Skip to content

Add a device-side allocator - #631

Open
maleadt wants to merge 2 commits into
mainfrom
tb/device-heap
Open

Add a device-side allocator#631
maleadt wants to merge 2 commits into
mainfrom
tb/device-heap

Conversation

@maleadt

@maleadt maleadt commented Sep 6, 2026

Copy link
Copy Markdown
Member

Implement GPUCompiler's malloc hook so Julia objects that survive optimization can be
allocated in oneAPI kernels. This includes boxed values, Refs passed to non-inlined
functions, and allocations on exception paths.

The motivating failure is GPUCompiler #906:
preserving inferred invoke specializations exposed allocations inside DomainError
constructors, causing kernels such as sqrt.(::oneArray{ComplexF32}) to fail compilation
with an unresolved gpu_malloc. GPUCompiler #908
now 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) and
terminates 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 KernelState argument through device code. Before SPIR-V
argument 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.

@github-actions

github-actions Bot commented Sep 6, 2026

Copy link
Copy Markdown
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

maleadt commented Sep 6, 2026

Copy link
Copy Markdown
Member Author

@michel2323 I figured that landing an allocator first is more important, instead of bundling several features in #621.

@michel2323

michel2323 commented Sep 6, 2026

Copy link
Copy Markdown
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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants