Can I get an identity of an F# quotation?

Question

F# quotation is a wonderful feature, it allows us to treat F# expression as normal F# value. In my context, I use F# quotation to code Gpu kernel, and compile it into Gpu bitcode module.

There is one problem. I don't want to compile Gpu kernel each time, I would like to cache the compiled Gpu bitcode module. Thus I need a key, or identity from an F# quotation value. I would like to have a cache system like:

let compile : Expr -> GpuModule

let cache = ConcurrentDictionary<Key, GpuModule>()

let jitCompile (expr:Expr) =
    let key = getQuotationKey(expr)
    cache.GetOrAdd(key, fun key -> compile expr)

There is one solution, to use the quotation expr instance as the key. But look at this piece of code:

open Microsoft.FSharp.Quotations

let foo (expr:Expr) =
    printfn "%O" expr

[<EntryPoint>]
let main argv = 

    for i = 1 to 10 do
        foo <@ fun x y -> x + y @>

    0

If I check the compiled IL code, I got these IL instructions:

IL_0000: nop
IL_0001: ldc.i4.1
IL_0002: stloc.0
IL_0003: br IL_00a2
// loop start (head: IL_00a2)
    IL_0008: ldtoken '<StartupCode$ConsoleApplication2>.$Program'
    IL_000d: call class [mscorlib]System.Type [mscorlib]System.Type::GetTypeFromHandle(valuetype [mscorlib]System.RuntimeTypeHandle)
    IL_0012: ldc.i4.5
    IL_0013: newarr [mscorlib]System.Type
    IL_0018: dup
    IL_0019: ldc.i4.0
    IL_001a: ldtoken [mscorlib]System.Int32
    IL_001f: call class [mscorlib]System.Type [mscorlib]System.Type::GetTypeFromHandle(valuetype [mscorlib]System.RuntimeTypeHandle)
    IL_0024: stelem.any [mscorlib]System.Type
    IL_0029: dup
    IL_002a: ldc.i4.1
    IL_002b: ldtoken [FSharp.Core]Microsoft.FSharp.Core.Operators
    IL_0030: call class [mscorlib]System.Type [mscorlib]System.Type::GetTypeFromHandle(valuetype [mscorlib]System.RuntimeTypeHandle)
    IL_0035: stelem.any [mscorlib]System.Type
    IL_003a: dup
    IL_003b: ldc.i4.2
    IL_003c: ldtoken [mscorlib]System.Tuple`2
    IL_0041: call class [mscorlib]System.Type [mscorlib]System.Type::GetTypeFromHandle(valuetype [mscorlib]System.RuntimeTypeHandle)
    IL_0046: stelem.any [mscorlib]System.Type
    IL_004b: dup
    IL_004c: ldc.i4.3
    IL_004d: ldtoken [mscorlib]System.String
    IL_0052: call class [mscorlib]System.Type [mscorlib]System.Type::GetTypeFromHandle(valuetype [mscorlib]System.RuntimeTypeHandle)
    IL_0057: stelem.any [mscorlib]System.Type
    IL_005c: dup
    IL_005d: ldc.i4.4
    IL_005e: ldtoken [mscorlib]System.Tuple`5
    IL_0063: call class [mscorlib]System.Type [mscorlib]System.Type::GetTypeFromHandle(valuetype [mscorlib]System.RuntimeTypeHandle)
    IL_0068: stelem.any [mscorlib]System.Type
    IL_006d: ldc.i4.0
    IL_006e: newarr [mscorlib]System.Type
    IL_0073: ldc.i4.0
    IL_0074: newarr [FSharp.Core]Microsoft.FSharp.Quotations.FSharpExpr
    IL_0079: ldc.i4 372
    IL_007e: newarr [mscorlib]System.Byte
    IL_0083: dup
    IL_0084: ldtoken field valuetype '<PrivateImplementationDetails$ConsoleApplication2>'/T1805_372Bytes@ Program::field1806@
    IL_0089: call void [mscorlib]System.Runtime.CompilerServices.RuntimeHelpers::InitializeArray(class [mscorlib]System.Array, valuetype [mscorlib]System.RuntimeFieldHandle)
    IL_008e: call class [FSharp.Core]Microsoft.FSharp.Quotations.FSharpExpr [FSharp.Core]Microsoft.FSharp.Quotations.FSharpExpr::Deserialize40(class [mscorlib]System.Type, class [mscorlib]System.Type[], class [mscorlib]System.Type[], class [FSharp.Core]Microsoft.FSharp.Quotations.FSharpExpr[], uint8[])
    IL_0093: call class [FSharp.Core]Microsoft.FSharp.Quotations.FSharpExpr`1<!!0> [FSharp.Core]Microsoft.FSharp.Quotations.FSharpExpr::Cast<class [FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<int32, class [FSharp.Core]Microsoft.FSharp.Core.FSharpFunc`2<int32, int32>>>(class [FSharp.Core]Microsoft.FSharp.Quotations.FSharpExpr)
    IL_0098: call void Program::foo(class [FSharp.Core]Microsoft.FSharp.Quotations.FSharpExpr)
    IL_009d: nop
    IL_009e: ldloc.0
    IL_009f: ldc.i4.1
    IL_00a0: add
    IL_00a1: stloc.0

    IL_00a2: ldloc.0
    IL_00a3: ldc.i4.s 11
    IL_00a5: blt IL_0008
// end loop

IL_00aa: ldc.i4.0
IL_00ab: ret

This is a large code, but basically it does these things in the loop:

