Add alloca intrinsic for per-workitem stack scratch

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>

Author: vchuravy

Project.toml

@@ -1,6 +1,6 @@
 name = "GPUCompiler"
 uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
-version = "1.22.4"
+version = "1.23.0"
 authors = ["Tim Besard <tim.besard@gmail.com>"]
 
 [workspace]

src/irgen.jl

@@ -151,6 +151,10 @@ function irgen(@nospecialize(job::CompilerJob))
         # the job's configured level, so device code can branch on it as a compile-time
         # constant that is part of the cache key (unlike reading the `-g` global directly).
         lower_debug_level!(job, mod)
+
+        # materialize `GPUCompiler.alloca` intrinsics as real entry-block allocas, before the
+        # optimizer runs so the slots can be promoted (see `lower_alloca!`).
+        lower_alloca!(job, mod)
     end
 
     return mod, compiled, gv_to_value
@@ -1215,6 +1219,125 @@ function lower_debug_level!(@nospecialize(job::CompilerJob), mod::LLVM.Module)
     return true
 end
 
+
+## stack allocation
+
+# device code can request a fixed-size, per-workitem stack scratch buffer via
+# `alloca(T, Val(N))`, returning a `Ptr{T}` to uninitialized storage for `N` elements of
+# `T`. this emits the `julia.gpu.alloca.<bytes>.<align>` intrinsic, which `lower_alloca!`
+# (run from `irgen`, before the optimizer) materializes as a real entry-block `alloca`.
+#
+# this exists because emitting an `alloca` directly through `llvmcall` is unsound/ineffective:
+# the `Ptr` round-trip through `ptrtoint`/`inttoptr` blocks SROA/mem2reg promotion, the target
+# stack address space (e.g. AS 5 on NVPTX/AMDGPU) isn't known at the front end, and the
+# LangRef lifetime of an `alloca` is tied to the (inlined) `llvmcall` wrapper. lowering it
+# ourselves lets us place the slot in the kernel entry block, in the datalayout's alloca
+# address space, early enough for the optimizer to promote it.
+
+function alloca_intr(mod::LLVM.Module, bytes::Integer, align::Integer)
+    name = "julia.gpu.alloca.$(bytes).$(align)"
+    intr = if haskey(functions(mod), name)
+        functions(mod)[name]
+    else
+        # returns an opaque pointer; intentionally *not* readnone/speculatable, as each call
+        # must yield a distinct slot and must not be hoisted or CSE'd.
+        LLVM.Function(mod, name, LLVM.FunctionType(LLVM.PointerType()))
+    end
+    return intr
+end
+
+# run-time equivalent: emits a call to the alloca intrinsic, returning a `Ptr{T}` to scratch
+# storage for `N` elements of `T` (materialized by `lower_alloca!`).
+function alloca_value(@nospecialize(T), N::Int)
+    isbitstype(T) ||
+        error("GPUCompiler.alloca only supports `isbits` element types, got $T")
+    N >= 0 || throw(ArgumentError("GPUCompiler.alloca count must be non-negative, got $N"))
+
+    bytes = sizeof(T) * N
+    align = Base.datatype_alignment(T)
+
+    # a zero-byte allocation has no storage to point at; hand back a null pointer rather than
+    # emitting a degenerate 0-element alloca.
+    if bytes == 0
+        return :(reinterpret(Ptr{$T}, C_NULL))
+    end
+
+    @dispose ctx=Context() begin
+        T_ptr = LLVM.PointerType()
+
+        # create function
+        llvm_f, _ = create_function(T_ptr)
+        mod = LLVM.parent(llvm_f)
+
+        # get intrinsic
+        intr = alloca_intr(mod, bytes, align)
+        intr_ft = function_type(intr)
+
+        # generate IR
+        @dispose builder=IRBuilder() begin
+            entry = BasicBlock(llvm_f, "entry")
+            position!(builder, entry)
+
+            ptr = call!(builder, intr_ft, intr, Value[], "alloca")
+
+            ret!(builder, ptr)
+        end
+
+        call_function(llvm_f, Ptr{T})
+    end
+end
+
+# device-facing accessor: a `Ptr{T}` to per-workitem stack scratch for `N` elements of `T`.
+# the storage is uninitialized and only valid within the calling kernel. `T` must be `isbits`
+# (an `alloca` of GC-tracked references would be unrooted). intended as a building block for
+# higher-level scratch abstractions (e.g. KernelAbstractions' `@private`).
+@inline @generated alloca(::Type{T}, ::Val{N}) where {T,N} = alloca_value(T, N)
+export alloca
+
+# replace every `julia.gpu.alloca.*` call with an entry-block alloca in the containing function
+function lower_alloca!(@nospecialize(job::CompilerJob), mod::LLVM.Module)
+    changed = false
+    prefix = "julia.gpu.alloca."
+
+    for intr in collect(functions(mod))
+        fn = LLVM.name(intr)
+        startswith(fn, prefix) || continue
+
+        bytes, align = parse.(Int, split(fn[length(prefix)+1:end], '.'))
+        slot_typ = LLVM.ArrayType(LLVM.Int8Type(), bytes)
+
+        for use in collect(uses(intr))
+            call = user(use)
+            @assert call isa LLVM.CallInst
+            f = LLVM.parent(LLVM.parent(call))
+
+            @dispose builder=IRBuilder() begin
+                # materialize the slot at the top of the entry block so that it is a static
+                # alloca (promotable, and allocated once rather than per loop iteration).
+                position!(builder, first(instructions(first(blocks(f)))))
+                slot = alloca!(builder, slot_typ, "alloca")
+                alignment!(slot, align)
+
+                # `alloca!` placed the slot in the datalayout's alloca address space; cast back
+                # to generic (AS 0) to match the `Ptr` the front end handed out.
+                ptr = if LLVM.addrspace(value_type(slot)) == 0
+                    slot
+                else
+                    addrspacecast!(builder, slot, LLVM.PointerType())
+                end
+                replace_uses!(call, ptr)
+            end
+            erase!(call)
+            changed = true
+        end
+
+        @assert isempty(uses(intr))
+        erase!(intr)
+    end
+
+    return changed
+end
+
 # convert kernel state argument from pass-by-value to pass-by-reference
 #
 # the kernel state argument is always passed by value to avoid codegen issues with byval.

