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

1 Introduction

• Dynamic production workloads vary tensor shapes, while GPU kernels perform best when compiled with full knowledge of their inputs.Static compilation must recompile when shapes change; JIT compilation tolerates dynamic shapes but adds compilation latency on the inference critical path.Source: §1 · p. 1

Contributions

• For statically known ranks, DLCB solves operator-derived shape constraints to separate hard-coded dimensions from runtime parameters and checks residual constraints before kernel launches.The generated host code also computes concrete dimension values needed as kernel launch arguments.Source: §1 · p. 2, Contributions · p. 2
• DLCB selects static kernels for fully known shapes, parameterized AOT kernels for known ranks with dynamic sizes, and cached JIT compilation when ranks are unknown.The paths share a unified interface; AOT kernels use launch-time dimension values and host-side checks.Source: §1 · p. 2

1 Introduction

• The PyTorch bridge partitions programs into supported subgraphs for AOT compilation and leaves unsupported subgraphs to run in PyTorch eager execution.Compiled binaries are spliced back into the broader execution graph.Source: §1 · p. 2

Contributions

• An automatic shape-generalization pass identifies dimensions that can become free variables, allowing compiled binaries to be reused across shapes without manual symbolic-dimension annotations.The pass operates on statically traced PyTorch programs.Source: §1 · p. 2, Contributions · p. 2

2.1 PyTorch-DLCB bridge

• The bridge partitions supported operators into compilable regions, lowers them to DLCB IR, and leaves unsupported operators in the PyTorch graph.For traced programs, shape generalization identifies dimensions safe to abstract; traced size-one broadcast dimensions can receive specialized paths guarded by runtime checks.Source: §2.1.1 · p. 3, §2.1.2 · p. 3, §2.1.3 · p. 3

2.2 Graph compilation

• Graph compilation incrementally adds operator-derived shape constraints, enabling contradictions such as incompatible operand shapes to be detected during graph construction.Optimization includes stride assignment, fusion, and pattern rewrites; the graph is then partitioned into vendor-library or custom-kernel thunks.Source: §2.2 · p. 4, §2.2.1 · p. 4, §2.2.2 · p. 4
• Coordinate-based inference maps each output coordinate to the required inputs, and kgen emits CUDA kernels that compute outputs by iterating over output coordinates.Dynamic dimensions left free by constraint solving become integer kernel parameters.Source: §2.2.3 · p. 4
• The host code generator turns residual expressions into an interpreted CPU program that resolves dimensions, verifies constraints, and supplies launch arguments before each kernel launch.At runtime, the interpreter fills output slots with resolved shapes, strides, and constants for the kernel launcher.Source: §2.2.4 · pp. 4–5

2.3 Runtime kernel launch

• Static thunks specialize on concrete shapes, dynamic AOT thunks reuse a binary for inputs with matching rank and element type, and JIT thunks compile on demand when rank is unknown.Dynamic AOT invocation runs the host program to resolve dimensions before memory allocation and kernel dispatch; JIT kernels are cached for subsequent matching shapes.Source: §2.3 · p. 6

3 Constraint Solving

• The fixed-point solver operates on equality constraints between dimension expressions represented as nodes in a DAG, and is also instantiated for coordinate-set constraints during kernel generation.Expression nodes track free variables and use content-addressed hashes for structural equality.Source: §3 · p. 6
• The solver repeatedly applies prioritized rewrites—including substitution, list unpacking, arithmetic inversion, and common-subexpression cancellation—until no rule applies.Contradictory equalities between integer literals cause failure.Source: §3 · pp. 6–7
• The rule set does not solve arbitrary equation systems, but the authors report it suffices for constraints generated by the supported operators’ regular patterns.They argue that a general-purpose solver would add unnecessary overhead and complexity for this bounded constraint vocabulary.Source: §3 · p. 7
• At a fixed point, dimensions are either compile-time integer literals, runtime named literals, or residual expressions evaluated by the host program at launch.Named literals become kernel parameters; residual constraints are compiled into interpreted host code.Source: §3 · p. 7

4 Dynamic Kernel Generation

• kgen derives input coordinates for arbitrary operator compositions through backward coordinate inference, then generates CUDA kernels that load required inputs and write outputs.This avoids implementing index arithmetic separately for each operator combination.Source: §4 · p. 7

4.1 Coordinate Representation

• Coordinates represent logical tensor positions through coordinate sets and arithmetic expressions, with bounded coordinate spaces supporting scalar or vector positions.Expression nodes are shared when structurally identical, so sub-expressions are computed only once in the generated kernel.Source: §4.1 · pp. 7–8

