.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. .. _relax-vm-arch: Relax Virtual Machine ===================== This document explains the Relax VM architecture in detail, covering the compilation pipeline from Relax IR to bytecode, the instruction set, the execution model, and the Python-level user interface. Overview -------- The end-to-end flow from model to execution is: 1. **Relax IR** — a high-level computational graph (``relax.Function`` inside an ``IRModule``). 2. **Compilation** — ``tvm.compile()`` applies the Relax transformation pipeline, then invokes ``VMCodeGen`` to translate each Relax function into bytecode instructions. 3. **Linking** — TIR functions are compiled to native kernels (via LLVM, CUDA, etc.); the bytecode, constant pool, and compiled kernels are packaged together into a ``VMExecutable``. 4. **Execution** — at runtime, a ``VirtualMachine`` loads the executable, initializes devices and memory allocators, and runs the bytecode. .. code-block:: text IRModule (Relax + TIR) │ ▼ relax_pipeline (FuseOps, LegalizeOps, ...) IRModule (optimized) │ ▼ VMCodeGen ExecBuilder (bytecode) + IRModule (TIR only) │ │ │ ▼ tirx.build() │ runtime.Module (native kernels) │ │ ▼ VMLink ▼ VMExecutable ◄───────── linked together │ ▼ VirtualMachine(exec, device) Runtime execution Compilation: From Relax IR to Bytecode -------------------------------------- Build entry point ~~~~~~~~~~~~~~~~~ The main entry point is ``tvm.compile()`` (which delegates to ``relax.build()`` in ``python/tvm/relax/vm_build.py``): .. code-block:: python import tvm from tvm import relax @tvm.script.ir_module class MyModule: @R.function def main(x: R.Tensor((3, 4), "float32")): return R.add(x, x) target = tvm.target.Target("llvm") ex = tvm.compile(MyModule, target) Internally, ``relax.build()`` performs these steps: 1. Apply the **Relax pipeline** (``relax.get_pipeline("default")``), which includes operator legalization, fusion, buffer planning, and other graph-level passes. 2. Create an ``ExecBuilder`` and run **VMCodeGen** (``src/relax/backend/vm/codegen_vm.cc``), which walks each ``relax.Function`` and emits bytecode instructions. The Relax functions are removed from the IRModule; only TIR functions remain. 3. Compile the remaining TIR functions to native code via ``tirx.build()``. 4. **Link** the bytecode executable with the compiled native module using ``VMLink``, producing a ``VMExecutable``. Two execution modes are supported: - ``exec_mode="bytecode"`` (default): Relax functions are interpreted by the VM's bytecode dispatch loop. - ``exec_mode="compiled"``: Relax functions are compiled into TIR functions (``VMTIRCodeGen``) that directly manipulate the register file, bypassing the interpreter loop. This avoids dispatch overhead but produces more code. Bytecode generation ~~~~~~~~~~~~~~~~~~~ The ``CodeGenVM`` class (``src/relax/backend/vm/codegen_vm.cc``) is an ``ExprFunctor`` that visits each Relax expression and emits instructions through the ``ExecBuilder``: - Each ``relax.Var`` is mapped to a register. - Function parameters occupy registers 0 through N-1. - Each binding in a ``SeqExpr`` generates one or more instructions; the result is stored in a new register. - Function calls (``R.call_tir``, ``R.call_packed``, operator calls) become ``Call`` instructions. - Conditional expressions (``relax.If``, written as Python ``if`` in TVMScript) become an ``If`` instruction followed by ``Goto`` to skip branches. - The function body ends with a ``Ret`` instruction. Instruction Set --------------- The VM uses a **register-based** architecture with an intentionally minimal instruction set. There are only four opcodes: .. list-table:: :header-rows: 1 :widths: 15 30 55 * - Opcode - Fields - Semantics * - ``Call`` - ``dst``, ``func_idx``, ``num_args``, ``args[]`` - Call function ``func_idx`` with the given arguments; store the result in register ``dst``. * - ``Ret`` - ``result`` - Return the value in register ``result`` to the caller. * - ``Goto`` - ``pc_offset`` - Jump forward or backward by ``pc_offset`` instructions. * - ``If`` - ``cond``, ``false_offset`` - If register ``cond`` is nonzero, fall through (pc++); otherwise jump by ``false_offset``. The VM itself performs **no mathematical computation**. All actual work — matrix multiplications, convolutions, elementwise operations — is carried out by compiled TIR kernels or external libraries (cuBLAS, cuDNN, etc.), dispatched through ``Call`` instructions. Instruction encoding ~~~~~~~~~~~~~~~~~~~~ Each instruction argument (``Instruction::Arg``) is a 64-bit word encoded as: - **Bits [63:56]** — ``ArgKind`` (8 bits): ``kRegister`` (0), ``kImmediate`` (1), ``kConstIdx`` (2), or ``kFuncIdx`` (3). - **Bits [55:0]** — value (56 bits, sign-extended). Two special register values exist: - ``kVoidRegister``: indicates "no destination" (the return value is discarded). - ``kVMRegister``: refers to the VM context pointer itself, passed as the first argument to closures. The instruction stream is stored as a flat ``vector`` (``instr_data``) with an offset table (``instr_offset``) for random access. Executable ---------- A ``VMExecutable`` (``include/tvm/runtime/vm/executable.h``) bundles everything needed for execution: - **Function table** (``func_table``): a ``vector`` describing every function. Each entry records the function's kind, name, instruction range (``start_instr`` to ``end_instr``), number of arguments, register file size, and parameter names. - **Constant pool** (``constants``): model weights, shape tuples, and other compile-time constants. - **Bytecode** (``instr_data`` + ``instr_offset``): the instruction stream. - **Imported modules**: the compiled TIR kernels and external libraries. Function kinds ~~~~~~~~~~~~~~ The VM recognizes three function kinds (``VMFuncInfo::FuncKind``): .. list-table:: :header-rows: 1 :widths: 20 80 * - Kind - Description * - ``kPackedFunc`` - An external C/C++ function looked up from imported modules or the global PackedFunc registry. Examples: ``vm.builtin.alloc_shape_heap``, ``vm.builtin.match_shape``. * - ``kVMFunc`` - A bytecode-interpreted Relax function. The VM interprets its instructions in ``RunLoop()``. * - ``kVMTIRFunc`` - A Relax function compiled to a TIR function (``exec_mode="compiled"``). Found in imports under the name ``__vmtir__``. Called directly with register file pointers, bypassing the interpreter loop. Serialization ~~~~~~~~~~~~~ The executable supports binary serialization for deployment: .. code-block:: python # Save ex.export_library("model.so") # Load loaded = tvm.runtime.load_module("model.so") vm = relax.VirtualMachine(loaded, tvm.cuda()) The binary format includes a magic number (``0xD225DE2F4214151E``), a version string (currently ``"0.14"``), followed by four sections: globals (the function table), memory scopes, constant pool, and bytecode. ``AsText()`` and ``AsPython()`` provide human-readable representations for debugging. Runtime Execution ----------------- VM initialization ~~~~~~~~~~~~~~~~~ At runtime, a ``VirtualMachine`` is created and initialized: .. code-block:: python from tvm.relax import VirtualMachine vm = VirtualMachine(exec_module, tvm.cuda()) Under the hood: 1. **LoadExecutable**: the bytecode and metadata are loaded from the ``VMExecutable``. 2. **Init**: devices and memory allocators are set up. Each device gets an ``Allocator`` (either ``NAIVE_ALLOCATOR`` or ``POOLED_ALLOCATOR``, defaulting to pooled). A CPU device is always added for shape computations. 3. **InitFuncPool**: the function pool is populated — ``kPackedFunc`` entries are resolved from imports or the global registry; ``kVMFunc`` and ``kVMTIRFunc`` entries are wrapped in ``VMClosure`` objects. 4. **Constant pool**: model constants are loaded and optionally transferred to the target device. The bytecode dispatch loop ~~~~~~~~~~~~~~~~~~~~~~~~~~ When a ``kVMFunc`` is invoked, the VM enters ``InvokeBytecode()``: 1. A new ``VMFrame`` is pushed onto the call stack. Each frame contains: - A **register file** (``vector``) — type-erased slots that can hold tensors, shapes, closures, or any TVM object. The size is determined at compile time (``VMFuncInfo::register_file_size``). - The **return program counter** — where to resume after the function returns. - The **caller's return register** — which register in the parent frame receives the result. 2. Function arguments are written to registers 0..N-1. 3. The program counter (``pc_``) is set to the function's ``start_instr``. 4. ``RunLoop()`` executes instructions until a ``Ret`` is encountered: - **Call**: resolve arguments (from registers, immediates, constant pool, or function pool), invoke the target function via ``InvokeClosurePacked()``, store the result in ``dst``. - **Ret**: read the return value from the specified register, write the result to the caller's return register, and return from ``RunLoop()`` (the frame is popped by an RAII guard when ``InvokeBytecode()`` exits). - **Goto**: adjust ``pc_`` by the offset. - **If**: check the condition register; if nonzero, fall through; otherwise jump by ``false_offset``. The dispatch loop is implemented in ``src/runtime/vm/vm.cc`` (``VirtualMachineImpl::RunLoop``). .. code-block:: text Frame Stack Register File (per frame) ┌─────────────┐ ┌────┬────┬────┬─────┬────┐ │ Frame 2 │ ───────► │ R0 │ R1 │ R2 │ ... │ Rn │ ├─────────────┤ └────┴────┴────┴─────┴────┘ │ Frame 1 │ ───────► [register file] ├─────────────┤ │ Frame 0 │ ───────► [register file] └─────────────┘ VMClosure and function dispatch ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Functions in the VM are stored in a ``func_pool_`` indexed by function table position. ``kVMFunc`` and ``kVMTIRFunc`` entries are wrapped as ``VMClosure`` objects, while ``kPackedFunc`` entries are stored as plain ``ffi::Function``. A ``VMClosure`` stores: - ``func_name``: the function's string name. - ``impl``: a ``ffi::Function`` that takes the VM context pointer as its first argument, followed by the actual parameters. When the VM encounters a ``Call`` instruction, it looks up the function in ``func_pool_`` by index and dispatches via ``InvokeClosurePacked()``. If the target is a ``VMClosure``, the VM pointer is prepended to the arguments and ``impl`` is invoked. If it is a plain ``ffi::Function``, it is called directly. ``VMClosure::BindLastArgs`` enables partial application — it creates a new function with some arguments pre-bound at the end, useful for implementing captured closures in Relax. Built-in operations ~~~~~~~~~~~~~~~~~~~ The VM relies on several built-in PackedFuncs (registered in ``src/runtime/vm/builtin.cc``) for runtime support: - ``vm.builtin.alloc_shape_heap``: allocate workspace for symbolic shape computations. - ``vm.builtin.match_shape``: validate tensor shapes against expected patterns at runtime, supporting assertions (``kAssertEqualToImm``, ``kAssertEqualToLoad``), storing symbolic dimensions to the shape heap (``kStoreToHeap``), or no-ops (``kNoOp``). - ``vm.builtin.make_shape``: construct shape tuples from immediates or heap-loaded values. - ``vm.builtin.match_prim_value``: validate primitive values (e.g., integers) against expected patterns. - ``vm.builtin.copy``: copy a value into a register. Used in several codegen scenarios: materializing non-register arguments (immediates, constants) into registers, ensuring each variable binding gets its own register, and merging results from if/else branches. Python Interface ---------------- Users interact with the VM through ``tvm.relax.VirtualMachine``: .. code-block:: python import tvm from tvm import relax import numpy as np # Compile ex = tvm.compile(MyModule, target="llvm") # Create VM vm = relax.VirtualMachine(ex, tvm.cpu()) # Direct invocation inp = tvm.runtime.tensor(np.random.rand(3, 4).astype("float32")) result = vm["main"](inp) # Stateful interface (useful for RPC) vm.set_input("main", inp) vm.invoke_stateful("main") output = vm.get_outputs("main") Key methods: - ``vm["func_name"](*args)`` — direct invocation, returns the result. - ``vm.set_input()`` / ``vm.invoke_stateful()`` / ``vm.get_outputs()`` — stateful interface that avoids sending output over the wire, useful for RPC-based remote execution. - ``vm.save_function(func_name, saved_name, *args)`` — pre-bind arguments for repeated calls, reducing dictionary lookup overhead during benchmarking. - ``vm.time_evaluator(func_name, dev)`` — returns a timing function following the same convention as ``tvm.runtime.Module.time_evaluator``. - ``vm.set_instrument(func)`` — register an instrumentation callback that is invoked before/after every ``Call`` instruction. The callback can return ``VMInstrumentReturnKind.SKIP_RUN`` to skip the call. Instrumentation ~~~~~~~~~~~~~~~ The VM supports observability via instrumentation: **Instrumentation** via ``set_instrument()``: .. code-block:: python def my_instrument(func, func_symbol, before_run, ret_value, *args): if before_run: print(f"About to call: {func_symbol}") return VMInstrumentReturnKind.NO_OP vm.set_instrument(my_instrument) vm["main"](inp) The instrument function is called before and after every ``Call`` instruction, receiving the function object, its symbol name, a flag indicating before/after, the return value (only valid after), and all arguments. Inspecting Bytecode ------------------- The executable provides text and Python representations of the compiled bytecode: .. code-block:: python ex = tvm.compile(MyModule, target="llvm") print(ex.as_text()) # Human-readable instruction listing print(ex.as_python()) # Equivalent Python program print(ex.stats()) # Summary statistics These are invaluable for debugging compilation issues — they show exactly which functions are called, in what order, and how registers are used. Source Code Map --------------- .. list-table:: :header-rows: 1 :widths: 45 55 * - Path - Contents * - ``include/tvm/runtime/vm/bytecode.h`` - Instruction, Opcode, and Arg definitions * - ``include/tvm/runtime/vm/executable.h`` - VMExecutable, VMFuncInfo, serialization * - ``include/tvm/runtime/vm/vm.h`` - VirtualMachine base class, VMClosure * - ``src/runtime/vm/vm.cc`` - VirtualMachineImpl, RunLoop, InvokeBytecode * - ``src/runtime/vm/executable.cc`` - Serialization/deserialization, text output * - ``src/runtime/vm/builtin.cc`` - Built-in operations (shape matching, allocation) * - ``src/relax/backend/vm/codegen_vm.cc`` - CodeGenVM: Relax IR → bytecode * - ``src/relax/backend/vm/codegen_vm_tir.cc`` - VMTIRCodeGen: Relax IR → compiled TIR * - ``python/tvm/runtime/vm.py`` - Python VirtualMachine wrapper * - ``python/tvm/relax/vm_build.py`` - ``relax.build()`` and VMExecutable Python class