Add a testsuite for inference caching (#567)

Author: vchuravy

src/jlgen.jl

@@ -749,3 +749,14 @@ function compile_method_instance(@nospecialize(job::CompilerJob))
 
     return llvm_mod, compiled
 end
+
+function CC.typeinf(interp::GPUInterpreter, frame::CC.InferenceState)
+    if CC.__measure_typeinf__[]
+        CC.Timings.enter_new_timer(frame)
+        v = CC._typeinf(interp, frame)
+        CC.Timings.exit_current_timer(frame)
+        return v
+    else
+        return CC._typeinf(interp, frame)
+    end
+end

test/Project.toml

@@ -2,6 +2,7 @@
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 SPIRV_LLVM_Translator_unified_jll = "85f0d8ed-5b39-5caa-b1ae-7472de402361"

test/native_tests.jl

@@ -535,6 +535,65 @@ end
                ["jl_invoke", "apply_iterate",
                 "inttoptr", "apply_type"]) broken=VERSION>=v"1.11.0-DEV.392"
 end
+end # testitem
+
+@testitem "native precompile" setup=[Precompile,] begin
+
+using Test
+precompile_test_harness("Inference caching") do load_path
+    # Write out the Native test setup as a micro package
+    create_standalone(load_path, "NativeCompiler", "native_testsetup.jl")
+
+    write(joinpath(load_path, "InferenceCaching.jl"), :(module InferenceCaching
+        import NativeCompiler
+        import GPUCompiler
+        using PrecompileTools
+
+        function kernel()
+            return
+        end
+
+        let
+            job, _ = NativeCompiler.create_job(kernel, ())
+            GPUCompiler.code_typed(job)
+        end
+        
+        # identity is foreign
+        @setup_workload begin
+            job, _ = NativeCompiler.create_job(identity, (Int,))
+            @compile_workload begin
+                GPUCompiler.code_typed(job)
+            end
+        end
+    end) |> string)
+
+    Base.compilecache(Base.PkgId("InferenceCaching"))
+    @eval let
+        import NativeCompiler
+
+        # Check that no cached entry is present
+        identity_mi = GPUCompiler.methodinstance(typeof(identity), Tuple{Int})
+
+        token = let
+            job, _ = NativeCompiler.create_job(identity, (Int,))
+            GPUCompiler.ci_cache_token(job)
+        end
+        ci = isdefined(identity_mi, :cache) ? identity_mi.cache : nothing
+        while ci !== nothing
+            @test ci.owner !== token
+            ci = isdefined(ci, :next) ? ci.next : nothing
+        end
+
+        using InferenceCaching
+
+        # Check that kernel survived
+        kernel_mi = GPUCompiler.methodinstance(typeof(InferenceCaching.kernel), Tuple{})
+        @test check_presence(kernel_mi, token)
+
+        # check that identity survived
+        @test check_presence(identity_mi, token)
+    end
+end
 
 ############################################################################################
 

test/precompile_testsetup.jl

@@ -0,0 +1,68 @@
+@testsetup module Precompile
+
+using Test
+using ReTestItems
+
+export precompile_test_harness, check_presence, create_standalone
+
+function precompile_test_harness(@nospecialize(f), testset::String)
+    @testset "$testset" begin
+        precompile_test_harness(f, true)
+    end
+end
+function precompile_test_harness(@nospecialize(f), separate::Bool)
+    load_path = mktempdir()
+    load_cache_path = separate ? mktempdir() : load_path
+    try
+        pushfirst!(LOAD_PATH, load_path)
+        pushfirst!(DEPOT_PATH, load_cache_path)
+        f(load_path)
+    finally
+        try
+            rm(load_path, force=true, recursive=true)
+        catch err
+            @show err
+        end
+        if separate
+            try
+                rm(load_cache_path, force=true, recursive=true)
+            catch err
+                @show err
+            end
+        end
+        filter!((≠)(load_path), LOAD_PATH)
+        separate && filter!((≠)(load_cache_path), DEPOT_PATH)
+    end
+    nothing
+end
+
+function check_presence(mi, token)
+    found = false
+    ci = isdefined(mi, :cache) ? mi.cache : nothing
+    while ci !== nothing
+        if ci.owner === token && ci.max_world == typemax(UInt)
+            found = true
+            break
+        end
+        ci = isdefined(ci, :next) ? ci.next : nothing
+    end
+    return found
+end
+
+function create_standalone(load_path, name::String, file)
+    cp(joinpath(@__DIR__, "runtime.jl"), joinpath(load_path, "runtime.jl"), force=true)
+
+    TS = include(file)
+    code = TS.code
+    if code.head == :begin
+        code.head = :block
+    end
+    @assert code.head == :block
+    code = Expr(:module, true, Symbol(name), code)
+
+    # Write out the test setup as a micro package
+    write(joinpath(load_path, "$name.jl"), string(code))
+    Base.compilecache(Base.PkgId(name))
+end
+
+end # testsetup

