Add a device-side allocator

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.

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.

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)
 

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"

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

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