High-Performance Portable GPU Primitives for Arbitrary Types and Operators in Julia

GPUs execute thousands of threads organized into warps—the atomic unit of execution—and blocks of warps scheduled across Streaming Multiprocessors (SMs) [11]. A memory hierarchy trades capacity for speed: global memory is large but slow, shared memory is block-scoped and fast, and registers are per-thread and fastest. Threads within a warp can exchange data directly via shuffle instructions without touching memory [12]; warps within a block synchronize through shared memory [13]; cross-block coordination requires global memory fences or atomics [14, 15].

Memory access performance depends critically on coalescing—contiguous, aligned warp accesses are serviced in a single transaction—and on issuing wide loads (e.g., 128-bit float4 in CUDA) to maximize bandwidth per transaction [16, 17].

These three mechanisms—warp-level primitives, memory fences, and vectorized memory access—are the low-level capabilities required for vendor-competitive parallel primitives, and form the focus of KernelIntrinsics.jl (Section IV). Peak performance further requires tuning block count and items-per-thread to hardware characteristics such as warp width and L2 cache size that vary across architectures.

Julia’s GPU ecosystem is built on a layered stack. At the lowest level, vendor-specific packages—CUDA.jl for NVIDIA, AMDGPU.jl for AMD ROCm, oneAPI.jl for Intel, and Metal.jl for Apple—provide full access to each platform’s programming model and vendor libraries. Above this, KernelAbstractions.jl [6] provides a backend-agnostic kernel language: a single kernel definition can be compiled to any supported backend through Julia’s multiple dispatch and JIT compilation.

KernelAbstractions.jl exposes a programming model in which the user defines kernels using the @kernel macro. Kernels are parameterized by a backend (e.g., CUDABackend(), ROCBackend()) and compiled at first invocation for the target device. The framework provides workgroup-level abstractions—including workgroup indices, local indices, and @synchronize barriers—that map to thread blocks, thread-local indices, and __syncthreads() on NVIDIA, with analogous mappings on other backends. Shared memory is allocated via @localmem, which maps to __shared__ memory on CUDA and local data share on AMD.

Listing 1 illustrates the simplicity of the programming model. This abstraction covers a substantial portion of the GPU programming surface, enabling portable kernels for many workloads. However, several capabilities essential for vendor-competitive performance on fundamental primitives are not exposed:

• •

  Warp-level shuffle operations. There is no portable abstraction for shfl_sync, shfl_down_sync, or their equivalents. Kernels requiring register-level data exchange within a warp—critical for fast reductions and scans—must fall back to vendor-specific intrinsics, breaking portability.

• •

  Ordered memory accesses. Fine-grained memory fences (threadfence(), threadfence_block()) that control the visibility ordering of memory operations across threads, blocks, or the device are not available. These are required for inter-block coordination protocols such as the decoupled look-back used in single-pass scan algorithms [18].

• •

  Vectorized loads and stores. There is no mechanism to emit 64-bit or 128-bit load/store instructions. Achieving peak memory bandwidth in copy and bandwidth-bound kernels requires this capability.

These gaps motivate our KernelIntrinsics.jl layer (Section IV), which provides portable, backend-dispatched implementations of all three capabilities, enabling KernelForge.jl to build vendor-competitive algorithms without sacrificing portability.

Beyond the copy operation—which serves as the practical bandwidth ceiling against which other primitives are measured (cf. Figure 1)—several fundamental operations form the backbone of GPU computation, serving a role analogous to BLAS [19] in linear algebra. We focus on three: mapreduce, scan, and matrix–vector/vector–matrix products.

The mapreduce operation transforms each element via a mapping function $f:T\to S$, then combines all transformed values using a commutative and associative operator $\mathop{\mathrm{op}}$ (commutativity is required here, in contrast to scan which requires only associativity) to compute $\mathop{\mathrm{op}}(f(\texttt{src}[1]),f(\texttt{src}[2]),\ldots,f(\texttt{src}[n]))$. Efficient GPU reduction proceeds hierarchically: threads first reduce within registers, then within warps using shuffle operations, then across warps within a block using shared memory, and finally across blocks. Implementations leveraging warp-level primitives achieve up to $7.8\times$ speedup over tree-based reductions through shared memory alone [20, 12].

