TTFX-related improvements (#856)

Author: maleadt

Project.toml

@@ -1,6 +1,6 @@
 name = "GPUCompiler"
 uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
-version = "1.22.4"
+version = "1.22.5"
 authors = ["Tim Besard <tim.besard@gmail.com>"]
 
 [workspace]
@@ -14,6 +14,7 @@ Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
+REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 Scratch = "6c6a2e73-6563-6170-7368-637461726353"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
@@ -34,6 +35,7 @@ Logging = "1"
 NVPTX_LLVM_Backend_jll = "22"
 PrecompileTools = "1"
 Preferences = "1"
+REPL = "1"
 Scratch = "1"
 Serialization = "1"
 TOML = "1"

src/jlgen.jl

@@ -984,17 +984,24 @@ function compile_method_instance(@nospecialize(job::CompilerJob))
             push!(code_instances, ci)
         end
     else
+        # `jl_get_llvm_cis` can report stale CIs that no longer cover the world
+        # this job was compiled in. The explicit cache lookup path already filters
+        # by world; do the same here before de-duplicating by MI.
+        filter!(code_instances) do ci
+            ci.min_world <= job.world <= ci.max_world
+        end
+
         # To avoid a clash in the compiled cache containing both with an interpreter token (like GPUCompiler.GPUCompilerCacheToken) and native,
         # prefer the non-native code-instance.
         # TODO: in the future we should migrate compiled to have the ci as the key, not the mi.
-        native_mis = Set{MethodInstance}()
+        owned_mis = Set{MethodInstance}()
         for ci in code_instances
             if ci.owner !== nothing
-                push!(native_mis, ci.def::MethodInstance)
+                push!(owned_mis, ci.def::MethodInstance)
             end
         end
         filter!(code_instances) do ci
-            return ci.owner !== nothing || in(ci.def, native_mis)
+            return ci.owner !== nothing || ci.def ∉ owned_mis
         end
     end
 

src/precompile.jl

@@ -1,5 +1,9 @@
 using PrecompileTools: @setup_workload, @compile_workload
 
+# Load REPL so that interactive use doesn't invalidate the precompiled compiler
+# pipeline (fixed on Julia 1.14, see JuliaLang/julia#61714).
+import REPL
+
 @setup_workload begin
     precompile_module = @eval module $(gensym())
         using ..GPUCompiler

test/native/precompile.jl

@@ -25,6 +25,10 @@ precompile_test_harness("Inference caching") do load_path
             return x*x
         end
 
+        function checked_convert(x)
+            return Int(x)
+        end
+
         let
             job, _ = NativeCompiler.Native.create_job(kernel, (Vector{Int}, Int))
             precompile(job)
@@ -48,6 +52,13 @@ precompile_test_harness("Inference caching") do load_path
                 precompile(job)
             end
         end
+
+        @setup_workload begin
+            job, _ = NativeCompiler.Native.create_job(checked_convert, (UInt,); validate=false)
+            @compile_workload begin
+                precompile(job)
+            end
+        end
     end) |> string)
 
     Base.compilecache(Base.PkgId("NativeBackend"), stderr, stdout)
@@ -78,6 +89,21 @@ precompile_test_harness("Inference caching") do load_path
         # check that identity survived
         @test check_presence(identity_mi, token) broken=(v"1.12.0-DEV.1268" <= VERSION < v"1.12.5" || v"1.13.0-" <= VERSION < v"1.13.0-beta3"|| v"1.14.0-" <= VERSION < v"1.14.0-DEV.1843")
 
+        # Recompiling a foreign method after loading precompiled owner-token CIs
+        # may also surface a native owner-less CI for the same MethodInstance.
+        # GPUCompiler should prefer the owner-token CI instead of recording both.
+        job, _ = NativeCompiler.Native.create_job(identity, (Int,))
+        JuliaContext() do ctx
+            _, meta = GPUCompiler.compile(:llvm, job)
+            @test haskey(meta.compiled, job.source)
+        end
+
+        job, _ = NativeCompiler.Native.create_job(NativeBackend.checked_convert, (UInt,); validate=false)
+        JuliaContext() do ctx
+            _, meta = GPUCompiler.compile(:llvm, job)
+            @test haskey(meta.compiled, job.source)
+        end
+
         GPUCompiler.clear_disk_cache!()
         @test GPUCompiler.disk_cache_enabled() == false