Can I suppress F# compiler making copy of function in IL code?

Question

I want to create a JIT GPU compiler. You give an F# function, and we JIT compile it. The key of JIT compiling is to be able to cache the compiling result. I tried to use the MethodInfo as the caching key, but it won't work. It seems that F# compiler will make a copy of the function instead of referencing the origin function. Is there a way to suppress this behavior?

Here is a test code, ideally, it should be just compiled twice, but it did it 4 times.

let compileGpuCode (m:MethodInfo) =
    printfn "JIT compiling..."
    printfn "Type  : %A" m.ReflectedType
    printfn "Method: %A" m
    printfn ""
    "fake gpu code"

let gpuCodeCache = ConcurrentDictionary<MethodInfo, string>()

let launchGpu (func:int -> int -> int) =
    let m = func.GetType().GetMethod("Invoke", [| typeof<int>; typeof<int> |])
    let gpuCode = gpuCodeCache.GetOrAdd(m, compileGpuCode)
    // launch gpuCode
    ()

let myGpuCode (a:int) (b:int) = a + 2 * b

[<Test>]
let testFSFuncReflection() =
    launchGpu (+)
    launchGpu (+)
    launchGpu myGpuCode
    launchGpu myGpuCode

Here is the output:

JIT compiling...
Type  : AleaTest.FS.Lab.Experiments+testFSFuncReflection@50
Method: Int32 Invoke(Int32, Int32)

JIT compiling...
Type  : AleaTest.FS.Lab.Experiments+testFSFuncReflection@51-1
Method: Int32 Invoke(Int32, Int32)

JIT compiling...
Type  : AleaTest.FS.Lab.Experiments+testFSFuncReflection@52-2
Method: Int32 Invoke(Int32, Int32)

JIT compiling...
Type  : AleaTest.FS.Lab.Experiments+testFSFuncReflection@53-3
Method: Int32 Invoke(Int32, Int32)

Answer 1

The F# compiler treats your code more as something like this:

launchGpu (fun a b -> myGpuCode a b)
launchGpu (fun a b -> myGpuCode a b)

When compiling this, it will generate a new class to represent the function on each of the lines. If you wrote your test as follows:

let f = myGpuCode
launchGpu f
launchGpu f

... it would generate just one class (for the one place where the function is referenced) and then share the same type in both of the calls - so this would work.

In this example, the compiler actually inlines myGpuCode because it is too short, but if you make it more complex, then it generates very simple Invoke function in both of the classes:

ldarg.1
ldarg.2
call int32 Test::myGpuCode(int32, int32)
ret

I'm sure there is a plenty of caveats, but you could just check if the body of the generated class contains the same IL and uses that as your key instead. Once you have the Invoke method, you can get the IL body using the following:

let m = func.GetType().GetMethod("Invoke", [| typeof<int>; typeof<int> |])
let body = m.GetMethodBody().GetILAsByteArray()

This will be the same for both of the classes - ideally, you could also analyze this to figure out if the code is just calling some other method.