diff --git a/Project.toml b/Project.toml index 2171c9fa..9429a68e 100644 --- a/Project.toml +++ b/Project.toml @@ -36,7 +36,7 @@ BFloat16s = "0.6" CUDACore = "6.3" CUDA_Compiler_jll = "0.5, 0.6" CUDA_Tile_jll = "13.3" -CompilerCaching = "0.4.6" +CompilerCaching = "0.5" DLFP8Types = "0.1.0" EnumX = "1.0" GPUArrays = "11" diff --git a/src/compiler/driver.jl b/src/compiler/driver.jl index f7885687..16aa4a7e 100644 --- a/src/compiler/driver.jl +++ b/src/compiler/driver.jl @@ -117,17 +117,9 @@ function get_inferred(cache::CacheView{K,V}, ci::Core.CodeInstance, if const_argtypes === nothing src = @something get_source(ci) else - src = @something get_source(ci, const_argtypes) - - # Extract the return type from a const-specialized entry. - cached = CC.traverse_analysis_results(ci) do @nospecialize(result) - result isa CompilerCaching.CachedResult{V} ? result : nothing - end - for entry in cached.const_entries - if entry.argtypes == const_argtypes - rettype = CC.widenconst(entry.rettype) - end - end + entry = @something specialization(cache, ci, const_argtypes) + src = @something get_source(entry) + rettype = CC.widenconst(entry.rettype) end ir = CC.inflate_ir(src, mi) return ir, rettype @@ -214,16 +206,20 @@ to resolve them on every phase. @inline function ensure_compiled(cache::CacheView{K,CuTileResults}, mi::Core.MethodInstance, const_argtypes::Union{Vector{Any}, Nothing}) where {K} - # Fast path: already cached - hit = const_argtypes === nothing ? lookup(cache, mi) : - lookup(cache, mi, const_argtypes) - hit !== nothing && return hit + if const_argtypes === nothing + hit = lookup(cache, mi) + hit !== nothing && return hit - # Slow path: run inference (and const-prop, if requested) and re-resolve. - ci = get_ci(cache, mi; const_argtypes) - res = const_argtypes === nothing ? results(cache, ci) : - results(cache, ci, const_argtypes) - return (ci, res) + ci = get_ci(cache, mi; const_argtypes) + return (ci, results(cache, ci)) + else + hit = lookup(cache, mi, const_argtypes) + hit !== nothing && return (hit[1], results(hit[2])) + + ci = get_ci(cache, mi; const_argtypes) + entry = @something specialization(cache, ci, const_argtypes) + return (ci, results(entry)) + end end # Cached wrappers around the driver's emit_* functions. These check/populate diff --git a/src/cuTile.jl b/src/cuTile.jl index 55ffc7db..51b79867 100644 --- a/src/cuTile.jl +++ b/src/cuTile.jl @@ -22,7 +22,8 @@ using EnumX public BFloat16 import CompilerCaching -using CompilerCaching: CacheView, method_instance, match_method_instance, typeinf!, results, lookup, get_source +using CompilerCaching: CacheView, method_instance, match_method_instance, typeinf!, results, lookup, + specialization, get_source # Shared definitions include("shapes.jl") diff --git a/test/host/inference_cache.jl b/test/host/inference_cache.jl new file mode 100644 index 00000000..98f3aca6 --- /dev/null +++ b/test/host/inference_cache.jl @@ -0,0 +1,33 @@ +@testset "Inference cache" begin + choose_type(flag, x) = flag ? x + one(x) : Float32(x) + world = Base.get_world_counter() + cache = ct.CacheView{ct.CuTileResults}(gensym(:inference_cache), world) + mi = ct.lookup_method_instance(choose_type, Tuple{Bool, Int32}; world) + + ci, generic = ct.ensure_compiled(cache, mi, nothing) + @test ct.ensure_compiled(cache, mi, nothing) === (ci, generic) + _, generic_rt = ct.get_inferred(cache, ci, mi) + @test generic_rt == Union{Int32, Float32} + + specialized = ct.CuTileResults[] + for (flag, expected_rt) in ((true, Int32), (false, Float32)) + argtypes = Any[ct.CC.Const(choose_type), ct.CC.Const(flag), Int32] + specialized_ci, res = ct.ensure_compiled(cache, mi, argtypes) + @test specialized_ci === ci + @test res !== generic + @test all(previous -> previous !== res, specialized) + push!(specialized, res) + + # Source and return type must come from the same specialization. + ir, rt = ct.get_inferred(cache, ci, mi; const_argtypes=argtypes) + @test ir isa ct.CC.IRCode + @test rt === expected_rt + structured = ct.emit_structured!(cache, mi, ci, res; const_argtypes=argtypes) + @test structured[2] === expected_rt + + # An equivalent argument vector must reuse the compiled results. + @test ct.ensure_compiled(cache, mi, copy(argtypes)) === (ci, res) + @test ct.emit_structured!(cache, mi, ci, res; + const_argtypes=argtypes) === structured + end +end