arXiv subjects

Alexander Collins

Publications and source records attributed to Alexander Collins.

4 recordsLinked to original sources

arXiv2026

DLCB: Ahead-of-Time Compilation for Dynamic Deep Learning

Deep learning workloads are increasingly deployed in settings where tensor shapes are not fully known at compile time. Batch sizes vary across requests, sequence lengths differ between inputs, and model architectures admit a range of spatial resolutions. This dynamism creates a fundamental tension: ahead-of-time compiled GPU kernels deliver peak performance but traditionally require fully static tensor shapes, while just-in-time compilation supports dynamic shapes at the cost of runtime compilation overhead. We present DLCB (Deep Learning Compiler Backend), a deep learning compiler that resolves this tension through a unified three-tier compilation strategy. For programs with fully static tensor shapes, DLCB generates static CUDA C++ kernels ahead of time. For programs whose tensor rank is statically known but whose dimension sizes are dynamic, DLCB generates AOT kernels parameterized by those dimensions, compiling once and executing across a range of shapes. For programs with tensors of unknown rank, DLCB falls back to JIT compilation at runtime. For AOT compilation with dynamic shape tensors, our approach builds a system of shape constraints from the input program, and uses a solver to reduce these to either a) fixed constraints that produce hard-coded constants in the generated kernel; b) unresolved constraints which become kernel parameters passed at launch time, along with interpreted host code to check that the constraint holds for a given set of inputs and to compute the values to pass. Our compiler and runtime are embedded in PyTorch, and use an automatic shape generalization pass to enable AOT compilation with dynamic shapes for supported subsets of PyTorch programs

cs.PL↗

arXiv2025

Tawa: Automatic Warp Specialization for Modern GPUs with Asynchronous References

Modern GPUs feature specialized hardware units that enable high-performance, asynchronous dataflow execution. However, the conventional SIMT programming model is fundamentally misaligned with this task-parallel hardware, creating a significant programmability gap. While hardware-level warp specialization is the key to unlocking peak performance, it forces developers to manually orchestrate complex, low-level communication and software pipelines--a process that is labor-intensive, error-prone, and unsustainable. To address this challenge, we present Tawa, an automated compiler that systematically generates high-performance, warp-specialized code from a high-level, tile-based program. Central to our approach is a novel IR abstraction, asynchronous references (aref), which expresses warp-level communication without exposing low-level hardware details. Using this abstraction, Tawa automatically partitions programs into producer-consumer roles and manages the intricate dataflow pipeline, relieving developers of invasive kernel rewriting. Evaluation on NVIDIA H100 GPUs across representative LLM kernels shows that Tawa delivers high hardware utilization, achieving up to 1.1$\times$ speedup over highly optimized cuBLAS GEMM kernels. For attention workloads, Tawa attains 1.2$\times$ speedup over Triton and matches the performance of the hand-optimized CUTLASS C++ FlashAttention-3 kernel with far less programming effort.

cs.LG↗

arXiv2022

Axon: A Language for Dynamic Shapes in Deep Learning Graphs

Axon is a language that enables shape and rank inference for tensors in a Deep Learning graphs. It aims to make shapes implicit and inferred, in a similar manner to how types are implicit and inferred in many functional programming languages. Tensor dimensions are represented by expressions consisting of symbolic variables, constants, and arithmetic operators. Tensor shapes can be expressed as either a sequence of these dimension expressions, as a symbolic variable, or as an appending of other shapes. This allows complex constraints on shapes to be expressed. Axon is functional in style, with a type system similar in to Standard ML, extended to include shape information. It provides a suite of built in operators over tensors, including pointwise arithmetic operators, maps, reduction, loops and user defined functions. We describe a shape inference algorithm based on constraint solving which infers information about shapes, from both shape information provided by the programmer and the structure of the program. This allows fully automatic inference of the shapes of tensors for complex Deep Learning graphs. This approach reduces programmer effort when specifying graphs, as tensor shapes are not explicit, allows composition of Deep Learning graphs while maintaining input and output tensor shape compatibility, and aids in automated error detection by identifying shape mismatches at runtime.

cs.PL↗

arXiv2020

Probabilistic Programming with CuPPL

Probabilistic Programming Languages (PPLs) are a powerful tool in machine learning, allowing highly expressive generative models to be expressed succinctly. They couple complex inference algorithms, implemented by the language, with an expressive modelling language that allows a user to implement any computable function as the generative model. Such languages are usually implemented on top of existing high level programming languages and do not make use of hardware accelerators. PPLs that do make use of accelerators exist, but restrict the expressivity of the language in order to do so. In this paper, we present a language and toolchain that generates highly efficient code for both CPUs and GPUs. The language is functional in style, and the tool chain is built on top of LLVM. Our implementation uses de-limited continuations on CPU to perform inference, and custom CUDA codes on GPU. We obtain significant speed ups across a suite of PPL workloads, compared to other state of the art approaches on CPU. Furthermore, our compiler can also generate efficient code that runs on CUDA GPUs.

cs.PL↗