Use CompilerCaching 0.5 specialization entries

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"

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

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")

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