Skip to content

Commit

Permalink
Fix precompilation harness.
Browse files Browse the repository at this point in the history
  • Loading branch information
maleadt committed Jan 13, 2025
1 parent 339330c commit 37d8298
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 37 deletions.
23 changes: 7 additions & 16 deletions test/helpers/precompile.jl
Original file line number Diff line number Diff line change
Expand Up @@ -4,27 +4,18 @@ function precompile_test_harness(@nospecialize(f), testset::String)
end
end
function precompile_test_harness(@nospecialize(f), separate::Bool)
load_path = mktempdir()
load_cache_path = separate ? mktempdir() : load_path
# XXX: clean-up may fail on Windows, because opened files are not deletable.
# fix this by running the harness in a separate process, such that the
# compilation cache files are not opened?
load_path = mktempdir(cleanup=true)
load_cache_path = separate ? mktempdir(cleanup=true) : 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)
popfirst!(DEPOT_PATH)
popfirst!(LOAD_PATH)
end
nothing
end
Expand Down
25 changes: 13 additions & 12 deletions test/native/precompile.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,11 @@

precompile_test_harness("Inference caching") do load_path
# Write out the Native test setup as a micro package
create_standalone(load_path, "TestCompiler", "native.jl")
create_standalone(load_path, "NativeCompiler", "native.jl")

write(joinpath(load_path, "InferenceCaching.jl"), :(module InferenceCaching
import TestCompiler
write(joinpath(load_path, "NativeBackend.jl"), :(
module NativeBackend
import NativeCompiler
using PrecompileTools

function kernel(A, x)
Expand All @@ -14,36 +15,36 @@ precompile_test_harness("Inference caching") do load_path
end

let
job, _ = TestCompiler.Native.create_job(kernel, (Vector{Int}, Int))
job, _ = NativeCompiler.Native.create_job(kernel, (Vector{Int}, Int))
precompile(job)
end

# identity is foreign
@setup_workload begin
job, _ = TestCompiler.Native.create_job(identity, (Int,))
job, _ = NativeCompiler.Native.create_job(identity, (Int,))
@compile_workload begin
precompile(job)
end
end
end) |> string)

Base.compilecache(Base.PkgId("InferenceCaching"))
Base.compilecache(Base.PkgId("NativeBackend"))
@eval let
import TestCompiler
import NativeCompiler

# Check that no cached entry is present
identity_mi = GPUCompiler.methodinstance(typeof(identity), Tuple{Int})

token = let
job, _ = TestCompiler.Native.create_job(identity, (Int,))
job, _ = NativeCompiler.Native.create_job(identity, (Int,))
GPUCompiler.ci_cache_token(job)
end
@test !check_presence(identity_mi, token)

using InferenceCaching
using NativeBackend

# Check that kernel survived
kernel_mi = GPUCompiler.methodinstance(typeof(InferenceCaching.kernel), Tuple{Vector{Int}, Int})
kernel_mi = GPUCompiler.methodinstance(typeof(NativeBackend.kernel), Tuple{Vector{Int}, Int})
@test check_presence(kernel_mi, token)

# check that identity survived
Expand All @@ -55,15 +56,15 @@ precompile_test_harness("Inference caching") do load_path
GPUCompiler.enable_disk_cache!()
@test GPUCompiler.disk_cache_enabled() == true

job, _ = TestCompiler.Native.create_job(InferenceCaching.kernel, (Vector{Int}, Int))
job, _ = NativeCompiler.Native.create_job(NativeBackend.kernel, (Vector{Int}, Int))
@assert job.source == kernel_mi
ci = GPUCompiler.ci_cache_lookup(GPUCompiler.ci_cache(job), job.source, job.world, job.world)
@assert ci !== nothing
@assert ci.inferred !== nothing
path = GPUCompiler.cache_file(ci, job.config)
@test path !== nothing
@test !ispath(path)
TestCompiler.Native.cached_execution(InferenceCaching.kernel, (Vector{Int}, Int))
NativeCompiler.Native.cached_execution(NativeBackend.kernel, (Vector{Int}, Int))
@test ispath(path)
GPUCompiler.clear_disk_cache!()
@test !ispath(path)
Expand Down
19 changes: 10 additions & 9 deletions test/ptx/precompile.jl
Original file line number Diff line number Diff line change
@@ -1,32 +1,33 @@
precompile_test_harness("Inference caching") do load_path
# Write out the PTX test helpers as a micro package
create_standalone(load_path, "TestCompiler", "ptx.jl")
create_standalone(load_path, "PTXCompiler", "ptx.jl")

write(joinpath(load_path, "InferenceCaching.jl"), :(module InferenceCaching
import TestCompiler
write(joinpath(load_path, "PTXBackend.jl"), :(
module PTXBackend
import PTXCompiler
using PrecompileTools

function kernel()
return
end

let
job, _ = TestCompiler.PTX.create_job(kernel, ())
job, _ = PTXCompiler.PTX.create_job(kernel, ())
precompile(job)
end

# identity is foreign
@setup_workload begin
job, _ = TestCompiler.PTX.create_job(identity, (Int,))
job, _ = PTXCompiler.PTX.create_job(identity, (Int,))
@compile_workload begin
precompile(job)
end
end
end) |> string)

Base.compilecache(Base.PkgId("InferenceCaching"))
Base.compilecache(Base.PkgId("PTXBackend"))
@eval let
import TestCompiler
import PTXCompiler

# Check that no cached entry is present
identity_mi = GPUCompiler.methodinstance(typeof(identity), Tuple{Int})
Expand All @@ -41,10 +42,10 @@ precompile_test_harness("Inference caching") do load_path
ci = isdefined(ci, :next) ? ci.next : nothing
end

using InferenceCaching
using PTXBackend

# Check that kernel survived
kernel_mi = GPUCompiler.methodinstance(typeof(InferenceCaching.kernel), Tuple{})
kernel_mi = GPUCompiler.methodinstance(typeof(PTXBackend.kernel), Tuple{})
@test check_presence(kernel_mi, token)

# check that identity survived
Expand Down

0 comments on commit 37d8298

Please sign in to comment.