Deprecate 3-arg link_libraries! in favor of lazy loading/linking

Project.toml

@@ -23,7 +23,7 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 [compat]
 ExprTools = "0.1"
 InteractiveUtils = "1"
-LLVM = "9.5"
+LLVM = "9.6"
 Libdl = "1"
 Logging = "1"
 PrecompileTools = "1"

src/GPUCompiler.jl

@@ -67,6 +67,8 @@ include("driver.jl")
 include("execution.jl")
 include("reflection.jl")
 
+include("deprecated.jl")
+
 include("precompile.jl")
 
 function __init__()

src/deprecated.jl

@@ -0,0 +1,45 @@
+# Deprecations scheduled for removal in the next major release.
+
+function defs(mod::LLVM.Module)
+    Base.depwarn("`GPUCompiler.defs(mod)` is deprecated; inline `filter(f -> !isdeclaration(f), collect(functions(mod)))`.",
+                 :defs)
+    filter(f -> !isdeclaration(f), collect(functions(mod)))
+end
+
+function decls(mod::LLVM.Module)
+    Base.depwarn("`GPUCompiler.decls(mod)` is deprecated; inline `filter(f -> isdeclaration(f) && !LLVM.isintrinsic(f), collect(functions(mod)))`.",
+                 :decls)
+    filter(f -> isdeclaration(f) && !LLVM.isintrinsic(f), collect(functions(mod)))
+end
+
+link_library!(mod::LLVM.Module, lib::LLVM.Module) = link_library!(mod, [lib])
+function link_library!(mod::LLVM.Module, libs::Vector{LLVM.Module})
+    Base.depwarn("`GPUCompiler.link_library!` is deprecated; call `LLVM.link!(mod, copy(lib))` directly, or `LLVM.link!(mod, lib; only_needed=true)` with a freshly-parsed library.",
+                 :link_library!)
+    libs = [copy(lib) for lib in libs]
+    for lib in libs
+        link!(mod, lib)
+    end
+end
+
+# no-op 3-arg fallback so downstream overrides that chain via
+# `invoke(GPUCompiler.link_libraries!, Tuple{CompilerJob, Module,
+# Vector{String}}, ...)` still resolve.
+link_libraries!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
+                undefined_fns::Vector{String}) = return
+
+# `true` when a downstream package has defined a 3-arg `link_libraries!`
+# override for `job`, i.e. the dispatched method isn't our fallback above.
+#
+# Uses the same `jl_gf_invoke_lookup` path as `Core._hasmethod` rather than
+# `which`, so it's safe to call from generated-function-adjacent contexts
+# where `Base.get_world_counter()` returns `typemax(UInt)` and reflection
+# queries like `which` / `methods` fail (see JuliaLang/julia#48611).
+# All this because Enzyme.jl calls GPUCompiler.jl from a generated function.
+function has_legacy_link_libraries(@nospecialize(job::CompilerJob))
+    tt = Tuple{typeof(link_libraries!), typeof(job),
+               LLVM.Module, Vector{String}}
+    world = ccall(:jl_get_tls_world_age, UInt, ())
+    m = ccall(:jl_gf_invoke_lookup, Any, (Any, Any, UInt), tt, nothing, world)
+    return m !== nothing && (m::Method).module !== @__MODULE__
+end

src/driver.jl

@@ -257,7 +257,7 @@ const __llvm_initialized = Ref(false)
                     merge!(compiled, dyn_meta.compiled)
                     if haskey(dyn_meta, :gv_to_value)
                         merge!(gv_to_value, dyn_meta.gv_to_value)
-                    end    
+                    end
                     @assert context(dyn_ir) == context(ir)
                     link!(ir, dyn_ir)
                     changed = true
@@ -303,20 +303,40 @@ const __llvm_initialized = Ref(false)
         # load the runtime outside of a timing block (because it recurses into the compiler)
         if !uses_julia_runtime(job)
             runtime = load_runtime(job)
-            runtime_fns = LLVM.name.(defs(runtime))
-            runtime_intrinsics = ["julia.gc_alloc_obj"]
         end
 
         @tracepoint "Library linking" begin
             # target-specific libraries
-            undefined_fns = LLVM.name.(decls(ir))
-            @tracepoint "target libraries" link_libraries!(job, ir, undefined_fns)
+            @tracepoint "target libraries" begin
+                if has_legacy_link_libraries(job)
+                    Base.depwarn(
+                        "3-arg `link_libraries!(job, mod, undefined_fns)` is deprecated; " *
+                        "migrate your override to the 2-arg form `link_libraries!(job, mod)`. " *
+                        "Instead of inspecting `undefined_fns` to decide what to link, " *
+                        "parse the library lazily with `parse(LLVM.Module, bytes; lazy=true)` " *
+                        "and link it with `LLVM.link!(mod, lib; only_needed=true)` — " *
+                        "the linker will then materialize only the referenced symbols.",
+                        :link_libraries!)
+                    undefined_fns = [LLVM.name(f) for f in functions(ir)
+                                     if isdeclaration(f) && !LLVM.isintrinsic(f)]
+                    link_libraries!(job, ir, undefined_fns)
+                else
+                    link_libraries!(job, ir)
+                end
+            end
 
             # GPU run-time library