4.2 Output Annotation and the Inference Problem

• kgen assigns each output tensor a free-variable coordinate set, then determines the input coordinates needed for each output in topological order.The resulting coordinate map can be evaluated independently by each thread.Source: §4.2 · p. 8

4.3 Coordinate Inference

• Coordinate inference propagates output-to-input coordinate transformations backward until rewrite rules resolve every graph node’s coordinate expression..Source: §4.3 · p. 8
• Operator-specific reverse transforms handle identity and broadcast indexing, reshape arithmetic, reduction loops, and gather’s index-dependent data access.For broadcast dimensions of size 1, the input coordinate maps to 0; reductions iterate across the reduced axis.Source: §4.3 · p. 8

4.4 CUDA C++ Code Generation

• The completed coordinate map is emitted as CUDA C++, compiled to a cubin at graph-compile time, and launched without recompilation at inference.The host program supplies unresolved dynamic dimensions; generated code evaluates coordinates, loads inputs, and writes each thread’s output.Source: §4.4 · p. 8

4.5 Dynamic Vectorization

• Dynamic vectorization selects a type-compatible width at compile time, then uses a runtime guard to choose vectorized execution or scalar fallback.The vectorized branch requires every input and output’s last dimension to be divisible by the width and contiguous; grouped coordinates become vector loads and stores.Source: §4.5 · p. 9

5.1 End-to-end inference throughput

• Dynamic shapes prevent constant-folding and common-subexpression elimination of shape- and stride-dependent expressions, leaving expensive integer division or modulo in address computations.For the same subgraph, dynamic kernels are slower than static counterparts.Source: §5.1 · p. 9
• A single runtime vector-length guard often misses TorchVision’s differing optimal widths; checking widths 1, 2, 4, and 8 adds overhead.Vectorization typically improves single-kernel performance by 20–25%.Source: §5.1 · p. 9
• Unknown compile-time dimensions prevent dynamic mode from identifying NCHW versus NHWC tensors and applying layout optimization.NHWC is generally faster for convolution kernels because it better leverages tensor cores.Source: §5.1 · p. 9
• If only a few dimensions are dynamic, future work could constant-fold more expressions and derive NCHW versus NHWC layout from shape variables.Source: §5.1 · p. 9
• 1.10x static and 0.92x dynamic speedup relative to eager; dynamic mode is slower than eager on average.Dynamic mode still beats eager on some models.Source: §5.1 · p. 9, Figure 2 · p. 10

5.2 Compilation cost

• 0.84 s median compile time versus 9.4 s for torch.compile—roughly a 10× reduction—and DLCB-dynamic compiles once without recompiling for new shapes.The dynamic compilation cost is incurred ahead of time at model load.Source: Figure 3 · p. 10, §5.2 · p. 11

Static-shape compilers

• Production-oriented static compilers such as XLA, TVM, Halide, Triton, Glow [16], and Relay [15] require concrete dimensions or shape-specific tuning.MLIR supports symbolic dimensions, while DLCB describes its constraint language as a lightweight CUDA-targeted analogue.Source: §6 · p. 11

Dynamic-shape systems

• DISC [24] handles dynamic shapes with a CPU shape-inference subgraph, whereas DLCB classifies dimensions at compile time and passes free dimensions as kernel parameters.Nimble dispatches to per-shape kernels with JIT fallback; DLCB’s AOT path avoids runtime recompilation for known rank.Source: §6 · p. 11

JIT compilation and PyTorch

• PyTorch 2 is DLCB’s state-of-the-art reference, while nvFuser generates fused CUDA kernels at runtime and recompiles for each distinct shape.Ansor and AutoTVM also require schedule re-search per shape; Proteus respecializes kernels via runtime JIT.Source: §6 · p. 11

Parameterized kernels

• PLUTO fixes tile sizes at compile time, while PPCG passes problem size N at runtime; DLCB generalizes runtime parameters to dynamic tensor dimensions.Tiramisu also parameterizes schedules by runtime problem sizes.Source: §6 · p. 11

• DLCB partitions dimensions into compile-time constants and runtime parameters, verifies residual shape constraints before launch, and falls back to JIT when rank is unknown.Source: §7 · p. 11
• Automatic shape generalization identifies safe dimensions to abstract in traced PyTorch programs, enabling parameterized AOT kernels for production inputs with varying shapes.Source: §7 · pp. 11–12
• DLCB’s dynamic-shape kernels approach fully-static compilation performance, while one dynamic binary avoids repeated compilation as deployment shapes vary.Source: §7 · p. 12