Fast GPU Linear Algebra via Compile Time Expression Fusion

Abstract: We describe the Bandicoot GPU linear algebra toolkit, a C++ based library that prioritises ease of use without compromising efficiency. Bandicoot's API is compatible with the popular Armadillo CPU linear algebra library, enabling easy transition for existing CPU-based codebases. Unlike other GPU-focused toolkits, Bandicoot uses template metaprogramming to generate fused GPU kernels directly at compile time, yielding efficient kernels that are often able to saturate memory bandwidth. This removes the need for runtime overhead or JIT infrastructure. Empirical results show that Bandicoot outperforms (sometimes by considerable margins) commonly-used linear algebra toolkits including PyTorch, TensorFlow, and JAX.

• The paper introduces Bandicoot, a C++ GPU linear algebra library that compiles fused expressions into optimized single kernels.
• Its compile-time fusion technique eliminates runtime overhead and intermediate data transfers, outperforming JIT-based frameworks.
• Empirical results on large-scale matrices demonstrate that Bandicoot saturates memory bandwidth, ensuring scalability and smooth integration with Armadillo-based pipelines.

Fast GPU Linear Algebra via Compile Time Expression Fusion

Introduction and Motivation

The paper "Fast GPU Linear Algebra via Compile Time Expression Fusion" (2604.22242) introduces Bandicoot, a C++ GPU linear algebra library designed to fuse compound expressions into efficient GPU kernels at compile time, rather than at runtime or via just-in-time (JIT) infrastructure. Given the dominance of GPUs in contemporary scientific computing and machine learning, achieving both high-level usability and hardware-efficient computation across heterogeneous architectures remains a central challenge. Existing solutions either require deep low-level programming skills (e.g., CUDA/cuBLAS) or incur significant runtime and memory overheads through high-level abstractions (MATLAB, PyTorch, TensorFlow, JAX). Bandicoot targets an intersection of portability, ease of use, and maximal hardware throughput by leveraging template metaprogramming to perform all expression fusion and kernel generation as part of the ahead-of-time (AOT) compilation step.

Design and User Experience

Bandicoot exposes an API closely matching the widely-used CPU-based Armadillo linear algebra toolkit. This compatibility enables straightforward migration of pre-existing Armadillo-centric computational pipelines to GPU acceleration.

A typical use case exemplifies Bandicoot's ease of integration: a line-by-line replacement of Armadillo headers and namespaces ports code from CPU to GPU, with automatic asynchronous dispatch of computation and minimal mandatory device management. GPU synchronization occurs only when required (such as matrix printing or explicit synchronization calls).

Figure 1: C++ program using Bandicoot for GPU execution, directly paralleling the Armadillo CPU-based alternative.

Bandicoot situates itself as a practical "drop-in" for high-level scientific applications and machine learning systems historically developed atop CPU-bound libraries. Despite this ease of use, developers need to adapt their computational paradigms to matrix-level manipulations (as opposed to scalar-level or elementwise loops), due to costly CPU-GPU memory transfers associated with fine-grained element access.

Compile-Time Expression Fusion and Kernel Generation

Fundamentally, Bandicoot diverges from standard GPU linear algebra libraries in its exploitation of C++ template metaprogramming to fuse entire compound mathematical expressions into optimized single kernels at compile time. Most other frameworks—such as PyTorch, TensorFlow, or ArrayFire—analyze and/or JIT-compile operations at runtime, incurring overhead and producing superfluous intermediate data movements and kernel launches, especially for composite expressions. In Bandicoot, the full expression is represented as an abstract syntax tree (AST) via C++ types, enabling the compiler to both fuse and optimize all operations without generating any unnecessary temporary storage or kernel launch overhead.

The internal architecture of Bandicoot is structured to maintain strict modularity. The device/backend layer (currently supporting CUDA and OpenCL, with others such as HIP/ROCm and Metal planned) is cleanly decoupled from the expression optimization stack and the high-level API.

Figure 2: Bandicoot's layered design, with user code compiled into backend-specific, fused GPU kernels, supporting portability and backend extensibility.

Kernel source code for fused operations is generated at compile time using skeleton templates parameterized by macros that encode the structure of the target expression. For example, in-memory AST traversal at compile time yields instantiations of code that precisely reflect the fused operations needed, embedding only the requisite kernels into the binary—no runtime selection or JIT infrastructure is needed.

Figure 3: Example of a Bandicoot compile-time-generated skeleton kernel for a simple copy operation.

