From ca3b1b7a5ee0c787e168b1082c5eeaec3a4b4f16 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 23 Sep 2026 18:12:59 +0200 Subject: [PATCH 1/2] Reject allocations of objects that contain references. There is no garbage collector on the device, and `gc_pool_alloc` does not provide the object header expected by Julia's runtime. Only statically typed uses of objects holding plain data work. Objects that reference other objects get GC orderings on their loads and stores, which not all back-ends can express, and typically come from GPU-incompatible code. After `AllocOpt`, `GPULowerGCFrame` now marks allocations of such objects, and `check_ir` reports them with a backtrace. --- src/optim.jl | 44 +++++++++++++++++++++++++++++--- src/relocation.jl | 5 ++-- src/validation.jl | 5 ++++ test/metal.jl | 64 +++++++++++++++++++++++++++++++++++++++++++++++ test/spirv.jl | 55 ++++++++++++++++++++++++++++++++++++++++ 5 files changed, 168 insertions(+), 5 deletions(-) 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..3b8ae3c37 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 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..4baaabb5e 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -1806,6 +1806,70 @@ 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 + 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 + 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 From 31b28f494c672ab0c61fa6332dafe7cb2e9a050c Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 18:02:34 +0200 Subject: [PATCH 2/2] Resolve small type tags when looking up referenced objects. Julia 1.10+ refers to some types by a small tag instead of their address: an offset into `jl_small_typeof`, below `jl_max_tags << 4`. Codegen emits such a tag as the type operand of `julia.gc_alloc_obj`, e.g., since 1.12 when allocating a `Core.svec` inline, where `check_allocation!` passed it to `referenced_object`. That treated every `inttoptr` constant as an object address and dereferenced the tag, crashing the compiler. Look small tags up in `jl_small_typeof` instead, as `jl_to_typeof` does, and treat unused table entries and non-constant operands as unknown. Larger constants are still object addresses, as embedded by 1.10's codegen or by baking relocations. Such allocations of a `SimpleVector` are now reported as allocations of an object with references. --- src/relocation.jl | 17 +++++++++++++++-- test/metal.jl | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/src/relocation.jl b/src/relocation.jl index 3b8ae3c37..bff9c12a9 100644 --- a/src/relocation.jl +++ b/src/relocation.jl @@ -1036,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/test/metal.jl b/test/metal.jl index 4baaabb5e..67c7b11f1 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -1849,6 +1849,14 @@ end 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}, @@ -1867,6 +1875,30 @@ end 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