Most deep learning systems ship as fixed inference engines: you bring a model, the engine runs it, and the optimizations are whatever the vendor froze into the release. Apache TVM is the opposite bet. Its README calls it an Open Machine Learning Compiler Framework built Python-first, so that quick customization of the compilation itself is the whole point, and universal deployment is the second principle, bringing models into minimum deployable modules. The lineage reads like a history of the field: part of the TIR and arithmetic simplification modules originate from Halide, the integer set analysis and loop transformation primitives come from Loopy, and the symbolic scan operator design was inspired by Theano. The current design, after several rounds of redesign, is a cross-level one, with TensorIR as the tensor-level representation and Relax as the graph-level representation, jointly optimizing computational graphs, tensor programs, and libraries. As always in this series, what follows is an educational tour of published source code.

The tree itself is living evidence of another redesign happening right now. On the main branch the old TIR namespace has been renamed to s_tir, so the public C++ headers split into include/tvm/relax, include/tvm/s_tir, include/tvm/te, include/tvm/topi, and include/tvm/runtime. Alongside them sits a self-contained sub-project called tirx, with its own CMakeLists.txt, its own include and src trees, and its own Python package tvm_tirx, which the Python code in the main package already imports side by side with the legacy namespaces. Reading this repository means reading both a mature compiler and the blueprint of its successor at the same time, which is a rare treat.

Overview of Apache TVM: frontends emit Relax programs, fifty-six C++ passes rewrite them, a dispatch stage hands subgraphs to vendor libraries, subgraphs lower into TensorIR where the schedule language, MetaSchedule, and dlight rules transform loops, and the Relax VM with Disco executes the result, all while the tirx sub-project reimagines the tensor layer.

Reading the overview from left to right:

• fx_translator.py is the PyTorch door at 1296 lines, with TorchFXImporter at line 31 and from_fx at 1184, flanked by a 2458-line exported_program_translator.py and a 246-line dynamo.py in the same folder.
• expr.h defines the graph-level objects in 552 lines: DataflowVarNode at 73, DataflowBlockNode at 240, and SeqExprNode at 265, with the 324-line block_builder.h declaring the C++ builder.
• fuse_ops.cc is the flagship of the 56 C++ pass files, 1588 lines with GraphCreator at 115, FunctionCreator at 403, OperatorFusor at 809, PatternBasedPartitioner at 1178, and CompositeFunctionAnnotator at 1366.
• cutlass.py is the dispatch stage that carves subgraphs out for external kernels, with matmul_patterns at 183, attention_patterns at 410, WorkspaceAnnotator at 530, and partition_for_cutlass at 574.
• stmt.h is the TensorIR statement vocabulary in 412 lines: MatchTensorRegionNode at 49, the schedulable SBlockNode at 112, and SBlockRealizeNode at 188.
• schedule.py is the 4042-line Python scheduling surface, with copy at 244, fuse at 678, split at 741, reorder at 949, compute_inline at 2259, decompose_reduction at 2452, and transform_layout at 3351.
• gemv.py is one of the dlight rules for LLM decode shapes, 693 lines with the GEMV class at 39 and its sch_inner_reduction, sch_outer_reduction, and fallback strategies at 90, 433, and 638.
• executable.cc implements the deployable artifact, 596 lines whose kind at line 50 returns relax.VMExecutable and whose SaveCodeSection at 314 and LoadCodeSection at 402 define the serialization format.

Why You Need This

The first reason is leverage. When a model misbehaves on your hardware, the usual options are to file an issue or to fork the engine. In TVM the graph-level optimizations, the tensor-level schedules, and even the code generation targets are ordinary code you can read and extend from Python, because the project states Python-first development as a principle and enforces it structurally: the passes are driven from python/tvm/relax/transform, the schedule primitives live in a 4042-line Python class, and the frontends accept PyTorch FX, exported programs, ONNX, TFLite, and StableHLO. When you want a transformation that nobody ships, like folding a batch norm into a preceding convolution for inference, the repository already shows the move in fold_batch_norm_to_conv2d_for_inference.py, and fast_math.py, fuse_transpose_matmul.py, and remove_redundant_reshape.py sit in the same folder as worked examples of the pattern.

