From 7eaba66a06e6a54fdae944ce2ab0cfd2dca9eb50 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 20 Oct 2026 15:55:50 +0200 Subject: [PATCH 1/4] Move intrinsic-lowering helpers out of the Metal back-end Other back-ends need to rewrite calls to LLVM intrinsics the same way Metal does, e.g., to lower math intrinsics to library calls. Generalize the vector min/max scalarization into `scalarize_vector_calls!`, which handles any elementwise intrinsic. --- src/metal.jl | 99 +---------------------------------- src/utils.jl | 141 ++++++++++++++++++++++++++++++++++++++++++++++++++ test/utils.jl | 17 ++++++ 3 files changed, 159 insertions(+), 98 deletions(-) diff --git a/src/metal.jl b/src/metal.jl index 7d5aca3c..ec1a42c1 100644 --- a/src/metal.jl +++ b/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 diff --git a/src/utils.jl b/src/utils.jl index 3fc15d79..512e70b0 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -522,3 +522,144 @@ function load_pointer_as_word!(builder::IRBuilder, load::LLVM.LoadInst, T_word:: erase!(load) return word end + +# 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]) + f = declare!(mod, name, ft) + call = call!(builder, ft, f, args) + # an existing definition may have been optimized to use a different convention + call.callconv = f.callconv + return call +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 + +""" + scalarize_vector_calls!(select, fun) + +Replace selected intrinsic calls with one scalar call per lane. `select(call, intrinsic)` +chooses elementwise intrinsics; scalar operands (such as `powi`'s exponent) are shared. +Fixed vectors and structures of fixed vectors are supported. Fast-math flags and debug +locations are preserved. Reductions and scalable vectors are not elementwise calls. +""" +function scalarize_vector_calls!(select, fun::LLVM.Function) + lower_intrinsic_calls!(fun) do builder, call, intrinsic + select(call, intrinsic) || return nothing + result_type = call.value_type + fields = result_type isa LLVM.StructType ? collect(result_type.elements) : + LLVMType[result_type] + all(typ -> typ isa LLVM.VectorType, fields) || return nothing + # capturing an abstractly-typed value would expose this code, and its callers, to + # invalidation by unrelated method definitions + lanes = (first(fields)::LLVM.VectorType).length + all(typ -> (typ::LLVM.VectorType).length == lanes, fields) || + error("Mismatched vector result lengths") + scalar_fields = LLVMType[typ.element_type for typ in fields] + scalar_result = result_type isa LLVM.StructType ? LLVM.StructType(scalar_fields) : + only(scalar_fields) + args = collect(call.arguments) + scalar_types = LLVMType[arg.value_type isa LLVM.VectorType ? + arg.value_type.element_type : arg.value_type for arg in args] + scalar_name = call.called_function.name + for typ in unique(vcat(fields, LLVMType[arg.value_type for arg in args])) + typ isa LLVM.VectorType || continue + scalar_name = replace(scalar_name, type_suffix(typ) => type_suffix(typ.element_type)) + end + scalar_type = LLVM.FunctionType(scalar_result, scalar_types) + scalar = declare!(fun.parent, scalar_name, scalar_type) + results = LLVM.Value[PoisonValue(typ) for typ in fields] + for lane in 0:lanes-1 + index = ConstantInt(LLVM.Int32Type(), lane) + lane_args = LLVM.Value[arg.value_type isa LLVM.VectorType ? + extract_element!(builder, arg, index) : arg for arg in args] + value = call!(builder, scalar_type, scalar, lane_args) + supports_fast_math(value) && (value.fast_math = call.fast_math) + for field in eachindex(fields) + element = result_type isa LLVM.StructType ? + extract_value!(builder, value, field - 1) : value + results[field] = insert_element!(builder, results[field], element, index) + end + end + result_type isa LLVM.StructType || return only(results) + result = PoisonValue(result_type) + for field in eachindex(fields) + result = insert_value!(builder, result, results[field], field - 1) + end + result + end +end + +@public call_declared!, lower_intrinsic_calls!, scalarize_vector_calls! diff --git a/test/utils.jl b/test/utils.jl index c36f2c40..294b3a35 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -308,3 +308,20 @@ end # Check that we can call this function from the CPU, to support deferred codegen for Enzyme. @test ccall("extern deferred_codegen", llvmcall, UInt, (UInt,), id) == id end + +@testset "elementwise scalarization" begin + JuliaContext() do ctx + @dispose mod=parse(LLVM.Module, """ + declare <2 x float> @llvm.powi.v2f32.i32(<2 x float>, i32) + define <2 x float> @test(<2 x float> %value, i32 %power) { + %result = call afn <2 x float> @llvm.powi.v2f32.i32(<2 x float> %value, i32 %power) + ret <2 x float> %result + } + """) begin + @test GPUCompiler.scalarize_vector_calls!((call, intrinsic) -> true, + mod.functions["test"]) + @test verify(mod) === nothing + @test count("call afn float @llvm.powi.f32.i32", string(mod)) == 2 + end + end +end From 815f28df98c40017c9d0fde3b16780d50f2006ef Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 20 Oct 2026 15:57:57 +0200 Subject: [PATCH 2/4] Add declarative device libraries Back-ends now describe the libraries that implement device functions with `device_libraries(job)`: Julia methods (`DeviceLibraryMethod`s, as used for the GPU runtime) or LLVM bitcode. Definitions are imported on demand, searching the libraries in order until no library provides a remaining reference, both before optimization (so library code is optimized along with the kernel) and after GC lowering (which introduces runtime calls). This replaces the separate runtime linking, and lets back-ends drop fixed-point linking loops over vendor libraries. `link_libraries!` keeps working. Julia libraries can be optimized on their own, for standalone numerical code, and bitcode libraries can be force-inlined, as AMDGPU does with the ROCm device libraries. --- README.md | 23 +++ src/GPUCompiler.jl | 3 + src/devicelib.jl | 342 +++++++++++++++++++++++++++++++++++++++++ src/driver.jl | 16 +- src/interface.jl | 4 +- src/optim.jl | 52 +------ src/rtlib.jl | 80 +++++----- test/device_library.jl | 204 ++++++++++++++++++++++++ test/helpers/metal.jl | 7 +- test/metal.jl | 3 +- 10 files changed, 628 insertions(+), 106 deletions(-) create mode 100644 src/devicelib.jl create mode 100644 test/device_library.jl diff --git a/README.md b/README.md index fc8988a2..f47ebb1e 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,29 @@ 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 before optimization, for everything the kernel references, so that library +code is optimized along with the kernel (as Clang does with `-mlink-builtin-bitcode`), +and again after GC lowering, which introduces calls to the runtime. + +Library selection 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 diff --git a/src/GPUCompiler.jl b/src/GPUCompiler.jl index 8231c94b..8a590397 100644 --- a/src/GPUCompiler.jl +++ b/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. diff --git a/src/devicelib.jl b/src/devicelib.jl new file mode 100644 index 00000000..47cb4813 --- /dev/null +++ b/src/devicelib.jl @@ -0,0 +1,342 @@ +# device libraries: definitions linked into kernels on demand + +""" + DeviceLibraryMethod(f, return_type, argument_types; llvm_name, + llvm_return_type=nothing, llvm_types=nothing) + +Describe an exported Julia specialization and its explicit LLVM ABI. Type generators +are evaluated in the compilation's LLVM context. +""" +function DeviceLibraryMethod(def::Function, return_type::Type, types::Tuple; + name=nameof(def), llvm_name="gpu_$name", + llvm_return_type=nothing, llvm_types=nothing) + Runtime.RuntimeMethodInstance(def, return_type, types, name, + llvm_return_type, llvm_types, llvm_name) +end + +""" + DeviceLibrary(name; methods, optimize=false, always_inline=optimize) + DeviceLibrary(name; bitcode, always_inline=false) + +Describe a library of device functions: Julia specializations (a tuple of +[`DeviceLibraryMethod`](@ref)s), or LLVM bitcode (a path or bytes). Bitcode is read once +per session, so files must not change while Julia is running. + +Libraries are linked on demand: only definitions the kernel references, and their +dependencies, are imported. Julia libraries are only compiled when used. + +`optimize=true` optimizes a Julia library on its own, and requires it to be standalone +device code without Julia runtime calls, address spaces, relocations, mutable globals or +indirect calls. `always_inline=true` forces imported definitions to be inlined; +otherwise the inliner decides. +""" +struct DeviceLibrary + name::Symbol + methods::Tuple + bitcode::Union{Nothing,String,Vector{UInt8}} + optimize::Bool + always_inline::Bool + # export manifest, computed on first use + exports::Base.RefValue{Union{Nothing,Set{String}}} + + function DeviceLibrary(name::Symbol; methods=nothing, bitcode=nothing, + optimize=false, always_inline=optimize) + (methods === nothing) != (bitcode === nothing) || + throw(ArgumentError("Specify exactly one of methods and bitcode")) + if methods !== nothing + methods isa Tuple || throw(ArgumentError("methods must be a tuple")) + all(method -> method isa Runtime.RuntimeMethodInstance, methods) || + throw(ArgumentError("Expected DeviceLibraryMethod descriptors")) + allunique(method.llvm_name for method in methods) || + throw(ArgumentError("Duplicate exports in $name")) + exports = Set(method.llvm_name for method in methods) + else + optimize && throw(ArgumentError("Only Julia libraries can be optimized")) + bitcode isa Union{String,Vector{UInt8}} || + throw(ArgumentError("bitcode must be a path or Vector{UInt8}")) + exports = nothing + end + new(name, something(methods, ()), bitcode, optimize, always_inline, + Ref{Union{Nothing,Set{String}}}(exports)) + end +end + +""" + runtime_library(mod::Module) + +Describe the target's implementations of GPUCompiler's runtime functions (e.g. +`signal_exception` or `malloc`) defined in `mod`, using GPUCompiler's signatures. +""" +function runtime_library(mod::Module) + methods = Tuple(DeviceLibraryMethod(getfield(mod, method.def), + method.return_type, method.types; + method.name, method.llvm_name, + method.llvm_return_type, method.llvm_types) + for method in runtime_methods() + if method.def isa Symbol && isdefined(mod, method.def)) + DeviceLibrary(nameof(mod); methods) +end + +""" + device_libraries(job) -> Tuple + +Libraries to resolve the job's external references with, in search order. The default +uses `runtime_module(job)`; back-ends overriding this hook should include +`runtime_library(MyRuntime)` themselves. GPUCompiler's own runtime is added for targets +that do not use Julia's runtime. + +The selection must only depend on the job's configuration, as it is not part of any cache +key, and Julia library methods must be resolvable through the job's method table. +""" +device_libraries(@nospecialize(job::CompilerJob)) = + uses_julia_runtime(job) ? () : (runtime_library(runtime_module(job)),) + +function selected_device_libraries(@nospecialize(job::CompilerJob)) + libraries = device_libraries(job) + libraries isa Tuple || error("device_libraries must return a tuple") + uses_julia_runtime(job) && return libraries + core = DeviceLibrary(:gpu_runtime; + methods=Tuple(method for method in runtime_methods() + if !(method.def isa Symbol))) + return (core, libraries...) +end + + +## bitcode libraries + +const bitcode_files = Dict{String,Vector{UInt8}}() +const bitcode_files_lock = ReentrantLock() + +function library_bytes(lib::DeviceLibrary) + lib.bitcode isa Vector{UInt8} && return lib.bitcode + path = lib.bitcode::String + Base.@lock bitcode_files_lock get!(() -> read(path), bitcode_files, path) +end + +# the names a library can define +function library_exports(lib::DeviceLibrary) + exports = lib.exports[] + exports === nothing || return exports + @dispose ctx=Context() mod=parse(LLVM.Module, MemoryBuffer(library_bytes(lib)); lazy=true) begin + exports = Set{String}() + for value in Iterators.flatten((mod.functions, mod.globals, mod.aliases)) + value isa LLVM.GlobalAlias || !isdeclaration(value) || continue + value.linkage in (LLVM.Linkage.Private, LLVM.Linkage.Internal, + LLVM.Linkage.AvailableExternally) && continue + push!(exports, value.name) + end + end + lib.exports[] = exports +end + + +## Julia libraries + +# Julia library artifacts are cached per session. +const device_libs = Dict{Any,RuntimeLibrary}() +const device_libs_lock = ReentrantLock() + +@locked function load_julia_library(lib::DeviceLibrary, @nospecialize(job::CompilerJob)) + config = runtime_config(job) + # a `Vector{Any}` key avoids compiling `Dict` methods for every library's tuple type + key = Any[lib.name, lib.optimize, lib.always_inline, config, + !supports_typed_pointers(context()), + lib.optimize ? job.config.opt_level : 0] + for method in lib.methods + push!(key, method.def, method.return_type, method.types, method.name, + method.llvm_return_type, method.llvm_types, method.llvm_name) + end + cached = Base.@lock device_libs_lock begin + cached = get(device_libs, key, nothing) + if cached === nothing || !device_library_valid(cached, job, lib.methods) + cached = build_device_library(job, config, lib) + device_libs[key] = cached + end + cached + end + parse(LLVM.Module, MemoryBuffer(cached.bytes)), cached.relocations +end + +function build_device_library(@nospecialize(job::CompilerJob), config::CompilerConfig, + lib::DeviceLibrary) + @dispose mod=LLVM.Module(string(lib.name)) begin + sources = MethodInstance[] + instances = CodeInstance[] + relocs = Relocations() + for method in lib.methods + source = runtime_method_instance(job, method) + source === nothing && error("No method for $(method.name) in library $(lib.name)") + push!(sources, source) + push!(instances, emit_function!(mod, relocs, config, source, method, job.world)) + end + if lib.optimize + standalone = CompilerJob(first(sources), + CompilerConfig(config; libraries=false, opt_level=job.config.opt_level), + job.world) + optimize!(standalone, mod, relocs; opt_level=job.config.opt_level) + prune_dead_relocations!(mod, relocs) + validate_device_library(mod, relocs, lib) + end + if lib.always_inline + for method in lib.methods + attrs = mod.functions[method.llvm_name].function_attributes + delete!(attrs, EnumAttribute("noinline")) + push!(attrs, EnumAttribute("alwaysinline")) + end + end + io = IOBuffer() + write(io, mod) + return RuntimeLibrary(take!(io), sources, instances, job.world, relocs) + end +end + +function validate_device_library(mod, relocs, lib) + verify(mod) + isempty(relocs) || error("Optimized library $(lib.name) contains relocations") + runtime_names = Set(method.llvm_name for method in runtime_methods()) + function check_type(typ, seen=Set{LLVMType}()) + typ in seen && return + push!(seen, typ) + if typ isa LLVM.PointerType + typ.addrspace in 10:13 && + error("Optimized library $(lib.name) contains Julia address spaces") + supports_typed_pointers(context()) && check_type(typ.element_type, seen) + elseif typ isa LLVM.FunctionType + check_type(typ.return_type, seen) + foreach(typ -> check_type(typ, seen), typ.parameters) + elseif typ isa Union{LLVM.ArrayType,LLVM.VectorType} + check_type(typ.element_type, seen) + elseif typ isa LLVM.StructType && !isopaque(typ) + foreach(typ -> check_type(typ, seen), typ.elements) + end + end + for global_var in mod.globals + (global_var.constant && global_var.linkage == LLVM.Linkage.Private) || + error("Optimized library $(lib.name) contains globals that are not private constants") + check_type(global_var.global_value_type) + end + for fun in mod.functions + check_type(fun.function_type) + if isdeclaration(fun) && !isempty(fun.uses) && + (startswith(fun.name, "julia.") || startswith(fun.name, "jl_") || + startswith(fun.name, "ijl_") || fun.name in runtime_names) + error("Optimized library $(lib.name) contains runtime operation $(fun.name)") + end + for block in fun.blocks, inst in block.instructions + check_type(inst.value_type) + foreach(value -> check_type(value.value_type), inst.operands) + inst isa LLVM.CallBase && inst.called_function === nothing && + !(inst.called_operand isa LLVM.InlineAsm) && + error("Optimized library $(lib.name) contains an indirect call") + end + end + return +end + + +## resolution + +# Whether a definition of type `actual` can be called through a declaration of type +# `expected`. Typed-pointer element types and struct names are not part of the ABI. +abi_compatible(expected::LLVMType, actual::LLVMType) = expected == actual +abi_compatible(expected::LLVM.PointerType, actual::LLVM.PointerType) = + expected.addrspace == actual.addrspace +function abi_compatible(expected::LLVM.FunctionType, actual::LLVM.FunctionType) + LLVM.isvararg(expected) == LLVM.isvararg(actual) && + abi_compatible(expected.return_type, actual.return_type) && + length(expected.parameters) == length(actual.parameters) && + all(abi_compatible(e, a) for (e, a) in zip(expected.parameters, actual.parameters)) +end +function abi_compatible(expected::LLVM.StructType, actual::LLVM.StructType) + expected == actual && return true + (isopaque(expected) || isopaque(actual)) && return false + ispacked(expected) == ispacked(actual) && + length(expected.elements) == length(actual.elements) && + all(abi_compatible(e, a) for (e, a) in zip(expected.elements, actual.elements)) +end +function abi_compatible(expected::T, actual::T) where {T<:Union{LLVM.ArrayType,LLVM.VectorType}} + length(expected) == length(actual) && + abi_compatible(expected.element_type, actual.element_type) +end + +needed_symbols(mod::LLVM.Module) = + Set(value.name for value in Iterators.flatten((mod.functions, mod.globals)) + if isdeclaration(value) && !isempty(value.uses)) + +# Import the definitions `mod` references from the job's device libraries, until no +# library provides any remaining reference. Libraries are searched in order. Returns the +# names of imported functions and aliases. +function link_device_libraries!(@nospecialize(job::CompilerJob), mod::LLVM.Module, + relocs::Relocations) + libraries = selected_device_libraries(job) + definitions() = Iterators.flatten((Iterators.filter(!isdeclaration, mod.functions), + mod.aliases)) + defined = Set(value.name for value in definitions()) + needed = needed_symbols(mod) + progress = true + while progress && !isempty(needed) + progress = false + for lib in libraries + names = intersect(needed, library_exports(lib)) + isempty(names) && continue + library, library_relocs = if lib.bitcode === nothing + load_julia_library(lib, job) + else + parse(LLVM.Module, MemoryBuffer(library_bytes(lib)); lazy=true), Relocations() + end + try + for name in names + expected = get(mod.functions, name, nothing) + actual = get(library.functions, name, nothing) + actual = something(actual, get(library.aliases, name, nothing), Some(nothing)) + (expected === nothing || actual === nothing) && continue + actual_type = LLVM.global_value_type(actual) + abi_compatible(expected.function_type, actual_type) || + error("Device library $(lib.name) defines $name as $actual_type, incompatible with its use as $(expected.function_type)") + end + if lib.always_inline && lib.bitcode !== nothing + for fun in library.functions + isdeclaration(fun) && continue + attrs = fun.function_attributes + haskey(attrs, :noinline) || push!(attrs, EnumAttribute("alwaysinline")) + end + end + catch + dispose(library) + rethrow() + end + # avoid warnings about mismatching triples and data layouts + library.triple = mod.triple + library.datalayout = mod.datalayout + link_relocatable!(mod, relocs, library, library_relocs; only_needed=true) + needed = needed_symbols(mod) + progress = true + end + end + return [value.name for value in definitions() if !(value.name in defined)] +end + +# Before optimization: import every referenced definition, so that it gets optimized +# along with the kernel. +function link_device_libraries_early!(@nospecialize(job::CompilerJob), mod::LLVM.Module, + relocs::Relocations) + @tracepoint "device libraries" link_device_libraries!(job, mod, relocs) + @tracepoint "target libraries" link_libraries!(job, mod) + return +end + +# After optimization: import definitions for references introduced by lowering passes. +# Definitions imported before optimization that are still around are reused (the +# optimizer does not change their signatures). +function link_device_libraries_late!(@nospecialize(job::CompilerJob), mod::LLVM.Module, + relocs::Relocations) + imported = link_device_libraries!(job, mod, relocs) + link_libraries!(job, mod) + for name in imported + value = something(get(mod.functions, name, nothing), get(mod.aliases, name, nothing)) + value.linkage = LLVM.Linkage.Internal + end + return +end + +@public DeviceLibrary, DeviceLibraryMethod, device_libraries, runtime_library diff --git a/src/driver.jl b/src/driver.jl index 13c3f81d..0ac2754b 100644 --- a/src/driver.jl +++ b/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 diff --git a/src/interface.jl b/src/interface.jl index 1b656eaf..223f0fcc 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -742,8 +742,8 @@ 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 # finalization of the module, before deferred codegen and optimization diff --git a/src/optim.jl b/src/optim.jl index b27f3a37..09abb2f6 100644 --- a/src/optim.jl +++ b/src/optim.jl @@ -61,8 +61,6 @@ function optimize!(@nospecialize(job::CompilerJob), mod::LLVM.Module, register!(pb, GPULowerPTLSPass(job)) register!(pb, GPULowerGCFramePass(job, relocs)) register!(pb, GPULinkRuntimePass(job, relocs)) - register!(pb, GPULinkLibrariesPass(job)) - register!(pb, GPUFinishRuntimeIntrinsicsPass(job)) register!(pb, AddKernelStatePass(job)) register!(pb, LowerKernelStatePass(job)) register!(pb, CleanupKernelStatePass(job)) @@ -327,12 +325,12 @@ function buildIntrinsicLoweringPipeline(mpm, @nospecialize(job::CompilerJob), op # 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 + # Use the registered pass because it owns `relocs`. + add!(mpm, "GPULinkRuntime") end # lower kernel state intrinsics @@ -488,47 +486,13 @@ struct LinkRuntime relocs::Relocations 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) + link_device_libraries_late!(self.job, mod, self.relocs) + finish_runtime_intrinsics!(self.job, mod) return true end GPULinkRuntimePass(job, relocs::Relocations) = ModulePass("GPULinkRuntime", LinkRuntime(job, relocs)) -struct LinkLibraries - job::CompilerJob -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) - return true -end -GPULinkLibrariesPass(job) = ModulePass("GPULinkLibraries", LinkLibraries(job)) - -struct FinishRuntimeIntrinsics - job::CompilerJob -end -function (self::FinishRuntimeIntrinsics)(mod::LLVM.Module) - return finish_runtime_intrinsics!(self.job, mod) -end -GPUFinishRuntimeIntrinsicsPass(job) = - ModulePass("GPUFinishRuntimeIntrinsics", FinishRuntimeIntrinsics(job)) - # lower object allocations to to PTX malloc # # this is a PoC implementation that is very simple: allocate, and never free. it also runs diff --git a/src/rtlib.jl b/src/rtlib.jl index 6c2d16a5..485bff41 100644 --- a/src/rtlib.jl +++ b/src/rtlib.jl @@ -73,7 +73,10 @@ function LLVM.call!(builder, rt::Runtime.RuntimeMethodInstance, args=LLVM.Value[ end end - call!(builder, ft, f, args) + call = call!(builder, ft, f, args) + # an existing definition may have been optimized to use a different convention + call.callconv = f.callconv + return call end @@ -83,10 +86,10 @@ end # required relocation metadata, this persists with the function's `CodeInstance`. # `runtime_libs` below provides the session cache on older Julia versions. mutable struct RuntimeFunctionResults - bitcode::Union{Nothing,Vector{UInt8}} - relocations::Relocations - RuntimeFunctionResults() = new(nothing, Relocations()) + artifacts::Dict{String,Tuple{Vector{UInt8},Relocations}} + RuntimeFunctionResults() = new(Dict{String,Tuple{Vector{UInt8},Relocations}}()) end +const runtime_function_results_lock = ReentrantLock() # Compile a single runtime function and link it into `mod`. The renamed bitcode is # memoized through `RuntimeFunctionResults`; the session-local `runtime_libs` cache @@ -99,10 +102,16 @@ function emit_function!(mod, relocs::Relocations, config::CompilerConfig, # Don't run a standalone inference walk on a miss: `compile_unhooked` below drives # inference itself. ci, res = runtime_function_results(rt_job) - if res !== nothing && res.bitcode !== nothing - link_relocatable!(mod, relocs, - parse(LLVM.Module, res.bitcode), - res.relocations) + cached = res === nothing ? nothing : + Base.@lock runtime_function_results_lock get(res.artifacts, name, nothing) + if cached !== nothing + bytes, relocations = cached + artifact = parse(LLVM.Module, bytes) + if artifact.functions[name].function_type != runtime_function_type(rt_job, method) + dispose(artifact) + error("Invalid ABI for cached device-library function '$name'") + end + link_relocatable!(mod, relocs, artifact, relocations) ci === nothing && (ci = runtime_code_instance(rt_job)) return ci::CodeInstance end @@ -112,11 +121,13 @@ function emit_function!(mod, relocs::Relocations, config::CompilerConfig, new_mod, meta = compile_unhooked(:llvm, rt_job; resolve_relocations=false) ft = meta.entry.function_type expected_ft = runtime_function_type(rt_job, method) - if ft.return_type != expected_ft.return_type - error("Invalid return type for runtime function '$(method.name)': expected $(expected_ft.return_type), got $(ft.return_type)") + if ft != expected_ft + dispose(new_mod) + error("Invalid ABI for device-library function '$(method.name)': expected $expected_ft, got $ft") end # recent Julia versions include prototypes for all runtime functions, even if unused + collect_cglobal_relocations!(rt_job, new_mod, meta.relocations) with_llvm_machine(config.target) do tm run!(StripDeadPrototypesPass(), new_mod, tm) end @@ -137,8 +148,9 @@ function emit_function!(mod, relocs::Relocations, config::CompilerConfig, ci === nothing && (ci = runtime_code_instance(rt_job)) if supports_relocatable_ir() res === nothing && (res = runtime_results(RuntimeFunctionResults, ci, rt_job.config)) - res.bitcode = take!(io) - res.relocations = meta.relocations + Base.@lock runtime_function_results_lock begin + res.artifacts[name] = (take!(io), meta.relocations) + end end link_relocatable!(mod, relocs, new_mod, meta.relocations) @@ -201,32 +213,12 @@ function runtime_config(@nospecialize(job::CompilerJob)) toplevel=false, only_entry=false, strip=false, name=nothing) end -# build the runtime library, returning its bitcode +# build the runtime library as a single artifact function build_runtime(@nospecialize(job::CompilerJob), config::CompilerConfig) - @dispose mod=LLVM.Module("GPUCompiler run-time library") begin - sources = MethodInstance[] - code_instances = CodeInstance[] - relocs = Relocations() - - for method in runtime_methods() - resolved = runtime_method_instance(job, method) - resolved === nothing && continue - source = resolved - push!(sources, source) - push!(code_instances, - emit_function!(mod, relocs, config, source, method, job.world)) - end - - # we cannot optimize the runtime library, because the code would then be optimized - # again during main compilation (and optimizing twice isn't safe). for example, - # optimization removes Julia address spaces, which would then lead to type - # mismatches when using functions from the runtime library from IR that has not been - # stripped of AS info. - - io = IOBuffer() - write(io, mod) - return take!(io), sources, code_instances, relocs - end + methods = [method for method in runtime_methods() if !(method.def isa Symbol)] + append!(methods, runtime_library(runtime_module(job)).methods) + sort!(methods; by=method -> method.name) + build_device_library(job, config, DeviceLibrary(:runtime; methods=Tuple(methods))) end # Runtime.methods is a Dict, but library layout and source validation require the same order. @@ -251,14 +243,15 @@ mutable struct RuntimeLibrary relocations::Relocations end -function runtime_library_valid(lib::RuntimeLibrary, @nospecialize(job::CompilerJob)) +function device_library_valid(lib::RuntimeLibrary, @nospecialize(job::CompilerJob), + @nospecialize(methods)) # Method-table changes and CI invalidations always advance the world counter. A runtime # already validated for this exact world therefore needs no per-function scan on the # common cache-hit path. job.world == lib.validated_world && return true i = 0 - for method in runtime_methods() + for method in methods resolved = runtime_method_instance(job, method) resolved === nothing && continue i += 1 @@ -271,9 +264,15 @@ function runtime_library_valid(lib::RuntimeLibrary, @nospecialize(job::CompilerJ return true end +runtime_library_valid(lib::RuntimeLibrary, @nospecialize(job::CompilerJob)) = + device_library_valid(lib, job, runtime_methods()) + const runtime_libs = Dict{Tuple{CompilerConfig, Bool}, RuntimeLibrary}() const runtime_libs_lock = ReentrantLock() +# Load the runtime library as a single module. This is not used by the compiler, which +# imports runtime functions on demand (see `link_device_libraries!`), but back-ends call it +# to populate the per-function caches, e.g., during precompilation. @locked function load_runtime(@nospecialize(job::CompilerJob)) config = runtime_config(job) key = (config, !supports_typed_pointers(context())) @@ -281,8 +280,7 @@ const runtime_libs_lock = ReentrantLock() cached = Base.@lock runtime_libs_lock begin cached = get(runtime_libs, key, nothing) if cached === nothing || !runtime_library_valid(cached, job) - bytes, sources, code_instances, relocations = build_runtime(job, config) - cached = RuntimeLibrary(bytes, sources, code_instances, job.world, relocations) + cached = build_runtime(job, config) runtime_libs[key] = cached end cached diff --git a/test/device_library.jl b/test/device_library.jl new file mode 100644 index 00000000..cd802ee6 --- /dev/null +++ b/test/device_library.jl @@ -0,0 +1,204 @@ +using Test + +module LibraryFixtures +using GPUCompiler, LLVM, LLVM.IR, LLVM.Build + +# a runtime whose allocator doesn't fold away +module Runtime + malloc(sz) = Ptr{Cvoid}(Int(0xDEADBEEF)) + signal_exception() = return + report_oom(sz) = return + report_exception(ex) = return + report_exception_name(ex) = return + report_exception_frame(idx, func, file, line) = return +end + +helper(value::Int32) = value + Int32(7) +vector_export(::Vector) = Int32(7) +exported(value::Int32) = helper(value) +caller(value::Int32) = ccall("extern test_export", llvmcall, Int32, (Int32,), value) +alias_caller(value::Int32) = + ccall("extern test_alias", llvmcall, Int32, (Int32,), value) +other_caller(value::Int32) = + ccall("extern test_export", llvmcall, Int32, (Int32,), value) + + ccall("extern other", llvmcall, Int32, (Int32,), value) +const julia_library = GPUCompiler.DeviceLibrary(:test; + methods=(GPUCompiler.DeviceLibraryMethod(exported, Int32, (Int32,); + llvm_name="test_export"), + GPUCompiler.DeviceLibraryMethod(exported, Int32, (Int32,); + llvm_name="test_alias"))) + +struct Params <: GPUCompiler.AbstractCompilerParams + libraries::Tuple +end +GPUCompiler.runtime_module(::CompilerJob{<:Any,Params}) = Runtime +GPUCompiler.device_libraries(job::CompilerJob{<:Any,Params}) = + (job.config.params.libraries..., GPUCompiler.runtime_library(Runtime)) +function job(func, types; libs=(julia_library,), target=:native, kwargs...) + source = GPUCompiler.methodinstance(typeof(func), Base.to_tuple_type(types)) + target = target === :native ? NativeCompilerTarget() : PTXCompilerTarget(; cap=v"7.0") + validate = target isa PTXCompilerTarget + CompilerJob(source, CompilerConfig(target, Params(libs); + kernel=false, validate, kwargs...)) +end +function ir(func, types=(Int32,); kwargs...) + JuliaContext() do ctx + mod, meta = GPUCompiler.compile(:llvm, job(func, types; kwargs...)) + try + @assert LLVM.verify(mod) === nothing + string(mod) + finally + LLVM.dispose(mod) + end + end +end + +bitcode(text) = JuliaContext() do ctx + @dispose mod=parse(LLVM.Module, text) begin + buffer = IOBuffer() + write(buffer, mod) + take!(buffer) + end +end +end + +@testset "Julia libraries" begin + # imported before optimization, and optimized with the caller + text = LibraryFixtures.ir(LibraryFixtures.caller) + @test occursin("define internal fastcc i32 @test_export", text) + @test contains(text, r"add i32 %\S+, 7") + @test occursin("define internal fastcc i32 @test_alias", + LibraryFixtures.ir(LibraryFixtures.alias_caller)) + + @test occursin("define", LibraryFixtures.ir(LibraryFixtures.caller; optimize=false)) + @test occursin("declare i32 @test_export", + LibraryFixtures.ir(LibraryFixtures.caller; libraries=false)) + + # equivalent descriptors share cached artifacts + vectors(T) = GPUCompiler.DeviceLibrary(:vectors; + methods=(GPUCompiler.DeviceLibraryMethod(LibraryFixtures.vector_export, Int32, + (UnionAll(T, Vector{T}),); + llvm_name="vector_export"),)) + job = LibraryFixtures.job(LibraryFixtures.caller, (Int32,)) + cached = JuliaContext() do ctx + GPUCompiler.load_julia_library(vectors(TypeVar(:T, Real)), job) + length(GPUCompiler.device_libs) + end + @test cached == JuliaContext() do ctx + GPUCompiler.load_julia_library(vectors(TypeVar(:S, Real)), job) + length(GPUCompiler.device_libs) + end + + # libraries follow method redefinitions + @eval LibraryFixtures helper(value::Int32) = value + Int32(19) + @test contains(LibraryFixtures.ir(LibraryFixtures.caller), r"add i32 %\S+, 19") +end + +@testset "optimized Julia libraries" begin + library = GPUCompiler.DeviceLibrary(:optimized; + methods=LibraryFixtures.julia_library.methods, optimize=true) + @test !occursin("@test_export", + LibraryFixtures.ir(LibraryFixtures.caller; libs=(library,))) + + @test_throws ArgumentError GPUCompiler.DeviceLibrary(:bad; bitcode=UInt8[], optimize=true) + @test_throws "Duplicate" GPUCompiler.DeviceLibrary(:bad; + methods=(first(LibraryFixtures.julia_library.methods), + first(LibraryFixtures.julia_library.methods))) +end + +@testset "bitcode libraries" begin + bitcode = LibraryFixtures.bitcode + upstream = GPUCompiler.DeviceLibrary(:upstream; bitcode=bitcode(""" + declare i32 @downstream(i32) + define i32 @test_export(i32 %value) { + %result = call i32 @downstream(i32 %value) + ret i32 %result + }""")) + downstream = GPUCompiler.DeviceLibrary(:downstream; bitcode=bitcode(""" + define i32 @downstream(i32 %value) { + %result = mul i32 %value, 37 + ret i32 %result + }""")) + + # dependencies are resolved regardless of the order of libraries + for libraries in ((upstream, downstream), (downstream, upstream)) + text = LibraryFixtures.ir(LibraryFixtures.caller; libs=libraries) + @test occursin("mul i32 %value, 37", text) + @test !occursin("declare", text) + end + + # unresolved references are left to validation + @test occursin("declare i32 @downstream", + LibraryFixtures.ir(LibraryFixtures.caller; libs=(upstream,))) + + # libraries can share linkonce definitions + shared(name, constant) = GPUCompiler.DeviceLibrary(Symbol(name); bitcode=bitcode(""" + define linkonce_odr i32 @shared() { ret i32 $constant } + define i32 @$name(i32 %value) { + %shared = call i32 @shared() + %result = add i32 %value, %shared + ret i32 %result + }""")) + text = LibraryFixtures.ir(LibraryFixtures.other_caller; + libs=(shared("test_export", 41), shared("other", 41))) + @test occursin("ret i32 41", text) + @test count("define internal fastcc i32 @shared", text) == 1 + @test !occursin("declare", text) + + # aliases are exports too + aliased = GPUCompiler.DeviceLibrary(:aliased; bitcode=bitcode(""" + define i32 @implementation(i32 %value) { + %result = mul i32 %value, 53 + ret i32 %result + } + @test_export = alias i32 (i32), i32 (i32)* @implementation""")) + text = LibraryFixtures.ir(LibraryFixtures.caller; libs=(aliased,)) + @test occursin("mul i32 %value, 53", text) + @test !occursin("declare", text) + mismatch = GPUCompiler.DeviceLibrary(:mismatch; bitcode=bitcode(""" + define float @implementation(float %value) { ret float %value } + @test_export = alias float (float), float (float)* @implementation""")) + @test_throws "incompatible" LibraryFixtures.ir(LibraryFixtures.caller; libs=(mismatch,)) + + # struct types are compared structurally + returns_struct = GPUCompiler.DeviceLibrary(:returns_struct; bitcode=bitcode(""" + %pair = type { i32, i32 } + declare %pair @make_pair(i32) + define i32 @test_export(i32 %value) { + %pair = call %pair @make_pair(i32 %value) + %result = extractvalue %pair %pair, 1 + ret i32 %result + }""")) + makes_struct = GPUCompiler.DeviceLibrary(:makes_struct; bitcode=bitcode(""" + %pair = type { i32, i32 } + define %pair @make_pair(i32 %value) { + %first = insertvalue %pair undef, i32 %value, 0 + %second = insertvalue %pair %first, i32 43, 1 + ret %pair %second + }""")) + text = LibraryFixtures.ir(LibraryFixtures.caller; libs=(returns_struct, makes_struct)) + @test occursin("i32 43", text) + @test !occursin("declare", text) + + # incompatible definitions are rejected + mismatch = GPUCompiler.DeviceLibrary(:mismatch; bitcode=bitcode(""" + define float @test_export(float %value) { ret float %value }""")) + @test_throws "incompatible" LibraryFixtures.ir(LibraryFixtures.caller; + libs=(mismatch,)) + + # forced inlining + noinline = GPUCompiler.DeviceLibrary(:noinline; bitcode=bitcode(""" + define i32 @test_export(i32 %value) noinline { + %result = mul i32 %value, 47 + ret i32 %result + }""")) + inline = GPUCompiler.DeviceLibrary(:inline; bitcode=bitcode(""" + define i32 @test_export(i32 %value) { + %result = mul i32 %value, 47 + ret i32 %result + }"""), always_inline=true) + @test occursin("call fastcc i32 @test_export", + LibraryFixtures.ir(LibraryFixtures.caller; libs=(noinline,))) + @test !occursin("@test_export", + LibraryFixtures.ir(LibraryFixtures.caller; libs=(inline,))) +end diff --git a/test/helpers/metal.jl b/test/helpers/metal.jl index 577d0262..6b77834c 100644 --- a/test/helpers/metal.jl +++ b/test/helpers/metal.jl @@ -14,9 +14,10 @@ module ThreadedRuntime signal_exception() = return malloc(sz) = C_NULL - function report_oom(sz) - ccall("extern julia.air.thread_position_in_grid.v3i32", llvmcall, - NTuple{3, VecElement{UInt32}}, ()) + @noinline function report_oom(sz) + position = ccall("extern julia.air.thread_position_in_grid.v3i32", llvmcall, + NTuple{3, VecElement{UInt32}}, ()) + parentmodule(parentmodule(@__MODULE__)).sink(position[1].value) return end report_exception(ex) = return diff --git a/test/metal.jl b/test/metal.jl index 658a79e6..ee5300bf 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -572,7 +572,8 @@ end Metal.code_llvm_threaded_runtime(io, mod.kernel, Tuple{Int}; kernel=true, dump_module=true) end - @test occursin(r"define .*@gpu_gc_pool_alloc\([^)]*<3 x i32>", ir) + @test !occursin("julia.air.thread_position_in_grid", ir) + @test occursin(r"call .*@gpu_report_oom\([^)]*<3 x i32>", ir) @test occursin("thread_position_in_grid", ir) @test !occursin(r"declare .*@gpu_gc_pool_alloc\(", ir) From 1fd96509214125ad9b5a1559d5d08dbe903b6c0f Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 20 Oct 2026 15:59:41 +0200 Subject: [PATCH 3/4] Lower unselectable operations late, through device libraries Some LLVM intrinsics have no lowering on GPU targets, e.g., math functions on vectors (#679) or any double-precision operation on hardware without it. Keeping them as intrinsics during optimization lets LLVM reason about them, so lower them afterwards: the new `lower_intrinsics!` hook runs after optimization, before GC lowering, and its library calls are resolved like any other reference. Julia libraries imported at that point are lowered by the same hook, so their ABI matches the lowered caller. `finish_imports!` lets back-ends prepare late imports, which are simplified once inlined. --- README.md | 14 +++++-- src/devicelib.jl | 48 +++++++++++++++------ src/driver.jl | 2 + src/interface.jl | 22 ++++++++++ src/optim.jl | 73 ++++++++++++++++++++++++++++---- src/ptx.jl | 3 ++ test/device_library.jl | 94 +++++++++++++++++++++++++++++++++++++++++- 7 files changed, 229 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index f47ebb1e..d3fde856 100644 --- a/README.md +++ b/README.md @@ -38,11 +38,17 @@ implementation of GPUCompiler's runtime, and is selected by default for 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 before optimization, for everything the kernel references, so that library -code is optimized along with the kernel (as Clang does with `-mlink-builtin-bitcode`), -and again after GC lowering, which introduces calls to the runtime. +This happens twice: -Library selection should only depend on the job's configuration, which +- 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. diff --git a/src/devicelib.jl b/src/devicelib.jl index 47cb4813..9c8804a1 100644 --- a/src/devicelib.jl +++ b/src/devicelib.jl @@ -132,15 +132,17 @@ end ## Julia libraries -# Julia library artifacts are cached per session. +# Julia library artifacts are cached per session. Artifacts linked before optimization +# keep the original ABI; those linked afterwards have been lowered by `lower_intrinsics!`. const device_libs = Dict{Any,RuntimeLibrary}() const device_libs_lock = ReentrantLock() -@locked function load_julia_library(lib::DeviceLibrary, @nospecialize(job::CompilerJob)) +@locked function load_julia_library(lib::DeviceLibrary, @nospecialize(job::CompilerJob); + lowered::Bool) config = runtime_config(job) # a `Vector{Any}` key avoids compiling `Dict` methods for every library's tuple type key = Any[lib.name, lib.optimize, lib.always_inline, config, - !supports_typed_pointers(context()), + !supports_typed_pointers(context()), lowered, lib.optimize ? job.config.opt_level : 0] for method in lib.methods push!(key, method.def, method.return_type, method.types, method.name, @@ -149,7 +151,7 @@ const device_libs_lock = ReentrantLock() cached = Base.@lock device_libs_lock begin cached = get(device_libs, key, nothing) if cached === nothing || !device_library_valid(cached, job, lib.methods) - cached = build_device_library(job, config, lib) + cached = build_device_library(job, config, lib; lowered) device_libs[key] = cached end cached @@ -158,7 +160,7 @@ const device_libs_lock = ReentrantLock() end function build_device_library(@nospecialize(job::CompilerJob), config::CompilerConfig, - lib::DeviceLibrary) + lib::DeviceLibrary; lowered::Bool=false) @dispose mod=LLVM.Module(string(lib.name)) begin sources = MethodInstance[] instances = CodeInstance[] @@ -177,6 +179,13 @@ function build_device_library(@nospecialize(job::CompilerJob), config::CompilerC prune_dead_relocations!(mod, relocs) validate_device_library(mod, relocs, lib) end + if lowered + lower_intrinsics!(CompilerJob(first(sources), config, job.world), mod) + for method in lib.methods + haskey(mod.functions, method.llvm_name) || + error("Lowering library $(lib.name) removed export $(method.llvm_name)") + end + end if lib.always_inline for method in lib.methods attrs = mod.functions[method.llvm_name].function_attributes @@ -267,7 +276,7 @@ needed_symbols(mod::LLVM.Module) = # library provides any remaining reference. Libraries are searched in order. Returns the # names of imported functions and aliases. function link_device_libraries!(@nospecialize(job::CompilerJob), mod::LLVM.Module, - relocs::Relocations) + relocs::Relocations; lowered::Bool) libraries = selected_device_libraries(job) definitions() = Iterators.flatten((Iterators.filter(!isdeclaration, mod.functions), mod.aliases)) @@ -280,7 +289,7 @@ function link_device_libraries!(@nospecialize(job::CompilerJob), mod::LLVM.Modul names = intersect(needed, library_exports(lib)) isempty(names) && continue library, library_relocs = if lib.bitcode === nothing - load_julia_library(lib, job) + load_julia_library(lib, job; lowered) else parse(LLVM.Module, MemoryBuffer(library_bytes(lib)); lazy=true), Relocations() end @@ -320,23 +329,36 @@ end # along with the kernel. function link_device_libraries_early!(@nospecialize(job::CompilerJob), mod::LLVM.Module, relocs::Relocations) - @tracepoint "device libraries" link_device_libraries!(job, mod, relocs) + @tracepoint "device libraries" link_device_libraries!(job, mod, relocs; lowered=false) @tracepoint "target libraries" link_libraries!(job, mod) return end -# After optimization: import definitions for references introduced by lowering passes. -# Definitions imported before optimization that are still around are reused (the -# optimizer does not change their signatures). +# After optimization: lower what the target cannot select, and import definitions for the +# references this introduces. Definitions imported before optimization that are still +# around are reused: the optimizer does not change their signatures, only their calling +# convention, which `call_declared!` follows. +function lower_and_link_device_libraries!(@nospecialize(job::CompilerJob), + mod::LLVM.Module, relocs::Relocations) + lower_intrinsics!(job, mod) + return link_device_libraries_late!(job, mod, relocs) +end + +# Import definitions for references introduced after optimization. Returns whether +# anything was imported. function link_device_libraries_late!(@nospecialize(job::CompilerJob), mod::LLVM.Module, relocs::Relocations) - imported = link_device_libraries!(job, mod, relocs) + imported = link_device_libraries!(job, mod, relocs; lowered=true) link_libraries!(job, mod) for name in imported value = something(get(mod.functions, name, nothing), get(mod.aliases, name, nothing)) value.linkage = LLVM.Linkage.Internal end - return + # repeat what imports before optimization went through + CPUFeatures(job)(mod) + lower_syncscopes!(job, mod) + isempty(imported) || finish_imports!(job, mod, imported) + return !isempty(imported) end @public DeviceLibrary, DeviceLibraryMethod, device_libraries, runtime_library diff --git a/src/driver.jl b/src/driver.jl index 0ac2754b..ecdd69b8 100644 --- a/src/driver.jl +++ b/src/driver.jl @@ -392,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. diff --git a/src/interface.jl b/src/interface.jl index 223f0fcc..736c9032 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -746,6 +746,28 @@ prepare_job!(@nospecialize(job::CompilerJob)) = return # 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 diff --git a/src/optim.jl b/src/optim.jl index 09abb2f6..b2ec4ee8 100644 --- a/src/optim.jl +++ b/src/optim.jl @@ -60,7 +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)) + 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)) @@ -319,17 +322,23 @@ 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 end # GC lowering can introduce calls to the runtime if job.config.libraries - # Use the registered pass because it owns `relocs`. add!(mpm, "GPULinkRuntime") end @@ -402,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) @@ -481,17 +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 + +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 +GPULowerIntrinsicsPass(state) = ModulePass("GPULowerIntrinsics", LowerIntrinsics(state)) + +struct LinkRuntime + state::LateImports end function (self::LinkRuntime)(mod::LLVM.Module) - link_device_libraries_late!(self.job, mod, self.relocs) - finish_runtime_intrinsics!(self.job, mod) + state = self.state + state.imported |= link_device_libraries_late!(state.job, mod, state.relocs) + finish_runtime_intrinsics!(state.job, mod) + return true +end +GPULinkRuntimePass(state) = ModulePass("GPULinkRuntime", LinkRuntime(state)) + +# 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::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 -GPULinkRuntimePass(job, relocs::Relocations) = - ModulePass("GPULinkRuntime", LinkRuntime(job, relocs)) +GPULateImportsCleanupPass(state, opt_level) = + ModulePass("GPULateImportsCleanup", LateImportsCleanup(state, opt_level)) # lower object allocations to to PTX malloc # diff --git a/src/ptx.jl b/src/ptx.jl index 4eb1c3da..80cda738 100644 --- a/src/ptx.jl +++ b/src/ptx.jl @@ -255,6 +255,9 @@ function finish_linked_module!(@nospecialize(job::CompilerJob{PTXCompilerTarget} return end +finish_imports!(@nospecialize(job::CompilerJob{PTXCompilerTarget}), mod::LLVM.Module, + imported) = finish_linked_module!(job, mod) + function optimize_module!(@nospecialize(job::CompilerJob{PTXCompilerTarget}), mod::LLVM.Module) with_llvm_machine(job.config.target) do tm diff --git a/test/device_library.jl b/test/device_library.jl index cd802ee6..48f25a1d 100644 --- a/test/device_library.jl +++ b/test/device_library.jl @@ -28,12 +28,39 @@ const julia_library = GPUCompiler.DeviceLibrary(:test; GPUCompiler.DeviceLibraryMethod(exported, Int32, (Int32,); llvm_name="test_alias"))) +# `lower_intrinsics!` replaces `llvm.exp.f32` with a call to `test_exp` +exp_caller(value::Float32) = ccall("llvm.exp.f32", llvmcall, Float32, (Float32,), value) +both_caller(value::Float32) = + ccall("extern test_exp", llvmcall, Float32, (Float32,), value) + exp_caller(value) +library_exp(value::Float32) = value * 2f0 + 3f0 +fma_exp(value::Float32) = fma(value, 2f0, 3f0) +mutable struct Box + value::Float32 +end +@noinline escape(box::Box) = (box.value += 1f0; box.value) +allocating_exp(value::Float32) = escape(Box(value)) +exp_library(f; kwargs...) = GPUCompiler.DeviceLibrary(:exp; + methods=(GPUCompiler.DeviceLibraryMethod(f, Float32, (Float32,); + llvm_name="test_exp"),), kwargs...) + struct Params <: GPUCompiler.AbstractCompilerParams libraries::Tuple end GPUCompiler.runtime_module(::CompilerJob{<:Any,Params}) = Runtime GPUCompiler.device_libraries(job::CompilerJob{<:Any,Params}) = (job.config.params.libraries..., GPUCompiler.runtime_library(Runtime)) +function GPUCompiler.lower_intrinsics!(::CompilerJob{<:Any,Params}, mod::LLVM.Module) + exp = LLVM.Intrinsic("llvm.exp") + for fun in mod.functions + isdeclaration(fun) && continue + GPUCompiler.lower_intrinsic_calls!(fun) do builder, call, intrinsic + intrinsic == exp || return nothing + GPUCompiler.call_declared!(builder, "test_exp", LLVM.FloatType(), + collect(LLVM.Value, call.arguments)) + end + end +end + function job(func, types; libs=(julia_library,), target=:native, kwargs...) source = GPUCompiler.methodinstance(typeof(func), Base.to_tuple_type(types)) target = target === :native ? NativeCompilerTarget() : PTXCompilerTarget(; cap=v"7.0") @@ -53,6 +80,10 @@ function ir(func, types=(Int32,); kwargs...) end end +asm(func, types; kwargs...) = JuliaContext() do ctx + first(GPUCompiler.compile(:asm, job(func, types; kwargs...))) +end + bitcode(text) = JuliaContext() do ctx @dispose mod=parse(LLVM.Module, text) begin buffer = IOBuffer() @@ -81,11 +112,11 @@ end llvm_name="vector_export"),)) job = LibraryFixtures.job(LibraryFixtures.caller, (Int32,)) cached = JuliaContext() do ctx - GPUCompiler.load_julia_library(vectors(TypeVar(:T, Real)), job) + GPUCompiler.load_julia_library(vectors(TypeVar(:T, Real)), job; lowered=false) length(GPUCompiler.device_libs) end @test cached == JuliaContext() do ctx - GPUCompiler.load_julia_library(vectors(TypeVar(:S, Real)), job) + GPUCompiler.load_julia_library(vectors(TypeVar(:S, Real)), job; lowered=false) length(GPUCompiler.device_libs) end @@ -202,3 +233,62 @@ end @test !occursin("@test_export", LibraryFixtures.ir(LibraryFixtures.caller; libs=(inline,))) end + +@testset "late imports" begin + exp_library = LibraryFixtures.exp_library + + # lowered intrinsics are imported after optimization + for library in (exp_library(LibraryFixtures.library_exp), + exp_library(LibraryFixtures.library_exp; optimize=true), + exp_library(LibraryFixtures.library_exp; optimize=true, always_inline=false)) + text = LibraryFixtures.ir(LibraryFixtures.exp_caller, (Float32,); + libs=(library,)) + @test !occursin("llvm.exp", text) + @test !occursin("declare float @test_exp", text) + @test occursin("3.000000e+00", text) + @test occursin("call fastcc float @test_exp", text) == !library.always_inline + end + text = LibraryFixtures.ir(LibraryFixtures.exp_caller, (Float32,); + libs=(exp_library(LibraryFixtures.library_exp),), + optimize=false) + @test !occursin("llvm.exp", text) + @test occursin("define", text) + + # an export can be needed both before and after optimization + text = LibraryFixtures.ir(LibraryFixtures.both_caller, (Float32,); + libs=(exp_library(LibraryFixtures.library_exp),)) + @test !occursin("llvm.exp", text) + @test !occursin("declare float @test_exp", text) + @test !occursin("unreachable", text) + @test count("call fastcc float @test_exp", text) == 2 + + # Julia libraries imported late are lowered like any other Julia code + library = exp_library(LibraryFixtures.allocating_exp) + text = LibraryFixtures.ir(LibraryFixtures.exp_caller, (Float32,); + libs=(library,), target=:ptx) + @test !occursin("@julia.", text) + @test occursin("gpu_gc_pool_alloc", text) + + # late imports are processed like early ones + text = LibraryFixtures.ir(LibraryFixtures.exp_caller, (Float32,); + libs=(exp_library(LibraryFixtures.fma_exp),), target=:ptx) + @test !occursin("@julia.cpu", text) + fence = GPUCompiler.DeviceLibrary(:fence; bitcode=LibraryFixtures.bitcode(""" + define float @test_exp(float %value) { + fence syncscope("workgroup") seq_cst + ret float %value + }""")) + text = LibraryFixtures.ir(LibraryFixtures.exp_caller, (Float32,); + libs=(fence,), target=:ptx) + @test occursin("syncscope(\"block\")", text) + + # late imports are internalized, including aliases (which need PTX 6.3) + if NVPTX_LLVM_Backend_jll.is_available() + aliased = GPUCompiler.DeviceLibrary(:aliased; bitcode=LibraryFixtures.bitcode(""" + define float @implementation(float %value) { ret float %value } + @test_exp = alias float (float), float (float)* @implementation""")) + text = LibraryFixtures.asm(LibraryFixtures.exp_caller, (Float32,); + libs=(aliased,), target=:ptx) + @test !occursin(".alias", text) + end +end From dbe5000c44a2b8c3d33f82704ae8e819f546718d Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 20 Oct 2026 16:01:40 +0200 Subject: [PATCH 4/4] PTX: Select precise libdevice implementations unless using fast math NVVMReflect folds `__CUDA_PREC_SQRT` and `__CUDA_PREC_DIV` to 0, so libdevice used its approximate square roots and divisions (e.g. in `__nv_hypotf`) even without fast math. Fold them to `!fastmath` instead, matching nvcc's `-prec-sqrt`/`-prec-div` defaults. Also accept a bare global as the reflect argument on LLVM < 17: with opaque pointers the zero-index GEP folds away there too. --- src/ptx.jl | 22 +++++++++------------- test/ptx.jl | 35 +++++++++++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 13 deletions(-) diff --git a/src/ptx.jl b/src/ptx.jl index 80cda738..c48b29a3 100644 --- a/src/ptx.jl +++ b/src/ptx.jl @@ -645,7 +645,9 @@ end # This matches LLVM's NVVMReflectPass: `__CUDA_ARCH` is derived from the target # capability, `__CUDA_FTZ` is read from the `nvvm-reflect-ftz` module flag, and every # other key folds to 0. Knobs like denormal flushing or FMAD contraction must be -# configured through module flags or LLVM fast-math flags, not here. +# configured through module flags or LLVM fast-math flags, not here. The exception are +# `__CUDA_PREC_SQRT` and `__CUDA_PREC_DIV`, which select precise libdevice implementations +# unless fast math is enabled, like nvcc does (LLVM defaults these to approximations). const NVVM_REFLECT_FUNCTION = "__nvvm_reflect" function nvvm_reflect!(@nospecialize(job::CompilerJob), mod::LLVM.Module) changed = false @@ -682,18 +684,10 @@ function nvvm_reflect!(@nospecialize(job::CompilerJob), mod::LLVM.Module) continue end - # decode the string argument - if LLVM.version() >= v"17" - sym = call.operands[1] - else - str = call.operands[1] - if !isa(str, LLVM.ConstantExpr) || str.opcode != LLVM.Opcode.GetElementPtr - @safe_error """Unrecognized format of __nvvm_reflect call: - $(string(call)) - Operand should be a GEP instruction, got a $(typeof(str)). Please file an issue.""" - continue - end - sym = str.operands[1] + # decode the string argument (a GEP only with typed pointers) + sym = call.operands[1] + if isa(sym, LLVM.ConstantExpr) && sym.opcode == LLVM.Opcode.GetElementPtr + sym = sym.operands[1] if isa(sym, LLVM.ConstantExpr) && sym.opcode == LLVM.Opcode.GetElementPtr # CUDA 11.0 or below sym = sym.operands[1] @@ -720,6 +714,8 @@ function nvvm_reflect!(@nospecialize(job::CompilerJob), mod::LLVM.Module) job.config.target.cap.major*100 + job.config.target.cap.minor*10) elseif reflect_arg == "__CUDA_FTZ" ConstantInt(reflect_typ, ftz_val) + elseif reflect_arg == "__CUDA_PREC_SQRT" || reflect_arg == "__CUDA_PREC_DIV" + ConstantInt(reflect_typ, !job.config.target.fastmath) else ConstantInt(reflect_typ, 0) end diff --git a/test/ptx.jl b/test/ptx.jl index 46e8e209..8f2018ef 100644 --- a/test/ptx.jl +++ b/test/ptx.jl @@ -1012,6 +1012,41 @@ end end end +@testset "nvvm_reflect" begin + reflect(key; kwargs...) = JuliaContext() do ctx + ir = if LLVM.version() >= v"17" + """ + @key = private unnamed_addr constant [$(sizeof(key)+1) x i8] c"$(key)\\00" + declare i32 @__nvvm_reflect(ptr) + define i32 @test() { + %value = call i32 @__nvvm_reflect(ptr @key) + ret i32 %value + }""" + else + """ + @key = private unnamed_addr constant [$(sizeof(key)+1) x i8] c"$(key)\\00" + declare i32 @__nvvm_reflect(i8*) + define i32 @test() { + %value = call i32 @__nvvm_reflect(i8* getelementptr ([$(sizeof(key)+1) x i8], [$(sizeof(key)+1) x i8]* @key, i64 0, i64 0)) + ret i32 %value + }""" + end + @dispose mod=parse(LLVM.Module, ir) begin + job, _ = PTX.create_job(identity, (Int,); cap=v"7.5", kwargs...) + GPUCompiler.nvvm_reflect!(job, mod) + ret = LLVM.terminator(only(mod.functions["test"].blocks)) + convert(Int, ret.operands[1]) + end + end + + @test reflect("__CUDA_ARCH") == 750 + @test reflect("__CUDA_PREC_SQRT") == 1 + @test reflect("__CUDA_PREC_DIV") == 1 + @test reflect("__CUDA_PREC_SQRT"; fastmath=true) == 0 + @test reflect("__CUDA_PREC_DIV"; fastmath=true) == 0 + @test reflect("__CUDA_UNKNOWN") == 0 +end + @testset "feature_set" begin # PTXCompilerTarget.feature_set controls the suffix on the LLVM CPU name, which is # what the NVPTX backend uses to flip `hasArchAccelFeatures()`. Verify it makes its