Using F# Units of Measure with System.Numerics.Vector<T>

Question

I struggle using F# units of measure in combination with the System.Numerics.Vector<'T> type. Let's have a look at a toy problem: Assume we have an array xs of type float<m>[] and for some reason we wanted to square all its components, resulting in an array of type float<m^2>[]. That works perfectly well with scalar code:

xs |> Array.map (fun x -> x * x) // float<m^2>[]

Now suppose we wanted to vectorize this operation by performing the multiplication in System.Numerics.Vector<float>.Count-sized chunks using SIMD, for instance like so:

open System.Numerics
let simdWidth = Vector<float>.Count
// fill with dummy data
let xs = Array.init (simdWidth * 10) (fun i -> float i * 1.0<m>)
// array to store the results
let rs: float<m^2> array = Array.zeroCreate (xs |> Array.length)
// number of SIMD operations required
let chunks = (xs |> Array.length) / simdWidth
// for simplicity, assume xs.Length % simdWidth = 0
for i = 0 to chunks - 1 do
    let v = Vector<_>(xs, i * simdWidth) // Vector<float<m>>, containing xs.[i .. i+simdWidth-1]
    let u = v * v                        // Vector<float<m>>; expected: Vector<float<m^2>>
    u.CopyTo(rs, i * simdWidth)          // units mismatch

I believe I understand why this happens: How could the F# compiler know, what System.Numerics.Vector<'T>.op_Multiply does and what arithmetic rules apply? It could literally be any operation. So how should it be able to deduce the correct units?

The question is: What the best way to make this work? How can we tell the compiler which rules apply?

Attempt 1: Remove all units of measure information from xs and add it back later on:

// remove all UoM from all arrays
let xsWoM = Array.map (fun x -> x / 1.0<m>) xs
// ...
// perform computation using xsWoM etc.
// ...
// add back units again
let xs = Array.map (fun x -> x * 1.0<m>) xsWoM

Problems: Performs unnecessary computations and/or copy operations, defeats the purpose of vectorizing the code for performance reasons. Also, largely defeats the purpose of using UoM to begin with.

Attempt 2: Use inline IL to change the return type of Vector<'T>.op_Multiply:

// reinterpret x to be of type 'b
let inline retype (x: 'a) : 'b = (# "" x: 'b #)
let inline (.*.) (u: Vector<float<'m>>) (v: Vector<float<'m>>): Vector<float<'m^2>> = u * v |> retype
// ...
let u = v .*. v // asserts type Vector<float<m^2>>

Problems: Doesn't require any additional operations, but uses a deprecated feature (inline IL) and isn't fully generic (only with respect to the unit of measure).

Does anyone have a better solution for this*?

*Note that the above example really is a toy problem to demonstrate the general issue. The real program solves a much more complicate initial value problem, involving many kinds of physical quantities.

Answer 1

The compiler is fine with figuring out how to apply unit rules for multiplication, the problem here is that you have a wrapped type. In your first example, when you write xs |> Array.map (fun x -> x * x), you're describing the multiplication in terms of the elements of the array rather than on the array directly.

When you have a Vector<float<m>>, the units are attached to the float rather than the Vector so when you try and multiply Vectors, the compiler won't treat that type as if it has any units.

Given the methods exposed by the class, I don't think there is an easy workaround using the Vector<'T> directly but there are options for wrapping the type.

Something like this could give you a unit-friendly vector:

type VectorWithUnits<'a, [<Measure>]'b> = 
    |VectorWithUnits of Vector<'a>

    static member inline (*) (a : VectorWithUnits<'a0, 'b0>, b : VectorWithUnits<'a0, 'b1>) 
        : VectorWithUnits<'a0, 'b0*'b1> =
        match a, b with
        |VectorWithUnits a, VectorWithUnits b -> VectorWithUnits <| a * b

In this case, the units are attached to the vector and multiplying vectors works as expected with the correct unit behaviour.

The problem is that now we can have separate and distinct unit of measure annotations on the Vector<'T> and on the float itself.

You can the turn an array of a specific type with units of measure into a set of the Vectors using:

let toFloatVectors (array : float<'m>[]) : VectorWithUnits<float,'m>[]  =
    let arrs = array |> Array.chunkBySize (Vector<float>.Count)
    arrs |> Array.map (Array.map (float) >> Vector >> VectorWithUnits)

and back:

let fromFloatVectors (vectors : VectorWithUnits<float,'m>[]) : float<'m>[] =
    let arr = Array.zeroCreate<float> (Array.length vectors)
    vectors |> Array.iteri (fun i uVec ->
        match uVec with
        |VectorWithUnits vec -> vec.CopyTo arr)
    arr |> Array.map (LanguagePrimitives.FloatWithMeasure<'m>)

A hacky alternative:

If you give up the generic type 'T you can make a float Vector behave properly through some fairly horrible boxing and runtime casting. This abuses the fact that units of measure are a compile time construct that no longer exist at runtime.

type FloatVectorWithUnits<[<Measure>]'b> = 
    |FloatVectorWithUnits of Vector<float<'b>>

    static member ( * ) (a : FloatVectorWithUnits<'b0>, b : FloatVectorWithUnits<'b1>) =
        match a, b with
        |FloatVectorWithUnits a, FloatVectorWithUnits b ->
            let c, d = box a :?> Vector<float<'b0*'b1>>, box b :?> Vector<float<'b0*'b1>>
            c * d |> FloatVectorWithUnits