• Load byte array for the quotation from some static field
• Setup type informations
• Call FSharp.Quotations.FSharpExpr::Deserialize40 to re-create the quotation object;

So from this observation, my questions are:

1. Although quotation are stored in one static field, but each time we write <@ ... @>, they will create a new Expr instance, even when the static field is same. So I cannot use the Expr instance as the key, it would be good to get the static field token and use that as the key. But I don't know how to get that information;
2. We saw there are lots of IL instructions to just re-create an quotation instance, even when they are the same quotation. This might have some performance issues, can F# compiler be optimized here?

Regards, Xiang.

@kvb gave a wonderful answer. Looks like we just need to fix the Var compare in the quotations (when var has a counterpart and has the same type). Follow his answer I made the following tests and it works:

let comparer =
    let rec compareQuots vs = function
        | ShapeLambda(v,e), ShapeLambda(v',e') ->
            compareQuots (vs |> Map.add v v') (e,e')
        | ShapeCombination(o,es), ShapeCombination(o',es') ->
            o = o' && (es.Length = es'.Length) && List.forall2 (fun q1 q2 -> compareQuots vs (q1, q2)) es es'
        | ShapeVar v, ShapeVar v' when Map.tryFind v vs = Some v' && v.Type = v'.Type ->
            true
        | _ -> false

    let rec hashQuot n vs = function
        | ShapeLambda(v,e) ->
            hashQuot (n+1) (vs |> Map.add v n) e
        | ShapeCombination(o,es) ->
            es |> List.fold (fun h e -> 31 * h + hashQuot n vs e) (o.GetHashCode())
        | ExprShape.ShapeVar v ->
            Map.find v vs

    { new System.Collections.Generic.IEqualityComparer<_> with 
        member __.Equals(q1,q2) = compareQuots Map.empty (q1,q2)
        member __.GetHashCode q = hashQuot 0 Map.empty q }

type Module = int

let mutable counter = 0

let compile (expr:Expr) =
    counter <- counter + 1
    printfn "Compiling #.%d module..." counter
    counter

let cache = ConcurrentDictionary<Expr, Module>(comparer)

let jitCompile (expr:Expr) =
    cache.GetOrAdd(expr, compile)

[<Test>]
let testJITCompile() =
    Assert.AreEqual(1, jitCompile <@ fun x y -> x + y @>)
    Assert.AreEqual(1, jitCompile <@ fun x y -> x + y @>)
    Assert.AreEqual(1, jitCompile <@ fun a b -> a + b @>)
    Assert.AreEqual(2, jitCompile <@ fun a b -> a + b + 1 @>)

    let combineExpr (expr:Expr<int -> int -> int>) =
        <@ fun (a:int) (b:int) -> ((%expr) a b) + 1 @> 

    // although (combineExpr <@ (+) @>) = <@ fun a b -> a + b + 1 @>
    // but they are treated as different expr.
    Assert.AreEqual(3, jitCompile (combineExpr <@ (+) @>))
    Assert.AreEqual(3, jitCompile (combineExpr <@ (+) @>))
    Assert.AreEqual(4, jitCompile (combineExpr <@ (-) @>))

Answer 1

Creating a new object each time through the loop doesn't necessarily mean that the object can't be used as a key, as long as the objects compare equal each time.

The real issue you'll have is that "the same" quotation means something different to you than to the F# compiler, particularly when it comes to variables in quotations. For example, you can verify that

<@ [1 + 1] @> = <@ [1 + 1] @>

evaluates to true, and

<@ fun x -> x @> = <@ fun y -> y @>

evaluates to false (which hopefully makes sense, since the lambdas are equivalent up to renaming, but not identical). Perhaps more surprisingly, you'll see that

<@ fun x -> x @> = <@ fun x -> x @>

also evaluates to false. This is because the variables in each quotation are treated as different variables which just happen to share the same name. You'll see the same behavior in your loop - each iteration's variable x is considered different.

However, all is not lost; all you need to do is to use a custom IEqualityComparer<Quotations.Expr>. I think something like this should work to identify any quotations that are identical modulo variable renaming:

let comparer = 
    let rec compareQuots vs = function
    | Quotations.ExprShape.ShapeLambda(v,e), Quotations.ExprShape.ShapeLambda(v',e') ->
        compareQuots (vs |> Map.add v v') (e,e')
    | Quotations.ExprShape.ShapeCombination(o,es), Quotations.ExprShape.ShapeCombination(o',es') ->
        o = o' && (es.Length = es'.Length) && List.forall2 (fun q1 q2 -> compareQuots vs (q1, q2)) es es'
    | Quotations.ExprShape.ShapeVar v, Quotations.ExprShape.ShapeVar v' when Map.tryFind v vs = Some v' && v.Type = v'.Type -> 
        true
    | _ -> false

    let rec hashQuot n vs = function
    | Quotations.ExprShape.ShapeLambda(v,e) -> 
        hashQuot (n+1) (vs |> Map.add v n) e
    | Quotations.ExprShape.ShapeCombination(o,es) -> 
        es |> List.fold (fun h e -> 31 * h + hashQuot n vs e) (o.GetHashCode())
    | Quotations.ExprShape.ShapeVar v -> 
        Map.find v vs

    { new System.Collections.Generic.IEqualityComparer<_> with 
        member __.Equals(q1,q2) = compareQuots Map.empty (q1,q2)
        member __.GetHashCode q = hashQuot 0 Map.empty q }

let cache = ConcurrentDictionary<Expr, Module>(comparer)