Device runtime: host-visible exceptions, a per-work-item malloc, and a compile-smoke grid

docs/src/kernels.md

@@ -62,3 +62,31 @@ 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.
 
+
+## Exceptions and Dynamic Allocation
+
+Kernels can throw. An exception on the device aborts the work-item that threw it and is
+reported on the host as a `KernelException` at the next synchronization — `synchronize()`,
+`oneAPI.@sync`, or copying data back with `Array` — after the reason was printed by the
+device:
+
+```julia
+julia> function kernel(a)
+           a[2] = 1f0     # bounds-checked
+           return
+       end;
+
+julia> @oneapi kernel(oneArray(Float32[0]));
+
+julia> synchronize()
+ERROR: Out-of-bounds array access.
+ERROR: KernelException: exception thrown during kernel execution on device Intel(R) Data Center GPU Max 1550
+```
+
+Exception objects that survive optimization, and any other Julia object that has to be
+heap-allocated inside a kernel (for example a `Ref` passed to a `@noinline` function), are
+allocated from a small per-work-item heap in private memory; objects never outlive the
+work-item that created them and are never freed. The heap is limited to
+`oneAPI.PRIVATE_HEAP_SIZE` bytes per work-item, and exhausting it is reported as a
+`KernelException` as well, rather than failing silently. Code on a hot path should not
+allocate.

src/compiler/compilation.jl

@@ -30,6 +30,10 @@ end
 
 GPUCompiler.runtime_module(::oneAPICompilerJob) = oneAPI
 
+# hidden first kernel argument carrying the exception flag and the private heap; see
+# src/device/runtime.jl, src/exceptions.jl and `add_private_heap!` below
+GPUCompiler.kernel_state_type(::oneAPICompilerJob) = KernelState
+
 GPUCompiler.method_table_view(job::oneAPICompilerJob) =
     GPUCompiler.StackedMethodTable(job.world, method_table, SPIRVIntrinsics.method_table)
 
@@ -64,6 +68,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)
+    # before the kernel state is turned into a by-reference argument below
+    job.config.kernel && add_private_heap!(mod, entry)
+
     entry = invoke(GPUCompiler.finish_ir!,
                    Tuple{CompilerJob{SPIRVCompilerTarget}, typeof(mod), typeof(entry)},
                    job, mod, entry)