The scan (prefix sum) operation relaxes the commutativity requirement, taking an associative operator $\mathop{\mathrm{op}}$ (not necessarily commutative) and computing a sequence of partial reductions: the output $\texttt{dst}[i]=\mathop{\mathrm{op}}(\texttt{src}[1],\texttt{src}[2],\ldots,\texttt{src}[i])$ contains the accumulated result up to position $i$ for all $1\leq i\leq n$. Scan is among the most important parallel primitives: it enables stream compaction, radix sort, histogram computation, and sparse matrix operations [21, 22]. The state-of-the-art is the single-pass parallel prefix algorithm with decoupled look-back [18], which achieves throughput approaching copy operations by requiring only approximately $2n$ data movement ($n$ inputs read, $n$ outputs written) and dissociating local computation from global prefix propagation latencies through strategic redundant work. This protocol requires inter-block communication through global memory with careful use of memory fences and status flags—precisely the capabilities absent from existing portable frameworks – see Section II-B.

Finally, we treat matrix–vector and vector–matrix products as distinct primitives. Given an $n\times p$ matrix $A$, a vector $x\in T^{n}$, a mapping function $f:T\times T\to S$, and an associative reduction operator $\mathop{\mathrm{op}}$ over $S$, the matrix–vector product computes a vector $y\in S^{p}$ where $y[j]=\mathop{\mathrm{op}}_{i=1}^{n}f(x[i],\,A[i,j])$, reducing over the rows of $A$. The vector–matrix product symmetrically computes $z\in S^{n}$ where $z[i]=\mathop{\mathrm{op}}_{j=1}^{p}f(A[i,j],\,x[j])$, reducing over the columns. Setting $f=\times$ and $\mathop{\mathrm{op}}=+$ recovers the standard BLAS GEMV operation, but the generalized formulation supports arbitrary algebraic structures—for instance, tropical semirings (where $f=+$ and $\mathop{\mathrm{op}}=\min$) used in shortest-path algorithms, or log-space operations for numerical stability. These two operations are not symmetric on GPUs: because arrays are stored in column-major order, matrix–vector products access $A$ with stride-one (coalesced) reads along columns, while vector–matrix products must read along rows, resulting in strided access patterns. Achieving high throughput for both orientations requires distinct kernel strategies, as we describe in Section V-C.

Performance portability in GPU computing has been addressed through various C++ frameworks, each embodying different design philosophies. RAJA [23] achieves portability by decoupling algorithm implementation from execution strategy through execution policies. Kokkos [24] offers a more encompassing framework that abstracts both execution and memory spaces. Additional frameworks pursuing similar goals include OpenACC, OpenMP, and SYCL. Despite their differences, these frameworks share a fundamental constraint: their dependence on ahead-of-time compilation hinders integration with dynamic languages such as Python and Julia.

Performance studies reveal a consistent pattern: portable frameworks typically incur overhead compared to hand-optimized vendor implementations. Martineau et al. [3] found 5–30% penalties versus architecture-specific code; Artigues et al. [4] reported 2–3$\times$ slowdowns on NVIDIA V100s compared to native CUDA. Davis et al. [25] confirmed that Kokkos performs moderately worse than CUDA across five proxy applications on V100, A100, and H100 GPUs. Sedova et al. [1] concluded that non-portable optimizations are essential to production molecular dynamics codes. These studies establish the conventional wisdom that a portability tax is unavoidable when using abstraction layers.

For 1-D parallel primitives (copy, reduce, scan), NVIDIA’s CUB [26] and Thrust libraries provide highly optimized CUDA implementations, achieving near-peak performance through careful exploitation of warp-level operations, shared memory banking patterns, and instruction-level parallelism. AMD’s rocPRIM and Intel’s oneDPL offer similar functionality for their respective platforms. The existence of multiple vendor-specific libraries with incompatible APIs illustrates the fundamental challenge: achieving peak performance currently requires separate implementations per vendor.

For matrix–vector operations, vendor libraries such as cuBLAS, rocBLAS, and oneMKL provide optimized GEMV routines, but only for standard numeric types with fixed arithmetic ($f=\times$, $\oplus=+$). Applications requiring non-standard algebraic structures—such as tropical semirings for shortest-path problems or log-space operations for numerical stability—cannot use these libraries and must resort to custom kernels, forfeiting vendor-level optimization. KernelForge.jl’s generalized formulation subsumes standard GEMV while extending it to arbitrary types and operators, without sacrificing performance on the standard numerical case.