test/native.jl

@@ -739,3 +739,44 @@ end
     @test !occursin("deferred_codegen", ir)
     @test occursin("call void @julia_kernel", ir)
 end
+
+@testset "stack allocation intrinsic" begin
+    mod = @eval module $(gensym())
+        import ..GPUCompiler
+
+        function scratch(x)
+            p = GPUCompiler.alloca(Float32, Val(8))
+            @inbounds unsafe_store!(p, x, 1)
+            @inbounds unsafe_store!(p, x, 8)
+            return @inbounds unsafe_load(p, 1) + unsafe_load(p, 8)
+        end
+
+        # zero-element scratch yields a (null) pointer without emitting an alloca
+        empty_scratch() = GPUCompiler.alloca(Float32, Val(0)) === reinterpret(Ptr{Float32}, C_NULL)
+    end
+
+    # the intrinsic is materialized as a single entry-block `alloca [8 x f32 = 32 x i8]`,
+    # and no `julia.gpu.alloca` call/declaration survives lowering.
+    @test @filecheck begin
+        @check_label "define float @{{(julia|j)_scratch_[0-9]+}}"
+        @check "alloca [32 x i8], align 4"
+        @check_not "julia.gpu.alloca"
+        Native.code_llvm(mod.scratch, Tuple{Float32}; optimize=false)
+    end
+
+    # once optimized the slot is promoted away entirely (result is x + x).
+    @test @filecheck begin
+        @check_label "define float @{{(julia|j)_scratch_[0-9]+}}"
+        @check_not "alloca"
+        @check_not "julia.gpu.alloca"
+        Native.code_llvm(mod.scratch, Tuple{Float32})
+    end
+
+    # a zero-byte allocation lowers to a null pointer rather than a degenerate alloca.
+    @test @filecheck begin
+        @check_label "define {{.*}}@{{(julia|j)_empty_scratch_[0-9]+}}"
+        @check_not "alloca"
+        @check_not "julia.gpu.alloca"
+        Native.code_llvm(mod.empty_scratch, Tuple{})
+    end
+end