LazyTensors: Dynamic Tensor Execution

• LazyTensors are a dynamic computation model that intercepts tensor operations at runtime, buffering them into a private computation graph without altering the user’s code semantics.
• They optimize performance by lowering the accumulated graph to XLA HLO IR, enabling global optimizations, fusion, and hardware-targeted code generation for significant speedups.
• The architecture integrates custom tensor types, a lowering library, and a runtime scheduler, supporting multiple host languages and devices such as CPUs, GPUs, and TPUs.

LazyTensors are a technique for combining the ergonomic flexibility of eager (define-by-run) tensor computation with the optimization capabilities of domain-specific compilers (DSCs) such as XLA. The core mechanism involves intercepting all tensor operations at runtime, recording them in an internal intermediate representation (IR) graph, and delaying (“lazily”) their materialization until a synchronization or observation barrier is encountered. This approach allows programmers to exploit the full expressiveness of host languages like Python and Swift, while enabling global graph-level optimizations, fusion, and hardware-targeted code generation typically restricted to static or subset languages (Suhan et al., 2021).

1. Motivation and Conceptual Foundations

Eager execution frameworks (e.g., PyTorch, NumPy) provide tight integration with host languages and dynamic model construction, but lack a global computation graph for whole-program optimization. Traditional approaches—such as autograd-style tracing, source-to-source translators (AutoGraph), or partial compiling (TorchScript)—necessitate a “language subset,” restricting the use of features like recursion, exceptions, and nontrivial control flow within compiled regions. This imposes “function coloring” throughout the codebase and burdens users with distinguishing between compilable and non-compilable code paths (Suhan et al., 2021).

LazyTensors eliminate the language subset problem by buffering all tensor-level operations into a private computation graph (a DAG), without altering semantics observable from the user’s perspective. User code continues to execute eagerly—for each operation, an object with the full Tensor API is returned—but underneath, the system accumulates a graph of tensor operations. Upon reaching a barrier (e.g., when a non-tensor result is observed or a synchronization API is called), the graph is lowered to XLA HLO IR, compiled, and executed, returning concrete tensor results.

2. Architectural Components

The LazyTensor system architecture is composed of three principal components:

A. Custom Tensor Types:

A host-language class (e.g., XlaTensor in PyTorch, LazyTensor in Swift) that mirrors the public API of the baseline eager Tensor class. All tensor methods (add, matmul, view, in-place updates, etc.) are overloaded to intercept and record operations.

B. Lowering Library:

A library, implemented in C++, maps each recorded high-level tensor op to XLA HLO instructions or an equivalent IR. This lowering step directly enables optimizations and accelerator-specific code generation.

C. Runtime Scheduler and Cache:

A runtime maintains a growing DAG of recorded tensor ops. On encountering a barrier, it finalizes and canonicalizes the current graph, checks for a cached compiled module, and if necessary invokes the XLA compiler (shape inference, fusion, scheduling, codegen). Device contexts are managed such that tensors are associated with the correct computational resource; input handles are mapped as parameters to the compiled executable. The system then retrieves results as full tensors and resets internal state for future graph accumulation.

These core mechanisms are shared across multiple host languages and frameworks. For example, more than 75% of the C++ runtime codebase is shared by PyTorch on Cloud TPU and Swift for TensorFlow, supporting CPUs, GPUs, and TPUs (Suhan et al., 2021).

3. Formal Model and Graph Representation

A LazyTensor computation forms a directed acyclic graph $G = (V, E)$, where each vertex $v \in V$ is annotated with:

• $op(v) \in \mathrm{Op}$: the tensor operation (e.g., aten::add, aten::conv2d)
• $shape(v) \in \mathbb{N}^d$: the static output shape
• $dev(v) \in \mathrm{Devices}$: device tag
• $const(v) \in \mathbb{R}$ (optionally): scalar constant

Edges $(u \to v)$ represent input dependencies between nodes. Roots of $G$ correspond to tensor outputs expected by the user. The system performs shape inference for each node by propagating up from inputs, applies a fusion pass on maximal subsets of element-wise operations, partitions the graph by device, then lowers the partitioned graphs to XLA HLO modules before feeding them into the backend compiler and runtime.

Execution semantics ensure that the compiled program behaves equivalently to the step-wise eager trace for any input tensors, meeting functional transparency requirements. This enables whole-graph reconstruction, optimization, and execution without deviating from the user’s programming abstractions (Suhan et al., 2021).

4. Host-Language Integration and API

PyTorch

The XlaTensor (Python subclass of torch.Tensor) overloads all tensor methods and intercepts operations to record them as LazyNodes. A C++ backend builds the IR, communicates with the XLA client, and manages device context arenas for tensor liveness. Barriers (such as mark_step()) must be invoked to flush and execute the current computation graph.

Swift for TensorFlow

A struct named LazyTensor, conforming to TensorProtocol, implements the appropriate operator overloads and captures user interactions in the same shared C++ backend. Barriers are handled via LazyTensorBarrier().

