Add alloca intrinsic for per-workitem stack scratch#859
Conversation
|
Conceptually fine, but I didn't realize we needed this.
Hmm, isn't this optimized away?
Yeah, I guess. How do we make sure Julia's
|
|
The crux is that llvmcall + alloca is UB and I have seen LLVM delete the alloca and replace it with poison. I can dig up the old PR I had for that. My overarching goal is to remove the StaticArrays dependency that Kernel abstraction currently forces and which you have complained about before since creates behavior where we have to hope for escape analysis to work. The rest is Claude obsessing over details. We have fixed alloca Julia to be emitted correctly |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #859 +/- ##
==========================================
- Coverage 79.77% 78.71% -1.06%
==========================================
Files 25 25
Lines 4796 4971 +175
==========================================
+ Hits 3826 3913 +87
- Misses 970 1058 +88 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
vchuravy
force-pushed
the
vc/alloca_intrinsic
branch
from
116c36b to
09e210a
Compare
Introduce `GPUCompiler.alloca(::Type{T}, ::Val{N})::Ptr{T}`, which hands device
code a fixed-size, per-workitem stack scratch buffer for `N` elements of `T`.
This is meant to replace abstractions like KernelAbstractions' `@private`
`MArray`-backed scratchpad with a direct stack allocation. Emitting the `alloca`
through `llvmcall` directly is unsound/ineffective: the `Ptr` round-trip through
`ptrtoint`/`inttoptr` blocks SROA/mem2reg promotion, the target stack address
space (e.g. addrspace 5 on NVPTX/AMDGPU) isn't known at the front end, and the
LangRef lifetime of the `alloca` is tied to the inlined `llvmcall` wrapper.
Instead, the front end emits a `julia.gpu.alloca.<bytes>.<align>` intrinsic that
`lower_alloca!` (run from `irgen`, before the optimizer) materializes as a real
entry-block `alloca` in the datalayout's alloca address space, cast back to
generic. Running before optimization lets the slot be promoted just like the
mutable stack allocations Julia already emits. `T` must be `isbits`.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
vchuravy
force-pushed
the
vc/alloca_intrinsic
branch
from
09e210a to
e13a8c5
Compare
Test that lowers correctly for all GPU targets
1. On Julia 1.10/1.11, LLVM runs in typed-pointers mode. PointerType() creates an opaque pointer in LLVM 15, causing assembly parsing errors on llvmcall since opaque pointers are disabled by default. Fix this by using PointerType(Int8Type()) when supports_typed_pointers() is true. 2. In lower_alloca!, if the pointer types or address spaces differ (which is common under typed pointers where the slot is [32 x i8]* but the call returns i8*), dynamically handle casting using bitcast! or addrspacecast! as appropriate. 3. On Julia 1.10/1.11, under optimize=false, the compiler compiles the llvmcall into a non-inlined helper function. Fix test checking by passing dump_module=true to code_llvm in the unoptimized checks to allow matching inside helper functions. 4. Dev the package in test/Project.toml so that Pkg loads the local repository package under Julia 1.10.
|
https://github.com/JuliaLang/julia/blob/3d710776cb31ce05775ee10290f709df1f1e939c/src/codegen.cpp#L2308 This is how base does it for the alloca address-ace |
- Encode size and alignment as constant call operands of a single `julia.gpu.alloca` declaration (matching `deferred_codegen`), instead of mangling them into the intrinsic name; `lower_alloca!` now guards with an O(1) name lookup like `lower_debug_level!`. - Replace the hand-rolled bitcast/addrspacecast selection with `pointercast!`, and hoist the IRBuilder out of the per-call loop. - Test the lowering on native (generic path) and GCN (nonzero alloca address space) only, dropping the identical PTX/SPIRV/Metal copies. - Don't export `alloca`; downstream users call it qualified. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
https://github.com/JuliaLang/julia/blob/3d710776cb31ce05775ee10290f709df1f1e939c/src/codegen.cpp#L2316-L2328 for the size of the alloys (we may also just want to do always power of 2) |
Materialize `julia.gpu.alloca` slots with an element type equal to the alignment (capped at 64 bits) rather than a flat `[bytes x i8]`. SROA takes a hint from the element type and will shred an i8 array into unaligned scalars; aligned integer chunks let it split into aligned pieces instead, which is much better for vectorization. Mirrors Julia's `emit_static_alloca`. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
The one thing I'm a little worried is using Ptr{} as a representation. Those addrspacecasts aren't free on amdgpu. addrspace(0) has an addressing cost when translating to private memory so it might be worth to check if the codegen of it is as expected |
|
I sadly can't get around it. We don't know the target during the generated function stage, but we must provide a stable return type so I can't use LLVMPtr |
|
Can't we make a type for this? Something like AllocaArray or something? Though I guess it needs more intrinsics then. Another option is to just remove the addrspacecast in a pass? Edit: In theory InferAddressSpaces is capable of removing these, but it deserves a double check (I'm not sure if our 64 bit indexing mucks that up) |
Open to all ideas, but I am essentially working around the fact here that Base Julia has no true support for alloca, and using StaticArrays is just "make a type" We can put the AS into the alloca intrinsic as a type parameter and then force each backend to use a device overlay |
The alloca intrinsic now takes a target address space and returns an
LLVMPtr{T,AS} instead of a Ptr{T}. lower_alloca! emits an addrspacecast
only when the requested AS differs from the datalayout's alloca AS. The
two-argument form defaults to AS 0 (generic).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
@gbaraldi I think I am happy with this. Since each backend will have to define |
Drop the two-argument convenience method so callers must always specify the target address space. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
3 participants
Summary
Adds
GPUCompiler.alloca(::Type{T}, ::Val{N})::Ptr{T}, a primitive that hands device code a fixed-size, per-workitem stack scratch buffer forNelements ofT. The intended use is to replace abstractions like KernelAbstractions'@privateMArray-backed scratchpad with a direct stack allocation (see companion KA PR).Why an intrinsic instead of
llvmcall+allocaEmitting the
allocadirectly throughllvmcallis unsound/ineffective:Ptrround-trips throughptrtoint/inttoptrat thellvmcallboundary, so SROA/mem2reg can never promote the slot to registers — which is the whole point of small@privatescratch.allocamust live in the datalayout's alloca address space (e.g. addrspace 5) and beaddrspacecastto generic; this can't be written portably across the PTX/GCN/Metal/SPIR-V back-ends from Julia.alloca's storage is freed when its function returns; in anllvmcallwrapper that's the wrapper, and correctness relies entirely on the inliner relocating a static entry-block alloca.How it works
julia.gpu.alloca.<bytes>.<align>intrinsic (modeled onjulia.gpu.debug_level).lower_alloca!, run fromirgenbefore the optimizer, materializes it as a real entry-blockallocain the datalayout's alloca address space (so SROA/mem2reg can promote it), cast back to generic for the returnedPtr. This produces the same IR shape Julia already emits for mutable stack allocations.Tmust beisbits; a zero-byte request returns a null pointer.Testing
Adds a native-target testset asserting an entry-block
alloca [N x i8], align <align>, full promotion when optimized, no surviving intrinsic, and the zero-byte → null path.Verified end-to-end (including execution) through the POCL/SPIR-V back-end via the companion KernelAbstractions change:
@private Float32 (4,)lowers toalloca [16 x i8], align 4in addrspace 0 (OpenCLFunction/private) and runs correctly.🤖 Generated with Claude Code