The second reason is that scheduling is a first-class, learnable language rather than buried heuristics. A TensorIR block, the SBlockNode in stmt.h, is a loop program annotated with what it reads, writes, and reduces, and the Schedule object exposes verbs that rewrite those loops while preserving the contract. Split, fuse, reorder, compute_inline, decompose_reduction, and transform_layout are the nouns of a grammar you can apply by hand in a REPL and watch the IR print out after each step. MetaSchedule then automates the search across 107 files of schedule_rule, search_strategy, task_scheduler, and cost_model machinery, with tune_tir at line 45 of tir_integration.py as the entry point. And for the LLM era, the dlight package adds hand-written rules tuned for known shapes: gpu/matmul.py at 782 lines classifies blocks by detecting iteration traits at 172 and builds index maps at 242, gemv.py handles the memory-bound decode case, and rmsnorm.py at 144 lines pattern-matches the rsqrt block at 60. These are the kernels that make small-model serving fast, and you can read every line of them.

The third reason is the deployment story. Compilation ends in a Relax VM executable whose format is small enough to understand from executable.cc alone: sections for globals, memory scopes, constants, and the instruction stream, with bytecode.h at 246 lines defining the opcodes. The 284-line vm_build.py ties the whole flow together with _vmcodegen at 51, _vmlink at 97, and build at 164. For multi-device serving, the Disco runtime in include/tvm/runtime/disco, with session.h at 389 lines, shards VM programs across workers, and the nn frontend’s llm subfolder ships KV cache and paged-attention plumbing in kv_cache.py, tree_attn.py, _page_kernels.py, _prefill_kernels.py, and _decode_kernels.py. This is why TVM keeps being the foundation for vertical LLM compilers: the pieces from graph ingest to distributed execution are all in one Apache-2.0 tree.

Detail view: torch, ONNX, and the native nn API feed Relax IR and its BlockBuilder, C++ and Python pass folders rewrite and fuse, the CUDA dispatch stage annotates subgraphs for cuBLAS, cuDNN, CUTLASS, and FlashInfer, TensorIR headers and their schedule machinery transform loops under MetaSchedule and dlight, and the VM, Disco, and tirx layers close the path to deployment.

The graph level is worth reading first because it is small. expr.h gives Relax first-class dataflow: a DataflowBlockNode at line 240 marks the region where intermediate values live and die inside a function, DataflowVarNode at 73 names those locals, and SeqExprNode at 265 binds dataflow blocks to the returns that escape them. The BlockBuilder class at line 108 of block_builder.py, 818 lines in total, is the Python-side constructor for all of this, and its emit methods record each call into the open dataflow block as you write it. Passes then rewrite this graph: among the 56 C++ files in src/relax/transform you will find kill_after_last_use.cc, static_plan_block_memory.cc, rewrite_cuda_graph.cc, and combine_parallel_matmul.cc, while fuse_ops.cc alone stacks five cooperating classes, from GraphCreator walking the expression at 115 to the PatternBasedPartitioner at 1178 that carves out matched subgraphs and the CompositeFunctionAnnotator at 1366 that labels them for the dispatch stage.

Ingestion and dispatch complete the front of the compiler. The torch folder offers three ways in: raw FX graphs through TorchFXImporter at line 31 of fx_translator.py, torch.export programs through the 2458-line exported_program_translator.py, and a dynamo backend in dynamo.py. The nn package is the native path, a PyTorch-like API in 879 lines of core.py where Module sits at 354 and export_tvm at 460 turns a composed module into an IRModule and parameter list, with Linear at 97, RMSNorm at 629, and Embedding at 916 in the 1199-line modules.py. When a subgraph should run as a vendor kernel instead of compiled loops, the backend/cuda folder annotates it: cutlass.py declares the fusion patterns with matmul_patterns at 183 and attention_patterns at 410, annotates workspaces through WorkspaceAnnotator at 530, and hands the partitioned module to codegen via partition_for_cutlass at 574. Siblings route to cuBLAS in cublas.py at 246 lines and FlashInfer at 345, and the same dispatch idea extends to ROCm, Metal, and generic GPU and CPU targets in neighboring backend folders. It is the same cutlass templates we toured earlier in this series, now consumed from a compiler that decides when to use them.

