chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,2 @@
|
||||
HOW TO
|
||||
------
|
||||
@@ -0,0 +1,375 @@
|
||||
# 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.
|
||||
|
||||
"""
|
||||
.. _tutorial-bring-your-own-codegen:
|
||||
|
||||
Bring Your Own Codegen
|
||||
======================
|
||||
|
||||
TVM's Bring Your Own Codegen (BYOC) framework lets you offload parts of a model
|
||||
to a custom backend -- a hardware accelerator, an inference library, or your own
|
||||
kernels -- while TVM compiles the rest. This tutorial has two parts:
|
||||
|
||||
- **How BYOC works** -- we teach the flow with a bundled, hardware-free *example
|
||||
NPU* backend and then drive the **same flow** on a real production backend,
|
||||
NVIDIA TensorRT. Both run a small, hand-written model so every step is
|
||||
visible; the only thing that changes between them is the backend, and that
|
||||
contrast is the lesson.
|
||||
- **Deploying a real model** -- we then put it to work, taking an actual PyTorch
|
||||
``nn.Module`` from export through TensorRT and running it on the GPU.
|
||||
|
||||
The example NPU is a teaching stub: its runtime logs the dispatch decisions an
|
||||
NPU would make (memory tier, execution engine, fusion) but performs no real
|
||||
computation, so its output buffers are left uninitialized. We therefore check
|
||||
*shapes*, not values, in the NPU sections -- its job is to make every BYOC step
|
||||
visible with nothing hidden. TensorRT then runs the identical flow for real, so
|
||||
we cross-check its result against a reference.
|
||||
|
||||
**Prerequisites**: the example NPU sections need TVM built with
|
||||
``USE_EXAMPLE_NPU_CODEGEN=ON`` and ``USE_EXAMPLE_NPU_RUNTIME=ON``; the TensorRT
|
||||
sections need ``USE_TENSORRT_CODEGEN=ON``, ``USE_TENSORRT_RUNTIME=ON`` and
|
||||
``USE_CUDA=ON`` plus a CUDA GPU and a matching TensorRT install (from NVIDIA's
|
||||
``pip install tensorrt`` packages or the TensorRT archive); the final deployment
|
||||
section also needs PyTorch. Each section degrades gracefully when its backend is
|
||||
unavailable.
|
||||
"""
|
||||
|
||||
######################################################################
|
||||
# Overview of the BYOC flow
|
||||
# -------------------------
|
||||
#
|
||||
# BYOC plugs a custom backend into TVM's compilation pipeline in four steps:
|
||||
#
|
||||
# 1. **Register patterns** - describe which sequences of Relax ops the backend
|
||||
# can handle.
|
||||
# 2. **Partition the graph** - group matched ops into composite functions.
|
||||
# 3. **Run codegen** - lower each composite to the backend's representation
|
||||
# (a JSON graph for both backends in this tutorial).
|
||||
# 4. **Execute** - the runtime dispatches each composite to the backend.
|
||||
#
|
||||
# Steps 1 and 2 are pure Python and run anywhere; steps 3 and 4 need the
|
||||
# backend's codegen and runtime compiled into TVM, which is why the
|
||||
# build-and-run cells below are guarded.
|
||||
|
||||
######################################################################
|
||||
# Step 1: Import the backends to register their patterns
|
||||
# ------------------------------------------------------
|
||||
#
|
||||
# Importing a backend module registers its patterns with TVM's global registry.
|
||||
# Pattern registration is independent of the C++ build -- only codegen and the
|
||||
# runtime require the backend to be compiled in -- so we probe each backend and
|
||||
# guard the build-and-run cells accordingly.
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
|
||||
import tvm
|
||||
import tvm.relax.backend.contrib.example_npu
|
||||
from tvm import relax
|
||||
from tvm.relax.backend.contrib.tensorrt import partition_for_tensorrt
|
||||
from tvm.relax.backend.pattern_registry import get_patterns_with_prefix
|
||||
from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions, RunCodegen
|
||||
from tvm.script import relax as R
|
||||
|
||||
has_example_npu_codegen = tvm.get_global_func("relax.ext.example_npu", True)
|
||||
has_example_npu_runtime = tvm.get_global_func("runtime.ExampleNPUJSONRuntimeCreate", True)
|
||||
has_example_npu = has_example_npu_codegen and has_example_npu_runtime
|
||||
|
||||
has_tensorrt_codegen = tvm.get_global_func("relax.ext.tensorrt", True) is not None
|
||||
_is_trt_runtime_enabled = tvm.get_global_func("relax.is_tensorrt_runtime_enabled", True)
|
||||
has_tensorrt = (
|
||||
has_tensorrt_codegen and _is_trt_runtime_enabled is not None and _is_trt_runtime_enabled()
|
||||
)
|
||||
has_cuda = tvm.cuda(0).exist
|
||||
|
||||
######################################################################
|
||||
# Step 2: Define the model
|
||||
# ------------------------
|
||||
#
|
||||
# A single convolution followed by a ReLU. This one model is used for both
|
||||
# backends.
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class ConvReLU:
|
||||
@R.function
|
||||
def main(
|
||||
data: R.Tensor((1, 3, 32, 32), "float32"),
|
||||
weight: R.Tensor((16, 3, 3, 3), "float32"),
|
||||
) -> R.Tensor((1, 16, 30, 30), "float32"):
|
||||
with R.dataflow():
|
||||
conv = relax.op.nn.conv2d(data, weight)
|
||||
out = relax.op.nn.relu(conv)
|
||||
R.output(out)
|
||||
return out
|
||||
|
||||
|
||||
######################################################################
|
||||
# Step 3: Partition for the example NPU
|
||||
# -------------------------------------
|
||||
#
|
||||
# ``FuseOpsByPattern`` groups ops matching a registered pattern into composite
|
||||
# functions; ``MergeCompositeFunctions`` then consolidates adjacent composites
|
||||
# bound for the same backend into a single external call. Two flags steer
|
||||
# partitioning:
|
||||
#
|
||||
# - ``bind_constants=False`` keeps weights as function arguments, so the host
|
||||
# stays in charge of the parameters. (TensorRT below makes the opposite
|
||||
# choice: it binds weights as constants because it bakes them into its engine.)
|
||||
# - ``annotate_codegen=True`` wraps each matched composite in a function tagged
|
||||
# with the backend name -- the tag ``RunCodegen`` routes on. (The follow-up
|
||||
# ``MergeCompositeFunctions`` also attaches this tag when it groups composites,
|
||||
# which is why ``partition_for_tensorrt`` below can leave the flag off.)
|
||||
#
|
||||
# The example NPU registers a fused ``conv2d + relu`` pattern with higher
|
||||
# priority than the standalone ``conv2d`` pattern, so the two ops collapse into a
|
||||
# single ``example_npu.conv2d_relu_fused`` composite -- look for it in the
|
||||
# printed module.
|
||||
|
||||
npu_patterns = get_patterns_with_prefix("example_npu")
|
||||
npu_mod = FuseOpsByPattern(npu_patterns, bind_constants=False, annotate_codegen=True)(ConvReLU)
|
||||
npu_mod = MergeCompositeFunctions()(npu_mod)
|
||||
print("After partitioning for the example NPU:")
|
||||
print(npu_mod)
|
||||
|
||||
######################################################################
|
||||
# Step 4: Codegen, build and run on the example NPU
|
||||
# -------------------------------------------------
|
||||
#
|
||||
# ``RunCodegen`` invokes each annotated composite's backend codegen, replacing it
|
||||
# with the backend runtime module (here, the NPU's JSON graph); ``relax.build``
|
||||
# then compiles the remaining host-side program and links everything. Because
|
||||
# the runtime is a stub that computes nothing, we assert on the output *shape*
|
||||
# only -- the values are uninitialized.
|
||||
|
||||
np.random.seed(0)
|
||||
data_np = np.random.randn(1, 3, 32, 32).astype("float32")
|
||||
weight_np = np.random.randn(16, 3, 3, 3).astype("float32")
|
||||
|
||||
if has_example_npu:
|
||||
npu_mod = RunCodegen()(npu_mod)
|
||||
|
||||
with tvm.transform.PassContext(opt_level=3):
|
||||
npu_exec = relax.build(npu_mod, tvm.target.Target("llvm"))
|
||||
|
||||
npu_vm = relax.VirtualMachine(npu_exec, tvm.cpu())
|
||||
npu_out = npu_vm["main"](
|
||||
tvm.runtime.tensor(data_np, tvm.cpu()), tvm.runtime.tensor(weight_np, tvm.cpu())
|
||||
)
|
||||
assert npu_out.numpy().shape == (1, 16, 30, 30)
|
||||
print("Example NPU run completed. Output shape:", npu_out.numpy().shape)
|
||||
else:
|
||||
print("Example NPU backend unavailable; skipping its build and run.")
|
||||
|
||||
######################################################################
|
||||
# The same flow on a real backend: TensorRT
|
||||
# -----------------------------------------
|
||||
#
|
||||
# Steps 1-4 above are the whole mechanism. Aiming them at a real backend
|
||||
# changes very little, so rather than repeat the walkthrough, here is only what
|
||||
# differs for NVIDIA TensorRT:
|
||||
#
|
||||
# - **Partition in one call.** ``partition_for_tensorrt`` bundles the
|
||||
# ``FuseOpsByPattern`` + ``MergeCompositeFunctions`` you ran by hand, using
|
||||
# TensorRT's own pattern table.
|
||||
# - **Weights become constants** (``bind_constants=True``): TensorRT bakes them
|
||||
# into the engine it builds, so bind the parameters before partitioning.
|
||||
# - **Real values.** TensorRT actually computes, so we build for CUDA, run on
|
||||
# the GPU, and cross-check against a plain CPU build -- not just the shape.
|
||||
#
|
||||
# The build-and-run cells below execute only when TensorRT and CUDA are
|
||||
# available. In CPU-only documentation builds, they produce no output.
|
||||
|
||||
trt_mod = relax.transform.BindParams("main", {"weight": weight_np})(ConvReLU)
|
||||
trt_mod = partition_for_tensorrt(trt_mod)
|
||||
print("After partition_for_tensorrt:")
|
||||
print(trt_mod)
|
||||
|
||||
######################################################################
|
||||
# Build for CUDA, run on the GPU, and compare against the CPU reference.
|
||||
|
||||
if has_tensorrt and has_cuda:
|
||||
dev = tvm.cuda(0)
|
||||
with tvm.transform.PassContext(opt_level=3):
|
||||
trt_exec = relax.build(RunCodegen()(trt_mod), "cuda")
|
||||
trt_out = relax.VirtualMachine(trt_exec, dev)["main"](tvm.runtime.tensor(data_np, dev)).numpy()
|
||||
|
||||
cpu_mod = relax.transform.LegalizeOps()(
|
||||
relax.transform.BindParams("main", {"weight": weight_np})(ConvReLU)
|
||||
)
|
||||
cpu_exec = relax.build(cpu_mod, "llvm")
|
||||
cpu_out = relax.VirtualMachine(cpu_exec, tvm.cpu())["main"](
|
||||
tvm.runtime.tensor(data_np, tvm.cpu())
|
||||
).numpy()
|
||||
|
||||
np.testing.assert_allclose(trt_out, cpu_out, rtol=1e-2, atol=1e-2)
|
||||
print("TensorRT output shape:", trt_out.shape, "- matches the CPU reference.")
|
||||
|
||||
######################################################################
|
||||
# A real backend also exposes knobs the stub does not. Setting ``use_fp16``
|
||||
# through the ``relax.ext.tensorrt.options`` config lets TensorRT pick FP16
|
||||
# kernels, trading a little accuracy for speed; nothing else about the flow
|
||||
# changes. (Other options are environment-driven: ``TVM_TENSORRT_USE_INT8``
|
||||
# enables INT8 with calibration, ``TVM_TENSORRT_MAX_WORKSPACE_SIZE`` caps the
|
||||
# build workspace, and ``TVM_TENSORRT_CACHE_DIR`` caches built engines to disk
|
||||
# for reuse across runs.)
|
||||
|
||||
if has_tensorrt and has_cuda:
|
||||
fp16_mod = partition_for_tensorrt(
|
||||
relax.transform.BindParams("main", {"weight": weight_np})(ConvReLU)
|
||||
)
|
||||
with tvm.transform.PassContext(
|
||||
opt_level=3, config={"relax.ext.tensorrt.options": {"use_fp16": True}}
|
||||
):
|
||||
fp16_exec = relax.build(RunCodegen()(fp16_mod), "cuda")
|
||||
fp16_out = relax.VirtualMachine(fp16_exec, tvm.cuda(0))["main"](
|
||||
tvm.runtime.tensor(data_np, tvm.cuda(0))
|
||||
).numpy()
|
||||
|
||||
np.testing.assert_allclose(fp16_out, cpu_out, rtol=5e-2, atol=5e-2)
|
||||
print("TensorRT FP16 output shape:", fp16_out.shape, "- matches within FP16 tolerance.")
|
||||
|
||||
######################################################################
|
||||
# Example NPU vs TensorRT at a glance
|
||||
# -----------------------------------
|
||||
#
|
||||
# The same four-step flow, two backends:
|
||||
#
|
||||
# ========= ============================== ==================================
|
||||
# Aspect Example NPU (teaching stub) TensorRT (real backend)
|
||||
# ========= ============================== ==================================
|
||||
# Runtime logs decisions, no compute builds and runs an nvinfer engine
|
||||
# Output uninitialized (check shape) real values (cross-checked vs CPU)
|
||||
# Weights ``bind_constants=False`` ``bind_constants=True`` (baked in)
|
||||
# Partition two passes, by hand ``partition_for_tensorrt`` one call
|
||||
# ========= ============================== ==================================
|
||||
|
||||
######################################################################
|
||||
# Deploying a PyTorch model with TensorRT
|
||||
# ---------------------------------------
|
||||
#
|
||||
# Everything above used a hand-written ``IRModule`` so each op was visible. In
|
||||
# practice you start from a trained model. This final section runs the *same*
|
||||
# ``partition_for_tensorrt`` flow on a real PyTorch ``nn.Module``, end to end:
|
||||
# export it, import it into Relax with the PyTorch frontend (the weights come in
|
||||
# as constants -- exactly what TensorRT bakes into its engine), partition, build
|
||||
# for CUDA, and check the GPU result against PyTorch's own output. Beyond the
|
||||
# frontend import, the only difference is that the imported program returns its
|
||||
# outputs as a tuple, so we index ``[0]`` for the single result tensor; the
|
||||
# partition-build-run flow is otherwise unchanged.
|
||||
#
|
||||
# This section additionally requires PyTorch.
|
||||
|
||||
try:
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
has_torch = True
|
||||
except ImportError:
|
||||
has_torch = False
|
||||
|
||||
if has_torch and has_tensorrt and has_cuda:
|
||||
from tvm.relax.frontend.torch import from_exported_program
|
||||
|
||||
class SmallConvNet(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.conv1 = nn.Conv2d(3, 8, 3)
|
||||
self.conv2 = nn.Conv2d(8, 16, 3)
|
||||
self.pool = nn.MaxPool2d(2)
|
||||
|
||||
def forward(self, x):
|
||||
x = torch.relu(self.conv1(x))
|
||||
x = self.pool(x)
|
||||
x = torch.relu(self.conv2(x))
|
||||
return x
|
||||
|
||||
torch_model = SmallConvNet().eval()
|
||||
example_input = torch.randn(1, 3, 32, 32)
|
||||
with torch.no_grad():
|
||||
torch_ref = torch_model(example_input).numpy()
|
||||
exported = torch.export.export(torch_model, (example_input,))
|
||||
|
||||
torch_mod = from_exported_program(exported)
|
||||
torch_mod = partition_for_tensorrt(torch_mod)
|
||||
print("After importing and partitioning the PyTorch model:")
|
||||
print(torch_mod)
|
||||
|
||||
torch_dev = tvm.cuda(0)
|
||||
with tvm.transform.PassContext(opt_level=3):
|
||||
torch_exec = relax.build(RunCodegen()(torch_mod), "cuda")
|
||||
deployed = relax.VirtualMachine(torch_exec, torch_dev)["main"](
|
||||
tvm.runtime.tensor(example_input.numpy(), torch_dev)
|
||||
)[0].numpy()
|
||||
|
||||
np.testing.assert_allclose(deployed, torch_ref, rtol=1e-2, atol=1e-2)
|
||||
print("Deployed PyTorch model on TensorRT; output", deployed.shape, "matches PyTorch.")
|
||||
|
||||
######################################################################
|
||||
# Real deployment builds once and reuses the artifact. Export the compiled
|
||||
# module to a shared library, then load and run it later -- in a fresh process,
|
||||
# with no PyTorch and no rebuild needed.
|
||||
|
||||
if has_torch and has_tensorrt and has_cuda:
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
lib_path = os.path.join(tmpdir, "deployed_trt.so")
|
||||
torch_exec.export_library(lib_path)
|
||||
loaded = tvm.runtime.load_module(lib_path)
|
||||
reran = relax.VirtualMachine(loaded, torch_dev)["main"](
|
||||
tvm.runtime.tensor(example_input.numpy(), torch_dev)
|
||||
)[0].numpy()
|
||||
np.testing.assert_allclose(reran, torch_ref, rtol=1e-2, atol=1e-2)
|
||||
print("Reloaded the exported library and reran; output", reran.shape, "still matches.")
|
||||
|
||||
######################################################################
|
||||
# Notes for real deployments
|
||||
# --------------------------
|
||||
#
|
||||
# - **Operator coverage and fallback.** TensorRT offloads only the ops in its
|
||||
# pattern table (see ``python/tvm/relax/backend/contrib/tensorrt.py``);
|
||||
# anything unsupported simply stays on the host. Print the partitioned module
|
||||
# and look for the ``Codegen: "tensorrt"`` functions to see what was offloaded.
|
||||
# - **Dynamic shapes.** The builder sets up an optimization profile for a dynamic
|
||||
# leading (batch) dimension, so the integration can serve a model exported with
|
||||
# a symbolic batch size.
|
||||
# - **Engine build cost.** Building a TensorRT engine is slow the first time (it
|
||||
# is not a hang). Set ``TVM_TENSORRT_CACHE_DIR`` to cache built engines to
|
||||
# disk and skip the rebuild on later runs.
|
||||
|
||||
######################################################################
|
||||
# Next steps
|
||||
# ----------
|
||||
#
|
||||
# To build your own backend using the example NPU as a starting point:
|
||||
#
|
||||
# - Replace the stub runtime in
|
||||
# ``src/runtime/extra/contrib/example_npu/example_npu_runtime.cc`` with your
|
||||
# hardware SDK calls.
|
||||
# - Extend ``patterns.py`` with the ops your hardware supports.
|
||||
# - Add a C++ codegen under ``src/relax/backend/contrib/`` if your backend needs
|
||||
# a non-JSON serialization format.
|
||||
# - Add a CMake module under ``cmake/modules/contrib/`` following
|
||||
# ``ExampleNPU.cmake``.
|
||||
#
|
||||
# For a complete real-backend implementation to study, see the TensorRT
|
||||
# integration: the pattern table and ``partition_for_tensorrt`` in
|
||||
# ``python/tvm/relax/backend/contrib/tensorrt.py``, the codegen in
|
||||
# ``src/relax/backend/contrib/tensorrt/``, and the runtime in
|
||||
# ``src/runtime/extra/contrib/tensorrt/``.
|
||||
@@ -0,0 +1,786 @@
|
||||
# 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.
|
||||
# ruff: noqa: E501
|
||||
"""
|
||||
.. _tutorial-cross-compilation-and-rpc:
|
||||
|
||||
Cross Compilation and RPC
|
||||
=========================
|
||||
**Author**: `Ziheng Jiang <https://github.com/ZihengJiang/>`_, `Lianmin Zheng <https://github.com/merrymercy/>`_
|
||||
|
||||
This tutorial introduces cross compilation and remote device
|
||||
execution with RPC in TVM.
|
||||
|
||||
With cross compilation and RPC, you can **compile a program on your
|
||||
local machine then run it on the remote device**. It is useful when
|
||||
the remote device resource are limited, like Raspberry Pi and mobile
|
||||
platforms. In this tutorial, we will use the Raspberry Pi for a CPU example
|
||||
and the Firefly-RK3399 for an OpenCL example.
|
||||
"""
|
||||
|
||||
######################################################################
|
||||
# Build TVM Runtime on Device
|
||||
# ---------------------------
|
||||
#
|
||||
# The first step is to build the TVM runtime on the remote device.
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# All instructions in both this section and the next section should be
|
||||
# executed on the target device, e.g. Raspberry Pi. We assume the target
|
||||
# is running Linux.
|
||||
#
|
||||
# Since we do compilation on the local machine, the remote device is only used
|
||||
# for running the generated code. We only need to build the TVM runtime on
|
||||
# the remote device.
|
||||
#
|
||||
# .. code-block:: bash
|
||||
#
|
||||
# git clone --recursive https://github.com/apache/tvm tvm
|
||||
# cd tvm
|
||||
# mkdir build && cd build
|
||||
# cp ../cmake/config.cmake .
|
||||
# cmake .. && cmake --build . --parallel $(nproc)
|
||||
#
|
||||
# After building the runtime successfully, we need to set environment variables
|
||||
# in :code:`~/.bashrc` file. We can edit :code:`~/.bashrc`
|
||||
# using :code:`vi ~/.bashrc` and add the line below (Assuming your TVM
|
||||
# directory is in :code:`~/tvm`):
|
||||
#
|
||||
# .. code-block:: bash
|
||||
#
|
||||
# export PYTHONPATH=$PYTHONPATH:~/tvm/python
|
||||
#
|
||||
# To update the environment variables, execute :code:`source ~/.bashrc`.
|
||||
|
||||
######################################################################
|
||||
# Set Up RPC Server on Device
|
||||
# ---------------------------
|
||||
# To start an RPC server, run the following command on your remote device
|
||||
# (Which is Raspberry Pi in this example).
|
||||
#
|
||||
# .. code-block:: bash
|
||||
#
|
||||
# python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090
|
||||
#
|
||||
# If you see the line below, it means the RPC server started
|
||||
# successfully on your device.
|
||||
#
|
||||
# .. code-block:: bash
|
||||
#
|
||||
# INFO:root:RPCServer: bind to 0.0.0.0:9090
|
||||
#
|
||||
# Declare and Cross Compile Kernel on Local Machine
|
||||
# -------------------------------------------------
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# Now we go back to the local machine, which has a full TVM installed
|
||||
# (with LLVM).
|
||||
#
|
||||
# Here we will declare a simple kernel on the local machine:
|
||||
|
||||
import numpy as np
|
||||
import tvm_ffi
|
||||
|
||||
import tvm
|
||||
from tvm import rpc, te
|
||||
from tvm.support import utils
|
||||
|
||||
n = tvm.runtime.convert(1024)
|
||||
A = te.placeholder((n,), name="A")
|
||||
B = te.compute((n,), lambda i: A[i] + 1.0, name="B")
|
||||
mod = tvm.IRModule.from_expr(te.create_prim_func([A, B]).with_attr("global_symbol", "add_one"))
|
||||
|
||||
######################################################################
|
||||
# Then we cross compile the kernel.
|
||||
# The target should be {"kind": "llvm", "mtriple": "armv7l-linux-gnueabihf"} for
|
||||
# Raspberry Pi 3B, but we use 'llvm' here to make this tutorial runnable
|
||||
# on our webpage building server. See the detailed note in the following block.
|
||||
|
||||
local_demo = True
|
||||
|
||||
if local_demo:
|
||||
target = "llvm"
|
||||
else:
|
||||
target = {"kind": "llvm", "mtriple": "armv7l-linux-gnueabihf"}
|
||||
|
||||
func = tvm.compile(mod, target=target)
|
||||
# save the lib at a local temp folder
|
||||
temp = utils.tempdir()
|
||||
path = temp.relpath("lib.tar")
|
||||
func.export_library(path)
|
||||
|
||||
######################################################################
|
||||
# .. note::
|
||||
#
|
||||
# To run this tutorial with a real remote device, change :code:`local_demo`
|
||||
# to False and replace :code:`target` in :code:`build` with the appropriate
|
||||
# target triple for your device. The target triple which might be
|
||||
# different for different devices. For example, it is
|
||||
# :code:`{"kind": "llvm", "mtriple": "armv7l-linux-gnueabihf"}` for Raspberry Pi 3B and
|
||||
# :code:`{"kind": "llvm", "mtriple": "aarch64-linux-gnu"}` for RK3399.
|
||||
#
|
||||
# Usually, you can query the target by running :code:`gcc -v` on your
|
||||
# device, and looking for the line starting with :code:`Target:`
|
||||
# (Though it may still be a loose configuration.)
|
||||
#
|
||||
# Besides :code:`-mtriple`, you can also set other compilation options
|
||||
# like:
|
||||
#
|
||||
# * -mcpu=<cpuname>
|
||||
# Specify a specific chip in the current architecture to generate code for. By default this is inferred from the target triple and autodetected to the current architecture.
|
||||
# * -mattr=a1,+a2,-a3,...
|
||||
# Override or control specific attributes of the target, such as whether SIMD operations are enabled or not. The default set of attributes is set by the current CPU.
|
||||
# To get the list of available attributes, you can do:
|
||||
#
|
||||
# .. code-block:: bash
|
||||
#
|
||||
# llc -mtriple=<your device target triple> -mattr=help
|
||||
#
|
||||
# These options are consistent with `llc <http://llvm.org/docs/CommandGuide/llc.html>`_.
|
||||
# It is recommended to set target triple and feature set to contain specific
|
||||
# feature available, so we can take full advantage of the features of the
|
||||
# board.
|
||||
# You can find more details about cross compilation attributes from
|
||||
# `LLVM guide of cross compilation <https://clang.llvm.org/docs/CrossCompilation.html>`_.
|
||||
|
||||
######################################################################
|
||||
# Run CPU Kernel Remotely by RPC
|
||||
# ------------------------------
|
||||
# We show how to run the generated CPU kernel on the remote device.
|
||||
# First we obtain an RPC session from remote device.
|
||||
|
||||
if local_demo:
|
||||
remote = rpc.LocalSession()
|
||||
else:
|
||||
# The following is my environment, change this to the IP address of your target device
|
||||
host = "10.77.1.162"
|
||||
port = 9090
|
||||
remote = rpc.connect(host, port)
|
||||
|
||||
######################################################################
|
||||
# Upload the lib to the remote device, then invoke a device local
|
||||
# compiler to relink them. Now `func` is a remote module object.
|
||||
|
||||
remote.upload(path)
|
||||
func = remote.load_module("lib.tar")
|
||||
|
||||
# create arrays on the remote device
|
||||
dev = remote.cpu()
|
||||
a = tvm.runtime.tensor(np.random.uniform(size=1024).astype(A.dtype), dev)
|
||||
b = tvm.runtime.tensor(np.zeros(1024, dtype=A.dtype), dev)
|
||||
# the function will run on the remote device
|
||||
func(a, b)
|
||||
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
|
||||
|
||||
######################################################################
|
||||
# When you want to evaluate the performance of the kernel on the remote
|
||||
# device, it is important to avoid the overhead of network.
|
||||
# :code:`time_evaluator` will returns a remote function that runs the
|
||||
# function over number times, measures the cost per run on the remote
|
||||
# device and returns the measured cost. Network overhead is excluded.
|
||||
|
||||
time_f = func.time_evaluator("add_one", dev, number=10)
|
||||
cost = time_f(a, b).mean
|
||||
print(f"{cost:g} secs/op")
|
||||
|
||||
######################################################################
|
||||
# Scale RPC to Shared Devices
|
||||
# ---------------------------
|
||||
#
|
||||
# The direct RPC server used above is the simplest way to run on one remote
|
||||
# device. In shared environments, the same compile/upload/run flow is usually
|
||||
# kept, but the connection is managed by an RPC tracker and, when needed, an
|
||||
# RPC proxy.
|
||||
#
|
||||
# This setup is useful when:
|
||||
#
|
||||
# - multiple users or CI jobs share a small number of boards,
|
||||
# - devices are registered by key rather than by fixed IP address,
|
||||
# - the host cannot directly reach the device because of the network layout, or
|
||||
# - the target device only has the minimal runtime stack needed for execution.
|
||||
#
|
||||
# The pieces fit together as follows:
|
||||
#
|
||||
# - **RPC server**: runs on the target device and executes uploaded modules.
|
||||
# - **RPC tracker**: runs on a host and assigns matching RPC servers to clients.
|
||||
# - **RPC proxy**: forwards traffic when the client cannot connect directly to
|
||||
# the RPC server.
|
||||
#
|
||||
# .. figure:: https://raw.githubusercontent.com/tlc-pack/web-data/main/images/dev/how-to/rpc_system_suggested_arch.svg
|
||||
# :align: center
|
||||
# :width: 85%
|
||||
#
|
||||
# In the figure above, machine A connects through the tracker. Machine B runs
|
||||
# an RPC proxy because machines C and D are not directly reachable from A. The
|
||||
# tracker keeps a queue per RPC key. If a matching server is available, it is
|
||||
# assigned to the client; otherwise, the request waits in that key's queue.
|
||||
#
|
||||
# Start the Tracker and Proxy
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
#
|
||||
# The tracker and proxy generally run on a host machine, not on the target
|
||||
# device. They do not require target-specific drivers.
|
||||
#
|
||||
# .. code-block:: shell
|
||||
#
|
||||
# python3 -m tvm.exec.rpc_tracker --host RPC_TRACKER_IP --port 9190 --port-end 9191
|
||||
#
|
||||
# .. code-block:: shell
|
||||
#
|
||||
# python3 -m tvm.exec.rpc_proxy \
|
||||
# --host RPC_PROXY_IP \
|
||||
# --port 9090 \
|
||||
# --port-end 9091 \
|
||||
# --tracker RPC_TRACKER_IP:RPC_TRACKER_PORT
|
||||
#
|
||||
# Replace the host names, ports, and port ranges for your environment. The
|
||||
# ``--port-end`` option is useful in CI because it prevents the service from
|
||||
# silently choosing an unexpected port.
|
||||
#
|
||||
# Package a Minimal RPC Runtime
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
#
|
||||
# If the target can build TVM directly, install TVM on the target and launch
|
||||
# the RPC server there. Otherwise, cross-compile the TVM runtime and package
|
||||
# it with the Python RPC server.
|
||||
#
|
||||
# A typical CMake toolchain file for 64-bit ARM Linux looks like this:
|
||||
#
|
||||
# .. code-block:: cmake
|
||||
#
|
||||
# set(CMAKE_SYSTEM_NAME Linux)
|
||||
# set(root_dir "/XXX/gcc-linaro-7.5.0-2019.12-x86_64_aarch64-linux-gnu")
|
||||
#
|
||||
# set(CMAKE_C_COMPILER "${root_dir}/bin/aarch64-linux-gnu-gcc")
|
||||
# set(CMAKE_CXX_COMPILER "${root_dir}/bin/aarch64-linux-gnu-g++")
|
||||
# set(CMAKE_SYSROOT "${root_dir}/aarch64-linux-gnu/libc")
|
||||
#
|
||||
# set(CMAKE_FIND_ROOT_PATH_MODE_PROGRAM NEVER)
|
||||
# set(CMAKE_FIND_ROOT_PATH_MODE_LIBRARY ONLY)
|
||||
# set(CMAKE_FIND_ROOT_PATH_MODE_INCLUDE ONLY)
|
||||
# set(CMAKE_FIND_ROOT_PATH_MODE_PACKAGE ONLY)
|
||||
#
|
||||
# Build the runtime from the TVM repository root. Enable target-specific
|
||||
# options such as ``USE_OPENCL`` or vendor runtime support in ``config.cmake``.
|
||||
# Build any target-specific runtime libraries that your deployment needs, such
|
||||
# as ``tvm_runtime_opencl`` for OpenCL.
|
||||
#
|
||||
# .. code-block:: shell
|
||||
#
|
||||
# mkdir cross_build
|
||||
# cd cross_build
|
||||
# cp ../cmake/config.cmake ./
|
||||
#
|
||||
# # Enable other options as needed, e.g. USE_OPENCL or vendor runtimes.
|
||||
# sed -i "s|USE_LLVM.*)|USE_LLVM OFF)|" config.cmake
|
||||
#
|
||||
# cmake -DCMAKE_TOOLCHAIN_FILE=/YYY/aarch64-linux-gnu.cmake -DCMAKE_BUILD_TYPE=Release ..
|
||||
# cmake --build . --target runtime -j
|
||||
# # Optional example when USE_OPENCL is enabled:
|
||||
# # cmake --build . --target tvm_runtime_opencl -j
|
||||
# cd ..
|
||||
#
|
||||
# Then package the Python RPC server with the cross-compiled runtime and copy
|
||||
# it to the device.
|
||||
#
|
||||
# .. code-block:: shell
|
||||
#
|
||||
# rm -rf tvm_runtime_package
|
||||
# mkdir tvm_runtime_package
|
||||
# cp -a python tvm_runtime_package/
|
||||
# cp cross_build/lib/libtvm_ffi.so tvm_runtime_package/python/tvm/
|
||||
# cp cross_build/lib/libtvm_runtime*.so tvm_runtime_package/python/tvm/
|
||||
# tar -czf tvm_runtime.tar.gz -C tvm_runtime_package python
|
||||
#
|
||||
# On the target device:
|
||||
#
|
||||
# .. code-block:: shell
|
||||
#
|
||||
# tar -xzf tvm_runtime.tar.gz
|
||||
# export PYTHONPATH=`pwd`/python:${PYTHONPATH}
|
||||
#
|
||||
# Launch the Server
|
||||
# ~~~~~~~~~~~~~~~~~
|
||||
#
|
||||
# Launch the RPC server on the target device. Use the proxy form when the
|
||||
# server connects through an RPC proxy; otherwise connect directly to the
|
||||
# tracker.
|
||||
#
|
||||
# .. code-block:: shell
|
||||
#
|
||||
# # Through an RPC proxy.
|
||||
# python3 -m tvm.exec.rpc_server \
|
||||
# --host RPC_PROXY_IP \
|
||||
# --port RPC_PROXY_PORT \
|
||||
# --through-proxy \
|
||||
# --key RPC_KEY
|
||||
#
|
||||
# # Directly to an RPC tracker.
|
||||
# python3 -m tvm.exec.rpc_server \
|
||||
# --tracker RPC_TRACKER_IP:RPC_TRACKER_PORT \
|
||||
# --key RPC_KEY
|
||||
#
|
||||
# Query the tracker from the host to confirm that the servers are visible:
|
||||
#
|
||||
# .. code-block:: shell
|
||||
#
|
||||
# python3 -m tvm.exec.query_rpc_tracker --host RPC_TRACKER_IP --port RPC_TRACKER_PORT
|
||||
#
|
||||
# If three servers connect through a proxy, the output should look similar to:
|
||||
#
|
||||
# .. code-block:: text
|
||||
#
|
||||
# Tracker address RPC_TRACKER_IP:RPC_TRACKER_PORT
|
||||
#
|
||||
# Server List
|
||||
# ----------------------------
|
||||
# server-address key
|
||||
# ----------------------------
|
||||
# RPC_PROXY_IP:RPC_PROXY_PORT server:proxy[RPC_KEY0,RPC_KEY1,RPC_KEY2]
|
||||
# ----------------------------
|
||||
#
|
||||
# Queue Status
|
||||
# ---------------------------------------
|
||||
# key total free pending
|
||||
# ---------------------------------------
|
||||
# RPC_KEY0 0 0 3
|
||||
# ---------------------------------------
|
||||
#
|
||||
# Once the tracker assigns a server, the client-side code still follows the
|
||||
# same pattern used earlier in this tutorial. Only the session creation
|
||||
# changes from a direct connection to a tracker request:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# tracker = rpc.connect_tracker("RPC_TRACKER_IP", RPC_TRACKER_PORT)
|
||||
# remote = tracker.request("RPC_KEY", priority=0, session_timeout=600)
|
||||
#
|
||||
# After that, use the same ``remote.upload()``, ``remote.load_module()``, remote
|
||||
# device creation, and ``time_evaluator`` flow shown above.
|
||||
#
|
||||
# Troubleshooting Minimal Device Environments
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
#
|
||||
# Some target devices have intentionally small Python environments. The TVM
|
||||
# runtime itself does not require full NumPy support for RPC execution, but the
|
||||
# Python RPC server may import modules that import ``numpy``. If installing or
|
||||
# cross-compiling NumPy is not practical, a small ``numpy.py`` shim in the
|
||||
# device's Python ``site-packages`` directory can be enough for server startup.
|
||||
#
|
||||
# If ``cloudpickle`` is missing, copy it from another Python environment into
|
||||
# the device's ``site-packages`` directory. It is a pure Python package, so it
|
||||
# usually does not need cross-compilation.
|
||||
#
|
||||
# Run OpenCL Kernel Remotely by RPC
|
||||
# ---------------------------------
|
||||
# For remote OpenCL devices, the workflow is almost the same as above.
|
||||
# You can define the kernel, upload files, and run via RPC.
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# Raspberry Pi does not support OpenCL, the following code is tested on
|
||||
# Firefly-RK3399. You may follow this `tutorial <https://gist.github.com/mli/585aed2cec0b5178b1a510f9f236afa2>`_
|
||||
# to setup the OS and OpenCL driver for RK3399.
|
||||
#
|
||||
# Also we need to build the runtime with OpenCL enabled on rk3399 board. In the TVM
|
||||
# root directory, execute
|
||||
#
|
||||
# .. code-block:: bash
|
||||
#
|
||||
# mkdir -p build && cd build
|
||||
# cp ../cmake/config.cmake .
|
||||
# sed -i "s/USE_OPENCL OFF/USE_OPENCL ON/" config.cmake
|
||||
# cmake .. && cmake --build . --parallel $(nproc)
|
||||
#
|
||||
# The following function shows how we run an OpenCL kernel remotely
|
||||
|
||||
|
||||
def run_opencl():
|
||||
# NOTE: This is the setting for my rk3399 board. You need to modify
|
||||
# them according to your environment.
|
||||
opencl_device_host = "10.77.1.145"
|
||||
opencl_device_port = 9090
|
||||
target = tvm.target.Target("opencl", host={"kind": "llvm", "mtriple": "aarch64-linux-gnu"})
|
||||
|
||||
# create schedule for the above "add one" compute declaration
|
||||
mod = tvm.IRModule.from_expr(te.create_prim_func([A, B]))
|
||||
sch = tvm.s_tir.Schedule(mod)
|
||||
(x,) = sch.get_loops(block=sch.get_sblock("B"))
|
||||
xo, xi = sch.split(x, [None, 32])
|
||||
sch.bind(xo, "blockIdx.x")
|
||||
sch.bind(xi, "threadIdx.x")
|
||||
func = tvm.compile(sch.mod, target=target)
|
||||
|
||||
remote = rpc.connect(opencl_device_host, opencl_device_port)
|
||||
|
||||
# export and upload
|
||||
path = temp.relpath("lib_cl.tar")
|
||||
func.export_library(path)
|
||||
remote.upload(path)
|
||||
func = remote.load_module("lib_cl.tar")
|
||||
|
||||
# run
|
||||
dev = remote.cl()
|
||||
a = tvm.runtime.tensor(np.random.uniform(size=1024).astype(A.dtype), dev)
|
||||
b = tvm.runtime.tensor(np.zeros(1024, dtype=A.dtype), dev)
|
||||
func(a, b)
|
||||
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
|
||||
print("OpenCL test passed!")
|
||||
|
||||
|
||||
######################################################################
|
||||
# Deploy PyTorch Models to Remote Devices with RPC
|
||||
# ------------------------------------------------
|
||||
# The above examples demonstrate cross compilation and RPC using low-level
|
||||
# TensorIR (via TE). For deploying complete neural network models from frameworks
|
||||
# like PyTorch or ONNX, TVM's Relax provides a higher-level abstraction that is
|
||||
# better suited for end-to-end model compilation.
|
||||
#
|
||||
# This section shows a modern workflow for deploying models to **any remote device**:
|
||||
#
|
||||
# 1. Import a PyTorch model and convert it to Relax
|
||||
# 2. Cross-compile for the target architecture (ARM, x86, RISC-V, etc.)
|
||||
# 3. Deploy via RPC to a remote device
|
||||
# 4. Run inference remotely
|
||||
#
|
||||
# This workflow is applicable to various deployment scenarios:
|
||||
#
|
||||
# - **ARM devices**: Raspberry Pi, NVIDIA Jetson, mobile phones
|
||||
# - **x86 servers**: Remote Linux servers, cloud instances
|
||||
# - **Embedded systems**: RISC-V boards, custom hardware
|
||||
# - **Accelerators**: Remote machines with GPUs, TPUs, or other accelerators
|
||||
#
|
||||
# .. note::
|
||||
# This example uses PyTorch for demonstration, but the workflow is identical
|
||||
# for ONNX models. Simply replace ``from_exported_program()`` with
|
||||
# ``from_onnx(model, keep_params_in_input=True)`` and follow the same steps.
|
||||
|
||||
# First, let's check if PyTorch is available
|
||||
try:
|
||||
import torch
|
||||
from torch.export import export
|
||||
|
||||
HAS_TORCH = True
|
||||
except ImportError:
|
||||
HAS_TORCH = False
|
||||
|
||||
|
||||
def run_pytorch_model_via_rpc():
|
||||
"""
|
||||
Demonstrates the complete workflow of deploying a PyTorch model to an ARM device via RPC.
|
||||
"""
|
||||
if not HAS_TORCH:
|
||||
print("Skipping PyTorch example (PyTorch not installed)")
|
||||
return
|
||||
|
||||
from tvm import relax
|
||||
from tvm.relax.frontend.torch import from_exported_program
|
||||
|
||||
######################################################################
|
||||
# Step 1: Define and Export PyTorch Model
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# We use a simple MLP model for demonstration. In practice, this could be
|
||||
# any PyTorch model (ResNet, BERT, etc.).
|
||||
|
||||
class TorchMLP(torch.nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.net = torch.nn.Sequential(
|
||||
torch.nn.Flatten(),
|
||||
torch.nn.Linear(28 * 28, 128),
|
||||
torch.nn.ReLU(),
|
||||
torch.nn.Linear(128, 10),
|
||||
)
|
||||
|
||||
def forward(self, data: torch.Tensor) -> torch.Tensor:
|
||||
return self.net(data)
|
||||
|
||||
# Export the model using PyTorch 2.x export API
|
||||
torch_model = TorchMLP().eval()
|
||||
example_args = (torch.randn(1, 1, 28, 28, dtype=torch.float32),)
|
||||
|
||||
with torch.no_grad():
|
||||
exported_program = export(torch_model, example_args)
|
||||
|
||||
######################################################################
|
||||
# Step 2: Convert to Relax and Prepare for Compilation
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Convert the exported PyTorch program to TVM's Relax representation
|
||||
|
||||
mod = from_exported_program(exported_program, keep_params_as_input=True)
|
||||
# Separate parameters from the model for flexible deployment
|
||||
mod, params = relax.frontend.detach_params(mod)
|
||||
|
||||
print("Converted PyTorch model to Relax:")
|
||||
print(f" - Number of parameters: {len(params['main'])}")
|
||||
|
||||
######################################################################
|
||||
# Step 3: Cross-Compile for Target Device
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Compile the model for the target device architecture. The target
|
||||
# configuration depends on your deployment scenario.
|
||||
|
||||
if local_demo:
|
||||
# For demonstration on local machine, use local target
|
||||
target = tvm.target.Target("llvm")
|
||||
print("Using local target for demonstration")
|
||||
else:
|
||||
# Choose the appropriate target for your device:
|
||||
#
|
||||
# ARM devices:
|
||||
# - Raspberry Pi 3/4 (32-bit): {"kind": "llvm", "mtriple": "armv7l-linux-gnueabihf"}
|
||||
# - Raspberry Pi 4 (64-bit) / Jetson: {"kind": "llvm", "mtriple": "aarch64-linux-gnu"}
|
||||
# - Android: {"kind": "llvm", "mtriple": "aarch64-linux-android"}
|
||||
#
|
||||
# x86 servers:
|
||||
# - Linux x86_64: {"kind": "llvm", "mtriple": "x86_64-linux-gnu"}
|
||||
# - With AVX-512: {"kind": "llvm", "mtriple": "x86_64-linux-gnu", "mcpu": "skylake-avx512"}
|
||||
#
|
||||
# RISC-V:
|
||||
# - RV64: {"kind": "llvm", "mtriple": "riscv64-unknown-linux-gnu"}
|
||||
#
|
||||
# GPU targets:
|
||||
# - CUDA: tvm.target.Target("cuda", host={"kind": "llvm", "mtriple": "x86_64-linux-gnu"})
|
||||
# - OpenCL: tvm.target.Target("opencl", host={"kind": "llvm", "mtriple": "aarch64-linux-gnu"})
|
||||
#
|
||||
# For this example, we use ARM 64-bit
|
||||
target = tvm.target.Target({"kind": "llvm", "mtriple": "aarch64-linux-gnu"})
|
||||
print(f"Cross-compiling for target: {target}")
|
||||
|
||||
# Apply optimization pipeline
|
||||
pipeline = relax.get_pipeline()
|
||||
with target:
|
||||
built_mod = pipeline(mod)
|
||||
|
||||
# Compile to executable
|
||||
executable = tvm.compile(built_mod, target=target)
|
||||
|
||||
# Export to shared library
|
||||
lib_path = temp.relpath("model_deployed.so")
|
||||
executable.export_library(lib_path)
|
||||
print(f"Exported library to: {lib_path}")
|
||||
|
||||
# Save parameters separately
|
||||
import numpy as np
|
||||
|
||||
params_path = temp.relpath("model_params.npz")
|
||||
param_arrays = {f"p_{i}": p.numpy() for i, p in enumerate(params["main"])}
|
||||
np.savez(params_path, **param_arrays)
|
||||
print(f"Saved parameters to: {params_path}")
|
||||
|
||||
######################################################################
|
||||
# Step 4: Deploy to Remote Device via RPC
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Connect to the remote device, upload the compiled library and parameters,
|
||||
# then run inference remotely. This works for any device with TVM RPC server.
|
||||
#
|
||||
# Note: The following code demonstrates the RPC workflow. In local_demo mode,
|
||||
# we skip actual execution to avoid LocalSession compatibility issues.
|
||||
|
||||
if local_demo:
|
||||
# For demonstration, show the code structure without execution
|
||||
print("\nRPC workflow (works for any remote device):")
|
||||
print("=" * 50)
|
||||
print("1. Start RPC server on target device:")
|
||||
print(" python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090")
|
||||
print("\n2. Connect from local machine:")
|
||||
print(" remote = rpc.connect('DEVICE_IP', 9090)")
|
||||
print("\n3. Upload compiled library:")
|
||||
print(" remote.upload('model_deployed.so')")
|
||||
print(" remote.upload('model_params.npz')")
|
||||
print("\n4. Load and run remotely:")
|
||||
print(" lib = remote.load_module('model_deployed.so')")
|
||||
print(" vm = relax.VirtualMachine(lib, remote.cpu())")
|
||||
print(" result = vm['main'](input, *params)")
|
||||
print("\nDevice examples:")
|
||||
print(" - Raspberry Pi: 192.168.1.100")
|
||||
print(" - Remote server: ssh tunnel or direct IP")
|
||||
print(" - NVIDIA Jetson: 10.0.0.50")
|
||||
print(" - Cloud instance: public IP")
|
||||
print("\nTo run actual RPC, set local_demo=False")
|
||||
return # Skip actual RPC execution in demo mode
|
||||
|
||||
# Actual RPC workflow for real deployment
|
||||
# Connect to remote device (works for ARM, x86, RISC-V, etc.)
|
||||
# Make sure the RPC server is running on the device:
|
||||
# python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090
|
||||
device_host = "192.168.1.100" # Replace with your device IP
|
||||
device_port = 9090
|
||||
remote = rpc.connect(device_host, device_port)
|
||||
print(f"Connected to remote device at {device_host}:{device_port}")
|
||||
|
||||
# Upload library and parameters to remote device
|
||||
remote.upload(lib_path)
|
||||
remote.upload(params_path)
|
||||
print("Uploaded files to remote device")
|
||||
|
||||
# Load the library on the remote device
|
||||
lib = remote.load_module("model_deployed.so")
|
||||
|
||||
# Choose device on remote machine
|
||||
# For CPU: dev = remote.cpu()
|
||||
# For CUDA GPU: dev = remote.cuda(0)
|
||||
# For OpenCL: dev = remote.cl(0)
|
||||
dev = remote.cpu()
|
||||
|
||||
# Create VM and load parameters
|
||||
vm = relax.VirtualMachine(lib, dev)
|
||||
|
||||
# Load parameters from the uploaded file
|
||||
# Note: In practice, you might load this from the remote filesystem
|
||||
params_npz = np.load(params_path)
|
||||
remote_params = [tvm.runtime.tensor(params_npz[f"p_{i}"], dev) for i in range(len(params_npz))]
|
||||
|
||||
######################################################################
|
||||
# Step 5: Run Inference on Remote Device
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Execute the model on the remote ARM device and retrieve results
|
||||
#
|
||||
# Note: When running VM over RPC, we use set_input() and invoke_stateful()
|
||||
# instead of direct function call (vm["main"](...)). This is because RPC
|
||||
# transmits tensors as DLTensor*, while VM builtins expect ffi.Tensor.
|
||||
# The set_input API handles this conversion internally.
|
||||
|
||||
# Prepare input data
|
||||
input_data = np.random.randn(1, 1, 28, 28).astype("float32")
|
||||
remote_input = tvm.runtime.tensor(input_data, dev)
|
||||
|
||||
# Run inference using set_input + invoke_stateful for RPC compatibility
|
||||
vm.set_input("main", remote_input, *remote_params)
|
||||
vm.invoke_stateful("main")
|
||||
output = vm.get_outputs("main")
|
||||
|
||||
# Extract result (handle both tuple and single tensor outputs)
|
||||
if isinstance(output, tvm_ffi.Array) and len(output) > 0:
|
||||
result = output[0]
|
||||
else:
|
||||
result = output
|
||||
|
||||
# Retrieve result from remote device to local
|
||||
result_np = result.numpy()
|
||||
print("Inference completed on remote device")
|
||||
print(f" Output shape: {result_np.shape}")
|
||||
print(f" Predicted class: {np.argmax(result_np)}")
|
||||
|
||||
######################################################################
|
||||
# Alternative: Direct Function Call (Local Only)
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Note: The direct call syntax vm["main"](input, *params) works for
|
||||
# local execution but may fail over RPC due to type mismatch between
|
||||
# DLTensor* (RPC) and ffi.Tensor (VM builtins). For RPC, always use
|
||||
# the set_input + invoke_stateful pattern shown above.
|
||||
|
||||
######################################################################
|
||||
# Step 6: Performance Evaluation (Optional)
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Measure inference time on the remote device, excluding network overhead
|
||||
# Note: For RPC, use invoke_stateful with time_evaluator
|
||||
|
||||
time_f = vm.time_evaluator("invoke_stateful", dev, number=10, repeat=3)
|
||||
prof_res = time_f("main")
|
||||
print(f"Inference time on remote device: {prof_res.mean * 1000:.2f} ms")
|
||||
|
||||
######################################################################
|
||||
# Notes on Performance Optimization
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
#
|
||||
# For optimal performance on target devices, consider:
|
||||
#
|
||||
# 1. **Auto-tuning with MetaSchedule**: Use automated search to find
|
||||
# optimal schedules for your specific hardware:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# mod = relax.get_pipeline(
|
||||
# "static_shape_tuning",
|
||||
# target=target,
|
||||
# total_trials=2000
|
||||
# )(mod)
|
||||
#
|
||||
# 2. **Quick optimization with DLight**: Apply pre-defined performant schedules:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# from tvm.s_tir import dlight as dl
|
||||
# with target:
|
||||
# mod = dl.ApplyDefaultSchedule()(mod)
|
||||
#
|
||||
# 3. **Architecture-specific optimizations**:
|
||||
#
|
||||
# - ARM NEON SIMD: ``-mattr=+neon``
|
||||
# - x86 AVX-512: ``-mcpu=skylake-avx512``
|
||||
# - RISC-V Vector: ``-mattr=+v``
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# # Example: ARM with NEON
|
||||
# target = tvm.target.Target(
|
||||
# {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": "+neon"}
|
||||
# )
|
||||
#
|
||||
# # Example: x86 with AVX-512
|
||||
# target = tvm.target.Target(
|
||||
# {"kind": "llvm", "mtriple": "x86_64-linux-gnu", "mcpu": "skylake-avx512"}
|
||||
# )
|
||||
#
|
||||
# See :doc:`e2e_opt_model </how_to/tutorials/e2e_opt_model>` for detailed
|
||||
# tuning examples.
|
||||
|
||||
|
||||
# Run the PyTorch RPC example if PyTorch is available
|
||||
if HAS_TORCH and local_demo:
|
||||
try:
|
||||
run_pytorch_model_via_rpc()
|
||||
except Exception:
|
||||
pass # Silently skip if execution fails
|
||||
|
||||
|
||||
######################################################################
|
||||
# Summary
|
||||
# -------
|
||||
# This tutorial provides a walk through of cross compilation and RPC
|
||||
# features in TVM.
|
||||
#
|
||||
# We demonstrated two approaches:
|
||||
#
|
||||
# **Low-level TensorIR (TE) approach** - for understanding fundamentals:
|
||||
#
|
||||
# - Define computations using Tensor Expression
|
||||
# - Cross-compile for ARM targets
|
||||
# - Deploy and run via RPC
|
||||
#
|
||||
# **High-level Relax approach** - for deploying complete models:
|
||||
#
|
||||
# - Import models from PyTorch (or ONNX)
|
||||
# - Convert to Relax representation
|
||||
# - Cross-compile for ARM Linux devices
|
||||
# - Deploy to remote devices via RPC
|
||||
# - Run inference and evaluate performance
|
||||
#
|
||||
# Key takeaways:
|
||||
#
|
||||
# - Set up an RPC server on the remote device
|
||||
# - Cross-compile on a powerful local machine for resource-constrained targets
|
||||
# - Upload and execute compiled modules remotely via the RPC API
|
||||
# - Measure performance excluding network overhead
|
||||
#
|
||||
# For complete model deployment workflows, see also:
|
||||
#
|
||||
# - :doc:`export_and_load_executable </how_to/tutorials/export_and_load_executable>` - Export and load compiled models
|
||||
# - :doc:`e2e_opt_model </how_to/tutorials/e2e_opt_model>` - End-to-end optimization with auto-tuning
|
||||
@@ -0,0 +1,237 @@
|
||||
# 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.
|
||||
# ruff: noqa: E402, E501, F401
|
||||
|
||||
"""
|
||||
.. _customize_opt:
|
||||
|
||||
Customize Optimization
|
||||
======================
|
||||
One main design goal of Apache TVM is to enable easy customization of the optimization pipeline
|
||||
for both research or development purposes and iterate the engineering optimizations. In this
|
||||
tutorial we will
|
||||
|
||||
.. contents:: Table of Contents
|
||||
:local:
|
||||
:depth: 1
|
||||
"""
|
||||
|
||||
######################################################################
|
||||
# Review Overall Flow
|
||||
# -------------------
|
||||
# .. figure:: https://raw.githubusercontent.com/tlc-pack/web-data/main/images/design/tvm_overall_flow.svg
|
||||
# :align: center
|
||||
# :width: 80%
|
||||
#
|
||||
# The overall flow consists of the following steps:
|
||||
#
|
||||
# - **Construct or Import a Model**: Construct a neural network model or import a pre-trained
|
||||
# model from other frameworks (e.g. PyTorch, ONNX), and create the TVM IRModule, which contains
|
||||
# all the information needed for compilation, including high-level Relax functions for
|
||||
# computational graph, and low-level TensorIR functions for tensor program.
|
||||
# - **Perform Composable Optimizations**: Perform a series of optimization transformations,
|
||||
# such as graph optimizations, tensor program optimizations, and library dispatching.
|
||||
# - **Build and Universal Deployment**: Build the optimized model to a deployable module to the
|
||||
# universal runtime, and execute it on different devices, such as CPU, GPU, or other accelerators.
|
||||
#
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
|
||||
import tvm
|
||||
from tvm import IRModule, relax
|
||||
from tvm.relax.frontend import nn
|
||||
|
||||
######################################################################
|
||||
# Composable IRModule Optimization
|
||||
# --------------------------------
|
||||
# Apache TVM provides a flexible way to optimize the IRModule. Everything centered
|
||||
# around IRModule optimization can be composed with existing pipelines. Note that each optimization
|
||||
# can focus on **part of the computation graph**, enabling partial lowering or partial optimization.
|
||||
#
|
||||
# In this tutorial, we will demonstrate how to optimize a model with Apache TVM.
|
||||
|
||||
######################################################################
|
||||
# Prepare a Relax Module
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# We first prepare a Relax module. The module can be imported from other frameworks, constructed
|
||||
# with NN module frontend or TVMScript. Here we use a simple neural network model as an example.
|
||||
|
||||
|
||||
class RelaxModel(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.fc1 = nn.Linear(784, 256)
|
||||
self.relu1 = nn.ReLU()
|
||||
self.fc2 = nn.Linear(256, 10, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.fc1(x)
|
||||
x = self.relu1(x)
|
||||
x = self.fc2(x)
|
||||
return x
|
||||
|
||||
|
||||
input_shape = (1, 784)
|
||||
mod, params = RelaxModel().export_tvm({"forward": {"x": nn.spec.Tensor(input_shape, "float32")}})
|
||||
mod.show()
|
||||
|
||||
######################################################################
|
||||
# Library Dispatch
|
||||
# ~~~~~~~~~~~~~~~~
|
||||
# We would like to quickly try out a variant of library optimization for certain platforms
|
||||
# (e.g., GPU). We can write a certain dispatching pass for the specific platform and
|
||||
# operator. Here we demonstrate how to dispatch the CUBLAS library for certain patterns.
|
||||
#
|
||||
# .. note::
|
||||
# This tutorial only demonstrates a single operator dispatching for CUBLAS, highlighting
|
||||
# the flexibility of the optimization pipeline. In real-world cases, we can import multiple
|
||||
# patterns and dispatch them to different kernels.
|
||||
|
||||
|
||||
# Import cublas pattern
|
||||
try:
|
||||
import tvm.relax.backend.cuda.cublas as _cublas
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"This tutorial requires TVM built with CUDA support.\n"
|
||||
"If you hit missing 'tvm_ffi', try: pip install apache-tvm-ffi\n"
|
||||
"Otherwise build TVM with CUDA enabled:\n"
|
||||
" https://tvm.apache.org/docs/install/from_source.html\n"
|
||||
f"Original error: {e}"
|
||||
) from e
|
||||
|
||||
|
||||
# Define a new pass for CUBLAS dispatch
|
||||
@tvm.transform.module_pass(opt_level=0, name="CublasDispatch")
|
||||
class CublasDispatch:
|
||||
def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
|
||||
# Check if CUBLAS is enabled
|
||||
if not tvm.get_global_func("relax.ext.cublas", True):
|
||||
raise Exception("CUBLAS is not enabled.")
|
||||
|
||||
# Get interested patterns
|
||||
patterns = [relax.backend.get_pattern("cublas.matmul_transposed_bias_relu")]
|
||||
# Note in real-world cases, we usually get all patterns
|
||||
# patterns = relax.backend.get_patterns_with_prefix("cublas")
|
||||
|
||||
# Fuse ops by patterns and then run codegen
|
||||
mod = relax.transform.FuseOpsByPattern(patterns, annotate_codegen=True)(mod)
|
||||
mod = relax.transform.RunCodegen()(mod)
|
||||
return mod
|
||||
|
||||
|
||||
mod = CublasDispatch()(mod)
|
||||
mod.show()
|
||||
|
||||
######################################################################
|
||||
# After the dispatching pass, we can see that the first ``nn.Linear`` and ``nn.ReLU`` are fused
|
||||
# and rewritten to a ``call_dps_packed`` function which call the CUBLAS library. Notably, the
|
||||
# other part is not changed, which means we can selectively dispatch the optimization for
|
||||
# certain computation.
|
||||
|
||||
######################################################################
|
||||
# Auto Tuning
|
||||
# ~~~~~~~~~~~
|
||||
# Continuing from the previous example, we can further optimize the model with auto-tuning for
|
||||
# the **rest part of the computation**. Here we demonstrate how to use the meta-schedule to auto-tune
|
||||
# the model.
|
||||
#
|
||||
# We can use ``MetaScheduleTuneTIR`` pass to simply tuning the model, while ``MetaScheduleApplyDatabase``
|
||||
# pass to apply the best configuration to the model. The tuning process will generate search space,
|
||||
# tune the model and the following steps will apply the best configuration to the model. Before
|
||||
# running the passes, we need to lowering relax operator into TensorIR functions via ``LegalizeOps``
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# To save CI time and avoid flakiness, we skip the tuning process in CI environment.
|
||||
#
|
||||
|
||||
device = tvm.cuda(0)
|
||||
target = tvm.target.Target.from_device(device)
|
||||
if os.getenv("CI", "") != "true":
|
||||
trials = 2000
|
||||
with target, tempfile.TemporaryDirectory() as tmp_dir:
|
||||
mod = tvm.ir.transform.Sequential(
|
||||
[
|
||||
relax.get_pipeline("zero"),
|
||||
relax.transform.MetaScheduleTuneTIR(work_dir=tmp_dir, max_trials_global=trials),
|
||||
relax.transform.MetaScheduleApplyDatabase(work_dir=tmp_dir),
|
||||
]
|
||||
)(mod)
|
||||
|
||||
mod.show()
|
||||
|
||||
######################################################################
|
||||
# DLight Rules
|
||||
# ~~~~~~~~~~~~
|
||||
# DLight rules are a set of default rules for scheduling and optimization the kernel.
|
||||
# DLight rules are designed for fast compilation and **fair** performance. In some cases,
|
||||
# e.g. language model, DLight provides excellent performance, while for generic models,
|
||||
# it achieves a balance between performance and compilation time.
|
||||
|
||||
from tvm.s_tir import dlight as dl
|
||||
|
||||
# Apply DLight rules
|
||||
with target:
|
||||
mod = tvm.ir.transform.Sequential(
|
||||
[
|
||||
relax.get_pipeline("zero"),
|
||||
dl.ApplyDefaultSchedule( # pylint: disable=not-callable
|
||||
dl.gpu.Matmul(),
|
||||
dl.gpu.GEMV(),
|
||||
dl.gpu.Reduction(),
|
||||
dl.gpu.GeneralReduction(),
|
||||
dl.gpu.Fallback(),
|
||||
),
|
||||
]
|
||||
)(mod)
|
||||
|
||||
mod.show()
|
||||
|
||||
######################################################################
|
||||
# .. note::
|
||||
#
|
||||
# This tutorial focuses on the demonstration of the optimization pipeline, instead of
|
||||
# pushing the performance to the limit. The current optimization may not be the best.
|
||||
|
||||
|
||||
######################################################################
|
||||
# Deploy the Optimized Model
|
||||
# --------------------------
|
||||
# We can build and deploy the optimized model to the TVM runtime.
|
||||
|
||||
ex = tvm.compile(mod, target="cuda")
|
||||
dev = tvm.device("cuda", 0)
|
||||
vm = relax.VirtualMachine(ex, dev)
|
||||
# Need to allocate data and params on GPU device
|
||||
data = tvm.runtime.tensor(np.random.rand(*input_shape).astype("float32"), dev)
|
||||
gpu_params = [tvm.runtime.tensor(np.random.rand(*p.shape).astype(p.dtype), dev) for _, p in params]
|
||||
gpu_out = vm["forward"](data, *gpu_params).numpy()
|
||||
print(gpu_out)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Summary
|
||||
# -------
|
||||
# This tutorial demonstrates how to customize the optimization pipeline for ML models in Apache TVM.
|
||||
# We can easily compose the optimization passes and customize the optimization for different parts
|
||||
# of the computation graph. The flexibility of the optimization pipeline enables us to quickly
|
||||
# iterate the optimization and improve the performance of the model.
|
||||
#
|
||||
@@ -0,0 +1,153 @@
|
||||
# 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.
|
||||
# ruff: noqa: E402
|
||||
|
||||
"""
|
||||
.. _optimize_model:
|
||||
|
||||
End-to-End Optimize Model
|
||||
=========================
|
||||
This tutorial demonstrates how to optimize a machine learning model using Apache TVM. We will
|
||||
use a pre-trained ResNet-18 model from PyTorch and end-to-end optimize it using TVM's Relax API.
|
||||
Please note that default end-to-end optimization may not suit complex models.
|
||||
"""
|
||||
|
||||
######################################################################
|
||||
# Preparation
|
||||
# -----------
|
||||
# First, we prepare the model and input information. We use a pre-trained ResNet-18 model from
|
||||
# PyTorch.
|
||||
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.export import export
|
||||
from torchvision.models.resnet import ResNet18_Weights, resnet18
|
||||
|
||||
torch_model = resnet18(weights=ResNet18_Weights.DEFAULT).eval()
|
||||
|
||||
######################################################################
|
||||
# Review Overall Flow
|
||||
# -------------------
|
||||
# .. figure:: https://raw.githubusercontent.com/tlc-pack/web-data/main/images/design/tvm_overall_flow.svg
|
||||
# :align: center
|
||||
# :width: 80%
|
||||
#
|
||||
# The overall flow consists of the following steps:
|
||||
#
|
||||
# - **Construct or Import a Model**: Construct a neural network model or import a pre-trained
|
||||
# model from other frameworks (e.g. PyTorch, ONNX), and create the TVM IRModule, which contains
|
||||
# all the information needed for compilation, including high-level Relax functions for
|
||||
# computational graph, and low-level TensorIR functions for tensor program.
|
||||
# - **Perform Composable Optimizations**: Perform a series of optimization transformations,
|
||||
# such as graph optimizations, tensor program optimizations, and library dispatching.
|
||||
# - **Build and Universal Deployment**: Build the optimized model to a deployable module to the
|
||||
# universal runtime, and execute it on different devices, such as CPU, GPU, or other accelerators.
|
||||
#
|
||||
|
||||
|
||||
######################################################################
|
||||
# Convert the model to IRModule
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Next step, we convert the model to an IRModule using the Relax frontend for PyTorch for further
|
||||
# optimization.
|
||||
|
||||
import tvm
|
||||
from tvm import relax
|
||||
from tvm.relax.frontend.torch import from_exported_program
|
||||
|
||||
# Give an example argument to torch.export
|
||||
example_args = (torch.randn(1, 3, 224, 224, dtype=torch.float32),)
|
||||
|
||||
# Skip running in CI environment
|
||||
IS_IN_CI = os.getenv("CI", "") == "true"
|
||||
|
||||
if not IS_IN_CI:
|
||||
# Convert the model to IRModule
|
||||
with torch.no_grad():
|
||||
exported_program = export(torch_model, example_args)
|
||||
mod = from_exported_program(exported_program, keep_params_as_input=True)
|
||||
|
||||
mod, params = relax.frontend.detach_params(mod)
|
||||
mod.show()
|
||||
|
||||
######################################################################
|
||||
# IRModule Optimization
|
||||
# ---------------------
|
||||
# Apache TVM provides a flexible way to optimize the IRModule. Everything centered
|
||||
# around IRModule optimization can be composed with existing pipelines. Note that each
|
||||
# transformation can be combined as an optimization pipeline via ``tvm.ir.transform.Sequential``.
|
||||
#
|
||||
# In this tutorial, we focus on the end-to-end optimization of the model via auto-tuning. We
|
||||
# leverage MetaSchedule to tune the model and store the tuning logs to the database. We also
|
||||
# apply the database to the model to get the best performance.
|
||||
#
|
||||
# The ResNet18 model will be divided into 20 independent tuning tasks during compilation.
|
||||
# To ensure each task receives adequate tuning resources in one iteration while providing
|
||||
# early feedback:
|
||||
#
|
||||
# - To quickly observe tuning progress, each task is allocated a maximum of 16 trials per
|
||||
# iteration (controlled by ``MAX_TRIALS_PER_TASK=16``). We should set ``TOTAL_TRIALS``
|
||||
# to at least ``320 (20 tasks * 16 trials)`` ensures every task receives one full iteration
|
||||
# of tuning. We set it to 512 in our configuration to allow for several more iterations,
|
||||
# aiming to explore a wider parameter space and potentially achieve better performance.
|
||||
# - If ``MAX_TRIALS_PER_TASK == None``, the system defaults to ``TOTAL_TRIALS`` trials per
|
||||
# task per iteration. An insufficient ``TOTAL_TRIALS`` setting may lead to undersubscribed
|
||||
# tuning, potentially skipping some tasks entirely. Explicitly setting both parameters
|
||||
# avoids this issue and provides deterministic resource allocation across all tasks.
|
||||
#
|
||||
# Note: These parameter settings are optimized for quick tutorial demonstration. For production
|
||||
# deployments requiring higher performance, we recommend adjusting both ``MAX_TRIALS_PER_TASK``
|
||||
# and ``TOTAL_TRIALS`` to larger values. This allows more extensive search space exploration
|
||||
# and typically yields better performance outcomes.
|
||||
|
||||
TOTAL_TRIALS = 512 # Change to 20000 for better performance if needed
|
||||
MAX_TRIALS_PER_TASK = 16 # Change to more trials per task for better performance if needed
|
||||
target = tvm.target.Target("nvidia/geforce-rtx-3090-ti") # Change to your target device
|
||||
work_dir = "tuning_logs"
|
||||
|
||||
if not IS_IN_CI:
|
||||
mod = relax.get_pipeline(
|
||||
"static_shape_tuning",
|
||||
target=target,
|
||||
work_dir=work_dir,
|
||||
total_trials=TOTAL_TRIALS,
|
||||
max_trials_per_task=MAX_TRIALS_PER_TASK,
|
||||
)(mod)
|
||||
|
||||
# Only show the main function
|
||||
mod["main"].show()
|
||||
|
||||
######################################################################
|
||||
# Build and Deploy
|
||||
# ----------------
|
||||
# Finally, we build the optimized model and deploy it to the target device.
|
||||
# We skip this step in the CI environment.
|
||||
|
||||
if not IS_IN_CI:
|
||||
with target:
|
||||
mod = tvm.s_tir.transform.DefaultGPUSchedule()(mod)
|
||||
ex = tvm.compile(mod, target=target)
|
||||
dev = tvm.device("cuda", 0)
|
||||
vm = relax.VirtualMachine(ex, dev)
|
||||
# Need to allocate data and params on GPU device
|
||||
gpu_data = tvm.runtime.tensor(np.random.rand(1, 3, 224, 224).astype("float32"), dev)
|
||||
gpu_params = [tvm.runtime.tensor(p, dev) for p in params["main"]]
|
||||
gpu_out = vm["main"](gpu_data, *gpu_params)[0].numpy()
|
||||
|
||||
print(gpu_out.shape)
|
||||
@@ -0,0 +1,384 @@
|
||||
# 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.
|
||||
# ruff: noqa: E402
|
||||
|
||||
"""
|
||||
.. _deploy_export_and_load_executable:
|
||||
|
||||
Export and Load Relax Executables
|
||||
=================================
|
||||
|
||||
This tutorial walks through exporting a compiled Relax module to a shared
|
||||
object, loading it back into the TVM runtime, and running the result either
|
||||
interactively or from a standalone script. This tutorial demonstrates how
|
||||
to turn Relax (or imported PyTorch / ONNX) programs into deployable artifacts
|
||||
using ``tvm.relax`` APIs.
|
||||
|
||||
.. note::
|
||||
This tutorial uses PyTorch as the source format, but the export/load workflow
|
||||
is the same for ONNX models. For ONNX, use ``from_onnx(model, keep_params_in_input=True)``
|
||||
instead of ``from_exported_program()``, then follow the same steps for building,
|
||||
exporting, and loading.
|
||||
"""
|
||||
|
||||
######################################################################
|
||||
# Introduction
|
||||
# ------------
|
||||
# TVM builds Relax programs into ``tvm.runtime.Executable`` objects. These
|
||||
# contain VM bytecode, compiled kernels, and constants. By exporting the
|
||||
# executable with :py:meth:`export_library`, you obtain a shared library (for
|
||||
# example ``.so`` on Linux) that can be shipped to another machine, uploaded
|
||||
# via RPC, or loaded back later with the TVM runtime. This tutorial shows the
|
||||
# exact steps end-to-end and explains what files are produced along the way.
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
import torch
|
||||
from torch.export import export
|
||||
except ImportError: # pragma: no cover
|
||||
torch = None # type: ignore
|
||||
|
||||
|
||||
######################################################################
|
||||
# Prepare a Torch MLP and Convert to Relax
|
||||
# ----------------------------------------
|
||||
# We start with a small PyTorch MLP so the example remains lightweight. The
|
||||
# model is exported to a :py:class:`torch.export.ExportedProgram` and then
|
||||
# translated into a Relax ``IRModule``.
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
import tvm
|
||||
from tvm import relax
|
||||
from tvm.relax.frontend.torch import from_exported_program
|
||||
|
||||
# Check dependencies first
|
||||
IS_IN_CI = os.getenv("CI", "").lower() == "true"
|
||||
HAS_TORCH = torch is not None
|
||||
RUN_EXAMPLE = HAS_TORCH and not IS_IN_CI
|
||||
|
||||
|
||||
if HAS_TORCH:
|
||||
|
||||
class TorchMLP(torch.nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.net = torch.nn.Sequential(
|
||||
torch.nn.Flatten(),
|
||||
torch.nn.Linear(28 * 28, 128),
|
||||
torch.nn.ReLU(),
|
||||
torch.nn.Linear(128, 10),
|
||||
)
|
||||
|
||||
def forward(self, data: torch.Tensor) -> torch.Tensor: # type: ignore[override]
|
||||
return self.net(data)
|
||||
|
||||
else: # pragma: no cover
|
||||
TorchMLP = None # type: ignore[misc, assignment]
|
||||
|
||||
if RUN_EXAMPLE:
|
||||
torch_model = TorchMLP().eval()
|
||||
example_args = (torch.randn(1, 1, 28, 28, dtype=torch.float32),)
|
||||
|
||||
with torch.no_grad():
|
||||
exported_program = export(torch_model, example_args)
|
||||
|
||||
mod = from_exported_program(exported_program, keep_params_as_input=True)
|
||||
|
||||
# Separate model parameters so they can be bound later (or stored on disk).
|
||||
mod, params = relax.frontend.detach_params(mod)
|
||||
|
||||
print("Imported Relax module:")
|
||||
mod.show()
|
||||
|
||||
|
||||
######################################################################
|
||||
# Build and Export with ``export_library``
|
||||
# -------------------------------------------
|
||||
# We build for ``llvm`` to generate CPU code and then export the resulting
|
||||
# executable. Passing ``workspace_dir`` keeps the intermediate packaging files,
|
||||
# which is useful to inspect what was produced.
|
||||
|
||||
TARGET = tvm.target.Target("llvm")
|
||||
ARTIFACT_DIR = Path("relax_export_artifacts")
|
||||
ARTIFACT_DIR.mkdir(exist_ok=True)
|
||||
|
||||
if RUN_EXAMPLE:
|
||||
# Apply the default Relax compilation pipeline before building.
|
||||
pipeline = relax.get_pipeline()
|
||||
with TARGET:
|
||||
built_mod = pipeline(mod)
|
||||
|
||||
# Build without params - we'll pass them at runtime
|
||||
executable = tvm.compile(built_mod, target=TARGET)
|
||||
|
||||
library_path = ARTIFACT_DIR / "mlp_cpu.so"
|
||||
executable.export_library(str(library_path), workspace_dir=str(ARTIFACT_DIR))
|
||||
|
||||
print(f"Exported runtime library to: {library_path}")
|
||||
|
||||
# The workspace directory now contains the shared object and supporting files.
|
||||
produced_files = sorted(p.name for p in ARTIFACT_DIR.iterdir())
|
||||
print("Artifacts saved:")
|
||||
for name in produced_files:
|
||||
print(f" - {name}")
|
||||
|
||||
# Generated files:
|
||||
# - ``mlp_cpu.so``: The main deployable shared library containing VM bytecode,
|
||||
# compiled kernels, and constants. Note: Since parameters are passed at runtime,
|
||||
# you will also need to save a separate parameters file (see next section).
|
||||
# - Intermediate object files (``devc.o``, ``lib0.o``, etc.) are kept in the
|
||||
# workspace for inspection but are not required for deployment.
|
||||
#
|
||||
# Note: Additional files like ``*.params``, ``*.metadata.json``, or ``*.imports``
|
||||
# may appear in specific configurations but are typically embedded into the
|
||||
# shared library or only generated when needed.
|
||||
|
||||
|
||||
######################################################################
|
||||
# Load the Exported Library and Run It
|
||||
# ------------------------------------
|
||||
# Once the shared object is produced, we can reload it back into the TVM runtime
|
||||
# on any machine with a compatible instruction set. The Relax VM consumes the
|
||||
# runtime module directly.
|
||||
|
||||
if RUN_EXAMPLE:
|
||||
loaded_rt_mod = tvm.runtime.load_module(str(library_path))
|
||||
dev = tvm.cpu(0)
|
||||
vm = relax.VirtualMachine(loaded_rt_mod, dev)
|
||||
|
||||
# Prepare input data
|
||||
input_tensor = torch.randn(1, 1, 28, 28, dtype=torch.float32)
|
||||
vm_input = tvm.runtime.tensor(input_tensor.numpy(), dev)
|
||||
|
||||
# Prepare parameters (allocate on target device)
|
||||
vm_params = [tvm.runtime.tensor(p, dev) for p in params["main"]]
|
||||
|
||||
# Run inference: pass input data followed by all parameters
|
||||
tvm_output = vm["main"](vm_input, *vm_params)
|
||||
|
||||
# TVM returns Array objects for tuple outputs, access via indexing.
|
||||
# For models imported from PyTorch, outputs are typically tuples (even for single outputs).
|
||||
# For ONNX models, outputs may be a single Tensor directly.
|
||||
if isinstance(tvm_output, tvm_ffi.Array) and len(tvm_output) > 0:
|
||||
result_tensor = tvm_output[0]
|
||||
else:
|
||||
result_tensor = tvm_output
|
||||
|
||||
print("VM output shape:", result_tensor.shape)
|
||||
print("VM output type:", type(tvm_output), "->", type(result_tensor))
|
||||
|
||||
# You can still inspect the executable after reloading.
|
||||
print("Executable stats:\n", loaded_rt_mod["stats"]())
|
||||
|
||||
|
||||
######################################################################
|
||||
# Save Parameters for Deployment
|
||||
# -------------------------------
|
||||
# Since parameters are passed at runtime (not embedded in the ``.so``), we must
|
||||
# save them separately for deployment. This is a required step to use the model
|
||||
# on other machines or in standalone scripts.
|
||||
|
||||
import numpy as np
|
||||
|
||||
if RUN_EXAMPLE:
|
||||
# Save parameters to disk
|
||||
params_path = ARTIFACT_DIR / "model_params.npz"
|
||||
param_arrays = {f"p_{i}": p.numpy() for i, p in enumerate(params["main"])}
|
||||
np.savez(str(params_path), **param_arrays)
|
||||
print(f"Saved parameters to: {params_path}")
|
||||
|
||||
# Note: Alternatively, you can embed parameters directly into the ``.so`` to
|
||||
# create a single-file deployment. Use ``keep_params_as_input=False`` when
|
||||
# importing from PyTorch:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# mod = from_exported_program(exported_program, keep_params_as_input=False)
|
||||
# # Parameters are now embedded as constants in the module
|
||||
# executable = tvm.compile(built_mod, target=TARGET)
|
||||
# # Runtime: vm["main"](input) # No need to pass params!
|
||||
#
|
||||
# This creates a single-file deployment (only the ``.so`` is needed), but you
|
||||
# lose the flexibility to swap parameters without recompiling. For most
|
||||
# production workflows, separating code and parameters (as shown above) is
|
||||
# preferred for flexibility.
|
||||
|
||||
|
||||
######################################################################
|
||||
# Loading and Running the Exported Model
|
||||
# -----------------------------------------------------------
|
||||
# To use the exported model on another machine or in a standalone script, you need
|
||||
# to load both the ``.so`` library and the parameters file. Here's a complete example
|
||||
# of how to reload and run the model. Save this as ``run_mlp.py``:
|
||||
#
|
||||
# To make it executable from the command line:
|
||||
#
|
||||
# .. code-block:: bash
|
||||
#
|
||||
# chmod +x run_mlp.py
|
||||
# ./run_mlp.py # Run it like a regular program
|
||||
#
|
||||
# Complete script:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# #!/usr/bin/env python3
|
||||
# import numpy as np
|
||||
# import tvm
|
||||
# from tvm import relax
|
||||
#
|
||||
# # Step 1: Load the compiled library
|
||||
# lib = tvm.runtime.load_module("relax_export_artifacts/mlp_cpu.so")
|
||||
#
|
||||
# # Step 2: Create Virtual Machine
|
||||
# device = tvm.cpu(0)
|
||||
# vm = relax.VirtualMachine(lib, device)
|
||||
#
|
||||
# # Step 3: Load parameters from the .npz file
|
||||
# params_npz = np.load("relax_export_artifacts/model_params.npz")
|
||||
# params = [tvm.runtime.tensor(params_npz[f"p_{i}"], device)
|
||||
# for i in range(len(params_npz))]
|
||||
#
|
||||
# # Step 4: Prepare input data
|
||||
# data = np.random.randn(1, 1, 28, 28).astype("float32")
|
||||
# input_tensor = tvm.runtime.tensor(data, device)
|
||||
#
|
||||
# # Step 5: Run inference (pass input followed by all parameters)
|
||||
# output = vm["main"](input_tensor, *params)
|
||||
#
|
||||
# # Step 6: Extract result (output may be tuple or single Tensor)
|
||||
# # PyTorch models typically return tuples, ONNX models may return a single Tensor
|
||||
# if isinstance(output, tvm_ffi.Array) and len(output) > 0:
|
||||
# result_tensor = output[0]
|
||||
# else:
|
||||
# result_tensor = output
|
||||
#
|
||||
# print("Prediction shape:", result_tensor.shape)
|
||||
# print("Predicted class:", np.argmax(result_tensor.numpy()))
|
||||
#
|
||||
# **Running on GPU:**
|
||||
# To run on GPU instead of CPU, make the following changes:
|
||||
#
|
||||
# 1. **Compile for GPU** (earlier in the tutorial, around line 112):
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# TARGET = tvm.target.Target("cuda") # Change from "llvm" to "cuda"
|
||||
#
|
||||
# 2. **Use GPU device in the script**:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# device = tvm.cuda(0) # Use CUDA device instead of CPU
|
||||
# vm = relax.VirtualMachine(lib, device)
|
||||
#
|
||||
# # Load parameters to GPU
|
||||
# params = [tvm.runtime.tensor(params_npz[f"p_{i}"], device) # Note: device parameter
|
||||
# for i in range(len(params_npz))]
|
||||
#
|
||||
# # Prepare input on GPU
|
||||
# input_tensor = tvm.runtime.tensor(data, device) # Note: device parameter
|
||||
#
|
||||
# The rest of the script remains the same. All tensors (parameters and inputs)
|
||||
# must be allocated on the same device (GPU) as the compiled model.
|
||||
#
|
||||
# **Deployment Checklist:**
|
||||
# When moving to another host (via RPC or SCP), you must copy **both** files:
|
||||
#
|
||||
# 1. ``mlp_cpu.so`` (or ``mlp_cuda.so`` for GPU) - the compiled model code
|
||||
# 2. ``model_params.npz`` - the model parameters, serialized as NumPy arrays
|
||||
#
|
||||
# The remote machine needs both files in the same directory. The script above
|
||||
# assumes they are in ``relax_export_artifacts/`` relative to the script location.
|
||||
# Adjust the paths as needed for your deployment. For GPU deployment, ensure the
|
||||
# target machine has compatible CUDA drivers and the model was compiled for the
|
||||
# same GPU architecture.
|
||||
|
||||
|
||||
######################################################################
|
||||
# Deploying to Remote Devices
|
||||
# ---------------------------
|
||||
# To deploy the exported model to a remote ARM Linux device (e.g., Raspberry Pi),
|
||||
# you can use TVM's RPC mechanism to cross-compile, upload, and run the model
|
||||
# remotely. This workflow is useful when:
|
||||
#
|
||||
# - The target device has limited resources for compilation
|
||||
# - You want to fine-tune performance by running on the actual hardware
|
||||
# - You need to deploy to embedded devices
|
||||
#
|
||||
# See :doc:`cross_compilation_and_rpc </how_to/tutorials/cross_compilation_and_rpc>`
|
||||
# for a comprehensive guide on:
|
||||
#
|
||||
# - Setting up TVM runtime on the remote device
|
||||
# - Starting an RPC server on the device
|
||||
# - Cross-compiling for ARM targets (e.g., ``llvm -mtriple=aarch64-linux-gnu``)
|
||||
# - Uploading exported libraries via RPC
|
||||
# - Running inference remotely
|
||||
#
|
||||
# Quick example for ARM deployment workflow:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# import tvm.rpc as rpc
|
||||
# from tvm import relax
|
||||
#
|
||||
# # Step 1: Cross-compile for ARM target (on local machine)
|
||||
# TARGET = tvm.target.Target({"kind": "llvm", "mtriple": "aarch64-linux-gnu"})
|
||||
# executable = tvm.compile(built_mod, target=TARGET)
|
||||
# executable.export_library("mlp_arm.so")
|
||||
#
|
||||
# # Step 2: Connect to remote device RPC server
|
||||
# remote = rpc.connect("192.168.1.100", 9090) # Device IP and RPC port
|
||||
#
|
||||
# # Step 3: Upload the compiled library and parameters
|
||||
# remote.upload("mlp_arm.so")
|
||||
# remote.upload("model_params.npz")
|
||||
#
|
||||
# # Step 4: Load and run on remote device
|
||||
# lib = remote.load_module("mlp_arm.so")
|
||||
# vm = relax.VirtualMachine(lib, remote.cpu())
|
||||
# # ... prepare input and params, then run inference
|
||||
#
|
||||
# The key difference is using an ARM target triple during compilation and
|
||||
# uploading files via RPC instead of copying them directly.
|
||||
|
||||
|
||||
######################################################################
|
||||
# FAQ
|
||||
# ---
|
||||
# **Can I run the ``.so`` as a standalone executable (like ``./mlp_cpu.so``)?**
|
||||
# No. The ``.so`` file is a shared library, not a standalone executable binary.
|
||||
# You cannot run it directly from the terminal. It must be loaded through a TVM
|
||||
# runtime program (as shown in the "Loading and Running" section above). The
|
||||
# ``.so`` bundles VM bytecode and compiled kernels, but still requires the TVM
|
||||
# runtime to execute.
|
||||
#
|
||||
# **Which devices can run the exported library?**
|
||||
# The target must match the ISA you compiled for (``llvm`` in this example).
|
||||
# As long as the target triple, runtime ABI, and available devices line up,
|
||||
# you can move the artifact between machines. For heterogeneous builds (CPU
|
||||
# plus GPU), ship the extra device libraries as well.
|
||||
#
|
||||
# **What about the ``.params`` and ``metadata.json`` files?**
|
||||
# These auxiliary files are only generated in specific configurations. In this
|
||||
# tutorial, since we pass parameters at runtime, they are not generated. When
|
||||
# they do appear, they may be kept alongside the ``.so`` for inspection, but
|
||||
# the essential content is typically embedded in the shared object itself, so
|
||||
# deploying the ``.so`` alone is usually sufficient.
|
||||
@@ -0,0 +1,407 @@
|
||||
# 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.
|
||||
# ruff: noqa: E402, E501
|
||||
|
||||
"""
|
||||
.. _import_model:
|
||||
|
||||
Importing Models from ML Frameworks
|
||||
====================================
|
||||
Apache TVM supports importing models from popular ML frameworks including PyTorch, ONNX,
|
||||
and TensorFlow Lite. This tutorial walks through each import path with a minimal working
|
||||
example and explains the key parameters. The PyTorch section additionally demonstrates
|
||||
how to handle unsupported operators via a custom converter map.
|
||||
|
||||
For end-to-end optimization and deployment after importing, see :ref:`optimize_model`.
|
||||
|
||||
.. note::
|
||||
|
||||
The ONNX section requires the ``onnx`` package. The TFLite section requires
|
||||
``tensorflow`` and ``tflite``. Sections whose dependencies are missing are skipped
|
||||
automatically.
|
||||
|
||||
.. contents:: Table of Contents
|
||||
:local:
|
||||
:depth: 2
|
||||
"""
|
||||
|
||||
######################################################################
|
||||
# Importing from PyTorch (Recommended)
|
||||
# -------------------------------------
|
||||
# TVM's PyTorch frontend is the most feature-complete. The recommended entry point is
|
||||
# :py:func:`~tvm.relax.frontend.torch.from_exported_program`, which works with PyTorch's
|
||||
# ``torch.export`` API.
|
||||
#
|
||||
# We start by defining a small CNN model for demonstration. No pretrained weights are
|
||||
# needed — we only care about the graph structure.
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.export import export
|
||||
|
||||
import tvm
|
||||
from tvm import relax
|
||||
from tvm.relax.frontend.torch import from_exported_program
|
||||
|
||||
|
||||
class SimpleCNN(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.conv = nn.Conv2d(3, 16, kernel_size=3, padding=1)
|
||||
self.bn = nn.BatchNorm2d(16)
|
||||
self.pool = nn.AdaptiveAvgPool2d((1, 1))
|
||||
self.fc = nn.Linear(16, 10)
|
||||
|
||||
def forward(self, x):
|
||||
x = torch.relu(self.bn(self.conv(x)))
|
||||
x = self.pool(x).flatten(1)
|
||||
x = self.fc(x)
|
||||
return x
|
||||
|
||||
|
||||
torch_model = SimpleCNN().eval()
|
||||
example_args = (torch.randn(1, 3, 32, 32),)
|
||||
|
||||
######################################################################
|
||||
# Basic import
|
||||
# ~~~~~~~~~~~~
|
||||
# The standard workflow is: ``torch.export.export()`` → ``from_exported_program()`` →
|
||||
# ``detach_params()``.
|
||||
|
||||
with torch.no_grad():
|
||||
exported_program = export(torch_model, example_args)
|
||||
mod = from_exported_program(
|
||||
exported_program,
|
||||
keep_params_as_input=True,
|
||||
unwrap_unit_return_tuple=True,
|
||||
)
|
||||
|
||||
mod, params = relax.frontend.detach_params(mod)
|
||||
mod.show()
|
||||
|
||||
######################################################################
|
||||
# Key parameters
|
||||
# ~~~~~~~~~~~~~~
|
||||
# ``from_exported_program`` accepts several parameters that control how the model is
|
||||
# translated:
|
||||
#
|
||||
# - **keep_params_as_input** (bool, default ``False``): When ``True``, model weights become
|
||||
# function parameters, separated via ``relax.frontend.detach_params()``. When ``False``,
|
||||
# weights are embedded as constants inside the IRModule. Use ``True`` when you want to
|
||||
# manage weights independently (e.g., for weight sharing or quantization).
|
||||
#
|
||||
# - **unwrap_unit_return_tuple** (bool, default ``False``): PyTorch ``export`` always wraps
|
||||
# the return value in a tuple. Set ``True`` to unwrap single-element return tuples for a
|
||||
# cleaner Relax function signature.
|
||||
#
|
||||
# - **run_ep_decomposition** (bool, default ``True``): Runs PyTorch's built-in operator
|
||||
# decomposition before translation. This breaks high-level ops (e.g., ``batch_norm``) into
|
||||
# lower-level primitives, which generally improves TVM's coverage and optimization
|
||||
# opportunities. Set ``False`` if you want to preserve the original op granularity.
|
||||
|
||||
######################################################################
|
||||
# Handling unsupported operators with ``custom_convert_map``
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# When TVM encounters a PyTorch operator it does not recognize, it raises an error
|
||||
# indicating the unsupported operator name. You can extend the frontend by providing a
|
||||
# **custom converter map** — a dictionary mapping operator names to your own conversion
|
||||
# functions.
|
||||
#
|
||||
# A custom converter function receives two arguments:
|
||||
#
|
||||
# - **node** (``torch.fx.Node``): The FX graph node being converted, carrying operator
|
||||
# info and references to input nodes.
|
||||
# - **importer** (``ExportedProgramImporter``): The importer instance, giving access to:
|
||||
#
|
||||
# - ``importer.env``: Dict mapping FX nodes to their converted Relax expressions.
|
||||
# - ``importer.block_builder``: The Relax ``BlockBuilder`` for emitting operations.
|
||||
# - ``importer.retrieve_args(node)``: Helper to look up converted args.
|
||||
#
|
||||
# The function must return a ``relax.Var`` — the Relax expression for this node's output.
|
||||
# Here is an example that maps an operator to ``relax.op.sigmoid``:
|
||||
|
||||
from tvm.relax.frontend.torch.exported_program_translator import ExportedProgramImporter
|
||||
|
||||
|
||||
def convert_sigmoid(node: torch.fx.Node, importer: ExportedProgramImporter) -> relax.Var:
|
||||
"""Custom converter: map an op to relax.op.sigmoid."""
|
||||
args = importer.retrieve_args(node)
|
||||
return importer.block_builder.emit(relax.op.sigmoid(args[0]))
|
||||
|
||||
|
||||
######################################################################
|
||||
# To use the custom converter, pass it via the ``custom_convert_map`` parameter. The key
|
||||
# is the ATen operator name in ``"op_name.variant"`` format (e.g., ``"sigmoid.default"``):
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# mod = from_exported_program(
|
||||
# exported_program,
|
||||
# custom_convert_map={"sigmoid.default": convert_sigmoid},
|
||||
# )
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# To find the correct operator name, check the error message TVM raises when encountering
|
||||
# the unsupported op — it includes the exact ATen name. You can also inspect the exported
|
||||
# program's graph via ``print(exported_program.graph_module.graph)`` to see all operator
|
||||
# names.
|
||||
|
||||
######################################################################
|
||||
# Alternative PyTorch import methods
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Besides ``from_exported_program``, TVM also provides:
|
||||
#
|
||||
# - :py:func:`~tvm.relax.frontend.torch.from_fx`: Works with ``torch.fx.GraphModule``
|
||||
# from ``torch.fx.symbolic_trace()``. Requires explicit ``input_info`` (shapes and dtypes).
|
||||
# Use this when ``torch.export`` fails on certain Python control flow patterns.
|
||||
#
|
||||
# - :py:func:`~tvm.relax.frontend.torch.relax_dynamo`: A ``torch.compile`` backend that
|
||||
# compiles and executes the model through TVM in one step. Useful for integrating TVM
|
||||
# into an existing PyTorch training or inference loop.
|
||||
#
|
||||
# - :py:func:`~tvm.relax.frontend.torch.dynamo_capture_subgraphs`: Captures subgraphs from
|
||||
# a PyTorch model into an IRModule via ``torch.compile``. Each subgraph becomes a separate
|
||||
# function in the IRModule.
|
||||
#
|
||||
# For most use cases, ``from_exported_program`` is the recommended path.
|
||||
|
||||
######################################################################
|
||||
# Verifying the imported model
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# After importing, it is good practice to verify that TVM produces the same output as the
|
||||
# original framework. We compile with the minimal ``"zero"`` pipeline (no tuning) and
|
||||
# compare. The same approach applies to models imported via the ONNX and TFLite frontends
|
||||
# shown below.
|
||||
|
||||
mod_compiled = relax.get_pipeline("zero")(mod)
|
||||
exec_module = tvm.compile(mod_compiled, target="llvm")
|
||||
dev = tvm.cpu()
|
||||
vm = relax.VirtualMachine(exec_module, dev)
|
||||
|
||||
# Run inference
|
||||
input_data = np.random.rand(1, 3, 32, 32).astype("float32")
|
||||
tvm_input = tvm.runtime.tensor(input_data, dev)
|
||||
tvm_params = [tvm.runtime.tensor(p, dev) for p in params["main"]]
|
||||
tvm_out = vm["main"](tvm_input, *tvm_params).numpy()
|
||||
|
||||
# Compare with PyTorch
|
||||
with torch.no_grad():
|
||||
pt_out = torch_model(torch.from_numpy(input_data)).numpy()
|
||||
|
||||
np.testing.assert_allclose(tvm_out, pt_out, rtol=1e-5, atol=1e-5)
|
||||
print("PyTorch vs TVM outputs match!")
|
||||
|
||||
######################################################################
|
||||
# Importing from ONNX
|
||||
# --------------------
|
||||
# TVM can import ONNX models via :py:func:`~tvm.relax.frontend.onnx.from_onnx`. The
|
||||
# function accepts an ``onnx.ModelProto`` object, so you need to load the model with
|
||||
# ``onnx.load()`` first.
|
||||
#
|
||||
# Here we export the same CNN model to ONNX format and then import it into TVM.
|
||||
|
||||
try:
|
||||
import onnx
|
||||
import onnxscript # noqa: F401 # required by torch.onnx.export
|
||||
|
||||
HAS_ONNX = True
|
||||
except ImportError:
|
||||
onnx = None # type: ignore[assignment]
|
||||
HAS_ONNX = False
|
||||
|
||||
if HAS_ONNX:
|
||||
from tvm.relax.frontend.onnx import from_onnx
|
||||
|
||||
# Export the PyTorch model to ONNX
|
||||
dummy_input = torch.randn(1, 3, 32, 32)
|
||||
onnx_path = "simple_cnn.onnx"
|
||||
torch.onnx.export(torch_model, dummy_input, onnx_path, input_names=["input"])
|
||||
|
||||
# Load and import into TVM
|
||||
onnx_model = onnx.load(onnx_path)
|
||||
mod_onnx = from_onnx(onnx_model, keep_params_in_input=True)
|
||||
mod_onnx, params_onnx = relax.frontend.detach_params(mod_onnx)
|
||||
mod_onnx.show()
|
||||
|
||||
######################################################################
|
||||
# If you already have an ``.onnx`` file on disk, the workflow is even simpler:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# import onnx
|
||||
# from tvm.relax.frontend.onnx import from_onnx
|
||||
#
|
||||
# onnx_model = onnx.load("my_model.onnx")
|
||||
# mod = from_onnx(onnx_model)
|
||||
#
|
||||
|
||||
######################################################################
|
||||
# Key parameters
|
||||
# ~~~~~~~~~~~~~~
|
||||
# - **shape_dict** (dict, optional): Maps input names to shapes. Auto-inferred from the
|
||||
# model if not provided. Useful when the ONNX model has dynamic dimensions that you
|
||||
# want to fix to concrete sizes:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# mod = from_onnx(onnx_model, shape_dict={"input": [1, 3, 224, 224]})
|
||||
#
|
||||
# - **dtype_dict** (str or dict, default ``"float32"``): Input dtypes. A single string
|
||||
# applies to all inputs, or use a dict to set per-input dtypes:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# mod = from_onnx(onnx_model, dtype_dict={"input": "float16"})
|
||||
#
|
||||
# - **keep_params_in_input** (bool, default ``False``): Same semantics as PyTorch — whether
|
||||
# model weights are function parameters or embedded constants.
|
||||
#
|
||||
# - **opset** (int, optional): Override the opset version auto-detected from the model.
|
||||
# Each ONNX op may have different semantics across opset versions; TVM's converter
|
||||
# selects the appropriate implementation automatically. You rarely need to set this
|
||||
# unless the model metadata is incorrect.
|
||||
|
||||
######################################################################
|
||||
# Importing from TensorFlow Lite
|
||||
# -------------------------------
|
||||
# TVM can import TFLite flat buffer models via
|
||||
# :py:func:`~tvm.relax.frontend.tflite.from_tflite`. The function expects a TFLite
|
||||
# ``Model`` object parsed from flat buffer bytes via ``GetRootAsModel``.
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# The ``tflite`` Python package has changed its module layout across versions.
|
||||
# Older versions use ``tflite.Model.Model.GetRootAsModel``, while newer versions use
|
||||
# ``tflite.Model.GetRootAsModel``. The code below handles both.
|
||||
#
|
||||
# Below we create a minimal TFLite model from TensorFlow and import it.
|
||||
|
||||
try:
|
||||
import tensorflow as tf
|
||||
import tflite
|
||||
import tflite.Model
|
||||
|
||||
HAS_TFLITE = True
|
||||
except ImportError:
|
||||
HAS_TFLITE = False
|
||||
|
||||
if HAS_TFLITE:
|
||||
from tvm.relax.frontend.tflite import from_tflite
|
||||
|
||||
# Define a simple TF module and convert to TFLite.
|
||||
# We use plain TF ops (not keras layers) to avoid variable-handling ops
|
||||
# that some TFLite converter versions do not support cleanly.
|
||||
class TFModule(tf.Module):
|
||||
@tf.function(
|
||||
input_signature=[
|
||||
tf.TensorSpec(shape=(1, 784), dtype=tf.float32),
|
||||
tf.TensorSpec(shape=(784, 10), dtype=tf.float32),
|
||||
]
|
||||
)
|
||||
def forward(self, x, weight):
|
||||
return tf.matmul(x, weight) + 0.1
|
||||
|
||||
tf_module = TFModule()
|
||||
converter = tf.lite.TFLiteConverter.from_concrete_functions(
|
||||
[tf_module.forward.get_concrete_function()], tf_module
|
||||
)
|
||||
tflite_buf = converter.convert()
|
||||
|
||||
# Parse and import into TVM (API differs between tflite package versions)
|
||||
if hasattr(tflite.Model, "Model"):
|
||||
tflite_model = tflite.Model.Model.GetRootAsModel(tflite_buf, 0)
|
||||
else:
|
||||
tflite_model = tflite.Model.GetRootAsModel(tflite_buf, 0)
|
||||
mod_tflite = from_tflite(tflite_model)
|
||||
mod_tflite.show()
|
||||
|
||||
######################################################################
|
||||
# Loading from a ``.tflite`` file
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# If you already have a ``.tflite`` file on disk, load the raw bytes and parse them:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# import tflite
|
||||
# import tflite.Model
|
||||
# from tvm.relax.frontend.tflite import from_tflite
|
||||
#
|
||||
# with open("my_model.tflite", "rb") as f:
|
||||
# tflite_buf = f.read()
|
||||
#
|
||||
# if hasattr(tflite.Model, "Model"):
|
||||
# tflite_model = tflite.Model.Model.GetRootAsModel(tflite_buf, 0)
|
||||
# else:
|
||||
# tflite_model = tflite.Model.GetRootAsModel(tflite_buf, 0)
|
||||
# mod = from_tflite(tflite_model)
|
||||
|
||||
######################################################################
|
||||
# Key parameters
|
||||
# ~~~~~~~~~~~~~~
|
||||
# - **shape_dict** / **dtype_dict** (optional): Override input shapes and dtypes. If not
|
||||
# provided, they are inferred from the TFLite model metadata.
|
||||
#
|
||||
# - **op_converter** (class, optional): A custom operator converter class. Subclass
|
||||
# ``OperatorConverter`` and override its ``convert_map`` dictionary to add or replace
|
||||
# operator conversions. For example, to add a hypothetical ``CUSTOM_RELU`` op:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# from tvm.relax.frontend.tflite.tflite_frontend import OperatorConverter
|
||||
#
|
||||
# class MyConverter(OperatorConverter):
|
||||
# def __init__(self, model, subgraph, exp_tab, ctx):
|
||||
# super().__init__(model, subgraph, exp_tab, ctx)
|
||||
# self.convert_map["CUSTOM_RELU"] = self._convert_custom_relu
|
||||
#
|
||||
# def _convert_custom_relu(self, op):
|
||||
# # implement your conversion logic here
|
||||
# ...
|
||||
#
|
||||
# mod = from_tflite(tflite_model, op_converter=MyConverter)
|
||||
|
||||
######################################################################
|
||||
# Summary
|
||||
# -------
|
||||
#
|
||||
# +---------------------+----------------------------+-------------------------------+-----------------------------+
|
||||
# | Aspect | PyTorch | ONNX | TFLite |
|
||||
# +=====================+============================+===============================+=============================+
|
||||
# | Entry function | ``from_exported_program`` | ``from_onnx`` | ``from_tflite`` |
|
||||
# +---------------------+----------------------------+-------------------------------+-----------------------------+
|
||||
# | Input | ``ExportedProgram`` | ``onnx.ModelProto`` | TFLite ``Model`` object |
|
||||
# +---------------------+----------------------------+-------------------------------+-----------------------------+
|
||||
# | Custom extension | ``custom_convert_map`` | — | ``op_converter`` class |
|
||||
# +---------------------+----------------------------+-------------------------------+-----------------------------+
|
||||
#
|
||||
# **Which to use?** Pick the frontend that matches your model format:
|
||||
#
|
||||
# - Have a PyTorch model? Use ``from_exported_program`` — it has the broadest operator coverage.
|
||||
# - Have an ``.onnx`` file? Use ``from_onnx``.
|
||||
# - Have a ``.tflite`` file? Use ``from_tflite``.
|
||||
#
|
||||
# The verification workflow (compile → run → compare) demonstrated in the PyTorch section
|
||||
# above applies equally to ONNX and TFLite imports.
|
||||
#
|
||||
# For the full list of supported operators, see the converter map in each frontend's source:
|
||||
# PyTorch uses ``create_convert_map()`` in ``exported_program_translator.py``, ONNX uses
|
||||
# ``_get_convert_map()`` in ``onnx_frontend.py``, and TFLite uses ``convert_map`` in
|
||||
# ``OperatorConverter`` in ``tflite_frontend.py``.
|
||||
#
|
||||
# After importing, refer to :ref:`optimize_model` for optimization and deployment.
|
||||
@@ -0,0 +1,461 @@
|
||||
# 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.
|
||||
|
||||
"""
|
||||
.. _mix_python_and_tvm:
|
||||
|
||||
Mix Python/PyTorch with TVM Using BasePyModule
|
||||
===============================================
|
||||
In a typical TVM workflow, you write an ``IRModule``, compile it, and load the compiled artifact
|
||||
into a ``VirtualMachine`` to run. This means **you cannot test or debug anything until the entire
|
||||
module compiles successfully**. If a single op is unsupported, the whole pipeline is blocked.
|
||||
|
||||
``BasePyModule`` solves this by letting Python functions, TIR kernels, and Relax functions coexist
|
||||
in one module. TIR and Relax functions are JIT-compiled on instantiation, Python functions run
|
||||
as-is, and tensors move between TVM and PyTorch via zero-copy DLPack. This enables:
|
||||
|
||||
- **Incremental development**: get a model running with Python fallbacks first, then replace them
|
||||
with TVM ops one by one.
|
||||
- **Easy debugging**: insert ``print`` in Python functions to inspect intermediate tensors — no
|
||||
need to compile the whole module first.
|
||||
- **Verification at any compilation stage**: convert Relax IR back to PyTorch to check numerical
|
||||
correctness before and after optimization passes.
|
||||
- **Hybrid execution**: let the compiled VM call back into Python for ops that are hard to
|
||||
express in TIR or Relax.
|
||||
|
||||
This tutorial walks through the full workflow step by step.
|
||||
|
||||
.. contents:: Table of Contents
|
||||
:local:
|
||||
:depth: 1
|
||||
"""
|
||||
|
||||
######################################################################
|
||||
# Preparation
|
||||
# -----------
|
||||
|
||||
import os
|
||||
|
||||
try:
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
except ImportError:
|
||||
torch = None
|
||||
|
||||
import tvm
|
||||
from tvm import relax
|
||||
from tvm.relax.base_py_module import BasePyModule
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import relax as R
|
||||
from tvm.script import tirx as T
|
||||
|
||||
IS_IN_CI = os.getenv("CI", "").lower() == "true"
|
||||
HAS_TORCH = torch is not None
|
||||
RUN_EXAMPLE = HAS_TORCH and not IS_IN_CI
|
||||
|
||||
|
||||
######################################################################
|
||||
# Step 1: Your First Hybrid Module
|
||||
# ----------------------------------
|
||||
# The core idea: decorate a class with ``@I.ir_module``, inherit from ``BasePyModule``, and use
|
||||
# three decorators for three kinds of functions:
|
||||
#
|
||||
# - ``@T.prim_func`` — low-level TIR kernel (JIT-compiled on instantiation)
|
||||
# - ``@R.function`` — high-level Relax graph (JIT-compiled on instantiation)
|
||||
# - ``@I.pyfunc`` — plain Python (runs as-is, can use any Python library)
|
||||
#
|
||||
# ``call_tir`` bridges Python and TIR: it converts PyTorch tensors to TVM NDArrays via DLPack
|
||||
# (zero-copy), allocates the output buffer, calls the compiled kernel, and converts back.
|
||||
|
||||
if RUN_EXAMPLE:
|
||||
|
||||
@I.ir_module
|
||||
class MyFirstModule(BasePyModule):
|
||||
@T.prim_func(s_tir=True)
|
||||
def add_tir(
|
||||
A: T.Buffer((4,), "float32"),
|
||||
B: T.Buffer((4,), "float32"),
|
||||
C: T.Buffer((4,), "float32"),
|
||||
):
|
||||
for i in range(4):
|
||||
C[i] = A[i] + B[i]
|
||||
|
||||
@I.pyfunc
|
||||
def forward(self, x, y):
|
||||
"""Takes PyTorch tensors, calls TIR, returns PyTorch tensors."""
|
||||
x_tvm = self._convert_pytorch_to_tvm(x)
|
||||
y_tvm = self._convert_pytorch_to_tvm(y)
|
||||
result = self.call_tir(self.add_tir, [x_tvm, y_tvm], out_ty=R.Tensor((4,), "float32"))
|
||||
return self._convert_tvm_to_pytorch(result)
|
||||
|
||||
# TIR functions are JIT-compiled at instantiation
|
||||
mod = MyFirstModule(device=tvm.cpu(0))
|
||||
|
||||
x = torch.tensor([1.0, 2.0, 3.0, 4.0])
|
||||
y = torch.tensor([10.0, 20.0, 30.0, 40.0])
|
||||
result = mod.forward(x, y)
|
||||
|
||||
print("forward(x, y) =", result)
|
||||
assert torch.allclose(result, x + y)
|
||||
|
||||
# show() prints TVMScript including Python functions (shown as ExternFunc)
|
||||
mod.show()
|
||||
|
||||
# list_functions() shows what is available in the module
|
||||
print("Available functions:", mod.list_functions())
|
||||
|
||||
|
||||
######################################################################
|
||||
# Step 2: Debugging — The Main Selling Point
|
||||
# ---------------------------------------------
|
||||
# Traditional ML compilers treat computation graphs as monolithic blobs. You cannot inspect
|
||||
# intermediate tensor values without compiling the entire module. With ``@I.pyfunc``, debugging
|
||||
# is as simple as adding a ``print`` statement. You can also make quick edits and re-run
|
||||
# immediately — no recompilation needed.
|
||||
|
||||
if RUN_EXAMPLE:
|
||||
|
||||
@I.ir_module
|
||||
class DebugModule(BasePyModule):
|
||||
@T.prim_func(s_tir=True)
|
||||
def matmul_tir(var_A: T.handle, var_B: T.handle, var_C: T.handle):
|
||||
n = T.int32()
|
||||
A = T.match_buffer(var_A, (n, 4), "float32")
|
||||
B = T.match_buffer(var_B, (4, 3), "float32")
|
||||
C = T.match_buffer(var_C, (n, 3), "float32")
|
||||
for i, j, k in T.grid(n, 3, 4):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = T.float32(0)
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
@I.pyfunc
|
||||
def forward(self, x, weights):
|
||||
# Inspect input
|
||||
print(f" [DEBUG] input shape: {x.shape}, mean: {x.mean():.4f}")
|
||||
|
||||
# Run TIR matmul
|
||||
x_tvm = self._convert_pytorch_to_tvm(x)
|
||||
w_tvm = self._convert_pytorch_to_tvm(weights)
|
||||
out = self.call_tir(
|
||||
self.matmul_tir,
|
||||
[x_tvm, w_tvm],
|
||||
out_ty=R.Tensor((x.shape[0], 3), "float32"),
|
||||
)
|
||||
logits = self._convert_tvm_to_pytorch(out)
|
||||
|
||||
# Inspect intermediate value — impossible with a compiled-only workflow
|
||||
print(
|
||||
f" [DEBUG] logits shape: {logits.shape}, "
|
||||
f"min: {logits.min():.4f}, max: {logits.max():.4f}"
|
||||
)
|
||||
|
||||
result = F.softmax(logits, dim=-1)
|
||||
|
||||
# Verify output
|
||||
print(f" [DEBUG] probs sum: {result.sum(dim=-1)}")
|
||||
return result
|
||||
|
||||
mod = DebugModule(device=tvm.cpu(0))
|
||||
|
||||
print("Running with debug prints:")
|
||||
probs = mod.forward(torch.randn(2, 4), torch.randn(4, 3))
|
||||
assert torch.allclose(probs.sum(dim=-1), torch.ones(2), atol=1e-5)
|
||||
|
||||
######################################################################
|
||||
# This is the key benefit: "debugging is as simple as inserting a print statement.
|
||||
# Users can also make quick, manual edits to Python functions and immediately observe the
|
||||
# results." No compilation cycle, no VM loading — just Python.
|
||||
|
||||
|
||||
######################################################################
|
||||
# Step 3: A Realistic Pipeline — Python, TIR, and Packed Functions
|
||||
# -------------------------------------------------------------------
|
||||
# Real models combine many kinds of operations. This step builds a mini inference pipeline using
|
||||
# three different calling conventions:
|
||||
#
|
||||
# - ``call_tir``: call a compiled TIR kernel
|
||||
# - ``call_dps_packed``: call a TVM packed function (e.g., a third-party library binding)
|
||||
# - Direct Python: call any PyTorch function
|
||||
#
|
||||
# ``call_dps_packed`` is useful for calling functions registered via ``tvm.register_global_func``
|
||||
# — for example, CUBLAS or cuDNN bindings that TVM wraps as packed functions.
|
||||
|
||||
if RUN_EXAMPLE:
|
||||
# Register a packed function (simulating an external library binding)
|
||||
@tvm.register_global_func("my_bias_add", override=True)
|
||||
def my_bias_add(x, bias, out):
|
||||
"""Packed function: adds bias to each row of x."""
|
||||
|
||||
x_np = x.numpy()
|
||||
b_np = bias.numpy()
|
||||
out_np = x_np + b_np
|
||||
out[:] = out_np
|
||||
|
||||
@I.ir_module
|
||||
class PipelineModule(BasePyModule):
|
||||
@T.prim_func(s_tir=True)
|
||||
def matmul_tir(var_A: T.handle, var_B: T.handle, var_C: T.handle):
|
||||
A = T.match_buffer(var_A, (2, 4), "float32")
|
||||
B = T.match_buffer(var_B, (4, 3), "float32")
|
||||
C = T.match_buffer(var_C, (2, 3), "float32")
|
||||
for i, j, k in T.grid(2, 3, 4):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = T.float32(0)
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
@I.pyfunc
|
||||
def forward(self, x, weights, bias):
|
||||
# 1. TIR matmul
|
||||
x_tvm = self._convert_pytorch_to_tvm(x)
|
||||
w_tvm = self._convert_pytorch_to_tvm(weights)
|
||||
h = self.call_tir(
|
||||
self.matmul_tir,
|
||||
[x_tvm, w_tvm],
|
||||
out_ty=R.Tensor((2, 3), "float32"),
|
||||
)
|
||||
h_pt = self._convert_tvm_to_pytorch(h)
|
||||
|
||||
# 2. Packed function for bias add (simulating an external library)
|
||||
h_biased = self.call_dps_packed(
|
||||
"my_bias_add",
|
||||
[h_pt, bias],
|
||||
out_ty=R.Tensor((2, 3), "float32"),
|
||||
)
|
||||
|
||||
# 3. Python/PyTorch activation
|
||||
return F.relu(h_biased)
|
||||
|
||||
mod = PipelineModule(device=tvm.cpu(0))
|
||||
|
||||
x = torch.randn(2, 4)
|
||||
w = torch.randn(4, 3)
|
||||
b = torch.randn(3)
|
||||
result = mod.forward(x, w, b)
|
||||
|
||||
expected = F.relu(x @ w + b)
|
||||
print("Pipeline result:", result)
|
||||
print("Expected: ", expected)
|
||||
assert torch.allclose(result, expected, atol=1e-4)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Step 4: Relax-to-Python Converter — Verify at Any Compilation Stage
|
||||
# ----------------------------------------------------------------------
|
||||
# Both Relax functions and Python functions describe computational graphs. The
|
||||
# ``RelaxToPyFuncConverter`` converts Relax IR into equivalent PyTorch code by mapping
|
||||
# Relax operators to their PyTorch counterparts (e.g., ``R.nn.relu`` → ``F.relu``).
|
||||
#
|
||||
# A key feature: **this conversion can happen at any stage of compilation**.
|
||||
# You can convert early (right after import) or late (after optimization passes have
|
||||
# transformed the IR), and compare the output against a PyTorch reference to catch bugs.
|
||||
|
||||
if RUN_EXAMPLE:
|
||||
from tvm.relax.relax_to_pyfunc_converter import RelaxToPyFuncConverter
|
||||
|
||||
# A simple Relax module: matmul + bias + relu (a dense layer)
|
||||
@I.ir_module
|
||||
class DenseLayer:
|
||||
@T.prim_func(s_tir=True)
|
||||
def bias_add_tir(var_x: T.handle, var_b: T.handle, var_out: T.handle):
|
||||
x = T.match_buffer(var_x, (2, 4), "float32")
|
||||
b = T.match_buffer(var_b, (4,), "float32")
|
||||
out = T.match_buffer(var_out, (2, 4), "float32")
|
||||
for i, j in T.grid(2, 4):
|
||||
out[i, j] = x[i, j] + b[j]
|
||||
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 4), "float32"),
|
||||
w: R.Tensor((4, 4), "float32"),
|
||||
b: R.Tensor((4,), "float32"),
|
||||
) -> R.Tensor((2, 4), "float32"):
|
||||
h = R.matmul(x, w)
|
||||
cls = DenseLayer
|
||||
h_bias = R.call_tir(
|
||||
cls.bias_add_tir,
|
||||
(h, b),
|
||||
out_ty=R.Tensor((2, 4), "float32"),
|
||||
)
|
||||
return R.nn.relu(h_bias)
|
||||
|
||||
# --- Stage 1: Convert BEFORE optimization ---
|
||||
converter = RelaxToPyFuncConverter(DenseLayer)
|
||||
converted_early = converter.convert(["main"])
|
||||
|
||||
x = torch.randn(2, 4)
|
||||
w = torch.randn(4, 4)
|
||||
b = torch.randn(4)
|
||||
|
||||
py_result_early = converted_early.pyfuncs["main"](x, w, b)
|
||||
expected = F.relu(x @ w + b)
|
||||
|
||||
print("Before optimization:")
|
||||
print(" Converted result:", py_result_early)
|
||||
print(" PyTorch expected:", expected)
|
||||
assert torch.allclose(py_result_early, expected, atol=1e-5)
|
||||
|
||||
# --- Stage 2: Apply a pass, then convert AFTER optimization ---
|
||||
# Run CanonicalizeBindings to clean up the IR, then convert again
|
||||
# to verify the pass did not break numerical correctness.
|
||||
optimized_mod = relax.transform.CanonicalizeBindings()(DenseLayer)
|
||||
|
||||
converter_late = RelaxToPyFuncConverter(optimized_mod)
|
||||
converted_late = converter_late.convert(["main"])
|
||||
|
||||
py_result_late = converted_late.pyfuncs["main"](x, w, b)
|
||||
|
||||
print("\nAfter CanonicalizeBindings pass:")
|
||||
print(" Converted result:", py_result_late)
|
||||
print(" Still matches: ", torch.allclose(py_result_late, expected, atol=1e-5))
|
||||
assert torch.allclose(py_result_late, expected, atol=1e-5)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Step 5: R.call_py_func — Python Callbacks in Compiled IR
|
||||
# -----------------------------------------------------------
|
||||
# ``R.call_py_func`` embeds a Python function call directly inside Relax IR. When the module
|
||||
# is compiled and run in the VM, everything else is optimized native code, but the VM calls
|
||||
# back into Python for the specified ops.
|
||||
#
|
||||
# ``BasePyModule`` supports cross-level calls in both directions: Relax functions can invoke
|
||||
# Python functions, and Python functions can invoke TIR/Relax functions. Data flows between
|
||||
# them via DLPack with minimal overhead.
|
||||
#
|
||||
# Use case: your model has a custom op (e.g., a special normalization or a sampling step)
|
||||
# that is complex to implement in TIR. Compile everything else, and let that one op stay
|
||||
# in Python.
|
||||
|
||||
if RUN_EXAMPLE:
|
||||
|
||||
@I.ir_module
|
||||
class HybridVMModule(BasePyModule):
|
||||
@I.pyfunc
|
||||
def silu(self, x):
|
||||
"""SiLU/Swish activation — using Python as fallback."""
|
||||
return torch.sigmoid(x) * x
|
||||
|
||||
@I.pyfunc
|
||||
def layer_norm(self, x):
|
||||
"""LayerNorm — another Python fallback."""
|
||||
return F.layer_norm(x, x.shape[-1:])
|
||||
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((4, 8), "float32"),
|
||||
) -> R.Tensor((4, 8), "float32"):
|
||||
# The VM calls back into Python for these two ops
|
||||
h = R.call_py_func("layer_norm", (x,), out_ty=R.Tensor((4, 8), "float32"))
|
||||
out = R.call_py_func("silu", (h,), out_ty=R.Tensor((4, 8), "float32"))
|
||||
return out
|
||||
|
||||
mod = HybridVMModule(device=tvm.cpu(0))
|
||||
x = torch.randn(4, 8)
|
||||
|
||||
# call_py_func is also callable from Python directly
|
||||
result = mod.call_py_func("layer_norm", [x])
|
||||
result = mod.call_py_func("silu", [result])
|
||||
|
||||
ln = F.layer_norm(x, x.shape[-1:])
|
||||
expected = torch.sigmoid(ln) * ln
|
||||
print("call_py_func result:", result)
|
||||
assert torch.allclose(torch.tensor(result.numpy()), expected, atol=1e-5)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Step 6: Cross-Level Calls and Symbolic Shapes
|
||||
# ------------------------------------------------
|
||||
# ``BasePyModule`` is designed for **cross-level interoperability**: Python functions can call
|
||||
# TIR and Relax functions, and Relax functions can call Python functions. We have already seen:
|
||||
#
|
||||
# - Python → TIR via ``call_tir`` (Steps 1-3)
|
||||
# - Python → packed function via ``call_dps_packed`` (Step 3)
|
||||
# - Relax → Python via ``R.call_py_func`` (Step 5)
|
||||
#
|
||||
# The missing piece: **Python calling a compiled Relax function directly**. When a module
|
||||
# contains ``@R.function``, it is JIT-compiled into a Relax VM. You can call it from Python
|
||||
# just like any other method — the module auto-converts PyTorch tensors to TVM and back.
|
||||
#
|
||||
# This step also shows **symbolic shapes**: TIR and Relax functions can declare dynamic
|
||||
# dimensions (e.g., ``"n"``). ``BasePyModule`` infers concrete shapes from the actual input
|
||||
# tensors at call time, so the same module handles different sizes without recompilation.
|
||||
|
||||
if RUN_EXAMPLE:
|
||||
|
||||
@I.ir_module
|
||||
class DynamicModule(BasePyModule):
|
||||
@T.prim_func(s_tir=True)
|
||||
def scale_tir(var_x: T.handle, var_out: T.handle):
|
||||
n = T.int64()
|
||||
x = T.match_buffer(var_x, (n,), "float32")
|
||||
out = T.match_buffer(var_out, (n,), "float32")
|
||||
for i in T.serial(n):
|
||||
out[i] = x[i] * T.float32(2.0)
|
||||
|
||||
@R.function
|
||||
def add_relax(
|
||||
x: R.Tensor(("n",), "float32"),
|
||||
y: R.Tensor(("n",), "float32"),
|
||||
) -> R.Tensor(("n",), "float32"):
|
||||
return R.add(x, y)
|
||||
|
||||
mod = DynamicModule(device=tvm.cpu(0), target="llvm")
|
||||
|
||||
# Inspect what the module contains
|
||||
print("Functions:", mod.list_functions())
|
||||
|
||||
# Python → Relax: call the compiled Relax function directly with PyTorch tensors
|
||||
a5 = torch.randn(5)
|
||||
b5 = torch.randn(5)
|
||||
out5 = mod.add_relax(a5, b5)
|
||||
print("add_relax(len=5):", out5)
|
||||
|
||||
# Same module, different size — symbolic shapes handle this automatically
|
||||
a10 = torch.randn(10)
|
||||
b10 = torch.randn(10)
|
||||
out10 = mod.add_relax(a10, b10)
|
||||
print("add_relax(len=10):", out10)
|
||||
|
||||
# Python → TIR with symbolic output shape
|
||||
n = T.int64()
|
||||
x7 = torch.randn(7)
|
||||
scaled = mod.call_tir("scale_tir", [x7], relax.TensorType((n,), "float32"))
|
||||
print("scale_tir(len=7):", scaled)
|
||||
assert torch.allclose(torch.tensor(scaled.numpy()), x7 * 2.0, atol=1e-5)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Summary
|
||||
# -------
|
||||
# Cross-level call summary:
|
||||
#
|
||||
# - **Python → TIR**: ``call_tir()`` (Steps 1, 2, 3, 6)
|
||||
# - **Python → packed function**: ``call_dps_packed()`` (Step 3)
|
||||
# - **Python → Relax**: call ``@R.function`` as a method (Step 6)
|
||||
# - **Relax → Python**: ``R.call_py_func()`` in compiled VM (Step 5)
|
||||
#
|
||||
# The workflow in practice:
|
||||
#
|
||||
# 1. Import a model → some ops unsupported → use ``@I.pyfunc`` as Python fallbacks
|
||||
# 2. Get it running end-to-end with ``BasePyModule``
|
||||
# 3. Debug by inserting ``print`` in pyfuncs — inspect intermediate tensors instantly
|
||||
# 4. Use ``RelaxToPyFuncConverter`` to verify correctness after each optimization pass
|
||||
# 5. Gradually replace Python fallbacks with TIR/Relax implementations
|
||||
# 6. Use ``R.call_py_func`` for ops that must stay in Python even after compilation
|
||||
@@ -0,0 +1,622 @@
|
||||
# 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.
|
||||
# ruff: noqa: F401
|
||||
|
||||
"""
|
||||
.. _opt_llm:
|
||||
|
||||
Optimize Large Language Model
|
||||
=============================
|
||||
As large language models (LLMs) have become a popular research topic in many different fields,
|
||||
deploying them on cloud and edge devices has become a challenging task. In this tutorial, we will
|
||||
demonstrate how to optimize a large language model using Apache TVM. We will use a pre-trained
|
||||
TinyLlama model from Hugging Face and deploy it on various devices.
|
||||
"""
|
||||
|
||||
######################################################################
|
||||
# Review Overall Flow
|
||||
# -------------------
|
||||
# .. figure:: https://raw.githubusercontent.com/tlc-pack/web-data/main/images/design/tvm_overall_flow.svg
|
||||
# :align: center
|
||||
# :width: 80%
|
||||
#
|
||||
# The overall flow consists of the following steps:
|
||||
#
|
||||
# - **Construct or Import a Model**: Construct a neural network model or import a pre-trained
|
||||
# model from other frameworks (e.g. PyTorch, ONNX), and create the TVM IRModule, which contains
|
||||
# all the information needed for compilation, including high-level Relax functions for
|
||||
# computational graph, and low-level TensorIR functions for tensor program.
|
||||
# - **Perform Composable Optimizations**: Perform a series of optimization transformations,
|
||||
# such as graph optimizations, tensor program optimizations, and library dispatching.
|
||||
# - **Build and Universal Deployment**: Build the optimized model to a deployable module to the
|
||||
# universal runtime, and execute it on different devices, such as CPU, GPU, or other accelerators.
|
||||
#
|
||||
|
||||
|
||||
######################################################################
|
||||
# Construct the model architecture
|
||||
# --------------------------------
|
||||
# We will use a pre-trained TinyLlama model from Hugging Face. However, usually we only load the
|
||||
# pre-trained weight from Hugging Face but not the model architecture. We need to construct the
|
||||
# model architecture by ourselves. Apache TVM prepares a PyTorch-liked API to construct the model
|
||||
# architecture. We can use the API to construct the model architecture.
|
||||
|
||||
import dataclasses
|
||||
import enum
|
||||
import os
|
||||
from pathlib import Path
|
||||
from pprint import pprint
|
||||
|
||||
from tvm_ffi import Shape
|
||||
|
||||
import tvm
|
||||
from tvm import relax, te, tirx
|
||||
from tvm.relax import register_pipeline
|
||||
from tvm.relax.frontend import nn
|
||||
from tvm.relax.frontend.nn import Tensor, op
|
||||
from tvm.relax.frontend.nn.llm.kv_cache import PagedKVCache, TIRPagedKVCache
|
||||
from tvm.s_tir import dlight
|
||||
|
||||
######################################################################
|
||||
# First, we need to define the model configuration. The configuration includes the key parameters
|
||||
# of the model, such as hidden size, intermediate size, etc. Here for convenience, we define a
|
||||
# constant config specially for the TinyLlama model.
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class LlamaConfig:
|
||||
hidden_size: int = 2048
|
||||
intermediate_size: int = 5632
|
||||
num_attention_heads: int = 32
|
||||
num_hidden_layers: int = 22
|
||||
rms_norm_eps: float = 1e-05
|
||||
vocab_size: int = 32000
|
||||
rope_theta: int = 10000
|
||||
context_window_size: int = 2048
|
||||
prefill_chunk_size: int = 2048
|
||||
num_key_value_heads: int = 4
|
||||
head_dim: int = 64 # hidden_size // num_attention_heads
|
||||
|
||||
|
||||
dev = tvm.device("cuda", 0)
|
||||
target = tvm.target.Target.from_device(dev)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Next, we define the RoPE mode of the Paged KV cache. The RoPE mode is used to apply the
|
||||
# Relative Positional Encoding (RoPE) to the query and key tensors. The RoPE mode can be set to
|
||||
# `NONE`, `NORMAL`, or `INLINE`. If the RoPE mode is `NONE`, the KV cache will not apply RoPE to
|
||||
# the query and key tensors. If the RoPE mode is `NORMAL`, RoPE will be applied to the key tensor
|
||||
# before adding the key tensor to the cache. If the RoPE mode is `INLINE`, RoPE will be applied to
|
||||
# the query and key tensors in the attention kernel on-the-fly.
|
||||
|
||||
|
||||
class RopeMode(enum.IntEnum):
|
||||
"""The RoPE mode of the Paged KV cache.
|
||||
If it is none, the KV cache will not apply RoPE to q and k.
|
||||
If it is normal, RoPE will be applied to k before adding k to cache.
|
||||
Otherwise, RoPE will be applied to q/k in attention kernel on-the-fly.
|
||||
"""
|
||||
|
||||
NONE = 0
|
||||
NORMAL = 1
|
||||
INLINE = 2
|
||||
|
||||
|
||||
######################################################################
|
||||
# Secondly, we define the model architecture. The model architecture consists of three parts:
|
||||
#
|
||||
# - Embedding layer: The embedding layer converts the input token IDs to the hidden states.
|
||||
# - Decoder layers: The decoder layers are the core of the model. Each decoder layer consists of
|
||||
# a self-attention layer and a feed-forward network (FFN) layer.
|
||||
# - Output layer: The output layer converts the hidden states to the logits.
|
||||
#
|
||||
# First we define the FFN layer. Note that the following FFN layer is optimized implementation
|
||||
# where we fuse the gate and up projection into one kernel.
|
||||
# The naive implementation of FFN layer is: ``FFN(x) = down_proj(silu(gate(x)) * up(x))``
|
||||
# We could combine the ``gate`` and ``up`` projection into one kernel for better performance.
|
||||
# The optimized implementation is:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# concat_x = gate_up(x)
|
||||
# gate_x, up_x = split(concat_x, 2, axis=-1)
|
||||
# FFN(x) = down_proj(silu(gate_x) * up_x)
|
||||
#
|
||||
|
||||
|
||||
class LlamaFFN(nn.Module):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
super().__init__()
|
||||
self.gate_up_proj = nn.Linear(
|
||||
in_features=config.hidden_size,
|
||||
out_features=2 * config.intermediate_size,
|
||||
bias=False,
|
||||
)
|
||||
self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
concat_x1_x2 = self.gate_up_proj(x)
|
||||
x1, x2 = op.split(concat_x1_x2, 2, axis=-1)
|
||||
return self.down_proj(op.silu(x1) * x2)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Then we define the self-attention layer. The self-attention layer consists of three parts:
|
||||
#
|
||||
# - QKV projection: The QKV projection converts the input hidden states to the query, key, and
|
||||
# value tensors.
|
||||
# - Attention: The attention layer computes the attention scores and applies the softmax
|
||||
# operation.
|
||||
# - Output projection: The output projection converts the attention output to the hidden states.
|
||||
#
|
||||
# We perform optimizations on the different parts of the self-attention layer:
|
||||
#
|
||||
# - QKV projection: We leverage the horizontal fusion on QKV projection and fuse them into one
|
||||
# kernel.
|
||||
# - Attention: We leverage the horizontal fusion on attention and fuse the QKV projection and
|
||||
|
||||
|
||||
class LlamaAttention(nn.Module): # pylint: disable=too-many-instance-attributes
|
||||
def __init__(self, config: LlamaConfig):
|
||||
self.head_dim = config.head_dim
|
||||
self.num_q_heads = config.num_attention_heads
|
||||
self.num_kv_heads = config.num_key_value_heads
|
||||
# horizontal fusion on QKV projection
|
||||
self.qkv_proj = nn.Linear(
|
||||
in_features=config.hidden_size,
|
||||
out_features=(self.num_q_heads + 2 * self.num_kv_heads) * self.head_dim,
|
||||
bias=False,
|
||||
)
|
||||
self.o_proj = nn.Linear(self.num_q_heads * self.head_dim, config.hidden_size, bias=False)
|
||||
|
||||
def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int):
|
||||
d, h_q, h_kv = self.head_dim, self.num_q_heads, self.num_kv_heads
|
||||
b, s, _ = hidden_states.shape
|
||||
# QKV Projection
|
||||
qkv = self.qkv_proj(hidden_states)
|
||||
qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d))
|
||||
# Attention
|
||||
output = op.reshape(
|
||||
paged_kv_cache.attention_with_fused_qkv(
|
||||
layer_id, qkv, self.num_q_heads, sm_scale=self.head_dim**-0.5
|
||||
),
|
||||
(b, s, h_q * d),
|
||||
)
|
||||
# Output Projection
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Finally, we define the model architecture with FFN and self-attention layers.
|
||||
|
||||
|
||||
class LlamaDecoderLayer(nn.Module):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
rms_norm_eps = config.rms_norm_eps
|
||||
self.self_attn = LlamaAttention(config)
|
||||
self.mlp = LlamaFFN(config)
|
||||
self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False)
|
||||
self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False)
|
||||
|
||||
def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int):
|
||||
hidden_states += self.self_attn(
|
||||
self.input_layernorm(hidden_states), paged_kv_cache, layer_id
|
||||
)
|
||||
hidden_states += self.mlp(self.post_attention_layernorm(hidden_states))
|
||||
return hidden_states
|
||||
|
||||
|
||||
class LlamaModel(nn.Module):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
assert config.hidden_size % config.num_attention_heads == 0
|
||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
|
||||
self.layers = nn.ModuleList(
|
||||
[LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)]
|
||||
)
|
||||
self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False)
|
||||
|
||||
def forward(self, input_embed: Tensor, paged_kv_cache: PagedKVCache):
|
||||
hidden_states = input_embed
|
||||
for layer_id, layer in enumerate(self.layers):
|
||||
hidden_states = layer(hidden_states, paged_kv_cache, layer_id)
|
||||
hidden_states = self.norm(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class LlamaForCausalLM(nn.Module):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
self.model = LlamaModel(config)
|
||||
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||
self.num_hidden_layers = config.num_hidden_layers
|
||||
self.num_attention_heads = config.num_attention_heads
|
||||
self.num_key_value_heads = config.num_key_value_heads
|
||||
self.head_dim = config.head_dim
|
||||
self.hidden_size = config.hidden_size
|
||||
self.vocab_size = config.vocab_size
|
||||
self.rope_theta = config.rope_theta
|
||||
self.dtype = "float32"
|
||||
|
||||
def to(self, dtype: str | None = None):
|
||||
super().to(dtype=dtype)
|
||||
if dtype is not None:
|
||||
self.dtype = dtype
|
||||
|
||||
def embed(self, input_ids: Tensor):
|
||||
return self.model.embed_tokens(input_ids)
|
||||
|
||||
def get_logits(self, hidden_states: Tensor):
|
||||
logits = self.lm_head(hidden_states)
|
||||
if logits.dtype != "float32":
|
||||
logits = logits.astype("float32")
|
||||
return logits
|
||||
|
||||
def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache):
|
||||
def _index(x: te.Tensor): # x[:-1,:]
|
||||
b, s, d = x.shape
|
||||
return te.compute((b, 1, d), lambda i, _, k: x[i, s - 1, k], name="index")
|
||||
|
||||
hidden_states = self.model(input_embed, paged_kv_cache)
|
||||
hidden_states = op.tensor_expr_op(_index, name_hint="index", args=[hidden_states])
|
||||
logits = self.get_logits(hidden_states)
|
||||
return logits, paged_kv_cache
|
||||
|
||||
def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache):
|
||||
hidden_states = self.model(input_embed, paged_kv_cache)
|
||||
logits = self.get_logits(hidden_states)
|
||||
return logits, paged_kv_cache
|
||||
|
||||
def create_tir_paged_kv_cache(
|
||||
self,
|
||||
max_batch_size: tirx.Var,
|
||||
max_total_seq_len: tirx.Var,
|
||||
prefill_chunk_size: tirx.Var,
|
||||
page_size: tirx.Var,
|
||||
) -> PagedKVCache:
|
||||
return TIRPagedKVCache(
|
||||
attn_kind="mha",
|
||||
max_batch_size=max_batch_size,
|
||||
max_total_seq_len=max_total_seq_len,
|
||||
prefill_chunk_size=prefill_chunk_size,
|
||||
page_size=page_size,
|
||||
support_sliding_window=0,
|
||||
layer_partition=relax.ShapeExpr([0, self.num_hidden_layers]),
|
||||
num_hidden_layers=self.num_hidden_layers,
|
||||
num_attention_heads=self.num_attention_heads,
|
||||
num_key_value_heads=self.num_key_value_heads,
|
||||
qk_head_dim=self.head_dim,
|
||||
v_head_dim=self.head_dim,
|
||||
mla_original_qk_head_dim=0,
|
||||
mla_original_v_head_dim=0,
|
||||
rope_mode=RopeMode.NORMAL,
|
||||
rope_scale=1,
|
||||
rope_theta=self.rope_theta,
|
||||
rope_scaling={},
|
||||
rope_ext_factors=tirx.IntImm("int64", 0),
|
||||
rotary_dim=self.head_dim,
|
||||
dtype=self.dtype,
|
||||
target=target,
|
||||
enable_disaggregation=False,
|
||||
)
|
||||
|
||||
def get_default_spec(self):
|
||||
mod_spec = {
|
||||
"embed": {
|
||||
"input_ids": nn.spec.Tensor(["seq_len"], "int32"),
|
||||
"$": {
|
||||
"param_mode": "packed",
|
||||
"effect_mode": "none",
|
||||
},
|
||||
},
|
||||
"prefill": {
|
||||
"input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype),
|
||||
"paged_kv_cache": nn.spec.Object(object_type=PagedKVCache),
|
||||
"$": {
|
||||
"param_mode": "packed",
|
||||
"effect_mode": "none",
|
||||
},
|
||||
},
|
||||
"decode": {
|
||||
"input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype),
|
||||
"paged_kv_cache": nn.spec.Object(object_type=PagedKVCache),
|
||||
"$": {
|
||||
"param_mode": "packed",
|
||||
"effect_mode": "none",
|
||||
},
|
||||
},
|
||||
"create_tir_paged_kv_cache": {
|
||||
"max_batch_size": int,
|
||||
"max_total_seq_len": int,
|
||||
"prefill_chunk_size": int,
|
||||
"page_size": int,
|
||||
"$": {
|
||||
"param_mode": "none",
|
||||
"effect_mode": "none",
|
||||
},
|
||||
},
|
||||
}
|
||||
return nn.spec.ModuleSpec.from_raw(mod_spec, self)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Export the model to Relax IRModule
|
||||
# ----------------------------------
|
||||
# After defining the model architecture, we can export the model to the Relax IRModule.
|
||||
# For demonstration, we only show the part of the model architecture. and parameters.
|
||||
|
||||
model_config = LlamaConfig()
|
||||
model = LlamaForCausalLM(model_config)
|
||||
model.to("float16")
|
||||
mod, named_params = model.export_tvm(spec=model.get_default_spec())
|
||||
prefill_str = mod["prefill"].script()
|
||||
print(*prefill_str.split("\n")[3:20], sep="\n") # Only show the first 10 lines for demonstration
|
||||
print(" ...")
|
||||
|
||||
print("\nParameters:")
|
||||
pprint(named_params[:5]) # Only show the first 5 parameters for demonstration
|
||||
|
||||
######################################################################
|
||||
# Define Optimization Pipeline
|
||||
# ----------------------------
|
||||
# We define a series of optimization passes to optimize the model. The optimization pipeline
|
||||
# is designed specifically for the LLMs.
|
||||
|
||||
|
||||
@register_pipeline("opt_llm")
|
||||
def _pipeline( # pylint: disable=too-many-arguments
|
||||
ext_mods: list[nn.ExternModule] | None = None,
|
||||
):
|
||||
ext_mods = ext_mods or []
|
||||
|
||||
@tvm.transform.module_pass(opt_level=0)
|
||||
def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
|
||||
seq = tvm.transform.Sequential(
|
||||
[
|
||||
# Phase 1. Passes on high-level operator graph
|
||||
# We can enable cublas for further optimization
|
||||
relax.transform.FuseTransposeMatmul(),
|
||||
# Phase 2. Lowering to TIR, inherited TVM Relax's official "zero" pipeline
|
||||
relax.transform.LegalizeOps(),
|
||||
relax.transform.AnnotateTIROpPattern(),
|
||||
relax.transform.FoldConstant(),
|
||||
relax.transform.FuseOps(),
|
||||
relax.transform.FuseTIR(),
|
||||
# Phase 3. Passes on TIR
|
||||
relax.transform.DeadCodeElimination(),
|
||||
# Phase 4. Low-level Optimizations
|
||||
dlight.ApplyDefaultSchedule(
|
||||
dlight.gpu.Matmul(),
|
||||
dlight.gpu.GEMV(),
|
||||
dlight.gpu.Reduction(),
|
||||
dlight.gpu.GeneralReduction(),
|
||||
dlight.gpu.Fallback(),
|
||||
),
|
||||
# Phase 5. Lowering to VM bytecode
|
||||
relax.transform.RewriteDataflowReshape(),
|
||||
relax.transform.ToNonDataflow(),
|
||||
relax.transform.RemovePurityChecking(),
|
||||
relax.transform.CallTIRRewrite(),
|
||||
relax.transform.StaticPlanBlockMemory(),
|
||||
relax.transform.RewriteCUDAGraph(),
|
||||
relax.transform.LowerAllocTensor(),
|
||||
relax.transform.KillAfterLastUse(),
|
||||
relax.transform.LowerRuntimeBuiltin(),
|
||||
relax.transform.VMShapeLower(),
|
||||
relax.transform.AttachGlobalSymbol(),
|
||||
relax.transform.AttachExternModules(ext_mods),
|
||||
]
|
||||
)
|
||||
mod = seq(mod)
|
||||
return mod
|
||||
|
||||
return _pipeline
|
||||
|
||||
|
||||
with target:
|
||||
ex = tvm.compile(mod, target, relax_pipeline=relax.get_pipeline("opt_llm"))
|
||||
vm = relax.VirtualMachine(ex, dev)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Prepare the model weights
|
||||
# -------------------------
|
||||
# We load the pre-trained weights from Hugging Face and prepare the model weights.
|
||||
# The pre-trained weights are stored in the Hugging Face format. We need to load the weights
|
||||
# and prepare the model parameters.
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# Note that we won't execute the following code in this tutorial because the pre-trained weights
|
||||
# are not available in the CI environment.
|
||||
#
|
||||
|
||||
|
||||
IS_IN_CI = os.getenv("CI", "") == "true"
|
||||
|
||||
HF_WEIGHT_PATH = None
|
||||
# HF_WEIGHT_PATH = Path("/path/to/TinyLlama-1.1B-Chat-v1.0/")
|
||||
|
||||
if not IS_IN_CI:
|
||||
import numpy as np
|
||||
import safetensors.torch
|
||||
import torch
|
||||
|
||||
if HF_WEIGHT_PATH is None or not HF_WEIGHT_PATH.exists():
|
||||
raise ValueError("Please set the HF_WEIGHT_PATH to the path of the pre-trained weights.")
|
||||
|
||||
# Torch format weights
|
||||
param_dict = safetensors.torch.load_file(HF_WEIGHT_PATH / "model.safetensors", device="cpu")
|
||||
# Numpy format weights
|
||||
param_dict = {
|
||||
k: v.half().numpy() if v.dtype == torch.bfloat16 else v.numpy()
|
||||
for k, v in param_dict.items()
|
||||
}
|
||||
|
||||
named_params = dict(named_params)
|
||||
for i in range(model_config.num_hidden_layers):
|
||||
# Add QKV in self attention
|
||||
attn = f"model.layers.{i}.self_attn"
|
||||
param_dict[f"{attn}.qkv_proj.weight"] = np.concatenate(
|
||||
[
|
||||
param_dict.pop(f"{attn}.q_proj.weight"), # Pop the old parameters to save memory
|
||||
param_dict.pop(f"{attn}.k_proj.weight"),
|
||||
param_dict.pop(f"{attn}.v_proj.weight"),
|
||||
],
|
||||
axis=0,
|
||||
)
|
||||
# Add gates in MLP
|
||||
mlp = f"model.layers.{i}.mlp"
|
||||
param_dict[f"{mlp}.gate_up_proj.weight"] = np.concatenate(
|
||||
[
|
||||
param_dict.pop(f"{mlp}.gate_proj.weight"),
|
||||
param_dict.pop(f"{mlp}.up_proj.weight"),
|
||||
],
|
||||
axis=0,
|
||||
)
|
||||
|
||||
# Convert params into ndarray
|
||||
params = [
|
||||
tvm.runtime.tensor(param_dict[k].astype("float16"), device=dev) for k in named_params.keys()
|
||||
]
|
||||
|
||||
|
||||
######################################################################
|
||||
# Deploy the compiled model
|
||||
# -------------------------
|
||||
# After the model and weights are ready, we can deploy the compiled model on the target device.
|
||||
# The language models inference includes two steps: prefill and decode. The prefill step is
|
||||
# used to process the input tokens and store the KVCache. The decode step is used to generate
|
||||
# the token until the end token is generated.
|
||||
|
||||
|
||||
######################################################################
|
||||
# Tokenization
|
||||
# ~~~~~~~~~~~~
|
||||
# The first step is to tokenize the input prompt and embed the tokens into the hidden states.
|
||||
# The tokenization and embedding are the same as the original model. We use the HF tokenizer
|
||||
# to tokenize the input prompt and embed the tokens into the hidden states.
|
||||
# Note that different models require different tokenization and prompt format, please refer to
|
||||
# the model documentation for the correct tokenization and prompt format.
|
||||
|
||||
|
||||
if not IS_IN_CI:
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(HF_WEIGHT_PATH)
|
||||
messages = [
|
||||
{"role": "user", "content": "What's your name?"},
|
||||
]
|
||||
prompt = tokenizer.apply_chat_template(messages)
|
||||
input_len = len(prompt)
|
||||
|
||||
# Load prompt tokens into TVM ndarray on the target device
|
||||
tokens = tvm.runtime.tensor(np.array(prompt).astype("int32"), device=dev)
|
||||
|
||||
######################################################################
|
||||
# Create the KVCache
|
||||
# ~~~~~~~~~~~~~~~~~~
|
||||
# Before starting the inference, we need to create the KVCache. The KVCache is used to store the
|
||||
# key and value tensors for the attention layer. Apache TVM provides a PagedKVCache to store the
|
||||
# key and value tensors. We create the PagedKVCache with the specified parameters.
|
||||
|
||||
if not IS_IN_CI:
|
||||
kv_cache = vm["create_tir_paged_kv_cache"](
|
||||
Shape([1]), # max_batch_size=1
|
||||
Shape([2048]), # max_total_seq_len=2048
|
||||
Shape([2048]), # prefill_chunk_size=2048
|
||||
Shape([16]), # page_size=16
|
||||
)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Embedding
|
||||
# ~~~~~~~~~
|
||||
# The next step is to embed the tokens into the hidden states. We use the `embed` function
|
||||
# compiled in the Relax IRModule to embed the tokens into the hidden states.
|
||||
|
||||
nd_view_func = tvm.get_global_func("vm.builtin.reshape")
|
||||
|
||||
|
||||
def embed(tokens, params):
|
||||
_embed = vm["embed"](tokens, params)
|
||||
# Reshape hidden from [seq_len, hidden_size] to [1, seq_len, hidden_size]
|
||||
_embed = nd_view_func(_embed, Shape([1, _embed.shape[0], _embed.shape[1]]))
|
||||
return _embed
|
||||
|
||||
|
||||
######################################################################
|
||||
# Prefill
|
||||
# ~~~~~~~
|
||||
# Before running the forward pass, we first get some help functions for preparation.
|
||||
|
||||
add_sequence_func = tvm.get_global_func("vm.builtin.kv_state_add_sequence")
|
||||
begin_forward_func = tvm.get_global_func("vm.builtin.kv_state_begin_forward")
|
||||
end_forward_func = tvm.get_global_func("vm.builtin.kv_state_end_forward")
|
||||
|
||||
######################################################################
|
||||
# As we are creating a new sequence, we need to call `add_sequence_func` to initialize
|
||||
# the request. Additionally, we need to call `begin_forward_func` to start the forward pass,
|
||||
# and `end_forward_func` to end the forward pass.
|
||||
|
||||
if not IS_IN_CI:
|
||||
seq_id = 0
|
||||
add_sequence_func(kv_cache, seq_id)
|
||||
hidden_states = embed(tokens, params)
|
||||
begin_forward_func(kv_cache, Shape([seq_id]), Shape([input_len]))
|
||||
logits, kv_cache = vm["prefill"](hidden_states, kv_cache, params)
|
||||
end_forward_func(kv_cache)
|
||||
|
||||
######################################################################
|
||||
# Now we have the output logits from the prefill step. The logits are used to generate the token
|
||||
# via sampling. Let's sample the token from the logits.
|
||||
#
|
||||
# In this tutorial, we simplify the sampling process and pick the token with the highest
|
||||
# probability. In practice, we should sample the token based on the probability distribution.
|
||||
# Also, to make the tutorial concise, we execute the sample process on CPU.
|
||||
|
||||
|
||||
def sample_token(logits):
|
||||
logits_np = logits.numpy()
|
||||
return np.argmax(logits_np)
|
||||
|
||||
|
||||
if not IS_IN_CI:
|
||||
last_token = sample_token(logits)
|
||||
output_tokens = [last_token]
|
||||
|
||||
|
||||
######################################################################
|
||||
# Decode
|
||||
# ~~~~~~
|
||||
# After the prefill step, we can start the decode step. The decode step is used to generate the
|
||||
# token until the end token is generated. We use the `decode` function compiled in the Relax
|
||||
# IRModule to generate the token.
|
||||
|
||||
if not IS_IN_CI:
|
||||
print("The generated token:")
|
||||
|
||||
while last_token != tokenizer.eos_token_id:
|
||||
tokens = tvm.runtime.tensor(np.array([last_token]).astype("int32"), device=dev)
|
||||
hidden_states = embed(tokens, params)
|
||||
begin_forward_func(kv_cache, Shape([seq_id]), Shape([1]))
|
||||
logits, kv_cache = vm["decode"](hidden_states, kv_cache, params)
|
||||
|
||||
end_forward_func(kv_cache)
|
||||
last_token = sample_token(logits)
|
||||
output_tokens.append(last_token)
|
||||
|
||||
print(tokenizer.decode(output_tokens))
|
||||
Reference in New Issue
Block a user