diff --git a/README.md b/README.md index fc8988a2..d3fde856 100644 --- a/README.md +++ b/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 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..9c8804a1 --- /dev/null +++ b/src/devicelib.jl @@ -0,0 +1,364 @@ +# 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. 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); + 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()), lowered, + 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; lowered) + 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; lowered::Bool=false) + @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 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 + 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; lowered::Bool) + 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; lowered) + 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; lowered=false) + @tracepoint "target libraries" link_libraries!(job, mod) + return +end + +# 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; 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 + # 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 13c3f81d..ecdd69b8 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 @@ -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. diff --git a/src/interface.jl b/src/interface.jl index 1b656eaf..736c9032 100644 --- a/src/interface.jl +++ b/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 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/optim.jl b/src/optim.jl index b27f3a37..b2ec4ee8 100644 --- a/src/optim.jl +++ b/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 # diff --git a/src/ptx.jl b/src/ptx.jl index 4eb1c3da..c48b29a3 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 @@ -642,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 @@ -679,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] @@ -717,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/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/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/device_library.jl b/test/device_library.jl new file mode 100644 index 00000000..48f25a1d --- /dev/null +++ b/test/device_library.jl @@ -0,0 +1,294 @@ +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"))) + +# `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") + 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 + +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() + 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; lowered=false) + length(GPUCompiler.device_libs) + end + @test cached == JuliaContext() do ctx + GPUCompiler.load_julia_library(vectors(TypeVar(:S, Real)), job; lowered=false) + 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 + +@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 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) 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 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