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" 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/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/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" diff --git a/src/device/runtime.jl b/src/device/runtime.jl index 2f48b0f1..c7551b9b 100644 --- a/src/device/runtime.jl +++ b/src/device/runtime.jl @@ -1,36 +1,78 @@ # 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))) + # `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) + @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..3efcdbf1 100644 --- a/test/execution.jl +++ b/test/execution.jl @@ -742,3 +742,127 @@ 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. 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() + 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]) + @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