MetaTensor: Tensor Fusion in a Meta-Tracing Compiler

Abstract

Python tensor compilers capture a graph with a separate frontend, and a fused region ends wherever that frontend stops. Branches, exceptions and in-place writes in the host program therefore split fusion. We show that the JIT compiler of a dynamic language already provides what such a frontend reimplements, namely tracing, guards, specialization and deoptimization. MetaTensor adds tensor fusion to PyPy’s meta-tracing JIT compiler as one pass in its trace optimizer, with TensorPyPy as its tensor library. The pass keeps each recorded tensor operation as a virtual, an object the compiler never allocates, so the operand links of virtuals form the computation graph inside the optimizer. Forcing a virtual emits one fused Triton kernel, and the rest of the trace becomes host machine code that launches it. A guard saves live tensor virtuals in its deoptimization state instead of forcing them, and two decoders rebuild them when the guard fails, so fused regions span host control flow. On the 72 of 161 models of the PyTorch 2 benchmark suites that MetaTensor runs from upstream source, it is 1.12× faster than TorchInductor, 1.09× faster than the better of TorchInductor with and without CUDA graphs, and 1.04× faster than JAX/XLA in geometric mean. The gain comes from elementwise work fused into convolution kernels, a fused attention kernel and a shorter host path. On seven TorchBench models, ending fused regions at every guard costs 2.47×, and building the same graph from runtime objects costs 1.55×, or 1.49× and 1.11× without the one elementwise-only model.