Add declarative device libraries and late intrinsic lowering

README.md

@@ -26,6 +26,35 @@ as [CUDA.jl](https://github.com/JuliaGPU/CUDA.jl), [AMDGPU.jl](https://github.co
 
 ## Development
 
+### Device libraries
+
+Back-ends describe the libraries that implement functions called from device code by
+overriding `GPUCompiler.device_libraries(job)`, returning a tuple of `DeviceLibrary`s.
+A library either exports Julia methods, `DeviceLibrary(:name; methods=(DeviceLibraryMethod(f,
+rettyp, argtyps; llvm_name), ...))`, or contains LLVM bitcode, `DeviceLibrary(:name;
+bitcode=path_or_bytes)`. `runtime_library(MyRuntime)` describes the target's
+implementation of GPUCompiler's runtime, and is selected by default for
+`runtime_module(job)`.
+
+Definitions are imported on demand, searching the libraries in order, until no library
+provides any remaining reference; unresolved references are reported by IR validation.
+This happens twice:
+
+- before optimization, for everything the kernel references, so that library code is
+  optimized along with the kernel (as Clang does with `-mlink-builtin-bitcode`);
+- after optimization, once `lower_intrinsics!(job, mod)` has lowered operations the
+  target cannot select (e.g. LLVM math intrinsics) to library calls, and again after GC
+  lowering, which introduces calls to the runtime. These late imports are processed by
+  `finish_imports!(job, mod, imported)`. Julia libraries imported at this point have been
+  lowered by `lower_intrinsics!` too.
+
+Library selection and lowering should only depend on the job's configuration, which
+identifies cached kernels. Compiled kernels are not invalidated when a library changes:
+Julia library methods are recompiled when redefined, but cached kernels keep using the
+previous version, and bitcode files are read once per session.
+
+### Testing
+
 GPUCompiler.jl's own test suite runs on GitHub Actions. In addition, Buildkite runs the test
 suites of the GPU back-ends that build on GPUCompiler (CUDA.jl, AMDGPU.jl, Metal.jl, oneAPI.jl
 and OpenCL.jl) against the version of GPUCompiler being tested. These back-end tests run on

src/GPUCompiler.jl

@@ -80,6 +80,7 @@ include("determinism.jl")
 include("optim.jl")
 include("validation.jl")
 include("rtlib.jl")
+include("devicelib.jl")
 include("mcgen.jl")
 include("debug.jl")
 include("driver.jl")
@@ -94,6 +95,8 @@ include("precompile.jl")
 function __init__()
     STDERR_HAS_COLOR[] = get(stderr, :color, false)
     empty!(session_results_cache)
+    empty!(device_libs)
+    empty!(bitcode_files)
 
     @static if !HAS_INTEGRATED_CACHE
         # CodeInstances created by GPUCompiler's precompile workload are process-local.

src/driver.jl

@@ -335,21 +335,7 @@ const __llvm_initialized = Ref(false)
     entrypoints = filter(((job′, _),) -> job′ === job || job′.config.kernel, jobs)
 
     if job.config.toplevel && job.config.libraries
-        # load the runtime outside of a timing block (because it recurses into the compiler)
-        if !uses_julia_runtime(job)
-            runtime, runtime_relocs = load_runtime(job)
-        end
-
-        @tracepoint "Library linking" begin
-            # target-specific libraries
-            @tracepoint "target libraries" link_libraries!(job, ir)
-
-            # GPU run-time library
-            if !uses_julia_runtime(job)
-                @tracepoint "runtime library" link_relocatable!(
-                    ir, relocations, runtime, runtime_relocs; only_needed=true)
-            end
-        end
+        @tracepoint "Library linking" link_device_libraries_early!(job, ir, relocations)
     end
 
     @tracepoint "IR post-processing" begin
@@ -406,6 +392,8 @@ const __llvm_initialized = Ref(false)
                         end
                     end
                 end
+            elseif job.config.libraries
+                lower_and_link_device_libraries!(job, ir, relocations)
             end
 
             # Runtime linking during optimization can add relocations.

src/interface.jl

@@ -742,10 +742,32 @@ end
 # priming the cache with entries that cannot be easily inferred.
 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.
+# extension point used to link-in external bitcode files, called before optimization and
+# again afterwards for references introduced late. prefer `device_libraries`.
 link_libraries!(@nospecialize(job::CompilerJob), mod::LLVM.Module) = return
 
+"""
+    lower_intrinsics!(job, mod)
+
+Lower operations that the target cannot select (e.g. LLVM math intrinsics) to calls into
+device libraries. This runs after optimization, on the whole module, after which the
+definitions for any new references are imported. It also runs on Julia device libraries
+imported at that point, so it should only depend on `runtime_config(job)`. Names of
+defined functions must be preserved. Not called when `libraries=false`.
+"""
+lower_intrinsics!(@nospecialize(job::CompilerJob), mod::LLVM.Module) = return
+
+"""
+    finish_imports!(job, mod, imported)
+
+Prepare definitions imported from device libraries after optimization, whose names are
+listed in `imported`. Definitions imported before optimization are processed by
+`finish_linked_module!` instead.
+"""
+finish_imports!(@nospecialize(job::CompilerJob), mod::LLVM.Module, imported) = return
+
+@public lower_intrinsics!, finish_imports!
+
 # finalization of the module, before deferred codegen and optimization
 finish_module!(@nospecialize(job::CompilerJob), mod::LLVM.Module, entry::LLVM.Function) =
     entry

src/metal.jl

@@ -2588,86 +2588,6 @@ end
 #
 # we don't have a proper back-end, so we're missing out on intrinsics-related functionality.
 
-# The function `name` of `mod`, declaring it with type `ft` (or creating it with `create()`)
-# if it doesn't exist yet. An existing function must have that type.
-function declare!(create, mod::LLVM.Module, name::String, ft::LLVM.FunctionType)
-    f = get!(create, mod.functions, name)
-    f.function_type == ft ||
-        error("Conflicting declarations of $name: $(f.function_type) and $ft")
-    return f
-end
-declare!(mod::LLVM.Module, name::String, ft::LLVM.FunctionType) =
-    declare!(() -> LLVM.Function(mod, name, ft), mod, name, ft)
-
-# Call the function `name`, declaring it for the types of `args` and return type `T_ret`.
-function call_declared!(builder::IRBuilder, name::String, T_ret::LLVMType,
-                        args::Vector{<:LLVM.Value})
-    mod = builder.insert_block.parent.parent
-    ft = LLVM.FunctionType(T_ret, LLVMType[arg.value_type for arg in args])
-    return call!(builder, ft, declare!(mod, name, ft), args)
-end
-
-# the suffix LLVM and AIR use to mangle overloaded intrinsics on `typ`, e.g. `v4f32`
-function type_suffix(@nospecialize(typ::LLVMType))
-    typ isa LLVM.IntegerType && return "i$(typ.width)"
-    typ isa LLVM.HalfType && return "f16"
-    typ isa LLVM.BFloatType && return "bf16"
-    typ isa LLVM.FloatType && return "f32"
-    typ isa LLVM.DoubleType && return "f64"
-    typ isa LLVM.VectorType && return "v$(typ.length)$(type_suffix(typ.element_type))"
-    error("Unsupported intrinsic type: $typ")
-end
-
-# the Julia floating-point type of an LLVM one (the C API lacks getPrimitiveSizeInBits)
-function julia_float_type(typ::LLVMType)
-    typ isa LLVM.HalfType && return Float16
-    typ isa LLVM.FloatType && return Float32
-    typ isa LLVM.DoubleType && return Float64
-    error("Unsupported floating-point type: $typ")
-end
-
-# the intrinsic that `inst` calls, if any
-# The tables below map intrinsics by name, but calls are matched on the intrinsic's ID:
-# the `name` of an `Intrinsic` cannot name an overloaded intrinsic (LLVM asserts that it
-# isn't), and IDs are only known at run time.
-# Intrinsics that the running LLVM doesn't know (e.g. `llvm.vector.reduce.fminimum` before
-# LLVM 17) can't be called, so they are left out.
-intrinsic_table(table) = Dict(intr => val for (intr, val) in
-                              ((tryparse(LLVM.Intrinsic, name), val) for (name, val) in table)
-                              if intr !== nothing)
-
-function called_intrinsic(inst::LLVM.Instruction)
-    inst isa LLVM.CallBase || return nothing
-    callee = inst.called_function
-    return callee === nothing ? nothing : callee.intrinsic
-end
-
-# Lower the calls to intrinsics in `fun`: `lower(builder, call, intrinsic)` is called for each,
-# with `builder` positioned at the call, and returns the value that replaces the call,
-# `:erase` to remove it, or `nothing` to keep it. The calls are collected first, so `lower`
-# can emit code, but it must not erase other instructions.
-function lower_intrinsic_calls!(lower, fun::LLVM.Function)
-    calls = LLVM.CallBase[]
-    for bb in fun.blocks, inst in bb.instructions
-        inst isa LLVM.CallBase || continue
-        called_intrinsic(inst) === nothing || push!(calls, inst)
-    end
-    isempty(calls) && return false
-    changed = false
-    @dispose builder=IRBuilder() begin
-        for call in calls
-            position!(builder, LLVM.before(call))
-            builder.debug_location = call.debug_location
-            new = lower(builder, call, called_intrinsic(call))
-            new === nothing && continue
-            new === :erase || replace_uses!(call, new)
-            erase!(call)
-            changed = true
-        end
-    end
-    return changed
-end
-
 # AIR has no vector floating-point min/max intrinsic; only the scalar `air.fmin`/`air.fmax`
 # exist. Julia's NaN-propagating `min`/`max` lower to `llvm.minimum`/`llvm.maximum` (and the
 # non-propagating `llvm.minnum`/`llvm.maxnum`), which LLVM's vectorizers can widen to vector
@@ -2678,24 +2598,7 @@ end
 # form, and semantically exact.
 function scalarize_vector_minmax!(fun::LLVM.Function)
     minmax = LLVM.Intrinsic.(["llvm.minnum", "llvm.maxnum", "llvm.minimum", "llvm.maximum"])
-    mod = fun.parent
-    return lower_intrinsic_calls!(fun) do builder, call, intr
-        vecty = call.value_type
-        (intr in minmax && vecty isa LLVM.VectorType) || return nothing
-        # the scalar overload of the same intrinsic, e.g. llvm.minimum.v4f32 -> llvm.minimum.f32
-        scalar_f = LLVM.Function(mod, intr, LLVMType[vecty.element_type])
-        scalar_ft = scalar_f.function_type
-        arg0, arg1 = call.arguments
-        res = PoisonValue(vecty)
-        for i in 0:Int(vecty.length)-1
-            idx = ConstantInt(LLVM.Int32Type(), i)
-            a = extract_element!(builder, arg0, idx)
-            b = extract_element!(builder, arg1, idx)
-            s = call!(builder, scalar_ft, scalar_f, LLVM.Value[a, b])
-            res = insert_element!(builder, res, s, idx)
-        end
-        res
-    end
+    scalarize_vector_calls!((call, intrinsic) -> intrinsic in minmax, fun)
 end
 
 # AIR has no vector reductions either, and Apple's back-end rejects `llvm.vector.reduce.*` or

src/optim.jl

@@ -60,9 +60,10 @@ function optimize!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
             register!(pb, GPULowerCPUFeaturesPass(job))
             register!(pb, GPULowerPTLSPass(job))
             register!(pb, GPULowerGCFramePass(job, relocs))
-            register!(pb, GPULinkRuntimePass(job, relocs))
-            register!(pb, GPULinkLibrariesPass(job))
-            register!(pb, GPUFinishRuntimeIntrinsicsPass(job))
+            imports = LateImports(job, relocs, false)
+            register!(pb, GPULowerIntrinsicsPass(imports))
+            register!(pb, GPULinkRuntimePass(imports))
+            register!(pb, GPULateImportsCleanupPass(imports, opt_level))
             register!(pb, AddKernelStatePass(job))
             register!(pb, LowerKernelStatePass(job))
             register!(pb, CleanupKernelStatePass(job))
@@ -321,18 +322,24 @@ end
 function buildIntrinsicLoweringPipeline(mpm, @nospecialize(job::CompilerJob), opt_level)
     add!(mpm, RemoveNIPass())
 
+    # lower operations the target cannot select, now that the optimizer is done with them,
+    # and import their implementations. this happens before GC lowering, so that it also
+    # applies to Julia code imported here.
+    # (use the registered passes, which share their state)
+    if job.config.libraries
+        add!(mpm, "GPULowerIntrinsics")
+    end
+
     # lower GC intrinsics
     if !uses_julia_runtime(job)
         add!(mpm, FunctionPassManager()) do fpm
-            # Use the registered pass because it owns `relocs`.
             add!(fpm, "GPULowerGCFrame")
         end
-        if job.config.libraries
-            # Use the registered pass because it owns `relocs`; the others only capture `job`.
-            add!(mpm, "GPULinkRuntime")
-            add!(mpm, GPULinkLibrariesPass(job))
-            add!(mpm, GPUFinishRuntimeIntrinsicsPass(job))
-        end
+    end
+
+    # GC lowering can introduce calls to the runtime
+    if job.config.libraries
+        add!(mpm, "GPULinkRuntime")
     end
 
     # lower kernel state intrinsics
@@ -404,6 +411,9 @@ function buildIntrinsicLoweringPipeline(mpm, @nospecialize(job::CompilerJob), op
     # Julia's operand bundles confuse the inliner, so repeat here now they are gone.
     # FIXME: we should fix the inliner so that inlined code gets optimized early-on
     add!(mpm, AlwaysInlinerPass())
+    if job.config.libraries
+        add!(mpm, "GPULateImportsCleanup")
+    end
 end
 
 function buildCleanupPipeline(mpm, @nospecialize(job::CompilerJob), opt_level)
@@ -483,51 +493,62 @@ function (self::CPUFeatures)(mod::LLVM.Module)
 end
 GPULowerCPUFeaturesPass(job) = ModulePass("GPULowerCPUFeatures", CPUFeatures(job))
 
-struct LinkRuntime
+# state shared by the passes that import device-library definitions after optimization
+mutable struct LateImports
     job::CompilerJob
     relocs::Relocations
+    # whether any definitions were imported, which then need to be optimized
+    imported::Bool
 end
-function (self::LinkRuntime)(mod::LLVM.Module)
-    self.job.config.libraries || return false
-    uses_julia_runtime(self.job) && return false
-
-    # GC lowering can introduce new calls to GPU runtime functions after the runtime
-    # was linked initially. Link again now so those calls resolve to definitions before
-    # later intrinsic-lowering passes inspect or rewrite the runtime call graph.
-    runtime, runtime_relocs = load_runtime(self.job)
-    # `RemoveNIPass` stripped non-integral address spaces from `mod`'s datalayout, but the
-    # cached runtime kept them; align it (as with target libraries) to avoid a warning.
-    runtime.triple = mod.triple
-    runtime.datalayout = mod.datalayout
-    link_relocatable!(mod, self.relocs, runtime, runtime_relocs; only_needed=true)
+
+struct LowerIntrinsics
+    state::LateImports
+end
+function (self::LowerIntrinsics)(mod::LLVM.Module)
+    state = self.state
+    state.imported |= lower_and_link_device_libraries!(state.job, mod, state.relocs)
     return true
 end
-GPULinkRuntimePass(job, relocs::Relocations) =
-    ModulePass("GPULinkRuntime", LinkRuntime(job, relocs))
+GPULowerIntrinsicsPass(state) = ModulePass("GPULowerIntrinsics", LowerIntrinsics(state))
 
-struct LinkLibraries
-    job::CompilerJob
+struct LinkRuntime
+    state::LateImports
 end
-function (self::LinkLibraries)(mod::LLVM.Module)
-    self.job.config.libraries || return false
-
-    # Runtime functions materialized by GC lowering can introduce target-library
-    # references after the normal library-linking phase has run. Give back-ends a
-    # second chance to resolve those references before they lower their own runtime
-    # intrinsics. This is a no-op for back-ends without a link_libraries! override.
-    link_libraries!(self.job, mod)
+function (self::LinkRuntime)(mod::LLVM.Module)
+    state = self.state
+    state.imported |= link_device_libraries_late!(state.job, mod, state.relocs)
+    finish_runtime_intrinsics!(state.job, mod)
     return true
 end
-GPULinkLibrariesPass(job) = ModulePass("GPULinkLibraries", LinkLibraries(job))
+GPULinkRuntimePass(state) = ModulePass("GPULinkRuntime", LinkRuntime(state))
 
-struct FinishRuntimeIntrinsics
-    job::CompilerJob
+# Definitions imported after optimization did not get optimized along with the kernel;
+# simplify them after they have been inlined.
+struct LateImportsCleanup
+    state::LateImports
+    opt_level::Int
 end
-function (self::FinishRuntimeIntrinsics)(mod::LLVM.Module)
-    return finish_runtime_intrinsics!(self.job, mod)
+function (self::LateImportsCleanup)(mod::LLVM.Module)
+    self.state.imported && self.opt_level >= 1 || return false
+    job = self.state.job
+    @dispose pb=PassBuilder() begin
+        tti = llvm_targetinfo(job.config.target)
+        tti === nothing || LLVM.target_transform_info!(pb, tti)
+        add!(pb, FunctionPassManager()) do fpm
+            add!(fpm, SROAPass())
+            add!(fpm, instcombine_pass(job))
+            add!(fpm, SimplifyCFGPass())
+            add!(fpm, self.opt_level >= 2 ? GVNPass() : EarlyCSEPass())
+            add!(fpm, DCEPass())
+        end
+        with_llvm_machine(job.config.target) do tm
+            run!(pb, mod, tm)
+        end
+    end
+    return true
 end
-GPUFinishRuntimeIntrinsicsPass(job) =
-    ModulePass("GPUFinishRuntimeIntrinsics", FinishRuntimeIntrinsics(job))
+GPULateImportsCleanupPass(state, opt_level) =
+    ModulePass("GPULateImportsCleanup", LateImportsCleanup(state, opt_level))
 
 # lower object allocations to to PTX malloc
 #