@@ -92,6 +99,39 @@ function GPUCompiler.finish_ir!(job::oneAPICompilerJob, mod::LLVM.Module,
     return entry
 end
 
+# Give the kernel a private heap for `malloc` (src/device/runtime.jl): allocate the arena in
+# the entry block of the kernel, initialize its cursor, and thread the pointer into the kernel
+# state that the entry passes on to every device function. Runs after the kernel-state passes,
+# so the state is the entry's first (by-value) parameter; only kernels whose code reaches
+# `gpu_malloc` pay for the arena.
+function add_private_heap!(mod::LLVM.Module, entry::LLVM.Function)
+    haskey(functions(mod), "gpu_malloc") || return false
+    T_state = convert(LLVMType, KernelState)
+    params = parameters(entry)
+    (isempty(params) || value_type(params[1]) != T_state) && return false
+    state = params[1]
+    isempty(uses(state)) && return false
+
+    users = LLVM.Value[user(use) for use in uses(state)]
+    @dispose builder = IRBuilder() begin
+        position!(builder, first(instructions(first(blocks(entry)))))
+        T_i8 = LLVM.Int8Type()
+        T_i32 = LLVM.Int32Type()
+        arena = alloca!(builder, LLVM.ArrayType(T_i8, HEAP_HEADER + PRIVATE_HEAP_SIZE), "private_heap")
+        alignment!(arena, 16)
+        store!(builder, ConstantInt(T_i32, 0), arena)
+        heap_field = findfirst(==(:heap), fieldnames(KernelState)) - 1
+        new_state = insert_value!(builder, state, arena, heap_field, "state_with_heap")
+        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
+
 # 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,12 @@ abstract type AbstractKernel{F,TT} end
         end
     end
 
+    # the kernel state is the hidden first argument of every compiled kernel (see
+    # `GPUCompiler.kernel_state_type`); `onecall` itself stays agnostic so that it can
+    # also launch kernels that were not compiled by us
+    pushfirst!(call_t, KernelState)
+    pushfirst!(call_args, :(kernel.state))
+
     # finalize types
     call_tt = Base.to_tuple_type(call_t)
 
@@ -209,6 +215,8 @@ end
 struct HostKernel{F,TT} <: AbstractKernel{F,TT}
     f::F
     fun::ZeKernel
+    # for the context and device `fun` was linked against
+    state::KernelState
 end
 
 # Upper bound on the spill (scratch) memory a single work-group may require, in bytes.
@@ -297,7 +305,7 @@ function zefunction(f::F, tt::TT=Tuple{}; kwargs...) where {F,TT}
         # about world age here, as GPUCompiler already does and will return a different object
         h = hash(fun, hash(f, hash(tt)))
         get!(_kernel_instances, h) do
-            HostKernel{F,tt}(f, fun)
+            HostKernel{F, tt}(f, fun, kernel_state(ctx, dev))
         end::HostKernel{F,tt}
     end
 end
@@ -346,7 +354,10 @@ end
     spill > s.scratch_hwm && scratch_hedge!(s, spill)
 
     append_launch!(s.list, kernel, groups)
-    oneL0.sync_each_submission() && oneL0.synchronize(s.list)
+    if oneL0.sync_each_submission()
+        oneL0.synchronize(s.list)
+        check_exceptions(s.ctx, s.dev)
+    end
     return
 end
 
@@ -358,6 +369,8 @@ end
     execute!(queue) do list
         append_launch!(list, kernel, groups)
     end
+    oneL0.sync_each_submission() && check_exceptions(queue.context, queue.device)
+    return
 end
 
 # Slow path of the scratch hedge, firing once per (stream, spill tier): retire in-flight

src/context.jl

@@ -405,6 +405,10 @@ function oneL0.synchronize(s::oneStream)
         oneL0.synchronize(q)
         s.mkl_dirty = false
     end
+    # every user-facing synchronization funnels through here, so this is where a device
+    # exception surfaces (src/exceptions.jl); `synchronize_all_streams` deliberately does
+    # not check, it runs from finalizers
+    check_exceptions(s.ctx, s.dev)
     return
 end
 

src/device/quirks.jl

@@ -29,6 +29,22 @@ end
 @device_override @noinline Core.throw_inexacterror(f::Symbol, ::Type{T}, val) where {T} =
     @print_and_throw "Inexact conversion"
 
+# Base also constructs these exceptions directly, without a throw helper that could be
+# overridden: `Int32(::Float32)` and `round(Int, ::Float64)` (float.jl), `x^y` for `Complex`
+# (`_cpow`), and the local `throw1`/`throw2` closures of `exponent` (math.jl) that
+# `sqrt(::Complex)` reaches through `ssqs`. Left alone, such a throw allocates the exception
+# object on the device heap and signals the host without printing anything; replacing the
+# constructors covers every such site at once with a printed reason (at the cost of its
+# specificity) and keeps the heap out of it. Unlike Base's `@nospecialize`d inner
+# constructors these are specialized and inlined: a `@noinline` callee taking `Any` would
+# have to box its (e.g. `Float32`) argument, which is itself an allocation.
+@device_override @inline Core.InexactError(f::Symbol, args...) =
+    @print_and_throw "Inexact conversion"
+@device_override @inline Core.DomainError(val) =
+    @print_and_throw "Argument outside the domain of the function"
+@device_override @inline Core.DomainError(val, msg::AbstractString) =
+    @print_and_throw "Argument outside the domain of the function"
+
 # abstractarray.jl
 @device_override @noinline Base.throw_boundserror(A, I) =
     @print_and_throw "Out-of-bounds array access"
@@ -38,7 +54,8 @@ 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 built from a string; this one prints and keeps the
+# device heap out of the hot path.
 import LinearAlgebra
 @device_override function Base.setindex!(D::LinearAlgebra.Diagonal, v, i::Int, j::Int)
     @boundscheck checkbounds(D, i, j)
@@ -51,7 +68,7 @@ import LinearAlgebra
 end
 
 # number.jl
-# XXX: remove when we have malloc
+# Base's version throws a BoundsError; same reasoning as above.
 @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,9 +1,50 @@
 # device runtime libraries
+#
+# GPUCompiler resolves the back-end runtime by name in this module (`runtime_module`):
+# `signal_exception`, `report_*` and `malloc` below are compiled into the runtime library
+# that is linked into every kernel, with `malloc` becoming the `gpu_malloc` symbol that
+# `gc_pool_alloc` — and so every heap allocation that survives optimization — calls.
 
 
-## Julia library
+## kernel state
+
+# Passed by value as the hidden first argument of every kernel and threaded by GPUCompiler
+# to every device function that calls `kernel_state()`. The exception flag is filled in on
+# the host (`kernel_state` in src/exceptions.jl, once per linked kernel); the heap pointer
+# is patched in on the device, in the kernel's entry block (`add_private_heap!` in
+# src/compiler/compilation.jl), because it points at private memory.
+struct KernelState
+    # host USM, 16 bytes, read by the host after synchronization and cleared by it:
+    #   [1] nonzero after `signal_exception`, [2] nonzero after `report_oom`
+    exception_flag::LLVMPtr{Int32, AS.CrossWorkgroup}
+    # per-work-item private memory: a `HEAP_HEADER`-byte header whose first word is the
+    # bump cursor, followed by `PRIVATE_HEAP_SIZE` allocatable bytes; null in kernels that
+    # do not allocate
+    heap::Ptr{UInt8}
+end
+
+# Bytes of private memory set aside for dynamic allocations, per work-item. Julia objects
+# allocated in a kernel never outlive the work-item that created them — exception objects
+# on a throw path, boxes handed to a `@noinline` callee — so private memory is the right
+# place for them, and the only one: address space 0, which Julia's boxed objects live in
+# after GPUCompiler strips its address spaces, is private memory to SPIR-V and Intel's
+# compiler, and a store through an address-space-0 pointer into global memory is silently
+# lost. The arena is only materialized in kernels whose code calls `malloc`.
+const PRIVATE_HEAP_SIZE = 1024
+const HEAP_HEADER = 16     # keeps the first allocation 16-byte aligned
+
+@inline @generated kernel_state() = GPUCompiler.kernel_state_value(KernelState)
+
+
+## exceptions
 
 function signal_exception()
+    unsafe_store!(kernel_state().exception_flag, Int32(1), 1)
+    return
+end
+
+function report_oom(sz)
+    unsafe_store!(kernel_state().exception_flag, Int32(1), 2)
     return
 end
 
@@ -15,8 +56,6 @@ function report_exception(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.
@@ -31,6 +70,21 @@ function report_exception_frame(idx, func, file, line)
 end
 
 
-## SPIRV libraries
+## dynamic memory allocation
 
-# TODO
+# Bump allocator over the work-item's private arena: nothing is ever freed (the arena dies
+# with the work-item), and exhaustion returns null, which `gc_pool_alloc` turns into
+# `report_oom` + `OutOfMemoryError`, i.e. a loud `KernelException` on the host rather than
+# a silent failure. The cursor is private to the work-item, so no atomics are needed.
+function malloc(sz::Csize_t)
+    heap = kernel_state().heap
+    heap == C_NULL && return C_NULL
+    sz > PRIVATE_HEAP_SIZE && return C_NULL
+    bytes = (UInt32(sz) + UInt32(15)) & ~UInt32(15)
+    cursor = convert(Ptr{UInt32}, heap)
+    old = unsafe_load(cursor)
+    new = old + bytes
+    new > PRIVATE_HEAP_SIZE && return C_NULL
+    unsafe_store!(cursor, new)
+    return Ptr{Cvoid}(heap + HEAP_HEADER + old)
+end

src/exceptions.jl

@@ -0,0 +1,83 @@
+# device exceptions
+#
+# Every kernel receives a `KernelState` (src/device/runtime.jl) that names a host-visible
+# flag owned by the (context, device) it was linked for. The device runtime sets the flag
+# when a kernel throws; the host reads it whenever it synchronizes a stream and raises a
+# `KernelException`.
+
+export KernelException
+
+
+## exception type
+
+"""
+    KernelException
+
+Thrown on the host, at the next synchronization, after a kernel threw an exception on the
+device. The device runtime prints the reason (e.g. `ERROR: Out-of-bounds array access.`)
+to the standard output of the process at the time of the throw; this exception only carries
+the device and whether the work-item ran out of private heap memory along the way.
+"""
+struct KernelException <: Exception
+    dev::ZeDevice
+    oom::Bool
+end
+
+function Base.showerror(io::IO, err::KernelException)
+    name = oneL0.properties(err.dev).name
+    print(io, "KernelException: exception thrown during kernel execution on device $name")
+    if err.oom
+        print(
+            io, " (a work-item allocated more than the $(PRIVATE_HEAP_SIZE) bytes of ",
+            "private heap memory available to it)"
+        )
+    end
+    return
+end
+
+
+## exception flags
+
+# one 16-byte host USM buffer per (context, device): [1] exception, [2] oom (Int32 each)
+const exception_flags = Dict{Tuple{ZeContext, ZeDevice}, oneL0.HostBuffer}()
+const exception_flags_lock = ReentrantLock()
+
+function exception_flag(ctx::ZeContext, dev::ZeDevice)
+    return Base.@lock exception_flags_lock get!(exception_flags, (ctx, dev)) do
+        flag = oneL0.host_alloc(ctx, 16, 16)
+        # a pointer embedded in the kernel state is an indirect access as far as Level Zero
+        # is concerned, so the buffer needs explicit residency
+        oneL0.make_resident(ctx, dev, flag)
+        p = convert(Ptr{Int32}, flag)
+        unsafe_store!(p, Int32(0), 1)
+        unsafe_store!(p, Int32(0), 2)
+        flag
+    end
+end
+
+# the kernel state for kernels linked against `ctx`/`dev`; built once per `HostKernel`.
+# The private heap pointer is filled in on the device (`add_private_heap!`).
+function kernel_state(ctx::ZeContext, dev::ZeDevice)
+    flag = exception_flag(ctx, dev)
+    return KernelState(
+        reinterpret(LLVMPtr{Int32, AS.CrossWorkgroup}, pointer(flag)),
+        C_NULL
+    )
+end
+
+
+## host-side check
+
+# Called after a stream has been synchronized. Clears the flag words so the exception is
+# reported once; the clear is an atomic swap so two tasks synchronizing the same device
+# cannot both report one throw.
+function check_exceptions(ctx::ZeContext, dev::ZeDevice)
+    flag = Base.@lock exception_flags_lock get(exception_flags, (ctx, dev), nothing)
+    flag === nothing && return
+    p = convert(Ptr{Int32}, flag)
+    unsafe_load(p, 1) == 0 && return
+    thrown = Core.Intrinsics.atomic_pointerswap(p, Int32(0), :sequentially_consistent)
+    oom = Core.Intrinsics.atomic_pointerswap(p + sizeof(Int32), Int32(0), :sequentially_consistent)
+    thrown == 0 && return
+    throw(KernelException(dev, oom != 0))
+end

src/oneAPI.jl

@@ -51,6 +51,7 @@ include("context.jl")
 include("memory.jl")
 include("pool.jl")
 include("array.jl")
+include("exceptions.jl")
 
 # compiler implementation
 include("compiler/compilation.jl")