Both implementations remain backend-agnostic: While XLA is the reference, the approach is compatible with other compilers (Glow, TVM, LLVM-IR), provided a suitable lowering is implemented. Device management ensures that graph compilation and execution are tailored per device resource. The system’s language-agnostic design is further substantiated by the sharing of ~85% of the Swift for TensorFlow LazyTensor codebase with the Python front-end (Suhan et al., 2021).

5. Addressing the Language Subset Problem

Traditional tracing or scripted-front-ends force explicit demarcation of compilable functions, fragmenting the user codebase and precluding seamless interleaving of arbitrary host-language logic. LazyTensor’s design principle dictates that “the only gateway between host language and DSC is the Tensor op interface.” All host-language constructs—such as exceptions, loops, lists, dynamic dispatch, and recursion—are permitted without constraint.

Whenever a program step produces a non-tensor result (Boolean for an if-statement, string for I/O, Python or Swift integer, etc.), LazyTensor detects this transition and forcibly materializes the computation graph via a blocking barrier. As a result, no special annotations are required, and all tensor-compatible operations are automatically captured and compiled as appropriate. This resolves the language subset problem without compromising usability or expressivity (Suhan et al., 2021).

6. Performance Characteristics and Benchmarks

Code Reuse

The architecture yields significant code reuse: Swift for TensorFlow’s LazyTensor codebase comprises approximately 66 K SLoC, with around 85% shared between Swift and Python implementations, confirming language-independence and reusability.

Empirical Performance

• HuggingFace RoBERTa-base (Cloud TPU v3-8, 4 chips) vs Native PyTorch (4×V100 GPU):
  • Native PyTorch (4×V100): 48 batch, 3-epoch WikiText-103, eval-perplexity 3.14, 133.4 min
  • PyTorch + LazyTensor (4×TPUv3): 128 batch, 3.25 perplexity, 38.3 min (∼3.5× speedup)
  • Speedup arises from larger batch sizes, higher arithmetic intensity on TPU, fused global graphs, and reduced host-device communications.
• ResNet-50 on ImageNet (Swift for TensorFlow with TPUv3 pods):

| Cores | 90-epoch Time | Throughput (examples/sec) | Per-core ex/sec | |-------|---------------|---------------------------|-----------------| | 16 | 189 min | 10,164 | 635 | | 32 | 96 min | 20,015 | 625 | | 128 | 25 min | 77,726 | 607 |

Near-linear scaling is maintained as cores increase, with per-chip throughput remaining within 5% across an 8× increase.

• Small Model (WordSeg) on GPU:

| Operation | Eager (ms) | Lazy (ms) | Ratio | |---------------------|------------|-----------|-------| | Score 4 | 96 | 65 | 0.68 | | Score 8 | 182 | 118 | 0.65 | | Score 14 | 282 | 165 | 0.59 | | ScoreGradient 4 | 425 | 281 | 0.66 |

Even for models fitting within GPU cache, 30–40% reductions in latency are observed (Suhan et al., 2021).

7. Limitations and Prospects for Future Development

Compile Overhead: JIT compilation overhead can be substantial for small or highly interactive workloads. While the cost is incurred once per unique graph, dynamic-shaped models (e.g., Mask-RCNN on COCO) that alter shapes frequently experience recurrent recompilation latency.

Shape Constraints: XLA backend requires statically known shapes for compilation. Models with dynamic shape requirements may need to fall back to eager execution or pay the compilation cost for each new shape instance.

Control-Flow Limitations: The LazyTensor IR does not natively support control-flow (loops, if-statements) over tensor conditions; such constructs are unrolled at record time. Future work may extend the IR and lowering passes to support XLA’s while/cond operators via deeper integration.

Manual Barriers: Presently, API users must invoke explicit synchronization barriers (mark_step/LazyTensorBarrier). Automatic inference of execution boundaries or pipeline segmentation for long-running code remains open for future enhancements.

IR Memory Usage: The accumulation of large or complex computation DAGs can induce in-memory overheads. Strategies for IR compression, arena-based allocation, and JIT-style IR compaction are identified areas of improvement.

Graph Reuse and Mixed Precision: Users could further expedite workloads by providing “graph shape reuse” hints across iterations. While mixed-precision support exists via environment toggles, more granular profile-guided specialization or device-specific layout optimization is a potential target (Suhan et al., 2021).

A plausible implication is that LazyTensors provide a pathway for interoperable, highly optimized tensor programming across an expanding array of host languages and hardware targets, supporting the increasing demands of modern ML and scientific computing workloads.

---

For comparison, template-based C++ libraries such as TLoops also enable tensor algebra to be written in a style close to analytic notation while emitting high-performance CUDA or C kernels. Unlike LazyTensor’s runtime graph capture, TLoops leverages compile-time metaprogramming to generate explicit kernel code via expression templates, achieving near hand-tuned performance on CPUs and GPUs (Lewis et al., 2018). However, TLoops lacks dynamic shape generality, restricts operations to a fixed set of templates, and is less agnostic with respect to high-level language features or dynamic model architectures.

Both approaches exemplify the tradeoffs between static and dynamic tensor computation models, and highlight the ongoing evolution of software abstractions for high-performance, hardware-agnostic tensor execution.