Add a device-side allocator - #631
Conversation
|
Your PR no longer requires formatting changes. Thank you for your contribution! |
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.
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.
fd02bac to
914dedb
Compare
|
@michel2323 I figured that landing an allocator first is more important, instead of bundling several features in #621. |
|
Thank you! I'm OOO until next week, but I'll try to read it before. I'm glad you're taking a look. Obviously, I heavily used Claude and I did that bundled PR along benchmarking a code. |
`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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Fy9hCZkRFuf9fbRQNagCxf
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Fy9hCZkRFuf9fbRQNagCxf
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Fy9hCZkRFuf9fbRQNagCxf
|
CI failure is JuliaGPU/GPUCompiler.jl#924 |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #631 +/- ##
==========================================
+ Coverage 81.31% 81.50% +0.19%
==========================================
Files 50 50
Lines 3569 3607 +38
==========================================
+ Hits 2902 2940 +38
Misses 667 667 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
I've merged it already in order to move forwards with a GPUCompiler change depending on this, but feel free to still review; I'll address in follow-up PRs then. |
Implement GPUCompiler's
mallochook so Julia objects that survive optimization can beallocated in oneAPI kernels. This includes boxed values,
Refs passed to non-inlinedfunctions, and allocations on exception paths.
The motivating failure is GPUCompiler #906:
preserving inferred invoke specializations exposed allocations inside
DomainErrorconstructors, causing kernels such as
sqrt.(::oneArray{ComplexF32})to fail compilationwith an unresolved
gpu_malloc. GPUCompiler #908now proposes a null-returning fallback for backends without an allocator. This PR supplies
actual storage, so surviving allocations can execute instead of always taking the OOM path.
Each work-item gets a 1 KiB private arena. Allocations advance a cursor in 16-byte increments;
nothing is freed individually, and the arena is reclaimed when the work-item exits.
Exhaustion prints
ERROR: Out of dynamic GPU memory (trying to allocate N bytes)andterminates that work-item. It does not raise a host-side exception, so output may be incomplete.
Objects in the arena must not be shared between work-items or retained across launches.
Private memory matches the pointers Julia emits for boxed objects: address space 0 maps
to SPIR-V private memory. A USM buffer cannot back those pointers on Intel GPUs; the
earlier investigation in #621 reported
lost stores through such pointers on a Max 1550.
GPUCompiler threads a hidden
KernelStateargument through device code. Before SPIR-Vargument lowering, the compiler reserves the arena in the kernel entry block, initializes
its cursor and capacity, and puts its pointer in the state. The pass looks for reads of
the heap field, so it also handles inlined allocator calls. Kernels without those reads
reserve no arena. Device compilation may optimize away heap storage, but there is no
general guarantee that allocations are free.
This extracts the allocator approach from #621 without its exception mailbox, constructor
overrides, or broader compilation tests. Host-visible exception reporting remains separate.
The existing error-printing overrides remain useful because ordinary device exceptions
still do not report their reason to the host.