The Julia GPU ecosystem provides comprehensive GPU support through several packages. CUDA.jl offers both high-level array abstractions and low-level kernel programming for NVIDIA GPUs. KernelAbstractions.jl [6] extends portability to multiple backends (CUDA, ROCm, oneAPI, Metal) through unified kernel syntax and backend dispatch. AcceleratedKernels.jl [10] builds atop KernelAbstractions.jl to provide portable implementations of common operations. However, existing Julia GPU libraries prioritize generality over optimization: CUDA.jl’s implementations favor simplicity, while AcceleratedKernels.jl focuses on broad portability rather than matching vendor library performance.

KernelForge.jl differs from the frameworks above in two respects. First, unlike C++ portability layers that compile ahead of time, it exploits Julia’s JIT compilation to generate code specialized to the concrete data type, operator, and problem size at the call site—eliminating the abstraction overhead that accounts for the 5–30% portability tax reported in existing studies. Second, unlike Julia GPU libraries that prioritize API breadth, KernelForge.jl targets a narrow set of primitives—scan, mapreduce, and matrix–vector products—and optimizes them through a two-layer architecture: KernelIntrinsics.jl confines all vendor-specific functionality (warp shuffles, memory fences, vectorized loads) to a thin extension layer, while the algorithmic layer builds exclusively on these portable abstractions. On the A40, this design matches CUB and cuBLAS throughput on all three primitives; on the MI300X, it delivers competitive performance on mapreduce and scan relative to the available Julia baselines, and matches rocBLAS on matrix–vector operations for wide-matrix configurations.

A shuffle operation within a warp enables a given lane to read a value directly from the register of another lane, as determined by a source index or mask. Vendor implementations, however, define shuffle intrinsics only for a narrow set of types—typically 32-bit primitives such as Float32 and Int32.

CUDA.jl partially addresses this limitation through a recursion mechanism that extends shuffling to types like Int64 or complex numbers, but support for more composite types such as tuples or quaternions remains unavailable.

KernelIntrinsics.jl generalizes this approach by leveraging Julia’s @generated functions to recursively decompose any composite type at compile time. The key observation is that once shuffle intrinsics are defined for a single 32-bit primitive (e.g., UInt32), any composite type can be supported by decomposing it into its constituent primitive fields and shuffling each independently. We formalize this through a recursive definition: a Bitstype is either a concrete primitive of at most 64 bits, a tuple of Bitstypes, or a struct whose fields are all Bitstypes. Using @generated functions, the compiler statically unrolls the recursion over struct fields and tuple elements, producing specialized, zero-overhead code for each concrete type. KernelIntrinsics.jl implements shuffle operations over arbitrary Bitstypes under this definition, fully abstracting away vendor-specific type restrictions with no runtime cost.

This generalized shuffle mechanism enables fast warp-level communication for arbitrarily complex data types entirely through registers, without resorting to shared memory. KernelForge.jl builds on this capability to expose a high-level interface: the user can simply write mapreduce(f, op, src) or scan(f, op, src) and obtain correct, efficient results even when the elements of src are complex structures, tuples, or nested compositions thereof.

When a thread stores an element, the write is not instantly visible to all other threads. The value is first written to the L1 cache, where it is visible only to threads within the same workgroup. It is then propagated to the L2 cache, which is shared across workgroups. This creates a potential communication hazard: a thread in a different workgroup may load a stale value from its own L1 cache before the updated value has arrived in the L2 cache.

Consider a producer-consumer pattern where one workgroup writes data to global memory and then sets a flag to signal completion. A release semantic on the flag store guarantees that all prior writes are visible to any thread that subsequently observes the flag. A matching acquire semantic on the corresponding load guarantees that all subsequent reads observe up-to-date values. Together, release–acquire pairs establish a happens-before relationship between workgroups without requiring a full system-wide fence.

KernelIntrinsics.jl exposes this pattern through the @access macro: @access flag[i] = 0x01 emits a release store, while x = @access flag[i] emits an acquire load. On NVIDIA GPUs, these lower directly to PTX memory ordering annotations — for instance, st.release.gpu.global for the store and fence.acq_rel.gpu for a GPU-scoped fence. On AMD GPUs, the equivalent semantics are expressed through GCN scope bits (sc1/sc0) combined with explicit cache-flush and invalidation instructions (buffer_wbl2, buffer_inv). Both lowerings are verified at the assembly level in KernelIntrinsics.jl’s test suite, confirming that Julia’s compilation pipeline introduces no unintended fences or reorderings. For backends where fine-grained ordering annotations are unavailable, @access falls back to a full memory fence, which is less precise but preserves the same correctness guarantees with no changes to the algorithmic layer.

