diff --git a/src/optim.jl b/src/optim.jl index 3b01e23d9..5c11faca8 100644 --- a/src/optim.jl +++ b/src/optim.jl @@ -59,7 +59,7 @@ function optimize!(@nospecialize(job::CompilerJob), mod::LLVM.Module, register!(pb, GPULowerCPUFeaturesPass(job)) register!(pb, GPULowerPTLSPass(job)) - register!(pb, GPULowerGCFramePass(job)) + register!(pb, GPULowerGCFramePass(job, relocs)) register!(pb, GPULinkRuntimePass(job, relocs)) register!(pb, GPULinkLibrariesPass(job)) register!(pb, GPUFinishRuntimeIntrinsicsPass(job)) @@ -323,7 +323,8 @@ function buildIntrinsicLoweringPipeline(mpm, @nospecialize(job::CompilerJob), op # lower GC intrinsics if !uses_julia_runtime(job) add!(mpm, NewPMFunctionPassManager()) do fpm - add!(fpm, GPULowerGCFramePass(job)) + # 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`. @@ -537,6 +538,7 @@ GPUFinishRuntimeIntrinsicsPass(job) = # lower-level intrinsics which then can be lowered to architecture-specific code. struct LowerGCFrame job::CompilerJob + relocs::Relocations end function (self::LowerGCFrame)(fun::LLVM.Function) mod = LLVM.parent(fun) @@ -555,10 +557,13 @@ function (self::LowerGCFrame)(fun::LLVM.Function) # decode the call ops = arguments(call) sz = ops[2] + typ = ops[3] # replace with PTX alloc_obj @dispose builder=IRBuilder() begin position!(builder, call) + debuglocation!(builder, call) + check_allocation!(builder, typ, self.relocs) ptr = call!(builder, Runtime.get(:gc_pool_alloc), [sz]) replace_uses!(call, ptr) end @@ -589,7 +594,40 @@ function (self::LowerGCFrame)(fun::LLVM.Function) return changed end -GPULowerGCFramePass(job) = NewPMFunctionPass("GPULowerGCFrame", LowerGCFrame(job)) +GPULowerGCFramePass(job, relocs::Relocations) = + NewPMFunctionPass("GPULowerGCFrame", LowerGCFrame(job, relocs)) + +# only allocations of objects without references to other heap objects are supported. +# +# there is no garbage collector on the device, and `gc_pool_alloc` only provides the object's +# memory (without the header Julia's runtime expects, e.g., for `typeof`), so what works is a +# statically-typed use of an object that holds plain data (an escaping `Ref{Int}`, a mutable +# struct with bits fields, a boxed bits value). references are a different matter: Julia puts +# GC orderings on their loads and stores, which not every back-end can express, and they +# typically come from code that isn't GPU compatible, e.g., closures boxing captured variables. +# allocations that the optimizer could not remove are checked here, after Julia's `AllocOpt`, +# marking the unsupported ones for `check_ir` to report with a backtrace. +const UNSUPPORTED_ALLOCATION_MARKER = "gpu_unsupported_allocation" +function check_allocation!(builder::IRBuilder, typ::LLVM.Value, relocs::Relocations) + ref = referenced_object(typ, relocs) + ref === nothing && return + T = something(ref) + (T isa DataType && isconcretetype(T)) || return + Base.datatype_pointerfree(T) && return + + # emit a call to a marker, passing the name of the type + bb = position(builder) + mod = LLVM.parent(LLVM.parent(bb)) + name = globalstring_ptr!(builder, string(T), "allocated_type") + marker_type = LLVM.FunctionType(LLVM.VoidType(), [value_type(name)]) + marker = if haskey(functions(mod), UNSUPPORTED_ALLOCATION_MARKER) + functions(mod)[UNSUPPORTED_ALLOCATION_MARKER] + else + LLVM.Function(mod, UNSUPPORTED_ALLOCATION_MARKER, marker_type) + end + call!(builder, marker_type, marker, [name]) + return +end # lower the `julia.ptls_states` intrinsic by removing it, since it is GPU incompatible. # diff --git a/src/relocation.jl b/src/relocation.jl index 2371fa83f..bff9c12a9 100644 --- a/src/relocation.jl +++ b/src/relocation.jl @@ -1018,8 +1018,9 @@ end function referenced_object(value, relocs::Relocations) # This is best-effort: optimized shapes fall back to the unknown-binding error path. - while value isa ConstantExpr && - opcode(value) in (LLVM.API.LLVMBitCast, LLVM.API.LLVMAddrSpaceCast) + while (value isa ConstantExpr && + opcode(value) in (LLVM.API.LLVMBitCast, LLVM.API.LLVMAddrSpaceCast)) || + value isa LLVM.BitCastInst || value isa LLVM.AddrSpaceCastInst value = first(operands(value)) end if value isa LLVM.LoadInst @@ -1035,8 +1036,21 @@ function referenced_object(value, relocs::Relocations) end end elseif value isa ConstantExpr && opcode(value) == LLVM.API.LLVMIntToPtr - ptr = Ptr{Cvoid}(convert(Int, first(operands(value)))) - return Some(Base.unsafe_pointer_to_objref(ptr)) + addr = first(operands(value)) + addr isa ConstantInt || return nothing + addr = convert(UInt, addr) + addr < UInt(64 << 4) && return small_typeof(addr) # jl_max_tags << 4 + return Some(Base.unsafe_pointer_to_objref(Ptr{Cvoid}(addr))) end return nothing end + +# Codegen refers to some types by their small tag instead of their address, e.g., in the +# type operand of `julia.gc_alloc_obj`. Such a tag cannot be a heap address; like +# `jl_to_typeof`, resolve it through Julia's table, whose unused entries are null. +function small_typeof(tag::UInt) + table = cglobal(:jl_small_typeof, Ptr{Cvoid}) + ptr = unsafe_load(table, tag รท sizeof(Ptr{Cvoid}) + 1) + ptr == C_NULL && return nothing + return Some(Base.unsafe_pointer_to_objref(ptr)) +end diff --git a/src/validation.jl b/src/validation.jl index c07b93c34..dc553d921 100644 --- a/src/validation.jl +++ b/src/validation.jl @@ -142,6 +142,7 @@ const LAZY_FUNCTION = "call to a lazy-initialized function" const DELAYED_BINDING = "use of an undefined name" const DYNAMIC_CALL = "dynamic function invocation" const UNKNOWN_INTRINSIC = "call to an unknown LLVM intrinsic" +const UNSUPPORTED_ALLOCATION = "allocation of an object with references" function Base.showerror(io::IO, err::InvalidIRError) print(io, "InvalidIRError: compiling ", err.job.source, " resulted in invalid LLVM IR") @@ -155,6 +156,8 @@ function Base.showerror(io::IO, err::InvalidIRError) printstyled(io, " (use of '", meta, "')"; color=:red) elseif kind == STATIC_ASSERTION printstyled(io, " (", meta, ")"; color=:red) + elseif kind == UNSUPPORTED_ALLOCATION + printstyled(io, " (", meta, ")"; color=:red) end end Base.show_backtrace(io, bt) @@ -268,6 +271,8 @@ function check_ir!(job, errors::Vector{IRError}, inst::LLVM.CallInst, relocs::Re push!(errors, (STATIC_ASSERTION, bt, static_assert_message(inst))) elseif is_unknown_intrinsic_error(inst) push!(errors, (UNKNOWN_INTRINSIC, bt, nothing)) + elseif fn == UNSUPPORTED_ALLOCATION_MARKER + push!(errors, (UNSUPPORTED_ALLOCATION, bt, static_assert_message(inst))) elseif fn == "jl_get_binding_or_error" || fn == "ijl_get_binding_or_error" try m, sym = arguments(inst) diff --git a/test/metal.jl b/test/metal.jl index 18e7ea282..67c7b11f1 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -1806,6 +1806,102 @@ end end end + +@testset "unsupported allocations" begin + # without a garbage collector, device code can only allocate objects that do not + # reference other objects + mod = @eval module $(gensym()) + using ..GPUCompiler + import ..ExternalAllocatorRuntime, ..Metal + struct Params <: GPUCompiler.AbstractCompilerParams + table::Bool # use the `:table` relocation strategy, as Metal.jl does + end + GPUCompiler.runtime_module(::CompilerJob{<:Any,Params}) = ExternalAllocatorRuntime + GPUCompiler.relocation_lowering(job::CompilerJob{<:Any,Params}) = + job.config.params.table ? (:table) : (:bake) + GPUCompiler.kernel_state_type(job::CompilerJob{<:Any,Params}) = + job.config.params.table ? Metal.TableKernelState : Nothing + GPUCompiler.isintrinsic(job::CompilerJob{MetalCompilerTarget,Params}, fn::String) = + fn == "test_malloc" || + @invoke GPUCompiler.isintrinsic(job::CompilerJob{MetalCompilerTarget}, fn::String) + + mutable struct Counter + x::Int + end + mutable struct Holder + c::Counter + end + @noinline bump!(c::Counter) = (c.x += 1; return) + @noinline bump!(h::Holder) = bump!(h.c) + + # an escaping object with plain data + function plain(out, x) + c = Counter(x) + bump!(c) + unsafe_store!(out, c.x) + return + end + + # one that references another object + function nested(out, x) + h = Holder(Counter(x)) + bump!(h) + unsafe_store!(out, h.c.x) + return + end + + # one whose type codegen refers to by its small tag instead of its address + @noinline svec_length(s::Core.SimpleVector) = length(s) + function svec(out, x) + s = Core.svec(x) + unsafe_store!(out, svec_length(s)) + return + end + end + function compile(f, table) + source = methodinstance(typeof(f), Tuple{Core.LLVMPtr{Int,1},Int}, + Base.get_world_counter()) + target = MetalCompilerTarget(; macos=v"12.2", metal=v"3.0", air=v"3.0") + job = CompilerJob(source, CompilerConfig(target, mod.Params(table); kernel=true)) + JuliaContext() do ctx + GPUCompiler.compile(:llvm, job) + end + end + + # the allocated type is found through relocations of either strategy + for table in (false, true) + @test compile(mod.plain, table) isa Tuple + @test_throws_message(InvalidIRError, compile(mod.nested, table)) do msg + occursin("unsupported allocation of an object with references", msg) && + occursin("Holder)", msg) && occursin("[2] nested", msg) + end + # older versions call `jl_f_svec` instead of allocating inline + if VERSION >= v"1.12" + @test_throws_message(InvalidIRError, compile(mod.svec, table)) do msg + occursin("unsupported allocation of an object with references", msg) && + occursin("(Core.SimpleVector)", msg) && occursin("[1] svec", msg) + end + end + end + + # small tags are resolved on every version, while addresses are still dereferenced + Context() do ctx + relocs = GPUCompiler.Relocations() + ref(addr) = GPUCompiler.referenced_object( + const_inttoptr(ConstantInt(UInt64(addr)), LLVM.PointerType(LLVM.Int8Type())), + relocs) + # the header of an object holds its type's small tag + obj = Core.svec(1, 2) + ptr = ccall(:jl_value_ptr, Ptr{UInt}, (Any,), obj) + tag = GC.@preserve obj unsafe_load(ptr - sizeof(UInt)) & ~UInt(15) + @test tag < 64 << 4 + @test something(ref(tag)) === Core.SimpleVector + @test ref(63 << 4) === nothing # unused tag + addr = ccall(:jl_value_ptr, Ptr{Cvoid}, (Any,), Core.SimpleVector) + @test something(ref(UInt(addr))) === Core.SimpleVector + end +end + @testset "atomic demotion" begin # Demote every LLVM atomic load/store, whatever its ordering or value type. Context() do ctx diff --git a/test/spirv.jl b/test/spirv.jl index 785d009c5..28d6a997d 100644 --- a/test/spirv.jl +++ b/test/spirv.jl @@ -363,6 +363,61 @@ end end end +@testset "unsupported allocations" begin + # without a garbage collector, device code can only allocate objects that do not + # reference other objects + mod = @eval module $(gensym()) + using ..GPUCompiler + import ..ExternalAllocatorRuntime + struct Params <: GPUCompiler.AbstractCompilerParams end + GPUCompiler.runtime_module(::CompilerJob{<:Any,Params}) = ExternalAllocatorRuntime + GPUCompiler.isintrinsic(job::CompilerJob{SPIRVCompilerTarget,Params}, fn::String) = + fn == "test_malloc" || + @invoke GPUCompiler.isintrinsic(job::CompilerJob{SPIRVCompilerTarget}, fn::String) + + mutable struct Counter + x::Int + end + mutable struct Holder + c::Counter + end + @noinline bump!(c::Counter) = (c.x += 1; return) + @noinline bump!(h::Holder) = bump!(h.c) + + # an escaping object with plain data + function plain(out, x) + c = Counter(x) + bump!(c) + unsafe_store!(out, c.x) + return + end + + # one that references another object + function nested(out, x) + h = Holder(Counter(x)) + bump!(h) + unsafe_store!(out, h.c.x) + return + end + end + function compile(f) + source = methodinstance(typeof(f), Tuple{Core.LLVMPtr{Int,1},Int}, + Base.get_world_counter()) + target = SPIRVCompilerTarget(; backend, validate=true) + job = CompilerJob(source, CompilerConfig(target, mod.Params(); kernel=true)) + JuliaContext() do ctx + GPUCompiler.compile(:llvm, job) + end + end + + @test compile(mod.plain) isa Tuple + @test_throws_message(InvalidIRError, compile(mod.nested)) do msg + occursin("unsupported allocation of an object with references", msg) && + occursin("Holder)", msg) && occursin("[2] nested", msg) + end +end + + @testset "atomic demotion" begin # Julia's `unordered` heap-reference accesses and `release` type-tag stores cannot be # expressed in SPIR-V when they involve pointers (OpAtomicLoad/OpAtomicStore take scalars