From 46edf96857a0fbf986495a0de110abe0a3910a8e Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Tue, 16 Apr 2024 06:59:41 -0400 Subject: [PATCH] Add a testsuite for inference caching (#567) --- src/jlgen.jl | 11 ++++++ test/Project.toml | 1 + test/native_tests.jl | 59 +++++++++++++++++++++++++++++++ test/precompile_testsetup.jl | 68 ++++++++++++++++++++++++++++++++++++ test/ptx_tests.jl | 56 +++++++++++++++++++++++++++++ test/ptx_testsetup.jl | 1 - test/runtests.jl | 9 ++++- 7 files changed, 203 insertions(+), 2 deletions(-) create mode 100644 test/precompile_testsetup.jl diff --git a/src/jlgen.jl b/src/jlgen.jl index 518924b2..c09af674 100644 --- a/src/jlgen.jl +++ b/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 diff --git a/test/Project.toml b/test/Project.toml index ba189bc9..c5be5973 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/native_tests.jl b/test/native_tests.jl index d9f3ab9f..ad753bde 100644 --- a/test/native_tests.jl +++ b/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 ############################################################################################ diff --git a/test/precompile_testsetup.jl b/test/precompile_testsetup.jl new file mode 100644 index 00000000..ee3bec06 --- /dev/null +++ b/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 diff --git a/test/ptx_tests.jl b/test/ptx_tests.jl index 10b40db8..e0a692c4 100644 --- a/test/ptx_tests.jl +++ b/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 ############################################################################################ diff --git a/test/ptx_testsetup.jl b/test/ptx_testsetup.jl index 7d7db554..ed5026f1 100644 --- a/test/ptx_testsetup.jl +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index 57b75fbc..9411dbc7 100644 --- a/test/runtests.jl +++ b/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