├── .gitignore ├── LICENSE ├── Makefile ├── README.md ├── ci ├── build_docs.sh ├── docker_bash.sh ├── docs_jenkinsfile.groovy └── update_site.py ├── docs ├── _static │ └── img │ │ ├── empty-thumb.png │ │ ├── overview.svg │ │ ├── tvm-logo-small.png │ │ └── tvm-logo-square.png ├── conf.py ├── deep_dive │ └── tensor_ir │ │ ├── abstraction.rst │ │ ├── index.rst │ │ ├── learning.rst │ │ └── tutorials │ │ ├── README.rst │ │ ├── creation.py │ │ └── transformation.py ├── get_started │ ├── install.rst │ ├── overview.md │ └── tutorials │ │ ├── README.rst │ │ ├── ir_module.py │ │ └── quick_start.py ├── index.rst └── reference │ ├── api │ ├── dlight.rst │ ├── error.rst │ ├── index.rst │ ├── instrument.rst │ ├── ir.rst │ ├── meta_schedule.rst │ ├── relax │ │ ├── analysis.rst │ │ ├── block_builder.rst │ │ ├── frontend.rst │ │ ├── op.rst │ │ ├── relax.rst │ │ └── transform.rst │ ├── rpc.rst │ ├── runtime │ │ ├── disco.rst │ │ ├── ndarray.rst │ │ ├── profiling.rst │ │ ├── relax_vm.rst │ │ └── runtime.rst │ ├── target.rst │ ├── te.rst │ ├── tir │ │ ├── analysis.rst │ │ ├── schedule.rst │ │ ├── stmt_functor.rst │ │ ├── tir.rst │ │ └── transform.rst │ ├── topi.rst │ └── transform.rst │ └── publications.rst ├── pylintrc └── requirements.txt /.gitignore: -------------------------------------------------------------------------------- 1 | # Sphinx documentation 2 | _build/ 3 | _staging/ 4 | 5 | # Vscode 6 | .vscode/ 7 | 8 | # Byte-compiled / optimized / DLL files 9 | __pycache__/ 10 | *.py[cod] 11 | *$py.class 12 | 13 | # C extensions 14 | *.so 15 | 16 | # Distribution / packaging 17 | .Python 18 | build/ 19 | develop-eggs/ 20 | dist/ 21 | downloads/ 22 | eggs/ 23 | .eggs/ 24 | lib/ 25 | lib64/ 26 | parts/ 27 | sdist/ 28 | var/ 29 | wheels/ 30 | share/python-wheels/ 31 | *.egg-info/ 32 | .installed.cfg 33 | *.egg 34 | MANIFEST 35 | 36 | # PyInstaller 37 | # Usually these files are written by a python script from a template 38 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 39 | *.manifest 40 | *.spec 41 | 42 | # Installer logs 43 | pip-log.txt 44 | pip-delete-this-directory.txt 45 | 46 | # Unit test / coverage reports 47 | htmlcov/ 48 | .tox/ 49 | .nox/ 50 | .coverage 51 | .coverage.* 52 | .cache 53 | nosetests.xml 54 | coverage.xml 55 | *.cover 56 | *.py,cover 57 | .hypothesis/ 58 | .pytest_cache/ 59 | cover/ 60 | 61 | # Translations 62 | *.mo 63 | *.pot 64 | 65 | # Django stuff: 66 | *.log 67 | local_settings.py 68 | db.sqlite3 69 | db.sqlite3-journal 70 | 71 | # Flask stuff: 72 | instance/ 73 | .webassets-cache 74 | 75 | # Scrapy stuff: 76 | .scrapy 77 | 78 | # Sphinx documentation 79 | docs/_build/ 80 | 81 | # PyBuilder 82 | .pybuilder/ 83 | target/ 84 | 85 | # Jupyter Notebook 86 | .ipynb_checkpoints 87 | 88 | # IPython 89 | profile_default/ 90 | ipython_config.py 91 | 92 | # pyenv 93 | # For a library or package, you might want to ignore these files since the code is 94 | # intended to run in multiple environments; otherwise, check them in: 95 | # .python-version 96 | 97 | # pipenv 98 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 99 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 100 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 101 | # install all needed dependencies. 102 | #Pipfile.lock 103 | 104 | # poetry 105 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 106 | # This is especially recommended for binary packages to ensure reproducibility, and is more 107 | # commonly ignored for libraries. 108 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 109 | #poetry.lock 110 | 111 | # pdm 112 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 113 | #pdm.lock 114 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 115 | # in version control. 116 | # https://pdm.fming.dev/#use-with-ide 117 | .pdm.toml 118 | 119 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 120 | __pypackages__/ 121 | 122 | # Celery stuff 123 | celerybeat-schedule 124 | celerybeat.pid 125 | 126 | # SageMath parsed files 127 | *.sage.py 128 | 129 | # Environments 130 | .env 131 | .venv 132 | env/ 133 | venv/ 134 | ENV/ 135 | env.bak/ 136 | venv.bak/ 137 | 138 | # Spyder project settings 139 | .spyderproject 140 | .spyproject 141 | 142 | # Rope project settings 143 | .ropeproject 144 | 145 | # mkdocs documentation 146 | /site 147 | 148 | # mypy 149 | .mypy_cache/ 150 | .dmypy.json 151 | dmypy.json 152 | 153 | # Pyre type checker 154 | .pyre/ 155 | 156 | # pytype static type analyzer 157 | .pytype/ 158 | 159 | # Cython debug symbols 160 | cython_debug/ 161 | 162 | # PyCharm 163 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 164 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 165 | # and can be added to the global gitignore or merged into this file. For a more nuclear 166 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 167 | #.idea/ 168 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line, and also 5 | # from the environment for the first two. 6 | SPHINXOPTS ?= 7 | SPHINXBUILD ?= python -m sphinx 8 | SOURCEDIR = docs 9 | BUILDDIR = _build 10 | STAGINGDIR = _staging 11 | 12 | .PHONY: help Makefile 13 | 14 | # Put it first so that "make" without argument is like "make help". 15 | help: 16 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 17 | 18 | # Catch-all target: route all unknown targets to Sphinx using the new 19 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 20 | # %: Makefile 21 | # @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 22 | 23 | clean: 24 | rm -rf $(BUILDDIR) 25 | rm -rf $(STAGINGDIR) 26 | 27 | staging: 28 | # Prepare the staging directory. Sphinx gallery automatically 29 | # writes new .rst files into the current directory. This can 30 | # cause issues when switching branches. By sequestering the 31 | # auto-generated files into the staging directory, they can be 32 | # removed without knowing the exact directory. 33 | 34 | mkdir -p $(STAGINGDIR) 35 | 36 | # Remove any symlinks that currently exist 37 | find $(STAGINGDIR) -type l -exec rm {} \; 38 | 39 | # Reproduce the directory structure 40 | find ${SOURCEDIR} \ 41 | -path ./$(BUILDDIR) -prune -o -path ./$(STAGINGDIR) -prune -o \ 42 | \( -name "*.rst" -o -name "*.md" \) \ 43 | -printf "$(STAGINGDIR)/%h\n" \ 44 | | sort | uniq | xargs mkdir -p 45 | 46 | # Symlink all .rst/.md files into the staging directory 47 | find ${SOURCEDIR} \ 48 | -path ./$(BUILDDIR) -prune -o -path ./$(STAGINGDIR) -prune -o \ 49 | \( -name "*.rst" -o -name "*.md" \) \ 50 | -exec ln -s $(PWD)/{} $(STAGINGDIR)/{} \; 51 | 52 | ln -s $(PWD)/$(SOURCEDIR)/conf.py $(STAGINGDIR)/${SOURCEDIR}/conf.py 53 | ln -s $(PWD)/$(SOURCEDIR)/_static $(STAGINGDIR)/${SOURCEDIR}/_static 54 | 55 | html: staging 56 | $(SPHINXBUILD) -M html $(STAGINGDIR)/docs $(BUILDDIR) $(SPHINXOPTS) 57 | @echo 58 | @echo "Build finished. The HTML pages are in $(BUILDDIR)." 59 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # TVM Unity Documentation 2 | 3 | The documentation was built upon [Sphinx](https://www.sphinx-doc.org/en/master/). 4 | 5 | ## Dependencies 6 | 7 | Run the following command in this directory to install dependencies first: 8 | 9 | ```bash 10 | pip install -r requirements.txt 11 | ``` 12 | 13 | ## Build the Documentation 14 | 15 | Then you can build the documentation by running: 16 | 17 | ```bash 18 | make html 19 | ``` 20 | 21 | ## View the Documentation 22 | 23 | Run the following command to start a simple HTTP server: 24 | 25 | ```bash 26 | python -m http.server -d _build/html 8000 27 | ``` 28 | 29 | Then you can view the documentation in your browser at `http://localhost:8000` (the port can be customized by changing the port number). -------------------------------------------------------------------------------- /ci/build_docs.sh: -------------------------------------------------------------------------------- 1 | pip install -r ./requirements.txt 2 | make clean && make html -------------------------------------------------------------------------------- /ci/docker_bash.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # Licensed to the Apache Software Foundation (ASF) under one 4 | # or more contributor license agreements. See the NOTICE file 5 | # distributed with this work for additional information 6 | # regarding copyright ownership. The ASF licenses this file 7 | # to you under the Apache License, Version 2.0 (the 8 | # "License"); you may not use this file except in compliance 9 | # with the License. You may obtain a copy of the License at 10 | # 11 | # http://www.apache.org/licenses/LICENSE-2.0 12 | # 13 | # Unless required by applicable law or agreed to in writing, 14 | # software distributed under the License is distributed on an 15 | # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 16 | # KIND, either express or implied. See the License for the 17 | # specific language governing permissions and limitations 18 | # under the License. 19 | 20 | # 21 | # Start a bash, mount REPO_MOUNT_POINT to be current directory. 22 | # 23 | # Usage: docker_bash.sh [-i|--interactive] [--net=host] [-t|--tty] 24 | # [--mount MOUNT_DIR] [--repo-mount-point REPO_MOUNT_POINT] 25 | # [--dry-run] [--name NAME] [--privileged] 26 | # [--] [COMMAND] 27 | # 28 | # Usage: docker_bash.sh 29 | # Starts an interactive session 30 | # 31 | # Usage2: docker_bash.sh [-i] [COMMAND] 32 | # Execute command in the docker image, default non-interactive 33 | # With -i, execute interactively. 34 | # 35 | 36 | set -euo pipefail 37 | 38 | function show_usage() { 39 | cat < [--] [COMMAND] 44 | 45 | -h, --help 46 | 47 | Display this help message. 48 | 49 | -i, --interactive 50 | 51 | Start the docker session in interactive mode. 52 | 53 | -t, --tty 54 | 55 | Start the docker session with a pseudo terminal (tty). 56 | 57 | --net=host 58 | 59 | Expose servers run into the container to the host, passing the 60 | "--net=host" argument through to docker. On MacOS, this is 61 | instead passed as "-p 8888:8888" since the host networking driver 62 | isn't supported. 63 | 64 | --mount MOUNT_DIR 65 | 66 | Expose MOUNT_DIR as an additional mount point inside the docker 67 | container. The mount point inside the container is the same as 68 | the folder location outside the container. This option can be 69 | specified multiple times. 70 | 71 | --repo-mount-point REPO_MOUNT_POINT 72 | 73 | The directory inside the docker container at which the TVM 74 | repository should be mounted, and is used as the workspace inside 75 | the docker container. 76 | 77 | If unspecified, the mount location depends on the environment. If 78 | running inside Jenkins, the mount location will be /workspace. 79 | Otherwise, the mount location of the repository will be the same 80 | as the external location of the repository, to maintain 81 | compatibility with git-worktree. 82 | 83 | --no-gpu 84 | 85 | Do not use GPU device drivers even if using an CUDA Docker image 86 | 87 | --dry-run 88 | 89 | Print the docker command to be run, but do not execute it. 90 | 91 | --env 92 | 93 | Pass an environment variable through to the container. 94 | 95 | --name 96 | 97 | Set the name of the docker container, and the hostname that will 98 | appear inside the container. 99 | 100 | --privileged 101 | 102 | Give extended privileges to this container. 103 | 104 | DOCKER_IMAGE_NAME 105 | 106 | The name of the docker container to be run. This can be an 107 | explicit name of a docker image (e.g. "tlcpack/ci-gpu:v0.76") or 108 | can be a shortcut as defined in the TVM Jenkinsfile 109 | (e.g. "ci_gpu"). 110 | 111 | COMMAND 112 | 113 | The command to be run inside the docker container. If this is set 114 | to "bash", the --interactive, --tty and --net=host flags are set. 115 | If no command is specified, defaults to "bash". If the command 116 | contains dash-prefixed arguments, the command should be preceded 117 | by -- to indicate arguments that are not intended for bash.sh. 118 | 119 | EOF 120 | } 121 | 122 | 123 | ################################# 124 | ### Start of argument parsing ### 125 | ################################# 126 | 127 | SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)" 128 | REPO_DIR="$(dirname "${SCRIPT_DIR}")" 129 | 130 | DRY_RUN=false 131 | INTERACTIVE=false 132 | TTY=false 133 | USE_NET_HOST=false 134 | USE_GPU=true 135 | DOCKER_IMAGE_NAME= 136 | COMMAND=bash 137 | MOUNT_DIRS=( ) 138 | CONTAINER_NAME= 139 | 140 | # TODO(Lunderberg): Remove this if statement and always set to 141 | # "${REPO_DIR}". The consistent directory for Jenkins is currently 142 | # necessary to allow cmake build commands to run in CI after the build 143 | # steps. 144 | # TODO(https://github.com/apache/tvm/issues/11952): 145 | # Figure out a better way to keep the same path 146 | # between build and testing stages. 147 | if [[ -n "${JENKINS_HOME:-}" ]]; then 148 | REPO_MOUNT_POINT=/workspace 149 | else 150 | REPO_MOUNT_POINT="${REPO_DIR}" 151 | fi 152 | 153 | 154 | function parse_error() { 155 | echo "$@" >&2 156 | show_usage >&2 157 | exit 1 158 | } 159 | 160 | # Handle joined flags, such as interpreting -ih as -i -h. Either rewrites 161 | # the current argument if it is a joined argument, or shifts all arguments 162 | # otherwise. Should be called as "eval $break_joined_flag" where joined 163 | # flags are possible. Can't use a function definition, because it needs 164 | # to overwrite the parent scope's behavior. 165 | break_joined_flag='if (( ${#1} == 2 )); then shift; else set -- -"${1#-i}" "${@:2}"; fi' 166 | 167 | DOCKER_ENV=( ) 168 | DOCKER_FLAGS=( ) 169 | 170 | while (( $# )); do 171 | case "$1" in 172 | -h|--help) 173 | show_usage 174 | exit 0 175 | ;; 176 | 177 | -i*|--interactive) 178 | INTERACTIVE=true 179 | eval $break_joined_flag 180 | ;; 181 | 182 | -t*|--tty) 183 | TTY=true 184 | eval $break_joined_flag 185 | ;; 186 | 187 | --net=host) 188 | USE_NET_HOST=true 189 | shift 190 | ;; 191 | 192 | --net) 193 | DOCKER_FLAGS+=( --net "$2" ) 194 | shift 2 195 | ;; 196 | 197 | --mount) 198 | if [[ -n "$2" ]]; then 199 | MOUNT_DIRS+=("$2") 200 | shift 2 201 | else 202 | parse_error 'ERROR: --mount requires a non-empty argument' 203 | fi 204 | ;; 205 | 206 | --mount=?*) 207 | MOUNT_DIRS+=("${1#*=}") 208 | shift 209 | ;; 210 | 211 | --name) 212 | if [[ -n "$2" ]]; then 213 | CONTAINER_NAME="$2" 214 | shift 2 215 | else 216 | parse_error 'ERROR: --name requires a non empty argument' 217 | fi 218 | ;; 219 | 220 | --privileged) 221 | DOCKER_FLAGS+=( "--privileged" ) 222 | shift 1 223 | ;; 224 | 225 | --env) 226 | DOCKER_ENV+=( --env "$2" ) 227 | shift 2 228 | ;; 229 | 230 | --volume) 231 | DOCKER_FLAGS+=( --volume "$2" ) 232 | shift 2 233 | ;; 234 | 235 | --dry-run) 236 | DRY_RUN=true 237 | shift 238 | ;; 239 | 240 | --no-gpu) 241 | USE_GPU=false 242 | shift 243 | ;; 244 | 245 | --repo-mount-point) 246 | if [[ -n "$2" ]]; then 247 | REPO_MOUNT_POINT="$2" 248 | shift 2 249 | else 250 | parse_error 'ERROR: --repo-mount-point requires a non-empty argument' 251 | fi 252 | ;; 253 | 254 | --repo-mount-point=?*) 255 | REPO_MOUNT_POINT="${1#*=}" 256 | shift 257 | ;; 258 | 259 | --) 260 | shift 261 | COMMAND=( "$@" ) 262 | break 263 | ;; 264 | 265 | -*|--*) 266 | echo "Error: Unknown flag: $1" >&2 267 | echo " If this flag is intended to be passed to the" >&2 268 | echo " docker command, please add -- before the docker" >&2 269 | echo " command (e.g. docker_bash.sh ci_gpu -- build -j2)" >&2 270 | show_usage >&2 271 | exit 1 272 | ;; 273 | 274 | *) 275 | # First positional argument is the image name, all 276 | # remaining below to the COMMAND. 277 | if [[ -z "${DOCKER_IMAGE_NAME}" ]]; then 278 | DOCKER_IMAGE_NAME=$1 279 | shift 280 | else 281 | COMMAND=( "$@" ) 282 | break 283 | fi 284 | ;; 285 | esac 286 | done 287 | 288 | if [[ -z "${DOCKER_IMAGE_NAME}" ]]; then 289 | echo "Error: Missing DOCKER_IMAGE_NAME" >&2 290 | show_usage >&2 291 | fi 292 | 293 | if [[ ${COMMAND[@]+"${COMMAND[@]}"} = bash ]]; then 294 | INTERACTIVE=true 295 | TTY=true 296 | USE_NET_HOST=true 297 | fi 298 | 299 | 300 | 301 | ############################### 302 | ### End of argument parsing ### 303 | ############################### 304 | 305 | DOCKER_MOUNT=( ) 306 | DOCKER_DEVICES=( ) 307 | 308 | # Set up working directories 309 | 310 | DOCKER_FLAGS+=( --workdir "${REPO_MOUNT_POINT}" ) 311 | DOCKER_MOUNT+=( --volume "${REPO_DIR}":"${REPO_MOUNT_POINT}" 312 | --volume "${SCRIPT_DIR}":/docker 313 | ) 314 | 315 | # Set up CI-specific environment variables 316 | DOCKER_ENV+=( --env CI_BUILD_HOME="${REPO_MOUNT_POINT}" 317 | --env CI_BUILD_USER="$(id -u -n)" 318 | --env CI_BUILD_UID="$(id -u)" 319 | --env CI_BUILD_GROUP="$(id -g -n)" 320 | --env CI_BUILD_GID="$(id -g)" 321 | --env CI_PYTEST_ADD_OPTIONS="${CI_PYTEST_ADD_OPTIONS:-}" 322 | --env CI_IMAGE_NAME="${DOCKER_IMAGE_NAME}" 323 | ) 324 | 325 | # Remove the container once it finishes running (--rm). 326 | DOCKER_FLAGS+=(--rm) 327 | 328 | 329 | # Expose services running in container to the host. 330 | if $USE_NET_HOST; then 331 | if [[ $(uname) == "Darwin" ]]; then 332 | # Docker's host networking driver isn't supported on macOS. 333 | # Use default bridge network and expose port for jupyter notebook. 334 | DOCKER_FLAGS+=( -p 8888:8888 ) 335 | else 336 | DOCKER_FLAGS+=(--net=host) 337 | fi 338 | fi 339 | 340 | # Set up interactive sessions 341 | if ${INTERACTIVE}; then 342 | DOCKER_FLAGS+=( --interactive ) 343 | fi 344 | 345 | if ${TTY}; then 346 | DOCKER_FLAGS+=( --tty ) 347 | fi 348 | 349 | # Setup the docker name and the hostname inside the container 350 | if [[ ! -z "${CONTAINER_NAME}" ]]; then 351 | DOCKER_FLAGS+=( --name ${CONTAINER_NAME} --hostname ${CONTAINER_NAME}) 352 | fi 353 | 354 | # Expose external directories to the docker container 355 | for MOUNT_DIR in ${MOUNT_DIRS[@]+"${MOUNT_DIRS[@]}"}; do 356 | DOCKER_MOUNT+=( --volume "${MOUNT_DIR}:${MOUNT_DIR}" ) 357 | done 358 | 359 | # Use nvidia-docker for GPU container. If nvidia-docker is not 360 | # available, fall back to using "--gpus all" flag, requires docker 361 | # version 19.03 or higher. 362 | if [[ "$USE_GPU" == "true" ]] && [[ "${DOCKER_IMAGE_NAME}" == *"gpu"* || "${DOCKER_IMAGE_NAME}" == *"cuda"* ]]; then 363 | if type nvidia-docker 1> /dev/null 2> /dev/null; then 364 | DOCKER_BINARY=nvidia-docker 365 | else 366 | DOCKER_BINARY=docker 367 | DOCKER_FLAGS+=( --gpus all ) 368 | fi 369 | 370 | # nvidia-docker treats Vulkan as a graphics API, so we need to 371 | # request passthrough of graphics APIs. This could also be set in 372 | # the Dockerfile. 373 | DOCKER_ENV+=( --env NVIDIA_DRIVER_CAPABILITIES=compute,graphics,utility ) 374 | 375 | # But as of nvidia-docker version 2.6.0-1, we still need to pass 376 | # through the nvidia icd files ourselves. 377 | ICD_SEARCH_LOCATIONS=( 378 | # https://github.com/KhronosGroup/Vulkan-Loader/blob/master/loader/LoaderAndLayerInterface.md#icd-discovery-on-linux 379 | /usr/local/etc/vulkan/icd.d 380 | /usr/local/share/vulkan/icd.d 381 | /etc/vulkan/icd.d 382 | /usr/share/vulkan/icd.d 383 | # https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md#icd-installation 384 | /etc/glvnd/egl_vendor.d 385 | /usr/share/glvnd/egl_vendor.d 386 | ) 387 | for filename in $(find "${ICD_SEARCH_LOCATIONS[@]}" -name "*nvidia*.json" 2> /dev/null); do 388 | DOCKER_MOUNT+=( --volume "${filename}":"${filename}":ro ) 389 | done 390 | 391 | else 392 | DOCKER_BINARY=docker 393 | fi 394 | 395 | # Pass any restrictions of allowed CUDA devices from the host to the 396 | # docker container. 397 | if [[ -n ${CUDA_VISIBLE_DEVICES:-} ]]; then 398 | DOCKER_ENV+=( --env CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" ) 399 | fi 400 | 401 | 402 | # Set TVM import path inside the docker image 403 | if [[ "${DOCKER_IMAGE_NAME}" == *"ci"* ]]; then 404 | DOCKER_ENV+=( --env PYTHONPATH="${REPO_MOUNT_POINT}"/python ) 405 | fi 406 | 407 | 408 | 409 | # If the Vitis-AI docker image is selected, expose the Xilinx FPGA 410 | # devices and required volumes containing e.g. DSA's and overlays 411 | if [[ "${DOCKER_IMAGE_NAME}" == *"demo_vitis_ai"* && -d "/dev/shm" && -d "/opt/xilinx/dsa" && -d "/opt/xilinx/overlaybins" ]]; then 412 | DOCKER_MOUNT+=( --volume /dev/shm:/dev/shm 413 | --volume /opt/xilinx/dsa:/opt/xilinx/dsa 414 | --volume /opt/xilinx/overlaybins:/opt/xilinx/overlaybins 415 | ) 416 | 417 | XCLMGMT_DRIVER="$(find /dev -name xclmgmt\*)" 418 | for DRIVER in "${XCLMGMT_DRIVER}"; do 419 | DOCKER_DEVICES+=( --device="${DRIVER}" ) 420 | done 421 | 422 | RENDER_DRIVER="$(find /dev/dri -name renderD\*)" 423 | for DRIVER in "${RENDER_DRIVER}"; do 424 | DOCKER_DEVICES+=( --device="${DRIVER}" ) 425 | done 426 | fi 427 | 428 | # Add ROCm devices and set ROCM_ENABLED=1 which is used in the with_the_same_user script 429 | # to add the user to the video group 430 | if [[ "${DOCKER_IMAGE_NAME}" == *"rocm"* && -d "/dev/dri" ]]; then 431 | DOCKER_DEVICES+=( --device=/dev/kfd --device=/dev/dri ) 432 | DOCKER_ENV+=( --env ROCM_ENABLED=1 ) 433 | fi 434 | 435 | # When running from a git worktree, also mount the original git dir. 436 | if [ -f "${REPO_DIR}/.git" ]; then 437 | git_dir=$(cd ${REPO_DIR} && git rev-parse --git-common-dir) 438 | if [ "${git_dir}" != "${REPO_DIR}/.git" ]; then 439 | DOCKER_MOUNT+=( --volume "${git_dir}:${git_dir}" ) 440 | fi 441 | fi 442 | 443 | # Print arguments. 444 | echo "REPO_DIR: ${REPO_DIR}" 445 | echo "DOCKER CONTAINER NAME: ${DOCKER_IMAGE_NAME}" 446 | echo "" 447 | 448 | echo Running \'${COMMAND[@]+"${COMMAND[@]}"}\' inside ${DOCKER_IMAGE_NAME}... 449 | 450 | DOCKER_CMD=(${DOCKER_BINARY} run 451 | ${DOCKER_FLAGS[@]+"${DOCKER_FLAGS[@]}"} 452 | ${DOCKER_ENV[@]+"${DOCKER_ENV[@]}"} 453 | ${DOCKER_MOUNT[@]+"${DOCKER_MOUNT[@]}"} 454 | ${DOCKER_DEVICES[@]+"${DOCKER_DEVICES[@]}"} 455 | "${DOCKER_IMAGE_NAME}" 456 | ${COMMAND[@]+"${COMMAND[@]}"} 457 | ) 458 | 459 | if ${DRY_RUN}; then 460 | echo ${DOCKER_CMD[@]+"${DOCKER_CMD[@]}"} 461 | else 462 | ${DOCKER_CMD[@]+"${DOCKER_CMD[@]}"} 463 | fi 464 | -------------------------------------------------------------------------------- /ci/docs_jenkinsfile.groovy: -------------------------------------------------------------------------------- 1 | #!groovy 2 | // -*- mode: groovy -*- 3 | 4 | // Licensed to the Apache Software Foundation (ASF) under one 5 | // or more contributor license agreements. See the NOTICE file 6 | // distributed with this work for additional information 7 | // regarding copyright ownership. The ASF licenses this file 8 | // to you under the Apache License, Version 2.0 (the 9 | // "License"); you may not use this file except in compliance 10 | // with the License. You may obtain a copy of the License at 11 | // 12 | // http://www.apache.org/licenses/LICENSE-2.0 13 | // 14 | // Unless required by applicable law or agreed to in writing, 15 | // software distributed under the License is distributed on an 16 | // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 17 | // KIND, either express or implied. See the License for the 18 | // specific language governing permissions and limitations 19 | // under the License. 20 | 21 | // Jenkins pipeline 22 | // See documents at https://jenkins.io/doc/book/pipeline/jenkinsfile/ 23 | 24 | import org.jenkinsci.plugins.pipeline.modeldefinition.Utils 25 | 26 | // NOTE: these lines are scanned by docker/dev_common.sh. Please update the regex as needed. --> 27 | ci_gpu = 'tlcpackstaging/ci_gpu:20240428-060115-0b09ed018' 28 | // <--- End of regex-scanned config. 29 | 30 | // Parameters to allow overriding (in Jenkins UI), the images 31 | // to be used by a given build. When provided, they take precedence 32 | // over default values above. 33 | properties([ 34 | parameters([ 35 | string(name: 'ci_gpu_param', defaultValue: ''), 36 | ]) 37 | ]) 38 | 39 | // command to start a docker container 40 | docker_run = 'ci/docker_bash.sh' 41 | // timeout in minutes 42 | max_time = 30 43 | 44 | def per_exec_ws(folder) { 45 | return "workspace/exec_${env.EXECUTOR_NUMBER}/" + folder 46 | } 47 | 48 | // initialize source codes 49 | def init_git() { 50 | checkout scm 51 | // Add more info about job node 52 | sh ( 53 | script: "echo NODE_NAME=${env.NODE_NAME}", 54 | label: 'Show executor node info', 55 | ) 56 | retry(5) { 57 | timeout(time: 2, unit: 'MINUTES') { 58 | sh (script: 'git submodule update --init --recursive -f', label: 'Update git submodules') 59 | } 60 | } 61 | } 62 | 63 | stage('Prepare') { 64 | node('CPU-SMALL') { 65 | // When something is provided in ci_*_param, use it, otherwise default with ci_* 66 | ci_gpu = params.ci_gpu_param ?: ci_gpu 67 | 68 | sh (script: """ 69 | echo "Docker images being used in this build:" 70 | echo " ci_gpu = ${ci_gpu}" 71 | """, label: 'Docker image names') 72 | } 73 | } 74 | 75 | stage('Build') { 76 | timeout(time: max_time, unit: 'MINUTES') { 77 | node('GPU') { 78 | ws(per_exec_ws('mlc-docs/build')) { 79 | init_git() 80 | sh (script: "${docker_run} ${ci_gpu} nvidia-smi", label: 'Check GPU info') 81 | sh (script: "${docker_run} ${ci_gpu} ./ci/build_docs.sh", label: 'Build docs') 82 | if (env.BRANCH_NAME == 'main') { 83 | withCredentials([string( 84 | credentialsId: 'MLC_ACCESS_TOKEN', 85 | variable: 'GITHUB_TOKEN', 86 | )]) { 87 | sh (script: 'if [ ! -d docs-site ]; then git clone https://$GITHUB_TOKEN@github.com/mlc-ai/docs docs-site; fi') 88 | sh (script: 'python ci/update_site.py --site-path docs-site --source-path _build/html', label: 'Depoly') 89 | } 90 | } 91 | } 92 | } 93 | } 94 | } 95 | -------------------------------------------------------------------------------- /ci/update_site.py: -------------------------------------------------------------------------------- 1 | """Update the site by importing content from the folder.""" 2 | import argparse 3 | import logging 4 | import os 5 | import subprocess 6 | from datetime import datetime 7 | 8 | 9 | def py_str(cstr): 10 | return cstr.decode("utf-8") 11 | 12 | 13 | # list of files to skip 14 | skip_list = { 15 | ".gitignore", 16 | ".nojekyll", 17 | "CNAME", 18 | "README.md", 19 | "LICENSE", 20 | } 21 | 22 | 23 | def main(): 24 | logging.basicConfig(level=logging.WARNING) 25 | parser = argparse.ArgumentParser(description="Deploy a built html to the root.") 26 | parser.add_argument("--dry-run", action="store_true") 27 | parser.add_argument("--site-path", type=str) 28 | parser.add_argument("--source-path", type=str, default="_build/html") 29 | 30 | args = parser.parse_args() 31 | 32 | def run_cmd(cmd): 33 | proc = subprocess.Popen( 34 | cmd, cwd=args.site_path, stdout=subprocess.PIPE, stderr=subprocess.STDOUT 35 | ) 36 | (out, _) = proc.communicate() 37 | if proc.returncode != 0: 38 | msg = f"cmd error: {cmd}" 39 | msg += py_str(out) 40 | raise RuntimeError(msg) 41 | return py_str(out) 42 | 43 | run_cmd(["git", "config", "user.name", "mlc-bot"]) 44 | run_cmd( 45 | ["git", "config", "user.email", "106439794+mlc-bot@users.noreply.github.com"] 46 | ) 47 | run_cmd(["git", "fetch"]) 48 | run_cmd(["git", "checkout", "-B", "gh-pages", "origin/gh-pages"]) 49 | files = run_cmd(["git", "ls-files"]) 50 | skip_set = set(skip_list) 51 | 52 | for fname in files.split("\n"): 53 | fname = fname.strip() 54 | if fname and fname not in skip_set: 55 | if not args.dry_run: 56 | run_cmd(["rm", "-rf", fname]) 57 | print(f"Remove {fname}") 58 | 59 | os.system(f"cp -rf {args.source_path}/* {args.site_path}") 60 | run_cmd(["git", "add", "--all"]) 61 | 62 | if not args.dry_run: 63 | try: 64 | run_cmd(["git", "commit", "-am", f" Update at {datetime.now()}"]) 65 | except RuntimeError: 66 | pass 67 | run_cmd(["git", "push", "origin", "gh-pages"]) 68 | print("Finish updating and push to origin/gh-pages ...") 69 | 70 | 71 | if __name__ == "__main__": 72 | main() 73 | -------------------------------------------------------------------------------- /docs/_static/img/empty-thumb.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mlc-ai/docs/8919ed0e828a13644f87584c5e3ae9611d6f4739/docs/_static/img/empty-thumb.png -------------------------------------------------------------------------------- /docs/_static/img/tvm-logo-small.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mlc-ai/docs/8919ed0e828a13644f87584c5e3ae9611d6f4739/docs/_static/img/tvm-logo-small.png -------------------------------------------------------------------------------- /docs/_static/img/tvm-logo-square.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mlc-ai/docs/8919ed0e828a13644f87584c5e3ae9611d6f4739/docs/_static/img/tvm-logo-square.png -------------------------------------------------------------------------------- /docs/conf.py: -------------------------------------------------------------------------------- 1 | # pylint: disable=invalid-name, exec-used, redefined-builtin 2 | # -*- coding: utf-8 -*- 3 | import inspect 4 | import os 5 | from functools import partial 6 | from hashlib import md5 7 | from importlib import import_module 8 | from pathlib import Path 9 | import re 10 | from typing import List 11 | from unittest.mock import patch 12 | 13 | import sphinx_rtd_theme 14 | import tlcpack_sphinx_addon 15 | import tvm 16 | 17 | # -- General configuration ------------------------------------------------ 18 | 19 | # General information about the project. 20 | 21 | project = "Apache TVM Unity" 22 | author = "Apache Software Foundation" 23 | copyright = f"2020 - 2023, {author}" 24 | github_doc_root = "https://github.com/mlc-ai/docs/" 25 | 26 | # Version information. 27 | curr_path = Path(__file__).expanduser().absolute().parent 28 | # Can't use curr_path.parent, because sphinx_gallery requires a relative path. 29 | home_path = Path(os.pardir, os.pardir) if "_staging" in str(curr_path) else Path(os.pardir) 30 | home_path = os.path.join(home_path, "docs") 31 | 32 | version = tvm.__version__ 33 | release = version 34 | 35 | extensions = [ 36 | "myst_parser", 37 | "sphinx_gallery.gen_gallery", 38 | "sphinx_tabs.tabs", 39 | "sphinx_toolbox.collapse", 40 | "sphinxcontrib.httpdomain", 41 | "sphinx.ext.autodoc", 42 | "sphinx.ext.autosummary", 43 | "sphinx.ext.napoleon", 44 | ] 45 | 46 | source_suffix = { 47 | '.rst': 'restructuredtext', 48 | '.md': 'markdown', 49 | } 50 | 51 | language = "en" 52 | 53 | exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] 54 | 55 | # This line is used for bypass the issue of sidebar toc 56 | exclude_patterns.append("**/tutorials/index.rst") 57 | exclude_patterns.append("**/tutorials/README.rst") 58 | 59 | # The name of the Pygments (syntax highlighting) style to use. 60 | pygments_style = "sphinx" 61 | 62 | # A list of ignored prefixes for module index sorting. 63 | # If true, `todo` and `todoList` produce output, else they produce nothing. 64 | todo_include_todos = False 65 | 66 | 67 | myst_enable_extensions = [ 68 | "amsmath", 69 | "attrs_inline", 70 | "colon_fence", 71 | "deflist", 72 | "dollarmath", 73 | "fieldlist", 74 | "html_admonition", 75 | "html_image", 76 | "replacements", 77 | "smartquotes", 78 | "strikethrough", 79 | "substitution", 80 | "tasklist", 81 | ] 82 | 83 | 84 | # -- Options for customization -------------------------------------------- 85 | 86 | 87 | def monkey_patch(module_name, func_name): 88 | """Helper function for monkey-patching library functions. 89 | 90 | Used to modify a few sphinx-gallery behaviors to make the "Open in Colab" button work correctly. 91 | Should be called as a decorator with arguments. Note this behaves differently from unittest's 92 | @mock.patch, as our monkey_patch decorator should be placed on the new version of the function. 93 | """ 94 | module = import_module(module_name) 95 | original_func = getattr(module, func_name) 96 | 97 | def decorator(function): 98 | updated_func = partial(function, real_func=original_func) 99 | setattr(module, func_name, updated_func) 100 | return updated_func 101 | 102 | return decorator 103 | 104 | 105 | # -- Options for Colab integration -------------------------------------------- 106 | 107 | CURRENT_FILE_CONF = None 108 | 109 | 110 | @monkey_patch("sphinx_gallery.py_source_parser", "split_code_and_text_blocks") 111 | def split_code_and_text_blocks(source_file, real_func, return_node=False): 112 | """Monkey-patch split_code_and_text_blocks to expose sphinx-gallery's file-level config. 113 | 114 | It's kinda gross, but we need access to file_conf to detect the requires_cuda flag. 115 | """ 116 | global CURRENT_FILE_CONF 117 | 118 | res = real_func(source_file, return_node) 119 | CURRENT_FILE_CONF = res[0] 120 | return res 121 | 122 | 123 | # This header replaces the default sphinx-gallery one in sphinx_gallery/gen_rst.py. 124 | COLAB_HTML_HEADER = """ 125 | .. DO NOT EDIT. THIS FILE WAS AUTOMATICALLY GENERATED BY 126 | .. TVM'S MONKEY-PATCHED VERSION OF SPHINX-GALLERY. TO MAKE 127 | .. CHANGES, EDIT THE SOURCE PYTHON FILE: 128 | .. "{python_file}" 129 | 130 | .. only:: html 131 | 132 | .. note:: 133 | :class: sphx-glr-download-link-note 134 | 135 | This tutorial can be used interactively with Google Colab! You can also click 136 | :ref:`here ` to run the Jupyter notebook locally. 137 | 138 | .. image:: {button_svg} 139 | :align: center 140 | :target: {colab_url} 141 | :width: 300px 142 | 143 | .. rst-class:: sphx-glr-example-title 144 | 145 | .. _sphx_glr_{ref_name}: 146 | 147 | """ 148 | 149 | # Google Colab allows opening .ipynb files on GitHub by appending a GitHub path to this base URL. 150 | COLAB_URL_BASE = "https://colab.research.google.com/github" 151 | # The GitHub path where the site is automatically deployed by tvm-bot. 152 | IPYTHON_GITHUB_BASE = "mlc-ai/docs/blob/gh-pages/_downloads/" 153 | # The SVG image of the "Open in Colab" button. 154 | BUTTON = ( 155 | "https://raw.githubusercontent.com/tlc-pack/web-data/main/images/utilities/colab_button.svg" 156 | ) 157 | 158 | 159 | @monkey_patch("sphinx_gallery.gen_rst", "save_rst_example") 160 | def save_rst_example( 161 | example_rst, 162 | example_file, 163 | time_elapsed, 164 | memory_used, 165 | gallery_conf, 166 | real_func, 167 | language="Python", 168 | ): 169 | """Monkey-patch save_rst_example to include the "Open in Colab" button.""" 170 | assert language == "Python", f"Only Python examples are supported, but got {language}." 171 | # The url is the md5 hash of the notebook path. 172 | example_fname = os.path.relpath(example_file, gallery_conf["src_dir"]) 173 | ref_fname = example_fname.replace(os.path.sep, "_") 174 | notebook_path = example_fname[:-2] + "ipynb" 175 | digest = md5(notebook_path.encode()).hexdigest() 176 | 177 | # Fixed documentation versions must link to different (earlier) .ipynb notebooks. 178 | colab_url = f"{COLAB_URL_BASE}/{IPYTHON_GITHUB_BASE}" 179 | if "dev" not in version: 180 | colab_url += version + "/" 181 | colab_url += digest + "/" + os.path.basename(notebook_path) 182 | 183 | new_header = COLAB_HTML_HEADER.format( 184 | python_file=example_fname, 185 | ref_name=ref_fname, 186 | colab_url=colab_url, 187 | button_svg=BUTTON, 188 | ) 189 | with patch("sphinx_gallery.gen_rst.EXAMPLE_HEADER", new_header): 190 | real_func(example_rst, example_file, time_elapsed, memory_used, gallery_conf) 191 | 192 | 193 | INSTALL_TVM_UNITY_DEV = """ 194 | %%shell 195 | # Installs the nightly dev build of Apache TVM Unity from PyPI. If you wish to build 196 | # from source, see https://mlc.ai/docs/get_started/install.html#option-2-build-from-source 197 | pip install mlc-ai-nightly --pre -f https://mlc.ai/wheels""" 198 | 199 | INSTALL_TVM_UNITY_CUDA_DEV = """\ 200 | %%shell 201 | # Installs the nightly dev build of Apache TVM Unity from PyPI, with CUDA enabled. 202 | # To use this, you must request a Google Colab instance with a GPU by going to Runtime -> 203 | # Change runtime type -> Hardware accelerator -> GPU. 204 | # If you wish to build from source, see see 205 | # https://mlc.ai/docs/get_started/install.html#option-2-build-from-source 206 | pip install mlc-ai-nightly-cu118 --pre -f https://mlc.ai/wheels""" 207 | 208 | 209 | @monkey_patch("sphinx_gallery.gen_rst", "jupyter_notebook") 210 | def jupyter_notebook(script_blocks, gallery_conf, target_dir, real_func): 211 | """Monkey-patch sphinx-gallery to add a TVM import block to each IPython notebook. 212 | 213 | If we had only one import block, we could skip the patching and just set first_notebook_cell. 214 | However, how we import TVM depends on if we are using a fixed or dev version, and whether we 215 | will use the GPU. 216 | 217 | Tutorials requiring a CUDA-enabled build of TVM should use the flag: 218 | # sphinx_gallery_requires_cuda = True 219 | """ 220 | 221 | requires_cuda = CURRENT_FILE_CONF.get("requires_cuda", False) 222 | 223 | new_conf = { 224 | **gallery_conf, 225 | "first_notebook_cell": INSTALL_TVM_UNITY_CUDA_DEV 226 | if requires_cuda 227 | else INSTALL_TVM_UNITY_DEV, 228 | } 229 | return real_func(script_blocks, new_conf, target_dir) 230 | 231 | 232 | # -- Options for HTML output ---------------------------------------------- 233 | 234 | 235 | def fixup_tutorials(original_url: str) -> str: 236 | if "tutorials" in original_url: 237 | # The `index.rst` is omitted from the URL in the sidebar 238 | assert not original_url.endswith("index.rst") 239 | return re.sub(r"tutorials/(.*)\.rst", "tutorials/\\1.py", original_url) 240 | else: 241 | # do nothing for normal non-tutorial .rst files 242 | return original_url 243 | 244 | 245 | # The theme is set by the make target 246 | 247 | html_theme = "sphinx_rtd_theme" 248 | html_theme_path = [sphinx_rtd_theme.get_html_theme_path()] 249 | 250 | templates_path = [] 251 | 252 | html_static_path = [] 253 | 254 | footer_copyright = "© 2023 Apache Software Foundation | All rights reserved" 255 | footer_note = " ".join( 256 | """ 257 | Copyright © 2023 The Apache Software Foundation. Apache TVM, Apache, the Apache feather, 258 | and the Apache TVM project logo are either trademarks or registered trademarks of 259 | the Apache Software Foundation.""".split( 260 | "\n" 261 | ) 262 | ).strip() 263 | 264 | header_logo = "https://tvm.apache.org/assets/images/logo.svg" 265 | html_logo = "_static/img/tvm-logo-small.png" 266 | html_favicon = "_static/img/tvm-logo-square.png" 267 | 268 | html_theme_options = { 269 | "logo_only": True, 270 | } 271 | 272 | header_links = [ 273 | ("Github", "https://github.com/apache/tvm/tree/unity/"), 274 | ("MLC-LLM", "https://mlc.ai/mlc-llm/"), 275 | ("MLC-Tutorial", "https://mlc.ai/"), 276 | ] 277 | 278 | header_dropdown = { 279 | "name": "ASF", 280 | "items": [ 281 | ("Apache Homepage", "https://apache.org/"), 282 | ("License", "https://www.apache.org/licenses/"), 283 | ("Sponsorship", "https://www.apache.org/foundation/sponsorship.html"), 284 | ("Security", "https://www.apache.org/security/"), 285 | ("Thanks", "https://www.apache.org/foundation/thanks.html"), 286 | ("Events", "https://www.apache.org/events/current-event"), 287 | ], 288 | } 289 | 290 | html_context = { 291 | "footer_copyright": footer_copyright, 292 | "footer_note": footer_note, 293 | "header_links": header_links, 294 | "header_dropdown": header_dropdown, 295 | "display_github": True, 296 | "github_user": "mlc-ai", 297 | "github_repo": "docs", 298 | "github_version": "main/docs/", 299 | "theme_vcs_pageview_mode": "edit", 300 | "edit_link_hook_fn": fixup_tutorials, 301 | # "header_logo": "/path/to/logo", 302 | # "header_logo_link": "", 303 | # "version_selecter": "", 304 | } 305 | 306 | 307 | # add additional overrides 308 | templates_path += [tlcpack_sphinx_addon.get_templates_path()] 309 | html_static_path += [tlcpack_sphinx_addon.get_static_path()] 310 | 311 | 312 | # Sphinx-Gallery Settings 313 | examples_dirs = [ 314 | f"{home_path}/get_started/tutorials", 315 | f"{home_path}/deep_dive/tensor_ir/tutorials/", 316 | ] 317 | 318 | gallery_dirs = [ 319 | "get_started/tutorials/", 320 | "deep_dive/tensor_ir/tutorials/", 321 | ] 322 | 323 | sphinx_gallery_conf = { 324 | "examples_dirs": examples_dirs, 325 | "gallery_dirs": gallery_dirs, 326 | "backreferences_dir": "gen_modules/backreferences", 327 | "filename_pattern": r".*\.py", 328 | "ignore_pattern": r"__init__\.py", 329 | "show_signature": False, 330 | "download_all_examples": False, 331 | "promote_jupyter_magic": True, 332 | } 333 | 334 | # -- Options for Autodoc ---------------------------------------------- 335 | 336 | add_module_names = True 337 | autodoc_default_options = { 338 | "member-order": "alphabetical", 339 | } 340 | 341 | tvm_class_name_rewrite_map = { 342 | "tvm.tir": ["Var", "Call"], 343 | "tvm.relax": ["Var", "Call"], 344 | "tvm.relax.frontend.nn": ["Module"], 345 | } 346 | 347 | 348 | def distinguish_class_name(name: str, lines: List[str]): 349 | """Distinguish the docstring of type annotations. 350 | 351 | In the whole TVM, there are many classes with the same name but in different modules, 352 | e.g. ``tir.Var``, ``relax.Var``. This function is used to distinguish them in the docstring, 353 | by adding the module name as prefix. 354 | 355 | To be specific, this function will check the current object name, and if it in the specific 356 | module with specific name, it will add the module name as prefix to the class name to prevent 357 | the confusion. Further, we only add the prefix to those standalone class name, but skip 358 | the pattern of `xx.Var`, `Var.xx` and `xx.Var.xx`. 359 | 360 | Parameters 361 | ---------- 362 | name : str 363 | The full name of the object in the doc. 364 | 365 | lines : list 366 | The docstring lines, need to be modified inplace. 367 | """ 368 | remap = {} 369 | for module_name in tvm_class_name_rewrite_map: 370 | if name.startswith(module_name): 371 | short_name = module_name[4:] if module_name.startswith("tvm.") else module_name 372 | for class_name in tvm_class_name_rewrite_map[module_name]: 373 | remap.update({class_name: f"{short_name}.{class_name}"}) 374 | 375 | for k, v in remap.items(): 376 | for i in range(len(lines)): 377 | lines[i] = re.sub(rf"(? None: 40 | for i in range(128): 41 | with T.block("C"): 42 | vi = T.axis.spatial(128, i) 43 | C[vi] = A[vi] + B[vi] 44 | 45 | Key Elements of Tensor Programs 46 | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 47 | 48 | The demonstrated primitive tensor function calculates the element-wise sum of two vectors. 49 | The function: 50 | 51 | - Accepts three **multi-dimensional buffers** as parameters, and generates one **multi-dimensional 52 | buffer** as output. 53 | - Incorporates a solitary **loop nest** ``i`` that facilitates the computation. 54 | - Features a singular **compute statement** that calculates the element-wise sum of the two 55 | vectors. 56 | 57 | Extra Structure in TensorIR 58 | ~~~~~~~~~~~~~~~~~~~~~~~~~~~ 59 | Crucially, we are unable to execute arbitrary transformations on the program, as certain 60 | computations rely on the loop's sequence. Fortunately, the majority of primitive tensor 61 | functions we focus on possess favorable properties, such as independence among loop iterations. 62 | For instance, the aforementioned program includes block and iteration annotations: 63 | 64 | - The **block annotation** ``with T.block("C")`` signifies that the block is the fundamental 65 | computation unit designated for scheduling. A block may encompass a single computation 66 | statement, multiple computation statements with loops, or opaque intrinsics such as Tensor 67 | Core instructions. 68 | - The **iteration annotation** ``T.axis.spatial``, indicating that variable ``vi`` is mapped 69 | to ``i``, and all iterations are independent. 70 | 71 | While this information isn't crucial for *executing* the specific program, it proves useful when 72 | transforming the program. Consequently, we can confidently parallelize or reorder loops associated 73 | with ``vi``, provided we traverse all the index elements from 0 to 128. 74 | -------------------------------------------------------------------------------- /docs/deep_dive/tensor_ir/index.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | .. _tensor-ir: 19 | 20 | TensorIR 21 | ======== 22 | TensorIR is one of the core abstraction in Apache TVM Unity stack, which is used to 23 | represent and optimize the primitive tensor functions. 24 | 25 | .. toctree:: 26 | :maxdepth: 2 27 | 28 | abstraction 29 | learning 30 | tutorials/creation 31 | tutorials/transformation 32 | -------------------------------------------------------------------------------- /docs/deep_dive/tensor_ir/learning.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | .. _tir-learning: 19 | 20 | Understand TensorIR Abstraction 21 | =============================== 22 | TensorIR is the tensor program abstraction in Apache TVM, which is one of the standard 23 | machine learning compilation frameworks. The principal objective of tensor program abstraction 24 | is to depict loops and associated hardware acceleration options, including threading, the 25 | application of specialized hardware instructions, and memory access. 26 | 27 | To help our explanations, let us use the following sequence of tensor computations as 28 | a motivating example. Specifically, for two :math:`128 \times 128` matrices ``A`` and ``B``, let us perform the 29 | following two steps of tensor computations. 30 | 31 | .. math:: 32 | 33 | Y_{i, j} &= \sum_k A_{i, k} \times B_{k, j} \\ 34 | C_{i, j} &= \mathbb{relu}(Y_{i, j}) = \mathbb{max}(Y_{i, j}, 0) 35 | 36 | 37 | The above computations resemble a typical primitive tensor function commonly seen in neural networks, 38 | a linear layer with relu activation. We use TensorIR to depict the above computations as follows. 39 | 40 | Before we invoke TensorIR, let's use native Python codes with NumPy to show the computation: 41 | 42 | .. code:: python 43 | 44 | def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray): 45 | Y = np.empty((128, 128), dtype="float32") 46 | for i in range(128): 47 | for j in range(128): 48 | for k in range(128): 49 | if k == 0: 50 | Y[i, j] = 0 51 | Y[i, j] = Y[i, j] + A[i, k] * B[k, j] 52 | for i in range(128): 53 | for j in range(128): 54 | C[i, j] = max(Y[i, j], 0) 55 | 56 | With the low-level NumPy example in mind, now we are ready to introduce TensorIR. The code block 57 | below shows a TensorIR implementation of ``mm_relu``. The particular code is implemented in a 58 | language called TVMScript, which is a domain-specific dialect embedded in python AST. 59 | 60 | .. code:: python 61 | 62 | @tvm.script.ir_module 63 | class MyModule: 64 | @T.prim_func 65 | def mm_relu(A: T.Buffer((128, 128), "float32"), 66 | B: T.Buffer((128, 128), "float32"), 67 | C: T.Buffer((128, 128), "float32")): 68 | Y = T.alloc_buffer((128, 128), dtype="float32") 69 | for i, j, k in T.grid(128, 128, 128): 70 | with T.block("Y"): 71 | vi = T.axis.spatial(128, i) 72 | vj = T.axis.spatial(128, j) 73 | vk = T.axis.reduce(128, k) 74 | with T.init(): 75 | Y[vi, vj] = T.float32(0) 76 | Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj] 77 | for i, j in T.grid(128, 128): 78 | with T.block("C"): 79 | vi = T.axis.spatial(128, i) 80 | vj = T.axis.spatial(128, j) 81 | C[vi, vj] = T.max(Y[vi, vj], T.float32(0)) 82 | 83 | 84 | Next, let's invest the elements in the above TensorIR program. 85 | 86 | Function Parameters and Buffers 87 | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 88 | **The function parameters correspond to the same set of parameters on the numpy function.** 89 | 90 | .. code:: python 91 | 92 | # TensorIR 93 | def mm_relu(A: T.Buffer[(128, 128), "float32"], 94 | B: T.Buffer[(128, 128), "float32"], 95 | C: T.Buffer[(128, 128), "float32"]): 96 | ... 97 | # NumPy 98 | def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray): 99 | ... 100 | 101 | Here ``A``, ``B``, and ``C`` takes a type named ``T.Buffer``, which with shape 102 | argument ``(128, 128)`` and data type ``float32``. This additional information 103 | helps possible MLC process to generate code that specializes in the shape and data 104 | type. 105 | 106 | **Similarly, TensorIR also uses a buffer type in intermediate result allocation.** 107 | 108 | .. code:: python 109 | 110 | # TensorIR 111 | Y = T.alloc_buffer((128, 128), dtype="float32") 112 | # NumPy 113 | Y = np.empty((128, 128), dtype="float32") 114 | 115 | Loop Iterations 116 | ~~~~~~~~~~~~~~~ 117 | **There are also direct correspondence of loop iterations.** 118 | 119 | ``T.grid`` is a syntactic sugar in TensorIR for us to write multiple nested iterators. 120 | 121 | .. code:: python 122 | 123 | # TensorIR with `T.grid` 124 | for i, j, k in T.grid(128, 128, 128): 125 | ... 126 | # TensorIR with `range` 127 | for i in range(128): 128 | for j in range(128): 129 | for k in range(128): 130 | ... 131 | # NumPy 132 | for i in range(128): 133 | for j in range(128): 134 | for k in range(128): 135 | ... 136 | 137 | Computational Block 138 | ~~~~~~~~~~~~~~~~~~~ 139 | A significant distinction lies in computational statements: 140 | **TensorIR incorporates an additional construct termed** ``T.block``. 141 | 142 | .. code:: python 143 | 144 | # TensorIR 145 | with T.block("Y"): 146 | vi = T.axis.spatial(128, i) 147 | vj = T.axis.spatial(128, j) 148 | vk = T.axis.reduce(128, k) 149 | with T.init(): 150 | Y[vi, vj] = T.float32(0) 151 | Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj] 152 | # NumPy 153 | vi, vj, vk = i, j, k 154 | if vk == 0: 155 | Y[vi, vj] = 0 156 | Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj] 157 | 158 | A **block** represents a fundamental computation unit within TensorIR. Importantly, 159 | a block encompasses more information than standard NumPy code. It comprises a set of block axes 160 | ``(vi, vj, vk)`` and the computations delineated around them. 161 | 162 | .. code:: python 163 | 164 | vi = T.axis.spatial(128, i) 165 | vj = T.axis.spatial(128, j) 166 | vk = T.axis.reduce(128, k) 167 | 168 | The above three lines declare the **key properties** about block axes in the following syntax. 169 | 170 | .. code:: python 171 | 172 | [block_axis] = T.axis.[axis_type]([axis_range], [mapped_value]) 173 | 174 | These three lines convey the following details: 175 | 176 | - They specify the binding of ``vi``, ``vj``, ``vk`` (in this instance, to ``i``, ``j``, ``k``). 177 | - They declare the original range intended for ``vi``, ``vj``, ``vk`` 178 | (the 128 in ``T.axis.spatial(128, i)``). 179 | - They announce the properties of the iterators (spatial, reduce). 180 | 181 | Block Axis Properties 182 | ~~~~~~~~~~~~~~~~~~~~~ 183 | Let's delve deeper into the properties of the block axis. These properties signify the axis's 184 | relationship to the computation in progress. The block comprises three axes ``vi``, ``vj``, and 185 | ``vk``, meanwhile the block reads the buffer ``A[vi, vk]``, ``B[vk, vj]`` and writs the buffer 186 | ``Y[vi, vj]``. Strictly speaking, the block performs (reduction) updates to Y, which we label 187 | as write for the time being, as we don't require the value of Y from another block. 188 | 189 | Significantly, for a fixed value of ``vi`` and ``vj``, the computation block yields a point 190 | value at a spatial location of ``Y`` (``Y[vi, vj]``) that is independent of other locations in ``Y`` 191 | (with different ``vi``, ``vj`` values). We can refer to ``vi``, ``vj`` as **spatial axes** since 192 | they directly correspond to the start of a spatial region of buffers that the block writes to. 193 | The axes involved in reduction (``vk``) are designated as **reduce axes**. 194 | 195 | Why Extra Information in Block 196 | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 197 | One crucial observation is that the additional information (block axis range and their properties) 198 | makes the block to be **self-contained** when it comes to the iterations that it is supposed to 199 | carry out independent from the external loop-nest ``i, j, k``. 200 | 201 | The block axis information also provides additional properties that help us to validate the correctness of the 202 | external loops that are used to carry out the computation. For example, the above code block will result in an 203 | error because the loop expects an iterator of size 128, but we only bound it to a for loop of size 127. 204 | 205 | .. code:: python 206 | 207 | # wrong program due to loop and block iteration mismatch 208 | for i in range(127): 209 | with T.block("C"): 210 | vi = T.axis.spatial(128, i) 211 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^ 212 | error here due to iterator size mismatch 213 | ... 214 | 215 | Sugars for Block Axes Binding 216 | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 217 | In situations where each of the block axes is directly mapped to an outer loop iterator, 218 | we can use ``T.axis.remap`` to declare the block axis in a single line. 219 | 220 | .. code:: python 221 | 222 | # SSR means the properties of each axes are "spatial", "spatial", "reduce" 223 | vi, vj, vk = T.axis.remap("SSR", [i, j, k]) 224 | 225 | which is equivalent to 226 | 227 | .. code:: python 228 | 229 | vi = T.axis.spatial(range_of_i, i) 230 | vj = T.axis.spatial(range_of_j, j) 231 | vk = T.axis.reduce (range_of_k, k) 232 | 233 | So we can also write the programs as follows. 234 | 235 | .. code:: python 236 | 237 | @tvm.script.ir_module 238 | class MyModuleWithAxisRemapSugar: 239 | @T.prim_func 240 | def mm_relu(A: T.Buffer((128, 128), "float32"), 241 | B: T.Buffer((128, 128), "float32"), 242 | C: T.Buffer((128, 128), "float32")): 243 | Y = T.alloc_buffer((128, 128), dtype="float32") 244 | for i, j, k in T.grid(128, 128, 128): 245 | with T.block("Y"): 246 | vi, vj, vk = T.axis.remap("SSR", [i, j, k]) 247 | with T.init(): 248 | Y[vi, vj] = T.float32(0) 249 | Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj] 250 | for i, j in T.grid(128, 128): 251 | with T.block("C"): 252 | vi, vj = T.axis.remap("SS", [i, j]) 253 | C[vi, vj] = T.max(Y[vi, vj], T.float32(0)) 254 | 255 | -------------------------------------------------------------------------------- /docs/deep_dive/tensor_ir/tutorials/README.rst: -------------------------------------------------------------------------------- 1 | Deep Dive: TensorIR 2 | ------------------- 3 | -------------------------------------------------------------------------------- /docs/deep_dive/tensor_ir/tutorials/creation.py: -------------------------------------------------------------------------------- 1 | # Licensed to the Apache Software Foundation (ASF) under one 2 | # or more contributor license agreements. See the NOTICE file 3 | # distributed with this work for additional information 4 | # regarding copyright ownership. The ASF licenses this file 5 | # to you under the Apache License, Version 2.0 (the 6 | # "License"); you may not use this file except in compliance 7 | # with the License. You may obtain a copy of the License at 8 | # 9 | # http://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, 12 | # software distributed under the License is distributed on an 13 | # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | # KIND, either express or implied. See the License for the 15 | # specific language governing permissions and limitations 16 | # under the License. 17 | 18 | """ 19 | .. _tir-creation: 20 | 21 | TensorIR Creation 22 | ----------------- 23 | In this section, we will introduce the methods to write a TensorIR function 24 | in Apache TVM Unity. This tutorial presumes familiarity with the fundamental concepts of TensorIR. 25 | If not already acquainted, please refer to :ref:`tir-learning` initially. 26 | 27 | .. note:: 28 | 29 | This tutorial concentrates on the construction of **standalone** TensorIR functions. The 30 | techniques presented here are not requisite for end users to compile Relax models. For specifics 31 | on converting Relax Models to TensorIR functions, please refer to 🚧 TODO: Add link. 32 | 33 | """ 34 | 35 | ###################################################################### 36 | # Create TensorIR using TVMScript 37 | # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 38 | # The most straightforward way to create a TensorIR function via TVMScript. 39 | # TVMScript is a TVM Python dialect that represents TensorIR in TVM. 40 | # 41 | # .. important:: 42 | # 43 | # While TVMScript employs Python syntax and AST, ensuring full compatibility 44 | # with Python tools like auto-completion and linting, it is not a native Python 45 | # language and cannot be executed by a Python interpreter. 46 | # 47 | # More precisely, the decorator **@tvm.script** extracts the Python AST from 48 | # the decorated function, subsequently parsing it into TensorIR. 49 | # 50 | # Standard Format 51 | # *************** 52 | # Let's take an example of ``mm_relu`` from :ref:`tir-learning`. Here is the complete 53 | # format of the ir_module and in TVMScript: 54 | 55 | 56 | import numpy as np 57 | import tvm 58 | from tvm.script import ir as I 59 | from tvm.script import tir as T 60 | 61 | 62 | @I.ir_module 63 | class MyModule: 64 | @T.prim_func 65 | def mm_relu( 66 | A: T.Buffer((128, 128), "float32"), 67 | B: T.Buffer((128, 128), "float32"), 68 | C: T.Buffer((128, 128), "float32"), 69 | ): 70 | Y = T.alloc_buffer((128, 128), dtype="float32") 71 | for i in range(128): 72 | for j in range(128): 73 | for k in range(128): 74 | with T.block("Y"): 75 | vi = T.axis.spatial(128, i) 76 | vj = T.axis.spatial(128, j) 77 | vk = T.axis.reduce(128, k) 78 | T.reads(A[vi, vk], B[vk, vj]) 79 | T.writes(Y[vi, vj]) 80 | with T.init(): 81 | Y[vi, vj] = T.float32(0) 82 | Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj] 83 | for i in range(128): 84 | for j in range(128): 85 | with T.block("C"): 86 | vi = T.axis.spatial(128, i) 87 | vj = T.axis.spatial(128, j) 88 | T.reads(Y[vi, vj]) 89 | T.writes(C[vi, vj]) 90 | C[vi, vj] = T.max(Y[vi, vj], T.float32(0)) 91 | 92 | 93 | ###################################################################### 94 | # Concise with Syntactic Sugar 95 | # **************************** 96 | # For ease of writing, we can employ the following syntactic sugar to 97 | # streamline the code: 98 | # 99 | # - Utilize ``T.grid`` to condense nested loops; 100 | # - Employ ``T.axis.remap`` to abbreviate block iterator annotations; 101 | # - Exclude ``T.reads`` and ``T.writes`` for blocks whose content can 102 | # be inferred from the block body; 103 | 104 | 105 | @I.ir_module 106 | class ConciseModule: 107 | @T.prim_func 108 | def mm_relu( 109 | A: T.Buffer((128, 128), "float32"), 110 | B: T.Buffer((128, 128), "float32"), 111 | C: T.Buffer((128, 128), "float32"), 112 | ): 113 | Y = T.alloc_buffer((128, 128), dtype="float32") 114 | for i, j, k in T.grid(128, 128, 128): 115 | with T.block("Y"): 116 | vi, vj, vk = T.axis.remap("SSR", [i, j, k]) 117 | with T.init(): 118 | Y[vi, vj] = T.float32(0) 119 | Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj] 120 | for i, j in T.grid(128, 128): 121 | with T.block("C"): 122 | vi, vj = T.axis.remap("SS", [i, j]) 123 | C[vi, vj] = T.max(Y[vi, vj], T.float32(0)) 124 | 125 | 126 | ###################################################################### 127 | # We can use the following code to verify that the two modules are equivalent: 128 | 129 | print(tvm.ir.structural_equal(MyModule, ConciseModule)) 130 | 131 | ###################################################################### 132 | # Interactive with Python Variables 133 | # ********************************* 134 | # Despite TVMScript not being executed by a Python interpreter, limited 135 | # interaction with Python is feasible. For instance, Python variables can 136 | # be used to ascertain the shape and data type of a TensorIR. 137 | 138 | # Python variables 139 | M = N = K = 128 140 | dtype = "float32" 141 | 142 | 143 | # IRModule in TVMScript 144 | @I.ir_module 145 | class ConciseModuleFromPython: 146 | @T.prim_func 147 | def mm_relu( 148 | A: T.Buffer((M, K), dtype), 149 | B: T.Buffer((K, N), dtype), 150 | C: T.Buffer((M, N), dtype), 151 | ): 152 | Y = T.alloc_buffer((M, N), dtype) 153 | for i, j, k in T.grid(M, N, K): 154 | with T.block("Y"): 155 | vi, vj, vk = T.axis.remap("SSR", [i, j, k]) 156 | with T.init(): 157 | Y[vi, vj] = T.cast(T.float32(0), dtype) 158 | Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj] 159 | for i, j in T.grid(M, N): 160 | with T.block("C"): 161 | vi, vj = T.axis.remap("SS", [i, j]) 162 | C[vi, vj] = T.max(Y[vi, vj], T.cast(T.float32(0), dtype)) 163 | 164 | 165 | ###################################################################### 166 | # Check the equivalence: 167 | 168 | print(tvm.ir.structural_equal(ConciseModule, ConciseModuleFromPython)) 169 | 170 | 171 | ###################################################################### 172 | # TensorIR Function with Dynamic Shapes 173 | # ************************************* 174 | # Despite TVMScript not being executed by a Python interpreter, limited 175 | # interaction with Python is feasible. For instance, Python variables can 176 | # be used to ascertain the shape and data type of a TensorIR. 177 | 178 | 179 | @I.ir_module 180 | class DynamicShapeModule: 181 | @T.prim_func 182 | def mm_relu(a: T.handle, b: T.handle, c: T.handle): 183 | # Dynamic shape definition 184 | M, N, K = T.int32(), T.int32(), T.int32() 185 | 186 | # Bind the input buffers with the dynamic shapes 187 | A = T.match_buffer(a, [M, K], dtype) 188 | B = T.match_buffer(b, [K, N], dtype) 189 | C = T.match_buffer(c, [M, N], dtype) 190 | Y = T.alloc_buffer((M, N), dtype) 191 | for i, j, k in T.grid(M, N, K): 192 | with T.block("Y"): 193 | vi, vj, vk = T.axis.remap("SSR", [i, j, k]) 194 | with T.init(): 195 | Y[vi, vj] = T.cast(T.float32(0), dtype) 196 | Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj] 197 | for i, j in T.grid(M, N): 198 | with T.block("C"): 199 | vi, vj = T.axis.remap("SS", [i, j]) 200 | C[vi, vj] = T.max(Y[vi, vj], T.cast(T.float32(0), dtype)) 201 | 202 | 203 | ###################################################################### 204 | # Now let's check the runtime dynamic shape inference: 205 | 206 | 207 | def evaluate_dynamic_shape(lib: tvm.runtime.Module, m: int, n: int, k: int): 208 | A = tvm.nd.array(np.random.uniform(size=(m, k)).astype("float32")) 209 | B = tvm.nd.array(np.random.uniform(size=(k, n)).astype("float32")) 210 | C = tvm.nd.array(np.zeros((m, n), dtype="float32")) 211 | lib(A, B, C) 212 | return C.numpy() 213 | 214 | 215 | # Compile lib only once 216 | dyn_shape_lib = tvm.build(DynamicShapeModule, target="llvm") 217 | # Able to handle different shapes 218 | print(evaluate_dynamic_shape(dyn_shape_lib, m=4, n=4, k=4)) 219 | print(evaluate_dynamic_shape(dyn_shape_lib, m=64, n=64, k=128)) 220 | 221 | ###################################################################### 222 | # Create TensorIR using Tensor Expression 223 | # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 224 | # Often, the specifics of TensorIR are disregarded in favor of expressing the computation more 225 | # succinctly, leading to the pragmatic generation of TensorIR. This is where Tensor Expression 226 | # (TE) becomes relevant. 227 | # 228 | # Tensor Expression (TE) serves as a domain-specific language delineating a sequence of 229 | # computations through an expression-like API. 230 | # 231 | # .. note:: 232 | # 233 | # Tensor Expression comprises two components within the TVM stack: the expression and the 234 | # schedule. The expression is the domain-specific language embodying the computation pattern, 235 | # precisely what we're addressing in this section. Conversely, the TE schedule is the legacy 236 | # scheduling method, has been superseded by the TensorIR schedule in the TVM Unity stack. 237 | # 238 | # Create Static-Shape Functions 239 | # ***************************** 240 | # We use the same example of ``mm_relu`` from the last subsection to demonstrate the 241 | # TE creation method. 242 | 243 | from tvm import te 244 | 245 | A = te.placeholder((128, 128), "float32", name="A") 246 | B = te.placeholder((128, 128), "float32", name="B") 247 | k = te.reduce_axis((0, 128), "k") 248 | Y = te.compute((128, 128), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="Y") 249 | C = te.compute((128, 128), lambda i, j: te.max(Y[i, j], 0), name="C") 250 | 251 | ###################################################################### 252 | # Here ``te.compute`` takes the signature ``te.compute(output_shape, fcompute)``. 253 | # And the fcompute function describes how we want to compute the value of each 254 | # element ``Y[i, j]`` for a given index: 255 | # 256 | # .. code:: python 257 | # 258 | # lambda i, j: te.sum(A[i, k] * B[k, j], axis=k) 259 | # 260 | # The aforementioned lambda expression encapsulates the computation: 261 | # :math:`Y_{i, j} = \sum_k A_{i, k} \times B_{k, j}`. Upon defining the computation, 262 | # we can formulate a TensorIR function by incorporating the pertinent parameters of interest. 263 | # In this specific instance, we aim to construct a function with two input parameters **A, B** 264 | # and one output parameter **C**. 265 | 266 | te_func = te.create_prim_func([A, B, C]).with_attr({"global_symbol": "mm_relu"}) 267 | TEModule = tvm.IRModule({"mm_relu": te_func}) 268 | TEModule.show() 269 | 270 | ###################################################################### 271 | # Create Dynamic-Shape Functions 272 | # ****************************** 273 | # We can also create a dynamic-shape function using Tensor Expression. The only difference 274 | # is that we need to specify the shape of the input tensors as symbolic variables. 275 | 276 | # Declare symbolic variables 277 | M, N, K = te.var("m"), te.var("n"), te.var("k") 278 | A = te.placeholder((M, N), "float32", name="A") 279 | B = te.placeholder((K, N), "float32", name="B") 280 | k = te.reduce_axis((0, K), "k") 281 | Y = te.compute((M, N), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="Y") 282 | C = te.compute((M, N), lambda i, j: te.max(Y[i, j], 0), name="C") 283 | 284 | dyn_te_func = te.create_prim_func([A, B, C]).with_attr({"global_symbol": "mm_relu"}) 285 | DynamicTEModule = tvm.IRModule({"mm_relu": dyn_te_func}) 286 | DynamicTEModule.show() 287 | -------------------------------------------------------------------------------- /docs/deep_dive/tensor_ir/tutorials/transformation.py: -------------------------------------------------------------------------------- 1 | # Licensed to the Apache Software Foundation (ASF) under one 2 | # or more contributor license agreements. See the NOTICE file 3 | # distributed with this work for additional information 4 | # regarding copyright ownership. The ASF licenses this file 5 | # to you under the Apache License, Version 2.0 (the 6 | # "License"); you may not use this file except in compliance 7 | # with the License. You may obtain a copy of the License at 8 | # 9 | # http://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, 12 | # software distributed under the License is distributed on an 13 | # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | # KIND, either express or implied. See the License for the 15 | # specific language governing permissions and limitations 16 | # under the License. 17 | 18 | """ 19 | .. _tir-transform: 20 | 21 | Transformation 22 | -------------- 23 | In this section, we will get to the main ingredients of all MLC flows - transformations 24 | of primitive tensor functions. 25 | """ 26 | 27 | ###################################################################### 28 | # In the :ref:`previous section `, we have given an example of how to write 29 | # ``mm_relu`` using TensorIR. In practice, there can be multiple ways to implement 30 | # the same functionality, and each implementation can result in different performance. 31 | # 32 | # .. note:: 33 | # This tutorial primarily illustrates the application of TensorIR Transformation, 34 | # rather than delving into optimization techniques. 35 | # 36 | # First, let's take a look at the implementation of ``mm_relu`` in the previous section: 37 | 38 | import tvm 39 | from tvm.script import ir as I 40 | from tvm.script import tir as T 41 | 42 | 43 | @I.ir_module 44 | class MyModule: 45 | @T.prim_func 46 | def main( 47 | A: T.Buffer((128, 128), "float32"), 48 | B: T.Buffer((128, 128), "float32"), 49 | C: T.Buffer((128, 128), "float32"), 50 | ): 51 | T.func_attr({"tir.noalias": T.bool(True)}) 52 | Y = T.alloc_buffer((128, 128)) 53 | for i, j, k in T.grid(128, 128, 128): 54 | with T.block("Y"): 55 | vi, vj, vk = T.axis.remap("SSR", [i, j, k]) 56 | with T.init(): 57 | Y[vi, vj] = T.float32(0) 58 | Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj] 59 | for i, j in T.grid(128, 128): 60 | with T.block("C"): 61 | vi, vj = T.axis.remap("SS", [i, j]) 62 | C[vi, vj] = T.max(Y[vi, vj], T.float32(0)) 63 | 64 | 65 | ###################################################################### 66 | # Before we transform the function, let's first evaluate the performance of the 67 | # original implementation. 68 | 69 | import numpy as np 70 | 71 | a_np = np.random.uniform(size=(128, 128)).astype("float32") 72 | b_np = np.random.uniform(size=(128, 128)).astype("float32") 73 | c_np = a_np @ b_np 74 | 75 | a_nd = tvm.nd.array(a_np) 76 | b_nd = tvm.nd.array(b_np) 77 | c_nd = tvm.nd.array(np.zeros((128, 128), dtype="float32")) 78 | 79 | 80 | def evaluate(mod: tvm.IRModule): 81 | lib = tvm.build(mod, target="llvm") 82 | # check correctness 83 | lib(a_nd, b_nd, c_nd) 84 | np.testing.assert_allclose(c_nd.numpy(), c_np, rtol=1e-5) 85 | # evaluate performance 86 | f_timer = lib.time_evaluator("main", tvm.cpu()) 87 | print(f_timer(a_nd, b_nd, c_nd)) 88 | 89 | 90 | evaluate(MyModule) 91 | 92 | ###################################################################### 93 | # Initialization Schedule 94 | # *********************** 95 | # We initiate the process of code transformation by establishing a Schedule helper class, 96 | # utilizing the provided **MyModule** as input. 97 | 98 | sch = tvm.tir.Schedule(MyModule) 99 | 100 | ###################################################################### 101 | # Loop Tiling 102 | # *********** 103 | # Subsequently, we execute the requisite operations to acquire a reference to 104 | # block **Y** and its associated loops. 105 | 106 | block_Y = sch.get_block("Y") 107 | i, j, k = sch.get_loops(block_Y) 108 | 109 | ###################################################################### 110 | # We now proceed to execute the transformations. The initial modification involves 111 | # splitting loop ``j`` into two separate loops, with the inner loop possessing a 112 | # length of 4. It is crucial to understand that the transformation process is procedural; 113 | # thus, inadvertent execution of the block twice will yield an error stating the 114 | # non-existence of variable ``j``. 115 | 116 | j0, j1 = sch.split(j, factors=[None, 8]) 117 | 118 | ###################################################################### 119 | # The outcome of the transformation can be examined, as it is retained within ``sch.mod``. 120 | 121 | sch.mod.show() 122 | 123 | ###################################################################### 124 | # Following the initial transformation phase, two supplementary loops, ``j_0`` and ``j_1``, 125 | # have been generated with respective ranges of 32 and 4. The subsequent 126 | # action involves reordering these two loops. 127 | 128 | sch.reorder(j0, k, j1) 129 | sch.mod.show() 130 | evaluate(sch.mod) 131 | 132 | ###################################################################### 133 | # Leverage Localities 134 | # ******************* 135 | # Subsequently, we will execute two additional transformation steps to achieve a different 136 | # variant. First, we employ a primitive known as **reverse_compute_at** to relocate block 137 | # **C** to an inner loop of **Y**. 138 | 139 | block_C = sch.get_block("C") 140 | sch.reverse_compute_at(block_C, j0) 141 | sch.mod.show() 142 | 143 | ###################################################################### 144 | # Rewrite Reduction 145 | # ***************** 146 | # Until now, the reduction initialization and update step have been maintained together 147 | # within a single block body. This amalgamated form facilitates loop transformations, 148 | # as the outer loops ``i``, ``j`` of initialization and updates generally need to remain 149 | # synchronized. 150 | # 151 | # Following the loop transformations, we can segregate the initialization of Y's elements 152 | # from the reduction update via the **decompose_reduction** primitive. 153 | 154 | sch.decompose_reduction(block_Y, k) 155 | sch.mod.show() 156 | evaluate(sch.mod) 157 | 158 | ###################################################################### 159 | # Trace the Transformation 160 | # ************************ 161 | # TensorIR schedule is a procedural language, and the transformation is executed in a 162 | # step-by-step manner. We can trace the transformation by printing the schedule or the 163 | # history of the schedule. 164 | # 165 | # We've already see the schedule by printing ``sch.mod``. We can also print the history 166 | # of the schedule by ``sch.trace``. 167 | 168 | sch.trace.show() 169 | 170 | ###################################################################### 171 | # Alternatively, we can output the IRModule in conjunction with the historical trace. 172 | 173 | sch.show() 174 | -------------------------------------------------------------------------------- /docs/get_started/install.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | .. _install: 19 | 20 | Installing Apache TVM 21 | ===================== 22 | 23 | .. contents:: Table of Contents 24 | :local: 25 | :depth: 2 26 | 27 | Option 1. Prebuilt Package 28 | -------------------------- 29 | To help our community to use Apache TVM, a nightly prebuilt developer package is provided by 30 | `MLC community `_ 31 | 32 | Please visit the installation page for installation instructions: https://mlc.ai/package/. 33 | 34 | Option 2. Build from Source 35 | --------------------------- 36 | While it is generally recommended to always use the prebuilt TVM, if you require more customization, 37 | you may need to build it from source. 38 | 39 | Step 1. Install Dependencies 40 | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 41 | TVM requires the following dependencies: 42 | 43 | - CMake (>= 3.24.0) 44 | - LLVM (recommended >= 15) 45 | - Git 46 | - A recent C++ compiler supporting C++ 17, at the minimum 47 | - GCC 7.1 48 | - Clang 5.0 49 | - Apple Clang 9.3 50 | - Visual Studio 2019 (v16.7) 51 | - Python (>= 3.8) 52 | - (Optional) Conda (Strongly Recommended) 53 | 54 | 55 | For Ubuntu/Debian users, the following APT Repository may help: 56 | 57 | - CMake: https://apt.kitware.com 58 | - LLVM: https://apt.llvm.org 59 | 60 | Step 2. Get Source from Github 61 | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 62 | First, You can also choose to clone the source repo from Github. The code of Apache TVM is hosted 63 | under the `Apache TVM `_ 64 | 65 | .. code:: bash 66 | 67 | git clone https://github.com/apache/tvm --recursive 68 | 69 | .. note:: 70 | It's important to use the ``--recursive`` flag when cloning the TVM repository, which will 71 | automatically clone the submodules. If you forget to use this flag, you can manually clone the submodules 72 | by running ``git submodule update --init --recursive`` in the root directory of the TVM repository. 73 | 74 | 75 | Step 3. Configure and Build 76 | ~~~~~~~~~~~~~~~~~~~~~~~~~~~ 77 | Create a build directory and run CMake to configure the build. The following example shows how to build 78 | 79 | .. code:: bash 80 | 81 | cd tvm 82 | rm -rf build && mkdir build && cd build 83 | # Specify the build configuration via CMake options 84 | cp ../cmake/config.cmake . 85 | 86 | Then you can configure the build by running CMake: 87 | 88 | .. code:: bash 89 | 90 | # controls default compilation flags (Candidates: Release, Debug, RelWithDebInfo) 91 | echo "set(CMAKE_BUILD_TYPE RelWithDebInfo)" >> config.cmake 92 | 93 | # LLVM is a must dependency 94 | echo "set(USE_LLVM \"llvm-config --ignore-libllvm --link-static\")" >> config.cmake 95 | echo "set(HIDE_PRIVATE_SYMBOLS ON)" >> config.cmake 96 | 97 | # GPU SDKs, turn on if needed 98 | echo "set(USE_CUDA OFF)" >> config.cmake 99 | echo "set(USE_METAL OFF)" >> config.cmake 100 | echo "set(USE_VULKAN OFF)" >> config.cmake 101 | echo "set(USE_OPENCL OFF)" >> config.cmake 102 | 103 | # cuBLAS, cuDNN, cutlass support, turn on if needed 104 | echo "set(USE_CUBLAS OFF)" >> config.cmake 105 | echo "set(USE_CUDNN OFF)" >> config.cmake 106 | echo "set(USE_CUTLASS OFF)" >> config.cmake 107 | 108 | .. note:: 109 | 110 | ``HIDE_PRIVATE_SYMBOLS`` is a configuration option that enables the ``-fvisibility=hidden`` flag. 111 | This flag helps prevent potential symbol conflicts between TVM and PyTorch. These conflicts arise 112 | due to the frameworks shipping LLVMs of different versions. 113 | 114 | Once config.cmake is edited accordingly, kick off build with the commands below: 115 | 116 | .. code:: bash 117 | 118 | cmake .. && cmake --build . --parallel $(nproc) 119 | 120 | A success build should produce ``libtvm`` and ``libtvm_runtime`` under ``build/`` directory. 121 | 122 | .. tabs :: 123 | 124 | .. code-tab :: bash Install via environment variable 125 | 126 | export TVM_HOME=/path/to/tvm 127 | export PYTHONPATH=$TVM_HOME/python:$PYTHONPATH 128 | 129 | .. code-tab :: bash Install via pip local project 130 | 131 | cd /path-to-tvm/python 132 | pip install -e . 133 | 134 | Step 4. Validate Installation 135 | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 136 | Please the following code to validate the TVM installation: 137 | 138 | .. code:: bash 139 | 140 | python -c "import tvm; print(tvm.__file__)" 141 | 142 | If the installation is successful, you should see the path to the TVM Python package printed out. 143 | 144 | Also, please verify you installed the versions 145 | 146 | .. code:: bash 147 | 148 | python -c "import tvm.relax; print(\"OK\")" 149 | 150 | If the installation is successful, you should see ``OK`` printed out. -------------------------------------------------------------------------------- /docs/get_started/overview.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | # Overview 19 | 20 | Apache TVM is a machine learning compilation framework, following the principle of **Python-first development** 21 | and **universal deployment**. It takes in pre-trained machine learning models, 22 | 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 23 | and more. 24 | 25 | ## Key Principle 26 | 27 | - **Python-first**: the optimization process is fully customizable in Python. 28 | It is easy to customize the optimization pipeline without recompiling the TVM stack. 29 | - **Composable**: the optimization process is composable. It is easy to compose 30 | new optimization passes, libraries and codegen to the existing pipeline. 31 | 32 | ## Key Goals 33 | 34 | - **Optimize** performance of ML workloads, composing libraries and codegen. 35 | - **Deploy** ML workloads to a diverse set of new environments, including new runtime and new hardware. 36 | - **Continuously improve and customize** ML deployment pipeline in Python by quickly customizing library dispatching, 37 | bringing in customized operators and code generation. 38 | 39 | ## Key Flow 40 | 41 | Here is a typical flow of using TVM to deploy a machine learning model. For a runnable example, please refer to 42 | [Quick Start](#quick_start) 43 | 44 | 1. **Import/construct an ML model** 45 | 46 | TVM supports importing models from various frameworks, such as PyTorch, TensorFlow for generic ML models. Meanwhile, we can create models directly using Relax frontend for scenarios of large language models. 47 | 48 | 2. **Perform composable optimization** transformations via `pipelines` 49 | 50 | The pipeline encapsulates a collection of transformations to achieve two goals: 51 | 52 | - **Graph Optimizations**: such as operator fusion, and layout rewrites. 53 | - **Tensor Program Optimization**: Map the operators to low-level implementations (both library or codegen) 54 | 55 | :::{note} 56 | The twos are goals but not the stages of the pipeline. The two optimizations are performed 57 | **at the same level**, or separately in two stages. 58 | ::: 59 | 60 | 3. **Build and universal deploy** 61 | 62 | 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. 63 | -------------------------------------------------------------------------------- /docs/get_started/tutorials/README.rst: -------------------------------------------------------------------------------- 1 | Get Started 2 | ----------- 3 | -------------------------------------------------------------------------------- /docs/get_started/tutorials/ir_module.py: -------------------------------------------------------------------------------- 1 | # Licensed to the Apache Software Foundation (ASF) under one 2 | # or more contributor license agreements. See the NOTICE file 3 | # distributed with this work for additional information 4 | # regarding copyright ownership. The ASF licenses this file 5 | # to you under the Apache License, Version 2.0 (the 6 | # "License"); you may not use this file except in compliance 7 | # with the License. You may obtain a copy of the License at 8 | # 9 | # http://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, 12 | # software distributed under the License is distributed on an 13 | # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | # KIND, either express or implied. See the License for the 15 | # specific language governing permissions and limitations 16 | # under the License. 17 | 18 | """ 19 | .. _ir_module: 20 | 21 | IRModule 22 | ======== 23 | This tutorial presents the core abstraction of Apache TVM Unity, the IRModule. 24 | The IRModule encompasses the **entirety** of the ML models, incorporating the 25 | computational graph, tensor programs, and potential calls to external libraries. 26 | 27 | .. contents:: Table of Contents 28 | :local: 29 | :depth: 1 30 | """ 31 | 32 | import numpy as np 33 | import tvm 34 | from tvm import relax 35 | 36 | ###################################################################### 37 | # Create IRModule 38 | # --------------- 39 | # IRModules can be initialized in various ways. We demonstrate a few of them 40 | # below. 41 | 42 | import torch 43 | from torch import fx, nn 44 | from tvm.relax.frontend.torch import from_fx 45 | 46 | ###################################################################### 47 | # Import from existing models 48 | # ~~~~~~~~~~~~~~~~~~~~~~~~~~~ 49 | # The most common way to initialize an IRModule is to import from an existing 50 | # model. Apache TVM Unity accommodates imports from a range of frameworks, 51 | # such as PyTorch and ONNX. This tutorial solely demonstrates the import process 52 | # from PyTorch. For other frameworks, please refer to the 🚧 corresponding tutorial. 53 | 54 | 55 | # Create a dummy model 56 | class TorchModel(nn.Module): 57 | def __init__(self): 58 | super(TorchModel, self).__init__() 59 | self.fc1 = nn.Linear(784, 256) 60 | self.relu1 = nn.ReLU() 61 | self.fc2 = nn.Linear(256, 10) 62 | 63 | def forward(self, x): 64 | x = self.fc1(x) 65 | x = self.relu1(x) 66 | x = self.fc2(x) 67 | return x 68 | 69 | 70 | # Give the input shape and data type 71 | input_info = [((1, 784), "float32")] 72 | 73 | # Convert the model to IRModule 74 | with torch.no_grad(): 75 | torch_fx_model = fx.symbolic_trace(TorchModel()) 76 | mod_from_torch = from_fx(torch_fx_model, input_info, keep_params_as_input=True) 77 | 78 | mod_from_torch, params_from_torch = relax.frontend.detach_params(mod_from_torch) 79 | # Print the IRModule 80 | mod_from_torch.show() 81 | 82 | ###################################################################### 83 | # Write with Relax NN Module 84 | # ~~~~~~~~~~~~~~~~~~~~~~~~~~ 85 | # Apache TVM Unity also provides a set of PyTorch-liked APIs, to help users 86 | # write the IRModule directly. For details, please refer to the 🚧 corresponding tutorial. 87 | 88 | from tvm.relax.frontend import nn 89 | 90 | 91 | class RelaxModel(nn.Module): 92 | def __init__(self): 93 | super(RelaxModel, self).__init__() 94 | self.fc1 = nn.Linear(784, 256) 95 | self.relu1 = nn.ReLU() 96 | self.fc2 = nn.Linear(256, 10) 97 | 98 | def forward(self, x): 99 | x = self.fc1(x) 100 | x = self.relu1(x) 101 | x = self.fc2(x) 102 | return x 103 | 104 | 105 | mod_from_relax, params_from_relax = RelaxModel().export_tvm( 106 | {"forward": {"x": nn.spec.Tensor((1, 784), "float32")}} 107 | ) 108 | mod_from_relax.show() 109 | 110 | ###################################################################### 111 | # Create via TVMScript 112 | # ~~~~~~~~~~~~~~~~~~~~ 113 | # TVMScript is a Python-based DSL for IRModules. We are able to 114 | # directly output the IRModule in the TVMScript syntax, or alternatively, 115 | # parse the TVMScript to obtain an IRModule. 116 | 117 | from tvm.script import ir as I 118 | from tvm.script import relax as R 119 | 120 | 121 | @I.ir_module 122 | class TVMScriptModule: 123 | @R.function 124 | def main( 125 | x: R.Tensor((1, 784), dtype="float32"), 126 | fc1_weight: R.Tensor((256, 784), dtype="float32"), 127 | fc1_bias: R.Tensor((256,), dtype="float32"), 128 | fc2_weight: R.Tensor((10, 256), dtype="float32"), 129 | fc2_bias: R.Tensor((10,), dtype="float32"), 130 | ) -> R.Tensor((1, 10), dtype="float32"): 131 | R.func_attr({"num_input": 1}) 132 | with R.dataflow(): 133 | permute_dims = R.permute_dims(fc1_weight, axes=None) 134 | matmul = R.matmul(x, permute_dims, out_dtype="void") 135 | add = R.add(matmul, fc1_bias) 136 | relu = R.nn.relu(add) 137 | permute_dims1 = R.permute_dims(fc2_weight, axes=None) 138 | matmul1 = R.matmul(relu, permute_dims1, out_dtype="void") 139 | add1 = R.add(matmul1, fc2_bias) 140 | gv = add1 141 | R.output(gv) 142 | return gv 143 | 144 | 145 | mod_from_script = TVMScriptModule 146 | mod_from_script.show() 147 | 148 | ###################################################################### 149 | # Attributes of an IRModule 150 | # ------------------------- 151 | # An IRModule is a collection of functions, indexed by GlobalVars. 152 | 153 | mod = mod_from_torch 154 | print(mod.get_global_vars()) 155 | 156 | ###################################################################### 157 | # We can access the functions in the IRModule by indexing with the GlobalVars 158 | # or their names 159 | 160 | # index by global var name 161 | print(mod["main"]) 162 | # index by global var, and checking they are the same function 163 | (gv,) = mod.get_global_vars() 164 | assert mod[gv] == mod["main"] 165 | 166 | ###################################################################### 167 | # Transformations on IRModules 168 | # ---------------------------- 169 | # Transformations are the import component of Apache TVM Unity. One transformation 170 | # takes in an IRModule and outputs another IRModule. We can apply a sequence of 171 | # transformations to an IRModule to obtain a new IRModule. That is the common way to 172 | # optimize a model. 173 | # 174 | # In this getting started tutorial, we only demonstrate how to apply transformations 175 | # to an IRModule. For details of each transformation, please refer to the 176 | # :ref:`Transformation API Reference ` 177 | 178 | ###################################################################### 179 | # We first apply **LegalizeOps** transformation to the IRModule. This transformation 180 | # will convert the Relax module into a mixed stage, with both Relax and TensorIR function 181 | # within the same module. Meanwhile, the Relax operators will be converted into ``call_tir``. 182 | 183 | mod = mod_from_torch 184 | mod = relax.transform.LegalizeOps()(mod) 185 | mod.show() 186 | 187 | ###################################################################### 188 | # After the transformation, there are much more functions inside the module. Let's print 189 | # the global vars again. 190 | 191 | print(mod.get_global_vars()) 192 | 193 | ###################################################################### 194 | # Next, Apache TVM Unity provides a set of default transformation pipelines for users, 195 | # to simplify the transformation process. We can then apply the default pipeline to the module. 196 | # The default **zero** pipeline contains very fundamental transformations, including: 197 | # 198 | # - **LegalizeOps**: This transform converts the Relax operators into `call_tir` functions 199 | # with the corresponding TensorIR Functions. After this transform, the IRModule will 200 | # contain both Relax functions and TensorIR functions. 201 | # - **AnnotateTIROpPattern**: This transform annotates the pattern of the TensorIR functions, 202 | # preparing them for subsequent operator fusion. 203 | # - **FoldConstant**: This pass performs constant folding, optimizing operations 204 | # involving constants. 205 | # - **FuseOps and FuseTIR**: These two passes work together to fuse operators based on the 206 | # patterns annotated in the previous step (AnnotateTIROpPattern). These passes transform 207 | # both Relax functions and TensorIR functions. 208 | # 209 | # .. note:: 210 | # 211 | # Here, we have applied **LegalizeOps** twice in the flow. The second time is useless but 212 | # harmless. 213 | # 214 | # Every passes can be duplicated in the flow, since we ensure the passes can handle all legal 215 | # IRModule inputs. This design can help users to construct their own pipeline. 216 | 217 | mod = relax.get_pipeline("zero")(mod) 218 | mod.show() 219 | 220 | ###################################################################### 221 | # Deploy the IRModule Universally 222 | # ------------------------------- 223 | # After the optimization, we can compile the model into a TVM runtime module. 224 | # Notably, Apache TVM Unity provides the ability of universal deployment, which means 225 | # we can deploy the same IRModule on different backends, including CPU, GPU, and other emerging 226 | # backends. 227 | # 228 | # Deploy on CPU 229 | # ~~~~~~~~~~~~~ 230 | # We can deploy the IRModule on CPU by specifying the target as ``llvm``. 231 | 232 | exec = relax.build(mod, target="llvm") 233 | dev = tvm.cpu() 234 | vm = relax.VirtualMachine(exec, dev) 235 | 236 | raw_data = np.random.rand(1, 784).astype("float32") 237 | data = tvm.nd.array(raw_data, dev) 238 | cpu_out = vm["main"](data, *params_from_torch["main"]).numpy() 239 | print(cpu_out) 240 | 241 | ###################################################################### 242 | # Deploy on GPU 243 | # ~~~~~~~~~~~~~ 244 | # Besides, CPU backend, we can also deploy the IRModule on GPU. GPU requires 245 | # programs containing extra information, such as the thread bindings and shared memory 246 | # allocations. We need a further transformation to generate the GPU programs. 247 | # 248 | # We use ``DLight`` to generate the GPU programs. In this tutorial, we won't go into 249 | # the details of ``DLight``. For more information, please refer to the 🚧 corresponding tutorial. 250 | # 251 | 252 | from tvm import dlight as dl 253 | 254 | with tvm.target.Target("cuda"): 255 | gpu_mod = dl.ApplyDefaultSchedule( 256 | dl.gpu.Matmul(), 257 | dl.gpu.Fallback(), 258 | )(mod) 259 | 260 | ###################################################################### 261 | # Now we can compile the IRModule on GPU, the similar way as we did on CPU. 262 | 263 | exec = relax.build(gpu_mod, target="cuda") 264 | dev = tvm.device("cuda", 0) 265 | vm = relax.VirtualMachine(exec, dev) 266 | # Need to allocate data and params on GPU device 267 | data = tvm.nd.array(raw_data, dev) 268 | gpu_params = [tvm.nd.array(p, dev) for p in params_from_torch["main"]] 269 | gpu_out = vm["main"](data, *gpu_params).numpy() 270 | print(gpu_out) 271 | 272 | # Check the correctness of the results 273 | assert np.allclose(cpu_out, gpu_out, atol=1e-3) 274 | 275 | ###################################################################### 276 | # Deploy on Other Backends 277 | # ~~~~~~~~~~~~~~~~~~~~~~~~ 278 | # Apache TVM Unity also supports other backends, such as different kinds of GPUs 279 | # (Metal, ROCm, Vulkan and OpenCL), different kinds of CPUs (x86, ARM), and other 280 | # emerging backends (e.g., WebAssembly). The deployment process is similar to the 281 | # GPU backend. 282 | -------------------------------------------------------------------------------- /docs/get_started/tutorials/quick_start.py: -------------------------------------------------------------------------------- 1 | # Licensed to the Apache Software Foundation (ASF) under one 2 | # or more contributor license agreements. See the NOTICE file 3 | # distributed with this work for additional information 4 | # regarding copyright ownership. The ASF licenses this file 5 | # to you under the Apache License, Version 2.0 (the 6 | # "License"); you may not use this file except in compliance 7 | # with the License. You may obtain a copy of the License at 8 | # 9 | # http://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, 12 | # software distributed under the License is distributed on an 13 | # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | # KIND, either express or implied. See the License for the 15 | # specific language governing permissions and limitations 16 | # under the License. 17 | 18 | """ 19 | .. _quick_start: 20 | 21 | Quick Start 22 | =========== 23 | 24 | This tutorial is for people who are new to Apache TVM. Taking an simple example 25 | to show how to use Apache TVM to compile a simple neural network. 26 | 27 | .. contents:: Table of Contents 28 | :local: 29 | :depth: 2 30 | 31 | """ 32 | 33 | ################################################################################ 34 | # Overview 35 | # -------- 36 | # Apache TVM is a machine learning compilation framework, following the principle of 37 | # **Python-first development** and **universal deployment**. It takes in pre-trained 38 | # machine learning models, compiles and generates deployable modules that can be embedded 39 | # and run everywhere. 40 | # Apache TVM also enables customizing optimization processes to introduce new optimizations, 41 | # libraries, codegen and more. 42 | # 43 | # Apache TVM can help to: 44 | # 45 | # - **Optimize** performance of ML workloads, composing libraries and codegen. 46 | # - **Deploy** ML workloads to a diverse set of new environments, including new runtime and new 47 | # hardware. 48 | # - **Continuously improve and customize** ML deployment pipeline in Python by quickly customizing 49 | # library dispatching, bringing in customized operators and code generation. 50 | 51 | ################################################################################ 52 | # Overall Flow 53 | # ------------ 54 | # Then we will show the overall flow of using Apache TVM to compile a neural network model, 55 | # showing how to optimize, deploy and run the model. 56 | # The overall flow is illustrated as the figure: 57 | # 58 | # .. figure:: ../../_static/img/overview.svg 59 | # :align: center 60 | # :width: 80% 61 | # 62 | # The overall flow consists of the following steps: 63 | # 64 | # - **Construct or Import a Model**: Construct a neural network model or import a pre-trained 65 | # model from other frameworks (e.g. PyTorch, ONNX), and create the TVM IRModule, which contains 66 | # all the information needed for compilation, including high-level Relax functions for 67 | # computational graph, and low-level TensorIR functions for tensor program. 68 | # - **Perform Composable Optimizations**: Perform a series of optimization transformations, 69 | # such as graph optimizations, tensor program optimizations, and library dispatching. 70 | # - **Build and Universal Deployment**: Build the optimized model to a deployable module to the 71 | # universal runtime, and execute it on different devices, such as CPU, GPU, or other accelerators. 72 | 73 | ################################################################################ 74 | # Construct or Import a Model 75 | # ~~~~~~~~~~~~~~~~~~~~~~~~~~~ 76 | # Before we get started, let's construct a neural network model first. 77 | # In this tutorial, to make things simple, we will defined a two-layer MLP networks 78 | # directly in this script with TVM Relax frontend, which is a similar API to PyTorch. 79 | # 80 | 81 | import tvm 82 | from tvm import relax 83 | from tvm.relax.frontend import nn 84 | 85 | 86 | class MLPModel(nn.Module): 87 | def __init__(self): 88 | super(MLPModel, self).__init__() 89 | self.fc1 = nn.Linear(784, 256) 90 | self.relu1 = nn.ReLU() 91 | self.fc2 = nn.Linear(256, 10) 92 | 93 | def forward(self, x): 94 | x = self.fc1(x) 95 | x = self.relu1(x) 96 | x = self.fc2(x) 97 | return x 98 | 99 | 100 | ################################################################################ 101 | # Then we can export the model to TVM IRModule, which is the central intermediate representation 102 | # in TVM. 103 | 104 | mod, param_spec = MLPModel().export_tvm( 105 | spec={"forward": {"x": nn.spec.Tensor((1, 784), "float32")}} 106 | ) 107 | mod.show() 108 | 109 | ################################################################################ 110 | # Perform Optimization Transformations 111 | # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 112 | # Apache TVM leverage ``pipeline`` to transform and optimize program. 113 | # The pipeline encapsulates a collection of transformation that gets two goals (at the same level): 114 | # 115 | # - **Model optimizations**: such as operator fusion, layout rewrites. 116 | # - **Tensor program optimization**: Map the operators to low-level implementations 117 | # (both library or codegen) 118 | # 119 | # .. note:: 120 | # The twos are goals but not the stages of the pipeline. The two optimizations are performed 121 | # **at the same level**, or separately in two stages. 122 | # 123 | # .. note:: 124 | # In this tutorial we only demonstrate the overall flow, by leverage ``zero`` optimization 125 | # pipeline, instead of optimizing for any specific target. 126 | 127 | mod = relax.get_pipeline("zero")(mod) 128 | 129 | 130 | ################################################################################ 131 | # Build and Universal Deployment 132 | # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 133 | # After the optimization, we can build the model to a deployable module and run it on 134 | # different devices. 135 | 136 | 137 | import numpy as np 138 | 139 | target = tvm.target.Target("llvm") 140 | ex = relax.build(mod, target) 141 | device = tvm.cpu() 142 | vm = relax.VirtualMachine(ex, device) 143 | data = np.random.rand(1, 784).astype("float32") 144 | tvm_data = tvm.nd.array(data, device=device) 145 | params = [np.random.rand(*param.shape).astype("float32") for _, param in param_spec] 146 | params = [tvm.nd.array(param, device=device) for param in params] 147 | print(vm["forward"](tvm_data, *params).numpy()) 148 | 149 | ################################################################################ 150 | # Our goal is to bring machine learning to the application with any language of interest, 151 | # with the minimum runtime support. 152 | # 153 | # - Each function in IRModule becomes a runnable function in the runtime. For example in LLM 154 | # cases, we can call ``prefill`` and ``decode`` functions directly. 155 | # 156 | # .. code-block:: Python 157 | # 158 | # prefill_logits = vm["prefill"](inputs, weight, kv_cache) 159 | # decoded_logits = vm["decode"](inputs, weight, kv_cache) 160 | # 161 | # - TVM runtime comes with native data structures, such as NDArray, can also have zero 162 | # copy exchange with existing ecosystem (DLPack exchange with PyTorch) 163 | # 164 | # .. code-block:: Python 165 | # 166 | # # Convert PyTorch tensor to TVM NDArray 167 | # x_tvm = tvm.nd.from_dlpack(x_torch.to_dlpack()) 168 | # # Convert TVM NDArray to PyTorch tensor 169 | # x_torch = torch.from_dlpack(x_tvm.to_dlpack()) 170 | # 171 | # - TVM runtime works in non-python environments, so it works on settings such as mobile 172 | # 173 | # .. code-block:: C++ 174 | # 175 | # // C++ snippet 176 | # runtime::Module vm = ex.GetFunction("load_executable")(); 177 | # vm.GetFunction("init")(...); 178 | # NDArray out = vm.GetFunction("prefill")(data, weight, kv_cache); 179 | # 180 | # .. code-block:: Java 181 | # 182 | # // Java snippet 183 | # Module vm = ex.getFunction("load_executable").invoke(); 184 | # vm.getFunction("init").pushArg(...).invoke; 185 | # NDArray out = vm.getFunction("prefill").pushArg(data).pushArg(weight).pushArg(kv_cache).invoke(); 186 | # 187 | 188 | ################################################################################ 189 | # Read next 190 | # --------- 191 | # This tutorial demonstrates the overall flow of using Apache TVM to compile a neural network model. 192 | # For more advanced or specific topics, please refer to the following tutorials: 193 | # 194 | # - If you want to convert a PyTorch model to TVM, please refer to TODOLINK 195 | # - If you want to optimize LLM models with TVM, please refer to TODOLINK 196 | # -------------------------------------------------------------------------------- /docs/index.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | 🚧 Apache TVM Unity Documentation 19 | ================================= 20 | 21 | TVM Unity is developed under Apache TVM's Github Repo in a separate `unity` branch. The documentation will be host here until it is merged into the main branch. 22 | 23 | .. toctree:: 24 | :maxdepth: 1 25 | :caption: Get Started 26 | 27 | get_started/overview.md 28 | get_started/install 29 | get_started/tutorials/quick_start 30 | get_started/tutorials/ir_module 31 | 32 | 33 | .. The Deep Dive content is comprehensive 34 | .. we maintain a ``maxdepth`` of 2 to display more information on the main page. 35 | 36 | .. toctree:: 37 | :maxdepth: 2 38 | :caption: Deep Dive 39 | 40 | deep_dive/tensor_ir/index 41 | 42 | .. toctree:: 43 | :maxdepth: 1 44 | :caption: References 45 | 46 | reference/api/index 47 | reference/publications 48 | 49 | -------------------------------------------------------------------------------- /docs/reference/api/dlight.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.dlight 19 | ---------- 20 | .. automodule:: tvm.dlight 21 | :members: 22 | :imported-members: 23 | -------------------------------------------------------------------------------- /docs/reference/api/error.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.error 19 | --------- 20 | .. automodule:: tvm.error 21 | :members: 22 | :imported-members: 23 | -------------------------------------------------------------------------------- /docs/reference/api/index.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | API Reference 19 | ============= 20 | 21 | .. toctree:: 22 | :maxdepth: 1 23 | :caption: tvm 24 | 25 | error 26 | ir 27 | instrument 28 | transform 29 | target 30 | 31 | .. toctree:: 32 | :maxdepth: 1 33 | :caption: tvm.runtime 34 | 35 | runtime/runtime 36 | runtime/ndarray 37 | runtime/relax_vm 38 | runtime/disco 39 | runtime/profiling 40 | 41 | .. toctree:: 42 | :maxdepth: 1 43 | :caption: tvm.te 44 | 45 | te 46 | topi 47 | 48 | .. toctree:: 49 | :maxdepth: 1 50 | :caption: tvm.tir 51 | 52 | tir/tir 53 | tir/analysis 54 | tir/schedule 55 | tir/stmt_functor 56 | tir/transform 57 | 58 | .. toctree:: 59 | :maxdepth: 1 60 | :caption: tvm.meta_schedule 61 | 62 | meta_schedule 63 | 64 | .. toctree:: 65 | :maxdepth: 1 66 | :caption: tvm.dlight 67 | 68 | dlight 69 | 70 | .. toctree:: 71 | :maxdepth: 1 72 | :caption: tvm.relax 73 | 74 | relax/relax 75 | relax/analysis 76 | relax/block_builder 77 | relax/frontend 78 | relax/op 79 | relax/transform 80 | 81 | .. toctree:: 82 | :maxdepth: 1 83 | :caption: Misc 84 | 85 | rpc 86 | -------------------------------------------------------------------------------- /docs/reference/api/instrument.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.instrument 19 | -------------- 20 | .. automodule:: tvm.instrument 21 | :members: 22 | :imported-members: 23 | -------------------------------------------------------------------------------- /docs/reference/api/ir.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.ir 19 | ------ 20 | .. automodule:: tvm.ir 21 | :members: 22 | :imported-members: 23 | -------------------------------------------------------------------------------- /docs/reference/api/meta_schedule.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.meta_schedule 19 | ----------------- 20 | .. automodule:: tvm.meta_schedule 21 | :members: 22 | :imported-members: 23 | -------------------------------------------------------------------------------- /docs/reference/api/relax/analysis.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.relax.analysis 19 | ------------------ 20 | .. automodule:: tvm.relax.analysis 21 | :members: 22 | :imported-members: 23 | 24 | -------------------------------------------------------------------------------- /docs/reference/api/relax/block_builder.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.relax.block_builder 19 | ----------------------- 20 | .. automodule:: tvm.relax.block_builder 21 | :members: 22 | -------------------------------------------------------------------------------- /docs/reference/api/relax/frontend.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.relax.frontend 19 | ------------------ 20 | .. automodule:: tvm.relax.frontend 21 | :members: 22 | :imported-members: 23 | 24 | tvm.relax.frontend.nn 25 | ********************* 26 | .. automodule:: tvm.relax.frontend.nn 27 | :members: 28 | :imported-members: 29 | :exclude-members: BlockBuilder 30 | :noindex: 31 | 32 | tvm.relax.frontend.onnx 33 | *********************** 34 | .. automodule:: tvm.relax.frontend.onnx 35 | :members: 36 | :imported-members: 37 | 38 | tvm.relax.frontend.stablehlo 39 | **************************** 40 | .. automodule:: tvm.relax.frontend.stablehlo 41 | :members: 42 | :imported-members: 43 | 44 | tvm.relax.frontend.torch 45 | ************************ 46 | .. automodule:: tvm.relax.frontend.torch 47 | :members: 48 | :imported-members: 49 | -------------------------------------------------------------------------------- /docs/reference/api/relax/op.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.relax.op 19 | ------------ 20 | 21 | tvm.relax.op 22 | ************ 23 | .. automodule:: tvm.relax.op 24 | :members: 25 | :imported-members: 26 | 27 | tvm.relax.op.nn 28 | *************** 29 | .. automodule:: tvm.relax.op.nn 30 | :members: 31 | :imported-members: 32 | 33 | tvm.relax.op.builtin 34 | ******************** 35 | .. automodule:: tvm.relax.op.builtin 36 | :members: 37 | :imported-members: 38 | 39 | tvm.relax.op.ccl 40 | **************** 41 | .. automodule:: tvm.relax.op.ccl 42 | :members: 43 | :imported-members: 44 | 45 | tvm.relax.op.distributed 46 | ************************ 47 | .. automodule:: tvm.relax.op.distributed 48 | :members: 49 | :imported-members: 50 | 51 | tvm.relax.op.grad 52 | ***************** 53 | .. automodule:: tvm.relax.op.grad 54 | :members: 55 | :imported-members: 56 | 57 | tvm.relax.op.image 58 | ****************** 59 | .. automodule:: tvm.relax.op.image 60 | :members: 61 | :imported-members: 62 | 63 | tvm.relax.op.memory 64 | ******************* 65 | .. automodule:: tvm.relax.op.memory 66 | :members: 67 | :imported-members: 68 | 69 | tvm.relax.op.op_attrs 70 | ********************* 71 | .. automodule:: tvm.relax.op.op_attrs 72 | :members: 73 | 74 | -------------------------------------------------------------------------------- /docs/reference/api/relax/relax.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.relax 19 | --------- 20 | .. automodule:: tvm.relax 21 | :members: 22 | :imported-members: 23 | :exclude-members: BlockBuilder, Span, GlobalVar, SourceName, TupleType, Type, FuncType 24 | -------------------------------------------------------------------------------- /docs/reference/api/relax/transform.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | .. _api-relax-transformation: 19 | 20 | tvm.relax.transform 21 | ------------------- 22 | .. automodule:: tvm.relax.transform 23 | :members: 24 | :imported-members: 25 | -------------------------------------------------------------------------------- /docs/reference/api/rpc.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.rpc 19 | ------- 20 | .. automodule:: tvm.rpc 21 | :members: 22 | :imported-members: 23 | -------------------------------------------------------------------------------- /docs/reference/api/runtime/disco.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.runtime.disco 19 | ----------------- 20 | .. automodule:: tvm.runtime.disco 21 | :members: 22 | :imported-members: -------------------------------------------------------------------------------- /docs/reference/api/runtime/ndarray.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.runtime.ndarray 19 | ------------------- 20 | .. automodule:: tvm.runtime.ndarray 21 | :members: 22 | -------------------------------------------------------------------------------- /docs/reference/api/runtime/profiling.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.runtime.profiling 19 | --------------------- 20 | .. automodule:: tvm.runtime.profiling 21 | :members: -------------------------------------------------------------------------------- /docs/reference/api/runtime/relax_vm.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.runtime.relax_vm 19 | -------------------- 20 | .. automodule:: tvm.runtime.relax_vm 21 | :members: -------------------------------------------------------------------------------- /docs/reference/api/runtime/runtime.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.runtime 19 | ----------- 20 | .. automodule:: tvm.runtime 21 | :members: 22 | :imported-members: 23 | :exclude-members: NDArray 24 | -------------------------------------------------------------------------------- /docs/reference/api/target.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.target 19 | ---------- 20 | .. automodule:: tvm.target 21 | :members: 22 | :imported-members: 23 | -------------------------------------------------------------------------------- /docs/reference/api/te.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.te 19 | ------ 20 | .. Exclude the ops imported from tir. 21 | 22 | .. automodule:: tvm.te 23 | :members: 24 | :imported-members: 25 | :exclude-members: Stage, Schedule 26 | 27 | -------------------------------------------------------------------------------- /docs/reference/api/tir/analysis.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.tir.analysis 19 | ---------------- 20 | .. automodule:: tvm.tir.analysis.analysis 21 | :members: 22 | -------------------------------------------------------------------------------- /docs/reference/api/tir/schedule.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.tir.schedule 19 | ----------------- 20 | .. automodule:: tvm.tir.schedule 21 | :members: 22 | :imported-members: 23 | -------------------------------------------------------------------------------- /docs/reference/api/tir/stmt_functor.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.tir.stmt_functor 19 | -------------------- 20 | .. automodule:: tvm.tir.stmt_functor 21 | :members: 22 | -------------------------------------------------------------------------------- /docs/reference/api/tir/tir.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.tir 19 | ------- 20 | .. automodule:: tvm.tir 21 | :members: 22 | :imported-members: 23 | :exclude-members: PrimExpr, const, StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError 24 | 25 | -------------------------------------------------------------------------------- /docs/reference/api/tir/transform.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | 19 | tvm.tir.transform 20 | ----------------- 21 | .. automodule:: tvm.tir.transform 22 | :members: 23 | :imported-members: 24 | -------------------------------------------------------------------------------- /docs/reference/api/topi.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.topi 19 | -------- 20 | 21 | tvm.topi 22 | ~~~~~~~~ 23 | .. automodule:: tvm.topi 24 | :members: 25 | :imported-members: 26 | :exclude-members: PrimExpr, Cast 27 | 28 | tvm.topi.nn 29 | ~~~~~~~~~~~ 30 | .. automodule:: tvm.topi.nn 31 | :members: 32 | :imported-members: 33 | 34 | tvm.topi.image 35 | ~~~~~~~~~~~~~~ 36 | .. automodule:: tvm.topi.image 37 | :members: 38 | :imported-members: 39 | 40 | tvm.topi.sparse 41 | ~~~~~~~~~~~~~~~ 42 | .. automodule:: tvm.topi.sparse 43 | :members: 44 | :imported-members: 45 | -------------------------------------------------------------------------------- /docs/reference/api/transform.rst: -------------------------------------------------------------------------------- 1 | .. Licensed to the Apache Software Foundation (ASF) under one 2 | or more contributor license agreements. See the NOTICE file 3 | distributed with this work for additional information 4 | regarding copyright ownership. The ASF licenses this file 5 | to you under the Apache License, Version 2.0 (the 6 | "License"); you may not use this file except in compliance 7 | with the License. You may obtain a copy of the License at 8 | 9 | .. http://www.apache.org/licenses/LICENSE-2.0 10 | 11 | .. Unless required by applicable law or agreed to in writing, 12 | software distributed under the License is distributed on an 13 | "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 | KIND, either express or implied. See the License for the 15 | specific language governing permissions and limitations 16 | under the License. 17 | 18 | tvm.transform 19 | ------------- 20 | .. automodule:: tvm.transform 21 | :members: 22 | :imported-members: 23 | -------------------------------------------------------------------------------- /docs/reference/publications.rst: -------------------------------------------------------------------------------- 1 | .. _Publications: 2 | 3 | Publications 4 | ============ 5 | 6 | The underlying compiler techniques employed by Apache TVM are outlined in the following papers: 7 | 8 | 9 | .. code:: bibtex 10 | 11 | @inproceedings{tensorir, 12 | author = {Feng, Siyuan and Hou, Bohan and Jin, Hongyi and Lin, Wuwei and Shao, Junru and Lai, Ruihang and Ye, Zihao and Zheng, Lianmin and Yu, Cody Hao and Yu, Yong and Chen, Tianqi}, 13 | title = {TensorIR: An Abstraction for Automatic Tensorized Program Optimization}, 14 | year = {2023}, 15 | isbn = {9781450399166}, 16 | publisher = {Association for Computing Machinery}, 17 | address = {New York, NY, USA}, 18 | url = {https://doi.org/10.1145/3575693.3576933}, 19 | doi = {10.1145/3575693.3576933}, 20 | booktitle = {Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2}, 21 | pages = {804--817}, 22 | numpages = {14}, 23 | keywords = {Tensor Computation, Machine Learning Compiler, Deep Neural Network}, 24 | location = {Vancouver, BC, Canada}, 25 | series = {ASPLOS 2023} 26 | } 27 | 28 | @inproceedings{metaschedule, 29 | author = {Shao, Junru and Zhou, Xiyou and Feng, Siyuan and Hou, Bohan and Lai, Ruihang and Jin, Hongyi and Lin, Wuwei and Masuda, Masahiro and Yu, Cody Hao and Chen, Tianqi}, 30 | booktitle = {Advances in Neural Information Processing Systems}, 31 | editor = {S. Koyejo and S. Mohamed and A. Agarwal and D. Belgrave and K. Cho and A. Oh}, 32 | pages = {35783--35796}, 33 | publisher = {Curran Associates, Inc.}, 34 | title = {Tensor Program Optimization with Probabilistic Programs}, 35 | url = {https://proceedings.neurips.cc/paper_files/paper/2022/file/e894eafae43e68b4c8dfdacf742bcbf3-Paper-Conference.pdf}, 36 | volume = {35}, 37 | year = {2022} 38 | } 39 | 40 | @inproceedings{tvm, 41 | author = {Tianqi Chen and Thierry Moreau and Ziheng Jiang and Lianmin Zheng and Eddie Yan and Haichen Shen and Meghan Cowan and Leyuan Wang and Yuwei Hu and Luis Ceze and Carlos Guestrin and Arvind Krishnamurthy}, 42 | title = {{TVM}: An Automated {End-to-End} Optimizing Compiler for Deep Learning}, 43 | booktitle = {13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)}, 44 | year = {2018}, 45 | isbn = {978-1-939133-08-3}, 46 | address = {Carlsbad, CA}, 47 | pages = {578--594}, 48 | url = {https://www.usenix.org/conference/osdi18/presentation/chen}, 49 | publisher = {USENIX Association}, 50 | month = oct, 51 | } -------------------------------------------------------------------------------- /pylintrc: -------------------------------------------------------------------------------- 1 | disable=missing-docstring, -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | sphinx-tabs==3.4.5 2 | sphinx-rtd-theme==2.0.0 3 | sphinx==7.0.0 4 | sphinx-toolbox==3.5.0 5 | tlcpack-sphinx-addon==0.2.3 6 | sphinxcontrib_httpdomain==1.8.1 7 | sphinxcontrib-napoleon==0.7 8 | sphinx_gallery==0.15.0 9 | myst-parser==2.0.0 10 | matplotlib 11 | --pre 12 | --find-links https://mlc.ai/wheels 13 | mlc-ai-nightly-cu121 --------------------------------------------------------------------------------