The tensor level shows what the s_tir rename actually bought. stmt.h keeps the schedulable loop program deliberately tiny: an SBlockNode at 112 declares its iter vars, its reads and writes as MatchTensorRegionNode at 49, and an SBlockRealizeNode at 188 instantiates it under a loop nest. The Python Schedule at 4042 lines is the user surface, and the C++ machinery that executes each primitive lives in src/s_tir/schedule, where concrete_schedule.cc pairs with instruction.cc and state.cc to record schedule steps as replayable instructions, a design that makes tuning results portable. MetaSchedule wraps that recording ability in search machinery, and compile_tir at line 249 of tir_integration.py shows the handoff from tuned schedule to lowered code across the 39 transform files in src/s_tir/transform. One detail in dlight/gpu/matmul.py tells the migration story better than any changelog: its import block between lines 24 and 33 pulls from tvm.s_tir and tvm_tirx simultaneously, so the legacy namespace and the successor package are literally coexisting inside one rule file.

The runtime and the next generation close the tour. executable.cc serializes a compiled module into sections, with Stats at line 64 for introspection and WriteToFile at 205 for shipping, and the VM interpreter in include/tvm/runtime/vm, headed by vm.h at 218 lines, executes the bytecode. Disco extends one device to many with session.h at 389 lines and its worker protocol. Then there is tirx, the sub-project that is easy to mistake for a stray directory: it has its own CMakeLists.txt, its own pyproject.toml pinning mlc-z3-static at 4.16.0 or newer for z3-backed symbolic reasoning, its own tvm_tirx Python package with a backend/cuda/codegen tree, and its own src split into backend, codegen, ir, runtime, script, support, sym, and target. The AGENTS.md file instructs builders to configure it with TVM_TIRX_USE_LLVM and TVM_TIRX_USE_CUDA flags and to develop with PYTHONPATH rather than editable installs, and it expects commit tags like [REFACTOR][IR]. The legacy layers are all still here too, te and topi under include/tvm, and the old tvmc command-line tool has been removed from the tree, python/tvm/driver now holding only build_module.py.

From Install to a Compiled Model

Build from source with LLVM and CUDA enabled, then work from the Python tree as the repository’s own agent guidance recommends:

git clone https://github.com/apache/tvm.git
cd tvm && git submodule update --init --recursive
mkdir build
cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=RelWithDebInfo \
  -DTVM_TIRX_USE_LLVM=ON -DTVM_TIRX_USE_CUDA=ON
cmake --build build --parallel
export PYTHONPATH="$(pwd)/python:$(pwd)/.local/python"

The nn frontend lets you compose a model in PyTorch style and export it straight to Relax, then compile to a VM executable:

from tvm.relax.frontend.nn import Module, Tensor, Linear, ReLU, spec
import tvm
from tvm import relax

class Tiny(Module):
    def __init__(self):
        super().__init__()
        self.dense = Linear(16, 16)
        self.act = ReLU()

    def forward(self, x: Tensor) -> Tensor:
        return self.act(self.dense(x))

# export_tvm returns an IRModule plus a parameter list,
# exactly what test_frontend_nn_exporter.py exercises
mod, params = Tiny().export_tvm(
    spec={"forward": {"x": spec.Tensor((3, 16), "float32")}}
)

with tvm.target.Target("llvm"):
    ex = relax.build(mod, "llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())

Honest limits: this tree is mid-migration, and the coexistence of tvm.s_tir, tvm_tirx, te, and topi means documentation and reality do not always agree; the version number is derived from the git tag at build time by setuptools_scm rather than stamped in the source. Building takes a real C++ toolchain and the submodules, the AGENTS.md explicitly warns that editable installs can make one worktree silently import another’s code, and the tirx sub-project is clearly moving target material. None of that changes the essential fact: if you want to understand or bend how models get compiled, the entire path from a PyTorch graph to a serializable virtual machine is open source here, readable in an afternoon.

Enjoyed this post? Never miss out on future posts by following us