From 441c8edef24c0e114018123fe90ff9e6d6e8def6 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sun, 6 Sep 2026 11:59:42 +0200 Subject: [PATCH 1/5] Add a per-work-item device heap Provide GPUCompiler's malloc hook so boxed Julia objects and allocations on exception paths can execute in kernels. Preserving inferred invoke specializations exposes these allocations, as reported in GPUCompiler.jl#906. Use a 1 KiB private arena per work-item because Julia's boxed-object pointers map to SPIR-V private memory and cannot address a global USM heap. Initialize the arena through the kernel state before SPIR-V argument lowering, detecting heap-field reads even when malloc has been inlined. Non-allocating kernels do not reserve an arena. Round allocations to 16 bytes and return null on exhaustion or size overflow. Report exhaustion before GPUCompiler terminates the work-item. Cover boxed values, independent heaps, allocation lifetime, fresh arenas across launches, alignment, failed requests, OOM output, and the motivating math kernels. --- src/compiler/compilation.jl | 54 +++++++++++++++++ src/compiler/execution.jl | 4 ++ src/device/runtime.jl | 74 +++++++++++++++++------ test/execution.jl | 117 ++++++++++++++++++++++++++++++++++++ 4 files changed, 232 insertions(+), 17 deletions(-) diff --git a/src/compiler/compilation.jl b/src/compiler/compilation.jl index a8c117a8..92557313 100644 --- a/src/compiler/compilation.jl +++ b/src/compiler/compilation.jl @@ -30,6 +30,8 @@ end GPUCompiler.runtime_module(::oneAPICompilerJob) = oneAPI +GPUCompiler.kernel_state_type(::oneAPICompilerJob) = KernelState + GPUCompiler.method_table_view(job::oneAPICompilerJob) = GPUCompiler.StackedMethodTable(job.world, method_table, SPIRVIntrinsics.method_table) @@ -64,6 +66,9 @@ end # finish_ir! runs later in the pipeline, after optimizations that create nested insertvalue function GPUCompiler.finish_ir!(job::oneAPICompilerJob, mod::LLVM.Module, entry::LLVM.Function) + # Initialize the heap before SPIR-V lowering converts the state to a reference. + job.config.kernel && add_heap!(mod, entry) + entry = invoke(GPUCompiler.finish_ir!, Tuple{CompilerJob{SPIRVCompilerTarget}, typeof(mod), typeof(entry)}, job, mod, entry) @@ -92,6 +97,55 @@ function GPUCompiler.finish_ir!(job::oneAPICompilerJob, mod::LLVM.Module, return entry end +# Reserve a private arena only when device code reads the heap pointer. At this point, +# GPUCompiler has threaded the state through callees as a leading by-value argument. +function add_heap!(mod::LLVM.Module, entry::LLVM.Function) + T_state = convert(LLVMType, KernelState) + heap_field = Base.fieldindex(KernelState, :heap) - 1 + uses_heap(mod, T_state, heap_field) || return false + + params = parameters(entry) + if isempty(params) || value_type(params[1]) != T_state + error("kernel `$(LLVM.name(entry))` allocates but has no kernel state to hold the heap") + end + state = params[1] + users = LLVM.Value[user(use) for use in uses(state)] + + T_size = convert(LLVMType, Csize_t) + T_heap = LLVM.StructType([T_size, T_size, LLVM.ArrayType(LLVM.Int8Type(), HEAP_SIZE)]) + T_ptr = convert(LLVMType, fieldtype(KernelState, :heap)) + @dispose builder = IRBuilder() begin + position!(builder, first(instructions(first(blocks(entry))))) + + heap = alloca!(builder, T_heap, "heap") + alignment!(heap, HEAP_ALIGNMENT) + store!(builder, ConstantInt(T_size, 0), struct_gep!(builder, T_heap, heap, 0)) + store!(builder, ConstantInt(T_size, HEAP_SIZE), struct_gep!(builder, T_heap, heap, 1)) + + # Replace the original uses, excluding the insertvalue that constructs the state. + ptr = pointercast!(builder, heap, T_ptr) + new_state = insert_value!(builder, state, ptr, heap_field, "state") + for u in users + ops = operands(u) + for i in 1:length(ops) + ops[i] == state && (ops[i] = new_state) + end + end + end + + return true +end + +# Inspect field reads rather than calls to malloc, which may already have been inlined. +function uses_heap(mod::LLVM.Module, T_state::LLVMType, heap_field::Integer) + for f in functions(mod), bb in blocks(f), inst in instructions(bb) + inst isa LLVM.ExtractValueInst || continue + value_type(operands(inst)[1]) == T_state || continue + unsafe_load(LLVM.API.LLVMGetIndices(inst)) == heap_field && return true + end + return false +end + # Flatten nested insertvalue instructions # This works around a bug in Intel's SPIR-V runtime where OpCompositeInsert # with nested array indices corrupts adjacent struct fields. diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index bd874d67..12d9f4a8 100644 --- a/src/compiler/execution.jl +++ b/src/compiler/execution.jl @@ -195,6 +195,10 @@ abstract type AbstractKernel{F,TT} end end end + # Match GPUCompiler's hidden state argument. Keep onecall usable for foreign kernels. + pushfirst!(call_t, KernelState) + pushfirst!(call_args, :(KernelState())) + # finalize types call_tt = Base.to_tuple_type(call_t) diff --git a/src/device/runtime.jl b/src/device/runtime.jl index 2f48b0f1..d7c83e0b 100644 --- a/src/device/runtime.jl +++ b/src/device/runtime.jl @@ -1,36 +1,76 @@ # device runtime libraries -## Julia library +## kernel state +# GPUCompiler passes this as a hidden kernel argument and forwards it to device callees. +struct KernelState + # Initialized on the device by add_heap!; the host passes a null pointer. + heap::LLVMPtr{UInt8, AS.Function} +end + +KernelState() = KernelState(reinterpret(LLVMPtr{UInt8, AS.Function}, C_NULL)) + +@inline @generated kernel_state() = GPUCompiler.kernel_state_value(KernelState) + + +## dynamic memory allocation + +# Julia's boxed objects use address-space-0 pointers, which SPIR-V maps to private +# memory. A global (USM) allocation cannot back those pointers on Intel GPUs. +# Use a per-work-item bump allocator: objects remain valid until the work-item exits, +# and must not be shared with other work-items or retained across launches. +# +# add_heap! reserves the arena in the kernel entry block. Its header contains the cursor +# and capacity in bytes, keeping the runtime independent of the compiler's chosen size. + +# bytes of private memory reserved per work-item for dynamic allocations +const HEAP_SIZE = 1024 + +# alignment of every allocation; the largest Julia's codegen assumes for heap objects +const HEAP_ALIGNMENT = 16 + +# Two 64-bit words keep the payload aligned to HEAP_ALIGNMENT. +const HEAP_HEADER = 2 * sizeof(Csize_t) + +function malloc(sz::Csize_t) + heap = kernel_state().heap + heap == reinterpret(LLVMPtr{UInt8, AS.Function}, C_NULL) && return C_NULL + + header = reinterpret(LLVMPtr{Csize_t, AS.Function}, heap) + cursor = unsafe_load(header, 1, Val(sizeof(Csize_t))) + capacity = unsafe_load(header, 2, Val(sizeof(Csize_t))) + + bytes = (sz + Csize_t(HEAP_ALIGNMENT - 1)) & ~Csize_t(HEAP_ALIGNMENT - 1) + bytes < sz && return C_NULL # alignment rounding overflowed + bytes > capacity - cursor && return C_NULL # gc_pool_alloc reports exhaustion + + unsafe_store!(header, cursor + bytes, 1, Val(sizeof(Csize_t))) + return reinterpret(Ptr{Cvoid}, heap + HEAP_HEADER + cursor) +end + +function report_oom(sz) + @println("ERROR: Out of dynamic GPU memory (trying to allocate ", sz, " bytes)") + return +end + + +## exceptions + +# SPIR-V has no way to abort a kernel, and the exception is not reported to the host: the +# work-item that threw simply exits (see `lower_unreachable_control_flow!` in GPUCompiler). function signal_exception() return end function report_exception(ex) - # @cuprintf(""" - # ERROR: a %s was thrown during kernel execution. - # Run Julia on debug level 2 for device stack traces. - # """, ex) return end -report_oom(sz) = return #@cuprintf("ERROR: Out of dynamic GPU memory (trying to allocate %i bytes)\n", sz) - function report_exception_name(ex) - # @cuprintf(""" - # ERROR: a %s was thrown during kernel execution. - # Stacktrace: - # """, ex) return end function report_exception_frame(idx, func, file, line) - # @cuprintf(" [%i] %s at %s:%i\n", idx, func, file, line) return end - - -## SPIRV libraries - -# TODO diff --git a/test/execution.jl b/test/execution.jl index 1f44c128..03c6f9f6 100644 --- a/test/execution.jl +++ b/test/execution.jl @@ -742,3 +742,120 @@ end end for _ in 1:2]) @test all(results) end + +############################################################################################ + +# Keep allocation consumers at top level so kernels do not capture test state. + +@noinline heap_consume(r::Base.RefValue{Float32}) = r[] + 1.0f0 + +struct HeapAnyBox + x::Any +end +@noinline heap_consume(b::HeapAnyBox) = (b.x::Float32) * 2.0f0 + +@testset "device heap" begin + # Keep objects alive across a call so allocation survives Julia/LLVM optimization. + function ref_kernel(a) + i = get_global_id() + @inbounds a[i] = heap_consume(Ref(a[i])) + return + end + a = oneArray(Float32[41]) + @oneapi ref_kernel(a) + @test Array(a) == [42] + + # so is a struct whose `Any` field boxes its value + function anybox_kernel(a) + i = get_global_id() + @inbounds a[i] = heap_consume(HeapAnyBox(a[i])) + return + end + a = oneArray(Float32[21]) + @oneapi anybox_kernel(a) + @test Array(a) == [42] + + # every work-item has its own heap + n = 4096 + a = oneArray(Float32.(1:n)) + @oneapi items = 256 groups = n ÷ 256 ref_kernel(a) + @test Array(a) == Float32.(2:(n + 1)) + + # objects stay valid across later allocations by the same work-item + function select_kernel(a, idx) + i = get_global_id() + refs = ntuple(j -> Ref(a[i] * j), Val(4)) + @inbounds a[i] = heap_consume(refs[idx]) + return + end + a = oneArray(Float32[1, 2, 3, 4]) + @oneapi items = 4 select_kernel(a, 3) + @test Array(a) == Float32[4, 7, 10, 13] + + # nothing is freed: a work-item that allocates more than the heap holds runs out of + # memory, which is reported, and exits without writing its result + function loop_kernel(a, n) + i = get_global_id() + @inbounds x = a[i] + for _ in 1:n + x = heap_consume(Ref(x)) + end + @inbounds a[i] = x + return + end + fits = oneAPI.HEAP_SIZE ÷ oneAPI.HEAP_ALIGNMENT + a = oneArray(Float32.(1:256)) + @oneapi items = 256 loop_kernel(a, fits) + @test Array(a) == Float32.(1:256) .+ fits + # A later launch starts with an empty heap, even after using the entire arena. + @oneapi items = 256 loop_kernel(a, fits) + @test Array(a) == Float32.(1:256) .+ 2 * fits + a = oneArray(Float32[1]) + _, out = @grab_output begin + @oneapi loop_kernel(a, fits + 1) + synchronize() + end + @test occursin("Out of dynamic GPU memory", out) + @test Array(a) == [1] + + # Failed requests must not consume space. Exercise the size arithmetic directly, + # including overflow when rounding up and allocations with different sizes. + function allocation_kernel(out, sizes) + for i in eachindex(sizes) + ptr = oneAPI.malloc(sizes[i]) + out[i] = UInt(ptr) + end + return + end + sizes = oneArray( + Csize_t[ + typemax(Csize_t), oneAPI.HEAP_SIZE + 1, 1, 17, + oneAPI.HEAP_SIZE - 3 * oneAPI.HEAP_ALIGNMENT, 1, + ] + ) + out = oneAPI.zeros(UInt, length(sizes)) + @oneapi allocation_kernel(out, sizes) + ptrs = Array(out) + @test ptrs[[1, 2, 6]] == [0, 0, 0] + @test all(!iszero, ptrs[3:5]) + @test all(p -> p % oneAPI.HEAP_ALIGNMENT == 0, ptrs[3:5]) + @test ptrs[4] - ptrs[3] == oneAPI.HEAP_ALIGNMENT + @test ptrs[5] - ptrs[4] == 2 * oneAPI.HEAP_ALIGNMENT + + # These valid inputs must compile even if exception paths retain boxed arguments + # (GPUCompiler.jl#906, exposed by preserving inferred invoke specializations). + powers = Float32[1, 2, 4, 8] + @test Array(exponent.(oneArray(powers))) == exponent.(powers) + z = ComplexF32[1 + 2im, -3 + 4im, 0, 2 - 3im] + @test Array(sqrt.(oneArray(z))) ≈ sqrt.(z) + + # only kernels that allocate carry a heap + function plain_kernel(a) + i = get_global_id() + @inbounds a[i] += 1.0f0 + return + end + T = Tuple{oneDeviceVector{Float32, oneAPI.AS.CrossWorkgroup}} + @test occursin(r"%heap\d* = alloca", sprint(io -> oneAPI.code_llvm(io, ref_kernel, T; kernel = true))) + @test !occursin(r"%heap\d* = alloca", sprint(io -> oneAPI.code_llvm(io, plain_kernel, T; kernel = true))) +end From 914dedb1007e1a6cb36d591bc09e6876fd051a0b Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sun, 6 Sep 2026 12:00:04 +0200 Subject: [PATCH 2/5] Document device allocation limits and failure behavior Explain the per-work-item heap's size, alignment, lifetime, and cumulative usage in loops. Make clear that exhaustion exits the work-item without a host exception and can leave kernel output incomplete. Avoid promising that the device compiler eliminates the storage or its performance cost. Keep the existing scalar-indexing and Diagonal error overrides: their printed diagnostics remain useful even with an allocator, since ordinary device exceptions still do not report their reason to the host. --- docs/src/kernels.md | 20 ++++++++++++++++++++ src/device/quirks.jl | 4 ++-- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/docs/src/kernels.md b/docs/src/kernels.md index 07e50c7a..a9eeb46f 100644 --- a/docs/src/kernels.md +++ b/docs/src/kernels.md @@ -62,3 +62,23 @@ and correspond to the standard OpenCL built-in functions. Note that the indices are 1-based, so they can be used to index Julia arrays directly. See [Device Intrinsics](device.md) for the full list. + +## Dynamic Memory Allocation + +Kernels can allocate Julia objects, such as a `Ref` passed to a `@noinline` function or a +boxed value in an `Any` field. Allocations that survive optimization use a 1 KiB heap +private to each work-item. Each allocation is rounded up to 16 bytes, and memory is only +reclaimed when the work-item exits. Allocated objects must not be shared with other +work-items or retained across kernel launches. + +When the heap is exhausted, the work-item prints an error and exits without completing +its work. This does not raise a host-side exception, and kernel output may be incomplete: + +``` +ERROR: Out of dynamic GPU memory (trying to allocate 4 bytes) +``` + +Kernels without remaining allocations do not reserve an arena. The device compiler may +optimize away some heap storage, but allocations can increase private-memory use and +reduce performance. Avoid repeated allocations in loops: even short-lived objects consume +heap space until the work-item exits. diff --git a/src/device/quirks.jl b/src/device/quirks.jl index 987922a5..1d151f1b 100644 --- a/src/device/quirks.jl +++ b/src/device/quirks.jl @@ -38,7 +38,7 @@ end @print_and_throw "sincos(x) is only defined for finite x." # diagonal.jl -# XXX: remove when we have malloc +# Base's version throws an ArgumentError; this one prints the reason import LinearAlgebra @device_override function Base.setindex!(D::LinearAlgebra.Diagonal, v, i::Int, j::Int) @boundscheck checkbounds(D, i, j) @@ -51,7 +51,7 @@ import LinearAlgebra end # number.jl -# XXX: remove when we have malloc +# Base's version throws a BoundsError; this one prints the reason @device_override @inline function Base.getindex(x::Number, I::Integer...) @boundscheck all(isone, I) || @print_and_throw "Out-of-bounds access of scalar value" From e5472b5d4f3e12b04cebbc12661e8d87f2273e6e Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sun, 6 Sep 2026 14:15:54 +0200 Subject: [PATCH 3/5] Avoid a checked integer conversion in the device allocator `LLVMPtr + UInt64` converts the offset through `Int`, which drags an `InexactError` throw path, its printed message and a baked host symbol pointer into every allocating kernel. The cursor is below the capacity, so reinterpret it instead. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Fy9hCZkRFuf9fbRQNagCxf --- src/device/runtime.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/device/runtime.jl b/src/device/runtime.jl index d7c83e0b..c7551b9b 100644 --- a/src/device/runtime.jl +++ b/src/device/runtime.jl @@ -46,7 +46,9 @@ function malloc(sz::Csize_t) bytes > capacity - cursor && return C_NULL # gc_pool_alloc reports exhaustion unsafe_store!(header, cursor + bytes, 1, Val(sizeof(Csize_t))) - return reinterpret(Ptr{Cvoid}, heap + HEAP_HEADER + cursor) + # `cursor` fits an `Int` (it is below the capacity); reinterpret rather than convert, + # as the checked conversion would drag an `InexactError` throw path into every kernel + return reinterpret(Ptr{Cvoid}, heap + HEAP_HEADER + (cursor % Int)) end function report_oom(sz) From 44ac4ca9cb6bac4a745362c49c662de3c9ca8775 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sun, 6 Sep 2026 14:15:54 +0200 Subject: [PATCH 4/5] Select boxes with branches in the device heap test Under `--check-bounds=yes`, as on CI, indexing a tuple of boxed objects with a run-time index emits a bounds check that reads the tuple type's field count through a host pointer, which yields garbage on the device and makes every work-item take the silent bounds-error exit. Pick the box with branches instead; the test still covers objects staying valid across later allocations. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Fy9hCZkRFuf9fbRQNagCxf --- test/execution.jl | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/test/execution.jl b/test/execution.jl index 03c6f9f6..3efcdbf1 100644 --- a/test/execution.jl +++ b/test/execution.jl @@ -781,11 +781,18 @@ end @oneapi items = 256 groups = n ÷ 256 ref_kernel(a) @test Array(a) == Float32.(2:(n + 1)) - # objects stay valid across later allocations by the same work-item + # objects stay valid across later allocations by the same work-item. Selected with + # branches rather than indexed from a tuple: under `--check-bounds=yes` the bounds check + # of a dynamic tuple index reads the field count through a host pointer, which cannot + # work on the device. function select_kernel(a, idx) i = get_global_id() - refs = ntuple(j -> Ref(a[i] * j), Val(4)) - @inbounds a[i] = heap_consume(refs[idx]) + r1 = Ref(a[i] * 1) + r2 = Ref(a[i] * 2) + r3 = Ref(a[i] * 3) + r4 = Ref(a[i] * 4) + r = idx == 1 ? r1 : idx == 2 ? r2 : idx == 3 ? r3 : r4 + @inbounds a[i] = heap_consume(r) return end a = oneArray(Float32[1, 2, 3, 4]) From 399e0c5af6e1a82603b8e32450397ffbafa61659 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sun, 6 Sep 2026 16:50:24 +0200 Subject: [PATCH 5/5] Require GPUCompiler 2.5.4 Kernels that read boxed fields, which the device heap now makes possible, contain `unordered` heap-reference accesses on Julia 1.12+ that the Khronos translator turns into invalid pointer-typed atomics. GPUCompiler 2.5.4 demotes them for SPIR-V targets. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Fy9hCZkRFuf9fbRQNagCxf --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 640cac57..60204c72 100644 --- a/Project.toml +++ b/Project.toml @@ -39,7 +39,7 @@ Adapt = "4" CEnum = "0.4, 0.5" ExprTools = "0.1" GPUArrays = "11.5.14" -GPUCompiler = "2" +GPUCompiler = "2.5.4" GPUToolbox = "0.1, 0.2, 0.3, 1, 3" KernelAbstractions = "0.9.39" LLVM = "6, 7, 8, 9"