chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,306 @@
|
||||
.. 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.
|
||||
|
||||
.. _codegen-arch:
|
||||
|
||||
Code Generation
|
||||
===============
|
||||
|
||||
Code generation is the final stage of the TVM compilation pipeline — it translates TIR
|
||||
``PrimFunc``\ s into executable code for a target device. This document explains how TIR
|
||||
functions become native CPU instructions, GPU kernels, or source code strings, covering the
|
||||
target dispatch mechanism, the two codegen families (LLVM and Source), and the runtime module
|
||||
system that wraps the generated code.
|
||||
|
||||
|
||||
Where Codegen Fits
|
||||
------------------
|
||||
|
||||
When a user calls ``tvm.compile()``, the compilation proceeds in two phases:
|
||||
|
||||
1. **Relax phase**: the Relax pipeline optimizes and fuses the computational graph, then
|
||||
``VMCodeGen`` translates Relax functions into VM bytecode (see :ref:`relax-vm-arch`).
|
||||
2. **TIR phase**: TIR ``PrimFunc``\ s (the actual compute kernels) are compiled to native code.
|
||||
|
||||
The TIR phase is handled internally by ``tirx.build()`` (called from ``relax.build()``).
|
||||
It performs these steps:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
TIR PrimFuncs (in IRModule)
|
||||
│
|
||||
▼ TIR pipeline ← lowering passes (flatten buffers, lower intrinsics, etc.)
|
||||
TIR PrimFuncs (lowered)
|
||||
│
|
||||
▼ split_host_device_mods() ← separate host and device functions
|
||||
Host IRModule + Device IRModule(s)
|
||||
│ │
|
||||
▼ ▼
|
||||
codegen_build() codegen_build() ← target-specific code generation
|
||||
│ │
|
||||
▼ ▼
|
||||
Host Module Device Module(s)
|
||||
│ │
|
||||
▼ import_module() │
|
||||
Host Module ◄─────────────┘ ← device modules imported into host
|
||||
│
|
||||
▼ (returned to relax.build for linking with VM bytecode)
|
||||
|
||||
|
||||
Target Dispatch
|
||||
---------------
|
||||
|
||||
The core dispatch logic lives in ``codegen::Build()`` (``src/target/codegen.cc``), which is
|
||||
called from the Python-side ``codegen_build()`` in ``tirx/build.py``. It selects the correct
|
||||
backend based on the ``Target`` object:
|
||||
|
||||
.. code-block:: cpp
|
||||
|
||||
ffi::Module Build(IRModule mod, Target target) {
|
||||
std::string build_f_name = "target.build." + target->kind->name;
|
||||
const auto bf = tvm::ffi::Function::GetGlobal(build_f_name);
|
||||
return (*bf)(mod, target).cast<ffi::Module>();
|
||||
}
|
||||
|
||||
Each backend registers its build function via FFI:
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 25 30 45
|
||||
|
||||
* - FFI Key
|
||||
- Backend
|
||||
- Codegen Class
|
||||
* - ``target.build.llvm``
|
||||
- CPU (x86, ARM, etc.)
|
||||
- ``CodeGenCPU`` (→ LLVM IR → machine code)
|
||||
* - ``target.build.cuda``
|
||||
- NVIDIA GPU
|
||||
- ``CodeGenCUDA`` (→ CUDA C → PTX/cubin)
|
||||
* - ``target.build.rocm``
|
||||
- AMD GPU
|
||||
- ``CodeGenAMDGPU`` (→ LLVM IR → AMDGPU ISA)
|
||||
* - ``target.build.nvptx``
|
||||
- NVIDIA PTX
|
||||
- ``CodeGenNVPTX`` (→ LLVM IR → PTX)
|
||||
* - ``target.build.metal``
|
||||
- Apple GPU
|
||||
- ``CodeGenMetal`` (→ Metal Shading Language)
|
||||
* - ``target.build.opencl``
|
||||
- OpenCL devices
|
||||
- ``CodeGenOpenCL`` (→ OpenCL C)
|
||||
* - ``target.build.vulkan``
|
||||
- Vulkan devices
|
||||
- ``CodeGenSPIRV`` (→ SPIR-V binary)
|
||||
* - ``target.build.webgpu``
|
||||
- WebGPU
|
||||
- ``CodeGenWebGPU`` (→ WGSL)
|
||||
* - ``target.build.c``
|
||||
- C host code
|
||||
- ``CodeGenCHost`` (→ C source)
|
||||
|
||||
|
||||
Two Codegen Families
|
||||
--------------------
|
||||
|
||||
TVM has two families of code generators, corresponding to two fundamentally different strategies
|
||||
for producing executable code:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
LLVM Family Source Family
|
||||
────────── ─────────────
|
||||
TIR → LLVM IR → machine code TIR → source string → external compiler
|
||||
(in-process, JIT or AOT) (CUDA C, OpenCL C, Metal, WGSL)
|
||||
|
||||
LLVM family
|
||||
~~~~~~~~~~~
|
||||
|
||||
``CodeGenLLVM`` (``src/target/llvm/codegen_llvm.h``) translates TIR directly to LLVM IR using
|
||||
the LLVM C++ API. The generated ``llvm::Module`` is then compiled to native code by LLVM's
|
||||
backend (x86, ARM, NVPTX, AMDGPU, etc.).
|
||||
|
||||
**Inheritance**:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
CodeGenLLVM (base)
|
||||
├── CodeGenCPU ← x86, ARM (target.build.llvm)
|
||||
│ └── CodeGenHexagon
|
||||
├── CodeGenNVPTX ← NVIDIA PTX via LLVM (target.build.nvptx)
|
||||
└── CodeGenAMDGPU ← AMD GPU via LLVM (target.build.rocm)
|
||||
|
||||
``CodeGenLLVM`` inherits from both ``ExprFunctor<llvm::Value*(const Expr&)>`` and
|
||||
``StmtFunctor<void(const Stmt&)>``. Each TIR node type has a corresponding visitor:
|
||||
|
||||
- **Expressions** (``VisitExpr_``) convert TIR expressions to LLVM ``Value``\ s:
|
||||
arithmetic ops → LLVM binary instructions, ``BufferLoad`` → load with pointer arithmetic,
|
||||
``Cast`` → LLVM type conversions, ``Call`` → intrinsic or extern function calls.
|
||||
- **Statements** (``VisitStmt_``) emit LLVM IR side effects:
|
||||
``BufferStore`` → store instructions, ``For`` → loop basic blocks with branches,
|
||||
``IfThenElse`` → conditional branches, ``AllocBuffer`` → stack or heap allocation.
|
||||
|
||||
The key methods on ``CodeGenLLVM`` are:
|
||||
|
||||
- ``Create(LLVMTarget*)`` — factory that returns a target-specific subclass.
|
||||
- ``Init(...)`` — set up the LLVM context, module, and builder.
|
||||
- ``DeclareFunction(gvar, f)`` / ``AddFunction(gvar, f)`` — forward-declare then compile a
|
||||
``PrimFunc`` to LLVM IR.
|
||||
- ``Finish()`` — return the completed ``llvm::Module``.
|
||||
|
||||
Source family
|
||||
~~~~~~~~~~~~~
|
||||
|
||||
``CodeGenC`` (``src/target/source/codegen_c.h``) generates C-like source code as text. Each
|
||||
target subclass overrides methods to emit target-specific syntax.
|
||||
|
||||
**Inheritance**:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
CodeGenC (base)
|
||||
├── CodeGenCUDA ← CUDA C (target.build.cuda)
|
||||
├── CodeGenOpenCL ← OpenCL C (target.build.opencl)
|
||||
├── CodeGenMetal ← Metal Shading Language (target.build.metal)
|
||||
├── CodeGenWebGPU ← WGSL (target.build.webgpu)
|
||||
└── CodeGenCHost ← C host code (target.build.c)
|
||||
|
||||
``CodeGenC`` also uses the visitor pattern (``ExprFunctor`` and ``StmtFunctor``), but outputs to
|
||||
``std::ostream`` instead of constructing LLVM IR. Subclasses override target-specific methods:
|
||||
|
||||
- ``PrintStorageScope(scope, os)`` — emit memory qualifiers (e.g., ``__shared__`` for CUDA,
|
||||
``__local`` for OpenCL).
|
||||
- ``BindThreadIndex(iv)`` — emit thread index bindings (e.g., ``threadIdx.x``, ``blockIdx.y``).
|
||||
- ``PrintType(dtype, os)`` — emit target-specific type names (e.g., ``half`` for float16).
|
||||
- ``PrintVecBinaryOp(...)`` — emit vectorized operations in target syntax.
|
||||
|
||||
For CUDA, the build flow (``BuildCUDA`` in ``src/target/opt/build_cuda_on.cc``) is:
|
||||
|
||||
1. ``CodeGenCUDA`` generates CUDA C source.
|
||||
2. An optional post-processing callback (``tvm_callback_cuda_postproc``) transforms the source.
|
||||
3. A Python callback (``tvm_callback_cuda_compile``) compiles the source to PTX or cubin via
|
||||
NVRTC or NVCC.
|
||||
4. The result is wrapped in a ``CUDAModule``.
|
||||
|
||||
Design choice
|
||||
~~~~~~~~~~~~~
|
||||
|
||||
Why two families?
|
||||
|
||||
- **LLVM family** produces higher-quality code — LLVM applies its own optimization passes
|
||||
(instruction selection, register allocation, vectorization). Best for CPU targets where TVM
|
||||
has full control over the compilation.
|
||||
- **Source family** is more portable — it generates human-readable source that can be compiled
|
||||
by vendor toolchains (NVCC, Metal compiler, etc.). This is necessary for GPU targets where
|
||||
the vendor compiler handles device-specific optimizations and the runtime compilation model
|
||||
(e.g., NVRTC for CUDA, runtime shader compilation for Metal/OpenCL).
|
||||
|
||||
|
||||
Host/Device Split
|
||||
-----------------
|
||||
|
||||
When compiling for GPU targets, TIR functions are split into two categories:
|
||||
|
||||
- **Host functions** — run on the CPU. They set up kernel launch parameters (grid/block
|
||||
dimensions), allocate memory, and invoke device kernels. Compiled with ``target.build.llvm``
|
||||
or ``target.build.c``.
|
||||
- **Device functions** — the actual compute kernels that run on the GPU. Compiled with the
|
||||
target-specific codegen (``target.build.cuda``, etc.).
|
||||
|
||||
``split_host_device_mods()`` (``python/tvm/tirx/build.py``) separates functions by their
|
||||
``target`` attribute: functions whose target kind is ``"llvm"`` or ``"c"`` go to the host
|
||||
module; all others go to device modules grouped by target.
|
||||
|
||||
After compilation, device modules are imported into the host module via ``import_module()``,
|
||||
forming a module tree. At runtime, the host module dispatches to the imported device module
|
||||
when a device kernel is called.
|
||||
|
||||
|
||||
Runtime Modules
|
||||
---------------
|
||||
|
||||
Each codegen produces a ``runtime.Module`` — the container that holds the generated code and
|
||||
exposes it as callable ``PackedFunc``\ s.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 20 35 45
|
||||
|
||||
* - Module Type
|
||||
- How Code Is Stored
|
||||
- How Code Is Executed
|
||||
* - ``LLVMModule``
|
||||
- LLVM IR (in-memory ``llvm::Module``)
|
||||
- JIT-compiled on first call (MCJIT or ORC). Function pointers cached for subsequent calls.
|
||||
* - ``CUDAModule``
|
||||
- PTX or cubin binary
|
||||
- Loaded via CUDA driver API (``cuModuleLoad``). Kernels launched via ``cuLaunchKernel``.
|
||||
* - ``CSourceModule``
|
||||
- C source string
|
||||
- Not directly executable. Used as a build artifact for AOT compilation.
|
||||
* - ``DeviceSourceModule``
|
||||
- Device source string (OpenCL C, Metal, WGSL)
|
||||
- Compiled at runtime by the device driver (e.g., ``clCreateProgramWithSource``).
|
||||
|
||||
All module types implement the same interface: ``GetFunction(name)`` returns a ``PackedFunc``
|
||||
that can be called from Python or C++. The VM and other runtime components use this interface
|
||||
to invoke compiled kernels without knowing which backend produced them.
|
||||
|
||||
The module tree is serializable via ``export_library()``, which packs the host module and all
|
||||
imported device modules into a single shared library (``.so`` / ``.dll`` / ``.dylib``) or
|
||||
a tar archive for deployment.
|
||||
|
||||
|
||||
Source Code Map
|
||||
---------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 50 50
|
||||
|
||||
* - Path
|
||||
- Contents
|
||||
* - ``python/tvm/tirx/build.py``
|
||||
- ``tirx.build()``: TIR compilation entry, host/device split, module linking
|
||||
* - ``src/target/codegen.cc``
|
||||
- ``codegen::Build()``: target dispatch via ``"target.build.<kind>"``
|
||||
* - ``src/target/llvm/codegen_llvm.h``
|
||||
- ``CodeGenLLVM``: TIR → LLVM IR base class
|
||||
* - ``src/target/llvm/codegen_cpu.h``
|
||||
- ``CodeGenCPU``: CPU-specific LLVM codegen (x86, ARM)
|
||||
* - ``src/target/llvm/codegen_nvptx.cc``
|
||||
- ``CodeGenNVPTX``: NVIDIA PTX via LLVM
|
||||
* - ``src/target/llvm/codegen_amdgpu.cc``
|
||||
- ``CodeGenAMDGPU``: AMD GPU via LLVM
|
||||
* - ``src/target/llvm/llvm_module.cc``
|
||||
- ``LLVMModuleNode``: runtime module with JIT compilation
|
||||
* - ``src/target/source/codegen_c.h``
|
||||
- ``CodeGenC``: TIR → C-like source base class
|
||||
* - ``src/target/source/codegen_cuda.h``
|
||||
- ``CodeGenCUDA``: TIR → CUDA C
|
||||
* - ``src/target/source/codegen_opencl.h``
|
||||
- ``CodeGenOpenCL``: TIR → OpenCL C
|
||||
* - ``src/target/source/codegen_metal.h``
|
||||
- ``CodeGenMetal``: TIR → Metal Shading Language
|
||||
* - ``src/target/source/codegen_c_host.h``
|
||||
- ``CodeGenCHost``: TIR → C host code
|
||||
* - ``src/target/opt/build_cuda_on.cc``
|
||||
- ``BuildCUDA``: CUDA build flow (codegen → compile → module)
|
||||
* - ``src/target/spirv/codegen_spirv.h``
|
||||
- ``CodeGenSPIRV``: TIR → SPIR-V for Vulkan
|
||||
* - ``src/target/source/codegen_webgpu.h``
|
||||
- ``CodeGenWebGPU``: TIR → WGSL
|
||||
@@ -0,0 +1,245 @@
|
||||
.. 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.
|
||||
|
||||
|
||||
.. _tvm-target-specific-overview:
|
||||
|
||||
Device/Target Interactions
|
||||
==========================
|
||||
|
||||
This documented is intended for developers interested in understanding
|
||||
how the TVM framework interacts with specific device APIs, or who
|
||||
may want to implement support for a new API or new hardware.
|
||||
|
||||
There are three main aspects that must be implemented for any new
|
||||
runtime environment.
|
||||
|
||||
* The :ref:`DeviceAPI <tvm-target-specific-device-api>` class gives a
|
||||
handle to a specific device, and the API used to interact with it.
|
||||
It defines a common interface for querying device parameters
|
||||
(e.g. memory available, number of threads, etc.) and for performing
|
||||
simple actions (e.g. copying memory from the host, or between
|
||||
buffers on the device).
|
||||
|
||||
* The :ref:`Target <tvm-target-specific-target>` class contains a
|
||||
description of the device on which a function will run. It is
|
||||
exposed both to the target code generators and to the optimization
|
||||
passes.
|
||||
|
||||
* The :ref:`target code generators <tvm-target-specific-codegen>`
|
||||
construct a :ref:`Module <tvm-runtime-system-module>` consisting of
|
||||
one or more :ref:`PackedFunc <tvm-runtime-system-packed-func>`, from
|
||||
an IRModule.
|
||||
|
||||
.. _tvm-target-specific-device-api:
|
||||
|
||||
DeviceAPI
|
||||
---------
|
||||
|
||||
The ``DeviceAPI`` represents a handle to a specific hardware device
|
||||
API. (e.g. ``CUDADeviceAPI`` handles all interactions through the
|
||||
CUDA framework.) Most ``DeviceAPI`` methods accept a ``device_id``
|
||||
parameter to specify which device should be accessed. In Python,
|
||||
these are typically accessed using the :py:func:`tvm.runtime.device`
|
||||
function, which returns a handle to a specific device, accessed
|
||||
through a specific API. (e.g. ``tvm.runtime.device('cuda',0)`` gives
|
||||
access to physical device ``0``, accessed through the CUDA API.)
|
||||
|
||||
.. _device_api.h: https://github.com/apache/tvm/blob/main/include/tvm/runtime/device_api.h
|
||||
|
||||
* Attribute queries - ``GetAttr`` allows different
|
||||
device-specific parameters to be queried, such as the device name,
|
||||
number of threads, etc. The parameters that can be queried are
|
||||
defined in ``enum DeviceAttrKind`` in `device_api.h`_. Not all
|
||||
query-able parameters are supported by all devices. If a parameter
|
||||
cannot be queried (e.g. ``kMaxClockRate`` on Vulkan), or if a
|
||||
parameter isn't applicable (e.g. ``kWarpSize`` on CPU), then those
|
||||
queries should return ``nullptr``.
|
||||
|
||||
* Setting active device - ``SetDevice`` should set a
|
||||
particular device as being active. If a ``PackedFunc`` generated by
|
||||
the target-specific code gen requires execution on a device, it
|
||||
should run on the active device.
|
||||
|
||||
* Memory management - Utilities for allocating and deallocating memory
|
||||
on the device.
|
||||
|
||||
* Allocate data space - ``AllocDataSpace`` and ``FreeDataSpace``
|
||||
allocate and free space on the device. These allocations can be
|
||||
provided as inputs and outputs to an operator and make up the
|
||||
primary data flow of the operator graph. It must be possible to
|
||||
transfer data from the host to/from a data space. The return
|
||||
value is an opaque ``void*``. While some implementations return a
|
||||
memory address, this is not required, and the ``void*`` may be an
|
||||
opaque handle that is interpretable only by the device backend
|
||||
that generated it. The ``void*`` is used as an argument to other
|
||||
backend-specific functions, such as ``CopyDataFromTo``.
|
||||
|
||||
* Allocate work space - ``AllocWorkspace`` and ``FreeWorkspace``
|
||||
allocate and free space on the device. Unlike data space, these
|
||||
are used for storage of intermediate values within an operator
|
||||
definition, and are not required to be transferable to/from the
|
||||
host device. If a ``DeviceAPI`` subclass does not implement these
|
||||
methods, they will default to calling the corresponding
|
||||
``DataSpace`` functions.
|
||||
|
||||
* Copy data - ``CopyDataFromTo`` should copy data from one location
|
||||
to another. The type of copy is determined by the ``dev_from``
|
||||
and ``dev_to`` parameters. Implementations should support copying
|
||||
memory from CPU to device, from device to CPU, and from one buffer
|
||||
to another on a single device. If the source or destination
|
||||
locations are on the CPU, the corresponding ``void*`` points to a
|
||||
CPU address that can be passed into ``memcpy``. If the source or
|
||||
destinations locations are on the device, the corresponding
|
||||
``void*`` was previously generated by either ``AllocDataSpace`` or
|
||||
``AllocWorkspace``.
|
||||
|
||||
These copies are queued to execute on a specific
|
||||
``TVMStreamHandle``. However, implementations should not assume
|
||||
that CPU buffers remains valid or accessible after the call to
|
||||
``CopyDataFromTo`` completes.
|
||||
|
||||
|
||||
* Execution stream management - utilities for handling
|
||||
``TVMStreamHandle``, which represents parallel streams of execution
|
||||
used to execute commands.
|
||||
|
||||
* Create stream - ``CreateStream`` and ``FreeStream`` should
|
||||
allocate/free a handle to a stream of execution. If a device
|
||||
implements only a single queue of commands, then ``CreateStream``
|
||||
should return ``nullptr``.
|
||||
|
||||
* Set active stream - ``SetStream`` should set a stream as being
|
||||
active. While active, if a ``PackedFunc`` generated by the
|
||||
target-specific code gen requires execution on a device, the work
|
||||
should be submitted to the active stream.
|
||||
|
||||
* Synchronize to CPU - ``StreamSync`` should synchronize a stream of
|
||||
execution to the CPU. The call to ``StreamSync`` should return
|
||||
once all memory transfers and computations submitted prior to the
|
||||
``StreamSync`` call have completed.
|
||||
|
||||
* Synchronize between streams - ``SyncStreamFromTo`` should
|
||||
introduce a synchronization barrier between the source and
|
||||
destination stream. That is, the destination stream may not
|
||||
proceed beyond commands currently queued until the source stream
|
||||
has completed all commands that are currently queued.
|
||||
|
||||
|
||||
In order to be usable by the TVM framework, the new DeviceAPI should
|
||||
then be registered with the following steps.
|
||||
|
||||
#. Create a function that instantiates the new DeviceAPI, and returns
|
||||
a pointer to it::
|
||||
|
||||
FooDeviceAPI* FooDeviceAPI::Global() {
|
||||
static FooDeviceAPI inst;
|
||||
return &inst;
|
||||
}
|
||||
|
||||
#. Register the function to the tvm registry::
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def("device_api.foo", FooDeviceAPI::Global);
|
||||
}
|
||||
|
||||
.. _base.h: https://github.com/apache/tvm/blob/main/include/tvm/runtime/base.h
|
||||
|
||||
#. Add an entry for the new DeviceAPI to the ``TVMDeviceExtType`` enum
|
||||
in `base.h`_. The value should be an unused value greater
|
||||
than ``DLDeviceType::kDLExtDev``, but less than
|
||||
``DeviceAPIManager::kMaxDeviceAPI``.
|
||||
|
||||
#. Add a case in ``DeviceName`` in `device_api.h`_ to convert from the
|
||||
enum value to a string representation. This string representation
|
||||
should match the name given to ``GlobalDef().def``.
|
||||
|
||||
#. Add entries to the ``_DEVICE_TYPE_TO_NAME`` and ``_DEVICE_NAME_TO_TYPE`` dictionaries of
|
||||
:py:class:`tvm.runtime.Device` for the new enum value.
|
||||
|
||||
|
||||
.. _tvm-target-specific-target:
|
||||
|
||||
Target Definition
|
||||
-----------------
|
||||
|
||||
The ``Target`` object is a lookup table of properties about a physical
|
||||
device, its hardware/driver limits, and its capabilities. The
|
||||
``Target`` is accessible both during optimization and code generation
|
||||
stages. While the same ``Target`` class is used for all runtime
|
||||
targets, each runtime target may need to add target-specific options.
|
||||
|
||||
.. _target_kind.cc: https://github.com/apache/tvm/blob/main/src/target/target_kind.cc
|
||||
|
||||
In `target_kind.cc`_, add a new declaration of
|
||||
``TVM_REGISTER_TARGET_KIND``, passing a string name of the new target,
|
||||
and the ``TVMDeviceExtType`` or ``DLDeviceType`` enum value for the
|
||||
device on which that target should run. Typically, the target name
|
||||
and the device name will match (e.g., the ``"cuda"`` target runs on
|
||||
the ``kDLCUDA`` device). There are exceptions, such as when multiple
|
||||
different code generation targets can run on the same physical device
|
||||
(e.g., the ``"llvm"`` and ``"c"`` targets both run on the ``kDLCPU``
|
||||
device type).
|
||||
|
||||
All options for a specific target kind are added with the
|
||||
``add_attr_option`` function, with optional default values. A `Target`
|
||||
parser can be added with ``set_target_parser`` to process
|
||||
any parameters that are dynamically based on other parameters or
|
||||
queried from device properties.
|
||||
|
||||
This argument definition defines a parser that can unpack a string
|
||||
description of a target. This is done in the ``Target::Target(const
|
||||
String&)`` constructor in C++, which accepts a JSON-formatted string
|
||||
and is typically called using the :py:class:`tvm.target.Target` python
|
||||
object. For example, ``tvm.target.Target('{"kind": "cuda",
|
||||
"max_num_threads": 1024}')`` will create a ``cuda`` target, while
|
||||
overriding the default maximum number of threads.
|
||||
|
||||
In a code generator, the target properties can be accessed using
|
||||
``target->GetAttr<T>(param_name)`` in C++, or with the
|
||||
``target.attrs`` dictionary in Python.
|
||||
|
||||
|
||||
.. _tvm-target-specific-codegen:
|
||||
|
||||
Target Code Generators
|
||||
----------------------
|
||||
|
||||
The code generators take an optimized ``IRModule`` and converts it
|
||||
into an executable representation. Each code generator must be
|
||||
registered in order to be used by the TVM framework. This is done by
|
||||
registering a function named ``"target.build.foo"``, where ``foo`` is
|
||||
the same name as was used in the ``TVM_REGISTER_TARGET_KIND``
|
||||
definition above. ::
|
||||
|
||||
tvm::runtime::Module GeneratorFooCode(IRModule mod, Target target);
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def("target.build.foo", GeneratorFooCode);
|
||||
}
|
||||
|
||||
The code generator takes two arguments. The first is the ``IRModule``
|
||||
to compile, and the second is the ``Target`` that describes the device
|
||||
on which the code should run. Because the environment performing the
|
||||
compilation is not necessarily the same as the environment that will
|
||||
be executing the code, code generators should not perform any
|
||||
attribute lookups on the device itself, and should instead access
|
||||
parameters stored in the ``Target``.
|
||||
|
||||
Each function in the input ``IRModule`` should be accessible by name
|
||||
in the output ``runtime::Module``.
|
||||
@@ -0,0 +1,361 @@
|
||||
.. 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.
|
||||
|
||||
.. _external-library-dispatch:
|
||||
|
||||
External Library Dispatch (BYOC)
|
||||
================================
|
||||
|
||||
When deploying models, certain operator patterns (e.g., matmul + bias + relu) can be executed
|
||||
more efficiently by vendor-optimized libraries such as cuBLAS, CUTLASS, cuDNN, or DNNL. TVM's
|
||||
**BYOC (Bring Your Own Codegen)** mechanism identifies these patterns in a Relax module and
|
||||
offloads them to external backends, while keeping the rest of the computation on TVM's own
|
||||
generated kernels.
|
||||
|
||||
This document explains the BYOC pipeline: how patterns are registered, how subgraphs are
|
||||
matched and extracted, how backend code generators are invoked, and how the externally compiled
|
||||
code is executed at runtime.
|
||||
|
||||
|
||||
Overview
|
||||
--------
|
||||
|
||||
The BYOC pipeline consists of four stages:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
IRModule (high-level Relax IR)
|
||||
│
|
||||
▼ FuseOpsByPattern ← match high-level ops, create composite functions
|
||||
IRModule (with Composite + Codegen attributes)
|
||||
│
|
||||
▼ RunCodegen ← invoke backend codegen via FFI
|
||||
IRModule (with call_dps_packed to ExternFunc)
|
||||
+ external runtime Modules
|
||||
│
|
||||
▼ LegalizeOps + FuseOps + ... ← compile remaining ops normally
|
||||
│
|
||||
▼ VM compilation ← link external modules into executable
|
||||
Deployable artifact
|
||||
|
||||
Each stage is a Relax transformation pass that operates on the ``IRModule``:
|
||||
|
||||
1. **FuseOpsByPattern** — matches operator subgraphs against registered patterns and groups them
|
||||
into composite functions annotated with ``Composite`` and ``Codegen`` attributes.
|
||||
2. **MergeCompositeFunctions** (optional) — merges multiple composite functions targeting the same
|
||||
backend when inter-operator dependencies allow.
|
||||
3. **RunCodegen** — finds all functions with a ``Codegen`` attribute, invokes the corresponding
|
||||
backend code generator via FFI, and replaces the original calls with ``call_dps_packed``
|
||||
to externally compiled functions.
|
||||
4. **Linking** — the resulting external ``runtime.Module``\ s are attached to the ``IRModule``
|
||||
as the ``external_mods`` attribute and bundled into the final executable during
|
||||
``relax.build()``.
|
||||
|
||||
|
||||
Pattern Registration
|
||||
--------------------
|
||||
|
||||
Each backend registers the operator patterns it supports in a **global pattern registry**
|
||||
(``python/tvm/relax/backend/pattern_registry.py``). The registry is a static table that maps
|
||||
pattern names to ``FusionPattern`` objects.
|
||||
|
||||
Registering patterns
|
||||
~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.relax.backend.pattern_registry import register_patterns
|
||||
from tvm.relax.backend.patterns import make_matmul_pattern
|
||||
|
||||
register_patterns([
|
||||
(
|
||||
"cublas.matmul", # pattern name (prefix = backend)
|
||||
*make_matmul_pattern( # returns (DFPattern, annotation_patterns)
|
||||
with_bias=False,
|
||||
),
|
||||
_check_matmul, # check function
|
||||
),
|
||||
(
|
||||
"cublas.matmul_bias_relu",
|
||||
*make_matmul_pattern(
|
||||
with_bias=True,
|
||||
activation="relax.nn.relu",
|
||||
),
|
||||
_check_matmul,
|
||||
),
|
||||
# ... more patterns
|
||||
])
|
||||
|
||||
Each entry is a tuple of ``(name, pattern, annotation_patterns, check_func)`` that gets
|
||||
converted to a ``FusionPattern`` object. The name prefix (e.g., ``"cublas"``) identifies the
|
||||
backend; ``get_patterns_with_prefix("cublas")`` retrieves all patterns for that backend.
|
||||
|
||||
Patterns registered later have **higher priority** — when a subgraph matches multiple patterns,
|
||||
the highest-priority match wins.
|
||||
|
||||
Pattern templates
|
||||
~~~~~~~~~~~~~~~~~
|
||||
|
||||
``python/tvm/relax/backend/patterns.py`` provides reusable templates for common patterns:
|
||||
|
||||
- ``make_matmul_pattern(with_bias, activation, transposed_rhs)`` — matmul with optional bias
|
||||
and activation fusion
|
||||
- ``make_conv2d_pattern(with_bias, activation)`` — 2D convolution
|
||||
- ``make_attention_pattern()`` — multi-head attention
|
||||
- ``make_residual_block_pattern()`` — residual connections
|
||||
- ``make_layer_norm_pattern()`` / ``make_rms_norm_pattern()`` — normalization layers
|
||||
|
||||
Each template returns ``(DFPattern, Mapping[str, DFPattern])`` — the main pattern and its
|
||||
annotation sub-patterns.
|
||||
|
||||
Check functions
|
||||
~~~~~~~~~~~~~~~
|
||||
|
||||
The check function validates whether a matched subgraph can actually be handled by the backend.
|
||||
It receives a ``PatternCheckContext`` and returns ``True`` to accept or ``False`` to reject.
|
||||
|
||||
Typical checks include:
|
||||
|
||||
- **Data type support**: verify the operand dtypes are supported (e.g., cuBLAS supports
|
||||
float16, float32, int8, bfloat16, float8 for matmul).
|
||||
- **Shape constraints**: verify reduction axes are constant, batch dimensions are compatible.
|
||||
- **Leaking intermediates**: reject if an intermediate result is used outside the fused group
|
||||
(via ``has_leaking_intermediate_variables()``).
|
||||
|
||||
|
||||
Partitioning
|
||||
------------
|
||||
|
||||
After patterns are registered, a backend provides a **partition function** that applies
|
||||
``FuseOpsByPattern`` to an ``IRModule``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# python/tvm/relax/backend/cuda/cublas.py
|
||||
def partition_for_cublas(mod, bind_constants=False):
|
||||
patterns = get_patterns_with_prefix("cublas")
|
||||
return transform.FuseOpsByPattern(
|
||||
patterns, bind_constants=bind_constants, annotate_codegen=True
|
||||
)(mod)
|
||||
|
||||
With ``annotate_codegen=True``, each matched subgraph is wrapped in a two-level function
|
||||
structure:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
# Outer function — tagged for the codegen backend
|
||||
@R.function
|
||||
def fused_relax_matmul_cublas0(args...):
|
||||
R.func_attr({"Codegen": "cublas", "global_symbol": "fused_relax_matmul_cublas0"})
|
||||
...
|
||||
# Inner function — identifies the specific pattern
|
||||
@R.function(private=True)
|
||||
def composite(args...):
|
||||
R.func_attr({"Composite": "cublas.matmul_bias_relu"})
|
||||
lv0 = R.matmul(x, w)
|
||||
lv1 = R.add(lv0, bias)
|
||||
lv2 = R.nn.relu(lv1)
|
||||
return lv2
|
||||
...
|
||||
|
||||
The outer function carries the ``Codegen`` attribute that ``RunCodegen`` uses to dispatch to the
|
||||
right backend. The inner function carries the ``Composite`` attribute that the backend codegen
|
||||
uses to identify which operation to emit.
|
||||
|
||||
MergeCompositeFunctions
|
||||
~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
When ``annotate_codegen=False``, ``FuseOpsByPattern`` only creates inner functions with
|
||||
``Composite`` attributes. A separate ``MergeCompositeFunctions`` pass then groups multiple
|
||||
composite functions targeting the same backend into a single outer function with ``Codegen``
|
||||
and ``global_symbol`` attributes.
|
||||
|
||||
This is useful when multiple sequential operations should be sent to the same backend as a
|
||||
single unit (e.g., a sequence of cuBLAS matmuls that share intermediate results). The pass
|
||||
checks that merging does not create cyclic dependencies between groups.
|
||||
|
||||
|
||||
Code Generation
|
||||
---------------
|
||||
|
||||
``RunCodegen`` (``src/relax/transform/run_codegen.cc``) is the pass that triggers backend
|
||||
code generation:
|
||||
|
||||
1. Scan the module for all functions with a ``Codegen`` attribute.
|
||||
2. Group them by backend target name.
|
||||
3. For each backend, look up the registered codegen function via FFI key
|
||||
``"relax.ext.<backend>"`` (e.g., ``"relax.ext.cublas"``).
|
||||
4. Call the codegen function, which returns an array of compiled ``runtime.Module``\ s.
|
||||
5. Replace the original function calls with ``call_dps_packed(ExternFunc(...), args)``.
|
||||
6. Attach the compiled modules to the ``IRModule`` as the ``external_mods`` attribute.
|
||||
|
||||
Codegen registration
|
||||
~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Each backend registers a codegen function via TVM's FFI mechanism:
|
||||
|
||||
.. code-block:: cpp
|
||||
|
||||
// src/relax/backend/contrib/cublas/codegen.cc
|
||||
ffi::Array<ffi::Module> CublasCompiler(
|
||||
ffi::Array<Function> functions,
|
||||
ffi::Map<ffi::String, ffi::Any> options,
|
||||
ffi::Map<Constant, ffi::String> constant_names) {
|
||||
ffi::Array<ffi::Module> compiled_functions;
|
||||
for (const auto& func : functions) {
|
||||
CublasJSONSerializer serializer(constant_names, AnalyzeVar2Value(func));
|
||||
serializer.serialize(func);
|
||||
auto graph_json = serializer.GetJSON();
|
||||
auto names = serializer.GetConstantNames();
|
||||
const auto pf = ffi::Function::GetGlobalRequired("runtime.CublasJSONRuntimeCreate");
|
||||
compiled_functions.push_back(
|
||||
pf(GetExtSymbol(func), graph_json, names).cast<ffi::Module>());
|
||||
}
|
||||
return compiled_functions;
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def("relax.ext.cublas", CublasCompiler);
|
||||
}
|
||||
|
||||
The codegen function receives:
|
||||
|
||||
- ``functions``: the Relax functions with ``Codegen`` attribute to compile.
|
||||
- ``options``: backend-specific compilation options.
|
||||
- ``constant_names``: mapping from constant values to their names (for weight handling).
|
||||
|
||||
It returns an array of ``runtime.Module`` objects — one per function — that contain the
|
||||
externally compiled code.
|
||||
|
||||
Codegen strategies
|
||||
~~~~~~~~~~~~~~~~~~
|
||||
|
||||
TVM provides two base classes for implementing backend codegens:
|
||||
|
||||
- **JSONSerializer** (``src/relax/backend/contrib/codegen_json/codegen_json.h``): serializes the
|
||||
composite function into a JSON graph representation. At runtime, a backend-specific JSON
|
||||
runtime module interprets the graph and dispatches to library calls. Used by cuBLAS, cuDNN,
|
||||
and most backends.
|
||||
|
||||
- **CSourceCodegen** (``src/relax/backend/contrib/codegen_c/codegen_c.h``): generates C/CUDA
|
||||
source code that is compiled and linked. Used when the backend requires ahead-of-time
|
||||
compilation.
|
||||
|
||||
|
||||
Runtime Execution
|
||||
-----------------
|
||||
|
||||
After ``RunCodegen``, the original high-level function calls are replaced with:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
R.call_dps_packed(ExternFunc("fused_relax_matmul_cublas0"), (x, w, bias), ...)
|
||||
|
||||
At runtime, ``call_dps_packed`` invokes the externally compiled function through the
|
||||
``PackedFunc`` interface. The external ``runtime.Module``\ s (produced by the codegen) are
|
||||
imported into the final executable during ``relax.build()`` and are available via the module's
|
||||
function lookup mechanism.
|
||||
|
||||
For JSON-based backends (cuBLAS, cuDNN), the runtime module deserializes the JSON graph and
|
||||
dispatches each node to the corresponding library API call. For source-based backends, the
|
||||
compiled native code is called directly.
|
||||
|
||||
|
||||
Adding a New Backend
|
||||
--------------------
|
||||
|
||||
To add support for a new external library:
|
||||
|
||||
1. **Define patterns** in ``python/tvm/relax/backend/<target>/``:
|
||||
|
||||
- Create DFPatterns using templates from ``patterns.py`` or custom patterns.
|
||||
- Write check functions to validate dtypes, shapes, and other constraints.
|
||||
- Register patterns with ``register_patterns()``.
|
||||
- Provide a ``partition_for_<backend>(mod)`` convenience function.
|
||||
|
||||
2. **Implement codegen** in ``src/relax/backend/contrib/<target>/``:
|
||||
|
||||
- Subclass ``JSONSerializer`` or ``CSourceCodegen``.
|
||||
- Implement the visitor that converts composite functions to the target format.
|
||||
- Register the codegen function as ``"relax.ext.<target>"``.
|
||||
|
||||
3. **Implement runtime** (for JSON-based backends):
|
||||
|
||||
- Create a JSON runtime module that interprets the serialized graph and dispatches
|
||||
to the library's API calls.
|
||||
- Register the runtime constructor as ``"runtime.<Target>JSONRuntimeCreate"``.
|
||||
|
||||
|
||||
Supported Backends
|
||||
------------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 15 25 60
|
||||
|
||||
* - Backend
|
||||
- Patterns
|
||||
- Operations
|
||||
* - cuBLAS
|
||||
- ``cublas.*``
|
||||
- Matmul (with bias, activation, transpose, dequantize variants)
|
||||
* - CUTLASS
|
||||
- ``cutlass.*``
|
||||
- Matmul, conv2d, attention, residual blocks, decode matmul
|
||||
* - cuDNN
|
||||
- ``cudnn.*``
|
||||
- Conv2d (NHWC/NCHW), stacked attention
|
||||
* - DNNL
|
||||
- ``dnnl.*``
|
||||
- Matmul, conv2d (x86 CPU). Codegen exists at C++ level; patterns are
|
||||
defined in tests rather than pre-registered.
|
||||
|
||||
|
||||
Source Code Map
|
||||
---------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 50 50
|
||||
|
||||
* - Path
|
||||
- Contents
|
||||
* - ``python/tvm/relax/backend/pattern_registry.py``
|
||||
- Pattern registry API (register_patterns, get_patterns_with_prefix)
|
||||
* - ``python/tvm/relax/backend/patterns.py``
|
||||
- Reusable pattern templates (make_matmul_pattern, etc.)
|
||||
* - ``python/tvm/relax/backend/cuda/cublas.py``
|
||||
- cuBLAS patterns and partition_for_cublas
|
||||
* - ``python/tvm/relax/backend/cuda/cutlass.py``
|
||||
- CUTLASS patterns and partition_for_cutlass
|
||||
* - ``python/tvm/relax/backend/cuda/cudnn.py``
|
||||
- cuDNN patterns and partition_for_cudnn
|
||||
* - ``src/relax/backend/pattern_registry.cc``
|
||||
- Pattern registry C++ implementation
|
||||
* - ``src/relax/transform/run_codegen.cc``
|
||||
- RunCodegen pass (CodeGenRunner)
|
||||
* - ``src/relax/transform/merge_composite_functions.cc``
|
||||
- MergeCompositeFunctions pass
|
||||
* - ``src/relax/backend/contrib/cublas/codegen.cc``
|
||||
- cuBLAS codegen (JSONSerializer-based)
|
||||
* - ``src/relax/backend/contrib/cutlass/codegen.cc``
|
||||
- CUTLASS codegen
|
||||
* - ``src/relax/backend/contrib/codegen_json/codegen_json.h``
|
||||
- JSONSerializer base class
|
||||
* - ``src/relax/backend/contrib/codegen_c/codegen_c.h``
|
||||
- CSourceCodegen base class
|
||||
@@ -0,0 +1,388 @@
|
||||
.. 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.
|
||||
|
||||
.. _fusion-arch:
|
||||
|
||||
Operator Fusion
|
||||
===============
|
||||
|
||||
Operator fusion is one of the most impactful optimizations in TVM. Instead of launching one kernel
|
||||
per operator (e.g., conv2d, bias_add, relu), fusion merges multiple operators into a single kernel,
|
||||
eliminating intermediate memory allocations and kernel launch overhead.
|
||||
|
||||
TVM provides two complementary fusion mechanisms:
|
||||
|
||||
- **Automatic fusion** (``FuseOps`` + ``FuseTIR``): groups operators based on their computational
|
||||
patterns using a post-dominator analysis algorithm.
|
||||
- **Pattern-based fusion** (``FuseOpsByPattern``): groups operators that match user-defined
|
||||
dataflow patterns, typically for offloading to external backends (cuBLAS, CUTLASS, DNNL, etc.).
|
||||
|
||||
Both produce the same output: Relax functions marked with ``Primitive=True`` that are later
|
||||
lowered to fused TIR kernels or dispatched to external libraries.
|
||||
|
||||
Overview
|
||||
--------
|
||||
|
||||
Fusion involves three passes:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
IRModule (after LegalizeOps)
|
||||
│
|
||||
▼ AnnotateTIROpPattern ← label each op (elementwise, reduce, etc.)
|
||||
IRModule (annotated)
|
||||
│
|
||||
▼ FuseOps ← group ops into fused Relax functions
|
||||
IRModule (with fused functions marked Primitive=True)
|
||||
│
|
||||
▼ FuseTIR ← merge TIR PrimFuncs inside each group
|
||||
IRModule (fused TIR kernels)
|
||||
|
||||
In the compilation pipeline, these passes appear in the backend-specific ``legalize_passes``
|
||||
phase. For example, the CUDA pipeline (``python/tvm/relax/backend/cuda/pipeline.py``) runs:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
LegalizeOps() # lower Relax ops to call_tir
|
||||
AnnotateTIROpPattern() # annotate pattern kinds
|
||||
FoldConstant()
|
||||
FuseOps() # group ops
|
||||
FuseTIR() # merge TIR functions
|
||||
|
||||
|
||||
Operator Pattern Classification
|
||||
-------------------------------
|
||||
|
||||
Before fusion, ``AnnotateTIROpPattern`` analyzes each TIR function in the module and assigns
|
||||
an ``OpPatternKind``. The fusion algorithm uses these pattern kinds to decide which operators
|
||||
can be fused together.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 20 10 70
|
||||
|
||||
* - Pattern Kind
|
||||
- Value
|
||||
- Description
|
||||
* - ``kElemWise``
|
||||
- 0
|
||||
- Elementwise: one-to-one input/output mapping (e.g., ``add``, ``relu``, ``exp``).
|
||||
* - ``kBroadcast``
|
||||
- 1
|
||||
- Broadcasting: output axes map to input axes in order, but some input axes may be
|
||||
broadcast (e.g., ``bias_add``). Note: ``transpose`` is **not** broadcast because axes
|
||||
are reordered.
|
||||
* - ``kInjective``
|
||||
- 2
|
||||
- Injective: each output element depends on a single input element, but the mapping may
|
||||
be non-trivial (e.g., ``reshape``, ``concatenate``, ``transpose``).
|
||||
* - ``kCommReduce``
|
||||
- 3
|
||||
- Communicative reduction: output elements aggregate over input elements
|
||||
(e.g., ``sum``, ``max``, ``mean``).
|
||||
* - ``kOutEWiseFusable``
|
||||
- 4
|
||||
- Complex operation whose output can accept elementwise followers, but cannot chain
|
||||
with another complex op (e.g., ``conv2d``, ``matmul``, ``dense``).
|
||||
* - ``kTuple``
|
||||
- 7
|
||||
- Tuple node. Can fuse into subsequent injective ops but is treated specially.
|
||||
* - ``kOpaque``
|
||||
- 8
|
||||
- Opaque: cannot be fused (e.g., external function calls, operations with side effects).
|
||||
|
||||
These kinds form an ordering: lower values are "simpler" and more fusable. The fusion algorithm
|
||||
uses ``CombinePattern(lhs, rhs) = max(lhs, rhs)`` when merging patterns along a path.
|
||||
|
||||
|
||||
FuseOps: Automatic Fusion
|
||||
-------------------------
|
||||
|
||||
``FuseOps`` (``src/relax/transform/fuse_ops.cc``) groups bindings in a dataflow block into
|
||||
new Relax functions. It operates only within ``DataflowBlock``\ s — if your module doesn't have
|
||||
any, run ``ConvertToDataflow`` first.
|
||||
|
||||
Algorithm
|
||||
~~~~~~~~~
|
||||
|
||||
The fusion algorithm addresses diamond-shaped dataflow branches, where a single producer
|
||||
(e.g., conv2d) has multiple consumers that eventually reconverge:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
conv2d
|
||||
/ | \
|
||||
/ | \
|
||||
op op op
|
||||
\ | /
|
||||
\ | /
|
||||
elemwise add
|
||||
|
||||
At the point of ``conv2d``, we don't know if all future paths will merge. The algorithm uses
|
||||
**post-dominator analysis** to resolve this:
|
||||
|
||||
1. **Build forward graph**: construct an ``IndexedForwardGraph`` from the dataflow block.
|
||||
Each node has an ``OpPatternKind`` and a list of forward edges.
|
||||
|
||||
2. **Build post-dominator tree**: compute the immediate post-dominator of each node using
|
||||
Least Common Ancestor (LCA) on the DAG. The post-dominator of a node is the closest
|
||||
downstream node where **all** future paths converge.
|
||||
|
||||
3. **Fuse groups**: for each node in topological order, check if it can be fused with its
|
||||
immediate post-dominator:
|
||||
|
||||
- **CheckPath**: verify that all paths from the node to its post-dominator satisfy the
|
||||
fusion conditions (pattern compatibility, depth limits, argument limits).
|
||||
- **CommitFuse**: mark all intermediate nodes as belonging to the same group using a
|
||||
Union-Find data structure.
|
||||
|
||||
4. **Create grouped functions**: extract each group into a new ``relax.Function`` with the
|
||||
attribute ``Primitive=True``. Replace the original bindings with a call to the grouped
|
||||
function.
|
||||
|
||||
Fusion rules
|
||||
~~~~~~~~~~~~
|
||||
|
||||
The key fusion decisions depend on the ``OpPatternKind`` of the source, the path, and the
|
||||
post-dominator. The algorithm runs in three phases (via ``GraphPartitioner::RunFuse``) so that
|
||||
higher-complexity ops get a chance to fuse first:
|
||||
|
||||
- **Phase 0**: ``kOutEWiseFusable`` ops (e.g., ``conv2d``) can fuse with their elementwise
|
||||
post-dominator if all intermediate ops are broadcast or simpler. This enables patterns like
|
||||
conv2d + bias_add + relu. Two ``kOutEWiseFusable`` ops cannot fuse together.
|
||||
- **Phase 1**: ``kInjective`` and ``kTuple`` ops can fuse only when all paths to the
|
||||
post-dominator are injective or simpler. This is deferred to phase 1 so that
|
||||
``kOutEWiseFusable`` groups are finalized first.
|
||||
- **Phase 2**: fuse injective ops into intermediate tuple nodes that have already been absorbed
|
||||
by subsequent injective groups.
|
||||
|
||||
``kElemWise`` / ``kBroadcast`` ops are processed in **every** phase (not restricted to one):
|
||||
they can fuse into a post-dominator that is injective or reduction. The sink (final node) may
|
||||
also be a ``kOutEWiseFusable`` group that was formed in phase 0 — this is how elementwise
|
||||
producers merge into an existing conv2d fusion group.
|
||||
|
||||
Additional constraints:
|
||||
|
||||
- **Reduction** (``kCommReduce``) ops never initiate fusion — they act as sinks only. Elementwise
|
||||
and broadcast producers can fuse *into* a reduction, but a reduction cannot fuse forward.
|
||||
- **Opaque** ops are fusion barriers.
|
||||
- A group cannot exceed ``kMaxFusedOps`` (256) nodes or the maximum function argument count.
|
||||
|
||||
Example
|
||||
~~~~~~~
|
||||
|
||||
Given two elementwise ops (``add``, ``exp``) and one injective op (``squeeze``).
|
||||
The examples below are simplified pseudocode — real TVMScript would reference TIR functions
|
||||
via ``cls.func_name``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Before FuseOps (simplified)
|
||||
@R.function
|
||||
def main(x: R.Tensor((10, 20), "float32")):
|
||||
with R.dataflow():
|
||||
lv0 = R.call_tir(add, (x, const_1), out_ty=R.Tensor((10, 20), "float32"))
|
||||
lv1 = R.call_tir(exp, (lv0,), out_ty=R.Tensor((10, 20), "float32"))
|
||||
gv = R.call_tir(squeeze, (lv1,), out_ty=R.Tensor((10, 20), "float32"))
|
||||
R.output(gv)
|
||||
return gv
|
||||
|
||||
After ``FuseOps``, all three are grouped into a single function:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# After FuseOps
|
||||
@R.function(private=True)
|
||||
def fused_add_exp_squeeze(x, p0):
|
||||
R.func_attr({"Primitive": True})
|
||||
with R.dataflow():
|
||||
lv0 = R.call_tir(add, (x, p0), ...)
|
||||
lv1 = R.call_tir(exp, (lv0,), ...)
|
||||
gv = R.call_tir(squeeze, (lv1,), ...)
|
||||
R.output(gv)
|
||||
return gv
|
||||
|
||||
@R.function
|
||||
def main(x: R.Tensor((10, 20), "float32")):
|
||||
with R.dataflow():
|
||||
gv = fused_add_exp_squeeze(x, const_1)
|
||||
R.output(gv)
|
||||
return gv
|
||||
|
||||
|
||||
FuseTIR: Merging TIR Functions
|
||||
------------------------------
|
||||
|
||||
``FuseTIR`` (``src/relax/transform/fuse_tir.cc``) takes the grouped Relax functions produced by
|
||||
``FuseOps`` and merges their internal TIR ``PrimFunc``\ s into a single TIR function.
|
||||
|
||||
Before ``FuseTIR``, a fused group still contains multiple ``R.call_tir`` calls to separate
|
||||
TIR functions. ``FuseTIR`` inlines and merges them:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
Before FuseTIR:
|
||||
fused_add_exp_squeeze:
|
||||
call_tir(add, ...) → separate TIR PrimFunc
|
||||
call_tir(exp, ...) → separate TIR PrimFunc
|
||||
call_tir(squeeze, ...) → separate TIR PrimFunc
|
||||
|
||||
After FuseTIR:
|
||||
fused_add_exp_squeeze: → single merged TIR PrimFunc
|
||||
|
||||
The merged function eliminates intermediate buffers — the output of ``add`` is directly consumed
|
||||
by ``exp`` without writing to and reading from global memory. This is the core performance benefit
|
||||
of fusion.
|
||||
|
||||
Internally, ``FuseTIR`` uses a ``SymbolicMatcher`` to align symbolic shape variables across the
|
||||
TIR functions being merged, ensuring that dimensions are correctly mapped when combining buffer
|
||||
accesses.
|
||||
|
||||
|
||||
FuseOpsByPattern: Pattern-Based Fusion
|
||||
--------------------------------------
|
||||
|
||||
While ``FuseOps`` makes fusion decisions automatically based on operator patterns,
|
||||
``FuseOpsByPattern`` lets you specify exactly which operator combinations to fuse using
|
||||
the Relax :ref:`Dataflow Pattern Language (DPL) <relax-dpl>`.
|
||||
|
||||
This is primarily used for **backend-specific dispatch**: identifying operator subgraphs that
|
||||
should be offloaded to external libraries like cuBLAS, CUTLASS, cuDNN, or DNNL.
|
||||
|
||||
FusionPattern
|
||||
~~~~~~~~~~~~~
|
||||
|
||||
A ``FusionPattern`` (``python/tvm/relax/transform/transform.py``) defines what to match:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.relax.dpl import wildcard, is_op
|
||||
from tvm.relax.transform import FusionPattern
|
||||
|
||||
# Match: matmul(x, w) + bias
|
||||
x = wildcard()
|
||||
w = wildcard()
|
||||
bias = wildcard()
|
||||
matmul = is_op("relax.matmul")(x, w)
|
||||
out = is_op("relax.add")(matmul, bias)
|
||||
|
||||
pattern = FusionPattern(
|
||||
name="cutlass.matmul_bias",
|
||||
pattern=out,
|
||||
annotation_patterns={"matmul": matmul, "bias": bias},
|
||||
check=my_check_function, # optional validation
|
||||
)
|
||||
|
||||
Fields:
|
||||
|
||||
- ``name``: pattern identifier, typically prefixed with the backend name (e.g.,
|
||||
``"cutlass.matmul_bias"``).
|
||||
- ``pattern``: a DFPattern describing the subgraph to match. See the
|
||||
:ref:`DPL deep dive <relax-dpl>` for the full pattern language.
|
||||
- ``annotation_patterns``: a mapping of names to sub-patterns within the main pattern. These
|
||||
are extracted during matching and made available to the ``check`` function and
|
||||
``attrs_getter``.
|
||||
- ``check``: an optional ``Callable[[PatternCheckContext], bool]`` that validates whether
|
||||
a match should be accepted. Receives the matched expression, annotated sub-expressions,
|
||||
variable usages, and binding information.
|
||||
- ``attrs_getter``: an optional function that extracts attributes (e.g., transpose flags,
|
||||
data types) from the matched expressions to annotate the grouped function.
|
||||
|
||||
Applying patterns
|
||||
~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.relax.transform import FuseOpsByPattern
|
||||
|
||||
mod = FuseOpsByPattern(
|
||||
patterns=[pattern1, pattern2, ...], # ordered by priority
|
||||
bind_constants=True,
|
||||
annotate_codegen=False,
|
||||
)(mod)
|
||||
|
||||
Key parameters:
|
||||
|
||||
- ``patterns``: a list of ``FusionPattern`` objects, ordered by priority. Higher-priority
|
||||
patterns come first — if a subgraph matches multiple patterns, the first match wins.
|
||||
- ``bind_constants``: if ``True``, constants used by the matched subgraph are captured inside
|
||||
the grouped function.
|
||||
- ``annotate_codegen``: if ``True``, wraps each composite function with an outer function
|
||||
annotated with ``"Codegen"`` and ``"global_symbol"`` attributes for external backend dispatch.
|
||||
The ``"Codegen"`` value is derived from the pattern name prefix (e.g., ``"dnnl"`` from
|
||||
``"dnnl.conv2d_relu"``).
|
||||
|
||||
PatternCheckContext
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
The ``check`` function receives a ``PatternCheckContext`` with:
|
||||
|
||||
- ``matched_expr``: the root expression matched by the pattern.
|
||||
- ``annotated_expr``: a mapping from annotation pattern names to their matched expressions.
|
||||
- ``matched_bindings``: variable-to-value bindings within the matched subgraph.
|
||||
- ``var_usages``: a mapping from variable definitions to all their uses in the function.
|
||||
- ``value_to_bound_var``: reverse mapping from values to the variables they are bound to.
|
||||
|
||||
This context enables sophisticated validation logic, such as checking that an intermediate
|
||||
result is not used outside the fused group, or verifying data type compatibility.
|
||||
|
||||
|
||||
How Backends Use Fusion
|
||||
-----------------------
|
||||
|
||||
The default backend pipelines (CUDA, ROCm, CPU, etc.) all include ``FuseOps`` + ``FuseTIR``
|
||||
in their ``legalize_passes`` phase for automatic fusion, as shown in the `Overview`_ above.
|
||||
|
||||
For external library dispatch (cuBLAS, CUTLASS, cuDNN, DNNL), ``FuseOpsByPattern`` is used
|
||||
separately. These are **not** included in the default pipeline — users add them explicitly
|
||||
when building a custom compilation flow. The typical sequence is:
|
||||
|
||||
1. **Pattern-based dispatch** (``FuseOpsByPattern``): identify subgraphs that should be
|
||||
offloaded to external libraries. For example, CUTLASS patterns match
|
||||
matmul+bias+activation combinations (``python/tvm/relax/backend/cuda/cutlass.py``).
|
||||
Functions marked by patterns are annotated with ``Composite`` and optionally ``Codegen``
|
||||
attributes. See :ref:`external-library-dispatch` for the full BYOC pipeline.
|
||||
|
||||
2. **Automatic fusion** (``FuseOps`` + ``FuseTIR``): remaining operators that were not
|
||||
matched by backend patterns are fused automatically based on their pattern kinds.
|
||||
|
||||
|
||||
Source Code Map
|
||||
---------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 50 50
|
||||
|
||||
* - Path
|
||||
- Contents
|
||||
* - ``src/relax/transform/fuse_ops.cc``
|
||||
- FuseOps and FuseOpsByPattern implementation
|
||||
* - ``src/relax/analysis/graph_partitioner.h``
|
||||
- IndexedForwardGraph, DominatorTree, GraphPartitioner (Union-Find)
|
||||
* - ``src/relax/transform/fuse_tir.cc``
|
||||
- FuseTIR implementation, SymbolicMatcher
|
||||
* - ``include/tvm/relax/op_attr_types.h``
|
||||
- ``OpPatternKind`` enum definition
|
||||
* - ``python/tvm/relax/transform/transform.py``
|
||||
- Python API: FuseOps, FuseTIR, FuseOpsByPattern, FusionPattern
|
||||
* - ``python/tvm/relax/dpl/``
|
||||
- Dataflow Pattern Language (DFPattern, is_op, wildcard, etc.)
|
||||
* - ``python/tvm/relax/backend/cuda/cutlass.py``
|
||||
- Example: CUTLASS fusion patterns
|
||||
* - ``python/tvm/relax/backend/cuda/cublas.py``
|
||||
- Example: cuBLAS fusion patterns
|
||||
@@ -0,0 +1,436 @@
|
||||
.. 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.
|
||||
|
||||
Design and Architecture
|
||||
=======================
|
||||
|
||||
This document is intended for developers who want to understand the architecture of Apache TVM and/or actively develop on the project.
|
||||
This page is organized as follows:
|
||||
|
||||
- The `Overall Flow`_ gives an overview of the steps that TVM takes to turn a high level description of a model into a deployable module.
|
||||
To get started, please read this section first.
|
||||
- Brief introduction to the key components of the TVM stack. Feel free to also check out the :ref:`TensorIR Deep Dive <tensor-ir-deep-dive>`
|
||||
and :ref:`Relax Deep Dive <relax-deep-dive>` for more details about the two major components in the TVM stack.
|
||||
|
||||
This guide provides a few complementary views of the architecture.
|
||||
First, we review a single end-to-end compilation flow and discuss the key data structures and the transformations.
|
||||
This runtime-based view focuses on the interactions of each components when running the compiler.
|
||||
Then we will review the logical modules of the codebase and their relationship. This part provides a static overarching view of the design.
|
||||
|
||||
Overall Flow
|
||||
------------
|
||||
|
||||
In this guide, we will study an example compilation flow in the compiler. The figure below shows the flow. At a high-level, it contains several steps:
|
||||
|
||||
- **Model Creation**: Create the IRModule to be optimized and compiled, which contains a collection of functions that internally represent the model.
|
||||
Users can manually construct IRModule via NNModule, TVMScript, or import a pre-trained model from Relax frontend.
|
||||
- **Transformation**: The compiler transforms an IRModule to another functionally equivalent or approximately
|
||||
equivalent(e.g. in the case of quantization) IRModule. Many of the transformations are target (backend) independent.
|
||||
We also allow target to affect the configuration of the transformation pipeline.
|
||||
- **Target Translation**: The compiler translates(codegen) the IRModule to an executable format specified by the target.
|
||||
The target translation result is encapsulated as a `runtime.Module` that can be exported, loaded, and executed on the target runtime environment.
|
||||
- **Runtime Execution**: the user loads back a `runtime.Module` and runs the compiled functions in the supported runtime environment.
|
||||
|
||||
|
||||
.. figure:: https://raw.githubusercontent.com/tlc-pack/web-data/main/images/design/tvm_overall_flow.svg
|
||||
:align: center
|
||||
:width: 80%
|
||||
|
||||
|
||||
Key data structures
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
One of the best ways to design and understand a complex system is to identify the key data structures and APIs that
|
||||
manipulate (transform) these data structures. Once we identified the key data structures, we can then breakdown a system into logical
|
||||
components that either define a collection of key data structures or transformations among the data structures.
|
||||
|
||||
**IRModule** is the primary data structure used across the entire stack. An IRModule (intermediate representation module)
|
||||
contains a collection of functions. Currently, we support two primary variants of functions.
|
||||
|
||||
- **relax::Function** is a high-level functional program representation. A relax.Function represents high-level graph structure,
|
||||
usually corresponds to an end-to-end model or a sub-graph of the overall model. You can view a relax.Function as a computational
|
||||
graph with additional support for control-flow, and complex data structures.
|
||||
- **tirx::PrimFunc** is a low-level program representation that contains elements including loop-nest choices, multi-dimensional load/store,
|
||||
threading, and vector/tensor instructions. It is usually used to represent an operator program that executes a (possibly-fused) layer in a model.
|
||||
|
||||
During the compilation and transformation, all relax operators are lowered to ``tirx::PrimFunc`` or ``TVM PackedFunc``, which can be executed directly
|
||||
on the target device, while the calls to relax operators are lowered to calls to low-level functions (e.g. ``R.call_tir`` or ``R.call_dps_packed``).
|
||||
|
||||
Transformations
|
||||
~~~~~~~~~~~~~~~
|
||||
|
||||
Now that we have covered the key data structures, let us talk about the transformations. Each transformation could serve one of the following purposes:
|
||||
|
||||
- optimization: transform a program to an equivalent, possibly more optimized version.
|
||||
- lowering: transform a program to a lower-level representation that is closer to the target.
|
||||
|
||||
relax transformations
|
||||
^^^^^^^^^^^^^^^^^^^^^
|
||||
relax transformations contain a collection of passes that apply to relax functions. The optimizations include common graph-level
|
||||
optimizations such as constant folding and dead-code elimination for operators, and backend-specific optimizations such as library dispatch.
|
||||
|
||||
TensorIR transformations
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
- **TensorIR schedule**: TensorIR schedules are designed to optimize the TensorIR functions for a specific target, with user-guided instructions and control how the target code is generated.
|
||||
For CPU targets, a TensorIR PrimFunc can generate valid code and execute on the target device without schedule but with very-low performance. However, for GPU targets, the schedule is essential
|
||||
for generating valid code with thread bindings. For more details, please refer to the :ref:`TensorIR Transformation <tirx-transform>` section. Additionally, we provides ``MetaSchedule`` to
|
||||
automate the search of TensorIR schedule.
|
||||
- **Lowering Passes**: These passes usually perform after the schedule is applied, transforming a TensorIR PrimFunc into another functionally equivalent PrimFunc, but closer to the
|
||||
target-specific representation. For example, there are passes to flatten multi-dimensional access to one-dimensional pointer access, to expand the intrinsics into target-specific ones,
|
||||
and to decorate the function entry to meet the runtime calling convention.
|
||||
|
||||
Many low-level optimizations can be handled in the target phase by the LLVM,
|
||||
CUDA C, and other target compilers. As a result, we leave low-level
|
||||
optimizations such as register allocation to the downstream compilers and only
|
||||
focus on optimizations that are not covered by them.
|
||||
|
||||
cross-level transformations
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
Apache TVM enables cross-level optimization of end-to-end models. As the IRModule includes both Relax and TensorIR functions, the cross-level transformations are designed to mutate
|
||||
the IRModule by applying different transformations to these two types of functions.
|
||||
|
||||
For example, ``relax.LegalizeOps`` pass mutates the IRModule by lowering relax operators, adding corresponding TensorIR PrimFunc into the IRModule, and replacing the relax operators
|
||||
with calls to the lowered TensorIR PrimFunc. Another example is the operator fusion pipeline
|
||||
(``relax.FuseOps`` + ``relax.FuseTIR``), which fuses multiple consecutive tensor operations into a
|
||||
single kernel. See :ref:`fusion-arch` for a detailed explanation of the fusion algorithm, operator
|
||||
pattern classification, and pattern-based fusion for external backends.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
fusion
|
||||
|
||||
Target Translation
|
||||
~~~~~~~~~~~~~~~~~~
|
||||
|
||||
The target translation phase transforms an IRModule to the corresponding target executable format.
|
||||
For backends such as x86 and ARM, we use the LLVM IRBuilder to build in-memory LLVM IR.
|
||||
We can also generate source-level languages such as CUDA C and OpenCL.
|
||||
Finally, we support direct translations of a Relax function (sub-graph) to specific targets via external code generators.
|
||||
See :ref:`codegen-arch` for how TIR functions are compiled to native code through the LLVM and
|
||||
Source codegen families.
|
||||
See :ref:`external-library-dispatch` for the full BYOC (Bring Your Own Codegen) pipeline that
|
||||
offloads operator subgraphs to vendor libraries like cuBLAS, CUTLASS, and cuDNN.
|
||||
It is important that the final code generation phase is as lightweight as possible. Vast majority of transformations
|
||||
and lowering should be performed before the target translation phase.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
codegen
|
||||
external_library_dispatch
|
||||
|
||||
We also provide a Target structure to specify the compilation target.
|
||||
The transformations before the target translation phase can also be affected by the target — for example,
|
||||
a target's vector length would change the vectorization behavior.
|
||||
|
||||
|
||||
Runtime Execution
|
||||
~~~~~~~~~~~~~~~~~
|
||||
|
||||
The main goal of TVM's runtime is to provide a minimal API for loading and executing the compiled artifact in a language of their choice, including Python, C++, Rust, Go, Java, and JavaScript. The code snippet below shows such an example in Python:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import tvm
|
||||
# Example runtime execution program in python, with type annotated
|
||||
mod: tvm.runtime.Module = tvm.runtime.load_module("compiled_artifact.so")
|
||||
arr: tvm.runtime.Tensor = tvm.runtime.tensor([1, 2, 3], device=tvm.cuda(0))
|
||||
fun: tvm_ffi.Function = mod["addone"]
|
||||
fun(arr)
|
||||
print(arr.numpy())
|
||||
|
||||
|
||||
:py:class:`tvm.runtime.Module` encapsulates the result of compilation. A runtime.Module contains a GetFunction method to obtain :py:class:`tvm_ffi.Function` instances by name.
|
||||
|
||||
:py:class:`tvm_ffi.Function` is a type-erased function interface for both the generated functions. A tvm_ffi.Function can take arguments and return values with the
|
||||
following types: POD types(int, float), string, tvm_ffi.Function, runtime.Module, runtime.Tensor, and other sub-classes of runtime.Object.
|
||||
|
||||
:py:class:`tvm.runtime.Module` and :py:class:`tvm_ffi.Function` are powerful mechanisms to modularize the runtime. For example, to get the above `addone` function on CUDA, we can use LLVM to generate the host-side code to compute the launching parameters(e.g. size of the thread groups) and then call into another tvm_ffi.Function from a CUDAModule that is backed by the CUDA driver API. The same mechanism can be used for OpenCL kernels.
|
||||
|
||||
The above example only deals with a simple `addone` function. The code snippet below gives an example of an end-to-end model execution using the Relax Virtual Machine, which is built on the same runtime.Module and tvm_ffi.Function interface:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import tvm
|
||||
from tvm import relax
|
||||
# Load the compiled artifact
|
||||
mod: tvm.runtime.Module = tvm.runtime.load_module("resnet18.so")
|
||||
# Create a VM instance on cuda(0)
|
||||
vm = relax.VirtualMachine(mod, tvm.cuda(0))
|
||||
data: tvm.runtime.Tensor = get_input_data()
|
||||
# Run the model — vm["main"] returns a PackedFunc
|
||||
result = vm["main"](data).numpy()
|
||||
|
||||
The main take away is that runtime.Module and runtime.PackedFunc are sufficient to encapsulate both operator level programs (such as addone), as well as the end-to-end models.
|
||||
|
||||
Summary and Discussions
|
||||
~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
In summary, the key data structures in the compilation flows are:
|
||||
|
||||
- IRModule: contains relax.Function and tirx.PrimFunc
|
||||
- runtime.Module: contains runtime.PackedFunc
|
||||
|
||||
Most parts of the compilation are transformations among the key data structures.
|
||||
|
||||
- relax/transform and tirx/transform are deterministic rule-based transformations
|
||||
- meta-schedule contains the search-based transformations
|
||||
|
||||
Finally, the compilation flow example is only a typical use-case of the TVM stack.
|
||||
We expose these key data structures and transformations to python and C++ APIs. As a result, you can use TVM just like the way you use numpy,
|
||||
except that the data structure of interest changes from the numpy.ndarray to tvm.IRModule. Here are some example use-cases:
|
||||
|
||||
- Directly construct IRModule using the python API.
|
||||
- Compose a custom set of transformations(e.g. customize quantization).
|
||||
- Manipulate the IR directly using TVM's python API.
|
||||
|
||||
|
||||
tvm/support
|
||||
-----------
|
||||
The support module contains the most common utilities for the infrastructure, such as generic arena allocator, socket, and logging.
|
||||
|
||||
|
||||
tvm/runtime
|
||||
-----------
|
||||
|
||||
The runtime serves as the foundation of the TVM stack. It provides the mechanism to load and execute compiled artifacts.
|
||||
The runtime defines a stable standard set of C APIs to interface with frontend languages such as Python and Rust.
|
||||
|
||||
`runtime::Object` is one of the primary data structures in TVM runtime besides the `ffi::Function`.
|
||||
It is a reference-counted base class with a type index to support runtime type checking and downcasting.
|
||||
The object system allows the developer to introduce new data structures to the runtime, such as Array, Map, and new IR data structures.
|
||||
|
||||
Besides deployment use-cases, the compiler itself also makes heavy use of TVM's runtime mechanism.
|
||||
All of the IR data structures are subclasses of `runtime::Object`, as a result, they can be directly accessed and manipulated from the Python frontend.
|
||||
We use the PackedFunc mechanism to expose various APIs to the frontend.
|
||||
|
||||
Runtime support for different hardware backends are defined in subdirectories of runtime(e.g. runtime/opencl).
|
||||
These hardware-specific runtime modules define APIs for device memory allocation and device function serialization.
|
||||
|
||||
`runtime/rpc` implements an RPC support for PackedFunc. We can use the RPC mechanism to send a cross-compiled library to a remote
|
||||
device and benchmark the execution performance. The rpc infrastructure enables data collection from a wide range of hardware backends
|
||||
for learning-based optimizations.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
runtime
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
introduction_to_module_serialization
|
||||
|
||||
Relax Virtual Machine
|
||||
~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Relax defines *what* to compute — it is a graph-level IR that describes the operators and dataflow
|
||||
of a model. The Relax Virtual Machine (VM) handles *how* to run it — it is the runtime component
|
||||
that executes the compiled result. The VM uses a register-based interpreter with only four opcodes
|
||||
(``Call``, ``Ret``, ``Goto``, ``If``) and performs no mathematical computation itself — it
|
||||
orchestrates control flow while dispatching actual work to compiled TIR kernels or external
|
||||
libraries.
|
||||
|
||||
See :ref:`relax-vm-arch` for the full architecture documentation, including the compilation
|
||||
pipeline, instruction set details, execution model, and Python interface.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
relax_vm
|
||||
|
||||
Disco: Distributed Runtime
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Disco is TVM's distributed runtime for executing models across multiple devices. When a model is
|
||||
too large to fit on a single GPU, the ``relax.distributed`` module annotates how tensors should be
|
||||
partitioned and placed across a mesh of devices at compile time. Disco then takes over at runtime:
|
||||
it manages a group of workers, dispatches the compiled program to all of them simultaneously, and
|
||||
coordinates inter-device communication through collective operations such as allreduce, allgather,
|
||||
broadcast, and scatter.
|
||||
|
||||
The central abstraction is the ``Session``, which owns the workers and exposes a SPMD-style
|
||||
programming interface. Every object that lives on workers is represented by a ``DRef`` — a
|
||||
distributed reference that maps to a concrete value on each worker. When the controller invokes a
|
||||
``DPackedFunc`` through the session, all workers execute the same PackedFunc call synchronously, each
|
||||
operating on its own local shard. Compiled VM modules can be loaded into a session as ``DModule``
|
||||
objects and called in the same fashion. The session also provides collective primitives backed by
|
||||
NCCL or RCCL, so that workers can exchange partial results without routing data through the
|
||||
controller.
|
||||
|
||||
Three session backends cover different deployment topologies. ``ThreadedSession`` spawns workers as
|
||||
threads within a single process — this is the most common choice for multi-GPU inference on a
|
||||
single machine. ``ProcessSession`` launches workers as separate OS processes connected by pipes,
|
||||
providing stronger isolation. ``SocketSession`` extends the model to multi-node clusters by
|
||||
connecting workers across machines via TCP sockets.
|
||||
|
||||
tvm/node
|
||||
--------
|
||||
The node module adds additional features on top of the `runtime::Object` for IR data structures.
|
||||
The main features include reflection, serialization, structural equivalence, and hashing.
|
||||
|
||||
Thanks to the node module, we can directly access any field of the TVM's IRNode by their name in Python.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
x = tvm.tirx.Var("x", "int32")
|
||||
y = tvm.tirx.Add(x, x)
|
||||
# a and b are fields of a tirx.Add node
|
||||
# we can directly use the field name to access the IR structures
|
||||
assert y.a == x
|
||||
|
||||
We can also serialize arbitrary IR node into a JSON format, and load them back.
|
||||
The ability to save/store, and inspect an IR node provides a foundation for making the compiler more accessible.
|
||||
|
||||
tvm/ir
|
||||
------
|
||||
The `tvm/ir` folder contains the unified data structure and interfaces across all IR function variants.
|
||||
The components in `tvm/ir` are shared by `tvm/relax` and `tvm/tirx`, notable ones include
|
||||
|
||||
- IRModule
|
||||
- Type
|
||||
- PassContext and Pass
|
||||
- Op
|
||||
|
||||
Different variants of functions(e.g. relax.Function and tirx.PrimFunc) can co-exist in an IRModule.
|
||||
While these variants may not have the same content representation, they use the same data structure to represent types.
|
||||
As a consequence, we use the same data structure to represent function (type) signatures of these variants.
|
||||
The unified type system allows one function variant to call another function
|
||||
once we clearly define the calling convention. This opens doors for future cross-function-variant optimizations.
|
||||
|
||||
We also provide a unified PassContext for configuring the pass behavior, and common composite passes to execute a pass pipeline.
|
||||
The following code snippet gives an example of PassContext configuration.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# configure the behavior of the tirx.UnrollLoop pass
|
||||
with tvm.transform.PassContext(config={"tirx.UnrollLoop": { "auto_max_step": 10 }}):
|
||||
# code affected by the pass context
|
||||
|
||||
|
||||
Op is the common class to represent all system-defined primitive operator/intrinsics.
|
||||
Developers can register new Ops as well as their additional attributes(e.g. whether the Op is elementwise) to the system.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
pass_infra
|
||||
|
||||
tvm/script (TVMScript)
|
||||
----------------------
|
||||
|
||||
TVMScript is a Python-based DSL for writing TVM IR. It allows users to define ``IRModule``\ s
|
||||
— containing both Relax functions and TIR ``PrimFunc``\ s — using familiar Python syntax with
|
||||
three import aliases: ``I`` (module-level), ``T`` (TIR), and ``R`` (Relax). Although TVMScript
|
||||
uses Python syntax, it is not executed by the Python interpreter — decorators like
|
||||
``@I.ir_module``, ``@T.prim_func``, and ``@R.function`` extract the Python AST and transform
|
||||
it into TVM IR through a parser and IR builder pipeline.
|
||||
|
||||
TVMScript also supports **roundtrip**: any ``IRModule`` can be printed back to TVMScript via
|
||||
``mod.script()`` and re-parsed to produce a structurally equivalent module. See
|
||||
:ref:`tvmscript-arch` for the full architecture documentation, including the parser dispatch
|
||||
mechanism, IR builder frame stack, printer pipeline, and syntax reference.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
tvmscript
|
||||
|
||||
|
||||
tvm/target
|
||||
----------
|
||||
The target module contains all the code generators that translate an IRModule to a target runtime.Module.
|
||||
It also provides a common `Target` class that describes the target.
|
||||
|
||||
Targets can be constructed from a registered tag, a configuration dictionary, or a tag with attribute overrides:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.target import Target
|
||||
|
||||
# From a registered tag
|
||||
target = Target("nvidia/nvidia-a100")
|
||||
|
||||
# From a config dictionary
|
||||
target = Target({"kind": "cuda", "arch": "sm_80"})
|
||||
|
||||
# From a tag with attribute overrides
|
||||
target = Target({"tag": "nvidia/nvidia-a100", "l2_cache_size_bytes": 12345})
|
||||
|
||||
Use ``Target.list_kinds()`` to see all available target kinds, and ``target.attrs`` to inspect
|
||||
target attributes.
|
||||
|
||||
The compilation pipeline can be customized according to the target by querying the attribute information
|
||||
in the target and builtin information registered to each target id(cuda, opencl).
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
device_target_interactions
|
||||
|
||||
tvm/relax
|
||||
---------
|
||||
|
||||
Relax is the high-level IR used to represent the computational graph of a model. Various optimizations are defined in ``relax.transform``.
|
||||
Note that Relax usually works closely with the TensorIR IRModule, most of the transformations are applied on both Relax and TensorIR functions
|
||||
in the IRModule. Please refer to the :ref:`Relax Deep Dive <relax-deep-dive>` for more details.
|
||||
|
||||
tvm/tirx
|
||||
--------
|
||||
|
||||
``tirx`` contains the core IR definitions and lowering infrastructure
|
||||
for TensorIR (split from the former ``tir`` module). ``tirx::PrimFunc``
|
||||
represents low-level tensor functions that can be transformed by tirx passes.
|
||||
|
||||
The tirx module includes:
|
||||
|
||||
- IR data structures (PrimFunc, Buffer, SBlock, expressions, statements).
|
||||
- Analysis passes in ``tirx/analysis``.
|
||||
- Transformation and lowering passes in ``tirx/transform``.
|
||||
|
||||
tvm/s_tir
|
||||
---------
|
||||
|
||||
``s_tir`` (Schedulable TIR, split from the former ``tir`` module) contains
|
||||
schedule primitives and auto-tuning tools that operate on ``tirx::PrimFunc``:
|
||||
|
||||
- Schedule primitives to control code generation (tiling, vectorization, thread
|
||||
binding) in ``s_tir/schedule``.
|
||||
- Builtin tensor intrinsics in ``s_tir/tensor_intrin``.
|
||||
- MetaSchedule for automated performance tuning.
|
||||
- DLight for pre-defined, high-performance schedules.
|
||||
|
||||
Please refer to the :ref:`TensorIR Deep Dive <tensor-ir-deep-dive>` for more details.
|
||||
|
||||
tvm/arith
|
||||
---------
|
||||
|
||||
This module is closely tied to TensorIR. One of the key problems in the low-level code generation is the analysis of the indices'
|
||||
arithmetic properties — the positiveness, variable bound, and the integer set that describes the iterator space. arith module provides
|
||||
a collection of tools that do (primarily integer) analysis. A TensorIR pass can use these analyses to simplify and optimize the code.
|
||||
|
||||
tvm/te and tvm/topi
|
||||
-------------------
|
||||
|
||||
TE stands for Tensor Expression. TE is a domain-specific language (DSL) for describing tensor computations. Importantly, a tensor expression
|
||||
itself is not a self-contained function that can be stored into IRModule. We can use ``te.create_prim_func`` to convert a tensor expression to a ``tirx::PrimFunc``
|
||||
and then integrate it into the IRModule.
|
||||
|
||||
While possible to construct operators directly via TensorIR or tensor expressions (TE) for each use case, it is tedious to do so.
|
||||
`topi` (Tensor operator inventory) provides a set of pre-defined operators defined by numpy and found in common deep learning workloads.
|
||||
@@ -0,0 +1,193 @@
|
||||
.. 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.
|
||||
|
||||
Introduction to Module Serialization
|
||||
====================================
|
||||
|
||||
When to deploy TVM runtime module, no matter whether it is CPU or GPU, TVM only needs one single dynamic
|
||||
shared library. The key is our unified module serialization mechanism. This document will introduce TVM module
|
||||
serialization format standard and implementation details.
|
||||
|
||||
|
||||
*************
|
||||
Serialization
|
||||
*************
|
||||
|
||||
The entrance API is ``export_library`` of ``tvm.runtime.Module``.
|
||||
Inside this function, we will do the following steps:
|
||||
|
||||
1. Collect all DSO modules (LLVM modules and C modules)
|
||||
|
||||
2. Once we have DSO modules, we will call ``save`` function to save them into files.
|
||||
|
||||
3. Next, we will check whether we have imported modules, such as CUDA,
|
||||
OpenCL or anything else. We don't restrict the module type here.
|
||||
Once we have imported modules, we will create one file named ``devc.o`` / ``dev.cc``
|
||||
(so that we could embed the binary blob data of import modules into one dynamic shared library),
|
||||
then call function ``_PackImportsToLLVM`` or ``_PackImportsToC`` to do module serialization.
|
||||
|
||||
4. Finally, we call ``fcompile`` which invokes ``_cc.create_shared`` to get
|
||||
dynamic shared library.
|
||||
|
||||
.. note::
|
||||
1. For C source modules, we will compile them and link them together with the DSO module.
|
||||
|
||||
2. Use ``_PackImportsToLLVM`` or ``_PackImportsToC`` depends on whether we enable LLVM in TVM.
|
||||
They achieve the same goal in fact.
|
||||
|
||||
***************************************************
|
||||
Under the Hood of Serialization and Format Standard
|
||||
***************************************************
|
||||
|
||||
As said before, we will do the serialization work in the ``_PackImportsToLLVM`` or ``_PackImportsToC``.
|
||||
They both call ``SerializeModule`` to serialize the runtime module. In ``SerializeModule``
|
||||
function, we firstly construct one helper class ``ModuleSerializer``. It will take ``module`` to do some
|
||||
initialization work, like marking module index. Then we could use its ``SerializeModule`` to serialize module.
|
||||
|
||||
For better understanding, let us dig the implementation of this class a little deeper.
|
||||
|
||||
The following code is used to construct ``ModuleSerializer``:
|
||||
|
||||
.. code:: c++
|
||||
|
||||
explicit ModuleSerializer(runtime::Module mod) : mod_(mod) {
|
||||
Init();
|
||||
}
|
||||
private:
|
||||
void Init() {
|
||||
CreateModuleIndex();
|
||||
CreateImportTree();
|
||||
}
|
||||
|
||||
In ``CreateModuleIndex()``, We will inspect module import relationship
|
||||
using DFS and create index for them. Note the root module is fixed at
|
||||
location 0. In our example, we have module relationship like this:
|
||||
|
||||
.. code:: c++
|
||||
|
||||
llvm_mod:imports
|
||||
- cuda_mod
|
||||
|
||||
So LLVM module will have index 0, CUDA module will have index 1.
|
||||
|
||||
After constructing module index, we will try to construct import tree (``CreateImportTree()``),
|
||||
which will be used to restore module import relationship when we load
|
||||
the exported library back. In our design, we use CSR format to store
|
||||
import tree, each row is parent index, the child indices correspond to its children
|
||||
index. In code, we use ``import_tree_row_ptr_`` and
|
||||
``import_tree_child_indices_`` to represent them.
|
||||
|
||||
After initialization, we could serialize module using ``SerializeModule`` function.
|
||||
In its function logic, we will assume the serialization format like this:
|
||||
|
||||
.. code:: c++
|
||||
|
||||
binary_blob_size
|
||||
binary_blob_type_key
|
||||
binary_blob_logic
|
||||
binary_blob_type_key
|
||||
binary_blob_logic
|
||||
...
|
||||
_import_tree
|
||||
_import_tree_logic
|
||||
|
||||
``binary_blob_size`` is the number of blobs we will have in this
|
||||
serialization step. There will be three blobs in our example which
|
||||
are created for LLVM module, CUDA module, and ``_import_tree``, respectively.
|
||||
|
||||
``binary_blob_type_key`` is the blob type key of module. For LLVM / C module, whose
|
||||
blob type key is ``_lib``. For CUDA module, it is ``cuda``, which could be got by ``module->type_key()``.
|
||||
|
||||
``binary_blob_logic`` is the logic handling of blob. For most of blob (like CUDA, OpenCL), we will call
|
||||
``SaveToBinary`` function to serialize blob into binary. However, like LLVM / C module, we will only write
|
||||
``_lib`` to indicate this is a DSO module.
|
||||
|
||||
.. note::
|
||||
Whether or not it is required to implement the SaveToBinary virtual function depends on
|
||||
how the module is used. For example, if the module has information we need when we load
|
||||
the dynamic shared library back, we should do. Like CUDA module, we need its binary data
|
||||
passing to GPU driver when we load the dynamic shared library, so we should implement
|
||||
``SaveToBinary`` to serialize its binary data. But for host module (like DSO), we don't
|
||||
need other information when we load the dynamic shared library, so we don't need to implement
|
||||
``SaveToBinary``. However, if in the future, we want to record some meta information of DSO module,
|
||||
we could implement ``SaveToBinary`` for DSO module too.
|
||||
|
||||
Finally, we will write one key ``_import_tree`` unless our module only
|
||||
has one DSO module and it is in the root. It is used to reconstruct the
|
||||
module import relationship when we load the exported library back as said
|
||||
before. The ``import_tree_logic`` is just to write ``import_tree_row_ptr_``
|
||||
and ``import_tree_child_indices_`` into stream.
|
||||
|
||||
After this step, we will pack it into a symbol
|
||||
``runtime::symbol::tvm_ffi_library_bin`` that can be recovered in the dynamic
|
||||
library.
|
||||
|
||||
Now, we complete the serialization part. As you have seen, we could
|
||||
support arbitrary modules to import ideally.
|
||||
|
||||
****************
|
||||
Deserialization
|
||||
****************
|
||||
|
||||
The entrance API is ``tvm.runtime.load_module``. This function
|
||||
actually calls ``_LoadFromFile``. If we dig it a little deeper, this is
|
||||
``Module::LoadFromFile``. In our example, the file is ``deploy.so``,
|
||||
according to the function logic, we will call ``module.loadfile_so`` in
|
||||
``dso_library.cc``. The key is here:
|
||||
|
||||
.. code:: c++
|
||||
|
||||
// Load the imported modules
|
||||
const char* library_bin = reinterpret_cast<const char*>(
|
||||
lib->GetSymbol(runtime::symbol::tvm_ffi_library_bin));
|
||||
Module root_mod;
|
||||
if (library_bin != nullptr) {
|
||||
root_mod = ProcessLibraryBin(library_bin, lib);
|
||||
} else {
|
||||
// Only have one single DSO Module
|
||||
root_mod = Module(n);
|
||||
}
|
||||
|
||||
As said before, we will pack the blob into the symbol
|
||||
``runtime::symbol::tvm_ffi_library_bin``. During deserialization part, we will
|
||||
inspect it. If we have ``runtime::symbol::tvm_ffi_library_bin``, we will call ``ProcessLibraryBin``,
|
||||
whose logic like this:
|
||||
|
||||
.. code:: c++
|
||||
|
||||
READ(blob_size)
|
||||
READ(blob_type_key)
|
||||
for (size_t i = 0; i < blob_size; i++) {
|
||||
if (blob_type_key == "_lib") {
|
||||
// construct dso module using lib
|
||||
} else if (blob_type_key == "_import_tree") {
|
||||
// READ(_import_tree_row_ptr)
|
||||
// READ(_import_tree_child_indices)
|
||||
} else {
|
||||
// call module.loadbinary_blob_type_key, such as module.loadbinary_cuda
|
||||
// to restore.
|
||||
}
|
||||
}
|
||||
// Using _import_tree_row_ptr and _import_tree_child_indices to
|
||||
// restore module import relationship. The first module is the
|
||||
// root module according to our invariance as said before.
|
||||
return root_module;
|
||||
|
||||
After this, we will set the ``ctx_address`` to be the ``root_module`` so
|
||||
that allow lookup of symbol from root (so all symbols are visible).
|
||||
|
||||
Finally, we complete the deserialization part.
|
||||
@@ -0,0 +1,669 @@
|
||||
.. 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.
|
||||
|
||||
.. _pass-infra:
|
||||
|
||||
Pass Infrastructure
|
||||
===================
|
||||
|
||||
Both Relax and TVM IR contain a series of optimization passes which improve performance metrics
|
||||
of models such as mean inference, memory footprint, or power consumption for
|
||||
specific devices. There is a suite of standard optimizations as well as machine
|
||||
learning-specific optimizations including constant folding, dead code
|
||||
elimination, operator layout alteration, operator fusion, buffer handling, and
|
||||
loop transformation, etc. Each of these passes is structured as a ir-to-ir
|
||||
transformation using the analysis result collected during and/or before traversal.
|
||||
|
||||
However, as TVM evolves quickly, the need for a more systematic and efficient
|
||||
way to manage these passes is becoming apparent. In addition, a generic
|
||||
framework that manages the passes across different layers of the TVM stack (e.g.
|
||||
Relax and TensorIR) paves the way for developers to quickly prototype and plug the
|
||||
implemented passes into the system.
|
||||
|
||||
This doc describes the design of such an infra that takes the advantage of the
|
||||
way production compilers are used to manage the optimization passes and the style
|
||||
modern deep learning frameworks adopted to build up layers.
|
||||
|
||||
For example, many existing production compilers, such as GCC and LLVM, employ
|
||||
pass managers to effectively manage the execution of passes. Initially managing
|
||||
passes is straightforward as the number of passes is small, but mature compilers
|
||||
will contain hundreds of individual passes. Often external users will want to
|
||||
have custom passes correctly scheduled without having to modify a single
|
||||
handcrafted pass order.
|
||||
|
||||
Similarly, modern deep learning frameworks, such as PyTorch, also have
|
||||
the tendency to enable pass-style layer construction scheme through
|
||||
`Sequential`_. With such constructs, these modern frameworks are able to
|
||||
conveniently add modules/layers to their containers and build up neural
|
||||
networks easily.
|
||||
|
||||
The design of the TVM pass infra is largely inspired by the hierarchical
|
||||
pass manager used in LLVM and the block-style containers used in the popular
|
||||
deep learning frameworks. The major goals of the pass infra include:
|
||||
|
||||
#) enabling better programmatic orchestration of optimizations. This allows
|
||||
users to flexibly customize and build their own optimization pipelines.
|
||||
|
||||
#) providing a user-friendly way to debug optimization passes.
|
||||
|
||||
#) alleviating developers from manually and respectively resolving the
|
||||
dependencies between passes.
|
||||
|
||||
#) simplifying the implementation of new passes for developers. For example, we
|
||||
allow users to implement a pass in Python and let the pass infra manipulate
|
||||
its execution.
|
||||
|
||||
The Design
|
||||
----------
|
||||
|
||||
We focus on ease of extension for users, making it possible for users to quickly
|
||||
add new passes without loss of backward compatibility. The design contains both
|
||||
the backend and the frontend. The former implements the main logic of the pass
|
||||
infra. The latter provides simple APIs for users to interact with, i.e.,
|
||||
allowing users to quickly create their own optimization pipelines.
|
||||
|
||||
C++ Backend
|
||||
~~~~~~~~~~~
|
||||
|
||||
We provide a ``PassInfo`` object to contain the basic information needed by
|
||||
a pass. ``name`` is the pass name, ``opt_level`` indicates at which optimization
|
||||
level the pass will be enabled, and ``required`` represents the passes that are
|
||||
required to execute a certain pass (see `include/tvm/ir/transform.h`_ for
|
||||
more details). For example, during registration of a pass (will be covered in
|
||||
later), the pass developers can specify the name of the pass, the optimization
|
||||
level it will be performed at, and/or the passes that are required.
|
||||
``opt_level`` could be used to help the pass infra identify if a certain pass
|
||||
needs to be executed when running under a user-provided optimization level. The
|
||||
``required`` field can be used by the pass infra to resolve pass dependencies.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
class PassInfoNode : public Object {
|
||||
int opt_level;
|
||||
ffi::String name;
|
||||
bool traceable;
|
||||
ffi::Array<ffi::String> required;
|
||||
};
|
||||
|
||||
PassContext
|
||||
^^^^^^^^^^^
|
||||
|
||||
``PassContext`` carries useful information for an optimization pass. For
|
||||
example, it contains the error reporting system so optimization authors can
|
||||
provide diagnostics about why an optimization fails. ``PassContext`` is also
|
||||
designed to replace the old ``BuildConfig`` which was used to help users
|
||||
configure the compilation options, including optimization level and
|
||||
required/disabled passes, etc. For instance, we may have a configuration which
|
||||
performs all passes at ``opt_level=3`` with some disabled passes using
|
||||
``disabled_pass=xx`` provided by ``PassContext``. Now we could glob all passes
|
||||
at ``opt_level=3`` and exclude those in the disabled pass list. ``PassContext``
|
||||
also provides a way to instrument all passes. See section :ref:`pass_instrument_cpp_backend`.
|
||||
|
||||
This class is designed for users to conveniently write the Python ``with``
|
||||
syntax to perform optimizations under a certain configuration. In addition, the
|
||||
users can obtain the context that is available within a certain program scope in
|
||||
a thread-safe way through ``PassContext::Current()``, since a thread-local store
|
||||
``PassContextThreadLocalStore`` is used to hold the created pass context
|
||||
objects. Examples will be provided later to show how we can use both the C++ and
|
||||
Python APIs to create a compilation pipeline using pass context.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
class PassContextNode : public Object {
|
||||
public:
|
||||
int opt_level{2};
|
||||
ffi::Array<ffi::String> required_pass;
|
||||
ffi::Array<ffi::String> disabled_pass;
|
||||
mutable ffi::Optional<DiagnosticContext> diag_ctx;
|
||||
ffi::Map<ffi::String, Any> config;
|
||||
ffi::Array<instrument::PassInstrument> instruments;
|
||||
};
|
||||
|
||||
class PassContext : public ObjectRef {
|
||||
public:
|
||||
TVM_DLL static PassContext Create();
|
||||
TVM_DLL static PassContext Current();
|
||||
TVM_DLL void InstrumentEnterPassContext();
|
||||
TVM_DLL void InstrumentExitPassContext();
|
||||
TVM_DLL bool InstrumentBeforePass(const IRModule& mod, const PassInfo& info) const;
|
||||
TVM_DLL void InstrumentAfterPass(const IRModule& mod, const PassInfo& info) const;
|
||||
/* Other fields are omitted. */
|
||||
|
||||
private:
|
||||
// The entry of a pass context scope.
|
||||
TVM_DLL void EnterWithScope();
|
||||
// The exit of a pass context scope.
|
||||
TVM_DLL void ExitWithScope();
|
||||
|
||||
// Classes to get the Python `with` like syntax.
|
||||
friend class tvm::With<PassContext>;
|
||||
};
|
||||
|
||||
struct PassContextThreadLocalEntry {
|
||||
/*! \brief The default pass context. */
|
||||
PassContext default_context;
|
||||
/*! \brief The current pass context. */
|
||||
std::stack<PassContext> context_stack;
|
||||
PassContextThreadLocalEntry() {
|
||||
default_context = PassContext(ffi::make_object<PassContextNode>());
|
||||
}
|
||||
};
|
||||
|
||||
Pass Constructs
|
||||
^^^^^^^^^^^^^^^
|
||||
|
||||
The pass infra is designed in a hierarchical manner, and it could work at
|
||||
different granularities of Relax/TensorIR programs. A pure virtual class ``PassNode`` is
|
||||
introduced to serve as the base of the different optimization passes. This class
|
||||
contains several virtual methods that must be implemented by the
|
||||
subclasses at the level of modules, functions, or sequences of passes.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
class PassNode : Object {
|
||||
virtual PassInfo Info() const = 0;
|
||||
virtual Module operator()(const IRModule& mod
|
||||
const PassContext& pass_ctx) const = 0;
|
||||
};
|
||||
|
||||
The functor shows how a pass must be realized, i.e. it always works on a
|
||||
:py:class:`IRModule` under a certain context. All passes are designed in a ``Module`` to ``Module``
|
||||
manner. Therefore, optimizations governed by the pass infra will
|
||||
always update the whole module.
|
||||
|
||||
Several subclasses have been created to implement different types of
|
||||
optimization passes, e.g., function-level passes, module-level passes, and
|
||||
sequential passes. Each subclass itself could act as a pass manager. For
|
||||
instance, they could collect the required passes and execute them or build
|
||||
a dependency graph based on the given metadata. The full definition of them
|
||||
can be found in `src/ir/transform.cc`_.
|
||||
|
||||
Module-Level Passes
|
||||
^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Module level passes are geared mainly for global and inter-procedural
|
||||
optimizations (IPO), which are similar to the module pass used in LLVM. Some
|
||||
typical passes in Relax that need the global picture of a module, such as
|
||||
A-normal form conversion and lambda lifting, etc., fall into this set. At this
|
||||
level, users can even add and/or delete functions in a module. Note that all
|
||||
passes
|
||||
|
||||
.. code:: c++
|
||||
|
||||
class ModulePassNode : PassNode {
|
||||
PassInfo pass_info;
|
||||
std::function<Module(Module, PassContext)> pass_func;
|
||||
Module operator()(const Module& mod, const PassContext& pass_ctx) const final;
|
||||
// Other members/methods are omitted
|
||||
};
|
||||
|
||||
``pass_info`` maintains the information needed by a module-level pass.
|
||||
``pass_func`` sketches the real optimization. For example, we may need to
|
||||
perform dead code elimination on the module. We could implement the algorithm in
|
||||
the ``pass_func`` and let it run on a module. It will then remove the dead code
|
||||
including the unused functions in the module. Note that this field is designed
|
||||
as a packed function, which enables the implementation of the optimization in
|
||||
both C++ and Python.
|
||||
|
||||
Function-Level Passes
|
||||
^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Function-level passes are used to implement various intra-function level
|
||||
optimizations for a given Relax/TensorIR module. It fetches one function at a time from
|
||||
the function list of a module for optimization and yields a rewritten Relax
|
||||
``Function`` or TensorIR ``PrimFunc``. Most of passes can be classified into this category, such as
|
||||
common subexpression elimination and inference simplification in Relax as well as vectorization
|
||||
and flattening storage in TensorIR, etc.
|
||||
|
||||
Note that the scope of passes at this level is either a Relax function or a TensorIR primitive function.
|
||||
Therefore, we cannot add or delete a function through these passes as they are not aware of
|
||||
the global information.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
class FunctionPassNode : PassNode {
|
||||
PassInfo pass_info;
|
||||
std::function<Function(Function, Module, PassContext)> pass_func;
|
||||
Module operator()(const Module& mod, const PassContext& pass_ctx) const final;
|
||||
bool SkipFunction(const Function& func) const;
|
||||
// Other members/methods are omitted...
|
||||
};
|
||||
|
||||
``pass_info`` is identical to what we just described in the module pass.
|
||||
``pass_func`` takes a function for optimization, it also needs a module as we
|
||||
may use it for reporting errors. A function could be annotated with
|
||||
"SkipOptimization" so that it will be ignored during optimization.
|
||||
|
||||
Sequential Passes
|
||||
^^^^^^^^^^^^^^^^^
|
||||
|
||||
``SequentialPass`` is similar to Pytorch ``nn.Sequential`` that contains a host
|
||||
of passes for execution.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
class SequentialPassNode : PassNode {
|
||||
PassInfo pass_info;
|
||||
// Passes need to be executed.
|
||||
ffi::Array<Pass> passes;
|
||||
bool PassEnabled(const PassInfo& info) const;
|
||||
Module operator()(const Module& mod, const PassContext& pass_ctx) const final;
|
||||
};
|
||||
|
||||
The following code shows how individual passes in a sequential pass are invoked.
|
||||
Essentially, we sequentially execute each pass in a sequential pass using the
|
||||
order that they were appended to the pass list.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
Module SequentialNode::operator()(const Module& module,
|
||||
const PassContext& pass_ctx) const {
|
||||
Module mod = module;
|
||||
for (const Pass& pass : passes) {
|
||||
TVM_FFI_ICHECK(pass.defined()) << "Found undefined pass for optimization.";
|
||||
const PassInfo& pass_info = pass->Info();
|
||||
if (!PassEnabled(pass_info)) continue;
|
||||
for (const auto& it : pass_info->required) {
|
||||
mod = GetPass(it)(std::move(mod), pass_ctx);
|
||||
}
|
||||
mod = pass(mod, pass_ctx);
|
||||
}
|
||||
return mod;
|
||||
}
|
||||
|
||||
Upon the invocation of a pass, we first check if this pass is enabled. This is
|
||||
done by first checking if the pass is explicitly disabled by a user, followed by
|
||||
inspecting if it is specified as a required pass by the user. If it is still
|
||||
undetermined whether this pass is enabled, its ``opt_level`` will be checked.
|
||||
This pass will be enabled and therefore executed only when its optimization
|
||||
level is not less than the configured optimization level in the pass context.
|
||||
|
||||
To execute the pass, we need first to retrieve the registered pass in the TVM
|
||||
packed function registry using the pass name. This is possible because every
|
||||
pass is registered with an API endpoint as we will show later.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
Pass GetPass(const std::string& pass_name) {
|
||||
std::string fpass_name = "relax.transform." + pass_name;
|
||||
const std::optional<tvm::ffi::Function> f = tvm::ffi::Function::GetGlobal(fpass_name);
|
||||
TVM_FFI_ICHECK(f.has_value()) << "Cannot find " << fpass_name
|
||||
<< "to create the pass " << pass_name;
|
||||
return (*f)();
|
||||
}
|
||||
|
||||
Some helper functions are provided to create each type of these aforementioned
|
||||
passes. These helpers are also exposed to the Python frontend for users to
|
||||
favorably use Python APIs to create a specific pass object.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
Pass CreateFunctionPass(
|
||||
std::function<Function(Function, IRModule, PassContext)> pass_func,
|
||||
int opt_level,
|
||||
ffi::String name,
|
||||
ffi::Array<ffi::String> required,
|
||||
bool traceable = false);
|
||||
|
||||
Pass CreatePrimFuncPass(
|
||||
std::function<PrimFunc(PrimFunc, IRModule, PassContext)> pass_func,
|
||||
int opt_level,
|
||||
ffi::String name,
|
||||
ffi::Array<ffi::String> required,
|
||||
bool traceable = false);
|
||||
|
||||
Pass CreateModulePass(
|
||||
std::function<IRModule(IRModule, PassContext)> pass_func,
|
||||
int opt_level,
|
||||
ffi::String name,
|
||||
ffi::Array<ffi::String> required,
|
||||
bool traceable = false);
|
||||
|
||||
Pass Sequential(tvm::ffi::Array<Pass> passes, PassInfo pass_info);
|
||||
|
||||
Pass Registration
|
||||
^^^^^^^^^^^^^^^^^
|
||||
|
||||
We've covered the concept of different level of passes and the context used for
|
||||
compilation. It would be interesting to see how easily users can register
|
||||
a pass. Let's take const folding as an example. This pass has already been
|
||||
implemented to fold constants in a Relax function (found in
|
||||
`src/relax/transform/fold_constant.cc`_).
|
||||
|
||||
An API was provided to perform the ``Expr`` to ``Expr`` transformation.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
Expr FoldConstant(const Expr& expr);
|
||||
|
||||
In order to register this pass to the pass infra, we first need to decide at
|
||||
which level this pass will be performed. As const folding happens on individual
|
||||
functions, we should intuitively create a ``FunctionPass`` for it through
|
||||
``CreateFunctionPass``. The ``pass_func`` is returned as a packed function that
|
||||
invokes the ``Expr`` to ``Expr`` API on each function in a `IRModule`. ``{}``
|
||||
indicates that no prerequisite is required for this pass. Otherwise, the pass
|
||||
developer has to identify and list them.
|
||||
|
||||
Meanwhile, a pass API endpoint is registered with the name
|
||||
``"relax.transform.FoldConstant"``. This pass, therefore, becomes an entry in the
|
||||
registry that can be accessed by both C++ (e.g. the ``GetPass`` above) and
|
||||
Python when needed.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
namespace transform {
|
||||
|
||||
Pass FoldConstant() {
|
||||
auto pass_func =
|
||||
[=](Function f, IRModule m, PassContext pc) { return ConstantFolder::Fold(f, m); };
|
||||
return CreateFunctionPass(pass_func, 0, "FoldConstant", {});
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def("relax.transform.FoldConstant", FoldConstant);
|
||||
}
|
||||
|
||||
} // namespace transform
|
||||
|
||||
To allow other C++ modules to apply this pass, we declare a free function in
|
||||
`include/tvm/relax/transform.h`_ as the following:
|
||||
|
||||
.. code:: c++
|
||||
|
||||
TVM_DLL Pass FoldConstant();
|
||||
|
||||
.. _pass_instrument_cpp_backend:
|
||||
|
||||
Pass Instrument
|
||||
^^^^^^^^^^^^^^^
|
||||
|
||||
Pass Instrument is a mechanism to analyze the pass itself. For example,
|
||||
we can use the infrastructure to know how much time and memory a pass requires
|
||||
or how a pass can transform the IR module.
|
||||
|
||||
We introduce four instrument points in the life-cycle of ``PassContext``.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
TVM_DLL void InstrumentEnterPassContext();
|
||||
TVM_DLL void InstrumentExitPassContext();
|
||||
TVM_DLL bool InstrumentBeforePass(const IRModule& mod, const PassInfo& info) const;
|
||||
TVM_DLL void InstrumentAfterPass(const IRModule& mod, const PassInfo& info) const;
|
||||
|
||||
``InstrumentEnterPassContext`` is called immediately when entering the scope
|
||||
of the ``PassContext`` instance.
|
||||
|
||||
``InstrumentExitPassContext`` is called when leaving the scope of ``PassContext``,
|
||||
or exceptions occur during the execution of passes.
|
||||
This method is also called when instruments is being overridden by ``override_instruments`` in :py:class:`tvm.transform.PassContext`.
|
||||
See :ref:`pass_instrument_overriden`.
|
||||
|
||||
``InstrumentBeforePass`` is called before execution.
|
||||
``InstrumentAfterPass`` is called after execution if the pass should be run. The behavior is like:
|
||||
|
||||
.. code:: c++
|
||||
|
||||
if (pass_ctx.InstrumentBeforePass(ir_module, pass_info)) {
|
||||
new_ir_module = run_pass(ir_module, pass_ctx);
|
||||
pass_ctx.InstrumentAfterPass(new_ir_module, pass_info);
|
||||
return new_ir_module;
|
||||
}
|
||||
|
||||
The ``PassInstrument`` interface allow you to run arbitrary code inside above four methods.
|
||||
Multiple ``PassInstrument`` instances can be registed into a single
|
||||
``PassContext``. ``PassInstrument`` instances are called sequentially in the order of
|
||||
``instruments`` argument passed to ``PassContext``.
|
||||
|
||||
``PassInstrument`` provides following interfaces:
|
||||
|
||||
.. code:: c++
|
||||
|
||||
namespace instrument {
|
||||
|
||||
class PassInstrumentNode : public Object {
|
||||
public:
|
||||
ffi::String name;
|
||||
virtual void EnterPassContext() const = 0;
|
||||
virtual void ExitPassContext() const = 0;
|
||||
virtual bool ShouldRun(const IRModule& mod, const transform::PassInfo& info) const = 0;
|
||||
virtual void RunBeforePass(const IRModule& mod, const transform::PassInfo& info) const = 0;
|
||||
virtual void RunAfterPass(const IRModule& mod, const transform::PassInfo& info) const = 0;
|
||||
/* Other fields are omitted. */
|
||||
};
|
||||
|
||||
class PassInstrument : public ObjectRef {
|
||||
public:
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PassInstrument, ObjectRef, PassInstrumentNode);
|
||||
};
|
||||
|
||||
} // namespace instrument
|
||||
|
||||
Python frontend are provided to implement ``PassInstrument`` quickly. See :ref:`pass_instrument_py_frontend`.
|
||||
|
||||
Within a ``PassContext``, the call sequence of a ``PassInstrument`` instance is like:
|
||||
|
||||
::
|
||||
|
||||
with PassContext(instruments=[pi]) # pi = a PassInstrument implementation.
|
||||
pi.EnterPassContext()
|
||||
|
||||
if pi.ShouldRun(Pass1):
|
||||
pi.RunBeforePass()
|
||||
Pass1()
|
||||
pi.RunAfterPass()
|
||||
|
||||
if pi.ShouldRun(Pass2):
|
||||
pi.RunBeforePass()
|
||||
Pass2()
|
||||
pi.RunAfterPass()
|
||||
|
||||
pi.ExitPassContext()
|
||||
|
||||
Here is a brief introduction of relations between ``PassInstrument`` interfaces
|
||||
and ``PassContext`` methods. See (`src/ir/transform.cc`_) for more details.
|
||||
|
||||
- ``InstrumentEnterPassContext``
|
||||
|
||||
* ``EnterPassContext()`` is executed in the order of ``instruments`` passed to the ``PassContext``.
|
||||
* When an exception raises, ``PassContext`` disable the pass instrumentation
|
||||
by clearing all registered ``PassInstrument`` instances.
|
||||
* Then ``PassContext`` execute ``ExitPassContext()`` method of each ``PassInstrument``
|
||||
instances which successfully finished ``EnterPassContext()``
|
||||
* For example, if ``PassInstrument`` A, B, and C are registered to a ``PassContext``
|
||||
and A finished ``EnterPassContext()`` while B throws an exception, then C
|
||||
is never executed; ``ExitPassContext()`` of A is executed.
|
||||
|
||||
- ``InstrumentExitPassContext``
|
||||
|
||||
* ``ExitPassContext()`` of each ``PassInstrument`` instances are executed in
|
||||
the order of ``instruments`` passed to the ``PassContext``.
|
||||
* While an exception occurs, ``instruments`` is cleared.
|
||||
* ``PassInstrument`` Instances registered after the one throwing exceptions do not execute ``ExitPassContext``.
|
||||
|
||||
- ``InstrumentBeforePass``
|
||||
|
||||
* ``ShouldRun`` is executed if the pass is not listed as a required pass.
|
||||
* ``RunBeforePass`` is executed in the order of ``instruments`` if the pass is not blocked by ``ShouldRun``.
|
||||
* Note that ``InstrumentBeforePass`` returns a boolean indicating whether or not the pass should be run.
|
||||
* When an exception occur, it is thrown immediately.
|
||||
We rely on Python Context Manager to exit ``PassContext`` safely
|
||||
(meaning ``ExitPassContext`` of each instruments will be run. For C++, please refer to `include/tvm/support/with.h`_.)
|
||||
|
||||
- ``InstrumentAfterPass``
|
||||
|
||||
* ``RunAfterPass`` is executed in the order of ``instruments`` passed to the ``PassContext``.
|
||||
* When an exception occur, it is thrown immediately.
|
||||
We rely on Python Context Manager or ``With`` class(`include/tvm/support/with.h`_) to exit ``PassContext`` safely
|
||||
|
||||
Built-in Instrument
|
||||
^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
There are several built-in instruments.
|
||||
|
||||
- PassTimingInstrument (see `src/ir/instrument.cc`_)
|
||||
|
||||
* Profile the execution time of passes.
|
||||
|
||||
- PrintBeforeAll (see `python/tvm/ir/instrument.py`_)
|
||||
|
||||
* Print the IR module and pass info before each pass executes.
|
||||
|
||||
- PrintAfterAll (see `python/tvm/ir/instrument.py`_)
|
||||
|
||||
* Print the IR module and pass info after each pass executes.
|
||||
|
||||
- PassPrintingInstrument (see `python/tvm/ir/instrument.py`_)
|
||||
|
||||
* Selectively print the IR module before or after specific named passes.
|
||||
|
||||
- DumpIR (see `python/tvm/ir/instrument.py`_)
|
||||
|
||||
* Dump the IR module to files after each pass executes.
|
||||
|
||||
Python Frontend
|
||||
~~~~~~~~~~~~~~~
|
||||
|
||||
Only some simple APIs are needed for the frontend side. For example, we can
|
||||
provide users the following APIs to create and execute a pass (full
|
||||
implementation is provided in `python/tvm/relax/transform/transform.py`_ and
|
||||
`python/tvm/ir/transform.py`_). The backend
|
||||
receives the information and decides which function it should use to create
|
||||
a Pass object.
|
||||
|
||||
PassContext
|
||||
^^^^^^^^^^^
|
||||
|
||||
Python frontend provides a wrapper for the ``PassContext`` to enable the
|
||||
``with`` syntax by overriding ``__enter__`` and ``__exit__``. A ``current``
|
||||
static method is offered for users to get the context that is in use under
|
||||
a certain scope.
|
||||
|
||||
.. code:: python
|
||||
|
||||
@tvm_ffi.register_object("transform.PassContext")
|
||||
class PassContext(tvm.runtime.Object):
|
||||
def __enter__(self):
|
||||
_transform.EnterPassContext(self)
|
||||
return self
|
||||
|
||||
def __exit__(self, ptype, value, trace, config):
|
||||
_transform.ExitPassContext(self)
|
||||
|
||||
@staticmethod
|
||||
def current():
|
||||
"""Return the current pass context."""
|
||||
return _transform.GetCurrentPassContext()
|
||||
|
||||
A ``PassContext`` is used to configure the compilation options, including the
|
||||
optimization level and required/disabled passes. It can also take a dictionary
|
||||
of configs so that different passes can conveniently fetch the passed data, such
|
||||
as fallback device info and step/depth for loop unrolling, etc. In order to
|
||||
enable fetching the required config, the key must be registered through
|
||||
``TVM_REGISTER_PASS_CONFIG_OPTION``. For example, the following is used by the
|
||||
loop unrolling pass
|
||||
|
||||
.. code:: c++
|
||||
|
||||
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.UnrollLoop", UnrollLoopConfig);
|
||||
|
||||
Please refer to `src/tirx/transform/unroll_loop.cc`_ for more details.
|
||||
|
||||
.. _pass_instrument_py_frontend:
|
||||
|
||||
Pass Instrument
|
||||
^^^^^^^^^^^^^^^
|
||||
|
||||
One can implement a ``PassInstrument`` by using the ``pass_instrument``
|
||||
decorator(`python/tvm/ir/instrument.py`_) on a class implementing following methods.
|
||||
Note that it is recommended to use the ``pass_instrument`` decorator to implement
|
||||
``PassInstrument``, instead of overriding or subclassing.
|
||||
|
||||
- ``enter_pass_ctx``
|
||||
|
||||
* This method is run when entering ``PassContext``.
|
||||
|
||||
- ``exit_pass_ctx``
|
||||
|
||||
* This method is run when exiting ``PassContext``.
|
||||
|
||||
- ``should_run``
|
||||
|
||||
* This method is run before a pass is executed, returning a boolean
|
||||
indicating whether or not the pass should be run.
|
||||
|
||||
- ``run_before_pass``
|
||||
|
||||
* If a pass should be run, this method is run just before pass execution.
|
||||
|
||||
- ``run_after_pass``
|
||||
|
||||
* This method is run right after a pass has been executed.
|
||||
|
||||
``PassInstrument`` instances can be registered through ``instruments`` argument in
|
||||
:py:class:`tvm.transform.PassContext`.
|
||||
|
||||
See `python/tvm/ir/instrument.py`_ for examples of how to implement ``PassInstrument`` with Python APIs.
|
||||
|
||||
.. _pass_instrument_overriden:
|
||||
|
||||
Override Instruments in Current PassContext
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
``override_instruments`` method is provided to override the ``instruments`` of current ``PassContext``.
|
||||
For example, if passes are run without explicitly creating a new ``PassContext``,
|
||||
one can still register ``PassInstrument`` into the global ``PassContext`` by:
|
||||
|
||||
.. code:: python
|
||||
|
||||
cur_pass_ctx = tvm.transform.PassContext.current()
|
||||
# override PassInstrument instances
|
||||
cur_pass_ctx.override_instruments([pass_inst])
|
||||
mod = pass_seq(mod)
|
||||
result = pass_inst.get_result()
|
||||
|
||||
Note that when ``override_instruments`` is called, the ``exit_pass_ctx`` method of
|
||||
old ``PassInstrument`` instances are called. Then the ``enter_pass_ctx`` method of
|
||||
new ``PassInstrument`` are called.
|
||||
|
||||
.. _Sequential: https://pytorch.org/docs/stable/nn.html?highlight=sequential#torch.nn.Sequential
|
||||
|
||||
.. _Block: https://pytorch.org/docs/stable/generated/torch.nn.Module.html
|
||||
|
||||
.. _include/tvm/ir/transform.h: https://github.com/apache/tvm/blob/main/include/tvm/ir/transform.h
|
||||
|
||||
.. _include/tvm/support/with.h: https://github.com/apache/tvm/blob/main/include/tvm/support/with.h
|
||||
|
||||
.. _src/relax/ir/transform.cc: https://github.com/apache/tvm/blob/main/src/relax/ir/transform.cc
|
||||
|
||||
.. _src/ir/transform.cc: https://github.com/apache/tvm/blob/main/src/ir/transform.cc
|
||||
|
||||
.. _src/ir/instrument.cc: https://github.com/apache/tvm/blob/main/src/ir/instrument.cc
|
||||
|
||||
.. _src/relax/transform/fold_constant.cc: https://github.com/apache/tvm/blob/main/src/relax/transform/fold_constant.cc
|
||||
|
||||
.. _python/tvm/relax/transform/transform.py: https://github.com/apache/tvm/blob/main/python/tvm/relax/transform/transform.py
|
||||
|
||||
.. _include/tvm/relax/transform.h: https://github.com/apache/tvm/blob/main/include/tvm/relax/transform.h
|
||||
|
||||
.. _python/tvm/ir/transform.py: https://github.com/apache/tvm/blob/main/python/tvm/ir/transform.py
|
||||
|
||||
.. _python/tvm/ir/instrument.py: https://github.com/apache/tvm/blob/main/python/tvm/ir/instrument.py
|
||||
|
||||
.. _src/tirx/transform/unroll_loop.cc: https://github.com/apache/tvm/blob/main/src/tirx/transform/unroll_loop.cc
|
||||
|
||||
.. _use pass infra: https://github.com/apache/tvm/blob/main/docs/how_to/tutorials/customize_opt.py
|
||||
@@ -0,0 +1,427 @@
|
||||
.. 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<ExecWord>`` (``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<VMFuncInfo>`` 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__<func_name>``. 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<ffi::Any>``) — 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
|
||||
@@ -0,0 +1,281 @@
|
||||
.. 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.
|
||||
|
||||
.. _tvm-runtime-system:
|
||||
|
||||
TVM Runtime System
|
||||
==================
|
||||
|
||||
TVM supports multiple programming languages for the compiler stack development and deployment.
|
||||
In this note, we explain the key elements of the TVM runtime.
|
||||
|
||||
.. image:: https://tvm.apache.org/images/release/tvm_flexible.png
|
||||
|
||||
We need to satisfy quite a few interesting requirements:
|
||||
|
||||
- Deployment: invoke the compiled function from python/javascript/c++ language.
|
||||
- Debug: define a function in python and call that from a compiled function.
|
||||
- Link: write driver code to call device specific code (CUDA) and call it from compiled host function.
|
||||
- Prototype: define an IR pass from python and call that from C++ backend.
|
||||
- Expose: compiler stack developed in c++ to front-end (i.e, python)
|
||||
- Experiment: ship a compiled function to an embedded device to directly run there.
|
||||
|
||||
We want to be able to define a function from any language and call from another.
|
||||
We also want the runtime core to be minimal to deploy to embedded devices.
|
||||
|
||||
.. _tvm-runtime-system-packed-func:
|
||||
|
||||
PackedFunc
|
||||
----------
|
||||
|
||||
`PackedFunc`_ is a simple but elegant solution we find to solve the
|
||||
challenges listed. A single ``PackedFunc`` object represents a
|
||||
function call whose caller and callee may be in different languages.
|
||||
|
||||
The following code block provides an example in C++
|
||||
|
||||
.. _PackedFunc: https://github.com/apache/tvm/blob/main/3rdparty/tvm-ffi/include/tvm/ffi/function.h
|
||||
|
||||
.. code:: c
|
||||
|
||||
#include <tvm/ffi/function.h>
|
||||
|
||||
void MyAdd(ffi::PackedArgs args, ffi::Any* rv) {
|
||||
// automatically convert arguments to desired type.
|
||||
int a = args[0].cast<int>();
|
||||
int b = args[1].cast<int>();
|
||||
// automatically assign value return to rv
|
||||
*rv = a + b;
|
||||
}
|
||||
|
||||
void CallPacked() {
|
||||
PackedFunc myadd = PackedFunc(MyAdd);
|
||||
// get back 3
|
||||
int c = myadd(1, 2);
|
||||
}
|
||||
|
||||
In the above codeblock, we defined a PackedFunc MyAdd. It takes two arguments
|
||||
: ``args`` represents input arguments and ``rv`` represents return value.
|
||||
The function is type-erased, which means that the function signature does not restrict which input type to pass in or type to return.
|
||||
Under the hood, when we call a PackedFunc, it packs the input arguments to ffi::PackedArgs on stack,
|
||||
and gets the result back via ffi::Any.
|
||||
|
||||
Thanks to template tricks in C++, we can call a PackedFunc just like a normal function. Because of its type-erased nature, we can call a PackedFunc from dynamic languages like python, without additional glue code for each new type function created.
|
||||
The following example registers PackedFunc in C++ and calls from python.
|
||||
|
||||
.. code:: c
|
||||
|
||||
// register a global packed function in c++
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def_packed("myadd", MyAdd);
|
||||
}
|
||||
|
||||
.. code:: python
|
||||
|
||||
import tvm
|
||||
|
||||
myadd = tvm.get_global_func("myadd")
|
||||
# prints 3
|
||||
print(myadd(1, 2))
|
||||
|
||||
Most of the magic of PackedFunc lies in ``ffi::PackedArgs`` and ``ffi::Any`` structure.
|
||||
We restrict a list of possible types which can be passed.
|
||||
Here are the common ones:
|
||||
|
||||
- int, float and string
|
||||
- PackedFunc itself
|
||||
- Module for compiled modules
|
||||
- DLTensor* for tensor object exchange
|
||||
- TVM Object to represent any object in IR
|
||||
|
||||
The restriction makes the implementation simple without the need of serialization.
|
||||
Despite being minimum, the PackedFunc is sufficient for the use-case of deep learning deployment as
|
||||
most functions only take DLTensor or numbers.
|
||||
|
||||
Since one PackedFunc can take another PackedFunc as an argument,
|
||||
we can pass functions from python (as PackedFunc) to C++.
|
||||
|
||||
.. code:: c
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def_packed("callhello", [](ffi::PackedArgs args, ffi::Any* rv) {
|
||||
ffi::Function f = args[0].cast<ffi::Function>();
|
||||
f("hello world");
|
||||
});
|
||||
}
|
||||
|
||||
.. code:: python
|
||||
|
||||
import tvm
|
||||
|
||||
def callback(msg):
|
||||
print(msg)
|
||||
|
||||
# convert to PackedFunc
|
||||
f = tvm.runtime.convert(callback)
|
||||
callhello = tvm.get_global_func("callhello")
|
||||
# prints hello world
|
||||
callhello(f)
|
||||
|
||||
TVM provides a `minimum C API`_,
|
||||
which allows us to embed the PackedFunc into any languages. Besides python, so far we supported
|
||||
`java`_ and `javascript`_.
|
||||
This philosophy of embedded API is very like Lua, except that we don't have a new language but use C++.
|
||||
|
||||
.. _minimum C API: https://github.com/apache/tvm/blob/main/3rdparty/tvm-ffi/include/tvm/ffi/c_api.h
|
||||
.. _java: https://github.com/apache/tvm/tree/main/jvm
|
||||
.. _javascript: https://github.com/apache/tvm/tree/main/web
|
||||
|
||||
|
||||
One fun fact about PackedFunc is that we use it for both compiler and deployment stack.
|
||||
|
||||
- All compiler pass functions of TVM are exposed to frontend as PackedFunc
|
||||
- The compiled module also returns the compiled function as PackedFunc
|
||||
|
||||
To keep the runtime minimum, we isolated the IR Object support from the deployment runtime. The resulting runtime takes around 200K - 600K depending on how many runtime driver modules (e.g., CUDA) get included.
|
||||
|
||||
The overhead of calling into PackedFunc vs. a normal function is small, as it is only saving a few values on the stack.
|
||||
So it is OK as long as we don't wrap small functions.
|
||||
In summary, the PackedFunc is the universal glue in TVM where we use it extensively to support our compiler and deployment.
|
||||
|
||||
.. _tvm-runtime-system-module:
|
||||
|
||||
Module
|
||||
------
|
||||
|
||||
Since TVM supports multiple types of devices, we need to support different type of drivers.
|
||||
We have to use the driver API to load the kernel, set up the argument in packed format and perform kernel launch.
|
||||
We also need to patch up the driver API so that the exposed functions are threadsafe.
|
||||
So we often need to implement these driver glues in C++ and expose them to the user.
|
||||
We can certainly not do it for each type of functions, so again PackedFunc is our answer.
|
||||
|
||||
TVM defines the compiled object as `Module`_.
|
||||
The user can get the compiled function from Module as PackedFunc.
|
||||
The generated compiled code can dynamically get function from Module in runtime. It caches the function handle in the first call and reuses in subsequent calls. We use this to link device code and callback into any PackedFunc(e.g., python) from generated code.
|
||||
|
||||
.. _Module: https://github.com/apache/tvm/blob/main/3rdparty/tvm-ffi/include/tvm/ffi/extra/module.h
|
||||
|
||||
The ModuleNode is an abstract class that can be implemented by each type of device.
|
||||
So far we support modules for CUDA, Metal, OpenCL and loading dynamic shared libraries. This abstraction makes introduction
|
||||
of new device easy, and we do not need to redo the host code generation for each type of device.
|
||||
|
||||
Remote Deployment
|
||||
-----------------
|
||||
|
||||
The PackedFunc and Module system also makes it easy to ship the function into remote devices directly.
|
||||
Under the hood, we have an RPCModule that serializes the arguments to do the data movement and launches the computation on the remote.
|
||||
|
||||
.. image:: https://tvm.apache.org/images/release/tvm_rpc.png
|
||||
|
||||
The RPC server itself is minimum and can be bundled into the runtime. We can start a minimum TVM
|
||||
RPC server on iPhone/android/raspberry pi or even the browser. The cross compilation on server and shipping of the module for testing can be done in the same script. Checkout
|
||||
:ref:`tutorial-cross-compilation-and-rpc` for more details.
|
||||
|
||||
|
||||
This instant feedback gives us a lot of advantages. For example, to test the correctness of generated code on iPhone, we no longer have to write test-cases in swift/objective-c from scratch -- We can use RPC to execute on iPhone, copy the result back and do verification on the host via numpy. We can also do the profiling using the same script.
|
||||
|
||||
TVM Object and Compiler Stack
|
||||
-----------------------------
|
||||
|
||||
As we mentioned earlier, we build compiler stack API on top of the PackedFunc runtime system.
|
||||
We faced a constant changing of the compiler API for the need of research. We need a new language object or IR node whenever we want to test out new primitives.
|
||||
However, we don't want to change our API from time to time. Besides that, we also want to
|
||||
|
||||
- be able to serialize any language object and IRs
|
||||
- be able to explore, print, and manipulate the IR objects in front-end language to do quick prototyping.
|
||||
|
||||
We introduced a base class, called `Object`_ to solve this problem.
|
||||
All the language object in the compiler stack is a subclass of ``Object``. Each object contains a string type_key that uniquely identifies
|
||||
the type of object. We choose string instead of int as type key so new ``Object`` class can be added in the decentralized fashion without
|
||||
adding the code back to the central repo. To ease the speed of dispatching, we allocate an integer type_index at runtime for each type_key.
|
||||
|
||||
.. _Object: https://github.com/apache/tvm/blob/main/include/tvm/runtime/object.h
|
||||
|
||||
Since usually one ``Object`` could be referenced in multiple places in the language, we use a shared_ptr to keep
|
||||
track of reference. We use ``ObjectRef`` class to represent a reference to the ``Object``.
|
||||
We can roughly view ``ObjectRef`` class as shared_ptr to the ``Object`` container.
|
||||
We can also define subclass ``ObjectRef`` to hold each subtypes of ``Object``. Each subclass of ``Object`` needs to define the
|
||||
RegisterReflection function.
|
||||
|
||||
|
||||
Each ``Object`` subclass will override this to register its members. Here is an example implementation of IntImmNode.
|
||||
|
||||
.. code:: c
|
||||
|
||||
class IntImmNode : public PrimExprNode {
|
||||
public:
|
||||
/*! \brief the Internal value. */
|
||||
int64_t value;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<IntImmNode>().def_ro("value", &IntImmNode::value);
|
||||
}
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.IntImm", IntImmNode, PrimExprNode);
|
||||
};
|
||||
// in cc file
|
||||
TVM_FFI_STATIC_INIT_BLOCK() { IntImmNode::RegisterReflection(); }
|
||||
|
||||
The RegisterReflection gives us a reflection API to register each member of the object.
|
||||
We can use this function to visit the node and serialize any language object recursively.
|
||||
It also allows us to get members of an object easily in front-end language.
|
||||
For example, we can access the value field of the IntImmNode.
|
||||
|
||||
.. code:: python
|
||||
|
||||
import tvm
|
||||
|
||||
x = tvm.tirx.IntImm("int32", 1)
|
||||
# access the value field of IntImmNode
|
||||
print(x.value)
|
||||
|
||||
New ``Object`` can be added to C++ without changing the front-end runtime, making it easy to make extensions to the compiler stack.
|
||||
Note that this is not the fastest way to expose members to front-end language, but might be one of the simplest
|
||||
approaches possible. We also find that it fits our purposes as we mainly use python for testing and prototyping and still use c++
|
||||
to do the heavy lifting job.
|
||||
|
||||
Implementation Details
|
||||
----------------------
|
||||
|
||||
Each argument in PackedFunc contains a union value `TVMValue`_
|
||||
and a type code. This design allows the dynamically typed language to convert to the corresponding type directly, and statically typed language to
|
||||
do runtime type checking during conversion.
|
||||
|
||||
.. _TVMValue: https://github.com/apache/tvm/blob/main/3rdparty/tvm-ffi/include/tvm/ffi/c_api.h
|
||||
|
||||
The relevant files are
|
||||
|
||||
- `function.h`_ for C++ PackedFunc API
|
||||
- `c_api.h`_ for C API.
|
||||
|
||||
.. _function.h: https://github.com/apache/tvm/blob/main/3rdparty/tvm-ffi/include/tvm/ffi/function.h
|
||||
.. _c_api.h: https://github.com/apache/tvm/blob/main/3rdparty/tvm-ffi/include/tvm/ffi/c_api.h
|
||||
|
||||
To support extension types, we used a registry system to register type related information, like support of any
|
||||
in C++. See the ``tvm-ffi`` subproject under ``3rdparty/tvm-ffi/`` for more details on the FFI type system.
|
||||
|
||||
|
||||
Runtime-Specific Information
|
||||
============================
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:glob:
|
||||
|
||||
runtimes/*
|
||||
@@ -0,0 +1,259 @@
|
||||
.. 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.
|
||||
|
||||
.. _tvm-runtime-vulkan:
|
||||
|
||||
Vulkan Runtime
|
||||
==============
|
||||
|
||||
TVM supports using Vulkan compute shaders to execute queries. Each
|
||||
computational kernel is compiled into a SPIR-V shader, which can then
|
||||
be called using the TVM interface.
|
||||
|
||||
.. _tvm-runtime-vulkan-features:
|
||||
|
||||
Vulkan Features, Limits
|
||||
-----------------------
|
||||
|
||||
.. _Required Limits: https://www.khronos.org/registry/vulkan/specs/1.2-extensions/html/vkspec.html#limits-minmax
|
||||
|
||||
Since different Vulkan implementations may enable different optional
|
||||
features or have different physical limits, the code generation must
|
||||
know which features are available to use. These correspond to
|
||||
specific Vulkan capabilities/limits as in
|
||||
:ref:`Vulkan Capabilities Table <tvm-table-vulkan-capabilities>`.
|
||||
If unspecified, TVM assumes that a capability is not available, or
|
||||
that a limit is the minimum guaranteed by the Vulkan spec in the
|
||||
`Required Limits`_ section.
|
||||
|
||||
These parameters can be either explicitly specific when defining a
|
||||
:ref:`Target <tvm-target-specific-target>`, or can be queried from a
|
||||
device. To query from a device, the special parameter
|
||||
``-from_device=N`` can be used to query all vulkan device parameters
|
||||
from device id ``N``. Any additional parameters explicitly specified
|
||||
will override the parameters queried from the device.
|
||||
|
||||
.. _VkSubgroupFeatureFlagBits: https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VkSubgroupFeatureFlagBits.html
|
||||
|
||||
.. list-table:: Vulkan Capabilities
|
||||
:name: tvm-runtime-table-vulkan-capabilities
|
||||
:header-rows: 1
|
||||
|
||||
* - Target Parameter
|
||||
- Required Vulkan Version/Extension
|
||||
- Parameter Queried
|
||||
- Default Value
|
||||
|
||||
* - ``supported_subgroup_operations``
|
||||
- Vulkan 1.1+
|
||||
- ``VkPhysicalDeviceSubgroupProperties::supportedOperations``
|
||||
- 0 (interpreted as `VkSubgroupFeatureFlagBits`_)
|
||||
|
||||
* - ``max_push_constants_size``
|
||||
-
|
||||
- ``VkPhysicalDeviceLimits::maxPushConstantsSize``
|
||||
- 128 bytes
|
||||
|
||||
* - ``max_uniform_buffer_range``
|
||||
-
|
||||
- ``VkPhysicalDeviceLimits::maxUniformBufferRange``
|
||||
- 16384 bytes
|
||||
|
||||
|
||||
* - ``max_storage_buffer_range``
|
||||
-
|
||||
- ``VkPhysicalDeviceLimits::maxStorageBufferRange``
|
||||
- 2\ :sup:`27`\ bytes
|
||||
|
||||
|
||||
* - ``max_per_stage_descriptor_storage_buffer``
|
||||
-
|
||||
- ``VkPhysicalDeviceLimits::maxPerStageDescriptorStorageBuffers``
|
||||
- 4
|
||||
|
||||
|
||||
* - ``supports_storage_buffer_storage_class``
|
||||
- VK_KHR_storage_buffer_storage_class
|
||||
-
|
||||
- false
|
||||
|
||||
|
||||
* - ``supports_storage_buffer_8bit_access``
|
||||
- VK_KHR_8bit_storage
|
||||
- ``VkPhysicalDevice8BitStorageFeaturesKHR::storageBuffer8BitAccess``
|
||||
- false
|
||||
|
||||
|
||||
* - ``supports_storage_buffer_16bit_access``
|
||||
- VK_KHR_16bit_storage
|
||||
- ``VkPhysicalDevice16BitStorageFeaturesKHR::storageBuffer16BitAccess``
|
||||
- false
|
||||
|
||||
|
||||
* - ``supports_float16``
|
||||
- VK_KHR_shader_float16_int8
|
||||
- ``VkPhysicalDeviceShaderFloat16Int8FeaturesKHR::shaderFloat16``
|
||||
- false
|
||||
|
||||
|
||||
* - ``supports_float64``
|
||||
-
|
||||
- ``VkPhysicalDeviceFeatures::shaderFloat64``
|
||||
- false
|
||||
|
||||
|
||||
* - ``supports_int8``
|
||||
- VK_KHR_shader_float16_int8
|
||||
- ``VkPhysicalDeviceShaderFloat16Int8FeaturesKHR::shaderInt8``
|
||||
- false
|
||||
|
||||
|
||||
* - ``supports_int16``
|
||||
-
|
||||
- ``VkPhysicalDeviceFeatures::shaderInt16``
|
||||
- false
|
||||
|
||||
|
||||
* - ``supports_int64``
|
||||
-
|
||||
- ``VkPhysicalDeviceFeatures::shaderInt64``
|
||||
- false
|
||||
|
||||
|
||||
|
||||
As of May 2021, not all Vulkan implementations are supported. For
|
||||
example, support for 64-bit integers is required. If a Vulkan target
|
||||
is not supported, an error message should be issued during SPIR-V code
|
||||
generation. Efforts are also underway to remove these requirements
|
||||
and support additional Vulkan implementations.
|
||||
|
||||
|
||||
.. _tvm-runtime-vulkan-spirv-capabilities:
|
||||
|
||||
SPIR-V Capabilities
|
||||
-------------------
|
||||
|
||||
Some of the device-specific capabilities also correspond to SPIR-V
|
||||
capabilities or extensions that must be declared in the shader, or a
|
||||
minimum SPIR-V version required in order to use a feature. The
|
||||
TVM-generated shaders will declare the minimum set of
|
||||
extensions/capabilities and the minimum allowed version of SPIR-V
|
||||
that are needed to execute the compiled graph.
|
||||
|
||||
If the shader generation requires a capability or extension that is
|
||||
not enabled in the ``Target``, an exception will be raised.
|
||||
|
||||
|
||||
.. list-table:: Vulkan Capabilities
|
||||
:name: tvm-table-vulkan-capabilities
|
||||
:header-rows: 1
|
||||
|
||||
* - Target Parameter
|
||||
- Required SPIR-V Version/Extension
|
||||
- Declared Capability
|
||||
|
||||
* - ``supported_subgroup_operations``
|
||||
- SPIR-V 1.3+
|
||||
- Varies, see `VkSubgroupFeatureFlagBits`_
|
||||
|
||||
* - ``supports_storage_buffer_storage_class``
|
||||
- SPV_KHR_storage_buffer_storage_class
|
||||
-
|
||||
|
||||
* - ``supports_storage_buffer_8bit_access``
|
||||
- SPV_KHR_8bit_storage
|
||||
- StorageBuffer8BitAccess
|
||||
|
||||
* - ``supports_storage_buffer_16bit_access``
|
||||
- SPV_KHR_16bit_storage
|
||||
- StorageBuffer16BitAccess
|
||||
|
||||
* - ``supports_float16``
|
||||
-
|
||||
- Float16
|
||||
|
||||
|
||||
* - ``supports_float64``
|
||||
-
|
||||
- Float64
|
||||
|
||||
|
||||
* - ``supports_int8``
|
||||
-
|
||||
- Int8
|
||||
|
||||
|
||||
* - ``supports_int16``
|
||||
-
|
||||
- Int16
|
||||
|
||||
|
||||
* - ``supports_int64``
|
||||
-
|
||||
- Int64
|
||||
|
||||
|
||||
Vulkan-Specific Environment Variables
|
||||
-------------------------------------
|
||||
|
||||
Both the SPIR-V code generation and the Vulkan runtime have
|
||||
environment variables that can modify some of the runtime behavior.
|
||||
These are intended for debugging purposes, both to more easily test
|
||||
specific code paths, and to output more information as needed. All
|
||||
boolean flags are true if the environment variable is set to a
|
||||
non-zero integer. An unset variable, the integer zero, or an empty
|
||||
string are all false boolean flags.
|
||||
|
||||
.. _VK_KHR_push_descriptor: https://khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VK_KHR_push_descriptor.html
|
||||
|
||||
.. _VK_KHR_descriptor_update_template: https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VK_KHR_descriptor_update_template.html
|
||||
|
||||
.. _VK_KHR_dedicated_allocation: https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VK_KHR_dedicated_allocation.html
|
||||
|
||||
.. _VkMemoryDedicatedRequirements: https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VkMemoryDedicatedRequirements.html
|
||||
|
||||
.. _Vulkan validation layers: https://github.com/KhronosGroup/Vulkan-LoaderAndValidationLayers/blob/master/layers/README.md
|
||||
|
||||
.. _spvValidate: https://github.com/KhronosGroup/SPIRV-Tools#validator
|
||||
|
||||
|
||||
* ``TVM_VULKAN_DISABLE_PUSH_DESCRIPTOR`` - A boolean flag. If true,
|
||||
TVM will explicitly allocate descriptors, and will not use the
|
||||
`VK_KHR_push_descriptor`_ or `VK_KHR_descriptor_update_template`_
|
||||
extensions. If false, TVM will decide whether to use these
|
||||
extensions based on their availability.
|
||||
|
||||
* ``TVM_VULKAN_DISABLE_DEDICATED_ALLOCATION`` - A boolean flag. If
|
||||
true, TVM will not mark memory allocations as being dedicated
|
||||
allocations, and will not use the `VK_KHR_dedicated_allocation`_
|
||||
extension. If false, TVM will decide whether memory allocations
|
||||
should be marked as dedicated based on the
|
||||
`VkMemoryDedicatedRequirements`_ for that buffer.
|
||||
|
||||
* ``TVM_VULKAN_ENABLE_VALIDATION_LAYERS`` - A boolean flag. If true,
|
||||
TVM will enable `Vulkan validation layers`_ that the device
|
||||
supports. If false, no validation layers are enabled.
|
||||
|
||||
* ``TVM_VULKAN_DISABLE_SHADER_VALIDATION`` - A boolean flag. If true,
|
||||
the SPIR-V shader validation done with `spvValidate`_ is skipped.
|
||||
If false (default), all SPIR-V shaders generated by TVM are
|
||||
validated with `spvValidate`_.
|
||||
|
||||
* ``TVM_VULKAN_DEBUG_SHADER_SAVEPATH`` - A path to a directory. If
|
||||
set to a non-empty string, the Vulkan codegen will save TIR, binary
|
||||
SPIR-V, and disassembled SPIR-V shaders to this directory, to be
|
||||
used for debugging purposes.
|
||||
@@ -0,0 +1,575 @@
|
||||
.. 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.
|
||||
|
||||
.. _tvmscript-arch:
|
||||
|
||||
TVMScript
|
||||
=========
|
||||
|
||||
TVMScript is a Python-based domain-specific language (DSL) for writing TVM IR. It lets users
|
||||
define ``IRModule``\ s — containing both Relax functions and TIR ``PrimFunc``\ s — using
|
||||
familiar Python syntax. Although TVMScript *looks* like Python, it is **not executed by the
|
||||
Python interpreter**. Instead, Python decorators extract the AST from the source code and
|
||||
transform it into TVM IR through a dedicated parser and IR builder pipeline.
|
||||
|
||||
TVMScript serves two roles in the TVM stack:
|
||||
|
||||
- **Authoring**: users write TIR kernels and Relax programs directly in TVMScript.
|
||||
- **Roundtrip**: every ``IRModule`` can be printed back to TVMScript via ``mod.script()`` and
|
||||
re-parsed to produce an equivalent module. This makes TVMScript the primary tool for
|
||||
inspecting, debugging, and serializing IR.
|
||||
|
||||
|
||||
Overview
|
||||
--------
|
||||
|
||||
The TVMScript system has three components:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
Parsing (Python source → TVM IR):
|
||||
|
||||
Python source (TVMScript)
|
||||
│
|
||||
▼ ast.parse + convert
|
||||
│
|
||||
Doc AST (mirror of Python AST)
|
||||
│
|
||||
▼ Parser (dispatch by token: ir / tirx / relax)
|
||||
│
|
||||
▼ IR Builder (frame stack)
|
||||
│
|
||||
TVM IR (IRModule, PrimFunc, relax.Function)
|
||||
|
||||
|
||||
Printing (TVM IR → Python source):
|
||||
|
||||
TVM IR
|
||||
│
|
||||
▼ IRDocsifier (C++, dispatch by token + type)
|
||||
│
|
||||
Doc tree (ExprDoc, StmtDoc, ...)
|
||||
│
|
||||
▼ DocToPythonScript
|
||||
│
|
||||
TVMScript text
|
||||
|
||||
- **Parser** (Python): reads Python source, converts it to a ``Doc AST`` (a mirror of
|
||||
Python's ``ast`` module), then walks the tree using dialect-specific handlers that call
|
||||
into the IR builder.
|
||||
- **IR Builder** (Python + C++): provides a frame-stack API where each ``with`` block or
|
||||
decorator pushes a frame. When the frame exits, the constructed IR is finalized. The builder
|
||||
is shared across dialects — TIR and Relax each register their own frame types.
|
||||
- **Printer** (C++): converts TVM IR objects to a ``Doc`` tree (an intermediate representation
|
||||
of Python syntax), then formats the tree into valid TVMScript text.
|
||||
|
||||
|
||||
Decorators
|
||||
----------
|
||||
|
||||
TVMScript uses three import aliases by convention:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.script import ir as I # module-level constructs
|
||||
from tvm.script import tirx as T # TIR constructs
|
||||
from tvm.script import relax as R # Relax constructs
|
||||
|
||||
The primary decorators are:
|
||||
|
||||
- ``@I.ir_module``: marks a Python class as an ``IRModule``. Each method inside becomes a
|
||||
function in the module.
|
||||
- ``@T.prim_func``: marks a function as a TIR ``PrimFunc``.
|
||||
- ``@R.function``: marks a function as a ``relax.Function``.
|
||||
|
||||
These can be composed:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@I.ir_module
|
||||
class MyModule:
|
||||
@T.prim_func
|
||||
def add_kernel(A: T.Buffer((128,), "float32"),
|
||||
B: T.Buffer((128,), "float32"),
|
||||
C: T.Buffer((128,), "float32")):
|
||||
for i in range(128):
|
||||
with T.sblock("compute"):
|
||||
vi = T.axis.spatial(128, i)
|
||||
C[vi] = A[vi] + B[vi]
|
||||
|
||||
@R.function
|
||||
def main(x: R.Tensor((128,), "float32"),
|
||||
y: R.Tensor((128,), "float32")) -> R.Tensor((128,), "float32"):
|
||||
with R.dataflow():
|
||||
out = R.call_tir(cls.add_kernel, (x, y),
|
||||
out_ty=R.Tensor((128,), "float32"))
|
||||
R.output(out)
|
||||
return out
|
||||
|
||||
When Python encounters ``@I.ir_module``, the decorator does **not** execute the class body.
|
||||
Instead, it calls ``tvm.script.parse()`` which extracts the source code of the class,
|
||||
builds a Doc AST, and hands it to the parser.
|
||||
|
||||
|
||||
Parser Architecture
|
||||
-------------------
|
||||
|
||||
The parser lives in ``python/tvm/script/parser/``.
|
||||
|
||||
Dispatch mechanism
|
||||
~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Different IR dialects (TIR, Relax) need different handling for the same Python syntax. For
|
||||
example, ``if ... else`` inside ``@T.prim_func`` creates a TIR ``If`` branch, while the same
|
||||
syntax inside ``@R.function`` creates a Relax ``If`` node with different semantics.
|
||||
|
||||
The parser maintains a **dispatch token** stack (``["default"]`` initially). When it encounters
|
||||
a decorated function, it inspects the decorator to determine the token — ``"tirx"`` for
|
||||
``@T.prim_func``, ``"relax"`` for ``@R.function`` — and pushes it onto the stack.
|
||||
|
||||
Each AST node type is dispatched via a virtual table:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
ParseVTable[(token, node_type)] → handler function
|
||||
|
||||
Lookup order:
|
||||
1. (current_token, node_type) e.g. ("tirx", "For")
|
||||
2. ("default", node_type) e.g. ("default", "For")
|
||||
3. generic_visit fallback
|
||||
|
||||
Dialect-specific parsers (``parser/tirx/parser.py``, ``parser/relax/parser.py``) register
|
||||
handlers using ``@dispatch.register(token, type_name)`` decorators.
|
||||
|
||||
Parse flow
|
||||
~~~~~~~~~~
|
||||
|
||||
The entry point is ``parse(program, extra_vars)``:
|
||||
|
||||
1. **Source extraction**: the program's source code is extracted (from a class, function, or
|
||||
string) and converted to a Doc AST via Python's ``ast`` module.
|
||||
|
||||
2. **AST walking**: the ``Parser`` (a subclass of ``doc.NodeVisitor``) walks the Doc AST.
|
||||
For each node, it looks up the handler in the dispatch table.
|
||||
|
||||
3. **Expression evaluation**: expressions like ``T.grid(128, 128)`` are evaluated by the
|
||||
``ExprEvaluator``, which resolves names against the variable table and the ``T.``/``R.``
|
||||
module namespaces.
|
||||
|
||||
4. **Value binding**: assignment statements (``A = T.match_buffer(...)`` in TIR,
|
||||
``lv = R.add(x, y)`` in Relax) go through dialect-specific ``bind_*_value()`` functions
|
||||
that register the resulting TVM objects in the parser's ``VarTable``.
|
||||
|
||||
5. **Scoping**: the ``VarTable`` maintains a stack of frames. Entering a ``with`` block,
|
||||
``for`` loop, or function body pushes a new frame; exiting pops it. This ensures variables
|
||||
are scoped correctly.
|
||||
|
||||
Variable table
|
||||
~~~~~~~~~~~~~~
|
||||
|
||||
The ``VarTable`` is the parser's symbol table:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
VarTable
|
||||
├── frames: [VarTableFrame, ...] ← stack of scopes
|
||||
└── name2value: {str: [Any, ...]} ← name → value stack (for shadowing)
|
||||
|
||||
When a name is looked up, the most recent binding wins. When a frame is popped, all bindings
|
||||
introduced in that frame are removed.
|
||||
|
||||
|
||||
IR Builder Architecture
|
||||
-----------------------
|
||||
|
||||
The IR builder (``python/tvm/script/ir_builder/``, backed by C++ in ``src/script/ir_builder/``)
|
||||
provides a frame-stack API for constructing IR incrementally.
|
||||
|
||||
Frame stack
|
||||
~~~~~~~~~~~
|
||||
|
||||
The core idea: each IR scope (module, function, block, loop) is a **frame**. Frames are pushed
|
||||
on ``__enter__`` and popped on ``__exit__``. When a frame exits, it finalizes the IR it
|
||||
represents and attaches it to the parent frame.
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
IRBuilder (thread-local singleton)
|
||||
└── frame stack:
|
||||
├── IRModuleFrame ← @I.ir_module
|
||||
│ ├── PrimFuncFrame ← @T.prim_func
|
||||
│ │ ├── ForFrame ← T.grid(...) / T.serial(...)
|
||||
│ │ │ └── SBlockFrame ← T.sblock(...)
|
||||
│ │ └── ...
|
||||
│ └── FunctionFrame ← @R.function
|
||||
│ └── BindingBlockFrame ← R.dataflow()
|
||||
└── ...
|
||||
|
||||
This design means the parser never needs to build a complete IR tree in memory — it
|
||||
constructs IR top-down by entering and exiting frames, and each frame handles its own
|
||||
finalization.
|
||||
|
||||
TIR builder
|
||||
~~~~~~~~~~~
|
||||
|
||||
The TIR builder (``ir_builder/tirx/ir.py``) provides functions that map directly to TVMScript
|
||||
syntax. Key categories:
|
||||
|
||||
**Function and block**:
|
||||
|
||||
- ``T.prim_func()`` → ``PrimFuncFrame``
|
||||
- ``T.sblock(name)`` → ``SBlockFrame`` (spatial block)
|
||||
- ``T.init()`` → ``BlockInitFrame`` (reduction initialization)
|
||||
- ``T.reads(...)``, ``T.writes(...)`` → declare buffer access regions
|
||||
|
||||
**Loops**:
|
||||
|
||||
- ``T.grid(*extents)`` → ``ForFrame`` returning loop variables
|
||||
- ``T.serial(start, stop)``, ``T.parallel(...)``, ``T.vectorized(...)``,
|
||||
``T.unroll(...)``, ``T.thread_binding(...)`` → loop with specific iterator type
|
||||
|
||||
**Block axes**:
|
||||
|
||||
- ``T.axis.spatial(dom, binding)`` — spatial iteration axis
|
||||
- ``T.axis.reduce(dom, binding)`` — reduction axis
|
||||
- ``T.axis.remap(kinds, bindings)`` — shorthand for multiple axes
|
||||
|
||||
**Buffers**:
|
||||
|
||||
- ``T.match_buffer(param, shape, dtype)`` — match function parameter to buffer
|
||||
- ``T.alloc_buffer(shape, dtype)`` — allocate intermediate buffer
|
||||
- ``T.Buffer(shape, dtype)`` — buffer type annotation in function signatures
|
||||
|
||||
Relax builder
|
||||
~~~~~~~~~~~~~
|
||||
|
||||
The Relax builder (``ir_builder/relax/ir.py``) provides:
|
||||
|
||||
**Function and dataflow**:
|
||||
|
||||
- ``R.function()`` → ``FunctionFrame``
|
||||
- ``R.dataflow()`` → ``BindingBlockFrame``
|
||||
- ``R.output(*vars)`` → expose variables from a dataflow block
|
||||
|
||||
**Emit**:
|
||||
|
||||
- ``R.emit(value)`` → emit a binding, returns a ``Var``
|
||||
- ``R.emit_match_cast(value, ty)`` → emit with type assertion
|
||||
|
||||
**Type annotations**:
|
||||
|
||||
- ``R.Tensor(shape, dtype)`` — tensor type
|
||||
- ``R.Tuple(*fields)`` — tuple type
|
||||
- ``R.Shape(values)`` — shape type
|
||||
- ``R.Any()`` — any Relax value type
|
||||
|
||||
**Calling conventions**:
|
||||
|
||||
- ``R.call_tir(func, args, out_ty)`` — call a TIR function
|
||||
- ``R.call_packed(name, *args)`` — call a PackedFunc
|
||||
- ``R.call_dps_packed(func, *args)`` — call using destination-passing style
|
||||
|
||||
**Operators**: the ``R`` module also re-exports all Relax operators
|
||||
(``R.add``, ``R.matmul``, ``R.nn.conv2d``, etc.) so they can be used directly in TVMScript.
|
||||
|
||||
|
||||
Printer Architecture
|
||||
--------------------
|
||||
|
||||
The printer converts TVM IR back to TVMScript text. It is implemented primarily in C++
|
||||
(``src/script/printer/``) for performance.
|
||||
|
||||
Doc tree
|
||||
~~~~~~~~
|
||||
|
||||
The printer does **not** generate text directly. Instead, it first builds a ``Doc`` tree — an
|
||||
intermediate representation that mirrors Python syntax:
|
||||
|
||||
- **Expression docs**: ``IdDoc``, ``AttrAccessDoc``, ``CallDoc``, ``IndexDoc``,
|
||||
``OperationDoc``, ``LiteralDoc``, ``TupleDoc``, ``ListDoc``, etc.
|
||||
- **Statement docs**: ``AssignDoc``, ``ForDoc``, ``IfDoc``, ``ScopeDoc`` (``with`` blocks),
|
||||
``FunctionDoc``, ``ClassDoc``, ``ReturnDoc``, ``CommentDoc``, etc.
|
||||
|
||||
For example, ``T.axis.spatial(128, i)`` is represented as:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
CallDoc(
|
||||
callee=AttrAccessDoc(AttrAccessDoc(IdDoc("T"), "axis"), "spatial"),
|
||||
args=[LiteralDoc(128), IdDoc("i")]
|
||||
)
|
||||
|
||||
IRDocsifier
|
||||
~~~~~~~~~~~
|
||||
|
||||
The ``IRDocsifier`` (``include/tvm/script/printer/ir_docsifier.h``) is the main dispatcher.
|
||||
It maintains:
|
||||
|
||||
- A dispatch table mapping ``(token, type_index)`` pairs to converter functions.
|
||||
- A frame stack for tracking the current scope (similar to the builder's frame stack).
|
||||
- A variable-to-name mapping to produce readable names.
|
||||
|
||||
Each IR dialect registers its own converters:
|
||||
|
||||
- ``src/script/printer/tirx/`` — converts PrimFunc, Buffer, SBlock, loops, expressions.
|
||||
- ``src/script/printer/relax/`` — converts relax.Function, bindings, types, operators.
|
||||
- ``src/script/printer/ir/`` — converts IRModule, shared types.
|
||||
|
||||
The final step calls ``DocToPythonScript()`` (``src/script/printer/doc_printer/python_doc_printer.cc``)
|
||||
to format the Doc tree into properly indented Python text.
|
||||
|
||||
Roundtrip guarantee
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
For any ``IRModule`` constructed through the compiler:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
text = mod.script() # IR → TVMScript text
|
||||
reparsed = tvm.script.from_source(text) # text → IR
|
||||
tvm.ir.assert_structural_equal(mod, reparsed)
|
||||
|
||||
This roundtrip property is relied upon by testing infrastructure and serialization workflows.
|
||||
Note that the printed text may differ from hand-written TVMScript — the printer uses canonical
|
||||
forms (e.g., explicit ``R.emit`` calls, fully qualified buffer annotations) that are not required
|
||||
in hand-written code.
|
||||
|
||||
|
||||
Supported Python Syntax
|
||||
-----------------------
|
||||
|
||||
TVMScript supports a subset of Python syntax. The table below summarizes what is supported
|
||||
and how each construct is interpreted:
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 25 15 60
|
||||
|
||||
* - Python Syntax
|
||||
- TIR
|
||||
- Relax
|
||||
* - ``for i in range(n)``
|
||||
- Serial loop nest
|
||||
- Not supported (no Relax-level ``for`` handler)
|
||||
* - ``with T.sblock(...)``
|
||||
- Spatial block scope
|
||||
- N/A
|
||||
* - ``with R.dataflow()``
|
||||
- N/A
|
||||
- Dataflow block
|
||||
* - ``if ... else``
|
||||
- TIR ``If`` branch (PrimExpr condition) or static eval (Python bool)
|
||||
- Relax ``If`` node (plain Python ``if cond:`` syntax)
|
||||
* - ``while``
|
||||
- ``T.While`` loop
|
||||
- Not supported
|
||||
* - ``x = expr``
|
||||
- Variable binding
|
||||
- Emit binding (implicit ``R.emit``)
|
||||
* - ``x: T.Buffer(...)``
|
||||
- Buffer annotation
|
||||
- N/A
|
||||
* - ``x: R.Tensor(...)``
|
||||
- N/A
|
||||
- Struct info annotation
|
||||
* - ``return``
|
||||
- Not used
|
||||
- Function return value
|
||||
* - ``A[i, j]``
|
||||
- Buffer load
|
||||
- Not applicable (use operators)
|
||||
* - ``A[i, j] = expr``
|
||||
- Buffer store
|
||||
- Not applicable
|
||||
* - Arithmetic (``+``, ``-``, etc.)
|
||||
- PrimExpr operations
|
||||
- Calls to Relax operators
|
||||
* - Function calls
|
||||
- ``T.*`` intrinsics
|
||||
- ``R.*`` operators or ``call_tir`` / ``call_packed``
|
||||
|
||||
**Not supported**: ``class`` definitions (except for ``@I.ir_module``), ``try/except``,
|
||||
``yield``, ``async/await``, list comprehensions, ``lambda``, ``import``, and ``global``
|
||||
statements.
|
||||
|
||||
|
||||
TIR Syntax Reference
|
||||
---------------------
|
||||
|
||||
Function definition
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@T.prim_func
|
||||
def func_name(a: T.handle, b: T.handle):
|
||||
A = T.match_buffer(a, (m, n), "float32")
|
||||
B = T.match_buffer(b, (m,), "float32")
|
||||
# function body
|
||||
|
||||
- ``T.handle`` — opaque handle parameter (matched to a buffer inside the function).
|
||||
- ``T.Buffer(shape, dtype)`` — can also be used directly in the signature:
|
||||
``def func(A: T.Buffer((128,), "float32"))``.
|
||||
|
||||
Block and axes
|
||||
~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock("block_name"):
|
||||
vi = T.axis.spatial(128, i)
|
||||
vj = T.axis.reduce(128, j)
|
||||
T.reads(A[vi, vj])
|
||||
T.writes(B[vi])
|
||||
# compute
|
||||
|
||||
- ``T.axis.spatial`` / ``T.axis.reduce`` / ``T.axis.scan`` — declare axis variables with
|
||||
their iteration domain and binding to outer loop variables.
|
||||
- ``T.axis.remap("SR", [i, j])`` — shorthand: ``S`` = spatial, ``R`` = reduce.
|
||||
- ``T.reads(...)``, ``T.writes(...)`` — declare buffer regions accessed by this block.
|
||||
|
||||
Loop types
|
||||
~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
for i in T.serial(0, 128): # sequential
|
||||
for i in T.parallel(0, 128): # parallel
|
||||
for i in T.vectorized(0, 128): # vectorized
|
||||
for i in T.unroll(0, 128): # unrolled
|
||||
for i in T.thread_binding(0, 128, thread="threadIdx.x"): # GPU thread
|
||||
|
||||
Buffer operations
|
||||
~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
C = T.alloc_buffer((128, 128), "float32") # intermediate buffer
|
||||
val = A[i, j] # buffer load
|
||||
B[i] = val + 1.0 # buffer store
|
||||
|
||||
Common intrinsics
|
||||
~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
T.exp(x), T.log(x), T.sqrt(x), T.tanh(x), ... # math functions
|
||||
T.cast(x, "float16") # type cast
|
||||
T.if_then_else(cond, true_val, false_val) # conditional expression
|
||||
T.min(a, b), T.max(a, b) # min/max
|
||||
T.call_extern("func_name", *args) # external function call
|
||||
T.call_packed("func_name", *args) # packed function call
|
||||
T.tvm_storage_sync("shared") # GPU memory fence
|
||||
|
||||
|
||||
Relax Syntax Reference
|
||||
-----------------------
|
||||
|
||||
Function definition
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@R.function
|
||||
def main(x: R.Tensor((128, 128), "float32"),
|
||||
y: R.Tensor((128,), "float32")) -> R.Tensor((128, 128), "float32"):
|
||||
# function body
|
||||
return result
|
||||
|
||||
- ``R.Tensor(shape, dtype)`` — tensor type annotation.
|
||||
- ``R.Tuple(...)``, ``R.Shape(...)``, ``R.Any()`` — other Relax type annotations.
|
||||
- ``R.function(private=True)`` — marks the function as module-private.
|
||||
- ``R.function(pure=False)`` — marks the function as having side effects.
|
||||
|
||||
Dataflow blocks
|
||||
~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
with R.dataflow():
|
||||
lv0 = R.add(x, y)
|
||||
lv1 = R.nn.relu(lv0)
|
||||
R.output(lv1)
|
||||
|
||||
Variables inside a ``R.dataflow()`` block are local to that block. ``R.output(...)`` exposes
|
||||
variables to the outer scope.
|
||||
|
||||
Calling TIR functions
|
||||
~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
out = R.call_tir(cls.my_kernel, (x, y), out_ty=R.Tensor((128,), "float32"))
|
||||
|
||||
- ``cls.my_kernel`` — references a TIR ``PrimFunc`` in the same module.
|
||||
- ``out_ty`` — the type (shape and dtype) of the output tensor.
|
||||
|
||||
Control flow
|
||||
~~~~~~~~~~~~
|
||||
|
||||
Relax ``if`` uses plain Python ``if`` syntax. The condition must be a Relax variable with
|
||||
boolean type. Both branches are required.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@R.function
|
||||
def f(cond: R.Tensor((), "bool"), x: R.Tensor((128,), "float32")):
|
||||
if cond:
|
||||
result = R.add(x, x)
|
||||
else:
|
||||
result = R.multiply(x, x)
|
||||
return result
|
||||
|
||||
|
||||
Source Code Map
|
||||
---------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 50 50
|
||||
|
||||
* - Path
|
||||
- Contents
|
||||
* - ``python/tvm/script/parser/core/``
|
||||
- Core parser: dispatch, expression evaluator, variable table, Doc AST
|
||||
* - ``python/tvm/script/parser/tirx/``
|
||||
- TIR-specific parser handlers and value binding
|
||||
* - ``python/tvm/script/parser/relax/``
|
||||
- Relax-specific parser handlers and value binding
|
||||
* - ``python/tvm/script/parser/ir/``
|
||||
- ``@I.ir_module`` entry point and module-level parsing
|
||||
* - ``python/tvm/script/ir_builder/base.py``
|
||||
- IRBuilder base class and frame stack mechanism
|
||||
* - ``python/tvm/script/ir_builder/tirx/``
|
||||
- TIR frame types and builder functions (``T.*``)
|
||||
* - ``python/tvm/script/ir_builder/relax/``
|
||||
- Relax frame types and builder functions (``R.*``)
|
||||
* - ``python/tvm/script/ir_builder/ir/``
|
||||
- IRModule builder (``I.*``)
|
||||
* - ``src/script/printer/``
|
||||
- C++ printer: Doc tree, IRDocsifier, Python code generation
|
||||
* - ``src/script/printer/tirx/``
|
||||
- TIR-specific IR-to-Doc converters
|
||||
* - ``src/script/printer/relax/``
|
||||
- Relax-specific IR-to-Doc converters
|
||||
* - ``src/script/ir_builder/``
|
||||
- C++ backend for frame stack and IR construction
|
||||
* - ``include/tvm/script/printer/``
|
||||
- C++ headers: Doc classes, IRDocsifier, dispatch functor
|
||||
* - ``include/tvm/script/ir_builder/``
|
||||
- C++ headers: builder base, dialect-specific frame types
|
||||
Reference in New Issue
Block a user