For the expression $Y = X + 3$, the macro expansion defines in-kernel element access as a direct mapping from the AST, with all operators present as inlined, fused source. This methodology ensures single-pass computation over matrices, minimizes PCIe/memory transaction counts, and saturates memory bandwidth in bandwidth-limited scenarios.

Empirical Evaluation and Performance

The paper provides a comprehensive empirical comparison of Bandicoot against prominent GPU libraries—including PyTorch, TensorFlow, JAX, and ArrayFire—over a suite of expressions encompassing basic matrix operations, compound algebraic expressions, and neural network activation functions. Benchmarks are conducted on square matrices of dimension $10{,}000 \times 10{,}000$ (each $\sim 384$ GB), utilizing an NVIDIA RTX 4090 GPU with repeated trials to ensure statistical significance.

Figure 4: Average runtime for 50 trials per expression; Bandicoot matches or exceeds competitors for nearly all expressions due to compile-time fusion and minimal overhead.

Key results demonstrated include:

• Bandicoot achieves the lowest or equal-lowest runtime in most test expressions.
• Absence of runtime/JIT overhead is an explicit advantage, as metrics are evaluated in kernel execution only, excluding warmup/JIT compilation (which Bandicoot does not need beyond initial GPU kernel driver initialization).
• ArrayFire, despite being C++ and high-performance, still incurs overhead due to its internal JIT compilation instead of leveraging compile-time fusions.

Bandwidth-bound operations are further examined by scaling the number of source matrices being added together—an archetypal memory-limited scenario. Only Bandicoot maintains high throughput for arbitrary numbers of matrices. Competing frameworks degrade sharply beyond low operand counts due to multiple kernel launches and intermediate temporaries.

Figure 5: Matrix addition throughput (GB/s) versus number of matrices, showing Bandicoot alone approaches hardware memory bandwidth and scales to large numbers of operands without throughput collapse.

Notably, Bandicoot is able to saturate the observed hardware memory bandwidth ceiling (measured at approximately 868 GB/s for the target GPU), maintaining peak throughput even as the complexity of the operation increases.

Theoretical and Practical Implications

Bandicoot's template metaprogramming-based expression fusion strategy achieves several impactful advances:

• Static expression fusion eliminates the need for runtime expression parsing, graph tracing, or JIT infrastructure, yielding both smaller binaries (only kernels in use are compiled) and reduced runtime memory/compute overhead.
• Vendor/backend independence is assured by modular backend design, with current CUDA and OpenCL support and plans for expansion to Vulkan, HIP/ROCm, Apple Metal, and others.
• Immediate transposability for existing Armadillo-based codes lowers adoption barriers across scientific and ML workloads.
• Sustained peak memory-bound performance for arbitrary complex expressions has direct implications for high-throughput scientific computing, large-scale ML workloads, and real-time AI inference serving.
• The compile-time approach sidesteps constraints inherent to languages lacking sufficient metaprogramming sophistication or to projects requiring runtime dynamism (e.g., user-input program generation or runtime symbolic computation), but is generally superior for the static, strongly-typed scientific computing regimes for which C++ excels.

Future Perspectives and AI Relevance

Expected extensions include full backend parity with all major vendor stacks (including Apple Metal and AMD ROCm) and support for low-precision formats (e.g., bfloat16, fp8, quantized types) critical to efficient deep neural network training and inference at scale. Expanding the operator/factorization set to reach full Armadillo parity (potentially leveraging libraries like MAGMA) will further amplify Bandicoot's reach into broader computational fields.

From the perspective of AI systems, the ability to deploy custom, bandwidth-saturating kernels for arbitrary model architectures and linear algebraic backbones is particularly salient for large-scale model inference and training, where kernel launch overheads and memory-bound subexpressions often dominate resource utilization. The static fusion paradigm is especially valuable in settings—such as edge inference, data center deployment, or tightly resource-constrained AI systems—where volatility and unpredictability of runtime environments impede the efficacy of JIT or traced compilation models.

Conclusion

The Bandicoot library implements an avant-garde, compile-time kernel fusion approach for GPU linear algebra, leveraging C++ template metaprogramming to statically generate hardware-optimized, fused kernels for arbitrary expressions. Empirical evaluation substantiates that Bandicoot either matches or surpasses existing frameworks in execution speed and, uniquely, maintains saturating memory throughput for compound matrix expressions of arbitrary size. The methodologically distinctive compile-time fusion technique demonstrated here elucidates a compelling direction for high-performance, portable, and user-friendly linear algebra in scientific and AI workloads, with immediate impact and clear growth trajectories as hardware and compiler capabilities advance.