GPU kernel performance is limited by either compute throughput or memory bandwidth [27, 28]. For large problems that saturate memory bandwidth, vectorized loads—where each thread issues wide transactions (e.g., 128-bit loads of four Float32 values) rather than scalar ones—are essential to maximize bandwidth utilization, particularly for problems that fit within the L2 cache. KernelForge.jl defines its kernels with a static parameter Nitem that controls how many elements each thread processes, specialized at compile time via Julia’s Val type with zero overhead. The optimal Nitem is not necessarily dictated by the vector load width: the scan kernel uses 16 Float32 values per thread, since processing more elements sequentially amortizes synchronization cost across lanes and warps. Figure 1 illustrates this on a copy kernel, where 128-bit loads maximize bandwidth and KernelForge.jl matches or exceeds CUDA.jl’s internal libcuda implementation.

Vectorized loads require contiguous memory access. For strided subarrays, KernelForge.jl falls back to loading individual elements into a tuple, forgoing vectorized loads but preserving the multi-item-per-thread structure and its synchronization benefits.

Vectorized loads and stores provide substantial performance gains for problems that fit within the L2 cache, but they impose strict alignment requirements. The PTX ISA mandates that the address of any memory access be aligned to a multiple of the access size [29]. For instance, a ld.v4.f32 instruction loads 16 bytes and therefore requires the starting address to be 16-byte aligned. In practice, this means that a thread issuing a 128-bit load of Float32 elements can only begin reading at element indices that are multiples of 4 (using 0-based indexing). A load starting at index 2 or 3 would violate this constraint and result in undefined behavior or a hardware fault.

More generally, for a vectorized load of Nitem elements of size $s$ bytes, the byte address must be a multiple of $\texttt{Nitem}\times s$. This creates a critical issue for matrix operations: since a matrix may have a number of rows $n$ that is not a multiple of Nitem, the starting address of the second column (and subsequent columns) is not guaranteed to satisfy the alignment constraint, even if the first column is properly aligned.

A natural solution is to decompose each misaligned load at runtime into a sequence of smaller, properly aligned loads. For example, if Nitem $=4$ and the starting index $i$ has a misalignment of 2 (i.e., $(i-1)\bmod\texttt{Nitem}=2$), the load can be split into a scalar load at $i$, a vectorized 2-element load at $i+1$, and a scalar load at $i+3$. KernelIntrinsics.jl formalizes this through the function vload_pattern, which is a @generated function that emits an optimal load sequence for a statically known alignment pattern, expressed as a tuple of integers summing to Nitem (e.g., $(1,2,1)$ in the example above). The outer vload function then dispatches at runtime to the appropriate vload_pattern specialization via a compile-time-generated switch table, so that only the branch selection occurs at runtime while each load sequence itself is fully specialized.

Crucially, this mechanism is entirely hidden from the user. As shown in Listing 2, a vectorized copy kernel in KernelForge.jl requires only a call to vload and vstore!—alignment handling, pattern decomposition, and specialization are managed transparently by the library and at compile time using multiple dispatch.

The mapreduce primitive (Section II-C) is implemented with a fixed grid of blocks (e.g., 100 blocks of 256 threads, tuned per architecture). Each thread strides across the input array with stride equal to the total thread count, accumulating a partial result in registers. This fixed-grid strategy processes arbitrarily large arrays without growing the launch configuration.

Within each block, partial results are reduced hierarchically: first across threads within each warp using shuffle operations, then across warps via shared memory, yielding one partial value per block. Warp-level shuffles are provided by KernelIntrinsics.jl (Section IV-A) and support arbitrary composite types at no runtime cost.

Inter-block aggregation is handled without a second kernel launch. Each block writes its partial result to global memory and raises a UInt8 completion flag, initialized to zero before launch. A release-semantic store, issued via the @access macro, guarantees that the partial result is visible to all other blocks before the flag is observed. A designated block spins on all flags using acquire-semantic loads—also via @access— and performs the final reduction once all partial results are confirmed visible. This single-launch design avoids the two-kernel sequence used by CUDA.jl.

The implementation selects among three paths based on input dimensionality. One-dimensional arrays use the strategy above. Two-dimensional arrays are dispatched to the matrix-vector and vector-matrix kernels (Section V-C), which are already optimized for coalesced access and warp-level parallelism along both axes. Higher-dimensional arrays and arbitrary iterators fall through to a general reduction kernel, at a modest performance cost relative to the specialized paths.

The scan primitive computes an inclusive or exclusive prefix reduction over an input array using any associative operator. It shares the same building blocks as mapreduce—vectorized loads, warp-level shuffles, and release/acquire flag synchronization via @access—and achieves optimal single-pass throughput via the decoupled lookback algorithm of Merrill and Garland [18].

