Reject allocations of objects that contain references

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.
 #

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

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)

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

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