Add a device-side allocator - #631
Open
maleadt wants to merge 2 commits into
Open
Conversation
Contributor
|
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.
maleadt
force-pushed
the
tb/device-heap
branch
from
September 6, 2026 10:05
fd02bac to
914dedb
Compare
Member
Author
|
@michel2323 I figured that landing an allocator first is more important, instead of bundling several features in #621. |
Member
|
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. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.