The input is partitioned into tiles of $256\times\texttt{Nitem}$ elements, one tile per thread block. Each thread loads Nitem elements using vload (Section IV-C), accumulates a local prefix entirely in registers, then participates in a warp-shuffle reduction and a shared-memory exchange to compute the tile’s aggregate. Global memory is therefore read exactly once per element. Upon completion, a designated thread stores the tile aggregate and raises a partial status flag via a release store (@access flag[i] = PARTIAL).

Global prefix propagation proceeds via the lookback phase. A single warp scans backwards over the preceding 32 tiles (64 on AMD wavefront-64 hardware), spinning on each status flag via acquire loads (@access flag[j]) until at least a partial flag is visible. If all 32 (or 64) preceding tiles have raised only partial flags, the warp reduces their aggregates via shuffle-based reduction and continues looking back at the next group of tiles. As soon as the warp encounters a tile with a prefix flag—indicating that the complete inclusive prefix through that tile is already available—it can immediately finalize the current tile’s inclusive prefix without inspecting any earlier tiles. The finalized prefix is written to global memory and the tile raises its own prefix flag via a release store, unblocking all subsequent tiles. The correctness of this scheme relies on the release–acquire ordering established by @access (Section IV-B): without it, a tile could observe a raised flag while the associated aggregate value remains stale in a remote L1 cache.

Once the inclusive prefix is known, each thread computes its final output values entirely in registers and writes them back to the destination array using vectorized stores, so global memory is written exactly once per element.

The resulting implementation is correct for any associative operator, including noncommutative cases such as quaternion multiplication, and achieves throughput competitive with CUB on standard types (Section VII).

Matrix-vector and vector-matrix products are central to neural network inference and training, and are typically served by highly optimized proprietary libraries such as cuBLAS. These libraries, however, are vendor-locked and restricted to standard numeric arithmetic. KernelForge.jl provides open-source implementations that match cuBLAS throughput on CUDA for standard types, while supporting arbitrary element types and associative operator pairs. This generality enables use cases beyond the standard $(\times,+)$ semiring— such as tropical semiring operations for shortest-path computations or log-space accumulation for numerical stability—without sacrificing performance on the standard numerical case.

Because matrices are stored in column-major order, horizontal reductions (matrix-vector products) and vertical reductions (vector-matrix products) access memory along different axes and require distinct kernels with different coalescing strategies (Section II-C). The optimal thread organization also depends on matrix shape. For tall, narrow matrices, each column resembles an independent 1-D reduction: the same fixed-grid block striding used in mapreduce applies directly, with a small number of blocks (default: 100, tunable per architecture) assigned per column. For wide, short matrices, work must instead be distributed across both dimensions to keep all threads occupied; Figure 2 illustrates this layout for the vector-matrix case. KernelForge.jl selects the appropriate strategy at kernel launch based on matrix shape, with dispatch resolved statically via Julia’s Val type so that each path is independently specialized by the compiler with no runtime overhead.

We evaluate KernelForge.jl on the AMD MI300X against AMDGPU.jl and AcceleratedKernels.jl (cf. Tables I–VIII). Unlike the NVIDIA setting, where CUDA.@profile isolates individual kernel segments, ROCm timing relies on HIP events (AMDGPU.HIP.HIPEvent), which measure device-side elapsed time including kernel launch overhead and any on-device allocations. CPU-side pipeline timing is avoided, as it exhibits substantially higher variance on ROCm—in the worst cases, the standard deviation can exceed the mean when temporary or destination buffers are allocated within the timed region.

To mitigate this, KernelForge.jl and AcceleratedKernels.jl are benchmarked with pre-allocated temporary buffers—supplied via KernelForge’s get_allocation interface or AcceleratedKernels.jl’s temp keyword argument—which both expose uniformly across all primitives. AMDGPU.jl does not provide an equivalent facility and is therefore benchmarked with internal allocation; its reported times should be interpreted as upper bounds on what that library could achieve with a pre-allocation interface, rather than as a measure of its kernel throughput alone. Consequently, comparisons against AMDGPU.jl primarily demonstrate the practical advantage of KernelForge’s allocation-forwarding design over libraries that lack this facility, not algorithmic superiority per se. As with CUDA.jl dispatching to cuBLAS, AMDGPU.jl dispatches LinearAlgebra.mul! to rocblas_gemv internally, making it the effective rocBLAS baseline for matrix–vector and vector–matrix primitives.