-            if !uses_julia_runtime(job) && any(fn -> fn in runtime_fns ||
-                                                        fn in runtime_intrinsics,
-                                                undefined_fns)
-                @tracepoint "runtime library" link_library!(ir, runtime)
+            if !uses_julia_runtime(job)
+                # Calls to `gc_pool_alloc` are inserted after linking, so spoof
+                # it here so that lazy linking pulls it in, if needed.
+                if haskey(functions(ir), "julia.gc_alloc_obj")
+                    rt = Runtime.get(:gc_pool_alloc)
+                    if !haskey(functions(ir), rt.llvm_name)
+                        LLVM.Function(ir, rt.llvm_name, convert(LLVM.FunctionType, rt))
+                    end
+                end
+
+                @tracepoint "runtime library" link!(ir, runtime; only_needed=true)
             end
         end
     end

src/interface.jl

@@ -360,8 +360,7 @@ prepare_job!(@nospecialize(job::CompilerJob)) = return
 
 # early extension point used to link-in external bitcode files.
 # this is typically used by downstream packages to link vendor libraries.
-link_libraries!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
-                undefined_fns::Vector{String}) = return
+link_libraries!(@nospecialize(job::CompilerJob), mod::LLVM.Module) = return
 
 # finalization of the module, before deferred codegen and optimization
 finish_module!(@nospecialize(job::CompilerJob), mod::LLVM.Module, entry::LLVM.Function) =

src/rtlib.jl

@@ -1,16 +1,5 @@
 # compiler support for working with run-time libraries
 
-link_library!(mod::LLVM.Module, lib::LLVM.Module) = link_library!(mod, [lib])
-function link_library!(mod::LLVM.Module, libs::Vector{LLVM.Module})
-    # linking is destructive, so copy the libraries
-    libs = [copy(lib) for lib in libs]
-
-    for lib in libs
-        link!(mod, lib)
-    end
-end
-
-
 #
 # GPU run-time library
 #
@@ -126,60 +115,38 @@ function build_runtime(@nospecialize(job::CompilerJob))
     mod
 end
 
-const runtime_lock = ReentrantLock()
-
-const runtime_cache = Dict{String, Vector{UInt8}}()
-
 @locked function load_runtime(@nospecialize(job::CompilerJob))
     global compile_cache
     if compile_cache === nothing    # during precompilation
         return build_runtime(job)
     end
 
-    lock(runtime_lock) do
-        slug = runtime_slug(job)
-        if !supports_typed_pointers(context())
-            slug *= "-opaque"
-        end
-        name = "runtime_$(slug).bc"
-        path = joinpath(compile_cache, name)
-
-        # cache the runtime library on disk and in memory
-        lib = if haskey(runtime_cache, slug)
-            parse(LLVM.Module, runtime_cache[slug])
-        elseif ispath(path)
-            runtime_cache[slug] = open(path) do io
-                read(io)
-            end
-            parse(LLVM.Module, runtime_cache[slug])
-        end
-
-        if lib === nothing
-            @debug "Building the GPU runtime library at $path"
-            mkpath(compile_cache)
-            lib = build_runtime(job)
-
-            # atomic write to disk
-            temp_path, io = mktemp(dirname(path); cleanup=false)
-            write(io, lib)
-            close(io)
-            @static if VERSION >= v"1.12.0-DEV.1023"
-                mv(temp_path, path; force=true)
-            else
-                Base.rename(temp_path, path, force=true)
-            end
+    slug = runtime_slug(job)
+    if !supports_typed_pointers(context())
+        slug *= "-opaque"
+    end
+    name = "runtime_$(slug).bc"
+    path = joinpath(compile_cache, name)
+
+    if !ispath(path)
+        @debug "Building the GPU runtime library at $path"
+        mkpath(compile_cache)
+        lib = build_runtime(job)
+
+        # atomic write to disk
+        temp_path, io = mktemp(dirname(path); cleanup=false)
+        write(io, lib)
+        close(io)
+        @static if VERSION >= v"1.12.0-DEV.1023"
+            mv(temp_path, path; force=true)
+        else
+            Base.rename(temp_path, path, force=true)
         end
-
-        return lib
     end
+
+    return parse(LLVM.Module, MemoryBufferFile(path); lazy=true)
 end
 
 # remove the existing cache
 # NOTE: call this function from global scope, so any change triggers recompilation.
-function reset_runtime()
-    lock(runtime_lock) do
-        rm(compile_cache; recursive=true, force=true)
-    end
-
-    return
-end
+reset_runtime() = rm(compile_cache; recursive=true, force=true)

src/utils.jl

@@ -1,7 +1,3 @@
-defs(mod::LLVM.Module)  = filter(f -> !isdeclaration(f), collect(functions(mod)))
-decls(mod::LLVM.Module) = filter(f ->  isdeclaration(f) && !LLVM.isintrinsic(f),
-                                 collect(functions(mod)))
-
 ## debug verification
 
 should_verify() = ccall(:jl_is_debugbuild, Cint, ()) == 1 ||