test/ptx_tests.jl

@@ -321,7 +321,63 @@ end
 end
 
 end
+end # testitem
 
+@testitem "PTX precompile" setup=[Precompile,] begin
+precompile_test_harness("Inference caching") do load_path
+    # Write out the PTX test setup as a micro package
+    create_standalone(load_path, "PTXCompiler", "ptx_testsetup.jl")
+
+    write(joinpath(load_path, "InferenceCaching.jl"), :(module InferenceCaching
+        import PTXCompiler
+        import GPUCompiler
+        using PrecompileTools
+
+        function kernel()
+            return
+        end
+
+        let
+            job, _ = PTXCompiler.create_job(kernel, ())
+            GPUCompiler.code_typed(job)
+        end
+        
+        # identity is foreign
+        @setup_workload begin
+            job, _ = PTXCompiler.create_job(identity, (Int,))
+            @compile_workload begin
+                GPUCompiler.code_typed(job)
+            end
+        end
+    end) |> string)
+
+    Base.compilecache(Base.PkgId("InferenceCaching"))
+    @eval let
+        import PTXCompiler
+
+        # Check that no cached entry is present
+        identity_mi = GPUCompiler.methodinstance(typeof(identity), Tuple{Int})
+
+        token = let
+            job, _ = PTXCompiler.create_job(identity, (Int,))
+            GPUCompiler.ci_cache_token(job)
+        end
+        ci = isdefined(identity_mi, :cache) ? identity_mi.cache : nothing
+        while ci !== nothing
+            @test ci.owner !== token
+            ci = isdefined(ci, :next) ? ci.next : nothing
+        end
+
+        using InferenceCaching
+
+        # Check that kernel survived
+        kernel_mi = GPUCompiler.methodinstance(typeof(InferenceCaching.kernel), Tuple{})
+        @test check_presence(kernel_mi, token)
+
+        # check that identity survived
+        @test check_presence(identity_mi, token)
+    end
+end
 
 ############################################################################################
 

test/ptx_testsetup.jl

@@ -9,7 +9,6 @@ include("runtime.jl")
 struct CompilerParams <: AbstractCompilerParams end
 
 PTXCompilerJob = CompilerJob{PTXCompilerTarget,CompilerParams}
-GPUCompiler.runtime_module(::PTXCompilerJob) = TestRuntime
 
 struct PTXKernelState
     data::Int64

test/runtests.jl

@@ -22,14 +22,21 @@ runtests(GPUCompiler; nworkers=min(Sys.CPU_THREADS,4), nworker_threads=1,
         end
     end
 
-    if ti.name in ["PTX", "GCN"] && Sys.isapple() && VERSION >= v"1.10-"
+    if ti.name in ["PTX", "GCN", "PTX precompile"] && Sys.isapple() && VERSION >= v"1.10-"
         # support for AMDGPU and NVTX on macOS has been removed from Julia's LLVM build
         return false
     end
 
+
     if ti.name in ["SPIRV"] && !(SPIRV_LLVM_Translator_unified_jll.is_available() && SPIRV_Tools_jll.is_available())
         # SPIRV needs it's tools to be available
         return false
     end
+
+    if ti.name in ["PTX precompile", "native precompile"] && VERSION < v"1.11-"
+        # precompile needs v1.11
+        return false
+    end
+
     true
 end