chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,2 @@
|
||||
doxygen
|
||||
modules
|
||||
@@ -0,0 +1,5 @@
|
||||
[rstcheck]
|
||||
report_level = warning
|
||||
ignore_directives = automodule, autosummary, currentmodule, toctree, ifconfig, tab-set, collapse, tabs, dropdown, doxygenclass, doxygenfunction, doxygenstruct, doxygennamespace, figure, code-block, graphviz
|
||||
ignore_roles = ref, cpp:class, cpp:func, py:func, c:macro, doc, any
|
||||
ignore_languages = cpp, python
|
||||
+2342
File diff suppressed because it is too large
Load Diff
+250
@@ -0,0 +1,250 @@
|
||||
# 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.
|
||||
|
||||
# Makefile for Sphinx documentation
|
||||
#
|
||||
|
||||
# You can set these variables from the command line.
|
||||
SPHINXOPTS =
|
||||
SPHINXBUILD = python3 -m sphinx
|
||||
PAPER =
|
||||
BUILDDIR = _build
|
||||
STAGINGDIR = _staging
|
||||
|
||||
# Internal variables.
|
||||
PAPEROPT_a4 = -D latex_paper_size=a4
|
||||
PAPEROPT_letter = -D latex_paper_size=letter
|
||||
ALLSPHINXOPTS = -d $(PWD)/$(BUILDDIR)/doctrees $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) .
|
||||
# the i18n builder cannot share the environment and doctrees with the others
|
||||
I18NSPHINXOPTS = $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) .
|
||||
|
||||
.PHONY: help clean html dirhtml singlehtml pickle json htmlhelp qthelp devhelp epub latex latexpdf text man changes linkcheck doctest coverage gettext staging
|
||||
|
||||
help:
|
||||
@echo "Please use \`make <target>' where <target> is one of"
|
||||
@echo " html to make standalone HTML files"
|
||||
@echo " dirhtml to make HTML files named index.html in directories"
|
||||
@echo " singlehtml to make a single large HTML file"
|
||||
@echo " pickle to make pickle files"
|
||||
@echo " json to make JSON files"
|
||||
@echo " htmlhelp to make HTML files and a HTML help project"
|
||||
@echo " qthelp to make HTML files and a qthelp project"
|
||||
@echo " applehelp to make an Apple Help Book"
|
||||
@echo " devhelp to make HTML files and a Devhelp project"
|
||||
@echo " epub to make an epub"
|
||||
@echo " latex to make LaTeX files, you can set PAPER=a4 or PAPER=letter"
|
||||
@echo " latexpdf to make LaTeX files and run them through pdflatex"
|
||||
@echo " latexpdfja to make LaTeX files and run them through platex/dvipdfmx"
|
||||
@echo " text to make text files"
|
||||
@echo " man to make manual pages"
|
||||
@echo " texinfo to make Texinfo files"
|
||||
@echo " info to make Texinfo files and run them through makeinfo"
|
||||
@echo " gettext to make PO message catalogs"
|
||||
@echo " changes to make an overview of all changed/added/deprecated items"
|
||||
@echo " xml to make Docutils-native XML files"
|
||||
@echo " pseudoxml to make pseudoxml-XML files for display purposes"
|
||||
@echo " linkcheck to check all external links for integrity"
|
||||
@echo " doctest to run all doctests embedded in the documentation (if enabled)"
|
||||
@echo " coverage to run coverage check of the documentation (if enabled)"
|
||||
|
||||
clean:
|
||||
rm -rf $(BUILDDIR)
|
||||
rm -rf $(STAGINGDIR)
|
||||
|
||||
# TODO(Lunderberg): Remove these lines once the CI steps have
|
||||
# propagated.
|
||||
|
||||
# Remove folders that have since been relocated into
|
||||
# $(STAGINGDIR). This allows `task_python_docs.sh` to
|
||||
# run, even if a commit that predates $(STAGINGDIR) was
|
||||
# previously run on that node.
|
||||
rm -rf gen_modules
|
||||
rm -rf user_tutorials
|
||||
rm -rf tutorials
|
||||
|
||||
staging:
|
||||
# Prepare the staging directory. Sphinx gallery automatically
|
||||
# writes new .rst files into the current directory. This can
|
||||
# cause issues when switching branches. By sequestering the
|
||||
# auto-generated files into the staging directory, they can be
|
||||
# removed without knowing the exact directory.
|
||||
|
||||
mkdir -p $(STAGINGDIR)
|
||||
|
||||
# Remove any symlinks that currently exist
|
||||
find $(STAGINGDIR) -type l -exec rm {} \;
|
||||
|
||||
# Reproduce the directory structure
|
||||
find . \
|
||||
-path ./$(BUILDDIR) -prune -o -path ./$(STAGINGDIR) -prune -o \
|
||||
-name "*.rst" \
|
||||
-printf "$(STAGINGDIR)/%h\n" \
|
||||
| sort | uniq | xargs mkdir -p
|
||||
|
||||
# Symlink all .rst files into the staging directory
|
||||
find . \
|
||||
-path ./$(BUILDDIR) -prune -o -path ./$(STAGINGDIR) -prune -o \
|
||||
-name "*.rst" \
|
||||
-exec ln -s $(PWD)/{} $(STAGINGDIR)/{} \;
|
||||
|
||||
ln -s $(PWD)/conf.py $(STAGINGDIR)/conf.py
|
||||
ln -s $(PWD)/_static $(STAGINGDIR)/_static
|
||||
|
||||
|
||||
html: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b html $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/html
|
||||
@echo
|
||||
@echo "Build finished. The HTML pages are in $(BUILDDIR)/html."
|
||||
|
||||
htmldepoly: html
|
||||
python3 $(PWD)/download_3rdparty_embeds.py -v
|
||||
@echo "Replaced external URLs with local files."
|
||||
|
||||
dirhtml: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b dirhtml $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/dirhtml
|
||||
@echo
|
||||
@echo "Build finished. The HTML pages are in $(BUILDDIR)/dirhtml."
|
||||
|
||||
singlehtml: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b singlehtml $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/singlehtml
|
||||
@echo
|
||||
@echo "Build finished. The HTML page is in $(BUILDDIR)/singlehtml."
|
||||
|
||||
pickle: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b pickle $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/pickle
|
||||
@echo
|
||||
@echo "Build finished; now you can process the pickle files."
|
||||
|
||||
json: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b json $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/json
|
||||
@echo
|
||||
@echo "Build finished; now you can process the JSON files."
|
||||
|
||||
htmlhelp: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b htmlhelp $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/htmlhelp
|
||||
@echo
|
||||
@echo "Build finished; now you can run HTML Help Workshop with the" \
|
||||
".hhp project file in $(BUILDDIR)/htmlhelp."
|
||||
|
||||
qthelp: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b qthelp $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/qthelp
|
||||
@echo
|
||||
@echo "Build finished; now you can run "qcollectiongenerator" with the" \
|
||||
".qhcp project file in $(BUILDDIR)/qthelp, like this:"
|
||||
@echo "# qcollectiongenerator $(BUILDDIR)/qthelp/rabit.qhcp"
|
||||
@echo "To view the help file:"
|
||||
@echo "# assistant -collectionFile $(BUILDDIR)/qthelp/rabit.qhc"
|
||||
|
||||
applehelp: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b applehelp $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/applehelp
|
||||
@echo
|
||||
@echo "Build finished. The help book is in $(BUILDDIR)/applehelp."
|
||||
@echo "N.B. You won't be able to view it unless you put it in" \
|
||||
"~/Library/Documentation/Help or install it in your application" \
|
||||
"bundle."
|
||||
|
||||
devhelp: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b devhelp $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/devhelp
|
||||
@echo
|
||||
@echo "Build finished."
|
||||
@echo "To view the help file:"
|
||||
@echo "# mkdir -p $$HOME/.local/share/devhelp/rabit"
|
||||
@echo "# ln -s $(BUILDDIR)/devhelp $$HOME/.local/share/devhelp/rabit"
|
||||
@echo "# devhelp"
|
||||
|
||||
epub: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b epub $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/epub
|
||||
@echo
|
||||
@echo "Build finished. The epub file is in $(BUILDDIR)/epub."
|
||||
|
||||
latex: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/latex
|
||||
@echo
|
||||
@echo "Build finished; the LaTeX files are in $(BUILDDIR)/latex."
|
||||
@echo "Run \`make' in that directory to run these through (pdf)latex" \
|
||||
"(use \`make latexpdf' here to do that automatically)."
|
||||
|
||||
latexpdf: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/latex
|
||||
@echo "Running LaTeX files through pdflatex..."
|
||||
$(MAKE) -C $(BUILDDIR)/latex all-pdf
|
||||
@echo "pdflatex finished; the PDF files are in $(BUILDDIR)/latex."
|
||||
|
||||
latexpdfja: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/latex
|
||||
@echo "Running LaTeX files through platex and dvipdfmx..."
|
||||
$(MAKE) -C $(BUILDDIR)/latex all-pdf-ja
|
||||
@echo "pdflatex finished; the PDF files are in $(BUILDDIR)/latex."
|
||||
|
||||
text: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b text $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/text
|
||||
@echo
|
||||
@echo "Build finished. The text files are in $(BUILDDIR)/text."
|
||||
|
||||
man: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b man $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/man
|
||||
@echo
|
||||
@echo "Build finished. The manual pages are in $(BUILDDIR)/man."
|
||||
|
||||
texinfo: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b texinfo $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/texinfo
|
||||
@echo
|
||||
@echo "Build finished. The Texinfo files are in $(BUILDDIR)/texinfo."
|
||||
@echo "Run \`make' in that directory to run these through makeinfo" \
|
||||
"(use \`make info' here to do that automatically)."
|
||||
|
||||
info: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b texinfo $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/texinfo
|
||||
@echo "Running Texinfo files through makeinfo..."
|
||||
make -C $(BUILDDIR)/texinfo info
|
||||
@echo "makeinfo finished; the Info files are in $(BUILDDIR)/texinfo."
|
||||
|
||||
gettext: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b gettext $(I18NSPHINXOPTS) $(PWD)/$(BUILDDIR)/locale
|
||||
@echo
|
||||
@echo "Build finished. The message catalogs are in $(BUILDDIR)/locale."
|
||||
|
||||
changes: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b changes $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/changes
|
||||
@echo
|
||||
@echo "The overview file is in $(BUILDDIR)/changes."
|
||||
|
||||
linkcheck: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b linkcheck $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/linkcheck
|
||||
@echo
|
||||
@echo "Link check complete; look for any errors in the above output " \
|
||||
"or in $(BUILDDIR)/linkcheck/output.txt."
|
||||
|
||||
doctest: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b doctest $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/doctest
|
||||
@echo "Testing of doctests in the sources finished, look at the " \
|
||||
"results in $(BUILDDIR)/doctest/output.txt."
|
||||
|
||||
coverage: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b coverage $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/coverage
|
||||
@echo "Testing of coverage in the sources finished, look at the " \
|
||||
"results in $(BUILDDIR)/coverage/python.txt."
|
||||
|
||||
xml: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b xml $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/xml
|
||||
@echo
|
||||
@echo "Build finished. The XML files are in $(BUILDDIR)/xml."
|
||||
|
||||
pseudoxml: staging
|
||||
cd $(STAGINGDIR) && $(SPHINXBUILD) -b pseudoxml $(ALLSPHINXOPTS) $(PWD)/$(BUILDDIR)/pseudoxml
|
||||
@echo
|
||||
@echo "Build finished. The pseudo-XML files are in $(BUILDDIR)/pseudoxml."
|
||||
+173
@@ -0,0 +1,173 @@
|
||||
<!--- Licensed to the Apache Software Foundation (ASF) under one -->
|
||||
<!--- or more contributor license agreements. See the NOTICE file -->
|
||||
<!--- distributed with this work for additional information -->
|
||||
<!--- regarding copyright ownership. The ASF licenses this file -->
|
||||
<!--- to you under the Apache License, Version 2.0 (the -->
|
||||
<!--- "License"); you may not use this file except in compliance -->
|
||||
<!--- with the License. You may obtain a copy of the License at -->
|
||||
|
||||
<!--- http://www.apache.org/licenses/LICENSE-2.0 -->
|
||||
|
||||
<!--- Unless required by applicable law or agreed to in writing, -->
|
||||
<!--- software distributed under the License is distributed on an -->
|
||||
<!--- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -->
|
||||
<!--- KIND, either express or implied. See the License for the -->
|
||||
<!--- specific language governing permissions and limitations -->
|
||||
<!--- under the License. -->
|
||||
|
||||
# TVM Documentation
|
||||
|
||||
This folder contains the source of TVM's documentation, hosted at https://tvm.apache.org/docs
|
||||
|
||||
## Build Locally
|
||||
|
||||
### With Docker (recommended)
|
||||
|
||||
1. Build TVM and the docs inside the [tlcpack/ci-gpu image](https://hub.docker.com/r/tlcpack/ci-gpu) using the [`ci.py`](../tests/scripts/ci.py) script.
|
||||
|
||||
```bash
|
||||
# If this runs into errors, try cleaning your 'build' directory
|
||||
python tests/scripts/ci.py docs
|
||||
|
||||
# See other doc building options
|
||||
python tests/scripts/ci.py docs --help
|
||||
```
|
||||
|
||||
2. Serve the docs and visit http://localhost:8000 in your browser
|
||||
|
||||
```bash
|
||||
# Run an HTTP server you can visit to view the docs in your browser
|
||||
python tests/scripts/ci.py serve-docs
|
||||
```
|
||||
|
||||
### Native
|
||||
|
||||
1. [Build TVM](https://tvm.apache.org/docs/install/from_source.html) first in the repo root folder, then make it importable:
|
||||
|
||||
```bash
|
||||
export TVM_HOME=/path-to-tvm
|
||||
export TVM_LIBRARY_PATH=$TVM_HOME/build
|
||||
pip install --target=$TVM_HOME/python $TVM_HOME/3rdparty/tvm-ffi
|
||||
export PYTHONPATH=$TVM_HOME/python:$PYTHONPATH
|
||||
```
|
||||
|
||||
`docs/conf.py` unconditionally imports `tvm` at startup, so the build will
|
||||
fail immediately if TVM is not importable.
|
||||
|
||||
2. Install dependencies
|
||||
|
||||
```bash
|
||||
# Pillow on Ubuntu may require libjpeg-dev from apt
|
||||
./docker/bash.sh ci_gpu -c \
|
||||
'python3 -m pip install --quiet sphinx-book-theme==1.1.4 && python3 -m pip freeze' > frozen-requirements.txt
|
||||
|
||||
pip install -r frozen-requirements.txt
|
||||
```
|
||||
|
||||
3. Generate the docs
|
||||
|
||||
```bash
|
||||
# TVM_TUTORIAL_EXEC_PATTERN=none skips the tutorial execution to the build
|
||||
# work on most environments (e.g. MacOS).
|
||||
export TVM_TUTORIAL_EXEC_PATTERN=none
|
||||
|
||||
cd docs
|
||||
make html
|
||||
```
|
||||
|
||||
4. Run an HTTP server and visit http://localhost:8000 in your browser
|
||||
|
||||
```bash
|
||||
cd _build/html && python3 -m http.server
|
||||
```
|
||||
|
||||
## Only Execute Specified Tutorials
|
||||
|
||||
The document build process will execute all the tutorials in the sphinx gallery.
|
||||
This will cause failure in some cases when certain machines do not have necessary
|
||||
environment. You can set `TVM_TUTORIAL_EXEC_PATTERN` to only execute
|
||||
the path that matches the regular expression pattern.
|
||||
|
||||
For example, to only build tutorials under `/get_started/tutorials`, run
|
||||
|
||||
```bash
|
||||
python tests/scripts/ci.py docs --tutorial-pattern=/get_started/tutorials
|
||||
```
|
||||
|
||||
To only build one specific file, do
|
||||
|
||||
```bash
|
||||
# The slash \ is used to get . in regular expression
|
||||
python tests/scripts/ci.py docs --tutorial-pattern=file_name\.py
|
||||
```
|
||||
|
||||
## Helper Scripts
|
||||
|
||||
The following script mirrors the CI docs pipeline: it runs a sphinx pre-check
|
||||
(when not running locally) and then performs a full `make htmldepoly` build,
|
||||
including tutorial execution. You will need a GPU CI environment.
|
||||
|
||||
```bash
|
||||
tests/scripts/task_python_docs.sh
|
||||
```
|
||||
|
||||
To build docs locally without executing tutorials (fastest local iteration):
|
||||
|
||||
```bash
|
||||
cd docs && TVM_TUTORIAL_EXEC_PATTERN=none make html
|
||||
```
|
||||
|
||||
Note: the sphinx pre-check (warning validation) only runs in CI (`IS_LOCAL=0`).
|
||||
`python tests/scripts/ci.py docs` always sets `IS_LOCAL=1` and skips the
|
||||
pre-check regardless of other flags.
|
||||
|
||||
To run the full build including tutorial executions:
|
||||
|
||||
```bash
|
||||
python tests/scripts/ci.py docs --full
|
||||
```
|
||||
|
||||
## Define the Order of Tutorials
|
||||
|
||||
You can define the order of tutorials with `subsection_order` and
|
||||
`within_subsection_order` in [`conf.py`](conf.py).
|
||||
By default, the tutorials within one subsection are sorted by filename.
|
||||
|
||||
## Google Colab Integration
|
||||
|
||||
All the TVM tutorials can be opened and used interactively in Google Colab by
|
||||
clicking the button at the top of the page. To do this, `sphinx-gallery` builds
|
||||
`.ipynb` files from each tutorial, which are automatically deployed to the
|
||||
[apache/tvm-site](https://github.com/apache/tvm-site/tree/asf-site) repo's
|
||||
`asf-site` branch by [@tvm-bot](https://github.com/tvm-bot).
|
||||
|
||||
To make sure your tutorial runs correctly on Colab, any non-Python parts of
|
||||
the tutorial (e.g. dependency installations) should be prefixed by an
|
||||
[IPython magic command](https://ipython.readthedocs.io/en/stable/interactive/magics.html).
|
||||
These will not be included in the built `HTML` file. For example, to install
|
||||
Pytorch in your tutorial, add a ReStructured Text block like the following:
|
||||
|
||||
```python
|
||||
######################################################################
|
||||
# To run this tutorial, we must install PyTorch:
|
||||
#
|
||||
# .. code-block:: bash
|
||||
#
|
||||
# %%shell
|
||||
# pip install torch
|
||||
#
|
||||
```
|
||||
|
||||
### Interactive Bash Scripts
|
||||
|
||||
In stock IPython, the `%%bash` magic command should be used to run shell
|
||||
commands. However, this command does not give real-time output - the
|
||||
tutorial's user will not see any output until the entire cell finishes
|
||||
running. When running commands that take several minutes (e.g. installing
|
||||
dependencies), this is annoying.
|
||||
|
||||
Luckily, Google Colab has the `%%shell` magic command that does the same
|
||||
thing as `%%bash`, but gives output in real time. This command is specific
|
||||
to Colab, and its [source code](https://github.com/googlecolab/colabtools)
|
||||
is public. Thus, `%%shell` should be used instead of `%%bash` when writing
|
||||
TVM tutorials.
|
||||
Vendored
+2
@@ -0,0 +1,2 @@
|
||||
The logo file in this repo is an exception due to the need of sphinx.
|
||||
By default we avoid to put large binary blobs into this repo.
|
||||
BIN
Binary file not shown.
|
After Width: | Height: | Size: 104 KiB |
Vendored
BIN
Binary file not shown.
|
After Width: | Height: | Size: 6.5 KiB |
Vendored
BIN
Binary file not shown.
|
After Width: | Height: | Size: 3.4 KiB |
@@ -0,0 +1,306 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
.. _codegen-arch:
|
||||
|
||||
Code Generation
|
||||
===============
|
||||
|
||||
Code generation is the final stage of the TVM compilation pipeline — it translates TIR
|
||||
``PrimFunc``\ s into executable code for a target device. This document explains how TIR
|
||||
functions become native CPU instructions, GPU kernels, or source code strings, covering the
|
||||
target dispatch mechanism, the two codegen families (LLVM and Source), and the runtime module
|
||||
system that wraps the generated code.
|
||||
|
||||
|
||||
Where Codegen Fits
|
||||
------------------
|
||||
|
||||
When a user calls ``tvm.compile()``, the compilation proceeds in two phases:
|
||||
|
||||
1. **Relax phase**: the Relax pipeline optimizes and fuses the computational graph, then
|
||||
``VMCodeGen`` translates Relax functions into VM bytecode (see :ref:`relax-vm-arch`).
|
||||
2. **TIR phase**: TIR ``PrimFunc``\ s (the actual compute kernels) are compiled to native code.
|
||||
|
||||
The TIR phase is handled internally by ``tirx.build()`` (called from ``relax.build()``).
|
||||
It performs these steps:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
TIR PrimFuncs (in IRModule)
|
||||
│
|
||||
▼ TIR pipeline ← lowering passes (flatten buffers, lower intrinsics, etc.)
|
||||
TIR PrimFuncs (lowered)
|
||||
│
|
||||
▼ split_host_device_mods() ← separate host and device functions
|
||||
Host IRModule + Device IRModule(s)
|
||||
│ │
|
||||
▼ ▼
|
||||
codegen_build() codegen_build() ← target-specific code generation
|
||||
│ │
|
||||
▼ ▼
|
||||
Host Module Device Module(s)
|
||||
│ │
|
||||
▼ import_module() │
|
||||
Host Module ◄─────────────┘ ← device modules imported into host
|
||||
│
|
||||
▼ (returned to relax.build for linking with VM bytecode)
|
||||
|
||||
|
||||
Target Dispatch
|
||||
---------------
|
||||
|
||||
The core dispatch logic lives in ``codegen::Build()`` (``src/target/codegen.cc``), which is
|
||||
called from the Python-side ``codegen_build()`` in ``tirx/build.py``. It selects the correct
|
||||
backend based on the ``Target`` object:
|
||||
|
||||
.. code-block:: cpp
|
||||
|
||||
ffi::Module Build(IRModule mod, Target target) {
|
||||
std::string build_f_name = "target.build." + target->kind->name;
|
||||
const auto bf = tvm::ffi::Function::GetGlobal(build_f_name);
|
||||
return (*bf)(mod, target).cast<ffi::Module>();
|
||||
}
|
||||
|
||||
Each backend registers its build function via FFI:
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 25 30 45
|
||||
|
||||
* - FFI Key
|
||||
- Backend
|
||||
- Codegen Class
|
||||
* - ``target.build.llvm``
|
||||
- CPU (x86, ARM, etc.)
|
||||
- ``CodeGenCPU`` (→ LLVM IR → machine code)
|
||||
* - ``target.build.cuda``
|
||||
- NVIDIA GPU
|
||||
- ``CodeGenCUDA`` (→ CUDA C → PTX/cubin)
|
||||
* - ``target.build.rocm``
|
||||
- AMD GPU
|
||||
- ``CodeGenAMDGPU`` (→ LLVM IR → AMDGPU ISA)
|
||||
* - ``target.build.nvptx``
|
||||
- NVIDIA PTX
|
||||
- ``CodeGenNVPTX`` (→ LLVM IR → PTX)
|
||||
* - ``target.build.metal``
|
||||
- Apple GPU
|
||||
- ``CodeGenMetal`` (→ Metal Shading Language)
|
||||
* - ``target.build.opencl``
|
||||
- OpenCL devices
|
||||
- ``CodeGenOpenCL`` (→ OpenCL C)
|
||||
* - ``target.build.vulkan``
|
||||
- Vulkan devices
|
||||
- ``CodeGenSPIRV`` (→ SPIR-V binary)
|
||||
* - ``target.build.webgpu``
|
||||
- WebGPU
|
||||
- ``CodeGenWebGPU`` (→ WGSL)
|
||||
* - ``target.build.c``
|
||||
- C host code
|
||||
- ``CodeGenCHost`` (→ C source)
|
||||
|
||||
|
||||
Two Codegen Families
|
||||
--------------------
|
||||
|
||||
TVM has two families of code generators, corresponding to two fundamentally different strategies
|
||||
for producing executable code:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
LLVM Family Source Family
|
||||
────────── ─────────────
|
||||
TIR → LLVM IR → machine code TIR → source string → external compiler
|
||||
(in-process, JIT or AOT) (CUDA C, OpenCL C, Metal, WGSL)
|
||||
|
||||
LLVM family
|
||||
~~~~~~~~~~~
|
||||
|
||||
``CodeGenLLVM`` (``src/target/llvm/codegen_llvm.h``) translates TIR directly to LLVM IR using
|
||||
the LLVM C++ API. The generated ``llvm::Module`` is then compiled to native code by LLVM's
|
||||
backend (x86, ARM, NVPTX, AMDGPU, etc.).
|
||||
|
||||
**Inheritance**:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
CodeGenLLVM (base)
|
||||
├── CodeGenCPU ← x86, ARM (target.build.llvm)
|
||||
│ └── CodeGenHexagon
|
||||
├── CodeGenNVPTX ← NVIDIA PTX via LLVM (target.build.nvptx)
|
||||
└── CodeGenAMDGPU ← AMD GPU via LLVM (target.build.rocm)
|
||||
|
||||
``CodeGenLLVM`` inherits from both ``ExprFunctor<llvm::Value*(const Expr&)>`` and
|
||||
``StmtFunctor<void(const Stmt&)>``. Each TIR node type has a corresponding visitor:
|
||||
|
||||
- **Expressions** (``VisitExpr_``) convert TIR expressions to LLVM ``Value``\ s:
|
||||
arithmetic ops → LLVM binary instructions, ``BufferLoad`` → load with pointer arithmetic,
|
||||
``Cast`` → LLVM type conversions, ``Call`` → intrinsic or extern function calls.
|
||||
- **Statements** (``VisitStmt_``) emit LLVM IR side effects:
|
||||
``BufferStore`` → store instructions, ``For`` → loop basic blocks with branches,
|
||||
``IfThenElse`` → conditional branches, ``AllocBuffer`` → stack or heap allocation.
|
||||
|
||||
The key methods on ``CodeGenLLVM`` are:
|
||||
|
||||
- ``Create(LLVMTarget*)`` — factory that returns a target-specific subclass.
|
||||
- ``Init(...)`` — set up the LLVM context, module, and builder.
|
||||
- ``DeclareFunction(gvar, f)`` / ``AddFunction(gvar, f)`` — forward-declare then compile a
|
||||
``PrimFunc`` to LLVM IR.
|
||||
- ``Finish()`` — return the completed ``llvm::Module``.
|
||||
|
||||
Source family
|
||||
~~~~~~~~~~~~~
|
||||
|
||||
``CodeGenC`` (``src/target/source/codegen_c.h``) generates C-like source code as text. Each
|
||||
target subclass overrides methods to emit target-specific syntax.
|
||||
|
||||
**Inheritance**:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
CodeGenC (base)
|
||||
├── CodeGenCUDA ← CUDA C (target.build.cuda)
|
||||
├── CodeGenOpenCL ← OpenCL C (target.build.opencl)
|
||||
├── CodeGenMetal ← Metal Shading Language (target.build.metal)
|
||||
├── CodeGenWebGPU ← WGSL (target.build.webgpu)
|
||||
└── CodeGenCHost ← C host code (target.build.c)
|
||||
|
||||
``CodeGenC`` also uses the visitor pattern (``ExprFunctor`` and ``StmtFunctor``), but outputs to
|
||||
``std::ostream`` instead of constructing LLVM IR. Subclasses override target-specific methods:
|
||||
|
||||
- ``PrintStorageScope(scope, os)`` — emit memory qualifiers (e.g., ``__shared__`` for CUDA,
|
||||
``__local`` for OpenCL).
|
||||
- ``BindThreadIndex(iv)`` — emit thread index bindings (e.g., ``threadIdx.x``, ``blockIdx.y``).
|
||||
- ``PrintType(dtype, os)`` — emit target-specific type names (e.g., ``half`` for float16).
|
||||
- ``PrintVecBinaryOp(...)`` — emit vectorized operations in target syntax.
|
||||
|
||||
For CUDA, the build flow (``BuildCUDA`` in ``src/target/opt/build_cuda_on.cc``) is:
|
||||
|
||||
1. ``CodeGenCUDA`` generates CUDA C source.
|
||||
2. An optional post-processing callback (``tvm_callback_cuda_postproc``) transforms the source.
|
||||
3. A Python callback (``tvm_callback_cuda_compile``) compiles the source to PTX or cubin via
|
||||
NVRTC or NVCC.
|
||||
4. The result is wrapped in a ``CUDAModule``.
|
||||
|
||||
Design choice
|
||||
~~~~~~~~~~~~~
|
||||
|
||||
Why two families?
|
||||
|
||||
- **LLVM family** produces higher-quality code — LLVM applies its own optimization passes
|
||||
(instruction selection, register allocation, vectorization). Best for CPU targets where TVM
|
||||
has full control over the compilation.
|
||||
- **Source family** is more portable — it generates human-readable source that can be compiled
|
||||
by vendor toolchains (NVCC, Metal compiler, etc.). This is necessary for GPU targets where
|
||||
the vendor compiler handles device-specific optimizations and the runtime compilation model
|
||||
(e.g., NVRTC for CUDA, runtime shader compilation for Metal/OpenCL).
|
||||
|
||||
|
||||
Host/Device Split
|
||||
-----------------
|
||||
|
||||
When compiling for GPU targets, TIR functions are split into two categories:
|
||||
|
||||
- **Host functions** — run on the CPU. They set up kernel launch parameters (grid/block
|
||||
dimensions), allocate memory, and invoke device kernels. Compiled with ``target.build.llvm``
|
||||
or ``target.build.c``.
|
||||
- **Device functions** — the actual compute kernels that run on the GPU. Compiled with the
|
||||
target-specific codegen (``target.build.cuda``, etc.).
|
||||
|
||||
``split_host_device_mods()`` (``python/tvm/tirx/build.py``) separates functions by their
|
||||
``target`` attribute: functions whose target kind is ``"llvm"`` or ``"c"`` go to the host
|
||||
module; all others go to device modules grouped by target.
|
||||
|
||||
After compilation, device modules are imported into the host module via ``import_module()``,
|
||||
forming a module tree. At runtime, the host module dispatches to the imported device module
|
||||
when a device kernel is called.
|
||||
|
||||
|
||||
Runtime Modules
|
||||
---------------
|
||||
|
||||
Each codegen produces a ``runtime.Module`` — the container that holds the generated code and
|
||||
exposes it as callable ``PackedFunc``\ s.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 20 35 45
|
||||
|
||||
* - Module Type
|
||||
- How Code Is Stored
|
||||
- How Code Is Executed
|
||||
* - ``LLVMModule``
|
||||
- LLVM IR (in-memory ``llvm::Module``)
|
||||
- JIT-compiled on first call (MCJIT or ORC). Function pointers cached for subsequent calls.
|
||||
* - ``CUDAModule``
|
||||
- PTX or cubin binary
|
||||
- Loaded via CUDA driver API (``cuModuleLoad``). Kernels launched via ``cuLaunchKernel``.
|
||||
* - ``CSourceModule``
|
||||
- C source string
|
||||
- Not directly executable. Used as a build artifact for AOT compilation.
|
||||
* - ``DeviceSourceModule``
|
||||
- Device source string (OpenCL C, Metal, WGSL)
|
||||
- Compiled at runtime by the device driver (e.g., ``clCreateProgramWithSource``).
|
||||
|
||||
All module types implement the same interface: ``GetFunction(name)`` returns a ``PackedFunc``
|
||||
that can be called from Python or C++. The VM and other runtime components use this interface
|
||||
to invoke compiled kernels without knowing which backend produced them.
|
||||
|
||||
The module tree is serializable via ``export_library()``, which packs the host module and all
|
||||
imported device modules into a single shared library (``.so`` / ``.dll`` / ``.dylib``) or
|
||||
a tar archive for deployment.
|
||||
|
||||
|
||||
Source Code Map
|
||||
---------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 50 50
|
||||
|
||||
* - Path
|
||||
- Contents
|
||||
* - ``python/tvm/tirx/build.py``
|
||||
- ``tirx.build()``: TIR compilation entry, host/device split, module linking
|
||||
* - ``src/target/codegen.cc``
|
||||
- ``codegen::Build()``: target dispatch via ``"target.build.<kind>"``
|
||||
* - ``src/target/llvm/codegen_llvm.h``
|
||||
- ``CodeGenLLVM``: TIR → LLVM IR base class
|
||||
* - ``src/target/llvm/codegen_cpu.h``
|
||||
- ``CodeGenCPU``: CPU-specific LLVM codegen (x86, ARM)
|
||||
* - ``src/target/llvm/codegen_nvptx.cc``
|
||||
- ``CodeGenNVPTX``: NVIDIA PTX via LLVM
|
||||
* - ``src/target/llvm/codegen_amdgpu.cc``
|
||||
- ``CodeGenAMDGPU``: AMD GPU via LLVM
|
||||
* - ``src/target/llvm/llvm_module.cc``
|
||||
- ``LLVMModuleNode``: runtime module with JIT compilation
|
||||
* - ``src/target/source/codegen_c.h``
|
||||
- ``CodeGenC``: TIR → C-like source base class
|
||||
* - ``src/target/source/codegen_cuda.h``
|
||||
- ``CodeGenCUDA``: TIR → CUDA C
|
||||
* - ``src/target/source/codegen_opencl.h``
|
||||
- ``CodeGenOpenCL``: TIR → OpenCL C
|
||||
* - ``src/target/source/codegen_metal.h``
|
||||
- ``CodeGenMetal``: TIR → Metal Shading Language
|
||||
* - ``src/target/source/codegen_c_host.h``
|
||||
- ``CodeGenCHost``: TIR → C host code
|
||||
* - ``src/target/opt/build_cuda_on.cc``
|
||||
- ``BuildCUDA``: CUDA build flow (codegen → compile → module)
|
||||
* - ``src/target/spirv/codegen_spirv.h``
|
||||
- ``CodeGenSPIRV``: TIR → SPIR-V for Vulkan
|
||||
* - ``src/target/source/codegen_webgpu.h``
|
||||
- ``CodeGenWebGPU``: TIR → WGSL
|
||||
@@ -0,0 +1,245 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
|
||||
.. _tvm-target-specific-overview:
|
||||
|
||||
Device/Target Interactions
|
||||
==========================
|
||||
|
||||
This documented is intended for developers interested in understanding
|
||||
how the TVM framework interacts with specific device APIs, or who
|
||||
may want to implement support for a new API or new hardware.
|
||||
|
||||
There are three main aspects that must be implemented for any new
|
||||
runtime environment.
|
||||
|
||||
* The :ref:`DeviceAPI <tvm-target-specific-device-api>` class gives a
|
||||
handle to a specific device, and the API used to interact with it.
|
||||
It defines a common interface for querying device parameters
|
||||
(e.g. memory available, number of threads, etc.) and for performing
|
||||
simple actions (e.g. copying memory from the host, or between
|
||||
buffers on the device).
|
||||
|
||||
* The :ref:`Target <tvm-target-specific-target>` class contains a
|
||||
description of the device on which a function will run. It is
|
||||
exposed both to the target code generators and to the optimization
|
||||
passes.
|
||||
|
||||
* The :ref:`target code generators <tvm-target-specific-codegen>`
|
||||
construct a :ref:`Module <tvm-runtime-system-module>` consisting of
|
||||
one or more :ref:`PackedFunc <tvm-runtime-system-packed-func>`, from
|
||||
an IRModule.
|
||||
|
||||
.. _tvm-target-specific-device-api:
|
||||
|
||||
DeviceAPI
|
||||
---------
|
||||
|
||||
The ``DeviceAPI`` represents a handle to a specific hardware device
|
||||
API. (e.g. ``CUDADeviceAPI`` handles all interactions through the
|
||||
CUDA framework.) Most ``DeviceAPI`` methods accept a ``device_id``
|
||||
parameter to specify which device should be accessed. In Python,
|
||||
these are typically accessed using the :py:func:`tvm.runtime.device`
|
||||
function, which returns a handle to a specific device, accessed
|
||||
through a specific API. (e.g. ``tvm.runtime.device('cuda',0)`` gives
|
||||
access to physical device ``0``, accessed through the CUDA API.)
|
||||
|
||||
.. _device_api.h: https://github.com/apache/tvm/blob/main/include/tvm/runtime/device_api.h
|
||||
|
||||
* Attribute queries - ``GetAttr`` allows different
|
||||
device-specific parameters to be queried, such as the device name,
|
||||
number of threads, etc. The parameters that can be queried are
|
||||
defined in ``enum DeviceAttrKind`` in `device_api.h`_. Not all
|
||||
query-able parameters are supported by all devices. If a parameter
|
||||
cannot be queried (e.g. ``kMaxClockRate`` on Vulkan), or if a
|
||||
parameter isn't applicable (e.g. ``kWarpSize`` on CPU), then those
|
||||
queries should return ``nullptr``.
|
||||
|
||||
* Setting active device - ``SetDevice`` should set a
|
||||
particular device as being active. If a ``PackedFunc`` generated by
|
||||
the target-specific code gen requires execution on a device, it
|
||||
should run on the active device.
|
||||
|
||||
* Memory management - Utilities for allocating and deallocating memory
|
||||
on the device.
|
||||
|
||||
* Allocate data space - ``AllocDataSpace`` and ``FreeDataSpace``
|
||||
allocate and free space on the device. These allocations can be
|
||||
provided as inputs and outputs to an operator and make up the
|
||||
primary data flow of the operator graph. It must be possible to
|
||||
transfer data from the host to/from a data space. The return
|
||||
value is an opaque ``void*``. While some implementations return a
|
||||
memory address, this is not required, and the ``void*`` may be an
|
||||
opaque handle that is interpretable only by the device backend
|
||||
that generated it. The ``void*`` is used as an argument to other
|
||||
backend-specific functions, such as ``CopyDataFromTo``.
|
||||
|
||||
* Allocate work space - ``AllocWorkspace`` and ``FreeWorkspace``
|
||||
allocate and free space on the device. Unlike data space, these
|
||||
are used for storage of intermediate values within an operator
|
||||
definition, and are not required to be transferable to/from the
|
||||
host device. If a ``DeviceAPI`` subclass does not implement these
|
||||
methods, they will default to calling the corresponding
|
||||
``DataSpace`` functions.
|
||||
|
||||
* Copy data - ``CopyDataFromTo`` should copy data from one location
|
||||
to another. The type of copy is determined by the ``dev_from``
|
||||
and ``dev_to`` parameters. Implementations should support copying
|
||||
memory from CPU to device, from device to CPU, and from one buffer
|
||||
to another on a single device. If the source or destination
|
||||
locations are on the CPU, the corresponding ``void*`` points to a
|
||||
CPU address that can be passed into ``memcpy``. If the source or
|
||||
destinations locations are on the device, the corresponding
|
||||
``void*`` was previously generated by either ``AllocDataSpace`` or
|
||||
``AllocWorkspace``.
|
||||
|
||||
These copies are queued to execute on a specific
|
||||
``TVMStreamHandle``. However, implementations should not assume
|
||||
that CPU buffers remains valid or accessible after the call to
|
||||
``CopyDataFromTo`` completes.
|
||||
|
||||
|
||||
* Execution stream management - utilities for handling
|
||||
``TVMStreamHandle``, which represents parallel streams of execution
|
||||
used to execute commands.
|
||||
|
||||
* Create stream - ``CreateStream`` and ``FreeStream`` should
|
||||
allocate/free a handle to a stream of execution. If a device
|
||||
implements only a single queue of commands, then ``CreateStream``
|
||||
should return ``nullptr``.
|
||||
|
||||
* Set active stream - ``SetStream`` should set a stream as being
|
||||
active. While active, if a ``PackedFunc`` generated by the
|
||||
target-specific code gen requires execution on a device, the work
|
||||
should be submitted to the active stream.
|
||||
|
||||
* Synchronize to CPU - ``StreamSync`` should synchronize a stream of
|
||||
execution to the CPU. The call to ``StreamSync`` should return
|
||||
once all memory transfers and computations submitted prior to the
|
||||
``StreamSync`` call have completed.
|
||||
|
||||
* Synchronize between streams - ``SyncStreamFromTo`` should
|
||||
introduce a synchronization barrier between the source and
|
||||
destination stream. That is, the destination stream may not
|
||||
proceed beyond commands currently queued until the source stream
|
||||
has completed all commands that are currently queued.
|
||||
|
||||
|
||||
In order to be usable by the TVM framework, the new DeviceAPI should
|
||||
then be registered with the following steps.
|
||||
|
||||
#. Create a function that instantiates the new DeviceAPI, and returns
|
||||
a pointer to it::
|
||||
|
||||
FooDeviceAPI* FooDeviceAPI::Global() {
|
||||
static FooDeviceAPI inst;
|
||||
return &inst;
|
||||
}
|
||||
|
||||
#. Register the function to the tvm registry::
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def("device_api.foo", FooDeviceAPI::Global);
|
||||
}
|
||||
|
||||
.. _base.h: https://github.com/apache/tvm/blob/main/include/tvm/runtime/base.h
|
||||
|
||||
#. Add an entry for the new DeviceAPI to the ``TVMDeviceExtType`` enum
|
||||
in `base.h`_. The value should be an unused value greater
|
||||
than ``DLDeviceType::kDLExtDev``, but less than
|
||||
``DeviceAPIManager::kMaxDeviceAPI``.
|
||||
|
||||
#. Add a case in ``DeviceName`` in `device_api.h`_ to convert from the
|
||||
enum value to a string representation. This string representation
|
||||
should match the name given to ``GlobalDef().def``.
|
||||
|
||||
#. Add entries to the ``_DEVICE_TYPE_TO_NAME`` and ``_DEVICE_NAME_TO_TYPE`` dictionaries of
|
||||
:py:class:`tvm.runtime.Device` for the new enum value.
|
||||
|
||||
|
||||
.. _tvm-target-specific-target:
|
||||
|
||||
Target Definition
|
||||
-----------------
|
||||
|
||||
The ``Target`` object is a lookup table of properties about a physical
|
||||
device, its hardware/driver limits, and its capabilities. The
|
||||
``Target`` is accessible both during optimization and code generation
|
||||
stages. While the same ``Target`` class is used for all runtime
|
||||
targets, each runtime target may need to add target-specific options.
|
||||
|
||||
.. _target_kind.cc: https://github.com/apache/tvm/blob/main/src/target/target_kind.cc
|
||||
|
||||
In `target_kind.cc`_, add a new declaration of
|
||||
``TVM_REGISTER_TARGET_KIND``, passing a string name of the new target,
|
||||
and the ``TVMDeviceExtType`` or ``DLDeviceType`` enum value for the
|
||||
device on which that target should run. Typically, the target name
|
||||
and the device name will match (e.g., the ``"cuda"`` target runs on
|
||||
the ``kDLCUDA`` device). There are exceptions, such as when multiple
|
||||
different code generation targets can run on the same physical device
|
||||
(e.g., the ``"llvm"`` and ``"c"`` targets both run on the ``kDLCPU``
|
||||
device type).
|
||||
|
||||
All options for a specific target kind are added with the
|
||||
``add_attr_option`` function, with optional default values. A `Target`
|
||||
parser can be added with ``set_target_parser`` to process
|
||||
any parameters that are dynamically based on other parameters or
|
||||
queried from device properties.
|
||||
|
||||
This argument definition defines a parser that can unpack a string
|
||||
description of a target. This is done in the ``Target::Target(const
|
||||
String&)`` constructor in C++, which accepts a JSON-formatted string
|
||||
and is typically called using the :py:class:`tvm.target.Target` python
|
||||
object. For example, ``tvm.target.Target('{"kind": "cuda",
|
||||
"max_num_threads": 1024}')`` will create a ``cuda`` target, while
|
||||
overriding the default maximum number of threads.
|
||||
|
||||
In a code generator, the target properties can be accessed using
|
||||
``target->GetAttr<T>(param_name)`` in C++, or with the
|
||||
``target.attrs`` dictionary in Python.
|
||||
|
||||
|
||||
.. _tvm-target-specific-codegen:
|
||||
|
||||
Target Code Generators
|
||||
----------------------
|
||||
|
||||
The code generators take an optimized ``IRModule`` and converts it
|
||||
into an executable representation. Each code generator must be
|
||||
registered in order to be used by the TVM framework. This is done by
|
||||
registering a function named ``"target.build.foo"``, where ``foo`` is
|
||||
the same name as was used in the ``TVM_REGISTER_TARGET_KIND``
|
||||
definition above. ::
|
||||
|
||||
tvm::runtime::Module GeneratorFooCode(IRModule mod, Target target);
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def("target.build.foo", GeneratorFooCode);
|
||||
}
|
||||
|
||||
The code generator takes two arguments. The first is the ``IRModule``
|
||||
to compile, and the second is the ``Target`` that describes the device
|
||||
on which the code should run. Because the environment performing the
|
||||
compilation is not necessarily the same as the environment that will
|
||||
be executing the code, code generators should not perform any
|
||||
attribute lookups on the device itself, and should instead access
|
||||
parameters stored in the ``Target``.
|
||||
|
||||
Each function in the input ``IRModule`` should be accessible by name
|
||||
in the output ``runtime::Module``.
|
||||
@@ -0,0 +1,361 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
.. _external-library-dispatch:
|
||||
|
||||
External Library Dispatch (BYOC)
|
||||
================================
|
||||
|
||||
When deploying models, certain operator patterns (e.g., matmul + bias + relu) can be executed
|
||||
more efficiently by vendor-optimized libraries such as cuBLAS, CUTLASS, cuDNN, or DNNL. TVM's
|
||||
**BYOC (Bring Your Own Codegen)** mechanism identifies these patterns in a Relax module and
|
||||
offloads them to external backends, while keeping the rest of the computation on TVM's own
|
||||
generated kernels.
|
||||
|
||||
This document explains the BYOC pipeline: how patterns are registered, how subgraphs are
|
||||
matched and extracted, how backend code generators are invoked, and how the externally compiled
|
||||
code is executed at runtime.
|
||||
|
||||
|
||||
Overview
|
||||
--------
|
||||
|
||||
The BYOC pipeline consists of four stages:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
IRModule (high-level Relax IR)
|
||||
│
|
||||
▼ FuseOpsByPattern ← match high-level ops, create composite functions
|
||||
IRModule (with Composite + Codegen attributes)
|
||||
│
|
||||
▼ RunCodegen ← invoke backend codegen via FFI
|
||||
IRModule (with call_dps_packed to ExternFunc)
|
||||
+ external runtime Modules
|
||||
│
|
||||
▼ LegalizeOps + FuseOps + ... ← compile remaining ops normally
|
||||
│
|
||||
▼ VM compilation ← link external modules into executable
|
||||
Deployable artifact
|
||||
|
||||
Each stage is a Relax transformation pass that operates on the ``IRModule``:
|
||||
|
||||
1. **FuseOpsByPattern** — matches operator subgraphs against registered patterns and groups them
|
||||
into composite functions annotated with ``Composite`` and ``Codegen`` attributes.
|
||||
2. **MergeCompositeFunctions** (optional) — merges multiple composite functions targeting the same
|
||||
backend when inter-operator dependencies allow.
|
||||
3. **RunCodegen** — finds all functions with a ``Codegen`` attribute, invokes the corresponding
|
||||
backend code generator via FFI, and replaces the original calls with ``call_dps_packed``
|
||||
to externally compiled functions.
|
||||
4. **Linking** — the resulting external ``runtime.Module``\ s are attached to the ``IRModule``
|
||||
as the ``external_mods`` attribute and bundled into the final executable during
|
||||
``relax.build()``.
|
||||
|
||||
|
||||
Pattern Registration
|
||||
--------------------
|
||||
|
||||
Each backend registers the operator patterns it supports in a **global pattern registry**
|
||||
(``python/tvm/relax/backend/pattern_registry.py``). The registry is a static table that maps
|
||||
pattern names to ``FusionPattern`` objects.
|
||||
|
||||
Registering patterns
|
||||
~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.relax.backend.pattern_registry import register_patterns
|
||||
from tvm.relax.backend.patterns import make_matmul_pattern
|
||||
|
||||
register_patterns([
|
||||
(
|
||||
"cublas.matmul", # pattern name (prefix = backend)
|
||||
*make_matmul_pattern( # returns (DFPattern, annotation_patterns)
|
||||
with_bias=False,
|
||||
),
|
||||
_check_matmul, # check function
|
||||
),
|
||||
(
|
||||
"cublas.matmul_bias_relu",
|
||||
*make_matmul_pattern(
|
||||
with_bias=True,
|
||||
activation="relax.nn.relu",
|
||||
),
|
||||
_check_matmul,
|
||||
),
|
||||
# ... more patterns
|
||||
])
|
||||
|
||||
Each entry is a tuple of ``(name, pattern, annotation_patterns, check_func)`` that gets
|
||||
converted to a ``FusionPattern`` object. The name prefix (e.g., ``"cublas"``) identifies the
|
||||
backend; ``get_patterns_with_prefix("cublas")`` retrieves all patterns for that backend.
|
||||
|
||||
Patterns registered later have **higher priority** — when a subgraph matches multiple patterns,
|
||||
the highest-priority match wins.
|
||||
|
||||
Pattern templates
|
||||
~~~~~~~~~~~~~~~~~
|
||||
|
||||
``python/tvm/relax/backend/patterns.py`` provides reusable templates for common patterns:
|
||||
|
||||
- ``make_matmul_pattern(with_bias, activation, transposed_rhs)`` — matmul with optional bias
|
||||
and activation fusion
|
||||
- ``make_conv2d_pattern(with_bias, activation)`` — 2D convolution
|
||||
- ``make_attention_pattern()`` — multi-head attention
|
||||
- ``make_residual_block_pattern()`` — residual connections
|
||||
- ``make_layer_norm_pattern()`` / ``make_rms_norm_pattern()`` — normalization layers
|
||||
|
||||
Each template returns ``(DFPattern, Mapping[str, DFPattern])`` — the main pattern and its
|
||||
annotation sub-patterns.
|
||||
|
||||
Check functions
|
||||
~~~~~~~~~~~~~~~
|
||||
|
||||
The check function validates whether a matched subgraph can actually be handled by the backend.
|
||||
It receives a ``PatternCheckContext`` and returns ``True`` to accept or ``False`` to reject.
|
||||
|
||||
Typical checks include:
|
||||
|
||||
- **Data type support**: verify the operand dtypes are supported (e.g., cuBLAS supports
|
||||
float16, float32, int8, bfloat16, float8 for matmul).
|
||||
- **Shape constraints**: verify reduction axes are constant, batch dimensions are compatible.
|
||||
- **Leaking intermediates**: reject if an intermediate result is used outside the fused group
|
||||
(via ``has_leaking_intermediate_variables()``).
|
||||
|
||||
|
||||
Partitioning
|
||||
------------
|
||||
|
||||
After patterns are registered, a backend provides a **partition function** that applies
|
||||
``FuseOpsByPattern`` to an ``IRModule``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# python/tvm/relax/backend/cuda/cublas.py
|
||||
def partition_for_cublas(mod, bind_constants=False):
|
||||
patterns = get_patterns_with_prefix("cublas")
|
||||
return transform.FuseOpsByPattern(
|
||||
patterns, bind_constants=bind_constants, annotate_codegen=True
|
||||
)(mod)
|
||||
|
||||
With ``annotate_codegen=True``, each matched subgraph is wrapped in a two-level function
|
||||
structure:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
# Outer function — tagged for the codegen backend
|
||||
@R.function
|
||||
def fused_relax_matmul_cublas0(args...):
|
||||
R.func_attr({"Codegen": "cublas", "global_symbol": "fused_relax_matmul_cublas0"})
|
||||
...
|
||||
# Inner function — identifies the specific pattern
|
||||
@R.function(private=True)
|
||||
def composite(args...):
|
||||
R.func_attr({"Composite": "cublas.matmul_bias_relu"})
|
||||
lv0 = R.matmul(x, w)
|
||||
lv1 = R.add(lv0, bias)
|
||||
lv2 = R.nn.relu(lv1)
|
||||
return lv2
|
||||
...
|
||||
|
||||
The outer function carries the ``Codegen`` attribute that ``RunCodegen`` uses to dispatch to the
|
||||
right backend. The inner function carries the ``Composite`` attribute that the backend codegen
|
||||
uses to identify which operation to emit.
|
||||
|
||||
MergeCompositeFunctions
|
||||
~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
When ``annotate_codegen=False``, ``FuseOpsByPattern`` only creates inner functions with
|
||||
``Composite`` attributes. A separate ``MergeCompositeFunctions`` pass then groups multiple
|
||||
composite functions targeting the same backend into a single outer function with ``Codegen``
|
||||
and ``global_symbol`` attributes.
|
||||
|
||||
This is useful when multiple sequential operations should be sent to the same backend as a
|
||||
single unit (e.g., a sequence of cuBLAS matmuls that share intermediate results). The pass
|
||||
checks that merging does not create cyclic dependencies between groups.
|
||||
|
||||
|
||||
Code Generation
|
||||
---------------
|
||||
|
||||
``RunCodegen`` (``src/relax/transform/run_codegen.cc``) is the pass that triggers backend
|
||||
code generation:
|
||||
|
||||
1. Scan the module for all functions with a ``Codegen`` attribute.
|
||||
2. Group them by backend target name.
|
||||
3. For each backend, look up the registered codegen function via FFI key
|
||||
``"relax.ext.<backend>"`` (e.g., ``"relax.ext.cublas"``).
|
||||
4. Call the codegen function, which returns an array of compiled ``runtime.Module``\ s.
|
||||
5. Replace the original function calls with ``call_dps_packed(ExternFunc(...), args)``.
|
||||
6. Attach the compiled modules to the ``IRModule`` as the ``external_mods`` attribute.
|
||||
|
||||
Codegen registration
|
||||
~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Each backend registers a codegen function via TVM's FFI mechanism:
|
||||
|
||||
.. code-block:: cpp
|
||||
|
||||
// src/relax/backend/contrib/cublas/codegen.cc
|
||||
ffi::Array<ffi::Module> CublasCompiler(
|
||||
ffi::Array<Function> functions,
|
||||
ffi::Map<ffi::String, ffi::Any> options,
|
||||
ffi::Map<Constant, ffi::String> constant_names) {
|
||||
ffi::Array<ffi::Module> compiled_functions;
|
||||
for (const auto& func : functions) {
|
||||
CublasJSONSerializer serializer(constant_names, AnalyzeVar2Value(func));
|
||||
serializer.serialize(func);
|
||||
auto graph_json = serializer.GetJSON();
|
||||
auto names = serializer.GetConstantNames();
|
||||
const auto pf = ffi::Function::GetGlobalRequired("runtime.CublasJSONRuntimeCreate");
|
||||
compiled_functions.push_back(
|
||||
pf(GetExtSymbol(func), graph_json, names).cast<ffi::Module>());
|
||||
}
|
||||
return compiled_functions;
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def("relax.ext.cublas", CublasCompiler);
|
||||
}
|
||||
|
||||
The codegen function receives:
|
||||
|
||||
- ``functions``: the Relax functions with ``Codegen`` attribute to compile.
|
||||
- ``options``: backend-specific compilation options.
|
||||
- ``constant_names``: mapping from constant values to their names (for weight handling).
|
||||
|
||||
It returns an array of ``runtime.Module`` objects — one per function — that contain the
|
||||
externally compiled code.
|
||||
|
||||
Codegen strategies
|
||||
~~~~~~~~~~~~~~~~~~
|
||||
|
||||
TVM provides two base classes for implementing backend codegens:
|
||||
|
||||
- **JSONSerializer** (``src/relax/backend/contrib/codegen_json/codegen_json.h``): serializes the
|
||||
composite function into a JSON graph representation. At runtime, a backend-specific JSON
|
||||
runtime module interprets the graph and dispatches to library calls. Used by cuBLAS, cuDNN,
|
||||
and most backends.
|
||||
|
||||
- **CSourceCodegen** (``src/relax/backend/contrib/codegen_c/codegen_c.h``): generates C/CUDA
|
||||
source code that is compiled and linked. Used when the backend requires ahead-of-time
|
||||
compilation.
|
||||
|
||||
|
||||
Runtime Execution
|
||||
-----------------
|
||||
|
||||
After ``RunCodegen``, the original high-level function calls are replaced with:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
R.call_dps_packed(ExternFunc("fused_relax_matmul_cublas0"), (x, w, bias), ...)
|
||||
|
||||
At runtime, ``call_dps_packed`` invokes the externally compiled function through the
|
||||
``PackedFunc`` interface. The external ``runtime.Module``\ s (produced by the codegen) are
|
||||
imported into the final executable during ``relax.build()`` and are available via the module's
|
||||
function lookup mechanism.
|
||||
|
||||
For JSON-based backends (cuBLAS, cuDNN), the runtime module deserializes the JSON graph and
|
||||
dispatches each node to the corresponding library API call. For source-based backends, the
|
||||
compiled native code is called directly.
|
||||
|
||||
|
||||
Adding a New Backend
|
||||
--------------------
|
||||
|
||||
To add support for a new external library:
|
||||
|
||||
1. **Define patterns** in ``python/tvm/relax/backend/<target>/``:
|
||||
|
||||
- Create DFPatterns using templates from ``patterns.py`` or custom patterns.
|
||||
- Write check functions to validate dtypes, shapes, and other constraints.
|
||||
- Register patterns with ``register_patterns()``.
|
||||
- Provide a ``partition_for_<backend>(mod)`` convenience function.
|
||||
|
||||
2. **Implement codegen** in ``src/relax/backend/contrib/<target>/``:
|
||||
|
||||
- Subclass ``JSONSerializer`` or ``CSourceCodegen``.
|
||||
- Implement the visitor that converts composite functions to the target format.
|
||||
- Register the codegen function as ``"relax.ext.<target>"``.
|
||||
|
||||
3. **Implement runtime** (for JSON-based backends):
|
||||
|
||||
- Create a JSON runtime module that interprets the serialized graph and dispatches
|
||||
to the library's API calls.
|
||||
- Register the runtime constructor as ``"runtime.<Target>JSONRuntimeCreate"``.
|
||||
|
||||
|
||||
Supported Backends
|
||||
------------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 15 25 60
|
||||
|
||||
* - Backend
|
||||
- Patterns
|
||||
- Operations
|
||||
* - cuBLAS
|
||||
- ``cublas.*``
|
||||
- Matmul (with bias, activation, transpose, dequantize variants)
|
||||
* - CUTLASS
|
||||
- ``cutlass.*``
|
||||
- Matmul, conv2d, attention, residual blocks, decode matmul
|
||||
* - cuDNN
|
||||
- ``cudnn.*``
|
||||
- Conv2d (NHWC/NCHW), stacked attention
|
||||
* - DNNL
|
||||
- ``dnnl.*``
|
||||
- Matmul, conv2d (x86 CPU). Codegen exists at C++ level; patterns are
|
||||
defined in tests rather than pre-registered.
|
||||
|
||||
|
||||
Source Code Map
|
||||
---------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 50 50
|
||||
|
||||
* - Path
|
||||
- Contents
|
||||
* - ``python/tvm/relax/backend/pattern_registry.py``
|
||||
- Pattern registry API (register_patterns, get_patterns_with_prefix)
|
||||
* - ``python/tvm/relax/backend/patterns.py``
|
||||
- Reusable pattern templates (make_matmul_pattern, etc.)
|
||||
* - ``python/tvm/relax/backend/cuda/cublas.py``
|
||||
- cuBLAS patterns and partition_for_cublas
|
||||
* - ``python/tvm/relax/backend/cuda/cutlass.py``
|
||||
- CUTLASS patterns and partition_for_cutlass
|
||||
* - ``python/tvm/relax/backend/cuda/cudnn.py``
|
||||
- cuDNN patterns and partition_for_cudnn
|
||||
* - ``src/relax/backend/pattern_registry.cc``
|
||||
- Pattern registry C++ implementation
|
||||
* - ``src/relax/transform/run_codegen.cc``
|
||||
- RunCodegen pass (CodeGenRunner)
|
||||
* - ``src/relax/transform/merge_composite_functions.cc``
|
||||
- MergeCompositeFunctions pass
|
||||
* - ``src/relax/backend/contrib/cublas/codegen.cc``
|
||||
- cuBLAS codegen (JSONSerializer-based)
|
||||
* - ``src/relax/backend/contrib/cutlass/codegen.cc``
|
||||
- CUTLASS codegen
|
||||
* - ``src/relax/backend/contrib/codegen_json/codegen_json.h``
|
||||
- JSONSerializer base class
|
||||
* - ``src/relax/backend/contrib/codegen_c/codegen_c.h``
|
||||
- CSourceCodegen base class
|
||||
@@ -0,0 +1,388 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
.. _fusion-arch:
|
||||
|
||||
Operator Fusion
|
||||
===============
|
||||
|
||||
Operator fusion is one of the most impactful optimizations in TVM. Instead of launching one kernel
|
||||
per operator (e.g., conv2d, bias_add, relu), fusion merges multiple operators into a single kernel,
|
||||
eliminating intermediate memory allocations and kernel launch overhead.
|
||||
|
||||
TVM provides two complementary fusion mechanisms:
|
||||
|
||||
- **Automatic fusion** (``FuseOps`` + ``FuseTIR``): groups operators based on their computational
|
||||
patterns using a post-dominator analysis algorithm.
|
||||
- **Pattern-based fusion** (``FuseOpsByPattern``): groups operators that match user-defined
|
||||
dataflow patterns, typically for offloading to external backends (cuBLAS, CUTLASS, DNNL, etc.).
|
||||
|
||||
Both produce the same output: Relax functions marked with ``Primitive=True`` that are later
|
||||
lowered to fused TIR kernels or dispatched to external libraries.
|
||||
|
||||
Overview
|
||||
--------
|
||||
|
||||
Fusion involves three passes:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
IRModule (after LegalizeOps)
|
||||
│
|
||||
▼ AnnotateTIROpPattern ← label each op (elementwise, reduce, etc.)
|
||||
IRModule (annotated)
|
||||
│
|
||||
▼ FuseOps ← group ops into fused Relax functions
|
||||
IRModule (with fused functions marked Primitive=True)
|
||||
│
|
||||
▼ FuseTIR ← merge TIR PrimFuncs inside each group
|
||||
IRModule (fused TIR kernels)
|
||||
|
||||
In the compilation pipeline, these passes appear in the backend-specific ``legalize_passes``
|
||||
phase. For example, the CUDA pipeline (``python/tvm/relax/backend/cuda/pipeline.py``) runs:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
LegalizeOps() # lower Relax ops to call_tir
|
||||
AnnotateTIROpPattern() # annotate pattern kinds
|
||||
FoldConstant()
|
||||
FuseOps() # group ops
|
||||
FuseTIR() # merge TIR functions
|
||||
|
||||
|
||||
Operator Pattern Classification
|
||||
-------------------------------
|
||||
|
||||
Before fusion, ``AnnotateTIROpPattern`` analyzes each TIR function in the module and assigns
|
||||
an ``OpPatternKind``. The fusion algorithm uses these pattern kinds to decide which operators
|
||||
can be fused together.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 20 10 70
|
||||
|
||||
* - Pattern Kind
|
||||
- Value
|
||||
- Description
|
||||
* - ``kElemWise``
|
||||
- 0
|
||||
- Elementwise: one-to-one input/output mapping (e.g., ``add``, ``relu``, ``exp``).
|
||||
* - ``kBroadcast``
|
||||
- 1
|
||||
- Broadcasting: output axes map to input axes in order, but some input axes may be
|
||||
broadcast (e.g., ``bias_add``). Note: ``transpose`` is **not** broadcast because axes
|
||||
are reordered.
|
||||
* - ``kInjective``
|
||||
- 2
|
||||
- Injective: each output element depends on a single input element, but the mapping may
|
||||
be non-trivial (e.g., ``reshape``, ``concatenate``, ``transpose``).
|
||||
* - ``kCommReduce``
|
||||
- 3
|
||||
- Communicative reduction: output elements aggregate over input elements
|
||||
(e.g., ``sum``, ``max``, ``mean``).
|
||||
* - ``kOutEWiseFusable``
|
||||
- 4
|
||||
- Complex operation whose output can accept elementwise followers, but cannot chain
|
||||
with another complex op (e.g., ``conv2d``, ``matmul``, ``dense``).
|
||||
* - ``kTuple``
|
||||
- 7
|
||||
- Tuple node. Can fuse into subsequent injective ops but is treated specially.
|
||||
* - ``kOpaque``
|
||||
- 8
|
||||
- Opaque: cannot be fused (e.g., external function calls, operations with side effects).
|
||||
|
||||
These kinds form an ordering: lower values are "simpler" and more fusable. The fusion algorithm
|
||||
uses ``CombinePattern(lhs, rhs) = max(lhs, rhs)`` when merging patterns along a path.
|
||||
|
||||
|
||||
FuseOps: Automatic Fusion
|
||||
-------------------------
|
||||
|
||||
``FuseOps`` (``src/relax/transform/fuse_ops.cc``) groups bindings in a dataflow block into
|
||||
new Relax functions. It operates only within ``DataflowBlock``\ s — if your module doesn't have
|
||||
any, run ``ConvertToDataflow`` first.
|
||||
|
||||
Algorithm
|
||||
~~~~~~~~~
|
||||
|
||||
The fusion algorithm addresses diamond-shaped dataflow branches, where a single producer
|
||||
(e.g., conv2d) has multiple consumers that eventually reconverge:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
conv2d
|
||||
/ | \
|
||||
/ | \
|
||||
op op op
|
||||
\ | /
|
||||
\ | /
|
||||
elemwise add
|
||||
|
||||
At the point of ``conv2d``, we don't know if all future paths will merge. The algorithm uses
|
||||
**post-dominator analysis** to resolve this:
|
||||
|
||||
1. **Build forward graph**: construct an ``IndexedForwardGraph`` from the dataflow block.
|
||||
Each node has an ``OpPatternKind`` and a list of forward edges.
|
||||
|
||||
2. **Build post-dominator tree**: compute the immediate post-dominator of each node using
|
||||
Least Common Ancestor (LCA) on the DAG. The post-dominator of a node is the closest
|
||||
downstream node where **all** future paths converge.
|
||||
|
||||
3. **Fuse groups**: for each node in topological order, check if it can be fused with its
|
||||
immediate post-dominator:
|
||||
|
||||
- **CheckPath**: verify that all paths from the node to its post-dominator satisfy the
|
||||
fusion conditions (pattern compatibility, depth limits, argument limits).
|
||||
- **CommitFuse**: mark all intermediate nodes as belonging to the same group using a
|
||||
Union-Find data structure.
|
||||
|
||||
4. **Create grouped functions**: extract each group into a new ``relax.Function`` with the
|
||||
attribute ``Primitive=True``. Replace the original bindings with a call to the grouped
|
||||
function.
|
||||
|
||||
Fusion rules
|
||||
~~~~~~~~~~~~
|
||||
|
||||
The key fusion decisions depend on the ``OpPatternKind`` of the source, the path, and the
|
||||
post-dominator. The algorithm runs in three phases (via ``GraphPartitioner::RunFuse``) so that
|
||||
higher-complexity ops get a chance to fuse first:
|
||||
|
||||
- **Phase 0**: ``kOutEWiseFusable`` ops (e.g., ``conv2d``) can fuse with their elementwise
|
||||
post-dominator if all intermediate ops are broadcast or simpler. This enables patterns like
|
||||
conv2d + bias_add + relu. Two ``kOutEWiseFusable`` ops cannot fuse together.
|
||||
- **Phase 1**: ``kInjective`` and ``kTuple`` ops can fuse only when all paths to the
|
||||
post-dominator are injective or simpler. This is deferred to phase 1 so that
|
||||
``kOutEWiseFusable`` groups are finalized first.
|
||||
- **Phase 2**: fuse injective ops into intermediate tuple nodes that have already been absorbed
|
||||
by subsequent injective groups.
|
||||
|
||||
``kElemWise`` / ``kBroadcast`` ops are processed in **every** phase (not restricted to one):
|
||||
they can fuse into a post-dominator that is injective or reduction. The sink (final node) may
|
||||
also be a ``kOutEWiseFusable`` group that was formed in phase 0 — this is how elementwise
|
||||
producers merge into an existing conv2d fusion group.
|
||||
|
||||
Additional constraints:
|
||||
|
||||
- **Reduction** (``kCommReduce``) ops never initiate fusion — they act as sinks only. Elementwise
|
||||
and broadcast producers can fuse *into* a reduction, but a reduction cannot fuse forward.
|
||||
- **Opaque** ops are fusion barriers.
|
||||
- A group cannot exceed ``kMaxFusedOps`` (256) nodes or the maximum function argument count.
|
||||
|
||||
Example
|
||||
~~~~~~~
|
||||
|
||||
Given two elementwise ops (``add``, ``exp``) and one injective op (``squeeze``).
|
||||
The examples below are simplified pseudocode — real TVMScript would reference TIR functions
|
||||
via ``cls.func_name``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Before FuseOps (simplified)
|
||||
@R.function
|
||||
def main(x: R.Tensor((10, 20), "float32")):
|
||||
with R.dataflow():
|
||||
lv0 = R.call_tir(add, (x, const_1), out_ty=R.Tensor((10, 20), "float32"))
|
||||
lv1 = R.call_tir(exp, (lv0,), out_ty=R.Tensor((10, 20), "float32"))
|
||||
gv = R.call_tir(squeeze, (lv1,), out_ty=R.Tensor((10, 20), "float32"))
|
||||
R.output(gv)
|
||||
return gv
|
||||
|
||||
After ``FuseOps``, all three are grouped into a single function:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# After FuseOps
|
||||
@R.function(private=True)
|
||||
def fused_add_exp_squeeze(x, p0):
|
||||
R.func_attr({"Primitive": True})
|
||||
with R.dataflow():
|
||||
lv0 = R.call_tir(add, (x, p0), ...)
|
||||
lv1 = R.call_tir(exp, (lv0,), ...)
|
||||
gv = R.call_tir(squeeze, (lv1,), ...)
|
||||
R.output(gv)
|
||||
return gv
|
||||
|
||||
@R.function
|
||||
def main(x: R.Tensor((10, 20), "float32")):
|
||||
with R.dataflow():
|
||||
gv = fused_add_exp_squeeze(x, const_1)
|
||||
R.output(gv)
|
||||
return gv
|
||||
|
||||
|
||||
FuseTIR: Merging TIR Functions
|
||||
------------------------------
|
||||
|
||||
``FuseTIR`` (``src/relax/transform/fuse_tir.cc``) takes the grouped Relax functions produced by
|
||||
``FuseOps`` and merges their internal TIR ``PrimFunc``\ s into a single TIR function.
|
||||
|
||||
Before ``FuseTIR``, a fused group still contains multiple ``R.call_tir`` calls to separate
|
||||
TIR functions. ``FuseTIR`` inlines and merges them:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
Before FuseTIR:
|
||||
fused_add_exp_squeeze:
|
||||
call_tir(add, ...) → separate TIR PrimFunc
|
||||
call_tir(exp, ...) → separate TIR PrimFunc
|
||||
call_tir(squeeze, ...) → separate TIR PrimFunc
|
||||
|
||||
After FuseTIR:
|
||||
fused_add_exp_squeeze: → single merged TIR PrimFunc
|
||||
|
||||
The merged function eliminates intermediate buffers — the output of ``add`` is directly consumed
|
||||
by ``exp`` without writing to and reading from global memory. This is the core performance benefit
|
||||
of fusion.
|
||||
|
||||
Internally, ``FuseTIR`` uses a ``SymbolicMatcher`` to align symbolic shape variables across the
|
||||
TIR functions being merged, ensuring that dimensions are correctly mapped when combining buffer
|
||||
accesses.
|
||||
|
||||
|
||||
FuseOpsByPattern: Pattern-Based Fusion
|
||||
--------------------------------------
|
||||
|
||||
While ``FuseOps`` makes fusion decisions automatically based on operator patterns,
|
||||
``FuseOpsByPattern`` lets you specify exactly which operator combinations to fuse using
|
||||
the Relax :ref:`Dataflow Pattern Language (DPL) <relax-dpl>`.
|
||||
|
||||
This is primarily used for **backend-specific dispatch**: identifying operator subgraphs that
|
||||
should be offloaded to external libraries like cuBLAS, CUTLASS, cuDNN, or DNNL.
|
||||
|
||||
FusionPattern
|
||||
~~~~~~~~~~~~~
|
||||
|
||||
A ``FusionPattern`` (``python/tvm/relax/transform/transform.py``) defines what to match:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.relax.dpl import wildcard, is_op
|
||||
from tvm.relax.transform import FusionPattern
|
||||
|
||||
# Match: matmul(x, w) + bias
|
||||
x = wildcard()
|
||||
w = wildcard()
|
||||
bias = wildcard()
|
||||
matmul = is_op("relax.matmul")(x, w)
|
||||
out = is_op("relax.add")(matmul, bias)
|
||||
|
||||
pattern = FusionPattern(
|
||||
name="cutlass.matmul_bias",
|
||||
pattern=out,
|
||||
annotation_patterns={"matmul": matmul, "bias": bias},
|
||||
check=my_check_function, # optional validation
|
||||
)
|
||||
|
||||
Fields:
|
||||
|
||||
- ``name``: pattern identifier, typically prefixed with the backend name (e.g.,
|
||||
``"cutlass.matmul_bias"``).
|
||||
- ``pattern``: a DFPattern describing the subgraph to match. See the
|
||||
:ref:`DPL deep dive <relax-dpl>` for the full pattern language.
|
||||
- ``annotation_patterns``: a mapping of names to sub-patterns within the main pattern. These
|
||||
are extracted during matching and made available to the ``check`` function and
|
||||
``attrs_getter``.
|
||||
- ``check``: an optional ``Callable[[PatternCheckContext], bool]`` that validates whether
|
||||
a match should be accepted. Receives the matched expression, annotated sub-expressions,
|
||||
variable usages, and binding information.
|
||||
- ``attrs_getter``: an optional function that extracts attributes (e.g., transpose flags,
|
||||
data types) from the matched expressions to annotate the grouped function.
|
||||
|
||||
Applying patterns
|
||||
~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.relax.transform import FuseOpsByPattern
|
||||
|
||||
mod = FuseOpsByPattern(
|
||||
patterns=[pattern1, pattern2, ...], # ordered by priority
|
||||
bind_constants=True,
|
||||
annotate_codegen=False,
|
||||
)(mod)
|
||||
|
||||
Key parameters:
|
||||
|
||||
- ``patterns``: a list of ``FusionPattern`` objects, ordered by priority. Higher-priority
|
||||
patterns come first — if a subgraph matches multiple patterns, the first match wins.
|
||||
- ``bind_constants``: if ``True``, constants used by the matched subgraph are captured inside
|
||||
the grouped function.
|
||||
- ``annotate_codegen``: if ``True``, wraps each composite function with an outer function
|
||||
annotated with ``"Codegen"`` and ``"global_symbol"`` attributes for external backend dispatch.
|
||||
The ``"Codegen"`` value is derived from the pattern name prefix (e.g., ``"dnnl"`` from
|
||||
``"dnnl.conv2d_relu"``).
|
||||
|
||||
PatternCheckContext
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
The ``check`` function receives a ``PatternCheckContext`` with:
|
||||
|
||||
- ``matched_expr``: the root expression matched by the pattern.
|
||||
- ``annotated_expr``: a mapping from annotation pattern names to their matched expressions.
|
||||
- ``matched_bindings``: variable-to-value bindings within the matched subgraph.
|
||||
- ``var_usages``: a mapping from variable definitions to all their uses in the function.
|
||||
- ``value_to_bound_var``: reverse mapping from values to the variables they are bound to.
|
||||
|
||||
This context enables sophisticated validation logic, such as checking that an intermediate
|
||||
result is not used outside the fused group, or verifying data type compatibility.
|
||||
|
||||
|
||||
How Backends Use Fusion
|
||||
-----------------------
|
||||
|
||||
The default backend pipelines (CUDA, ROCm, CPU, etc.) all include ``FuseOps`` + ``FuseTIR``
|
||||
in their ``legalize_passes`` phase for automatic fusion, as shown in the `Overview`_ above.
|
||||
|
||||
For external library dispatch (cuBLAS, CUTLASS, cuDNN, DNNL), ``FuseOpsByPattern`` is used
|
||||
separately. These are **not** included in the default pipeline — users add them explicitly
|
||||
when building a custom compilation flow. The typical sequence is:
|
||||
|
||||
1. **Pattern-based dispatch** (``FuseOpsByPattern``): identify subgraphs that should be
|
||||
offloaded to external libraries. For example, CUTLASS patterns match
|
||||
matmul+bias+activation combinations (``python/tvm/relax/backend/cuda/cutlass.py``).
|
||||
Functions marked by patterns are annotated with ``Composite`` and optionally ``Codegen``
|
||||
attributes. See :ref:`external-library-dispatch` for the full BYOC pipeline.
|
||||
|
||||
2. **Automatic fusion** (``FuseOps`` + ``FuseTIR``): remaining operators that were not
|
||||
matched by backend patterns are fused automatically based on their pattern kinds.
|
||||
|
||||
|
||||
Source Code Map
|
||||
---------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 50 50
|
||||
|
||||
* - Path
|
||||
- Contents
|
||||
* - ``src/relax/transform/fuse_ops.cc``
|
||||
- FuseOps and FuseOpsByPattern implementation
|
||||
* - ``src/relax/analysis/graph_partitioner.h``
|
||||
- IndexedForwardGraph, DominatorTree, GraphPartitioner (Union-Find)
|
||||
* - ``src/relax/transform/fuse_tir.cc``
|
||||
- FuseTIR implementation, SymbolicMatcher
|
||||
* - ``include/tvm/relax/op_attr_types.h``
|
||||
- ``OpPatternKind`` enum definition
|
||||
* - ``python/tvm/relax/transform/transform.py``
|
||||
- Python API: FuseOps, FuseTIR, FuseOpsByPattern, FusionPattern
|
||||
* - ``python/tvm/relax/dpl/``
|
||||
- Dataflow Pattern Language (DFPattern, is_op, wildcard, etc.)
|
||||
* - ``python/tvm/relax/backend/cuda/cutlass.py``
|
||||
- Example: CUTLASS fusion patterns
|
||||
* - ``python/tvm/relax/backend/cuda/cublas.py``
|
||||
- Example: cuBLAS fusion patterns
|
||||
@@ -0,0 +1,436 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
Design and Architecture
|
||||
=======================
|
||||
|
||||
This document is intended for developers who want to understand the architecture of Apache TVM and/or actively develop on the project.
|
||||
This page is organized as follows:
|
||||
|
||||
- The `Overall Flow`_ gives an overview of the steps that TVM takes to turn a high level description of a model into a deployable module.
|
||||
To get started, please read this section first.
|
||||
- Brief introduction to the key components of the TVM stack. Feel free to also check out the :ref:`TensorIR Deep Dive <tensor-ir-deep-dive>`
|
||||
and :ref:`Relax Deep Dive <relax-deep-dive>` for more details about the two major components in the TVM stack.
|
||||
|
||||
This guide provides a few complementary views of the architecture.
|
||||
First, we review a single end-to-end compilation flow and discuss the key data structures and the transformations.
|
||||
This runtime-based view focuses on the interactions of each components when running the compiler.
|
||||
Then we will review the logical modules of the codebase and their relationship. This part provides a static overarching view of the design.
|
||||
|
||||
Overall Flow
|
||||
------------
|
||||
|
||||
In this guide, we will study an example compilation flow in the compiler. The figure below shows the flow. At a high-level, it contains several steps:
|
||||
|
||||
- **Model Creation**: Create the IRModule to be optimized and compiled, which contains a collection of functions that internally represent the model.
|
||||
Users can manually construct IRModule via NNModule, TVMScript, or import a pre-trained model from Relax frontend.
|
||||
- **Transformation**: The compiler transforms an IRModule to another functionally equivalent or approximately
|
||||
equivalent(e.g. in the case of quantization) IRModule. Many of the transformations are target (backend) independent.
|
||||
We also allow target to affect the configuration of the transformation pipeline.
|
||||
- **Target Translation**: The compiler translates(codegen) the IRModule to an executable format specified by the target.
|
||||
The target translation result is encapsulated as a `runtime.Module` that can be exported, loaded, and executed on the target runtime environment.
|
||||
- **Runtime Execution**: the user loads back a `runtime.Module` and runs the compiled functions in the supported runtime environment.
|
||||
|
||||
|
||||
.. figure:: https://raw.githubusercontent.com/tlc-pack/web-data/main/images/design/tvm_overall_flow.svg
|
||||
:align: center
|
||||
:width: 80%
|
||||
|
||||
|
||||
Key data structures
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
One of the best ways to design and understand a complex system is to identify the key data structures and APIs that
|
||||
manipulate (transform) these data structures. Once we identified the key data structures, we can then breakdown a system into logical
|
||||
components that either define a collection of key data structures or transformations among the data structures.
|
||||
|
||||
**IRModule** is the primary data structure used across the entire stack. An IRModule (intermediate representation module)
|
||||
contains a collection of functions. Currently, we support two primary variants of functions.
|
||||
|
||||
- **relax::Function** is a high-level functional program representation. A relax.Function represents high-level graph structure,
|
||||
usually corresponds to an end-to-end model or a sub-graph of the overall model. You can view a relax.Function as a computational
|
||||
graph with additional support for control-flow, and complex data structures.
|
||||
- **tirx::PrimFunc** is a low-level program representation that contains elements including loop-nest choices, multi-dimensional load/store,
|
||||
threading, and vector/tensor instructions. It is usually used to represent an operator program that executes a (possibly-fused) layer in a model.
|
||||
|
||||
During the compilation and transformation, all relax operators are lowered to ``tirx::PrimFunc`` or ``TVM PackedFunc``, which can be executed directly
|
||||
on the target device, while the calls to relax operators are lowered to calls to low-level functions (e.g. ``R.call_tir`` or ``R.call_dps_packed``).
|
||||
|
||||
Transformations
|
||||
~~~~~~~~~~~~~~~
|
||||
|
||||
Now that we have covered the key data structures, let us talk about the transformations. Each transformation could serve one of the following purposes:
|
||||
|
||||
- optimization: transform a program to an equivalent, possibly more optimized version.
|
||||
- lowering: transform a program to a lower-level representation that is closer to the target.
|
||||
|
||||
relax transformations
|
||||
^^^^^^^^^^^^^^^^^^^^^
|
||||
relax transformations contain a collection of passes that apply to relax functions. The optimizations include common graph-level
|
||||
optimizations such as constant folding and dead-code elimination for operators, and backend-specific optimizations such as library dispatch.
|
||||
|
||||
TensorIR transformations
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
- **TensorIR schedule**: TensorIR schedules are designed to optimize the TensorIR functions for a specific target, with user-guided instructions and control how the target code is generated.
|
||||
For CPU targets, a TensorIR PrimFunc can generate valid code and execute on the target device without schedule but with very-low performance. However, for GPU targets, the schedule is essential
|
||||
for generating valid code with thread bindings. For more details, please refer to the :ref:`TensorIR Transformation <tirx-transform>` section. Additionally, we provides ``MetaSchedule`` to
|
||||
automate the search of TensorIR schedule.
|
||||
- **Lowering Passes**: These passes usually perform after the schedule is applied, transforming a TensorIR PrimFunc into another functionally equivalent PrimFunc, but closer to the
|
||||
target-specific representation. For example, there are passes to flatten multi-dimensional access to one-dimensional pointer access, to expand the intrinsics into target-specific ones,
|
||||
and to decorate the function entry to meet the runtime calling convention.
|
||||
|
||||
Many low-level optimizations can be handled in the target phase by the LLVM,
|
||||
CUDA C, and other target compilers. As a result, we leave low-level
|
||||
optimizations such as register allocation to the downstream compilers and only
|
||||
focus on optimizations that are not covered by them.
|
||||
|
||||
cross-level transformations
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
Apache TVM enables cross-level optimization of end-to-end models. As the IRModule includes both Relax and TensorIR functions, the cross-level transformations are designed to mutate
|
||||
the IRModule by applying different transformations to these two types of functions.
|
||||
|
||||
For example, ``relax.LegalizeOps`` pass mutates the IRModule by lowering relax operators, adding corresponding TensorIR PrimFunc into the IRModule, and replacing the relax operators
|
||||
with calls to the lowered TensorIR PrimFunc. Another example is the operator fusion pipeline
|
||||
(``relax.FuseOps`` + ``relax.FuseTIR``), which fuses multiple consecutive tensor operations into a
|
||||
single kernel. See :ref:`fusion-arch` for a detailed explanation of the fusion algorithm, operator
|
||||
pattern classification, and pattern-based fusion for external backends.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
fusion
|
||||
|
||||
Target Translation
|
||||
~~~~~~~~~~~~~~~~~~
|
||||
|
||||
The target translation phase transforms an IRModule to the corresponding target executable format.
|
||||
For backends such as x86 and ARM, we use the LLVM IRBuilder to build in-memory LLVM IR.
|
||||
We can also generate source-level languages such as CUDA C and OpenCL.
|
||||
Finally, we support direct translations of a Relax function (sub-graph) to specific targets via external code generators.
|
||||
See :ref:`codegen-arch` for how TIR functions are compiled to native code through the LLVM and
|
||||
Source codegen families.
|
||||
See :ref:`external-library-dispatch` for the full BYOC (Bring Your Own Codegen) pipeline that
|
||||
offloads operator subgraphs to vendor libraries like cuBLAS, CUTLASS, and cuDNN.
|
||||
It is important that the final code generation phase is as lightweight as possible. Vast majority of transformations
|
||||
and lowering should be performed before the target translation phase.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
codegen
|
||||
external_library_dispatch
|
||||
|
||||
We also provide a Target structure to specify the compilation target.
|
||||
The transformations before the target translation phase can also be affected by the target — for example,
|
||||
a target's vector length would change the vectorization behavior.
|
||||
|
||||
|
||||
Runtime Execution
|
||||
~~~~~~~~~~~~~~~~~
|
||||
|
||||
The main goal of TVM's runtime is to provide a minimal API for loading and executing the compiled artifact in a language of their choice, including Python, C++, Rust, Go, Java, and JavaScript. The code snippet below shows such an example in Python:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import tvm
|
||||
# Example runtime execution program in python, with type annotated
|
||||
mod: tvm.runtime.Module = tvm.runtime.load_module("compiled_artifact.so")
|
||||
arr: tvm.runtime.Tensor = tvm.runtime.tensor([1, 2, 3], device=tvm.cuda(0))
|
||||
fun: tvm_ffi.Function = mod["addone"]
|
||||
fun(arr)
|
||||
print(arr.numpy())
|
||||
|
||||
|
||||
:py:class:`tvm.runtime.Module` encapsulates the result of compilation. A runtime.Module contains a GetFunction method to obtain :py:class:`tvm_ffi.Function` instances by name.
|
||||
|
||||
:py:class:`tvm_ffi.Function` is a type-erased function interface for both the generated functions. A tvm_ffi.Function can take arguments and return values with the
|
||||
following types: POD types(int, float), string, tvm_ffi.Function, runtime.Module, runtime.Tensor, and other sub-classes of runtime.Object.
|
||||
|
||||
:py:class:`tvm.runtime.Module` and :py:class:`tvm_ffi.Function` are powerful mechanisms to modularize the runtime. For example, to get the above `addone` function on CUDA, we can use LLVM to generate the host-side code to compute the launching parameters(e.g. size of the thread groups) and then call into another tvm_ffi.Function from a CUDAModule that is backed by the CUDA driver API. The same mechanism can be used for OpenCL kernels.
|
||||
|
||||
The above example only deals with a simple `addone` function. The code snippet below gives an example of an end-to-end model execution using the Relax Virtual Machine, which is built on the same runtime.Module and tvm_ffi.Function interface:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import tvm
|
||||
from tvm import relax
|
||||
# Load the compiled artifact
|
||||
mod: tvm.runtime.Module = tvm.runtime.load_module("resnet18.so")
|
||||
# Create a VM instance on cuda(0)
|
||||
vm = relax.VirtualMachine(mod, tvm.cuda(0))
|
||||
data: tvm.runtime.Tensor = get_input_data()
|
||||
# Run the model — vm["main"] returns a PackedFunc
|
||||
result = vm["main"](data).numpy()
|
||||
|
||||
The main take away is that runtime.Module and runtime.PackedFunc are sufficient to encapsulate both operator level programs (such as addone), as well as the end-to-end models.
|
||||
|
||||
Summary and Discussions
|
||||
~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
In summary, the key data structures in the compilation flows are:
|
||||
|
||||
- IRModule: contains relax.Function and tirx.PrimFunc
|
||||
- runtime.Module: contains runtime.PackedFunc
|
||||
|
||||
Most parts of the compilation are transformations among the key data structures.
|
||||
|
||||
- relax/transform and tirx/transform are deterministic rule-based transformations
|
||||
- meta-schedule contains the search-based transformations
|
||||
|
||||
Finally, the compilation flow example is only a typical use-case of the TVM stack.
|
||||
We expose these key data structures and transformations to python and C++ APIs. As a result, you can use TVM just like the way you use numpy,
|
||||
except that the data structure of interest changes from the numpy.ndarray to tvm.IRModule. Here are some example use-cases:
|
||||
|
||||
- Directly construct IRModule using the python API.
|
||||
- Compose a custom set of transformations(e.g. customize quantization).
|
||||
- Manipulate the IR directly using TVM's python API.
|
||||
|
||||
|
||||
tvm/support
|
||||
-----------
|
||||
The support module contains the most common utilities for the infrastructure, such as generic arena allocator, socket, and logging.
|
||||
|
||||
|
||||
tvm/runtime
|
||||
-----------
|
||||
|
||||
The runtime serves as the foundation of the TVM stack. It provides the mechanism to load and execute compiled artifacts.
|
||||
The runtime defines a stable standard set of C APIs to interface with frontend languages such as Python and Rust.
|
||||
|
||||
`runtime::Object` is one of the primary data structures in TVM runtime besides the `ffi::Function`.
|
||||
It is a reference-counted base class with a type index to support runtime type checking and downcasting.
|
||||
The object system allows the developer to introduce new data structures to the runtime, such as Array, Map, and new IR data structures.
|
||||
|
||||
Besides deployment use-cases, the compiler itself also makes heavy use of TVM's runtime mechanism.
|
||||
All of the IR data structures are subclasses of `runtime::Object`, as a result, they can be directly accessed and manipulated from the Python frontend.
|
||||
We use the PackedFunc mechanism to expose various APIs to the frontend.
|
||||
|
||||
Runtime support for different hardware backends are defined in subdirectories of runtime(e.g. runtime/opencl).
|
||||
These hardware-specific runtime modules define APIs for device memory allocation and device function serialization.
|
||||
|
||||
`runtime/rpc` implements an RPC support for PackedFunc. We can use the RPC mechanism to send a cross-compiled library to a remote
|
||||
device and benchmark the execution performance. The rpc infrastructure enables data collection from a wide range of hardware backends
|
||||
for learning-based optimizations.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
runtime
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
introduction_to_module_serialization
|
||||
|
||||
Relax Virtual Machine
|
||||
~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Relax defines *what* to compute — it is a graph-level IR that describes the operators and dataflow
|
||||
of a model. The Relax Virtual Machine (VM) handles *how* to run it — it is the runtime component
|
||||
that executes the compiled result. The VM uses a register-based interpreter with only four opcodes
|
||||
(``Call``, ``Ret``, ``Goto``, ``If``) and performs no mathematical computation itself — it
|
||||
orchestrates control flow while dispatching actual work to compiled TIR kernels or external
|
||||
libraries.
|
||||
|
||||
See :ref:`relax-vm-arch` for the full architecture documentation, including the compilation
|
||||
pipeline, instruction set details, execution model, and Python interface.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
relax_vm
|
||||
|
||||
Disco: Distributed Runtime
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Disco is TVM's distributed runtime for executing models across multiple devices. When a model is
|
||||
too large to fit on a single GPU, the ``relax.distributed`` module annotates how tensors should be
|
||||
partitioned and placed across a mesh of devices at compile time. Disco then takes over at runtime:
|
||||
it manages a group of workers, dispatches the compiled program to all of them simultaneously, and
|
||||
coordinates inter-device communication through collective operations such as allreduce, allgather,
|
||||
broadcast, and scatter.
|
||||
|
||||
The central abstraction is the ``Session``, which owns the workers and exposes a SPMD-style
|
||||
programming interface. Every object that lives on workers is represented by a ``DRef`` — a
|
||||
distributed reference that maps to a concrete value on each worker. When the controller invokes a
|
||||
``DPackedFunc`` through the session, all workers execute the same PackedFunc call synchronously, each
|
||||
operating on its own local shard. Compiled VM modules can be loaded into a session as ``DModule``
|
||||
objects and called in the same fashion. The session also provides collective primitives backed by
|
||||
NCCL or RCCL, so that workers can exchange partial results without routing data through the
|
||||
controller.
|
||||
|
||||
Three session backends cover different deployment topologies. ``ThreadedSession`` spawns workers as
|
||||
threads within a single process — this is the most common choice for multi-GPU inference on a
|
||||
single machine. ``ProcessSession`` launches workers as separate OS processes connected by pipes,
|
||||
providing stronger isolation. ``SocketSession`` extends the model to multi-node clusters by
|
||||
connecting workers across machines via TCP sockets.
|
||||
|
||||
tvm/node
|
||||
--------
|
||||
The node module adds additional features on top of the `runtime::Object` for IR data structures.
|
||||
The main features include reflection, serialization, structural equivalence, and hashing.
|
||||
|
||||
Thanks to the node module, we can directly access any field of the TVM's IRNode by their name in Python.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
x = tvm.tirx.Var("x", "int32")
|
||||
y = tvm.tirx.Add(x, x)
|
||||
# a and b are fields of a tirx.Add node
|
||||
# we can directly use the field name to access the IR structures
|
||||
assert y.a == x
|
||||
|
||||
We can also serialize arbitrary IR node into a JSON format, and load them back.
|
||||
The ability to save/store, and inspect an IR node provides a foundation for making the compiler more accessible.
|
||||
|
||||
tvm/ir
|
||||
------
|
||||
The `tvm/ir` folder contains the unified data structure and interfaces across all IR function variants.
|
||||
The components in `tvm/ir` are shared by `tvm/relax` and `tvm/tirx`, notable ones include
|
||||
|
||||
- IRModule
|
||||
- Type
|
||||
- PassContext and Pass
|
||||
- Op
|
||||
|
||||
Different variants of functions(e.g. relax.Function and tirx.PrimFunc) can co-exist in an IRModule.
|
||||
While these variants may not have the same content representation, they use the same data structure to represent types.
|
||||
As a consequence, we use the same data structure to represent function (type) signatures of these variants.
|
||||
The unified type system allows one function variant to call another function
|
||||
once we clearly define the calling convention. This opens doors for future cross-function-variant optimizations.
|
||||
|
||||
We also provide a unified PassContext for configuring the pass behavior, and common composite passes to execute a pass pipeline.
|
||||
The following code snippet gives an example of PassContext configuration.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# configure the behavior of the tirx.UnrollLoop pass
|
||||
with tvm.transform.PassContext(config={"tirx.UnrollLoop": { "auto_max_step": 10 }}):
|
||||
# code affected by the pass context
|
||||
|
||||
|
||||
Op is the common class to represent all system-defined primitive operator/intrinsics.
|
||||
Developers can register new Ops as well as their additional attributes(e.g. whether the Op is elementwise) to the system.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
pass_infra
|
||||
|
||||
tvm/script (TVMScript)
|
||||
----------------------
|
||||
|
||||
TVMScript is a Python-based DSL for writing TVM IR. It allows users to define ``IRModule``\ s
|
||||
— containing both Relax functions and TIR ``PrimFunc``\ s — using familiar Python syntax with
|
||||
three import aliases: ``I`` (module-level), ``T`` (TIR), and ``R`` (Relax). Although TVMScript
|
||||
uses Python syntax, it is not executed by the Python interpreter — decorators like
|
||||
``@I.ir_module``, ``@T.prim_func``, and ``@R.function`` extract the Python AST and transform
|
||||
it into TVM IR through a parser and IR builder pipeline.
|
||||
|
||||
TVMScript also supports **roundtrip**: any ``IRModule`` can be printed back to TVMScript via
|
||||
``mod.script()`` and re-parsed to produce a structurally equivalent module. See
|
||||
:ref:`tvmscript-arch` for the full architecture documentation, including the parser dispatch
|
||||
mechanism, IR builder frame stack, printer pipeline, and syntax reference.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
tvmscript
|
||||
|
||||
|
||||
tvm/target
|
||||
----------
|
||||
The target module contains all the code generators that translate an IRModule to a target runtime.Module.
|
||||
It also provides a common `Target` class that describes the target.
|
||||
|
||||
Targets can be constructed from a registered tag, a configuration dictionary, or a tag with attribute overrides:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.target import Target
|
||||
|
||||
# From a registered tag
|
||||
target = Target("nvidia/nvidia-a100")
|
||||
|
||||
# From a config dictionary
|
||||
target = Target({"kind": "cuda", "arch": "sm_80"})
|
||||
|
||||
# From a tag with attribute overrides
|
||||
target = Target({"tag": "nvidia/nvidia-a100", "l2_cache_size_bytes": 12345})
|
||||
|
||||
Use ``Target.list_kinds()`` to see all available target kinds, and ``target.attrs`` to inspect
|
||||
target attributes.
|
||||
|
||||
The compilation pipeline can be customized according to the target by querying the attribute information
|
||||
in the target and builtin information registered to each target id(cuda, opencl).
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
device_target_interactions
|
||||
|
||||
tvm/relax
|
||||
---------
|
||||
|
||||
Relax is the high-level IR used to represent the computational graph of a model. Various optimizations are defined in ``relax.transform``.
|
||||
Note that Relax usually works closely with the TensorIR IRModule, most of the transformations are applied on both Relax and TensorIR functions
|
||||
in the IRModule. Please refer to the :ref:`Relax Deep Dive <relax-deep-dive>` for more details.
|
||||
|
||||
tvm/tirx
|
||||
--------
|
||||
|
||||
``tirx`` contains the core IR definitions and lowering infrastructure
|
||||
for TensorIR (split from the former ``tir`` module). ``tirx::PrimFunc``
|
||||
represents low-level tensor functions that can be transformed by tirx passes.
|
||||
|
||||
The tirx module includes:
|
||||
|
||||
- IR data structures (PrimFunc, Buffer, SBlock, expressions, statements).
|
||||
- Analysis passes in ``tirx/analysis``.
|
||||
- Transformation and lowering passes in ``tirx/transform``.
|
||||
|
||||
tvm/s_tir
|
||||
---------
|
||||
|
||||
``s_tir`` (Schedulable TIR, split from the former ``tir`` module) contains
|
||||
schedule primitives and auto-tuning tools that operate on ``tirx::PrimFunc``:
|
||||
|
||||
- Schedule primitives to control code generation (tiling, vectorization, thread
|
||||
binding) in ``s_tir/schedule``.
|
||||
- Builtin tensor intrinsics in ``s_tir/tensor_intrin``.
|
||||
- MetaSchedule for automated performance tuning.
|
||||
- DLight for pre-defined, high-performance schedules.
|
||||
|
||||
Please refer to the :ref:`TensorIR Deep Dive <tensor-ir-deep-dive>` for more details.
|
||||
|
||||
tvm/arith
|
||||
---------
|
||||
|
||||
This module is closely tied to TensorIR. One of the key problems in the low-level code generation is the analysis of the indices'
|
||||
arithmetic properties — the positiveness, variable bound, and the integer set that describes the iterator space. arith module provides
|
||||
a collection of tools that do (primarily integer) analysis. A TensorIR pass can use these analyses to simplify and optimize the code.
|
||||
|
||||
tvm/te and tvm/topi
|
||||
-------------------
|
||||
|
||||
TE stands for Tensor Expression. TE is a domain-specific language (DSL) for describing tensor computations. Importantly, a tensor expression
|
||||
itself is not a self-contained function that can be stored into IRModule. We can use ``te.create_prim_func`` to convert a tensor expression to a ``tirx::PrimFunc``
|
||||
and then integrate it into the IRModule.
|
||||
|
||||
While possible to construct operators directly via TensorIR or tensor expressions (TE) for each use case, it is tedious to do so.
|
||||
`topi` (Tensor operator inventory) provides a set of pre-defined operators defined by numpy and found in common deep learning workloads.
|
||||
@@ -0,0 +1,193 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
Introduction to Module Serialization
|
||||
====================================
|
||||
|
||||
When to deploy TVM runtime module, no matter whether it is CPU or GPU, TVM only needs one single dynamic
|
||||
shared library. The key is our unified module serialization mechanism. This document will introduce TVM module
|
||||
serialization format standard and implementation details.
|
||||
|
||||
|
||||
*************
|
||||
Serialization
|
||||
*************
|
||||
|
||||
The entrance API is ``export_library`` of ``tvm.runtime.Module``.
|
||||
Inside this function, we will do the following steps:
|
||||
|
||||
1. Collect all DSO modules (LLVM modules and C modules)
|
||||
|
||||
2. Once we have DSO modules, we will call ``save`` function to save them into files.
|
||||
|
||||
3. Next, we will check whether we have imported modules, such as CUDA,
|
||||
OpenCL or anything else. We don't restrict the module type here.
|
||||
Once we have imported modules, we will create one file named ``devc.o`` / ``dev.cc``
|
||||
(so that we could embed the binary blob data of import modules into one dynamic shared library),
|
||||
then call function ``_PackImportsToLLVM`` or ``_PackImportsToC`` to do module serialization.
|
||||
|
||||
4. Finally, we call ``fcompile`` which invokes ``_cc.create_shared`` to get
|
||||
dynamic shared library.
|
||||
|
||||
.. note::
|
||||
1. For C source modules, we will compile them and link them together with the DSO module.
|
||||
|
||||
2. Use ``_PackImportsToLLVM`` or ``_PackImportsToC`` depends on whether we enable LLVM in TVM.
|
||||
They achieve the same goal in fact.
|
||||
|
||||
***************************************************
|
||||
Under the Hood of Serialization and Format Standard
|
||||
***************************************************
|
||||
|
||||
As said before, we will do the serialization work in the ``_PackImportsToLLVM`` or ``_PackImportsToC``.
|
||||
They both call ``SerializeModule`` to serialize the runtime module. In ``SerializeModule``
|
||||
function, we firstly construct one helper class ``ModuleSerializer``. It will take ``module`` to do some
|
||||
initialization work, like marking module index. Then we could use its ``SerializeModule`` to serialize module.
|
||||
|
||||
For better understanding, let us dig the implementation of this class a little deeper.
|
||||
|
||||
The following code is used to construct ``ModuleSerializer``:
|
||||
|
||||
.. code:: c++
|
||||
|
||||
explicit ModuleSerializer(runtime::Module mod) : mod_(mod) {
|
||||
Init();
|
||||
}
|
||||
private:
|
||||
void Init() {
|
||||
CreateModuleIndex();
|
||||
CreateImportTree();
|
||||
}
|
||||
|
||||
In ``CreateModuleIndex()``, We will inspect module import relationship
|
||||
using DFS and create index for them. Note the root module is fixed at
|
||||
location 0. In our example, we have module relationship like this:
|
||||
|
||||
.. code:: c++
|
||||
|
||||
llvm_mod:imports
|
||||
- cuda_mod
|
||||
|
||||
So LLVM module will have index 0, CUDA module will have index 1.
|
||||
|
||||
After constructing module index, we will try to construct import tree (``CreateImportTree()``),
|
||||
which will be used to restore module import relationship when we load
|
||||
the exported library back. In our design, we use CSR format to store
|
||||
import tree, each row is parent index, the child indices correspond to its children
|
||||
index. In code, we use ``import_tree_row_ptr_`` and
|
||||
``import_tree_child_indices_`` to represent them.
|
||||
|
||||
After initialization, we could serialize module using ``SerializeModule`` function.
|
||||
In its function logic, we will assume the serialization format like this:
|
||||
|
||||
.. code:: c++
|
||||
|
||||
binary_blob_size
|
||||
binary_blob_type_key
|
||||
binary_blob_logic
|
||||
binary_blob_type_key
|
||||
binary_blob_logic
|
||||
...
|
||||
_import_tree
|
||||
_import_tree_logic
|
||||
|
||||
``binary_blob_size`` is the number of blobs we will have in this
|
||||
serialization step. There will be three blobs in our example which
|
||||
are created for LLVM module, CUDA module, and ``_import_tree``, respectively.
|
||||
|
||||
``binary_blob_type_key`` is the blob type key of module. For LLVM / C module, whose
|
||||
blob type key is ``_lib``. For CUDA module, it is ``cuda``, which could be got by ``module->type_key()``.
|
||||
|
||||
``binary_blob_logic`` is the logic handling of blob. For most of blob (like CUDA, OpenCL), we will call
|
||||
``SaveToBinary`` function to serialize blob into binary. However, like LLVM / C module, we will only write
|
||||
``_lib`` to indicate this is a DSO module.
|
||||
|
||||
.. note::
|
||||
Whether or not it is required to implement the SaveToBinary virtual function depends on
|
||||
how the module is used. For example, if the module has information we need when we load
|
||||
the dynamic shared library back, we should do. Like CUDA module, we need its binary data
|
||||
passing to GPU driver when we load the dynamic shared library, so we should implement
|
||||
``SaveToBinary`` to serialize its binary data. But for host module (like DSO), we don't
|
||||
need other information when we load the dynamic shared library, so we don't need to implement
|
||||
``SaveToBinary``. However, if in the future, we want to record some meta information of DSO module,
|
||||
we could implement ``SaveToBinary`` for DSO module too.
|
||||
|
||||
Finally, we will write one key ``_import_tree`` unless our module only
|
||||
has one DSO module and it is in the root. It is used to reconstruct the
|
||||
module import relationship when we load the exported library back as said
|
||||
before. The ``import_tree_logic`` is just to write ``import_tree_row_ptr_``
|
||||
and ``import_tree_child_indices_`` into stream.
|
||||
|
||||
After this step, we will pack it into a symbol
|
||||
``runtime::symbol::tvm_ffi_library_bin`` that can be recovered in the dynamic
|
||||
library.
|
||||
|
||||
Now, we complete the serialization part. As you have seen, we could
|
||||
support arbitrary modules to import ideally.
|
||||
|
||||
****************
|
||||
Deserialization
|
||||
****************
|
||||
|
||||
The entrance API is ``tvm.runtime.load_module``. This function
|
||||
actually calls ``_LoadFromFile``. If we dig it a little deeper, this is
|
||||
``Module::LoadFromFile``. In our example, the file is ``deploy.so``,
|
||||
according to the function logic, we will call ``module.loadfile_so`` in
|
||||
``dso_library.cc``. The key is here:
|
||||
|
||||
.. code:: c++
|
||||
|
||||
// Load the imported modules
|
||||
const char* library_bin = reinterpret_cast<const char*>(
|
||||
lib->GetSymbol(runtime::symbol::tvm_ffi_library_bin));
|
||||
Module root_mod;
|
||||
if (library_bin != nullptr) {
|
||||
root_mod = ProcessLibraryBin(library_bin, lib);
|
||||
} else {
|
||||
// Only have one single DSO Module
|
||||
root_mod = Module(n);
|
||||
}
|
||||
|
||||
As said before, we will pack the blob into the symbol
|
||||
``runtime::symbol::tvm_ffi_library_bin``. During deserialization part, we will
|
||||
inspect it. If we have ``runtime::symbol::tvm_ffi_library_bin``, we will call ``ProcessLibraryBin``,
|
||||
whose logic like this:
|
||||
|
||||
.. code:: c++
|
||||
|
||||
READ(blob_size)
|
||||
READ(blob_type_key)
|
||||
for (size_t i = 0; i < blob_size; i++) {
|
||||
if (blob_type_key == "_lib") {
|
||||
// construct dso module using lib
|
||||
} else if (blob_type_key == "_import_tree") {
|
||||
// READ(_import_tree_row_ptr)
|
||||
// READ(_import_tree_child_indices)
|
||||
} else {
|
||||
// call module.loadbinary_blob_type_key, such as module.loadbinary_cuda
|
||||
// to restore.
|
||||
}
|
||||
}
|
||||
// Using _import_tree_row_ptr and _import_tree_child_indices to
|
||||
// restore module import relationship. The first module is the
|
||||
// root module according to our invariance as said before.
|
||||
return root_module;
|
||||
|
||||
After this, we will set the ``ctx_address`` to be the ``root_module`` so
|
||||
that allow lookup of symbol from root (so all symbols are visible).
|
||||
|
||||
Finally, we complete the deserialization part.
|
||||
@@ -0,0 +1,669 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
.. _pass-infra:
|
||||
|
||||
Pass Infrastructure
|
||||
===================
|
||||
|
||||
Both Relax and TVM IR contain a series of optimization passes which improve performance metrics
|
||||
of models such as mean inference, memory footprint, or power consumption for
|
||||
specific devices. There is a suite of standard optimizations as well as machine
|
||||
learning-specific optimizations including constant folding, dead code
|
||||
elimination, operator layout alteration, operator fusion, buffer handling, and
|
||||
loop transformation, etc. Each of these passes is structured as a ir-to-ir
|
||||
transformation using the analysis result collected during and/or before traversal.
|
||||
|
||||
However, as TVM evolves quickly, the need for a more systematic and efficient
|
||||
way to manage these passes is becoming apparent. In addition, a generic
|
||||
framework that manages the passes across different layers of the TVM stack (e.g.
|
||||
Relax and TensorIR) paves the way for developers to quickly prototype and plug the
|
||||
implemented passes into the system.
|
||||
|
||||
This doc describes the design of such an infra that takes the advantage of the
|
||||
way production compilers are used to manage the optimization passes and the style
|
||||
modern deep learning frameworks adopted to build up layers.
|
||||
|
||||
For example, many existing production compilers, such as GCC and LLVM, employ
|
||||
pass managers to effectively manage the execution of passes. Initially managing
|
||||
passes is straightforward as the number of passes is small, but mature compilers
|
||||
will contain hundreds of individual passes. Often external users will want to
|
||||
have custom passes correctly scheduled without having to modify a single
|
||||
handcrafted pass order.
|
||||
|
||||
Similarly, modern deep learning frameworks, such as PyTorch, also have
|
||||
the tendency to enable pass-style layer construction scheme through
|
||||
`Sequential`_. With such constructs, these modern frameworks are able to
|
||||
conveniently add modules/layers to their containers and build up neural
|
||||
networks easily.
|
||||
|
||||
The design of the TVM pass infra is largely inspired by the hierarchical
|
||||
pass manager used in LLVM and the block-style containers used in the popular
|
||||
deep learning frameworks. The major goals of the pass infra include:
|
||||
|
||||
#) enabling better programmatic orchestration of optimizations. This allows
|
||||
users to flexibly customize and build their own optimization pipelines.
|
||||
|
||||
#) providing a user-friendly way to debug optimization passes.
|
||||
|
||||
#) alleviating developers from manually and respectively resolving the
|
||||
dependencies between passes.
|
||||
|
||||
#) simplifying the implementation of new passes for developers. For example, we
|
||||
allow users to implement a pass in Python and let the pass infra manipulate
|
||||
its execution.
|
||||
|
||||
The Design
|
||||
----------
|
||||
|
||||
We focus on ease of extension for users, making it possible for users to quickly
|
||||
add new passes without loss of backward compatibility. The design contains both
|
||||
the backend and the frontend. The former implements the main logic of the pass
|
||||
infra. The latter provides simple APIs for users to interact with, i.e.,
|
||||
allowing users to quickly create their own optimization pipelines.
|
||||
|
||||
C++ Backend
|
||||
~~~~~~~~~~~
|
||||
|
||||
We provide a ``PassInfo`` object to contain the basic information needed by
|
||||
a pass. ``name`` is the pass name, ``opt_level`` indicates at which optimization
|
||||
level the pass will be enabled, and ``required`` represents the passes that are
|
||||
required to execute a certain pass (see `include/tvm/ir/transform.h`_ for
|
||||
more details). For example, during registration of a pass (will be covered in
|
||||
later), the pass developers can specify the name of the pass, the optimization
|
||||
level it will be performed at, and/or the passes that are required.
|
||||
``opt_level`` could be used to help the pass infra identify if a certain pass
|
||||
needs to be executed when running under a user-provided optimization level. The
|
||||
``required`` field can be used by the pass infra to resolve pass dependencies.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
class PassInfoNode : public Object {
|
||||
int opt_level;
|
||||
ffi::String name;
|
||||
bool traceable;
|
||||
ffi::Array<ffi::String> required;
|
||||
};
|
||||
|
||||
PassContext
|
||||
^^^^^^^^^^^
|
||||
|
||||
``PassContext`` carries useful information for an optimization pass. For
|
||||
example, it contains the error reporting system so optimization authors can
|
||||
provide diagnostics about why an optimization fails. ``PassContext`` is also
|
||||
designed to replace the old ``BuildConfig`` which was used to help users
|
||||
configure the compilation options, including optimization level and
|
||||
required/disabled passes, etc. For instance, we may have a configuration which
|
||||
performs all passes at ``opt_level=3`` with some disabled passes using
|
||||
``disabled_pass=xx`` provided by ``PassContext``. Now we could glob all passes
|
||||
at ``opt_level=3`` and exclude those in the disabled pass list. ``PassContext``
|
||||
also provides a way to instrument all passes. See section :ref:`pass_instrument_cpp_backend`.
|
||||
|
||||
This class is designed for users to conveniently write the Python ``with``
|
||||
syntax to perform optimizations under a certain configuration. In addition, the
|
||||
users can obtain the context that is available within a certain program scope in
|
||||
a thread-safe way through ``PassContext::Current()``, since a thread-local store
|
||||
``PassContextThreadLocalStore`` is used to hold the created pass context
|
||||
objects. Examples will be provided later to show how we can use both the C++ and
|
||||
Python APIs to create a compilation pipeline using pass context.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
class PassContextNode : public Object {
|
||||
public:
|
||||
int opt_level{2};
|
||||
ffi::Array<ffi::String> required_pass;
|
||||
ffi::Array<ffi::String> disabled_pass;
|
||||
mutable ffi::Optional<DiagnosticContext> diag_ctx;
|
||||
ffi::Map<ffi::String, Any> config;
|
||||
ffi::Array<instrument::PassInstrument> instruments;
|
||||
};
|
||||
|
||||
class PassContext : public ObjectRef {
|
||||
public:
|
||||
TVM_DLL static PassContext Create();
|
||||
TVM_DLL static PassContext Current();
|
||||
TVM_DLL void InstrumentEnterPassContext();
|
||||
TVM_DLL void InstrumentExitPassContext();
|
||||
TVM_DLL bool InstrumentBeforePass(const IRModule& mod, const PassInfo& info) const;
|
||||
TVM_DLL void InstrumentAfterPass(const IRModule& mod, const PassInfo& info) const;
|
||||
/* Other fields are omitted. */
|
||||
|
||||
private:
|
||||
// The entry of a pass context scope.
|
||||
TVM_DLL void EnterWithScope();
|
||||
// The exit of a pass context scope.
|
||||
TVM_DLL void ExitWithScope();
|
||||
|
||||
// Classes to get the Python `with` like syntax.
|
||||
friend class tvm::With<PassContext>;
|
||||
};
|
||||
|
||||
struct PassContextThreadLocalEntry {
|
||||
/*! \brief The default pass context. */
|
||||
PassContext default_context;
|
||||
/*! \brief The current pass context. */
|
||||
std::stack<PassContext> context_stack;
|
||||
PassContextThreadLocalEntry() {
|
||||
default_context = PassContext(ffi::make_object<PassContextNode>());
|
||||
}
|
||||
};
|
||||
|
||||
Pass Constructs
|
||||
^^^^^^^^^^^^^^^
|
||||
|
||||
The pass infra is designed in a hierarchical manner, and it could work at
|
||||
different granularities of Relax/TensorIR programs. A pure virtual class ``PassNode`` is
|
||||
introduced to serve as the base of the different optimization passes. This class
|
||||
contains several virtual methods that must be implemented by the
|
||||
subclasses at the level of modules, functions, or sequences of passes.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
class PassNode : Object {
|
||||
virtual PassInfo Info() const = 0;
|
||||
virtual Module operator()(const IRModule& mod
|
||||
const PassContext& pass_ctx) const = 0;
|
||||
};
|
||||
|
||||
The functor shows how a pass must be realized, i.e. it always works on a
|
||||
:py:class:`IRModule` under a certain context. All passes are designed in a ``Module`` to ``Module``
|
||||
manner. Therefore, optimizations governed by the pass infra will
|
||||
always update the whole module.
|
||||
|
||||
Several subclasses have been created to implement different types of
|
||||
optimization passes, e.g., function-level passes, module-level passes, and
|
||||
sequential passes. Each subclass itself could act as a pass manager. For
|
||||
instance, they could collect the required passes and execute them or build
|
||||
a dependency graph based on the given metadata. The full definition of them
|
||||
can be found in `src/ir/transform.cc`_.
|
||||
|
||||
Module-Level Passes
|
||||
^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Module level passes are geared mainly for global and inter-procedural
|
||||
optimizations (IPO), which are similar to the module pass used in LLVM. Some
|
||||
typical passes in Relax that need the global picture of a module, such as
|
||||
A-normal form conversion and lambda lifting, etc., fall into this set. At this
|
||||
level, users can even add and/or delete functions in a module. Note that all
|
||||
passes
|
||||
|
||||
.. code:: c++
|
||||
|
||||
class ModulePassNode : PassNode {
|
||||
PassInfo pass_info;
|
||||
std::function<Module(Module, PassContext)> pass_func;
|
||||
Module operator()(const Module& mod, const PassContext& pass_ctx) const final;
|
||||
// Other members/methods are omitted
|
||||
};
|
||||
|
||||
``pass_info`` maintains the information needed by a module-level pass.
|
||||
``pass_func`` sketches the real optimization. For example, we may need to
|
||||
perform dead code elimination on the module. We could implement the algorithm in
|
||||
the ``pass_func`` and let it run on a module. It will then remove the dead code
|
||||
including the unused functions in the module. Note that this field is designed
|
||||
as a packed function, which enables the implementation of the optimization in
|
||||
both C++ and Python.
|
||||
|
||||
Function-Level Passes
|
||||
^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Function-level passes are used to implement various intra-function level
|
||||
optimizations for a given Relax/TensorIR module. It fetches one function at a time from
|
||||
the function list of a module for optimization and yields a rewritten Relax
|
||||
``Function`` or TensorIR ``PrimFunc``. Most of passes can be classified into this category, such as
|
||||
common subexpression elimination and inference simplification in Relax as well as vectorization
|
||||
and flattening storage in TensorIR, etc.
|
||||
|
||||
Note that the scope of passes at this level is either a Relax function or a TensorIR primitive function.
|
||||
Therefore, we cannot add or delete a function through these passes as they are not aware of
|
||||
the global information.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
class FunctionPassNode : PassNode {
|
||||
PassInfo pass_info;
|
||||
std::function<Function(Function, Module, PassContext)> pass_func;
|
||||
Module operator()(const Module& mod, const PassContext& pass_ctx) const final;
|
||||
bool SkipFunction(const Function& func) const;
|
||||
// Other members/methods are omitted...
|
||||
};
|
||||
|
||||
``pass_info`` is identical to what we just described in the module pass.
|
||||
``pass_func`` takes a function for optimization, it also needs a module as we
|
||||
may use it for reporting errors. A function could be annotated with
|
||||
"SkipOptimization" so that it will be ignored during optimization.
|
||||
|
||||
Sequential Passes
|
||||
^^^^^^^^^^^^^^^^^
|
||||
|
||||
``SequentialPass`` is similar to Pytorch ``nn.Sequential`` that contains a host
|
||||
of passes for execution.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
class SequentialPassNode : PassNode {
|
||||
PassInfo pass_info;
|
||||
// Passes need to be executed.
|
||||
ffi::Array<Pass> passes;
|
||||
bool PassEnabled(const PassInfo& info) const;
|
||||
Module operator()(const Module& mod, const PassContext& pass_ctx) const final;
|
||||
};
|
||||
|
||||
The following code shows how individual passes in a sequential pass are invoked.
|
||||
Essentially, we sequentially execute each pass in a sequential pass using the
|
||||
order that they were appended to the pass list.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
Module SequentialNode::operator()(const Module& module,
|
||||
const PassContext& pass_ctx) const {
|
||||
Module mod = module;
|
||||
for (const Pass& pass : passes) {
|
||||
TVM_FFI_ICHECK(pass.defined()) << "Found undefined pass for optimization.";
|
||||
const PassInfo& pass_info = pass->Info();
|
||||
if (!PassEnabled(pass_info)) continue;
|
||||
for (const auto& it : pass_info->required) {
|
||||
mod = GetPass(it)(std::move(mod), pass_ctx);
|
||||
}
|
||||
mod = pass(mod, pass_ctx);
|
||||
}
|
||||
return mod;
|
||||
}
|
||||
|
||||
Upon the invocation of a pass, we first check if this pass is enabled. This is
|
||||
done by first checking if the pass is explicitly disabled by a user, followed by
|
||||
inspecting if it is specified as a required pass by the user. If it is still
|
||||
undetermined whether this pass is enabled, its ``opt_level`` will be checked.
|
||||
This pass will be enabled and therefore executed only when its optimization
|
||||
level is not less than the configured optimization level in the pass context.
|
||||
|
||||
To execute the pass, we need first to retrieve the registered pass in the TVM
|
||||
packed function registry using the pass name. This is possible because every
|
||||
pass is registered with an API endpoint as we will show later.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
Pass GetPass(const std::string& pass_name) {
|
||||
std::string fpass_name = "relax.transform." + pass_name;
|
||||
const std::optional<tvm::ffi::Function> f = tvm::ffi::Function::GetGlobal(fpass_name);
|
||||
TVM_FFI_ICHECK(f.has_value()) << "Cannot find " << fpass_name
|
||||
<< "to create the pass " << pass_name;
|
||||
return (*f)();
|
||||
}
|
||||
|
||||
Some helper functions are provided to create each type of these aforementioned
|
||||
passes. These helpers are also exposed to the Python frontend for users to
|
||||
favorably use Python APIs to create a specific pass object.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
Pass CreateFunctionPass(
|
||||
std::function<Function(Function, IRModule, PassContext)> pass_func,
|
||||
int opt_level,
|
||||
ffi::String name,
|
||||
ffi::Array<ffi::String> required,
|
||||
bool traceable = false);
|
||||
|
||||
Pass CreatePrimFuncPass(
|
||||
std::function<PrimFunc(PrimFunc, IRModule, PassContext)> pass_func,
|
||||
int opt_level,
|
||||
ffi::String name,
|
||||
ffi::Array<ffi::String> required,
|
||||
bool traceable = false);
|
||||
|
||||
Pass CreateModulePass(
|
||||
std::function<IRModule(IRModule, PassContext)> pass_func,
|
||||
int opt_level,
|
||||
ffi::String name,
|
||||
ffi::Array<ffi::String> required,
|
||||
bool traceable = false);
|
||||
|
||||
Pass Sequential(tvm::ffi::Array<Pass> passes, PassInfo pass_info);
|
||||
|
||||
Pass Registration
|
||||
^^^^^^^^^^^^^^^^^
|
||||
|
||||
We've covered the concept of different level of passes and the context used for
|
||||
compilation. It would be interesting to see how easily users can register
|
||||
a pass. Let's take const folding as an example. This pass has already been
|
||||
implemented to fold constants in a Relax function (found in
|
||||
`src/relax/transform/fold_constant.cc`_).
|
||||
|
||||
An API was provided to perform the ``Expr`` to ``Expr`` transformation.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
Expr FoldConstant(const Expr& expr);
|
||||
|
||||
In order to register this pass to the pass infra, we first need to decide at
|
||||
which level this pass will be performed. As const folding happens on individual
|
||||
functions, we should intuitively create a ``FunctionPass`` for it through
|
||||
``CreateFunctionPass``. The ``pass_func`` is returned as a packed function that
|
||||
invokes the ``Expr`` to ``Expr`` API on each function in a `IRModule`. ``{}``
|
||||
indicates that no prerequisite is required for this pass. Otherwise, the pass
|
||||
developer has to identify and list them.
|
||||
|
||||
Meanwhile, a pass API endpoint is registered with the name
|
||||
``"relax.transform.FoldConstant"``. This pass, therefore, becomes an entry in the
|
||||
registry that can be accessed by both C++ (e.g. the ``GetPass`` above) and
|
||||
Python when needed.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
namespace transform {
|
||||
|
||||
Pass FoldConstant() {
|
||||
auto pass_func =
|
||||
[=](Function f, IRModule m, PassContext pc) { return ConstantFolder::Fold(f, m); };
|
||||
return CreateFunctionPass(pass_func, 0, "FoldConstant", {});
|
||||
}
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def("relax.transform.FoldConstant", FoldConstant);
|
||||
}
|
||||
|
||||
} // namespace transform
|
||||
|
||||
To allow other C++ modules to apply this pass, we declare a free function in
|
||||
`include/tvm/relax/transform.h`_ as the following:
|
||||
|
||||
.. code:: c++
|
||||
|
||||
TVM_DLL Pass FoldConstant();
|
||||
|
||||
.. _pass_instrument_cpp_backend:
|
||||
|
||||
Pass Instrument
|
||||
^^^^^^^^^^^^^^^
|
||||
|
||||
Pass Instrument is a mechanism to analyze the pass itself. For example,
|
||||
we can use the infrastructure to know how much time and memory a pass requires
|
||||
or how a pass can transform the IR module.
|
||||
|
||||
We introduce four instrument points in the life-cycle of ``PassContext``.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
TVM_DLL void InstrumentEnterPassContext();
|
||||
TVM_DLL void InstrumentExitPassContext();
|
||||
TVM_DLL bool InstrumentBeforePass(const IRModule& mod, const PassInfo& info) const;
|
||||
TVM_DLL void InstrumentAfterPass(const IRModule& mod, const PassInfo& info) const;
|
||||
|
||||
``InstrumentEnterPassContext`` is called immediately when entering the scope
|
||||
of the ``PassContext`` instance.
|
||||
|
||||
``InstrumentExitPassContext`` is called when leaving the scope of ``PassContext``,
|
||||
or exceptions occur during the execution of passes.
|
||||
This method is also called when instruments is being overridden by ``override_instruments`` in :py:class:`tvm.transform.PassContext`.
|
||||
See :ref:`pass_instrument_overriden`.
|
||||
|
||||
``InstrumentBeforePass`` is called before execution.
|
||||
``InstrumentAfterPass`` is called after execution if the pass should be run. The behavior is like:
|
||||
|
||||
.. code:: c++
|
||||
|
||||
if (pass_ctx.InstrumentBeforePass(ir_module, pass_info)) {
|
||||
new_ir_module = run_pass(ir_module, pass_ctx);
|
||||
pass_ctx.InstrumentAfterPass(new_ir_module, pass_info);
|
||||
return new_ir_module;
|
||||
}
|
||||
|
||||
The ``PassInstrument`` interface allow you to run arbitrary code inside above four methods.
|
||||
Multiple ``PassInstrument`` instances can be registed into a single
|
||||
``PassContext``. ``PassInstrument`` instances are called sequentially in the order of
|
||||
``instruments`` argument passed to ``PassContext``.
|
||||
|
||||
``PassInstrument`` provides following interfaces:
|
||||
|
||||
.. code:: c++
|
||||
|
||||
namespace instrument {
|
||||
|
||||
class PassInstrumentNode : public Object {
|
||||
public:
|
||||
ffi::String name;
|
||||
virtual void EnterPassContext() const = 0;
|
||||
virtual void ExitPassContext() const = 0;
|
||||
virtual bool ShouldRun(const IRModule& mod, const transform::PassInfo& info) const = 0;
|
||||
virtual void RunBeforePass(const IRModule& mod, const transform::PassInfo& info) const = 0;
|
||||
virtual void RunAfterPass(const IRModule& mod, const transform::PassInfo& info) const = 0;
|
||||
/* Other fields are omitted. */
|
||||
};
|
||||
|
||||
class PassInstrument : public ObjectRef {
|
||||
public:
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PassInstrument, ObjectRef, PassInstrumentNode);
|
||||
};
|
||||
|
||||
} // namespace instrument
|
||||
|
||||
Python frontend are provided to implement ``PassInstrument`` quickly. See :ref:`pass_instrument_py_frontend`.
|
||||
|
||||
Within a ``PassContext``, the call sequence of a ``PassInstrument`` instance is like:
|
||||
|
||||
::
|
||||
|
||||
with PassContext(instruments=[pi]) # pi = a PassInstrument implementation.
|
||||
pi.EnterPassContext()
|
||||
|
||||
if pi.ShouldRun(Pass1):
|
||||
pi.RunBeforePass()
|
||||
Pass1()
|
||||
pi.RunAfterPass()
|
||||
|
||||
if pi.ShouldRun(Pass2):
|
||||
pi.RunBeforePass()
|
||||
Pass2()
|
||||
pi.RunAfterPass()
|
||||
|
||||
pi.ExitPassContext()
|
||||
|
||||
Here is a brief introduction of relations between ``PassInstrument`` interfaces
|
||||
and ``PassContext`` methods. See (`src/ir/transform.cc`_) for more details.
|
||||
|
||||
- ``InstrumentEnterPassContext``
|
||||
|
||||
* ``EnterPassContext()`` is executed in the order of ``instruments`` passed to the ``PassContext``.
|
||||
* When an exception raises, ``PassContext`` disable the pass instrumentation
|
||||
by clearing all registered ``PassInstrument`` instances.
|
||||
* Then ``PassContext`` execute ``ExitPassContext()`` method of each ``PassInstrument``
|
||||
instances which successfully finished ``EnterPassContext()``
|
||||
* For example, if ``PassInstrument`` A, B, and C are registered to a ``PassContext``
|
||||
and A finished ``EnterPassContext()`` while B throws an exception, then C
|
||||
is never executed; ``ExitPassContext()`` of A is executed.
|
||||
|
||||
- ``InstrumentExitPassContext``
|
||||
|
||||
* ``ExitPassContext()`` of each ``PassInstrument`` instances are executed in
|
||||
the order of ``instruments`` passed to the ``PassContext``.
|
||||
* While an exception occurs, ``instruments`` is cleared.
|
||||
* ``PassInstrument`` Instances registered after the one throwing exceptions do not execute ``ExitPassContext``.
|
||||
|
||||
- ``InstrumentBeforePass``
|
||||
|
||||
* ``ShouldRun`` is executed if the pass is not listed as a required pass.
|
||||
* ``RunBeforePass`` is executed in the order of ``instruments`` if the pass is not blocked by ``ShouldRun``.
|
||||
* Note that ``InstrumentBeforePass`` returns a boolean indicating whether or not the pass should be run.
|
||||
* When an exception occur, it is thrown immediately.
|
||||
We rely on Python Context Manager to exit ``PassContext`` safely
|
||||
(meaning ``ExitPassContext`` of each instruments will be run. For C++, please refer to `include/tvm/support/with.h`_.)
|
||||
|
||||
- ``InstrumentAfterPass``
|
||||
|
||||
* ``RunAfterPass`` is executed in the order of ``instruments`` passed to the ``PassContext``.
|
||||
* When an exception occur, it is thrown immediately.
|
||||
We rely on Python Context Manager or ``With`` class(`include/tvm/support/with.h`_) to exit ``PassContext`` safely
|
||||
|
||||
Built-in Instrument
|
||||
^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
There are several built-in instruments.
|
||||
|
||||
- PassTimingInstrument (see `src/ir/instrument.cc`_)
|
||||
|
||||
* Profile the execution time of passes.
|
||||
|
||||
- PrintBeforeAll (see `python/tvm/ir/instrument.py`_)
|
||||
|
||||
* Print the IR module and pass info before each pass executes.
|
||||
|
||||
- PrintAfterAll (see `python/tvm/ir/instrument.py`_)
|
||||
|
||||
* Print the IR module and pass info after each pass executes.
|
||||
|
||||
- PassPrintingInstrument (see `python/tvm/ir/instrument.py`_)
|
||||
|
||||
* Selectively print the IR module before or after specific named passes.
|
||||
|
||||
- DumpIR (see `python/tvm/ir/instrument.py`_)
|
||||
|
||||
* Dump the IR module to files after each pass executes.
|
||||
|
||||
Python Frontend
|
||||
~~~~~~~~~~~~~~~
|
||||
|
||||
Only some simple APIs are needed for the frontend side. For example, we can
|
||||
provide users the following APIs to create and execute a pass (full
|
||||
implementation is provided in `python/tvm/relax/transform/transform.py`_ and
|
||||
`python/tvm/ir/transform.py`_). The backend
|
||||
receives the information and decides which function it should use to create
|
||||
a Pass object.
|
||||
|
||||
PassContext
|
||||
^^^^^^^^^^^
|
||||
|
||||
Python frontend provides a wrapper for the ``PassContext`` to enable the
|
||||
``with`` syntax by overriding ``__enter__`` and ``__exit__``. A ``current``
|
||||
static method is offered for users to get the context that is in use under
|
||||
a certain scope.
|
||||
|
||||
.. code:: python
|
||||
|
||||
@tvm_ffi.register_object("transform.PassContext")
|
||||
class PassContext(tvm.runtime.Object):
|
||||
def __enter__(self):
|
||||
_transform.EnterPassContext(self)
|
||||
return self
|
||||
|
||||
def __exit__(self, ptype, value, trace, config):
|
||||
_transform.ExitPassContext(self)
|
||||
|
||||
@staticmethod
|
||||
def current():
|
||||
"""Return the current pass context."""
|
||||
return _transform.GetCurrentPassContext()
|
||||
|
||||
A ``PassContext`` is used to configure the compilation options, including the
|
||||
optimization level and required/disabled passes. It can also take a dictionary
|
||||
of configs so that different passes can conveniently fetch the passed data, such
|
||||
as fallback device info and step/depth for loop unrolling, etc. In order to
|
||||
enable fetching the required config, the key must be registered through
|
||||
``TVM_REGISTER_PASS_CONFIG_OPTION``. For example, the following is used by the
|
||||
loop unrolling pass
|
||||
|
||||
.. code:: c++
|
||||
|
||||
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.UnrollLoop", UnrollLoopConfig);
|
||||
|
||||
Please refer to `src/tirx/transform/unroll_loop.cc`_ for more details.
|
||||
|
||||
.. _pass_instrument_py_frontend:
|
||||
|
||||
Pass Instrument
|
||||
^^^^^^^^^^^^^^^
|
||||
|
||||
One can implement a ``PassInstrument`` by using the ``pass_instrument``
|
||||
decorator(`python/tvm/ir/instrument.py`_) on a class implementing following methods.
|
||||
Note that it is recommended to use the ``pass_instrument`` decorator to implement
|
||||
``PassInstrument``, instead of overriding or subclassing.
|
||||
|
||||
- ``enter_pass_ctx``
|
||||
|
||||
* This method is run when entering ``PassContext``.
|
||||
|
||||
- ``exit_pass_ctx``
|
||||
|
||||
* This method is run when exiting ``PassContext``.
|
||||
|
||||
- ``should_run``
|
||||
|
||||
* This method is run before a pass is executed, returning a boolean
|
||||
indicating whether or not the pass should be run.
|
||||
|
||||
- ``run_before_pass``
|
||||
|
||||
* If a pass should be run, this method is run just before pass execution.
|
||||
|
||||
- ``run_after_pass``
|
||||
|
||||
* This method is run right after a pass has been executed.
|
||||
|
||||
``PassInstrument`` instances can be registered through ``instruments`` argument in
|
||||
:py:class:`tvm.transform.PassContext`.
|
||||
|
||||
See `python/tvm/ir/instrument.py`_ for examples of how to implement ``PassInstrument`` with Python APIs.
|
||||
|
||||
.. _pass_instrument_overriden:
|
||||
|
||||
Override Instruments in Current PassContext
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
``override_instruments`` method is provided to override the ``instruments`` of current ``PassContext``.
|
||||
For example, if passes are run without explicitly creating a new ``PassContext``,
|
||||
one can still register ``PassInstrument`` into the global ``PassContext`` by:
|
||||
|
||||
.. code:: python
|
||||
|
||||
cur_pass_ctx = tvm.transform.PassContext.current()
|
||||
# override PassInstrument instances
|
||||
cur_pass_ctx.override_instruments([pass_inst])
|
||||
mod = pass_seq(mod)
|
||||
result = pass_inst.get_result()
|
||||
|
||||
Note that when ``override_instruments`` is called, the ``exit_pass_ctx`` method of
|
||||
old ``PassInstrument`` instances are called. Then the ``enter_pass_ctx`` method of
|
||||
new ``PassInstrument`` are called.
|
||||
|
||||
.. _Sequential: https://pytorch.org/docs/stable/nn.html?highlight=sequential#torch.nn.Sequential
|
||||
|
||||
.. _Block: https://pytorch.org/docs/stable/generated/torch.nn.Module.html
|
||||
|
||||
.. _include/tvm/ir/transform.h: https://github.com/apache/tvm/blob/main/include/tvm/ir/transform.h
|
||||
|
||||
.. _include/tvm/support/with.h: https://github.com/apache/tvm/blob/main/include/tvm/support/with.h
|
||||
|
||||
.. _src/relax/ir/transform.cc: https://github.com/apache/tvm/blob/main/src/relax/ir/transform.cc
|
||||
|
||||
.. _src/ir/transform.cc: https://github.com/apache/tvm/blob/main/src/ir/transform.cc
|
||||
|
||||
.. _src/ir/instrument.cc: https://github.com/apache/tvm/blob/main/src/ir/instrument.cc
|
||||
|
||||
.. _src/relax/transform/fold_constant.cc: https://github.com/apache/tvm/blob/main/src/relax/transform/fold_constant.cc
|
||||
|
||||
.. _python/tvm/relax/transform/transform.py: https://github.com/apache/tvm/blob/main/python/tvm/relax/transform/transform.py
|
||||
|
||||
.. _include/tvm/relax/transform.h: https://github.com/apache/tvm/blob/main/include/tvm/relax/transform.h
|
||||
|
||||
.. _python/tvm/ir/transform.py: https://github.com/apache/tvm/blob/main/python/tvm/ir/transform.py
|
||||
|
||||
.. _python/tvm/ir/instrument.py: https://github.com/apache/tvm/blob/main/python/tvm/ir/instrument.py
|
||||
|
||||
.. _src/tirx/transform/unroll_loop.cc: https://github.com/apache/tvm/blob/main/src/tirx/transform/unroll_loop.cc
|
||||
|
||||
.. _use pass infra: https://github.com/apache/tvm/blob/main/docs/how_to/tutorials/customize_opt.py
|
||||
@@ -0,0 +1,427 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
.. _relax-vm-arch:
|
||||
|
||||
Relax Virtual Machine
|
||||
=====================
|
||||
|
||||
This document explains the Relax VM architecture in detail, covering the compilation pipeline
|
||||
from Relax IR to bytecode, the instruction set, the execution model, and the Python-level user
|
||||
interface.
|
||||
|
||||
Overview
|
||||
--------
|
||||
|
||||
The end-to-end flow from model to execution is:
|
||||
|
||||
1. **Relax IR** — a high-level computational graph (``relax.Function`` inside an ``IRModule``).
|
||||
2. **Compilation** — ``tvm.compile()`` applies the Relax transformation pipeline, then invokes
|
||||
``VMCodeGen`` to translate each Relax function into bytecode instructions.
|
||||
3. **Linking** — TIR functions are compiled to native kernels (via LLVM, CUDA, etc.); the bytecode,
|
||||
constant pool, and compiled kernels are packaged together into a ``VMExecutable``.
|
||||
4. **Execution** — at runtime, a ``VirtualMachine`` loads the executable, initializes devices and
|
||||
memory allocators, and runs the bytecode.
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
IRModule (Relax + TIR)
|
||||
│
|
||||
▼ relax_pipeline (FuseOps, LegalizeOps, ...)
|
||||
IRModule (optimized)
|
||||
│
|
||||
▼ VMCodeGen
|
||||
ExecBuilder (bytecode) + IRModule (TIR only)
|
||||
│ │
|
||||
│ ▼ tirx.build()
|
||||
│ runtime.Module (native kernels)
|
||||
│ │
|
||||
▼ VMLink ▼
|
||||
VMExecutable ◄───────── linked together
|
||||
│
|
||||
▼ VirtualMachine(exec, device)
|
||||
Runtime execution
|
||||
|
||||
|
||||
Compilation: From Relax IR to Bytecode
|
||||
--------------------------------------
|
||||
|
||||
Build entry point
|
||||
~~~~~~~~~~~~~~~~~
|
||||
|
||||
The main entry point is ``tvm.compile()`` (which delegates to ``relax.build()`` in
|
||||
``python/tvm/relax/vm_build.py``):
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import tvm
|
||||
from tvm import relax
|
||||
|
||||
@tvm.script.ir_module
|
||||
class MyModule:
|
||||
@R.function
|
||||
def main(x: R.Tensor((3, 4), "float32")):
|
||||
return R.add(x, x)
|
||||
|
||||
target = tvm.target.Target("llvm")
|
||||
ex = tvm.compile(MyModule, target)
|
||||
|
||||
Internally, ``relax.build()`` performs these steps:
|
||||
|
||||
1. Apply the **Relax pipeline** (``relax.get_pipeline("default")``), which includes operator
|
||||
legalization, fusion, buffer planning, and other graph-level passes.
|
||||
2. Create an ``ExecBuilder`` and run **VMCodeGen** (``src/relax/backend/vm/codegen_vm.cc``),
|
||||
which walks each ``relax.Function`` and emits bytecode instructions. The Relax functions are
|
||||
removed from the IRModule; only TIR functions remain.
|
||||
3. Compile the remaining TIR functions to native code via ``tirx.build()``.
|
||||
4. **Link** the bytecode executable with the compiled native module using ``VMLink``, producing
|
||||
a ``VMExecutable``.
|
||||
|
||||
Two execution modes are supported:
|
||||
|
||||
- ``exec_mode="bytecode"`` (default): Relax functions are interpreted by the VM's bytecode
|
||||
dispatch loop.
|
||||
- ``exec_mode="compiled"``: Relax functions are compiled into TIR functions (``VMTIRCodeGen``)
|
||||
that directly manipulate the register file, bypassing the interpreter loop. This avoids
|
||||
dispatch overhead but produces more code.
|
||||
|
||||
Bytecode generation
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
The ``CodeGenVM`` class (``src/relax/backend/vm/codegen_vm.cc``) is an ``ExprFunctor`` that visits
|
||||
each Relax expression and emits instructions through the ``ExecBuilder``:
|
||||
|
||||
- Each ``relax.Var`` is mapped to a register.
|
||||
- Function parameters occupy registers 0 through N-1.
|
||||
- Each binding in a ``SeqExpr`` generates one or more instructions; the result is stored in a
|
||||
new register.
|
||||
- Function calls (``R.call_tir``, ``R.call_packed``, operator calls) become ``Call`` instructions.
|
||||
- Conditional expressions (``relax.If``, written as Python ``if`` in TVMScript) become an ``If``
|
||||
instruction followed by ``Goto`` to skip branches.
|
||||
- The function body ends with a ``Ret`` instruction.
|
||||
|
||||
|
||||
Instruction Set
|
||||
---------------
|
||||
|
||||
The VM uses a **register-based** architecture with an intentionally minimal instruction set.
|
||||
There are only four opcodes:
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 15 30 55
|
||||
|
||||
* - Opcode
|
||||
- Fields
|
||||
- Semantics
|
||||
* - ``Call``
|
||||
- ``dst``, ``func_idx``, ``num_args``, ``args[]``
|
||||
- Call function ``func_idx`` with the given arguments; store the result in register ``dst``.
|
||||
* - ``Ret``
|
||||
- ``result``
|
||||
- Return the value in register ``result`` to the caller.
|
||||
* - ``Goto``
|
||||
- ``pc_offset``
|
||||
- Jump forward or backward by ``pc_offset`` instructions.
|
||||
* - ``If``
|
||||
- ``cond``, ``false_offset``
|
||||
- If register ``cond`` is nonzero, fall through (pc++); otherwise jump by ``false_offset``.
|
||||
|
||||
The VM itself performs **no mathematical computation**. All actual work — matrix multiplications,
|
||||
convolutions, elementwise operations — is carried out by compiled TIR kernels or external
|
||||
libraries (cuBLAS, cuDNN, etc.), dispatched through ``Call`` instructions.
|
||||
|
||||
Instruction encoding
|
||||
~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Each instruction argument (``Instruction::Arg``) is a 64-bit word encoded as:
|
||||
|
||||
- **Bits [63:56]** — ``ArgKind`` (8 bits): ``kRegister`` (0), ``kImmediate`` (1), ``kConstIdx`` (2),
|
||||
or ``kFuncIdx`` (3).
|
||||
- **Bits [55:0]** — value (56 bits, sign-extended).
|
||||
|
||||
Two special register values exist:
|
||||
|
||||
- ``kVoidRegister``: indicates "no destination" (the return value is discarded).
|
||||
- ``kVMRegister``: refers to the VM context pointer itself, passed as the first argument to
|
||||
closures.
|
||||
|
||||
The instruction stream is stored as a flat ``vector<ExecWord>`` (``instr_data``) with an offset
|
||||
table (``instr_offset``) for random access.
|
||||
|
||||
|
||||
Executable
|
||||
----------
|
||||
|
||||
A ``VMExecutable`` (``include/tvm/runtime/vm/executable.h``) bundles everything needed for
|
||||
execution:
|
||||
|
||||
- **Function table** (``func_table``): a ``vector<VMFuncInfo>`` describing every function. Each
|
||||
entry records the function's kind, name, instruction range (``start_instr`` to ``end_instr``),
|
||||
number of arguments, register file size, and parameter names.
|
||||
- **Constant pool** (``constants``): model weights, shape tuples, and other compile-time constants.
|
||||
- **Bytecode** (``instr_data`` + ``instr_offset``): the instruction stream.
|
||||
- **Imported modules**: the compiled TIR kernels and external libraries.
|
||||
|
||||
Function kinds
|
||||
~~~~~~~~~~~~~~
|
||||
|
||||
The VM recognizes three function kinds (``VMFuncInfo::FuncKind``):
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 20 80
|
||||
|
||||
* - Kind
|
||||
- Description
|
||||
* - ``kPackedFunc``
|
||||
- An external C/C++ function looked up from imported modules or the global PackedFunc
|
||||
registry. Examples: ``vm.builtin.alloc_shape_heap``, ``vm.builtin.match_shape``.
|
||||
* - ``kVMFunc``
|
||||
- A bytecode-interpreted Relax function. The VM interprets its instructions in ``RunLoop()``.
|
||||
* - ``kVMTIRFunc``
|
||||
- A Relax function compiled to a TIR function (``exec_mode="compiled"``). Found in
|
||||
imports under the name ``__vmtir__<func_name>``. Called directly with register file
|
||||
pointers, bypassing the interpreter loop.
|
||||
|
||||
Serialization
|
||||
~~~~~~~~~~~~~
|
||||
|
||||
The executable supports binary serialization for deployment:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Save
|
||||
ex.export_library("model.so")
|
||||
|
||||
# Load
|
||||
loaded = tvm.runtime.load_module("model.so")
|
||||
vm = relax.VirtualMachine(loaded, tvm.cuda())
|
||||
|
||||
The binary format includes a magic number (``0xD225DE2F4214151E``), a version string
|
||||
(currently ``"0.14"``), followed by four sections: globals (the function table), memory scopes,
|
||||
constant pool, and bytecode. ``AsText()`` and ``AsPython()`` provide human-readable representations
|
||||
for debugging.
|
||||
|
||||
|
||||
Runtime Execution
|
||||
-----------------
|
||||
|
||||
VM initialization
|
||||
~~~~~~~~~~~~~~~~~
|
||||
|
||||
At runtime, a ``VirtualMachine`` is created and initialized:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.relax import VirtualMachine
|
||||
|
||||
vm = VirtualMachine(exec_module, tvm.cuda())
|
||||
|
||||
Under the hood:
|
||||
|
||||
1. **LoadExecutable**: the bytecode and metadata are loaded from the ``VMExecutable``.
|
||||
2. **Init**: devices and memory allocators are set up. Each device gets an ``Allocator``
|
||||
(either ``NAIVE_ALLOCATOR`` or ``POOLED_ALLOCATOR``, defaulting to pooled). A CPU device
|
||||
is always added for shape computations.
|
||||
3. **InitFuncPool**: the function pool is populated — ``kPackedFunc`` entries are resolved from
|
||||
imports or the global registry; ``kVMFunc`` and ``kVMTIRFunc`` entries are wrapped in
|
||||
``VMClosure`` objects.
|
||||
4. **Constant pool**: model constants are loaded and optionally transferred to the target device.
|
||||
|
||||
The bytecode dispatch loop
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
When a ``kVMFunc`` is invoked, the VM enters ``InvokeBytecode()``:
|
||||
|
||||
1. A new ``VMFrame`` is pushed onto the call stack. Each frame contains:
|
||||
|
||||
- A **register file** (``vector<ffi::Any>``) — type-erased slots that can hold tensors,
|
||||
shapes, closures, or any TVM object. The size is determined at compile time
|
||||
(``VMFuncInfo::register_file_size``).
|
||||
- The **return program counter** — where to resume after the function returns.
|
||||
- The **caller's return register** — which register in the parent frame receives the result.
|
||||
|
||||
2. Function arguments are written to registers 0..N-1.
|
||||
3. The program counter (``pc_``) is set to the function's ``start_instr``.
|
||||
4. ``RunLoop()`` executes instructions until a ``Ret`` is encountered:
|
||||
|
||||
- **Call**: resolve arguments (from registers, immediates, constant pool, or function pool),
|
||||
invoke the target function via ``InvokeClosurePacked()``, store the result in ``dst``.
|
||||
- **Ret**: read the return value from the specified register, write the result to the
|
||||
caller's return register, and return from ``RunLoop()`` (the frame is popped by an RAII
|
||||
guard when ``InvokeBytecode()`` exits).
|
||||
- **Goto**: adjust ``pc_`` by the offset.
|
||||
- **If**: check the condition register; if nonzero, fall through; otherwise jump by
|
||||
``false_offset``.
|
||||
|
||||
The dispatch loop is implemented in ``src/runtime/vm/vm.cc`` (``VirtualMachineImpl::RunLoop``).
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
Frame Stack Register File (per frame)
|
||||
┌─────────────┐ ┌────┬────┬────┬─────┬────┐
|
||||
│ Frame 2 │ ───────► │ R0 │ R1 │ R2 │ ... │ Rn │
|
||||
├─────────────┤ └────┴────┴────┴─────┴────┘
|
||||
│ Frame 1 │ ───────► [register file]
|
||||
├─────────────┤
|
||||
│ Frame 0 │ ───────► [register file]
|
||||
└─────────────┘
|
||||
|
||||
VMClosure and function dispatch
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Functions in the VM are stored in a ``func_pool_`` indexed by function table position.
|
||||
``kVMFunc`` and ``kVMTIRFunc`` entries are wrapped as ``VMClosure`` objects, while ``kPackedFunc``
|
||||
entries are stored as plain ``ffi::Function``. A ``VMClosure`` stores:
|
||||
|
||||
- ``func_name``: the function's string name.
|
||||
- ``impl``: a ``ffi::Function`` that takes the VM context pointer as its first argument, followed
|
||||
by the actual parameters.
|
||||
|
||||
When the VM encounters a ``Call`` instruction, it looks up the function in ``func_pool_`` by
|
||||
index and dispatches via ``InvokeClosurePacked()``. If the target is a ``VMClosure``, the VM
|
||||
pointer is prepended to the arguments and ``impl`` is invoked. If it is a plain
|
||||
``ffi::Function``, it is called directly.
|
||||
|
||||
``VMClosure::BindLastArgs`` enables partial application — it creates a new function with
|
||||
some arguments pre-bound at the end, useful for implementing captured closures in Relax.
|
||||
|
||||
Built-in operations
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
The VM relies on several built-in PackedFuncs (registered in ``src/runtime/vm/builtin.cc``)
|
||||
for runtime support:
|
||||
|
||||
- ``vm.builtin.alloc_shape_heap``: allocate workspace for symbolic shape computations.
|
||||
- ``vm.builtin.match_shape``: validate tensor shapes against expected patterns at runtime,
|
||||
supporting assertions (``kAssertEqualToImm``, ``kAssertEqualToLoad``), storing symbolic
|
||||
dimensions to the shape heap (``kStoreToHeap``), or no-ops (``kNoOp``).
|
||||
- ``vm.builtin.make_shape``: construct shape tuples from immediates or heap-loaded values.
|
||||
- ``vm.builtin.match_prim_value``: validate primitive values (e.g., integers) against expected
|
||||
patterns.
|
||||
- ``vm.builtin.copy``: copy a value into a register. Used in several codegen scenarios:
|
||||
materializing non-register arguments (immediates, constants) into registers, ensuring each
|
||||
variable binding gets its own register, and merging results from if/else branches.
|
||||
|
||||
|
||||
Python Interface
|
||||
----------------
|
||||
|
||||
Users interact with the VM through ``tvm.relax.VirtualMachine``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import tvm
|
||||
from tvm import relax
|
||||
import numpy as np
|
||||
|
||||
# Compile
|
||||
ex = tvm.compile(MyModule, target="llvm")
|
||||
|
||||
# Create VM
|
||||
vm = relax.VirtualMachine(ex, tvm.cpu())
|
||||
|
||||
# Direct invocation
|
||||
inp = tvm.runtime.tensor(np.random.rand(3, 4).astype("float32"))
|
||||
result = vm["main"](inp)
|
||||
|
||||
# Stateful interface (useful for RPC)
|
||||
vm.set_input("main", inp)
|
||||
vm.invoke_stateful("main")
|
||||
output = vm.get_outputs("main")
|
||||
|
||||
Key methods:
|
||||
|
||||
- ``vm["func_name"](*args)`` — direct invocation, returns the result.
|
||||
- ``vm.set_input()`` / ``vm.invoke_stateful()`` / ``vm.get_outputs()`` — stateful interface
|
||||
that avoids sending output over the wire, useful for RPC-based remote execution.
|
||||
- ``vm.save_function(func_name, saved_name, *args)`` — pre-bind arguments for repeated calls,
|
||||
reducing dictionary lookup overhead during benchmarking.
|
||||
- ``vm.time_evaluator(func_name, dev)`` — returns a timing function following the same convention
|
||||
as ``tvm.runtime.Module.time_evaluator``.
|
||||
- ``vm.set_instrument(func)`` — register an instrumentation callback that is invoked before/after
|
||||
every ``Call`` instruction. The callback can return ``VMInstrumentReturnKind.SKIP_RUN`` to
|
||||
skip the call.
|
||||
|
||||
Instrumentation
|
||||
~~~~~~~~~~~~~~~
|
||||
|
||||
The VM supports observability via instrumentation:
|
||||
|
||||
**Instrumentation** via ``set_instrument()``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def my_instrument(func, func_symbol, before_run, ret_value, *args):
|
||||
if before_run:
|
||||
print(f"About to call: {func_symbol}")
|
||||
return VMInstrumentReturnKind.NO_OP
|
||||
|
||||
vm.set_instrument(my_instrument)
|
||||
vm["main"](inp)
|
||||
|
||||
The instrument function is called before and after every ``Call`` instruction, receiving the
|
||||
function object, its symbol name, a flag indicating before/after, the return value (only valid
|
||||
after), and all arguments.
|
||||
|
||||
|
||||
Inspecting Bytecode
|
||||
-------------------
|
||||
|
||||
The executable provides text and Python representations of the compiled bytecode:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
ex = tvm.compile(MyModule, target="llvm")
|
||||
print(ex.as_text()) # Human-readable instruction listing
|
||||
print(ex.as_python()) # Equivalent Python program
|
||||
print(ex.stats()) # Summary statistics
|
||||
|
||||
These are invaluable for debugging compilation issues — they show exactly which functions
|
||||
are called, in what order, and how registers are used.
|
||||
|
||||
|
||||
Source Code Map
|
||||
---------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 45 55
|
||||
|
||||
* - Path
|
||||
- Contents
|
||||
* - ``include/tvm/runtime/vm/bytecode.h``
|
||||
- Instruction, Opcode, and Arg definitions
|
||||
* - ``include/tvm/runtime/vm/executable.h``
|
||||
- VMExecutable, VMFuncInfo, serialization
|
||||
* - ``include/tvm/runtime/vm/vm.h``
|
||||
- VirtualMachine base class, VMClosure
|
||||
* - ``src/runtime/vm/vm.cc``
|
||||
- VirtualMachineImpl, RunLoop, InvokeBytecode
|
||||
* - ``src/runtime/vm/executable.cc``
|
||||
- Serialization/deserialization, text output
|
||||
* - ``src/runtime/vm/builtin.cc``
|
||||
- Built-in operations (shape matching, allocation)
|
||||
* - ``src/relax/backend/vm/codegen_vm.cc``
|
||||
- CodeGenVM: Relax IR → bytecode
|
||||
* - ``src/relax/backend/vm/codegen_vm_tir.cc``
|
||||
- VMTIRCodeGen: Relax IR → compiled TIR
|
||||
* - ``python/tvm/runtime/vm.py``
|
||||
- Python VirtualMachine wrapper
|
||||
* - ``python/tvm/relax/vm_build.py``
|
||||
- ``relax.build()`` and VMExecutable Python class
|
||||
@@ -0,0 +1,281 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
.. _tvm-runtime-system:
|
||||
|
||||
TVM Runtime System
|
||||
==================
|
||||
|
||||
TVM supports multiple programming languages for the compiler stack development and deployment.
|
||||
In this note, we explain the key elements of the TVM runtime.
|
||||
|
||||
.. image:: https://tvm.apache.org/images/release/tvm_flexible.png
|
||||
|
||||
We need to satisfy quite a few interesting requirements:
|
||||
|
||||
- Deployment: invoke the compiled function from python/javascript/c++ language.
|
||||
- Debug: define a function in python and call that from a compiled function.
|
||||
- Link: write driver code to call device specific code (CUDA) and call it from compiled host function.
|
||||
- Prototype: define an IR pass from python and call that from C++ backend.
|
||||
- Expose: compiler stack developed in c++ to front-end (i.e, python)
|
||||
- Experiment: ship a compiled function to an embedded device to directly run there.
|
||||
|
||||
We want to be able to define a function from any language and call from another.
|
||||
We also want the runtime core to be minimal to deploy to embedded devices.
|
||||
|
||||
.. _tvm-runtime-system-packed-func:
|
||||
|
||||
PackedFunc
|
||||
----------
|
||||
|
||||
`PackedFunc`_ is a simple but elegant solution we find to solve the
|
||||
challenges listed. A single ``PackedFunc`` object represents a
|
||||
function call whose caller and callee may be in different languages.
|
||||
|
||||
The following code block provides an example in C++
|
||||
|
||||
.. _PackedFunc: https://github.com/apache/tvm/blob/main/3rdparty/tvm-ffi/include/tvm/ffi/function.h
|
||||
|
||||
.. code:: c
|
||||
|
||||
#include <tvm/ffi/function.h>
|
||||
|
||||
void MyAdd(ffi::PackedArgs args, ffi::Any* rv) {
|
||||
// automatically convert arguments to desired type.
|
||||
int a = args[0].cast<int>();
|
||||
int b = args[1].cast<int>();
|
||||
// automatically assign value return to rv
|
||||
*rv = a + b;
|
||||
}
|
||||
|
||||
void CallPacked() {
|
||||
PackedFunc myadd = PackedFunc(MyAdd);
|
||||
// get back 3
|
||||
int c = myadd(1, 2);
|
||||
}
|
||||
|
||||
In the above codeblock, we defined a PackedFunc MyAdd. It takes two arguments
|
||||
: ``args`` represents input arguments and ``rv`` represents return value.
|
||||
The function is type-erased, which means that the function signature does not restrict which input type to pass in or type to return.
|
||||
Under the hood, when we call a PackedFunc, it packs the input arguments to ffi::PackedArgs on stack,
|
||||
and gets the result back via ffi::Any.
|
||||
|
||||
Thanks to template tricks in C++, we can call a PackedFunc just like a normal function. Because of its type-erased nature, we can call a PackedFunc from dynamic languages like python, without additional glue code for each new type function created.
|
||||
The following example registers PackedFunc in C++ and calls from python.
|
||||
|
||||
.. code:: c
|
||||
|
||||
// register a global packed function in c++
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def_packed("myadd", MyAdd);
|
||||
}
|
||||
|
||||
.. code:: python
|
||||
|
||||
import tvm
|
||||
|
||||
myadd = tvm.get_global_func("myadd")
|
||||
# prints 3
|
||||
print(myadd(1, 2))
|
||||
|
||||
Most of the magic of PackedFunc lies in ``ffi::PackedArgs`` and ``ffi::Any`` structure.
|
||||
We restrict a list of possible types which can be passed.
|
||||
Here are the common ones:
|
||||
|
||||
- int, float and string
|
||||
- PackedFunc itself
|
||||
- Module for compiled modules
|
||||
- DLTensor* for tensor object exchange
|
||||
- TVM Object to represent any object in IR
|
||||
|
||||
The restriction makes the implementation simple without the need of serialization.
|
||||
Despite being minimum, the PackedFunc is sufficient for the use-case of deep learning deployment as
|
||||
most functions only take DLTensor or numbers.
|
||||
|
||||
Since one PackedFunc can take another PackedFunc as an argument,
|
||||
we can pass functions from python (as PackedFunc) to C++.
|
||||
|
||||
.. code:: c
|
||||
|
||||
TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::GlobalDef().def_packed("callhello", [](ffi::PackedArgs args, ffi::Any* rv) {
|
||||
ffi::Function f = args[0].cast<ffi::Function>();
|
||||
f("hello world");
|
||||
});
|
||||
}
|
||||
|
||||
.. code:: python
|
||||
|
||||
import tvm
|
||||
|
||||
def callback(msg):
|
||||
print(msg)
|
||||
|
||||
# convert to PackedFunc
|
||||
f = tvm.runtime.convert(callback)
|
||||
callhello = tvm.get_global_func("callhello")
|
||||
# prints hello world
|
||||
callhello(f)
|
||||
|
||||
TVM provides a `minimum C API`_,
|
||||
which allows us to embed the PackedFunc into any languages. Besides python, so far we supported
|
||||
`java`_ and `javascript`_.
|
||||
This philosophy of embedded API is very like Lua, except that we don't have a new language but use C++.
|
||||
|
||||
.. _minimum C API: https://github.com/apache/tvm/blob/main/3rdparty/tvm-ffi/include/tvm/ffi/c_api.h
|
||||
.. _java: https://github.com/apache/tvm/tree/main/jvm
|
||||
.. _javascript: https://github.com/apache/tvm/tree/main/web
|
||||
|
||||
|
||||
One fun fact about PackedFunc is that we use it for both compiler and deployment stack.
|
||||
|
||||
- All compiler pass functions of TVM are exposed to frontend as PackedFunc
|
||||
- The compiled module also returns the compiled function as PackedFunc
|
||||
|
||||
To keep the runtime minimum, we isolated the IR Object support from the deployment runtime. The resulting runtime takes around 200K - 600K depending on how many runtime driver modules (e.g., CUDA) get included.
|
||||
|
||||
The overhead of calling into PackedFunc vs. a normal function is small, as it is only saving a few values on the stack.
|
||||
So it is OK as long as we don't wrap small functions.
|
||||
In summary, the PackedFunc is the universal glue in TVM where we use it extensively to support our compiler and deployment.
|
||||
|
||||
.. _tvm-runtime-system-module:
|
||||
|
||||
Module
|
||||
------
|
||||
|
||||
Since TVM supports multiple types of devices, we need to support different type of drivers.
|
||||
We have to use the driver API to load the kernel, set up the argument in packed format and perform kernel launch.
|
||||
We also need to patch up the driver API so that the exposed functions are threadsafe.
|
||||
So we often need to implement these driver glues in C++ and expose them to the user.
|
||||
We can certainly not do it for each type of functions, so again PackedFunc is our answer.
|
||||
|
||||
TVM defines the compiled object as `Module`_.
|
||||
The user can get the compiled function from Module as PackedFunc.
|
||||
The generated compiled code can dynamically get function from Module in runtime. It caches the function handle in the first call and reuses in subsequent calls. We use this to link device code and callback into any PackedFunc(e.g., python) from generated code.
|
||||
|
||||
.. _Module: https://github.com/apache/tvm/blob/main/3rdparty/tvm-ffi/include/tvm/ffi/extra/module.h
|
||||
|
||||
The ModuleNode is an abstract class that can be implemented by each type of device.
|
||||
So far we support modules for CUDA, Metal, OpenCL and loading dynamic shared libraries. This abstraction makes introduction
|
||||
of new device easy, and we do not need to redo the host code generation for each type of device.
|
||||
|
||||
Remote Deployment
|
||||
-----------------
|
||||
|
||||
The PackedFunc and Module system also makes it easy to ship the function into remote devices directly.
|
||||
Under the hood, we have an RPCModule that serializes the arguments to do the data movement and launches the computation on the remote.
|
||||
|
||||
.. image:: https://tvm.apache.org/images/release/tvm_rpc.png
|
||||
|
||||
The RPC server itself is minimum and can be bundled into the runtime. We can start a minimum TVM
|
||||
RPC server on iPhone/android/raspberry pi or even the browser. The cross compilation on server and shipping of the module for testing can be done in the same script. Checkout
|
||||
:ref:`tutorial-cross-compilation-and-rpc` for more details.
|
||||
|
||||
|
||||
This instant feedback gives us a lot of advantages. For example, to test the correctness of generated code on iPhone, we no longer have to write test-cases in swift/objective-c from scratch -- We can use RPC to execute on iPhone, copy the result back and do verification on the host via numpy. We can also do the profiling using the same script.
|
||||
|
||||
TVM Object and Compiler Stack
|
||||
-----------------------------
|
||||
|
||||
As we mentioned earlier, we build compiler stack API on top of the PackedFunc runtime system.
|
||||
We faced a constant changing of the compiler API for the need of research. We need a new language object or IR node whenever we want to test out new primitives.
|
||||
However, we don't want to change our API from time to time. Besides that, we also want to
|
||||
|
||||
- be able to serialize any language object and IRs
|
||||
- be able to explore, print, and manipulate the IR objects in front-end language to do quick prototyping.
|
||||
|
||||
We introduced a base class, called `Object`_ to solve this problem.
|
||||
All the language object in the compiler stack is a subclass of ``Object``. Each object contains a string type_key that uniquely identifies
|
||||
the type of object. We choose string instead of int as type key so new ``Object`` class can be added in the decentralized fashion without
|
||||
adding the code back to the central repo. To ease the speed of dispatching, we allocate an integer type_index at runtime for each type_key.
|
||||
|
||||
.. _Object: https://github.com/apache/tvm/blob/main/include/tvm/runtime/object.h
|
||||
|
||||
Since usually one ``Object`` could be referenced in multiple places in the language, we use a shared_ptr to keep
|
||||
track of reference. We use ``ObjectRef`` class to represent a reference to the ``Object``.
|
||||
We can roughly view ``ObjectRef`` class as shared_ptr to the ``Object`` container.
|
||||
We can also define subclass ``ObjectRef`` to hold each subtypes of ``Object``. Each subclass of ``Object`` needs to define the
|
||||
RegisterReflection function.
|
||||
|
||||
|
||||
Each ``Object`` subclass will override this to register its members. Here is an example implementation of IntImmNode.
|
||||
|
||||
.. code:: c
|
||||
|
||||
class IntImmNode : public PrimExprNode {
|
||||
public:
|
||||
/*! \brief the Internal value. */
|
||||
int64_t value;
|
||||
|
||||
static void RegisterReflection() {
|
||||
namespace refl = tvm::ffi::reflection;
|
||||
refl::ObjectDef<IntImmNode>().def_ro("value", &IntImmNode::value);
|
||||
}
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.IntImm", IntImmNode, PrimExprNode);
|
||||
};
|
||||
// in cc file
|
||||
TVM_FFI_STATIC_INIT_BLOCK() { IntImmNode::RegisterReflection(); }
|
||||
|
||||
The RegisterReflection gives us a reflection API to register each member of the object.
|
||||
We can use this function to visit the node and serialize any language object recursively.
|
||||
It also allows us to get members of an object easily in front-end language.
|
||||
For example, we can access the value field of the IntImmNode.
|
||||
|
||||
.. code:: python
|
||||
|
||||
import tvm
|
||||
|
||||
x = tvm.tirx.IntImm("int32", 1)
|
||||
# access the value field of IntImmNode
|
||||
print(x.value)
|
||||
|
||||
New ``Object`` can be added to C++ without changing the front-end runtime, making it easy to make extensions to the compiler stack.
|
||||
Note that this is not the fastest way to expose members to front-end language, but might be one of the simplest
|
||||
approaches possible. We also find that it fits our purposes as we mainly use python for testing and prototyping and still use c++
|
||||
to do the heavy lifting job.
|
||||
|
||||
Implementation Details
|
||||
----------------------
|
||||
|
||||
Each argument in PackedFunc contains a union value `TVMValue`_
|
||||
and a type code. This design allows the dynamically typed language to convert to the corresponding type directly, and statically typed language to
|
||||
do runtime type checking during conversion.
|
||||
|
||||
.. _TVMValue: https://github.com/apache/tvm/blob/main/3rdparty/tvm-ffi/include/tvm/ffi/c_api.h
|
||||
|
||||
The relevant files are
|
||||
|
||||
- `function.h`_ for C++ PackedFunc API
|
||||
- `c_api.h`_ for C API.
|
||||
|
||||
.. _function.h: https://github.com/apache/tvm/blob/main/3rdparty/tvm-ffi/include/tvm/ffi/function.h
|
||||
.. _c_api.h: https://github.com/apache/tvm/blob/main/3rdparty/tvm-ffi/include/tvm/ffi/c_api.h
|
||||
|
||||
To support extension types, we used a registry system to register type related information, like support of any
|
||||
in C++. See the ``tvm-ffi`` subproject under ``3rdparty/tvm-ffi/`` for more details on the FFI type system.
|
||||
|
||||
|
||||
Runtime-Specific Information
|
||||
============================
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:glob:
|
||||
|
||||
runtimes/*
|
||||
@@ -0,0 +1,259 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
.. _tvm-runtime-vulkan:
|
||||
|
||||
Vulkan Runtime
|
||||
==============
|
||||
|
||||
TVM supports using Vulkan compute shaders to execute queries. Each
|
||||
computational kernel is compiled into a SPIR-V shader, which can then
|
||||
be called using the TVM interface.
|
||||
|
||||
.. _tvm-runtime-vulkan-features:
|
||||
|
||||
Vulkan Features, Limits
|
||||
-----------------------
|
||||
|
||||
.. _Required Limits: https://www.khronos.org/registry/vulkan/specs/1.2-extensions/html/vkspec.html#limits-minmax
|
||||
|
||||
Since different Vulkan implementations may enable different optional
|
||||
features or have different physical limits, the code generation must
|
||||
know which features are available to use. These correspond to
|
||||
specific Vulkan capabilities/limits as in
|
||||
:ref:`Vulkan Capabilities Table <tvm-table-vulkan-capabilities>`.
|
||||
If unspecified, TVM assumes that a capability is not available, or
|
||||
that a limit is the minimum guaranteed by the Vulkan spec in the
|
||||
`Required Limits`_ section.
|
||||
|
||||
These parameters can be either explicitly specific when defining a
|
||||
:ref:`Target <tvm-target-specific-target>`, or can be queried from a
|
||||
device. To query from a device, the special parameter
|
||||
``-from_device=N`` can be used to query all vulkan device parameters
|
||||
from device id ``N``. Any additional parameters explicitly specified
|
||||
will override the parameters queried from the device.
|
||||
|
||||
.. _VkSubgroupFeatureFlagBits: https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VkSubgroupFeatureFlagBits.html
|
||||
|
||||
.. list-table:: Vulkan Capabilities
|
||||
:name: tvm-runtime-table-vulkan-capabilities
|
||||
:header-rows: 1
|
||||
|
||||
* - Target Parameter
|
||||
- Required Vulkan Version/Extension
|
||||
- Parameter Queried
|
||||
- Default Value
|
||||
|
||||
* - ``supported_subgroup_operations``
|
||||
- Vulkan 1.1+
|
||||
- ``VkPhysicalDeviceSubgroupProperties::supportedOperations``
|
||||
- 0 (interpreted as `VkSubgroupFeatureFlagBits`_)
|
||||
|
||||
* - ``max_push_constants_size``
|
||||
-
|
||||
- ``VkPhysicalDeviceLimits::maxPushConstantsSize``
|
||||
- 128 bytes
|
||||
|
||||
* - ``max_uniform_buffer_range``
|
||||
-
|
||||
- ``VkPhysicalDeviceLimits::maxUniformBufferRange``
|
||||
- 16384 bytes
|
||||
|
||||
|
||||
* - ``max_storage_buffer_range``
|
||||
-
|
||||
- ``VkPhysicalDeviceLimits::maxStorageBufferRange``
|
||||
- 2\ :sup:`27`\ bytes
|
||||
|
||||
|
||||
* - ``max_per_stage_descriptor_storage_buffer``
|
||||
-
|
||||
- ``VkPhysicalDeviceLimits::maxPerStageDescriptorStorageBuffers``
|
||||
- 4
|
||||
|
||||
|
||||
* - ``supports_storage_buffer_storage_class``
|
||||
- VK_KHR_storage_buffer_storage_class
|
||||
-
|
||||
- false
|
||||
|
||||
|
||||
* - ``supports_storage_buffer_8bit_access``
|
||||
- VK_KHR_8bit_storage
|
||||
- ``VkPhysicalDevice8BitStorageFeaturesKHR::storageBuffer8BitAccess``
|
||||
- false
|
||||
|
||||
|
||||
* - ``supports_storage_buffer_16bit_access``
|
||||
- VK_KHR_16bit_storage
|
||||
- ``VkPhysicalDevice16BitStorageFeaturesKHR::storageBuffer16BitAccess``
|
||||
- false
|
||||
|
||||
|
||||
* - ``supports_float16``
|
||||
- VK_KHR_shader_float16_int8
|
||||
- ``VkPhysicalDeviceShaderFloat16Int8FeaturesKHR::shaderFloat16``
|
||||
- false
|
||||
|
||||
|
||||
* - ``supports_float64``
|
||||
-
|
||||
- ``VkPhysicalDeviceFeatures::shaderFloat64``
|
||||
- false
|
||||
|
||||
|
||||
* - ``supports_int8``
|
||||
- VK_KHR_shader_float16_int8
|
||||
- ``VkPhysicalDeviceShaderFloat16Int8FeaturesKHR::shaderInt8``
|
||||
- false
|
||||
|
||||
|
||||
* - ``supports_int16``
|
||||
-
|
||||
- ``VkPhysicalDeviceFeatures::shaderInt16``
|
||||
- false
|
||||
|
||||
|
||||
* - ``supports_int64``
|
||||
-
|
||||
- ``VkPhysicalDeviceFeatures::shaderInt64``
|
||||
- false
|
||||
|
||||
|
||||
|
||||
As of May 2021, not all Vulkan implementations are supported. For
|
||||
example, support for 64-bit integers is required. If a Vulkan target
|
||||
is not supported, an error message should be issued during SPIR-V code
|
||||
generation. Efforts are also underway to remove these requirements
|
||||
and support additional Vulkan implementations.
|
||||
|
||||
|
||||
.. _tvm-runtime-vulkan-spirv-capabilities:
|
||||
|
||||
SPIR-V Capabilities
|
||||
-------------------
|
||||
|
||||
Some of the device-specific capabilities also correspond to SPIR-V
|
||||
capabilities or extensions that must be declared in the shader, or a
|
||||
minimum SPIR-V version required in order to use a feature. The
|
||||
TVM-generated shaders will declare the minimum set of
|
||||
extensions/capabilities and the minimum allowed version of SPIR-V
|
||||
that are needed to execute the compiled graph.
|
||||
|
||||
If the shader generation requires a capability or extension that is
|
||||
not enabled in the ``Target``, an exception will be raised.
|
||||
|
||||
|
||||
.. list-table:: Vulkan Capabilities
|
||||
:name: tvm-table-vulkan-capabilities
|
||||
:header-rows: 1
|
||||
|
||||
* - Target Parameter
|
||||
- Required SPIR-V Version/Extension
|
||||
- Declared Capability
|
||||
|
||||
* - ``supported_subgroup_operations``
|
||||
- SPIR-V 1.3+
|
||||
- Varies, see `VkSubgroupFeatureFlagBits`_
|
||||
|
||||
* - ``supports_storage_buffer_storage_class``
|
||||
- SPV_KHR_storage_buffer_storage_class
|
||||
-
|
||||
|
||||
* - ``supports_storage_buffer_8bit_access``
|
||||
- SPV_KHR_8bit_storage
|
||||
- StorageBuffer8BitAccess
|
||||
|
||||
* - ``supports_storage_buffer_16bit_access``
|
||||
- SPV_KHR_16bit_storage
|
||||
- StorageBuffer16BitAccess
|
||||
|
||||
* - ``supports_float16``
|
||||
-
|
||||
- Float16
|
||||
|
||||
|
||||
* - ``supports_float64``
|
||||
-
|
||||
- Float64
|
||||
|
||||
|
||||
* - ``supports_int8``
|
||||
-
|
||||
- Int8
|
||||
|
||||
|
||||
* - ``supports_int16``
|
||||
-
|
||||
- Int16
|
||||
|
||||
|
||||
* - ``supports_int64``
|
||||
-
|
||||
- Int64
|
||||
|
||||
|
||||
Vulkan-Specific Environment Variables
|
||||
-------------------------------------
|
||||
|
||||
Both the SPIR-V code generation and the Vulkan runtime have
|
||||
environment variables that can modify some of the runtime behavior.
|
||||
These are intended for debugging purposes, both to more easily test
|
||||
specific code paths, and to output more information as needed. All
|
||||
boolean flags are true if the environment variable is set to a
|
||||
non-zero integer. An unset variable, the integer zero, or an empty
|
||||
string are all false boolean flags.
|
||||
|
||||
.. _VK_KHR_push_descriptor: https://khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VK_KHR_push_descriptor.html
|
||||
|
||||
.. _VK_KHR_descriptor_update_template: https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VK_KHR_descriptor_update_template.html
|
||||
|
||||
.. _VK_KHR_dedicated_allocation: https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VK_KHR_dedicated_allocation.html
|
||||
|
||||
.. _VkMemoryDedicatedRequirements: https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VkMemoryDedicatedRequirements.html
|
||||
|
||||
.. _Vulkan validation layers: https://github.com/KhronosGroup/Vulkan-LoaderAndValidationLayers/blob/master/layers/README.md
|
||||
|
||||
.. _spvValidate: https://github.com/KhronosGroup/SPIRV-Tools#validator
|
||||
|
||||
|
||||
* ``TVM_VULKAN_DISABLE_PUSH_DESCRIPTOR`` - A boolean flag. If true,
|
||||
TVM will explicitly allocate descriptors, and will not use the
|
||||
`VK_KHR_push_descriptor`_ or `VK_KHR_descriptor_update_template`_
|
||||
extensions. If false, TVM will decide whether to use these
|
||||
extensions based on their availability.
|
||||
|
||||
* ``TVM_VULKAN_DISABLE_DEDICATED_ALLOCATION`` - A boolean flag. If
|
||||
true, TVM will not mark memory allocations as being dedicated
|
||||
allocations, and will not use the `VK_KHR_dedicated_allocation`_
|
||||
extension. If false, TVM will decide whether memory allocations
|
||||
should be marked as dedicated based on the
|
||||
`VkMemoryDedicatedRequirements`_ for that buffer.
|
||||
|
||||
* ``TVM_VULKAN_ENABLE_VALIDATION_LAYERS`` - A boolean flag. If true,
|
||||
TVM will enable `Vulkan validation layers`_ that the device
|
||||
supports. If false, no validation layers are enabled.
|
||||
|
||||
* ``TVM_VULKAN_DISABLE_SHADER_VALIDATION`` - A boolean flag. If true,
|
||||
the SPIR-V shader validation done with `spvValidate`_ is skipped.
|
||||
If false (default), all SPIR-V shaders generated by TVM are
|
||||
validated with `spvValidate`_.
|
||||
|
||||
* ``TVM_VULKAN_DEBUG_SHADER_SAVEPATH`` - A path to a directory. If
|
||||
set to a non-empty string, the Vulkan codegen will save TIR, binary
|
||||
SPIR-V, and disassembled SPIR-V shaders to this directory, to be
|
||||
used for debugging purposes.
|
||||
@@ -0,0 +1,575 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
.. _tvmscript-arch:
|
||||
|
||||
TVMScript
|
||||
=========
|
||||
|
||||
TVMScript is a Python-based domain-specific language (DSL) for writing TVM IR. It lets users
|
||||
define ``IRModule``\ s — containing both Relax functions and TIR ``PrimFunc``\ s — using
|
||||
familiar Python syntax. Although TVMScript *looks* like Python, it is **not executed by the
|
||||
Python interpreter**. Instead, Python decorators extract the AST from the source code and
|
||||
transform it into TVM IR through a dedicated parser and IR builder pipeline.
|
||||
|
||||
TVMScript serves two roles in the TVM stack:
|
||||
|
||||
- **Authoring**: users write TIR kernels and Relax programs directly in TVMScript.
|
||||
- **Roundtrip**: every ``IRModule`` can be printed back to TVMScript via ``mod.script()`` and
|
||||
re-parsed to produce an equivalent module. This makes TVMScript the primary tool for
|
||||
inspecting, debugging, and serializing IR.
|
||||
|
||||
|
||||
Overview
|
||||
--------
|
||||
|
||||
The TVMScript system has three components:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
Parsing (Python source → TVM IR):
|
||||
|
||||
Python source (TVMScript)
|
||||
│
|
||||
▼ ast.parse + convert
|
||||
│
|
||||
Doc AST (mirror of Python AST)
|
||||
│
|
||||
▼ Parser (dispatch by token: ir / tirx / relax)
|
||||
│
|
||||
▼ IR Builder (frame stack)
|
||||
│
|
||||
TVM IR (IRModule, PrimFunc, relax.Function)
|
||||
|
||||
|
||||
Printing (TVM IR → Python source):
|
||||
|
||||
TVM IR
|
||||
│
|
||||
▼ IRDocsifier (C++, dispatch by token + type)
|
||||
│
|
||||
Doc tree (ExprDoc, StmtDoc, ...)
|
||||
│
|
||||
▼ DocToPythonScript
|
||||
│
|
||||
TVMScript text
|
||||
|
||||
- **Parser** (Python): reads Python source, converts it to a ``Doc AST`` (a mirror of
|
||||
Python's ``ast`` module), then walks the tree using dialect-specific handlers that call
|
||||
into the IR builder.
|
||||
- **IR Builder** (Python + C++): provides a frame-stack API where each ``with`` block or
|
||||
decorator pushes a frame. When the frame exits, the constructed IR is finalized. The builder
|
||||
is shared across dialects — TIR and Relax each register their own frame types.
|
||||
- **Printer** (C++): converts TVM IR objects to a ``Doc`` tree (an intermediate representation
|
||||
of Python syntax), then formats the tree into valid TVMScript text.
|
||||
|
||||
|
||||
Decorators
|
||||
----------
|
||||
|
||||
TVMScript uses three import aliases by convention:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.script import ir as I # module-level constructs
|
||||
from tvm.script import tirx as T # TIR constructs
|
||||
from tvm.script import relax as R # Relax constructs
|
||||
|
||||
The primary decorators are:
|
||||
|
||||
- ``@I.ir_module``: marks a Python class as an ``IRModule``. Each method inside becomes a
|
||||
function in the module.
|
||||
- ``@T.prim_func``: marks a function as a TIR ``PrimFunc``.
|
||||
- ``@R.function``: marks a function as a ``relax.Function``.
|
||||
|
||||
These can be composed:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@I.ir_module
|
||||
class MyModule:
|
||||
@T.prim_func
|
||||
def add_kernel(A: T.Buffer((128,), "float32"),
|
||||
B: T.Buffer((128,), "float32"),
|
||||
C: T.Buffer((128,), "float32")):
|
||||
for i in range(128):
|
||||
with T.sblock("compute"):
|
||||
vi = T.axis.spatial(128, i)
|
||||
C[vi] = A[vi] + B[vi]
|
||||
|
||||
@R.function
|
||||
def main(x: R.Tensor((128,), "float32"),
|
||||
y: R.Tensor((128,), "float32")) -> R.Tensor((128,), "float32"):
|
||||
with R.dataflow():
|
||||
out = R.call_tir(cls.add_kernel, (x, y),
|
||||
out_ty=R.Tensor((128,), "float32"))
|
||||
R.output(out)
|
||||
return out
|
||||
|
||||
When Python encounters ``@I.ir_module``, the decorator does **not** execute the class body.
|
||||
Instead, it calls ``tvm.script.parse()`` which extracts the source code of the class,
|
||||
builds a Doc AST, and hands it to the parser.
|
||||
|
||||
|
||||
Parser Architecture
|
||||
-------------------
|
||||
|
||||
The parser lives in ``python/tvm/script/parser/``.
|
||||
|
||||
Dispatch mechanism
|
||||
~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Different IR dialects (TIR, Relax) need different handling for the same Python syntax. For
|
||||
example, ``if ... else`` inside ``@T.prim_func`` creates a TIR ``If`` branch, while the same
|
||||
syntax inside ``@R.function`` creates a Relax ``If`` node with different semantics.
|
||||
|
||||
The parser maintains a **dispatch token** stack (``["default"]`` initially). When it encounters
|
||||
a decorated function, it inspects the decorator to determine the token — ``"tirx"`` for
|
||||
``@T.prim_func``, ``"relax"`` for ``@R.function`` — and pushes it onto the stack.
|
||||
|
||||
Each AST node type is dispatched via a virtual table:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
ParseVTable[(token, node_type)] → handler function
|
||||
|
||||
Lookup order:
|
||||
1. (current_token, node_type) e.g. ("tirx", "For")
|
||||
2. ("default", node_type) e.g. ("default", "For")
|
||||
3. generic_visit fallback
|
||||
|
||||
Dialect-specific parsers (``parser/tirx/parser.py``, ``parser/relax/parser.py``) register
|
||||
handlers using ``@dispatch.register(token, type_name)`` decorators.
|
||||
|
||||
Parse flow
|
||||
~~~~~~~~~~
|
||||
|
||||
The entry point is ``parse(program, extra_vars)``:
|
||||
|
||||
1. **Source extraction**: the program's source code is extracted (from a class, function, or
|
||||
string) and converted to a Doc AST via Python's ``ast`` module.
|
||||
|
||||
2. **AST walking**: the ``Parser`` (a subclass of ``doc.NodeVisitor``) walks the Doc AST.
|
||||
For each node, it looks up the handler in the dispatch table.
|
||||
|
||||
3. **Expression evaluation**: expressions like ``T.grid(128, 128)`` are evaluated by the
|
||||
``ExprEvaluator``, which resolves names against the variable table and the ``T.``/``R.``
|
||||
module namespaces.
|
||||
|
||||
4. **Value binding**: assignment statements (``A = T.match_buffer(...)`` in TIR,
|
||||
``lv = R.add(x, y)`` in Relax) go through dialect-specific ``bind_*_value()`` functions
|
||||
that register the resulting TVM objects in the parser's ``VarTable``.
|
||||
|
||||
5. **Scoping**: the ``VarTable`` maintains a stack of frames. Entering a ``with`` block,
|
||||
``for`` loop, or function body pushes a new frame; exiting pops it. This ensures variables
|
||||
are scoped correctly.
|
||||
|
||||
Variable table
|
||||
~~~~~~~~~~~~~~
|
||||
|
||||
The ``VarTable`` is the parser's symbol table:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
VarTable
|
||||
├── frames: [VarTableFrame, ...] ← stack of scopes
|
||||
└── name2value: {str: [Any, ...]} ← name → value stack (for shadowing)
|
||||
|
||||
When a name is looked up, the most recent binding wins. When a frame is popped, all bindings
|
||||
introduced in that frame are removed.
|
||||
|
||||
|
||||
IR Builder Architecture
|
||||
-----------------------
|
||||
|
||||
The IR builder (``python/tvm/script/ir_builder/``, backed by C++ in ``src/script/ir_builder/``)
|
||||
provides a frame-stack API for constructing IR incrementally.
|
||||
|
||||
Frame stack
|
||||
~~~~~~~~~~~
|
||||
|
||||
The core idea: each IR scope (module, function, block, loop) is a **frame**. Frames are pushed
|
||||
on ``__enter__`` and popped on ``__exit__``. When a frame exits, it finalizes the IR it
|
||||
represents and attaches it to the parent frame.
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
IRBuilder (thread-local singleton)
|
||||
└── frame stack:
|
||||
├── IRModuleFrame ← @I.ir_module
|
||||
│ ├── PrimFuncFrame ← @T.prim_func
|
||||
│ │ ├── ForFrame ← T.grid(...) / T.serial(...)
|
||||
│ │ │ └── SBlockFrame ← T.sblock(...)
|
||||
│ │ └── ...
|
||||
│ └── FunctionFrame ← @R.function
|
||||
│ └── BindingBlockFrame ← R.dataflow()
|
||||
└── ...
|
||||
|
||||
This design means the parser never needs to build a complete IR tree in memory — it
|
||||
constructs IR top-down by entering and exiting frames, and each frame handles its own
|
||||
finalization.
|
||||
|
||||
TIR builder
|
||||
~~~~~~~~~~~
|
||||
|
||||
The TIR builder (``ir_builder/tirx/ir.py``) provides functions that map directly to TVMScript
|
||||
syntax. Key categories:
|
||||
|
||||
**Function and block**:
|
||||
|
||||
- ``T.prim_func()`` → ``PrimFuncFrame``
|
||||
- ``T.sblock(name)`` → ``SBlockFrame`` (spatial block)
|
||||
- ``T.init()`` → ``BlockInitFrame`` (reduction initialization)
|
||||
- ``T.reads(...)``, ``T.writes(...)`` → declare buffer access regions
|
||||
|
||||
**Loops**:
|
||||
|
||||
- ``T.grid(*extents)`` → ``ForFrame`` returning loop variables
|
||||
- ``T.serial(start, stop)``, ``T.parallel(...)``, ``T.vectorized(...)``,
|
||||
``T.unroll(...)``, ``T.thread_binding(...)`` → loop with specific iterator type
|
||||
|
||||
**Block axes**:
|
||||
|
||||
- ``T.axis.spatial(dom, binding)`` — spatial iteration axis
|
||||
- ``T.axis.reduce(dom, binding)`` — reduction axis
|
||||
- ``T.axis.remap(kinds, bindings)`` — shorthand for multiple axes
|
||||
|
||||
**Buffers**:
|
||||
|
||||
- ``T.match_buffer(param, shape, dtype)`` — match function parameter to buffer
|
||||
- ``T.alloc_buffer(shape, dtype)`` — allocate intermediate buffer
|
||||
- ``T.Buffer(shape, dtype)`` — buffer type annotation in function signatures
|
||||
|
||||
Relax builder
|
||||
~~~~~~~~~~~~~
|
||||
|
||||
The Relax builder (``ir_builder/relax/ir.py``) provides:
|
||||
|
||||
**Function and dataflow**:
|
||||
|
||||
- ``R.function()`` → ``FunctionFrame``
|
||||
- ``R.dataflow()`` → ``BindingBlockFrame``
|
||||
- ``R.output(*vars)`` → expose variables from a dataflow block
|
||||
|
||||
**Emit**:
|
||||
|
||||
- ``R.emit(value)`` → emit a binding, returns a ``Var``
|
||||
- ``R.emit_match_cast(value, ty)`` → emit with type assertion
|
||||
|
||||
**Type annotations**:
|
||||
|
||||
- ``R.Tensor(shape, dtype)`` — tensor type
|
||||
- ``R.Tuple(*fields)`` — tuple type
|
||||
- ``R.Shape(values)`` — shape type
|
||||
- ``R.Any()`` — any Relax value type
|
||||
|
||||
**Calling conventions**:
|
||||
|
||||
- ``R.call_tir(func, args, out_ty)`` — call a TIR function
|
||||
- ``R.call_packed(name, *args)`` — call a PackedFunc
|
||||
- ``R.call_dps_packed(func, *args)`` — call using destination-passing style
|
||||
|
||||
**Operators**: the ``R`` module also re-exports all Relax operators
|
||||
(``R.add``, ``R.matmul``, ``R.nn.conv2d``, etc.) so they can be used directly in TVMScript.
|
||||
|
||||
|
||||
Printer Architecture
|
||||
--------------------
|
||||
|
||||
The printer converts TVM IR back to TVMScript text. It is implemented primarily in C++
|
||||
(``src/script/printer/``) for performance.
|
||||
|
||||
Doc tree
|
||||
~~~~~~~~
|
||||
|
||||
The printer does **not** generate text directly. Instead, it first builds a ``Doc`` tree — an
|
||||
intermediate representation that mirrors Python syntax:
|
||||
|
||||
- **Expression docs**: ``IdDoc``, ``AttrAccessDoc``, ``CallDoc``, ``IndexDoc``,
|
||||
``OperationDoc``, ``LiteralDoc``, ``TupleDoc``, ``ListDoc``, etc.
|
||||
- **Statement docs**: ``AssignDoc``, ``ForDoc``, ``IfDoc``, ``ScopeDoc`` (``with`` blocks),
|
||||
``FunctionDoc``, ``ClassDoc``, ``ReturnDoc``, ``CommentDoc``, etc.
|
||||
|
||||
For example, ``T.axis.spatial(128, i)`` is represented as:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
CallDoc(
|
||||
callee=AttrAccessDoc(AttrAccessDoc(IdDoc("T"), "axis"), "spatial"),
|
||||
args=[LiteralDoc(128), IdDoc("i")]
|
||||
)
|
||||
|
||||
IRDocsifier
|
||||
~~~~~~~~~~~
|
||||
|
||||
The ``IRDocsifier`` (``include/tvm/script/printer/ir_docsifier.h``) is the main dispatcher.
|
||||
It maintains:
|
||||
|
||||
- A dispatch table mapping ``(token, type_index)`` pairs to converter functions.
|
||||
- A frame stack for tracking the current scope (similar to the builder's frame stack).
|
||||
- A variable-to-name mapping to produce readable names.
|
||||
|
||||
Each IR dialect registers its own converters:
|
||||
|
||||
- ``src/script/printer/tirx/`` — converts PrimFunc, Buffer, SBlock, loops, expressions.
|
||||
- ``src/script/printer/relax/`` — converts relax.Function, bindings, types, operators.
|
||||
- ``src/script/printer/ir/`` — converts IRModule, shared types.
|
||||
|
||||
The final step calls ``DocToPythonScript()`` (``src/script/printer/doc_printer/python_doc_printer.cc``)
|
||||
to format the Doc tree into properly indented Python text.
|
||||
|
||||
Roundtrip guarantee
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
For any ``IRModule`` constructed through the compiler:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
text = mod.script() # IR → TVMScript text
|
||||
reparsed = tvm.script.from_source(text) # text → IR
|
||||
tvm.ir.assert_structural_equal(mod, reparsed)
|
||||
|
||||
This roundtrip property is relied upon by testing infrastructure and serialization workflows.
|
||||
Note that the printed text may differ from hand-written TVMScript — the printer uses canonical
|
||||
forms (e.g., explicit ``R.emit`` calls, fully qualified buffer annotations) that are not required
|
||||
in hand-written code.
|
||||
|
||||
|
||||
Supported Python Syntax
|
||||
-----------------------
|
||||
|
||||
TVMScript supports a subset of Python syntax. The table below summarizes what is supported
|
||||
and how each construct is interpreted:
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 25 15 60
|
||||
|
||||
* - Python Syntax
|
||||
- TIR
|
||||
- Relax
|
||||
* - ``for i in range(n)``
|
||||
- Serial loop nest
|
||||
- Not supported (no Relax-level ``for`` handler)
|
||||
* - ``with T.sblock(...)``
|
||||
- Spatial block scope
|
||||
- N/A
|
||||
* - ``with R.dataflow()``
|
||||
- N/A
|
||||
- Dataflow block
|
||||
* - ``if ... else``
|
||||
- TIR ``If`` branch (PrimExpr condition) or static eval (Python bool)
|
||||
- Relax ``If`` node (plain Python ``if cond:`` syntax)
|
||||
* - ``while``
|
||||
- ``T.While`` loop
|
||||
- Not supported
|
||||
* - ``x = expr``
|
||||
- Variable binding
|
||||
- Emit binding (implicit ``R.emit``)
|
||||
* - ``x: T.Buffer(...)``
|
||||
- Buffer annotation
|
||||
- N/A
|
||||
* - ``x: R.Tensor(...)``
|
||||
- N/A
|
||||
- Struct info annotation
|
||||
* - ``return``
|
||||
- Not used
|
||||
- Function return value
|
||||
* - ``A[i, j]``
|
||||
- Buffer load
|
||||
- Not applicable (use operators)
|
||||
* - ``A[i, j] = expr``
|
||||
- Buffer store
|
||||
- Not applicable
|
||||
* - Arithmetic (``+``, ``-``, etc.)
|
||||
- PrimExpr operations
|
||||
- Calls to Relax operators
|
||||
* - Function calls
|
||||
- ``T.*`` intrinsics
|
||||
- ``R.*`` operators or ``call_tir`` / ``call_packed``
|
||||
|
||||
**Not supported**: ``class`` definitions (except for ``@I.ir_module``), ``try/except``,
|
||||
``yield``, ``async/await``, list comprehensions, ``lambda``, ``import``, and ``global``
|
||||
statements.
|
||||
|
||||
|
||||
TIR Syntax Reference
|
||||
---------------------
|
||||
|
||||
Function definition
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@T.prim_func
|
||||
def func_name(a: T.handle, b: T.handle):
|
||||
A = T.match_buffer(a, (m, n), "float32")
|
||||
B = T.match_buffer(b, (m,), "float32")
|
||||
# function body
|
||||
|
||||
- ``T.handle`` — opaque handle parameter (matched to a buffer inside the function).
|
||||
- ``T.Buffer(shape, dtype)`` — can also be used directly in the signature:
|
||||
``def func(A: T.Buffer((128,), "float32"))``.
|
||||
|
||||
Block and axes
|
||||
~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock("block_name"):
|
||||
vi = T.axis.spatial(128, i)
|
||||
vj = T.axis.reduce(128, j)
|
||||
T.reads(A[vi, vj])
|
||||
T.writes(B[vi])
|
||||
# compute
|
||||
|
||||
- ``T.axis.spatial`` / ``T.axis.reduce`` / ``T.axis.scan`` — declare axis variables with
|
||||
their iteration domain and binding to outer loop variables.
|
||||
- ``T.axis.remap("SR", [i, j])`` — shorthand: ``S`` = spatial, ``R`` = reduce.
|
||||
- ``T.reads(...)``, ``T.writes(...)`` — declare buffer regions accessed by this block.
|
||||
|
||||
Loop types
|
||||
~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
for i in T.serial(0, 128): # sequential
|
||||
for i in T.parallel(0, 128): # parallel
|
||||
for i in T.vectorized(0, 128): # vectorized
|
||||
for i in T.unroll(0, 128): # unrolled
|
||||
for i in T.thread_binding(0, 128, thread="threadIdx.x"): # GPU thread
|
||||
|
||||
Buffer operations
|
||||
~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
C = T.alloc_buffer((128, 128), "float32") # intermediate buffer
|
||||
val = A[i, j] # buffer load
|
||||
B[i] = val + 1.0 # buffer store
|
||||
|
||||
Common intrinsics
|
||||
~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
T.exp(x), T.log(x), T.sqrt(x), T.tanh(x), ... # math functions
|
||||
T.cast(x, "float16") # type cast
|
||||
T.if_then_else(cond, true_val, false_val) # conditional expression
|
||||
T.min(a, b), T.max(a, b) # min/max
|
||||
T.call_extern("func_name", *args) # external function call
|
||||
T.call_packed("func_name", *args) # packed function call
|
||||
T.tvm_storage_sync("shared") # GPU memory fence
|
||||
|
||||
|
||||
Relax Syntax Reference
|
||||
-----------------------
|
||||
|
||||
Function definition
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@R.function
|
||||
def main(x: R.Tensor((128, 128), "float32"),
|
||||
y: R.Tensor((128,), "float32")) -> R.Tensor((128, 128), "float32"):
|
||||
# function body
|
||||
return result
|
||||
|
||||
- ``R.Tensor(shape, dtype)`` — tensor type annotation.
|
||||
- ``R.Tuple(...)``, ``R.Shape(...)``, ``R.Any()`` — other Relax type annotations.
|
||||
- ``R.function(private=True)`` — marks the function as module-private.
|
||||
- ``R.function(pure=False)`` — marks the function as having side effects.
|
||||
|
||||
Dataflow blocks
|
||||
~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
with R.dataflow():
|
||||
lv0 = R.add(x, y)
|
||||
lv1 = R.nn.relu(lv0)
|
||||
R.output(lv1)
|
||||
|
||||
Variables inside a ``R.dataflow()`` block are local to that block. ``R.output(...)`` exposes
|
||||
variables to the outer scope.
|
||||
|
||||
Calling TIR functions
|
||||
~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
out = R.call_tir(cls.my_kernel, (x, y), out_ty=R.Tensor((128,), "float32"))
|
||||
|
||||
- ``cls.my_kernel`` — references a TIR ``PrimFunc`` in the same module.
|
||||
- ``out_ty`` — the type (shape and dtype) of the output tensor.
|
||||
|
||||
Control flow
|
||||
~~~~~~~~~~~~
|
||||
|
||||
Relax ``if`` uses plain Python ``if`` syntax. The condition must be a Relax variable with
|
||||
boolean type. Both branches are required.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@R.function
|
||||
def f(cond: R.Tensor((), "bool"), x: R.Tensor((128,), "float32")):
|
||||
if cond:
|
||||
result = R.add(x, x)
|
||||
else:
|
||||
result = R.multiply(x, x)
|
||||
return result
|
||||
|
||||
|
||||
Source Code Map
|
||||
---------------
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 50 50
|
||||
|
||||
* - Path
|
||||
- Contents
|
||||
* - ``python/tvm/script/parser/core/``
|
||||
- Core parser: dispatch, expression evaluator, variable table, Doc AST
|
||||
* - ``python/tvm/script/parser/tirx/``
|
||||
- TIR-specific parser handlers and value binding
|
||||
* - ``python/tvm/script/parser/relax/``
|
||||
- Relax-specific parser handlers and value binding
|
||||
* - ``python/tvm/script/parser/ir/``
|
||||
- ``@I.ir_module`` entry point and module-level parsing
|
||||
* - ``python/tvm/script/ir_builder/base.py``
|
||||
- IRBuilder base class and frame stack mechanism
|
||||
* - ``python/tvm/script/ir_builder/tirx/``
|
||||
- TIR frame types and builder functions (``T.*``)
|
||||
* - ``python/tvm/script/ir_builder/relax/``
|
||||
- Relax frame types and builder functions (``R.*``)
|
||||
* - ``python/tvm/script/ir_builder/ir/``
|
||||
- IRModule builder (``I.*``)
|
||||
* - ``src/script/printer/``
|
||||
- C++ printer: Doc tree, IRDocsifier, Python code generation
|
||||
* - ``src/script/printer/tirx/``
|
||||
- TIR-specific IR-to-Doc converters
|
||||
* - ``src/script/printer/relax/``
|
||||
- Relax-specific IR-to-Doc converters
|
||||
* - ``src/script/ir_builder/``
|
||||
- C++ backend for frame stack and IR construction
|
||||
* - ``include/tvm/script/printer/``
|
||||
- C++ headers: Doc classes, IRDocsifier, dispatch functor
|
||||
* - ``include/tvm/script/ir_builder/``
|
||||
- C++ headers: builder base, dialect-specific frame types
|
||||
+740
@@ -0,0 +1,740 @@
|
||||
# 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
|
||||
|
||||
#
|
||||
# documentation build configuration file, created by
|
||||
# sphinx-quickstart on Thu Jul 23 19:40:08 2015.
|
||||
#
|
||||
# This file is execfile()d with the current directory set to its
|
||||
# containing dir.
|
||||
#
|
||||
# Note that not all possible configuration values are present in this
|
||||
# autogenerated file.
|
||||
#
|
||||
# All configuration values have a default; values that are commented out
|
||||
# serve to show the default.
|
||||
import gc
|
||||
import inspect
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from functools import partial
|
||||
from hashlib import md5
|
||||
from importlib import import_module
|
||||
from pathlib import Path
|
||||
from textwrap import dedent, indent
|
||||
from unittest.mock import patch
|
||||
|
||||
# If extensions (or modules to document with autodoc) are in another directory,
|
||||
# add these directories to sys.path here. If the directory is relative to the
|
||||
# documentation root, use os.path.abspath to make it absolute, like shown here.
|
||||
curr_path = Path(__file__).expanduser().absolute().parent
|
||||
if curr_path.name == "_staging":
|
||||
# Can't use curr_path.parent, because sphinx_gallery requires a relative path.
|
||||
tvm_path = Path(os.pardir, os.pardir)
|
||||
else:
|
||||
tvm_path = Path(os.pardir)
|
||||
|
||||
sys.path.insert(0, str(tvm_path.resolve() / "python"))
|
||||
sys.path.insert(0, str(tvm_path.resolve() / "docs"))
|
||||
|
||||
# -- General configuration ------------------------------------------------
|
||||
|
||||
# General information about the project.
|
||||
project = "tvm"
|
||||
author = "Apache Software Foundation"
|
||||
copyright = f"2020 - 2026, {author}"
|
||||
github_doc_root = "https://github.com/apache/tvm/tree/main/docs/"
|
||||
|
||||
os.environ["TVM_BUILD_DOC"] = "1"
|
||||
|
||||
|
||||
# Version information.
|
||||
import tvm
|
||||
from tvm import te, testing, topi
|
||||
|
||||
# The version is derived from the Git tag by setuptools_scm at build time and exposed
|
||||
# as tvm.__version__ (see [tool.setuptools_scm] in pyproject.toml).
|
||||
version = tvm.__version__
|
||||
release = version
|
||||
|
||||
|
||||
def monkey_patch(module_name, func_name):
|
||||
"""Helper function for monkey-patching library functions.
|
||||
|
||||
Used to modify a few sphinx-gallery behaviors to make the "Open in Colab" button work correctly.
|
||||
Should be called as a decorator with arguments. Note this behaves differently from unittest's
|
||||
@mock.patch, as our monkey_patch decorator should be placed on the new version of the function.
|
||||
"""
|
||||
module = import_module(module_name)
|
||||
original_func = getattr(module, func_name)
|
||||
|
||||
def decorator(function):
|
||||
updated_func = partial(function, real_func=original_func)
|
||||
setattr(module, func_name, updated_func)
|
||||
return updated_func
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
CURRENT_FILE_CONF = None
|
||||
|
||||
|
||||
@monkey_patch("sphinx_gallery.py_source_parser", "split_code_and_text_blocks")
|
||||
def split_code_and_text_blocks(source_file, return_node, real_func):
|
||||
"""Monkey-patch split_code_and_text_blocks to expose sphinx-gallery's file-level config.
|
||||
|
||||
It's kinda gross, but we need access to file_conf to detect the requires_cuda flag.
|
||||
"""
|
||||
global CURRENT_FILE_CONF
|
||||
file_conf, blocks, node = real_func(source_file, return_node)
|
||||
CURRENT_FILE_CONF = file_conf
|
||||
return (file_conf, blocks, node)
|
||||
|
||||
|
||||
# This header replaces the default sphinx-gallery one in sphinx_gallery/gen_rst.py.
|
||||
# Colab button has been temporarily disabled due to prebuilt packages unavailability.
|
||||
COLAB_HTML_HEADER = """
|
||||
.. DO NOT EDIT. THIS FILE WAS AUTOMATICALLY GENERATED BY
|
||||
.. TVM'S MONKEY-PATCHED VERSION OF SPHINX-GALLERY. TO MAKE
|
||||
.. CHANGES, EDIT THE SOURCE PYTHON FILE:
|
||||
.. "{python_file}"
|
||||
|
||||
.. only:: html
|
||||
|
||||
.. note::
|
||||
:class: sphx-glr-download-link-note
|
||||
|
||||
You can click :ref:`here <sphx_glr_download_{ref_name}>` to run the Jupyter notebook locally.
|
||||
|
||||
.. rst-class:: sphx-glr-example-title
|
||||
|
||||
.. _sphx_glr_{ref_name}:
|
||||
|
||||
"""
|
||||
|
||||
# Google Colab allows opening .ipynb files on GitHub by appending a GitHub path to this base URL.
|
||||
COLAB_URL_BASE = "https://colab.research.google.com/github"
|
||||
|
||||
# The GitHub path where the site is automatically deployed by tvm-bot.
|
||||
IPYTHON_GITHUB_BASE = "apache/tvm-site/blob/asf-site/docs/_downloads/"
|
||||
|
||||
# The SVG image of the "Open in Colab" button.
|
||||
BUTTON = (
|
||||
"https://raw.githubusercontent.com/tlc-pack/web-data/main/images/utilities/colab_button.svg"
|
||||
)
|
||||
|
||||
|
||||
@monkey_patch("sphinx_gallery.gen_rst", "save_rst_example")
|
||||
def save_rst_example(
|
||||
example_rst, example_file, time_elapsed, memory_used, gallery_conf, language, real_func
|
||||
):
|
||||
"""Monkey-patch save_rst_example to customize the tutorial header.
|
||||
|
||||
Note: Colab button has been temporarily disabled. The colab_url and button_svg
|
||||
are still generated but not used in the header template.
|
||||
"""
|
||||
|
||||
# The url is the md5 hash of the notebook path.
|
||||
example_fname = os.path.relpath(example_file, gallery_conf["src_dir"])
|
||||
ref_fname = example_fname.replace(os.path.sep, "_")
|
||||
notebook_path = example_fname[:-2] + "ipynb"
|
||||
digest = md5(notebook_path.encode()).hexdigest()
|
||||
|
||||
# Fixed documentation versions must link to different (earlier) .ipynb notebooks.
|
||||
# Note: colab_url is generated but not currently used in the header template.
|
||||
colab_url = f"{COLAB_URL_BASE}/{IPYTHON_GITHUB_BASE}"
|
||||
if "dev" not in version:
|
||||
colab_url += version + "/"
|
||||
colab_url += digest + "/" + os.path.basename(notebook_path)
|
||||
|
||||
new_header = COLAB_HTML_HEADER.format(
|
||||
python_file=example_fname, ref_name=ref_fname, colab_url=colab_url, button_svg=BUTTON
|
||||
)
|
||||
with patch("sphinx_gallery.gen_rst.EXAMPLE_HEADER", new_header):
|
||||
real_func(
|
||||
example_rst, example_file, time_elapsed, memory_used, gallery_conf, language=language
|
||||
)
|
||||
|
||||
|
||||
INCLUDE_DIRECTIVE_RE = re.compile(r"^([ \t]*)\.\. include::\s*(.+)\n", flags=re.M)
|
||||
COMMENT_DIRECTIVE_RE = re.compile(r"^\.\.(?: .*)?\n(?:(?: .*)?\n)*", flags=re.M)
|
||||
ADMONITION_DIRECTIVE_RE = re.compile(r"^\.\. admonition:: *(.*)\n((?:(?: .*)?\n)*)\n", flags=re.M)
|
||||
|
||||
|
||||
@monkey_patch("sphinx_gallery.notebook", "rst2md")
|
||||
def rst2md(text, gallery_conf, target_dir, heading_levels, real_func):
|
||||
"""Monkey-patch rst2md to support comments and some include directives.
|
||||
|
||||
Currently, only include directives without any parameters are supported. Also, note that in
|
||||
reStructuredText any unrecognized explicit markup block is treated as a comment (see
|
||||
https://docutils.sourceforge.io/docs/ref/rst/restructuredtext.html#comments).
|
||||
|
||||
For callouts, we only replace generic "admonition" directives. All others should be replaced by
|
||||
sphinx-gallery's rst2md. Note that the "alert" and "alert-info" tags are support in most IPython
|
||||
notebooks, but they render kinda funky on Colab.
|
||||
"""
|
||||
|
||||
def load_include(match):
|
||||
full_path = os.path.join(target_dir, match.group(2))
|
||||
with open(full_path) as f:
|
||||
lines = f.read()
|
||||
indented = indent(lines, match.group(1)) + "\n"
|
||||
return indented
|
||||
|
||||
text = re.sub(INCLUDE_DIRECTIVE_RE, load_include, text)
|
||||
|
||||
# Replace generic, titled admonitions with indented text. Other admonitions (e.g. .. note::)
|
||||
# will be handled by sphinx-gallery's
|
||||
def rewrite_generic_admonition(match):
|
||||
title, text = match.groups()
|
||||
stripped_text = dedent(text).strip()
|
||||
return f'<div class="alert alert-info"><h4>{title}</h4><p>{stripped_text}</p></div>'
|
||||
|
||||
text = re.sub(ADMONITION_DIRECTIVE_RE, rewrite_generic_admonition, text)
|
||||
|
||||
# Call the real function, and then strip any remaining directives (i.e. comments)
|
||||
text = real_func(text, gallery_conf, target_dir, heading_levels)
|
||||
text = re.sub(COMMENT_DIRECTIVE_RE, "", text)
|
||||
return text
|
||||
|
||||
|
||||
def install_request_hook(gallery_conf, fname):
|
||||
testing.utils.install_request_hook(tvm_path.resolve() / "tests" / "python" / "request_hook.py")
|
||||
|
||||
|
||||
INSTALL_TVM_DEV = """\
|
||||
%%shell
|
||||
# Installs the latest dev build of TVM from PyPI. If you wish to build
|
||||
# from source, see https://tvm.apache.org/docs/install/from_source.html
|
||||
pip install apache-tvm --pre"""
|
||||
|
||||
INSTALL_TVM_FIXED = f"""\
|
||||
%%shell
|
||||
# Installs TVM version {version} from PyPI. If you wish to build
|
||||
# from source, see https://tvm.apache.org/docs/install/from_source.html
|
||||
pip install apache-tvm=={version}"""
|
||||
|
||||
INSTALL_TVM_CUDA_DEV = """\
|
||||
%%markdown
|
||||
> **Note:** This tutorial requires a CUDA-enabled build of TVM.
|
||||
> Pre-built CUDA wheels are not currently available on PyPI.
|
||||
> Please build TVM from source with CUDA enabled before running this notebook:
|
||||
> https://tvm.apache.org/docs/install/from_source.html"""
|
||||
|
||||
INSTALL_TVM_CUDA_FIXED = f"""\
|
||||
%%markdown
|
||||
> **Note:** This tutorial requires a CUDA-enabled build of TVM (version {version}).
|
||||
> Pre-built CUDA wheels are not currently available on PyPI.
|
||||
> Please build TVM from source with CUDA enabled before running this notebook:
|
||||
> https://tvm.apache.org/docs/install/from_source.html"""
|
||||
|
||||
|
||||
@monkey_patch("sphinx_gallery.gen_rst", "jupyter_notebook")
|
||||
def jupyter_notebook(script_blocks, gallery_conf, target_dir, real_func):
|
||||
"""Monkey-patch sphinx-gallery to add a TVM import block to each IPython notebook.
|
||||
|
||||
If we had only one import block, we could skip the patching and just set first_notebook_cell.
|
||||
However, how we import TVM depends on if we are using a fixed or dev version, and whether we
|
||||
will use the GPU.
|
||||
|
||||
Tutorials requiring a CUDA-enabled build of TVM should use the flag:
|
||||
# sphinx_gallery_requires_cuda = True
|
||||
"""
|
||||
|
||||
requires_cuda = CURRENT_FILE_CONF.get("requires_cuda", False)
|
||||
fixed_version = "dev" not in version
|
||||
|
||||
if fixed_version and requires_cuda:
|
||||
install_block = INSTALL_TVM_CUDA_FIXED
|
||||
elif fixed_version and not requires_cuda:
|
||||
install_block = INSTALL_TVM_FIXED
|
||||
elif not fixed_version and requires_cuda:
|
||||
install_block = INSTALL_TVM_CUDA_DEV
|
||||
else:
|
||||
install_block = INSTALL_TVM_DEV
|
||||
|
||||
new_conf = {**gallery_conf, "first_notebook_cell": install_block}
|
||||
return real_func(script_blocks, new_conf, target_dir)
|
||||
|
||||
|
||||
# Add any Sphinx extension module names here, as strings. They can be
|
||||
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom ones
|
||||
extensions = [
|
||||
"sphinx.ext.autodoc",
|
||||
"sphinx.ext.autosummary",
|
||||
"sphinx.ext.intersphinx",
|
||||
"sphinx.ext.napoleon",
|
||||
"sphinx.ext.mathjax",
|
||||
"sphinx_gallery.gen_gallery",
|
||||
"autodocsumm",
|
||||
]
|
||||
|
||||
# Add any paths that contain templates here, relative to this directory.
|
||||
templates_path = ["_templates"]
|
||||
|
||||
# The suffix(es) of source filenames.
|
||||
# You can specify multiple suffix as a list of string:
|
||||
# source_suffix = ['.rst', '.md']
|
||||
source_suffix = [".rst", ".md"]
|
||||
|
||||
# The encoding of source files.
|
||||
# source_encoding = 'utf-8-sig'
|
||||
|
||||
# generate autosummary even if no references
|
||||
autosummary_generate = True
|
||||
|
||||
# The main toctree document.
|
||||
main_doc = "index"
|
||||
|
||||
# The language for content autogenerated by Sphinx. Refer to documentation
|
||||
# for a list of supported languages.
|
||||
#
|
||||
# This is also used if you do content translation via gettext catalogs.
|
||||
# Usually you set "language" from the command line for these cases.
|
||||
language = None
|
||||
|
||||
# There are two options for replacing |today|: either, you set today to some
|
||||
# non-false value, then it is used:
|
||||
# today = ''
|
||||
# Else, today_fmt is used as the format for a strftime call.
|
||||
# today_fmt = '%B %d, %Y'
|
||||
|
||||
# List of patterns, relative to source directory, that match files and
|
||||
# directories to ignore when looking for source files.
|
||||
exclude_patterns = ["_build", "_staging"]
|
||||
|
||||
# The TIRx API pages autodoc modules (tvm.tirx, tvm.backend.cuda) that re-export
|
||||
# common IR types (PrimExpr, Op, ...) from several modules, which makes a handful
|
||||
# of autodoc'd cross references ambiguous. Silence that specific, benign category.
|
||||
suppress_warnings = ["ref.python"]
|
||||
|
||||
# The reST default role (used for this markup: `text`) to use for all
|
||||
# documents.
|
||||
# default_role = None
|
||||
|
||||
# If true, '()' will be appended to :func: etc. cross-reference text.
|
||||
# add_function_parentheses = True
|
||||
|
||||
# If true, the current module name will be prepended to all description
|
||||
# unit titles (such as .. function::).
|
||||
# add_module_names = True
|
||||
|
||||
# If true, sectionauthor and moduleauthor directives will be shown in the
|
||||
# output. They are ignored by default.
|
||||
# show_authors = False
|
||||
|
||||
# The name of the Pygments (syntax highlighting) style to use.
|
||||
pygments_style = "sphinx"
|
||||
|
||||
# A list of ignored prefixes for module index sorting.
|
||||
# modindex_common_prefix = []
|
||||
|
||||
# If true, keep warnings as "system message" paragraphs in the built documents.
|
||||
# keep_warnings = False
|
||||
|
||||
# If true, `todo` and `todoList` produce output, else they produce nothing.
|
||||
todo_include_todos = False
|
||||
|
||||
# -- Options for HTML output ----------------------------------------------
|
||||
|
||||
html_theme = "sphinx_book_theme"
|
||||
|
||||
html_title = "Apache TVM"
|
||||
|
||||
# Add any paths that contain custom static files (such as style sheets) here,
|
||||
# relative to this directory. They are copied after the builtin static files,
|
||||
# so a file named "default.css" will overwrite the builtin "default.css".
|
||||
html_static_path = ["_static"]
|
||||
|
||||
html_logo = "_static/img/tvm-logo-small.png"
|
||||
|
||||
html_favicon = "_static/img/tvm-logo-square.png"
|
||||
|
||||
# The Apache trademark/copyright footer is rendered through sphinx-book-theme's
|
||||
# ``extra_footer`` hook (see footer_html below). This mirrors how the tvm-ffi docs
|
||||
# (3rdparty/tvm-ffi/docs/conf.py) preserve the ASF menu under the book theme.
|
||||
footer_dropdown = {
|
||||
"name": "ASF",
|
||||
"items": [
|
||||
("Apache Homepage", "https://apache.org/"),
|
||||
("License", "https://www.apache.org/licenses/"),
|
||||
("Sponsorship", "https://www.apache.org/foundation/sponsorship.html"),
|
||||
("Security", "https://tvm.apache.org/docs/reference/security.html"),
|
||||
("Thanks", "https://www.apache.org/foundation/thanks.html"),
|
||||
("Events", "https://www.apache.org/events/current-event"),
|
||||
],
|
||||
}
|
||||
|
||||
footer_note = " ".join(
|
||||
"""
|
||||
Copyright © 2026 The Apache Software Foundation. Apache TVM, Apache, the Apache feather,
|
||||
and the Apache TVM project logo are either trademarks or registered trademarks of
|
||||
the Apache Software Foundation.""".split("\n")
|
||||
).strip()
|
||||
|
||||
|
||||
def footer_html() -> str:
|
||||
"""Build the extra footer: ASF dropdown and the Apache trademark note.
|
||||
|
||||
The copyright line is rendered natively by sphinx-book-theme (from the ``copyright``
|
||||
config value), so it is intentionally not repeated here.
|
||||
"""
|
||||
dropdown_items = ""
|
||||
for item_name, item_url in footer_dropdown["items"]:
|
||||
dropdown_items += f'<li><a class="dropdown-item" href="{item_url}" target="_blank" style="font-size: 0.9em;">{item_name}</a></li>\n'
|
||||
|
||||
return f"""
|
||||
<div class="footer-container" style="margin: 5px 0; font-size: 0.9em; color: #6c757d; text-align: right;">
|
||||
<div class="footer-line1" style="display: flex; justify-content: flex-end; align-items: center; gap: 0.9em; margin-bottom: 3px; flex-wrap: wrap;">
|
||||
<div class="footer-dropdown">
|
||||
<div class="dropdown">
|
||||
<button class="btn btn-link dropdown-toggle" type="button" id="footerDropdown" data-bs-toggle="dropdown"
|
||||
aria-expanded="false" style="font-size: 0.9em; color: #6c757d; text-decoration: none; padding: 0; border: none; background: none;">
|
||||
{footer_dropdown["name"]}
|
||||
</button>
|
||||
<ul class="dropdown-menu dropdown-menu-end" aria-labelledby="footerDropdown" style="font-size: 0.9em;">
|
||||
{dropdown_items} </ul>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="footer-line2" style="font-size: 0.9em; color: #6c757d; text-align: right;">
|
||||
{footer_note}
|
||||
</div>
|
||||
</div>
|
||||
"""
|
||||
|
||||
|
||||
html_theme_options = {
|
||||
"repository_url": "https://github.com/apache/tvm",
|
||||
"repository_branch": "main",
|
||||
"path_to_docs": "docs/",
|
||||
"use_repository_button": True,
|
||||
"use_edit_page_button": True,
|
||||
"use_source_button": True,
|
||||
"use_issues_button": True,
|
||||
"show_toc_level": 2,
|
||||
"show_navbar_depth": 1,
|
||||
"icon_links": [
|
||||
{
|
||||
"name": "TVM Homepage",
|
||||
"url": "https://tvm.apache.org/",
|
||||
"icon": "fa-solid fa-house",
|
||||
"type": "fontawesome",
|
||||
},
|
||||
{
|
||||
"name": "Community",
|
||||
"url": "https://tvm.apache.org/community",
|
||||
"icon": "fa-solid fa-users",
|
||||
"type": "fontawesome",
|
||||
},
|
||||
{
|
||||
"name": "Download",
|
||||
"url": "https://tvm.apache.org/download",
|
||||
"icon": "fa-solid fa-box-open",
|
||||
"type": "fontawesome",
|
||||
},
|
||||
],
|
||||
"extra_footer": footer_html(),
|
||||
}
|
||||
|
||||
|
||||
# Output file base name for HTML help builder.
|
||||
htmlhelp_basename = project + "doc"
|
||||
|
||||
# -- Options for LaTeX output ---------------------------------------------
|
||||
latex_elements = {}
|
||||
|
||||
# Grouping the document tree into LaTeX files. List of tuples
|
||||
# (source start file, target name, title,
|
||||
# author, documentclass [howto, manual, or own class]).
|
||||
latex_documents = [
|
||||
(main_doc, f"{project}.tex", project, author, "manual"),
|
||||
]
|
||||
|
||||
intersphinx_mapping = {
|
||||
"python": (f"https://docs.python.org/{sys.version_info.major}", None),
|
||||
# "numpy": ("https://numpy.org/doc/stable", None),
|
||||
# "scipy": ("https://docs.scipy.org/doc/scipy", None),
|
||||
# "matplotlib": ("https://matplotlib.org/", None),
|
||||
}
|
||||
|
||||
from sphinx_gallery.sorting import ExplicitOrder
|
||||
|
||||
examples_dirs = [
|
||||
tvm_path.joinpath("docs", "get_started", "tutorials"),
|
||||
tvm_path.joinpath("docs", "how_to", "tutorials"),
|
||||
tvm_path.joinpath("docs", "deep_dive", "relax", "tutorials"),
|
||||
tvm_path.joinpath("docs", "deep_dive", "tensor_ir", "tutorials"),
|
||||
]
|
||||
|
||||
gallery_dirs = [
|
||||
"get_started/tutorials/",
|
||||
"how_to/tutorials/",
|
||||
"deep_dive/relax/tutorials/",
|
||||
"deep_dive/tensor_ir/tutorials/",
|
||||
]
|
||||
|
||||
# Explicitly define the order within a subsection.
|
||||
# The listed files are sorted according to the list.
|
||||
# The unlisted files are sorted by filenames.
|
||||
# The unlisted files always appear after listed files.
|
||||
within_subsection_order = {}
|
||||
|
||||
|
||||
class WithinSubsectionOrder:
|
||||
def __init__(self, src_dir):
|
||||
self.src_dir = src_dir.split("/")[-1]
|
||||
|
||||
def __call__(self, filename):
|
||||
# If the order is provided, use the provided order
|
||||
if (
|
||||
self.src_dir in within_subsection_order
|
||||
and filename in within_subsection_order[self.src_dir]
|
||||
):
|
||||
index = within_subsection_order[self.src_dir].index(filename)
|
||||
assert index < 1e10
|
||||
return f"\0{index:010d}"
|
||||
|
||||
# Otherwise, sort by filename
|
||||
return filename
|
||||
|
||||
|
||||
# When running the tutorials on GPUs we are dependent on the Python garbage collector
|
||||
# collecting TVM packed function closures for any device memory to also be released. This
|
||||
# is not a good setup for machines with lots of CPU ram but constrained GPU ram, so force
|
||||
# a gc after each example.
|
||||
def force_gc(gallery_conf, fname):
|
||||
gc.collect()
|
||||
|
||||
|
||||
filename_pattern_default = ".*"
|
||||
|
||||
sphinx_gallery_conf = {
|
||||
"backreferences_dir": "gen_modules/backreferences",
|
||||
"doc_module": ("tvm", "numpy"),
|
||||
"reference_url": {
|
||||
"tvm": None,
|
||||
# "matplotlib": "https://matplotlib.org/",
|
||||
# "numpy": "https://numpy.org/doc/stable",
|
||||
},
|
||||
"examples_dirs": examples_dirs,
|
||||
"within_subsection_order": WithinSubsectionOrder,
|
||||
"gallery_dirs": gallery_dirs,
|
||||
"filename_pattern": os.environ.get("TVM_TUTORIAL_EXEC_PATTERN", filename_pattern_default),
|
||||
"download_all_examples": False,
|
||||
"min_reported_time": 60,
|
||||
"expected_failing_examples": [],
|
||||
"reset_modules": ("matplotlib", "seaborn", force_gc, install_request_hook),
|
||||
"promote_jupyter_magic": True,
|
||||
# Drop the "Gallery generated by Sphinx-Gallery" signature line on generated pages.
|
||||
"show_signature": False,
|
||||
}
|
||||
|
||||
autodoc_default_options = {
|
||||
"member-order": "bysource",
|
||||
}
|
||||
|
||||
# Maps the original namespace to list of potential modules
|
||||
# that we can import alias from.
|
||||
tvm_alias_check_map = {
|
||||
"tvm.te": ["tvm.tirx"],
|
||||
"tvm.tirx": ["tvm.ir", "tvm.runtime"],
|
||||
}
|
||||
|
||||
|
||||
def update_alias_docstring(name, obj, lines):
|
||||
"""Update the docstring of alias functions.
|
||||
|
||||
This function checks if the obj is an alias of another documented object
|
||||
in a different module.
|
||||
|
||||
If it is an alias, then it will append the alias information to the docstring.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
The full name of the object in the doc.
|
||||
|
||||
obj : object
|
||||
The original object.
|
||||
|
||||
lines : list
|
||||
The docstring lines, need to be modified inplace.
|
||||
"""
|
||||
arr = name.rsplit(".", 1)
|
||||
if len(arr) != 2:
|
||||
return
|
||||
target_mod, target_name = arr
|
||||
|
||||
if target_mod not in tvm_alias_check_map:
|
||||
return
|
||||
if not hasattr(obj, "__module__"):
|
||||
return
|
||||
obj_mod = obj.__module__
|
||||
|
||||
for amod in tvm_alias_check_map[target_mod]:
|
||||
if not obj_mod.startswith(amod):
|
||||
continue
|
||||
|
||||
if hasattr(sys.modules[amod], target_name):
|
||||
obj_type = ":py:func" if callable(obj) else ":py:class"
|
||||
lines.append(f".. rubric:: Alias of {obj_type}:`{amod}.{target_name}`")
|
||||
|
||||
|
||||
tvm_class_name_rewrite_map = {
|
||||
"tvm.tirx": ["Var", "Call"],
|
||||
"tvm.relax": ["Var", "Call", "StringImm"],
|
||||
"tvm.relax.frontend.nn": ["Module"],
|
||||
}
|
||||
|
||||
# When documenting modules under these prefixes, prefer types from the mapped module
|
||||
# to resolve ambiguous cross-references (e.g. Var exists in both tvm.tirx and tvm.relax).
|
||||
tvm_module_type_preference = {
|
||||
"tvm.s_tir": "tvm.tirx",
|
||||
}
|
||||
|
||||
|
||||
def distinguish_class_name(name: str, lines: list[str]):
|
||||
"""Distinguish the docstring of type annotations.
|
||||
|
||||
In the whole TVM, there are many classes with the same name but in different modules,
|
||||
e.g. ``tirx.Var``, ``relax.Var``. This function is used to distinguish them in the docstring,
|
||||
by adding the module name as prefix.
|
||||
|
||||
To be specific, this function will check the current object name, and if it in the specific
|
||||
module with specific name, it will add the module name as prefix to the class name to prevent
|
||||
the confusion. Further, we only add the prefix to those standalone class name, but skip
|
||||
the pattern of `xx.Var`, `Var.xx` and `xx.Var.xx`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
The full name of the object in the doc.
|
||||
|
||||
lines : list
|
||||
The docstring lines, need to be modified inplace.
|
||||
"""
|
||||
remap = {}
|
||||
for module_name in tvm_class_name_rewrite_map:
|
||||
if name.startswith(module_name):
|
||||
short_name = module_name[4:] if module_name.startswith("tvm.") else module_name
|
||||
for class_name in tvm_class_name_rewrite_map[module_name]:
|
||||
remap.update({class_name: f"{short_name}.{class_name}"})
|
||||
|
||||
for k, v in remap.items():
|
||||
for i in range(len(lines)):
|
||||
lines[i] = re.sub(rf"(?<!\.)\b{k}\b(?!\.)", v, lines[i])
|
||||
|
||||
|
||||
def process_docstring(app, what, name, obj, options, lines):
|
||||
"""Sphinx callback to process docstring"""
|
||||
if callable(obj) or inspect.isclass(obj):
|
||||
update_alias_docstring(name, obj, lines)
|
||||
distinguish_class_name(name, lines)
|
||||
|
||||
|
||||
def strip_ipython_magic(app, docname, source):
|
||||
"""Prevents IPython magic commands from being rendered in HTML files.
|
||||
|
||||
TODO rework this function to remove IPython magic commands from include directives too.
|
||||
"""
|
||||
for i in range(len(source)):
|
||||
source[i] = re.sub(r"%%.*\n\s*", "", source[i])
|
||||
|
||||
|
||||
def _patch_python_domain_find_obj():
|
||||
"""Patch PythonDomain.find_obj to resolve ambiguous cross-references.
|
||||
|
||||
Sphinx's ``warn-missing-reference`` event is only fired for unresolved
|
||||
references. Ambiguous short names such as ``StringImm`` already have
|
||||
multiple matches at ``PythonDomain.find_obj`` time, so the disambiguation
|
||||
needs to happen here instead.
|
||||
"""
|
||||
from sphinx.domains.python import PythonDomain
|
||||
|
||||
if getattr(PythonDomain.find_obj, "_tvm_patched", False):
|
||||
return
|
||||
|
||||
_original_find_obj = PythonDomain.find_obj
|
||||
|
||||
def _common_prefix_len(lhs: str, rhs: str) -> int:
|
||||
count = 0
|
||||
for lpart, rpart in zip(lhs.split("."), rhs.split(".")):
|
||||
if lpart != rpart:
|
||||
break
|
||||
count += 1
|
||||
return count
|
||||
|
||||
def _dedup_find_obj(self, env, modname, classname, name, objtype, searchmode=0):
|
||||
matches = _original_find_obj(self, env, modname, classname, name, objtype, searchmode)
|
||||
if len(matches) <= 1:
|
||||
return matches
|
||||
|
||||
short_name = name.rsplit(".", 1)[-1]
|
||||
|
||||
# Prefer a single canonical (non-aliased) entry if Sphinx already found one.
|
||||
canonical_matches = [match for match in matches if not match[1].aliased]
|
||||
if len(canonical_matches) == 1:
|
||||
return canonical_matches
|
||||
|
||||
# Use TVM's module context for the known short names we rewrite in docstrings.
|
||||
if modname:
|
||||
candidate_modules = sorted(
|
||||
(
|
||||
module_name
|
||||
for module_name, class_names in tvm_class_name_rewrite_map.items()
|
||||
if short_name in class_names and modname.startswith(module_name)
|
||||
),
|
||||
key=len,
|
||||
reverse=True,
|
||||
)
|
||||
for module_name in candidate_modules:
|
||||
target_name = f"{module_name}.{short_name}"
|
||||
context_matches = [match for match in matches if match[0] == target_name]
|
||||
if len(context_matches) == 1:
|
||||
return context_matches
|
||||
|
||||
# Fall back to the unique match that best shares the current module prefix.
|
||||
match_scores = {match[0]: _common_prefix_len(modname, match[0]) for match in matches}
|
||||
best_score = max(match_scores.values())
|
||||
if best_score > 1:
|
||||
best_matches = [match for match in matches if match_scores[match[0]] == best_score]
|
||||
if len(best_matches) == 1:
|
||||
return best_matches
|
||||
|
||||
# Check module type preference for cross-module resolution
|
||||
# (e.g. tvm.s_tir.analysis uses types from tvm.tirx).
|
||||
for prefix, preferred_mod in tvm_module_type_preference.items():
|
||||
if modname.startswith(prefix):
|
||||
preferred = [m for m in matches if m[0].startswith(preferred_mod + ".")]
|
||||
if len(preferred) >= 1:
|
||||
return preferred[:1]
|
||||
|
||||
return matches
|
||||
|
||||
_dedup_find_obj._tvm_patched = True
|
||||
PythonDomain.find_obj = _dedup_find_obj
|
||||
|
||||
|
||||
_patch_python_domain_find_obj()
|
||||
|
||||
|
||||
def setup(app):
|
||||
app.connect("source-read", strip_ipython_magic)
|
||||
app.connect("autodoc-process-docstring", process_docstring)
|
||||
@@ -0,0 +1,249 @@
|
||||
.. 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.
|
||||
|
||||
.. _ci_guide:
|
||||
|
||||
Using TVM's CI
|
||||
==============
|
||||
|
||||
.. contents::
|
||||
:local:
|
||||
|
||||
TVM uses a combination of Jenkins and GitHub Actions for continuous integration (CI).
|
||||
|
||||
- **Jenkins** runs Linux CI tests on
|
||||
`branches <https://ci.tlcpack.ai/job/tvm/>`_ and
|
||||
`pull requests <https://ci.tlcpack.ai/job/tvm/view/change-requests/>`_ through a
|
||||
build configuration specified in `Jenkinsfile templates <https://github.com/apache/tvm/blob/main/ci/jenkins/templates/>`_.
|
||||
Jenkins is the primary CI step that is codified to block merging.
|
||||
- **GitHub Actions** runs `linting <https://github.com/apache/tvm/blob/main/.github/workflows/lint.yml>`_
|
||||
(via pre-commit hooks) on pushes and pull requests, as well as minimal
|
||||
Windows and macOS build-and-test jobs.
|
||||
|
||||
This page describes how contributors and committers can use TVM's CI to verify their code. You can
|
||||
read more about the design of TVM CI in the `tlc-pack/ci <https://github.com/tlc-pack/ci>`_ repo.
|
||||
|
||||
For Contributors
|
||||
----------------
|
||||
|
||||
A standard Jenkins CI run looks something like this viewed in `Jenkins' BlueOcean viewer <https://ci.tlcpack.ai/blue/organizations/jenkins/tvm/activity>`_.
|
||||
CI runs usually take a couple hours to complete and pull requests (PRs) cannot be merged before CI
|
||||
has successfully completed. To diagnose failing steps, click through to the failing
|
||||
pipeline stage then to the failing step to see the output logs. For GitHub Actions jobs (lint,
|
||||
Windows, macOS), check the "Actions" tab on the pull request or repository page.
|
||||
|
||||
.. image:: https://github.com/tlc-pack/web-data/raw/main/images/contribute/ci.png
|
||||
:width: 800
|
||||
:alt: The Jenkins UI for a CI run
|
||||
|
||||
|
||||
Debugging Failures
|
||||
^^^^^^^^^^^^^^^^^^
|
||||
|
||||
When CI fails for some reason, there are several methods to diagnose the issue.
|
||||
|
||||
Jenkins Logs
|
||||
""""""""""""
|
||||
|
||||
.. |pytest| replace:: ``pytest``
|
||||
.. _pytest: https://docs.pytest.org/en/stable/
|
||||
|
||||
The first place to look for a failure is in the CI logs, follow the red Xs on
|
||||
the failing job to view the logs. Note:
|
||||
|
||||
* Jenkins does not display the full log by default, at the top of the log viewer
|
||||
is a button "Show complete log" which will take you to a plaintext version of the log
|
||||
* |pytest|_ failures are summarized at the bottom of the log but you will likely
|
||||
need to scroll up to view the actual failure.
|
||||
|
||||
Reproduce Failures
|
||||
""""""""""""""""""
|
||||
|
||||
Most TVM Python tests run under |pytest|_ and can be run as described in :ref:`pr-testing`.
|
||||
|
||||
|
||||
Reporting Issues
|
||||
^^^^^^^^^^^^^^^^
|
||||
|
||||
Issues with CI should be `reported on GitHub <https://github.com/apache/tvm/issues/new?assignees=&labels=&template=ci-problem.md&title=%5BCI+Problem%5D+>`_
|
||||
with a link to the relevant jobs, commits, or PRs.
|
||||
|
||||
|
||||
|
||||
For Maintainers
|
||||
---------------
|
||||
|
||||
This section discusses processes ran by TVM Maintainers.
|
||||
|
||||
|
||||
Procedures for Keeping CI Green
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
This section talks about common procedures used to keep CI passing.
|
||||
|
||||
Broken CI due to Simultaneous Merge
|
||||
"""""""""""""""""""""""""""""""""""
|
||||
|
||||
Developers rely on the TVM CI to get signal on their PRs before merging. Occasionally, two
|
||||
different PRs can pass CI individually but break ``main`` when both land. This in turn causes an
|
||||
error to show up on an unrelated PR that is based on the broken commit(s). Broken commits can be
|
||||
identified `through GitHub <https://github.com/apache/tvm/commits/main>`_ via the commit status icon
|
||||
or via `Jenkins <https://ci.tlcpack.ai/blue/organizations/jenkins/tvm/activity?branch=main>`_.
|
||||
|
||||
In these situations it is ultimately the responsibility of the TVM Committer who merged the PR to
|
||||
fix CI (others are encouraged to help). Typical responses to this situation are:
|
||||
1. revert the offending commit
|
||||
2. submit a forward fix to address the issue.
|
||||
|
||||
It is up to the committer and commit author which option to choose. A broken CI affects all TVM
|
||||
developers and should be fixed as soon as possible, while a revert may be especially painful for the
|
||||
author of the offending PR when that PR is large.
|
||||
|
||||
|
||||
Dealing with Flakiness
|
||||
^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
If you notice a failure on your PR that seems unrelated to your change, you should
|
||||
search `recent GitHub issues related to flaky tests <https://github.com/apache/tvm/issues?q=is%3Aissue+%5BCI+Problem%5D+Flaky>`_ and
|
||||
`file a new issue <https://github.com/apache/tvm/issues/new?assignees=&labels=&template=ci-problem.md&title=%5BCI+Problem%5D>`_
|
||||
if you don't see any reports of the failure. If a certain test or class of tests affects
|
||||
several PRs or commits on ``main`` with flaky failures, the test should be disabled via
|
||||
`pytest's @xfail decorator <https://docs.pytest.org/en/stable/how-to/skipping.html#xfail-mark-test-functions-as-expected-to-fail>`_ with `strict=False <https://docs.pytest.org/en/stable/how-to/skipping.html#strict-parameter>`_ and the relevant issue linked in the
|
||||
disabling PR.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@pytest.mark.xfail(strict=False, reason="Flaky test: https://github.com/apache/tvm/issues/1234")
|
||||
def test_something_flaky():
|
||||
pass
|
||||
|
||||
Then submit a PR as usual
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
git add <test file>
|
||||
git commit -m'[skip ci][ci] Disable flaky test: ``<test_name>``
|
||||
|
||||
See #<issue number>
|
||||
'
|
||||
gh pr create
|
||||
|
||||
|
||||
Skipping CI
|
||||
^^^^^^^^^^^
|
||||
|
||||
For reverts and trivial forward fixes, adding ``[skip ci]`` to the revert's
|
||||
PR title will cause Jenkins CI to shortcut and only run lint. Committers should
|
||||
take care that they only merge CI-skipped PRs to fix a failure on ``main`` and
|
||||
not in cases where the submitter wants to shortcut CI to merge a change faster.
|
||||
The PR title is checked when the build is first run (specifically during the Jenkins lint
|
||||
step, so changes after that has run do not affect CI and will require the job to
|
||||
be re-triggered by another ``git push``). Note that GitHub Actions lint always
|
||||
runs independently via pre-commit hooks.
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# Revert HEAD commit, make sure to insert '[skip ci]' at the beginning of
|
||||
# the commit subject
|
||||
git revert HEAD
|
||||
git checkout -b my_fix
|
||||
# After you have pushed your branch, create a PR as usual.
|
||||
git push my_repo
|
||||
# Example: Skip CI on a branch with an existing PR
|
||||
# Adding this commit to an existing branch will cause a new CI run where
|
||||
# Jenkins is skipped
|
||||
git commit --allow-empty --message "[skip ci] Trigger skipped CI"
|
||||
git push my_repo
|
||||
|
||||
|
||||
Docker Images
|
||||
^^^^^^^^^^^^^
|
||||
|
||||
Each CI job runs most of its work inside a Docker container, built from files
|
||||
in the `docker/ <https://github.com/apache/tvm/tree/main/docker>`_ folder.
|
||||
|
||||
|
||||
Updating a Docker Image Tag
|
||||
"""""""""""""""""""""""""""
|
||||
|
||||
To update a tag, a new image needs to be built and uploaded to Docker Hub, then
|
||||
the image tags in `docker-images.ini <https://github.com/apache/tvm/tree/main/ci/jenkins/docker-images.ini>`_
|
||||
need to be updated to match the image tags on Docker Hub.
|
||||
|
||||
Docker images are built automatically nightly via the `tvm-docker <https://ci.tlcpack.ai/job/tvm-docker/>`_,
|
||||
which uploads the built images to https://hub.docker.com/u/tlcpackstaging once
|
||||
they have passed CI. Post-merge CI runs on ``main`` build Docker images ad-hoc
|
||||
and upload them to the ``tlcpackstaging`` Docker Hub account as well. There is an
|
||||
auto-promotion process for ``tlcpackstaging`` Docker images to be moved to the
|
||||
``tlcpack`` account. This means that image tags from ``tlcpackstaging`` can be
|
||||
used in CI and they will be automatically moved to ``tlcpack`` after a successful
|
||||
post-merge CI run on ``main``. So the steps to update the image are:
|
||||
|
||||
1. Merge a PR that changes the Dockerfiles under ``docker/`` or scripts in ``docker/install``.
|
||||
2. Do either of:
|
||||
|
||||
a. Wait for the post-merge CI build from the PR to complete and upload the newly built image to the `tlcpackstaging <https://hub.docker.com/u/tlcpackstaging>`_ Docker Hub.
|
||||
b. Wait for the nightly Docker image build to complete and upload the newly built image to the `tlcpackstaging <https://hub.docker.com/u/tlcpackstaging>`_ Docker Hub.
|
||||
|
||||
3. Find the newly uploaded image tag on the `tlcpackstaging <https://hub.docker.com/u/tlcpackstaging>`_ Docker Hub, for example ``20221208-070144-22ff38dff`` and update the tag in ``ci/jenkins/docker-images.ini`` to use the tlcpackstaging tag but under the tlcpack account, e.g. ``tlcpack/ci-arm:20221208-070144-22ff38dff``. Send in a PR with these changes and wait for it to run through CI to ensure the new images are valid.
|
||||
4. Merge the ``docker-images.ini`` update PR. Once post-merge CI finishes running on ``main`` the ``tlcpackstaging`` tag will be re-uploaded to ``tlcpack`` automatically.
|
||||
|
||||
Adding a New Docker Image
|
||||
"""""""""""""""""""""""""
|
||||
|
||||
New docker images can be added to test TVM on a variety of platforms. Here are the steps for adding
|
||||
a new CI image:
|
||||
|
||||
1. Define the ``docker/Dockerfile.ci_foo`` and associated scripts in ``docker/install``. Create a PR containing only these changes (no ``Jenkinsfile`` changes).
|
||||
|
||||
Example: https://github.com/apache/tvm/pull/12230/files
|
||||
|
||||
2. A committer verifies the image builds locally and then reviews/approves this PR.
|
||||
3. A committer creates the ci-foo repos in https://hub.docker.com/u/tlcpack and https://hub.docker.com/u/tlcpackstaging.
|
||||
4. Create a PR to create an ECR repo for the image in tlcpack/ci: https://github.com/tlc-pack/ci/pull/46/files
|
||||
5. A committer creates and gets merged a PR to add the image to the ``Jenkinsfile``
|
||||
|
||||
Example: https://github.com/apache/tvm/pull/12369/files.
|
||||
|
||||
**NOTE**: The PR must be opened from a branch in apache/tvm, not from a branch in a forked repo.
|
||||
|
||||
6. A committer adds this image to the daily docker rebuild/validation run in tlcpack.
|
||||
|
||||
Example: https://github.com/tlc-pack/tlcpack/pull/131
|
||||
|
||||
|
||||
``ci-docker-staging``
|
||||
^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
The `ci-docker-staging <https://github.com/apache/tvm/tree/ci-docker-staging>`_
|
||||
branch is typically used to test updates to Docker images and ``Jenkinsfile`` changes. When
|
||||
running a build for a normal PR from a forked repository, Jenkins uses the code
|
||||
from the PR except for the ``Jenkinsfile`` itself, which comes from the base branch.
|
||||
When branches are built, the ``Jenkinsfile`` in the branch is used, so a committer
|
||||
with write access must push PRs to a branch in apache/tvm to properly test
|
||||
``Jenkinsfile`` changes. If your PR makes changes to the ``Jenkinsfile``, make sure
|
||||
to @ a `committer <https://github.com/apache/tvm/tree/main/CONTRIBUTORS.md>`_
|
||||
and ask them to push your PR as a branch to test the changes.
|
||||
|
||||
|
||||
CI Monitoring Rotation
|
||||
^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Some tests are also flaky and occasionally fail for reasons unrelated to the PR. The
|
||||
`CI monitoring rotation <https://github.com/apache/tvm/wiki/CI-Monitoring-Runbook>`_ watches for these failures and
|
||||
disables tests as necessary. It is the responsibility of those who wrote the test to ultimately fix
|
||||
and re-enable the test.
|
||||
@@ -0,0 +1,185 @@
|
||||
.. 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.
|
||||
|
||||
.. _code_guide:
|
||||
|
||||
Code Guide and Tips
|
||||
===================
|
||||
|
||||
.. contents::
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
This is a document used to record tips in TVM codebase for reviewers and contributors.
|
||||
Most of them are summarized through lessons during the contributing and process.
|
||||
|
||||
|
||||
C++ Code Styles
|
||||
---------------
|
||||
- Use the Google C/C++ style.
|
||||
- The public facing functions are documented in doxygen format.
|
||||
- Favor concrete type declaration over ``auto`` as long as it is short.
|
||||
- Favor passing by const reference (e.g. ``const Expr&``) over passing by value.
|
||||
Except when the function consumes the value by copy constructor or move,
|
||||
pass by value is better than pass by const reference in such cases.
|
||||
- Favor ``const`` member function when possible.
|
||||
|
||||
We use ``clang-format`` to enforce the code style. Because different version
|
||||
of clang-format might change by its version, it is recommended to use the same
|
||||
version of the clang-format as the main one.
|
||||
You can use pre-commit hooks to run formatting checks:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
# Run clang-format on all files
|
||||
pre-commit run clang-format --all-files
|
||||
|
||||
# Run all linters (Python + C++ + custom checks)
|
||||
pre-commit run --all-files
|
||||
|
||||
|
||||
clang-format is also not perfect, when necessary, you can use disble clang-format on certain code regions.
|
||||
|
||||
.. code :: c
|
||||
|
||||
// clang-format off
|
||||
void Test() {
|
||||
// clang-format will be disabled in this region.
|
||||
}
|
||||
// clang-format on
|
||||
|
||||
|
||||
Because clang-format may not recognize macros, it is recommended to use macro like normal function styles.
|
||||
|
||||
|
||||
.. code :: c
|
||||
|
||||
#define MACRO_IMPL { custom impl; }
|
||||
#define MACRO_FUNC(x)
|
||||
|
||||
// not preferred, because clang-format might recognize it as types.
|
||||
virtual void Func1() MACRO_IMPL
|
||||
|
||||
// preferred
|
||||
virtual void Func2() MACRO_IMPL;
|
||||
|
||||
void Func3() {
|
||||
// preferred
|
||||
MACRO_FUNC(xyz);
|
||||
}
|
||||
|
||||
|
||||
Python Code Styles
|
||||
------------------
|
||||
- The functions and classes are documented in `numpydoc <https://numpydoc.readthedocs.io/en/latest/>`_ format.
|
||||
- Check your code style using ``pre-commit run --all-files``
|
||||
- Stick to language features in ``python 3.10``
|
||||
|
||||
- For functions with early returns, prefer ``if``/``elif``/``else``
|
||||
chains for functions with parallel and short bodies to the
|
||||
conditions, such as functions that apply a simple mapping to the
|
||||
arguments. For more procedural functions, especially where the
|
||||
final ``else`` block would be much longer than the ``if`` and
|
||||
``elif`` blocks, prefer having the final ``else`` case unindented.
|
||||
|
||||
The pylint check ``no-else-return`` is disabled to allow for this
|
||||
distinction. See further discussion `here
|
||||
<https://github.com/apache/tvm/pull/11327>`_.
|
||||
|
||||
.. code:: python
|
||||
|
||||
# All cases have bodies with similar flow control. While this could
|
||||
# be expressed as a sequence of if conditions, a reader would need to
|
||||
# inspect the body of each condition to know that only one conditional
|
||||
# body may be reached.
|
||||
def sign(x):
|
||||
if x > 0:
|
||||
return "+"
|
||||
elif x < 0:
|
||||
return "-"
|
||||
else:
|
||||
return ""
|
||||
|
||||
# The initial special case is an early return for a special case,
|
||||
# followed by a more general method. Using an else block for the
|
||||
# condition would add unnecessary indentation for the remainder of the
|
||||
# function.
|
||||
def num_unique_subsets(values):
|
||||
if len(values)==0:
|
||||
return 1
|
||||
|
||||
# Longer, more general solution here
|
||||
...
|
||||
|
||||
Writing Python Tests
|
||||
--------------------
|
||||
We use `pytest <https://docs.pytest.org/en/stable/>`_ for all Python testing. Regular tests live
|
||||
under ``tests/python``, while environment-specific nightly tests live under ``tests/nightly/python``.
|
||||
See :doc:`testing` for details on running tests, target parametrization,
|
||||
and the target-specific marks used by CI.
|
||||
|
||||
If you want your test to run over a variety of targets, parametrize over ``target`` with ``@pytest.mark.parametrize``, tag GPU targets with ``@pytest.mark.gpu`` so the CI routes them to GPU nodes, and skip a target that is unavailable on the current machine with :py:func:`tvm.testing.device_enabled`. For example:
|
||||
|
||||
.. code:: python
|
||||
|
||||
@pytest.mark.parametrize("target", ["llvm", pytest.param("cuda", marks=pytest.mark.gpu)])
|
||||
def test_mytest(target):
|
||||
if not tvm.testing.device_enabled(target):
|
||||
pytest.skip(f"{target} not enabled")
|
||||
dev = tvm.device(target)
|
||||
...
|
||||
|
||||
will run ``test_mytest`` with ``target="llvm"`` and ``target="cuda"``, skipping any target whose device is not present. If you only want to test against a single target, drop the parametrization and hardcode the target. Mark GPU tests with ``@pytest.mark.gpu`` so the CI can select them, and skip when the required feature is unavailable with ``@pytest.mark.skipif``. For example, CUDA tests use:
|
||||
|
||||
.. code:: python
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not tvm.testing.env.has_cuda(), reason="need cuda")
|
||||
def test_mycudatest():
|
||||
...
|
||||
|
||||
The ``tvm.testing.env`` module exposes a ``has_*()`` probe for each runtime and hardware feature (e.g. ``has_cuda()``, ``has_rocm()``, ``has_vulkan()``, ``has_llvm()``). To skip a test when an optional Python package is missing, use ``pytest.importorskip("package_name")``.
|
||||
|
||||
|
||||
Network Resources
|
||||
-----------------
|
||||
|
||||
In CI, downloading files from the Internet is a big source of flaky test failures (e.g. remote
|
||||
server can go down or be slow), so try to avoid using the network at all during tests. In some cases
|
||||
this isn't a reasonable proposition (e.g. the docs tutorials which need to download models).
|
||||
|
||||
New network downloads are rejected by the CI `request hook
|
||||
<https://github.com/apache/tvm/blob/main/tests/python/request_hook.py>`_. Prefer
|
||||
checked-in fixtures or generated data. If a download is unavoidable, arrange a stable
|
||||
project-managed mirror with the maintainers and add an explicit ``URL_MAP`` entry.
|
||||
|
||||
|
||||
Handle Integer Constant Expression
|
||||
----------------------------------
|
||||
We often need to handle constant integer expressions in TVM. Before we do so, the first question we want to ask is that is it really necessary to get a constant integer. If symbolic expression also works and let the logic flow, we should use symbolic expression as much as possible. So the generated code works for shapes that are not known ahead of time.
|
||||
|
||||
Note that in some cases we cannot know certain information, e.g. sign of symbolic variable, it is ok to make assumptions in certain cases. While adding precise support if the variable is constant.
|
||||
|
||||
If we do have to get constant integer expression, we should get the constant value using type ``int64_t`` instead of ``int``, to avoid potential integer overflow. We can always reconstruct an integer with the corresponding expression type via ``make_const``. The following code gives an example.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
Expr CalculateExpr(Expr value) {
|
||||
int64_t int_value = GetConstInt<int64_t>(value);
|
||||
int_value = CalculateExprInInt64(int_value);
|
||||
return make_const(value.type(), int_value);
|
||||
}
|
||||
@@ -0,0 +1,221 @@
|
||||
.. 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.
|
||||
|
||||
.. _code_review_guide:
|
||||
|
||||
|
||||
Code Reviews
|
||||
============
|
||||
|
||||
.. contents::
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
Open source code is maintained by a community with diverse backgrounds, interests, and goals.
|
||||
Hence it is important to provide clear, documented and maintainable code and processes. Code reviews are a
|
||||
shepherding process used to collectively spot potential problems, improve quality of the code, and educate both contributors
|
||||
and reviewers about the code base and its assumptions. It is also one mechanism to ensure there are multiple people who can
|
||||
maintain a related piece of code together. Contributors are encouraged to polish the code to a reviewable
|
||||
state before requesting reviews. This is especially important for committer candidates, as committers are expected
|
||||
to participate in not only writing code but also reviewing it.
|
||||
|
||||
This document is a living guideline for code review in open source. Please also take sometime to read
|
||||
:ref:`community_guide` about the general development process.
|
||||
|
||||
Building Trust
|
||||
--------------
|
||||
|
||||
First and foremost, we are building a community that based on trust, which takes time
|
||||
and effort to both build and maintain. We expect our community members to work together in a
|
||||
constructive way and work together with common sense. Although we all have different sets of backgrounds,
|
||||
interests and goals we must work together to find solutions that work for the larger community.
|
||||
Trust-based collaboration is also a key tenant of the Apache way and an important factor to consider in growing the community,
|
||||
and promoting members to official roles.
|
||||
|
||||
Community Participation
|
||||
-----------------------
|
||||
|
||||
Everyone is welcomed to comment on PRs. We encourage committers to wait for some period of time(e.g. three days)
|
||||
before merging PR that contains a major architecture change. The goal is to give people time to speak up and
|
||||
express interest in reviewing and participate.
|
||||
|
||||
Remembering that we are all coming from different backgrounds is important here. For example some community members
|
||||
work in different time zones, only work on open source during work hours, or may be traveling or having other events
|
||||
going on in their lives. An important part of working in a large project is ensuring there is collective understanding,
|
||||
so no one person is a bottleneck. While it is important to allow time for participation in code review we also can not
|
||||
block all changes on all reviewers. Remember that helping people land PRs is a great way to encourage broader
|
||||
participation, especially for those who volunteer their time to contribute.
|
||||
|
||||
Part of this is trusting and communicating with fellow maintainers that if changes need to be applied in the future
|
||||
that PR authors will later follow through on their promises. It is the responsibility of committers to listen to all
|
||||
feedback whether from PMC members or new contributors and consider what actions need to be taken.
|
||||
|
||||
Read the code carefully
|
||||
-----------------------
|
||||
|
||||
Sometimes we may quickly read through the code and only pick up on a selective aspects of the code. These type of comments
|
||||
are usually helpful and should be welcomed in the community. However, they are only part of performing code review and
|
||||
should be part of more comprehensive feedback. A good and careful code review is a large time investment and sometimes
|
||||
can be longer than writing the actual contribution.
|
||||
|
||||
For example receiving only highly critical feedback on minor aspects of your PR rarely feels good, and it can be discouraging
|
||||
if your time and effort was not reciprocated during review. Practicing empathy when acting both as a contributor and committer
|
||||
is important and can help make you a more effective code reviewer and contributor.
|
||||
|
||||
We expect that all committers carefully read and understand the code before signing off. There is a lot of trust involved when
|
||||
a committer hits the merge button. In the meantime, we acknowledge that sometimes problems slip through, in that case, the
|
||||
merger is responsible for ensuring the correct follow up actions are taken.
|
||||
|
||||
Be Respectful
|
||||
-------------
|
||||
|
||||
- To everyone who are making comments: making constructive comment will help new contributors to land their PRs
|
||||
timely and help us welcome new members to the community.
|
||||
|
||||
- To authors: reviewers should spend significant time reading the code, and a careful review could be as time intensive
|
||||
as writing the code from scratch. Respectfully address review comments and reciprocate the review by helping review
|
||||
others changes in the future.
|
||||
|
||||
Most importantly focus on having a constructive conversation, and try to assume best intentions when interacting as a reviewer.
|
||||
If there is something in the process not working, consider getting some face time with the other contributors and discussing
|
||||
how to improve the process or communication.
|
||||
|
||||
Factors to Consider about Code Quality
|
||||
--------------------------------------
|
||||
|
||||
High quality code is critical to the long term success of the project. There are many factors of code quality to consider
|
||||
during a code review:
|
||||
|
||||
- F0: Overall architecture. This includes the definition of public modules, key data structures and public interfaces.
|
||||
Good architectural choices are critical to the success of the project in the long run.
|
||||
- F1: Architectural consistency. There are usually multiple ways to implement a new feature. We must ensure new
|
||||
features are consistent with previous overall architectural choices and interact well with the existing code.
|
||||
Every new feature increases the complexity of the project, and a consistent design ideally minimizes the increase
|
||||
in complexity bought by a new feature, making it easier to maintain code in the long run.
|
||||
- F2: Code robustness and test coverage. Ensure code runs correctly in all possible settings(platforms), ensure
|
||||
test coverage of the new feature. Clear error messages for user facing errors.
|
||||
- F3: User facing API documentation: documentation of public user facing APIs and key module interfaces are mandatory.
|
||||
This includes the API, data structures that appears in the public interface (i.e., `include/tvm` and user facing python APIs).
|
||||
We generally encourage well documented code and include some form of documentations for internal APIs that are used in
|
||||
multiple places, see also F4.
|
||||
- F4: Code readability. Readability involves multiple aspects: instructive and consistent function names, clear implementation
|
||||
of the overall flow, descriptive comments for complex code logic and internal functions. Readable code is easier to maintain.
|
||||
|
||||
Architectural design and consistency are the most important factors since they are likely to introduce the most long term technical debt.
|
||||
As a result, committers should most carefully consider these factors before merging the code.
|
||||
|
||||
Test coverage and API documentation are expected for code contributions.
|
||||
|
||||
Code readability is relatively a subjective matter compared to the other ones.
|
||||
Different people have different thoughts on how to best write code. Reviewers should make constructive and actionable comments.
|
||||
In the meantime, code review should not be used as a way to get others to write code exactly the way you would.
|
||||
Conversely you should also consider that what you may easily understand, or find acceptable might not work for the larger
|
||||
community or other members. Use your judgment on what is appropriate based on the content and the scope of the contribution
|
||||
and where the contributor is coming from.
|
||||
|
||||
We follow common :ref:`code_guide` when writing code. Style guides help ensure that code is readable and maintainable by others,
|
||||
long after the original author has moved on. Style guides are more than about code formatting — they also pertain
|
||||
to the correct way to document code, variable naming, and other conventions that are not enforced by automatic formatters.
|
||||
|
||||
Consensus Building
|
||||
------------------
|
||||
|
||||
Disagreements can happen during code reviews. We encourage building consensus among the people involved. We are working together
|
||||
and building trust with each other in OSS. The nature of OSS means sometimes we make compromises on less significant issues to
|
||||
make steady progress and welcome broader participation in the community. Compromise unfortunately means sometimes the world will
|
||||
not be exactly as we would like, this true even for leaders of the community.
|
||||
|
||||
- Be civil and build consensus through constructive technical-based conversations.
|
||||
- A committer who owns the area can serve as a shepherd to drive the discussion by taking all the conversations into consideration,
|
||||
and suggest a resolution with to move forward.
|
||||
- Because a lot of trust is involved on the committer(shepherd), they should read the PR carefully before sign off. Additionally,
|
||||
the merger should also take the responsibility to followup in case there are problems caused by the merge.
|
||||
|
||||
Consistency
|
||||
-----------
|
||||
|
||||
A final remark is that we are all human and its hard to always be perfectly consistent. If contributors feel that you didn't apply these guidelines
|
||||
in a consistent way it is important to listen and hear folks out. We will constantly have to iterate on processes and guidelines as we evolve as a community.
|
||||
Our goal is to strive to be consistent and objective but all of us are unfortunately human and imperfect and will need to adjust and learn.
|
||||
|
||||
Additional Recommendations
|
||||
--------------------------
|
||||
|
||||
Deliberate on API and Data Structures
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
A minimum and stable API is critical to the project’s life. A good API makes a huge difference. Always think very carefully about all the aspects including naming, argument definitions and behavior.
|
||||
|
||||
When possible, pay more attention still to the proposed API design during code reviews.
|
||||
Remember, it is easier to improve code implementation, but it is extremely hard to change an API once accepted.
|
||||
We should treat data structures that are shared across modules(e.g. AST) in the same way.
|
||||
If/when uncertain, start a conversation with more developers before committing.
|
||||
|
||||
Here are some useful principles for designing APIs:
|
||||
|
||||
- Be consistent with existing well-known package’s APIs if the features overlap.
|
||||
For example, tensor operation APIs should always be consistent with the numpy API.
|
||||
- Be consistent with existing APIs in the same project.
|
||||
For example, we should use the same argument ordering across all the optimization passes,
|
||||
so there is no "surprise" when using them.
|
||||
- Think about whether the API will change in the future.
|
||||
For example, we will have more options like loop_unrolling and device placement policy
|
||||
as we add more optimizations in build. We can package optimization knobs into a build
|
||||
configuration object. In this way, the build API is stable over time, even though it may be enriched.
|
||||
- Write documentation. Documentation is mandatory for APIs and sometimes writing documents helps
|
||||
us to think further about the design as well as whether we need to add further clarifications.
|
||||
- Minimum. Think about how many lines of code a user has to write to use the API.
|
||||
Remove layers of abstraction when possible.
|
||||
|
||||
Minimize Dependencies
|
||||
~~~~~~~~~~~~~~~~~~~~~
|
||||
Always be cautious in introducing dependencies. While it is important to reuse code and avoid reinventing the wheel,
|
||||
dependencies can increase burden of users in deployment. A good design principle is that a feature or function
|
||||
should only have a dependency if/when a user actually use it.
|
||||
|
||||
Concise Implementation
|
||||
~~~~~~~~~~~~~~~~~~~~~~
|
||||
Some basic principles applied here: favor vectorized array code over loops, use existing APIs that solve the problem.
|
||||
|
||||
Document Lessons in Code Reviews
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
When you find there are some common or recurring lessons that can be summarized,
|
||||
add it to the :ref:`code_guide`.
|
||||
It is always good to refer to the guideline document when requesting changes,
|
||||
so the lessons can be shared to all the community.
|
||||
|
||||
|
||||
Learn from other Code Reviews
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
There can be multiple reviewers reviewing the same changes. Many times other reviewers
|
||||
may spot things you did not find. Try to learn from other code reviews, when possible, document these lessons.
|
||||
|
||||
Approve and Request Changes Explicitly
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
The contributor and code owner can request code reviews from multiple reviewers.
|
||||
Remember to approve changes when your comments are addressed in a code review.
|
||||
To do so -- please click on changes tab in the pull request, then select approve,
|
||||
or comment on the code and click request changes.
|
||||
Code owner can decide if the code can be merged in case by case if some of the reviewers
|
||||
did not respond in time(e.g. a week) and existing reviews are sufficient.
|
||||
|
||||
Reviewers
|
||||
~~~~~~~~~
|
||||
Reviewers should strive to leave timely feedback on pull requests for which their
|
||||
review was requested. Reviewing code is an important part of the project's health
|
||||
and should be considered a regular responsibility for contributors. Automated
|
||||
tooling helps out in this regard, as PRs with no activity for a set amount of
|
||||
time will get a bot comment pinging the relevant parties.
|
||||
@@ -0,0 +1,121 @@
|
||||
.. 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.
|
||||
|
||||
.. _committer_guide:
|
||||
|
||||
Committer Guide
|
||||
===============
|
||||
|
||||
.. contents::
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
This is an evolving document to provide some helpful tips for committers.
|
||||
Most of them are lessons learned during development.
|
||||
We welcome every committer to contribute to this document.
|
||||
See the :ref:`community_guide` for an overview of
|
||||
the committership and the general development process.
|
||||
|
||||
Community First
|
||||
---------------
|
||||
The collective effort of the community moves the project forward and
|
||||
makes the project awesome for everyone.
|
||||
When we make a decision, it is always helpful to keep the community in mind.
|
||||
Here are some example questions that we can ask:
|
||||
|
||||
- How can I encourage new contributors to get more involved in the project?
|
||||
- Can I help to save my fellow committers' time?
|
||||
- Have I enabled the rest of the community to participate the
|
||||
design proposals?
|
||||
|
||||
|
||||
Public Archive Principle
|
||||
------------------------
|
||||
While private channels such as face to face discussion are useful for development,
|
||||
they also create barriers for the broader community's participation.
|
||||
The Apache way of development requires all decisions
|
||||
to be made in public channels, which are archived and accessible to everyone.
|
||||
As a result, any contributor can keep up with the development by watching the
|
||||
archives and join the development anytime.
|
||||
|
||||
While this principle applies to every contributor,
|
||||
it is especially important for committers.
|
||||
Here are some example applications of this principle:
|
||||
|
||||
- When getting a project-related question from a personal channel,
|
||||
encourage the person to open a public thread in the discuss forum,
|
||||
so others in the community can benefit from the answer.
|
||||
- After an in-person discussion, send a summary to public channels
|
||||
(as an RFC or a discuss thread).
|
||||
|
||||
|
||||
Independent Project Management
|
||||
------------------------------
|
||||
|
||||
Everyone is presumed to be wearing their Apache committer hat when participating in the project.
|
||||
That is, committers should act - in the context of the project activities - in the best interests of the project.
|
||||
Separating your hat between committer and any other roles you may have is important in all aspects.
|
||||
|
||||
In the context of project participation, it can be helpful to state which hat you are wearing in cases where that
|
||||
can cause confusion, especially in cases where you are not wearing committer hat. Two examples:
|
||||
|
||||
- "Wearing [foo] hat: [message when serving as foo's role and not as committer]".
|
||||
- "Wearing Apache TVM hat: [messages when serving as committer]".
|
||||
|
||||
Shepherd a Pull Request
|
||||
-----------------------
|
||||
|
||||
Here are some tips to shepherd a pull request.
|
||||
You can also take a look at the :ref:`code_review_guide`.
|
||||
|
||||
- Assign the PR to yourself, so that other committers
|
||||
know that the PR has already been tended to.
|
||||
- Make use of the status label to indicate the current status.
|
||||
- Check if an RFC needs to be sent.
|
||||
- If the contributor has not requested a reviewer, kindly
|
||||
ask the contributor to do so.
|
||||
If the PR comes from a new contributor,
|
||||
help the contributor to request reviewers
|
||||
and ask the contributor to do so next time.
|
||||
- Moderate the reviews, ask reviewers to approve explicitly.
|
||||
- Mark the PR as accepted and acknowledge the contributor/reviewers.
|
||||
- Merge the PR :)
|
||||
|
||||
|
||||
Time Management
|
||||
---------------
|
||||
There are many things that a committer can do, such as
|
||||
moderating discussions, pull request reviews and
|
||||
code contributions.
|
||||
|
||||
Working on an open source project can be rewarding,
|
||||
but also be a bit overwhelming sometimes.
|
||||
A little bit of time management might be helpful to alleviate the problem.
|
||||
For example, some committers have a "community day" in a week
|
||||
when they actively manage outstanding PRs,
|
||||
but watch the community less frequently in the rest of the time.
|
||||
|
||||
Remember that your merit will never go away, so please
|
||||
take your time and pace when contributing to the project :)
|
||||
|
||||
|
||||
Broad Collaboration
|
||||
-------------------
|
||||
Sometimes, we tend to only interact with people we know.
|
||||
However, broad collaborations are necessary to the success of the project.
|
||||
Try to keep that in mind, shepherd PRs for, and request code reviews from
|
||||
community members who you do not interact physically.
|
||||
@@ -0,0 +1,64 @@
|
||||
.. 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.
|
||||
|
||||
.. _community_guide:
|
||||
|
||||
TVM Community Guidelines
|
||||
========================
|
||||
|
||||
.. contents::
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
|
||||
TVM adopts the Apache style model and governs by merit. We believe that it is important to create an inclusive community where everyone can use, contribute to, and influence the direction of the project. See `CONTRIBUTORS.md <https://github.com/apache/tvm/blob/main/CONTRIBUTORS.md>`_ for the current list of contributors.
|
||||
|
||||
General Development Process
|
||||
---------------------------
|
||||
Everyone in the community is welcomed to send patches, documents, and propose new directions to the project. The key guideline here is to enable everyone in the community to get involved and participate the decision and development. When major changes are proposed, an RFC should be sent to allow discussion by the community. We encourage public discussion, archivable channels such as issues, discuss forum and mailing-list, so that everyone in the community can participate and review the process later.
|
||||
|
||||
Code reviews are one of the key ways to ensure the quality of the code. High-quality code reviews prevent technical debt for long-term and are crucial to the success of the project. A pull request needs to be reviewed before it gets merged. A committer who has the expertise of the corresponding area would moderate the pull request and the merge the code when it is ready. The corresponding committer could request multiple reviewers who are familiar with the area of the code. We encourage contributors to request code reviews themselves and help review each other's code -- remember everyone is volunteering their time to the community, high-quality code review itself costs as much as the actual code contribution, you could get your code quickly reviewed if you do others the same favor.
|
||||
|
||||
The community should strive to reach a consensus on technical decisions through discussion. We expect committers and PMCs to moderate technical discussions in a diplomatic way, and provide suggestions with clear technical reasoning when necessary.
|
||||
|
||||
Strategy Decision Process
|
||||
-------------------------
|
||||
It takes lazy 2/3 majority (at least 3 votes and twice as many +1 votes as -1 votes) of binding decisions to make the following
|
||||
strategic decisions in the TVM community:
|
||||
|
||||
- Adoption of a guidance-level community strategy to enable new directions or overall project evolution.
|
||||
- Establishment of a new module in the project.
|
||||
- Adoption of a new codebase: When the codebase for an existing, released product is to be replaced with an alternative codebase.
|
||||
If such a vote fails to gain approval, the existing code base will continue. This also covers the creation of new sub-projects within the project.
|
||||
|
||||
All these decisions are made after community conversations that get captured as part of the summary.
|
||||
|
||||
|
||||
Committers
|
||||
----------
|
||||
Committers are individuals who are granted the write access to the project. A committer is usually responsible for a certain area or several areas of the code where they oversee the code review process. The area of contribution can take all forms, including code contributions and code reviews, documents, education, and outreach. Committers are essential for a high quality and healthy project. The community actively look for new committers from contributors. Here is a list of useful traits that help the community to recognize potential committers:
|
||||
|
||||
- Sustained contribution to the project, demonstrated by discussion over RFCs, code reviews and proposals of new features, and other development activities. Being familiar with, and being able to take ownership on one or several areas of the project.
|
||||
- Quality of contributions: High-quality, readable code contributions indicated by pull requests that can be merged without a substantial code review. History of creating clean, maintainable code and including good test cases. Informative code reviews to help other contributors that adhere to a good standard.
|
||||
- Community involvement: active participation in the discussion forum, promote the projects via tutorials, talks and outreach. We encourage committers to collaborate broadly, e.g. do code reviews and discuss designs with community members that they do not interact physically.
|
||||
|
||||
The `Project Management Committee (PMC) <https://projects.apache.org/committee.html?tvm>`_ consists group of active committers that moderate the discussion, manage the project release, and proposes new committer/PMC members. Potential candidates are usually proposed via an internal discussion among PMCs, followed by a consensus approval, (i.e. at least 3 +1 votes, and no vetoes). Any veto must be accompanied by reasoning. PMCs should serve the community by upholding the community practices and guidelines TVM a better community for everyone. PMCs should strive to only nominate new candidates outside of their own organization.
|
||||
|
||||
|
||||
Reviewers
|
||||
---------
|
||||
Reviewers are individuals who actively contributed to the project and are willing to participate in the code review of new contributions. We identify reviewers from active contributors. The committers should explicitly solicit reviews from reviewers. High-quality code reviews prevent technical debt for long-term and are crucial to the success of the project. A pull request to the project has to be reviewed by at least one reviewer in order to be merged.
|
||||
@@ -0,0 +1,257 @@
|
||||
.. 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.
|
||||
|
||||
.. _doc_guide:
|
||||
|
||||
Documentation
|
||||
=============
|
||||
|
||||
.. contents::
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
TVM documentation loosely follows the `formal documentation style described by
|
||||
Divio <https://documentation.divio.com>`_. This system has been chosen because
|
||||
it is a "simple, comprehensive and nearly universally-applicable scheme. It is
|
||||
proven in practice across a wide variety of fields and applications."
|
||||
|
||||
This document describes the organization of TVM documentation, and how to write
|
||||
new documentation. See `docs/README.md <https://github.com/apache/tvm/tree/main/docs#build-locally>`_
|
||||
for instructions on building the docs.
|
||||
|
||||
The Four Document Types
|
||||
***********************
|
||||
|
||||
Introductory Tutorials
|
||||
----------------------
|
||||
|
||||
These are step by step guides to introduce new users to a project. An
|
||||
introductory tutorial is designed to get a user engaged with the software
|
||||
without necessarily explaining why the software works the way it does. Those
|
||||
explanations can be saved for other document types. An introductory tutorial
|
||||
focuses on a successful first experience. These are the most important docs to
|
||||
turning newcomers into new users and developers. A fully end-to-end
|
||||
tutorial — from installing TVM and supporting ML software, to creating and
|
||||
training a model, to compiling to different architectures — will give a new
|
||||
user the opportunity to use TVM in the most efficient way possible. A tutorial
|
||||
teaches a beginner something they need to know. This is in contrast with a
|
||||
how-to, which is meant to be an answer to a question that a user with some
|
||||
experience would ask.
|
||||
|
||||
Tutorials need to be repeatable and reliable, because the lack of success means
|
||||
a user will look for other solutions.
|
||||
|
||||
How-to Guides
|
||||
-------------
|
||||
|
||||
These are step by step guides on how to solve particular problems. The user can
|
||||
ask meaningful questions, and the documents provide answers. An examples of
|
||||
this type of document might be, "how do I compile an optimized model for ARM
|
||||
architecture?" or "how do I compile and optimize a PyTorch model?" These
|
||||
documents should be open enough that a user could see how to apply it to a new
|
||||
use case. Practical usability is more important than completeness. The title
|
||||
should tell the user what problem the how-to is solving.
|
||||
|
||||
How are tutorials different from how-tos? A tutorial is oriented towards the
|
||||
new developer, and focuses on successfully introducing them to the software and
|
||||
community. A how-to, in contrast, focuses on accomplishing a specific task
|
||||
within the context of basic understanding. A tutorial helps to on-board and
|
||||
assumes no prior knowledge. A how-to assumes minimum knowledge, and is meant to
|
||||
guide someone to accomplish a specific task.
|
||||
|
||||
Reference
|
||||
---------
|
||||
|
||||
Reference documentation describes how the software is configured and operated.
|
||||
APIs, key functions, commands, and interfaces are all candidates for reference
|
||||
documentation. These are the technical manuals that let users build their own
|
||||
interfaces and programs. They are information oriented, focused on lists and
|
||||
descriptions. You can assume that the audience has a grasp on how the software
|
||||
works and is looking for specific answers to specific questions. Ideally, the
|
||||
reference documentation should have the same structure as the code base and be
|
||||
generated automatically as much as possible.
|
||||
|
||||
Architecture Guides
|
||||
-------------------
|
||||
|
||||
Architecture Guides are explanations are background material on a topic. These
|
||||
documents help to illuminate and understand the application environment. Why
|
||||
are things the way they are? What were the design decisions, what alternatives
|
||||
were considered, what are the RFCs describing the existing system? This
|
||||
includes academic papers and links to publications relevant to the software.
|
||||
Within these documents you can explore contradictory and conflicting position,
|
||||
and help the reader make sense of how and why the software was built the way it
|
||||
is. It's not the place for how-tos and descriptions on how to accomplish tasks.
|
||||
They instead focus on higher level concepts that help with the understanding of
|
||||
the project. Generally these are written by the architects and developers of
|
||||
the project, but can useful to help both users and developers to have a deeper
|
||||
understanding of why the software works the way it does, and how to contribute
|
||||
to it in ways that are consistent with the underlying design principles.
|
||||
|
||||
Special considerations for TVM
|
||||
------------------------------
|
||||
|
||||
The TVM community has some special considerations that require deviation from
|
||||
the simple docs style outlined by Divio. The first consideration is that there
|
||||
is frequently overlap between the user and developer communities. Many projects
|
||||
document the developer and user experience with separate systems, but it is
|
||||
appropriate to consider both in this system, with differentiations where
|
||||
appropriate. As a result the tutorials and how-tos will be divided between
|
||||
"User Guides" that focus on the user experience, and "Developer Guides" that
|
||||
focus on the developer experience.
|
||||
|
||||
The next consideration is that there are special topics within the TVM
|
||||
community that benefit from additional attention. Special "Topic Guides" can be
|
||||
created to index existing material, and provide context on how to navigate that
|
||||
material most effectively.
|
||||
|
||||
To facilitate newcomers, a special "Getting Started" section with installation
|
||||
instructions, a overview of why to use TVM, and other first-experience
|
||||
documents will be produced.
|
||||
|
||||
|
||||
Technical Details
|
||||
*****************
|
||||
|
||||
We use the `Sphinx <https://www.sphinx-doc.org>`_ for the main documentation.
|
||||
Sphinx supports both reStructuredText and markdown. When possible, we
|
||||
encourage reStructuredText as it has richer features. Note that the
|
||||
Python doc-string and tutorials allow you to embed reStructuredText syntax.
|
||||
|
||||
See
|
||||
`docs/README.md <https://github.com/apache/tvm/tree/main/docs#build-locally>`_
|
||||
for instructions on building the docs.
|
||||
|
||||
|
||||
Python Reference Documentation
|
||||
------------------------------
|
||||
|
||||
We use the `numpydoc <https://numpydoc.readthedocs.io/en/latest/>`_ format to
|
||||
document the function and classes. The following snippet gives an example
|
||||
docstring. We always document all the public functions, and when necessary,
|
||||
provide a usage example of the features we support (as shown below).
|
||||
|
||||
.. code:: python
|
||||
|
||||
def myfunction(arg1, arg2, arg3=3):
|
||||
"""Briefly describe my function.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
arg1 : Type1
|
||||
Description of arg1
|
||||
|
||||
arg2 : Type2
|
||||
Description of arg2
|
||||
|
||||
arg3 : Type3, optional
|
||||
Description of arg3
|
||||
|
||||
Returns
|
||||
-------
|
||||
rv1 : RType1
|
||||
Description of return type one
|
||||
|
||||
Examples
|
||||
--------
|
||||
.. code:: python
|
||||
|
||||
# Example usage of myfunction
|
||||
x = myfunction(1, 2)
|
||||
"""
|
||||
return rv1
|
||||
|
||||
Be careful to leave blank lines between sections of your documents. In the
|
||||
above case, there has to be a blank line before ``Parameters``, ``Returns`` and
|
||||
``Examples`` in order for the doc to be built correctly. To add a new function to
|
||||
the docs, we need to add the `sphinx.autodoc
|
||||
<https://www.sphinx-doc.org/en/master/usage/extensions/autodoc.html>`_ rules to
|
||||
`docs/reference/api/python <https://github.com/apache/tvm/tree/main/docs/reference/api/python>`_.
|
||||
You can refer to the existing files under this folder on how to add the
|
||||
functions.
|
||||
|
||||
C++ Reference Documentation
|
||||
---------------------------
|
||||
|
||||
We use the doxygen format to document c++ functions. The following snippet
|
||||
shows an example of c++ docstring.
|
||||
|
||||
.. code:: c++
|
||||
|
||||
/*!
|
||||
* \brief Description of my function
|
||||
* \param arg1 Description of arg1
|
||||
* \param arg2 Description of arg2
|
||||
* \returns describe return value
|
||||
*/
|
||||
int myfunction(int arg1, int arg2) {
|
||||
// When necessary, also add comment to clarify internal logics
|
||||
}
|
||||
|
||||
Besides documenting function usages, we also highly recommend contributors to
|
||||
add comments about code logics to improve readability.
|
||||
|
||||
Sphinx Gallery How-Tos
|
||||
----------------------
|
||||
|
||||
We use `sphinx-gallery <https://sphinx-gallery.github.io/>`_ to build many
|
||||
Python how-tos. You can find the source code under `docs/how_to/tutorials
|
||||
<https://github.com/apache/tvm/tree/main/docs/how_to/tutorials>`_.
|
||||
One thing that worth noting is that the comment blocks are written in
|
||||
reStructuredText instead of markdown so be aware of the syntax.
|
||||
|
||||
The how-to code will run on our build server to generate the document page. So
|
||||
we may have a restriction like not being able to access a remote Raspberry Pi,
|
||||
in such case add a flag variable to the tutorial (e.g. ``use_rasp``) and allow
|
||||
users to easily switch to the real device by changing one flag. Then use the
|
||||
existing environment to demonstrate the usage.
|
||||
|
||||
If you add a new categorization of how-to, you will need to add references to
|
||||
`conf.py <https://github.com/apache/tvm/tree/main/docs/conf.py>`_ and the
|
||||
`top-level docs index <https://github.com/apache/tvm/tree/main/docs/index.rst>`_
|
||||
(how-to entries are registered there directly, not in a separate how-to index).
|
||||
|
||||
Refer to Another Location in the Document
|
||||
-----------------------------------------
|
||||
Please use sphinx's ``:ref:`` markup to refer to another location in the same doc.
|
||||
|
||||
.. code-block:: rst
|
||||
|
||||
.. _document-my-section-tag
|
||||
|
||||
My Section
|
||||
----------
|
||||
|
||||
You can use :ref:`document-my-section-tag` to refer to My Section.
|
||||
|
||||
Documents with Images / Figures
|
||||
-------------------------------
|
||||
reStructuredText's `figure <https://docutils.sourceforge.io/docs/ref/rst/directives.html#figure>`_
|
||||
and `image <https://docutils.sourceforge.io/docs/ref/rst/directives.html#image>`_
|
||||
elements allow a document to include an image URL.
|
||||
|
||||
Image files created for TVM documentation should reside in the `<https://github.com/tlc-pack/web-data>`_
|
||||
repository, while the `.rst` files *using* those images should reside in the main TVM repostitory
|
||||
(`<https://github.com/apache/tvm>`_).
|
||||
|
||||
This will require two GitHub Pull Requests, one for the image files and another for the `.rst` files.
|
||||
Discussion between the contributor and reviewers may be necessary to coordinate the review process.
|
||||
|
||||
*IMPORTANT NOTE:* When using two Pull Requests as described above, please merge the
|
||||
Pull Request in `<https://github.com/tlc-pack/web-data>`_ *before* merging
|
||||
the Pull Request in `<https://github.com/apache/tvm>`_.
|
||||
This helps ensure that all URL links in TVM's online documentation are valid.
|
||||
@@ -0,0 +1,97 @@
|
||||
.. 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.
|
||||
|
||||
.. _error-handling-guide:
|
||||
|
||||
Error Handling Guide
|
||||
====================
|
||||
|
||||
.. contents::
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
TVM contains structured error classes to indicate specific types of error.
|
||||
Please raise a specific error type when possible, so that users can
|
||||
write code to handle a specific error category if necessary.
|
||||
You can directly raise the specific error object in python.
|
||||
In other languages like c++, you simply add ``<ErrorType>:`` prefix to
|
||||
the error message(see below).
|
||||
|
||||
.. note::
|
||||
|
||||
Please refer to :py:mod:`tvm.error` for the list of errors.
|
||||
|
||||
Raise a Specific Error in C++
|
||||
-----------------------------
|
||||
You can add ``<ErrorType>:`` prefix to your error message to
|
||||
raise an error of the corresponding type.
|
||||
Note that you do not have to add a new type
|
||||
:py:class:`tvm.error.InternalError` will be raised by default when
|
||||
there is no error type prefix in the message.
|
||||
This mechanism works for both ``LOG(FATAL)`` and ``TVM_FFI_ICHECK`` macros.
|
||||
The following code gives an example on how to do so.
|
||||
|
||||
.. code:: c
|
||||
|
||||
// src/api_test.cc
|
||||
void ErrorTest(int x, int y) {
|
||||
TVM_FFI_ICHECK_EQ(x, y) << "ValueError: expect x and y to be equal."
|
||||
if (x == 1) {
|
||||
LOG(FATAL) << "InternalError: cannot reach here";
|
||||
}
|
||||
}
|
||||
|
||||
When a C++ function registered via the FFI raises an error with a typed prefix,
|
||||
the TVM FFI system will automatically map it to the corresponding Python exception
|
||||
class. For example, a ``ValueError:`` prefix in the error message will raise a Python
|
||||
``ValueError``, and an ``InternalError:`` prefix will raise ``tvm.error.InternalError``.
|
||||
|
||||
TVM's FFI system combines both the Python and C++ stacktraces into a single message,
|
||||
and generates the corresponding error class automatically.
|
||||
|
||||
|
||||
How to choose an Error Type
|
||||
---------------------------
|
||||
You can go through the error types are listed below, try to use common
|
||||
sense and also refer to the choices in the existing code.
|
||||
We try to keep a reasonable amount of error types.
|
||||
If you feel there is a need to add a new error type, do the following steps:
|
||||
|
||||
- Send a RFC proposal with a description and usage examples in the current codebase.
|
||||
- Add the new error type to :py:mod:`tvm.error` with clear documents.
|
||||
- Update the list in this file to include the new error type.
|
||||
- Change the code to use the new error type.
|
||||
|
||||
We also recommend to use less abstraction when creating the short error messages.
|
||||
The code is more readable in this way, and also opens path to craft specific
|
||||
error messages when necessary.
|
||||
|
||||
.. code:: python
|
||||
|
||||
def preferred():
|
||||
# Very clear about what is being raised and what is the error message.
|
||||
raise OpNotImplemented("Operator relu is not implemented in the ONNX frontend")
|
||||
|
||||
def _op_not_implemented(op_name):
|
||||
return OpNotImplemented("Operator {} is not implemented.").format(op_name)
|
||||
|
||||
def not_preferred():
|
||||
# Introduces another level of indirection.
|
||||
raise _op_not_implemented("relu")
|
||||
|
||||
If we need to introduce a wrapper function that constructs multi-line error messages,
|
||||
please put wrapper in the same file so other developers can look up the implementation easily.
|
||||
@@ -0,0 +1,143 @@
|
||||
.. 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.
|
||||
|
||||
.. _git-howto:
|
||||
|
||||
|
||||
Git Usage Tips
|
||||
==============
|
||||
|
||||
.. contents::
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
Here are some tips for git workflow.
|
||||
|
||||
How to resolve a conflict with ``main``
|
||||
---------------------------------------
|
||||
|
||||
- First rebase to most recent main
|
||||
|
||||
.. code:: bash
|
||||
|
||||
# The first two steps can be skipped after you do it once.
|
||||
git remote add upstream [url to tvm repo]
|
||||
git fetch upstream
|
||||
git rebase upstream/main
|
||||
|
||||
|
||||
- The git may show some conflicts it cannot merge, say ``conflicted.py``.
|
||||
|
||||
- Manually modify the file to resolve the conflict.
|
||||
- After you resolved the conflict, mark it as resolved by
|
||||
|
||||
.. code:: bash
|
||||
|
||||
git add conflicted.py
|
||||
|
||||
- Then you can continue rebase by
|
||||
|
||||
.. code:: bash
|
||||
|
||||
git rebase --continue
|
||||
|
||||
- Finally push to your fork, you may need to force push here.
|
||||
|
||||
.. code:: bash
|
||||
|
||||
git push --force
|
||||
|
||||
|
||||
How to combine multiple commits into one
|
||||
----------------------------------------
|
||||
|
||||
Sometimes we want to combine multiple commits, especially when later commits are only fixes to previous ones,
|
||||
to create a PR with set of meaningful commits. You can do it by following steps.
|
||||
|
||||
- Before doing so, configure the default editor of git if you haven't done so before.
|
||||
|
||||
.. code:: bash
|
||||
|
||||
git config core.editor the-editor-you-like
|
||||
|
||||
- Assume we want to merge last 3 commits, type the following commands
|
||||
|
||||
.. code:: bash
|
||||
|
||||
git rebase -i HEAD~3
|
||||
|
||||
- It will pop up an text editor. Set the first commit as ``pick``, and change later ones to ``squash``.
|
||||
- After you saved the file, it will pop up another text editor to ask you modify the combined commit message.
|
||||
- Push the changes to your fork, you need to force push.
|
||||
|
||||
.. code:: bash
|
||||
|
||||
git push --force
|
||||
|
||||
|
||||
Reset to the most recent main branch
|
||||
------------------------------------
|
||||
|
||||
You can always use git reset to reset your version to the most recent main.
|
||||
Note that **all your local changes will get lost**.
|
||||
So only do it when you do not have local changes or when your pull request just get merged.
|
||||
|
||||
.. code:: bash
|
||||
|
||||
git fetch origin main
|
||||
git reset --hard FETCH_HEAD
|
||||
|
||||
|
||||
Recover a Previous Commit after Reset
|
||||
-------------------------------------
|
||||
Sometimes we could mistakenly reset a branch to a wrong commit.
|
||||
When that happens, you can use the following command to show the list
|
||||
of recent commits
|
||||
|
||||
.. code:: bash
|
||||
|
||||
git reflog
|
||||
|
||||
Once you get the right hashtag, you can use git reset again to change
|
||||
the head to the right commit.
|
||||
|
||||
|
||||
Apply only k-Latest Commits on to the main
|
||||
------------------------------------------
|
||||
|
||||
Sometimes it is useful to only apply your k-latest changes on top of the main.
|
||||
This usually happens when you have other m-commits that are already merged
|
||||
before these k-commits. Directly rebase against the main might cause merge conflicts
|
||||
on these first m-commits (which can be safely discarded).
|
||||
|
||||
You can instead use the following command
|
||||
|
||||
.. code:: bash
|
||||
|
||||
# k is the concrete number
|
||||
# Put HEAD~2 for the last 1 commit.
|
||||
git rebase --onto upstream/main HEAD~k
|
||||
|
||||
You can then force push to the main. Note that the above command will discard
|
||||
all the commits before tha last k ones.
|
||||
|
||||
|
||||
What is the consequence of force push
|
||||
-------------------------------------
|
||||
|
||||
The previous two tips requires force push, this is because we altered the path of the commits.
|
||||
It is fine to force push to your own fork, as long as the commits changed are only yours.
|
||||
@@ -0,0 +1,53 @@
|
||||
.. 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.
|
||||
|
||||
Contributor Guide
|
||||
=================
|
||||
|
||||
TVM has been developed by community members.
|
||||
Everyone is welcomed to contribute.
|
||||
We value all forms of contributions, including, but not limited to:
|
||||
|
||||
- Code reviewing of the existing patches.
|
||||
- Documentation and usage examples
|
||||
- Community participation in forums and issues.
|
||||
- Code readability and developer guide
|
||||
|
||||
- We welcome contributions that add code comments
|
||||
to improve readability
|
||||
- We also welcome contributions to docs to explain the
|
||||
design choices of the internal.
|
||||
|
||||
- Test cases to make the codebase more robust
|
||||
- Tutorials, blog posts, talks that promote the project.
|
||||
|
||||
Here are guidelines for contributing to various aspect of the project:
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
community
|
||||
pull_request
|
||||
code_review
|
||||
committer_guide
|
||||
document
|
||||
code_guide
|
||||
testing
|
||||
git_howto
|
||||
ci
|
||||
release_process
|
||||
error_handling
|
||||
@@ -0,0 +1,267 @@
|
||||
.. 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.
|
||||
|
||||
Submit a Pull Request
|
||||
=====================
|
||||
|
||||
.. contents::
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
Guidelines
|
||||
----------
|
||||
|
||||
- We recommend authors send well scoped PRs that are easy to review and revert in case there is a problem. As such, authors should avoid merging multiple unrelated changes into a single PR
|
||||
- Before you submit a PR, please rebase your code on the most recent version of ``main``, you can do it by
|
||||
running
|
||||
|
||||
.. code:: bash
|
||||
|
||||
git remote add upstream [url to tvm repo]
|
||||
git fetch upstream
|
||||
git rebase upstream/main
|
||||
|
||||
- Make sure code passes lint checks
|
||||
|
||||
.. code:: bash
|
||||
|
||||
# Run all lint checks via pre-commit hooks
|
||||
pre-commit run --all-files
|
||||
|
||||
# Run specific linters individually
|
||||
pre-commit run ruff-check --all-files # Python lint
|
||||
pre-commit run ruff-format --all-files # Python format
|
||||
pre-commit run clang-format --all-files # C++ format
|
||||
|
||||
- Add test-cases to cover the new features or bugfix the patch introduces.
|
||||
- Document the code you wrote, see more at :ref:`doc_guide`
|
||||
- `Create a pull request <https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request>`_ and fix the problems reported by CI checks.
|
||||
- Request code reviews from other contributors and improve your patch according
|
||||
to their reviews by ``@``-ing them in your pull request. Use relevant topic
|
||||
tags in PR titles to make the affected area clear (e.g. ``[microTVM] Add a
|
||||
cool change`` and not ``a cool change for microTVM``).
|
||||
Please see the Commit Message Guideline below on the guidelines about the tags
|
||||
in a PR/commit title and how to write good PR/commit messages.
|
||||
|
||||
- To get your code reviewed quickly, we encourage you to help review others' code so they can do the favor in return.
|
||||
- Code review is a shepherding process that helps to improve contributor's code quality.
|
||||
We should treat it proactively, to improve the code as much as possible before the review.
|
||||
We highly value patches that can get in without extensive reviews.
|
||||
- The detailed guidelines and summarizes useful lessons.
|
||||
|
||||
- The PR can be merged after the reviewers approve the pull request.
|
||||
|
||||
Commit Message Guideline
|
||||
------------------------
|
||||
|
||||
Apache TVM uses the GitHub (GH) platform for patch submission and code review
|
||||
via Pull Requests (PRs). The final commit (title and body) that is merged into
|
||||
the Apache TVM main tree is composed of the PR's title and body and must be kept
|
||||
updated and reflecting the new changes in the code as per the reviews and
|
||||
discussions.
|
||||
|
||||
Although these guidelines apply essentially to the PRs’ title and body messages,
|
||||
because GH auto-generates the PR’s title and body from the commits on a given
|
||||
branch, it’s recommended to follow these guidelines right from the beginning,
|
||||
when preparing commits in general to be submitted to the Apache TVM project.
|
||||
This will ease the creation of a new PR, avoiding rework, and also will help the
|
||||
review.
|
||||
|
||||
The rules below will help to achieve uniformity that has several benefits, both
|
||||
for review and for the code base maintenance as a whole, helping you to write
|
||||
commit messages with a good quality suitable for the Apache TVM project,
|
||||
allowing fast log searches, bisecting, and so on.
|
||||
|
||||
*PR/commit title*:
|
||||
|
||||
- Guarantee a title exists (enforced);
|
||||
- Don’t use GitHub usernames in the title, like @username (enforced);
|
||||
- A tag must be present as a hint about what component(s) of the code
|
||||
the PRs / commits “touch” (enforced). For example [BugFix], [CI], [microTVM],
|
||||
and [TVMC]. Tags go between square brackets and appear first in the title. If
|
||||
more than one tag exist, multiple brackets should be used, like [BugFix][CI].
|
||||
The case recommended for tags, in geral, is the upper camel case. For example,
|
||||
prefer the forms [Fix], [BugFix], and [Docker] instead of [fix], [bug_fix],
|
||||
and [docker]. Acronyms should be kept as such so, for example, use [CI] and
|
||||
[TVMC] instead of [ci] and [tvmc]. Tags help reviewers to identify the PRs
|
||||
they can/want to review and also help the release folks when generating the
|
||||
release notes;
|
||||
- Use an imperative mood. Avoid titles like “Added operator X” and “Updated
|
||||
image Y in the CI”, instead use the forms “Add feature X” and “Update image Y
|
||||
in the CI” instead;
|
||||
- Observe proper use of caps at the beginning (uppercase for the first letter)
|
||||
and for acronyms, like, for instance, TVM, FVP, OpenCL. Hence instead of
|
||||
“fix tvm use of opencl library”, write it as “Fix TVM use of OpenCL library”;
|
||||
- Do not put a period at the end of the title.
|
||||
|
||||
*PR/commit body*:
|
||||
|
||||
- Guarantee a body exists (enforced);
|
||||
- Don’t use GitHub usernames in body text, like @username (enforced);
|
||||
- Avoid “bullet” commit message bodies: “bullet” commit message bodies are not
|
||||
bad per se, but “bullet” commit messages without any description or
|
||||
explanation is likely as bad as commits without any description, rationale,
|
||||
or explanation in the body.
|
||||
|
||||
For minor deviations from these guidelines, the community will normally favor
|
||||
reminding the contributor of this policy over reverting or blocking a commit /
|
||||
PR.
|
||||
|
||||
Commits and PRs without a title and/or a body are not considered minor
|
||||
deviations from these guidelines and hence must be avoided.
|
||||
|
||||
Most importantly, the contents of the commit message, especially the body,
|
||||
should be written to convey the intention of the change, so it should avoid
|
||||
being vague. For example, commits with a title like “Fix”, “Cleanup”, and
|
||||
“Fix flaky test” and without any body text should be avoided. Also, for the
|
||||
review, it will leave the reviewer wondering about what exactly was fixed or
|
||||
changed and why the change is necessary, slowing the review.
|
||||
|
||||
Below is an example that can be used as a model:
|
||||
|
||||
::
|
||||
|
||||
[microTVM] Zephyr: Remove zephyr_board option from build, flash, and open_transport methods
|
||||
|
||||
Currently it’s necessary to pass the board type via ‘zephyr_board’ option to
|
||||
the Project API build, flash, and open_transport methods.
|
||||
|
||||
However, since the board type is already configured when the project is
|
||||
created (i.e. when the generate_project method is called), it’s possible to
|
||||
avoid this redundancy by obtaining the board type from the project
|
||||
configuration files.
|
||||
|
||||
This commit adds code to obtain the board type from the project CMake files,
|
||||
removing this option from build, flash, and open_transport methods, so it’s
|
||||
only necessary to specify the ‘zephyr_board’ option when calling
|
||||
generate_project.
|
||||
|
||||
This commit also moves the ‘verbose’ and ‘west_cmd’ options from ‘build’
|
||||
method to ‘generate_project’, reducing further the number of required options
|
||||
when building a project, since the ‘build’ method is usually called more often
|
||||
than the ‘generate_project’.
|
||||
|
||||
After a new PR is created and the review starts it’s common that reviewers will
|
||||
request changes. Usually the author will address the reviewers’ comments and
|
||||
push additional commits on top of the initial ones. For these additional commits
|
||||
there is no recommendation regarding the commit messages. However if the
|
||||
additional commits render the PR title and/or body outdated then it's the
|
||||
author's responsibility to keep the PR title and body in sync with new changes
|
||||
in the code and updated the PR title and body accordingly (remember that the PR
|
||||
title and body will be used to compose the final commit message that will land
|
||||
in the main tree).
|
||||
|
||||
Committers will seek to fix any issues with the commit message prior to
|
||||
committing but they retain the right to inform the author of the rules and
|
||||
encourage them to follow them in future. Also, they retain the right to ask to
|
||||
the author to update the PR title and/or body when they are not correctly
|
||||
updated or fixed.
|
||||
|
||||
CI Environment
|
||||
--------------
|
||||
We use Docker images to create stable CI environments that can be deployed to multiple machines.
|
||||
Follow the steps in `this issue template <https://github.com/apache/tvm/issues/new?assignees=&labels=&template=ci-image.md&title=%5BCI+Image%5D+>`_
|
||||
to update a CI Docker image.
|
||||
|
||||
.. _pr-testing:
|
||||
|
||||
Testing
|
||||
-------
|
||||
Even though we have hooks to run unit tests automatically for each pull request, it's always recommended to run unit tests
|
||||
locally beforehand to reduce reviewers' burden and speedup review process.
|
||||
|
||||
Docker (recommended)
|
||||
^^^^^^^^^^^^^^^^^^^^
|
||||
``tests/scripts/ci.py`` replicates the CI environment locally and provides a user-friendly interface.
|
||||
The same Docker images and scripts used in CI are used directly to run tests. It also deposits builds
|
||||
in different folders so you can maintain multiple test environments without rebuilding from scratch
|
||||
each time (e.g. you can test a change in CPU and GPU while retaining incremental rebuilds).
|
||||
|
||||
.. code:: bash
|
||||
|
||||
# see all available platforms
|
||||
python tests/scripts/ci.py --help
|
||||
python tests/scripts/ci.py cpu --help
|
||||
|
||||
# run the CPU build in the ci_cpu docker container (build will be left in
|
||||
# the build-cpu/ folder)
|
||||
# note: the CPU and GPU Docker images are quite large and may take some
|
||||
# time to download on their first use
|
||||
python tests/scripts/ci.py cpu
|
||||
|
||||
# run the CPU build in the ci_cpu docker container and then run unittests
|
||||
python tests/scripts/ci.py cpu --unittest
|
||||
|
||||
# quickly iterate by running a specific test and skipping the rebuild each time
|
||||
python tests/scripts/ci.py cpu --skip-build --tests tests/python/s_tir/schedule/test_tir_schedule_rolling_buffer.py::test_upscale
|
||||
|
||||
# run the CPU build and drop into a shell in the container
|
||||
python tests/scripts/ci.py cpu --interactive
|
||||
|
||||
We regularly update our docker images and, over time, stale images may unnecessarily consume disk
|
||||
space. You can remove stale images that aren't used in the presently checked-out branch plus any
|
||||
other worktrees using the following command:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
docker/clear-stale-images.sh
|
||||
|
||||
Consult the ``--help`` for more options.
|
||||
|
||||
C++ (local)
|
||||
^^^^^^^^^^^
|
||||
|
||||
Running the C++ tests requires installation of gtest, following the instructions in
|
||||
:ref:`install-from-source-cpp-tests`
|
||||
|
||||
|
||||
.. code:: bash
|
||||
|
||||
# assume you are in tvm source root
|
||||
TVM_ROOT=`pwd`
|
||||
|
||||
./tests/scripts/task_cpp_unittest.sh
|
||||
|
||||
Python (local)
|
||||
^^^^^^^^^^^^^^
|
||||
Necessary dependencies:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
pip install --user pytest pytest-xdist Cython
|
||||
|
||||
If you want to run all tests:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
# build tvm (see install-from-source for CMake build instructions)
|
||||
cd build && cmake .. && cmake --build . --parallel $(nproc) && cd ..
|
||||
|
||||
python -m pytest -vvs -n auto tests/python
|
||||
|
||||
If you want to run a single test:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
# let python know where to find tvm related libraries
|
||||
export PYTHONPATH=python
|
||||
rm -rf python/tvm/*.pyc python/tvm/*/*.pyc python/tvm/*/*/*.pyc
|
||||
|
||||
python -m pytest -v tests/python/tirx-transform/test_tir_transform_storage_rewrite.py
|
||||
|
||||
# Additionally if you want to run a single test, for example test_add_one_2d inside a file.
|
||||
python -m pytest -v -k "test_add_one_2d" tests/python/relax/test_frontend_tflite.py
|
||||
@@ -0,0 +1,314 @@
|
||||
.. 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.
|
||||
|
||||
.. _release_process:
|
||||
|
||||
Release Process
|
||||
===============
|
||||
|
||||
.. contents::
|
||||
:depth: 2
|
||||
:local:
|
||||
|
||||
The release manager role in TVM means you are responsible for a few different things:
|
||||
|
||||
- Preparing release notes
|
||||
- Preparing your setup
|
||||
- Preparing for release candidates
|
||||
|
||||
- Cutting a release branch
|
||||
- Informing the community of timing
|
||||
- Making any necessary code changes in that branch (versions are derived from Git tags, so there are no manual version-number edits)
|
||||
|
||||
- Running the voting process for a release
|
||||
|
||||
- Creating release candidates
|
||||
- Calling votes and triaging issues
|
||||
|
||||
- Finalizing and posting a release:
|
||||
|
||||
- Updating the TVM website
|
||||
- Finalizing release notes
|
||||
- Announcing the release
|
||||
|
||||
|
||||
Versioning
|
||||
----------
|
||||
|
||||
TVM's version is derived automatically from the most recent Git tag by
|
||||
`setuptools_scm <https://setuptools-scm.readthedocs.io/>`_ at build time (configured
|
||||
under ``[tool.setuptools_scm]`` in ``pyproject.toml``). There are **no version numbers
|
||||
to edit by hand** in ``pyproject.toml``, ``python/tvm/libinfo.py`` or
|
||||
``include/tvm/runtime/base.h``; releasing is driven entirely by pushing Git tags:
|
||||
|
||||
- ``main`` carries a ``vMAJOR.MINOR.devN`` tag (e.g. ``v0.7.dev0``). Commits after it
|
||||
are versioned ``0.7.devN`` where ``N`` is the number of commits since the tag. This
|
||||
is why the next dev tag must be pushed on ``main`` when a release branch is cut:
|
||||
without it, follow-up commits would keep deriving their version from the previous
|
||||
cycle's tag.
|
||||
- A release branch (e.g. ``v0.6``) is tagged ``v0.6.0.rc0`` for a candidate (wheel
|
||||
version ``0.6.0rc0``) and ``v0.6.0`` for the formal release (wheel version
|
||||
``0.6.0``). Release wheels are built on the exact tag, so the version is the tag
|
||||
itself.
|
||||
|
||||
The legacy ``version.py`` stamping script has been removed. ``web/package.json`` (npm,
|
||||
a separate ecosystem) is no longer auto-stamped; bump it by hand when starting a new dev
|
||||
cycle or cutting a release. ``docs/conf.py`` reads ``tvm.__version__`` directly.
|
||||
|
||||
|
||||
Prepare the Release Notes
|
||||
-------------------------
|
||||
|
||||
Release note contains new features, improvement, bug fixes, known issues and deprecation, etc. TVM provides `monthly dev report <https://discuss.tvm.apache.org/search?q=TVM%20Monthly%20%23Announcement>`_ collects developing progress each month. It could be helpful to who writes the release notes.
|
||||
|
||||
It is recommended to open a GitHub issue to collect feedbacks for the release note draft before cutting the release branch. See the scripts in ``tests/scripts/release`` for some starting points.
|
||||
|
||||
|
||||
Prepare the Release Candidate
|
||||
-----------------------------
|
||||
|
||||
There may be some code changes necessary on the release branch before the release (for
|
||||
example cherry-picked fixes). Version numbers are derived from the release tags (see
|
||||
`Versioning`_), so there are no version numbers to update by hand.
|
||||
|
||||
|
||||
Prepare the GPG Key
|
||||
-------------------
|
||||
|
||||
You can skip this section if you have already uploaded your key.
|
||||
|
||||
After generating the gpg key, you need to upload your key to a public key server. Please refer to https://www.apache.org/dev/openpgp.html#generate-key for details.
|
||||
|
||||
If you want to do the release on another machine, you can transfer your gpg key to that machine via the ``gpg --export`` and ``gpg --import`` commands.
|
||||
|
||||
The last step is to update the KEYS file with your code signing key https://www.apache.org/dev/openpgp.html#export-public-key. Check in the changes to the TVM main branch, as well as ASF SVN,
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# the --depth=files will avoid checkout existing folders
|
||||
svn co --depth=files "https://dist.apache.org/repos/dist/dev/tvm" svn-tvm
|
||||
cd svn-tvm
|
||||
# edit KEYS file
|
||||
svn ci --username $ASF_USERNAME --password "$ASF_PASSWORD" -m "Update KEYS"
|
||||
# update downloads.apache.org (note that only PMC members can update the dist/release directory)
|
||||
svn rm --username $ASF_USERNAME --password "$ASF_PASSWORD" https://dist.apache.org/repos/dist/release/tvm/KEYS -m "Update KEYS"
|
||||
svn cp --username $ASF_USERNAME --password "$ASF_PASSWORD" https://dist.apache.org/repos/dist/dev/tvm/KEYS https://dist.apache.org/repos/dist/release/tvm/ -m "Update KEYS"
|
||||
|
||||
|
||||
Cut a Release Candidate
|
||||
-----------------------
|
||||
|
||||
To cut a release candidate for the ``v0.6`` release:
|
||||
|
||||
#. On the ``main`` commit that should be the last one included in the release, push the
|
||||
**next** dev-cycle tag ``v0.7.dev0``. This tag is what makes subsequent ``main``
|
||||
commits versioned ``0.7.devN`` (see `Versioning`_), so it must be pushed *before*
|
||||
branching.
|
||||
#. Cut the release branch off that same commit. Branches are named with the base
|
||||
release version without the patch, e.g. ``v0.6`` for the ``v0.6.0`` release.
|
||||
#. Push the first release-candidate tag ``v0.6.0.rc0`` on the release branch. CI then
|
||||
builds the candidate wheel (version ``0.6.0rc0``) for PyPI/TestPyPI testing. Keep
|
||||
this tag on a ``v0.6`` branch commit that is **not** also the ``v0.7.dev0``-tagged
|
||||
branch point: when two tags share one commit, which one ``setuptools_scm`` picks is
|
||||
fragile, so put the candidate tag on a release-prep commit on the branch.
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
git clone https://github.com/apache/tvm.git
|
||||
cd tvm/
|
||||
|
||||
# 1. Tag the next dev cycle on main (drives main's 0.7.devN version),
|
||||
# on the last commit to be included in the release.
|
||||
git checkout <last-commit-for-release>
|
||||
git tag v0.7.dev0
|
||||
git push origin refs/tags/v0.7.dev0
|
||||
|
||||
# 2. Cut the release branch off that same commit.
|
||||
git branch v0.6
|
||||
git push --set-upstream origin v0.6
|
||||
|
||||
# 3. Tag the first release candidate on the release branch. Keep this tag on a
|
||||
# release-prep commit, NOT the v0.7.dev0-tagged branch point, so the two tags
|
||||
# never share a commit.
|
||||
git checkout v0.6
|
||||
# ... make any release-prep commits (release notes, etc.) here ...
|
||||
git tag v0.6.0.rc0
|
||||
git push origin refs/tags/v0.6.0.rc0
|
||||
|
||||
The wheel/distribution version is derived from these tags by ``setuptools_scm`` at build
|
||||
time, so no source files need editing and you no longer run ``version.py`` to stamp the
|
||||
version (see `Versioning`_).
|
||||
|
||||
Go to the GitHub repositories "releases" tab and click "Draft a new release",
|
||||
|
||||
- Verify the release by checking the version numbers and ensuring that TVM can build and run the unit tests.
|
||||
- Provide the release tag in the form of ``v1.0.0.rc0`` where 0 means it's the first release candidate. The tag must match this pattern ``v[0-9]+\.[0-9]+\.[0-9]+\.rc[0-9]`` exactly!
|
||||
- Select the commit by clicking Target: branch > Recent commits > $commit_hash
|
||||
- Copy and paste release note draft into the description box
|
||||
- Select "This is a pre-release"
|
||||
- Click "Publish release"
|
||||
|
||||
Notice that one can still apply changes to the branch after the cut, while the tag is fixed. If any change is required for this release, a new tag has to be created.
|
||||
|
||||
Remove previous release candidate (if applied),
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
git push --delete origin v0.6.0.rc1
|
||||
|
||||
Create source code artifacts,
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# Replace v0.6.0 with the relevant version
|
||||
git clone git@github.com:apache/tvm.git apache-tvm-src-v0.6.0
|
||||
cd apache-tvm-src-v0.6.0
|
||||
git checkout v0.6
|
||||
git submodule update --init --recursive
|
||||
git checkout v0.6.0.rc0
|
||||
rm -rf .DS_Store
|
||||
find . -name ".git*" -print0 | xargs -0 rm -rf
|
||||
cd ..
|
||||
brew install gnu-tar
|
||||
gtar -czvf apache-tvm-src-v0.6.0.rc0.tar.gz apache-tvm-src-v0.6.0
|
||||
|
||||
Use your GPG key to sign the created artifact. First make sure your GPG is set to use the correct private key,
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
$ cat ~/.gnupg/gpg.conf
|
||||
default-key F42xxxxxxxxxxxxxxx
|
||||
|
||||
Create GPG signature as well as the hash of the file,
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
gpg --armor --output apache-tvm-src-v0.6.0.rc0.tar.gz.asc --detach-sig apache-tvm-src-v0.6.0.rc0.tar.gz
|
||||
shasum -a 512 apache-tvm-src-v0.6.0.rc0.tar.gz > apache-tvm-src-v0.6.0.rc0.tar.gz.sha512
|
||||
|
||||
|
||||
Update TVM Version on ``main``
|
||||
------------------------------
|
||||
|
||||
The next dev-cycle tag pushed on ``main`` during the cut (step 1 above, e.g. ``v0.7.dev0``)
|
||||
is what gives ``main`` its ``0.7.devN`` version — ``setuptools_scm`` derives it from that
|
||||
tag, so there are **no source version numbers to bump** (the old two-commit ``[Dont Squash]``
|
||||
bump and the ``python version.py`` stamping step are no longer needed). Make sure that tag
|
||||
sits on the ``main`` commit immediately after the last one included in the release branch;
|
||||
it is required so that nightly/dev packages built from ``main`` carry the correct
|
||||
``0.7.devN`` version.
|
||||
|
||||
Upload the Release Candidate
|
||||
----------------------------
|
||||
|
||||
Edit the release page on GitHub and upload the artifacts created by the previous steps.
|
||||
|
||||
The release manager also needs to upload the artifacts to ASF SVN,
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
# the --depth=files will avoid checkout existing folders
|
||||
svn co --depth=files "https://dist.apache.org/repos/dist/dev/tvm" svn-tvm
|
||||
cd svn-tvm
|
||||
mkdir tvm-v0.6.0-rc0
|
||||
# copy files into it
|
||||
svn add tvm-0.6.0-rc0
|
||||
svn ci --username $ASF_USERNAME --password "$ASF_PASSWORD" -m "Add RC"
|
||||
|
||||
|
||||
Cherry-Picking
|
||||
--------------
|
||||
After a release branch has been cut but before the release has been voted on, the release manager may cherry-pick commits from ``main``. Since release branches are protected on GitHub, to merge this fixes into the release branch (e.g. ``v0.11``), the release manager must file a PR with the cherry-picked changes against the release branch. The PR should roughly match the original one from ``main`` with extra details on why the commit is being cherry-picked. The community then goes through a normal review and merge process for these PRs. Note that these PRs against the release branches must be `signed <https://docs.github.com/en/authentication/managing-commit-signature-verification/signing-commits>`_.
|
||||
|
||||
|
||||
Call a Vote on the Release Candidate
|
||||
------------------------------------
|
||||
|
||||
The first voting takes place on the Apache TVM developers list (dev@tvm.apache.org). To get more attention, one can create a GitHub issue start with "[VOTE]" instead, it will be mirrored to dev@ automatically. Look at past voting threads to see how this proceeds. The email should follow this format.
|
||||
|
||||
- Provide the link to the draft of the release notes in the email
|
||||
- Provide the link to the release candidate artifacts
|
||||
- Make sure the email is in text format and the links are correct
|
||||
|
||||
For the dev@ vote, there must be at least 3 binding +1 votes and more +1 votes than -1 votes. Once the vote is done, you should also send out a summary email with the totals, with a subject that looks something like [VOTE][RESULT] ....
|
||||
|
||||
In ASF, votes are open at least 72 hours (3 days). If you don't get enough number of binding votes within that time, you cannot close the voting deadline. You need to extend it.
|
||||
|
||||
If the vote fails, the community needs to modify the release accordingly: create a new release candidate and re-run the voting process.
|
||||
|
||||
|
||||
Post the Release
|
||||
----------------
|
||||
|
||||
After the vote passes, to upload the binaries to Apache mirrors, you move the binaries from dev directory (this should be where they are voted) to release directory. This "moving" is the only way you can add stuff to the actual release directory. (Note: only PMC can move to release directory)
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
export SVN_EDITOR=vim
|
||||
svn mkdir https://dist.apache.org/repos/dist/release/tvm
|
||||
svn mv https://dist.apache.org/repos/dist/dev/tvm/tvm-v0.6.0-rc2 https://dist.apache.org/repos/dist/release/tvm/tvm-v0.6.0
|
||||
|
||||
# If you've added your signing key to the KEYS file, also update the release copy.
|
||||
svn co --depth=files "https://dist.apache.org/repos/dist/release/tvm" svn-tvm
|
||||
curl "https://dist.apache.org/repos/dist/dev/tvm/KEYS" > svn-tvm/KEYS
|
||||
(cd svn-tvm && svn ci --username $ASF_USERNAME --password "$ASF_PASSWORD" -m"Update KEYS")
|
||||
|
||||
Remember to create a new release TAG (v0.6.0 in this case) on GitHub and remove the pre-release candidate TAG.
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
git push --delete origin v0.6.0.rc2
|
||||
|
||||
|
||||
Update the TVM Website
|
||||
----------------------
|
||||
|
||||
The website repository is located at `https://github.com/apache/tvm-site <https://github.com/apache/tvm-site>`_. Modify the download page to include the release artifacts as well as the GPG signature and SHA hash. Since TVM's docs are continually updated, upload a fixed version of the release docs. If CI has deleted the docs from the release by the time you go to update the website, you can restart the CI build for the release branch on Jenkins. See the example code below for a starting point.
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
git clone https://github.com/apache/tvm-site.git
|
||||
pushd tvm-site
|
||||
git checkout asf-site
|
||||
pushd docs
|
||||
|
||||
# make release docs directory
|
||||
mkdir v0.9.0
|
||||
pushd v0.9.0
|
||||
|
||||
# download the release docs from CI
|
||||
# find this URL by inspecting the CI logs for the most recent build of the release branch
|
||||
curl -LO https://tvm-jenkins-artifacts-prod.s3.us-west-2.amazonaws.com/tvm/v0.9.0/1/docs/docs.tgz
|
||||
tar xf docs.tgz
|
||||
rm docs.tgz
|
||||
|
||||
# add the docs and push
|
||||
git add .
|
||||
git commit -m "Add v0.9.0 docs"
|
||||
git push
|
||||
|
||||
|
||||
Afterwards, modify the `downloads page <https://tvm.apache.org/download>`_ to support the latest release. An example of how to do this is `here <https://github.com/apache/tvm-site/pull/38>`_.
|
||||
|
||||
Post the Announcement
|
||||
---------------------
|
||||
|
||||
Send out an announcement email to announce@apache.org, and dev@tvm.apache.org. The announcement should include the link to release note and download page.
|
||||
|
||||
Patch Releases
|
||||
--------------
|
||||
Patch releases should be reserved for critical bug fixes. Patch releases must go through the same process as normal releases, with the option at the release manager's discretion of a shortened release candidate voting window of 24 hours to ensure that fixes are delivered quickly. A patch release is cut purely by tagging the release base branch (e.g. ``v0.11``): push ``v0.11.1.rc0`` for the candidate and ``v0.11.1`` for the formal patch release; no source version numbers need bumping.
|
||||
@@ -0,0 +1,312 @@
|
||||
.. 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.
|
||||
|
||||
Testing TVM
|
||||
===========
|
||||
|
||||
This page describes how to write and run Python tests for TVM,
|
||||
including the target parametrization utilities used by CI.
|
||||
|
||||
Python Target Parametrization
|
||||
-----------------------------
|
||||
|
||||
Summary
|
||||
~~~~~~~
|
||||
|
||||
For any supported runtime, TVM should produce numerically
|
||||
correct results. Therefore, when writing unit tests that validate
|
||||
the numeric output, these unit tests should be run on all supported
|
||||
runtimes. Since this is a very common use case, TVM has helper
|
||||
functions to parametrize unit tests such that they will run on all
|
||||
targets that are enabled and have a compatible device.
|
||||
|
||||
A single Python function in the test suite can expand to several
|
||||
parameterized unit tests, each of which tests a single target device.
|
||||
In order for a test to be run, all of the following must be true.
|
||||
|
||||
- The test exists in a file or directory that has been passed to
|
||||
`pytest`.
|
||||
|
||||
- The pytest marks applied to the function, either explicitly or
|
||||
through target parametrization, must be compatible with the
|
||||
expression passed to pytest's `-m` argument.
|
||||
|
||||
- For parametrized tests using the `target` fixture, the target must
|
||||
appear in the environment variable `TVM_TEST_TARGETS`.
|
||||
|
||||
- For parametrized tests using the `target` fixture, the build
|
||||
configuration in `config.cmake` must enable the corresponding
|
||||
runtime.
|
||||
|
||||
Unit-Test File Contents
|
||||
~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. _pytest-marks: https://docs.pytest.org/en/stable/how-to/mark.html
|
||||
|
||||
The recommended way to run a test on multiple targets is to parametrize
|
||||
over ``target`` with ``@pytest.mark.parametrize``. Tag each GPU target
|
||||
with ``pytest.mark.gpu`` so the CI routes it to a GPU node, skip a target
|
||||
that cannot run on the current machine with
|
||||
:py:func:`tvm.testing.device_enabled`, and obtain its device with
|
||||
``tvm.device(target)``. The function is run once per target, the
|
||||
success/failure of each is reported separately, and a target whose device
|
||||
is disabled in ``config.cmake`` or absent from the machine is reported as
|
||||
skipped.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"target",
|
||||
["llvm", pytest.param("cuda", marks=pytest.mark.gpu)],
|
||||
)
|
||||
def test_function(target):
|
||||
if not tvm.testing.device_enabled(target):
|
||||
pytest.skip(f"{target} not enabled")
|
||||
dev = tvm.device(target)
|
||||
# Test code goes here
|
||||
|
||||
For a test that only applies to a single target, omit the parametrization
|
||||
and gate the test with ``@pytest.mark.skipif`` (plus ``@pytest.mark.gpu``
|
||||
for a GPU target):
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(
|
||||
not tvm.testing.device_enabled("cuda"), reason="cuda not enabled"
|
||||
)
|
||||
def test_function():
|
||||
target = "cuda"
|
||||
dev = tvm.device(target)
|
||||
# Test code goes here
|
||||
|
||||
To exclude a target, leave it out of the parametrize list. To mark a
|
||||
target as expected to fail, wrap it with
|
||||
``pytest.param("target", marks=pytest.mark.xfail(reason=...))``.
|
||||
|
||||
Additional parameters can be combined with the target parametrization by
|
||||
stacking ``@pytest.mark.parametrize`` decorators, or by listing tuples of
|
||||
arguments. Tag the GPU rows with ``pytest.mark.gpu`` and skip in the body
|
||||
as above:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@pytest.mark.parametrize("target,impl", [
|
||||
("llvm", cpu_implementation),
|
||||
pytest.param("cuda", gpu_implementation_small_batch, marks=pytest.mark.gpu),
|
||||
pytest.param("cuda", gpu_implementation_large_batch, marks=pytest.mark.gpu),
|
||||
])
|
||||
def test_function(target, impl):
|
||||
if not tvm.testing.device_enabled(target):
|
||||
pytest.skip(f"{target} not enabled")
|
||||
dev = tvm.device(target)
|
||||
# Test code goes here
|
||||
|
||||
|
||||
Tests gate on hardware and carry metadata using
|
||||
`pytest marks <pytest-marks>`_. The most frequently applied
|
||||
marks are as follows.
|
||||
|
||||
- ``@pytest.mark.gpu`` - Tags a function as using GPU
|
||||
capabilities. This has no effect on its own, but can be paired with
|
||||
the command-line arguments ``-m gpu`` or ``-m 'not gpu'`` to restrict
|
||||
which tests pytest will execute. Apply it to any test that needs a
|
||||
GPU so that the CI runs it only on GPU nodes.
|
||||
|
||||
- ``@pytest.mark.skipif(not tvm.testing.env.has_X(), reason=...)`` -
|
||||
Skips a test when a required runtime or hardware feature is not
|
||||
available. The :py:mod:`tvm.testing.env` module exposes one memoized
|
||||
probe per capability (e.g. ``has_cuda()``, ``has_rocm()``,
|
||||
``has_vulkan()``, ``has_gpu()``, ``has_llvm()``), each of which
|
||||
returns ``False`` when the runtime is disabled in ``config.cmake`` or
|
||||
no compatible device is present. Pair it with ``@pytest.mark.gpu``
|
||||
for tests that use the GPU::
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not tvm.testing.env.has_cuda(), reason="need cuda")
|
||||
def test_cuda_vectorize_add():
|
||||
# Test code goes here
|
||||
|
||||
- ``pytest.importorskip("package_name")`` - Skips a test (or the whole
|
||||
module, when called at import time) if an optional Python package is
|
||||
not installed. Use this instead of a ``skipif`` for package
|
||||
dependencies.
|
||||
|
||||
Tests that execute on a local GPU must put the complete live-device
|
||||
lifetime in a small callback passed to
|
||||
:py:func:`tvm.testing.run_with_gpu_lock`. Target construction and
|
||||
compilation remain outside so that pytest-xdist workers can compile in
|
||||
parallel. Device creation, allocation, execution, synchronization,
|
||||
host conversion, result checks, and child-process teardown remain inside
|
||||
the callback so no device-backed object outlives the lock.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not tvm.testing.env.has_cuda(), reason="need cuda")
|
||||
def test_cuda_add_one():
|
||||
target = tvm.target.Target("cuda -arch=sm_90")
|
||||
executable = tvm.compile(make_add_one_module(), target)
|
||||
host_input = np.arange(16, dtype="float32")
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
device_input = tvm.runtime.tensor(host_input, dev)
|
||||
device_output = tvm.runtime.empty(host_input.shape, "float32", dev)
|
||||
executable(device_input, device_output)
|
||||
dev.sync()
|
||||
tvm.testing.assert_allclose(device_output.numpy(), host_input + 1)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
The wrapper uses the existing :py:class:`tvm_ffi.utils.FileLock` with a
|
||||
persistent machine-local path. A process exit releases the kernel lock;
|
||||
the remaining file is not stale ownership. Test startup must never
|
||||
delete or rotate it, because another process could then lock a different
|
||||
inode. Set ``TVM_TEST_LOCK_DIR`` only when all cooperating processes need
|
||||
an explicitly configured shared machine-local directory. The default
|
||||
temporary path coordinates processes running as the same user. Multi-user
|
||||
runners sharing a GPU must use one administrator-provisioned directory and
|
||||
persistent lock file that every contender can write, or enforce exclusivity
|
||||
through the runner. A per-user lock path cannot protect a GPU shared across
|
||||
users because each user would lock a different file.
|
||||
|
||||
There also exists a ``tvm.testing.enabled_targets()`` that returns
|
||||
all targets that are enabled and runnable on the current machine,
|
||||
based on the environment variable ``TVM_TEST_TARGETS``, the build
|
||||
configuration, and the physical hardware present. Some legacy tests
|
||||
explicitly loop over the targets returned from ``enabled_targets()``,
|
||||
but this style should not be used for new tests. The pytest output
|
||||
for this style silently skips runtimes that are disabled in
|
||||
``config.cmake``, or do not have a device on which they can run. In
|
||||
addition, the test halts on the first target to fail, which is
|
||||
ambiguous as to whether the error occurs on a particular target, or on
|
||||
every target.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Old style, do not use.
|
||||
def test_function():
|
||||
for target, dev in tvm.testing.enabled_targets():
|
||||
# Test code goes here
|
||||
|
||||
|
||||
|
||||
Running Locally
|
||||
~~~~~~~~~~~~~~~
|
||||
|
||||
To run the Python unit tests locally, use the command ``pytest`` in
|
||||
the ``${TVM_HOME}`` directory.
|
||||
|
||||
- Environment variables
|
||||
- ``TVM_TEST_TARGETS`` should be a semicolon-separated list of
|
||||
targets to run. If unset, will default to the targets defined in
|
||||
``tvm.testing.DEFAULT_TEST_TARGETS``.
|
||||
|
||||
Note: If ``TVM_TEST_TARGETS`` does not contain any targets that
|
||||
are both enabled, and have an accessible device of that type,
|
||||
then the tests will fall back to running on the ``llvm`` target
|
||||
only.
|
||||
|
||||
- ``TVM_LIBRARY_PATH`` should be a path to the ``libtvm.so``
|
||||
library. This can be used, for example, to run tests using a
|
||||
debug build. If unset, will search for ``libtvm.so`` relative to
|
||||
the TVM source directory.
|
||||
|
||||
- Command-line arguments
|
||||
|
||||
- Passing a path to a folder or file will run only the unit tests
|
||||
in that folder or file. This can be useful, for example, to
|
||||
avoid running tests located in ``tests/python/contrib`` on a
|
||||
system without a specific backend installed.
|
||||
|
||||
- The ``-m`` argument only runs unit tests that are tagged with a
|
||||
specific pytest marker. The most frequent usage is to use
|
||||
``-m gpu`` to run only tests that are marked with
|
||||
``@pytest.mark.gpu`` and use a GPU to run. It can also be used
|
||||
to run only tests that do not use a GPU, by passing ``not gpu``
|
||||
as the marker expression to ``-m``.
|
||||
|
||||
Note: This filtering takes place after the selection of targets
|
||||
based on the ``TVM_TEST_TARGETS`` environment variable. Even if
|
||||
``-m gpu`` is specified, if ``TVM_TEST_TARGETS`` does not
|
||||
contain GPU targets, no GPU tests will be run.
|
||||
|
||||
Running in a Local Docker Container
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. _tlcpack: https://hub.docker.com/u/tlcpack
|
||||
|
||||
The ``docker/bash.sh`` script can be used to run unit tests inside the
|
||||
same docker image as is used by the CI. The first argument should
|
||||
specify which docker image to run (e.g. ``docker/bash.sh ci_gpu``).
|
||||
Allowed image names are defined in ``ci/jenkins/data.py`` in the TVM source directory,
|
||||
and map to images at `tlcpack`_.
|
||||
|
||||
If no additional arguments are given, the docker image will be loaded
|
||||
with an interactive bash session. If a script is passed as an
|
||||
optional argument (e.g. ``docker/bash.sh ci_gpu tests/scripts/task_python_unittest.sh``), then that script will be
|
||||
executed inside the docker image.
|
||||
|
||||
Note: The docker images contain all system dependencies, but do not
|
||||
include the ``build/config.cmake`` configuration file for those
|
||||
systems. The TVM source directory is used as the home directory of
|
||||
the docker image, and so this will default to using the same
|
||||
config/build directories as the local config. One solution is to
|
||||
maintain separate ``build_local`` and ``build_docker`` directories,
|
||||
and make a symlink from ``build`` to the appropriate folder when
|
||||
entering/exiting docker.
|
||||
|
||||
Running in CI
|
||||
~~~~~~~~~~~~~
|
||||
|
||||
Everything in the CI starts from the task definitions present in the
|
||||
Jenkinsfile. This includes defining which docker image gets used,
|
||||
what the compile-time configuration is, and which tests are included
|
||||
in which stages.
|
||||
|
||||
- Docker images
|
||||
|
||||
Each task of the Jenkinsfile (e.g. 'BUILD: CPU') makes calls to
|
||||
``docker/bash.sh``. The argument following the call to
|
||||
docker/bash.sh defines the docker image in CI, just as it does
|
||||
locally.
|
||||
|
||||
- Compile-time configuration
|
||||
|
||||
The docker image does not have the ``config.cmake`` file built into
|
||||
it, so this is the first step in each of the ``BUILD`` tasks. This
|
||||
is done using the ``tests/scripts/task_config_build_*.sh`` scripts.
|
||||
Which script is used depends on the build being tested, and is
|
||||
specified in the Jenkinsfile.
|
||||
|
||||
Each ``BUILD`` task concludes by packing a library for use in later
|
||||
tests.
|
||||
|
||||
- Which tests run
|
||||
|
||||
The ``Unit Test`` stage of the Jenkinsfile determines how ``pytest``
|
||||
is called. Each task starts by unpacking a
|
||||
compiled library that was previous compiled in the ``BUILD`` stage,
|
||||
then runs a test script
|
||||
(e.g. ``tests/scripts/task_python_unittest.sh``). These scripts set
|
||||
the files/folders and command-line options that are passed to
|
||||
``pytest``.
|
||||
|
||||
Several of these scripts include the ``-m gpu`` option, which
|
||||
restricts the tests to only run tests that include the
|
||||
``@pytest.mark.gpu`` mark.
|
||||
@@ -0,0 +1,73 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
.. _relax-abstraction:
|
||||
|
||||
Graph Abstraction for ML Models
|
||||
-------------------------------
|
||||
Graph abstraction is a key technique used in machine learning (ML) compilers
|
||||
to represent and reason about the structure and data flow of ML models. By
|
||||
abstracting the model into a graph representation, the compiler can perform
|
||||
various optimizations to improve performance and efficiency. This tutorial will
|
||||
cover the basics of graph abstraction, its key elements of Relax IR, and how it enables optimization in ML compilers.
|
||||
|
||||
What is Graph Abstraction?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
Graph abstraction is the process of representing an ML model as a directed graph,
|
||||
where the nodes represent computational operations (e.g., matrix multiplication,
|
||||
convolution) and the edges represent the flow of data between these operations.
|
||||
This abstraction allows the compiler to analyze the dependencies and
|
||||
relationships between different parts of the model.
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm.script import relax as R
|
||||
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((1, 784), dtype="float32"),
|
||||
weight: R.Tensor((784, 256), dtype="float32"),
|
||||
bias: R.Tensor((256,), dtype="float32"),
|
||||
) -> R.Tensor((1, 256), dtype="float32"):
|
||||
with R.dataflow():
|
||||
lv0 = R.matmul(x, weight)
|
||||
lv1 = R.add(lv0, bias)
|
||||
gv = R.nn.relu(lv1)
|
||||
R.output(gv)
|
||||
return gv
|
||||
|
||||
Key Features of Relax
|
||||
~~~~~~~~~~~~~~~~~~~~~
|
||||
Relax, the graph representation utilized in Apache TVM,
|
||||
facilitates end-to-end optimization of ML models through several crucial
|
||||
features:
|
||||
|
||||
- **First-class symbolic shape**: Relax employs symbolic shapes to represent
|
||||
tensor dimensions, enabling global tracking of dynamic shape relationships
|
||||
across tensor operators and function calls.
|
||||
|
||||
- **Multi-level abstractions**: Relax supports cross-level abstractions, from
|
||||
high-level neural network layers to low-level tensor operations, enabling
|
||||
optimizations that span different hierarchies within the model.
|
||||
|
||||
- **Composable transformations**: Relax offers a framework for composable
|
||||
transformations that can be selectively applied to different model components.
|
||||
This includes capabilities such as partial lowering and partial specialization,
|
||||
providing flexible customization and optimization options.
|
||||
|
||||
These features collectively empower Relax to offer a powerful and adaptable approach
|
||||
to ML model optimization within the Apache TVM ecosystem.
|
||||
@@ -0,0 +1,557 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
.. _relax-dpl:
|
||||
|
||||
Dataflow Pattern Language (DPL)
|
||||
===============================
|
||||
The Dataflow Pattern Language (DPL) is Relax's built-in facility for
|
||||
**pattern matching and rewriting** on computation graphs. It lets you describe a
|
||||
sub-graph structure you are looking for, search for it inside a Relax function,
|
||||
and optionally replace it with a new structure -- all without hand-writing a
|
||||
full IR visitor.
|
||||
|
||||
DPL is used throughout the TVM stack:
|
||||
|
||||
- **Operator fusion** -- ``FuseOpsByPattern`` groups matched operators into a
|
||||
single fused function.
|
||||
- **Backend dispatch** -- CUTLASS, cuBLAS, cuDNN and other backends register
|
||||
patterns so the compiler can route sub-graphs to optimized library kernels.
|
||||
- **Custom graph transforms** -- users write their own patterns and rewriters
|
||||
to perform project-specific optimizations.
|
||||
|
||||
The typical workflow has three steps:
|
||||
|
||||
1. **Build a pattern** that describes the sub-graph shape (e.g. ``matmul`` followed
|
||||
by ``add``).
|
||||
2. **Match** the pattern against Relax IR to locate all occurrences.
|
||||
3. **Rewrite** each match into a replacement expression.
|
||||
|
||||
The public API lives in ``tvm.relax.dpl`` (source: ``python/tvm/relax/dpl/``).
|
||||
|
||||
|
||||
Building Patterns
|
||||
-----------------
|
||||
A *pattern* is a lightweight description of what an expression should look like.
|
||||
Patterns are built by combining small building blocks.
|
||||
|
||||
Basic Patterns
|
||||
~~~~~~~~~~~~~~
|
||||
The most common leaf patterns are:
|
||||
|
||||
- ``wildcard()`` -- matches any expression.
|
||||
- ``is_op("relax.add")`` -- matches a specific Relax operator.
|
||||
- ``is_const()`` -- matches any constant value.
|
||||
- ``is_var(name)`` -- matches a ``Var`` node (optionally with a given name).
|
||||
- ``is_dfv(name)`` -- matches a ``DataflowVar`` node.
|
||||
- ``is_gv(name)`` -- matches a ``GlobalVar``.
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm.relax.dpl import wildcard, is_op, is_const
|
||||
|
||||
# Match any relax.add call, regardless of arguments
|
||||
add_pattern = is_op("relax.add")(wildcard(), wildcard())
|
||||
|
||||
Call Patterns
|
||||
~~~~~~~~~~~~~
|
||||
Calling a pattern as a function produces a ``CallPattern``. The callee is the
|
||||
pattern itself, and the positional arguments are patterns for each operand:
|
||||
|
||||
.. code:: python
|
||||
|
||||
x = wildcard()
|
||||
w = wildcard()
|
||||
|
||||
# Match: relax.matmul(x, w)
|
||||
matmul = is_op("relax.matmul")(x, w)
|
||||
|
||||
For operators with variadic arguments, pass ``varg_default_wildcard=True`` so
|
||||
that extra arguments are matched by implicit wildcards:
|
||||
|
||||
.. code:: python
|
||||
|
||||
# Match relax.concat with any number of inputs
|
||||
concat = is_op("relax.concat")(wildcard(), varg_default_wildcard=True)
|
||||
|
||||
DPL also provides specialized helpers for common call patterns:
|
||||
|
||||
- ``is_call_tir(func_name, args)`` -- matches ``R.call_tir(func_name, (args...,))``.
|
||||
- ``is_call_dps_packed(func_name, args)`` -- matches ``R.call_dps_packed``.
|
||||
- ``is_call_packed(func_name, args)`` -- matches ``R.call_packed``.
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm.relax.dpl import is_call_tir, wildcard
|
||||
|
||||
# Match a call_tir that calls the function "decode"
|
||||
decode = is_call_tir("decode", args=[wildcard(), wildcard()])
|
||||
|
||||
Tuple Patterns
|
||||
~~~~~~~~~~~~~~
|
||||
``TuplePattern`` matches a Relax tuple with a fixed number of fields.
|
||||
It supports indexing with ``[]`` to create ``TupleGetItemPattern``:
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm.relax.dpl import is_tuple, wildcard
|
||||
|
||||
a, b = wildcard(), wildcard()
|
||||
tup = is_tuple([a, b])
|
||||
|
||||
# Match: getting the first element from the tuple
|
||||
first = tup[0]
|
||||
|
||||
Constraints
|
||||
~~~~~~~~~~~
|
||||
Any pattern can be further narrowed by attaching constraints:
|
||||
|
||||
- ``.has_dtype(dtype)`` -- the matched expression must have the given data type.
|
||||
- ``.has_shape(shape)`` -- the matched expression must have the given shape.
|
||||
- ``.has_attr(attrs)`` -- the matched call must carry the given attributes.
|
||||
- ``.has_ty(ty)`` -- the matched expression must have the given type.
|
||||
|
||||
.. code:: python
|
||||
|
||||
# Match a float16 matmul
|
||||
fp16_matmul = is_op("relax.matmul")(wildcard(), wildcard()).has_dtype("float16")
|
||||
|
||||
Logical Combinators
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
Patterns can be combined with logical operators:
|
||||
|
||||
- ``pat_a | pat_b`` -- match if **either** pattern matches (``OrPattern``).
|
||||
- ``pat_a & pat_b`` -- match if **both** patterns match (``AndPattern``).
|
||||
- ``~pat`` -- match anything **except** ``pat`` (``NotPattern``).
|
||||
|
||||
.. code:: python
|
||||
|
||||
# Match either relu or gelu activation
|
||||
activation = is_op("relax.nn.relu")(wildcard()) | is_op("relax.nn.gelu")(wildcard())
|
||||
|
||||
Sequence Patterns
|
||||
~~~~~~~~~~~~~~~~~
|
||||
When a pattern spans multiple bindings inside a ``DataflowBlock``, use
|
||||
*sequence operators* to express producer-consumer relationships:
|
||||
|
||||
- ``a ^ b`` (``used_by``) -- ``a`` is used by ``b`` (``a`` may also be used
|
||||
elsewhere).
|
||||
- ``a >> b`` (``only_used_by``) -- ``a`` is **only** used by ``b`` (no other
|
||||
consumers).
|
||||
|
||||
These return a ``PatternSeq`` that can be chained:
|
||||
|
||||
.. code:: python
|
||||
|
||||
x = wildcard()
|
||||
matmul = is_op("relax.matmul")(x, wildcard())
|
||||
add = is_op("relax.add")(matmul, wildcard())
|
||||
|
||||
# matmul result is exclusively consumed by the add
|
||||
seq = matmul >> add
|
||||
|
||||
High-level Helpers
|
||||
~~~~~~~~~~~~~~~~~~
|
||||
``make_fused_bias_activation_pattern`` builds a common
|
||||
``op -> optional bias -> optional activation`` chain in one call:
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm.relax.dpl import make_fused_bias_activation_pattern
|
||||
|
||||
# conv2d + bias + relu
|
||||
pattern = make_fused_bias_activation_pattern(
|
||||
"relax.nn.conv2d",
|
||||
with_bias=True,
|
||||
activation="relax.nn.relu",
|
||||
)
|
||||
|
||||
|
||||
Matching Without Rewriting
|
||||
--------------------------
|
||||
Sometimes you only need to **detect** a structure without replacing it.
|
||||
Every ``DFPattern`` exposes two matching methods:
|
||||
|
||||
- ``pattern.match(expr)`` -- returns ``True`` if the pattern matches.
|
||||
- ``pattern.extract_matched_expr(expr)`` -- returns a
|
||||
``dict[DFPattern, Expr]`` mapping each sub-pattern to the concrete
|
||||
expression it matched, or ``None`` on failure.
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm.relax.dpl import wildcard, is_op
|
||||
|
||||
x = wildcard()
|
||||
y = wildcard()
|
||||
add_pat = is_op("relax.add")(x, y)
|
||||
|
||||
# Assume `expr` is a Relax expression: R.add(a, b)
|
||||
if add_pat.match(expr):
|
||||
matched = add_pat.extract_matched_expr(expr)
|
||||
# matched[x] -> the expression that matched `x`
|
||||
# matched[y] -> the expression that matched `y`
|
||||
|
||||
When matching across variable bindings (e.g., ``lv0 = ...; lv1 = f(lv0)``),
|
||||
the matcher needs a ``var2val`` map so it can see through binding
|
||||
boundaries. Use ``tvm.relax.analysis.get_var2val(func)`` to build one:
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm.relax.analysis import get_var2val
|
||||
|
||||
var2val = get_var2val(func)
|
||||
matched = pattern.extract_matched_expr(expr, var2val=var2val)
|
||||
|
||||
|
||||
Rewriting Matched Patterns
|
||||
--------------------------
|
||||
|
||||
``rewrite_call``
|
||||
~~~~~~~~~~~~~~~~
|
||||
``rewrite_call`` is the simplest rewrite API. It walks every expression in a
|
||||
function, and when the pattern matches, it calls your callback to produce a
|
||||
replacement.
|
||||
|
||||
.. code:: python
|
||||
|
||||
rewrite_call(pattern, rewriter, func) -> Function
|
||||
|
||||
The callback signature is:
|
||||
|
||||
.. code:: python
|
||||
|
||||
def rewriter(
|
||||
matched_expr: Expr,
|
||||
matchings: dict[DFPattern, Expr],
|
||||
) -> Expr:
|
||||
...
|
||||
|
||||
**Example -- replace** ``reshape(reshape(x, s1), s2)`` **with**
|
||||
``reshape(x, s2)``:
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm import relax
|
||||
from tvm.relax.dpl import wildcard, is_op, rewrite_call
|
||||
|
||||
inp = wildcard()
|
||||
shape1, shape2 = wildcard(), wildcard()
|
||||
inner = is_op("relax.reshape")(inp, shape1)
|
||||
outer = is_op("relax.reshape")(inner, shape2)
|
||||
|
||||
def rewriter(expr, matchings):
|
||||
# Keep the original input but use the outermost target shape
|
||||
return relax.op.reshape(matchings[inp], matchings[outer].args[1])
|
||||
|
||||
new_func = rewrite_call(outer, rewriter, func)
|
||||
|
||||
``rewrite_call`` is best for **local, single-expression** rewrites.
|
||||
|
||||
|
||||
``rewrite_bindings`` with ``PatternContext``
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
When a rewrite involves **multiple bindings** across a ``DataflowBlock``
|
||||
(e.g., merging three separate matmuls into one), use ``rewrite_bindings``
|
||||
together with ``PatternContext``.
|
||||
|
||||
``PatternContext`` enables topological (graph-level) matching on an entire
|
||||
dataflow block rather than on individual expressions.
|
||||
|
||||
.. code:: python
|
||||
|
||||
rewrite_bindings(ctx, rewriter, func) -> Function
|
||||
|
||||
The callback receives *variables* rather than expressions:
|
||||
|
||||
.. code:: python
|
||||
|
||||
def rewriter(
|
||||
matchings: dict[DFPattern, Var],
|
||||
bindings: dict[Var, Expr],
|
||||
) -> dict[Var, Expr]:
|
||||
...
|
||||
|
||||
- ``matchings[pat]`` returns the **bound variable** (``Var``) whose right-hand
|
||||
side matched ``pat``. The ``Var`` itself carries ``ty`` and can be
|
||||
used directly in new expressions.
|
||||
- ``bindings`` maps each ``Var`` to its bound ``Expr`` (the right-hand side),
|
||||
useful when you need to inspect the original expression.
|
||||
|
||||
**Example -- merge three parallel matmuls into one**:
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm.script import relax as R
|
||||
from tvm.relax.dpl import wildcard, is_op, rewrite_bindings, PatternContext
|
||||
|
||||
with PatternContext() as ctx:
|
||||
inp_pat = wildcard()
|
||||
w1, w2, w3 = wildcard(), wildcard(), wildcard()
|
||||
|
||||
matmul1 = is_op("relax.matmul")(inp_pat, w1)
|
||||
matmul2 = is_op("relax.matmul")(inp_pat, w2)
|
||||
matmul3 = is_op("relax.matmul")(inp_pat, w3)
|
||||
|
||||
def rewriter(matchings, _bindings):
|
||||
inp = matchings[inp_pat]
|
||||
W1 = matchings[w1]
|
||||
W2 = matchings[w2]
|
||||
W3 = matchings[w3]
|
||||
width = W1.ty.shape[1]
|
||||
|
||||
concat_w = R.concat([W1, W2, W3], axis=1)
|
||||
merged = R.matmul(inp, concat_w)
|
||||
|
||||
return {
|
||||
matchings[matmul1]: R.strided_slice(
|
||||
merged, axes=[2], begin=[0], end=[width],
|
||||
),
|
||||
matchings[matmul2]: R.strided_slice(
|
||||
merged, axes=[2], begin=[width], end=[width * 2],
|
||||
),
|
||||
matchings[matmul3]: R.strided_slice(
|
||||
merged, axes=[2], begin=[width * 2], end=[width * 3],
|
||||
),
|
||||
}
|
||||
|
||||
new_func = rewrite_bindings(ctx, rewriter, func)
|
||||
|
||||
|
||||
Declarative Rewriting with ``@R.rewriter``
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
For straightforward one-to-one replacements you can declare the pattern and
|
||||
its replacement as two Relax functions in a single ``IRModule``. The
|
||||
``@R.rewriter`` decorator turns the module into a ``PatternMatchingRewriter``
|
||||
object that can be applied directly.
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm.script import relax as R
|
||||
|
||||
@R.rewriter
|
||||
class RewriteAddToPackedCall:
|
||||
@R.function
|
||||
def pattern(
|
||||
A: R.Tensor([16], "float32"),
|
||||
B: R.Tensor([16], "float32"),
|
||||
):
|
||||
C = R.add(A, B)
|
||||
return C
|
||||
|
||||
@R.function
|
||||
def replacement(
|
||||
A: R.Tensor([16], "float32"),
|
||||
B: R.Tensor([16], "float32"),
|
||||
):
|
||||
C = R.call_pure_packed(
|
||||
"my_fast_add",
|
||||
A,
|
||||
B,
|
||||
ty_args=R.Tensor([16], "float32"),
|
||||
)
|
||||
return C
|
||||
|
||||
# Apply to an IRModule or a single function
|
||||
rewritten_mod = RewriteAddToPackedCall(mod)
|
||||
|
||||
Composing Rewriters
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
Multiple ``PatternMatchingRewriter`` objects can be combined with the ``|``
|
||||
operator so they run as a single pass:
|
||||
|
||||
.. code:: python
|
||||
|
||||
combined = rewriter_a | rewriter_b
|
||||
result = combined(mod)
|
||||
|
||||
The left-hand rewriter is tried first; the right-hand rewriter only applies to
|
||||
bindings that were **not** already modified by the left.
|
||||
|
||||
|
||||
Using DPL in Compiler Passes
|
||||
-----------------------------
|
||||
The most common way DPL appears in the TVM codebase is through the
|
||||
``FuseOpsByPattern`` pass, which uses ``FusionPattern`` objects to drive
|
||||
operator fusion.
|
||||
|
||||
``FusionPattern``
|
||||
~~~~~~~~~~~~~~~~~
|
||||
A ``FusionPattern`` bundles four pieces of information:
|
||||
|
||||
- ``name`` -- a string label (e.g., ``"cutlass.matmul"``).
|
||||
- ``pattern`` -- a ``DFPattern`` that describes the sub-graph to match.
|
||||
- ``annotation_patterns`` -- a ``dict[str, DFPattern]`` that names interesting
|
||||
sub-patterns so the check function can inspect them.
|
||||
- ``check`` -- an optional ``Callable[[PatternCheckContext], bool]`` that
|
||||
performs additional validation after a structural match succeeds.
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm.relax.dpl import wildcard, is_op
|
||||
from tvm.relax.transform import FusionPattern
|
||||
|
||||
x = wildcard()
|
||||
w = wildcard()
|
||||
matmul = is_op("relax.matmul")(x, w)
|
||||
bias = wildcard()
|
||||
add = is_op("relax.add")(matmul, bias)
|
||||
|
||||
pattern = FusionPattern(
|
||||
name="my_backend.matmul_bias",
|
||||
pattern=add,
|
||||
annotation_patterns={"matmul": matmul, "bias": bias, "lhs": x, "rhs": w},
|
||||
check=my_check_fn,
|
||||
)
|
||||
|
||||
``PatternCheckContext``
|
||||
~~~~~~~~~~~~~~~~~~~~~~~
|
||||
When ``FuseOpsByPattern`` finds a structural match, it calls the ``check``
|
||||
function with a ``PatternCheckContext`` that provides:
|
||||
|
||||
- ``matched_expr`` -- the root expression of the match.
|
||||
- ``annotated_expr`` -- a ``dict[str, Expr]`` resolved from the
|
||||
``annotation_patterns``.
|
||||
- ``matched_bindings`` -- a ``dict[Var, Expr]`` of bindings being fused.
|
||||
- ``var_usages`` -- a ``dict[Var, Sequence[Var]]`` of variable use chains.
|
||||
- ``value_to_bound_var`` -- a ``dict[Expr, Var]`` mapping values back to
|
||||
their bound variables.
|
||||
|
||||
Use the check function to enforce constraints that cannot be expressed
|
||||
structurally (dtype restrictions, shape compatibility, attribute values, etc.):
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm.relax.transform import PatternCheckContext
|
||||
|
||||
def my_check_fn(ctx: PatternCheckContext) -> bool:
|
||||
matmul_expr = ctx.annotated_expr["matmul"]
|
||||
# Only accept float16 output
|
||||
if matmul_expr.ty.dtype != "float16":
|
||||
return False
|
||||
return True
|
||||
|
||||
``FuseOpsByPattern``
|
||||
~~~~~~~~~~~~~~~~~~~~
|
||||
``FuseOpsByPattern`` is a module-level pass that takes a list of
|
||||
``FusionPattern`` (or equivalent tuples) and groups every match into a fused
|
||||
sub-function.
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm.relax.dpl import wildcard, is_op
|
||||
from tvm.relax.transform import FuseOpsByPattern
|
||||
|
||||
# 1. Define the pattern
|
||||
w = wildcard()
|
||||
x = wildcard()
|
||||
wT = is_op("relax.permute_dims")(w)
|
||||
o = is_op("relax.matmul")(x, wT)
|
||||
annotations = {"o": o, "w": w, "x": x, "wT": wT}
|
||||
|
||||
def check(ctx):
|
||||
transpose_call = ctx.annotated_expr["wT"]
|
||||
ndim = transpose_call.args[0].ty.ndim
|
||||
if ndim == -1:
|
||||
return False
|
||||
if ndim == 2 and transpose_call.attrs.axes is None:
|
||||
return True
|
||||
axes = list(range(ndim))
|
||||
axes[-1], axes[-2] = axes[-2], axes[-1]
|
||||
return list(transpose_call.attrs.axes) == axes
|
||||
|
||||
# 2. Run the pass
|
||||
mod = FuseOpsByPattern(
|
||||
[("transpose_matmul_fuse", o, annotations, check)],
|
||||
bind_constants=False,
|
||||
)(mod)
|
||||
|
||||
When ``annotate_codegen=True``, each fused function is additionally wrapped
|
||||
with ``Codegen`` and ``global_symbol`` attributes, which is how backends like
|
||||
CUTLASS and cuBLAS register themselves for external code generation.
|
||||
|
||||
|
||||
Quick Reference
|
||||
---------------
|
||||
|
||||
**Pattern construction**
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 35 65
|
||||
|
||||
* - API
|
||||
- Description
|
||||
* - ``wildcard()``
|
||||
- Match any expression
|
||||
* - ``is_op(op_name)``
|
||||
- Match a Relax operator by name
|
||||
* - ``is_const()``
|
||||
- Match any constant
|
||||
* - ``is_var(name)`` / ``is_dfv(name)`` / ``is_gv(name)``
|
||||
- Match ``Var`` / ``DataflowVar`` / ``GlobalVar``
|
||||
* - ``is_tuple(fields)``
|
||||
- Match a tuple with given field patterns
|
||||
* - ``is_call_tir(name, args)``
|
||||
- Match ``R.call_tir``
|
||||
* - ``is_call_dps_packed(name, args)``
|
||||
- Match ``R.call_dps_packed``
|
||||
* - ``is_call_packed(name, args)``
|
||||
- Match ``R.call_packed``
|
||||
* - ``make_fused_bias_activation_pattern(...)``
|
||||
- Build ``op + bias + activation`` chain
|
||||
* - ``.has_dtype()`` / ``.has_shape()`` / ``.has_attr()`` / ``.has_ty()``
|
||||
- Attach constraints
|
||||
* - ``|`` / ``&`` / ``~``
|
||||
- Or / And / Not combinators
|
||||
* - ``^`` / ``>>``
|
||||
- used_by / only_used_by (sequence)
|
||||
|
||||
**Matching and rewriting**
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 35 65
|
||||
|
||||
* - API
|
||||
- Description
|
||||
* - ``pattern.match(expr)``
|
||||
- Returns ``True`` if pattern matches
|
||||
* - ``pattern.extract_matched_expr(expr)``
|
||||
- Returns ``dict[DFPattern, Expr]`` or ``None``
|
||||
* - ``rewrite_call(pattern, rewriter, func)``
|
||||
- Rewrite individual expressions
|
||||
* - ``rewrite_bindings(ctx, rewriter, func)``
|
||||
- Rewrite across bindings in a ``DataflowBlock``
|
||||
* - ``PatternMatchingRewriter.from_module(mod)``
|
||||
- Declarative rewriter from ``IRModule``
|
||||
* - ``@R.rewriter``
|
||||
- Decorator shorthand for ``from_module``
|
||||
|
||||
**Pass integration**
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 35 65
|
||||
|
||||
* - API
|
||||
- Description
|
||||
* - ``FusionPattern(name, pattern, annotations, check)``
|
||||
- Bundle pattern with metadata for ``FuseOpsByPattern``
|
||||
* - ``PatternCheckContext``
|
||||
- Runtime context passed to check functions
|
||||
* - ``FuseOpsByPattern(patterns, ...)``
|
||||
- Module pass that fuses matched sub-graphs
|
||||
@@ -0,0 +1,35 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
.. _relax-deep-dive:
|
||||
|
||||
Relax
|
||||
=====
|
||||
Relax is a high-level abstraction for graph optimization and transformation in Apache TVM stack.
|
||||
Additionally, Apache TVM combines Relax and TensorIR together for cross-level
|
||||
optimization. Hence, Relax is usually working closely with TensorIR for representing and optimizing
|
||||
the whole IRModule
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
abstraction
|
||||
learning
|
||||
dpl
|
||||
tutorials/relax_creation
|
||||
tutorials/relax_transformation
|
||||
@@ -0,0 +1,276 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
.. _relax-learning:
|
||||
|
||||
Understand Relax Abstraction
|
||||
============================
|
||||
Relax is a graph abstraction used in Apache TVM, which
|
||||
helps to end-to-end optimize ML models. The principal objective of Relax
|
||||
is to depict the structure and data flow of ML models, including the
|
||||
dependencies and relationships between different parts of the model, as
|
||||
well as how to execute the model on hardware.
|
||||
|
||||
End to End Model Execution
|
||||
--------------------------
|
||||
|
||||
In this chapter, we will use the following model as an example. This is
|
||||
a two-layer neural network that consists of two linear operations with
|
||||
relu activation.
|
||||
|
||||
.. image:: /_static/img/e2e_fashionmnist_mlp_model.png
|
||||
:width: 85%
|
||||
:align: center
|
||||
|
||||
|
||||
High-Level Operations Representation
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Let us begin by reviewing a Numpy implementation of the model.
|
||||
|
||||
.. code:: python
|
||||
|
||||
def numpy_mlp(data, w0, b0, w1, b1):
|
||||
lv0 = data @ w0 + b0
|
||||
lv1 = np.maximum(lv0, 0)
|
||||
lv2 = lv1 @ w1 + b1
|
||||
return lv2
|
||||
|
||||
The above example code shows the high-level array operations to perform the end-to-end model
|
||||
execution. Of course, we can rewrite the above code using Relax as follows:
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm.script import relax as R
|
||||
|
||||
@R.function
|
||||
def relax_mlp(
|
||||
data: R.Tensor(("n", 784), dtype="float32"),
|
||||
w0: R.Tensor((784, 128), dtype="float32"),
|
||||
b0: R.Tensor((128,), dtype="float32"),
|
||||
w1: R.Tensor((128, 10), dtype="float32"),
|
||||
b1: R.Tensor((10,), dtype="float32"),
|
||||
) -> R.Tensor(("n", 10), dtype="float32"):
|
||||
with R.dataflow():
|
||||
lv0 = R.matmul(data, w0) + b0
|
||||
lv1 = R.nn.relu(lv0)
|
||||
lv2 = R.matmul(lv1, w1) + b1
|
||||
R.output(lv2)
|
||||
return lv2
|
||||
|
||||
Low-Level Integration
|
||||
~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
However, again from the pov of machine learning compilation (MLC), we would like to see
|
||||
through the details under the hood of these array computations.
|
||||
|
||||
For the purpose of illustrating details under the hood, we will again write examples in low-level numpy:
|
||||
|
||||
We will use a loop instead of array functions when necessary to demonstrate the possible loop computations.
|
||||
When possible, we always explicitly allocate arrays via numpy.empty and pass them around.
|
||||
The code block below shows a low-level numpy implementation of the same model.
|
||||
|
||||
.. code:: python
|
||||
|
||||
def lnumpy_linear(X: np.ndarray, W: np.ndarray, B: np.ndarray, Z: np.ndarray):
|
||||
n, m, K = X.shape[0], W.shape[1], X.shape[1]
|
||||
Y = np.empty((n, m), dtype="float32")
|
||||
for i in range(n):
|
||||
for j in range(m):
|
||||
for k in range(K):
|
||||
if k == 0:
|
||||
Y[i, j] = 0
|
||||
Y[i, j] = Y[i, j] + X[i, k] * W[k, j]
|
||||
|
||||
for i in range(n):
|
||||
for j in range(m):
|
||||
Z[i, j] = Y[i, j] + B[j]
|
||||
|
||||
|
||||
def lnumpy_relu0(X: np.ndarray, Y: np.ndarray):
|
||||
n, m = X.shape
|
||||
for i in range(n):
|
||||
for j in range(m):
|
||||
Y[i, j] = np.maximum(X[i, j], 0)
|
||||
|
||||
def lnumpy_mlp(data, w0, b0, w1, b1):
|
||||
n = data.shape[0]
|
||||
lv0 = np.empty((n, 128), dtype="float32")
|
||||
lnumpy_linear(data, w0, b0, lv0)
|
||||
|
||||
lv1 = np.empty((n, 128), dtype="float32")
|
||||
lnumpy_relu0(lv0, lv1)
|
||||
|
||||
out = np.empty((n, 10), dtype="float32")
|
||||
lnumpy_linear(lv1, w1, b1, out)
|
||||
return out
|
||||
|
||||
With the low-level NumPy example in mind, now we are ready to introduce an Relax abstraction
|
||||
for the end-to-end model execution. The code block below shows a TVMScript implementation of the model.
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script import relax as R
|
||||
|
||||
@I.ir_module
|
||||
class Module:
|
||||
@T.prim_func(private=True)
|
||||
def linear(x: T.handle, w: T.handle, b: T.handle, z: T.handle):
|
||||
M, N, K = T.int64(), T.int64(), T.int64()
|
||||
X = T.match_buffer(x, (M, K), "float32")
|
||||
W = T.match_buffer(w, (K, N), "float32")
|
||||
B = T.match_buffer(b, (N,), "float32")
|
||||
Z = T.match_buffer(z, (M, N), "float32")
|
||||
Y = T.alloc_buffer((M, N), "float32")
|
||||
for i, j, k in T.grid(M, N, K):
|
||||
with T.sblock("Y"):
|
||||
v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
Y[v_i, v_j] = T.float32(0.0)
|
||||
Y[v_i, v_j] = Y[v_i, v_j] + X[v_i, v_k] * W[v_k, v_j]
|
||||
for i, j in T.grid(M, N):
|
||||
with T.sblock("Z"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
Z[v_i, v_j] = Y[v_i, v_j] + B[v_j]
|
||||
|
||||
@T.prim_func(private=True)
|
||||
def relu(x: T.handle, y: T.handle):
|
||||
M, N = T.int64(), T.int64()
|
||||
X = T.match_buffer(x, (M, N), "float32")
|
||||
Y = T.match_buffer(y, (M, N), "float32")
|
||||
for i, j in T.grid(M, N):
|
||||
with T.sblock("Y"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
Y[v_i, v_j] = T.max(X[v_i, v_j], T.float32(0.0))
|
||||
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor(("n", 784), dtype="float32"),
|
||||
w0: R.Tensor((784, 256), dtype="float32"),
|
||||
b0: R.Tensor((256,), dtype="float32"),
|
||||
w1: R.Tensor((256, 10), dtype="float32"),
|
||||
b1: R.Tensor((10,), dtype="float32")
|
||||
) -> R.Tensor(("n", 10), dtype="float32"):
|
||||
cls = Module
|
||||
n = T.int64()
|
||||
with R.dataflow():
|
||||
lv = R.call_tir(cls.linear, (x, w0, b0), out_ty=R.Tensor((n, 256), dtype="float32"))
|
||||
lv1 = R.call_tir(cls.relu, (lv,), out_ty=R.Tensor((n, 256), dtype="float32"))
|
||||
lv2 = R.call_tir(cls.linear, (lv1, w1, b1), out_ty=R.Tensor((n, 10), dtype="float32"))
|
||||
R.output(lv2)
|
||||
return lv2
|
||||
|
||||
The above code contains kinds of functions: the primitive tensor functions (``T.prim_func``) and a
|
||||
``R.function`` (relax function). Relax function is a new type of abstraction representing
|
||||
high-level neural network executions.
|
||||
|
||||
Note that the above relax module natively supports symbolic shapes, see the ``"n"`` in the
|
||||
tensor shapes in ``main`` function and ``M``, ``N``, ``K`` in the ``linear`` function. This is
|
||||
a key feature of Relax abstraction, which enables the compiler to track dynamic shape relations
|
||||
globally across tensor operators and function calls.
|
||||
|
||||
Again it is helpful to see the TVMScript code and low-level numpy code side-by-side and check the
|
||||
corresponding elements, and we are going to walk through each of them in detail. Since we already
|
||||
learned about primitive tensor functions, we are going to focus on the high-level execution part.
|
||||
|
||||
Key Elements of Relax
|
||||
---------------------
|
||||
This section will introduce the key elements of Relax abstraction and how it enables optimization
|
||||
in ML compilers.
|
||||
|
||||
Type
|
||||
~~~~
|
||||
Type is the Relax representation of expression type information. It can
|
||||
be ``TensorType``, ``TupleType``, etc. In the above example, we use ``TensorType``
|
||||
(short in ``R.Tensor`` in TVMScript) to represent the shape and dtype of the tensor of the inputs,
|
||||
outputs, and intermediate results.
|
||||
|
||||
R.call_tir
|
||||
~~~~~~~~~~
|
||||
The ``R.call_tir`` function is a new abstraction in Relax that allows calling primitive tensor
|
||||
functions in the same IRModule. This is a key feature of Relax that enables cross-level
|
||||
abstractions, from high-level neural network layers to low-level tensor operations.
|
||||
Taking one line from the above code as an example:
|
||||
|
||||
.. code:: python
|
||||
|
||||
lv = R.call_tir(cls.linear, (x, w0, b0), out_ty=R.Tensor((n, 256), dtype="float32"))
|
||||
|
||||
To explain what does ``R.call_tir`` work, let us review an equivalent low-level numpy
|
||||
implementation of the operation, as follows:
|
||||
|
||||
.. code:: python
|
||||
|
||||
lv0 = np.empty((n, 256), dtype="float32")
|
||||
lnumpy_linear(x, w0, b0, lv0)
|
||||
|
||||
Specifically, ``call_tir`` allocates an output tensor res, then pass the inputs and the output
|
||||
to the prim_func. After executing prim_func the result is populated in res, then we can return
|
||||
the result.
|
||||
|
||||
This convention is called **destination passing**, The idea is that input and output are explicitly
|
||||
allocated outside and passed to the low-level primitive function. This style is commonly used
|
||||
in low-level library designs, so higher-level frameworks can handle that memory allocation
|
||||
decision. Note that not all tensor operations can be presented in this style (specifically,
|
||||
there are operations whose output shape depends on the input). Nevertheless, in common practice,
|
||||
it is usually helpful to write the low-level function in this style when possible.
|
||||
|
||||
Dataflow Block
|
||||
~~~~~~~~~~~~~~
|
||||
Another important element in a relax function is the R.dataflow() scope annotation.
|
||||
|
||||
.. code:: python
|
||||
|
||||
with R.dataflow():
|
||||
lv = R.call_tir(cls.linear, (x, w0, b0), out_ty=R.Tensor((n, 256), dtype="float32"))
|
||||
lv1 = R.call_tir(cls.relu, (lv,), out_ty=R.Tensor((n, 256), dtype="float32"))
|
||||
lv2 = R.call_tir(cls.linear, (lv1, w1, b1), out_ty=R.Tensor((n, 10), dtype="float32"))
|
||||
R.output(lv2)
|
||||
|
||||
Before we talk about the dataflow block, let us first introduce the concept of **pure** and
|
||||
**side-effect**. A function is **pure** or **side-effect free** if:
|
||||
|
||||
- it only reads from its inputs and returns the result via its output
|
||||
- it will not change other parts of the program (such as incrementing a global counter).
|
||||
|
||||
For example, all ``R.call_tir`` functions are pure functions, as they only read from their inputs
|
||||
and write the output to another new allocated tensor. However, the **inplace operations** are not
|
||||
pure functions, in other words, they are side-effect functions, because they will change the existing
|
||||
intermediate or input tensors.
|
||||
|
||||
A dataflow block is a way for us to mark the computational graph regions of the program.
|
||||
Specifically, within a dataflow block, all the operations need to be **side-effect free**.
|
||||
Outside a dataflow block, the operations can contain side-effect.
|
||||
|
||||
.. note::
|
||||
|
||||
A common question that arises is why we need to manually mark dataflow blocks instead of
|
||||
automatically inferring them. There are two main reasons for this approach:
|
||||
|
||||
- Automatic inference of dataflow blocks can be challenging and imprecise, particularly
|
||||
when dealing with calls to packed functions (such as cuBLAS integrations). By manually
|
||||
marking dataflow blocks, we enable the compiler to accurately understand and optimize
|
||||
the program's dataflow.
|
||||
- Many optimizations can only be applied within dataflow blocks. For instance, fusion
|
||||
optimization is limited to operations within a single dataflow block. If the compiler
|
||||
were to incorrectly infer dataflow boundaries, it might miss crucial optimization
|
||||
opportunities, potentially impacting the program's performance.
|
||||
|
||||
By allowing manual marking of dataflow blocks, we ensure that the compiler has the most
|
||||
accurate information to work with, leading to more effective optimizations.
|
||||
@@ -0,0 +1,2 @@
|
||||
Deep Dive: Relax
|
||||
----------------
|
||||
@@ -0,0 +1,285 @@
|
||||
# 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
|
||||
|
||||
"""
|
||||
.. _relax-creation:
|
||||
|
||||
Relax Creation
|
||||
==============
|
||||
This tutorial demonstrates how to create Relax functions and programs.
|
||||
We'll cover various ways to define Relax functions, including using TVMScript,
|
||||
and relax NNModule API.
|
||||
"""
|
||||
|
||||
|
||||
######################################################################
|
||||
# Create Relax programs using TVMScript
|
||||
# -------------------------------------
|
||||
# TVMScript is a domain-specific language for representing Apache TVM's
|
||||
# intermediate representation (IR). It is a Python dialect that can be used
|
||||
# to define an IRModule, which contains both TensorIR and Relax functions.
|
||||
#
|
||||
# In this section, we will show how to define a simple MLP model with only
|
||||
# high-level Relax operators using TVMScript.
|
||||
|
||||
from tvm import relax, topi
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import relax as R
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
@I.ir_module
|
||||
class RelaxModule:
|
||||
@R.function
|
||||
def forward(
|
||||
data: R.Tensor(("n", 784), dtype="float32"),
|
||||
w0: R.Tensor((128, 784), dtype="float32"),
|
||||
b0: R.Tensor((128,), dtype="float32"),
|
||||
w1: R.Tensor((10, 128), dtype="float32"),
|
||||
b1: R.Tensor((10,), dtype="float32"),
|
||||
) -> R.Tensor(("n", 10), dtype="float32"):
|
||||
with R.dataflow():
|
||||
lv0 = R.matmul(data, R.permute_dims(w0)) + b0
|
||||
lv1 = R.nn.relu(lv0)
|
||||
lv2 = R.matmul(lv1, R.permute_dims(w1)) + b1
|
||||
R.output(lv2)
|
||||
return lv2
|
||||
|
||||
|
||||
RelaxModule.show()
|
||||
|
||||
######################################################################
|
||||
# Relax is not only a graph-level IR, but also supports cross-level
|
||||
# representation and transformation. To be specific, we can directly call
|
||||
# TensorIR functions in Relax function.
|
||||
|
||||
|
||||
@I.ir_module
|
||||
class RelaxModuleWithTIR:
|
||||
@T.prim_func(s_tir=True)
|
||||
def relu(x: T.handle, y: T.handle):
|
||||
n = T.int64()
|
||||
m = T.int64()
|
||||
X = T.match_buffer(x, (n, m), "float32")
|
||||
Y = T.match_buffer(y, (n, m), "float32")
|
||||
for i, j in T.grid(n, m):
|
||||
with T.sblock("relu"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
Y[vi, vj] = T.max(X[vi, vj], T.float32(0))
|
||||
|
||||
@R.function
|
||||
def forward(
|
||||
data: R.Tensor(("n", 784), dtype="float32"),
|
||||
w0: R.Tensor((128, 784), dtype="float32"),
|
||||
b0: R.Tensor((128,), dtype="float32"),
|
||||
w1: R.Tensor((10, 128), dtype="float32"),
|
||||
b1: R.Tensor((10,), dtype="float32"),
|
||||
) -> R.Tensor(("n", 10), dtype="float32"):
|
||||
n = T.int64()
|
||||
cls = RelaxModuleWithTIR
|
||||
with R.dataflow():
|
||||
lv0 = R.matmul(data, R.permute_dims(w0)) + b0
|
||||
lv1 = R.call_tir(cls.relu, lv0, R.Tensor((n, 128), dtype="float32"))
|
||||
lv2 = R.matmul(lv1, R.permute_dims(w1)) + b1
|
||||
R.output(lv2)
|
||||
return lv2
|
||||
|
||||
|
||||
RelaxModuleWithTIR.show()
|
||||
|
||||
######################################################################
|
||||
# .. note::
|
||||
#
|
||||
# You may notice that the printed output is different from the written
|
||||
# TVMScript code. This is because we print the IRModule in a standard
|
||||
# format, while we support syntax sugar for the input
|
||||
#
|
||||
# For example, we can combine multiple operators into a single line, as
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# lv0 = R.matmul(data, R.permute_dims(w0)) + b0
|
||||
#
|
||||
# However, the normalized expression requires only one operation in one
|
||||
# binding. So the printed output is different from the written TVMScript code,
|
||||
# as
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# lv: R.Tensor((784, 128), dtype="float32") = R.permute_dims(w0, axes=None)
|
||||
# lv1: R.Tensor((n, 128), dtype="float32") = R.matmul(data, lv)
|
||||
# lv0: R.Tensor((n, 128), dtype="float32") = R.add(lv1, b0)
|
||||
#
|
||||
|
||||
######################################################################
|
||||
# Create Relax programs using NNModule API
|
||||
# ----------------------------------------
|
||||
# Besides TVMScript, we also provide a PyTorch-like API for defining neural networks.
|
||||
# It is designed to be more intuitive and easier to use than TVMScript.
|
||||
#
|
||||
# In this section, we will show how to define the same MLP model using
|
||||
# Relax NNModule API.
|
||||
|
||||
from tvm.relax.frontend import nn
|
||||
|
||||
|
||||
class NNModule(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.fc1 = nn.Linear(784, 128)
|
||||
self.relu1 = nn.ReLU()
|
||||
self.fc2 = nn.Linear(128, 10)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.fc1(x)
|
||||
x = self.relu1(x)
|
||||
x = self.fc2(x)
|
||||
return x
|
||||
|
||||
|
||||
######################################################################
|
||||
# After we define the NNModule, we can export it to TVM IRModule via
|
||||
# ``export_tvm``.
|
||||
|
||||
mod, params = NNModule().export_tvm({"forward": {"x": nn.spec.Tensor(("n", 784), "float32")}})
|
||||
mod.show()
|
||||
|
||||
######################################################################
|
||||
# We can also insert customized function calls into the NNModule, such as
|
||||
# Tensor Expression(TE), TensorIR functions or other TVM packed functions.
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def tir_linear(x: T.handle, w: T.handle, b: T.handle, z: T.handle):
|
||||
M = T.int64()
|
||||
N = T.int64()
|
||||
K = T.int64()
|
||||
X = T.match_buffer(x, (M, K), "float32")
|
||||
W = T.match_buffer(w, (N, K), "float32")
|
||||
B = T.match_buffer(b, (N,), "float32")
|
||||
Z = T.match_buffer(z, (M, N), "float32")
|
||||
for i, j, k in T.grid(M, N, K):
|
||||
with T.sblock("linear"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
Z[vi, vj] = 0
|
||||
Z[vi, vj] = Z[vi, vj] + X[vi, vk] * W[vj, vk]
|
||||
for i, j in T.grid(M, N):
|
||||
with T.sblock("add"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
Z[vi, vj] = Z[vi, vj] + B[vj]
|
||||
|
||||
|
||||
class NNModuleWithTIR(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.fc1 = nn.Linear(784, 128)
|
||||
self.fc2 = nn.Linear(128, 10)
|
||||
|
||||
def forward(self, x):
|
||||
n = x.shape[0]
|
||||
# We can call external functions using nn.extern
|
||||
x = nn.extern(
|
||||
"env.linear",
|
||||
[x, self.fc1.weight, self.fc1.bias],
|
||||
out=nn.Tensor.placeholder((n, 128), "float32"),
|
||||
)
|
||||
# We can also call TensorIR via Tensor Expression API in TOPI
|
||||
x = nn.tensor_expr_op(topi.nn.relu, "relu", [x])
|
||||
# We can also call other TVM packed functions
|
||||
x = nn.tensor_ir_op(
|
||||
tir_linear,
|
||||
"tir_linear",
|
||||
[x, self.fc2.weight, self.fc2.bias],
|
||||
out=nn.Tensor.placeholder((n, 10), "float32"),
|
||||
)
|
||||
return x
|
||||
|
||||
|
||||
mod, params = NNModuleWithTIR().export_tvm(
|
||||
{"forward": {"x": nn.spec.Tensor(("n", 784), "float32")}}
|
||||
)
|
||||
mod.show()
|
||||
|
||||
|
||||
######################################################################
|
||||
# Create Relax programs using Block Builder API
|
||||
# ---------------------------------------------
|
||||
# In addition to the above APIs, we also provide a Block Builder API for
|
||||
# creating Relax programs. It is a IR builder API, which is more
|
||||
# low-level and widely used in TVM's internal logic, e.g writing a
|
||||
# customized pass.
|
||||
|
||||
bb = relax.BlockBuilder()
|
||||
n = T.int64()
|
||||
x = relax.Var("x", R.Tensor((n, 784), "float32"))
|
||||
fc1_weight = relax.Var("fc1_weight", R.Tensor((128, 784), "float32"))
|
||||
fc1_bias = relax.Var("fc1_bias", R.Tensor((128,), "float32"))
|
||||
fc2_weight = relax.Var("fc2_weight", R.Tensor((10, 128), "float32"))
|
||||
fc2_bias = relax.Var("fc2_bias", R.Tensor((10,), "float32"))
|
||||
with bb.function("forward", [x, fc1_weight, fc1_bias, fc2_weight, fc2_bias]):
|
||||
with bb.dataflow():
|
||||
lv0 = bb.emit(relax.op.matmul(x, relax.op.permute_dims(fc1_weight)) + fc1_bias)
|
||||
lv1 = bb.emit(relax.op.nn.relu(lv0))
|
||||
gv = bb.emit(relax.op.matmul(lv1, relax.op.permute_dims(fc2_weight)) + fc2_bias)
|
||||
bb.emit_output(gv)
|
||||
bb.emit_func_output(gv)
|
||||
|
||||
mod = bb.get()
|
||||
mod.show()
|
||||
|
||||
######################################################################
|
||||
# Also, Block Builder API supports building cross-level IRModule with both
|
||||
# Relax functions, TensorIR functions and other TVM packed functions.
|
||||
|
||||
bb = relax.BlockBuilder()
|
||||
with bb.function("forward", [x, fc1_weight, fc1_bias, fc2_weight, fc2_bias]):
|
||||
with bb.dataflow():
|
||||
lv0 = bb.emit(
|
||||
relax.call_dps_packed(
|
||||
"env.linear",
|
||||
[x, fc1_weight, fc1_bias],
|
||||
out_ty=relax.TensorType((n, 128), "float32"),
|
||||
)
|
||||
)
|
||||
lv1 = bb.emit_te(topi.nn.relu, lv0)
|
||||
tir_gv = bb.add_func(tir_linear, "tir_linear")
|
||||
gv = bb.emit(
|
||||
relax.call_tir(
|
||||
tir_gv,
|
||||
[lv1, fc2_weight, fc2_bias],
|
||||
out_ty=relax.TensorType((n, 10), "float32"),
|
||||
)
|
||||
)
|
||||
bb.emit_output(gv)
|
||||
bb.emit_func_output(gv)
|
||||
mod = bb.get()
|
||||
mod.show()
|
||||
|
||||
######################################################################
|
||||
# Note that the Block Builder API is not as user-friendly as the above APIs,
|
||||
# but it is lowest-level API and works closely with the IR definition. We
|
||||
# recommend using the above APIs for users who only want to define and
|
||||
# transform a ML model. But for those who want to build more complex
|
||||
# transformations, the Block Builder API is a more flexible choice.
|
||||
|
||||
######################################################################
|
||||
# Summary
|
||||
# -------
|
||||
# This tutorial demonstrates how to create Relax programs using TVMScript,
|
||||
# NNModule API, Block Builder API and PackedFunc API for different use cases.
|
||||
@@ -0,0 +1,142 @@
|
||||
# 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
|
||||
|
||||
"""
|
||||
.. _relax-transform:
|
||||
|
||||
Transformation
|
||||
--------------
|
||||
In this section, we will dive into the transformation of Relax programs.
|
||||
Transformations is one of the key ingredients of the compilation flows
|
||||
for optimizing and integrating with hardware backends.
|
||||
"""
|
||||
|
||||
######################################################################
|
||||
# Let's first create a simple Relax program as what we have done in
|
||||
# the :ref:`previous section <relax-creation>`.
|
||||
|
||||
import tvm
|
||||
from tvm import IRModule, relax
|
||||
from tvm.relax.frontend import nn
|
||||
|
||||
|
||||
class NNModule(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.fc1 = nn.Linear(784, 128)
|
||||
self.relu1 = nn.ReLU()
|
||||
self.fc2 = nn.Linear(128, 10)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.fc1(x)
|
||||
x = self.relu1(x)
|
||||
x = self.fc2(x)
|
||||
return x
|
||||
|
||||
|
||||
origin_mod, params = NNModule().export_tvm(
|
||||
{"forward": {"x": nn.spec.Tensor(("n", 784), "float32")}}
|
||||
)
|
||||
origin_mod.show()
|
||||
|
||||
######################################################################
|
||||
# Apply transformations
|
||||
# ~~~~~~~~~~~~~~~~~~~~~
|
||||
# Passes are the main way to apply transformations to the program.
|
||||
# We can apply passes to the program. As first step, let's apply
|
||||
# a built-in pass ``LegalizeOps`` to lower the high-level operators
|
||||
# into low-level operators.
|
||||
|
||||
mod = tvm.relax.transform.LegalizeOps()(origin_mod)
|
||||
mod.show()
|
||||
|
||||
######################################################################
|
||||
# As we can see from the output, the high-level operators (aka ``relax.op``) in the program
|
||||
# are replaced by their corresponding low-level operators (aka ``relax.call_tir``).
|
||||
#
|
||||
# Then let's trying to apply the operator fusion, which is a wide-used optimization technique
|
||||
# in ML compilers. Note that in relax, fusion optimizations are done with the collaboration of
|
||||
# a set of passes. We can apply them in a sequence.
|
||||
|
||||
mod = tvm.ir.transform.Sequential(
|
||||
[
|
||||
tvm.relax.transform.AnnotateTIROpPattern(),
|
||||
tvm.relax.transform.FuseOps(),
|
||||
tvm.relax.transform.FuseTIR(),
|
||||
]
|
||||
)(mod)
|
||||
mod.show()
|
||||
|
||||
######################################################################
|
||||
# As result, we can see that the ``matmul``, ``add`` and ``relu`` operators are fused
|
||||
# into one kernel (aka one ``call_tir``).
|
||||
#
|
||||
# For all built-in passes, please refer to :py:class:`relax.transform`.
|
||||
#
|
||||
# Custom Passes
|
||||
# ~~~~~~~~~~~~~
|
||||
# We can also define our own passes. Let's take an example of rewriting the ``relu``
|
||||
# operator to ``gelu`` operator.
|
||||
#
|
||||
# First, we need to write a Relax IR Mutator to do the rewriting.
|
||||
|
||||
from tvm.relax.expr_functor import PyExprMutator, mutator
|
||||
|
||||
|
||||
@mutator
|
||||
class ReluRewriter(PyExprMutator):
|
||||
def __init__(self, mod):
|
||||
super().__init__(mod)
|
||||
|
||||
def visit_call_(self, call: relax.Call) -> relax.Expr:
|
||||
# visit the relax.Call expr, and only handle the case when op is relax.nn.relu
|
||||
if call.op.name == "relax.nn.relu":
|
||||
return relax.op.nn.gelu(call.args[0])
|
||||
|
||||
return super().visit_call_(call)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Then we can write a pass to apply the mutator to the whole module.
|
||||
|
||||
|
||||
@tvm.transform.module_pass(opt_level=0, name="ReluToGelu")
|
||||
class ReluToGelu: # pylint: disable=too-few-public-methods
|
||||
def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
|
||||
"""IRModule-level transformation"""
|
||||
rewriter = ReluRewriter(mod)
|
||||
for g_var, func in mod.functions_items():
|
||||
if isinstance(func, relax.Function):
|
||||
func = rewriter.visit_expr(func)
|
||||
rewriter.builder_.update_func(g_var, func)
|
||||
return rewriter.builder_.get()
|
||||
|
||||
|
||||
mod = ReluToGelu()(origin_mod)
|
||||
mod.show()
|
||||
|
||||
######################################################################
|
||||
# The printed output shows that the ``relax.nn.relu`` operator is
|
||||
# rewritten to ``relax.nn.gelu`` operator.
|
||||
#
|
||||
# For the details of the mutator, please refer to :py:class:`relax.expr_functor.PyExprMutator`.
|
||||
#
|
||||
# Summary
|
||||
# ~~~~~~~
|
||||
# In this section, we have shown how to apply transformations to the Relax program.
|
||||
# We have also shown how to define and apply custom transformations.
|
||||
@@ -0,0 +1,72 @@
|
||||
.. 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.
|
||||
|
||||
.. _tir-abstraction:
|
||||
|
||||
Tensor Program Abstraction
|
||||
--------------------------
|
||||
Before we dive into the details of TensorIR, let's first introduce what is a primitive tensor
|
||||
function. Primitive tensor functions are functions that correspond to a single "unit" of
|
||||
computational operation. For example, a convolution operation can be a primitive tensor function,
|
||||
and a fused convolution + relu operation can also be a primitive tensor function.
|
||||
Usually, a typical abstraction for primitive tensor function implementation contains the following
|
||||
elements: multi-dimensional buffers, loop nests that drive the tensor computations, and finally,
|
||||
the compute statements themselves.
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm.script import tirx as T
|
||||
|
||||
@T.prim_func
|
||||
def main(
|
||||
A: T.Buffer((128,), "float32"),
|
||||
B: T.Buffer((128,), "float32"),
|
||||
C: T.Buffer((128,), "float32"),
|
||||
) -> None:
|
||||
for i in range(128):
|
||||
with T.sblock("C"):
|
||||
vi = T.axis.spatial(128, i)
|
||||
C[vi] = A[vi] + B[vi]
|
||||
|
||||
Key Elements of Tensor Programs
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
The demonstrated primitive tensor function calculates the element-wise sum of two vectors.
|
||||
The function:
|
||||
|
||||
- Accepts three **multi-dimensional buffers** as parameters, and generates one **multi-dimensional
|
||||
buffer** as output.
|
||||
- Incorporates a solitary **loop nest** ``i`` that facilitates the computation.
|
||||
- Features a singular **compute statement** that calculates the element-wise sum of the two
|
||||
vectors.
|
||||
|
||||
Extra Structure in TensorIR
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
Crucially, we are unable to execute arbitrary transformations on the program, as certain
|
||||
computations rely on the loop's sequence. Fortunately, the majority of primitive tensor
|
||||
functions we focus on possess favorable properties, such as independence among loop iterations.
|
||||
For instance, the aforementioned program includes block and iteration annotations:
|
||||
|
||||
- The **block annotation** ``with T.sblock("C")`` signifies that the block is the fundamental
|
||||
computation unit designated for scheduling. A block may encompass a single computation
|
||||
statement, multiple computation statements with loops, or opaque intrinsics such as Tensor
|
||||
Core instructions.
|
||||
- The **iteration annotation** ``T.axis.spatial``, indicating that variable ``vi`` is mapped
|
||||
to ``i``, and all iterations are independent.
|
||||
|
||||
While this information isn't crucial for *executing* the specific program, it proves useful when
|
||||
transforming the program. Consequently, we can confidently parallelize or reorder loops associated
|
||||
with ``vi``, provided we traverse all the index elements from 0 to 128.
|
||||
@@ -0,0 +1,43 @@
|
||||
.. 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.
|
||||
|
||||
.. _tensor-ir-deep-dive:
|
||||
|
||||
TensorIR
|
||||
========
|
||||
TensorIR is one of the core abstractions in the Apache TVM stack, used to
|
||||
represent and optimize primitive tensor functions.
|
||||
|
||||
The TensorIR codebase consists of two modules (split from the former ``tir``):
|
||||
|
||||
- **tirx** — Core IR definitions and lowering (PrimFunc, Buffer, SBlock,
|
||||
expressions, statements, lowering passes).
|
||||
- **s_tir** (Schedulable TIR) — Schedule primitives, MetaSchedule, DLight,
|
||||
and tensor intrinsics.
|
||||
|
||||
In TVMScript, both modules are accessed via
|
||||
``from tvm.script import tirx as T``.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
abstraction
|
||||
learning
|
||||
tutorials/tir_creation
|
||||
tutorials/tir_transformation
|
||||
tutorials/dlight_gpu_scheduling
|
||||
tutorials/meta_schedule
|
||||
@@ -0,0 +1,255 @@
|
||||
.. 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.
|
||||
|
||||
.. _tirx-learning:
|
||||
|
||||
Understand TensorIR Abstraction
|
||||
===============================
|
||||
TensorIR is the tensor program abstraction in Apache TVM, which is one of the standard
|
||||
machine learning compilation frameworks. The principal objective of tensor program abstraction
|
||||
is to depict loops and associated hardware acceleration options, including threading, the
|
||||
application of specialized hardware instructions, and memory access.
|
||||
|
||||
To help our explanations, let us use the following sequence of tensor computations as
|
||||
a motivating example. Specifically, for two :math:`128 \times 128` matrices ``A`` and ``B``, let us perform the
|
||||
following two steps of tensor computations.
|
||||
|
||||
.. math::
|
||||
|
||||
Y_{i, j} &= \sum_k A_{i, k} \times B_{k, j} \\
|
||||
C_{i, j} &= \mathbb{relu}(Y_{i, j}) = \mathbb{max}(Y_{i, j}, 0)
|
||||
|
||||
|
||||
The above computations resemble a typical primitive tensor function commonly seen in neural networks,
|
||||
a linear layer with relu activation. We use TensorIR to depict the above computations as follows.
|
||||
|
||||
Before we invoke TensorIR, let's use native Python codes with NumPy to show the computation:
|
||||
|
||||
.. code:: python
|
||||
|
||||
def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):
|
||||
Y = np.empty((128, 128), dtype="float32")
|
||||
for i in range(128):
|
||||
for j in range(128):
|
||||
for k in range(128):
|
||||
if k == 0:
|
||||
Y[i, j] = 0
|
||||
Y[i, j] = Y[i, j] + A[i, k] * B[k, j]
|
||||
for i in range(128):
|
||||
for j in range(128):
|
||||
C[i, j] = max(Y[i, j], 0)
|
||||
|
||||
With the low-level NumPy example in mind, now we are ready to introduce TensorIR. The code block
|
||||
below shows a TensorIR implementation of ``mm_relu``. The particular code is implemented in a
|
||||
language called TVMScript, which is a domain-specific dialect embedded in python AST.
|
||||
|
||||
.. code:: python
|
||||
|
||||
from tvm.script import tirx as T
|
||||
|
||||
@tvm.script.ir_module
|
||||
class MyModule:
|
||||
@T.prim_func
|
||||
def mm_relu(A: T.Buffer((128, 128), "float32"),
|
||||
B: T.Buffer((128, 128), "float32"),
|
||||
C: T.Buffer((128, 128), "float32")):
|
||||
Y = T.alloc_buffer((128, 128), dtype="float32")
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock("Y"):
|
||||
vi = T.axis.spatial(128, i)
|
||||
vj = T.axis.spatial(128, j)
|
||||
vk = T.axis.reduce(128, k)
|
||||
with T.init():
|
||||
Y[vi, vj] = T.float32(0)
|
||||
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock("C"):
|
||||
vi = T.axis.spatial(128, i)
|
||||
vj = T.axis.spatial(128, j)
|
||||
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
|
||||
|
||||
|
||||
Next, let's invest the elements in the above TensorIR program.
|
||||
|
||||
Function Parameters and Buffers
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
**The function parameters correspond to the same set of parameters on the numpy function.**
|
||||
|
||||
.. code:: python
|
||||
|
||||
# TensorIR
|
||||
def mm_relu(A: T.Buffer((128, 128), "float32"),
|
||||
B: T.Buffer((128, 128), "float32"),
|
||||
C: T.Buffer((128, 128), "float32")):
|
||||
...
|
||||
# NumPy
|
||||
def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):
|
||||
...
|
||||
|
||||
Here ``A``, ``B``, and ``C`` takes a type named ``T.Buffer``, which with shape
|
||||
argument ``(128, 128)`` and data type ``float32``. This additional information
|
||||
helps possible MLC process to generate code that specializes in the shape and data
|
||||
type.
|
||||
|
||||
**Similarly, TensorIR also uses a buffer type in intermediate result allocation.**
|
||||
|
||||
.. code:: python
|
||||
|
||||
# TensorIR
|
||||
Y = T.alloc_buffer((128, 128), dtype="float32")
|
||||
# NumPy
|
||||
Y = np.empty((128, 128), dtype="float32")
|
||||
|
||||
Loop Iterations
|
||||
~~~~~~~~~~~~~~~
|
||||
**There are also direct correspondence of loop iterations.**
|
||||
|
||||
``T.grid`` is a syntactic sugar in TensorIR for us to write multiple nested iterators.
|
||||
|
||||
.. code:: python
|
||||
|
||||
# TensorIR with `T.grid`
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
...
|
||||
# TensorIR with `range`
|
||||
for i in range(128):
|
||||
for j in range(128):
|
||||
for k in range(128):
|
||||
...
|
||||
# NumPy
|
||||
for i in range(128):
|
||||
for j in range(128):
|
||||
for k in range(128):
|
||||
...
|
||||
|
||||
Computational Block
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
A significant distinction lies in computational statements:
|
||||
**TensorIR incorporates an additional construct termed** ``T.sblock``.
|
||||
|
||||
.. code:: python
|
||||
|
||||
# TensorIR
|
||||
with T.sblock("Y"):
|
||||
vi = T.axis.spatial(128, i)
|
||||
vj = T.axis.spatial(128, j)
|
||||
vk = T.axis.reduce(128, k)
|
||||
with T.init():
|
||||
Y[vi, vj] = T.float32(0)
|
||||
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
# NumPy
|
||||
vi, vj, vk = i, j, k
|
||||
if vk == 0:
|
||||
Y[vi, vj] = 0
|
||||
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
A **block** represents a fundamental computation unit within TensorIR. Importantly,
|
||||
a block encompasses more information than standard NumPy code. It comprises a set of block axes
|
||||
``(vi, vj, vk)`` and the computations delineated around them.
|
||||
|
||||
.. code:: python
|
||||
|
||||
vi = T.axis.spatial(128, i)
|
||||
vj = T.axis.spatial(128, j)
|
||||
vk = T.axis.reduce(128, k)
|
||||
|
||||
The above three lines declare the **key properties** about block axes in the following syntax.
|
||||
|
||||
.. code:: python
|
||||
|
||||
[block_axis] = T.axis.[axis_type]([axis_range], [mapped_value])
|
||||
|
||||
These three lines convey the following details:
|
||||
|
||||
- They specify the binding of ``vi``, ``vj``, ``vk`` (in this instance, to ``i``, ``j``, ``k``).
|
||||
- They declare the original range intended for ``vi``, ``vj``, ``vk``
|
||||
(the 128 in ``T.axis.spatial(128, i)``).
|
||||
- They announce the properties of the iterators (spatial, reduce).
|
||||
|
||||
Block Axis Properties
|
||||
~~~~~~~~~~~~~~~~~~~~~
|
||||
Let's delve deeper into the properties of the block axis. These properties signify the axis's
|
||||
relationship to the computation in progress. The block comprises three axes ``vi``, ``vj``, and
|
||||
``vk``, meanwhile the block reads the buffer ``A[vi, vk]``, ``B[vk, vj]`` and writes the buffer
|
||||
``Y[vi, vj]``. Strictly speaking, the block performs (reduction) updates to Y, which we label
|
||||
as write for the time being, as we don't require the value of Y from another block.
|
||||
|
||||
Significantly, for a fixed value of ``vi`` and ``vj``, the computation block yields a point
|
||||
value at a spatial location of ``Y`` (``Y[vi, vj]``) that is independent of other locations in ``Y``
|
||||
(with different ``vi``, ``vj`` values). We can refer to ``vi``, ``vj`` as **spatial axes** since
|
||||
they directly correspond to the start of a spatial region of buffers that the block writes to.
|
||||
The axes involved in reduction (``vk``) are designated as **reduce axes**.
|
||||
|
||||
Why Extra Information in Block
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
One crucial observation is that the additional information (block axis range and their properties)
|
||||
makes the block to be **self-contained** when it comes to the iterations that it is supposed to
|
||||
carry out independent from the external loop-nest ``i, j, k``.
|
||||
|
||||
The block axis information also provides additional properties that help us to validate the correctness of the
|
||||
external loops that are used to carry out the computation. For example, the above code block will result in an
|
||||
error because the loop expects an iterator of size 128, but we only bound it to a for loop of size 127.
|
||||
|
||||
.. code:: python
|
||||
|
||||
# wrong program due to loop and block iteration mismatch
|
||||
for i in range(127):
|
||||
with T.sblock("C"):
|
||||
vi = T.axis.spatial(128, i)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
error here due to iterator size mismatch
|
||||
...
|
||||
|
||||
Sugars for Block Axes Binding
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
In situations where each of the block axes is directly mapped to an outer loop iterator,
|
||||
we can use ``T.axis.remap`` to declare the block axis in a single line.
|
||||
|
||||
.. code:: python
|
||||
|
||||
# SSR means the properties of each axes are "spatial", "spatial", "reduce"
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
|
||||
which is equivalent to
|
||||
|
||||
.. code:: python
|
||||
|
||||
vi = T.axis.spatial(range_of_i, i)
|
||||
vj = T.axis.spatial(range_of_j, j)
|
||||
vk = T.axis.reduce (range_of_k, k)
|
||||
|
||||
So we can also write the programs as follows.
|
||||
|
||||
.. code:: python
|
||||
|
||||
@tvm.script.ir_module
|
||||
class MyModuleWithAxisRemapSugar:
|
||||
@T.prim_func
|
||||
def mm_relu(A: T.Buffer((128, 128), "float32"),
|
||||
B: T.Buffer((128, 128), "float32"),
|
||||
C: T.Buffer((128, 128), "float32")):
|
||||
Y = T.alloc_buffer((128, 128), dtype="float32")
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock("Y"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
Y[vi, vj] = T.float32(0)
|
||||
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock("C"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
|
||||
@@ -0,0 +1,2 @@
|
||||
Deep Dive: TensorIR
|
||||
-------------------
|
||||
@@ -0,0 +1,316 @@
|
||||
# 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
|
||||
|
||||
"""
|
||||
.. _dlight_gpu_scheduling:
|
||||
|
||||
DLight: Rule-Based GPU Scheduling
|
||||
==================================
|
||||
TIR functions produced by Relax legalization need GPU-specific scheduling — thread binding,
|
||||
loop tiling, shared memory usage — before they can run efficiently on a GPU. There are two
|
||||
main approaches in TVM:
|
||||
|
||||
- **MetaSchedule**: explores a search space to find the best schedule. High quality, but
|
||||
compilation takes minutes to hours.
|
||||
- **DLight**: applies pre-defined scheduling rules deterministically. No tuning required,
|
||||
compilation completes in seconds. Performance is excellent for well-known patterns
|
||||
(e.g., GEMM, GEMV in LLM workloads) and fair for the rest.
|
||||
|
||||
This tutorial covers how DLight works, what rules are available, how to diagnose scheduling
|
||||
quality, and how to write custom rules.
|
||||
|
||||
.. contents:: Table of Contents
|
||||
:local:
|
||||
:depth: 1
|
||||
"""
|
||||
|
||||
######################################################################
|
||||
# Prepare a Model
|
||||
# ---------------
|
||||
# We build a small model with ``nn.Module`` that is rich enough to trigger multiple DLight
|
||||
# rules: ``Linear`` layers produce GEMM (matrix multiplication) kernels, ``LayerNorm``
|
||||
# produces a general-reduction kernel, and ``ReLU`` is a simple elementwise op.
|
||||
|
||||
import tvm
|
||||
from tvm import relax, tirx
|
||||
from tvm.relax.frontend import nn
|
||||
from tvm.s_tir import dlight as dl
|
||||
|
||||
|
||||
class DemoModel(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.fc1 = nn.Linear(768, 768)
|
||||
self.relu = nn.ReLU()
|
||||
self.norm = nn.LayerNorm(768)
|
||||
self.fc2 = nn.Linear(768, 256)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.norm(self.relu(self.fc1(x)))
|
||||
return self.fc2(x)
|
||||
|
||||
|
||||
mod, params = DemoModel().export_tvm({"forward": {"x": nn.spec.Tensor((1, 768), "float32")}})
|
||||
|
||||
######################################################################
|
||||
# Legalize Relax operators into TIR functions so that DLight has concrete kernels to schedule.
|
||||
|
||||
device = tvm.cuda(0)
|
||||
target = tvm.target.Target.from_device(device)
|
||||
with target:
|
||||
mod = relax.get_pipeline("zero")(mod)
|
||||
|
||||
######################################################################
|
||||
# At this point every TIR function in ``mod`` is **unscheduled** — it has no thread bindings
|
||||
# and would not run efficiently on a GPU. Let's see what functions we have:
|
||||
for gv, func in mod.functions_items():
|
||||
if isinstance(func, tirx.PrimFunc):
|
||||
print(f" {gv.name_hint}")
|
||||
|
||||
######################################################################
|
||||
# Basic Usage: ApplyDefaultSchedule
|
||||
# ---------------------------------
|
||||
# ``ApplyDefaultSchedule`` is an ``IRModule`` pass. It iterates over every TIR function in the
|
||||
# module and tries the given rules **in order**. For each function the first rule whose
|
||||
# ``apply()`` returns a non-``None`` schedule wins; subsequent rules are skipped.
|
||||
# After scheduling, the function is marked with ``tirx.is_scheduled`` so it won't be
|
||||
# scheduled again by a later ``ApplyDefaultSchedule`` call.
|
||||
|
||||
######################################################################
|
||||
# Here we use a common subset of rules. The full catalog (including ``LowBatchGEMV``,
|
||||
# ``Transpose``, ``RMSNorm``) is listed in the next section.
|
||||
|
||||
with target:
|
||||
scheduled_mod = dl.ApplyDefaultSchedule(
|
||||
dl.gpu.Matmul(), # GEMM: dense matrix multiplication
|
||||
dl.gpu.GEMV(), # matrix-vector products
|
||||
dl.gpu.Reduction(), # simple reductions (sum, max, ...)
|
||||
dl.gpu.GeneralReduction(), # compound reductions (softmax, layer norm, ...)
|
||||
dl.gpu.Fallback(), # catch-all for anything unmatched above
|
||||
)(mod)
|
||||
|
||||
scheduled_mod.show()
|
||||
|
||||
######################################################################
|
||||
# Compared with the unscheduled IR, you can now see thread bindings
|
||||
# (``blockIdx.x``, ``threadIdx.x``, ...) and loop transformations in each TIR function.
|
||||
|
||||
######################################################################
|
||||
# Rule Catalog
|
||||
# ------------
|
||||
# DLight ships a set of GPU scheduling rules. Each rule is a subclass of
|
||||
# ``ScheduleRule`` and implements an ``apply(func, target, tunable)`` method that returns
|
||||
# a ``Schedule`` if the rule matches, or ``None`` to pass.
|
||||
#
|
||||
# The built-in GPU rules, roughly from most specific to most general:
|
||||
#
|
||||
# .. list-table::
|
||||
# :header-rows: 1
|
||||
# :widths: 20 40 40
|
||||
#
|
||||
# * - Rule
|
||||
# - Pattern
|
||||
# - Typical operators
|
||||
# * - ``Matmul``
|
||||
# - GEMM index pattern ``C[S,I,J] += A[S,I,K] * B[S,J,K]``
|
||||
# - ``nn.Linear``, batched matmul
|
||||
# * - ``GEMV``
|
||||
# - Matrix-vector multiply (one dimension is 1)
|
||||
# - single-batch decode in attention
|
||||
# * - ``LowBatchGEMV``
|
||||
# - Low-batch GEMM scheduled with a GEMV strategy
|
||||
# - small-batch decode
|
||||
# * - ``Reduction``
|
||||
# - Simple accumulation ``X[...] += Y[...]``
|
||||
# - sum, max, argmax
|
||||
# * - ``GeneralReduction``
|
||||
# - Spatial dims followed by reduction dims (``S* R*``)
|
||||
# - softmax, layer norm, RMS norm
|
||||
# * - ``Transpose``
|
||||
# - Read/write indices are permutations of each other
|
||||
# - 2-D transpose
|
||||
# * - ``RMSNorm``
|
||||
# - Contains an ``rsqrt`` operation
|
||||
# - RMS normalization
|
||||
# * - ``Fallback``
|
||||
# - Any function (always matches)
|
||||
# - generic catch-all
|
||||
#
|
||||
# **Rule order matters.** ``ApplyDefaultSchedule`` stops at the first match, so:
|
||||
#
|
||||
# - Put **specialized** rules first (``Matmul``, ``GEMV``) — they have strict matching
|
||||
# conditions but produce high-quality schedules.
|
||||
# - Put **general** rules later (``GeneralReduction``, ``Fallback``) — they match broadly
|
||||
# but with less optimal schedules.
|
||||
# - If you put ``Fallback`` first, it would "steal" every function and no specialized
|
||||
# rule would ever run.
|
||||
|
||||
######################################################################
|
||||
# Diagnosing Schedule Quality
|
||||
# ---------------------------
|
||||
# A common question is: *which rule scheduled which function?* ``ApplyDefaultSchedule``
|
||||
# does not log this directly, but you can figure it out by applying rules one at a time.
|
||||
#
|
||||
# **Step 1**: Apply each rule individually and record which functions it claims.
|
||||
|
||||
from collections import OrderedDict
|
||||
|
||||
rules = OrderedDict(
|
||||
[
|
||||
("Matmul", dl.gpu.Matmul()),
|
||||
("GEMV", dl.gpu.GEMV()),
|
||||
("LowBatchGEMV", dl.gpu.LowBatchGEMV()),
|
||||
("Reduction", dl.gpu.Reduction()),
|
||||
("GeneralReduction", dl.gpu.GeneralReduction()),
|
||||
("Transpose", dl.gpu.Transpose()),
|
||||
("RMSNorm", dl.gpu.RMSNorm()),
|
||||
]
|
||||
)
|
||||
|
||||
rule_assignment = {}
|
||||
for rule_name, rule in rules.items():
|
||||
with target:
|
||||
test_mod = dl.ApplyDefaultSchedule(rule)(mod)
|
||||
for gv, func in test_mod.functions_items():
|
||||
if isinstance(func, tirx.PrimFunc) and gv.name_hint not in rule_assignment:
|
||||
if "tirx.is_scheduled" in func.attrs and func.attrs["tirx.is_scheduled"] == 1:
|
||||
rule_assignment[gv.name_hint] = rule_name
|
||||
|
||||
######################################################################
|
||||
# **Step 2**: Functions not claimed by any specialized rule will fall through to ``Fallback``.
|
||||
|
||||
all_tir_funcs = [
|
||||
gv.name_hint for gv, func in mod.functions_items() if isinstance(func, tirx.PrimFunc)
|
||||
]
|
||||
fallback_funcs = [name for name in all_tir_funcs if name not in rule_assignment]
|
||||
|
||||
print("Rule assignments:")
|
||||
for name, rule_name in sorted(rule_assignment.items()):
|
||||
print(f" {name:40s} -> {rule_name}")
|
||||
if fallback_funcs:
|
||||
print("Handled by Fallback (may have suboptimal performance):")
|
||||
for name in sorted(fallback_funcs):
|
||||
print(f" {name}")
|
||||
|
||||
######################################################################
|
||||
# If an important kernel lands in the Fallback bucket, you have three options:
|
||||
#
|
||||
# 1. Write a **custom DLight rule** for it (see below).
|
||||
# 2. Use **MetaSchedule** to auto-tune that specific function.
|
||||
# 3. Manually schedule it with the ``tvm.s_tir.Schedule`` API.
|
||||
|
||||
######################################################################
|
||||
# DLight vs MetaSchedule
|
||||
# ----------------------
|
||||
# The two systems are complementary, not competing:
|
||||
#
|
||||
# .. list-table::
|
||||
# :header-rows: 1
|
||||
# :widths: 20 40 40
|
||||
#
|
||||
# * -
|
||||
# - DLight
|
||||
# - MetaSchedule
|
||||
# * - Mechanism
|
||||
# - Deterministic rule matching
|
||||
# - Search-space exploration
|
||||
# * - Compile time
|
||||
# - Seconds
|
||||
# - Minutes to hours
|
||||
# * - Performance
|
||||
# - Excellent on known patterns, fair otherwise
|
||||
# - Near-optimal with sufficient search budget
|
||||
# * - Best for
|
||||
# - Default path, rapid iteration, CI
|
||||
# - Hot-spot tuning in production
|
||||
#
|
||||
# A practical workflow:
|
||||
#
|
||||
# 1. Run ``ApplyDefaultSchedule`` with the full rule set to cover all functions.
|
||||
# 2. Profile the compiled model to identify hot-spot kernels.
|
||||
# 3. Use ``MetaScheduleTuneTIR`` to auto-tune only those kernels.
|
||||
#
|
||||
# Note that ``MetaScheduleTuneTIR`` does **not** automatically skip functions already
|
||||
# scheduled by DLight — it processes every ``PrimFunc`` in the module. In practice this
|
||||
# is harmless (tuning an already-scheduled function simply re-explores its space), but if
|
||||
# you want to avoid the extra search cost, filter the module or use ``MetaScheduleTuneIRMod``
|
||||
# with ``op_names`` to target specific functions.
|
||||
|
||||
######################################################################
|
||||
# Writing a Custom Rule
|
||||
# ---------------------
|
||||
# You can extend DLight by writing your own ``ScheduleRule``. The simplest way is
|
||||
# ``ScheduleRule.from_callable``, which wraps a plain function into a rule **instance**.
|
||||
|
||||
from tvm import s_tir
|
||||
from tvm.s_tir.dlight.analysis import normalize_prim_func
|
||||
from tvm.s_tir.dlight.base.schedule_rule import ScheduleRule
|
||||
|
||||
|
||||
@ScheduleRule.from_callable("MyTileAndBind")
|
||||
def my_tile_and_bind(func: tirx.PrimFunc, target: tvm.target.Target, tunable: bool):
|
||||
"""A minimal rule: for single-block injective functions, tile and bind to GPU threads."""
|
||||
if not isinstance(func, tirx.PrimFunc):
|
||||
return None
|
||||
sch = s_tir.Schedule(func)
|
||||
# Use normalize_prim_func to get block info with correct spatial/reduction classification.
|
||||
# This is the same analysis used by built-in DLight rules.
|
||||
block_infos = normalize_prim_func(sch)
|
||||
if block_infos is None or len(block_infos) != 1:
|
||||
return None # only handle single-block functions
|
||||
info = block_infos[0]
|
||||
if not info.is_injective():
|
||||
return None # skip reductions — dom_kind() uses iter_type, not loop kind
|
||||
loops = sch.get_loops(info.block_rv)
|
||||
if len(loops) == 0:
|
||||
return None
|
||||
fused = sch.fuse(*loops)
|
||||
bx, tx = sch.split(fused, factors=[None, 256])
|
||||
sch.bind(bx, "blockIdx.x")
|
||||
sch.bind(tx, "threadIdx.x")
|
||||
return sch
|
||||
|
||||
|
||||
######################################################################
|
||||
# Insert the custom rule into the rule chain. Note that ``from_callable`` returns an
|
||||
# **instance**, so pass it directly — do not call ``my_tile_and_bind()`` again.
|
||||
|
||||
with target:
|
||||
custom_mod = dl.ApplyDefaultSchedule(
|
||||
dl.gpu.Matmul(),
|
||||
dl.gpu.GeneralReduction(),
|
||||
my_tile_and_bind, # our custom rule, tried before Fallback
|
||||
dl.gpu.Fallback(),
|
||||
)(mod)
|
||||
|
||||
custom_mod.show()
|
||||
|
||||
######################################################################
|
||||
# To build a production-quality rule, subclass ``ScheduleRule`` directly and implement
|
||||
# ``apply()`` with full analysis logic (see ``tvm.s_tir.dlight.gpu.Matmul`` for an example).
|
||||
|
||||
######################################################################
|
||||
# Summary
|
||||
# -------
|
||||
# - **DLight** provides fast, deterministic GPU scheduling via rule matching.
|
||||
# - Rules are tried in order; the first match wins. Put specialized rules before general ones.
|
||||
# - Use the **single-rule probing** technique to diagnose which rule handles each function.
|
||||
# - Combine DLight with MetaSchedule: DLight for baseline coverage, MetaSchedule for hot-spot tuning.
|
||||
# - Extend DLight by writing custom ``ScheduleRule`` implementations.
|
||||
#
|
||||
# For DLight's role in the broader optimization pipeline, see :ref:`customize_opt`.
|
||||
@@ -0,0 +1,307 @@
|
||||
# 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
|
||||
|
||||
"""
|
||||
.. _meta_schedule_deep_dive:
|
||||
|
||||
MetaSchedule: Search-Based Auto-Tuning
|
||||
=======================================
|
||||
MetaSchedule is TVM's search-based auto-tuning framework, located in
|
||||
``python/tvm/s_tir/meta_schedule/``. It explores different TIR schedules
|
||||
(loop tiling, vectorization, thread binding, etc.) and measures them on real
|
||||
hardware to find the fastest implementation for each operator.
|
||||
|
||||
While **DLight** (see :ref:`dlight_gpu_scheduling`) provides rule-based scheduling with zero
|
||||
search time, MetaSchedule trades compilation time for better performance by searching over
|
||||
the space of possible schedules.
|
||||
|
||||
.. contents:: Table of Contents
|
||||
:local:
|
||||
:depth: 1
|
||||
"""
|
||||
|
||||
######################################################################
|
||||
# Architecture Overview
|
||||
# ---------------------
|
||||
# A MetaSchedule tuning session involves the following components:
|
||||
#
|
||||
# - **ExtractedTask**: A unique TIR workload extracted from a Relax IRModule,
|
||||
# with a ``task_name`` and ``weight`` (call frequency in the graph).
|
||||
# - **TuneContext**: Container holding all resources for a single tuning task
|
||||
# (module, target, space generator, search strategy, etc.).
|
||||
# - **SpaceGenerator** (default: ``PostOrderApply``): Generates the design space
|
||||
# of possible schedules by applying ``ScheduleRule`` instances to each block.
|
||||
# - **SearchStrategy** (default: ``EvolutionarySearch``): Explores the design
|
||||
# space using an evolutionary algorithm guided by a cost model.
|
||||
# - **CostModel** (default: ``XGBModel``): Predicts schedule performance using
|
||||
# XGBoost, reducing the number of actual hardware measurements needed.
|
||||
# Alternatives include ``MLPModel`` (neural network) and ``RandomModel``
|
||||
# (baseline).
|
||||
# - **Builder** / **Runner**: Compile and execute candidates on real hardware to
|
||||
# obtain measured run times.
|
||||
# - **Database** (default: ``JSONDatabase``): Persistently stores tuning records
|
||||
# (schedule traces + measured run times) for later retrieval.
|
||||
# - **TaskScheduler** (default: ``GradientBasedScheduler``): Allocates tuning
|
||||
# budget across multiple tasks based on their weights and estimated improvement
|
||||
# potential.
|
||||
#
|
||||
# The tuning loop works as follows:
|
||||
#
|
||||
# 1. The **TaskScheduler** picks a task to tune.
|
||||
# 2. The **SpaceGenerator** produces candidate schedules from the design space.
|
||||
# 3. The **SearchStrategy** selects candidates (guided by the **CostModel**),
|
||||
# sends them to the **Builder** and **Runner** for measurement.
|
||||
# 4. Measured results are committed to the **Database** and used to update the
|
||||
# **CostModel** for the next iteration.
|
||||
# 5. Repeat until the trial budget is exhausted.
|
||||
|
||||
######################################################################
|
||||
# Prepare a Model
|
||||
# ---------------
|
||||
# We reuse a simple model to demonstrate MetaSchedule APIs.
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import tvm
|
||||
from tvm import relax
|
||||
from tvm.relax.frontend import nn
|
||||
|
||||
|
||||
class DemoModel(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.fc1 = nn.Linear(784, 256)
|
||||
self.relu = nn.ReLU()
|
||||
self.fc2 = nn.Linear(256, 10, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.fc1(x)
|
||||
x = self.relu(x)
|
||||
x = self.fc2(x)
|
||||
return x
|
||||
|
||||
|
||||
input_shape = (1, 784)
|
||||
mod, params = DemoModel().export_tvm({"forward": {"x": nn.spec.Tensor(input_shape, "float32")}})
|
||||
|
||||
device = tvm.cuda(0)
|
||||
target = tvm.target.Target.from_device(device)
|
||||
|
||||
######################################################################
|
||||
# User-Facing Entry Points
|
||||
# ------------------------
|
||||
# MetaSchedule provides several levels of API, from high-level transforms to
|
||||
# low-level tuning functions.
|
||||
#
|
||||
# Transform-Based API (Recommended)
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# These are Relax passes that can be composed into a ``Sequential`` pipeline:
|
||||
#
|
||||
# - **MetaScheduleTuneIRMod**: Tunes an entire IRModule. Supports ``op_names``
|
||||
# for selective operator tuning.
|
||||
# - **MetaScheduleTuneTIR**: Tunes all TIR functions individually (no
|
||||
# ``op_names`` filtering).
|
||||
# - **MetaScheduleApplyDatabase**: Applies the best schedules from the tuning
|
||||
# database. Only replaces functions that have records; the rest are left
|
||||
# unchanged.
|
||||
#
|
||||
# Here is a typical tune-and-apply pipeline:
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# To save CI time and avoid flakiness, we skip the tuning process in CI.
|
||||
|
||||
if os.getenv("CI", "") != "true":
|
||||
with target, tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tuned_mod = tvm.ir.transform.Sequential(
|
||||
[
|
||||
relax.get_pipeline("zero"),
|
||||
relax.transform.MetaScheduleTuneTIR(
|
||||
work_dir=tmp_dir,
|
||||
max_trials_global=300,
|
||||
),
|
||||
relax.transform.MetaScheduleApplyDatabase(work_dir=tmp_dir),
|
||||
]
|
||||
)(mod)
|
||||
|
||||
tuned_mod.show()
|
||||
|
||||
######################################################################
|
||||
# Inspecting Tunable Tasks
|
||||
# ------------------------
|
||||
# Before tuning, use ``extract_tasks`` to see what MetaSchedule will tune:
|
||||
|
||||
from tvm.s_tir.meta_schedule.relax_integration import extract_tasks
|
||||
|
||||
with target:
|
||||
legalized_mod = relax.get_pipeline("zero")(mod)
|
||||
|
||||
tasks = extract_tasks(legalized_mod, target)
|
||||
for i, task in enumerate(tasks):
|
||||
print(f"Task {i}: {task.task_name} (weight={task.weight})")
|
||||
|
||||
######################################################################
|
||||
# Each ``ExtractedTask`` has:
|
||||
#
|
||||
# - ``task_name``: Derived from the PrimFunc name (e.g., ``"fused_matmul_add_relu"``).
|
||||
# - ``weight``: How many ``call_tir`` sites invoke this workload. The task
|
||||
# scheduler uses weights to allocate more budget to frequently-called operators.
|
||||
# - ``dispatched``: List of candidate TIR modules for this workload.
|
||||
|
||||
######################################################################
|
||||
# Selective Operator Tuning
|
||||
# -------------------------
|
||||
# ``MetaScheduleTuneIRMod`` accepts an ``op_names`` parameter to tune only
|
||||
# operators whose task name contains any of the given strings:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# with target:
|
||||
# mod = tvm.ir.transform.Sequential([
|
||||
# relax.transform.MetaScheduleTuneIRMod(
|
||||
# params={},
|
||||
# work_dir="./tuning_logs",
|
||||
# max_trials_global=300,
|
||||
# op_names=["matmul"], # Only tune matmul-related operators
|
||||
# ),
|
||||
# relax.transform.MetaScheduleApplyDatabase(work_dir="./tuning_logs"),
|
||||
# ])(mod)
|
||||
#
|
||||
# Operators without tuning records are left unscheduled -- you can apply DLight or
|
||||
# other rule-based schedules to cover them afterward.
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# ``MetaScheduleTuneTIR`` does not support ``op_names`` filtering. Use
|
||||
# ``MetaScheduleTuneIRMod`` when you need selective tuning.
|
||||
|
||||
######################################################################
|
||||
# Database
|
||||
# --------
|
||||
# When using a fixed ``work_dir``, tuning results are persisted in two
|
||||
# newline-delimited JSON files:
|
||||
#
|
||||
# - ``database_workload.json``: One line per unique workload (structural hash +
|
||||
# serialized IRModule).
|
||||
# - ``database_tuning_record.json``: One line per tuning record (workload index +
|
||||
# schedule trace + measured run times).
|
||||
#
|
||||
# Records are appended incrementally as tuning progresses.
|
||||
#
|
||||
# Resumption Semantics
|
||||
# ~~~~~~~~~~~~~~~~~~~~
|
||||
# When you re-run tuning with the same ``work_dir``, existing records are loaded
|
||||
# and used as warm-start seeds for the evolutionary search. The tuner does
|
||||
# **not** skip already-seen workloads entirely -- it starts from a better initial
|
||||
# population, so re-runs are faster than starting from scratch but still consume
|
||||
# trials.
|
||||
#
|
||||
# Once tuning is done, subsequent compilations only need
|
||||
# ``MetaScheduleApplyDatabase``:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# with target:
|
||||
# mod = relax.transform.MetaScheduleApplyDatabase(
|
||||
# work_dir="./tuning_logs"
|
||||
# )(mod)
|
||||
#
|
||||
# Database Implementations
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# MetaSchedule ships several database backends:
|
||||
#
|
||||
# - **JSONDatabase**: Persistent file-based storage (default). Created
|
||||
# automatically when you pass ``work_dir``.
|
||||
# - **MemoryDatabase**: In-memory, non-persistent. Useful for testing.
|
||||
# - **UnionDatabase**: Queries all sub-databases and returns the globally best
|
||||
# record.
|
||||
# - **OrderedUnionDatabase**: Queries sub-databases in order; returns from the
|
||||
# first one that has a match.
|
||||
# - **ScheduleFnDatabase**: Wraps a user-provided scheduling function.
|
||||
|
||||
######################################################################
|
||||
# Cross-Model Database Reuse
|
||||
# --------------------------
|
||||
# MetaSchedule identifies workloads by their structural hash. If two models
|
||||
# contain operators with the same shape, dtype, and computation, they share the
|
||||
# same hash and can reuse tuning records.
|
||||
#
|
||||
# module_equality Options
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# - ``"structural"`` (default): Exact structural match. Safe but strict.
|
||||
# - ``"anchor-block"``: Match based on the dominant compute block, ignoring
|
||||
# surrounding context. More permissive -- enables sharing across fused operators
|
||||
# that have the same core computation but different fusion boundaries.
|
||||
#
|
||||
# ``OrderedUnionDatabase`` enables a layered lookup strategy: check a local
|
||||
# database first, then fall back to a shared team database:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# from tvm.s_tir.meta_schedule.database import JSONDatabase, OrderedUnionDatabase
|
||||
#
|
||||
# local_db = JSONDatabase(work_dir="./my_tuning_logs")
|
||||
# shared_db = JSONDatabase(work_dir="/shared/tuning_db")
|
||||
# combined_db = OrderedUnionDatabase(local_db, shared_db)
|
||||
#
|
||||
# with target, combined_db:
|
||||
# mod = relax.transform.MetaScheduleApplyDatabase()(mod)
|
||||
|
||||
######################################################################
|
||||
# Key Parameters Reference
|
||||
# ------------------------
|
||||
#
|
||||
# .. list-table::
|
||||
# :header-rows: 1
|
||||
# :widths: 25 75
|
||||
#
|
||||
# * - Parameter
|
||||
# - Description
|
||||
# * - ``max_trials_global``
|
||||
# - Total trial budget shared across all tasks. Set proportional to the
|
||||
# number of tasks (e.g., 200-500 trials per task for good results).
|
||||
# * - ``max_trials_per_task``
|
||||
# - Per-task trial cap. Defaults to ``max_trials_global`` if not set.
|
||||
# * - ``op_names``
|
||||
# - List of strings to filter tasks by name (substring match).
|
||||
# ``MetaScheduleTuneIRMod`` only.
|
||||
# * - ``work_dir``
|
||||
# - Directory for database files and logs. Use a fixed path to enable
|
||||
# persistence and resumption.
|
||||
# * - ``cost_model``
|
||||
# - ``"xgb"`` (XGBoost, default), ``"mlp"`` (neural network), or
|
||||
# ``"random"`` (baseline). Only available via ``tune_relax``.
|
||||
# * - ``runner``
|
||||
# - ``"local"`` (default) or an ``RPCRunner`` instance for remote devices.
|
||||
# Only available via ``tune_relax``.
|
||||
# * - ``module_equality``
|
||||
# - ``"structural"`` (default) or ``"anchor-block"`` for more permissive
|
||||
# cross-model matching. Only available via ``tune_relax``.
|
||||
|
||||
######################################################################
|
||||
# Summary
|
||||
# -------
|
||||
# - **MetaSchedule** finds high-quality TIR schedules by searching over the
|
||||
# design space and measuring on real hardware.
|
||||
# - Use ``MetaScheduleTuneTIR`` for full-module tuning, or
|
||||
# ``MetaScheduleTuneIRMod`` with ``op_names`` for selective tuning.
|
||||
# - Tuning records persist in ``work_dir`` and can be reused across runs and
|
||||
# models with the same operator shapes.
|
||||
# - Combine with DLight: use DLight for fast baseline coverage, then MetaSchedule
|
||||
# for hot-spot tuning (see :ref:`dlight_gpu_scheduling`).
|
||||
@@ -0,0 +1,289 @@
|
||||
# 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
|
||||
|
||||
"""
|
||||
.. _tir-creation:
|
||||
|
||||
TensorIR Creation
|
||||
-----------------
|
||||
In this section, we will introduce the methods to write a TensorIR function
|
||||
in Apache TVM. This tutorial presumes familiarity with the fundamental concepts of TensorIR.
|
||||
If not already acquainted, please refer to :ref:`tirx-learning` initially.
|
||||
|
||||
.. note::
|
||||
|
||||
This tutorial concentrates on the construction of **standalone** TensorIR functions. The
|
||||
techniques presented here are not requisite for end users to compile Relax models.
|
||||
|
||||
"""
|
||||
|
||||
######################################################################
|
||||
# Create TensorIR using TVMScript
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# The most straightforward way to create a TensorIR function via TVMScript.
|
||||
# TVMScript is a TVM Python dialect that represents TensorIR in TVM.
|
||||
#
|
||||
# .. important::
|
||||
#
|
||||
# While TVMScript employs Python syntax and AST, ensuring full compatibility
|
||||
# with Python tools like auto-completion and linting, it is not a native Python
|
||||
# language and cannot be executed by a Python interpreter.
|
||||
#
|
||||
# More precisely, the decorator **@tvm.script** extracts the Python AST from
|
||||
# the decorated function, subsequently parsing it into TensorIR.
|
||||
#
|
||||
# Standard Format
|
||||
# ***************
|
||||
# Let's take an example of ``mm_relu`` from :ref:`tirx-learning`. Here is the complete
|
||||
# format of the ir_module and in TVMScript:
|
||||
|
||||
import numpy as np
|
||||
import tvm_ffi
|
||||
|
||||
import tvm
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
@I.ir_module
|
||||
class MyModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def mm_relu(
|
||||
A: T.Buffer((128, 128), "float32"),
|
||||
B: T.Buffer((128, 128), "float32"),
|
||||
C: T.Buffer((128, 128), "float32"),
|
||||
):
|
||||
Y = T.alloc_buffer((128, 128), dtype="float32")
|
||||
for i in range(128):
|
||||
for j in range(128):
|
||||
for k in range(128):
|
||||
with T.sblock("Y"):
|
||||
vi = T.axis.spatial(128, i)
|
||||
vj = T.axis.spatial(128, j)
|
||||
vk = T.axis.reduce(128, k)
|
||||
T.reads(A[vi, vk], B[vk, vj])
|
||||
T.writes(Y[vi, vj])
|
||||
with T.init():
|
||||
Y[vi, vj] = T.float32(0)
|
||||
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
for i in range(128):
|
||||
for j in range(128):
|
||||
with T.sblock("C"):
|
||||
vi = T.axis.spatial(128, i)
|
||||
vj = T.axis.spatial(128, j)
|
||||
T.reads(Y[vi, vj])
|
||||
T.writes(C[vi, vj])
|
||||
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
|
||||
|
||||
|
||||
######################################################################
|
||||
# Concise with Syntactic Sugar
|
||||
# ****************************
|
||||
# For ease of writing, we can employ the following syntactic sugar to
|
||||
# streamline the code:
|
||||
#
|
||||
# - Utilize ``T.grid`` to condense nested loops;
|
||||
# - Employ ``T.axis.remap`` to abbreviate block iterator annotations;
|
||||
# - Exclude ``T.reads`` and ``T.writes`` for blocks whose content can
|
||||
# be inferred from the block body;
|
||||
|
||||
|
||||
@I.ir_module
|
||||
class ConciseModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def mm_relu(
|
||||
A: T.Buffer((128, 128), "float32"),
|
||||
B: T.Buffer((128, 128), "float32"),
|
||||
C: T.Buffer((128, 128), "float32"),
|
||||
):
|
||||
Y = T.alloc_buffer((128, 128), dtype="float32")
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock("Y"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
Y[vi, vj] = T.float32(0)
|
||||
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock("C"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
|
||||
|
||||
|
||||
######################################################################
|
||||
# We can use the following code to verify that the two modules are equivalent:
|
||||
|
||||
print(tvm_ffi.structural_equal(MyModule, ConciseModule))
|
||||
|
||||
######################################################################
|
||||
# Interactive with Python Variables
|
||||
# *********************************
|
||||
# Despite TVMScript not being executed by a Python interpreter, limited
|
||||
# interaction with Python is feasible. For instance, Python variables can
|
||||
# be used to ascertain the shape and data type of a TensorIR.
|
||||
|
||||
# Python variables
|
||||
M = N = K = 128
|
||||
dtype = "float32"
|
||||
|
||||
|
||||
# IRModule in TVMScript
|
||||
@I.ir_module
|
||||
class ConciseModuleFromPython:
|
||||
@T.prim_func(s_tir=True)
|
||||
def mm_relu(
|
||||
A: T.Buffer((M, K), dtype),
|
||||
B: T.Buffer((K, N), dtype),
|
||||
C: T.Buffer((M, N), dtype),
|
||||
):
|
||||
Y = T.alloc_buffer((M, N), dtype)
|
||||
for i, j, k in T.grid(M, N, K):
|
||||
with T.sblock("Y"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
Y[vi, vj] = T.cast(T.float32(0), dtype)
|
||||
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
for i, j in T.grid(M, N):
|
||||
with T.sblock("C"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
C[vi, vj] = T.max(Y[vi, vj], T.cast(T.float32(0), dtype))
|
||||
|
||||
|
||||
######################################################################
|
||||
# Check the equivalence:
|
||||
|
||||
print(tvm_ffi.structural_equal(ConciseModule, ConciseModuleFromPython))
|
||||
|
||||
|
||||
######################################################################
|
||||
# TensorIR Function with Dynamic Shapes
|
||||
# *************************************
|
||||
# Despite TVMScript not being executed by a Python interpreter, limited
|
||||
# interaction with Python is feasible. For instance, Python variables can
|
||||
# be used to ascertain the shape and data type of a TensorIR.
|
||||
|
||||
|
||||
@I.ir_module
|
||||
class DynamicShapeModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def mm_relu(a: T.handle, b: T.handle, c: T.handle):
|
||||
# Dynamic shape definition
|
||||
M = T.int32()
|
||||
N = T.int32()
|
||||
K = T.int32()
|
||||
|
||||
# Bind the input buffers with the dynamic shapes
|
||||
A = T.match_buffer(a, [M, K], dtype)
|
||||
B = T.match_buffer(b, [K, N], dtype)
|
||||
C = T.match_buffer(c, [M, N], dtype)
|
||||
Y = T.alloc_buffer((M, N), dtype)
|
||||
for i, j, k in T.grid(M, N, K):
|
||||
with T.sblock("Y"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
Y[vi, vj] = T.cast(T.float32(0), dtype)
|
||||
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
for i, j in T.grid(M, N):
|
||||
with T.sblock("C"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
C[vi, vj] = T.max(Y[vi, vj], T.cast(T.float32(0), dtype))
|
||||
|
||||
|
||||
######################################################################
|
||||
# Now let's check the runtime dynamic shape inference:
|
||||
|
||||
|
||||
def evaluate_dynamic_shape(lib: tvm.runtime.Module, m: int, n: int, k: int):
|
||||
A = tvm.runtime.tensor(np.random.uniform(size=(m, k)).astype("float32"))
|
||||
B = tvm.runtime.tensor(np.random.uniform(size=(k, n)).astype("float32"))
|
||||
C = tvm.runtime.tensor(np.zeros((m, n), dtype="float32"))
|
||||
lib(A, B, C)
|
||||
return C.numpy()
|
||||
|
||||
|
||||
# Compile lib only once
|
||||
dyn_shape_lib = tvm.compile(DynamicShapeModule, target="llvm")
|
||||
# Able to handle different shapes
|
||||
print(evaluate_dynamic_shape(dyn_shape_lib, m=4, n=4, k=4))
|
||||
print(evaluate_dynamic_shape(dyn_shape_lib, m=64, n=64, k=128))
|
||||
|
||||
######################################################################
|
||||
# Create TensorIR using Tensor Expression
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Often, the specifics of TensorIR are disregarded in favor of expressing the computation more
|
||||
# succinctly, leading to the pragmatic generation of TensorIR. This is where Tensor Expression
|
||||
# (TE) becomes relevant.
|
||||
#
|
||||
# Tensor Expression (TE) serves as a domain-specific language delineating a sequence of
|
||||
# computations through an expression-like API.
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# Tensor Expression comprises two components within the TVM stack: the expression and the
|
||||
# schedule. The expression is the domain-specific language embodying the computation pattern,
|
||||
# precisely what we're addressing in this section. Conversely, the TE schedule is the legacy
|
||||
# scheduling method, has been superseded by the TensorIR schedule in the current TVM stack.
|
||||
#
|
||||
# Create Static-Shape Functions
|
||||
# *****************************
|
||||
# We use the same example of ``mm_relu`` from the last subsection to demonstrate the
|
||||
# TE creation method.
|
||||
|
||||
from tvm import te
|
||||
|
||||
A = te.placeholder((128, 128), "float32", name="A")
|
||||
B = te.placeholder((128, 128), "float32", name="B")
|
||||
k = te.reduce_axis((0, 128), "k")
|
||||
Y = te.compute((128, 128), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="Y")
|
||||
C = te.compute((128, 128), lambda i, j: te.max(Y[i, j], 0), name="C")
|
||||
|
||||
######################################################################
|
||||
# Here ``te.compute`` takes the signature ``te.compute(output_shape, fcompute)``.
|
||||
# And the fcompute function describes how we want to compute the value of each
|
||||
# element ``Y[i, j]`` for a given index:
|
||||
#
|
||||
# .. code:: python
|
||||
#
|
||||
# lambda i, j: te.sum(A[i, k] * B[k, j], axis=k)
|
||||
#
|
||||
# The aforementioned lambda expression encapsulates the computation:
|
||||
# :math:`Y_{i, j} = \sum_k A_{i, k} \times B_{k, j}`. Upon defining the computation,
|
||||
# we can formulate a TensorIR function by incorporating the pertinent parameters of interest.
|
||||
# In this specific instance, we aim to construct a function with two input parameters **A, B**
|
||||
# and one output parameter **C**.
|
||||
|
||||
te_func = te.create_prim_func([A, B, C]).with_attr({"global_symbol": "mm_relu"})
|
||||
TEModule = tvm.IRModule({"mm_relu": te_func})
|
||||
TEModule.show()
|
||||
|
||||
######################################################################
|
||||
# Create Dynamic-Shape Functions
|
||||
# ******************************
|
||||
# We can also create a dynamic-shape function using Tensor Expression. The only difference
|
||||
# is that we need to specify the shape of the input tensors as symbolic variables.
|
||||
|
||||
# Declare symbolic variables
|
||||
M, N, K = te.var("m"), te.var("n"), te.var("k")
|
||||
A = te.placeholder((M, N), "float32", name="A")
|
||||
B = te.placeholder((K, N), "float32", name="B")
|
||||
k = te.reduce_axis((0, K), "k")
|
||||
Y = te.compute((M, N), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="Y")
|
||||
C = te.compute((M, N), lambda i, j: te.max(Y[i, j], 0), name="C")
|
||||
|
||||
dyn_te_func = te.create_prim_func([A, B, C]).with_attr({"global_symbol": "mm_relu"})
|
||||
DynamicTEModule = tvm.IRModule({"mm_relu": dyn_te_func})
|
||||
DynamicTEModule.show()
|
||||
@@ -0,0 +1,177 @@
|
||||
# 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
|
||||
|
||||
"""
|
||||
.. _tirx-transform:
|
||||
|
||||
Transformation
|
||||
--------------
|
||||
In this section, we will get to the main ingredients of the compilation flows -
|
||||
transformations of primitive tensor functions.
|
||||
"""
|
||||
|
||||
######################################################################
|
||||
# In the :ref:`previous section <tirx-learning>`, we have given an example of how to write
|
||||
# ``mm_relu`` using TensorIR. In practice, there can be multiple ways to implement
|
||||
# the same functionality, and each implementation can result in different performance.
|
||||
#
|
||||
# .. note::
|
||||
# This tutorial primarily illustrates the application of TensorIR Transformation,
|
||||
# rather than delving into optimization techniques.
|
||||
#
|
||||
# First, let's take a look at the implementation of ``mm_relu`` in the previous section:
|
||||
|
||||
import tvm
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
@I.ir_module
|
||||
class MyModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
A: T.Buffer((128, 128), "float32"),
|
||||
B: T.Buffer((128, 128), "float32"),
|
||||
C: T.Buffer((128, 128), "float32"),
|
||||
):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
with T.sblock("root"):
|
||||
T.reads()
|
||||
T.writes()
|
||||
Y = T.sblock_alloc_buffer((128, 128))
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock("Y"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
Y[vi, vj] = T.float32(0)
|
||||
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock("C"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
|
||||
|
||||
|
||||
######################################################################
|
||||
# Before we transform the function, let's first evaluate the performance of the
|
||||
# original implementation.
|
||||
|
||||
import numpy as np
|
||||
|
||||
a_np = np.random.uniform(size=(128, 128)).astype("float32")
|
||||
b_np = np.random.uniform(size=(128, 128)).astype("float32")
|
||||
c_np = a_np @ b_np
|
||||
|
||||
a_nd = tvm.runtime.tensor(a_np)
|
||||
b_nd = tvm.runtime.tensor(b_np)
|
||||
c_nd = tvm.runtime.tensor(np.zeros((128, 128), dtype="float32"))
|
||||
|
||||
|
||||
def evaluate(mod: tvm.IRModule):
|
||||
lib = tvm.tirx.build(mod, target="llvm")
|
||||
# check correctness
|
||||
lib(a_nd, b_nd, c_nd)
|
||||
np.testing.assert_allclose(c_nd.numpy(), c_np, rtol=1e-5)
|
||||
# evaluate performance
|
||||
f_timer = lib.time_evaluator("main", tvm.cpu())
|
||||
print(f_timer(a_nd, b_nd, c_nd))
|
||||
|
||||
|
||||
evaluate(MyModule)
|
||||
|
||||
######################################################################
|
||||
# Initialization Schedule
|
||||
# ***********************
|
||||
# We initiate the process of code transformation by establishing a Schedule helper class,
|
||||
# utilizing the provided **MyModule** as input.
|
||||
|
||||
sch = tvm.s_tir.Schedule(MyModule)
|
||||
|
||||
######################################################################
|
||||
# Loop Tiling
|
||||
# ***********
|
||||
# Subsequently, we execute the requisite operations to acquire a reference to
|
||||
# block **Y** and its associated loops.
|
||||
|
||||
block_Y = sch.get_sblock("Y")
|
||||
i, j, k = sch.get_loops(block_Y)
|
||||
|
||||
######################################################################
|
||||
# We now proceed to execute the transformations. The initial modification involves
|
||||
# splitting loop ``j`` into two separate loops, with the inner loop possessing a
|
||||
# length of 8. It is crucial to understand that the transformation process is procedural;
|
||||
# thus, inadvertent execution of the block twice will yield an error stating the
|
||||
# non-existence of variable ``j``.
|
||||
|
||||
j0, j1 = sch.split(j, factors=[None, 8])
|
||||
|
||||
######################################################################
|
||||
# The outcome of the transformation can be examined, as it is retained within ``sch.mod``.
|
||||
|
||||
sch.mod.show()
|
||||
|
||||
######################################################################
|
||||
# Following the initial transformation phase, two supplementary loops, ``j_0`` and ``j_1``,
|
||||
# have been generated with respective ranges of 16 and 8. The subsequent
|
||||
# action involves reordering these two loops.
|
||||
|
||||
sch.reorder(j0, k, j1)
|
||||
sch.mod.show()
|
||||
evaluate(sch.mod)
|
||||
|
||||
######################################################################
|
||||
# Leverage Localities
|
||||
# *******************
|
||||
# Subsequently, we will execute two additional transformation steps to achieve a different
|
||||
# variant. First, we employ a primitive known as **reverse_compute_at** to relocate block
|
||||
# **C** to an inner loop of **Y**.
|
||||
|
||||
block_C = sch.get_sblock("C")
|
||||
sch.reverse_compute_at(block_C, j0)
|
||||
sch.mod.show()
|
||||
|
||||
######################################################################
|
||||
# Rewrite Reduction
|
||||
# *****************
|
||||
# Until now, the reduction initialization and update step have been maintained together
|
||||
# within a single block body. This amalgamated form facilitates loop transformations,
|
||||
# as the outer loops ``i``, ``j`` of initialization and updates generally need to remain
|
||||
# synchronized.
|
||||
#
|
||||
# Following the loop transformations, we can segregate the initialization of Y's elements
|
||||
# from the reduction update via the **decompose_reduction** primitive.
|
||||
|
||||
sch.decompose_reduction(block_Y, k)
|
||||
sch.mod.show()
|
||||
evaluate(sch.mod)
|
||||
|
||||
######################################################################
|
||||
# Trace the Transformation
|
||||
# ************************
|
||||
# TensorIR schedule is a procedural language, and the transformation is executed in a
|
||||
# step-by-step manner. We can trace the transformation by printing the schedule or the
|
||||
# history of the schedule.
|
||||
#
|
||||
# We've already see the schedule by printing ``sch.mod``. We can also print the history
|
||||
# of the schedule by ``sch.trace``.
|
||||
|
||||
sch.trace.show()
|
||||
|
||||
######################################################################
|
||||
# Alternatively, we can output the IRModule in conjunction with the historical trace.
|
||||
|
||||
sch.show()
|
||||
@@ -0,0 +1,318 @@
|
||||
# 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.
|
||||
|
||||
# pylint: disable=redefined-outer-name, missing-module-docstring
|
||||
import argparse
|
||||
import hashlib
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
from html.parser import HTMLParser
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
|
||||
# NOTE: This script is called by the Makefile via `make htmldepoly`.
|
||||
# It's not called every time the docs are built on CI. However, it's
|
||||
# can be only called during deployment stage, instead of building the docs.
|
||||
# Also, we can download the resources manually before running this script to
|
||||
# avoid the overhead of downloading the resources every time the docs are built.
|
||||
|
||||
# Set to store unique external URLs found during processing
|
||||
BASE_URL = "https://tvm.apache.org"
|
||||
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
HTML_DIR = os.path.join(SCRIPT_DIR, "_build/html")
|
||||
|
||||
|
||||
class ExternalURLParser(HTMLParser):
|
||||
"""HTML Parser to find external URLs in HTML content."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.external_urls: list[str] = []
|
||||
self.base_domain = urlparse(BASE_URL).netloc
|
||||
# Tags and their attributes that might contain external resources
|
||||
self.tags_to_check = {
|
||||
"img": "src",
|
||||
"script": "src",
|
||||
"iframe": "src",
|
||||
"video": "src",
|
||||
"audio": "src",
|
||||
"link": "href",
|
||||
"source": "src",
|
||||
"embed": "src",
|
||||
}
|
||||
|
||||
def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
|
||||
"""Handle HTML start tags to find external URLs."""
|
||||
if tag not in self.tags_to_check:
|
||||
return
|
||||
|
||||
attr_name = self.tags_to_check[tag]
|
||||
for name, value in attrs:
|
||||
if name != attr_name or not value:
|
||||
continue
|
||||
|
||||
if value.startswith(("http://", "https://")):
|
||||
domain = urlparse(value).netloc
|
||||
if domain and domain != self.base_domain:
|
||||
self.external_urls.append(value)
|
||||
|
||||
|
||||
def detect_html_external_urls(html_content: str) -> list[str]:
|
||||
"""
|
||||
Detect third-party embedded resources in HTML content.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
html_content : str
|
||||
The HTML content to analyze
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[str]
|
||||
List of external URLs found in the HTML content
|
||||
"""
|
||||
parser = ExternalURLParser()
|
||||
parser.feed(html_content)
|
||||
return parser.external_urls
|
||||
|
||||
|
||||
def detect_css_external_urls(css_content: str) -> list[str]:
|
||||
"""
|
||||
Detect external URLs in CSS content.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
css_content : str
|
||||
The CSS content to analyze
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[str]
|
||||
List of external URLs found in the CSS content
|
||||
"""
|
||||
external_urls: list[str] = []
|
||||
# Regex to find URLs in CSS
|
||||
url_pattern = re.compile(r'url\(["\']?(.*?)["\']?\)')
|
||||
matches = url_pattern.findall(css_content)
|
||||
for match in matches:
|
||||
if match.startswith(("http://", "https://")) and not match.startswith(BASE_URL):
|
||||
external_urls.append(match)
|
||||
return external_urls
|
||||
|
||||
|
||||
def all_files_in_dir(path: str) -> list[str]:
|
||||
"""
|
||||
Get a list of all files in a directory and its subdirectories.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
path : str
|
||||
The root directory path to search
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[str]
|
||||
List of full paths to all files found
|
||||
"""
|
||||
return [os.path.join(root, file) for root, _, files in os.walk(path) for file in files]
|
||||
|
||||
|
||||
def detect_urls(files: list[str], verbose: bool = False) -> list[str]:
|
||||
"""
|
||||
Detect external URLs in the given HTML and CSS files.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
files : List[str]
|
||||
List of file paths to check for external URLs
|
||||
verbose : bool, optional
|
||||
Whether to print verbose output, by default False
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[str]
|
||||
List of external URLs found in the files
|
||||
"""
|
||||
|
||||
external_urls: set[str] = set()
|
||||
for file in files:
|
||||
f_detect: Callable[[str, str], list[str]] | None = None
|
||||
if file.endswith(".html"):
|
||||
f_detect = detect_html_external_urls
|
||||
elif file.endswith(".css"):
|
||||
f_detect = detect_css_external_urls
|
||||
else:
|
||||
continue
|
||||
with open(file, encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
urls = f_detect(content)
|
||||
if verbose:
|
||||
print(f"Processing {file}")
|
||||
exist_urls, new_urls = 0, 0
|
||||
for url in urls:
|
||||
if url in external_urls:
|
||||
exist_urls += 1
|
||||
else:
|
||||
new_urls += 1
|
||||
if verbose:
|
||||
print(f"Found new {url}")
|
||||
print(f"Found {exist_urls} existing resources and {new_urls} new resources")
|
||||
external_urls.update(urls)
|
||||
if verbose:
|
||||
print(f"Total {len(external_urls)} external resources")
|
||||
print("External resources:")
|
||||
print("\n".join(external_urls))
|
||||
|
||||
return list(external_urls)
|
||||
|
||||
|
||||
def download_external_urls(
|
||||
external_urls: list[str], verbose: bool = False
|
||||
) -> tuple[dict[str, str], list[str]]:
|
||||
"""
|
||||
Download external URLs and save them to docs/_static/downloads.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
external_urls : List[str]
|
||||
List of external URLs to download
|
||||
verbose : bool, optional
|
||||
Whether to print verbose output, by default False
|
||||
|
||||
Returns
|
||||
-------
|
||||
Tuple[Dict[str, str], List[str]]
|
||||
A tuple containing:
|
||||
- Dictionary mapping original URLs to their downloaded file paths
|
||||
- List of paths to all downloaded files (including source maps)
|
||||
"""
|
||||
download_dir = os.path.join(HTML_DIR, "_static/downloads")
|
||||
os.makedirs(download_dir, exist_ok=True)
|
||||
used_file_names: set[str] = set()
|
||||
downloaded_files: list[str] = []
|
||||
remap_urls: dict[str, str] = {}
|
||||
for url in external_urls:
|
||||
query = urlparse(url).query
|
||||
if url.startswith("https://fonts.googleapis.com/css2"):
|
||||
file_name = f"{hashlib.md5(url.encode()).hexdigest()}.css"
|
||||
elif query:
|
||||
raise ValueError(f"Unsupported URL with query: {url}")
|
||||
else:
|
||||
file_name = urlparse(url).path.split("/")[-1]
|
||||
if verbose:
|
||||
print(f"remapping {url} to {file_name}")
|
||||
if file_name in used_file_names:
|
||||
raise ValueError(f"File name {file_name} already exists")
|
||||
used_file_names.add(file_name)
|
||||
response = requests.get(url, timeout=30)
|
||||
body = response.content
|
||||
with open(os.path.join(download_dir, file_name), "wb") as f:
|
||||
f.write(body)
|
||||
remap_urls[url] = os.path.join(download_dir, file_name)
|
||||
downloaded_files.append(os.path.join(download_dir, file_name))
|
||||
|
||||
# Also download the sourceMappingURL
|
||||
if not url.startswith("https://fonts.googleapis.com/css2"):
|
||||
map_file_name = f"{file_name}.map"
|
||||
response = requests.get(f"{url}.map", timeout=30)
|
||||
if response.status_code == 200:
|
||||
body = response.content
|
||||
with open(os.path.join(download_dir, map_file_name), "wb") as f:
|
||||
f.write(body)
|
||||
if verbose:
|
||||
print(f"Downloaded {map_file_name} for {url}")
|
||||
downloaded_files.append(os.path.join(download_dir, map_file_name))
|
||||
|
||||
return remap_urls, downloaded_files
|
||||
|
||||
|
||||
def replace_urls_in_files(remap_urls: dict[str, str], verbose: bool = False):
|
||||
"""
|
||||
Replace external URLs with their downloaded versions in HTML/CSS files.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
remap_urls : Dict[str, str]
|
||||
Dictionary mapping original URLs to their downloaded file paths
|
||||
verbose : bool, optional
|
||||
Whether to print verbose output, by default False
|
||||
"""
|
||||
for root, _, files in os.walk(HTML_DIR):
|
||||
for file in files:
|
||||
if not (file.endswith(".html") or file.endswith(".css")):
|
||||
continue
|
||||
|
||||
file_path = os.path.join(root, file)
|
||||
if verbose:
|
||||
print(f"Processing {file_path}")
|
||||
|
||||
# Calculate relative path from current file to _static/downloads
|
||||
rel_path = os.path.relpath(
|
||||
os.path.join(HTML_DIR, "_static/downloads"), os.path.dirname(file_path)
|
||||
)
|
||||
|
||||
with open(file_path, encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
|
||||
new_content = content
|
||||
for original_url, new_path in remap_urls.items():
|
||||
relative_url = os.path.join(rel_path, os.path.basename(new_path))
|
||||
new_content = new_content.replace(original_url, relative_url)
|
||||
|
||||
if new_content != content:
|
||||
with open(file_path, "w", encoding="utf-8") as f:
|
||||
f.write(new_content)
|
||||
if verbose:
|
||||
print(f"Updated {file_path}")
|
||||
|
||||
|
||||
def download_and_replace_urls(files: list[str] | None = None, verbose: bool = False):
|
||||
"""
|
||||
Download external URLs found in files and replace them with local copies.
|
||||
Recursively processes any new external URLs found in downloaded content.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
files : Optional[List[str]], optional
|
||||
List of files to check for external URLs. If None, checks all files under HTML_DIR
|
||||
verbose : bool, optional
|
||||
Whether to print verbose output, by default False
|
||||
"""
|
||||
if files is None:
|
||||
files = all_files_in_dir(HTML_DIR)
|
||||
remap_urls = {}
|
||||
while True:
|
||||
external_urls = detect_urls(files, verbose=verbose)
|
||||
if not external_urls:
|
||||
break
|
||||
round_remap_urls, files = download_external_urls(external_urls, verbose=verbose)
|
||||
remap_urls.update(round_remap_urls)
|
||||
|
||||
replace_urls_in_files(remap_urls, verbose=verbose)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = argparse.ArgumentParser()
|
||||
args.add_argument("-v", "--verbose", action="store_true")
|
||||
args.add_argument("-p", "--path", type=str, default=None)
|
||||
args = args.parse_args()
|
||||
|
||||
if args.path is not None:
|
||||
HTML_DIR = args.path
|
||||
download_and_replace_urls(verbose=args.verbose)
|
||||
@@ -0,0 +1,73 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
|
||||
TVM Errors
|
||||
==========
|
||||
|
||||
TVM may raise errors from Python code, from C++ code reached through the
|
||||
FFI, or from generated runtime modules. Error messages usually include
|
||||
a Python stack trace, and may also include a C++ stack trace when the
|
||||
error crosses the TVM FFI boundary.
|
||||
|
||||
Some errors report invalid user input, unsupported operators, missing
|
||||
runtime features, or unavailable hardware. Others report a failed
|
||||
internal check, usually raised by ``TVM_FFI_ICHECK`` or
|
||||
``TVM_FFI_THROW`` in C++ code. Internal check failures often indicate
|
||||
that TVM reached a state that the implementation did not expect.
|
||||
|
||||
What to Check First
|
||||
-------------------
|
||||
|
||||
- Make sure the TVM Python package and native libraries come from the
|
||||
same build. A common symptom of a mismatched environment is importing
|
||||
Python files from one checkout while loading ``libtvm`` from another.
|
||||
- Check that the required runtime is enabled in ``config.cmake``. For
|
||||
example, CUDA tests and CUDA compilation require a TVM build with
|
||||
CUDA support enabled.
|
||||
- Check that the target hardware is available to the process. GPU
|
||||
tests may be skipped or fail if the device is not visible inside the
|
||||
current container or environment.
|
||||
- If the error occurs while importing or converting a model, reduce the
|
||||
input to the smallest model, operator, or shape that reproduces the
|
||||
issue.
|
||||
|
||||
Reporting an Issue
|
||||
------------------
|
||||
|
||||
Search the `Apache TVM Discuss Forum <https://discuss.tvm.apache.org/>`_
|
||||
and the `TVM issue tracker <https://github.com/apache/tvm/issues>`_
|
||||
for the exact error message first. If you do not find an existing
|
||||
report, include the following details when starting a new discussion or
|
||||
filing an issue:
|
||||
|
||||
- The TVM version or git commit hash.
|
||||
- The Python version, operating system, and hardware.
|
||||
- The target and runtime being used, such as LLVM, CUDA, Vulkan, or RPC.
|
||||
- The relevant build configuration from ``config.cmake``.
|
||||
- A minimal script, model, input shape, or IR module that reproduces the
|
||||
failure.
|
||||
- The full error message, including both Python and C++ stack traces
|
||||
when present.
|
||||
|
||||
Developer Notes
|
||||
---------------
|
||||
|
||||
For guidance on raising typed errors from TVM code, see
|
||||
:ref:`error-handling-guide`. That guide covers when to use specific
|
||||
error types, how C++ error prefixes map to Python exceptions, and how
|
||||
``TVM_FFI_ICHECK`` interacts with TVM's error handling.
|
||||
@@ -0,0 +1,19 @@
|
||||
.. 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.
|
||||
|
||||
Index
|
||||
=====
|
||||
@@ -0,0 +1,66 @@
|
||||
.. 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.
|
||||
|
||||
Overview
|
||||
========
|
||||
|
||||
Apache TVM is a machine learning compilation framework, following the principle of **Python-first development**
|
||||
and **universal deployment**. It takes in pre-trained machine learning models,
|
||||
compiles and generates deployable modules that can be embedded and run everywhere. Apache TVM also enables customizing optimization processes to introduce new optimizations, libraries, codegen
|
||||
and more.
|
||||
|
||||
Key Principle
|
||||
-------------
|
||||
|
||||
- **Python-first**: the optimization process is fully customizable in Python.
|
||||
It is easy to customize the optimization pipeline without recompiling the TVM stack.
|
||||
- **Composable**: the optimization process is composable. It is easy to compose
|
||||
new optimization passes, libraries and codegen to the existing pipeline.
|
||||
|
||||
Key Goals
|
||||
---------
|
||||
|
||||
- **Optimize** performance of ML workloads, composing libraries and codegen.
|
||||
- **Deploy** ML workloads to a diverse set of new environments, including new runtime and new hardware.
|
||||
- **Continuously improve and customize** ML deployment pipeline in Python by quickly customizing library dispatching,
|
||||
bringing in customized operators and code generation.
|
||||
|
||||
Key Flow
|
||||
--------
|
||||
|
||||
Here is a typical flow of using TVM to deploy a machine learning model. For a runnable example,
|
||||
please refer to :ref:`quick_start`
|
||||
|
||||
1. **Import/construct an ML model**
|
||||
|
||||
TVM supports importing models from various frameworks, such as PyTorch and ONNX for generic ML models. Meanwhile, we can create models directly using Relax frontend for scenarios of large language models.
|
||||
|
||||
2. **Perform composable optimization** transformations via ``pipelines``
|
||||
|
||||
The pipeline encapsulates a collection of transformations to achieve two goals:
|
||||
|
||||
- **Graph Optimizations**: such as operator fusion, and layout rewrites.
|
||||
- **Tensor Program Optimization**: Map the operators to low-level implementations (both library or codegen)
|
||||
|
||||
.. note::
|
||||
|
||||
The two are goals but not the stages of the pipeline. The two optimizations are performed
|
||||
**at the same level**, or separately in two stages.
|
||||
|
||||
3. **Build and universal deploy**
|
||||
|
||||
Apache TVM aims to provide a universal deployment solution to bring machine learning everywhere with every language with minimum runtime support. TVM runtime can work in non-Python environments, so it works on mobile, edge devices or even bare metal devices. Additionally, TVM runtime comes with native data structures, and can also have zero copy exchange with the existing ecosystem (PyTorch, TensorFlow, TensorRT, etc.) using DLPack support.
|
||||
@@ -0,0 +1,2 @@
|
||||
Get Started
|
||||
-----------
|
||||
@@ -0,0 +1,285 @@
|
||||
# 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
|
||||
|
||||
"""
|
||||
.. _ir_module:
|
||||
|
||||
IRModule
|
||||
========
|
||||
This tutorial presents the core abstraction of Apache TVM, the IRModule.
|
||||
The IRModule encompasses the **entirety** of the ML models, incorporating the
|
||||
computational graph, tensor programs, and potential calls to external libraries.
|
||||
|
||||
.. contents:: Table of Contents
|
||||
:local:
|
||||
:depth: 1
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
######################################################################
|
||||
# Create IRModule
|
||||
# ---------------
|
||||
# IRModules can be initialized in various ways. We demonstrate a few of them
|
||||
# below.
|
||||
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
|
||||
|
||||
######################################################################
|
||||
# Import from existing models
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# The most common way to initialize an IRModule is to import from an existing
|
||||
# model. Apache TVM accommodates imports from a range of frameworks,
|
||||
# such as PyTorch and ONNX. This tutorial solely demonstrates the import process
|
||||
# from PyTorch.
|
||||
|
||||
|
||||
# Create a dummy model
|
||||
class TorchModel(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.fc1 = nn.Linear(784, 256)
|
||||
self.relu1 = nn.ReLU()
|
||||
self.fc2 = nn.Linear(256, 10)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.fc1(x)
|
||||
x = self.relu1(x)
|
||||
x = self.fc2(x)
|
||||
return x
|
||||
|
||||
|
||||
# Give an example argument to torch.export
|
||||
example_args = (torch.randn(1, 784, dtype=torch.float32),)
|
||||
|
||||
# Convert the model to IRModule
|
||||
with torch.no_grad():
|
||||
exported_program = export(TorchModel().eval(), example_args)
|
||||
mod_from_torch = from_exported_program(
|
||||
exported_program, keep_params_as_input=True, unwrap_unit_return_tuple=True
|
||||
)
|
||||
|
||||
mod_from_torch, params_from_torch = relax.frontend.detach_params(mod_from_torch)
|
||||
# Print the IRModule
|
||||
mod_from_torch.show()
|
||||
|
||||
######################################################################
|
||||
# Write with Relax NN Module
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Apache TVM also provides a set of PyTorch-liked APIs, to help users
|
||||
# write the IRModule directly.
|
||||
|
||||
from tvm.relax.frontend import nn
|
||||
|
||||
|
||||
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)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.fc1(x)
|
||||
x = self.relu1(x)
|
||||
x = self.fc2(x)
|
||||
return x
|
||||
|
||||
|
||||
mod_from_relax, params_from_relax = RelaxModel().export_tvm(
|
||||
{"forward": {"x": nn.spec.Tensor((1, 784), "float32")}}
|
||||
)
|
||||
mod_from_relax.show()
|
||||
|
||||
######################################################################
|
||||
# Create via TVMScript
|
||||
# ~~~~~~~~~~~~~~~~~~~~
|
||||
# TVMScript is a Python-based DSL for IRModules. We are able to
|
||||
# directly output the IRModule in the TVMScript syntax, or alternatively,
|
||||
# parse the TVMScript to obtain an IRModule.
|
||||
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import relax as R
|
||||
|
||||
|
||||
@I.ir_module
|
||||
class TVMScriptModule:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((1, 784), dtype="float32"),
|
||||
fc1_weight: R.Tensor((256, 784), dtype="float32"),
|
||||
fc1_bias: R.Tensor((256,), dtype="float32"),
|
||||
fc2_weight: R.Tensor((10, 256), dtype="float32"),
|
||||
fc2_bias: R.Tensor((10,), dtype="float32"),
|
||||
) -> R.Tensor((1, 10), dtype="float32"):
|
||||
R.func_attr({"num_input": 1})
|
||||
with R.dataflow():
|
||||
permute_dims = R.permute_dims(fc1_weight, axes=None)
|
||||
matmul = R.matmul(x, permute_dims, out_dtype=None)
|
||||
add = R.add(matmul, fc1_bias)
|
||||
relu = R.nn.relu(add)
|
||||
permute_dims1 = R.permute_dims(fc2_weight, axes=None)
|
||||
matmul1 = R.matmul(relu, permute_dims1, out_dtype=None)
|
||||
add1 = R.add(matmul1, fc2_bias)
|
||||
gv = add1
|
||||
R.output(gv)
|
||||
return gv
|
||||
|
||||
|
||||
mod_from_script = TVMScriptModule
|
||||
mod_from_script.show()
|
||||
|
||||
######################################################################
|
||||
# Attributes of an IRModule
|
||||
# -------------------------
|
||||
# An IRModule is a collection of functions, indexed by GlobalVars.
|
||||
|
||||
mod = mod_from_torch
|
||||
print(mod.get_global_vars())
|
||||
|
||||
######################################################################
|
||||
# We can access the functions in the IRModule by indexing with the GlobalVars
|
||||
# or their names
|
||||
|
||||
# index by global var name
|
||||
print(mod["main"])
|
||||
# index by global var, and checking they are the same function
|
||||
(gv,) = mod.get_global_vars()
|
||||
assert mod[gv] == mod["main"]
|
||||
|
||||
######################################################################
|
||||
# Transformations on IRModules
|
||||
# ----------------------------
|
||||
# Transformations are the import component of Apache TVM. One transformation
|
||||
# takes in an IRModule and outputs another IRModule. We can apply a sequence of
|
||||
# transformations to an IRModule to obtain a new IRModule. That is the common way to
|
||||
# optimize a model.
|
||||
#
|
||||
# In this getting started tutorial, we only demonstrate how to apply transformations
|
||||
# to an IRModule. For details of each transformation, please refer to the
|
||||
# :ref:`Transformation API Reference <api-relax-transformation>`
|
||||
|
||||
######################################################################
|
||||
# We first apply **LegalizeOps** transformation to the IRModule. This transformation
|
||||
# will convert the Relax module into a mixed stage, with both Relax and TensorIR function
|
||||
# within the same module. Meanwhile, the Relax operators will be converted into ``call_tir``.
|
||||
|
||||
mod = mod_from_torch
|
||||
mod = relax.transform.LegalizeOps()(mod)
|
||||
mod.show()
|
||||
|
||||
######################################################################
|
||||
# After the transformation, there are much more functions inside the module. Let's print
|
||||
# the global vars again.
|
||||
|
||||
print(mod.get_global_vars())
|
||||
|
||||
######################################################################
|
||||
# Next, Apache TVM provides a set of default transformation pipelines for users,
|
||||
# to simplify the transformation process. We can then apply the default pipeline to the module.
|
||||
# The default **zero** pipeline contains very fundamental transformations, including:
|
||||
#
|
||||
# - **LegalizeOps**: This transform converts the Relax operators into `call_tir` functions
|
||||
# with the corresponding TensorIR Functions. After this transform, the IRModule will
|
||||
# contain both Relax functions and TensorIR functions.
|
||||
# - **AnnotateTIROpPattern**: This transform annotates the pattern of the TensorIR functions,
|
||||
# preparing them for subsequent operator fusion.
|
||||
# - **FoldConstant**: This pass performs constant folding, optimizing operations
|
||||
# involving constants.
|
||||
# - **FuseOps and FuseTIR**: These two passes work together to fuse operators based on the
|
||||
# patterns annotated in the previous step (AnnotateTIROpPattern). These passes transform
|
||||
# both Relax functions and TensorIR functions.
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# Here, we have applied **LegalizeOps** twice in the flow. The second time is useless but
|
||||
# harmless.
|
||||
#
|
||||
# Every passes can be duplicated in the flow, since we ensure the passes can handle all legal
|
||||
# IRModule inputs. This design can help users to construct their own pipeline.
|
||||
|
||||
mod = relax.get_pipeline("zero")(mod)
|
||||
mod.show()
|
||||
|
||||
######################################################################
|
||||
# Deploy the IRModule Universally
|
||||
# -------------------------------
|
||||
# After the optimization, we can compile the model into a TVM runtime module.
|
||||
# Notably, Apache TVM provides the ability of universal deployment, which means
|
||||
# we can deploy the same IRModule on different backends, including CPU, GPU, and other emerging
|
||||
# backends.
|
||||
#
|
||||
# Deploy on CPU
|
||||
# ~~~~~~~~~~~~~
|
||||
# We can deploy the IRModule on CPU by specifying the target as ``llvm``.
|
||||
|
||||
exec = tvm.compile(mod, target="llvm")
|
||||
dev = tvm.cpu()
|
||||
vm = relax.VirtualMachine(exec, dev)
|
||||
|
||||
raw_data = np.random.rand(1, 784).astype("float32")
|
||||
data = tvm.runtime.tensor(raw_data, dev)
|
||||
cpu_out = vm["main"](data, *params_from_torch["main"]).numpy()
|
||||
print(cpu_out)
|
||||
|
||||
######################################################################
|
||||
# Deploy on GPU
|
||||
# ~~~~~~~~~~~~~
|
||||
# Besides, CPU backend, we can also deploy the IRModule on GPU. GPU requires
|
||||
# programs containing extra information, such as the thread bindings and shared memory
|
||||
# allocations. We need a further transformation to generate the GPU programs.
|
||||
#
|
||||
# We use ``DLight`` to generate the GPU programs. In this tutorial, we won't go into
|
||||
# the details of ``DLight``.
|
||||
#
|
||||
|
||||
from tvm.s_tir import dlight as dl
|
||||
|
||||
with tvm.target.Target("cuda"):
|
||||
gpu_mod = dl.ApplyDefaultSchedule(
|
||||
dl.gpu.Matmul(),
|
||||
dl.gpu.Fallback(),
|
||||
)(mod)
|
||||
|
||||
######################################################################
|
||||
# Now we can compile the IRModule on GPU, the similar way as we did on CPU.
|
||||
|
||||
exec = tvm.compile(gpu_mod, target="cuda")
|
||||
dev = tvm.device("cuda", 0)
|
||||
vm = relax.VirtualMachine(exec, dev)
|
||||
# Need to allocate data and params on GPU device
|
||||
data = tvm.runtime.tensor(raw_data, dev)
|
||||
gpu_params = [tvm.runtime.tensor(p, dev) for p in params_from_torch["main"]]
|
||||
gpu_out = vm["main"](data, *gpu_params).numpy()
|
||||
print(gpu_out)
|
||||
|
||||
# Check the correctness of the results
|
||||
assert np.allclose(cpu_out, gpu_out, atol=1e-3)
|
||||
|
||||
######################################################################
|
||||
# Deploy on Other Backends
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Apache TVM also supports other backends, such as different kinds of GPUs
|
||||
# (Metal, ROCm, Vulkan and OpenCL), different kinds of CPUs (x86, ARM), and other
|
||||
# emerging backends (e.g., WebAssembly). The deployment process is similar to the
|
||||
# GPU backend.
|
||||
@@ -0,0 +1,194 @@
|
||||
# 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
|
||||
|
||||
"""
|
||||
.. _quick_start:
|
||||
|
||||
Quick Start
|
||||
===========
|
||||
|
||||
This tutorial is for people who are new to Apache TVM. Taking an simple example
|
||||
to show how to use Apache TVM to compile a simple neural network.
|
||||
|
||||
.. contents:: Table of Contents
|
||||
:local:
|
||||
:depth: 2
|
||||
|
||||
"""
|
||||
|
||||
################################################################################
|
||||
# Overview
|
||||
# --------
|
||||
# Apache TVM is a machine learning compilation framework, following the principle of
|
||||
# **Python-first development** and **universal deployment**. It takes in pre-trained
|
||||
# machine learning models, compiles and generates deployable modules that can be embedded
|
||||
# and run everywhere.
|
||||
# Apache TVM also enables customizing optimization processes to introduce new optimizations,
|
||||
# libraries, codegen and more.
|
||||
#
|
||||
# Apache TVM can help to:
|
||||
#
|
||||
# - **Optimize** performance of ML workloads, composing libraries and codegen.
|
||||
# - **Deploy** ML workloads to a diverse set of new environments, including new runtime and new
|
||||
# hardware.
|
||||
# - **Continuously improve and customize** ML deployment pipeline in Python by quickly customizing
|
||||
# library dispatching, bringing in customized operators and code generation.
|
||||
|
||||
################################################################################
|
||||
# Overall Flow
|
||||
# ------------
|
||||
# Then we will show the overall flow of using Apache TVM to compile a neural network model,
|
||||
# showing how to optimize, deploy and run the model.
|
||||
# The overall flow is illustrated as the figure:
|
||||
#
|
||||
# .. 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 or Import a Model
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Before we get started, let's construct a neural network model first.
|
||||
# In this tutorial, to make things simple, we will define a two-layer MLP network
|
||||
# directly in this script with the TVM Relax frontend, which is a similar API to PyTorch.
|
||||
#
|
||||
|
||||
import tvm
|
||||
from tvm import relax
|
||||
from tvm.relax.frontend import nn
|
||||
|
||||
|
||||
class MLPModel(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.fc1 = nn.Linear(784, 256)
|
||||
self.relu1 = nn.ReLU()
|
||||
self.fc2 = nn.Linear(256, 10)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.fc1(x)
|
||||
x = self.relu1(x)
|
||||
x = self.fc2(x)
|
||||
return x
|
||||
|
||||
|
||||
################################################################################
|
||||
# Then we can export the model to TVM IRModule, which is the central intermediate representation
|
||||
# in TVM.
|
||||
|
||||
mod, param_spec = MLPModel().export_tvm(
|
||||
spec={"forward": {"x": nn.spec.Tensor((1, 784), "float32")}}
|
||||
)
|
||||
mod.show()
|
||||
|
||||
################################################################################
|
||||
# Perform Optimization Transformations
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Apache TVM leverage ``pipeline`` to transform and optimize program.
|
||||
# The pipeline encapsulates a collection of transformation that gets two goals (at the same level):
|
||||
#
|
||||
# - **Model optimizations**: such as operator fusion, layout rewrites.
|
||||
# - **Tensor program optimization**: Map the operators to low-level implementations
|
||||
# (both library or codegen)
|
||||
#
|
||||
# .. note::
|
||||
# The twos are goals but not the stages of the pipeline. The two optimizations are performed
|
||||
# **at the same level**, or separately in two stages.
|
||||
#
|
||||
# .. note::
|
||||
# In this tutorial we only demonstrate the overall flow, by leverage ``zero`` optimization
|
||||
# pipeline, instead of optimizing for any specific target.
|
||||
|
||||
mod = relax.get_pipeline("zero")(mod)
|
||||
|
||||
|
||||
################################################################################
|
||||
# Build and Universal Deployment
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# After the optimization, we can build the model to a deployable module and run it on
|
||||
# different devices.
|
||||
|
||||
|
||||
import numpy as np
|
||||
|
||||
target = tvm.target.Target("llvm")
|
||||
ex = tvm.compile(mod, target)
|
||||
device = tvm.cpu()
|
||||
vm = relax.VirtualMachine(ex, device)
|
||||
data = np.random.rand(1, 784).astype("float32")
|
||||
tvm_data = tvm.runtime.tensor(data, device=device)
|
||||
params = [np.random.rand(*param.shape).astype("float32") for _, param in param_spec]
|
||||
params = [tvm.runtime.tensor(param, device=device) for param in params]
|
||||
print(vm["forward"](tvm_data, *params).numpy())
|
||||
|
||||
################################################################################
|
||||
# Our goal is to bring machine learning to the application with any language of interest,
|
||||
# with the minimum runtime support.
|
||||
#
|
||||
# - Each function in IRModule becomes a runnable function in the runtime. For example in LLM
|
||||
# cases, we can call ``prefill`` and ``decode`` functions directly.
|
||||
#
|
||||
# .. code-block:: Python
|
||||
#
|
||||
# prefill_logits = vm["prefill"](inputs, weight, kv_cache)
|
||||
# decoded_logits = vm["decode"](inputs, weight, kv_cache)
|
||||
#
|
||||
# - TVM runtime comes with native data structures, such as Tensor, can also have zero
|
||||
# copy exchange with existing ecosystem (DLPack exchange with PyTorch)
|
||||
#
|
||||
# .. code-block:: Python
|
||||
#
|
||||
# # Convert PyTorch tensor to TVM Tensor
|
||||
# x_tvm = tvm.runtime.from_dlpack(x_torch)
|
||||
# # Convert TVM Tensor to PyTorch tensor
|
||||
# x_torch = torch.from_dlpack(x_tvm)
|
||||
#
|
||||
# - TVM runtime works in non-python environments, so it works on settings such as mobile
|
||||
#
|
||||
# .. code-block:: C++
|
||||
#
|
||||
# // C++ snippet
|
||||
# runtime::Module vm = ex.GetFunction("load_executable")();
|
||||
# vm.GetFunction("init")(...);
|
||||
# Tensor out = vm.GetFunction("prefill")(data, weight, kv_cache);
|
||||
#
|
||||
# .. code-block:: Java
|
||||
#
|
||||
# // Java snippet
|
||||
# Module vm = ex.getFunction("load_executable").invoke();
|
||||
# vm.getFunction("init").pushArg(...).invoke;
|
||||
# Tensor out = vm.getFunction("prefill").pushArg(data).pushArg(weight).pushArg(kv_cache).invoke();
|
||||
#
|
||||
|
||||
################################################################################
|
||||
# Read next
|
||||
# ---------
|
||||
# This tutorial demonstrates the overall flow of using Apache TVM to compile a neural network model.
|
||||
# For more advanced or specific topics, please refer to the following tutorials
|
||||
#
|
||||
@@ -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))
|
||||
@@ -0,0 +1,96 @@
|
||||
.. 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.
|
||||
|
||||
Apache TVM Documentation
|
||||
========================
|
||||
|
||||
Welcome to the documentation for Apache TVM, a deep learning compiler that
|
||||
enables access to high-performance machine learning anywhere for everyone.
|
||||
TVM's diverse community of hardware vendors, compiler engineers and ML
|
||||
researchers work together to build a unified, programmable software stack, that
|
||||
enriches the entire ML technology ecosystem and make it accessible to the wider
|
||||
ML community. TVM empowers users to leverage community-driven ML-based
|
||||
optimizations to push the limits and amplify the reach of their research and
|
||||
development, which in turn raises the collective performance of all ML, while
|
||||
driving its costs down.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: Getting Started
|
||||
|
||||
get_started/overview
|
||||
install/index
|
||||
get_started/tutorials/quick_start
|
||||
get_started/tutorials/ir_module
|
||||
errors
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: How To
|
||||
|
||||
how_to/tutorials/import_model
|
||||
how_to/tutorials/e2e_opt_model
|
||||
how_to/tutorials/customize_opt
|
||||
how_to/tutorials/optimize_llm
|
||||
how_to/tutorials/cross_compilation_and_rpc
|
||||
how_to/tutorials/export_and_load_executable
|
||||
how_to/tutorials/mix_python_and_tvm_with_pymodule
|
||||
how_to/tutorials/bring_your_own_codegen
|
||||
|
||||
.. The Deep Dive content is comprehensive
|
||||
.. we maintain a ``maxdepth`` of 2 to display more information on the main page.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
:caption: Deep Dive
|
||||
|
||||
arch/index
|
||||
deep_dive/tensor_ir/index
|
||||
deep_dive/relax/index
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
:caption: TIRx
|
||||
|
||||
tirx/overview
|
||||
tirx/install
|
||||
tirx/native_basics
|
||||
tirx/layout
|
||||
tirx/tile_primitives
|
||||
tirx/arch/index
|
||||
tirx/api/index
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: API Reference
|
||||
|
||||
reference/api/python/index
|
||||
reference/api/links
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: About
|
||||
|
||||
contribute/index
|
||||
reference/publications
|
||||
reference/security
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: Index
|
||||
|
||||
genindex
|
||||
@@ -0,0 +1,83 @@
|
||||
.. 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.
|
||||
|
||||
.. _docker-images:
|
||||
|
||||
Docker Images
|
||||
=============
|
||||
We provide docker utility scripts to help developers to setup development environment.
|
||||
They are also helpful run through TVM demo and tutorials.
|
||||
We need `docker <https://docs.docker.com/engine/installation/>`_ and
|
||||
`NVIDIA Container Toolkit <https://github.com/NVIDIA/nvidia-container-toolkit>`_
|
||||
if we want to use CUDA.
|
||||
|
||||
Get a tvm source distribution or clone the GitHub repo to get the auxiliary scripts
|
||||
|
||||
.. code:: bash
|
||||
|
||||
git clone --recursive https://github.com/apache/tvm tvm
|
||||
|
||||
|
||||
We can then use the following command to launch a docker image.
|
||||
|
||||
.. code:: bash
|
||||
|
||||
/path/to/tvm/docker/bash.sh <image-name>
|
||||
|
||||
Here the image-name can be a local docker image name, e.g. ``tvm.ci_cpu`` after you have done the local build.
|
||||
|
||||
This auxiliary script does the following things:
|
||||
|
||||
- Mount current directory to ``/workspace``
|
||||
- Switch user to be the same user that calls the ``bash.sh`` (so you can read/write host system)
|
||||
- Use the host-side network on Linux. Use the bridge network and expose port 8888 on macOS,
|
||||
because host networking driver isn't supported. (so you can use ``jupyter notebook``)
|
||||
|
||||
|
||||
Then you can start a Jupyter notebook by typing
|
||||
|
||||
.. code:: bash
|
||||
|
||||
jupyter notebook
|
||||
|
||||
You might see an error ``OSError: [Errno 99] Cannot assign requested address`` when starting
|
||||
a Jupyter notebook on macOS. You can change the binding IP address by
|
||||
|
||||
.. code:: bash
|
||||
|
||||
jupyter notebook --ip=0.0.0.0
|
||||
|
||||
Note that on macOS, because ``bash.sh`` uses the Docker bridge network, Jupyter will be reportedly running
|
||||
at an URL like ``http://{container_hostname}:8888/?token=...``. You should replace the ``container_hostname``
|
||||
with ``localhost`` when pasting it into browser.
|
||||
|
||||
|
||||
Docker Source
|
||||
-------------
|
||||
Check out `the docker source <https://github.com/apache/tvm/tree/main/docker>`_ if you are interested in
|
||||
building your own docker images.
|
||||
|
||||
|
||||
Run the following command to build the docker image.
|
||||
|
||||
.. code:: bash
|
||||
|
||||
/path/to/tvm/docker/build.sh <image-name>
|
||||
|
||||
|
||||
You can find some un-official third party pre-built images at `<https://hub.docker.com/r/tlcpack/>`_.
|
||||
These images are used for test purposes and are NOT of the ASF release.
|
||||
@@ -0,0 +1,356 @@
|
||||
.. 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.
|
||||
|
||||
.. _install-from-source:
|
||||
|
||||
Install from Source
|
||||
===================
|
||||
This page gives instructions on how to build and install the TVM package from source.
|
||||
|
||||
.. contents:: Table of Contents
|
||||
:local:
|
||||
:depth: 2
|
||||
|
||||
.. _install-dependencies:
|
||||
|
||||
Step 1. Install Dependencies
|
||||
----------------------------
|
||||
|
||||
Apache TVM requires the following dependencies:
|
||||
|
||||
- CMake (>= 3.24.0)
|
||||
- LLVM (recommended >= 15)
|
||||
- Git
|
||||
- A recent C++ compiler supporting C++ 20, at the minimum
|
||||
- GCC 10
|
||||
- Clang 10
|
||||
- Apple Clang 14
|
||||
- Visual Studio 2022
|
||||
|
||||
Optional dependencies that use newer C++20 standard library facilities, such
|
||||
as ``std::format``, may require a newer standard library (for example GCC 13
|
||||
or newer on Linux).
|
||||
- Python (>= 3.10)
|
||||
- (Optional) Conda (Strongly Recommended)
|
||||
|
||||
System Dependencies (Non-Conda)
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
If you are not using Conda, TVM requires several system libraries. On Ubuntu/Debian systems, install them with:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
sudo apt update
|
||||
sudo apt install zlib1g-dev libxml2-dev
|
||||
|
||||
For other operating systems, please refer to your package manager documentation.
|
||||
|
||||
The easiest way to manage dependency is via conda, which maintains a set of toolchains
|
||||
including LLVM across platforms. To create the environment of those build dependencies,
|
||||
one may simply use:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
# make sure to start with a fresh environment
|
||||
conda env remove -n tvm-build-venv
|
||||
# create the conda environment with build dependency
|
||||
conda create -n tvm-build-venv -c conda-forge \
|
||||
"llvmdev>=15" \
|
||||
"cmake>=3.24" \
|
||||
git \
|
||||
python=3.11
|
||||
# enter the build environment
|
||||
conda activate tvm-build-venv
|
||||
|
||||
.. note::
|
||||
**For Frontend Contributors (TFLite):** If you plan to run or contribute to the frontend tests (e.g., ``test_frontend_tflite.py``), you must install ``tensorflow==2.19.0``.
|
||||
|
||||
.. code:: bash
|
||||
|
||||
pip install tensorflow==2.19.0
|
||||
|
||||
Python 3.11 is supported.
|
||||
|
||||
Step 2. Get Source from GitHub
|
||||
------------------------------
|
||||
You can also choose to clone the source repo from GitHub.
|
||||
|
||||
.. code:: bash
|
||||
|
||||
git clone --recursive https://github.com/apache/tvm tvm
|
||||
|
||||
.. note::
|
||||
It's important to use the ``--recursive`` flag when cloning the TVM repository, which will
|
||||
automatically clone the submodules. If you forget to use this flag, you can manually clone the submodules
|
||||
by running ``git submodule update --init --recursive`` in the root directory of the TVM repository.
|
||||
|
||||
Step 3. Configure and Build
|
||||
---------------------------
|
||||
Create a build directory and run CMake to configure the build. The following example shows how to build
|
||||
|
||||
.. code:: bash
|
||||
|
||||
cd tvm
|
||||
rm -rf build && mkdir build && cd build
|
||||
# Specify the build configuration via CMake options
|
||||
cp ../cmake/config.cmake .
|
||||
|
||||
We want to specifically tweak the following flags by appending them to the end of the configuration file:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
# controls default compilation flags (Candidates: Release, Debug, RelWithDebInfo)
|
||||
echo "set(CMAKE_BUILD_TYPE RelWithDebInfo)" >> config.cmake
|
||||
|
||||
# LLVM is a must dependency for compiler end
|
||||
echo "set(USE_LLVM \"llvm-config --ignore-libllvm --link-static\")" >> config.cmake
|
||||
echo "set(HIDE_PRIVATE_SYMBOLS ON)" >> config.cmake
|
||||
|
||||
# GPU SDKs, turn on if needed
|
||||
echo "set(USE_CUDA OFF)" >> config.cmake
|
||||
echo "set(USE_METAL OFF)" >> config.cmake
|
||||
echo "set(USE_VULKAN OFF)" >> config.cmake
|
||||
echo "set(USE_OPENCL OFF)" >> config.cmake
|
||||
|
||||
# cuBLAS, cuDNN, cutlass support, turn on if needed
|
||||
echo "set(USE_CUBLAS OFF)" >> config.cmake
|
||||
echo "set(USE_CUDNN OFF)" >> config.cmake
|
||||
echo "set(USE_CUTLASS OFF)" >> config.cmake
|
||||
|
||||
|
||||
.. note::
|
||||
``HIDE_PRIVATE_SYMBOLS`` is a configuration option that enables the ``-fvisibility=hidden`` flag.
|
||||
This flag helps prevent potential symbol conflicts between TVM and PyTorch. These conflicts arise due to
|
||||
the frameworks shipping LLVMs of different versions.
|
||||
|
||||
`CMAKE_BUILD_TYPE <https://cmake.org/cmake/help/latest/variable/CMAKE_BUILD_TYPE.html>`_ controls default compilation flag:
|
||||
|
||||
- ``Debug`` sets ``-O0 -g``
|
||||
- ``RelWithDebInfo`` sets ``-O2 -g -DNDEBUG`` (recommended)
|
||||
- ``Release`` sets ``-O3 -DNDEBUG``
|
||||
|
||||
Once ``config.cmake`` is edited accordingly, kick off build with the commands below:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
cmake .. && cmake --build . --parallel $(nproc)
|
||||
|
||||
.. note::
|
||||
``nproc`` may not be available on all systems, please replace it with the number of cores on your system
|
||||
|
||||
A success build should produce ``libtvm`` and ``libtvm_runtime`` under ``build/`` directory.
|
||||
|
||||
Apache TVM relies on the tvm-ffi package to support its python bindings.
|
||||
Therefore, after we finish the build, we need to install the tvm-ffi package.
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
cd 3rdparty/tvm-ffi; pip install .; cd ..
|
||||
|
||||
|
||||
Leaving the build environment ``tvm-build-venv``, there are two ways to install the successful build into your environment:
|
||||
|
||||
- Install via environment variable
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
export TVM_HOME=/path-to-tvm
|
||||
pip install --target=$TVM_HOME/python $TVM_HOME/3rdparty/tvm-ffi
|
||||
export PYTHONPATH=$TVM_HOME/python:$PYTHONPATH
|
||||
|
||||
- Install via pip local project
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
conda activate your-own-env
|
||||
conda install python # make sure python is installed
|
||||
export TVM_LIBRARY_PATH=/path-to-tvm/build
|
||||
pip install -e /path-to-tvm
|
||||
|
||||
Step 4. Validate Installation
|
||||
-----------------------------
|
||||
|
||||
Using a compiler infrastructure with multiple language bindings could be error-prone.
|
||||
Therefore, it is highly recommended to validate Apache TVM installation before use.
|
||||
|
||||
**Step 1. Locate TVM Python package.** The following command can help confirm that TVM is properly installed as a python package and provide the location of the TVM python package:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
>>> python -c "import tvm; print(tvm.__file__)"
|
||||
/some-path/lib/python3.11/site-packages/tvm/__init__.py
|
||||
|
||||
**Step 2. Confirm which TVM library is used.** When maintaining multiple build or installation of TVM, it becomes important to double check if the python package is using the proper ``libtvm`` with the following command:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
>>> python -c "import tvm; print(tvm.base._LOADED_LIBS['tvm_runtime'])"
|
||||
<CDLL '/some-path/lib/python3.11/site-packages/tvm/libtvm_runtime.dylib', handle 95ada510 at 0x1030e4e50>
|
||||
|
||||
**Step 3. Reflect TVM build option.** Sometimes when downstream application fails, it could likely be some mistakes with a wrong TVM commit, or wrong build flags. To find it out, the following commands will be helpful:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
>>> python -c "import tvm; print('\n'.join(f'{k}: {v}' for k, v in tvm.support.libinfo().items()))"
|
||||
... # Omitted less relevant options
|
||||
GIT_COMMIT_HASH: 4f6289590252a1cf45a4dc37bce55a25043b8338
|
||||
HIDE_PRIVATE_SYMBOLS: ON
|
||||
USE_LLVM: llvm-config --link-static
|
||||
LLVM_VERSION: 15.0.7
|
||||
USE_VULKAN: OFF
|
||||
USE_CUDA: OFF
|
||||
CUDA_VERSION: NOT-FOUND
|
||||
USE_OPENCL: OFF
|
||||
USE_METAL: ON
|
||||
USE_ROCM: OFF
|
||||
|
||||
|
||||
**Step 4. Check device detection.** Sometimes it could be helpful to understand if TVM could detect your device at all with the following commands:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
>>> python -c "import tvm; print(tvm.metal().exist)"
|
||||
True # or False
|
||||
>>> python -c "import tvm; print(tvm.cuda().exist)"
|
||||
False # or True
|
||||
>>> python -c "import tvm; print(tvm.vulkan().exist)"
|
||||
False # or True
|
||||
|
||||
Please note that the commands above verify the presence of an actual device on the local machine for the TVM runtime (not the compiler) to execute properly. However, TVM compiler can perform compilation tasks without requiring a physical device. As long as the necessary toolchain, such as NVCC, is available, TVM supports cross-compilation even in the absence of an actual device.
|
||||
|
||||
|
||||
Step 5. Extra Python Dependencies
|
||||
---------------------------------
|
||||
Building from source does not ensure the installation of all necessary Python dependencies.
|
||||
|
||||
Python Dependencies
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
The following commands can be used to install the extra Python dependencies:
|
||||
|
||||
* Necessary dependencies:
|
||||
|
||||
.. code:: bash
|
||||
|
||||
pip3 install numpy cython
|
||||
|
||||
* If you want to use RPC Tracker
|
||||
|
||||
.. code:: bash
|
||||
|
||||
pip3 install tornado
|
||||
|
||||
* If you want to use auto-tuning module
|
||||
|
||||
.. code:: bash
|
||||
|
||||
pip3 install tornado psutil 'xgboost>=1.1.0' cloudpickle
|
||||
|
||||
Windows-Specific Build Notes
|
||||
----------------------------
|
||||
|
||||
If you're building TVM on Windows, note these platform-specific considerations:
|
||||
|
||||
Path Conventions
|
||||
~~~~~~~~~~~~~~~~
|
||||
- Use forward slashes (``/``) in Python/CMake paths, not Windows backslashes
|
||||
- Example: ``python cmake/config.cmake`` not ``python cmake\\config.cmake``
|
||||
|
||||
Advanced Build Configuration
|
||||
----------------------------
|
||||
|
||||
Ccache
|
||||
~~~~~~
|
||||
On supported platforms, the `Ccache compiler wrapper <https://ccache.dev/>`_ may be helpful for
|
||||
reducing TVM's build time, especially when building with `cutlass <https://github.com/NVIDIA/cutlass>`_.
|
||||
There are several ways to enable CCache in TVM builds:
|
||||
|
||||
- Leave ``USE_CCACHE=AUTO`` in ``build/config.cmake``. CCache will be used if it is found.
|
||||
|
||||
- Ccache's Masquerade mode. This is typically enabled during the Ccache installation process.
|
||||
To have TVM use Ccache in masquerade, simply specify the appropriate C/C++ compiler
|
||||
paths when configuring TVM's build system. For example:
|
||||
``cmake -DCMAKE_CXX_COMPILER=/usr/lib/ccache/c++ ...``.
|
||||
|
||||
- Ccache as CMake's C++ compiler prefix. When configuring TVM's build system,
|
||||
set the CMake variable ``CMAKE_CXX_COMPILER_LAUNCHER`` to an appropriate value.
|
||||
E.g. ``cmake -DCMAKE_CXX_COMPILER_LAUNCHER=ccache ...``.
|
||||
|
||||
|
||||
Building on Windows
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
TVM support build via MSVC using CMake. You will need to obtain a visual studio compiler.
|
||||
The minimum required VS version is **Visual Studio Enterprise 2019** (NOTE: we test
|
||||
against GitHub Actions' `windows-latest Runner <https://github.com/actions/runner-images>`_, so see that page for full details.
|
||||
We recommend following :ref:`install-dependencies` to obtain necessary dependencies and
|
||||
get an activated tvm-build environment. Then you can run the following command to build
|
||||
|
||||
.. code:: bash
|
||||
|
||||
mkdir build
|
||||
cd build
|
||||
cmake ..
|
||||
cd ..
|
||||
|
||||
The above command generates the solution file under the build directory.
|
||||
You can then run the following command to build
|
||||
|
||||
.. code:: bash
|
||||
|
||||
cmake --build build --config Release -- /m
|
||||
|
||||
CUDA Configuration
|
||||
..................
|
||||
For CUDA support on Windows:
|
||||
|
||||
.. code-block:: batch
|
||||
|
||||
set CUDA_PATH=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8
|
||||
set PATH=%CUDA_PATH%\bin;%PATH%
|
||||
cmake .. -DUSE_CUDA=ON
|
||||
|
||||
CMake & Compiler Setup
|
||||
......................
|
||||
- Specify generator: ``cmake -G "Visual Studio 16 2019" -A x64 ..``
|
||||
- Ensure Python is in PATH or specify: ``-DPython_EXECUTABLE=C:\Python311\python.exe``
|
||||
|
||||
|
||||
Building ROCm support
|
||||
~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Currently, ROCm is supported only on linux, so all the instructions are written with linux in mind.
|
||||
|
||||
- Set ``set(USE_ROCM ON)``, set ROCM_PATH to the correct path.
|
||||
- You need to first install HIP runtime from ROCm. Make sure the installation system has ROCm installed in it.
|
||||
- Install LLVM (>= 15), and LLD, make sure ``ld.lld`` is available via command line.
|
||||
|
||||
.. _install-from-source-cpp-tests:
|
||||
|
||||
Enable C++ Tests
|
||||
~~~~~~~~~~~~~~~~
|
||||
We use `Google Test <https://github.com/google/googletest>`_ to drive the C++
|
||||
tests in TVM. The easiest way to install GTest is from source.
|
||||
|
||||
.. code:: bash
|
||||
|
||||
git clone https://github.com/google/googletest
|
||||
cd googletest
|
||||
mkdir build
|
||||
cd build
|
||||
cmake -DBUILD_SHARED_LIBS=ON ..
|
||||
make
|
||||
sudo make install
|
||||
|
||||
After installing GTest, the C++ tests can be built and started with ``./tests/scripts/task_cpp_unittest.sh``, or built via CMake with ``-DUSE_GTEST=ON`` and then running ``./build/cpptest``.
|
||||
@@ -0,0 +1,41 @@
|
||||
.. 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.
|
||||
|
||||
.. _installation:
|
||||
|
||||
Installing TVM
|
||||
==============
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
pypi
|
||||
from_source
|
||||
docker
|
||||
|
||||
For most Python users, the quickest way to get started is to
|
||||
:ref:`install the Apache TVM wheel from PyPI <install-from-pypi>`.
|
||||
|
||||
Visit the :ref:`install TVM from source <install-from-source>` page to install
|
||||
TVM from the source code. Installing from source gives you the maximum
|
||||
flexibility to configure the build effectively from the official source
|
||||
releases. If you are interested in deploying to mobile or embedded devices,
|
||||
you do not need to install the entire TVM stack on your device. Instead, you
|
||||
only need the runtime.
|
||||
|
||||
If you would like to quickly try out TVM or run some demo and tutorials, you
|
||||
can :ref:`install from Docker <docker-images>`.
|
||||
@@ -0,0 +1,51 @@
|
||||
.. 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.
|
||||
|
||||
.. _install-from-pypi:
|
||||
|
||||
Install from PyPI
|
||||
=================
|
||||
|
||||
For most Python users, the quickest way to get started is to install the Apache
|
||||
TVM wheel from PyPI:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
pip install apache-tvm
|
||||
|
||||
This installs the Python package, including modules such as ``tvm.tirx``, and
|
||||
is suitable for trying tutorials that do not require a custom build.
|
||||
|
||||
CUDA environments
|
||||
-----------------
|
||||
|
||||
Some CUDA workflows use NVIDIA's Python CUDA bindings for runtime compilation.
|
||||
Install the CUDA extra in the same environment as TVM when you need this path:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
pip install "apache-tvm[cuda]"
|
||||
|
||||
This extra installs Python-side CUDA bindings only. It does not make the PyPI
|
||||
wheel a CUDA-enabled TVM build, and it does not install NVIDIA drivers or a CUDA
|
||||
toolkit. If you need CUDA support in TVM itself, build TVM from source with
|
||||
``USE_CUDA=ON``.
|
||||
|
||||
For more details on installing the TIRx compiler and optional kernel library,
|
||||
visit the :doc:`TIRx installation </tirx/install>` page. If you need to
|
||||
customize TVM's build configuration, visit the
|
||||
:ref:`install TVM from source <install-from-source>` page instead.
|
||||
@@ -0,0 +1,26 @@
|
||||
.. 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.
|
||||
|
||||
Other APIs
|
||||
==========
|
||||
|
||||
This page contains links to API references that are built with different doc
|
||||
build system.
|
||||
|
||||
* `C++ doxygen API <doxygen/index.html>`_
|
||||
* `Typescript typedoc API <typedoc/index.html>`_
|
||||
* `Java Javadoc API <javadoc/index.html>`_
|
||||
@@ -0,0 +1,22 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.arith
|
||||
---------
|
||||
.. automodule:: tvm.arith
|
||||
:members:
|
||||
:imported-members:
|
||||
@@ -0,0 +1,123 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.contrib
|
||||
-----------
|
||||
.. automodule:: tvm.contrib
|
||||
|
||||
tvm.contrib.cblas
|
||||
~~~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.cblas
|
||||
:members:
|
||||
|
||||
|
||||
tvm.contrib.coreml_runtime
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.coreml_runtime
|
||||
:members:
|
||||
|
||||
|
||||
tvm.contrib.cublas
|
||||
~~~~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.cublas
|
||||
:members:
|
||||
|
||||
|
||||
tvm.contrib.cublaslt
|
||||
~~~~~~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.cublaslt
|
||||
:members:
|
||||
|
||||
|
||||
tvm.contrib.cudnn
|
||||
~~~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.cudnn
|
||||
:members:
|
||||
|
||||
|
||||
tvm.contrib.dlpack
|
||||
~~~~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.dlpack
|
||||
:members:
|
||||
|
||||
|
||||
tvm.contrib.dnnl
|
||||
~~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.dnnl
|
||||
:members:
|
||||
|
||||
|
||||
tvm.contrib.download
|
||||
~~~~~~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.download
|
||||
:members:
|
||||
|
||||
|
||||
tvm.contrib.hipblas
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.hipblas
|
||||
:members:
|
||||
|
||||
|
||||
tvm.contrib.mkl
|
||||
~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.mkl
|
||||
:members:
|
||||
|
||||
|
||||
tvm.contrib.nnpack
|
||||
~~~~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.nnpack
|
||||
:members:
|
||||
|
||||
|
||||
tvm.contrib.pickle_memoize
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.pickle_memoize
|
||||
:members:
|
||||
|
||||
|
||||
tvm.contrib.random
|
||||
~~~~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.random
|
||||
:members:
|
||||
|
||||
|
||||
tvm.contrib.thrust
|
||||
~~~~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.thrust
|
||||
:members:
|
||||
|
||||
|
||||
tvm.contrib.tvmjs
|
||||
~~~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.tvmjs
|
||||
:members:
|
||||
|
||||
|
||||
tvm.contrib.cutlass
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.cutlass
|
||||
:members:
|
||||
:imported-members:
|
||||
|
||||
|
||||
tvm.contrib.hexagon
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
.. automodule:: tvm.contrib.hexagon
|
||||
:members:
|
||||
:imported-members:
|
||||
@@ -0,0 +1,22 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.driver
|
||||
----------
|
||||
.. automodule:: tvm.driver
|
||||
|
||||
.. autofunction:: tvm.compile
|
||||
@@ -0,0 +1,23 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.error
|
||||
---------
|
||||
.. automodule:: tvm.error
|
||||
:members:
|
||||
:imported-members:
|
||||
:autosummary:
|
||||
@@ -0,0 +1,22 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.exec
|
||||
--------
|
||||
.. automodule:: tvm.exec
|
||||
:members:
|
||||
:imported-members:
|
||||
@@ -0,0 +1,100 @@
|
||||
.. 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.
|
||||
|
||||
Python API
|
||||
==========
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: tvm
|
||||
|
||||
error
|
||||
ir
|
||||
arith
|
||||
instrument
|
||||
transform
|
||||
target
|
||||
driver
|
||||
testing
|
||||
exec
|
||||
support
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: tvm.runtime
|
||||
|
||||
runtime/runtime
|
||||
runtime/vm
|
||||
runtime/disco
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: tvm.relax
|
||||
|
||||
relax/relax
|
||||
relax/analysis
|
||||
relax/backend
|
||||
relax/block_builder
|
||||
relax/distributed
|
||||
relax/frontend
|
||||
relax/op
|
||||
relax/testing
|
||||
relax/training
|
||||
relax/transform
|
||||
relax/dpl
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: tvm.s_tir
|
||||
|
||||
s_tir/analysis
|
||||
s_tir/schedule
|
||||
s_tir/transform
|
||||
s_tir/dlight
|
||||
s_tir/backend
|
||||
s_tir/tensor_intrin
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: tvm.script
|
||||
|
||||
script/script
|
||||
script/parser
|
||||
script/ir_builder
|
||||
script/printer
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: tvm.te
|
||||
|
||||
te
|
||||
topi
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: tvm.s_tir.meta_schedule
|
||||
|
||||
meta_schedule
|
||||
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
:caption: Misc
|
||||
|
||||
rpc
|
||||
contrib
|
||||
@@ -0,0 +1,23 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.ir.instrument
|
||||
------------------
|
||||
.. automodule:: tvm.ir.instrument
|
||||
:members:
|
||||
:imported-members:
|
||||
:exclude-members: Path
|
||||
@@ -0,0 +1,24 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.ir
|
||||
------
|
||||
.. automodule:: tvm.ir
|
||||
:members:
|
||||
:imported-members:
|
||||
:exclude-members: Expr
|
||||
:autosummary:
|
||||
@@ -0,0 +1,146 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.s_tir.meta_schedule
|
||||
------------------------
|
||||
|
||||
tvm.s_tir.meta_schedule
|
||||
***********************
|
||||
.. automodule:: tvm.s_tir.meta_schedule
|
||||
:members:
|
||||
:imported-members:
|
||||
:autosummary:
|
||||
|
||||
tvm.s_tir.meta_schedule.arg_info
|
||||
********************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.arg_info
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.builder
|
||||
*******************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.builder
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.cost_model
|
||||
**********************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.cost_model
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.database
|
||||
********************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.database
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.feature_extractor
|
||||
*****************************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.feature_extractor
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.measure_callback
|
||||
****************************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.measure_callback
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.mutator
|
||||
*******************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.mutator
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.postproc
|
||||
********************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.postproc
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.runner
|
||||
******************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.runner
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.schedule_rule
|
||||
*************************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.schedule_rule
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.search_strategy
|
||||
***************************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.search_strategy
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.space_generator
|
||||
***************************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.space_generator
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.task_scheduler
|
||||
**************************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.task_scheduler
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.tir_integration
|
||||
***************************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.tir_integration
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.relax_integration
|
||||
*****************************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.relax_integration
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.trace_apply
|
||||
************************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.trace_apply
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.schedule
|
||||
********************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.schedule
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.meta_schedule.post_optimization
|
||||
*****************************************
|
||||
.. automodule:: tvm.s_tir.meta_schedule.post_optimization
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
@@ -0,0 +1,22 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.relax.analysis
|
||||
------------------
|
||||
.. automodule:: tvm.relax.analysis
|
||||
:members:
|
||||
:imported-members:
|
||||
@@ -0,0 +1,67 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.relax.backend
|
||||
-----------------
|
||||
|
||||
tvm.relax.backend
|
||||
*****************
|
||||
.. automodule:: tvm.relax.backend
|
||||
:members:
|
||||
:imported-members:
|
||||
|
||||
tvm.relax.backend.cuda
|
||||
**********************
|
||||
.. automodule:: tvm.relax.backend.cuda
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.backend.rocm
|
||||
**********************
|
||||
.. automodule:: tvm.relax.backend.rocm
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.backend.metal
|
||||
***********************
|
||||
.. automodule:: tvm.relax.backend.metal
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.backend.adreno
|
||||
************************
|
||||
.. automodule:: tvm.relax.backend.adreno
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.backend.gpu_generic
|
||||
*****************************
|
||||
.. automodule:: tvm.relax.backend.gpu_generic
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.backend.cpu_generic
|
||||
*****************************
|
||||
.. automodule:: tvm.relax.backend.cpu_generic
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.backend.contrib
|
||||
*************************
|
||||
.. automodule:: tvm.relax.backend.contrib
|
||||
:members:
|
||||
:no-index:
|
||||
@@ -0,0 +1,21 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.relax.block_builder
|
||||
-----------------------
|
||||
.. automodule:: tvm.relax.block_builder
|
||||
:members:
|
||||
@@ -0,0 +1,28 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.relax.distributed
|
||||
---------------------
|
||||
.. automodule:: tvm.relax.distributed
|
||||
:members:
|
||||
:imported-members:
|
||||
|
||||
tvm.relax.distributed.transform
|
||||
*******************************
|
||||
.. automodule:: tvm.relax.distributed.transform
|
||||
:members:
|
||||
:no-index:
|
||||
@@ -0,0 +1,21 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.relax.dpl
|
||||
-------------
|
||||
.. automodule:: tvm.relax.dpl
|
||||
:members:
|
||||
@@ -0,0 +1,60 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.relax.frontend
|
||||
------------------
|
||||
.. automodule:: tvm.relax.frontend
|
||||
:members:
|
||||
:imported-members:
|
||||
|
||||
tvm.relax.frontend.nn
|
||||
*********************
|
||||
.. automodule:: tvm.relax.frontend.nn
|
||||
:members:
|
||||
:imported-members:
|
||||
:exclude-members: BlockBuilder
|
||||
:noindex:
|
||||
|
||||
tvm.relax.frontend.nn.llm
|
||||
*************************
|
||||
.. automodule:: tvm.relax.frontend.nn.llm
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.frontend.onnx
|
||||
***********************
|
||||
.. automodule:: tvm.relax.frontend.onnx
|
||||
:members:
|
||||
:imported-members:
|
||||
|
||||
tvm.relax.frontend.stablehlo
|
||||
****************************
|
||||
.. automodule:: tvm.relax.frontend.stablehlo
|
||||
:members:
|
||||
:imported-members:
|
||||
|
||||
tvm.relax.frontend.tflite
|
||||
*************************
|
||||
.. automodule:: tvm.relax.frontend.tflite
|
||||
:members:
|
||||
:imported-members:
|
||||
|
||||
tvm.relax.frontend.torch
|
||||
************************
|
||||
.. automodule:: tvm.relax.frontend.torch
|
||||
:members:
|
||||
:imported-members:
|
||||
@@ -0,0 +1,95 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.relax.op
|
||||
------------
|
||||
|
||||
tvm.relax.op
|
||||
************
|
||||
.. automodule:: tvm.relax.op
|
||||
:members:
|
||||
:imported-members:
|
||||
:exclude-members: Expr
|
||||
|
||||
tvm.relax.op.nn
|
||||
***************
|
||||
.. automodule:: tvm.relax.op.nn
|
||||
:members:
|
||||
:imported-members:
|
||||
:exclude-members: Expr
|
||||
|
||||
tvm.relax.op.builtin
|
||||
********************
|
||||
.. automodule:: tvm.relax.op.builtin
|
||||
:members:
|
||||
:imported-members:
|
||||
:exclude-members: Expr
|
||||
|
||||
tvm.relax.op.ccl
|
||||
****************
|
||||
.. automodule:: tvm.relax.op.ccl
|
||||
:members:
|
||||
:imported-members:
|
||||
:exclude-members: Expr
|
||||
|
||||
tvm.relax.op.distributed
|
||||
************************
|
||||
.. automodule:: tvm.relax.op.distributed
|
||||
:members:
|
||||
:imported-members:
|
||||
:exclude-members: Expr
|
||||
|
||||
tvm.relax.op.grad
|
||||
*****************
|
||||
.. automodule:: tvm.relax.op.grad
|
||||
:members:
|
||||
:imported-members:
|
||||
:exclude-members: Expr
|
||||
|
||||
tvm.relax.op.image
|
||||
******************
|
||||
.. automodule:: tvm.relax.op.image
|
||||
:members:
|
||||
:imported-members:
|
||||
:exclude-members: Expr
|
||||
|
||||
tvm.relax.op.memory
|
||||
*******************
|
||||
.. automodule:: tvm.relax.op.memory
|
||||
:members:
|
||||
:imported-members:
|
||||
:exclude-members: Expr
|
||||
|
||||
tvm.relax.op.vision
|
||||
*******************
|
||||
.. automodule:: tvm.relax.op.vision
|
||||
:members:
|
||||
:imported-members:
|
||||
:exclude-members: Expr
|
||||
|
||||
tvm.relax.op.vm
|
||||
***************
|
||||
.. automodule:: tvm.relax.op.vm
|
||||
:members:
|
||||
:imported-members:
|
||||
:exclude-members: Expr
|
||||
|
||||
tvm.relax.op.op_attrs
|
||||
*********************
|
||||
.. automodule:: tvm.relax.op.op_attrs
|
||||
:members:
|
||||
:exclude-members: Attrs
|
||||
@@ -0,0 +1,23 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.relax
|
||||
---------
|
||||
.. automodule:: tvm.relax
|
||||
:members:
|
||||
:imported-members:
|
||||
:exclude-members: BlockBuilder, Call, Span, GlobalVar, SourceName, TupleType, Type, FuncType
|
||||
@@ -0,0 +1,22 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.relax.testing
|
||||
-----------------
|
||||
.. automodule:: tvm.relax.testing
|
||||
:members:
|
||||
:no-index:
|
||||
@@ -0,0 +1,22 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.relax.training
|
||||
------------------
|
||||
.. automodule:: tvm.relax.training
|
||||
:members:
|
||||
:imported-members:
|
||||
@@ -0,0 +1,131 @@
|
||||
.. 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.
|
||||
|
||||
.. _api-relax-transformation:
|
||||
|
||||
tvm.relax.transform
|
||||
-------------------
|
||||
.. automodule:: tvm.relax.transform
|
||||
:members:
|
||||
:imported-members:
|
||||
|
||||
tvm.relax.transform.legalize_ops
|
||||
*********************************
|
||||
.. automodule:: tvm.relax.transform.legalize_ops
|
||||
:members:
|
||||
|
||||
tvm.relax.transform.legalize_ops.binary
|
||||
========================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.binary
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.transform.legalize_ops.ccl
|
||||
=====================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.ccl
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.transform.legalize_ops.create
|
||||
========================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.create
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.transform.legalize_ops.datatype
|
||||
==========================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.datatype
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.transform.legalize_ops.distributed
|
||||
==============================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.distributed
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.transform.legalize_ops.grad
|
||||
======================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.grad
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.transform.legalize_ops.image
|
||||
=======================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.image
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.transform.legalize_ops.index
|
||||
=======================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.index
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.transform.legalize_ops.linear_algebra
|
||||
================================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.linear_algebra
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.transform.legalize_ops.manipulate
|
||||
=============================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.manipulate
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.transform.legalize_ops.nn
|
||||
====================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.nn
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.transform.legalize_ops.qdq
|
||||
=====================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.qdq
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.transform.legalize_ops.search
|
||||
========================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.search
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.transform.legalize_ops.statistical
|
||||
==============================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.statistical
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.transform.legalize_ops.unary
|
||||
=======================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.unary
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.transform.legalize_ops.vision
|
||||
========================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.vision
|
||||
:members:
|
||||
:no-index:
|
||||
|
||||
tvm.relax.transform.legalize_ops.adreno
|
||||
========================================
|
||||
.. automodule:: tvm.relax.transform.legalize_ops.adreno
|
||||
:members:
|
||||
:no-index:
|
||||
@@ -0,0 +1,23 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.rpc
|
||||
-------
|
||||
.. automodule:: tvm.rpc
|
||||
:members:
|
||||
:imported-members:
|
||||
:autosummary:
|
||||
@@ -0,0 +1,22 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.runtime.disco
|
||||
-----------------
|
||||
.. automodule:: tvm.runtime.disco
|
||||
:members:
|
||||
:imported-members:
|
||||
@@ -0,0 +1,21 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.runtime
|
||||
-----------
|
||||
.. automodule:: tvm.runtime
|
||||
:members:
|
||||
@@ -0,0 +1,21 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.runtime.vm
|
||||
--------------------
|
||||
.. automodule:: tvm.runtime.vm
|
||||
:members:
|
||||
@@ -0,0 +1,21 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.s_tir.analysis
|
||||
------------------
|
||||
.. automodule:: tvm.s_tir.analysis
|
||||
:members:
|
||||
@@ -0,0 +1,42 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.s_tir.backend
|
||||
-----------------
|
||||
|
||||
tvm.s_tir.backend
|
||||
*****************
|
||||
.. automodule:: tvm.s_tir.backend
|
||||
:members:
|
||||
:imported-members:
|
||||
|
||||
tvm.s_tir.backend.adreno
|
||||
************************
|
||||
.. automodule:: tvm.s_tir.backend.adreno
|
||||
:members:
|
||||
:imported-members:
|
||||
|
||||
tvm.s_tir.backend.adreno.pipeline
|
||||
*********************************
|
||||
.. automodule:: tvm.s_tir.backend.adreno.pipeline
|
||||
:members:
|
||||
|
||||
tvm.s_tir.backend.adreno.transform
|
||||
***********************************
|
||||
.. automodule:: tvm.s_tir.backend.adreno.transform
|
||||
:members:
|
||||
:imported-members:
|
||||
@@ -0,0 +1,66 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.s_tir.dlight
|
||||
----------------
|
||||
|
||||
tvm.s_tir.dlight
|
||||
****************
|
||||
.. automodule:: tvm.s_tir.dlight
|
||||
:members:
|
||||
:imported-members:
|
||||
|
||||
tvm.s_tir.dlight.gpu
|
||||
********************
|
||||
.. automodule:: tvm.s_tir.dlight.gpu
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.dlight.adreno
|
||||
***********************
|
||||
.. automodule:: tvm.s_tir.dlight.adreno
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.dlight.cpu
|
||||
********************
|
||||
.. automodule:: tvm.s_tir.dlight.cpu
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.dlight.analysis
|
||||
*************************
|
||||
.. automodule:: tvm.s_tir.dlight.analysis
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.dlight.base
|
||||
*********************
|
||||
.. automodule:: tvm.s_tir.dlight.base
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.dlight.benchmark
|
||||
**************************
|
||||
.. automodule:: tvm.s_tir.dlight.benchmark
|
||||
:members:
|
||||
:imported-members:
|
||||
:noindex:
|
||||
@@ -0,0 +1,22 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.s_tir.schedule
|
||||
------------------
|
||||
.. automodule:: tvm.s_tir.schedule
|
||||
:members:
|
||||
:imported-members:
|
||||
@@ -0,0 +1,73 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.s_tir.tensor_intrin
|
||||
-----------------------
|
||||
|
||||
tvm.s_tir.tensor_intrin
|
||||
***********************
|
||||
.. automodule:: tvm.s_tir.tensor_intrin
|
||||
:members:
|
||||
:imported-members:
|
||||
|
||||
tvm.s_tir.tensor_intrin.cuda
|
||||
****************************
|
||||
.. automodule:: tvm.s_tir.tensor_intrin.cuda
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.tensor_intrin.arm_cpu
|
||||
*******************************
|
||||
.. automodule:: tvm.s_tir.tensor_intrin.arm_cpu
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.tensor_intrin.x86
|
||||
****************************
|
||||
.. automodule:: tvm.s_tir.tensor_intrin.x86
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.tensor_intrin.rocm
|
||||
*****************************
|
||||
.. automodule:: tvm.s_tir.tensor_intrin.rocm
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.tensor_intrin.metal
|
||||
*****************************
|
||||
.. automodule:: tvm.s_tir.tensor_intrin.metal
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.tensor_intrin.hexagon
|
||||
*******************************
|
||||
.. automodule:: tvm.s_tir.tensor_intrin.hexagon
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.tensor_intrin.riscv_cpu
|
||||
*********************************
|
||||
.. automodule:: tvm.s_tir.tensor_intrin.riscv_cpu
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
tvm.s_tir.tensor_intrin.dot_product_common
|
||||
******************************************
|
||||
.. automodule:: tvm.s_tir.tensor_intrin.dot_product_common
|
||||
:members:
|
||||
:noindex:
|
||||
@@ -0,0 +1,23 @@
|
||||
.. Licensed to the Apache Software Foundation (ASF) under one
|
||||
or more contributor license agreements. See the NOTICE file
|
||||
distributed with this work for additional information
|
||||
regarding copyright ownership. The ASF licenses this file
|
||||
to you under the Apache License, Version 2.0 (the
|
||||
"License"); you may not use this file except in compliance
|
||||
with the License. You may obtain a copy of the License at
|
||||
|
||||
.. http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
.. Unless required by applicable law or agreed to in writing,
|
||||
software distributed under the License is distributed on an
|
||||
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations
|
||||
under the License.
|
||||
|
||||
tvm.s_tir.transform
|
||||
-------------------
|
||||
.. automodule:: tvm.s_tir.transform
|
||||
:members:
|
||||
:imported-members:
|
||||
:no-index:
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user