You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/15 14:14:05 UTC
(tvm) branch tirx/standalone-package-cleanup updated (6787851cdd -> 58d620f53f)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a change to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
omit 6787851cdd [FIX][Docs] Include extracted core headers in C++ documentation
omit 8887114dc6 [FIX][Docs] Resolve extracted TIRx API namespaces
omit 10c1daef2d [FIX][Testing] Probe the loaded NCCL runtime capability
omit c90cc6d7e8 [FIX][IR] Keep the name-supply template constructor local to consumers
omit 97539acf4e [FIX][Runtime] Release the GIL while connecting Disco socket sessions
omit 6932e6edde [CI] Dump Python stacks when a test stops making progress
omit 8162368043 [FIX][CMake] Keep private visibility on owned library targets
omit a644010c08 [FIX][Metal] Restore the compiler callback after its owning test
omit a6fe2d8a27 [FIX][CMake] Normalize the packaged Z3 license path on Windows
omit 418ad1a5fd [FIX][CUDA] Bound embedded launch source literals for MSVC
omit 652020f489 [FIX][Python] Retain parent version and external-module headers
omit f9beb1517e [FIX][CUTLASS] Propagate owning runtime headers to the vendor binding
omit a5caabad43 [FIX][Web] Locate WebAssembly assets in the parent package
omit c35f1001db [FIX][Packaging] Require the FFI release with typed reflection attributes
omit e4eeae99f1 [STYLE] Format relocated include paths and printer imports
omit bf51c1d633 [FIX][Lint] Apply repository lint ownership to the extracted tree
omit 2453a3e630 [FIX][Packaging] Limit source distributions to their owned Python packages
omit ded2e085a6 [FIX][CI] Build both native partitions with the pinned FFI dependency
omit 320be3cbdc [FIX][Script] Synchronize the generated TIRx namespace reference
omit c47a1857dd [REFACTOR][Python] Format imports after the core namespace split
omit abe1148879 [FIX][Tests] Keep mixed printer coverage with parent dialect registration
omit ed86ccd4dc [REFACTOR][CMake] Build and install independent core and parent libraries
omit 38ffb38b7f [REFACTOR][Script] Emit imports from each dialect owning package
omit 19ced11be2 [REFACTOR][Native] Export core declarations and retain parent runtime consumers
omit 6bcdd9e7ac [REFACTOR][CI] Exercise both partitions and publish standalone core wheels
omit 9b74c5b87d [REFACTOR][Python] Load core and parent extensions from their owning packages
omit d0c5a163c5 [REFACTOR][Tests] Preserve fixture and dialect ownership across the package split
omit 3cdb284ca9 [REFACTOR][Packaging] Package TIRx with an independent version and dependency boundary
omit a82605f932 [REFACTOR][Python] Rename extracted core imports to tvm_tirx
omit baab6e3ee5 [REFACTOR][TIRx] Relocate standalone core sources and existing tests
add 66394cd89f [Fix][Relax][Frontend][Torch] Validate `dim` in the `squeeze` converter (#20322)
add c14bc80e8b [REFACTOR][CI] Share native Jenkins pipelines and organize CI helpers (#20735)
add aa34ac394f [CI] Require approving reviews without mandatory status checks (#20743)
new 8fe02db438 [REFACTOR][TIRx] Relocate standalone core sources and existing tests
new 08fadc36f8 [REFACTOR][Python] Rename extracted core imports to tvm_tirx
new 3a7bc102bf [REFACTOR][Packaging] Package TIRx with an independent version and dependency boundary
new 3fdcc3f579 [REFACTOR][Tests] Preserve fixture and dialect ownership across the package split
new 0bb1535c9c [REFACTOR][Python] Load core and parent extensions from their owning packages
new 350ea906a3 [REFACTOR][CI] Exercise both partitions and publish standalone core wheels
new 9f9c9fe550 [REFACTOR][Native] Export core declarations and retain parent runtime consumers
new 72ac640cf4 [REFACTOR][Script] Emit imports from each dialect owning package
new 760fae90d2 [REFACTOR][CMake] Build and install independent core and parent libraries
new 6c977552d0 [FIX][Tests] Keep mixed printer coverage with parent dialect registration
new 1614b98573 [REFACTOR][Python] Format imports after the core namespace split
new f393d43497 [FIX][Script] Synchronize the generated TIRx namespace reference
new fc53ebd4b7 [FIX][CI] Build both native partitions with the pinned FFI dependency
new a0c7813387 [FIX][Packaging] Limit source distributions to their owned Python packages
new 64ac09abda [FIX][Lint] Apply repository lint ownership to the extracted tree
new e7fff21d10 [STYLE] Format relocated include paths and printer imports
new 8c3a0874fc [FIX][Packaging] Require the FFI release with typed reflection attributes
new bb9ec8c2e1 [FIX][Web] Locate WebAssembly assets in the parent package
new ea1c75476d [FIX][CUTLASS] Propagate owning runtime headers to the vendor binding
new 739ffe23c1 [FIX][Python] Retain parent version and external-module headers
new 144ecab59d [FIX][CUDA] Bound embedded launch source literals for MSVC
new 8fb70273cd [FIX][CMake] Normalize the packaged Z3 license path on Windows
new 1d8feb6857 [FIX][Metal] Restore the compiler callback after its owning test
new ee7de53ff0 [FIX][CMake] Keep private visibility on owned library targets
new 5cd26dc7db [CI] Dump Python stacks when a test stops making progress
new 33920a5adc [FIX][Runtime] Release the GIL while connecting Disco socket sessions
new c5f78007c5 [FIX][IR] Keep the name-supply template constructor local to consumers
new 04f31ef669 [FIX][Testing] Probe the loaded NCCL runtime capability
new 78d2f04415 [FIX][Docs] Resolve extracted TIRx API namespaces
new dfb083b431 [FIX][Docs] Include extracted core headers in C++ documentation
new 0043f3ae27 [FIX][Python] Resolve the incoming squeeze regression through the core namespace
new 674b3ae17d [STYLE][CI] Format the native build dependency command
new 58d620f53f [FIX][Docs] Point wheel guidance to the standalone core configuration
This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version. This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:
* -- * -- B -- O -- O -- O (6787851cdd)
\
N -- N -- N refs/heads/tirx/standalone-package-cleanup (58d620f53f)
You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.
Any revisions marked "omit" are not gone; other references still
refer to them. Any revisions marked "discard" are gone forever.
The 33 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.asf.yaml | 10 +-
.github/actions/setup/action.yml | 6 +-
.github/workflows/nightly_docker_update.yml | 24 -
.github/workflows/publish_wheel.yml | 6 +-
.github/workflows/tvmbot.yml | 2 +-
ci/README.md | 16 +-
ci/{scripts => }/github/__init__.py | 0
ci/{scripts/jenkins => github}/git_utils.py | 0
ci/{scripts => }/github/github_tvmbot.py | 12 +-
ci/jenkins/README.md | 65 +-
ci/jenkins/data.py | 6 +-
ci/jenkins/generate.py | 226 -------
ci/jenkins/generated/arm_jenkinsfile.groovy | 491 +--------------
.../{wasm_jenkinsfile.groovy => common.groovy} | 314 ++++------
ci/jenkins/generated/cpu_jenkinsfile.groovy | 557 +----------------
ci/jenkins/generated/docker_jenkinsfile.groovy | 661 ++------------------
ci/jenkins/generated/gpu_jenkinsfile.groovy | 677 ++-------------------
ci/jenkins/generated/wasm_jenkinsfile.groovy | 492 +--------------
.../scripts}/determine_docker_images.py | 7 +-
.../scripts}/git_change_docker.sh | 0
.../jenkins => jenkins/scripts}/git_skip_ci.py | 51 +-
.../scripts}/git_skip_ci_globs.py | 2 +-
.../jenkins => jenkins/scripts}/http_utils.py | 0
ci/{scripts/jenkins => jenkins/scripts}/retry.sh | 0
ci/{scripts/jenkins => jenkins/scripts}/s3.py | 45 +-
.../scripts}/should_rebuild_docker.py | 19 +-
.../jenkins/scripts/show_node_info.sh | 0
ci/jenkins/templates/arm_jenkinsfile.groovy.j2 | 35 --
ci/jenkins/templates/cpu_jenkinsfile.groovy.j2 | 52 --
ci/jenkins/templates/docker_jenkinsfile.groovy.j2 | 202 ------
ci/jenkins/templates/gpu_jenkinsfile.groovy.j2 | 144 -----
ci/jenkins/templates/utils/Build.groovy.j2 | 28 -
ci/jenkins/templates/utils/Prepare.groovy.j2 | 276 ---------
ci/jenkins/templates/utils/Test.groovy.j2 | 6 -
ci/jenkins/templates/utils/base.groovy.j2 | 111 ----
ci/jenkins/templates/utils/macros.j2 | 169 -----
ci/jenkins/templates/wasm_jenkinsfile.groovy.j2 | 39 --
ci/scripts/jenkins/open_docker_update_pr.py | 214 -------
ci/{scripts/package => wheel}/README.md | 2 +-
.../package => wheel}/build-environment.yaml | 0
.../manylinux_build_libtvm_runtime_cuda.sh | 2 +-
.../windows_build_libtvm_runtime_cuda.bat | 4 +-
docker/bash.sh | 2 +-
docker/dev_common.sh | 2 +-
docs/contribute/ci.rst | 2 +-
docs/contribute/testing.rst | 2 +-
pyproject.toml | 2 +-
.../frontend/torch/base_fx_graph_translator.py | 22 +-
tests/python/ci/test_ci.py | 194 ++----
tests/python/ci/test_jenkins_helpers.py | 149 +++++
tests/python/ci/test_tvmbot.py | 6 +-
tests/python/ci/test_utils.py | 4 +-
tests/python/relax/test_frontend_from_fx.py | 36 ++
tests/scripts/release/gather_prs.py | 6 +-
tests/scripts/release/make_notes.py | 5 -
tests/scripts/setup-pytest-env.sh | 2 +-
tests/scripts/task_build.py | 70 ++-
57 files changed, 739 insertions(+), 4738 deletions(-)
delete mode 100644 .github/workflows/nightly_docker_update.yml
copy ci/{scripts => }/github/__init__.py (100%)
copy ci/{scripts/jenkins => github}/git_utils.py (100%)
copy ci/{scripts => }/github/github_tvmbot.py (98%)
delete mode 100644 ci/jenkins/generate.py
copy ci/jenkins/generated/{wasm_jenkinsfile.groovy => common.groovy} (54%)
copy ci/{scripts/jenkins => jenkins/scripts}/determine_docker_images.py (96%)
copy ci/{scripts/jenkins => jenkins/scripts}/git_change_docker.sh (100%)
copy ci/{scripts/jenkins => jenkins/scripts}/git_skip_ci.py (56%)
copy ci/{scripts/jenkins => jenkins/scripts}/git_skip_ci_globs.py (98%)
copy ci/{scripts/jenkins => jenkins/scripts}/http_utils.py (100%)
copy ci/{scripts/jenkins => jenkins/scripts}/retry.sh (100%)
copy ci/{scripts/jenkins => jenkins/scripts}/s3.py (83%)
copy ci/{scripts/jenkins => jenkins/scripts}/should_rebuild_docker.py (91%)
copy tests/scripts/task_show_node_info.sh => ci/jenkins/scripts/show_node_info.sh (100%)
delete mode 100644 ci/jenkins/templates/arm_jenkinsfile.groovy.j2
delete mode 100644 ci/jenkins/templates/cpu_jenkinsfile.groovy.j2
delete mode 100644 ci/jenkins/templates/docker_jenkinsfile.groovy.j2
delete mode 100644 ci/jenkins/templates/gpu_jenkinsfile.groovy.j2
delete mode 100644 ci/jenkins/templates/utils/Build.groovy.j2
delete mode 100644 ci/jenkins/templates/utils/Prepare.groovy.j2
delete mode 100644 ci/jenkins/templates/utils/Test.groovy.j2
delete mode 100644 ci/jenkins/templates/utils/base.groovy.j2
delete mode 100644 ci/jenkins/templates/utils/macros.j2
delete mode 100644 ci/jenkins/templates/wasm_jenkinsfile.groovy.j2
delete mode 100755 ci/scripts/jenkins/open_docker_update_pr.py
copy ci/{scripts/package => wheel}/README.md (94%)
copy ci/{scripts/package => wheel}/build-environment.yaml (100%)
copy ci/{scripts/package => wheel}/manylinux_build_libtvm_runtime_cuda.sh (97%)
copy ci/{scripts/package => wheel}/windows_build_libtvm_runtime_cuda.bat (97%)
create mode 100644 tests/python/ci/test_jenkins_helpers.py
(tvm) 30/33: [FIX][Docs] Include extracted core headers in C++ documentation
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit dfb083b4311092d44b34994156a6ea5bb03aaa40
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 10:03:12 2026 +0000
[FIX][Docs] Include extracted core headers in C++ documentation
Include both the parent and standalone core header roots in the existing Doxygen input inventory.
---
docs/Doxyfile | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/docs/Doxyfile b/docs/Doxyfile
index 42d9ea3d3f..bd7a36495b 100644
--- a/docs/Doxyfile
+++ b/docs/Doxyfile
@@ -764,7 +764,7 @@ WARN_LOGFILE =
# spaces.
# Note: If this tag is empty the current directory is searched.
-INPUT = include/tvm
+INPUT = include/tvm tirx/include/tvm
# This tag can be used to specify the character encoding of the source files
# that doxygen parses. Internally doxygen uses the UTF-8 encoding. Doxygen uses
(tvm) 01/33: [REFACTOR][TIRx] Relocate standalone core sources and existing tests
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 8fe02db4380ac498bf1a5370bd7b12bec6193545
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 00:46:29 2026 +0000
[REFACTOR][TIRx] Relocate standalone core sources and existing tests
---
{3rdparty => tirx/3rdparty}/compiler-rt/builtin_fp16.h | 0
{cmake => tirx/cmake}/modules/CUDA.cmake | 0
{cmake => tirx/cmake}/modules/ClangFlags.cmake | 0
{cmake => tirx/cmake}/modules/Git.cmake | 0
{cmake => tirx/cmake}/modules/LLVM.cmake | 0
{cmake => tirx/cmake}/modules/Logging.cmake | 0
{cmake => tirx/cmake}/modules/Metal.cmake | 0
{cmake => tirx/cmake}/modules/OpenCL.cmake | 0
{cmake => tirx/cmake}/modules/OpenMP.cmake | 0
{cmake => tirx/cmake}/modules/ROCM.cmake | 0
{cmake => tirx/cmake}/modules/Vulkan.cmake | 0
{cmake => tirx/cmake}/modules/contrib/Z3.cmake | 0
{cmake => tirx/cmake}/utils/CCache.cmake | 0
{cmake => tirx/cmake}/utils/FindCUDA.cmake | 0
{cmake => tirx/cmake}/utils/FindLLVM.cmake | 0
{cmake => tirx/cmake}/utils/FindOpenCL.cmake | 0
{cmake => tirx/cmake}/utils/FindROCM.cmake | 0
{cmake => tirx/cmake}/utils/FindVulkan.cmake | 0
{cmake => tirx/cmake}/utils/Library.cmake | 0
{cmake => tirx/cmake}/utils/Linker.cmake | 0
{cmake => tirx/cmake}/utils/Summary.cmake | 0
{cmake => tirx/cmake}/utils/Utils.cmake | 0
{include => tirx/include}/tvm/backend/cuda/attr.h | 0
{include => tirx/include}/tvm/backend/cuda/op/iket.h | 0
{include => tirx/include}/tvm/backend/cuda/op/tcgen05.h | 0
{include => tirx/include}/tvm/backend/cuda/op/tensormap.h | 0
{include => tirx/include}/tvm/backend/cuda/op_attr_types.h | 0
{include => tirx/include}/tvm/backend/metal/op/cooperative_tensor.h | 0
{include => tirx/include}/tvm/backend/opencl/op/memory.h | 0
{include => tirx/include}/tvm/codegen/codegen.h | 0
{include => tirx/include}/tvm/codegen/intrin_rule.h | 0
{include => tirx/include}/tvm/ir/analysis.h | 0
{include => tirx/include}/tvm/ir/attrs.h | 0
{include => tirx/include}/tvm/ir/base_expr.h | 0
{include => tirx/include}/tvm/ir/cow.h | 0
{include => tirx/include}/tvm/ir/env_func.h | 0
{include => tirx/include}/tvm/ir/expr.h | 0
{include => tirx/include}/tvm/ir/expr_functor.h | 0
{include => tirx/include}/tvm/ir/function.h | 0
{include => tirx/include}/tvm/ir/global_info.h | 0
{include => tirx/include}/tvm/ir/location.h | 0
{include => tirx/include}/tvm/ir/module.h | 0
{include => tirx/include}/tvm/ir/object_functor.h | 0
{include => tirx/include}/tvm/ir/op.h | 0
{include => tirx/include}/tvm/ir/prim/expr.h | 0
{include => tirx/include}/tvm/ir/prim/op.h | 0
{include => tirx/include}/tvm/ir/prim/vector_expr.h | 0
{include => tirx/include}/tvm/ir/scope_stack.h | 0
{include => tirx/include}/tvm/ir/stmt.h | 0
{include => tirx/include}/tvm/ir/transform.h | 0
{include => tirx/include}/tvm/ir/type.h | 0
{include => tirx/include}/tvm/ir/unique_name_supply.h | 0
{include => tirx/include}/tvm/ir/with_context.h | 0
{include => tirx/include}/tvm/runtime/base.h | 0
{include => tirx/include}/tvm/runtime/c_backend_api.h | 0
{include => tirx/include}/tvm/runtime/device_api.h | 0
{include => tirx/include}/tvm/runtime/logging.h | 0
{include => tirx/include}/tvm/runtime/memory_scope.h | 0
{include => tirx/include}/tvm/runtime/tensor.h | 0
{include => tirx/include}/tvm/runtime/timer.h | 0
{include => tirx/include}/tvm/script/ir_builder/base.h | 0
{include => tirx/include}/tvm/script/ir_builder/frame.h | 0
{include => tirx/include}/tvm/script/ir_builder/ir.h | 0
{include => tirx/include}/tvm/script/printer/config.h | 0
{include => tirx/include}/tvm/script/printer/doc.h | 0
{include => tirx/include}/tvm/script/printer/doc_translator.h | 0
{include => tirx/include}/tvm/script/printer/printer.h | 0
{include => tirx/include}/tvm/script/printer/support.h | 0
{include => tirx/include}/tvm/support/base64.h | 0
{include => tirx/include}/tvm/support/bytes_io.h | 0
{include => tirx/include}/tvm/support/cuda/nvtx.h | 0
{include => tirx/include}/tvm/support/env.h | 0
{include => tirx/include}/tvm/support/file_utils.h | 0
{include => tirx/include}/tvm/support/gpu_thread_tag.h | 0
{include => tirx/include}/tvm/support/io.h | 0
{include => tirx/include}/tvm/support/ordered_map.h | 0
{include => tirx/include}/tvm/support/serializer.h | 0
{include => tirx/include}/tvm/support/str_escape.h | 0
{include => tirx/include}/tvm/sym/analyzer.h | 0
{include => tirx/include}/tvm/sym/bound.h | 0
{include => tirx/include}/tvm/sym/int_set.h | 0
{include => tirx/include}/tvm/sym/iter_affine_map.h | 0
{include => tirx/include}/tvm/sym/pattern.h | 0
{include => tirx/include}/tvm/target/config_schema.h | 0
{include => tirx/include}/tvm/target/target.h | 0
{include => tirx/include}/tvm/target/target_kind.h | 0
{include => tirx/include}/tvm/target/target_tag_registry.h | 0
{include => tirx/include}/tvm/tirx/analysis.h | 0
{include => tirx/include}/tvm/tirx/expr.h | 0
{include => tirx/include}/tvm/tirx/function.h | 0
{include => tirx/include}/tvm/tirx/index_map.h | 0
{include => tirx/include}/tvm/tirx/layout.h | 0
{include => tirx/include}/tvm/tirx/op/abi.h | 0
{include => tirx/include}/tvm/tirx/op/annotation.h | 0
{include => tirx/include}/tvm/tirx/op/cpu.h | 0
{include => tirx/include}/tvm/tirx/op/gpu.h | 0
{include => tirx/include}/tvm/tirx/op/math.h | 0
{include => tirx/include}/tvm/tirx/op/memory.h | 0
{include => tirx/include}/tvm/tirx/op/region.h | 0
{include => tirx/include}/tvm/tirx/op/vector.h | 0
{include => tirx/include}/tvm/tirx/op_attr_types.h | 0
{include => tirx/include}/tvm/tirx/script/ir_builder/frame.h | 0
{include => tirx/include}/tvm/tirx/script/ir_builder/ir.h | 0
{include => tirx/include}/tvm/tirx/stmt.h | 0
{include => tirx/include}/tvm/tirx/stmt_functor.h | 0
{include => tirx/include}/tvm/tirx/stmt_functor_with_analyzer.h | 0
{include => tirx/include}/tvm/tirx/tile/context.h | 0
{include => tirx/include}/tvm/tirx/tile/dispatch.h | 0
{include => tirx/include}/tvm/tirx/tile/op.h | 0
{include => tirx/include}/tvm/tirx/tile/thread_scope.h | 0
{include => tirx/include}/tvm/tirx/transform.h | 0
{include => tirx/include}/tvm/tirx/type.h | 0
{licenses => tirx/licenses}/LICENSE.builtin_fp16.txt | 0
{python/tvm => tirx/python/tvm_tirx}/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/_autoload_backends.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/adreno/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/adreno/target_tags.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/config.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/axis.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/backend_config.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/codegen/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/codegen/header.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/codegen/registry.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/codegen/schema.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/codegen/types.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/codegen/utils.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/compiler.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/cpp/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/cpp/asm.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/cpp/builtins.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/cpp/descriptors.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/cpp/instrument.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/cpp/nvshmem.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/host.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/iket.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/lang/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/lang/alloc_pool.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/lang/clc.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/lang/pipeline.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/lang/tile_scheduler.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/launch/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/launch/_generated.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/launch/_impl.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/launch/generate.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/launch/table.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/layout.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/op.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/ptx/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/ptx/engine.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/ptx/gen_coverage.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/ptx/gen_helpers.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/ptx/gen_stubs.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/ptx/render.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/ptx/table.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/script.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/target_tags.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/thread_scope.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/tile/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/tile/exec_context.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/tile/instructions.py | 0
.../tvm => tirx/python/tvm_tirx}/backend/cuda/tile/patterns/__init__.py | 0
.../tvm => tirx/python/tvm_tirx}/backend/cuda/tile/patterns/async_copy.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/tile/patterns/common.py | 0
.../python/tvm_tirx}/backend/cuda/tile/patterns/descriptors.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/tile/patterns/ewise.py | 0
.../tvm => tirx/python/tvm_tirx}/backend/cuda/tile/patterns/ldmatrix.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/tile/patterns/memory.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/tile/patterns/mma.py | 0
.../tvm => tirx/python/tvm_tirx}/backend/cuda/tile/patterns/register.py | 0
.../python/tvm_tirx}/backend/cuda/tile/patterns/tcgen05_mma.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/tile/patterns/tma.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/tile/patterns/tmem.py | 0
.../tvm => tirx/python/tvm_tirx}/backend/cuda/tile/patterns/tmem_copy.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/cuda/transforms/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/hexagon/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/hexagon/linker.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/hexagon/target_tags.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/loader.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/metal/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/metal/op.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/metal/script.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/metal/target_tags.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/opencl/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/opencl/script.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/rocm/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/axis.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/layout.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/op.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/pipeline.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/script.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/target_tags.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/thread_scope.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/binary/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/binary/default.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/binary/utils.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/common.py | 0
.../tvm => tirx/python/tvm_tirx}/backend/trn/tile/compose_op/__init__.py | 0
.../python/tvm_tirx}/backend/trn/tile/compose_op/binary_chain.py | 0
.../python/tvm_tirx}/backend/trn/tile/compose_op/binary_reduce.py | 0
.../python/tvm_tirx}/backend/trn/tile/compose_op/unary_reduce.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/compose_op/utils.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/copy/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/copy/default.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/dim_utils.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/gemm/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/gemm/default.py | 0
.../python/tvm_tirx}/backend/trn/tile/instruction_generator.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/instructions.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/private_alloc.py | 0
.../tvm => tirx/python/tvm_tirx}/backend/trn/tile/reduction/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/reduction/utils.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/select/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/select/default.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/unary/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/unary/default.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/unary/utils.py | 0
.../python/tvm_tirx}/backend/trn/tile/unary/with_bias_scale.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/tile/workspace_utils.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/trn/transform/__init__.py | 0
.../tvm => tirx/python/tvm_tirx}/backend/trn/transform/naive_allocator.py | 0
.../python/tvm_tirx}/backend/trn/transform/private_tensor_alloc.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/vulkan/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/backend/webgpu/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/base.py | 0
{python/tvm => tirx/python/tvm_tirx}/codegen.py | 0
{python/tvm => tirx/python/tvm_tirx}/error.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/_constant.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/_ffi_api.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/_overload_prim_expr.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/_tensor_expr_overload.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/attrs.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/base.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/expr.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/function.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/global_info.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/json_compact.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/location.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/module.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/op.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/prim/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/prim/_ffi_api.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/prim/expr.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/prim/op.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/stmt.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/supply.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/transform/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/transform/_ffi_instrument_api.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/transform/_ffi_transform_api.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/transform/core.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/transform/instrument.py | 0
{python/tvm => tirx/python/tvm_tirx}/ir/type.py | 0
{python/tvm => tirx/python/tvm_tirx}/libinfo.py | 0
{python/tvm => tirx/python/tvm_tirx}/runtime/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/runtime/_ffi_api.py | 0
{python/tvm => tirx/python/tvm_tirx}/runtime/_tensor.py | 0
{python/tvm => tirx/python/tvm_tirx}/runtime/device.py | 0
{python/tvm => tirx/python/tvm_tirx}/runtime/executable.py | 0
{python/tvm => tirx/python/tvm_tirx}/runtime/memory_scope.py | 0
{python/tvm => tirx/python/tvm_tirx}/runtime/module.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/_ffi_api.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/ir_builder/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/ir_builder/_ffi_api.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/ir_builder/base.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/ir_builder/frame.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/ir_builder/ir.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/ir_builder/parser_protocol.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/ir_builder/stmt.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/ir_builder/tracing.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/parser/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/parser/entry.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/parser/inspect_source.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/parser/ir.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/parser/jit_support.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/parser/prescan.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/parser/protocol_registry.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/parser/transpile.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/printer/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/printer/_ffi_api.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/printer/doc.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/printer/doc_printer.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/printer/highlight.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/printer/scriptable.py | 0
{python/tvm => tirx/python/tvm_tirx}/script/tirx.pyi | 0
{python/tvm => tirx/python/tvm_tirx}/support/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/support/cc.py | 0
{python/tvm => tirx/python/tvm_tirx}/support/clang.py | 0
{python/tvm => tirx/python/tvm_tirx}/support/libinfo.py | 0
{python/tvm => tirx/python/tvm_tirx}/support/ndk.py | 0
{python/tvm => tirx/python/tvm_tirx}/support/nvcc.py | 0
{python/tvm => tirx/python/tvm_tirx}/support/popen_pool.py | 0
{python/tvm => tirx/python/tvm_tirx}/support/popen_worker.py | 0
{python/tvm => tirx/python/tvm_tirx}/support/rocm.py | 0
{python/tvm => tirx/python/tvm_tirx}/support/tar.py | 0
{python/tvm => tirx/python/tvm_tirx}/support/utils.py | 0
{python/tvm => tirx/python/tvm_tirx}/support/xcode.py | 0
{python/tvm => tirx/python/tvm_tirx}/sym/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/sym/_ffi_api.py | 0
{python/tvm => tirx/python/tvm_tirx}/sym/analyzer.py | 0
{python/tvm => tirx/python/tvm_tirx}/sym/bound.py | 0
{python/tvm => tirx/python/tvm_tirx}/sym/int_set.py | 0
{python/tvm => tirx/python/tvm_tirx}/sym/iter_affine_map.py | 0
{python/tvm => tirx/python/tvm_tirx}/sym/pattern.py | 0
{python/tvm => tirx/python/tvm_tirx}/sym/utils.py | 0
{python/tvm => tirx/python/tvm_tirx}/target/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/target/_ffi_api.py | 0
{python/tvm => tirx/python/tvm_tirx}/target/detect_target.py | 0
{python/tvm => tirx/python/tvm_tirx}/target/intrin.py | 0
{python/tvm => tirx/python/tvm_tirx}/target/tag.py | 0
{python/tvm => tirx/python/tvm_tirx}/target/tag_registry/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/target/tag_registry/arm_cpu.py | 0
{python/tvm => tirx/python/tvm_tirx}/target/tag_registry/aws_cpu.py | 0
{python/tvm => tirx/python/tvm_tirx}/target/tag_registry/registry.py | 0
{python/tvm => tirx/python/tvm_tirx}/target/tag_registry/riscv_cpu.py | 0
{python/tvm => tirx/python/tvm_tirx}/target/target.py | 0
{python/tvm => tirx/python/tvm_tirx}/testing/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/testing/_ffi_api.py | 0
{python/tvm => tirx/python/tvm_tirx}/testing/env.py | 0
{python/tvm => tirx/python/tvm_tirx}/testing/locking.py | 0
{python/tvm => tirx/python/tvm_tirx}/testing/popen_pool.py | 0
{python/tvm => tirx/python/tvm_tirx}/testing/utils.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/_compile.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/_ffi_api.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/_tensor_view.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/analysis/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/analysis/_ffi_api.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/analysis/analysis.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/backend/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/bench.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/compilation_pipeline.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/expr.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/function.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/layout.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/op/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/script/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/script/ir_builder/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/script/ir_builder/_ffi_api.py | 0
.../python/tvm_tirx}/tirx/script/ir_builder/external_kernel.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/script/ir_builder/frame.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/script/ir_builder/ir.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/script/ir_builder/op.py | 0
.../python/tvm_tirx}/tirx/script/ir_builder/parser_protocol.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/script/ir_builder/triton.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/script/ir_builder/utils.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/script/jit.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/tensor.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/thread_scope.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/tile/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/tile/common.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/tile/dispatch.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/tile/instruction.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/tile/views.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/transform/__init__.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/transform/_ffi_api.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/transform/common.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/transform/function_pass.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/transform/transform.py | 0
{python/tvm => tirx/python/tvm_tirx}/tirx/type.py | 0
{src => tirx/src}/backend/cuda/codegen/codegen_cuda.cc | 0
{src => tirx/src}/backend/cuda/codegen/codegen_cuda.h | 0
{src => tirx/src}/backend/cuda/codegen/codegen_cuda_host.cc | 0
{src => tirx/src}/backend/cuda/codegen/cuda_fallback_module.cc | 0
{src => tirx/src}/backend/cuda/codegen/cuda_fallback_module.h | 0
{src => tirx/src}/backend/cuda/codegen/intrin_rule_cuda.cc | 0
{src => tirx/src}/backend/cuda/codegen/literal/cuda_half_t.h | 0
{src => tirx/src}/backend/cuda/codegen/literal/cuda_int8_t.h | 0
{src => tirx/src}/backend/cuda/codegen/llvm/codegen_nvptx.cc | 0
{src => tirx/src}/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc | 0
{src => tirx/src}/backend/cuda/codegen/ptx.cc | 0
{src => tirx/src}/backend/cuda/codegen/target_kind.cc | 0
{src => tirx/src}/backend/cuda/launch/launch_attributes.h | 0
{src => tirx/src}/backend/cuda/launch/launch_config.h | 0
{src => tirx/src}/backend/cuda/launch/launch_fields.h | 0
{src => tirx/src}/backend/cuda/launch/launch_plan.h | 0
{src => tirx/src}/backend/cuda/launch/launch_source.h | 0
{src => tirx/src}/backend/cuda/layout/axis.cc | 0
{src => tirx/src}/backend/cuda/op/iket.cc | 0
{src => tirx/src}/backend/cuda/op/legacy.cc | 0
{src => tirx/src}/backend/cuda/op/target_builtin.cc | 0
{src => tirx/src}/backend/cuda/op/tcgen05.cc | 0
{src => tirx/src}/backend/cuda/op/tensormap.cc | 0
{src => tirx/src}/backend/cuda/runtime/cuda_device_api.cc | 0
{src => tirx/src}/backend/cuda/runtime/cuda_module.cc | 0
{src => tirx/src}/backend/cuda/tile/exec_analysis.cc | 0
{src => tirx/src}/backend/cuda/tile/exec_analysis.h | 0
{src => tirx/src}/backend/cuda/tile/exec_context.cc | 0
{src => tirx/src}/backend/cuda/tile/exec_context.h | 0
{src => tirx/src}/backend/cuda/tile/filter_canonical.cc | 0
{src => tirx/src}/backend/cuda/tile/filter_canonical.h | 0
{src => tirx/src}/backend/cuda/tile/thread_scope.cc | 0
{src => tirx/src}/backend/cuda/tile/thread_scope.h | 0
{src => tirx/src}/backend/cuda/transforms/bind_backend_config.cc | 0
{src => tirx/src}/backend/cuda/transforms/lower_iket.cc | 0
{src => tirx/src}/backend/hexagon/codegen/hexagon_fallback_module.cc | 0
{src => tirx/src}/backend/hexagon/codegen/hexagon_fallback_module.h | 0
{src => tirx/src}/backend/hexagon/codegen/llvm/codegen_hexagon.cc | 0
{src => tirx/src}/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc | 0
{src => tirx/src}/backend/hexagon/codegen/memory_scope.h | 0
{src => tirx/src}/backend/hexagon/codegen/target_kind.cc | 0
{src => tirx/src}/backend/metal/codegen/codegen_metal.cc | 0
{src => tirx/src}/backend/metal/codegen/codegen_metal.h | 0
{src => tirx/src}/backend/metal/codegen/intrin_rule_metal.cc | 0
{src => tirx/src}/backend/metal/codegen/metal_fallback_module.cc | 0
{src => tirx/src}/backend/metal/codegen/metal_fallback_module.h | 0
{src => tirx/src}/backend/metal/codegen/target_kind.cc | 0
{src => tirx/src}/backend/metal/op/cooperative_tensor.cc | 0
{src => tirx/src}/backend/metal/op/target_builtin.cc | 0
{src => tirx/src}/backend/metal/runtime/metal_common.h | 0
{src => tirx/src}/backend/metal/runtime/metal_device_api.mm | 0
{src => tirx/src}/backend/metal/runtime/metal_module.mm | 0
{src => tirx/src}/backend/opencl/codegen/codegen_opencl.cc | 0
{src => tirx/src}/backend/opencl/codegen/codegen_opencl.h | 0
{src => tirx/src}/backend/opencl/codegen/intrin_rule_opencl.cc | 0
{src => tirx/src}/backend/opencl/codegen/opencl_fallback_module.cc | 0
{src => tirx/src}/backend/opencl/codegen/opencl_fallback_module.h | 0
{src => tirx/src}/backend/opencl/codegen/target_kind.cc | 0
{src => tirx/src}/backend/opencl/memory_scope.h | 0
{src => tirx/src}/backend/opencl/op/memory.cc | 0
{src => tirx/src}/backend/opencl/runtime/opencl_common.h | 0
{src => tirx/src}/backend/opencl/runtime/opencl_device_api.cc | 0
{src => tirx/src}/backend/opencl/runtime/opencl_module.cc | 0
{src => tirx/src}/backend/opencl/runtime/opencl_wrapper/README.md | 0
{src => tirx/src}/backend/opencl/runtime/opencl_wrapper/opencl_wrapper.cc | 0
{src => tirx/src}/backend/opencl/runtime/source_utils.h | 0
{src => tirx/src}/backend/opencl/runtime/texture.h | 0
{src => tirx/src}/backend/rocm/codegen/llvm/codegen_amdgpu.cc | 0
{src => tirx/src}/backend/rocm/codegen/llvm/intrin_rule_rocm.cc | 0
{src => tirx/src}/backend/rocm/codegen/rocm_fallback_module.cc | 0
{src => tirx/src}/backend/rocm/codegen/rocm_fallback_module.h | 0
{src => tirx/src}/backend/rocm/codegen/target_kind.cc | 0
{src => tirx/src}/backend/rocm/runtime/rocm_common.h | 0
{src => tirx/src}/backend/rocm/runtime/rocm_device_api.cc | 0
{src => tirx/src}/backend/rocm/runtime/rocm_module.cc | 0
{src => tirx/src}/backend/trn/codegen/codegen_trn.cc | 0
{src => tirx/src}/backend/trn/codegen/codegen_trn.h | 0
{src => tirx/src}/backend/trn/codegen/target_kind.cc | 0
{src => tirx/src}/backend/trn/layout/axis.cc | 0
{src => tirx/src}/backend/trn/op/target_builtin.cc | 0
{src => tirx/src}/backend/trn/transform/lower_trainium_layout.cc | 0
{src => tirx/src}/backend/vulkan/codegen/build_vulkan.cc | 0
{src => tirx/src}/backend/vulkan/codegen/codegen_spirv.cc | 0
{src => tirx/src}/backend/vulkan/codegen/codegen_spirv.h | 0
{src => tirx/src}/backend/vulkan/codegen/intrin_rule_spirv.cc | 0
{src => tirx/src}/backend/vulkan/codegen/ir_builder.cc | 0
{src => tirx/src}/backend/vulkan/codegen/ir_builder.h | 0
{src => tirx/src}/backend/vulkan/codegen/spirv_support.cc | 0
{src => tirx/src}/backend/vulkan/codegen/spirv_support.h | 0
{src => tirx/src}/backend/vulkan/codegen/spirv_utils.cc | 0
{src => tirx/src}/backend/vulkan/codegen/spirv_utils.h | 0
{src => tirx/src}/backend/vulkan/codegen/target_kind.cc | 0
{src => tirx/src}/backend/vulkan/codegen/vulkan_fallback_module.cc | 0
{src => tirx/src}/backend/vulkan/codegen/vulkan_fallback_module.h | 0
{src => tirx/src}/backend/vulkan/runtime/README.md | 0
{src => tirx/src}/backend/vulkan/runtime/spirv_shader.h | 0
{src => tirx/src}/backend/vulkan/runtime/thread_map.h | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_amdrgp.cc | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_amdrgp.h | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_buffer.cc | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_buffer.h | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_common.cc | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_common.h | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_device.cc | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_device.h | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_device_api.cc | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_device_api.h | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_instance.cc | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_instance.h | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_module.cc | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_stream.cc | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_stream.h | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_timer.cc | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_timer.h | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_wrapped_func.cc | 0
{src => tirx/src}/backend/vulkan/runtime/vulkan_wrapped_func.h | 0
{src => tirx/src}/backend/webgpu/codegen/codegen_webgpu.cc | 0
{src => tirx/src}/backend/webgpu/codegen/codegen_webgpu.h | 0
{src => tirx/src}/backend/webgpu/codegen/intrin_rule_webgpu.cc | 0
{src => tirx/src}/backend/webgpu/codegen/target_kind.cc | 0
{src => tirx/src}/backend/webgpu/codegen/webgpu_fallback_module.cc | 0
{src => tirx/src}/backend/webgpu/codegen/webgpu_fallback_module.h | 0
{src => tirx/src}/codegen/build_common.h | 0
{src => tirx/src}/codegen/c/codegen_c.cc | 0
{src => tirx/src}/codegen/c/codegen_c.h | 0
{src => tirx/src}/codegen/c/codegen_c_host.cc | 0
{src => tirx/src}/codegen/c/codegen_c_host.h | 0
{src => tirx/src}/codegen/c/codegen_params.cc | 0
{src => tirx/src}/codegen/c/codegen_params.h | 0
{src => tirx/src}/codegen/c/codegen_source_base.cc | 0
{src => tirx/src}/codegen/c/codegen_source_base.h | 0
{src => tirx/src}/codegen/c/source_module.cc | 0
{src => tirx/src}/codegen/c/target_kind.cc | 0
{src => tirx/src}/codegen/canonicalizer/arm_mprofile.cc | 0
{src => tirx/src}/codegen/canonicalizer/arm_mprofile.h | 0
{src => tirx/src}/codegen/codegen.cc | 0
{src => tirx/src}/codegen/intrin_rule.cc | 0
{src => tirx/src}/codegen/llvm/canonicalizer/arm_aprofile.cc | 0
{src => tirx/src}/codegen/llvm/canonicalizer/arm_aprofile.h | 0
{src => tirx/src}/codegen/llvm/canonicalizer/canonicalize.cc | 0
{src => tirx/src}/codegen/llvm/canonicalizer/canonicalize.h | 0
{src => tirx/src}/codegen/llvm/codegen_aarch64.cc | 0
{src => tirx/src}/codegen/llvm/codegen_arm.cc | 0
{src => tirx/src}/codegen/llvm/codegen_blob.cc | 0
{src => tirx/src}/codegen/llvm/codegen_blob.h | 0
{src => tirx/src}/codegen/llvm/codegen_cpu.cc | 0
{src => tirx/src}/codegen/llvm/codegen_cpu.h | 0
{src => tirx/src}/codegen/llvm/codegen_llvm.cc | 0
{src => tirx/src}/codegen/llvm/codegen_llvm.h | 0
{src => tirx/src}/codegen/llvm/codegen_params.cc | 0
{src => tirx/src}/codegen/llvm/codegen_params.h | 0
{src => tirx/src}/codegen/llvm/codegen_x86_64.cc | 0
{src => tirx/src}/codegen/llvm/intrin_rule_llvm.cc | 0
{src => tirx/src}/codegen/llvm/intrin_rule_llvm.h | 0
{src => tirx/src}/codegen/llvm/llvm_instance.cc | 0
{src => tirx/src}/codegen/llvm/llvm_instance.h | 0
{src => tirx/src}/codegen/llvm/llvm_module.cc | 0
{src => tirx/src}/codegen/llvm/llvm_module.h | 0
{src => tirx/src}/codegen/llvm/target_kind.cc | 0
{src => tirx/src}/codegen/module_export.cc | 0
{src => tirx/src}/ir/attrs.cc | 0
{src => tirx/src}/ir/env_func.cc | 0
{src => tirx/src}/ir/expr.cc | 0
{src => tirx/src}/ir/expr_functor.cc | 0
{src => tirx/src}/ir/function.cc | 0
{src => tirx/src}/ir/global_info.cc | 0
{src => tirx/src}/ir/instrument.cc | 0
{src => tirx/src}/ir/location.cc | 0
{src => tirx/src}/ir/module.cc | 0
{src => tirx/src}/ir/op.cc | 0
{src => tirx/src}/ir/prim/builtin.cc | 0
{src => tirx/src}/ir/prim/const_fold.h | 0
{src => tirx/src}/ir/prim/deep_equal.cc | 0
{src => tirx/src}/ir/prim/expr.cc | 0
{src => tirx/src}/ir/prim/int_operator.h | 0
{src => tirx/src}/ir/prim/math.cc | 0
{src => tirx/src}/ir/prim/op.cc | 0
{src => tirx/src}/ir/prim/op_utils.h | 0
{src => tirx/src}/ir/prim/symbolic_limits.cc | 0
{src => tirx/src}/ir/prim/vector_expr.cc | 0
{src => tirx/src}/ir/seq_stmt_mutate.h | 0
{src => tirx/src}/ir/side_effect.cc | 0
{src => tirx/src}/ir/stmt.cc | 0
{src => tirx/src}/ir/subscript_expr.cc | 0
{src => tirx/src}/ir/transform.cc | 0
{src => tirx/src}/ir/type.cc | 0
{src => tirx/src}/ir/unique_name_supply.cc | 0
{src => tirx/src}/runtime/builtin_fp16.cc | 0
{src => tirx/src}/runtime/cpu_device_api.cc | 0
{src => tirx/src}/runtime/device_api.cc | 0
{src => tirx/src}/runtime/launch_param.h | 0
{src => tirx/src}/runtime/logging.cc | 0
{src => tirx/src}/runtime/memory_scope.cc | 0
{src => tirx/src}/runtime/metadata.cc | 0
{src => tirx/src}/runtime/metadata.h | 0
{src => tirx/src}/runtime/module.cc | 0
{src => tirx/src}/runtime/pack_args.h | 0
{src => tirx/src}/runtime/static_library.cc | 0
{src => tirx/src}/runtime/static_library.h | 0
{src => tirx/src}/runtime/tensor.cc | 0
{src => tirx/src}/runtime/thread_pool.cc | 0
{src => tirx/src}/runtime/threading_backend.cc | 0
{src => tirx/src}/runtime/threading_backend.h | 0
{src => tirx/src}/runtime/timer.cc | 0
{src => tirx/src}/runtime/workspace_pool.cc | 0
{src => tirx/src}/runtime/workspace_pool.h | 0
{src => tirx/src}/script/ir_builder/base.cc | 0
{src => tirx/src}/script/ir_builder/frame.cc | 0
{src => tirx/src}/script/ir_builder/ir.cc | 0
{src => tirx/src}/script/ir_builder/utils.h | 0
{src => tirx/src}/script/printer/config.cc | 0
{src => tirx/src}/script/printer/doc.cc | 0
{src => tirx/src}/script/printer/doc_printer.cc | 0
{src => tirx/src}/script/printer/doc_printer.h | 0
{src => tirx/src}/script/printer/doc_translator.cc | 0
{src => tirx/src}/script/printer/ir/call.cc | 0
{src => tirx/src}/script/printer/ir/ir.cc | 0
{src => tirx/src}/script/printer/ir/misc.cc | 0
{src => tirx/src}/script/printer/ir/prim_expr.cc | 0
{src => tirx/src}/script/printer/ir/prim_type.cc | 0
{src => tirx/src}/script/printer/ir/stmt.cc | 0
{src => tirx/src}/script/printer/ir/utils.cc | 0
{src => tirx/src}/script/printer/ir/utils.h | 0
{src => tirx/src}/script/printer/printer.cc | 0
{src => tirx/src}/script/printer/script_printer.cc | 0
{src => tirx/src}/script/printer/utils.h | 0
{src => tirx/src}/support/limits.h | 0
{src => tirx/src}/sym/analyzer.cc | 0
{src => tirx/src}/sym/bound_deducer.cc | 0
{src => tirx/src}/sym/canonical_simplify.cc | 0
{src => tirx/src}/sym/const_fold.h | 0
{src => tirx/src}/sym/const_int_bound.cc | 0
{src => tirx/src}/sym/constraint_extract.cc | 0
{src => tirx/src}/sym/constraint_extract.h | 0
{src => tirx/src}/sym/constraint_helpers.h | 0
{src => tirx/src}/sym/detect_linear_equation.cc | 0
{src => tirx/src}/sym/int_operator.h | 0
{src => tirx/src}/sym/int_set.cc | 0
{src => tirx/src}/sym/interval_set.h | 0
{src => tirx/src}/sym/iter_affine_map.cc | 0
{src => tirx/src}/sym/known_comparison_analyzer.cc | 0
{src => tirx/src}/sym/modular_set.cc | 0
{src => tirx/src}/sym/pattern_match.h | 0
{src => tirx/src}/sym/product_normal_form.h | 0
{src => tirx/src}/sym/rewrite_simplify.cc | 0
{src => tirx/src}/sym/rewrite_simplify.h | 0
{src => tirx/src}/sym/simplify_base.cc | 0
{src => tirx/src}/sym/simplify_base.h | 0
{src => tirx/src}/sym/z3_prover.cc | 0
{src => tirx/src}/target/target.cc | 0
{src => tirx/src}/target/target_kind.cc | 0
{src => tirx/src}/target/target_tag_registry.cc | 0
{src => tirx/src}/tirx/analysis/check_well_formed.cc | 0
{src => tirx/src}/tirx/analysis/expr_complexity.cc | 0
{src => tirx/src}/tirx/analysis/stmt_finding.cc | 0
{src => tirx/src}/tirx/analysis/undefined_vars.cc | 0
{src => tirx/src}/tirx/analysis/verify_memory.cc | 0
{src => tirx/src}/tirx/ir/data_type_rewriter.cc | 0
{src => tirx/src}/tirx/ir/data_type_rewriter.h | 0
{src => tirx/src}/tirx/ir/function.cc | 0
{src => tirx/src}/tirx/ir/index_map.cc | 0
{src => tirx/src}/tirx/ir/ir_mutator_with_analyzer.cc | 0
{src => tirx/src}/tirx/ir/ir_visitor_with_analyzer.cc | 0
{src => tirx/src}/tirx/ir/layout/axis_registry.cc | 0
{src => tirx/src}/tirx/ir/layout/axis_registry.h | 0
{src => tirx/src}/tirx/ir/layout/compose_layout.cc | 0
{src => tirx/src}/tirx/ir/layout/layout.cc | 0
{src => tirx/src}/tirx/ir/layout/tile_canonicalize.cc | 0
{src => tirx/src}/tirx/ir/layout/tile_core.cc | 0
{src => tirx/src}/tirx/ir/layout/tile_internal.h | 0
{src => tirx/src}/tirx/ir/layout/tile_tile_ops.cc | 0
{src => tirx/src}/tirx/ir/layout/utils.cc | 0
{src => tirx/src}/tirx/ir/layout/utils.h | 0
{src => tirx/src}/tirx/ir/specialize.cc | 0
{src => tirx/src}/tirx/ir/stmt.cc | 0
{src => tirx/src}/tirx/ir/stmt_functor.cc | 0
{src => tirx/src}/tirx/ir/tensor.cc | 0
{src => tirx/src}/tirx/ir/tensor_load.cc | 0
{src => tirx/src}/tirx/ir/transform.cc | 0
{src => tirx/src}/tirx/ir/type.cc | 0
{src => tirx/src}/tirx/op/abi.cc | 0
{src => tirx/src}/tirx/op/annotation.cc | 0
{src => tirx/src}/tirx/op/cpu.cc | 0
{src => tirx/src}/tirx/op/gpu.cc | 0
{src => tirx/src}/tirx/op/math.cc | 0
{src => tirx/src}/tirx/op/memory.cc | 0
{src => tirx/src}/tirx/op/region.cc | 0
{src => tirx/src}/tirx/op/vector.cc | 0
{src => tirx/src}/tirx/script/ir_builder/frame.cc | 0
{src => tirx/src}/tirx/script/ir_builder/ir.cc | 0
{src => tirx/src}/tirx/script/ir_builder/utils.h | 0
{src => tirx/src}/tirx/script/printer/block.cc | 0
{src => tirx/src}/tirx/script/printer/expr.cc | 0
{src => tirx/src}/tirx/script/printer/for_loop.cc | 0
{src => tirx/src}/tirx/script/printer/function.cc | 0
{src => tirx/src}/tirx/script/printer/ir.cc | 0
{src => tirx/src}/tirx/script/printer/stmt.cc | 0
{src => tirx/src}/tirx/script/printer/tensor.cc | 0
{src => tirx/src}/tirx/tile/context.cc | 0
{src => tirx/src}/tirx/tile/dispatch.cc | 0
{src => tirx/src}/tirx/tile/lower.cc | 0
{src => tirx/src}/tirx/tile/op.cc | 0
{src => tirx/src}/tirx/tile/thread_scope.cc | 0
{src => tirx/src}/tirx/transform/bind_target.cc | 0
{src => tirx/src}/tirx/transform/common_subexpr_elim.cc | 0
{src => tirx/src}/tirx/transform/decorate_device_scope.cc | 0
{src => tirx/src}/tirx/transform/dtype_conversion.cc | 0
{src => tirx/src}/tirx/transform/dtype_conversion.h | 0
{src => tirx/src}/tirx/transform/flatten_tensor.cc | 0
{src => tirx/src}/tirx/transform/flattened_tensor.h | 0
{src => tirx/src}/tirx/transform/force_narrow_index_to_i32.cc | 0
{src => tirx/src}/tirx/transform/function_utils.cc | 0
{src => tirx/src}/tirx/transform/hoist_if.cc | 0
{src => tirx/src}/tirx/transform/ir_utils.cc | 0
{src => tirx/src}/tirx/transform/ir_utils.h | 0
{src => tirx/src}/tirx/transform/lower_intrin.cc | 0
{src => tirx/src}/tirx/transform/lower_thread_allreduce.cc | 0
{src => tirx/src}/tirx/transform/lower_tirx.cc | 0
{src => tirx/src}/tirx/transform/lower_tirx_cleanup.cc | 0
{src => tirx/src}/tirx/transform/lower_tirx_dedup_tensormap.cc | 0
{src => tirx/src}/tirx/transform/lower_tvm_builtin.cc | 0
{src => tirx/src}/tirx/transform/make_packed_api.cc | 0
{src => tirx/src}/tirx/transform/narrow_datatype.cc | 0
{src => tirx/src}/tirx/transform/remove_no_op.cc | 0
{src => tirx/src}/tirx/transform/remove_no_op.h | 0
{src => tirx/src}/tirx/transform/renormalize_split_pattern.cc | 0
{src => tirx/src}/tirx/transform/skip_assert.cc | 0
{src => tirx/src}/tirx/transform/split_host_device.cc | 0
{src => tirx/src}/tirx/transform/stmt_simplify.cc | 0
{src => tirx/src}/tirx/transform/storage_rewrite.cc | 0
{src => tirx/src}/tirx/transform/tvm_ffi_binder.cc | 0
{src => tirx/src}/tirx/transform/tvm_ffi_binder.h | 0
{src => tirx/src}/tirx/transform/unroll_loop.cc | 0
{src => tirx/src}/tirx/transform/unsupported_dtype_legalize.cc | 0
{src => tirx/src}/tirx/transform/vectorize_loop.cc | 0
{tests => tirx/tests}/cpp/ir/expr_functor_test.cc | 0
{tests => tirx/tests}/cpp/ir/op_signature_test.cc | 0
{tests => tirx/tests}/cpp/ir/tensor_load_test.cc | 0
{tests => tirx/tests}/cpp/runtime/logging_test.cc | 0
{tests => tirx/tests}/cpp/runtime/memory_scope_test.cc | 0
{tests => tirx/tests}/cpp/sym/pattern_match_test.cc | 0
{tests => tirx/tests}/cpp/target/config_schema_test.cc | 0
{tests => tirx/tests}/cpp/target/target_test.cc | 0
{tests => tirx/tests}/cpp/tirx/stmt_functor_test.cc | 0
{tests => tirx/tests}/cpp/tirx/tensor_var_test.cc | 0
{tests => tirx/tests}/cpp/tirx/thread_scope_test.cc | 0
.../python/all-platform-minimal-test/test_validate_runtime_library.py | 0
{tests => tirx/tests}/python/backend/cuda/codegen/test_cuda.py | 0
{tests => tirx/tests}/python/backend/cuda/codegen/test_fastmath.py | 0
{tests => tirx/tests}/python/backend/cuda/codegen/test_fp4.py | 0
{tests => tirx/tests}/python/backend/cuda/codegen/test_fp8.py | 0
{tests => tirx/tests}/python/backend/cuda/test_cuda_host_bundle.py | 0
{tests => tirx/tests}/python/backend/cuda/test_cuda_module_export.py | 0
{tests => tirx/tests}/python/backend/hexagon/codegen/test_hexagon.py | 0
{tests => tirx/tests}/python/backend/metal/codegen/test_metal.py | 0
{tests => tirx/tests}/python/backend/metal/test_metal_compile_callback.py | 0
{tests => tirx/tests}/python/backend/metal/test_metal_module_export.py | 0
{tests => tirx/tests}/python/backend/opencl/codegen/test_opencl.py | 0
{tests => tirx/tests}/python/backend/opencl/test_opencl_module_export.py | 0
{tests => tirx/tests}/python/backend/rocm/codegen/test_rocm.py | 0
{tests => tirx/tests}/python/backend/rocm/test_rocm_module_export.py | 0
{tests => tirx/tests}/python/backend/vulkan/codegen/test_vulkan.py | 0
{tests => tirx/tests}/python/backend/vulkan/test_vulkan_module_export.py | 0
{tests => tirx/tests}/python/backend/webgpu/codegen/test_webgpu.py | 0
{tests => tirx/tests}/python/codegen/c/test_c.py | 0
{tests => tirx/tests}/python/codegen/gpu_common/test_allreduce.py | 0
{tests => tirx/tests}/python/codegen/gpu_common/test_device.py | 0
{tests => tirx/tests}/python/codegen/gpu_common/test_gpu_intrinsics.py | 0
{tests => tirx/tests}/python/codegen/llvm/test_aarch64.py | 0
{tests => tirx/tests}/python/codegen/llvm/test_arm.py | 0
{tests => tirx/tests}/python/codegen/llvm/test_llvm.py | 0
{tests => tirx/tests}/python/codegen/llvm/test_riscv.py | 0
{tests => tirx/tests}/python/codegen/llvm/test_scalable_vector.py | 0
{tests => tirx/tests}/python/codegen/llvm/test_x86.py | 0
{tests => tirx/tests}/python/codegen/test_argument_validation.py | 0
{tests => tirx/tests}/python/codegen/test_assert.py | 0
{tests => tirx/tests}/python/codegen/test_bool.py | 0
{tests => tirx/tests}/python/codegen/test_common.py | 0
{tests => tirx/tests}/python/codegen/test_extern.py | 0
{tests => tirx/tests}/python/codegen/test_static_init.py | 0
{tests => tirx/tests}/python/ir/test_ir_attrs.py | 0
{tests => tirx/tests}/python/ir/test_ir_json.py | 0
{tests => tirx/tests}/python/ir/test_ir_module.py | 0
{tests => tirx/tests}/python/ir/test_ir_prim_expr_deep_equal.py | 0
{tests => tirx/tests}/python/ir/test_ir_type.py | 0
{tests => tirx/tests}/python/ir/test_nonnullable_ir.py | 0
{tests => tirx/tests}/python/ir/test_roundtrip_runtime_module.py | 0
{tests => tirx/tests}/python/runtime/test_runtime_device_api.py | 0
{tests => tirx/tests}/python/runtime/test_runtime_memory_scope.py | 0
{tests => tirx/tests}/python/runtime/test_runtime_module.py | 0
{tests => tirx/tests}/python/runtime/test_runtime_tensor.py | 0
{tests => tirx/tests}/python/script/conftest.py | 0
{tests => tirx/tests}/python/script/minilang.py | 0
{tests => tirx/tests}/python/script/test_basic_usage.py | 0
{tests => tirx/tests}/python/script/test_builder_void_emission.py | 0
{tests => tirx/tests}/python/script/test_cross_function_call.py | 0
{tests => tirx/tests}/python/script/test_dynamic_shape.py | 0
{tests => tirx/tests}/python/script/test_error_handling.py | 0
{tests => tirx/tests}/python/script/test_ir_builder.py | 0
{tests => tirx/tests}/python/script/test_jit_specialization.py | 0
{tests => tirx/tests}/python/script/test_meta_programming.py | 0
{tests => tirx/tests}/python/script/test_parser_entry.py | 0
{tests => tirx/tests}/python/script/test_script_printer_doc.py | 0
{tests => tirx/tests}/python/script/test_script_printer_entry.py | 0
{tests => tirx/tests}/python/script/test_source_locations.py | 0
{tests => tirx/tests}/python/script/test_special_parser_protocol.py | 0
{tests => tirx/tests}/python/support/test_ccache.py | 0
{tests => tirx/tests}/python/support/test_nvcc.py | 0
{tests => tirx/tests}/python/support/test_popen_pool.py | 0
{tests => tirx/tests}/python/support/test_util.py | 0
{tests => tirx/tests}/python/sym/test_sym_analyzer_object.py | 0
{tests => tirx/tests}/python/sym/test_sym_const_int_bound.py | 0
{tests => tirx/tests}/python/sym/test_sym_deduce_bound.py | 0
{tests => tirx/tests}/python/sym/test_sym_detect_clip_bound.py | 0
{tests => tirx/tests}/python/sym/test_sym_detect_linear_equation.py | 0
{tests => tirx/tests}/python/sym/test_sym_intset.py | 0
{tests => tirx/tests}/python/sym/test_sym_iter_affine_map.py | 0
{tests => tirx/tests}/python/sym/test_sym_rewrite_simplify.py | 0
{tests => tirx/tests}/python/sym/test_sym_simplify.py | 0
{tests => tirx/tests}/python/sym/test_sym_z3.py | 0
{tests => tirx/tests}/python/target/test_llvm_features_info.py | 0
{tests => tirx/tests}/python/target/test_riscv_features.py | 0
{tests => tirx/tests}/python/target/test_target_parser_mprofile.py | 0
{tests => tirx/tests}/python/target/test_x86_features.py | 0
{tests => tirx/tests}/python/testing/test_env.py | 0
{tests => tirx/tests}/python/tirx/__init__.py | 0
.../tests}/python/tirx/analysis/test_tirx_analysis_check_well_defined.py | 0
.../tests}/python/tirx/analysis/test_tirx_analysis_undefined_vars.py | 0
{tests => tirx/tests}/python/tirx/base/test_tirx_host_func.py | 0
{tests => tirx/tests}/python/tirx/base/test_tirx_imm_values.py | 0
{tests => tirx/tests}/python/tirx/base/test_tirx_index_map.py | 0
{tests => tirx/tests}/python/tirx/base/test_tirx_intrin.py | 0
{tests => tirx/tests}/python/tirx/base/test_tirx_nodes.py | 0
{tests => tirx/tests}/python/tirx/base/test_tirx_op_types.py | 0
{tests => tirx/tests}/python/tirx/base/test_tirx_ptx_cp_async.py | 0
{tests => tirx/tests}/python/tirx/base/test_tirx_ptx_griddepcontrol.py | 0
{tests => tirx/tests}/python/tirx/base/test_tirx_ptx_mma_sp.py | 0
{tests => tirx/tests}/python/tirx/base/test_tirx_ptx_scalar_f32_math.py | 0
{tests => tirx/tests}/python/tirx/base/test_tirx_scalable_datatype.py | 0
{tests => tirx/tests}/python/tirx/base/test_tirx_specialize.py | 0
{tests => tirx/tests}/python/tirx/base/test_tirx_tensor_type.py | 0
{tests => tirx/tests}/python/tirx/base/test_tirx_texture_scope.py | 0
{tests => tirx/tests}/python/tirx/codegen/conftest.py | 0
{tests => tirx/tests}/python/tirx/codegen/test_codegen_ampere.py | 0
{tests => tirx/tests}/python/tirx/codegen/test_codegen_blackwell.py | 0
{tests => tirx/tests}/python/tirx/codegen/test_codegen_cuda.py | 0
{tests => tirx/tests}/python/tirx/codegen/test_codegen_dsmem.py | 0
{tests => tirx/tests}/python/tirx/codegen/test_codegen_hopper.py | 0
{tests => tirx/tests}/python/tirx/codegen/test_codegen_nki.py | 0
{tests => tirx/tests}/python/tirx/codegen/test_cuda_cta_reduce.py | 0
{tests => tirx/tests}/python/tirx/codegen/test_cuda_wait_until.py | 0
{tests => tirx/tests}/python/tirx/codegen/test_cuda_warp_reduce.py | 0
{tests => tirx/tests}/python/tirx/codegen/test_ptx_addr.py | 0
{tests => tirx/tests}/python/tirx/codegen/test_ptx_cvt.py | 0
{tests => tirx/tests}/python/tirx/codegen/test_ptx_dialect.py | 0
{tests => tirx/tests}/python/tirx/codegen/test_ptx_ld_st_ops.py | 0
.../tests}/python/tirx/codegen/test_tcgen05_descriptor_encoders.py | 0
{tests => tirx/tests}/python/tirx/conftest.py | 0
{tests => tirx/tests}/python/tirx/iket/iket_profile_workload.py | 0
.../tests}/python/tirx/iket/oracle/generate_iket_official_oracle.py | 0
.../python/tirx/iket/oracle/iket_official_cutlass_4_6_0_oracle.json | 0
{tests => tirx/tests}/python/tirx/iket/test_iket_orchestration.py | 0
{tests => tirx/tests}/python/tirx/iket/test_iket_profiler.py | 0
{tests => tirx/tests}/python/tirx/iket/verify_iket_official_patch.py | 0
{tests => tirx/tests}/python/tirx/layout_fixtures.py | 0
{tests => tirx/tests}/python/tirx/script/test_tirx_script_basic_usage.py | 0
.../tests}/python/tirx/script/test_tirx_script_dynamic_shape.py | 0
.../tests}/python/tirx/script/test_tirx_script_error_handling.py | 0
{tests => tirx/tests}/python/tirx/script/test_tirx_script_gpu.py | 0
.../tests}/python/tirx/script/test_tirx_script_jit_specialization.py | 0
.../tests}/python/tirx/script/test_tirx_script_meta_programming.py | 0
.../tests}/python/tirx/script/test_tirx_script_printer_highlight.py | 0
.../python/tirx/script/test_tirx_script_printer_invisible_path_info.py | 0
.../python/tirx/script/test_tirx_script_printer_structural_equal.py | 0
.../tests}/python/tirx/script/test_tirx_script_printer_underlining.py | 0
.../tests}/python/tirx/script/test_tirx_script_source_locations.py | 0
{tests => tirx/tests}/python/tirx/script/test_tirx_script_tensors.py | 0
{tests => tirx/tests}/python/tirx/script/test_tirx_script_tracing.py | 0
{tests => tirx/tests}/python/tirx/test_alloc_pool.py | 0
{tests => tirx/tests}/python/tirx/test_bench_utils.py | 0
{tests => tirx/tests}/python/tirx/test_compilation_pipeline.py | 0
{tests => tirx/tests}/python/tirx/test_control_flow.py | 0
{tests => tirx/tests}/python/tirx/test_cuda_arch_marker.py | 0
{tests => tirx/tests}/python/tirx/test_cuda_backend_config.py | 0
{tests => tirx/tests}/python/tirx/test_cuda_launch.py | 0
{tests => tirx/tests}/python/tirx/test_exec_context.py | 0
{tests => tirx/tests}/python/tirx/test_inline.py | 0
{tests => tirx/tests}/python/tirx/test_jit.py | 0
{tests => tirx/tests}/python/tirx/test_layout.py | 0
{tests => tirx/tests}/python/tirx/test_op.py | 0
{tests => tirx/tests}/python/tirx/test_op_namespace_cleanup.py | 0
{tests => tirx/tests}/python/tirx/test_roundtrip_namespaces.py | 0
{tests => tirx/tests}/python/tirx/test_tensor_data_reinfer_type.py | 0
{tests => tirx/tests}/python/tirx/test_thread_scope.py | 0
{tests => tirx/tests}/python/tirx/test_tile_namespace.py | 0
{tests => tirx/tests}/python/tirx/test_tirx_axis_enum.py | 0
{tests => tirx/tests}/python/tirx/test_verifier.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/conftest.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_add_sub_mul_div.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_cp_async.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_cp_async_bulk.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_cvt.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_ewise.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_explicit_patterns.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_fma.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_ld_st_cache_policy.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_ld_st_gmem_smem.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_ld_st_registers.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_ldmatrix_stmatrix.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_mma_sync.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_mov.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_pattern_groups.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_pattern_math.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_pattern_mma.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_pattern_transfers.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_tcgen05_cp.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_tcgen05_ld_st.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_tcgen05_mma.py | 0
{tests => tirx/tests}/python/tirx/tile/cuda/test_tcgen05_mma_codegen.py | 0
.../tests}/python/tirx/tile/cuda/test_tcgen05_mma_ws_split_n.py | 0
{tests => tirx/tests}/python/tirx/tile/test_dispatcher.py | 0
{tests => tirx/tests}/python/tirx/tile/trn/test_binary_trn.py | 0
{tests => tirx/tests}/python/tirx/tile/trn/test_compose_op_trn.py | 0
{tests => tirx/tests}/python/tirx/tile/trn/test_copy_trn.py | 0
{tests => tirx/tests}/python/tirx/tile/trn/test_gemm_trn.py | 0
{tests => tirx/tests}/python/tirx/tile/trn/test_private_alloc_trn.py | 0
{tests => tirx/tests}/python/tirx/tile/trn/test_reduction_trn.py | 0
{tests => tirx/tests}/python/tirx/tile/trn/test_select_trn.py | 0
{tests => tirx/tests}/python/tirx/tile/trn/test_unary_trn.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_bf16_legalize.py | 0
.../python/tirx/transform/test_tirx_transform_common_subexpr_elim.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_convert_ssa.py | 0
.../python/tirx/transform/test_tirx_transform_decorate_device_scope.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_flatten_tensor.py | 0
.../tirx/transform/test_tirx_transform_force_narrow_index_to_i32.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_fp8_legalize.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_function_pass.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_helpers.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_hoist_if.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_lower_intrin.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_lower_tirx.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_lower_tvm_builtin.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_make_packed_api.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_naive_allocator.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_narrow_datatype.py | 0
.../tirx/transform/test_tirx_transform_pointer_value_type_rewrite.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_remove_no_op.py | 0
.../tirx/transform/test_tirx_transform_renormalize_split_pattern.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_simplify.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_split_host_device.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_storage_rewrite.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_unroll_loop.py | 0
.../tests}/python/tirx/transform/test_tirx_transform_vectorize.py | 0
{tests => tirx/tests}/python/tirx/utils.py | 0
919 files changed, 0 insertions(+), 0 deletions(-)
diff --git a/3rdparty/compiler-rt/builtin_fp16.h b/tirx/3rdparty/compiler-rt/builtin_fp16.h
similarity index 100%
rename from 3rdparty/compiler-rt/builtin_fp16.h
rename to tirx/3rdparty/compiler-rt/builtin_fp16.h
diff --git a/cmake/modules/CUDA.cmake b/tirx/cmake/modules/CUDA.cmake
similarity index 100%
rename from cmake/modules/CUDA.cmake
rename to tirx/cmake/modules/CUDA.cmake
diff --git a/cmake/modules/ClangFlags.cmake b/tirx/cmake/modules/ClangFlags.cmake
similarity index 100%
rename from cmake/modules/ClangFlags.cmake
rename to tirx/cmake/modules/ClangFlags.cmake
diff --git a/cmake/modules/Git.cmake b/tirx/cmake/modules/Git.cmake
similarity index 100%
rename from cmake/modules/Git.cmake
rename to tirx/cmake/modules/Git.cmake
diff --git a/cmake/modules/LLVM.cmake b/tirx/cmake/modules/LLVM.cmake
similarity index 100%
rename from cmake/modules/LLVM.cmake
rename to tirx/cmake/modules/LLVM.cmake
diff --git a/cmake/modules/Logging.cmake b/tirx/cmake/modules/Logging.cmake
similarity index 100%
rename from cmake/modules/Logging.cmake
rename to tirx/cmake/modules/Logging.cmake
diff --git a/cmake/modules/Metal.cmake b/tirx/cmake/modules/Metal.cmake
similarity index 100%
rename from cmake/modules/Metal.cmake
rename to tirx/cmake/modules/Metal.cmake
diff --git a/cmake/modules/OpenCL.cmake b/tirx/cmake/modules/OpenCL.cmake
similarity index 100%
rename from cmake/modules/OpenCL.cmake
rename to tirx/cmake/modules/OpenCL.cmake
diff --git a/cmake/modules/OpenMP.cmake b/tirx/cmake/modules/OpenMP.cmake
similarity index 100%
rename from cmake/modules/OpenMP.cmake
rename to tirx/cmake/modules/OpenMP.cmake
diff --git a/cmake/modules/ROCM.cmake b/tirx/cmake/modules/ROCM.cmake
similarity index 100%
rename from cmake/modules/ROCM.cmake
rename to tirx/cmake/modules/ROCM.cmake
diff --git a/cmake/modules/Vulkan.cmake b/tirx/cmake/modules/Vulkan.cmake
similarity index 100%
rename from cmake/modules/Vulkan.cmake
rename to tirx/cmake/modules/Vulkan.cmake
diff --git a/cmake/modules/contrib/Z3.cmake b/tirx/cmake/modules/contrib/Z3.cmake
similarity index 100%
rename from cmake/modules/contrib/Z3.cmake
rename to tirx/cmake/modules/contrib/Z3.cmake
diff --git a/cmake/utils/CCache.cmake b/tirx/cmake/utils/CCache.cmake
similarity index 100%
rename from cmake/utils/CCache.cmake
rename to tirx/cmake/utils/CCache.cmake
diff --git a/cmake/utils/FindCUDA.cmake b/tirx/cmake/utils/FindCUDA.cmake
similarity index 100%
rename from cmake/utils/FindCUDA.cmake
rename to tirx/cmake/utils/FindCUDA.cmake
diff --git a/cmake/utils/FindLLVM.cmake b/tirx/cmake/utils/FindLLVM.cmake
similarity index 100%
rename from cmake/utils/FindLLVM.cmake
rename to tirx/cmake/utils/FindLLVM.cmake
diff --git a/cmake/utils/FindOpenCL.cmake b/tirx/cmake/utils/FindOpenCL.cmake
similarity index 100%
rename from cmake/utils/FindOpenCL.cmake
rename to tirx/cmake/utils/FindOpenCL.cmake
diff --git a/cmake/utils/FindROCM.cmake b/tirx/cmake/utils/FindROCM.cmake
similarity index 100%
rename from cmake/utils/FindROCM.cmake
rename to tirx/cmake/utils/FindROCM.cmake
diff --git a/cmake/utils/FindVulkan.cmake b/tirx/cmake/utils/FindVulkan.cmake
similarity index 100%
rename from cmake/utils/FindVulkan.cmake
rename to tirx/cmake/utils/FindVulkan.cmake
diff --git a/cmake/utils/Library.cmake b/tirx/cmake/utils/Library.cmake
similarity index 100%
rename from cmake/utils/Library.cmake
rename to tirx/cmake/utils/Library.cmake
diff --git a/cmake/utils/Linker.cmake b/tirx/cmake/utils/Linker.cmake
similarity index 100%
rename from cmake/utils/Linker.cmake
rename to tirx/cmake/utils/Linker.cmake
diff --git a/cmake/utils/Summary.cmake b/tirx/cmake/utils/Summary.cmake
similarity index 100%
rename from cmake/utils/Summary.cmake
rename to tirx/cmake/utils/Summary.cmake
diff --git a/cmake/utils/Utils.cmake b/tirx/cmake/utils/Utils.cmake
similarity index 100%
rename from cmake/utils/Utils.cmake
rename to tirx/cmake/utils/Utils.cmake
diff --git a/include/tvm/backend/cuda/attr.h b/tirx/include/tvm/backend/cuda/attr.h
similarity index 100%
rename from include/tvm/backend/cuda/attr.h
rename to tirx/include/tvm/backend/cuda/attr.h
diff --git a/include/tvm/backend/cuda/op/iket.h b/tirx/include/tvm/backend/cuda/op/iket.h
similarity index 100%
rename from include/tvm/backend/cuda/op/iket.h
rename to tirx/include/tvm/backend/cuda/op/iket.h
diff --git a/include/tvm/backend/cuda/op/tcgen05.h b/tirx/include/tvm/backend/cuda/op/tcgen05.h
similarity index 100%
rename from include/tvm/backend/cuda/op/tcgen05.h
rename to tirx/include/tvm/backend/cuda/op/tcgen05.h
diff --git a/include/tvm/backend/cuda/op/tensormap.h b/tirx/include/tvm/backend/cuda/op/tensormap.h
similarity index 100%
rename from include/tvm/backend/cuda/op/tensormap.h
rename to tirx/include/tvm/backend/cuda/op/tensormap.h
diff --git a/include/tvm/backend/cuda/op_attr_types.h b/tirx/include/tvm/backend/cuda/op_attr_types.h
similarity index 100%
rename from include/tvm/backend/cuda/op_attr_types.h
rename to tirx/include/tvm/backend/cuda/op_attr_types.h
diff --git a/include/tvm/backend/metal/op/cooperative_tensor.h b/tirx/include/tvm/backend/metal/op/cooperative_tensor.h
similarity index 100%
rename from include/tvm/backend/metal/op/cooperative_tensor.h
rename to tirx/include/tvm/backend/metal/op/cooperative_tensor.h
diff --git a/include/tvm/backend/opencl/op/memory.h b/tirx/include/tvm/backend/opencl/op/memory.h
similarity index 100%
rename from include/tvm/backend/opencl/op/memory.h
rename to tirx/include/tvm/backend/opencl/op/memory.h
diff --git a/include/tvm/codegen/codegen.h b/tirx/include/tvm/codegen/codegen.h
similarity index 100%
rename from include/tvm/codegen/codegen.h
rename to tirx/include/tvm/codegen/codegen.h
diff --git a/include/tvm/codegen/intrin_rule.h b/tirx/include/tvm/codegen/intrin_rule.h
similarity index 100%
rename from include/tvm/codegen/intrin_rule.h
rename to tirx/include/tvm/codegen/intrin_rule.h
diff --git a/include/tvm/ir/analysis.h b/tirx/include/tvm/ir/analysis.h
similarity index 100%
rename from include/tvm/ir/analysis.h
rename to tirx/include/tvm/ir/analysis.h
diff --git a/include/tvm/ir/attrs.h b/tirx/include/tvm/ir/attrs.h
similarity index 100%
rename from include/tvm/ir/attrs.h
rename to tirx/include/tvm/ir/attrs.h
diff --git a/include/tvm/ir/base_expr.h b/tirx/include/tvm/ir/base_expr.h
similarity index 100%
rename from include/tvm/ir/base_expr.h
rename to tirx/include/tvm/ir/base_expr.h
diff --git a/include/tvm/ir/cow.h b/tirx/include/tvm/ir/cow.h
similarity index 100%
rename from include/tvm/ir/cow.h
rename to tirx/include/tvm/ir/cow.h
diff --git a/include/tvm/ir/env_func.h b/tirx/include/tvm/ir/env_func.h
similarity index 100%
rename from include/tvm/ir/env_func.h
rename to tirx/include/tvm/ir/env_func.h
diff --git a/include/tvm/ir/expr.h b/tirx/include/tvm/ir/expr.h
similarity index 100%
rename from include/tvm/ir/expr.h
rename to tirx/include/tvm/ir/expr.h
diff --git a/include/tvm/ir/expr_functor.h b/tirx/include/tvm/ir/expr_functor.h
similarity index 100%
rename from include/tvm/ir/expr_functor.h
rename to tirx/include/tvm/ir/expr_functor.h
diff --git a/include/tvm/ir/function.h b/tirx/include/tvm/ir/function.h
similarity index 100%
rename from include/tvm/ir/function.h
rename to tirx/include/tvm/ir/function.h
diff --git a/include/tvm/ir/global_info.h b/tirx/include/tvm/ir/global_info.h
similarity index 100%
rename from include/tvm/ir/global_info.h
rename to tirx/include/tvm/ir/global_info.h
diff --git a/include/tvm/ir/location.h b/tirx/include/tvm/ir/location.h
similarity index 100%
rename from include/tvm/ir/location.h
rename to tirx/include/tvm/ir/location.h
diff --git a/include/tvm/ir/module.h b/tirx/include/tvm/ir/module.h
similarity index 100%
rename from include/tvm/ir/module.h
rename to tirx/include/tvm/ir/module.h
diff --git a/include/tvm/ir/object_functor.h b/tirx/include/tvm/ir/object_functor.h
similarity index 100%
rename from include/tvm/ir/object_functor.h
rename to tirx/include/tvm/ir/object_functor.h
diff --git a/include/tvm/ir/op.h b/tirx/include/tvm/ir/op.h
similarity index 100%
rename from include/tvm/ir/op.h
rename to tirx/include/tvm/ir/op.h
diff --git a/include/tvm/ir/prim/expr.h b/tirx/include/tvm/ir/prim/expr.h
similarity index 100%
rename from include/tvm/ir/prim/expr.h
rename to tirx/include/tvm/ir/prim/expr.h
diff --git a/include/tvm/ir/prim/op.h b/tirx/include/tvm/ir/prim/op.h
similarity index 100%
rename from include/tvm/ir/prim/op.h
rename to tirx/include/tvm/ir/prim/op.h
diff --git a/include/tvm/ir/prim/vector_expr.h b/tirx/include/tvm/ir/prim/vector_expr.h
similarity index 100%
rename from include/tvm/ir/prim/vector_expr.h
rename to tirx/include/tvm/ir/prim/vector_expr.h
diff --git a/include/tvm/ir/scope_stack.h b/tirx/include/tvm/ir/scope_stack.h
similarity index 100%
rename from include/tvm/ir/scope_stack.h
rename to tirx/include/tvm/ir/scope_stack.h
diff --git a/include/tvm/ir/stmt.h b/tirx/include/tvm/ir/stmt.h
similarity index 100%
rename from include/tvm/ir/stmt.h
rename to tirx/include/tvm/ir/stmt.h
diff --git a/include/tvm/ir/transform.h b/tirx/include/tvm/ir/transform.h
similarity index 100%
rename from include/tvm/ir/transform.h
rename to tirx/include/tvm/ir/transform.h
diff --git a/include/tvm/ir/type.h b/tirx/include/tvm/ir/type.h
similarity index 100%
rename from include/tvm/ir/type.h
rename to tirx/include/tvm/ir/type.h
diff --git a/include/tvm/ir/unique_name_supply.h b/tirx/include/tvm/ir/unique_name_supply.h
similarity index 100%
rename from include/tvm/ir/unique_name_supply.h
rename to tirx/include/tvm/ir/unique_name_supply.h
diff --git a/include/tvm/ir/with_context.h b/tirx/include/tvm/ir/with_context.h
similarity index 100%
rename from include/tvm/ir/with_context.h
rename to tirx/include/tvm/ir/with_context.h
diff --git a/include/tvm/runtime/base.h b/tirx/include/tvm/runtime/base.h
similarity index 100%
rename from include/tvm/runtime/base.h
rename to tirx/include/tvm/runtime/base.h
diff --git a/include/tvm/runtime/c_backend_api.h b/tirx/include/tvm/runtime/c_backend_api.h
similarity index 100%
rename from include/tvm/runtime/c_backend_api.h
rename to tirx/include/tvm/runtime/c_backend_api.h
diff --git a/include/tvm/runtime/device_api.h b/tirx/include/tvm/runtime/device_api.h
similarity index 100%
rename from include/tvm/runtime/device_api.h
rename to tirx/include/tvm/runtime/device_api.h
diff --git a/include/tvm/runtime/logging.h b/tirx/include/tvm/runtime/logging.h
similarity index 100%
rename from include/tvm/runtime/logging.h
rename to tirx/include/tvm/runtime/logging.h
diff --git a/include/tvm/runtime/memory_scope.h b/tirx/include/tvm/runtime/memory_scope.h
similarity index 100%
rename from include/tvm/runtime/memory_scope.h
rename to tirx/include/tvm/runtime/memory_scope.h
diff --git a/include/tvm/runtime/tensor.h b/tirx/include/tvm/runtime/tensor.h
similarity index 100%
rename from include/tvm/runtime/tensor.h
rename to tirx/include/tvm/runtime/tensor.h
diff --git a/include/tvm/runtime/timer.h b/tirx/include/tvm/runtime/timer.h
similarity index 100%
rename from include/tvm/runtime/timer.h
rename to tirx/include/tvm/runtime/timer.h
diff --git a/include/tvm/script/ir_builder/base.h b/tirx/include/tvm/script/ir_builder/base.h
similarity index 100%
rename from include/tvm/script/ir_builder/base.h
rename to tirx/include/tvm/script/ir_builder/base.h
diff --git a/include/tvm/script/ir_builder/frame.h b/tirx/include/tvm/script/ir_builder/frame.h
similarity index 100%
rename from include/tvm/script/ir_builder/frame.h
rename to tirx/include/tvm/script/ir_builder/frame.h
diff --git a/include/tvm/script/ir_builder/ir.h b/tirx/include/tvm/script/ir_builder/ir.h
similarity index 100%
rename from include/tvm/script/ir_builder/ir.h
rename to tirx/include/tvm/script/ir_builder/ir.h
diff --git a/include/tvm/script/printer/config.h b/tirx/include/tvm/script/printer/config.h
similarity index 100%
rename from include/tvm/script/printer/config.h
rename to tirx/include/tvm/script/printer/config.h
diff --git a/include/tvm/script/printer/doc.h b/tirx/include/tvm/script/printer/doc.h
similarity index 100%
rename from include/tvm/script/printer/doc.h
rename to tirx/include/tvm/script/printer/doc.h
diff --git a/include/tvm/script/printer/doc_translator.h b/tirx/include/tvm/script/printer/doc_translator.h
similarity index 100%
rename from include/tvm/script/printer/doc_translator.h
rename to tirx/include/tvm/script/printer/doc_translator.h
diff --git a/include/tvm/script/printer/printer.h b/tirx/include/tvm/script/printer/printer.h
similarity index 100%
rename from include/tvm/script/printer/printer.h
rename to tirx/include/tvm/script/printer/printer.h
diff --git a/include/tvm/script/printer/support.h b/tirx/include/tvm/script/printer/support.h
similarity index 100%
rename from include/tvm/script/printer/support.h
rename to tirx/include/tvm/script/printer/support.h
diff --git a/include/tvm/support/base64.h b/tirx/include/tvm/support/base64.h
similarity index 100%
rename from include/tvm/support/base64.h
rename to tirx/include/tvm/support/base64.h
diff --git a/include/tvm/support/bytes_io.h b/tirx/include/tvm/support/bytes_io.h
similarity index 100%
rename from include/tvm/support/bytes_io.h
rename to tirx/include/tvm/support/bytes_io.h
diff --git a/include/tvm/support/cuda/nvtx.h b/tirx/include/tvm/support/cuda/nvtx.h
similarity index 100%
rename from include/tvm/support/cuda/nvtx.h
rename to tirx/include/tvm/support/cuda/nvtx.h
diff --git a/include/tvm/support/env.h b/tirx/include/tvm/support/env.h
similarity index 100%
rename from include/tvm/support/env.h
rename to tirx/include/tvm/support/env.h
diff --git a/include/tvm/support/file_utils.h b/tirx/include/tvm/support/file_utils.h
similarity index 100%
rename from include/tvm/support/file_utils.h
rename to tirx/include/tvm/support/file_utils.h
diff --git a/include/tvm/support/gpu_thread_tag.h b/tirx/include/tvm/support/gpu_thread_tag.h
similarity index 100%
rename from include/tvm/support/gpu_thread_tag.h
rename to tirx/include/tvm/support/gpu_thread_tag.h
diff --git a/include/tvm/support/io.h b/tirx/include/tvm/support/io.h
similarity index 100%
rename from include/tvm/support/io.h
rename to tirx/include/tvm/support/io.h
diff --git a/include/tvm/support/ordered_map.h b/tirx/include/tvm/support/ordered_map.h
similarity index 100%
rename from include/tvm/support/ordered_map.h
rename to tirx/include/tvm/support/ordered_map.h
diff --git a/include/tvm/support/serializer.h b/tirx/include/tvm/support/serializer.h
similarity index 100%
rename from include/tvm/support/serializer.h
rename to tirx/include/tvm/support/serializer.h
diff --git a/include/tvm/support/str_escape.h b/tirx/include/tvm/support/str_escape.h
similarity index 100%
rename from include/tvm/support/str_escape.h
rename to tirx/include/tvm/support/str_escape.h
diff --git a/include/tvm/sym/analyzer.h b/tirx/include/tvm/sym/analyzer.h
similarity index 100%
rename from include/tvm/sym/analyzer.h
rename to tirx/include/tvm/sym/analyzer.h
diff --git a/include/tvm/sym/bound.h b/tirx/include/tvm/sym/bound.h
similarity index 100%
rename from include/tvm/sym/bound.h
rename to tirx/include/tvm/sym/bound.h
diff --git a/include/tvm/sym/int_set.h b/tirx/include/tvm/sym/int_set.h
similarity index 100%
rename from include/tvm/sym/int_set.h
rename to tirx/include/tvm/sym/int_set.h
diff --git a/include/tvm/sym/iter_affine_map.h b/tirx/include/tvm/sym/iter_affine_map.h
similarity index 100%
rename from include/tvm/sym/iter_affine_map.h
rename to tirx/include/tvm/sym/iter_affine_map.h
diff --git a/include/tvm/sym/pattern.h b/tirx/include/tvm/sym/pattern.h
similarity index 100%
rename from include/tvm/sym/pattern.h
rename to tirx/include/tvm/sym/pattern.h
diff --git a/include/tvm/target/config_schema.h b/tirx/include/tvm/target/config_schema.h
similarity index 100%
rename from include/tvm/target/config_schema.h
rename to tirx/include/tvm/target/config_schema.h
diff --git a/include/tvm/target/target.h b/tirx/include/tvm/target/target.h
similarity index 100%
rename from include/tvm/target/target.h
rename to tirx/include/tvm/target/target.h
diff --git a/include/tvm/target/target_kind.h b/tirx/include/tvm/target/target_kind.h
similarity index 100%
rename from include/tvm/target/target_kind.h
rename to tirx/include/tvm/target/target_kind.h
diff --git a/include/tvm/target/target_tag_registry.h b/tirx/include/tvm/target/target_tag_registry.h
similarity index 100%
rename from include/tvm/target/target_tag_registry.h
rename to tirx/include/tvm/target/target_tag_registry.h
diff --git a/include/tvm/tirx/analysis.h b/tirx/include/tvm/tirx/analysis.h
similarity index 100%
rename from include/tvm/tirx/analysis.h
rename to tirx/include/tvm/tirx/analysis.h
diff --git a/include/tvm/tirx/expr.h b/tirx/include/tvm/tirx/expr.h
similarity index 100%
rename from include/tvm/tirx/expr.h
rename to tirx/include/tvm/tirx/expr.h
diff --git a/include/tvm/tirx/function.h b/tirx/include/tvm/tirx/function.h
similarity index 100%
rename from include/tvm/tirx/function.h
rename to tirx/include/tvm/tirx/function.h
diff --git a/include/tvm/tirx/index_map.h b/tirx/include/tvm/tirx/index_map.h
similarity index 100%
rename from include/tvm/tirx/index_map.h
rename to tirx/include/tvm/tirx/index_map.h
diff --git a/include/tvm/tirx/layout.h b/tirx/include/tvm/tirx/layout.h
similarity index 100%
rename from include/tvm/tirx/layout.h
rename to tirx/include/tvm/tirx/layout.h
diff --git a/include/tvm/tirx/op/abi.h b/tirx/include/tvm/tirx/op/abi.h
similarity index 100%
rename from include/tvm/tirx/op/abi.h
rename to tirx/include/tvm/tirx/op/abi.h
diff --git a/include/tvm/tirx/op/annotation.h b/tirx/include/tvm/tirx/op/annotation.h
similarity index 100%
rename from include/tvm/tirx/op/annotation.h
rename to tirx/include/tvm/tirx/op/annotation.h
diff --git a/include/tvm/tirx/op/cpu.h b/tirx/include/tvm/tirx/op/cpu.h
similarity index 100%
rename from include/tvm/tirx/op/cpu.h
rename to tirx/include/tvm/tirx/op/cpu.h
diff --git a/include/tvm/tirx/op/gpu.h b/tirx/include/tvm/tirx/op/gpu.h
similarity index 100%
rename from include/tvm/tirx/op/gpu.h
rename to tirx/include/tvm/tirx/op/gpu.h
diff --git a/include/tvm/tirx/op/math.h b/tirx/include/tvm/tirx/op/math.h
similarity index 100%
rename from include/tvm/tirx/op/math.h
rename to tirx/include/tvm/tirx/op/math.h
diff --git a/include/tvm/tirx/op/memory.h b/tirx/include/tvm/tirx/op/memory.h
similarity index 100%
rename from include/tvm/tirx/op/memory.h
rename to tirx/include/tvm/tirx/op/memory.h
diff --git a/include/tvm/tirx/op/region.h b/tirx/include/tvm/tirx/op/region.h
similarity index 100%
rename from include/tvm/tirx/op/region.h
rename to tirx/include/tvm/tirx/op/region.h
diff --git a/include/tvm/tirx/op/vector.h b/tirx/include/tvm/tirx/op/vector.h
similarity index 100%
rename from include/tvm/tirx/op/vector.h
rename to tirx/include/tvm/tirx/op/vector.h
diff --git a/include/tvm/tirx/op_attr_types.h b/tirx/include/tvm/tirx/op_attr_types.h
similarity index 100%
rename from include/tvm/tirx/op_attr_types.h
rename to tirx/include/tvm/tirx/op_attr_types.h
diff --git a/include/tvm/tirx/script/ir_builder/frame.h b/tirx/include/tvm/tirx/script/ir_builder/frame.h
similarity index 100%
rename from include/tvm/tirx/script/ir_builder/frame.h
rename to tirx/include/tvm/tirx/script/ir_builder/frame.h
diff --git a/include/tvm/tirx/script/ir_builder/ir.h b/tirx/include/tvm/tirx/script/ir_builder/ir.h
similarity index 100%
rename from include/tvm/tirx/script/ir_builder/ir.h
rename to tirx/include/tvm/tirx/script/ir_builder/ir.h
diff --git a/include/tvm/tirx/stmt.h b/tirx/include/tvm/tirx/stmt.h
similarity index 100%
rename from include/tvm/tirx/stmt.h
rename to tirx/include/tvm/tirx/stmt.h
diff --git a/include/tvm/tirx/stmt_functor.h b/tirx/include/tvm/tirx/stmt_functor.h
similarity index 100%
rename from include/tvm/tirx/stmt_functor.h
rename to tirx/include/tvm/tirx/stmt_functor.h
diff --git a/include/tvm/tirx/stmt_functor_with_analyzer.h b/tirx/include/tvm/tirx/stmt_functor_with_analyzer.h
similarity index 100%
rename from include/tvm/tirx/stmt_functor_with_analyzer.h
rename to tirx/include/tvm/tirx/stmt_functor_with_analyzer.h
diff --git a/include/tvm/tirx/tile/context.h b/tirx/include/tvm/tirx/tile/context.h
similarity index 100%
rename from include/tvm/tirx/tile/context.h
rename to tirx/include/tvm/tirx/tile/context.h
diff --git a/include/tvm/tirx/tile/dispatch.h b/tirx/include/tvm/tirx/tile/dispatch.h
similarity index 100%
rename from include/tvm/tirx/tile/dispatch.h
rename to tirx/include/tvm/tirx/tile/dispatch.h
diff --git a/include/tvm/tirx/tile/op.h b/tirx/include/tvm/tirx/tile/op.h
similarity index 100%
rename from include/tvm/tirx/tile/op.h
rename to tirx/include/tvm/tirx/tile/op.h
diff --git a/include/tvm/tirx/tile/thread_scope.h b/tirx/include/tvm/tirx/tile/thread_scope.h
similarity index 100%
rename from include/tvm/tirx/tile/thread_scope.h
rename to tirx/include/tvm/tirx/tile/thread_scope.h
diff --git a/include/tvm/tirx/transform.h b/tirx/include/tvm/tirx/transform.h
similarity index 100%
rename from include/tvm/tirx/transform.h
rename to tirx/include/tvm/tirx/transform.h
diff --git a/include/tvm/tirx/type.h b/tirx/include/tvm/tirx/type.h
similarity index 100%
rename from include/tvm/tirx/type.h
rename to tirx/include/tvm/tirx/type.h
diff --git a/licenses/LICENSE.builtin_fp16.txt b/tirx/licenses/LICENSE.builtin_fp16.txt
similarity index 100%
rename from licenses/LICENSE.builtin_fp16.txt
rename to tirx/licenses/LICENSE.builtin_fp16.txt
diff --git a/python/tvm/__init__.py b/tirx/python/tvm_tirx/__init__.py
similarity index 100%
rename from python/tvm/__init__.py
rename to tirx/python/tvm_tirx/__init__.py
diff --git a/python/tvm/backend/__init__.py b/tirx/python/tvm_tirx/backend/__init__.py
similarity index 100%
rename from python/tvm/backend/__init__.py
rename to tirx/python/tvm_tirx/backend/__init__.py
diff --git a/python/tvm/backend/_autoload_backends.py b/tirx/python/tvm_tirx/backend/_autoload_backends.py
similarity index 100%
rename from python/tvm/backend/_autoload_backends.py
rename to tirx/python/tvm_tirx/backend/_autoload_backends.py
diff --git a/python/tvm/backend/adreno/__init__.py b/tirx/python/tvm_tirx/backend/adreno/__init__.py
similarity index 100%
rename from python/tvm/backend/adreno/__init__.py
rename to tirx/python/tvm_tirx/backend/adreno/__init__.py
diff --git a/python/tvm/backend/adreno/target_tags.py b/tirx/python/tvm_tirx/backend/adreno/target_tags.py
similarity index 100%
rename from python/tvm/backend/adreno/target_tags.py
rename to tirx/python/tvm_tirx/backend/adreno/target_tags.py
diff --git a/python/tvm/backend/config.py b/tirx/python/tvm_tirx/backend/config.py
similarity index 100%
rename from python/tvm/backend/config.py
rename to tirx/python/tvm_tirx/backend/config.py
diff --git a/python/tvm/backend/cuda/__init__.py b/tirx/python/tvm_tirx/backend/cuda/__init__.py
similarity index 100%
rename from python/tvm/backend/cuda/__init__.py
rename to tirx/python/tvm_tirx/backend/cuda/__init__.py
diff --git a/python/tvm/backend/cuda/axis.py b/tirx/python/tvm_tirx/backend/cuda/axis.py
similarity index 100%
rename from python/tvm/backend/cuda/axis.py
rename to tirx/python/tvm_tirx/backend/cuda/axis.py
diff --git a/python/tvm/backend/cuda/backend_config.py b/tirx/python/tvm_tirx/backend/cuda/backend_config.py
similarity index 100%
rename from python/tvm/backend/cuda/backend_config.py
rename to tirx/python/tvm_tirx/backend/cuda/backend_config.py
diff --git a/python/tvm/backend/cuda/codegen/__init__.py b/tirx/python/tvm_tirx/backend/cuda/codegen/__init__.py
similarity index 100%
rename from python/tvm/backend/cuda/codegen/__init__.py
rename to tirx/python/tvm_tirx/backend/cuda/codegen/__init__.py
diff --git a/python/tvm/backend/cuda/codegen/header.py b/tirx/python/tvm_tirx/backend/cuda/codegen/header.py
similarity index 100%
rename from python/tvm/backend/cuda/codegen/header.py
rename to tirx/python/tvm_tirx/backend/cuda/codegen/header.py
diff --git a/python/tvm/backend/cuda/codegen/registry.py b/tirx/python/tvm_tirx/backend/cuda/codegen/registry.py
similarity index 100%
rename from python/tvm/backend/cuda/codegen/registry.py
rename to tirx/python/tvm_tirx/backend/cuda/codegen/registry.py
diff --git a/python/tvm/backend/cuda/codegen/schema.py b/tirx/python/tvm_tirx/backend/cuda/codegen/schema.py
similarity index 100%
rename from python/tvm/backend/cuda/codegen/schema.py
rename to tirx/python/tvm_tirx/backend/cuda/codegen/schema.py
diff --git a/python/tvm/backend/cuda/codegen/types.py b/tirx/python/tvm_tirx/backend/cuda/codegen/types.py
similarity index 100%
rename from python/tvm/backend/cuda/codegen/types.py
rename to tirx/python/tvm_tirx/backend/cuda/codegen/types.py
diff --git a/python/tvm/backend/cuda/codegen/utils.py b/tirx/python/tvm_tirx/backend/cuda/codegen/utils.py
similarity index 100%
rename from python/tvm/backend/cuda/codegen/utils.py
rename to tirx/python/tvm_tirx/backend/cuda/codegen/utils.py
diff --git a/python/tvm/backend/cuda/compiler.py b/tirx/python/tvm_tirx/backend/cuda/compiler.py
similarity index 100%
rename from python/tvm/backend/cuda/compiler.py
rename to tirx/python/tvm_tirx/backend/cuda/compiler.py
diff --git a/python/tvm/backend/cuda/cpp/__init__.py b/tirx/python/tvm_tirx/backend/cuda/cpp/__init__.py
similarity index 100%
rename from python/tvm/backend/cuda/cpp/__init__.py
rename to tirx/python/tvm_tirx/backend/cuda/cpp/__init__.py
diff --git a/python/tvm/backend/cuda/cpp/asm.py b/tirx/python/tvm_tirx/backend/cuda/cpp/asm.py
similarity index 100%
rename from python/tvm/backend/cuda/cpp/asm.py
rename to tirx/python/tvm_tirx/backend/cuda/cpp/asm.py
diff --git a/python/tvm/backend/cuda/cpp/builtins.py b/tirx/python/tvm_tirx/backend/cuda/cpp/builtins.py
similarity index 100%
rename from python/tvm/backend/cuda/cpp/builtins.py
rename to tirx/python/tvm_tirx/backend/cuda/cpp/builtins.py
diff --git a/python/tvm/backend/cuda/cpp/descriptors.py b/tirx/python/tvm_tirx/backend/cuda/cpp/descriptors.py
similarity index 100%
rename from python/tvm/backend/cuda/cpp/descriptors.py
rename to tirx/python/tvm_tirx/backend/cuda/cpp/descriptors.py
diff --git a/python/tvm/backend/cuda/cpp/instrument.py b/tirx/python/tvm_tirx/backend/cuda/cpp/instrument.py
similarity index 100%
rename from python/tvm/backend/cuda/cpp/instrument.py
rename to tirx/python/tvm_tirx/backend/cuda/cpp/instrument.py
diff --git a/python/tvm/backend/cuda/cpp/nvshmem.py b/tirx/python/tvm_tirx/backend/cuda/cpp/nvshmem.py
similarity index 100%
rename from python/tvm/backend/cuda/cpp/nvshmem.py
rename to tirx/python/tvm_tirx/backend/cuda/cpp/nvshmem.py
diff --git a/python/tvm/backend/cuda/host.py b/tirx/python/tvm_tirx/backend/cuda/host.py
similarity index 100%
rename from python/tvm/backend/cuda/host.py
rename to tirx/python/tvm_tirx/backend/cuda/host.py
diff --git a/python/tvm/backend/cuda/iket.py b/tirx/python/tvm_tirx/backend/cuda/iket.py
similarity index 100%
rename from python/tvm/backend/cuda/iket.py
rename to tirx/python/tvm_tirx/backend/cuda/iket.py
diff --git a/python/tvm/backend/cuda/lang/__init__.py b/tirx/python/tvm_tirx/backend/cuda/lang/__init__.py
similarity index 100%
rename from python/tvm/backend/cuda/lang/__init__.py
rename to tirx/python/tvm_tirx/backend/cuda/lang/__init__.py
diff --git a/python/tvm/backend/cuda/lang/alloc_pool.py b/tirx/python/tvm_tirx/backend/cuda/lang/alloc_pool.py
similarity index 100%
rename from python/tvm/backend/cuda/lang/alloc_pool.py
rename to tirx/python/tvm_tirx/backend/cuda/lang/alloc_pool.py
diff --git a/python/tvm/backend/cuda/lang/clc.py b/tirx/python/tvm_tirx/backend/cuda/lang/clc.py
similarity index 100%
rename from python/tvm/backend/cuda/lang/clc.py
rename to tirx/python/tvm_tirx/backend/cuda/lang/clc.py
diff --git a/python/tvm/backend/cuda/lang/pipeline.py b/tirx/python/tvm_tirx/backend/cuda/lang/pipeline.py
similarity index 100%
rename from python/tvm/backend/cuda/lang/pipeline.py
rename to tirx/python/tvm_tirx/backend/cuda/lang/pipeline.py
diff --git a/python/tvm/backend/cuda/lang/tile_scheduler.py b/tirx/python/tvm_tirx/backend/cuda/lang/tile_scheduler.py
similarity index 100%
rename from python/tvm/backend/cuda/lang/tile_scheduler.py
rename to tirx/python/tvm_tirx/backend/cuda/lang/tile_scheduler.py
diff --git a/python/tvm/backend/cuda/launch/__init__.py b/tirx/python/tvm_tirx/backend/cuda/launch/__init__.py
similarity index 100%
rename from python/tvm/backend/cuda/launch/__init__.py
rename to tirx/python/tvm_tirx/backend/cuda/launch/__init__.py
diff --git a/python/tvm/backend/cuda/launch/_generated.py b/tirx/python/tvm_tirx/backend/cuda/launch/_generated.py
similarity index 100%
rename from python/tvm/backend/cuda/launch/_generated.py
rename to tirx/python/tvm_tirx/backend/cuda/launch/_generated.py
diff --git a/python/tvm/backend/cuda/launch/_impl.py b/tirx/python/tvm_tirx/backend/cuda/launch/_impl.py
similarity index 100%
rename from python/tvm/backend/cuda/launch/_impl.py
rename to tirx/python/tvm_tirx/backend/cuda/launch/_impl.py
diff --git a/python/tvm/backend/cuda/launch/generate.py b/tirx/python/tvm_tirx/backend/cuda/launch/generate.py
similarity index 100%
rename from python/tvm/backend/cuda/launch/generate.py
rename to tirx/python/tvm_tirx/backend/cuda/launch/generate.py
diff --git a/python/tvm/backend/cuda/launch/table.py b/tirx/python/tvm_tirx/backend/cuda/launch/table.py
similarity index 100%
rename from python/tvm/backend/cuda/launch/table.py
rename to tirx/python/tvm_tirx/backend/cuda/launch/table.py
diff --git a/python/tvm/backend/cuda/layout.py b/tirx/python/tvm_tirx/backend/cuda/layout.py
similarity index 100%
rename from python/tvm/backend/cuda/layout.py
rename to tirx/python/tvm_tirx/backend/cuda/layout.py
diff --git a/python/tvm/backend/cuda/op.py b/tirx/python/tvm_tirx/backend/cuda/op.py
similarity index 100%
rename from python/tvm/backend/cuda/op.py
rename to tirx/python/tvm_tirx/backend/cuda/op.py
diff --git a/python/tvm/backend/cuda/ptx/__init__.py b/tirx/python/tvm_tirx/backend/cuda/ptx/__init__.py
similarity index 100%
rename from python/tvm/backend/cuda/ptx/__init__.py
rename to tirx/python/tvm_tirx/backend/cuda/ptx/__init__.py
diff --git a/python/tvm/backend/cuda/ptx/engine.py b/tirx/python/tvm_tirx/backend/cuda/ptx/engine.py
similarity index 100%
rename from python/tvm/backend/cuda/ptx/engine.py
rename to tirx/python/tvm_tirx/backend/cuda/ptx/engine.py
diff --git a/python/tvm/backend/cuda/ptx/gen_coverage.py b/tirx/python/tvm_tirx/backend/cuda/ptx/gen_coverage.py
similarity index 100%
rename from python/tvm/backend/cuda/ptx/gen_coverage.py
rename to tirx/python/tvm_tirx/backend/cuda/ptx/gen_coverage.py
diff --git a/python/tvm/backend/cuda/ptx/gen_helpers.py b/tirx/python/tvm_tirx/backend/cuda/ptx/gen_helpers.py
similarity index 100%
rename from python/tvm/backend/cuda/ptx/gen_helpers.py
rename to tirx/python/tvm_tirx/backend/cuda/ptx/gen_helpers.py
diff --git a/python/tvm/backend/cuda/ptx/gen_stubs.py b/tirx/python/tvm_tirx/backend/cuda/ptx/gen_stubs.py
similarity index 100%
rename from python/tvm/backend/cuda/ptx/gen_stubs.py
rename to tirx/python/tvm_tirx/backend/cuda/ptx/gen_stubs.py
diff --git a/python/tvm/backend/cuda/ptx/render.py b/tirx/python/tvm_tirx/backend/cuda/ptx/render.py
similarity index 100%
rename from python/tvm/backend/cuda/ptx/render.py
rename to tirx/python/tvm_tirx/backend/cuda/ptx/render.py
diff --git a/python/tvm/backend/cuda/ptx/table.py b/tirx/python/tvm_tirx/backend/cuda/ptx/table.py
similarity index 100%
rename from python/tvm/backend/cuda/ptx/table.py
rename to tirx/python/tvm_tirx/backend/cuda/ptx/table.py
diff --git a/python/tvm/backend/cuda/script.py b/tirx/python/tvm_tirx/backend/cuda/script.py
similarity index 100%
rename from python/tvm/backend/cuda/script.py
rename to tirx/python/tvm_tirx/backend/cuda/script.py
diff --git a/python/tvm/backend/cuda/target_tags.py b/tirx/python/tvm_tirx/backend/cuda/target_tags.py
similarity index 100%
rename from python/tvm/backend/cuda/target_tags.py
rename to tirx/python/tvm_tirx/backend/cuda/target_tags.py
diff --git a/python/tvm/backend/cuda/thread_scope.py b/tirx/python/tvm_tirx/backend/cuda/thread_scope.py
similarity index 100%
rename from python/tvm/backend/cuda/thread_scope.py
rename to tirx/python/tvm_tirx/backend/cuda/thread_scope.py
diff --git a/python/tvm/backend/cuda/tile/__init__.py b/tirx/python/tvm_tirx/backend/cuda/tile/__init__.py
similarity index 100%
rename from python/tvm/backend/cuda/tile/__init__.py
rename to tirx/python/tvm_tirx/backend/cuda/tile/__init__.py
diff --git a/python/tvm/backend/cuda/tile/exec_context.py b/tirx/python/tvm_tirx/backend/cuda/tile/exec_context.py
similarity index 100%
rename from python/tvm/backend/cuda/tile/exec_context.py
rename to tirx/python/tvm_tirx/backend/cuda/tile/exec_context.py
diff --git a/python/tvm/backend/cuda/tile/instructions.py b/tirx/python/tvm_tirx/backend/cuda/tile/instructions.py
similarity index 100%
rename from python/tvm/backend/cuda/tile/instructions.py
rename to tirx/python/tvm_tirx/backend/cuda/tile/instructions.py
diff --git a/python/tvm/backend/cuda/tile/patterns/__init__.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/__init__.py
similarity index 100%
rename from python/tvm/backend/cuda/tile/patterns/__init__.py
rename to tirx/python/tvm_tirx/backend/cuda/tile/patterns/__init__.py
diff --git a/python/tvm/backend/cuda/tile/patterns/async_copy.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/async_copy.py
similarity index 100%
rename from python/tvm/backend/cuda/tile/patterns/async_copy.py
rename to tirx/python/tvm_tirx/backend/cuda/tile/patterns/async_copy.py
diff --git a/python/tvm/backend/cuda/tile/patterns/common.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py
similarity index 100%
rename from python/tvm/backend/cuda/tile/patterns/common.py
rename to tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py
diff --git a/python/tvm/backend/cuda/tile/patterns/descriptors.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/descriptors.py
similarity index 100%
rename from python/tvm/backend/cuda/tile/patterns/descriptors.py
rename to tirx/python/tvm_tirx/backend/cuda/tile/patterns/descriptors.py
diff --git a/python/tvm/backend/cuda/tile/patterns/ewise.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/ewise.py
similarity index 100%
rename from python/tvm/backend/cuda/tile/patterns/ewise.py
rename to tirx/python/tvm_tirx/backend/cuda/tile/patterns/ewise.py
diff --git a/python/tvm/backend/cuda/tile/patterns/ldmatrix.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/ldmatrix.py
similarity index 100%
rename from python/tvm/backend/cuda/tile/patterns/ldmatrix.py
rename to tirx/python/tvm_tirx/backend/cuda/tile/patterns/ldmatrix.py
diff --git a/python/tvm/backend/cuda/tile/patterns/memory.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/memory.py
similarity index 100%
rename from python/tvm/backend/cuda/tile/patterns/memory.py
rename to tirx/python/tvm_tirx/backend/cuda/tile/patterns/memory.py
diff --git a/python/tvm/backend/cuda/tile/patterns/mma.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/mma.py
similarity index 100%
rename from python/tvm/backend/cuda/tile/patterns/mma.py
rename to tirx/python/tvm_tirx/backend/cuda/tile/patterns/mma.py
diff --git a/python/tvm/backend/cuda/tile/patterns/register.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/register.py
similarity index 100%
rename from python/tvm/backend/cuda/tile/patterns/register.py
rename to tirx/python/tvm_tirx/backend/cuda/tile/patterns/register.py
diff --git a/python/tvm/backend/cuda/tile/patterns/tcgen05_mma.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tcgen05_mma.py
similarity index 100%
rename from python/tvm/backend/cuda/tile/patterns/tcgen05_mma.py
rename to tirx/python/tvm_tirx/backend/cuda/tile/patterns/tcgen05_mma.py
diff --git a/python/tvm/backend/cuda/tile/patterns/tma.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py
similarity index 100%
rename from python/tvm/backend/cuda/tile/patterns/tma.py
rename to tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py
diff --git a/python/tvm/backend/cuda/tile/patterns/tmem.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tmem.py
similarity index 100%
rename from python/tvm/backend/cuda/tile/patterns/tmem.py
rename to tirx/python/tvm_tirx/backend/cuda/tile/patterns/tmem.py
diff --git a/python/tvm/backend/cuda/tile/patterns/tmem_copy.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tmem_copy.py
similarity index 100%
rename from python/tvm/backend/cuda/tile/patterns/tmem_copy.py
rename to tirx/python/tvm_tirx/backend/cuda/tile/patterns/tmem_copy.py
diff --git a/python/tvm/backend/cuda/transforms/__init__.py b/tirx/python/tvm_tirx/backend/cuda/transforms/__init__.py
similarity index 100%
rename from python/tvm/backend/cuda/transforms/__init__.py
rename to tirx/python/tvm_tirx/backend/cuda/transforms/__init__.py
diff --git a/python/tvm/backend/hexagon/__init__.py b/tirx/python/tvm_tirx/backend/hexagon/__init__.py
similarity index 100%
rename from python/tvm/backend/hexagon/__init__.py
rename to tirx/python/tvm_tirx/backend/hexagon/__init__.py
diff --git a/python/tvm/backend/hexagon/linker.py b/tirx/python/tvm_tirx/backend/hexagon/linker.py
similarity index 100%
rename from python/tvm/backend/hexagon/linker.py
rename to tirx/python/tvm_tirx/backend/hexagon/linker.py
diff --git a/python/tvm/backend/hexagon/target_tags.py b/tirx/python/tvm_tirx/backend/hexagon/target_tags.py
similarity index 100%
rename from python/tvm/backend/hexagon/target_tags.py
rename to tirx/python/tvm_tirx/backend/hexagon/target_tags.py
diff --git a/python/tvm/backend/loader.py b/tirx/python/tvm_tirx/backend/loader.py
similarity index 100%
rename from python/tvm/backend/loader.py
rename to tirx/python/tvm_tirx/backend/loader.py
diff --git a/python/tvm/backend/metal/__init__.py b/tirx/python/tvm_tirx/backend/metal/__init__.py
similarity index 100%
rename from python/tvm/backend/metal/__init__.py
rename to tirx/python/tvm_tirx/backend/metal/__init__.py
diff --git a/python/tvm/backend/metal/op.py b/tirx/python/tvm_tirx/backend/metal/op.py
similarity index 100%
rename from python/tvm/backend/metal/op.py
rename to tirx/python/tvm_tirx/backend/metal/op.py
diff --git a/python/tvm/backend/metal/script.py b/tirx/python/tvm_tirx/backend/metal/script.py
similarity index 100%
rename from python/tvm/backend/metal/script.py
rename to tirx/python/tvm_tirx/backend/metal/script.py
diff --git a/python/tvm/backend/metal/target_tags.py b/tirx/python/tvm_tirx/backend/metal/target_tags.py
similarity index 100%
rename from python/tvm/backend/metal/target_tags.py
rename to tirx/python/tvm_tirx/backend/metal/target_tags.py
diff --git a/python/tvm/backend/opencl/__init__.py b/tirx/python/tvm_tirx/backend/opencl/__init__.py
similarity index 100%
rename from python/tvm/backend/opencl/__init__.py
rename to tirx/python/tvm_tirx/backend/opencl/__init__.py
diff --git a/python/tvm/backend/opencl/script.py b/tirx/python/tvm_tirx/backend/opencl/script.py
similarity index 100%
rename from python/tvm/backend/opencl/script.py
rename to tirx/python/tvm_tirx/backend/opencl/script.py
diff --git a/python/tvm/backend/rocm/__init__.py b/tirx/python/tvm_tirx/backend/rocm/__init__.py
similarity index 100%
rename from python/tvm/backend/rocm/__init__.py
rename to tirx/python/tvm_tirx/backend/rocm/__init__.py
diff --git a/python/tvm/backend/trn/__init__.py b/tirx/python/tvm_tirx/backend/trn/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/__init__.py
rename to tirx/python/tvm_tirx/backend/trn/__init__.py
diff --git a/python/tvm/backend/trn/axis.py b/tirx/python/tvm_tirx/backend/trn/axis.py
similarity index 100%
rename from python/tvm/backend/trn/axis.py
rename to tirx/python/tvm_tirx/backend/trn/axis.py
diff --git a/python/tvm/backend/trn/layout.py b/tirx/python/tvm_tirx/backend/trn/layout.py
similarity index 100%
rename from python/tvm/backend/trn/layout.py
rename to tirx/python/tvm_tirx/backend/trn/layout.py
diff --git a/python/tvm/backend/trn/op.py b/tirx/python/tvm_tirx/backend/trn/op.py
similarity index 100%
rename from python/tvm/backend/trn/op.py
rename to tirx/python/tvm_tirx/backend/trn/op.py
diff --git a/python/tvm/backend/trn/pipeline.py b/tirx/python/tvm_tirx/backend/trn/pipeline.py
similarity index 100%
rename from python/tvm/backend/trn/pipeline.py
rename to tirx/python/tvm_tirx/backend/trn/pipeline.py
diff --git a/python/tvm/backend/trn/script.py b/tirx/python/tvm_tirx/backend/trn/script.py
similarity index 100%
rename from python/tvm/backend/trn/script.py
rename to tirx/python/tvm_tirx/backend/trn/script.py
diff --git a/python/tvm/backend/trn/target_tags.py b/tirx/python/tvm_tirx/backend/trn/target_tags.py
similarity index 100%
rename from python/tvm/backend/trn/target_tags.py
rename to tirx/python/tvm_tirx/backend/trn/target_tags.py
diff --git a/python/tvm/backend/trn/thread_scope.py b/tirx/python/tvm_tirx/backend/trn/thread_scope.py
similarity index 100%
rename from python/tvm/backend/trn/thread_scope.py
rename to tirx/python/tvm_tirx/backend/trn/thread_scope.py
diff --git a/python/tvm/backend/trn/tile/__init__.py b/tirx/python/tvm_tirx/backend/trn/tile/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/tile/__init__.py
rename to tirx/python/tvm_tirx/backend/trn/tile/__init__.py
diff --git a/python/tvm/backend/trn/tile/binary/__init__.py b/tirx/python/tvm_tirx/backend/trn/tile/binary/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/tile/binary/__init__.py
rename to tirx/python/tvm_tirx/backend/trn/tile/binary/__init__.py
diff --git a/python/tvm/backend/trn/tile/binary/default.py b/tirx/python/tvm_tirx/backend/trn/tile/binary/default.py
similarity index 100%
rename from python/tvm/backend/trn/tile/binary/default.py
rename to tirx/python/tvm_tirx/backend/trn/tile/binary/default.py
diff --git a/python/tvm/backend/trn/tile/binary/utils.py b/tirx/python/tvm_tirx/backend/trn/tile/binary/utils.py
similarity index 100%
rename from python/tvm/backend/trn/tile/binary/utils.py
rename to tirx/python/tvm_tirx/backend/trn/tile/binary/utils.py
diff --git a/python/tvm/backend/trn/tile/common.py b/tirx/python/tvm_tirx/backend/trn/tile/common.py
similarity index 100%
rename from python/tvm/backend/trn/tile/common.py
rename to tirx/python/tvm_tirx/backend/trn/tile/common.py
diff --git a/python/tvm/backend/trn/tile/compose_op/__init__.py b/tirx/python/tvm_tirx/backend/trn/tile/compose_op/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/tile/compose_op/__init__.py
rename to tirx/python/tvm_tirx/backend/trn/tile/compose_op/__init__.py
diff --git a/python/tvm/backend/trn/tile/compose_op/binary_chain.py b/tirx/python/tvm_tirx/backend/trn/tile/compose_op/binary_chain.py
similarity index 100%
rename from python/tvm/backend/trn/tile/compose_op/binary_chain.py
rename to tirx/python/tvm_tirx/backend/trn/tile/compose_op/binary_chain.py
diff --git a/python/tvm/backend/trn/tile/compose_op/binary_reduce.py b/tirx/python/tvm_tirx/backend/trn/tile/compose_op/binary_reduce.py
similarity index 100%
rename from python/tvm/backend/trn/tile/compose_op/binary_reduce.py
rename to tirx/python/tvm_tirx/backend/trn/tile/compose_op/binary_reduce.py
diff --git a/python/tvm/backend/trn/tile/compose_op/unary_reduce.py b/tirx/python/tvm_tirx/backend/trn/tile/compose_op/unary_reduce.py
similarity index 100%
rename from python/tvm/backend/trn/tile/compose_op/unary_reduce.py
rename to tirx/python/tvm_tirx/backend/trn/tile/compose_op/unary_reduce.py
diff --git a/python/tvm/backend/trn/tile/compose_op/utils.py b/tirx/python/tvm_tirx/backend/trn/tile/compose_op/utils.py
similarity index 100%
rename from python/tvm/backend/trn/tile/compose_op/utils.py
rename to tirx/python/tvm_tirx/backend/trn/tile/compose_op/utils.py
diff --git a/python/tvm/backend/trn/tile/copy/__init__.py b/tirx/python/tvm_tirx/backend/trn/tile/copy/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/tile/copy/__init__.py
rename to tirx/python/tvm_tirx/backend/trn/tile/copy/__init__.py
diff --git a/python/tvm/backend/trn/tile/copy/default.py b/tirx/python/tvm_tirx/backend/trn/tile/copy/default.py
similarity index 100%
rename from python/tvm/backend/trn/tile/copy/default.py
rename to tirx/python/tvm_tirx/backend/trn/tile/copy/default.py
diff --git a/python/tvm/backend/trn/tile/dim_utils.py b/tirx/python/tvm_tirx/backend/trn/tile/dim_utils.py
similarity index 100%
rename from python/tvm/backend/trn/tile/dim_utils.py
rename to tirx/python/tvm_tirx/backend/trn/tile/dim_utils.py
diff --git a/python/tvm/backend/trn/tile/gemm/__init__.py b/tirx/python/tvm_tirx/backend/trn/tile/gemm/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/tile/gemm/__init__.py
rename to tirx/python/tvm_tirx/backend/trn/tile/gemm/__init__.py
diff --git a/python/tvm/backend/trn/tile/gemm/default.py b/tirx/python/tvm_tirx/backend/trn/tile/gemm/default.py
similarity index 100%
rename from python/tvm/backend/trn/tile/gemm/default.py
rename to tirx/python/tvm_tirx/backend/trn/tile/gemm/default.py
diff --git a/python/tvm/backend/trn/tile/instruction_generator.py b/tirx/python/tvm_tirx/backend/trn/tile/instruction_generator.py
similarity index 100%
rename from python/tvm/backend/trn/tile/instruction_generator.py
rename to tirx/python/tvm_tirx/backend/trn/tile/instruction_generator.py
diff --git a/python/tvm/backend/trn/tile/instructions.py b/tirx/python/tvm_tirx/backend/trn/tile/instructions.py
similarity index 100%
rename from python/tvm/backend/trn/tile/instructions.py
rename to tirx/python/tvm_tirx/backend/trn/tile/instructions.py
diff --git a/python/tvm/backend/trn/tile/private_alloc.py b/tirx/python/tvm_tirx/backend/trn/tile/private_alloc.py
similarity index 100%
rename from python/tvm/backend/trn/tile/private_alloc.py
rename to tirx/python/tvm_tirx/backend/trn/tile/private_alloc.py
diff --git a/python/tvm/backend/trn/tile/reduction/__init__.py b/tirx/python/tvm_tirx/backend/trn/tile/reduction/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/tile/reduction/__init__.py
rename to tirx/python/tvm_tirx/backend/trn/tile/reduction/__init__.py
diff --git a/python/tvm/backend/trn/tile/reduction/utils.py b/tirx/python/tvm_tirx/backend/trn/tile/reduction/utils.py
similarity index 100%
rename from python/tvm/backend/trn/tile/reduction/utils.py
rename to tirx/python/tvm_tirx/backend/trn/tile/reduction/utils.py
diff --git a/python/tvm/backend/trn/tile/select/__init__.py b/tirx/python/tvm_tirx/backend/trn/tile/select/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/tile/select/__init__.py
rename to tirx/python/tvm_tirx/backend/trn/tile/select/__init__.py
diff --git a/python/tvm/backend/trn/tile/select/default.py b/tirx/python/tvm_tirx/backend/trn/tile/select/default.py
similarity index 100%
rename from python/tvm/backend/trn/tile/select/default.py
rename to tirx/python/tvm_tirx/backend/trn/tile/select/default.py
diff --git a/python/tvm/backend/trn/tile/unary/__init__.py b/tirx/python/tvm_tirx/backend/trn/tile/unary/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/tile/unary/__init__.py
rename to tirx/python/tvm_tirx/backend/trn/tile/unary/__init__.py
diff --git a/python/tvm/backend/trn/tile/unary/default.py b/tirx/python/tvm_tirx/backend/trn/tile/unary/default.py
similarity index 100%
rename from python/tvm/backend/trn/tile/unary/default.py
rename to tirx/python/tvm_tirx/backend/trn/tile/unary/default.py
diff --git a/python/tvm/backend/trn/tile/unary/utils.py b/tirx/python/tvm_tirx/backend/trn/tile/unary/utils.py
similarity index 100%
rename from python/tvm/backend/trn/tile/unary/utils.py
rename to tirx/python/tvm_tirx/backend/trn/tile/unary/utils.py
diff --git a/python/tvm/backend/trn/tile/unary/with_bias_scale.py b/tirx/python/tvm_tirx/backend/trn/tile/unary/with_bias_scale.py
similarity index 100%
rename from python/tvm/backend/trn/tile/unary/with_bias_scale.py
rename to tirx/python/tvm_tirx/backend/trn/tile/unary/with_bias_scale.py
diff --git a/python/tvm/backend/trn/tile/workspace_utils.py b/tirx/python/tvm_tirx/backend/trn/tile/workspace_utils.py
similarity index 100%
rename from python/tvm/backend/trn/tile/workspace_utils.py
rename to tirx/python/tvm_tirx/backend/trn/tile/workspace_utils.py
diff --git a/python/tvm/backend/trn/transform/__init__.py b/tirx/python/tvm_tirx/backend/trn/transform/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/transform/__init__.py
rename to tirx/python/tvm_tirx/backend/trn/transform/__init__.py
diff --git a/python/tvm/backend/trn/transform/naive_allocator.py b/tirx/python/tvm_tirx/backend/trn/transform/naive_allocator.py
similarity index 100%
rename from python/tvm/backend/trn/transform/naive_allocator.py
rename to tirx/python/tvm_tirx/backend/trn/transform/naive_allocator.py
diff --git a/python/tvm/backend/trn/transform/private_tensor_alloc.py b/tirx/python/tvm_tirx/backend/trn/transform/private_tensor_alloc.py
similarity index 100%
rename from python/tvm/backend/trn/transform/private_tensor_alloc.py
rename to tirx/python/tvm_tirx/backend/trn/transform/private_tensor_alloc.py
diff --git a/python/tvm/backend/vulkan/__init__.py b/tirx/python/tvm_tirx/backend/vulkan/__init__.py
similarity index 100%
rename from python/tvm/backend/vulkan/__init__.py
rename to tirx/python/tvm_tirx/backend/vulkan/__init__.py
diff --git a/python/tvm/backend/webgpu/__init__.py b/tirx/python/tvm_tirx/backend/webgpu/__init__.py
similarity index 100%
rename from python/tvm/backend/webgpu/__init__.py
rename to tirx/python/tvm_tirx/backend/webgpu/__init__.py
diff --git a/python/tvm/base.py b/tirx/python/tvm_tirx/base.py
similarity index 100%
rename from python/tvm/base.py
rename to tirx/python/tvm_tirx/base.py
diff --git a/python/tvm/codegen.py b/tirx/python/tvm_tirx/codegen.py
similarity index 100%
rename from python/tvm/codegen.py
rename to tirx/python/tvm_tirx/codegen.py
diff --git a/python/tvm/error.py b/tirx/python/tvm_tirx/error.py
similarity index 100%
rename from python/tvm/error.py
rename to tirx/python/tvm_tirx/error.py
diff --git a/python/tvm/ir/__init__.py b/tirx/python/tvm_tirx/ir/__init__.py
similarity index 100%
rename from python/tvm/ir/__init__.py
rename to tirx/python/tvm_tirx/ir/__init__.py
diff --git a/python/tvm/ir/_constant.py b/tirx/python/tvm_tirx/ir/_constant.py
similarity index 100%
rename from python/tvm/ir/_constant.py
rename to tirx/python/tvm_tirx/ir/_constant.py
diff --git a/python/tvm/ir/_ffi_api.py b/tirx/python/tvm_tirx/ir/_ffi_api.py
similarity index 100%
rename from python/tvm/ir/_ffi_api.py
rename to tirx/python/tvm_tirx/ir/_ffi_api.py
diff --git a/python/tvm/ir/_overload_prim_expr.py b/tirx/python/tvm_tirx/ir/_overload_prim_expr.py
similarity index 100%
rename from python/tvm/ir/_overload_prim_expr.py
rename to tirx/python/tvm_tirx/ir/_overload_prim_expr.py
diff --git a/python/tvm/ir/_tensor_expr_overload.py b/tirx/python/tvm_tirx/ir/_tensor_expr_overload.py
similarity index 100%
rename from python/tvm/ir/_tensor_expr_overload.py
rename to tirx/python/tvm_tirx/ir/_tensor_expr_overload.py
diff --git a/python/tvm/ir/attrs.py b/tirx/python/tvm_tirx/ir/attrs.py
similarity index 100%
rename from python/tvm/ir/attrs.py
rename to tirx/python/tvm_tirx/ir/attrs.py
diff --git a/python/tvm/ir/base.py b/tirx/python/tvm_tirx/ir/base.py
similarity index 100%
rename from python/tvm/ir/base.py
rename to tirx/python/tvm_tirx/ir/base.py
diff --git a/python/tvm/ir/expr.py b/tirx/python/tvm_tirx/ir/expr.py
similarity index 100%
rename from python/tvm/ir/expr.py
rename to tirx/python/tvm_tirx/ir/expr.py
diff --git a/python/tvm/ir/function.py b/tirx/python/tvm_tirx/ir/function.py
similarity index 100%
rename from python/tvm/ir/function.py
rename to tirx/python/tvm_tirx/ir/function.py
diff --git a/python/tvm/ir/global_info.py b/tirx/python/tvm_tirx/ir/global_info.py
similarity index 100%
rename from python/tvm/ir/global_info.py
rename to tirx/python/tvm_tirx/ir/global_info.py
diff --git a/python/tvm/ir/json_compact.py b/tirx/python/tvm_tirx/ir/json_compact.py
similarity index 100%
rename from python/tvm/ir/json_compact.py
rename to tirx/python/tvm_tirx/ir/json_compact.py
diff --git a/python/tvm/ir/location.py b/tirx/python/tvm_tirx/ir/location.py
similarity index 100%
rename from python/tvm/ir/location.py
rename to tirx/python/tvm_tirx/ir/location.py
diff --git a/python/tvm/ir/module.py b/tirx/python/tvm_tirx/ir/module.py
similarity index 100%
rename from python/tvm/ir/module.py
rename to tirx/python/tvm_tirx/ir/module.py
diff --git a/python/tvm/ir/op.py b/tirx/python/tvm_tirx/ir/op.py
similarity index 100%
rename from python/tvm/ir/op.py
rename to tirx/python/tvm_tirx/ir/op.py
diff --git a/python/tvm/ir/prim/__init__.py b/tirx/python/tvm_tirx/ir/prim/__init__.py
similarity index 100%
rename from python/tvm/ir/prim/__init__.py
rename to tirx/python/tvm_tirx/ir/prim/__init__.py
diff --git a/python/tvm/ir/prim/_ffi_api.py b/tirx/python/tvm_tirx/ir/prim/_ffi_api.py
similarity index 100%
rename from python/tvm/ir/prim/_ffi_api.py
rename to tirx/python/tvm_tirx/ir/prim/_ffi_api.py
diff --git a/python/tvm/ir/prim/expr.py b/tirx/python/tvm_tirx/ir/prim/expr.py
similarity index 100%
rename from python/tvm/ir/prim/expr.py
rename to tirx/python/tvm_tirx/ir/prim/expr.py
diff --git a/python/tvm/ir/prim/op.py b/tirx/python/tvm_tirx/ir/prim/op.py
similarity index 100%
rename from python/tvm/ir/prim/op.py
rename to tirx/python/tvm_tirx/ir/prim/op.py
diff --git a/python/tvm/ir/stmt.py b/tirx/python/tvm_tirx/ir/stmt.py
similarity index 100%
rename from python/tvm/ir/stmt.py
rename to tirx/python/tvm_tirx/ir/stmt.py
diff --git a/python/tvm/ir/supply.py b/tirx/python/tvm_tirx/ir/supply.py
similarity index 100%
rename from python/tvm/ir/supply.py
rename to tirx/python/tvm_tirx/ir/supply.py
diff --git a/python/tvm/ir/transform/__init__.py b/tirx/python/tvm_tirx/ir/transform/__init__.py
similarity index 100%
rename from python/tvm/ir/transform/__init__.py
rename to tirx/python/tvm_tirx/ir/transform/__init__.py
diff --git a/python/tvm/ir/transform/_ffi_instrument_api.py b/tirx/python/tvm_tirx/ir/transform/_ffi_instrument_api.py
similarity index 100%
rename from python/tvm/ir/transform/_ffi_instrument_api.py
rename to tirx/python/tvm_tirx/ir/transform/_ffi_instrument_api.py
diff --git a/python/tvm/ir/transform/_ffi_transform_api.py b/tirx/python/tvm_tirx/ir/transform/_ffi_transform_api.py
similarity index 100%
rename from python/tvm/ir/transform/_ffi_transform_api.py
rename to tirx/python/tvm_tirx/ir/transform/_ffi_transform_api.py
diff --git a/python/tvm/ir/transform/core.py b/tirx/python/tvm_tirx/ir/transform/core.py
similarity index 100%
rename from python/tvm/ir/transform/core.py
rename to tirx/python/tvm_tirx/ir/transform/core.py
diff --git a/python/tvm/ir/transform/instrument.py b/tirx/python/tvm_tirx/ir/transform/instrument.py
similarity index 100%
rename from python/tvm/ir/transform/instrument.py
rename to tirx/python/tvm_tirx/ir/transform/instrument.py
diff --git a/python/tvm/ir/type.py b/tirx/python/tvm_tirx/ir/type.py
similarity index 100%
rename from python/tvm/ir/type.py
rename to tirx/python/tvm_tirx/ir/type.py
diff --git a/python/tvm/libinfo.py b/tirx/python/tvm_tirx/libinfo.py
similarity index 100%
rename from python/tvm/libinfo.py
rename to tirx/python/tvm_tirx/libinfo.py
diff --git a/python/tvm/runtime/__init__.py b/tirx/python/tvm_tirx/runtime/__init__.py
similarity index 100%
rename from python/tvm/runtime/__init__.py
rename to tirx/python/tvm_tirx/runtime/__init__.py
diff --git a/python/tvm/runtime/_ffi_api.py b/tirx/python/tvm_tirx/runtime/_ffi_api.py
similarity index 100%
rename from python/tvm/runtime/_ffi_api.py
rename to tirx/python/tvm_tirx/runtime/_ffi_api.py
diff --git a/python/tvm/runtime/_tensor.py b/tirx/python/tvm_tirx/runtime/_tensor.py
similarity index 100%
rename from python/tvm/runtime/_tensor.py
rename to tirx/python/tvm_tirx/runtime/_tensor.py
diff --git a/python/tvm/runtime/device.py b/tirx/python/tvm_tirx/runtime/device.py
similarity index 100%
rename from python/tvm/runtime/device.py
rename to tirx/python/tvm_tirx/runtime/device.py
diff --git a/python/tvm/runtime/executable.py b/tirx/python/tvm_tirx/runtime/executable.py
similarity index 100%
rename from python/tvm/runtime/executable.py
rename to tirx/python/tvm_tirx/runtime/executable.py
diff --git a/python/tvm/runtime/memory_scope.py b/tirx/python/tvm_tirx/runtime/memory_scope.py
similarity index 100%
rename from python/tvm/runtime/memory_scope.py
rename to tirx/python/tvm_tirx/runtime/memory_scope.py
diff --git a/python/tvm/runtime/module.py b/tirx/python/tvm_tirx/runtime/module.py
similarity index 100%
rename from python/tvm/runtime/module.py
rename to tirx/python/tvm_tirx/runtime/module.py
diff --git a/python/tvm/script/__init__.py b/tirx/python/tvm_tirx/script/__init__.py
similarity index 100%
rename from python/tvm/script/__init__.py
rename to tirx/python/tvm_tirx/script/__init__.py
diff --git a/python/tvm/script/_ffi_api.py b/tirx/python/tvm_tirx/script/_ffi_api.py
similarity index 100%
rename from python/tvm/script/_ffi_api.py
rename to tirx/python/tvm_tirx/script/_ffi_api.py
diff --git a/python/tvm/script/ir_builder/__init__.py b/tirx/python/tvm_tirx/script/ir_builder/__init__.py
similarity index 100%
rename from python/tvm/script/ir_builder/__init__.py
rename to tirx/python/tvm_tirx/script/ir_builder/__init__.py
diff --git a/python/tvm/script/ir_builder/_ffi_api.py b/tirx/python/tvm_tirx/script/ir_builder/_ffi_api.py
similarity index 100%
rename from python/tvm/script/ir_builder/_ffi_api.py
rename to tirx/python/tvm_tirx/script/ir_builder/_ffi_api.py
diff --git a/python/tvm/script/ir_builder/base.py b/tirx/python/tvm_tirx/script/ir_builder/base.py
similarity index 100%
rename from python/tvm/script/ir_builder/base.py
rename to tirx/python/tvm_tirx/script/ir_builder/base.py
diff --git a/python/tvm/script/ir_builder/frame.py b/tirx/python/tvm_tirx/script/ir_builder/frame.py
similarity index 100%
rename from python/tvm/script/ir_builder/frame.py
rename to tirx/python/tvm_tirx/script/ir_builder/frame.py
diff --git a/python/tvm/script/ir_builder/ir.py b/tirx/python/tvm_tirx/script/ir_builder/ir.py
similarity index 100%
rename from python/tvm/script/ir_builder/ir.py
rename to tirx/python/tvm_tirx/script/ir_builder/ir.py
diff --git a/python/tvm/script/ir_builder/parser_protocol.py b/tirx/python/tvm_tirx/script/ir_builder/parser_protocol.py
similarity index 100%
rename from python/tvm/script/ir_builder/parser_protocol.py
rename to tirx/python/tvm_tirx/script/ir_builder/parser_protocol.py
diff --git a/python/tvm/script/ir_builder/stmt.py b/tirx/python/tvm_tirx/script/ir_builder/stmt.py
similarity index 100%
rename from python/tvm/script/ir_builder/stmt.py
rename to tirx/python/tvm_tirx/script/ir_builder/stmt.py
diff --git a/python/tvm/script/ir_builder/tracing.py b/tirx/python/tvm_tirx/script/ir_builder/tracing.py
similarity index 100%
rename from python/tvm/script/ir_builder/tracing.py
rename to tirx/python/tvm_tirx/script/ir_builder/tracing.py
diff --git a/python/tvm/script/parser/__init__.py b/tirx/python/tvm_tirx/script/parser/__init__.py
similarity index 100%
rename from python/tvm/script/parser/__init__.py
rename to tirx/python/tvm_tirx/script/parser/__init__.py
diff --git a/python/tvm/script/parser/entry.py b/tirx/python/tvm_tirx/script/parser/entry.py
similarity index 100%
rename from python/tvm/script/parser/entry.py
rename to tirx/python/tvm_tirx/script/parser/entry.py
diff --git a/python/tvm/script/parser/inspect_source.py b/tirx/python/tvm_tirx/script/parser/inspect_source.py
similarity index 100%
rename from python/tvm/script/parser/inspect_source.py
rename to tirx/python/tvm_tirx/script/parser/inspect_source.py
diff --git a/python/tvm/script/parser/ir.py b/tirx/python/tvm_tirx/script/parser/ir.py
similarity index 100%
rename from python/tvm/script/parser/ir.py
rename to tirx/python/tvm_tirx/script/parser/ir.py
diff --git a/python/tvm/script/parser/jit_support.py b/tirx/python/tvm_tirx/script/parser/jit_support.py
similarity index 100%
rename from python/tvm/script/parser/jit_support.py
rename to tirx/python/tvm_tirx/script/parser/jit_support.py
diff --git a/python/tvm/script/parser/prescan.py b/tirx/python/tvm_tirx/script/parser/prescan.py
similarity index 100%
rename from python/tvm/script/parser/prescan.py
rename to tirx/python/tvm_tirx/script/parser/prescan.py
diff --git a/python/tvm/script/parser/protocol_registry.py b/tirx/python/tvm_tirx/script/parser/protocol_registry.py
similarity index 100%
rename from python/tvm/script/parser/protocol_registry.py
rename to tirx/python/tvm_tirx/script/parser/protocol_registry.py
diff --git a/python/tvm/script/parser/transpile.py b/tirx/python/tvm_tirx/script/parser/transpile.py
similarity index 100%
rename from python/tvm/script/parser/transpile.py
rename to tirx/python/tvm_tirx/script/parser/transpile.py
diff --git a/python/tvm/script/printer/__init__.py b/tirx/python/tvm_tirx/script/printer/__init__.py
similarity index 100%
rename from python/tvm/script/printer/__init__.py
rename to tirx/python/tvm_tirx/script/printer/__init__.py
diff --git a/python/tvm/script/printer/_ffi_api.py b/tirx/python/tvm_tirx/script/printer/_ffi_api.py
similarity index 100%
rename from python/tvm/script/printer/_ffi_api.py
rename to tirx/python/tvm_tirx/script/printer/_ffi_api.py
diff --git a/python/tvm/script/printer/doc.py b/tirx/python/tvm_tirx/script/printer/doc.py
similarity index 100%
rename from python/tvm/script/printer/doc.py
rename to tirx/python/tvm_tirx/script/printer/doc.py
diff --git a/python/tvm/script/printer/doc_printer.py b/tirx/python/tvm_tirx/script/printer/doc_printer.py
similarity index 100%
rename from python/tvm/script/printer/doc_printer.py
rename to tirx/python/tvm_tirx/script/printer/doc_printer.py
diff --git a/python/tvm/script/printer/highlight.py b/tirx/python/tvm_tirx/script/printer/highlight.py
similarity index 100%
rename from python/tvm/script/printer/highlight.py
rename to tirx/python/tvm_tirx/script/printer/highlight.py
diff --git a/python/tvm/script/printer/scriptable.py b/tirx/python/tvm_tirx/script/printer/scriptable.py
similarity index 100%
rename from python/tvm/script/printer/scriptable.py
rename to tirx/python/tvm_tirx/script/printer/scriptable.py
diff --git a/python/tvm/script/tirx.pyi b/tirx/python/tvm_tirx/script/tirx.pyi
similarity index 100%
rename from python/tvm/script/tirx.pyi
rename to tirx/python/tvm_tirx/script/tirx.pyi
diff --git a/python/tvm/support/__init__.py b/tirx/python/tvm_tirx/support/__init__.py
similarity index 100%
rename from python/tvm/support/__init__.py
rename to tirx/python/tvm_tirx/support/__init__.py
diff --git a/python/tvm/support/cc.py b/tirx/python/tvm_tirx/support/cc.py
similarity index 100%
rename from python/tvm/support/cc.py
rename to tirx/python/tvm_tirx/support/cc.py
diff --git a/python/tvm/support/clang.py b/tirx/python/tvm_tirx/support/clang.py
similarity index 100%
rename from python/tvm/support/clang.py
rename to tirx/python/tvm_tirx/support/clang.py
diff --git a/python/tvm/support/libinfo.py b/tirx/python/tvm_tirx/support/libinfo.py
similarity index 100%
rename from python/tvm/support/libinfo.py
rename to tirx/python/tvm_tirx/support/libinfo.py
diff --git a/python/tvm/support/ndk.py b/tirx/python/tvm_tirx/support/ndk.py
similarity index 100%
rename from python/tvm/support/ndk.py
rename to tirx/python/tvm_tirx/support/ndk.py
diff --git a/python/tvm/support/nvcc.py b/tirx/python/tvm_tirx/support/nvcc.py
similarity index 100%
rename from python/tvm/support/nvcc.py
rename to tirx/python/tvm_tirx/support/nvcc.py
diff --git a/python/tvm/support/popen_pool.py b/tirx/python/tvm_tirx/support/popen_pool.py
similarity index 100%
rename from python/tvm/support/popen_pool.py
rename to tirx/python/tvm_tirx/support/popen_pool.py
diff --git a/python/tvm/support/popen_worker.py b/tirx/python/tvm_tirx/support/popen_worker.py
similarity index 100%
rename from python/tvm/support/popen_worker.py
rename to tirx/python/tvm_tirx/support/popen_worker.py
diff --git a/python/tvm/support/rocm.py b/tirx/python/tvm_tirx/support/rocm.py
similarity index 100%
rename from python/tvm/support/rocm.py
rename to tirx/python/tvm_tirx/support/rocm.py
diff --git a/python/tvm/support/tar.py b/tirx/python/tvm_tirx/support/tar.py
similarity index 100%
rename from python/tvm/support/tar.py
rename to tirx/python/tvm_tirx/support/tar.py
diff --git a/python/tvm/support/utils.py b/tirx/python/tvm_tirx/support/utils.py
similarity index 100%
rename from python/tvm/support/utils.py
rename to tirx/python/tvm_tirx/support/utils.py
diff --git a/python/tvm/support/xcode.py b/tirx/python/tvm_tirx/support/xcode.py
similarity index 100%
rename from python/tvm/support/xcode.py
rename to tirx/python/tvm_tirx/support/xcode.py
diff --git a/python/tvm/sym/__init__.py b/tirx/python/tvm_tirx/sym/__init__.py
similarity index 100%
rename from python/tvm/sym/__init__.py
rename to tirx/python/tvm_tirx/sym/__init__.py
diff --git a/python/tvm/sym/_ffi_api.py b/tirx/python/tvm_tirx/sym/_ffi_api.py
similarity index 100%
rename from python/tvm/sym/_ffi_api.py
rename to tirx/python/tvm_tirx/sym/_ffi_api.py
diff --git a/python/tvm/sym/analyzer.py b/tirx/python/tvm_tirx/sym/analyzer.py
similarity index 100%
rename from python/tvm/sym/analyzer.py
rename to tirx/python/tvm_tirx/sym/analyzer.py
diff --git a/python/tvm/sym/bound.py b/tirx/python/tvm_tirx/sym/bound.py
similarity index 100%
rename from python/tvm/sym/bound.py
rename to tirx/python/tvm_tirx/sym/bound.py
diff --git a/python/tvm/sym/int_set.py b/tirx/python/tvm_tirx/sym/int_set.py
similarity index 100%
rename from python/tvm/sym/int_set.py
rename to tirx/python/tvm_tirx/sym/int_set.py
diff --git a/python/tvm/sym/iter_affine_map.py b/tirx/python/tvm_tirx/sym/iter_affine_map.py
similarity index 100%
rename from python/tvm/sym/iter_affine_map.py
rename to tirx/python/tvm_tirx/sym/iter_affine_map.py
diff --git a/python/tvm/sym/pattern.py b/tirx/python/tvm_tirx/sym/pattern.py
similarity index 100%
rename from python/tvm/sym/pattern.py
rename to tirx/python/tvm_tirx/sym/pattern.py
diff --git a/python/tvm/sym/utils.py b/tirx/python/tvm_tirx/sym/utils.py
similarity index 100%
rename from python/tvm/sym/utils.py
rename to tirx/python/tvm_tirx/sym/utils.py
diff --git a/python/tvm/target/__init__.py b/tirx/python/tvm_tirx/target/__init__.py
similarity index 100%
rename from python/tvm/target/__init__.py
rename to tirx/python/tvm_tirx/target/__init__.py
diff --git a/python/tvm/target/_ffi_api.py b/tirx/python/tvm_tirx/target/_ffi_api.py
similarity index 100%
rename from python/tvm/target/_ffi_api.py
rename to tirx/python/tvm_tirx/target/_ffi_api.py
diff --git a/python/tvm/target/detect_target.py b/tirx/python/tvm_tirx/target/detect_target.py
similarity index 100%
rename from python/tvm/target/detect_target.py
rename to tirx/python/tvm_tirx/target/detect_target.py
diff --git a/python/tvm/target/intrin.py b/tirx/python/tvm_tirx/target/intrin.py
similarity index 100%
rename from python/tvm/target/intrin.py
rename to tirx/python/tvm_tirx/target/intrin.py
diff --git a/python/tvm/target/tag.py b/tirx/python/tvm_tirx/target/tag.py
similarity index 100%
rename from python/tvm/target/tag.py
rename to tirx/python/tvm_tirx/target/tag.py
diff --git a/python/tvm/target/tag_registry/__init__.py b/tirx/python/tvm_tirx/target/tag_registry/__init__.py
similarity index 100%
rename from python/tvm/target/tag_registry/__init__.py
rename to tirx/python/tvm_tirx/target/tag_registry/__init__.py
diff --git a/python/tvm/target/tag_registry/arm_cpu.py b/tirx/python/tvm_tirx/target/tag_registry/arm_cpu.py
similarity index 100%
rename from python/tvm/target/tag_registry/arm_cpu.py
rename to tirx/python/tvm_tirx/target/tag_registry/arm_cpu.py
diff --git a/python/tvm/target/tag_registry/aws_cpu.py b/tirx/python/tvm_tirx/target/tag_registry/aws_cpu.py
similarity index 100%
rename from python/tvm/target/tag_registry/aws_cpu.py
rename to tirx/python/tvm_tirx/target/tag_registry/aws_cpu.py
diff --git a/python/tvm/target/tag_registry/registry.py b/tirx/python/tvm_tirx/target/tag_registry/registry.py
similarity index 100%
rename from python/tvm/target/tag_registry/registry.py
rename to tirx/python/tvm_tirx/target/tag_registry/registry.py
diff --git a/python/tvm/target/tag_registry/riscv_cpu.py b/tirx/python/tvm_tirx/target/tag_registry/riscv_cpu.py
similarity index 100%
rename from python/tvm/target/tag_registry/riscv_cpu.py
rename to tirx/python/tvm_tirx/target/tag_registry/riscv_cpu.py
diff --git a/python/tvm/target/target.py b/tirx/python/tvm_tirx/target/target.py
similarity index 100%
rename from python/tvm/target/target.py
rename to tirx/python/tvm_tirx/target/target.py
diff --git a/python/tvm/testing/__init__.py b/tirx/python/tvm_tirx/testing/__init__.py
similarity index 100%
rename from python/tvm/testing/__init__.py
rename to tirx/python/tvm_tirx/testing/__init__.py
diff --git a/python/tvm/testing/_ffi_api.py b/tirx/python/tvm_tirx/testing/_ffi_api.py
similarity index 100%
rename from python/tvm/testing/_ffi_api.py
rename to tirx/python/tvm_tirx/testing/_ffi_api.py
diff --git a/python/tvm/testing/env.py b/tirx/python/tvm_tirx/testing/env.py
similarity index 100%
rename from python/tvm/testing/env.py
rename to tirx/python/tvm_tirx/testing/env.py
diff --git a/python/tvm/testing/locking.py b/tirx/python/tvm_tirx/testing/locking.py
similarity index 100%
rename from python/tvm/testing/locking.py
rename to tirx/python/tvm_tirx/testing/locking.py
diff --git a/python/tvm/testing/popen_pool.py b/tirx/python/tvm_tirx/testing/popen_pool.py
similarity index 100%
rename from python/tvm/testing/popen_pool.py
rename to tirx/python/tvm_tirx/testing/popen_pool.py
diff --git a/python/tvm/testing/utils.py b/tirx/python/tvm_tirx/testing/utils.py
similarity index 100%
rename from python/tvm/testing/utils.py
rename to tirx/python/tvm_tirx/testing/utils.py
diff --git a/python/tvm/tirx/__init__.py b/tirx/python/tvm_tirx/tirx/__init__.py
similarity index 100%
rename from python/tvm/tirx/__init__.py
rename to tirx/python/tvm_tirx/tirx/__init__.py
diff --git a/python/tvm/tirx/_compile.py b/tirx/python/tvm_tirx/tirx/_compile.py
similarity index 100%
rename from python/tvm/tirx/_compile.py
rename to tirx/python/tvm_tirx/tirx/_compile.py
diff --git a/python/tvm/tirx/_ffi_api.py b/tirx/python/tvm_tirx/tirx/_ffi_api.py
similarity index 100%
rename from python/tvm/tirx/_ffi_api.py
rename to tirx/python/tvm_tirx/tirx/_ffi_api.py
diff --git a/python/tvm/tirx/_tensor_view.py b/tirx/python/tvm_tirx/tirx/_tensor_view.py
similarity index 100%
rename from python/tvm/tirx/_tensor_view.py
rename to tirx/python/tvm_tirx/tirx/_tensor_view.py
diff --git a/python/tvm/tirx/analysis/__init__.py b/tirx/python/tvm_tirx/tirx/analysis/__init__.py
similarity index 100%
rename from python/tvm/tirx/analysis/__init__.py
rename to tirx/python/tvm_tirx/tirx/analysis/__init__.py
diff --git a/python/tvm/tirx/analysis/_ffi_api.py b/tirx/python/tvm_tirx/tirx/analysis/_ffi_api.py
similarity index 100%
rename from python/tvm/tirx/analysis/_ffi_api.py
rename to tirx/python/tvm_tirx/tirx/analysis/_ffi_api.py
diff --git a/python/tvm/tirx/analysis/analysis.py b/tirx/python/tvm_tirx/tirx/analysis/analysis.py
similarity index 100%
rename from python/tvm/tirx/analysis/analysis.py
rename to tirx/python/tvm_tirx/tirx/analysis/analysis.py
diff --git a/python/tvm/tirx/backend/__init__.py b/tirx/python/tvm_tirx/tirx/backend/__init__.py
similarity index 100%
rename from python/tvm/tirx/backend/__init__.py
rename to tirx/python/tvm_tirx/tirx/backend/__init__.py
diff --git a/python/tvm/tirx/bench.py b/tirx/python/tvm_tirx/tirx/bench.py
similarity index 100%
rename from python/tvm/tirx/bench.py
rename to tirx/python/tvm_tirx/tirx/bench.py
diff --git a/python/tvm/tirx/compilation_pipeline.py b/tirx/python/tvm_tirx/tirx/compilation_pipeline.py
similarity index 100%
rename from python/tvm/tirx/compilation_pipeline.py
rename to tirx/python/tvm_tirx/tirx/compilation_pipeline.py
diff --git a/python/tvm/tirx/expr.py b/tirx/python/tvm_tirx/tirx/expr.py
similarity index 100%
rename from python/tvm/tirx/expr.py
rename to tirx/python/tvm_tirx/tirx/expr.py
diff --git a/python/tvm/tirx/function.py b/tirx/python/tvm_tirx/tirx/function.py
similarity index 100%
rename from python/tvm/tirx/function.py
rename to tirx/python/tvm_tirx/tirx/function.py
diff --git a/python/tvm/tirx/layout.py b/tirx/python/tvm_tirx/tirx/layout.py
similarity index 100%
rename from python/tvm/tirx/layout.py
rename to tirx/python/tvm_tirx/tirx/layout.py
diff --git a/python/tvm/tirx/op/__init__.py b/tirx/python/tvm_tirx/tirx/op/__init__.py
similarity index 100%
rename from python/tvm/tirx/op/__init__.py
rename to tirx/python/tvm_tirx/tirx/op/__init__.py
diff --git a/python/tvm/tirx/script/__init__.py b/tirx/python/tvm_tirx/tirx/script/__init__.py
similarity index 100%
rename from python/tvm/tirx/script/__init__.py
rename to tirx/python/tvm_tirx/tirx/script/__init__.py
diff --git a/python/tvm/tirx/script/ir_builder/__init__.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/__init__.py
similarity index 100%
rename from python/tvm/tirx/script/ir_builder/__init__.py
rename to tirx/python/tvm_tirx/tirx/script/ir_builder/__init__.py
diff --git a/python/tvm/tirx/script/ir_builder/_ffi_api.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/_ffi_api.py
similarity index 100%
rename from python/tvm/tirx/script/ir_builder/_ffi_api.py
rename to tirx/python/tvm_tirx/tirx/script/ir_builder/_ffi_api.py
diff --git a/python/tvm/tirx/script/ir_builder/external_kernel.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/external_kernel.py
similarity index 100%
rename from python/tvm/tirx/script/ir_builder/external_kernel.py
rename to tirx/python/tvm_tirx/tirx/script/ir_builder/external_kernel.py
diff --git a/python/tvm/tirx/script/ir_builder/frame.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/frame.py
similarity index 100%
rename from python/tvm/tirx/script/ir_builder/frame.py
rename to tirx/python/tvm_tirx/tirx/script/ir_builder/frame.py
diff --git a/python/tvm/tirx/script/ir_builder/ir.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py
similarity index 100%
rename from python/tvm/tirx/script/ir_builder/ir.py
rename to tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py
diff --git a/python/tvm/tirx/script/ir_builder/op.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/op.py
similarity index 100%
rename from python/tvm/tirx/script/ir_builder/op.py
rename to tirx/python/tvm_tirx/tirx/script/ir_builder/op.py
diff --git a/python/tvm/tirx/script/ir_builder/parser_protocol.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py
similarity index 100%
rename from python/tvm/tirx/script/ir_builder/parser_protocol.py
rename to tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py
diff --git a/python/tvm/tirx/script/ir_builder/triton.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/triton.py
similarity index 100%
rename from python/tvm/tirx/script/ir_builder/triton.py
rename to tirx/python/tvm_tirx/tirx/script/ir_builder/triton.py
diff --git a/python/tvm/tirx/script/ir_builder/utils.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/utils.py
similarity index 100%
rename from python/tvm/tirx/script/ir_builder/utils.py
rename to tirx/python/tvm_tirx/tirx/script/ir_builder/utils.py
diff --git a/python/tvm/tirx/script/jit.py b/tirx/python/tvm_tirx/tirx/script/jit.py
similarity index 100%
rename from python/tvm/tirx/script/jit.py
rename to tirx/python/tvm_tirx/tirx/script/jit.py
diff --git a/python/tvm/tirx/tensor.py b/tirx/python/tvm_tirx/tirx/tensor.py
similarity index 100%
rename from python/tvm/tirx/tensor.py
rename to tirx/python/tvm_tirx/tirx/tensor.py
diff --git a/python/tvm/tirx/thread_scope.py b/tirx/python/tvm_tirx/tirx/thread_scope.py
similarity index 100%
rename from python/tvm/tirx/thread_scope.py
rename to tirx/python/tvm_tirx/tirx/thread_scope.py
diff --git a/python/tvm/tirx/tile/__init__.py b/tirx/python/tvm_tirx/tirx/tile/__init__.py
similarity index 100%
rename from python/tvm/tirx/tile/__init__.py
rename to tirx/python/tvm_tirx/tirx/tile/__init__.py
diff --git a/python/tvm/tirx/tile/common.py b/tirx/python/tvm_tirx/tirx/tile/common.py
similarity index 100%
rename from python/tvm/tirx/tile/common.py
rename to tirx/python/tvm_tirx/tirx/tile/common.py
diff --git a/python/tvm/tirx/tile/dispatch.py b/tirx/python/tvm_tirx/tirx/tile/dispatch.py
similarity index 100%
rename from python/tvm/tirx/tile/dispatch.py
rename to tirx/python/tvm_tirx/tirx/tile/dispatch.py
diff --git a/python/tvm/tirx/tile/instruction.py b/tirx/python/tvm_tirx/tirx/tile/instruction.py
similarity index 100%
rename from python/tvm/tirx/tile/instruction.py
rename to tirx/python/tvm_tirx/tirx/tile/instruction.py
diff --git a/python/tvm/tirx/tile/views.py b/tirx/python/tvm_tirx/tirx/tile/views.py
similarity index 100%
rename from python/tvm/tirx/tile/views.py
rename to tirx/python/tvm_tirx/tirx/tile/views.py
diff --git a/python/tvm/tirx/transform/__init__.py b/tirx/python/tvm_tirx/tirx/transform/__init__.py
similarity index 100%
rename from python/tvm/tirx/transform/__init__.py
rename to tirx/python/tvm_tirx/tirx/transform/__init__.py
diff --git a/python/tvm/tirx/transform/_ffi_api.py b/tirx/python/tvm_tirx/tirx/transform/_ffi_api.py
similarity index 100%
rename from python/tvm/tirx/transform/_ffi_api.py
rename to tirx/python/tvm_tirx/tirx/transform/_ffi_api.py
diff --git a/python/tvm/tirx/transform/common.py b/tirx/python/tvm_tirx/tirx/transform/common.py
similarity index 100%
rename from python/tvm/tirx/transform/common.py
rename to tirx/python/tvm_tirx/tirx/transform/common.py
diff --git a/python/tvm/tirx/transform/function_pass.py b/tirx/python/tvm_tirx/tirx/transform/function_pass.py
similarity index 100%
rename from python/tvm/tirx/transform/function_pass.py
rename to tirx/python/tvm_tirx/tirx/transform/function_pass.py
diff --git a/python/tvm/tirx/transform/transform.py b/tirx/python/tvm_tirx/tirx/transform/transform.py
similarity index 100%
rename from python/tvm/tirx/transform/transform.py
rename to tirx/python/tvm_tirx/tirx/transform/transform.py
diff --git a/python/tvm/tirx/type.py b/tirx/python/tvm_tirx/tirx/type.py
similarity index 100%
rename from python/tvm/tirx/type.py
rename to tirx/python/tvm_tirx/tirx/type.py
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/tirx/src/backend/cuda/codegen/codegen_cuda.cc
similarity index 100%
rename from src/backend/cuda/codegen/codegen_cuda.cc
rename to tirx/src/backend/cuda/codegen/codegen_cuda.cc
diff --git a/src/backend/cuda/codegen/codegen_cuda.h b/tirx/src/backend/cuda/codegen/codegen_cuda.h
similarity index 100%
rename from src/backend/cuda/codegen/codegen_cuda.h
rename to tirx/src/backend/cuda/codegen/codegen_cuda.h
diff --git a/src/backend/cuda/codegen/codegen_cuda_host.cc b/tirx/src/backend/cuda/codegen/codegen_cuda_host.cc
similarity index 100%
rename from src/backend/cuda/codegen/codegen_cuda_host.cc
rename to tirx/src/backend/cuda/codegen/codegen_cuda_host.cc
diff --git a/src/backend/cuda/codegen/cuda_fallback_module.cc b/tirx/src/backend/cuda/codegen/cuda_fallback_module.cc
similarity index 100%
rename from src/backend/cuda/codegen/cuda_fallback_module.cc
rename to tirx/src/backend/cuda/codegen/cuda_fallback_module.cc
diff --git a/src/backend/cuda/codegen/cuda_fallback_module.h b/tirx/src/backend/cuda/codegen/cuda_fallback_module.h
similarity index 100%
rename from src/backend/cuda/codegen/cuda_fallback_module.h
rename to tirx/src/backend/cuda/codegen/cuda_fallback_module.h
diff --git a/src/backend/cuda/codegen/intrin_rule_cuda.cc b/tirx/src/backend/cuda/codegen/intrin_rule_cuda.cc
similarity index 100%
rename from src/backend/cuda/codegen/intrin_rule_cuda.cc
rename to tirx/src/backend/cuda/codegen/intrin_rule_cuda.cc
diff --git a/src/backend/cuda/codegen/literal/cuda_half_t.h b/tirx/src/backend/cuda/codegen/literal/cuda_half_t.h
similarity index 100%
rename from src/backend/cuda/codegen/literal/cuda_half_t.h
rename to tirx/src/backend/cuda/codegen/literal/cuda_half_t.h
diff --git a/src/backend/cuda/codegen/literal/cuda_int8_t.h b/tirx/src/backend/cuda/codegen/literal/cuda_int8_t.h
similarity index 100%
rename from src/backend/cuda/codegen/literal/cuda_int8_t.h
rename to tirx/src/backend/cuda/codegen/literal/cuda_int8_t.h
diff --git a/src/backend/cuda/codegen/llvm/codegen_nvptx.cc b/tirx/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
similarity index 100%
rename from src/backend/cuda/codegen/llvm/codegen_nvptx.cc
rename to tirx/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
diff --git a/src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc b/tirx/src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc
similarity index 100%
rename from src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc
rename to tirx/src/backend/cuda/codegen/llvm/intrin_rule_nvptx.cc
diff --git a/src/backend/cuda/codegen/ptx.cc b/tirx/src/backend/cuda/codegen/ptx.cc
similarity index 100%
rename from src/backend/cuda/codegen/ptx.cc
rename to tirx/src/backend/cuda/codegen/ptx.cc
diff --git a/src/backend/cuda/codegen/target_kind.cc b/tirx/src/backend/cuda/codegen/target_kind.cc
similarity index 100%
rename from src/backend/cuda/codegen/target_kind.cc
rename to tirx/src/backend/cuda/codegen/target_kind.cc
diff --git a/src/backend/cuda/launch/launch_attributes.h b/tirx/src/backend/cuda/launch/launch_attributes.h
similarity index 100%
rename from src/backend/cuda/launch/launch_attributes.h
rename to tirx/src/backend/cuda/launch/launch_attributes.h
diff --git a/src/backend/cuda/launch/launch_config.h b/tirx/src/backend/cuda/launch/launch_config.h
similarity index 100%
rename from src/backend/cuda/launch/launch_config.h
rename to tirx/src/backend/cuda/launch/launch_config.h
diff --git a/src/backend/cuda/launch/launch_fields.h b/tirx/src/backend/cuda/launch/launch_fields.h
similarity index 100%
rename from src/backend/cuda/launch/launch_fields.h
rename to tirx/src/backend/cuda/launch/launch_fields.h
diff --git a/src/backend/cuda/launch/launch_plan.h b/tirx/src/backend/cuda/launch/launch_plan.h
similarity index 100%
rename from src/backend/cuda/launch/launch_plan.h
rename to tirx/src/backend/cuda/launch/launch_plan.h
diff --git a/src/backend/cuda/launch/launch_source.h b/tirx/src/backend/cuda/launch/launch_source.h
similarity index 100%
rename from src/backend/cuda/launch/launch_source.h
rename to tirx/src/backend/cuda/launch/launch_source.h
diff --git a/src/backend/cuda/layout/axis.cc b/tirx/src/backend/cuda/layout/axis.cc
similarity index 100%
rename from src/backend/cuda/layout/axis.cc
rename to tirx/src/backend/cuda/layout/axis.cc
diff --git a/src/backend/cuda/op/iket.cc b/tirx/src/backend/cuda/op/iket.cc
similarity index 100%
rename from src/backend/cuda/op/iket.cc
rename to tirx/src/backend/cuda/op/iket.cc
diff --git a/src/backend/cuda/op/legacy.cc b/tirx/src/backend/cuda/op/legacy.cc
similarity index 100%
rename from src/backend/cuda/op/legacy.cc
rename to tirx/src/backend/cuda/op/legacy.cc
diff --git a/src/backend/cuda/op/target_builtin.cc b/tirx/src/backend/cuda/op/target_builtin.cc
similarity index 100%
rename from src/backend/cuda/op/target_builtin.cc
rename to tirx/src/backend/cuda/op/target_builtin.cc
diff --git a/src/backend/cuda/op/tcgen05.cc b/tirx/src/backend/cuda/op/tcgen05.cc
similarity index 100%
rename from src/backend/cuda/op/tcgen05.cc
rename to tirx/src/backend/cuda/op/tcgen05.cc
diff --git a/src/backend/cuda/op/tensormap.cc b/tirx/src/backend/cuda/op/tensormap.cc
similarity index 100%
rename from src/backend/cuda/op/tensormap.cc
rename to tirx/src/backend/cuda/op/tensormap.cc
diff --git a/src/backend/cuda/runtime/cuda_device_api.cc b/tirx/src/backend/cuda/runtime/cuda_device_api.cc
similarity index 100%
rename from src/backend/cuda/runtime/cuda_device_api.cc
rename to tirx/src/backend/cuda/runtime/cuda_device_api.cc
diff --git a/src/backend/cuda/runtime/cuda_module.cc b/tirx/src/backend/cuda/runtime/cuda_module.cc
similarity index 100%
rename from src/backend/cuda/runtime/cuda_module.cc
rename to tirx/src/backend/cuda/runtime/cuda_module.cc
diff --git a/src/backend/cuda/tile/exec_analysis.cc b/tirx/src/backend/cuda/tile/exec_analysis.cc
similarity index 100%
rename from src/backend/cuda/tile/exec_analysis.cc
rename to tirx/src/backend/cuda/tile/exec_analysis.cc
diff --git a/src/backend/cuda/tile/exec_analysis.h b/tirx/src/backend/cuda/tile/exec_analysis.h
similarity index 100%
rename from src/backend/cuda/tile/exec_analysis.h
rename to tirx/src/backend/cuda/tile/exec_analysis.h
diff --git a/src/backend/cuda/tile/exec_context.cc b/tirx/src/backend/cuda/tile/exec_context.cc
similarity index 100%
rename from src/backend/cuda/tile/exec_context.cc
rename to tirx/src/backend/cuda/tile/exec_context.cc
diff --git a/src/backend/cuda/tile/exec_context.h b/tirx/src/backend/cuda/tile/exec_context.h
similarity index 100%
rename from src/backend/cuda/tile/exec_context.h
rename to tirx/src/backend/cuda/tile/exec_context.h
diff --git a/src/backend/cuda/tile/filter_canonical.cc b/tirx/src/backend/cuda/tile/filter_canonical.cc
similarity index 100%
rename from src/backend/cuda/tile/filter_canonical.cc
rename to tirx/src/backend/cuda/tile/filter_canonical.cc
diff --git a/src/backend/cuda/tile/filter_canonical.h b/tirx/src/backend/cuda/tile/filter_canonical.h
similarity index 100%
rename from src/backend/cuda/tile/filter_canonical.h
rename to tirx/src/backend/cuda/tile/filter_canonical.h
diff --git a/src/backend/cuda/tile/thread_scope.cc b/tirx/src/backend/cuda/tile/thread_scope.cc
similarity index 100%
rename from src/backend/cuda/tile/thread_scope.cc
rename to tirx/src/backend/cuda/tile/thread_scope.cc
diff --git a/src/backend/cuda/tile/thread_scope.h b/tirx/src/backend/cuda/tile/thread_scope.h
similarity index 100%
rename from src/backend/cuda/tile/thread_scope.h
rename to tirx/src/backend/cuda/tile/thread_scope.h
diff --git a/src/backend/cuda/transforms/bind_backend_config.cc b/tirx/src/backend/cuda/transforms/bind_backend_config.cc
similarity index 100%
rename from src/backend/cuda/transforms/bind_backend_config.cc
rename to tirx/src/backend/cuda/transforms/bind_backend_config.cc
diff --git a/src/backend/cuda/transforms/lower_iket.cc b/tirx/src/backend/cuda/transforms/lower_iket.cc
similarity index 100%
rename from src/backend/cuda/transforms/lower_iket.cc
rename to tirx/src/backend/cuda/transforms/lower_iket.cc
diff --git a/src/backend/hexagon/codegen/hexagon_fallback_module.cc b/tirx/src/backend/hexagon/codegen/hexagon_fallback_module.cc
similarity index 100%
rename from src/backend/hexagon/codegen/hexagon_fallback_module.cc
rename to tirx/src/backend/hexagon/codegen/hexagon_fallback_module.cc
diff --git a/src/backend/hexagon/codegen/hexagon_fallback_module.h b/tirx/src/backend/hexagon/codegen/hexagon_fallback_module.h
similarity index 100%
rename from src/backend/hexagon/codegen/hexagon_fallback_module.h
rename to tirx/src/backend/hexagon/codegen/hexagon_fallback_module.h
diff --git a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc b/tirx/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
similarity index 100%
rename from src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
rename to tirx/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
diff --git a/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc b/tirx/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
similarity index 100%
rename from src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
rename to tirx/src/backend/hexagon/codegen/llvm/intrin_rule_hexagon.cc
diff --git a/src/backend/hexagon/codegen/memory_scope.h b/tirx/src/backend/hexagon/codegen/memory_scope.h
similarity index 100%
rename from src/backend/hexagon/codegen/memory_scope.h
rename to tirx/src/backend/hexagon/codegen/memory_scope.h
diff --git a/src/backend/hexagon/codegen/target_kind.cc b/tirx/src/backend/hexagon/codegen/target_kind.cc
similarity index 100%
rename from src/backend/hexagon/codegen/target_kind.cc
rename to tirx/src/backend/hexagon/codegen/target_kind.cc
diff --git a/src/backend/metal/codegen/codegen_metal.cc b/tirx/src/backend/metal/codegen/codegen_metal.cc
similarity index 100%
rename from src/backend/metal/codegen/codegen_metal.cc
rename to tirx/src/backend/metal/codegen/codegen_metal.cc
diff --git a/src/backend/metal/codegen/codegen_metal.h b/tirx/src/backend/metal/codegen/codegen_metal.h
similarity index 100%
rename from src/backend/metal/codegen/codegen_metal.h
rename to tirx/src/backend/metal/codegen/codegen_metal.h
diff --git a/src/backend/metal/codegen/intrin_rule_metal.cc b/tirx/src/backend/metal/codegen/intrin_rule_metal.cc
similarity index 100%
rename from src/backend/metal/codegen/intrin_rule_metal.cc
rename to tirx/src/backend/metal/codegen/intrin_rule_metal.cc
diff --git a/src/backend/metal/codegen/metal_fallback_module.cc b/tirx/src/backend/metal/codegen/metal_fallback_module.cc
similarity index 100%
rename from src/backend/metal/codegen/metal_fallback_module.cc
rename to tirx/src/backend/metal/codegen/metal_fallback_module.cc
diff --git a/src/backend/metal/codegen/metal_fallback_module.h b/tirx/src/backend/metal/codegen/metal_fallback_module.h
similarity index 100%
rename from src/backend/metal/codegen/metal_fallback_module.h
rename to tirx/src/backend/metal/codegen/metal_fallback_module.h
diff --git a/src/backend/metal/codegen/target_kind.cc b/tirx/src/backend/metal/codegen/target_kind.cc
similarity index 100%
rename from src/backend/metal/codegen/target_kind.cc
rename to tirx/src/backend/metal/codegen/target_kind.cc
diff --git a/src/backend/metal/op/cooperative_tensor.cc b/tirx/src/backend/metal/op/cooperative_tensor.cc
similarity index 100%
rename from src/backend/metal/op/cooperative_tensor.cc
rename to tirx/src/backend/metal/op/cooperative_tensor.cc
diff --git a/src/backend/metal/op/target_builtin.cc b/tirx/src/backend/metal/op/target_builtin.cc
similarity index 100%
rename from src/backend/metal/op/target_builtin.cc
rename to tirx/src/backend/metal/op/target_builtin.cc
diff --git a/src/backend/metal/runtime/metal_common.h b/tirx/src/backend/metal/runtime/metal_common.h
similarity index 100%
rename from src/backend/metal/runtime/metal_common.h
rename to tirx/src/backend/metal/runtime/metal_common.h
diff --git a/src/backend/metal/runtime/metal_device_api.mm b/tirx/src/backend/metal/runtime/metal_device_api.mm
similarity index 100%
rename from src/backend/metal/runtime/metal_device_api.mm
rename to tirx/src/backend/metal/runtime/metal_device_api.mm
diff --git a/src/backend/metal/runtime/metal_module.mm b/tirx/src/backend/metal/runtime/metal_module.mm
similarity index 100%
rename from src/backend/metal/runtime/metal_module.mm
rename to tirx/src/backend/metal/runtime/metal_module.mm
diff --git a/src/backend/opencl/codegen/codegen_opencl.cc b/tirx/src/backend/opencl/codegen/codegen_opencl.cc
similarity index 100%
rename from src/backend/opencl/codegen/codegen_opencl.cc
rename to tirx/src/backend/opencl/codegen/codegen_opencl.cc
diff --git a/src/backend/opencl/codegen/codegen_opencl.h b/tirx/src/backend/opencl/codegen/codegen_opencl.h
similarity index 100%
rename from src/backend/opencl/codegen/codegen_opencl.h
rename to tirx/src/backend/opencl/codegen/codegen_opencl.h
diff --git a/src/backend/opencl/codegen/intrin_rule_opencl.cc b/tirx/src/backend/opencl/codegen/intrin_rule_opencl.cc
similarity index 100%
rename from src/backend/opencl/codegen/intrin_rule_opencl.cc
rename to tirx/src/backend/opencl/codegen/intrin_rule_opencl.cc
diff --git a/src/backend/opencl/codegen/opencl_fallback_module.cc b/tirx/src/backend/opencl/codegen/opencl_fallback_module.cc
similarity index 100%
rename from src/backend/opencl/codegen/opencl_fallback_module.cc
rename to tirx/src/backend/opencl/codegen/opencl_fallback_module.cc
diff --git a/src/backend/opencl/codegen/opencl_fallback_module.h b/tirx/src/backend/opencl/codegen/opencl_fallback_module.h
similarity index 100%
rename from src/backend/opencl/codegen/opencl_fallback_module.h
rename to tirx/src/backend/opencl/codegen/opencl_fallback_module.h
diff --git a/src/backend/opencl/codegen/target_kind.cc b/tirx/src/backend/opencl/codegen/target_kind.cc
similarity index 100%
rename from src/backend/opencl/codegen/target_kind.cc
rename to tirx/src/backend/opencl/codegen/target_kind.cc
diff --git a/src/backend/opencl/memory_scope.h b/tirx/src/backend/opencl/memory_scope.h
similarity index 100%
rename from src/backend/opencl/memory_scope.h
rename to tirx/src/backend/opencl/memory_scope.h
diff --git a/src/backend/opencl/op/memory.cc b/tirx/src/backend/opencl/op/memory.cc
similarity index 100%
rename from src/backend/opencl/op/memory.cc
rename to tirx/src/backend/opencl/op/memory.cc
diff --git a/src/backend/opencl/runtime/opencl_common.h b/tirx/src/backend/opencl/runtime/opencl_common.h
similarity index 100%
rename from src/backend/opencl/runtime/opencl_common.h
rename to tirx/src/backend/opencl/runtime/opencl_common.h
diff --git a/src/backend/opencl/runtime/opencl_device_api.cc b/tirx/src/backend/opencl/runtime/opencl_device_api.cc
similarity index 100%
rename from src/backend/opencl/runtime/opencl_device_api.cc
rename to tirx/src/backend/opencl/runtime/opencl_device_api.cc
diff --git a/src/backend/opencl/runtime/opencl_module.cc b/tirx/src/backend/opencl/runtime/opencl_module.cc
similarity index 100%
rename from src/backend/opencl/runtime/opencl_module.cc
rename to tirx/src/backend/opencl/runtime/opencl_module.cc
diff --git a/src/backend/opencl/runtime/opencl_wrapper/README.md b/tirx/src/backend/opencl/runtime/opencl_wrapper/README.md
similarity index 100%
rename from src/backend/opencl/runtime/opencl_wrapper/README.md
rename to tirx/src/backend/opencl/runtime/opencl_wrapper/README.md
diff --git a/src/backend/opencl/runtime/opencl_wrapper/opencl_wrapper.cc b/tirx/src/backend/opencl/runtime/opencl_wrapper/opencl_wrapper.cc
similarity index 100%
rename from src/backend/opencl/runtime/opencl_wrapper/opencl_wrapper.cc
rename to tirx/src/backend/opencl/runtime/opencl_wrapper/opencl_wrapper.cc
diff --git a/src/backend/opencl/runtime/source_utils.h b/tirx/src/backend/opencl/runtime/source_utils.h
similarity index 100%
rename from src/backend/opencl/runtime/source_utils.h
rename to tirx/src/backend/opencl/runtime/source_utils.h
diff --git a/src/backend/opencl/runtime/texture.h b/tirx/src/backend/opencl/runtime/texture.h
similarity index 100%
rename from src/backend/opencl/runtime/texture.h
rename to tirx/src/backend/opencl/runtime/texture.h
diff --git a/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc b/tirx/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
similarity index 100%
rename from src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
rename to tirx/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
diff --git a/src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc b/tirx/src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc
similarity index 100%
rename from src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc
rename to tirx/src/backend/rocm/codegen/llvm/intrin_rule_rocm.cc
diff --git a/src/backend/rocm/codegen/rocm_fallback_module.cc b/tirx/src/backend/rocm/codegen/rocm_fallback_module.cc
similarity index 100%
rename from src/backend/rocm/codegen/rocm_fallback_module.cc
rename to tirx/src/backend/rocm/codegen/rocm_fallback_module.cc
diff --git a/src/backend/rocm/codegen/rocm_fallback_module.h b/tirx/src/backend/rocm/codegen/rocm_fallback_module.h
similarity index 100%
rename from src/backend/rocm/codegen/rocm_fallback_module.h
rename to tirx/src/backend/rocm/codegen/rocm_fallback_module.h
diff --git a/src/backend/rocm/codegen/target_kind.cc b/tirx/src/backend/rocm/codegen/target_kind.cc
similarity index 100%
rename from src/backend/rocm/codegen/target_kind.cc
rename to tirx/src/backend/rocm/codegen/target_kind.cc
diff --git a/src/backend/rocm/runtime/rocm_common.h b/tirx/src/backend/rocm/runtime/rocm_common.h
similarity index 100%
rename from src/backend/rocm/runtime/rocm_common.h
rename to tirx/src/backend/rocm/runtime/rocm_common.h
diff --git a/src/backend/rocm/runtime/rocm_device_api.cc b/tirx/src/backend/rocm/runtime/rocm_device_api.cc
similarity index 100%
rename from src/backend/rocm/runtime/rocm_device_api.cc
rename to tirx/src/backend/rocm/runtime/rocm_device_api.cc
diff --git a/src/backend/rocm/runtime/rocm_module.cc b/tirx/src/backend/rocm/runtime/rocm_module.cc
similarity index 100%
rename from src/backend/rocm/runtime/rocm_module.cc
rename to tirx/src/backend/rocm/runtime/rocm_module.cc
diff --git a/src/backend/trn/codegen/codegen_trn.cc b/tirx/src/backend/trn/codegen/codegen_trn.cc
similarity index 100%
rename from src/backend/trn/codegen/codegen_trn.cc
rename to tirx/src/backend/trn/codegen/codegen_trn.cc
diff --git a/src/backend/trn/codegen/codegen_trn.h b/tirx/src/backend/trn/codegen/codegen_trn.h
similarity index 100%
rename from src/backend/trn/codegen/codegen_trn.h
rename to tirx/src/backend/trn/codegen/codegen_trn.h
diff --git a/src/backend/trn/codegen/target_kind.cc b/tirx/src/backend/trn/codegen/target_kind.cc
similarity index 100%
rename from src/backend/trn/codegen/target_kind.cc
rename to tirx/src/backend/trn/codegen/target_kind.cc
diff --git a/src/backend/trn/layout/axis.cc b/tirx/src/backend/trn/layout/axis.cc
similarity index 100%
rename from src/backend/trn/layout/axis.cc
rename to tirx/src/backend/trn/layout/axis.cc
diff --git a/src/backend/trn/op/target_builtin.cc b/tirx/src/backend/trn/op/target_builtin.cc
similarity index 100%
rename from src/backend/trn/op/target_builtin.cc
rename to tirx/src/backend/trn/op/target_builtin.cc
diff --git a/src/backend/trn/transform/lower_trainium_layout.cc b/tirx/src/backend/trn/transform/lower_trainium_layout.cc
similarity index 100%
rename from src/backend/trn/transform/lower_trainium_layout.cc
rename to tirx/src/backend/trn/transform/lower_trainium_layout.cc
diff --git a/src/backend/vulkan/codegen/build_vulkan.cc b/tirx/src/backend/vulkan/codegen/build_vulkan.cc
similarity index 100%
rename from src/backend/vulkan/codegen/build_vulkan.cc
rename to tirx/src/backend/vulkan/codegen/build_vulkan.cc
diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc b/tirx/src/backend/vulkan/codegen/codegen_spirv.cc
similarity index 100%
rename from src/backend/vulkan/codegen/codegen_spirv.cc
rename to tirx/src/backend/vulkan/codegen/codegen_spirv.cc
diff --git a/src/backend/vulkan/codegen/codegen_spirv.h b/tirx/src/backend/vulkan/codegen/codegen_spirv.h
similarity index 100%
rename from src/backend/vulkan/codegen/codegen_spirv.h
rename to tirx/src/backend/vulkan/codegen/codegen_spirv.h
diff --git a/src/backend/vulkan/codegen/intrin_rule_spirv.cc b/tirx/src/backend/vulkan/codegen/intrin_rule_spirv.cc
similarity index 100%
rename from src/backend/vulkan/codegen/intrin_rule_spirv.cc
rename to tirx/src/backend/vulkan/codegen/intrin_rule_spirv.cc
diff --git a/src/backend/vulkan/codegen/ir_builder.cc b/tirx/src/backend/vulkan/codegen/ir_builder.cc
similarity index 100%
rename from src/backend/vulkan/codegen/ir_builder.cc
rename to tirx/src/backend/vulkan/codegen/ir_builder.cc
diff --git a/src/backend/vulkan/codegen/ir_builder.h b/tirx/src/backend/vulkan/codegen/ir_builder.h
similarity index 100%
rename from src/backend/vulkan/codegen/ir_builder.h
rename to tirx/src/backend/vulkan/codegen/ir_builder.h
diff --git a/src/backend/vulkan/codegen/spirv_support.cc b/tirx/src/backend/vulkan/codegen/spirv_support.cc
similarity index 100%
rename from src/backend/vulkan/codegen/spirv_support.cc
rename to tirx/src/backend/vulkan/codegen/spirv_support.cc
diff --git a/src/backend/vulkan/codegen/spirv_support.h b/tirx/src/backend/vulkan/codegen/spirv_support.h
similarity index 100%
rename from src/backend/vulkan/codegen/spirv_support.h
rename to tirx/src/backend/vulkan/codegen/spirv_support.h
diff --git a/src/backend/vulkan/codegen/spirv_utils.cc b/tirx/src/backend/vulkan/codegen/spirv_utils.cc
similarity index 100%
rename from src/backend/vulkan/codegen/spirv_utils.cc
rename to tirx/src/backend/vulkan/codegen/spirv_utils.cc
diff --git a/src/backend/vulkan/codegen/spirv_utils.h b/tirx/src/backend/vulkan/codegen/spirv_utils.h
similarity index 100%
rename from src/backend/vulkan/codegen/spirv_utils.h
rename to tirx/src/backend/vulkan/codegen/spirv_utils.h
diff --git a/src/backend/vulkan/codegen/target_kind.cc b/tirx/src/backend/vulkan/codegen/target_kind.cc
similarity index 100%
rename from src/backend/vulkan/codegen/target_kind.cc
rename to tirx/src/backend/vulkan/codegen/target_kind.cc
diff --git a/src/backend/vulkan/codegen/vulkan_fallback_module.cc b/tirx/src/backend/vulkan/codegen/vulkan_fallback_module.cc
similarity index 100%
rename from src/backend/vulkan/codegen/vulkan_fallback_module.cc
rename to tirx/src/backend/vulkan/codegen/vulkan_fallback_module.cc
diff --git a/src/backend/vulkan/codegen/vulkan_fallback_module.h b/tirx/src/backend/vulkan/codegen/vulkan_fallback_module.h
similarity index 100%
rename from src/backend/vulkan/codegen/vulkan_fallback_module.h
rename to tirx/src/backend/vulkan/codegen/vulkan_fallback_module.h
diff --git a/src/backend/vulkan/runtime/README.md b/tirx/src/backend/vulkan/runtime/README.md
similarity index 100%
rename from src/backend/vulkan/runtime/README.md
rename to tirx/src/backend/vulkan/runtime/README.md
diff --git a/src/backend/vulkan/runtime/spirv_shader.h b/tirx/src/backend/vulkan/runtime/spirv_shader.h
similarity index 100%
rename from src/backend/vulkan/runtime/spirv_shader.h
rename to tirx/src/backend/vulkan/runtime/spirv_shader.h
diff --git a/src/backend/vulkan/runtime/thread_map.h b/tirx/src/backend/vulkan/runtime/thread_map.h
similarity index 100%
rename from src/backend/vulkan/runtime/thread_map.h
rename to tirx/src/backend/vulkan/runtime/thread_map.h
diff --git a/src/backend/vulkan/runtime/vulkan_amdrgp.cc b/tirx/src/backend/vulkan/runtime/vulkan_amdrgp.cc
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_amdrgp.cc
rename to tirx/src/backend/vulkan/runtime/vulkan_amdrgp.cc
diff --git a/src/backend/vulkan/runtime/vulkan_amdrgp.h b/tirx/src/backend/vulkan/runtime/vulkan_amdrgp.h
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_amdrgp.h
rename to tirx/src/backend/vulkan/runtime/vulkan_amdrgp.h
diff --git a/src/backend/vulkan/runtime/vulkan_buffer.cc b/tirx/src/backend/vulkan/runtime/vulkan_buffer.cc
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_buffer.cc
rename to tirx/src/backend/vulkan/runtime/vulkan_buffer.cc
diff --git a/src/backend/vulkan/runtime/vulkan_buffer.h b/tirx/src/backend/vulkan/runtime/vulkan_buffer.h
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_buffer.h
rename to tirx/src/backend/vulkan/runtime/vulkan_buffer.h
diff --git a/src/backend/vulkan/runtime/vulkan_common.cc b/tirx/src/backend/vulkan/runtime/vulkan_common.cc
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_common.cc
rename to tirx/src/backend/vulkan/runtime/vulkan_common.cc
diff --git a/src/backend/vulkan/runtime/vulkan_common.h b/tirx/src/backend/vulkan/runtime/vulkan_common.h
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_common.h
rename to tirx/src/backend/vulkan/runtime/vulkan_common.h
diff --git a/src/backend/vulkan/runtime/vulkan_device.cc b/tirx/src/backend/vulkan/runtime/vulkan_device.cc
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_device.cc
rename to tirx/src/backend/vulkan/runtime/vulkan_device.cc
diff --git a/src/backend/vulkan/runtime/vulkan_device.h b/tirx/src/backend/vulkan/runtime/vulkan_device.h
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_device.h
rename to tirx/src/backend/vulkan/runtime/vulkan_device.h
diff --git a/src/backend/vulkan/runtime/vulkan_device_api.cc b/tirx/src/backend/vulkan/runtime/vulkan_device_api.cc
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_device_api.cc
rename to tirx/src/backend/vulkan/runtime/vulkan_device_api.cc
diff --git a/src/backend/vulkan/runtime/vulkan_device_api.h b/tirx/src/backend/vulkan/runtime/vulkan_device_api.h
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_device_api.h
rename to tirx/src/backend/vulkan/runtime/vulkan_device_api.h
diff --git a/src/backend/vulkan/runtime/vulkan_instance.cc b/tirx/src/backend/vulkan/runtime/vulkan_instance.cc
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_instance.cc
rename to tirx/src/backend/vulkan/runtime/vulkan_instance.cc
diff --git a/src/backend/vulkan/runtime/vulkan_instance.h b/tirx/src/backend/vulkan/runtime/vulkan_instance.h
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_instance.h
rename to tirx/src/backend/vulkan/runtime/vulkan_instance.h
diff --git a/src/backend/vulkan/runtime/vulkan_module.cc b/tirx/src/backend/vulkan/runtime/vulkan_module.cc
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_module.cc
rename to tirx/src/backend/vulkan/runtime/vulkan_module.cc
diff --git a/src/backend/vulkan/runtime/vulkan_stream.cc b/tirx/src/backend/vulkan/runtime/vulkan_stream.cc
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_stream.cc
rename to tirx/src/backend/vulkan/runtime/vulkan_stream.cc
diff --git a/src/backend/vulkan/runtime/vulkan_stream.h b/tirx/src/backend/vulkan/runtime/vulkan_stream.h
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_stream.h
rename to tirx/src/backend/vulkan/runtime/vulkan_stream.h
diff --git a/src/backend/vulkan/runtime/vulkan_timer.cc b/tirx/src/backend/vulkan/runtime/vulkan_timer.cc
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_timer.cc
rename to tirx/src/backend/vulkan/runtime/vulkan_timer.cc
diff --git a/src/backend/vulkan/runtime/vulkan_timer.h b/tirx/src/backend/vulkan/runtime/vulkan_timer.h
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_timer.h
rename to tirx/src/backend/vulkan/runtime/vulkan_timer.h
diff --git a/src/backend/vulkan/runtime/vulkan_wrapped_func.cc b/tirx/src/backend/vulkan/runtime/vulkan_wrapped_func.cc
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_wrapped_func.cc
rename to tirx/src/backend/vulkan/runtime/vulkan_wrapped_func.cc
diff --git a/src/backend/vulkan/runtime/vulkan_wrapped_func.h b/tirx/src/backend/vulkan/runtime/vulkan_wrapped_func.h
similarity index 100%
rename from src/backend/vulkan/runtime/vulkan_wrapped_func.h
rename to tirx/src/backend/vulkan/runtime/vulkan_wrapped_func.h
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc b/tirx/src/backend/webgpu/codegen/codegen_webgpu.cc
similarity index 100%
rename from src/backend/webgpu/codegen/codegen_webgpu.cc
rename to tirx/src/backend/webgpu/codegen/codegen_webgpu.cc
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.h b/tirx/src/backend/webgpu/codegen/codegen_webgpu.h
similarity index 100%
rename from src/backend/webgpu/codegen/codegen_webgpu.h
rename to tirx/src/backend/webgpu/codegen/codegen_webgpu.h
diff --git a/src/backend/webgpu/codegen/intrin_rule_webgpu.cc b/tirx/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
similarity index 100%
rename from src/backend/webgpu/codegen/intrin_rule_webgpu.cc
rename to tirx/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
diff --git a/src/backend/webgpu/codegen/target_kind.cc b/tirx/src/backend/webgpu/codegen/target_kind.cc
similarity index 100%
rename from src/backend/webgpu/codegen/target_kind.cc
rename to tirx/src/backend/webgpu/codegen/target_kind.cc
diff --git a/src/backend/webgpu/codegen/webgpu_fallback_module.cc b/tirx/src/backend/webgpu/codegen/webgpu_fallback_module.cc
similarity index 100%
rename from src/backend/webgpu/codegen/webgpu_fallback_module.cc
rename to tirx/src/backend/webgpu/codegen/webgpu_fallback_module.cc
diff --git a/src/backend/webgpu/codegen/webgpu_fallback_module.h b/tirx/src/backend/webgpu/codegen/webgpu_fallback_module.h
similarity index 100%
rename from src/backend/webgpu/codegen/webgpu_fallback_module.h
rename to tirx/src/backend/webgpu/codegen/webgpu_fallback_module.h
diff --git a/src/codegen/build_common.h b/tirx/src/codegen/build_common.h
similarity index 100%
rename from src/codegen/build_common.h
rename to tirx/src/codegen/build_common.h
diff --git a/src/codegen/c/codegen_c.cc b/tirx/src/codegen/c/codegen_c.cc
similarity index 100%
rename from src/codegen/c/codegen_c.cc
rename to tirx/src/codegen/c/codegen_c.cc
diff --git a/src/codegen/c/codegen_c.h b/tirx/src/codegen/c/codegen_c.h
similarity index 100%
rename from src/codegen/c/codegen_c.h
rename to tirx/src/codegen/c/codegen_c.h
diff --git a/src/codegen/c/codegen_c_host.cc b/tirx/src/codegen/c/codegen_c_host.cc
similarity index 100%
rename from src/codegen/c/codegen_c_host.cc
rename to tirx/src/codegen/c/codegen_c_host.cc
diff --git a/src/codegen/c/codegen_c_host.h b/tirx/src/codegen/c/codegen_c_host.h
similarity index 100%
rename from src/codegen/c/codegen_c_host.h
rename to tirx/src/codegen/c/codegen_c_host.h
diff --git a/src/codegen/c/codegen_params.cc b/tirx/src/codegen/c/codegen_params.cc
similarity index 100%
rename from src/codegen/c/codegen_params.cc
rename to tirx/src/codegen/c/codegen_params.cc
diff --git a/src/codegen/c/codegen_params.h b/tirx/src/codegen/c/codegen_params.h
similarity index 100%
rename from src/codegen/c/codegen_params.h
rename to tirx/src/codegen/c/codegen_params.h
diff --git a/src/codegen/c/codegen_source_base.cc b/tirx/src/codegen/c/codegen_source_base.cc
similarity index 100%
rename from src/codegen/c/codegen_source_base.cc
rename to tirx/src/codegen/c/codegen_source_base.cc
diff --git a/src/codegen/c/codegen_source_base.h b/tirx/src/codegen/c/codegen_source_base.h
similarity index 100%
rename from src/codegen/c/codegen_source_base.h
rename to tirx/src/codegen/c/codegen_source_base.h
diff --git a/src/codegen/c/source_module.cc b/tirx/src/codegen/c/source_module.cc
similarity index 100%
rename from src/codegen/c/source_module.cc
rename to tirx/src/codegen/c/source_module.cc
diff --git a/src/codegen/c/target_kind.cc b/tirx/src/codegen/c/target_kind.cc
similarity index 100%
rename from src/codegen/c/target_kind.cc
rename to tirx/src/codegen/c/target_kind.cc
diff --git a/src/codegen/canonicalizer/arm_mprofile.cc b/tirx/src/codegen/canonicalizer/arm_mprofile.cc
similarity index 100%
rename from src/codegen/canonicalizer/arm_mprofile.cc
rename to tirx/src/codegen/canonicalizer/arm_mprofile.cc
diff --git a/src/codegen/canonicalizer/arm_mprofile.h b/tirx/src/codegen/canonicalizer/arm_mprofile.h
similarity index 100%
rename from src/codegen/canonicalizer/arm_mprofile.h
rename to tirx/src/codegen/canonicalizer/arm_mprofile.h
diff --git a/src/codegen/codegen.cc b/tirx/src/codegen/codegen.cc
similarity index 100%
rename from src/codegen/codegen.cc
rename to tirx/src/codegen/codegen.cc
diff --git a/src/codegen/intrin_rule.cc b/tirx/src/codegen/intrin_rule.cc
similarity index 100%
rename from src/codegen/intrin_rule.cc
rename to tirx/src/codegen/intrin_rule.cc
diff --git a/src/codegen/llvm/canonicalizer/arm_aprofile.cc b/tirx/src/codegen/llvm/canonicalizer/arm_aprofile.cc
similarity index 100%
rename from src/codegen/llvm/canonicalizer/arm_aprofile.cc
rename to tirx/src/codegen/llvm/canonicalizer/arm_aprofile.cc
diff --git a/src/codegen/llvm/canonicalizer/arm_aprofile.h b/tirx/src/codegen/llvm/canonicalizer/arm_aprofile.h
similarity index 100%
rename from src/codegen/llvm/canonicalizer/arm_aprofile.h
rename to tirx/src/codegen/llvm/canonicalizer/arm_aprofile.h
diff --git a/src/codegen/llvm/canonicalizer/canonicalize.cc b/tirx/src/codegen/llvm/canonicalizer/canonicalize.cc
similarity index 100%
rename from src/codegen/llvm/canonicalizer/canonicalize.cc
rename to tirx/src/codegen/llvm/canonicalizer/canonicalize.cc
diff --git a/src/codegen/llvm/canonicalizer/canonicalize.h b/tirx/src/codegen/llvm/canonicalizer/canonicalize.h
similarity index 100%
rename from src/codegen/llvm/canonicalizer/canonicalize.h
rename to tirx/src/codegen/llvm/canonicalizer/canonicalize.h
diff --git a/src/codegen/llvm/codegen_aarch64.cc b/tirx/src/codegen/llvm/codegen_aarch64.cc
similarity index 100%
rename from src/codegen/llvm/codegen_aarch64.cc
rename to tirx/src/codegen/llvm/codegen_aarch64.cc
diff --git a/src/codegen/llvm/codegen_arm.cc b/tirx/src/codegen/llvm/codegen_arm.cc
similarity index 100%
rename from src/codegen/llvm/codegen_arm.cc
rename to tirx/src/codegen/llvm/codegen_arm.cc
diff --git a/src/codegen/llvm/codegen_blob.cc b/tirx/src/codegen/llvm/codegen_blob.cc
similarity index 100%
rename from src/codegen/llvm/codegen_blob.cc
rename to tirx/src/codegen/llvm/codegen_blob.cc
diff --git a/src/codegen/llvm/codegen_blob.h b/tirx/src/codegen/llvm/codegen_blob.h
similarity index 100%
rename from src/codegen/llvm/codegen_blob.h
rename to tirx/src/codegen/llvm/codegen_blob.h
diff --git a/src/codegen/llvm/codegen_cpu.cc b/tirx/src/codegen/llvm/codegen_cpu.cc
similarity index 100%
rename from src/codegen/llvm/codegen_cpu.cc
rename to tirx/src/codegen/llvm/codegen_cpu.cc
diff --git a/src/codegen/llvm/codegen_cpu.h b/tirx/src/codegen/llvm/codegen_cpu.h
similarity index 100%
rename from src/codegen/llvm/codegen_cpu.h
rename to tirx/src/codegen/llvm/codegen_cpu.h
diff --git a/src/codegen/llvm/codegen_llvm.cc b/tirx/src/codegen/llvm/codegen_llvm.cc
similarity index 100%
rename from src/codegen/llvm/codegen_llvm.cc
rename to tirx/src/codegen/llvm/codegen_llvm.cc
diff --git a/src/codegen/llvm/codegen_llvm.h b/tirx/src/codegen/llvm/codegen_llvm.h
similarity index 100%
rename from src/codegen/llvm/codegen_llvm.h
rename to tirx/src/codegen/llvm/codegen_llvm.h
diff --git a/src/codegen/llvm/codegen_params.cc b/tirx/src/codegen/llvm/codegen_params.cc
similarity index 100%
rename from src/codegen/llvm/codegen_params.cc
rename to tirx/src/codegen/llvm/codegen_params.cc
diff --git a/src/codegen/llvm/codegen_params.h b/tirx/src/codegen/llvm/codegen_params.h
similarity index 100%
rename from src/codegen/llvm/codegen_params.h
rename to tirx/src/codegen/llvm/codegen_params.h
diff --git a/src/codegen/llvm/codegen_x86_64.cc b/tirx/src/codegen/llvm/codegen_x86_64.cc
similarity index 100%
rename from src/codegen/llvm/codegen_x86_64.cc
rename to tirx/src/codegen/llvm/codegen_x86_64.cc
diff --git a/src/codegen/llvm/intrin_rule_llvm.cc b/tirx/src/codegen/llvm/intrin_rule_llvm.cc
similarity index 100%
rename from src/codegen/llvm/intrin_rule_llvm.cc
rename to tirx/src/codegen/llvm/intrin_rule_llvm.cc
diff --git a/src/codegen/llvm/intrin_rule_llvm.h b/tirx/src/codegen/llvm/intrin_rule_llvm.h
similarity index 100%
rename from src/codegen/llvm/intrin_rule_llvm.h
rename to tirx/src/codegen/llvm/intrin_rule_llvm.h
diff --git a/src/codegen/llvm/llvm_instance.cc b/tirx/src/codegen/llvm/llvm_instance.cc
similarity index 100%
rename from src/codegen/llvm/llvm_instance.cc
rename to tirx/src/codegen/llvm/llvm_instance.cc
diff --git a/src/codegen/llvm/llvm_instance.h b/tirx/src/codegen/llvm/llvm_instance.h
similarity index 100%
rename from src/codegen/llvm/llvm_instance.h
rename to tirx/src/codegen/llvm/llvm_instance.h
diff --git a/src/codegen/llvm/llvm_module.cc b/tirx/src/codegen/llvm/llvm_module.cc
similarity index 100%
rename from src/codegen/llvm/llvm_module.cc
rename to tirx/src/codegen/llvm/llvm_module.cc
diff --git a/src/codegen/llvm/llvm_module.h b/tirx/src/codegen/llvm/llvm_module.h
similarity index 100%
rename from src/codegen/llvm/llvm_module.h
rename to tirx/src/codegen/llvm/llvm_module.h
diff --git a/src/codegen/llvm/target_kind.cc b/tirx/src/codegen/llvm/target_kind.cc
similarity index 100%
rename from src/codegen/llvm/target_kind.cc
rename to tirx/src/codegen/llvm/target_kind.cc
diff --git a/src/codegen/module_export.cc b/tirx/src/codegen/module_export.cc
similarity index 100%
rename from src/codegen/module_export.cc
rename to tirx/src/codegen/module_export.cc
diff --git a/src/ir/attrs.cc b/tirx/src/ir/attrs.cc
similarity index 100%
rename from src/ir/attrs.cc
rename to tirx/src/ir/attrs.cc
diff --git a/src/ir/env_func.cc b/tirx/src/ir/env_func.cc
similarity index 100%
rename from src/ir/env_func.cc
rename to tirx/src/ir/env_func.cc
diff --git a/src/ir/expr.cc b/tirx/src/ir/expr.cc
similarity index 100%
rename from src/ir/expr.cc
rename to tirx/src/ir/expr.cc
diff --git a/src/ir/expr_functor.cc b/tirx/src/ir/expr_functor.cc
similarity index 100%
rename from src/ir/expr_functor.cc
rename to tirx/src/ir/expr_functor.cc
diff --git a/src/ir/function.cc b/tirx/src/ir/function.cc
similarity index 100%
rename from src/ir/function.cc
rename to tirx/src/ir/function.cc
diff --git a/src/ir/global_info.cc b/tirx/src/ir/global_info.cc
similarity index 100%
rename from src/ir/global_info.cc
rename to tirx/src/ir/global_info.cc
diff --git a/src/ir/instrument.cc b/tirx/src/ir/instrument.cc
similarity index 100%
rename from src/ir/instrument.cc
rename to tirx/src/ir/instrument.cc
diff --git a/src/ir/location.cc b/tirx/src/ir/location.cc
similarity index 100%
rename from src/ir/location.cc
rename to tirx/src/ir/location.cc
diff --git a/src/ir/module.cc b/tirx/src/ir/module.cc
similarity index 100%
rename from src/ir/module.cc
rename to tirx/src/ir/module.cc
diff --git a/src/ir/op.cc b/tirx/src/ir/op.cc
similarity index 100%
rename from src/ir/op.cc
rename to tirx/src/ir/op.cc
diff --git a/src/ir/prim/builtin.cc b/tirx/src/ir/prim/builtin.cc
similarity index 100%
rename from src/ir/prim/builtin.cc
rename to tirx/src/ir/prim/builtin.cc
diff --git a/src/ir/prim/const_fold.h b/tirx/src/ir/prim/const_fold.h
similarity index 100%
rename from src/ir/prim/const_fold.h
rename to tirx/src/ir/prim/const_fold.h
diff --git a/src/ir/prim/deep_equal.cc b/tirx/src/ir/prim/deep_equal.cc
similarity index 100%
rename from src/ir/prim/deep_equal.cc
rename to tirx/src/ir/prim/deep_equal.cc
diff --git a/src/ir/prim/expr.cc b/tirx/src/ir/prim/expr.cc
similarity index 100%
rename from src/ir/prim/expr.cc
rename to tirx/src/ir/prim/expr.cc
diff --git a/src/ir/prim/int_operator.h b/tirx/src/ir/prim/int_operator.h
similarity index 100%
rename from src/ir/prim/int_operator.h
rename to tirx/src/ir/prim/int_operator.h
diff --git a/src/ir/prim/math.cc b/tirx/src/ir/prim/math.cc
similarity index 100%
rename from src/ir/prim/math.cc
rename to tirx/src/ir/prim/math.cc
diff --git a/src/ir/prim/op.cc b/tirx/src/ir/prim/op.cc
similarity index 100%
rename from src/ir/prim/op.cc
rename to tirx/src/ir/prim/op.cc
diff --git a/src/ir/prim/op_utils.h b/tirx/src/ir/prim/op_utils.h
similarity index 100%
rename from src/ir/prim/op_utils.h
rename to tirx/src/ir/prim/op_utils.h
diff --git a/src/ir/prim/symbolic_limits.cc b/tirx/src/ir/prim/symbolic_limits.cc
similarity index 100%
rename from src/ir/prim/symbolic_limits.cc
rename to tirx/src/ir/prim/symbolic_limits.cc
diff --git a/src/ir/prim/vector_expr.cc b/tirx/src/ir/prim/vector_expr.cc
similarity index 100%
rename from src/ir/prim/vector_expr.cc
rename to tirx/src/ir/prim/vector_expr.cc
diff --git a/src/ir/seq_stmt_mutate.h b/tirx/src/ir/seq_stmt_mutate.h
similarity index 100%
rename from src/ir/seq_stmt_mutate.h
rename to tirx/src/ir/seq_stmt_mutate.h
diff --git a/src/ir/side_effect.cc b/tirx/src/ir/side_effect.cc
similarity index 100%
rename from src/ir/side_effect.cc
rename to tirx/src/ir/side_effect.cc
diff --git a/src/ir/stmt.cc b/tirx/src/ir/stmt.cc
similarity index 100%
rename from src/ir/stmt.cc
rename to tirx/src/ir/stmt.cc
diff --git a/src/ir/subscript_expr.cc b/tirx/src/ir/subscript_expr.cc
similarity index 100%
rename from src/ir/subscript_expr.cc
rename to tirx/src/ir/subscript_expr.cc
diff --git a/src/ir/transform.cc b/tirx/src/ir/transform.cc
similarity index 100%
rename from src/ir/transform.cc
rename to tirx/src/ir/transform.cc
diff --git a/src/ir/type.cc b/tirx/src/ir/type.cc
similarity index 100%
rename from src/ir/type.cc
rename to tirx/src/ir/type.cc
diff --git a/src/ir/unique_name_supply.cc b/tirx/src/ir/unique_name_supply.cc
similarity index 100%
rename from src/ir/unique_name_supply.cc
rename to tirx/src/ir/unique_name_supply.cc
diff --git a/src/runtime/builtin_fp16.cc b/tirx/src/runtime/builtin_fp16.cc
similarity index 100%
rename from src/runtime/builtin_fp16.cc
rename to tirx/src/runtime/builtin_fp16.cc
diff --git a/src/runtime/cpu_device_api.cc b/tirx/src/runtime/cpu_device_api.cc
similarity index 100%
rename from src/runtime/cpu_device_api.cc
rename to tirx/src/runtime/cpu_device_api.cc
diff --git a/src/runtime/device_api.cc b/tirx/src/runtime/device_api.cc
similarity index 100%
rename from src/runtime/device_api.cc
rename to tirx/src/runtime/device_api.cc
diff --git a/src/runtime/launch_param.h b/tirx/src/runtime/launch_param.h
similarity index 100%
rename from src/runtime/launch_param.h
rename to tirx/src/runtime/launch_param.h
diff --git a/src/runtime/logging.cc b/tirx/src/runtime/logging.cc
similarity index 100%
rename from src/runtime/logging.cc
rename to tirx/src/runtime/logging.cc
diff --git a/src/runtime/memory_scope.cc b/tirx/src/runtime/memory_scope.cc
similarity index 100%
rename from src/runtime/memory_scope.cc
rename to tirx/src/runtime/memory_scope.cc
diff --git a/src/runtime/metadata.cc b/tirx/src/runtime/metadata.cc
similarity index 100%
rename from src/runtime/metadata.cc
rename to tirx/src/runtime/metadata.cc
diff --git a/src/runtime/metadata.h b/tirx/src/runtime/metadata.h
similarity index 100%
rename from src/runtime/metadata.h
rename to tirx/src/runtime/metadata.h
diff --git a/src/runtime/module.cc b/tirx/src/runtime/module.cc
similarity index 100%
rename from src/runtime/module.cc
rename to tirx/src/runtime/module.cc
diff --git a/src/runtime/pack_args.h b/tirx/src/runtime/pack_args.h
similarity index 100%
rename from src/runtime/pack_args.h
rename to tirx/src/runtime/pack_args.h
diff --git a/src/runtime/static_library.cc b/tirx/src/runtime/static_library.cc
similarity index 100%
rename from src/runtime/static_library.cc
rename to tirx/src/runtime/static_library.cc
diff --git a/src/runtime/static_library.h b/tirx/src/runtime/static_library.h
similarity index 100%
rename from src/runtime/static_library.h
rename to tirx/src/runtime/static_library.h
diff --git a/src/runtime/tensor.cc b/tirx/src/runtime/tensor.cc
similarity index 100%
rename from src/runtime/tensor.cc
rename to tirx/src/runtime/tensor.cc
diff --git a/src/runtime/thread_pool.cc b/tirx/src/runtime/thread_pool.cc
similarity index 100%
rename from src/runtime/thread_pool.cc
rename to tirx/src/runtime/thread_pool.cc
diff --git a/src/runtime/threading_backend.cc b/tirx/src/runtime/threading_backend.cc
similarity index 100%
rename from src/runtime/threading_backend.cc
rename to tirx/src/runtime/threading_backend.cc
diff --git a/src/runtime/threading_backend.h b/tirx/src/runtime/threading_backend.h
similarity index 100%
rename from src/runtime/threading_backend.h
rename to tirx/src/runtime/threading_backend.h
diff --git a/src/runtime/timer.cc b/tirx/src/runtime/timer.cc
similarity index 100%
rename from src/runtime/timer.cc
rename to tirx/src/runtime/timer.cc
diff --git a/src/runtime/workspace_pool.cc b/tirx/src/runtime/workspace_pool.cc
similarity index 100%
rename from src/runtime/workspace_pool.cc
rename to tirx/src/runtime/workspace_pool.cc
diff --git a/src/runtime/workspace_pool.h b/tirx/src/runtime/workspace_pool.h
similarity index 100%
rename from src/runtime/workspace_pool.h
rename to tirx/src/runtime/workspace_pool.h
diff --git a/src/script/ir_builder/base.cc b/tirx/src/script/ir_builder/base.cc
similarity index 100%
rename from src/script/ir_builder/base.cc
rename to tirx/src/script/ir_builder/base.cc
diff --git a/src/script/ir_builder/frame.cc b/tirx/src/script/ir_builder/frame.cc
similarity index 100%
rename from src/script/ir_builder/frame.cc
rename to tirx/src/script/ir_builder/frame.cc
diff --git a/src/script/ir_builder/ir.cc b/tirx/src/script/ir_builder/ir.cc
similarity index 100%
rename from src/script/ir_builder/ir.cc
rename to tirx/src/script/ir_builder/ir.cc
diff --git a/src/script/ir_builder/utils.h b/tirx/src/script/ir_builder/utils.h
similarity index 100%
rename from src/script/ir_builder/utils.h
rename to tirx/src/script/ir_builder/utils.h
diff --git a/src/script/printer/config.cc b/tirx/src/script/printer/config.cc
similarity index 100%
rename from src/script/printer/config.cc
rename to tirx/src/script/printer/config.cc
diff --git a/src/script/printer/doc.cc b/tirx/src/script/printer/doc.cc
similarity index 100%
rename from src/script/printer/doc.cc
rename to tirx/src/script/printer/doc.cc
diff --git a/src/script/printer/doc_printer.cc b/tirx/src/script/printer/doc_printer.cc
similarity index 100%
rename from src/script/printer/doc_printer.cc
rename to tirx/src/script/printer/doc_printer.cc
diff --git a/src/script/printer/doc_printer.h b/tirx/src/script/printer/doc_printer.h
similarity index 100%
rename from src/script/printer/doc_printer.h
rename to tirx/src/script/printer/doc_printer.h
diff --git a/src/script/printer/doc_translator.cc b/tirx/src/script/printer/doc_translator.cc
similarity index 100%
rename from src/script/printer/doc_translator.cc
rename to tirx/src/script/printer/doc_translator.cc
diff --git a/src/script/printer/ir/call.cc b/tirx/src/script/printer/ir/call.cc
similarity index 100%
rename from src/script/printer/ir/call.cc
rename to tirx/src/script/printer/ir/call.cc
diff --git a/src/script/printer/ir/ir.cc b/tirx/src/script/printer/ir/ir.cc
similarity index 100%
rename from src/script/printer/ir/ir.cc
rename to tirx/src/script/printer/ir/ir.cc
diff --git a/src/script/printer/ir/misc.cc b/tirx/src/script/printer/ir/misc.cc
similarity index 100%
rename from src/script/printer/ir/misc.cc
rename to tirx/src/script/printer/ir/misc.cc
diff --git a/src/script/printer/ir/prim_expr.cc b/tirx/src/script/printer/ir/prim_expr.cc
similarity index 100%
rename from src/script/printer/ir/prim_expr.cc
rename to tirx/src/script/printer/ir/prim_expr.cc
diff --git a/src/script/printer/ir/prim_type.cc b/tirx/src/script/printer/ir/prim_type.cc
similarity index 100%
rename from src/script/printer/ir/prim_type.cc
rename to tirx/src/script/printer/ir/prim_type.cc
diff --git a/src/script/printer/ir/stmt.cc b/tirx/src/script/printer/ir/stmt.cc
similarity index 100%
rename from src/script/printer/ir/stmt.cc
rename to tirx/src/script/printer/ir/stmt.cc
diff --git a/src/script/printer/ir/utils.cc b/tirx/src/script/printer/ir/utils.cc
similarity index 100%
rename from src/script/printer/ir/utils.cc
rename to tirx/src/script/printer/ir/utils.cc
diff --git a/src/script/printer/ir/utils.h b/tirx/src/script/printer/ir/utils.h
similarity index 100%
rename from src/script/printer/ir/utils.h
rename to tirx/src/script/printer/ir/utils.h
diff --git a/src/script/printer/printer.cc b/tirx/src/script/printer/printer.cc
similarity index 100%
rename from src/script/printer/printer.cc
rename to tirx/src/script/printer/printer.cc
diff --git a/src/script/printer/script_printer.cc b/tirx/src/script/printer/script_printer.cc
similarity index 100%
rename from src/script/printer/script_printer.cc
rename to tirx/src/script/printer/script_printer.cc
diff --git a/src/script/printer/utils.h b/tirx/src/script/printer/utils.h
similarity index 100%
rename from src/script/printer/utils.h
rename to tirx/src/script/printer/utils.h
diff --git a/src/support/limits.h b/tirx/src/support/limits.h
similarity index 100%
rename from src/support/limits.h
rename to tirx/src/support/limits.h
diff --git a/src/sym/analyzer.cc b/tirx/src/sym/analyzer.cc
similarity index 100%
rename from src/sym/analyzer.cc
rename to tirx/src/sym/analyzer.cc
diff --git a/src/sym/bound_deducer.cc b/tirx/src/sym/bound_deducer.cc
similarity index 100%
rename from src/sym/bound_deducer.cc
rename to tirx/src/sym/bound_deducer.cc
diff --git a/src/sym/canonical_simplify.cc b/tirx/src/sym/canonical_simplify.cc
similarity index 100%
rename from src/sym/canonical_simplify.cc
rename to tirx/src/sym/canonical_simplify.cc
diff --git a/src/sym/const_fold.h b/tirx/src/sym/const_fold.h
similarity index 100%
rename from src/sym/const_fold.h
rename to tirx/src/sym/const_fold.h
diff --git a/src/sym/const_int_bound.cc b/tirx/src/sym/const_int_bound.cc
similarity index 100%
rename from src/sym/const_int_bound.cc
rename to tirx/src/sym/const_int_bound.cc
diff --git a/src/sym/constraint_extract.cc b/tirx/src/sym/constraint_extract.cc
similarity index 100%
rename from src/sym/constraint_extract.cc
rename to tirx/src/sym/constraint_extract.cc
diff --git a/src/sym/constraint_extract.h b/tirx/src/sym/constraint_extract.h
similarity index 100%
rename from src/sym/constraint_extract.h
rename to tirx/src/sym/constraint_extract.h
diff --git a/src/sym/constraint_helpers.h b/tirx/src/sym/constraint_helpers.h
similarity index 100%
rename from src/sym/constraint_helpers.h
rename to tirx/src/sym/constraint_helpers.h
diff --git a/src/sym/detect_linear_equation.cc b/tirx/src/sym/detect_linear_equation.cc
similarity index 100%
rename from src/sym/detect_linear_equation.cc
rename to tirx/src/sym/detect_linear_equation.cc
diff --git a/src/sym/int_operator.h b/tirx/src/sym/int_operator.h
similarity index 100%
rename from src/sym/int_operator.h
rename to tirx/src/sym/int_operator.h
diff --git a/src/sym/int_set.cc b/tirx/src/sym/int_set.cc
similarity index 100%
rename from src/sym/int_set.cc
rename to tirx/src/sym/int_set.cc
diff --git a/src/sym/interval_set.h b/tirx/src/sym/interval_set.h
similarity index 100%
rename from src/sym/interval_set.h
rename to tirx/src/sym/interval_set.h
diff --git a/src/sym/iter_affine_map.cc b/tirx/src/sym/iter_affine_map.cc
similarity index 100%
rename from src/sym/iter_affine_map.cc
rename to tirx/src/sym/iter_affine_map.cc
diff --git a/src/sym/known_comparison_analyzer.cc b/tirx/src/sym/known_comparison_analyzer.cc
similarity index 100%
rename from src/sym/known_comparison_analyzer.cc
rename to tirx/src/sym/known_comparison_analyzer.cc
diff --git a/src/sym/modular_set.cc b/tirx/src/sym/modular_set.cc
similarity index 100%
rename from src/sym/modular_set.cc
rename to tirx/src/sym/modular_set.cc
diff --git a/src/sym/pattern_match.h b/tirx/src/sym/pattern_match.h
similarity index 100%
rename from src/sym/pattern_match.h
rename to tirx/src/sym/pattern_match.h
diff --git a/src/sym/product_normal_form.h b/tirx/src/sym/product_normal_form.h
similarity index 100%
rename from src/sym/product_normal_form.h
rename to tirx/src/sym/product_normal_form.h
diff --git a/src/sym/rewrite_simplify.cc b/tirx/src/sym/rewrite_simplify.cc
similarity index 100%
rename from src/sym/rewrite_simplify.cc
rename to tirx/src/sym/rewrite_simplify.cc
diff --git a/src/sym/rewrite_simplify.h b/tirx/src/sym/rewrite_simplify.h
similarity index 100%
rename from src/sym/rewrite_simplify.h
rename to tirx/src/sym/rewrite_simplify.h
diff --git a/src/sym/simplify_base.cc b/tirx/src/sym/simplify_base.cc
similarity index 100%
rename from src/sym/simplify_base.cc
rename to tirx/src/sym/simplify_base.cc
diff --git a/src/sym/simplify_base.h b/tirx/src/sym/simplify_base.h
similarity index 100%
rename from src/sym/simplify_base.h
rename to tirx/src/sym/simplify_base.h
diff --git a/src/sym/z3_prover.cc b/tirx/src/sym/z3_prover.cc
similarity index 100%
rename from src/sym/z3_prover.cc
rename to tirx/src/sym/z3_prover.cc
diff --git a/src/target/target.cc b/tirx/src/target/target.cc
similarity index 100%
rename from src/target/target.cc
rename to tirx/src/target/target.cc
diff --git a/src/target/target_kind.cc b/tirx/src/target/target_kind.cc
similarity index 100%
rename from src/target/target_kind.cc
rename to tirx/src/target/target_kind.cc
diff --git a/src/target/target_tag_registry.cc b/tirx/src/target/target_tag_registry.cc
similarity index 100%
rename from src/target/target_tag_registry.cc
rename to tirx/src/target/target_tag_registry.cc
diff --git a/src/tirx/analysis/check_well_formed.cc b/tirx/src/tirx/analysis/check_well_formed.cc
similarity index 100%
rename from src/tirx/analysis/check_well_formed.cc
rename to tirx/src/tirx/analysis/check_well_formed.cc
diff --git a/src/tirx/analysis/expr_complexity.cc b/tirx/src/tirx/analysis/expr_complexity.cc
similarity index 100%
rename from src/tirx/analysis/expr_complexity.cc
rename to tirx/src/tirx/analysis/expr_complexity.cc
diff --git a/src/tirx/analysis/stmt_finding.cc b/tirx/src/tirx/analysis/stmt_finding.cc
similarity index 100%
rename from src/tirx/analysis/stmt_finding.cc
rename to tirx/src/tirx/analysis/stmt_finding.cc
diff --git a/src/tirx/analysis/undefined_vars.cc b/tirx/src/tirx/analysis/undefined_vars.cc
similarity index 100%
rename from src/tirx/analysis/undefined_vars.cc
rename to tirx/src/tirx/analysis/undefined_vars.cc
diff --git a/src/tirx/analysis/verify_memory.cc b/tirx/src/tirx/analysis/verify_memory.cc
similarity index 100%
rename from src/tirx/analysis/verify_memory.cc
rename to tirx/src/tirx/analysis/verify_memory.cc
diff --git a/src/tirx/ir/data_type_rewriter.cc b/tirx/src/tirx/ir/data_type_rewriter.cc
similarity index 100%
rename from src/tirx/ir/data_type_rewriter.cc
rename to tirx/src/tirx/ir/data_type_rewriter.cc
diff --git a/src/tirx/ir/data_type_rewriter.h b/tirx/src/tirx/ir/data_type_rewriter.h
similarity index 100%
rename from src/tirx/ir/data_type_rewriter.h
rename to tirx/src/tirx/ir/data_type_rewriter.h
diff --git a/src/tirx/ir/function.cc b/tirx/src/tirx/ir/function.cc
similarity index 100%
rename from src/tirx/ir/function.cc
rename to tirx/src/tirx/ir/function.cc
diff --git a/src/tirx/ir/index_map.cc b/tirx/src/tirx/ir/index_map.cc
similarity index 100%
rename from src/tirx/ir/index_map.cc
rename to tirx/src/tirx/ir/index_map.cc
diff --git a/src/tirx/ir/ir_mutator_with_analyzer.cc b/tirx/src/tirx/ir/ir_mutator_with_analyzer.cc
similarity index 100%
rename from src/tirx/ir/ir_mutator_with_analyzer.cc
rename to tirx/src/tirx/ir/ir_mutator_with_analyzer.cc
diff --git a/src/tirx/ir/ir_visitor_with_analyzer.cc b/tirx/src/tirx/ir/ir_visitor_with_analyzer.cc
similarity index 100%
rename from src/tirx/ir/ir_visitor_with_analyzer.cc
rename to tirx/src/tirx/ir/ir_visitor_with_analyzer.cc
diff --git a/src/tirx/ir/layout/axis_registry.cc b/tirx/src/tirx/ir/layout/axis_registry.cc
similarity index 100%
rename from src/tirx/ir/layout/axis_registry.cc
rename to tirx/src/tirx/ir/layout/axis_registry.cc
diff --git a/src/tirx/ir/layout/axis_registry.h b/tirx/src/tirx/ir/layout/axis_registry.h
similarity index 100%
rename from src/tirx/ir/layout/axis_registry.h
rename to tirx/src/tirx/ir/layout/axis_registry.h
diff --git a/src/tirx/ir/layout/compose_layout.cc b/tirx/src/tirx/ir/layout/compose_layout.cc
similarity index 100%
rename from src/tirx/ir/layout/compose_layout.cc
rename to tirx/src/tirx/ir/layout/compose_layout.cc
diff --git a/src/tirx/ir/layout/layout.cc b/tirx/src/tirx/ir/layout/layout.cc
similarity index 100%
rename from src/tirx/ir/layout/layout.cc
rename to tirx/src/tirx/ir/layout/layout.cc
diff --git a/src/tirx/ir/layout/tile_canonicalize.cc b/tirx/src/tirx/ir/layout/tile_canonicalize.cc
similarity index 100%
rename from src/tirx/ir/layout/tile_canonicalize.cc
rename to tirx/src/tirx/ir/layout/tile_canonicalize.cc
diff --git a/src/tirx/ir/layout/tile_core.cc b/tirx/src/tirx/ir/layout/tile_core.cc
similarity index 100%
rename from src/tirx/ir/layout/tile_core.cc
rename to tirx/src/tirx/ir/layout/tile_core.cc
diff --git a/src/tirx/ir/layout/tile_internal.h b/tirx/src/tirx/ir/layout/tile_internal.h
similarity index 100%
rename from src/tirx/ir/layout/tile_internal.h
rename to tirx/src/tirx/ir/layout/tile_internal.h
diff --git a/src/tirx/ir/layout/tile_tile_ops.cc b/tirx/src/tirx/ir/layout/tile_tile_ops.cc
similarity index 100%
rename from src/tirx/ir/layout/tile_tile_ops.cc
rename to tirx/src/tirx/ir/layout/tile_tile_ops.cc
diff --git a/src/tirx/ir/layout/utils.cc b/tirx/src/tirx/ir/layout/utils.cc
similarity index 100%
rename from src/tirx/ir/layout/utils.cc
rename to tirx/src/tirx/ir/layout/utils.cc
diff --git a/src/tirx/ir/layout/utils.h b/tirx/src/tirx/ir/layout/utils.h
similarity index 100%
rename from src/tirx/ir/layout/utils.h
rename to tirx/src/tirx/ir/layout/utils.h
diff --git a/src/tirx/ir/specialize.cc b/tirx/src/tirx/ir/specialize.cc
similarity index 100%
rename from src/tirx/ir/specialize.cc
rename to tirx/src/tirx/ir/specialize.cc
diff --git a/src/tirx/ir/stmt.cc b/tirx/src/tirx/ir/stmt.cc
similarity index 100%
rename from src/tirx/ir/stmt.cc
rename to tirx/src/tirx/ir/stmt.cc
diff --git a/src/tirx/ir/stmt_functor.cc b/tirx/src/tirx/ir/stmt_functor.cc
similarity index 100%
rename from src/tirx/ir/stmt_functor.cc
rename to tirx/src/tirx/ir/stmt_functor.cc
diff --git a/src/tirx/ir/tensor.cc b/tirx/src/tirx/ir/tensor.cc
similarity index 100%
rename from src/tirx/ir/tensor.cc
rename to tirx/src/tirx/ir/tensor.cc
diff --git a/src/tirx/ir/tensor_load.cc b/tirx/src/tirx/ir/tensor_load.cc
similarity index 100%
rename from src/tirx/ir/tensor_load.cc
rename to tirx/src/tirx/ir/tensor_load.cc
diff --git a/src/tirx/ir/transform.cc b/tirx/src/tirx/ir/transform.cc
similarity index 100%
rename from src/tirx/ir/transform.cc
rename to tirx/src/tirx/ir/transform.cc
diff --git a/src/tirx/ir/type.cc b/tirx/src/tirx/ir/type.cc
similarity index 100%
rename from src/tirx/ir/type.cc
rename to tirx/src/tirx/ir/type.cc
diff --git a/src/tirx/op/abi.cc b/tirx/src/tirx/op/abi.cc
similarity index 100%
rename from src/tirx/op/abi.cc
rename to tirx/src/tirx/op/abi.cc
diff --git a/src/tirx/op/annotation.cc b/tirx/src/tirx/op/annotation.cc
similarity index 100%
rename from src/tirx/op/annotation.cc
rename to tirx/src/tirx/op/annotation.cc
diff --git a/src/tirx/op/cpu.cc b/tirx/src/tirx/op/cpu.cc
similarity index 100%
rename from src/tirx/op/cpu.cc
rename to tirx/src/tirx/op/cpu.cc
diff --git a/src/tirx/op/gpu.cc b/tirx/src/tirx/op/gpu.cc
similarity index 100%
rename from src/tirx/op/gpu.cc
rename to tirx/src/tirx/op/gpu.cc
diff --git a/src/tirx/op/math.cc b/tirx/src/tirx/op/math.cc
similarity index 100%
rename from src/tirx/op/math.cc
rename to tirx/src/tirx/op/math.cc
diff --git a/src/tirx/op/memory.cc b/tirx/src/tirx/op/memory.cc
similarity index 100%
rename from src/tirx/op/memory.cc
rename to tirx/src/tirx/op/memory.cc
diff --git a/src/tirx/op/region.cc b/tirx/src/tirx/op/region.cc
similarity index 100%
rename from src/tirx/op/region.cc
rename to tirx/src/tirx/op/region.cc
diff --git a/src/tirx/op/vector.cc b/tirx/src/tirx/op/vector.cc
similarity index 100%
rename from src/tirx/op/vector.cc
rename to tirx/src/tirx/op/vector.cc
diff --git a/src/tirx/script/ir_builder/frame.cc b/tirx/src/tirx/script/ir_builder/frame.cc
similarity index 100%
rename from src/tirx/script/ir_builder/frame.cc
rename to tirx/src/tirx/script/ir_builder/frame.cc
diff --git a/src/tirx/script/ir_builder/ir.cc b/tirx/src/tirx/script/ir_builder/ir.cc
similarity index 100%
rename from src/tirx/script/ir_builder/ir.cc
rename to tirx/src/tirx/script/ir_builder/ir.cc
diff --git a/src/tirx/script/ir_builder/utils.h b/tirx/src/tirx/script/ir_builder/utils.h
similarity index 100%
rename from src/tirx/script/ir_builder/utils.h
rename to tirx/src/tirx/script/ir_builder/utils.h
diff --git a/src/tirx/script/printer/block.cc b/tirx/src/tirx/script/printer/block.cc
similarity index 100%
rename from src/tirx/script/printer/block.cc
rename to tirx/src/tirx/script/printer/block.cc
diff --git a/src/tirx/script/printer/expr.cc b/tirx/src/tirx/script/printer/expr.cc
similarity index 100%
rename from src/tirx/script/printer/expr.cc
rename to tirx/src/tirx/script/printer/expr.cc
diff --git a/src/tirx/script/printer/for_loop.cc b/tirx/src/tirx/script/printer/for_loop.cc
similarity index 100%
rename from src/tirx/script/printer/for_loop.cc
rename to tirx/src/tirx/script/printer/for_loop.cc
diff --git a/src/tirx/script/printer/function.cc b/tirx/src/tirx/script/printer/function.cc
similarity index 100%
rename from src/tirx/script/printer/function.cc
rename to tirx/src/tirx/script/printer/function.cc
diff --git a/src/tirx/script/printer/ir.cc b/tirx/src/tirx/script/printer/ir.cc
similarity index 100%
rename from src/tirx/script/printer/ir.cc
rename to tirx/src/tirx/script/printer/ir.cc
diff --git a/src/tirx/script/printer/stmt.cc b/tirx/src/tirx/script/printer/stmt.cc
similarity index 100%
rename from src/tirx/script/printer/stmt.cc
rename to tirx/src/tirx/script/printer/stmt.cc
diff --git a/src/tirx/script/printer/tensor.cc b/tirx/src/tirx/script/printer/tensor.cc
similarity index 100%
rename from src/tirx/script/printer/tensor.cc
rename to tirx/src/tirx/script/printer/tensor.cc
diff --git a/src/tirx/tile/context.cc b/tirx/src/tirx/tile/context.cc
similarity index 100%
rename from src/tirx/tile/context.cc
rename to tirx/src/tirx/tile/context.cc
diff --git a/src/tirx/tile/dispatch.cc b/tirx/src/tirx/tile/dispatch.cc
similarity index 100%
rename from src/tirx/tile/dispatch.cc
rename to tirx/src/tirx/tile/dispatch.cc
diff --git a/src/tirx/tile/lower.cc b/tirx/src/tirx/tile/lower.cc
similarity index 100%
rename from src/tirx/tile/lower.cc
rename to tirx/src/tirx/tile/lower.cc
diff --git a/src/tirx/tile/op.cc b/tirx/src/tirx/tile/op.cc
similarity index 100%
rename from src/tirx/tile/op.cc
rename to tirx/src/tirx/tile/op.cc
diff --git a/src/tirx/tile/thread_scope.cc b/tirx/src/tirx/tile/thread_scope.cc
similarity index 100%
rename from src/tirx/tile/thread_scope.cc
rename to tirx/src/tirx/tile/thread_scope.cc
diff --git a/src/tirx/transform/bind_target.cc b/tirx/src/tirx/transform/bind_target.cc
similarity index 100%
rename from src/tirx/transform/bind_target.cc
rename to tirx/src/tirx/transform/bind_target.cc
diff --git a/src/tirx/transform/common_subexpr_elim.cc b/tirx/src/tirx/transform/common_subexpr_elim.cc
similarity index 100%
rename from src/tirx/transform/common_subexpr_elim.cc
rename to tirx/src/tirx/transform/common_subexpr_elim.cc
diff --git a/src/tirx/transform/decorate_device_scope.cc b/tirx/src/tirx/transform/decorate_device_scope.cc
similarity index 100%
rename from src/tirx/transform/decorate_device_scope.cc
rename to tirx/src/tirx/transform/decorate_device_scope.cc
diff --git a/src/tirx/transform/dtype_conversion.cc b/tirx/src/tirx/transform/dtype_conversion.cc
similarity index 100%
rename from src/tirx/transform/dtype_conversion.cc
rename to tirx/src/tirx/transform/dtype_conversion.cc
diff --git a/src/tirx/transform/dtype_conversion.h b/tirx/src/tirx/transform/dtype_conversion.h
similarity index 100%
rename from src/tirx/transform/dtype_conversion.h
rename to tirx/src/tirx/transform/dtype_conversion.h
diff --git a/src/tirx/transform/flatten_tensor.cc b/tirx/src/tirx/transform/flatten_tensor.cc
similarity index 100%
rename from src/tirx/transform/flatten_tensor.cc
rename to tirx/src/tirx/transform/flatten_tensor.cc
diff --git a/src/tirx/transform/flattened_tensor.h b/tirx/src/tirx/transform/flattened_tensor.h
similarity index 100%
rename from src/tirx/transform/flattened_tensor.h
rename to tirx/src/tirx/transform/flattened_tensor.h
diff --git a/src/tirx/transform/force_narrow_index_to_i32.cc b/tirx/src/tirx/transform/force_narrow_index_to_i32.cc
similarity index 100%
rename from src/tirx/transform/force_narrow_index_to_i32.cc
rename to tirx/src/tirx/transform/force_narrow_index_to_i32.cc
diff --git a/src/tirx/transform/function_utils.cc b/tirx/src/tirx/transform/function_utils.cc
similarity index 100%
rename from src/tirx/transform/function_utils.cc
rename to tirx/src/tirx/transform/function_utils.cc
diff --git a/src/tirx/transform/hoist_if.cc b/tirx/src/tirx/transform/hoist_if.cc
similarity index 100%
rename from src/tirx/transform/hoist_if.cc
rename to tirx/src/tirx/transform/hoist_if.cc
diff --git a/src/tirx/transform/ir_utils.cc b/tirx/src/tirx/transform/ir_utils.cc
similarity index 100%
rename from src/tirx/transform/ir_utils.cc
rename to tirx/src/tirx/transform/ir_utils.cc
diff --git a/src/tirx/transform/ir_utils.h b/tirx/src/tirx/transform/ir_utils.h
similarity index 100%
rename from src/tirx/transform/ir_utils.h
rename to tirx/src/tirx/transform/ir_utils.h
diff --git a/src/tirx/transform/lower_intrin.cc b/tirx/src/tirx/transform/lower_intrin.cc
similarity index 100%
rename from src/tirx/transform/lower_intrin.cc
rename to tirx/src/tirx/transform/lower_intrin.cc
diff --git a/src/tirx/transform/lower_thread_allreduce.cc b/tirx/src/tirx/transform/lower_thread_allreduce.cc
similarity index 100%
rename from src/tirx/transform/lower_thread_allreduce.cc
rename to tirx/src/tirx/transform/lower_thread_allreduce.cc
diff --git a/src/tirx/transform/lower_tirx.cc b/tirx/src/tirx/transform/lower_tirx.cc
similarity index 100%
rename from src/tirx/transform/lower_tirx.cc
rename to tirx/src/tirx/transform/lower_tirx.cc
diff --git a/src/tirx/transform/lower_tirx_cleanup.cc b/tirx/src/tirx/transform/lower_tirx_cleanup.cc
similarity index 100%
rename from src/tirx/transform/lower_tirx_cleanup.cc
rename to tirx/src/tirx/transform/lower_tirx_cleanup.cc
diff --git a/src/tirx/transform/lower_tirx_dedup_tensormap.cc b/tirx/src/tirx/transform/lower_tirx_dedup_tensormap.cc
similarity index 100%
rename from src/tirx/transform/lower_tirx_dedup_tensormap.cc
rename to tirx/src/tirx/transform/lower_tirx_dedup_tensormap.cc
diff --git a/src/tirx/transform/lower_tvm_builtin.cc b/tirx/src/tirx/transform/lower_tvm_builtin.cc
similarity index 100%
rename from src/tirx/transform/lower_tvm_builtin.cc
rename to tirx/src/tirx/transform/lower_tvm_builtin.cc
diff --git a/src/tirx/transform/make_packed_api.cc b/tirx/src/tirx/transform/make_packed_api.cc
similarity index 100%
rename from src/tirx/transform/make_packed_api.cc
rename to tirx/src/tirx/transform/make_packed_api.cc
diff --git a/src/tirx/transform/narrow_datatype.cc b/tirx/src/tirx/transform/narrow_datatype.cc
similarity index 100%
rename from src/tirx/transform/narrow_datatype.cc
rename to tirx/src/tirx/transform/narrow_datatype.cc
diff --git a/src/tirx/transform/remove_no_op.cc b/tirx/src/tirx/transform/remove_no_op.cc
similarity index 100%
rename from src/tirx/transform/remove_no_op.cc
rename to tirx/src/tirx/transform/remove_no_op.cc
diff --git a/src/tirx/transform/remove_no_op.h b/tirx/src/tirx/transform/remove_no_op.h
similarity index 100%
rename from src/tirx/transform/remove_no_op.h
rename to tirx/src/tirx/transform/remove_no_op.h
diff --git a/src/tirx/transform/renormalize_split_pattern.cc b/tirx/src/tirx/transform/renormalize_split_pattern.cc
similarity index 100%
rename from src/tirx/transform/renormalize_split_pattern.cc
rename to tirx/src/tirx/transform/renormalize_split_pattern.cc
diff --git a/src/tirx/transform/skip_assert.cc b/tirx/src/tirx/transform/skip_assert.cc
similarity index 100%
rename from src/tirx/transform/skip_assert.cc
rename to tirx/src/tirx/transform/skip_assert.cc
diff --git a/src/tirx/transform/split_host_device.cc b/tirx/src/tirx/transform/split_host_device.cc
similarity index 100%
rename from src/tirx/transform/split_host_device.cc
rename to tirx/src/tirx/transform/split_host_device.cc
diff --git a/src/tirx/transform/stmt_simplify.cc b/tirx/src/tirx/transform/stmt_simplify.cc
similarity index 100%
rename from src/tirx/transform/stmt_simplify.cc
rename to tirx/src/tirx/transform/stmt_simplify.cc
diff --git a/src/tirx/transform/storage_rewrite.cc b/tirx/src/tirx/transform/storage_rewrite.cc
similarity index 100%
rename from src/tirx/transform/storage_rewrite.cc
rename to tirx/src/tirx/transform/storage_rewrite.cc
diff --git a/src/tirx/transform/tvm_ffi_binder.cc b/tirx/src/tirx/transform/tvm_ffi_binder.cc
similarity index 100%
rename from src/tirx/transform/tvm_ffi_binder.cc
rename to tirx/src/tirx/transform/tvm_ffi_binder.cc
diff --git a/src/tirx/transform/tvm_ffi_binder.h b/tirx/src/tirx/transform/tvm_ffi_binder.h
similarity index 100%
rename from src/tirx/transform/tvm_ffi_binder.h
rename to tirx/src/tirx/transform/tvm_ffi_binder.h
diff --git a/src/tirx/transform/unroll_loop.cc b/tirx/src/tirx/transform/unroll_loop.cc
similarity index 100%
rename from src/tirx/transform/unroll_loop.cc
rename to tirx/src/tirx/transform/unroll_loop.cc
diff --git a/src/tirx/transform/unsupported_dtype_legalize.cc b/tirx/src/tirx/transform/unsupported_dtype_legalize.cc
similarity index 100%
rename from src/tirx/transform/unsupported_dtype_legalize.cc
rename to tirx/src/tirx/transform/unsupported_dtype_legalize.cc
diff --git a/src/tirx/transform/vectorize_loop.cc b/tirx/src/tirx/transform/vectorize_loop.cc
similarity index 100%
rename from src/tirx/transform/vectorize_loop.cc
rename to tirx/src/tirx/transform/vectorize_loop.cc
diff --git a/tests/cpp/ir/expr_functor_test.cc b/tirx/tests/cpp/ir/expr_functor_test.cc
similarity index 100%
rename from tests/cpp/ir/expr_functor_test.cc
rename to tirx/tests/cpp/ir/expr_functor_test.cc
diff --git a/tests/cpp/ir/op_signature_test.cc b/tirx/tests/cpp/ir/op_signature_test.cc
similarity index 100%
rename from tests/cpp/ir/op_signature_test.cc
rename to tirx/tests/cpp/ir/op_signature_test.cc
diff --git a/tests/cpp/ir/tensor_load_test.cc b/tirx/tests/cpp/ir/tensor_load_test.cc
similarity index 100%
rename from tests/cpp/ir/tensor_load_test.cc
rename to tirx/tests/cpp/ir/tensor_load_test.cc
diff --git a/tests/cpp/runtime/logging_test.cc b/tirx/tests/cpp/runtime/logging_test.cc
similarity index 100%
rename from tests/cpp/runtime/logging_test.cc
rename to tirx/tests/cpp/runtime/logging_test.cc
diff --git a/tests/cpp/runtime/memory_scope_test.cc b/tirx/tests/cpp/runtime/memory_scope_test.cc
similarity index 100%
rename from tests/cpp/runtime/memory_scope_test.cc
rename to tirx/tests/cpp/runtime/memory_scope_test.cc
diff --git a/tests/cpp/sym/pattern_match_test.cc b/tirx/tests/cpp/sym/pattern_match_test.cc
similarity index 100%
rename from tests/cpp/sym/pattern_match_test.cc
rename to tirx/tests/cpp/sym/pattern_match_test.cc
diff --git a/tests/cpp/target/config_schema_test.cc b/tirx/tests/cpp/target/config_schema_test.cc
similarity index 100%
rename from tests/cpp/target/config_schema_test.cc
rename to tirx/tests/cpp/target/config_schema_test.cc
diff --git a/tests/cpp/target/target_test.cc b/tirx/tests/cpp/target/target_test.cc
similarity index 100%
rename from tests/cpp/target/target_test.cc
rename to tirx/tests/cpp/target/target_test.cc
diff --git a/tests/cpp/tirx/stmt_functor_test.cc b/tirx/tests/cpp/tirx/stmt_functor_test.cc
similarity index 100%
rename from tests/cpp/tirx/stmt_functor_test.cc
rename to tirx/tests/cpp/tirx/stmt_functor_test.cc
diff --git a/tests/cpp/tirx/tensor_var_test.cc b/tirx/tests/cpp/tirx/tensor_var_test.cc
similarity index 100%
rename from tests/cpp/tirx/tensor_var_test.cc
rename to tirx/tests/cpp/tirx/tensor_var_test.cc
diff --git a/tests/cpp/tirx/thread_scope_test.cc b/tirx/tests/cpp/tirx/thread_scope_test.cc
similarity index 100%
rename from tests/cpp/tirx/thread_scope_test.cc
rename to tirx/tests/cpp/tirx/thread_scope_test.cc
diff --git a/tests/python/all-platform-minimal-test/test_validate_runtime_library.py b/tirx/tests/python/all-platform-minimal-test/test_validate_runtime_library.py
similarity index 100%
rename from tests/python/all-platform-minimal-test/test_validate_runtime_library.py
rename to tirx/tests/python/all-platform-minimal-test/test_validate_runtime_library.py
diff --git a/tests/python/backend/cuda/codegen/test_cuda.py b/tirx/tests/python/backend/cuda/codegen/test_cuda.py
similarity index 100%
rename from tests/python/backend/cuda/codegen/test_cuda.py
rename to tirx/tests/python/backend/cuda/codegen/test_cuda.py
diff --git a/tests/python/backend/cuda/codegen/test_fastmath.py b/tirx/tests/python/backend/cuda/codegen/test_fastmath.py
similarity index 100%
rename from tests/python/backend/cuda/codegen/test_fastmath.py
rename to tirx/tests/python/backend/cuda/codegen/test_fastmath.py
diff --git a/tests/python/backend/cuda/codegen/test_fp4.py b/tirx/tests/python/backend/cuda/codegen/test_fp4.py
similarity index 100%
rename from tests/python/backend/cuda/codegen/test_fp4.py
rename to tirx/tests/python/backend/cuda/codegen/test_fp4.py
diff --git a/tests/python/backend/cuda/codegen/test_fp8.py b/tirx/tests/python/backend/cuda/codegen/test_fp8.py
similarity index 100%
rename from tests/python/backend/cuda/codegen/test_fp8.py
rename to tirx/tests/python/backend/cuda/codegen/test_fp8.py
diff --git a/tests/python/backend/cuda/test_cuda_host_bundle.py b/tirx/tests/python/backend/cuda/test_cuda_host_bundle.py
similarity index 100%
rename from tests/python/backend/cuda/test_cuda_host_bundle.py
rename to tirx/tests/python/backend/cuda/test_cuda_host_bundle.py
diff --git a/tests/python/backend/cuda/test_cuda_module_export.py b/tirx/tests/python/backend/cuda/test_cuda_module_export.py
similarity index 100%
rename from tests/python/backend/cuda/test_cuda_module_export.py
rename to tirx/tests/python/backend/cuda/test_cuda_module_export.py
diff --git a/tests/python/backend/hexagon/codegen/test_hexagon.py b/tirx/tests/python/backend/hexagon/codegen/test_hexagon.py
similarity index 100%
rename from tests/python/backend/hexagon/codegen/test_hexagon.py
rename to tirx/tests/python/backend/hexagon/codegen/test_hexagon.py
diff --git a/tests/python/backend/metal/codegen/test_metal.py b/tirx/tests/python/backend/metal/codegen/test_metal.py
similarity index 100%
rename from tests/python/backend/metal/codegen/test_metal.py
rename to tirx/tests/python/backend/metal/codegen/test_metal.py
diff --git a/tests/python/backend/metal/test_metal_compile_callback.py b/tirx/tests/python/backend/metal/test_metal_compile_callback.py
similarity index 100%
rename from tests/python/backend/metal/test_metal_compile_callback.py
rename to tirx/tests/python/backend/metal/test_metal_compile_callback.py
diff --git a/tests/python/backend/metal/test_metal_module_export.py b/tirx/tests/python/backend/metal/test_metal_module_export.py
similarity index 100%
rename from tests/python/backend/metal/test_metal_module_export.py
rename to tirx/tests/python/backend/metal/test_metal_module_export.py
diff --git a/tests/python/backend/opencl/codegen/test_opencl.py b/tirx/tests/python/backend/opencl/codegen/test_opencl.py
similarity index 100%
rename from tests/python/backend/opencl/codegen/test_opencl.py
rename to tirx/tests/python/backend/opencl/codegen/test_opencl.py
diff --git a/tests/python/backend/opencl/test_opencl_module_export.py b/tirx/tests/python/backend/opencl/test_opencl_module_export.py
similarity index 100%
rename from tests/python/backend/opencl/test_opencl_module_export.py
rename to tirx/tests/python/backend/opencl/test_opencl_module_export.py
diff --git a/tests/python/backend/rocm/codegen/test_rocm.py b/tirx/tests/python/backend/rocm/codegen/test_rocm.py
similarity index 100%
rename from tests/python/backend/rocm/codegen/test_rocm.py
rename to tirx/tests/python/backend/rocm/codegen/test_rocm.py
diff --git a/tests/python/backend/rocm/test_rocm_module_export.py b/tirx/tests/python/backend/rocm/test_rocm_module_export.py
similarity index 100%
rename from tests/python/backend/rocm/test_rocm_module_export.py
rename to tirx/tests/python/backend/rocm/test_rocm_module_export.py
diff --git a/tests/python/backend/vulkan/codegen/test_vulkan.py b/tirx/tests/python/backend/vulkan/codegen/test_vulkan.py
similarity index 100%
rename from tests/python/backend/vulkan/codegen/test_vulkan.py
rename to tirx/tests/python/backend/vulkan/codegen/test_vulkan.py
diff --git a/tests/python/backend/vulkan/test_vulkan_module_export.py b/tirx/tests/python/backend/vulkan/test_vulkan_module_export.py
similarity index 100%
rename from tests/python/backend/vulkan/test_vulkan_module_export.py
rename to tirx/tests/python/backend/vulkan/test_vulkan_module_export.py
diff --git a/tests/python/backend/webgpu/codegen/test_webgpu.py b/tirx/tests/python/backend/webgpu/codegen/test_webgpu.py
similarity index 100%
rename from tests/python/backend/webgpu/codegen/test_webgpu.py
rename to tirx/tests/python/backend/webgpu/codegen/test_webgpu.py
diff --git a/tests/python/codegen/c/test_c.py b/tirx/tests/python/codegen/c/test_c.py
similarity index 100%
rename from tests/python/codegen/c/test_c.py
rename to tirx/tests/python/codegen/c/test_c.py
diff --git a/tests/python/codegen/gpu_common/test_allreduce.py b/tirx/tests/python/codegen/gpu_common/test_allreduce.py
similarity index 100%
rename from tests/python/codegen/gpu_common/test_allreduce.py
rename to tirx/tests/python/codegen/gpu_common/test_allreduce.py
diff --git a/tests/python/codegen/gpu_common/test_device.py b/tirx/tests/python/codegen/gpu_common/test_device.py
similarity index 100%
rename from tests/python/codegen/gpu_common/test_device.py
rename to tirx/tests/python/codegen/gpu_common/test_device.py
diff --git a/tests/python/codegen/gpu_common/test_gpu_intrinsics.py b/tirx/tests/python/codegen/gpu_common/test_gpu_intrinsics.py
similarity index 100%
rename from tests/python/codegen/gpu_common/test_gpu_intrinsics.py
rename to tirx/tests/python/codegen/gpu_common/test_gpu_intrinsics.py
diff --git a/tests/python/codegen/llvm/test_aarch64.py b/tirx/tests/python/codegen/llvm/test_aarch64.py
similarity index 100%
rename from tests/python/codegen/llvm/test_aarch64.py
rename to tirx/tests/python/codegen/llvm/test_aarch64.py
diff --git a/tests/python/codegen/llvm/test_arm.py b/tirx/tests/python/codegen/llvm/test_arm.py
similarity index 100%
rename from tests/python/codegen/llvm/test_arm.py
rename to tirx/tests/python/codegen/llvm/test_arm.py
diff --git a/tests/python/codegen/llvm/test_llvm.py b/tirx/tests/python/codegen/llvm/test_llvm.py
similarity index 100%
rename from tests/python/codegen/llvm/test_llvm.py
rename to tirx/tests/python/codegen/llvm/test_llvm.py
diff --git a/tests/python/codegen/llvm/test_riscv.py b/tirx/tests/python/codegen/llvm/test_riscv.py
similarity index 100%
rename from tests/python/codegen/llvm/test_riscv.py
rename to tirx/tests/python/codegen/llvm/test_riscv.py
diff --git a/tests/python/codegen/llvm/test_scalable_vector.py b/tirx/tests/python/codegen/llvm/test_scalable_vector.py
similarity index 100%
rename from tests/python/codegen/llvm/test_scalable_vector.py
rename to tirx/tests/python/codegen/llvm/test_scalable_vector.py
diff --git a/tests/python/codegen/llvm/test_x86.py b/tirx/tests/python/codegen/llvm/test_x86.py
similarity index 100%
rename from tests/python/codegen/llvm/test_x86.py
rename to tirx/tests/python/codegen/llvm/test_x86.py
diff --git a/tests/python/codegen/test_argument_validation.py b/tirx/tests/python/codegen/test_argument_validation.py
similarity index 100%
rename from tests/python/codegen/test_argument_validation.py
rename to tirx/tests/python/codegen/test_argument_validation.py
diff --git a/tests/python/codegen/test_assert.py b/tirx/tests/python/codegen/test_assert.py
similarity index 100%
rename from tests/python/codegen/test_assert.py
rename to tirx/tests/python/codegen/test_assert.py
diff --git a/tests/python/codegen/test_bool.py b/tirx/tests/python/codegen/test_bool.py
similarity index 100%
rename from tests/python/codegen/test_bool.py
rename to tirx/tests/python/codegen/test_bool.py
diff --git a/tests/python/codegen/test_common.py b/tirx/tests/python/codegen/test_common.py
similarity index 100%
rename from tests/python/codegen/test_common.py
rename to tirx/tests/python/codegen/test_common.py
diff --git a/tests/python/codegen/test_extern.py b/tirx/tests/python/codegen/test_extern.py
similarity index 100%
rename from tests/python/codegen/test_extern.py
rename to tirx/tests/python/codegen/test_extern.py
diff --git a/tests/python/codegen/test_static_init.py b/tirx/tests/python/codegen/test_static_init.py
similarity index 100%
rename from tests/python/codegen/test_static_init.py
rename to tirx/tests/python/codegen/test_static_init.py
diff --git a/tests/python/ir/test_ir_attrs.py b/tirx/tests/python/ir/test_ir_attrs.py
similarity index 100%
rename from tests/python/ir/test_ir_attrs.py
rename to tirx/tests/python/ir/test_ir_attrs.py
diff --git a/tests/python/ir/test_ir_json.py b/tirx/tests/python/ir/test_ir_json.py
similarity index 100%
rename from tests/python/ir/test_ir_json.py
rename to tirx/tests/python/ir/test_ir_json.py
diff --git a/tests/python/ir/test_ir_module.py b/tirx/tests/python/ir/test_ir_module.py
similarity index 100%
rename from tests/python/ir/test_ir_module.py
rename to tirx/tests/python/ir/test_ir_module.py
diff --git a/tests/python/ir/test_ir_prim_expr_deep_equal.py b/tirx/tests/python/ir/test_ir_prim_expr_deep_equal.py
similarity index 100%
rename from tests/python/ir/test_ir_prim_expr_deep_equal.py
rename to tirx/tests/python/ir/test_ir_prim_expr_deep_equal.py
diff --git a/tests/python/ir/test_ir_type.py b/tirx/tests/python/ir/test_ir_type.py
similarity index 100%
rename from tests/python/ir/test_ir_type.py
rename to tirx/tests/python/ir/test_ir_type.py
diff --git a/tests/python/ir/test_nonnullable_ir.py b/tirx/tests/python/ir/test_nonnullable_ir.py
similarity index 100%
rename from tests/python/ir/test_nonnullable_ir.py
rename to tirx/tests/python/ir/test_nonnullable_ir.py
diff --git a/tests/python/ir/test_roundtrip_runtime_module.py b/tirx/tests/python/ir/test_roundtrip_runtime_module.py
similarity index 100%
rename from tests/python/ir/test_roundtrip_runtime_module.py
rename to tirx/tests/python/ir/test_roundtrip_runtime_module.py
diff --git a/tests/python/runtime/test_runtime_device_api.py b/tirx/tests/python/runtime/test_runtime_device_api.py
similarity index 100%
rename from tests/python/runtime/test_runtime_device_api.py
rename to tirx/tests/python/runtime/test_runtime_device_api.py
diff --git a/tests/python/runtime/test_runtime_memory_scope.py b/tirx/tests/python/runtime/test_runtime_memory_scope.py
similarity index 100%
rename from tests/python/runtime/test_runtime_memory_scope.py
rename to tirx/tests/python/runtime/test_runtime_memory_scope.py
diff --git a/tests/python/runtime/test_runtime_module.py b/tirx/tests/python/runtime/test_runtime_module.py
similarity index 100%
rename from tests/python/runtime/test_runtime_module.py
rename to tirx/tests/python/runtime/test_runtime_module.py
diff --git a/tests/python/runtime/test_runtime_tensor.py b/tirx/tests/python/runtime/test_runtime_tensor.py
similarity index 100%
rename from tests/python/runtime/test_runtime_tensor.py
rename to tirx/tests/python/runtime/test_runtime_tensor.py
diff --git a/tests/python/script/conftest.py b/tirx/tests/python/script/conftest.py
similarity index 100%
rename from tests/python/script/conftest.py
rename to tirx/tests/python/script/conftest.py
diff --git a/tests/python/script/minilang.py b/tirx/tests/python/script/minilang.py
similarity index 100%
rename from tests/python/script/minilang.py
rename to tirx/tests/python/script/minilang.py
diff --git a/tests/python/script/test_basic_usage.py b/tirx/tests/python/script/test_basic_usage.py
similarity index 100%
rename from tests/python/script/test_basic_usage.py
rename to tirx/tests/python/script/test_basic_usage.py
diff --git a/tests/python/script/test_builder_void_emission.py b/tirx/tests/python/script/test_builder_void_emission.py
similarity index 100%
rename from tests/python/script/test_builder_void_emission.py
rename to tirx/tests/python/script/test_builder_void_emission.py
diff --git a/tests/python/script/test_cross_function_call.py b/tirx/tests/python/script/test_cross_function_call.py
similarity index 100%
rename from tests/python/script/test_cross_function_call.py
rename to tirx/tests/python/script/test_cross_function_call.py
diff --git a/tests/python/script/test_dynamic_shape.py b/tirx/tests/python/script/test_dynamic_shape.py
similarity index 100%
rename from tests/python/script/test_dynamic_shape.py
rename to tirx/tests/python/script/test_dynamic_shape.py
diff --git a/tests/python/script/test_error_handling.py b/tirx/tests/python/script/test_error_handling.py
similarity index 100%
rename from tests/python/script/test_error_handling.py
rename to tirx/tests/python/script/test_error_handling.py
diff --git a/tests/python/script/test_ir_builder.py b/tirx/tests/python/script/test_ir_builder.py
similarity index 100%
rename from tests/python/script/test_ir_builder.py
rename to tirx/tests/python/script/test_ir_builder.py
diff --git a/tests/python/script/test_jit_specialization.py b/tirx/tests/python/script/test_jit_specialization.py
similarity index 100%
rename from tests/python/script/test_jit_specialization.py
rename to tirx/tests/python/script/test_jit_specialization.py
diff --git a/tests/python/script/test_meta_programming.py b/tirx/tests/python/script/test_meta_programming.py
similarity index 100%
rename from tests/python/script/test_meta_programming.py
rename to tirx/tests/python/script/test_meta_programming.py
diff --git a/tests/python/script/test_parser_entry.py b/tirx/tests/python/script/test_parser_entry.py
similarity index 100%
rename from tests/python/script/test_parser_entry.py
rename to tirx/tests/python/script/test_parser_entry.py
diff --git a/tests/python/script/test_script_printer_doc.py b/tirx/tests/python/script/test_script_printer_doc.py
similarity index 100%
rename from tests/python/script/test_script_printer_doc.py
rename to tirx/tests/python/script/test_script_printer_doc.py
diff --git a/tests/python/script/test_script_printer_entry.py b/tirx/tests/python/script/test_script_printer_entry.py
similarity index 100%
rename from tests/python/script/test_script_printer_entry.py
rename to tirx/tests/python/script/test_script_printer_entry.py
diff --git a/tests/python/script/test_source_locations.py b/tirx/tests/python/script/test_source_locations.py
similarity index 100%
rename from tests/python/script/test_source_locations.py
rename to tirx/tests/python/script/test_source_locations.py
diff --git a/tests/python/script/test_special_parser_protocol.py b/tirx/tests/python/script/test_special_parser_protocol.py
similarity index 100%
rename from tests/python/script/test_special_parser_protocol.py
rename to tirx/tests/python/script/test_special_parser_protocol.py
diff --git a/tests/python/support/test_ccache.py b/tirx/tests/python/support/test_ccache.py
similarity index 100%
rename from tests/python/support/test_ccache.py
rename to tirx/tests/python/support/test_ccache.py
diff --git a/tests/python/support/test_nvcc.py b/tirx/tests/python/support/test_nvcc.py
similarity index 100%
rename from tests/python/support/test_nvcc.py
rename to tirx/tests/python/support/test_nvcc.py
diff --git a/tests/python/support/test_popen_pool.py b/tirx/tests/python/support/test_popen_pool.py
similarity index 100%
rename from tests/python/support/test_popen_pool.py
rename to tirx/tests/python/support/test_popen_pool.py
diff --git a/tests/python/support/test_util.py b/tirx/tests/python/support/test_util.py
similarity index 100%
rename from tests/python/support/test_util.py
rename to tirx/tests/python/support/test_util.py
diff --git a/tests/python/sym/test_sym_analyzer_object.py b/tirx/tests/python/sym/test_sym_analyzer_object.py
similarity index 100%
rename from tests/python/sym/test_sym_analyzer_object.py
rename to tirx/tests/python/sym/test_sym_analyzer_object.py
diff --git a/tests/python/sym/test_sym_const_int_bound.py b/tirx/tests/python/sym/test_sym_const_int_bound.py
similarity index 100%
rename from tests/python/sym/test_sym_const_int_bound.py
rename to tirx/tests/python/sym/test_sym_const_int_bound.py
diff --git a/tests/python/sym/test_sym_deduce_bound.py b/tirx/tests/python/sym/test_sym_deduce_bound.py
similarity index 100%
rename from tests/python/sym/test_sym_deduce_bound.py
rename to tirx/tests/python/sym/test_sym_deduce_bound.py
diff --git a/tests/python/sym/test_sym_detect_clip_bound.py b/tirx/tests/python/sym/test_sym_detect_clip_bound.py
similarity index 100%
rename from tests/python/sym/test_sym_detect_clip_bound.py
rename to tirx/tests/python/sym/test_sym_detect_clip_bound.py
diff --git a/tests/python/sym/test_sym_detect_linear_equation.py b/tirx/tests/python/sym/test_sym_detect_linear_equation.py
similarity index 100%
rename from tests/python/sym/test_sym_detect_linear_equation.py
rename to tirx/tests/python/sym/test_sym_detect_linear_equation.py
diff --git a/tests/python/sym/test_sym_intset.py b/tirx/tests/python/sym/test_sym_intset.py
similarity index 100%
rename from tests/python/sym/test_sym_intset.py
rename to tirx/tests/python/sym/test_sym_intset.py
diff --git a/tests/python/sym/test_sym_iter_affine_map.py b/tirx/tests/python/sym/test_sym_iter_affine_map.py
similarity index 100%
rename from tests/python/sym/test_sym_iter_affine_map.py
rename to tirx/tests/python/sym/test_sym_iter_affine_map.py
diff --git a/tests/python/sym/test_sym_rewrite_simplify.py b/tirx/tests/python/sym/test_sym_rewrite_simplify.py
similarity index 100%
rename from tests/python/sym/test_sym_rewrite_simplify.py
rename to tirx/tests/python/sym/test_sym_rewrite_simplify.py
diff --git a/tests/python/sym/test_sym_simplify.py b/tirx/tests/python/sym/test_sym_simplify.py
similarity index 100%
rename from tests/python/sym/test_sym_simplify.py
rename to tirx/tests/python/sym/test_sym_simplify.py
diff --git a/tests/python/sym/test_sym_z3.py b/tirx/tests/python/sym/test_sym_z3.py
similarity index 100%
rename from tests/python/sym/test_sym_z3.py
rename to tirx/tests/python/sym/test_sym_z3.py
diff --git a/tests/python/target/test_llvm_features_info.py b/tirx/tests/python/target/test_llvm_features_info.py
similarity index 100%
rename from tests/python/target/test_llvm_features_info.py
rename to tirx/tests/python/target/test_llvm_features_info.py
diff --git a/tests/python/target/test_riscv_features.py b/tirx/tests/python/target/test_riscv_features.py
similarity index 100%
rename from tests/python/target/test_riscv_features.py
rename to tirx/tests/python/target/test_riscv_features.py
diff --git a/tests/python/target/test_target_parser_mprofile.py b/tirx/tests/python/target/test_target_parser_mprofile.py
similarity index 100%
rename from tests/python/target/test_target_parser_mprofile.py
rename to tirx/tests/python/target/test_target_parser_mprofile.py
diff --git a/tests/python/target/test_x86_features.py b/tirx/tests/python/target/test_x86_features.py
similarity index 100%
rename from tests/python/target/test_x86_features.py
rename to tirx/tests/python/target/test_x86_features.py
diff --git a/tests/python/testing/test_env.py b/tirx/tests/python/testing/test_env.py
similarity index 100%
rename from tests/python/testing/test_env.py
rename to tirx/tests/python/testing/test_env.py
diff --git a/tests/python/tirx/__init__.py b/tirx/tests/python/tirx/__init__.py
similarity index 100%
rename from tests/python/tirx/__init__.py
rename to tirx/tests/python/tirx/__init__.py
diff --git a/tests/python/tirx/analysis/test_tirx_analysis_check_well_defined.py b/tirx/tests/python/tirx/analysis/test_tirx_analysis_check_well_defined.py
similarity index 100%
rename from tests/python/tirx/analysis/test_tirx_analysis_check_well_defined.py
rename to tirx/tests/python/tirx/analysis/test_tirx_analysis_check_well_defined.py
diff --git a/tests/python/tirx/analysis/test_tirx_analysis_undefined_vars.py b/tirx/tests/python/tirx/analysis/test_tirx_analysis_undefined_vars.py
similarity index 100%
rename from tests/python/tirx/analysis/test_tirx_analysis_undefined_vars.py
rename to tirx/tests/python/tirx/analysis/test_tirx_analysis_undefined_vars.py
diff --git a/tests/python/tirx/base/test_tirx_host_func.py b/tirx/tests/python/tirx/base/test_tirx_host_func.py
similarity index 100%
rename from tests/python/tirx/base/test_tirx_host_func.py
rename to tirx/tests/python/tirx/base/test_tirx_host_func.py
diff --git a/tests/python/tirx/base/test_tirx_imm_values.py b/tirx/tests/python/tirx/base/test_tirx_imm_values.py
similarity index 100%
rename from tests/python/tirx/base/test_tirx_imm_values.py
rename to tirx/tests/python/tirx/base/test_tirx_imm_values.py
diff --git a/tests/python/tirx/base/test_tirx_index_map.py b/tirx/tests/python/tirx/base/test_tirx_index_map.py
similarity index 100%
rename from tests/python/tirx/base/test_tirx_index_map.py
rename to tirx/tests/python/tirx/base/test_tirx_index_map.py
diff --git a/tests/python/tirx/base/test_tirx_intrin.py b/tirx/tests/python/tirx/base/test_tirx_intrin.py
similarity index 100%
rename from tests/python/tirx/base/test_tirx_intrin.py
rename to tirx/tests/python/tirx/base/test_tirx_intrin.py
diff --git a/tests/python/tirx/base/test_tirx_nodes.py b/tirx/tests/python/tirx/base/test_tirx_nodes.py
similarity index 100%
rename from tests/python/tirx/base/test_tirx_nodes.py
rename to tirx/tests/python/tirx/base/test_tirx_nodes.py
diff --git a/tests/python/tirx/base/test_tirx_op_types.py b/tirx/tests/python/tirx/base/test_tirx_op_types.py
similarity index 100%
rename from tests/python/tirx/base/test_tirx_op_types.py
rename to tirx/tests/python/tirx/base/test_tirx_op_types.py
diff --git a/tests/python/tirx/base/test_tirx_ptx_cp_async.py b/tirx/tests/python/tirx/base/test_tirx_ptx_cp_async.py
similarity index 100%
rename from tests/python/tirx/base/test_tirx_ptx_cp_async.py
rename to tirx/tests/python/tirx/base/test_tirx_ptx_cp_async.py
diff --git a/tests/python/tirx/base/test_tirx_ptx_griddepcontrol.py b/tirx/tests/python/tirx/base/test_tirx_ptx_griddepcontrol.py
similarity index 100%
rename from tests/python/tirx/base/test_tirx_ptx_griddepcontrol.py
rename to tirx/tests/python/tirx/base/test_tirx_ptx_griddepcontrol.py
diff --git a/tests/python/tirx/base/test_tirx_ptx_mma_sp.py b/tirx/tests/python/tirx/base/test_tirx_ptx_mma_sp.py
similarity index 100%
rename from tests/python/tirx/base/test_tirx_ptx_mma_sp.py
rename to tirx/tests/python/tirx/base/test_tirx_ptx_mma_sp.py
diff --git a/tests/python/tirx/base/test_tirx_ptx_scalar_f32_math.py b/tirx/tests/python/tirx/base/test_tirx_ptx_scalar_f32_math.py
similarity index 100%
rename from tests/python/tirx/base/test_tirx_ptx_scalar_f32_math.py
rename to tirx/tests/python/tirx/base/test_tirx_ptx_scalar_f32_math.py
diff --git a/tests/python/tirx/base/test_tirx_scalable_datatype.py b/tirx/tests/python/tirx/base/test_tirx_scalable_datatype.py
similarity index 100%
rename from tests/python/tirx/base/test_tirx_scalable_datatype.py
rename to tirx/tests/python/tirx/base/test_tirx_scalable_datatype.py
diff --git a/tests/python/tirx/base/test_tirx_specialize.py b/tirx/tests/python/tirx/base/test_tirx_specialize.py
similarity index 100%
rename from tests/python/tirx/base/test_tirx_specialize.py
rename to tirx/tests/python/tirx/base/test_tirx_specialize.py
diff --git a/tests/python/tirx/base/test_tirx_tensor_type.py b/tirx/tests/python/tirx/base/test_tirx_tensor_type.py
similarity index 100%
rename from tests/python/tirx/base/test_tirx_tensor_type.py
rename to tirx/tests/python/tirx/base/test_tirx_tensor_type.py
diff --git a/tests/python/tirx/base/test_tirx_texture_scope.py b/tirx/tests/python/tirx/base/test_tirx_texture_scope.py
similarity index 100%
rename from tests/python/tirx/base/test_tirx_texture_scope.py
rename to tirx/tests/python/tirx/base/test_tirx_texture_scope.py
diff --git a/tests/python/tirx/codegen/conftest.py b/tirx/tests/python/tirx/codegen/conftest.py
similarity index 100%
rename from tests/python/tirx/codegen/conftest.py
rename to tirx/tests/python/tirx/codegen/conftest.py
diff --git a/tests/python/tirx/codegen/test_codegen_ampere.py b/tirx/tests/python/tirx/codegen/test_codegen_ampere.py
similarity index 100%
rename from tests/python/tirx/codegen/test_codegen_ampere.py
rename to tirx/tests/python/tirx/codegen/test_codegen_ampere.py
diff --git a/tests/python/tirx/codegen/test_codegen_blackwell.py b/tirx/tests/python/tirx/codegen/test_codegen_blackwell.py
similarity index 100%
rename from tests/python/tirx/codegen/test_codegen_blackwell.py
rename to tirx/tests/python/tirx/codegen/test_codegen_blackwell.py
diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py b/tirx/tests/python/tirx/codegen/test_codegen_cuda.py
similarity index 100%
rename from tests/python/tirx/codegen/test_codegen_cuda.py
rename to tirx/tests/python/tirx/codegen/test_codegen_cuda.py
diff --git a/tests/python/tirx/codegen/test_codegen_dsmem.py b/tirx/tests/python/tirx/codegen/test_codegen_dsmem.py
similarity index 100%
rename from tests/python/tirx/codegen/test_codegen_dsmem.py
rename to tirx/tests/python/tirx/codegen/test_codegen_dsmem.py
diff --git a/tests/python/tirx/codegen/test_codegen_hopper.py b/tirx/tests/python/tirx/codegen/test_codegen_hopper.py
similarity index 100%
rename from tests/python/tirx/codegen/test_codegen_hopper.py
rename to tirx/tests/python/tirx/codegen/test_codegen_hopper.py
diff --git a/tests/python/tirx/codegen/test_codegen_nki.py b/tirx/tests/python/tirx/codegen/test_codegen_nki.py
similarity index 100%
rename from tests/python/tirx/codegen/test_codegen_nki.py
rename to tirx/tests/python/tirx/codegen/test_codegen_nki.py
diff --git a/tests/python/tirx/codegen/test_cuda_cta_reduce.py b/tirx/tests/python/tirx/codegen/test_cuda_cta_reduce.py
similarity index 100%
rename from tests/python/tirx/codegen/test_cuda_cta_reduce.py
rename to tirx/tests/python/tirx/codegen/test_cuda_cta_reduce.py
diff --git a/tests/python/tirx/codegen/test_cuda_wait_until.py b/tirx/tests/python/tirx/codegen/test_cuda_wait_until.py
similarity index 100%
rename from tests/python/tirx/codegen/test_cuda_wait_until.py
rename to tirx/tests/python/tirx/codegen/test_cuda_wait_until.py
diff --git a/tests/python/tirx/codegen/test_cuda_warp_reduce.py b/tirx/tests/python/tirx/codegen/test_cuda_warp_reduce.py
similarity index 100%
rename from tests/python/tirx/codegen/test_cuda_warp_reduce.py
rename to tirx/tests/python/tirx/codegen/test_cuda_warp_reduce.py
diff --git a/tests/python/tirx/codegen/test_ptx_addr.py b/tirx/tests/python/tirx/codegen/test_ptx_addr.py
similarity index 100%
rename from tests/python/tirx/codegen/test_ptx_addr.py
rename to tirx/tests/python/tirx/codegen/test_ptx_addr.py
diff --git a/tests/python/tirx/codegen/test_ptx_cvt.py b/tirx/tests/python/tirx/codegen/test_ptx_cvt.py
similarity index 100%
rename from tests/python/tirx/codegen/test_ptx_cvt.py
rename to tirx/tests/python/tirx/codegen/test_ptx_cvt.py
diff --git a/tests/python/tirx/codegen/test_ptx_dialect.py b/tirx/tests/python/tirx/codegen/test_ptx_dialect.py
similarity index 100%
rename from tests/python/tirx/codegen/test_ptx_dialect.py
rename to tirx/tests/python/tirx/codegen/test_ptx_dialect.py
diff --git a/tests/python/tirx/codegen/test_ptx_ld_st_ops.py b/tirx/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
similarity index 100%
rename from tests/python/tirx/codegen/test_ptx_ld_st_ops.py
rename to tirx/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
diff --git a/tests/python/tirx/codegen/test_tcgen05_descriptor_encoders.py b/tirx/tests/python/tirx/codegen/test_tcgen05_descriptor_encoders.py
similarity index 100%
rename from tests/python/tirx/codegen/test_tcgen05_descriptor_encoders.py
rename to tirx/tests/python/tirx/codegen/test_tcgen05_descriptor_encoders.py
diff --git a/tests/python/tirx/conftest.py b/tirx/tests/python/tirx/conftest.py
similarity index 100%
rename from tests/python/tirx/conftest.py
rename to tirx/tests/python/tirx/conftest.py
diff --git a/tests/python/tirx/iket/iket_profile_workload.py b/tirx/tests/python/tirx/iket/iket_profile_workload.py
similarity index 100%
rename from tests/python/tirx/iket/iket_profile_workload.py
rename to tirx/tests/python/tirx/iket/iket_profile_workload.py
diff --git a/tests/python/tirx/iket/oracle/generate_iket_official_oracle.py b/tirx/tests/python/tirx/iket/oracle/generate_iket_official_oracle.py
similarity index 100%
rename from tests/python/tirx/iket/oracle/generate_iket_official_oracle.py
rename to tirx/tests/python/tirx/iket/oracle/generate_iket_official_oracle.py
diff --git a/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_0_oracle.json b/tirx/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_0_oracle.json
similarity index 100%
rename from tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_0_oracle.json
rename to tirx/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_0_oracle.json
diff --git a/tests/python/tirx/iket/test_iket_orchestration.py b/tirx/tests/python/tirx/iket/test_iket_orchestration.py
similarity index 100%
rename from tests/python/tirx/iket/test_iket_orchestration.py
rename to tirx/tests/python/tirx/iket/test_iket_orchestration.py
diff --git a/tests/python/tirx/iket/test_iket_profiler.py b/tirx/tests/python/tirx/iket/test_iket_profiler.py
similarity index 100%
rename from tests/python/tirx/iket/test_iket_profiler.py
rename to tirx/tests/python/tirx/iket/test_iket_profiler.py
diff --git a/tests/python/tirx/iket/verify_iket_official_patch.py b/tirx/tests/python/tirx/iket/verify_iket_official_patch.py
similarity index 100%
rename from tests/python/tirx/iket/verify_iket_official_patch.py
rename to tirx/tests/python/tirx/iket/verify_iket_official_patch.py
diff --git a/tests/python/tirx/layout_fixtures.py b/tirx/tests/python/tirx/layout_fixtures.py
similarity index 100%
rename from tests/python/tirx/layout_fixtures.py
rename to tirx/tests/python/tirx/layout_fixtures.py
diff --git a/tests/python/tirx/script/test_tirx_script_basic_usage.py b/tirx/tests/python/tirx/script/test_tirx_script_basic_usage.py
similarity index 100%
rename from tests/python/tirx/script/test_tirx_script_basic_usage.py
rename to tirx/tests/python/tirx/script/test_tirx_script_basic_usage.py
diff --git a/tests/python/tirx/script/test_tirx_script_dynamic_shape.py b/tirx/tests/python/tirx/script/test_tirx_script_dynamic_shape.py
similarity index 100%
rename from tests/python/tirx/script/test_tirx_script_dynamic_shape.py
rename to tirx/tests/python/tirx/script/test_tirx_script_dynamic_shape.py
diff --git a/tests/python/tirx/script/test_tirx_script_error_handling.py b/tirx/tests/python/tirx/script/test_tirx_script_error_handling.py
similarity index 100%
rename from tests/python/tirx/script/test_tirx_script_error_handling.py
rename to tirx/tests/python/tirx/script/test_tirx_script_error_handling.py
diff --git a/tests/python/tirx/script/test_tirx_script_gpu.py b/tirx/tests/python/tirx/script/test_tirx_script_gpu.py
similarity index 100%
rename from tests/python/tirx/script/test_tirx_script_gpu.py
rename to tirx/tests/python/tirx/script/test_tirx_script_gpu.py
diff --git a/tests/python/tirx/script/test_tirx_script_jit_specialization.py b/tirx/tests/python/tirx/script/test_tirx_script_jit_specialization.py
similarity index 100%
rename from tests/python/tirx/script/test_tirx_script_jit_specialization.py
rename to tirx/tests/python/tirx/script/test_tirx_script_jit_specialization.py
diff --git a/tests/python/tirx/script/test_tirx_script_meta_programming.py b/tirx/tests/python/tirx/script/test_tirx_script_meta_programming.py
similarity index 100%
rename from tests/python/tirx/script/test_tirx_script_meta_programming.py
rename to tirx/tests/python/tirx/script/test_tirx_script_meta_programming.py
diff --git a/tests/python/tirx/script/test_tirx_script_printer_highlight.py b/tirx/tests/python/tirx/script/test_tirx_script_printer_highlight.py
similarity index 100%
rename from tests/python/tirx/script/test_tirx_script_printer_highlight.py
rename to tirx/tests/python/tirx/script/test_tirx_script_printer_highlight.py
diff --git a/tests/python/tirx/script/test_tirx_script_printer_invisible_path_info.py b/tirx/tests/python/tirx/script/test_tirx_script_printer_invisible_path_info.py
similarity index 100%
rename from tests/python/tirx/script/test_tirx_script_printer_invisible_path_info.py
rename to tirx/tests/python/tirx/script/test_tirx_script_printer_invisible_path_info.py
diff --git a/tests/python/tirx/script/test_tirx_script_printer_structural_equal.py b/tirx/tests/python/tirx/script/test_tirx_script_printer_structural_equal.py
similarity index 100%
rename from tests/python/tirx/script/test_tirx_script_printer_structural_equal.py
rename to tirx/tests/python/tirx/script/test_tirx_script_printer_structural_equal.py
diff --git a/tests/python/tirx/script/test_tirx_script_printer_underlining.py b/tirx/tests/python/tirx/script/test_tirx_script_printer_underlining.py
similarity index 100%
rename from tests/python/tirx/script/test_tirx_script_printer_underlining.py
rename to tirx/tests/python/tirx/script/test_tirx_script_printer_underlining.py
diff --git a/tests/python/tirx/script/test_tirx_script_source_locations.py b/tirx/tests/python/tirx/script/test_tirx_script_source_locations.py
similarity index 100%
rename from tests/python/tirx/script/test_tirx_script_source_locations.py
rename to tirx/tests/python/tirx/script/test_tirx_script_source_locations.py
diff --git a/tests/python/tirx/script/test_tirx_script_tensors.py b/tirx/tests/python/tirx/script/test_tirx_script_tensors.py
similarity index 100%
rename from tests/python/tirx/script/test_tirx_script_tensors.py
rename to tirx/tests/python/tirx/script/test_tirx_script_tensors.py
diff --git a/tests/python/tirx/script/test_tirx_script_tracing.py b/tirx/tests/python/tirx/script/test_tirx_script_tracing.py
similarity index 100%
rename from tests/python/tirx/script/test_tirx_script_tracing.py
rename to tirx/tests/python/tirx/script/test_tirx_script_tracing.py
diff --git a/tests/python/tirx/test_alloc_pool.py b/tirx/tests/python/tirx/test_alloc_pool.py
similarity index 100%
rename from tests/python/tirx/test_alloc_pool.py
rename to tirx/tests/python/tirx/test_alloc_pool.py
diff --git a/tests/python/tirx/test_bench_utils.py b/tirx/tests/python/tirx/test_bench_utils.py
similarity index 100%
rename from tests/python/tirx/test_bench_utils.py
rename to tirx/tests/python/tirx/test_bench_utils.py
diff --git a/tests/python/tirx/test_compilation_pipeline.py b/tirx/tests/python/tirx/test_compilation_pipeline.py
similarity index 100%
rename from tests/python/tirx/test_compilation_pipeline.py
rename to tirx/tests/python/tirx/test_compilation_pipeline.py
diff --git a/tests/python/tirx/test_control_flow.py b/tirx/tests/python/tirx/test_control_flow.py
similarity index 100%
rename from tests/python/tirx/test_control_flow.py
rename to tirx/tests/python/tirx/test_control_flow.py
diff --git a/tests/python/tirx/test_cuda_arch_marker.py b/tirx/tests/python/tirx/test_cuda_arch_marker.py
similarity index 100%
rename from tests/python/tirx/test_cuda_arch_marker.py
rename to tirx/tests/python/tirx/test_cuda_arch_marker.py
diff --git a/tests/python/tirx/test_cuda_backend_config.py b/tirx/tests/python/tirx/test_cuda_backend_config.py
similarity index 100%
rename from tests/python/tirx/test_cuda_backend_config.py
rename to tirx/tests/python/tirx/test_cuda_backend_config.py
diff --git a/tests/python/tirx/test_cuda_launch.py b/tirx/tests/python/tirx/test_cuda_launch.py
similarity index 100%
rename from tests/python/tirx/test_cuda_launch.py
rename to tirx/tests/python/tirx/test_cuda_launch.py
diff --git a/tests/python/tirx/test_exec_context.py b/tirx/tests/python/tirx/test_exec_context.py
similarity index 100%
rename from tests/python/tirx/test_exec_context.py
rename to tirx/tests/python/tirx/test_exec_context.py
diff --git a/tests/python/tirx/test_inline.py b/tirx/tests/python/tirx/test_inline.py
similarity index 100%
rename from tests/python/tirx/test_inline.py
rename to tirx/tests/python/tirx/test_inline.py
diff --git a/tests/python/tirx/test_jit.py b/tirx/tests/python/tirx/test_jit.py
similarity index 100%
rename from tests/python/tirx/test_jit.py
rename to tirx/tests/python/tirx/test_jit.py
diff --git a/tests/python/tirx/test_layout.py b/tirx/tests/python/tirx/test_layout.py
similarity index 100%
rename from tests/python/tirx/test_layout.py
rename to tirx/tests/python/tirx/test_layout.py
diff --git a/tests/python/tirx/test_op.py b/tirx/tests/python/tirx/test_op.py
similarity index 100%
rename from tests/python/tirx/test_op.py
rename to tirx/tests/python/tirx/test_op.py
diff --git a/tests/python/tirx/test_op_namespace_cleanup.py b/tirx/tests/python/tirx/test_op_namespace_cleanup.py
similarity index 100%
rename from tests/python/tirx/test_op_namespace_cleanup.py
rename to tirx/tests/python/tirx/test_op_namespace_cleanup.py
diff --git a/tests/python/tirx/test_roundtrip_namespaces.py b/tirx/tests/python/tirx/test_roundtrip_namespaces.py
similarity index 100%
rename from tests/python/tirx/test_roundtrip_namespaces.py
rename to tirx/tests/python/tirx/test_roundtrip_namespaces.py
diff --git a/tests/python/tirx/test_tensor_data_reinfer_type.py b/tirx/tests/python/tirx/test_tensor_data_reinfer_type.py
similarity index 100%
rename from tests/python/tirx/test_tensor_data_reinfer_type.py
rename to tirx/tests/python/tirx/test_tensor_data_reinfer_type.py
diff --git a/tests/python/tirx/test_thread_scope.py b/tirx/tests/python/tirx/test_thread_scope.py
similarity index 100%
rename from tests/python/tirx/test_thread_scope.py
rename to tirx/tests/python/tirx/test_thread_scope.py
diff --git a/tests/python/tirx/test_tile_namespace.py b/tirx/tests/python/tirx/test_tile_namespace.py
similarity index 100%
rename from tests/python/tirx/test_tile_namespace.py
rename to tirx/tests/python/tirx/test_tile_namespace.py
diff --git a/tests/python/tirx/test_tirx_axis_enum.py b/tirx/tests/python/tirx/test_tirx_axis_enum.py
similarity index 100%
rename from tests/python/tirx/test_tirx_axis_enum.py
rename to tirx/tests/python/tirx/test_tirx_axis_enum.py
diff --git a/tests/python/tirx/test_verifier.py b/tirx/tests/python/tirx/test_verifier.py
similarity index 100%
rename from tests/python/tirx/test_verifier.py
rename to tirx/tests/python/tirx/test_verifier.py
diff --git a/tests/python/tirx/tile/cuda/conftest.py b/tirx/tests/python/tirx/tile/cuda/conftest.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/conftest.py
rename to tirx/tests/python/tirx/tile/cuda/conftest.py
diff --git a/tests/python/tirx/tile/cuda/test_add_sub_mul_div.py b/tirx/tests/python/tirx/tile/cuda/test_add_sub_mul_div.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_add_sub_mul_div.py
rename to tirx/tests/python/tirx/tile/cuda/test_add_sub_mul_div.py
diff --git a/tests/python/tirx/tile/cuda/test_cp_async.py b/tirx/tests/python/tirx/tile/cuda/test_cp_async.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_cp_async.py
rename to tirx/tests/python/tirx/tile/cuda/test_cp_async.py
diff --git a/tests/python/tirx/tile/cuda/test_cp_async_bulk.py b/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_cp_async_bulk.py
rename to tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py
diff --git a/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py b/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py
rename to tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py
diff --git a/tests/python/tirx/tile/cuda/test_cvt.py b/tirx/tests/python/tirx/tile/cuda/test_cvt.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_cvt.py
rename to tirx/tests/python/tirx/tile/cuda/test_cvt.py
diff --git a/tests/python/tirx/tile/cuda/test_ewise.py b/tirx/tests/python/tirx/tile/cuda/test_ewise.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_ewise.py
rename to tirx/tests/python/tirx/tile/cuda/test_ewise.py
diff --git a/tests/python/tirx/tile/cuda/test_explicit_patterns.py b/tirx/tests/python/tirx/tile/cuda/test_explicit_patterns.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_explicit_patterns.py
rename to tirx/tests/python/tirx/tile/cuda/test_explicit_patterns.py
diff --git a/tests/python/tirx/tile/cuda/test_fma.py b/tirx/tests/python/tirx/tile/cuda/test_fma.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_fma.py
rename to tirx/tests/python/tirx/tile/cuda/test_fma.py
diff --git a/tests/python/tirx/tile/cuda/test_ld_st_cache_policy.py b/tirx/tests/python/tirx/tile/cuda/test_ld_st_cache_policy.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_ld_st_cache_policy.py
rename to tirx/tests/python/tirx/tile/cuda/test_ld_st_cache_policy.py
diff --git a/tests/python/tirx/tile/cuda/test_ld_st_gmem_smem.py b/tirx/tests/python/tirx/tile/cuda/test_ld_st_gmem_smem.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_ld_st_gmem_smem.py
rename to tirx/tests/python/tirx/tile/cuda/test_ld_st_gmem_smem.py
diff --git a/tests/python/tirx/tile/cuda/test_ld_st_registers.py b/tirx/tests/python/tirx/tile/cuda/test_ld_st_registers.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_ld_st_registers.py
rename to tirx/tests/python/tirx/tile/cuda/test_ld_st_registers.py
diff --git a/tests/python/tirx/tile/cuda/test_ldmatrix_stmatrix.py b/tirx/tests/python/tirx/tile/cuda/test_ldmatrix_stmatrix.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_ldmatrix_stmatrix.py
rename to tirx/tests/python/tirx/tile/cuda/test_ldmatrix_stmatrix.py
diff --git a/tests/python/tirx/tile/cuda/test_mma_sync.py b/tirx/tests/python/tirx/tile/cuda/test_mma_sync.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_mma_sync.py
rename to tirx/tests/python/tirx/tile/cuda/test_mma_sync.py
diff --git a/tests/python/tirx/tile/cuda/test_mov.py b/tirx/tests/python/tirx/tile/cuda/test_mov.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_mov.py
rename to tirx/tests/python/tirx/tile/cuda/test_mov.py
diff --git a/tests/python/tirx/tile/cuda/test_pattern_groups.py b/tirx/tests/python/tirx/tile/cuda/test_pattern_groups.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_pattern_groups.py
rename to tirx/tests/python/tirx/tile/cuda/test_pattern_groups.py
diff --git a/tests/python/tirx/tile/cuda/test_pattern_math.py b/tirx/tests/python/tirx/tile/cuda/test_pattern_math.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_pattern_math.py
rename to tirx/tests/python/tirx/tile/cuda/test_pattern_math.py
diff --git a/tests/python/tirx/tile/cuda/test_pattern_mma.py b/tirx/tests/python/tirx/tile/cuda/test_pattern_mma.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_pattern_mma.py
rename to tirx/tests/python/tirx/tile/cuda/test_pattern_mma.py
diff --git a/tests/python/tirx/tile/cuda/test_pattern_transfers.py b/tirx/tests/python/tirx/tile/cuda/test_pattern_transfers.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_pattern_transfers.py
rename to tirx/tests/python/tirx/tile/cuda/test_pattern_transfers.py
diff --git a/tests/python/tirx/tile/cuda/test_tcgen05_cp.py b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_cp.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_tcgen05_cp.py
rename to tirx/tests/python/tirx/tile/cuda/test_tcgen05_cp.py
diff --git a/tests/python/tirx/tile/cuda/test_tcgen05_ld_st.py b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_ld_st.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_tcgen05_ld_st.py
rename to tirx/tests/python/tirx/tile/cuda/test_tcgen05_ld_st.py
diff --git a/tests/python/tirx/tile/cuda/test_tcgen05_mma.py b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_tcgen05_mma.py
rename to tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma.py
diff --git a/tests/python/tirx/tile/cuda/test_tcgen05_mma_codegen.py b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma_codegen.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_tcgen05_mma_codegen.py
rename to tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma_codegen.py
diff --git a/tests/python/tirx/tile/cuda/test_tcgen05_mma_ws_split_n.py b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma_ws_split_n.py
similarity index 100%
rename from tests/python/tirx/tile/cuda/test_tcgen05_mma_ws_split_n.py
rename to tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma_ws_split_n.py
diff --git a/tests/python/tirx/tile/test_dispatcher.py b/tirx/tests/python/tirx/tile/test_dispatcher.py
similarity index 100%
rename from tests/python/tirx/tile/test_dispatcher.py
rename to tirx/tests/python/tirx/tile/test_dispatcher.py
diff --git a/tests/python/tirx/tile/trn/test_binary_trn.py b/tirx/tests/python/tirx/tile/trn/test_binary_trn.py
similarity index 100%
rename from tests/python/tirx/tile/trn/test_binary_trn.py
rename to tirx/tests/python/tirx/tile/trn/test_binary_trn.py
diff --git a/tests/python/tirx/tile/trn/test_compose_op_trn.py b/tirx/tests/python/tirx/tile/trn/test_compose_op_trn.py
similarity index 100%
rename from tests/python/tirx/tile/trn/test_compose_op_trn.py
rename to tirx/tests/python/tirx/tile/trn/test_compose_op_trn.py
diff --git a/tests/python/tirx/tile/trn/test_copy_trn.py b/tirx/tests/python/tirx/tile/trn/test_copy_trn.py
similarity index 100%
rename from tests/python/tirx/tile/trn/test_copy_trn.py
rename to tirx/tests/python/tirx/tile/trn/test_copy_trn.py
diff --git a/tests/python/tirx/tile/trn/test_gemm_trn.py b/tirx/tests/python/tirx/tile/trn/test_gemm_trn.py
similarity index 100%
rename from tests/python/tirx/tile/trn/test_gemm_trn.py
rename to tirx/tests/python/tirx/tile/trn/test_gemm_trn.py
diff --git a/tests/python/tirx/tile/trn/test_private_alloc_trn.py b/tirx/tests/python/tirx/tile/trn/test_private_alloc_trn.py
similarity index 100%
rename from tests/python/tirx/tile/trn/test_private_alloc_trn.py
rename to tirx/tests/python/tirx/tile/trn/test_private_alloc_trn.py
diff --git a/tests/python/tirx/tile/trn/test_reduction_trn.py b/tirx/tests/python/tirx/tile/trn/test_reduction_trn.py
similarity index 100%
rename from tests/python/tirx/tile/trn/test_reduction_trn.py
rename to tirx/tests/python/tirx/tile/trn/test_reduction_trn.py
diff --git a/tests/python/tirx/tile/trn/test_select_trn.py b/tirx/tests/python/tirx/tile/trn/test_select_trn.py
similarity index 100%
rename from tests/python/tirx/tile/trn/test_select_trn.py
rename to tirx/tests/python/tirx/tile/trn/test_select_trn.py
diff --git a/tests/python/tirx/tile/trn/test_unary_trn.py b/tirx/tests/python/tirx/tile/trn/test_unary_trn.py
similarity index 100%
rename from tests/python/tirx/tile/trn/test_unary_trn.py
rename to tirx/tests/python/tirx/tile/trn/test_unary_trn.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_bf16_legalize.py b/tirx/tests/python/tirx/transform/test_tirx_transform_bf16_legalize.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_bf16_legalize.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_bf16_legalize.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_common_subexpr_elim.py b/tirx/tests/python/tirx/transform/test_tirx_transform_common_subexpr_elim.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_common_subexpr_elim.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_common_subexpr_elim.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_convert_ssa.py b/tirx/tests/python/tirx/transform/test_tirx_transform_convert_ssa.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_convert_ssa.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_convert_ssa.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_decorate_device_scope.py b/tirx/tests/python/tirx/transform/test_tirx_transform_decorate_device_scope.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_decorate_device_scope.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_decorate_device_scope.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_flatten_tensor.py b/tirx/tests/python/tirx/transform/test_tirx_transform_flatten_tensor.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_flatten_tensor.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_flatten_tensor.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_force_narrow_index_to_i32.py b/tirx/tests/python/tirx/transform/test_tirx_transform_force_narrow_index_to_i32.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_force_narrow_index_to_i32.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_force_narrow_index_to_i32.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_fp8_legalize.py b/tirx/tests/python/tirx/transform/test_tirx_transform_fp8_legalize.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_fp8_legalize.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_fp8_legalize.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_function_pass.py b/tirx/tests/python/tirx/transform/test_tirx_transform_function_pass.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_function_pass.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_function_pass.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_helpers.py b/tirx/tests/python/tirx/transform/test_tirx_transform_helpers.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_helpers.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_helpers.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_hoist_if.py b/tirx/tests/python/tirx/transform/test_tirx_transform_hoist_if.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_hoist_if.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_hoist_if.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_lower_intrin.py b/tirx/tests/python/tirx/transform/test_tirx_transform_lower_intrin.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_lower_intrin.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_lower_intrin.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_lower_tirx.py b/tirx/tests/python/tirx/transform/test_tirx_transform_lower_tirx.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_lower_tirx.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_lower_tirx.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_lower_tvm_builtin.py b/tirx/tests/python/tirx/transform/test_tirx_transform_lower_tvm_builtin.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_lower_tvm_builtin.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_lower_tvm_builtin.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_make_packed_api.py b/tirx/tests/python/tirx/transform/test_tirx_transform_make_packed_api.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_make_packed_api.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_make_packed_api.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_naive_allocator.py b/tirx/tests/python/tirx/transform/test_tirx_transform_naive_allocator.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_naive_allocator.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_naive_allocator.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_narrow_datatype.py b/tirx/tests/python/tirx/transform/test_tirx_transform_narrow_datatype.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_narrow_datatype.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_narrow_datatype.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_pointer_value_type_rewrite.py b/tirx/tests/python/tirx/transform/test_tirx_transform_pointer_value_type_rewrite.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_pointer_value_type_rewrite.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_pointer_value_type_rewrite.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_remove_no_op.py b/tirx/tests/python/tirx/transform/test_tirx_transform_remove_no_op.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_remove_no_op.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_remove_no_op.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_renormalize_split_pattern.py b/tirx/tests/python/tirx/transform/test_tirx_transform_renormalize_split_pattern.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_renormalize_split_pattern.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_renormalize_split_pattern.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_simplify.py b/tirx/tests/python/tirx/transform/test_tirx_transform_simplify.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_simplify.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_simplify.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_split_host_device.py b/tirx/tests/python/tirx/transform/test_tirx_transform_split_host_device.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_split_host_device.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_split_host_device.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_storage_rewrite.py b/tirx/tests/python/tirx/transform/test_tirx_transform_storage_rewrite.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_storage_rewrite.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_storage_rewrite.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_unroll_loop.py b/tirx/tests/python/tirx/transform/test_tirx_transform_unroll_loop.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_unroll_loop.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_unroll_loop.py
diff --git a/tests/python/tirx/transform/test_tirx_transform_vectorize.py b/tirx/tests/python/tirx/transform/test_tirx_transform_vectorize.py
similarity index 100%
rename from tests/python/tirx/transform/test_tirx_transform_vectorize.py
rename to tirx/tests/python/tirx/transform/test_tirx_transform_vectorize.py
diff --git a/tests/python/tirx/utils.py b/tirx/tests/python/tirx/utils.py
similarity index 100%
rename from tests/python/tirx/utils.py
rename to tirx/tests/python/tirx/utils.py
(tvm) 26/33: [FIX][Runtime] Release the GIL while connecting Disco socket sessions
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 33920a5adc3224d0d27b29329c8172f216305050
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 07:42:22 2026 +0000
[FIX][Runtime] Release the GIL while connecting Disco socket sessions
---
python/tvm/runtime/disco/session.py | 13 ++++++-------
1 file changed, 6 insertions(+), 7 deletions(-)
diff --git a/python/tvm/runtime/disco/session.py b/python/tvm/runtime/disco/session.py
index 6188600ca7..a658ebd165 100644
--- a/python/tvm/runtime/disco/session.py
+++ b/python/tvm/runtime/disco/session.py
@@ -615,13 +615,12 @@ class SocketSession(Session):
host: str,
port: int,
) -> None:
- self.__init_handle_by_constructor__(
- _ffi_api.SocketSession, # type: ignore # pylint: disable=no-member
- num_nodes,
- num_workers_per_node,
- num_groups,
- host,
- port,
+ # Connecting remote nodes can block. A regular FFI call releases the GIL
+ # so another Python thread can launch them; constructor calls retain it.
+ self.__move_handle_from__(
+ _ffi_api.SocketSession( # type: ignore # pylint: disable=no-member
+ num_nodes, num_workers_per_node, num_groups, host, port
+ )
)
(tvm) 03/33: [REFACTOR][Packaging] Package TIRx with an independent version and dependency boundary
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 3a7bc102bfb65bc072c48a9149a451a09d8495ec
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 01:20:54 2026 +0000
[REFACTOR][Packaging] Package TIRx with an independent version and dependency boundary
---
LICENSE | 20 ++---
pyproject.toml | 79 ++++++++-----------
LICENSE => tirx/LICENSE | 23 +-----
tirx/NOTICE | 5 ++
pyproject.toml => tirx/pyproject.toml | 141 +++++++---------------------------
5 files changed, 73 insertions(+), 195 deletions(-)
diff --git a/LICENSE b/LICENSE
index 94abb3512b..627aef16ea 100644
--- a/LICENSE
+++ b/LICENSE
@@ -201,10 +201,8 @@
limitations under the License.
------------------------------------------------------------------------------------
-This product bundles various third-party components under other open source licenses.
-This section summarizes those components and their licenses. See licenses/
-and 3rdparty/tvm-ffi/licenses/ for text of these licenses.
-
+This product bundles the following third-party components. See licenses/ and the
+respective component directories for their licenses.
Apache Software Foundation License 2.0
--------------------------------------
@@ -212,18 +210,14 @@ Apache Software Foundation License 2.0
3rdparty/OpenCL-Headers
3rdparty/cutlass_fpA_intB_gemm
3rdparty/tensorrt_llm
-3rdparty/tvm-ffi/3rdparty/dlpack
-
-
-MIT License
------------
-
-3rdparty/compiler-rt/builtin_fp16.h
-
BSD 3-Clause "New" or "Revised" License
---------------------------------------
3rdparty/cutlass
3rdparty/libflash_attn
-3rdparty/tvm-ffi/3rdparty/libbacktrace
+
+MIT License
+-----------
+
+src/runtime/extra/contrib/vllm
diff --git a/pyproject.toml b/pyproject.toml
index 5c7512032f..02badac363 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,11 +16,11 @@
# under the License.
[build-system]
-# mlc-z3-static ships the PIC static libz3 + headers consumed by USE_Z3=ON.
+# The parent links against the separately installed core package.
requires = [
"scikit-build-core>=0.11",
"setuptools-scm>=8",
- "mlc-z3-static>=4.16.0",
+ "apache-tvm-tirx==0.1.0a0",
]
build-backend = "scikit_build_core.build"
@@ -49,13 +49,7 @@ classifiers = [
"Topic :: Scientific/Engineering :: Artificial Intelligence",
"Topic :: Software Development :: Libraries :: Python Modules",
]
-dependencies = [
- "apache-tvm-ffi>=0.1.14.post1",
- "apache-tvm-ffi-orcjit==0.1.2",
- "ml_dtypes",
- "numpy",
- "typing_extensions",
-]
+dependencies = ["apache-tvm-tirx==0.1.0a0"]
[project.optional-dependencies]
torch = ["torch"]
@@ -104,51 +98,41 @@ metadata.version.provider = "scikit_build_core.metadata.setuptools_scm"
# Source distribution configuration
sdist.include = [
- # Build files
"/CMakeLists.txt",
"/pyproject.toml",
- "/cmake/**/*",
-
- # Source code
- "/src/**/*.cc",
- "/src/**/*.h",
- "/include/**/*.h",
-
- # Python source
- "/python/tvm/**/*.py",
- "/python/tvm/**/*.pyi",
-
- # Documentation and metadata
- "/docs/**/*",
+ "/cmake/**",
+ "/src/**",
+ "/include/**",
+ "/python/**",
+ "/licenses/**",
"/LICENSE",
- "/README.md",
"/NOTICE",
-
- # Tests
- "/tests/**/*",
-]
-
-sdist.exclude = [
- "**/.git",
- "**/.github",
- "**/__pycache__",
- "**/*.pyc",
- "build",
- "dist",
- "**/3rdparty/*/docs",
- "**/3rdparty/*/media",
- "**/3rdparty/*/examples",
- "**/3rdparty/*/test",
+ "/tests/**",
+ "/README.md",
+ "/apps/**",
+ "/web/**",
+ "/3rdparty/OpenCL-Headers/**",
+ "/3rdparty/cutlass/**",
+ "/3rdparty/cutlass_fpA_intB_gemm/**",
+ "/3rdparty/libflash_attn/**",
+ "/3rdparty/tensorrt_llm/**",
+ "!**/.git",
+ "!**/.git/**",
+ "!**/__pycache__/**",
+ "!**/*.pyc",
+ "!**/.pytest_cache/**",
+ "!**/build/**",
+ "!**/node_modules/**",
]
+sdist.exclude = ["**"]
# Logging
logging.level = "INFO"
[tool.scikit-build.cmake.define]
TVM_BUILD_PYTHON_MODULE = "ON"
+TVM_USE_INSTALLED_TIRX = "ON"
USE_CUDA = "OFF"
-# Statically link Z3 from the mlc-z3-static build dependency by default.
-USE_Z3 = "AUTO"
BUILD_TESTING = "OFF"
[tool.setuptools_scm]
@@ -176,6 +160,8 @@ markers = [
[tool.ruff]
include = [
"python/**/*.py",
+ "tirx/python/**/*.py",
+ "tirx/tests/**/*.py",
"tests/**/*.py",
"docs/**/*.py",
"apps/**/*.py",
@@ -217,9 +203,10 @@ unfixable = []
"python/tvm/relax/transform/legalize_ops/nn.py" = ["E501"]
# TIRx parser and lowering tests intentionally retain unused IR bindings.
"tests/python/tirx/**/*.py" = ["F841"]
+"tirx/tests/python/tirx/**/*.py" = ["F841"]
[tool.ruff.lint.isort]
-known-first-party = ["tvm"]
+known-first-party = ["tvm", "tvm_tirx"]
[tool.ruff.format]
quote-style = "double"
@@ -237,7 +224,7 @@ test-requires = ["pytest", "numpy"]
test-command = "pytest -vvs {project}/tests/python/all-platform-minimal-test"
[tool.cibuildwheel.linux]
-repair-wheel-command = "auditwheel repair --exclude libtvm_ffi.so --exclude libtvm_runtime_cuda.so --exclude 'libcuda.so.*' --exclude 'libcudart.so.*' --exclude 'libnvrtc.so.*' --exclude 'libnvrtc-builtins.so.*' -w {dest_dir} {wheel}"
+repair-wheel-command = "auditwheel repair --exclude libtvm_tirx.so --exclude libtvm_runtime.so --exclude libtvm_ffi.so --exclude libtvm_runtime_cuda.so --exclude 'libcuda.so.*' --exclude 'libcudart.so.*' --exclude 'libnvrtc.so.*' --exclude 'libnvrtc-builtins.so.*' -w {dest_dir} {wheel}"
[tool.cibuildwheel.macos]
# delocate breaks libtvm_runtime.dylib's ad-hoc signature (it links the excluded
@@ -245,7 +232,7 @@ repair-wheel-command = "auditwheel repair --exclude libtvm_ffi.so --exclude libt
repair-wheel-command = '''
set -euo pipefail
DELOCATE_TMP_DIR="$(mktemp -d)"
-delocate-wheel --ignore-missing-dependencies --exclude libtvm_ffi.dylib --require-archs {delocate_archs} -w "$DELOCATE_TMP_DIR" -v {wheel}
+delocate-wheel --ignore-missing-dependencies --exclude libtvm_tirx.dylib --exclude libtvm_runtime.dylib --exclude libtvm_ffi.dylib --require-archs {delocate_archs} -w "$DELOCATE_TMP_DIR" -v {wheel}
python -m pip install -q wheel
TMP_DIR="$(mktemp -d)"
python -m wheel unpack "$DELOCATE_TMP_DIR"/*.whl -d "$TMP_DIR"
@@ -256,4 +243,4 @@ rm -rf "$DELOCATE_TMP_DIR" "$TMP_DIR"
[tool.cibuildwheel.windows]
before-build = 'python -m pip install "delvewheel>=1.12.0"'
-repair-wheel-command = 'delvewheel repair --analyze-existing --ignore-existing --exclude "*tvm_ffi*.dll" --exclude tvm_runtime_cuda.dll --exclude nvcuda.dll --exclude cudart64_13.dll --exclude nvrtc64_130_0.dll -w {dest_dir} {wheel}'
+repair-wheel-command = 'delvewheel repair --analyze-existing --ignore-existing --exclude "*tvm_ffi*.dll" --exclude tvm_tirx.dll --exclude tvm_runtime.dll --exclude tvm_runtime_cuda.dll --exclude nvcuda.dll --exclude cudart64_13.dll --exclude nvrtc64_130_0.dll -w {dest_dir} {wheel}'
diff --git a/LICENSE b/tirx/LICENSE
similarity index 95%
copy from LICENSE
copy to tirx/LICENSE
index 94abb3512b..d959613f68 100644
--- a/LICENSE
+++ b/tirx/LICENSE
@@ -201,29 +201,10 @@
limitations under the License.
------------------------------------------------------------------------------------
-This product bundles various third-party components under other open source licenses.
-This section summarizes those components and their licenses. See licenses/
-and 3rdparty/tvm-ffi/licenses/ for text of these licenses.
-
-
-Apache Software Foundation License 2.0
---------------------------------------
-
-3rdparty/OpenCL-Headers
-3rdparty/cutlass_fpA_intB_gemm
-3rdparty/tensorrt_llm
-3rdparty/tvm-ffi/3rdparty/dlpack
-
+This product bundles the following third-party component. See licenses/ for its
+license text. External dependencies are distributed under their own licenses.
MIT License
-----------
3rdparty/compiler-rt/builtin_fp16.h
-
-
-BSD 3-Clause "New" or "Revised" License
----------------------------------------
-
-3rdparty/cutlass
-3rdparty/libflash_attn
-3rdparty/tvm-ffi/3rdparty/libbacktrace
diff --git a/tirx/NOTICE b/tirx/NOTICE
new file mode 100644
index 0000000000..35d3709888
--- /dev/null
+++ b/tirx/NOTICE
@@ -0,0 +1,5 @@
+Apache TVM
+Copyright 2019-2023 The Apache Software Foundation
+
+This product includes software developed at
+The Apache Software Foundation (http://www.apache.org/).
diff --git a/pyproject.toml b/tirx/pyproject.toml
similarity index 57%
copy from pyproject.toml
copy to tirx/pyproject.toml
index 5c7512032f..c07ec984de 100644
--- a/pyproject.toml
+++ b/tirx/pyproject.toml
@@ -19,18 +19,15 @@
# mlc-z3-static ships the PIC static libz3 + headers consumed by USE_Z3=ON.
requires = [
"scikit-build-core>=0.11",
- "setuptools-scm>=8",
"mlc-z3-static>=4.16.0",
+ "apache-tvm-ffi>=0.1.14.post1",
]
build-backend = "scikit_build_core.build"
[project]
-name = "apache-tvm"
-# The version is derived from the most recent Git tag by setuptools_scm at build
-# time (see [tool.setuptools_scm]); no manual version stamping is required.
-dynamic = ["version"]
-description = "Apache TVM: An End-to-End Deep Learning Compiler Stack"
-readme = "README.md"
+name = "apache-tvm-tirx"
+version = "0.1.0a0"
+description = "TVM TIRx compiler and runtime"
license = "Apache-2.0"
license-files = ["LICENSE"]
requires-python = ">=3.10"
@@ -55,15 +52,14 @@ dependencies = [
"ml_dtypes",
"numpy",
"typing_extensions",
+ "packaging",
]
[project.optional-dependencies]
torch = ["torch"]
cuda = ["cuda-bindings"]
-meta-schedule = ["xgboost"]
popen-pool = ["psutil", "cloudpickle"]
-rpc = ["tornado", "psutil", "cloudpickle"]
-all = ["xgboost", "psutil", "cloudpickle", "tornado"]
+all = ["psutil", "cloudpickle"]
[project.urls]
Homepage = "https://tvm.apache.org/"
@@ -75,6 +71,8 @@ Repository = "https://github.com/apache/tvm"
torch = ["torch"]
popen-pool = ["psutil", "cloudpickle"]
test = [
+ "scipy",
+ "z3-solver",
"pytest",
"pytest-xdist",
"pytest-cov",
@@ -95,51 +93,31 @@ wheel.py-api = "py3"
build-dir = "build"
# Wheel configuration
-wheel.packages = ["python/tvm"]
-wheel.install-dir = "tvm"
-
-# Derive the project version from the Git tag via setuptools_scm
-# (configured under [tool.setuptools_scm]).
-metadata.version.provider = "scikit_build_core.metadata.setuptools_scm"
+wheel.packages = ["python/tvm_tirx"]
+wheel.install-dir = "tvm_tirx"
# Source distribution configuration
sdist.include = [
- # Build files
"/CMakeLists.txt",
"/pyproject.toml",
- "/cmake/**/*",
-
- # Source code
- "/src/**/*.cc",
- "/src/**/*.h",
- "/include/**/*.h",
-
- # Python source
- "/python/tvm/**/*.py",
- "/python/tvm/**/*.pyi",
-
- # Documentation and metadata
- "/docs/**/*",
+ "/cmake/**",
+ "/src/**",
+ "/include/**",
+ "/python/**",
+ "/licenses/**",
"/LICENSE",
- "/README.md",
"/NOTICE",
-
- # Tests
- "/tests/**/*",
-]
-
-sdist.exclude = [
- "**/.git",
- "**/.github",
- "**/__pycache__",
- "**/*.pyc",
- "build",
- "dist",
- "**/3rdparty/*/docs",
- "**/3rdparty/*/media",
- "**/3rdparty/*/examples",
- "**/3rdparty/*/test",
+ "/tests/**",
+ "/3rdparty/compiler-rt/builtin_fp16.h",
+ "!**/.git",
+ "!**/.git/**",
+ "!**/__pycache__/**",
+ "!**/*.pyc",
+ "!**/.pytest_cache/**",
+ "!**/build/**",
+ "!**/node_modules/**",
]
+sdist.exclude = ["**"]
# Logging
logging.level = "INFO"
@@ -151,17 +129,6 @@ USE_CUDA = "OFF"
USE_Z3 = "AUTO"
BUILD_TESTING = "OFF"
-[tool.setuptools_scm]
-# Version comes from the most recent Git tag (vMAJOR.MINOR.devN or vMAJOR.MINOR.PATCH).
-# guess-next-dev reproduces the previous version.py behaviour for vMAJOR.MINOR.devN
-# tags (e.g. v0.26.dev0 + N commits -> 0.26.devN). local_scheme = "no-local-version"
-# drops the +g<hash> local segment: PyPI rejects local versions on upload, and
-# this matches the public version the old version.py stamped.
-version_file = "python/tvm/_version.py"
-version_scheme = "guess-next-dev"
-local_scheme = "no-local-version"
-fallback_version = "0.28.dev0"
-
[tool.pytest.ini_options]
testpaths = ["tests"]
addopts = "-v --tb=short"
@@ -173,68 +140,12 @@ markers = [
"cuda_arch(*arches, device='cuda0'): Require one of the exact CUDA runtime architectures",
]
-[tool.ruff]
-include = [
- "python/**/*.py",
- "tests/**/*.py",
- "docs/**/*.py",
- "apps/**/*.py",
- "ci/**/*.py",
- "conftest.py",
- "jvm/**/*.py",
- "web/tests/**/*.py",
-]
-line-length = 100
-indent-width = 4
-target-version = "py310"
-exclude = ["3rdparty", "build", "dist", ".venv", ".ruff_cache", "node_modules"]
-
-[tool.ruff.lint]
-select = [
- "E", # pycodestyle errors, https://docs.astral.sh/ruff/rules/#pycodestyle-e-w
- "F", # pyflakes, https://docs.astral.sh/ruff/rules/#pyflakes-f
- "I", # isort, https://docs.astral.sh/ruff/rules/#isort-i
- "UP", # pyupgrade, https://docs.astral.sh/ruff/rules/#pyupgrade-up
- "RUF", # ruff, https://docs.astral.sh/ruff/rules/#ruff-specific-rules-ruf
-]
-ignore = [
- # pyupgrade (style preference or low-value)
- "UP028", # yield-in-for-loop — sometimes clearer to keep
- "UP030", # format-literals — minor
- "UP031", # printf-string-formatting — style preference
- # ruff (preview rules, not yet stable)
- "RUF043", # pytest-raises-ambiguous-pattern
- "RUF059", # unused-unpacked-variable
-]
-fixable = ["ALL"]
-unfixable = []
-
-[tool.ruff.lint.per-file-ignores]
-"__init__.py" = ["E402", "F401", "F403", "F405"]
-"python/tvm/relax/op/nn/nn.py" = ["E501"]
-"docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py" = ["RUF003"]
-"python/tvm/relax/frontend/tflite/tflite_frontend.py" = ["E501"]
-"python/tvm/relax/transform/legalize_ops/nn.py" = ["E501"]
-# TIRx parser and lowering tests intentionally retain unused IR bindings.
-"tests/python/tirx/**/*.py" = ["F841"]
-
-[tool.ruff.lint.isort]
-known-first-party = ["tvm"]
-
-[tool.ruff.format]
-quote-style = "double"
-indent-style = "space"
-skip-magic-trailing-comma = false
-line-ending = "auto"
-docstring-code-format = false
-docstring-code-line-length = "dynamic"
-
[tool.cibuildwheel]
# Skip win32, i686 (32-bit), and musllinux wheels.
skip = "*-win32 *-manylinux_i686 *-musllinux*"
build-verbosity = 1
test-requires = ["pytest", "numpy"]
-test-command = "pytest -vvs {project}/tests/python/all-platform-minimal-test"
+test-command = "python -I -c \"import importlib.util; assert importlib.util.find_spec('tvm') is None\" && python -I -m pytest -vvs {package}/tests/python/all-platform-minimal-test {package}/tests/python/runtime/test_runtime_tensor.py {package}/tests/python/ir/test_ir_type.py {package}/tests/python/script/test_ir_builder.py {package}/tests/python/codegen/llvm/test_llvm.py::test_llvm_vadd_pipeline {package}/tests/python/codegen/gpu_common/test_allreduce.py::test_allreduce_sum_compile"
[tool.cibuildwheel.linux]
repair-wheel-command = "auditwheel repair --exclude libtvm_ffi.so --exclude libtvm_runtime_cuda.so --exclude 'libcuda.so.*' --exclude 'libcudart.so.*' --exclude 'libnvrtc.so.*' --exclude 'libnvrtc-builtins.so.*' -w {dest_dir} {wheel}"
(tvm) 09/33: [REFACTOR][CMake] Build and install independent core and parent libraries
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 760fae90d2ccbfe9287ec0b89ec2a77cc555b83e
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 01:37:33 2026 +0000
[REFACTOR][CMake] Build and install independent core and parent libraries
---
CMakeLists.txt | 836 +++++----------------
.../modules => cmake/modules/contrib}/CUDA.cmake | 78 --
.../modules => cmake/modules/contrib}/ROCM.cmake | 36 -
tests/cpp/CMakeLists.txt | 2 +-
CMakeLists.txt => tirx/CMakeLists.txt | 484 +++++-------
tirx/cmake/modules/CUDA.cmake | 73 +-
tirx/cmake/modules/ClangFlags.cmake | 2 +-
tirx/cmake/modules/Logging.cmake | 5 +-
tirx/cmake/modules/OpenCL.cmake | 11 +-
tirx/cmake/modules/ROCM.cmake | 31 -
tirx/cmake/modules/contrib/Z3.cmake | 12 +
.../Logging.cmake => tvm_tirxConfig.cmake.in} | 30 +-
{tests => tirx/tests}/cpp/CMakeLists.txt | 32 +-
13 files changed, 440 insertions(+), 1192 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index bfa62d6b9d..3e40018bb0 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -1,88 +1,99 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
cmake_minimum_required(VERSION 3.18)
project(tvm C CXX)
-
-# --- TVM version (no dependency on version.py) ---
-# When built via scikit-build-core the version is resolved by setuptools_scm and
-# passed in as SKBUILD_PROJECT_VERSION_FULL; bake it into the C++ TVM_VERSION macro.
-# A bare `cmake` build (no scikit-build-core) leaves the checked-in default in
-# include/tvm/runtime/base.h untouched. An explicit -DTVM_VERSION always wins.
-if(NOT DEFINED TVM_VERSION AND DEFINED SKBUILD_PROJECT_VERSION_FULL)
- set(TVM_VERSION "${SKBUILD_PROJECT_VERSION_FULL}")
+include(GNUInstallDirs)
+if(EXISTS "${CMAKE_BINARY_DIR}/config.cmake")
+ include("${CMAKE_BINARY_DIR}/config.cmake")
+elseif(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/config.cmake")
+ include("${CMAKE_CURRENT_SOURCE_DIR}/config.cmake")
endif()
-if(DEFINED TVM_VERSION)
- message(STATUS "TVM_VERSION=${TVM_VERSION}")
- add_compile_definitions(TVM_VERSION="${TVM_VERSION}")
+set(_installed_tirx ON)
+if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tirx/CMakeLists.txt")
+ set(_installed_tirx OFF)
endif()
-
-# Utility functions
-include(cmake/utils/Utils.cmake)
-include(cmake/utils/Summary.cmake)
-include(cmake/utils/Linker.cmake)
-include(cmake/utils/Library.cmake)
-include(cmake/utils/FindCUDA.cmake)
-include(cmake/utils/FindNCCL.cmake)
-include(cmake/utils/FindOpenCL.cmake)
-include(cmake/utils/FindVulkan.cmake)
-include(cmake/utils/FindLLVM.cmake)
-include(cmake/utils/FindROCM.cmake)
-include(cmake/utils/FindRCCL.cmake)
-include(cmake/utils/FindNVSHMEM.cmake)
-
-if(EXISTS ${CMAKE_BINARY_DIR}/config.cmake)
- include(${CMAKE_BINARY_DIR}/config.cmake)
+option(TVM_USE_INSTALLED_TIRX "Build against an installed apache-tvm-tirx package" ${_installed_tirx})
+option(TVM_BUILD_PYTHON_MODULE "Build Python module with scikit-build-core" OFF)
+if(TVM_BUILD_PYTHON_MODULE)
+ set(CMAKE_INSTALL_LIBDIR lib)
+endif()
+if(USE_CLML_GRAPH_EXECUTOR)
+ if(USE_CLML_GRAPH_EXECUTOR STREQUAL "ON")
+ set(USE_OPENCL "${CMAKE_CURRENT_SOURCE_DIR}/clml")
+ else()
+ set(USE_OPENCL "${USE_CLML_GRAPH_EXECUTOR}")
+ endif()
+endif()
+if(TVM_USE_INSTALLED_TIRX)
+ find_package(Python COMPONENTS Interpreter REQUIRED)
+ if(NOT tvm_tirx_DIR)
+ execute_process(COMMAND "${Python_EXECUTABLE}" -c
+ "import importlib.util; from pathlib import Path; print(Path(importlib.util.find_spec('tvm_tirx').origin).parent / 'lib/cmake/tvm_tirx')"
+ OUTPUT_VARIABLE tvm_tirx_DIR OUTPUT_STRIP_TRAILING_WHITESPACE
+ COMMAND_ERROR_IS_FATAL ANY)
+ endif()
+ find_package(tvm_tirx CONFIG REQUIRED)
+ if(NOT TVM_TIRX_VERSION STREQUAL "0.1.0a0")
+ message(FATAL_ERROR "TVM requires apache-tvm-tirx 0.1.0a0")
+ endif()
else()
- if(EXISTS ${CMAKE_SOURCE_DIR}/config.cmake)
- include(${CMAKE_SOURCE_DIR}/config.cmake)
+ if(NOT TVM_OPENCL_HEADERS_ROOT AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/OpenCL-Headers/CL/cl.h")
+ set(TVM_OPENCL_HEADERS_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/OpenCL-Headers")
endif()
+ add_subdirectory(tirx)
endif()
-
-# NOTE: do not modify this file to change option values.
-# You can create a config.cmake at build folder
-# and add set(OPTION VALUE) to override these build options.
-# Alernatively, use cmake -DOPTION=VALUE through command-line.
-tvm_option(USE_CUDA "Build with CUDA" OFF)
+foreach(_utility Utils Summary Linker Library FindCUDA FindROCM FindOpenCL CCache)
+ include("${TVM_TIRX_CMAKE_DIR}/utils/${_utility}.cmake")
+endforeach()
+foreach(_collective NCCL RCCL NVSHMEM)
+ include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils/Find${_collective}.cmake")
+endforeach()
+if(NOT TARGET tvm_ffi_header)
+ add_library(tvm_ffi_header ALIAS tvm_ffi::header)
+endif()
+# Keep existing integration modules' target names usable with an installed core.
+foreach(_library tvm_runtime tvm_runtime_static tvm_runtime_cuda tvm_runtime_rocm tvm_runtime_opencl tvm_runtime_vulkan tvm_runtime_metal)
+ if(NOT TARGET ${_library} AND TARGET tvm_tirx::${_library})
+ add_library(${_library} ALIAS tvm_tirx::${_library})
+ endif()
+endforeach()
+set(BUILD_FOR_ANDROID "${TVM_TIRX_BUILD_FOR_ANDROID}")
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_CUDA_STANDARD 20)
+set(CMAKE_CUDA_STANDARD_REQUIRED ON)
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+tvm_option(USE_CUDA "Build with CUDA" "${TVM_TIRX_USE_CUDA}")
tvm_option(USE_NCCL "Build with NCCL" OFF)
-tvm_option(USE_OPENCL "Build with OpenCL" OFF)
-tvm_option(USE_OPENCL_ENABLE_HOST_PTR "Enable OpenCL memory object access to host" OFF)
-tvm_option(USE_OPENCL_GTEST "Path to OpenCL specific gtest version for runtime cpp tests." /path/to/opencl/gtest)
-tvm_option(USE_VULKAN "Build with Vulkan" OFF)
-
-
-# Whether to use spirv-tools and SPIRV-Headers from Khronos GitHub or GitLab.
-#
-# Possible values:
-# - OFF: not to use
-# - /path/to/install: path to your khronis spirv-tools and SPIRV-Headers installation directory
-#
-tvm_option(USE_KHRONOS_SPIRV "Whether to use spirv-tools and SPIRV-Headers from Khronos GitHub or GitLab" OFF)
-tvm_option(USE_SPIRV_KHR_INTEGER_DOT_PRODUCT "whether enable SPIRV_KHR_DOT_PRODUCT" OFF)
-tvm_option(USE_METAL "Build with Metal" OFF)
-tvm_option(USE_ROCM "Build with ROCM" OFF)
+tvm_option(USE_ROCM "Build with ROCM" "${TVM_TIRX_USE_ROCM}")
tvm_option(USE_RCCL "Build with RCCL" OFF)
-tvm_option(ROCM_PATH "The path to rocm" /opt/rocm)
-tvm_option(USE_HEXAGON_QHL "Enable Hexagon QHL compiler lowering (requires external runtime wrappers)" OFF)
-
tvm_option(USE_RPC "Build with RPC" ON)
tvm_option(USE_THREADS "Build with thread support" ON)
-tvm_option(USE_LLVM "Build with LLVM, can be set to specific llvm-config path" OFF)
-tvm_option(USE_OPENMP "Build with OpenMP thread pool implementation" OFF)
tvm_option(TVM_LOG_BEFORE_THROW "Whether log before throw, for debugging purposes" OFF)
-tvm_option(USE_RTTI "Build with RTTI" ON)
tvm_option(USE_MSVC_MT "Build with MT" OFF)
tvm_option(INSTALL_DEV "Install compiler infrastructure" OFF)
tvm_option(HIDE_PRIVATE_SYMBOLS "Compile with -fvisibility=hidden." ON)
tvm_option(INDEX_DEFAULT_I64 "Defaults the index datatype to int64" ON)
tvm_option(USE_GTEST "Use GoogleTest for C++ sanity tests" AUTO)
-tvm_option(USE_CUSTOM_LOGGING "Use user-defined custom logging, tvm::runtime::detail::LogFatalImpl and tvm::runtime::detail::LogMessageImpl must be implemented" OFF)
tvm_option(USE_ALTERNATIVE_LINKER "Use 'mold' or 'lld' if found when invoking compiler to link artifact" AUTO)
-tvm_option(USE_CCACHE "Use ccache if found when invoking compiler" AUTO)
-
-# 3rdparty libraries
-tvm_option(COMPILER_RT_PATH "Path to COMPILER-RT" "3rdparty/compiler-rt")
-# Contrib library options
tvm_option(USE_BLAS "The blas library to be linked" none)
tvm_option(USE_AMX "Enable Intel AMX" OFF)
-tvm_option(USE_Z3 "Build with Z3 SMT solver support" AUTO)
tvm_option(USE_MKL "MKL root path when use MKL blas" OFF)
tvm_option(USE_DNNL "Enable DNNL codegen" OFF)
tvm_option(USE_CUDNN "Build with cuDNN" OFF)
@@ -108,190 +119,43 @@ tvm_option(USE_CLML "Build with CLML Codegen support" OFF)
tvm_option(USE_CLML_GRAPH_EXECUTOR "Build with CLML graph runtime" OFF)
tvm_option(USE_NVSHMEM "Build with NVSHMEM support" OFF)
-# Python package options
-tvm_option(TVM_BUILD_PYTHON_MODULE "Build Python module with scikit-build-core" OFF)
-
-# include directories
-include_directories(${CMAKE_INCLUDE_PATH})
-include_directories("include")
-include_directories(SYSTEM ${COMPILER_RT_PATH})
-
-# initial variables
-set(TVM_LINKER_LIBS "")
-set(TVM_RUNTIME_LINKER_LIBS "")
-set(TVM_RUNTIME_BACKEND_LIBS "")
-# Early target creation so contrib cmake files can call
-# target_link_libraries(tvm_runtime_extra PRIVATE <object_lib>) directly.
-add_library(tvm_runtime_extra SHARED)
-list(APPEND TVM_RUNTIME_BACKEND_LIBS tvm_runtime_extra)
-set_target_properties(tvm_runtime_extra PROPERTIES LINKER_LANGUAGE CXX)
-# INTERFACE target carrying compile definitions for OBJECT libs that build
-# into tvm_runtime_extra. On MSVC, TVM_RUNTIME_EXTRA_EXPORTS exports extra-owned
-# APIs while core-runtime APIs remain imports from tvm_runtime.dll.
-add_library(tvm_runtime_extra_defs INTERFACE)
-target_link_libraries(tvm_runtime_extra_defs INTERFACE tvm_ffi_header)
-target_compile_definitions(tvm_runtime_extra_defs
- INTERFACE TVM_RUNTIME_EXTRA_EXPORTS TVM_FFI_EXPORTS)
-
-
-# Check if this is being run on its own or as a subdirectory for another project
-# If we update to CMake 2.21+, we can use PROJECT_IS_TOP_LEVEL instead
-get_directory_property(IS_SUBPROJECT PARENT_DIRECTORY)
-
-if(NOT IS_SUBPROJECT AND NOT DEFINED "${CMAKE_EXPORT_COMPILE_COMMANDS}")
- # If not set manually, change the default to ON
- set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
-endif()
-
-if(TVM_LOG_BEFORE_THROW)
- # log error before throw as
- # when system have issues with stack trace
- add_definitions(-DTVM_LOG_BEFORE_THROW=1)
-endif()
-
-# Generic compilation options
+include_directories(include)
if(MSVC)
- add_definitions(-DWIN32_LEAN_AND_MEAN)
- add_definitions(-D_CRT_SECURE_NO_WARNINGS)
- add_definitions(-D_SCL_SECURE_NO_WARNINGS)
- add_definitions(-D_ENABLE_EXTENDED_ALIGNED_STORAGE)
- add_definitions(-DNOMINMAX)
- # regeneration does not work well with msbuild custom rules.
- set(CMAKE_SUPPRESS_REGENERATION ON)
- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /EHsc")
- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /MP")
- add_compile_options(/bigobj)
-
- # Use standard-conforming two-phase name resolution for templates.
- # This minimizes the differences between g++/clang builds on Linux,
- # and MSVC builds on Windows.
- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /permissive-")
-
- # MSVC already errors on undefined symbols, no additional flag needed.
- set(TVM_NO_UNDEFINED_SYMBOLS "")
-
+ add_compile_definitions(WIN32_LEAN_AND_MEAN _CRT_SECURE_NO_WARNINGS
+ _SCL_SECURE_NO_WARNINGS _ENABLE_EXTENDED_ALIGNED_STORAGE NOMINMAX)
+ add_compile_options(/bigobj /permissive- /EHsc)
if(USE_MSVC_MT)
- foreach(flag_var
- CMAKE_CXX_FLAGS CMAKE_CXX_FLAGS_DEBUG CMAKE_CXX_FLAGS_RELEASE
- CMAKE_CXX_FLAGS_MINSIZEREL CMAKE_CXX_FLAGS_RELWITHDEBINFO)
- if(${flag_var} MATCHES "/MD")
- string(REGEX REPLACE "/MD" "/MT" ${flag_var} "${${flag_var}}")
- endif(${flag_var} MATCHES "/MD")
- endforeach(flag_var)
- # Static linking. CMake behavior changed in 3.15 making this necessary.
- add_compile_options(/MT)
+ set(CMAKE_MSVC_RUNTIME_LIBRARY "MultiThreaded$<$<CONFIG:Debug>:Debug>")
endif()
- # Disable common MSVC warnings
- # Integer conversion warnings(e.g. int64 to int)
- add_compile_options(/wd4244)
- add_compile_options(/wd4267)
- # Signed unsigned constant comparison
- add_compile_options(/wd4018)
- # Aligned alloc may not met(need c++17)
- add_compile_options(/wd4316)
- # unreferenced local variables(usually in exception catch)
- add_compile_options(/wd4101)
- # always inline keyword not necessary
- add_compile_options(/wd4180)
- # DLL interface warning in c++
- add_compile_options(/wd4251)
- # destructor was implicitly defined as deleted
- add_compile_options(/wd4624)
- # unary minus operator applied to unsigned type, result still unsigned
- add_compile_options(/wd4146)
- # 'inline': used more than once
- add_compile_options(/wd4141)
- # unknown pragma
- add_compile_options(/wd4068)
-else(MSVC)
- set(WARNING_FLAG -Wall)
- if ("${CMAKE_BUILD_TYPE}" STREQUAL "Debug")
- message(STATUS "Build in Debug mode")
- set(CMAKE_C_FLAGS "-O0 -g ${WARNING_FLAG} -fPIC ${CMAKE_C_FLAGS}")
- set(CMAKE_CXX_FLAGS "-O0 -g ${WARNING_FLAG} -fPIC ${CMAKE_CXX_FLAGS}")
- set(CMAKE_CUDA_FLAGS "-O0 -g -Xcompiler=-Wall -Xcompiler=-fPIC ${CMAKE_CUDA_FLAGS}")
- else()
- set(CMAKE_C_FLAGS "-O2 ${WARNING_FLAG} -fPIC ${CMAKE_C_FLAGS}")
- set(CMAKE_CXX_FLAGS "-O2 ${WARNING_FLAG} -fPIC ${CMAKE_CXX_FLAGS}")
- set(CMAKE_CUDA_FLAGS "-O2 -Xcompiler=-Wall -Xcompiler=-fPIC ${CMAKE_CUDA_FLAGS}")
- set(TVM_VISIBILITY_FLAG "")
- if (HIDE_PRIVATE_SYMBOLS)
- message(STATUS "Hide private symbols...")
- set(TVM_VISIBILITY_FLAG "-fvisibility=hidden")
- endif(HIDE_PRIVATE_SYMBOLS)
- endif ()
- if (CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang" AND
- CMAKE_CXX_COMPILER_VERSION VERSION_GREATER 7.0)
- set(CMAKE_CXX_FLAGS "-faligned-new ${CMAKE_CXX_FLAGS}")
+else()
+ add_compile_options(-Wall)
+ if(HIDE_PRIVATE_SYMBOLS)
+ add_compile_options(-fvisibility=hidden)
endif()
-
- # ld option to warn if symbols are undefined (e.g. libtvm_runtime.so
- # using symbols only present in libtvm.so). Not needed for MSVC,
- # since this is already the default there.
- if(${CMAKE_SYSTEM_NAME} MATCHES "Darwin" OR ${CMAKE_SYSTEM_NAME} MATCHES "iOS")
+ if(APPLE)
set(TVM_NO_UNDEFINED_SYMBOLS "-Wl,-undefined,error")
else()
set(TVM_NO_UNDEFINED_SYMBOLS "-Wl,--no-undefined")
endif()
- message(STATUS "Forbidding undefined symbols in shared library, using ${TVM_NO_UNDEFINED_SYMBOLS} on platform ${CMAKE_SYSTEM_NAME}")
-
- # Detect if we're compiling for Android.
- set(TEST_FOR_ANDROID_CXX
- "#ifndef __ANDROID__"
- "#error"
- "#endif"
- "int main() {}")
- set(TEST_FOR_ANDROID_DIR
- "${CMAKE_BINARY_DIR}${CMAKE_FILES_DIRECTORY}/CMakeTmp")
- set(TEST_FOR_ANDROID_FILE "${TEST_FOR_ANDROID_DIR}/test_for_android.cc")
- string(REPLACE ";" "\n" TEST_FOR_ANDROID_CXX_TEXT "${TEST_FOR_ANDROID_CXX}")
- file(WRITE "${TEST_FOR_ANDROID_FILE}" "${TEST_FOR_ANDROID_CXX_TEXT}")
- try_compile(BUILD_FOR_ANDROID "${CMAKE_BINARY_DIR}${CMAKE_FILES_DIRECTORY}"
- "${TEST_FOR_ANDROID_FILE}")
- file(REMOVE "${TEST_FOR_ANDROID_FILE}")
- if(BUILD_FOR_ANDROID)
- message(STATUS "Building for Android")
- endif()
-endif(MSVC)
-
-list(APPEND TVM_RUNTIME_LINKER_LIBS ${CMAKE_DL_LIBS})
-
-# add source group
-tvm_file_glob(GLOB_RECURSE GROUP_SOURCE "src/*.cc")
-tvm_file_glob(GLOB_RECURSE GROUP_INCLUDE "src/*.h" "include/*.h")
-assign_source_group("Source" ${GROUP_SOURCE})
-assign_source_group("Include" ${GROUP_INCLUDE})
-
-# Source file lists
-# Core compiler sources, extended by the backend modules.
+endif()
+if(DEFINED TVM_VERSION)
+ add_compile_definitions(TVM_VERSION="${TVM_VERSION}")
+elseif(DEFINED SKBUILD_PROJECT_VERSION_FULL)
+ add_compile_definitions(TVM_VERSION="${SKBUILD_PROJECT_VERSION_FULL}")
+else()
+ add_compile_definitions(TVM_VERSION="0.28.dev0")
+endif()
+if(TVM_LOG_BEFORE_THROW)
+ add_compile_definitions(TVM_LOG_BEFORE_THROW=1)
+endif()
+set(TVM_LINKER_LIBS "")
+set(TVM_RUNTIME_LINKER_LIBS "")
+set(TVM_RUNTIME_BACKEND_LIBS "")
tvm_file_glob(GLOB_RECURSE COMPILER_SRCS
- src/ir/*.cc
- src/sym/*.cc
src/te/*.cc
- src/tirx/*.cc
src/s_tir/*.cc
src/topi/*.cc
src/driver/*.cc
- src/support/*.cc
- # Shared printer and TVMScript IR builder. Per-dialect printer/builder pieces
- # live under src/<dialect>/script/. List shared files explicitly so a dialect
- # implementation cannot be silently included through a recursive glob.
- src/script/ir_builder/base.cc
- src/script/ir_builder/ir.cc
- src/script/ir_builder/frame.cc
- src/script/printer/config.cc
- src/script/printer/doc.cc
- src/script/printer/doc_translator.cc
- src/script/printer/ir/call.cc
- src/script/printer/ir/utils.cc
- src/script/printer/ir/ir.cc
- src/script/printer/ir/misc.cc
- src/script/printer/ir/prim_type.cc
- src/script/printer/ir/prim_expr.cc
- src/script/printer/ir/stmt.cc
- src/script/printer/printer.cc
- src/script/printer/doc_printer.cc
- src/script/printer/script_printer.cc
src/relax/ir/*.cc
src/relax/op/*.cc
src/relax/analysis/*.cc
@@ -306,431 +170,139 @@ tvm_file_glob(GLOB_RECURSE COMPILER_SRCS
src/relax/op/distributed/*.cc
src/relax/script/*.cc
src/relax/testing/*.cc
- )
-
-tvm_file_glob(GLOB CODEGEN_SRCS
- src/target/*.cc
- src/codegen/c/*.cc
- src/codegen/*.cc
- src/codegen/canonicalizer/*.cc
- src/codegen/llvm/target_kind.cc
- src/codegen/llvm/canonicalizer/*.cc
- src/backend/cuda/codegen/*.cc
- src/backend/cuda/op/*.cc
- src/backend/cuda/transforms/*.cc
- src/backend/cuda/tile/*.cc
- src/backend/cuda/layout/*.cc
- src/backend/hexagon/codegen/*.cc
- src/backend/metal/codegen/*.cc
- src/backend/metal/op/*.cc
- src/backend/opencl/codegen/*.cc
- src/backend/opencl/op/*.cc
- src/backend/rocm/codegen/*.cc
- src/backend/trn/codegen/*.cc
- src/backend/trn/op/*.cc
- src/backend/trn/transform/*.cc
- src/backend/trn/layout/*.cc
- src/backend/vulkan/codegen/target_kind.cc
- src/backend/vulkan/codegen/vulkan_fallback_module.cc
- src/backend/webgpu/codegen/*.cc
- )
-
-list(APPEND COMPILER_SRCS ${CODEGEN_SRCS})
-# Core runtime sources, extended by the backend modules.
-tvm_file_glob(GLOB RUNTIME_SRCS
- src/runtime/*.cc
- src/runtime/rpc/minrpc/*.cc
)
-# VM-owned runtime implementations, including external-module constant loading.
-tvm_file_glob(GLOB RUNTIME_VM_SRCS src/runtime/vm/*.cc)
-list(APPEND RUNTIME_SRCS ${RUNTIME_VM_SRCS})
-# Note: src/runtime/extra/disco/** moves to libtvm_runtime_extra.
-# Note: src/backend/*/runtime sources move to per-backend DSOs.
-set(TVM_RUNTIME_EXT_OBJS "")
+# VM, RPC, Disco and optional integrations remain parent runtime extensions.
+add_library(tvm_runtime_extra SHARED)
+add_library(tvm_runtime_extra_defs INTERFACE)
+target_link_libraries(tvm_runtime_extra_defs INTERFACE
+ tvm_tirx::tvm_runtime tvm_tirx::private_headers tvm_ffi::header)
+target_compile_definitions(tvm_runtime_extra_defs INTERFACE TVM_RUNTIME_EXTRA_EXPORTS TVM_FFI_EXPORTS)
+target_link_libraries(tvm_runtime_extra PRIVATE tvm_runtime_extra_defs)
+tvm_file_glob(GLOB _vm_sources src/runtime/vm/*.cc src/runtime/rpc/minrpc/*.cc)
+target_sources(tvm_runtime_extra PRIVATE ${_vm_sources})
if(USE_RPC)
- message(STATUS "Build with RPC support...")
- tvm_file_glob(GLOB RUNTIME_RPC_SRCS src/runtime/rpc/*.cc)
- list(APPEND RUNTIME_SRCS ${RUNTIME_RPC_SRCS})
-endif(USE_RPC)
-
-# Package runtime rules
-if(NOT USE_RTTI)
+ tvm_file_glob(GLOB _rpc_sources src/runtime/rpc/*.cc)
+ target_sources(tvm_runtime_extra PRIVATE ${_rpc_sources})
+ if(WIN32)
+ target_link_libraries(tvm_runtime_extra PRIVATE ws2_32)
+ endif()
endif()
-
-if(INDEX_DEFAULT_I64)
- add_definitions(-DTVM_INDEX_DEFAULT_I64=1)
+if(USE_CUDA)
+ if(NOT TARGET tvm_tirx::tvm_runtime_cuda)
+ message(FATAL_ERROR "Parent CUDA support requires a CUDA-enabled core")
+ endif()
+ find_cuda(${USE_CUDA} ${USE_CUDNN})
+ enable_language(CUDA)
+ if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
+ set(CMAKE_CUDA_ARCHITECTURES native)
+ endif()
+ include_directories(SYSTEM ${CUDA_INCLUDE_DIRS})
+ tvm_file_glob(GLOB _vm_cuda_sources src/runtime/vm/cuda/*.cc)
+ target_sources(tvm_runtime_extra PRIVATE ${_vm_cuda_sources})
endif()
-
-if(USE_KALLOC_ALIGNMENT)
- message(STATUS "Build Alloc alignment set to ${USE_KALLOC_ALIGNMENT}")
- add_definitions(-DTVM_KALLOC_ALIGNMENT=${USE_KALLOC_ALIGNMENT})
-endif(USE_KALLOC_ALIGNMENT)
-
-# Caches the build.
-# Note that ccache-3.x doesn't support nvcc well, so CUDA kernels may never hit the cache and still
-# need to be re-compiled every time. Using ccache 4.0+ can resolve this issue.
-include(cmake/utils/CCache.cmake)
-
-include(CheckCXXCompilerFlag)
-set(CMAKE_CXX_STANDARD_REQUIRED ON)
-set(CMAKE_CXX_STANDARD 20)
-set(CMAKE_CUDA_STANDARD_REQUIRED ON)
-set(CMAKE_CUDA_STANDARD 20)
-
-# Module rules
-include(cmake/modules/CUDA.cmake)
-include(cmake/modules/contrib/CLML.cmake) # Must be before OpenCL.cmake
-include(cmake/modules/OpenCL.cmake)
-include(cmake/modules/OpenMP.cmake)
-include(cmake/modules/Vulkan.cmake)
-include(cmake/modules/Metal.cmake)
-include(cmake/modules/ROCM.cmake)
-include(cmake/modules/LLVM.cmake)
+if(USE_ROCM)
+ if(NOT TARGET tvm_tirx::tvm_runtime_rocm)
+ message(FATAL_ERROR "Parent ROCm support requires a ROCm-enabled core")
+ endif()
+ find_rocm(${USE_ROCM})
+ include_directories(SYSTEM ${ROCM_INCLUDE_DIRS})
+ add_compile_definitions(__HIP_PLATFORM_HCC__=1 __HIP_PLATFORM_AMD__=1)
+endif()
+if(USE_NVTX)
+ target_compile_definitions(tvm_runtime_extra_defs INTERFACE TVM_NVTX_ENABLED=1)
+endif()
+include(cmake/modules/contrib/CUDA.cmake)
+include(cmake/modules/contrib/ROCM.cmake)
include(cmake/modules/contrib/BLAS.cmake)
include(cmake/modules/contrib/DNNL.cmake)
include(cmake/modules/contrib/AMX.cmake)
include(cmake/modules/contrib/CUTLASS.cmake)
+include(cmake/modules/contrib/CLML.cmake)
include(cmake/modules/contrib/Random.cmake)
include(cmake/modules/contrib/Sort.cmake)
-include(cmake/modules/contrib/Z3.cmake)
include(cmake/modules/contrib/CoreML.cmake)
include(cmake/modules/contrib/TensorRT.cmake)
include(cmake/modules/contrib/NNAPI.cmake)
include(cmake/modules/contrib/ExampleNPU.cmake)
include(cmake/modules/contrib/vllm.cmake)
-include(cmake/modules/Git.cmake)
-
-include(cmake/modules/RuntimeExtra.cmake)
-add_subdirectory(3rdparty/tvm-ffi)
-include(GNUInstallDirs)
-
-if(USE_THREADS)
- message(STATUS "Build with thread support...")
- set(CMAKE_THREAD_PREFER_PTHREAD TRUE)
- set(THREADS_PREFER_PTHREAD_FLAG TRUE)
- find_package(Threads REQUIRED)
-endif()
-
-set(TVM_IS_DEBUG_BUILD OFF)
-if(CMAKE_BUILD_TYPE STREQUAL "Debug" OR CMAKE_BUILD_TYPE STREQUAL "RelWithDebInfo" OR CMAKE_CXX_FLAGS MATCHES "-g")
- set(TVM_IS_DEBUG_BUILD ON)
- set(FILE_PREFIX_MAP_FLAG "-ffile-prefix-map=..=${CMAKE_CURRENT_SOURCE_DIR}")
- CHECK_CXX_COMPILER_FLAG("${FILE_PREFIX_MAP_FLAG}" FILE_PREFIX_MAP_SUPPORTED)
-endif()
-
-# Core runtime objects are shared by the dynamic and static libraries.
-add_library(tvm_runtime_objs OBJECT ${RUNTIME_SRCS})
-target_link_libraries(tvm_runtime_objs PUBLIC tvm_ffi_header)
-if(TVM_VISIBILITY_FLAG)
- target_compile_options(tvm_runtime_objs PRIVATE "${TVM_VISIBILITY_FLAG}")
-endif()
-if(MSVC)
- # Runtime objects export only runtime/FFI symbols, never compiler symbols.
- target_compile_definitions(tvm_runtime_objs PRIVATE TVM_RUNTIME_EXPORTS TVM_FFI_EXPORTS)
+if(USE_CLML_GRAPH_EXECUTOR)
+ target_link_libraries(tvm_clml_objs PRIVATE tvm_tirx::tvm_runtime_opencl)
+ target_link_libraries(tvm_runtime_extra PRIVATE tvm_tirx::tvm_runtime_opencl)
endif()
-
-# Loaded with RTLD_GLOBAL so downstream runtime modules can resolve its symbols.
-add_library(tvm_runtime SHARED
- $<TARGET_OBJECTS:tvm_runtime_objs>
- ${TVM_RUNTIME_EXT_OBJS}
-)
-target_link_libraries(tvm_runtime PUBLIC tvm_ffi_shared)
-set_property(TARGET tvm_runtime APPEND PROPERTY LINK_OPTIONS "${TVM_NO_UNDEFINED_SYMBOLS}")
-set_property(TARGET tvm_runtime APPEND PROPERTY LINK_OPTIONS "${TVM_VISIBILITY_FLAG}")
-
-# Static consumers must retain registration objects from both archives, with
-# tvm_ffi_static before tvm_runtime_static so FFI type attributes initialize first:
-# -Wl,--whole-archive tvm_ffi_static tvm_runtime_static -Wl,--no-whole-archive
-# Use the corresponding archive-retention flags on other platforms.
-add_library(tvm_runtime_static STATIC
- $<TARGET_OBJECTS:tvm_runtime_objs>
- ${TVM_RUNTIME_EXT_OBJS}
-)
-target_link_libraries(tvm_runtime_static PUBLIC tvm_ffi_static)
-if(MSVC)
- # Static consumers reference ordinary symbols rather than DLL import stubs.
- target_compile_definitions(tvm_runtime_static INTERFACE
- TVM_RUNTIME_DLL= TVM_RUNTIME_DLL_EXPORT= TVM_FFI_DLL= TVM_FFI_DLL_EXPORT=)
-endif()
-
-foreach(runtime_target tvm_runtime tvm_runtime_static)
- target_include_directories(${runtime_target} PUBLIC "$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>")
- target_link_libraries(${runtime_target} PRIVATE ${TVM_RUNTIME_LINKER_LIBS})
- if(USE_THREADS)
- target_link_libraries(${runtime_target} PUBLIC Threads::Threads)
+# Retain the upstream runtime-extra assembly module and its target composition.
+include(cmake/modules/RuntimeExtra.cmake)
+target_link_libraries(tvm_runtime_extra PRIVATE ${TVM_RUNTIME_LINKER_LIBS} ${CMAKE_DL_LIBS})
+add_library(tvm_compiler SHARED ${COMPILER_SRCS})
+target_compile_definitions(tvm_compiler PRIVATE TVM_EXPORTS TVM_FFI_EXPORTS)
+target_link_libraries(tvm_compiler PUBLIC tvm_tirx::tvm_tirx tvm_runtime_extra)
+target_link_libraries(tvm_compiler PRIVATE tvm_tirx::private_headers ${TVM_LINKER_LIBS})
+foreach(_target tvm_compiler tvm_runtime_extra)
+ target_include_directories(${_target} PUBLIC
+ "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>"
+ "$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>")
+ target_link_options(${_target} PRIVATE ${TVM_NO_UNDEFINED_SYMBOLS})
+ set_target_properties(${_target} PROPERTIES
+ LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
+ ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
+ RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib")
+ if(MSVC)
+ foreach(_config RELEASE DEBUG RELWITHDEBINFO MINSIZEREL)
+ set_target_properties(${_target} PROPERTIES
+ LIBRARY_OUTPUT_DIRECTORY_${_config} "${CMAKE_BINARY_DIR}/lib"
+ ARCHIVE_OUTPUT_DIRECTORY_${_config} "${CMAKE_BINARY_DIR}/lib"
+ RUNTIME_OUTPUT_DIRECTORY_${_config} "${CMAKE_BINARY_DIR}/lib")
+ endforeach()
+ endif()
+ tvm_configure_target_library(${_target})
+ if(APPLE)
+ set_property(TARGET ${_target} APPEND PROPERTY INSTALL_RPATH "@loader_path/../../tvm_tirx/lib")
+ elseif(UNIX)
+ set_property(TARGET ${_target} APPEND PROPERTY INSTALL_RPATH "\$ORIGIN/../../tvm_tirx/lib")
endif()
+ tvm_ffi_add_apple_dsymutil(${_target})
endforeach()
-
-set_target_properties(tvm_runtime tvm_runtime_static PROPERTIES
- ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
- LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
- RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
-)
-if(MSVC)
- foreach(config_type Release RELEASE Debug DEBUG RelWithDebInfo RELWITHDEBINFO MinSizeRel MINSIZEREL)
- set_target_properties(tvm_runtime tvm_runtime_static PROPERTIES
- ARCHIVE_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib"
- LIBRARY_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib"
- RUNTIME_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib"
- )
- endforeach()
-endif()
-
-if(TVM_IS_DEBUG_BUILD AND FILE_PREFIX_MAP_SUPPORTED)
- target_compile_options(tvm_runtime PRIVATE $<$<COMPILE_LANGUAGE:CXX>:${FILE_PREFIX_MAP_FLAG}>)
- target_compile_options(tvm_runtime_objs PRIVATE $<$<COMPILE_LANGUAGE:CXX>:${FILE_PREFIX_MAP_FLAG}>)
-endif()
-
-tvm_ffi_add_apple_dsymutil(tvm_runtime)
-tvm_configure_target_library(tvm_runtime)
-add_custom_target(runtime DEPENDS tvm_runtime tvm_runtime_static)
-install(TARGETS tvm_runtime DESTINATION lib${LIB_SUFFIX})
-# Keep wheel contents unchanged; the archive is available to native consumers.
-if(NOT TVM_BUILD_PYTHON_MODULE)
- install(TARGETS tvm_runtime_static DESTINATION lib${LIB_SUFFIX})
-endif()
-add_library(tvm_objs OBJECT ${COMPILER_SRCS})
-target_link_libraries(tvm_objs PUBLIC tvm_ffi_header)
-if(TARGET tvm_llvm_header)
- target_link_libraries(tvm_objs PUBLIC tvm_llvm_header)
-endif()
-if(TVM_VISIBILITY_FLAG)
- target_compile_options(tvm_objs PRIVATE "${TVM_VISIBILITY_FLAG}")
-endif()
-if(MSVC)
- # Compiler objects import runtime symbols from the separate runtime DLL.
- target_compile_definitions(tvm_objs PRIVATE TVM_EXPORTS TVM_FFI_EXPORTS)
-endif()
-
-# Compiler library: built from compiler-only sources.
-# Links against tvm_runtime (and tvm_ffi_shared transitively). Loaded with
-# RTLD_LOCAL so compiler-internal symbols don't leak into the global namespace.
-add_library(tvm_compiler SHARED
- $<TARGET_OBJECTS:tvm_objs>
-)
-target_link_libraries(tvm_compiler PUBLIC tvm_runtime tvm_ffi_shared)
-target_include_directories(tvm_compiler PUBLIC "$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>")
-set_property(TARGET tvm_compiler APPEND PROPERTY LINK_OPTIONS "${TVM_NO_UNDEFINED_SYMBOLS}")
-set_property(TARGET tvm_compiler APPEND PROPERTY LINK_OPTIONS "${TVM_VISIBILITY_FLAG}")
-
-# Work around a GNU ld (binutils) relaxation bug that miscompiles
-# R_X86_64_GOTPCRELX relocations inside very large statically-linked archives.
-# When the full LLVM static libraries are linked into libtvm_compiler.so, the
-# library is large enough that ld can relax an indirect GOT call (LLVM built
-# with -fno-plt emits these) into a direct call with an incorrect displacement.
-# The call then targets read-only data instead of the intended function and
-# crashes at runtime with a SIGSEGV inside llvm::X86Subtarget during code
-# generation. Disabling linker relaxation keeps the GOT-indirect sequences and
-# avoids the miscompilation; it is harmless when LLVM is linked dynamically.
-# See binutils bug ld/25754.
-if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND NOT "${USE_LLVM}" MATCHES ${IS_FALSE_PATTERN})
- set_property(TARGET tvm_compiler APPEND PROPERTY LINK_OPTIONS "-Wl,--no-relax")
- # LLVM's --system-libs reports -lxml2, but TVM calls no libxml2 symbols. The
- # manylinux toolchain does not default to --as-needed, so set it explicitly:
- # the unused libxml2 is then not recorded as a NEEDED dependency, so auditwheel
- # does not vendor it (and its deps) into the wheel.
- set_property(TARGET tvm_compiler APPEND PROPERTY LINK_OPTIONS "-Wl,--as-needed")
-endif()
-
-set_target_properties(tvm_compiler PROPERTIES
- ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
- LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
- RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
-)
-if(MSVC)
- foreach(config_type Release RELEASE Debug DEBUG RelWithDebInfo RELWITHDEBINFO MinSizeRel MINSIZEREL)
- set_target_properties(tvm_compiler PROPERTIES
- ARCHIVE_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib"
- LIBRARY_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib"
- RUNTIME_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib"
- )
- endforeach()
+add_subdirectory(tests/cpp)
+if(GTEST_FOUND)
+ enable_testing()
+ include(CTest)
endif()
-
-if(USE_THREADS)
- target_link_libraries(tvm_compiler PUBLIC Threads::Threads)
+if(TARGET cpptest AND TARGET tirx_cpptest)
+ add_dependencies(cpptest tirx_cpptest)
endif()
-target_link_libraries(tvm_compiler PRIVATE ${TVM_LINKER_LIBS} ${TVM_RUNTIME_LINKER_LIBS})
-
-if(TVM_IS_DEBUG_BUILD)
- target_compile_options(tvm_compiler PRIVATE "${FILE_PREFIX_MAP_FLAG}")
- if(FILE_PREFIX_MAP_SUPPORTED)
- target_compile_options(tvm_compiler PRIVATE $<$<COMPILE_LANGUAGE:CXX>:${FILE_PREFIX_MAP_FLAG}>)
- target_compile_options(tvm_objs PRIVATE $<$<COMPILE_LANGUAGE:CXX>:${FILE_PREFIX_MAP_FLAG}>)
+add_custom_target(runtime DEPENDS tvm_runtime_extra)
+add_dependencies(runtime tvm_tirx::tvm_runtime)
+set(TVM_RUNTIME_BACKEND_LIBS tvm_runtime_extra)
+foreach(_backend cuda opencl vulkan metal rocm)
+ if(TARGET tvm_tirx::tvm_runtime_${_backend})
+ list(APPEND TVM_RUNTIME_BACKEND_LIBS tvm_tirx::tvm_runtime_${_backend})
endif()
-endif()
-
-tvm_ffi_add_apple_dsymutil(tvm_compiler)
-tvm_configure_target_library(tvm_compiler)
-install(TARGETS tvm_compiler DESTINATION lib${LIB_SUFFIX})
-
-# logging option for libbacktrace
-include(cmake/modules/Logging.cmake)
-
+endforeach()
if(USE_CPP_RPC)
- add_subdirectory("apps/cpp_rpc")
+ add_subdirectory(apps/cpp_rpc)
endif()
-
if(USE_CPP_RTVM)
- add_subdirectory("apps/cpp_rtvm")
+ add_subdirectory(apps/cpp_rtvm)
endif()
-
if(USE_IOS_RPC)
- add_subdirectory("apps/ios_rpc")
-endif()
-
-# Set flags for clang
-include(cmake/modules/ClangFlags.cmake)
-# Installation rules
-if (INSTALL_DEV)
- install(
- DIRECTORY "include/" DESTINATION "include"
- FILES_MATCHING
- PATTERN "*.h"
- )
-else(INSTALL_DEV)
- install(
- DIRECTORY "include/tvm/runtime" "include/tvm/support" DESTINATION "include/tvm"
- FILES_MATCHING
- PATTERN "*.h"
- )
-endif(INSTALL_DEV)
-
-include(CMakePackageConfigHelpers)
-set(PROJECT_CONFIG_CONTENT "@PACKAGE_INIT@\n")
-string(APPEND PROJECT_CONFIG_CONTENT "include(CMakeFindDependencyMacro)\n")
-string(APPEND PROJECT_CONFIG_CONTENT "find_dependency(Threads REQUIRED)\n")
-string(APPEND PROJECT_CONFIG_CONTENT
- "include(\"\${CMAKE_CURRENT_LIST_DIR}/${PROJECT_NAME}Targets.cmake\")")
-file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/temp_config_file.cmake" ${PROJECT_CONFIG_CONTENT})
-
-# install(EXPORT ${PROJECT_NAME}Targets
-# NAMESPACE ${PROJECT_NAME}::
-# DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/${PROJECT_NAME})
-
-# Create config for find_package()
-configure_package_config_file(
- "${CMAKE_CURRENT_BINARY_DIR}/temp_config_file.cmake" ${PROJECT_NAME}Config.cmake
- INSTALL_DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/${PROJECT_NAME}")
-
-install(
- FILES
- "${PROJECT_BINARY_DIR}/${PROJECT_NAME}Config.cmake"
- DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/${PROJECT_NAME}")
-
-find_and_set_linker(${USE_ALTERNATIVE_LINKER})
-
-if(${SUMMARIZE})
- print_summary()
-endif()
-
-dump_options_to_file("${TVM_ALL_OPTIONS}")
-
-if(USE_CUDA AND USE_CUTLASS)
- install(TARGETS fpA_intB_gemm EXPORT ${PROJECT_NAME}Targets DESTINATION lib${LIB_SUFFIX})
- install(TARGETS flash_attn EXPORT ${PROJECT_NAME}Targets DESTINATION lib${LIB_SUFFIX})
- # fpA_intB_gemm, fpA_intB_gemm_tvm, and flash_attn are linked by
- # tvm_runtime_extra (see RuntimeExtra.cmake); no link needed here.
-endif()
-
-if(USE_CUDA AND USE_NVTX)
- add_compile_definitions(TVM_NVTX_ENABLED=1)
+ add_subdirectory(apps/ios_rpc)
endif()
-
-# Tests inherit the final project linker flags and compile definitions.
-add_subdirectory(tests/cpp)
-if(GTEST_FOUND)
- # CTest must also be enabled at the project root for subdirectory discovery.
- enable_testing()
- include(CTest)
-endif()
-
-# Python package installation configuration
-# This section ensures that all necessary files are installed for the Python wheel
+install(TARGETS tvm_compiler tvm_runtime_extra
+ ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+ LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+ RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR})
+install(DIRECTORY include/ DESTINATION ${CMAKE_INSTALL_INCLUDEDIR} FILES_MATCHING PATTERN "*.h")
+# Parent minRPC compilation retains its own source and private header closure.
+install(DIRECTORY src/runtime/rpc/ DESTINATION src/runtime/rpc FILES_MATCHING PATTERN "*.h" PATTERN "*.cc")
+install(DIRECTORY src/support/ DESTINATION src/support FILES_MATCHING PATTERN "*.h")
+install(DIRECTORY src/runtime/extra/ DESTINATION src/runtime/extra FILES_MATCHING PATTERN "*.h" PATTERN "*.cuh")
if(TVM_BUILD_PYTHON_MODULE)
- message(STATUS "Configuring Python package installation")
-
- # Install compiled shared libraries into <project>/lib so that
- # tvm_ffi.libinfo.load_lib_ctypes(package="tvm", target_name=...) can find
- # them via the package RECORD or the project's lib/ fallback dir.
- install(TARGETS tvm_compiler DESTINATION "lib")
- install(TARGETS tvm_runtime DESTINATION "lib")
-
- # Install third-party compiled dependencies into the same lib/ dir.
- if(TARGET fpA_intB_gemm)
- tvm_configure_target_library(fpA_intB_gemm)
- install(TARGETS fpA_intB_gemm DESTINATION "lib")
- endif()
- if(TARGET flash_attn)
- tvm_configure_target_library(flash_attn)
- install(TARGETS flash_attn DESTINATION "lib")
- endif()
-
- # Install prebuilt extra runtime libraries into the same lib/ dir. This is how
- # the separately-built CUDA runtime (libtvm_runtime_cuda.so) is bundled: the
- # publishing flow builds it in a CUDA-enabled environment and passes its path
- # via TVM_PACKAGE_EXTRA_LIBS, so it ships through the normal CMake install
- # rather than a post-build wheel rewrite.
- foreach(_extra_lib IN LISTS TVM_PACKAGE_EXTRA_LIBS)
- if(NOT EXISTS "${_extra_lib}")
- message(FATAL_ERROR "TVM_PACKAGE_EXTRA_LIBS entry does not exist: ${_extra_lib}")
- endif()
- install(FILES "${_extra_lib}" DESTINATION "lib")
- endforeach()
-
- # Install minimal header files needed by Python extensions
- install(
- DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/include/tvm/runtime/"
- DESTINATION "include/tvm/runtime/"
- FILES_MATCHING
- PATTERN "*.h"
- )
-
- # Install minimal CMake configuration
- install(
- DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils/"
- DESTINATION "cmake/utils/"
- FILES_MATCHING
- PATTERN "*.cmake"
- )
-
- # Install CUTLASS headers only if available
+ install(DIRECTORY web/ DESTINATION web)
+ install(DIRECTORY licenses/ DESTINATION licenses)
+ install(FILES LICENSE NOTICE DESTINATION .)
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/cutlass/include")
- install(
- DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/cutlass/include/"
- DESTINATION "3rdparty/cutlass/include/"
- FILES_MATCHING
- PATTERN "*.h"
- PATTERN "*.hpp"
- )
+ install(DIRECTORY 3rdparty/cutlass/include/ DESTINATION 3rdparty/cutlass/include
+ FILES_MATCHING PATTERN "*.h" PATTERN "*.hpp")
endif()
-
- # Install minimal source files
- install(
- DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/src/runtime/"
- DESTINATION "src/runtime/"
- FILES_MATCHING
- PATTERN "*.cc"
- PATTERN "*.h"
- )
-
- # Install web package
- install(DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/web/" DESTINATION "web/")
-
- # Install licenses (required for distribution)
- install(
- DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/licenses/"
- DESTINATION "licenses/"
- )
-
- # Install essential metadata files
- install(FILES
- "${CMAKE_CURRENT_SOURCE_DIR}/README.md"
- "${CMAKE_CURRENT_SOURCE_DIR}/LICENSE"
- "${CMAKE_CURRENT_SOURCE_DIR}/NOTICE"
- DESTINATION "."
- )
-
- message(STATUS "Python package installation configured")
endif()
+find_and_set_linker(${USE_ALTERNATIVE_LINKER})
diff --git a/tirx/cmake/modules/CUDA.cmake b/cmake/modules/contrib/CUDA.cmake
similarity index 56%
copy from tirx/cmake/modules/CUDA.cmake
copy to cmake/modules/contrib/CUDA.cmake
index 4e2cdd508c..c4f6d68791 100644
--- a/tirx/cmake/modules/CUDA.cmake
+++ b/cmake/modules/contrib/CUDA.cmake
@@ -15,84 +15,6 @@
# specific language governing permissions and limitations
# under the License.
-# CUDA Module
-find_cuda(${USE_CUDA} ${USE_CUDNN})
-
-if(CUDA_FOUND)
- # always set the includedir when CUDA is available
- # avoid global retrigger of CMake
- include_directories(SYSTEM ${CUDA_INCLUDE_DIRS})
-endif(CUDA_FOUND)
-
-if(USE_CUDA)
- if(NOT CUDA_FOUND)
- message(FATAL_ERROR "Cannot find CUDA, USE_CUDA=" ${USE_CUDA})
- endif()
- message(STATUS "Build with CUDA ${CUDA_VERSION} support")
- enable_language(CUDA)
-
- # Ensure that include directives to NVCC are in the
- # `compile_commands.json`, as required by clangd.
- #
- # As of CMake 3.29.5 [0], if the NVCC version is 11 or higher, CMake
- # will generate a "options-file.rsp" containing the -I flags for
- # include directories, rather than providing them on the
- # command-line. This setting exists to work around the short
- # command-line length limits on Windows, but is enabled on all
- # platforms. If set, because include directories are not part of
- # the `compile_commands.json`, the clangd LSP cannot find the
- # include files.
- #
- # Furthermore, this override cannot be specified in a user's
- # `config.cmake` for TVM, because it must be set after CMake's
- # built-in CUDA support.
- #
- # [0] https://github.com/Kitware/CMake/commit/6377a438
- set(CMAKE_CUDA_USE_RESPONSE_FILE_FOR_INCLUDES 0)
-
- if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
- if(CMAKE_VERSION VERSION_LESS "3.24")
- message(FATAL_ERROR "CMAKE_CUDA_ARCHITECTURES not set. Please upgrade CMake to 3.24 to use native, or set CMAKE_CUDA_ARCHITECTURES manually")
- endif()
- message(STATUS "CMAKE_CUDA_ARCHITECTURES not set, using native")
- set(CMAKE_CUDA_ARCHITECTURES native)
- endif()
-endif(USE_CUDA)
-
-if(USE_CUDA)
- message(STATUS "Build cuda device runtime")
-
- # tvm_runtime_cuda links libcuda; without it the .so silently drops its
- # libcuda.so.1 dependency, so fail configure instead.
- if(NOT CUDA_CUDA_LIBRARY)
- message(FATAL_ERROR "USE_CUDA is on but libcuda was not found. "
- "Set -DCUDA_CUDA_LIBRARY=<path to libcuda.so> or make the driver stub discoverable.")
- endif()
-
- tvm_file_glob(GLOB RUNTIME_CUDA_SRCS src/backend/cuda/runtime/*.cc)
- tvm_file_glob(GLOB VM_CUDA_BUILTIN_SRC_CC src/runtime/vm/cuda/*.cc)
-
- add_library(tvm_runtime_cuda_objs OBJECT ${RUNTIME_CUDA_SRCS} ${VM_CUDA_BUILTIN_SRC_CC})
- target_link_libraries(tvm_runtime_cuda_objs PUBLIC tvm_ffi_header)
- # These sources compile into tvm_runtime_cuda.dll, so their TVM_RUNTIME_DLL /
- # TVM_FFI_DLL symbols must be dllexport on MSVC (e.g. GetCudaDeviceCount in
- # cuda_device_api.cc). Mirror tvm_runtime_objs; a no-op on non-MSVC platforms.
- target_compile_definitions(tvm_runtime_cuda_objs PRIVATE TVM_RUNTIME_EXPORTS TVM_FFI_EXPORTS)
- set_target_properties(tvm_runtime_cuda_objs PROPERTIES POSITION_INDEPENDENT_CODE ON)
- if(TVM_VISIBILITY_FLAG)
- target_compile_options(tvm_runtime_cuda_objs PRIVATE "${TVM_VISIBILITY_FLAG}")
- endif()
- add_library(tvm_runtime_cuda SHARED $<TARGET_OBJECTS:tvm_runtime_cuda_objs>)
- list(APPEND TVM_RUNTIME_BACKEND_LIBS tvm_runtime_cuda)
- target_link_libraries(tvm_runtime_cuda PUBLIC tvm_runtime PRIVATE ${CUDA_CUDART_LIBRARY} ${CUDA_CUDA_LIBRARY})
- tvm_configure_target_library(tvm_runtime_cuda RUNTIME_MODULE)
-
- if(USE_NVTX)
- message(STATUS "Build with NVTX support")
- target_link_libraries(tvm_runtime_cuda PRIVATE ${CUDA_NVTX_LIBRARY})
- endif()
-endif(USE_CUDA)
-
# Contrib sources gated by USE_CUDA go into libtvm_runtime_extra.
# See the RuntimeExtra assembly block in CMakeLists.txt.
diff --git a/tirx/cmake/modules/ROCM.cmake b/cmake/modules/contrib/ROCM.cmake
similarity index 64%
copy from tirx/cmake/modules/ROCM.cmake
copy to cmake/modules/contrib/ROCM.cmake
index 33abc75a2b..37ee066099 100644
--- a/tirx/cmake/modules/ROCM.cmake
+++ b/cmake/modules/contrib/ROCM.cmake
@@ -15,42 +15,6 @@
# specific language governing permissions and limitations
# under the License.
-# ROCM Module
-find_rocm(${USE_ROCM})
-
-if(ROCM_FOUND)
- # always set the includedir
- # avoid global retrigger of CMake
- include_directories(SYSTEM ${ROCM_INCLUDE_DIRS})
- add_definitions(-D__HIP_PLATFORM_HCC__=1)
- add_definitions(-D__HIP_PLATFORM_AMD__=1)
-endif(ROCM_FOUND)
-
-if(USE_ROCM)
- if(NOT ROCM_FOUND)
- message(FATAL_ERROR "Cannot find ROCM, USE_ROCM=" ${USE_ROCM})
- endif()
- message(STATUS "Build rocm device runtime")
-
- tvm_file_glob(GLOB RUNTIME_ROCM_SRCS src/backend/rocm/runtime/*.cc)
-
- set(_rocm_libs ${ROCM_HIPHCC_LIBRARY})
- if(ROCM_HSA_LIBRARY)
- list(APPEND _rocm_libs ${ROCM_HSA_LIBRARY})
- endif()
-
- add_library(tvm_runtime_rocm_objs OBJECT ${RUNTIME_ROCM_SRCS})
- target_link_libraries(tvm_runtime_rocm_objs PUBLIC tvm_ffi_header)
- set_target_properties(tvm_runtime_rocm_objs PROPERTIES POSITION_INDEPENDENT_CODE ON)
- if(TVM_VISIBILITY_FLAG)
- target_compile_options(tvm_runtime_rocm_objs PRIVATE "${TVM_VISIBILITY_FLAG}")
- endif()
- add_library(tvm_runtime_rocm SHARED $<TARGET_OBJECTS:tvm_runtime_rocm_objs>)
- list(APPEND TVM_RUNTIME_BACKEND_LIBS tvm_runtime_rocm)
- target_link_libraries(tvm_runtime_rocm PUBLIC tvm_runtime PRIVATE ${_rocm_libs})
- tvm_configure_target_library(tvm_runtime_rocm RUNTIME_MODULE)
-endif(USE_ROCM)
-
# HIPBLAS contrib goes into libtvm_runtime_extra.
if(USE_ROCM AND USE_HIPBLAS)
message(STATUS "Build with HIPBLAS support")
diff --git a/tests/cpp/CMakeLists.txt b/tests/cpp/CMakeLists.txt
index 35c9fadb9e..45be2dee75 100644
--- a/tests/cpp/CMakeLists.txt
+++ b/tests/cpp/CMakeLists.txt
@@ -16,7 +16,7 @@
# under the License.
# Enable ctest if gtest is available
-if(USE_GTEST)
+if(USE_GTEST AND NOT BUILD_TESTING STREQUAL "OFF")
# Check env var for backward compatibility. A better way to specify package
# locations is to use CMAKE_PREFIX_PATH or other standard CMake mechanism
# (see CMake documentation for `find_package`).
diff --git a/CMakeLists.txt b/tirx/CMakeLists.txt
similarity index 57%
copy from CMakeLists.txt
copy to tirx/CMakeLists.txt
index bfa62d6b9d..0b508905ee 100644
--- a/CMakeLists.txt
+++ b/tirx/CMakeLists.txt
@@ -1,32 +1,42 @@
-cmake_minimum_required(VERSION 3.18)
-project(tvm C CXX)
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
-# --- TVM version (no dependency on version.py) ---
-# When built via scikit-build-core the version is resolved by setuptools_scm and
-# passed in as SKBUILD_PROJECT_VERSION_FULL; bake it into the C++ TVM_VERSION macro.
-# A bare `cmake` build (no scikit-build-core) leaves the checked-in default in
-# include/tvm/runtime/base.h untouched. An explicit -DTVM_VERSION always wins.
-if(NOT DEFINED TVM_VERSION AND DEFINED SKBUILD_PROJECT_VERSION_FULL)
- set(TVM_VERSION "${SKBUILD_PROJECT_VERSION_FULL}")
-endif()
-if(DEFINED TVM_VERSION)
- message(STATUS "TVM_VERSION=${TVM_VERSION}")
- add_compile_definitions(TVM_VERSION="${TVM_VERSION}")
+cmake_minimum_required(VERSION 3.18)
+project(tvm_tirx C CXX)
+include(GNUInstallDirs)
+if(TVM_BUILD_PYTHON_MODULE)
+ set(CMAKE_INSTALL_LIBDIR lib)
endif()
+# The extracted package has an independent version; parent SCM tags do not apply.
+set(TVM_TIRX_VERSION "0.1.0a0")
+set(TVM_VERSION "${TVM_TIRX_VERSION}")
+add_compile_definitions(TVM_VERSION="${TVM_VERSION}")
+
# Utility functions
-include(cmake/utils/Utils.cmake)
-include(cmake/utils/Summary.cmake)
-include(cmake/utils/Linker.cmake)
-include(cmake/utils/Library.cmake)
-include(cmake/utils/FindCUDA.cmake)
-include(cmake/utils/FindNCCL.cmake)
-include(cmake/utils/FindOpenCL.cmake)
-include(cmake/utils/FindVulkan.cmake)
-include(cmake/utils/FindLLVM.cmake)
-include(cmake/utils/FindROCM.cmake)
-include(cmake/utils/FindRCCL.cmake)
-include(cmake/utils/FindNVSHMEM.cmake)
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils/Utils.cmake")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils/Summary.cmake")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils/Linker.cmake")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils/Library.cmake")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils/FindCUDA.cmake")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils/FindOpenCL.cmake")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils/FindVulkan.cmake")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils/FindLLVM.cmake")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils/FindROCM.cmake")
if(EXISTS ${CMAKE_BINARY_DIR}/config.cmake)
include(${CMAKE_BINARY_DIR}/config.cmake)
@@ -41,7 +51,6 @@ endif()
# and add set(OPTION VALUE) to override these build options.
# Alernatively, use cmake -DOPTION=VALUE through command-line.
tvm_option(USE_CUDA "Build with CUDA" OFF)
-tvm_option(USE_NCCL "Build with NCCL" OFF)
tvm_option(USE_OPENCL "Build with OpenCL" OFF)
tvm_option(USE_OPENCL_ENABLE_HOST_PTR "Enable OpenCL memory object access to host" OFF)
tvm_option(USE_OPENCL_GTEST "Path to OpenCL specific gtest version for runtime cpp tests." /path/to/opencl/gtest)
@@ -58,11 +67,9 @@ tvm_option(USE_KHRONOS_SPIRV "Whether to use spirv-tools and SPIRV-Headers from
tvm_option(USE_SPIRV_KHR_INTEGER_DOT_PRODUCT "whether enable SPIRV_KHR_DOT_PRODUCT" OFF)
tvm_option(USE_METAL "Build with Metal" OFF)
tvm_option(USE_ROCM "Build with ROCM" OFF)
-tvm_option(USE_RCCL "Build with RCCL" OFF)
tvm_option(ROCM_PATH "The path to rocm" /opt/rocm)
tvm_option(USE_HEXAGON_QHL "Enable Hexagon QHL compiler lowering (requires external runtime wrappers)" OFF)
-tvm_option(USE_RPC "Build with RPC" ON)
tvm_option(USE_THREADS "Build with thread support" ON)
tvm_option(USE_LLVM "Build with LLVM, can be set to specific llvm-config path" OFF)
tvm_option(USE_OPENMP "Build with OpenMP thread pool implementation" OFF)
@@ -78,35 +85,11 @@ tvm_option(USE_ALTERNATIVE_LINKER "Use 'mold' or 'lld' if found when invoking co
tvm_option(USE_CCACHE "Use ccache if found when invoking compiler" AUTO)
# 3rdparty libraries
-tvm_option(COMPILER_RT_PATH "Path to COMPILER-RT" "3rdparty/compiler-rt")
+tvm_option(COMPILER_RT_PATH "Path to COMPILER-RT" "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/compiler-rt")
# Contrib library options
-tvm_option(USE_BLAS "The blas library to be linked" none)
-tvm_option(USE_AMX "Enable Intel AMX" OFF)
tvm_option(USE_Z3 "Build with Z3 SMT solver support" AUTO)
-tvm_option(USE_MKL "MKL root path when use MKL blas" OFF)
-tvm_option(USE_DNNL "Enable DNNL codegen" OFF)
-tvm_option(USE_CUDNN "Build with cuDNN" OFF)
-tvm_option(USE_CUBLAS "Build with cuBLAS" OFF)
tvm_option(USE_NVTX "Build with NVTX" OFF)
-tvm_option(USE_CUTLASS "Build with CUTLASS" OFF)
-tvm_option(USE_THRUST "Build with Thrust" OFF)
-tvm_option(USE_CURAND "Build with cuRAND" OFF)
-tvm_option(USE_HIPBLAS "Build with ROCM:HIPBLAS" OFF)
-tvm_option(USE_SORT "Build with sort support" ON)
-tvm_option(USE_RANDOM "Build with random support" ON)
-tvm_option(USE_CPP_RPC "Build CPP RPC" OFF)
-tvm_option(USE_IOS_RPC "Build iOS RPC" OFF)
-tvm_option(USE_COREML "Build with coreml support" OFF)
-tvm_option(USE_TENSORRT_CODEGEN "Build with TensorRT Codegen support" OFF)
-tvm_option(USE_TENSORRT_RUNTIME "Build with TensorRT runtime" OFF)
-tvm_option(USE_NNAPI_CODEGEN "Build with NNAPI Codegen support" OFF)
-tvm_option(USE_NNAPI_RUNTIME "Build with NNAPI runtime" OFF)
-tvm_option(USE_EXAMPLE_NPU_CODEGEN "Build with Example NPU Codegen support" OFF)
-tvm_option(USE_EXAMPLE_NPU_RUNTIME "Build with Example NPU runtime" OFF)
tvm_option(SUMMARIZE "Print CMake option summary after configuring" OFF)
-tvm_option(USE_CLML "Build with CLML Codegen support" OFF)
-tvm_option(USE_CLML_GRAPH_EXECUTOR "Build with CLML graph runtime" OFF)
-tvm_option(USE_NVSHMEM "Build with NVSHMEM support" OFF)
# Python package options
tvm_option(TVM_BUILD_PYTHON_MODULE "Build Python module with scikit-build-core" OFF)
@@ -120,19 +103,24 @@ include_directories(SYSTEM ${COMPILER_RT_PATH})
set(TVM_LINKER_LIBS "")
set(TVM_RUNTIME_LINKER_LIBS "")
set(TVM_RUNTIME_BACKEND_LIBS "")
-# Early target creation so contrib cmake files can call
-# target_link_libraries(tvm_runtime_extra PRIVATE <object_lib>) directly.
-add_library(tvm_runtime_extra SHARED)
-list(APPEND TVM_RUNTIME_BACKEND_LIBS tvm_runtime_extra)
-set_target_properties(tvm_runtime_extra PROPERTIES LINKER_LANGUAGE CXX)
-# INTERFACE target carrying compile definitions for OBJECT libs that build
-# into tvm_runtime_extra. On MSVC, TVM_RUNTIME_EXTRA_EXPORTS exports extra-owned
-# APIs while core-runtime APIs remain imports from tvm_runtime.dll.
-add_library(tvm_runtime_extra_defs INTERFACE)
-target_link_libraries(tvm_runtime_extra_defs INTERFACE tvm_ffi_header)
-target_compile_definitions(tvm_runtime_extra_defs
- INTERFACE TVM_RUNTIME_EXTRA_EXPORTS TVM_FFI_EXPORTS)
-
+# Use one external FFI registry for both the core and its parent consumers.
+find_package(Python COMPONENTS Interpreter REQUIRED)
+if(NOT tvm_ffi_DIR)
+ execute_process(COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir
+ OUTPUT_VARIABLE tvm_ffi_DIR OUTPUT_STRIP_TRAILING_WHITESPACE
+ COMMAND_ERROR_IS_FATAL ANY)
+endif()
+if(NOT TARGET tvm_ffi::shared)
+ find_package(tvm_ffi CONFIG REQUIRED)
+ foreach(_ffi_target header shared static)
+ if(TARGET tvm_ffi::${_ffi_target})
+ set_target_properties(tvm_ffi::${_ffi_target} PROPERTIES IMPORTED_GLOBAL TRUE)
+ endif()
+ endforeach()
+endif()
+if(NOT TARGET tvm_ffi_header)
+ add_library(tvm_ffi_header ALIAS tvm_ffi::header)
+endif()
# Check if this is being run on its own or as a subdirectory for another project
# If we update to CMake 2.21+, we can use PROJECT_IS_TOP_LEVEL instead
@@ -267,11 +255,7 @@ assign_source_group("Include" ${GROUP_INCLUDE})
tvm_file_glob(GLOB_RECURSE COMPILER_SRCS
src/ir/*.cc
src/sym/*.cc
- src/te/*.cc
src/tirx/*.cc
- src/s_tir/*.cc
- src/topi/*.cc
- src/driver/*.cc
src/support/*.cc
# Shared printer and TVMScript IR builder. Per-dialect printer/builder pieces
# live under src/<dialect>/script/. List shared files explicitly so a dialect
@@ -292,20 +276,6 @@ tvm_file_glob(GLOB_RECURSE COMPILER_SRCS
src/script/printer/printer.cc
src/script/printer/doc_printer.cc
src/script/printer/script_printer.cc
- src/relax/ir/*.cc
- src/relax/op/*.cc
- src/relax/analysis/*.cc
- src/relax/transform/*.cc
- src/relax/backend/vm/*.cc
- src/relax/backend/adreno/*.cc
- src/relax/backend/task_extraction.cc
- src/relax/backend/pattern_registry.cc
- src/relax/utils.cc
- src/relax/distributed/*.cc
- src/relax/distributed/transform/*.cc
- src/relax/op/distributed/*.cc
- src/relax/script/*.cc
- src/relax/testing/*.cc
)
tvm_file_glob(GLOB CODEGEN_SRCS
@@ -339,21 +309,9 @@ list(APPEND COMPILER_SRCS ${CODEGEN_SRCS})
# Core runtime sources, extended by the backend modules.
tvm_file_glob(GLOB RUNTIME_SRCS
src/runtime/*.cc
- src/runtime/rpc/minrpc/*.cc
)
-# VM-owned runtime implementations, including external-module constant loading.
-tvm_file_glob(GLOB RUNTIME_VM_SRCS src/runtime/vm/*.cc)
-list(APPEND RUNTIME_SRCS ${RUNTIME_VM_SRCS})
-# Note: src/runtime/extra/disco/** moves to libtvm_runtime_extra.
-# Note: src/backend/*/runtime sources move to per-backend DSOs.
set(TVM_RUNTIME_EXT_OBJS "")
-if(USE_RPC)
- message(STATUS "Build with RPC support...")
- tvm_file_glob(GLOB RUNTIME_RPC_SRCS src/runtime/rpc/*.cc)
- list(APPEND RUNTIME_SRCS ${RUNTIME_RPC_SRCS})
-endif(USE_RPC)
-
# Package runtime rules
if(NOT USE_RTTI)
endif()
@@ -370,7 +328,7 @@ endif(USE_KALLOC_ALIGNMENT)
# Caches the build.
# Note that ccache-3.x doesn't support nvcc well, so CUDA kernels may never hit the cache and still
# need to be re-compiled every time. Using ccache 4.0+ can resolve this issue.
-include(cmake/utils/CCache.cmake)
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils/CCache.cmake")
include(CheckCXXCompilerFlag)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
@@ -379,30 +337,16 @@ set(CMAKE_CUDA_STANDARD_REQUIRED ON)
set(CMAKE_CUDA_STANDARD 20)
# Module rules
-include(cmake/modules/CUDA.cmake)
-include(cmake/modules/contrib/CLML.cmake) # Must be before OpenCL.cmake
-include(cmake/modules/OpenCL.cmake)
-include(cmake/modules/OpenMP.cmake)
-include(cmake/modules/Vulkan.cmake)
-include(cmake/modules/Metal.cmake)
-include(cmake/modules/ROCM.cmake)
-include(cmake/modules/LLVM.cmake)
-include(cmake/modules/contrib/BLAS.cmake)
-include(cmake/modules/contrib/DNNL.cmake)
-include(cmake/modules/contrib/AMX.cmake)
-include(cmake/modules/contrib/CUTLASS.cmake)
-include(cmake/modules/contrib/Random.cmake)
-include(cmake/modules/contrib/Sort.cmake)
-include(cmake/modules/contrib/Z3.cmake)
-include(cmake/modules/contrib/CoreML.cmake)
-include(cmake/modules/contrib/TensorRT.cmake)
-include(cmake/modules/contrib/NNAPI.cmake)
-include(cmake/modules/contrib/ExampleNPU.cmake)
-include(cmake/modules/contrib/vllm.cmake)
-include(cmake/modules/Git.cmake)
-
-include(cmake/modules/RuntimeExtra.cmake)
-add_subdirectory(3rdparty/tvm-ffi)
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/CUDA.cmake")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/OpenCL.cmake")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/OpenMP.cmake")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/Vulkan.cmake")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/Metal.cmake")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/ROCM.cmake")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/LLVM.cmake")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/contrib/Z3.cmake")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/Git.cmake")
+
include(GNUInstallDirs)
if(USE_THREADS)
@@ -435,41 +379,50 @@ add_library(tvm_runtime SHARED
$<TARGET_OBJECTS:tvm_runtime_objs>
${TVM_RUNTIME_EXT_OBJS}
)
-target_link_libraries(tvm_runtime PUBLIC tvm_ffi_shared)
+target_link_libraries(tvm_runtime PUBLIC tvm_ffi::shared)
set_property(TARGET tvm_runtime APPEND PROPERTY LINK_OPTIONS "${TVM_NO_UNDEFINED_SYMBOLS}")
set_property(TARGET tvm_runtime APPEND PROPERTY LINK_OPTIONS "${TVM_VISIBILITY_FLAG}")
# Static consumers must retain registration objects from both archives, with
-# tvm_ffi_static before tvm_runtime_static so FFI type attributes initialize first:
-# -Wl,--whole-archive tvm_ffi_static tvm_runtime_static -Wl,--no-whole-archive
+# tvm_ffi::static before tvm_runtime_static so FFI type attributes initialize first:
+# -Wl,--whole-archive tvm_ffi::static tvm_runtime_static -Wl,--no-whole-archive
# Use the corresponding archive-retention flags on other platforms.
+set(_ffi_static_dependency tvm_ffi::shared)
+if(TARGET tvm_ffi::static)
+ set(_ffi_static_dependency tvm_ffi::static)
+endif()
add_library(tvm_runtime_static STATIC
$<TARGET_OBJECTS:tvm_runtime_objs>
${TVM_RUNTIME_EXT_OBJS}
)
-target_link_libraries(tvm_runtime_static PUBLIC tvm_ffi_static)
+target_link_libraries(tvm_runtime_static PUBLIC ${_ffi_static_dependency})
if(MSVC)
# Static consumers reference ordinary symbols rather than DLL import stubs.
target_compile_definitions(tvm_runtime_static INTERFACE
- TVM_RUNTIME_DLL= TVM_RUNTIME_DLL_EXPORT= TVM_FFI_DLL= TVM_FFI_DLL_EXPORT=)
+ TVM_RUNTIME_DLL= TVM_RUNTIME_DLL_EXPORT=)
+ if(TARGET tvm_ffi::static)
+ target_compile_definitions(tvm_runtime_static INTERFACE TVM_FFI_DLL= TVM_FFI_DLL_EXPORT=)
+ endif()
endif()
-foreach(runtime_target tvm_runtime tvm_runtime_static)
- target_include_directories(${runtime_target} PUBLIC "$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>")
+set(_runtime_targets tvm_runtime tvm_runtime_static)
+
+foreach(runtime_target IN LISTS _runtime_targets)
+ target_include_directories(${runtime_target} PUBLIC "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>" "$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>")
target_link_libraries(${runtime_target} PRIVATE ${TVM_RUNTIME_LINKER_LIBS})
if(USE_THREADS)
target_link_libraries(${runtime_target} PUBLIC Threads::Threads)
endif()
endforeach()
-set_target_properties(tvm_runtime tvm_runtime_static PROPERTIES
+set_target_properties(${_runtime_targets} PROPERTIES
ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
)
if(MSVC)
foreach(config_type Release RELEASE Debug DEBUG RelWithDebInfo RELWITHDEBINFO MinSizeRel MINSIZEREL)
- set_target_properties(tvm_runtime tvm_runtime_static PROPERTIES
+ set_target_properties(${_runtime_targets} PROPERTIES
ARCHIVE_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib"
LIBRARY_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib"
RUNTIME_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib"
@@ -484,39 +437,41 @@ endif()
tvm_ffi_add_apple_dsymutil(tvm_runtime)
tvm_configure_target_library(tvm_runtime)
-add_custom_target(runtime DEPENDS tvm_runtime tvm_runtime_static)
+if(NOT IS_SUBPROJECT)
+ add_custom_target(runtime DEPENDS ${_runtime_targets})
+endif()
install(TARGETS tvm_runtime DESTINATION lib${LIB_SUFFIX})
# Keep wheel contents unchanged; the archive is available to native consumers.
-if(NOT TVM_BUILD_PYTHON_MODULE)
+if(TARGET tvm_runtime_static AND NOT TVM_BUILD_PYTHON_MODULE)
install(TARGETS tvm_runtime_static DESTINATION lib${LIB_SUFFIX})
endif()
-add_library(tvm_objs OBJECT ${COMPILER_SRCS})
-target_link_libraries(tvm_objs PUBLIC tvm_ffi_header)
+add_library(tvm_tirx_objs OBJECT ${COMPILER_SRCS})
+target_link_libraries(tvm_tirx_objs PUBLIC tvm_ffi_header)
if(TARGET tvm_llvm_header)
- target_link_libraries(tvm_objs PUBLIC tvm_llvm_header)
+ target_link_libraries(tvm_tirx_objs PUBLIC tvm_llvm_header)
endif()
if(TVM_VISIBILITY_FLAG)
- target_compile_options(tvm_objs PRIVATE "${TVM_VISIBILITY_FLAG}")
+ target_compile_options(tvm_tirx_objs PRIVATE "${TVM_VISIBILITY_FLAG}")
endif()
if(MSVC)
# Compiler objects import runtime symbols from the separate runtime DLL.
- target_compile_definitions(tvm_objs PRIVATE TVM_EXPORTS TVM_FFI_EXPORTS)
+ target_compile_definitions(tvm_tirx_objs PRIVATE TVM_TIRX_EXPORTS TVM_FFI_EXPORTS)
endif()
# Compiler library: built from compiler-only sources.
-# Links against tvm_runtime (and tvm_ffi_shared transitively). Loaded with
+# Links against tvm_runtime (and tvm_ffi::shared transitively). Loaded with
# RTLD_LOCAL so compiler-internal symbols don't leak into the global namespace.
-add_library(tvm_compiler SHARED
- $<TARGET_OBJECTS:tvm_objs>
+add_library(tvm_tirx SHARED
+ $<TARGET_OBJECTS:tvm_tirx_objs>
)
-target_link_libraries(tvm_compiler PUBLIC tvm_runtime tvm_ffi_shared)
-target_include_directories(tvm_compiler PUBLIC "$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>")
-set_property(TARGET tvm_compiler APPEND PROPERTY LINK_OPTIONS "${TVM_NO_UNDEFINED_SYMBOLS}")
-set_property(TARGET tvm_compiler APPEND PROPERTY LINK_OPTIONS "${TVM_VISIBILITY_FLAG}")
+target_link_libraries(tvm_tirx PUBLIC tvm_runtime tvm_ffi::shared)
+target_include_directories(tvm_tirx PUBLIC "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>" "$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>")
+set_property(TARGET tvm_tirx APPEND PROPERTY LINK_OPTIONS "${TVM_NO_UNDEFINED_SYMBOLS}")
+set_property(TARGET tvm_tirx APPEND PROPERTY LINK_OPTIONS "${TVM_VISIBILITY_FLAG}")
# Work around a GNU ld (binutils) relaxation bug that miscompiles
# R_X86_64_GOTPCRELX relocations inside very large statically-linked archives.
-# When the full LLVM static libraries are linked into libtvm_compiler.so, the
+# When the full LLVM static libraries are linked into libtvm_tirx.so, the
# library is large enough that ld can relax an indirect GOT call (LLVM built
# with -fno-plt emits these) into a direct call with an incorrect displacement.
# The call then targets read-only data instead of the intended function and
@@ -525,22 +480,22 @@ set_property(TARGET tvm_compiler APPEND PROPERTY LINK_OPTIONS "${TVM_VISIBILITY_
# avoids the miscompilation; it is harmless when LLVM is linked dynamically.
# See binutils bug ld/25754.
if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND NOT "${USE_LLVM}" MATCHES ${IS_FALSE_PATTERN})
- set_property(TARGET tvm_compiler APPEND PROPERTY LINK_OPTIONS "-Wl,--no-relax")
+ set_property(TARGET tvm_tirx APPEND PROPERTY LINK_OPTIONS "-Wl,--no-relax")
# LLVM's --system-libs reports -lxml2, but TVM calls no libxml2 symbols. The
# manylinux toolchain does not default to --as-needed, so set it explicitly:
# the unused libxml2 is then not recorded as a NEEDED dependency, so auditwheel
# does not vendor it (and its deps) into the wheel.
- set_property(TARGET tvm_compiler APPEND PROPERTY LINK_OPTIONS "-Wl,--as-needed")
+ set_property(TARGET tvm_tirx APPEND PROPERTY LINK_OPTIONS "-Wl,--as-needed")
endif()
-set_target_properties(tvm_compiler PROPERTIES
+set_target_properties(tvm_tirx PROPERTIES
ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
)
if(MSVC)
foreach(config_type Release RELEASE Debug DEBUG RelWithDebInfo RELWITHDEBINFO MinSizeRel MINSIZEREL)
- set_target_properties(tvm_compiler PROPERTIES
+ set_target_properties(tvm_tirx PROPERTIES
ARCHIVE_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib"
LIBRARY_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib"
RUNTIME_OUTPUT_DIRECTORY_${config_type} "${CMAKE_BINARY_DIR}/lib"
@@ -549,188 +504,105 @@ if(MSVC)
endif()
if(USE_THREADS)
- target_link_libraries(tvm_compiler PUBLIC Threads::Threads)
+ target_link_libraries(tvm_tirx PUBLIC Threads::Threads)
endif()
-target_link_libraries(tvm_compiler PRIVATE ${TVM_LINKER_LIBS} ${TVM_RUNTIME_LINKER_LIBS})
+target_link_libraries(tvm_tirx PRIVATE ${TVM_LINKER_LIBS} ${TVM_RUNTIME_LINKER_LIBS})
if(TVM_IS_DEBUG_BUILD)
- target_compile_options(tvm_compiler PRIVATE "${FILE_PREFIX_MAP_FLAG}")
+ target_compile_options(tvm_tirx PRIVATE "${FILE_PREFIX_MAP_FLAG}")
if(FILE_PREFIX_MAP_SUPPORTED)
- target_compile_options(tvm_compiler PRIVATE $<$<COMPILE_LANGUAGE:CXX>:${FILE_PREFIX_MAP_FLAG}>)
- target_compile_options(tvm_objs PRIVATE $<$<COMPILE_LANGUAGE:CXX>:${FILE_PREFIX_MAP_FLAG}>)
+ target_compile_options(tvm_tirx PRIVATE $<$<COMPILE_LANGUAGE:CXX>:${FILE_PREFIX_MAP_FLAG}>)
+ target_compile_options(tvm_tirx_objs PRIVATE $<$<COMPILE_LANGUAGE:CXX>:${FILE_PREFIX_MAP_FLAG}>)
endif()
endif()
-tvm_ffi_add_apple_dsymutil(tvm_compiler)
-tvm_configure_target_library(tvm_compiler)
-install(TARGETS tvm_compiler DESTINATION lib${LIB_SUFFIX})
+tvm_ffi_add_apple_dsymutil(tvm_tirx)
+tvm_configure_target_library(tvm_tirx)
+install(TARGETS tvm_tirx DESTINATION lib${LIB_SUFFIX})
# logging option for libbacktrace
-include(cmake/modules/Logging.cmake)
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/Logging.cmake")
-if(USE_CPP_RPC)
- add_subdirectory("apps/cpp_rpc")
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/ClangFlags.cmake")
+if(USE_CUDA AND USE_NVTX)
+ add_compile_definitions(TVM_NVTX_ENABLED=1)
+ target_compile_definitions(tvm_runtime INTERFACE TVM_NVTX_ENABLED=1)
endif()
-
-if(USE_CPP_RTVM)
- add_subdirectory("apps/cpp_rtvm")
+if(INDEX_DEFAULT_I64)
+ target_compile_definitions(tvm_runtime INTERFACE TVM_INDEX_DEFAULT_I64=1)
endif()
-
-if(USE_IOS_RPC)
- add_subdirectory("apps/ios_rpc")
+if(USE_KALLOC_ALIGNMENT)
+ target_compile_definitions(tvm_runtime INTERFACE TVM_KALLOC_ALIGNMENT=${USE_KALLOC_ALIGNMENT})
endif()
-
-# Set flags for clang
-include(cmake/modules/ClangFlags.cmake)
-# Installation rules
-if (INSTALL_DEV)
- install(
- DIRECTORY "include/" DESTINATION "include"
- FILES_MATCHING
- PATTERN "*.h"
- )
-else(INSTALL_DEV)
- install(
- DIRECTORY "include/tvm/runtime" "include/tvm/support" DESTINATION "include/tvm"
- FILES_MATCHING
- PATTERN "*.h"
- )
-endif(INSTALL_DEV)
-
-include(CMakePackageConfigHelpers)
-set(PROJECT_CONFIG_CONTENT "@PACKAGE_INIT@\n")
-string(APPEND PROJECT_CONFIG_CONTENT "include(CMakeFindDependencyMacro)\n")
-string(APPEND PROJECT_CONFIG_CONTENT "find_dependency(Threads REQUIRED)\n")
-string(APPEND PROJECT_CONFIG_CONTENT
- "include(\"\${CMAKE_CURRENT_LIST_DIR}/${PROJECT_NAME}Targets.cmake\")")
-file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/temp_config_file.cmake" ${PROJECT_CONFIG_CONTENT})
-
-# install(EXPORT ${PROJECT_NAME}Targets
-# NAMESPACE ${PROJECT_NAME}::
-# DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/${PROJECT_NAME})
-
-# Create config for find_package()
-configure_package_config_file(
- "${CMAKE_CURRENT_BINARY_DIR}/temp_config_file.cmake" ${PROJECT_NAME}Config.cmake
- INSTALL_DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/${PROJECT_NAME}")
-
-install(
- FILES
- "${PROJECT_BINARY_DIR}/${PROJECT_NAME}Config.cmake"
- DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/${PROJECT_NAME}")
-
-find_and_set_linker(${USE_ALTERNATIVE_LINKER})
-
-if(${SUMMARIZE})
- print_summary()
+if(TVM_DEBUG_WITH_ABI_CHANGE)
+ target_compile_definitions(tvm_runtime PUBLIC TVM_DEBUG_WITH_ABI_CHANGE)
+ target_compile_definitions(tvm_runtime_objs PRIVATE TVM_DEBUG_WITH_ABI_CHANGE)
+ target_compile_definitions(tvm_tirx_objs PRIVATE TVM_DEBUG_WITH_ABI_CHANGE)
endif()
-dump_options_to_file("${TVM_ALL_OPTIONS}")
-
-if(USE_CUDA AND USE_CUTLASS)
- install(TARGETS fpA_intB_gemm EXPORT ${PROJECT_NAME}Targets DESTINATION lib${LIB_SUFFIX})
- install(TARGETS flash_attn EXPORT ${PROJECT_NAME}Targets DESTINATION lib${LIB_SUFFIX})
- # fpA_intB_gemm, fpA_intB_gemm_tvm, and flash_attn are linked by
- # tvm_runtime_extra (see RuntimeExtra.cmake); no link needed here.
+get_target_property(_runtime_compile_definitions tvm_runtime INTERFACE_COMPILE_DEFINITIONS)
+if(_runtime_compile_definitions)
+ target_compile_definitions(tvm_runtime_static INTERFACE ${_runtime_compile_definitions})
endif()
-
-if(USE_CUDA AND USE_NVTX)
- add_compile_definitions(TVM_NVTX_ENABLED=1)
+add_library(tvm_tirx_private_headers INTERFACE)
+set_target_properties(tvm_tirx_private_headers PROPERTIES EXPORT_NAME private_headers)
+target_include_directories(tvm_tirx_private_headers INTERFACE
+ "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/src>"
+ "$<BUILD_INTERFACE:${COMPILER_RT_PATH}>"
+ "$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}/tvm_tirx/private>"
+ "$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}/tvm_tirx/compiler-rt>")
+set(_tirx_targets tvm_tirx ${_runtime_targets} tvm_tirx_private_headers ${TVM_RUNTIME_BACKEND_LIBS})
+foreach(_target IN LISTS _tirx_targets)
+ get_target_property(_export_name ${_target} EXPORT_NAME)
+ if(NOT _export_name)
+ set(_export_name ${_target})
+ endif()
+ add_library(tvm_tirx::${_export_name} ALIAS ${_target})
+endforeach()
+install(TARGETS ${_tirx_targets} EXPORT tvm_tirxTargets
+ ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+ LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+ RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR})
+install(DIRECTORY include/ DESTINATION ${CMAKE_INSTALL_INCLUDEDIR} FILES_MATCHING PATTERN "*.h")
+# CLML's retained OpenCL interface needs this bounded private include closure.
+foreach(_header backend/opencl/memory_scope.h
+ backend/opencl/runtime/opencl_common.h backend/opencl/runtime/texture.h
+ runtime/metadata.h runtime/pack_args.h runtime/workspace_pool.h)
+ get_filename_component(_dir "${_header}" DIRECTORY)
+ install(FILES "src/${_header}" DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}/tvm_tirx/private/${_dir}")
+endforeach()
+install(FILES 3rdparty/compiler-rt/builtin_fp16.h DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/tvm_tirx/compiler-rt)
+include(CMakePackageConfigHelpers)
+configure_package_config_file(cmake/tvm_tirxConfig.cmake.in
+ "${CMAKE_CURRENT_BINARY_DIR}/tvm_tirxConfig.cmake"
+ INSTALL_DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/tvm_tirx")
+install(EXPORT tvm_tirxTargets NAMESPACE tvm_tirx:: DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/tvm_tirx")
+install(FILES "${CMAKE_CURRENT_BINARY_DIR}/tvm_tirxConfig.cmake" DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/tvm_tirx")
+install(DIRECTORY cmake/utils cmake/modules DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/tvm_tirx" FILES_MATCHING PATTERN "*.cmake")
+set(TVM_TIRX_CMAKE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
+if(IS_SUBPROJECT)
+ set(TVM_TIRX_CMAKE_DIR "${TVM_TIRX_CMAKE_DIR}" PARENT_SCOPE)
+ foreach(_option USE_CUDA USE_ROCM BUILD_FOR_ANDROID)
+ set(TVM_TIRX_${_option} "${${_option}}" PARENT_SCOPE)
+ endforeach()
endif()
-
-# Tests inherit the final project linker flags and compile definitions.
add_subdirectory(tests/cpp)
if(GTEST_FOUND)
- # CTest must also be enabled at the project root for subdirectory discovery.
enable_testing()
include(CTest)
endif()
-
-# Python package installation configuration
-# This section ensures that all necessary files are installed for the Python wheel
-if(TVM_BUILD_PYTHON_MODULE)
- message(STATUS "Configuring Python package installation")
-
- # Install compiled shared libraries into <project>/lib so that
- # tvm_ffi.libinfo.load_lib_ctypes(package="tvm", target_name=...) can find
- # them via the package RECORD or the project's lib/ fallback dir.
- install(TARGETS tvm_compiler DESTINATION "lib")
- install(TARGETS tvm_runtime DESTINATION "lib")
-
- # Install third-party compiled dependencies into the same lib/ dir.
- if(TARGET fpA_intB_gemm)
- tvm_configure_target_library(fpA_intB_gemm)
- install(TARGETS fpA_intB_gemm DESTINATION "lib")
- endif()
- if(TARGET flash_attn)
- tvm_configure_target_library(flash_attn)
- install(TARGETS flash_attn DESTINATION "lib")
- endif()
-
- # Install prebuilt extra runtime libraries into the same lib/ dir. This is how
- # the separately-built CUDA runtime (libtvm_runtime_cuda.so) is bundled: the
- # publishing flow builds it in a CUDA-enabled environment and passes its path
- # via TVM_PACKAGE_EXTRA_LIBS, so it ships through the normal CMake install
- # rather than a post-build wheel rewrite.
+find_and_set_linker(${USE_ALTERNATIVE_LINKER})
+if(SUMMARIZE)
+ print_summary()
+endif()
+dump_options_to_file("${TVM_ALL_OPTIONS}")
+if(TVM_BUILD_PYTHON_MODULE AND NOT IS_SUBPROJECT)
foreach(_extra_lib IN LISTS TVM_PACKAGE_EXTRA_LIBS)
if(NOT EXISTS "${_extra_lib}")
message(FATAL_ERROR "TVM_PACKAGE_EXTRA_LIBS entry does not exist: ${_extra_lib}")
endif()
- install(FILES "${_extra_lib}" DESTINATION "lib")
+ install(FILES "${_extra_lib}" DESTINATION lib)
endforeach()
-
- # Install minimal header files needed by Python extensions
- install(
- DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/include/tvm/runtime/"
- DESTINATION "include/tvm/runtime/"
- FILES_MATCHING
- PATTERN "*.h"
- )
-
- # Install minimal CMake configuration
- install(
- DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils/"
- DESTINATION "cmake/utils/"
- FILES_MATCHING
- PATTERN "*.cmake"
- )
-
- # Install CUTLASS headers only if available
- if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/cutlass/include")
- install(
- DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/cutlass/include/"
- DESTINATION "3rdparty/cutlass/include/"
- FILES_MATCHING
- PATTERN "*.h"
- PATTERN "*.hpp"
- )
- endif()
-
- # Install minimal source files
- install(
- DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/src/runtime/"
- DESTINATION "src/runtime/"
- FILES_MATCHING
- PATTERN "*.cc"
- PATTERN "*.h"
- )
-
- # Install web package
- install(DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/web/" DESTINATION "web/")
-
- # Install licenses (required for distribution)
- install(
- DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/licenses/"
- DESTINATION "licenses/"
- )
-
- # Install essential metadata files
- install(FILES
- "${CMAKE_CURRENT_SOURCE_DIR}/README.md"
- "${CMAKE_CURRENT_SOURCE_DIR}/LICENSE"
- "${CMAKE_CURRENT_SOURCE_DIR}/NOTICE"
- DESTINATION "."
- )
-
- message(STATUS "Python package installation configured")
+ install(DIRECTORY licenses/ DESTINATION licenses)
+ install(FILES LICENSE NOTICE DESTINATION .)
endif()
diff --git a/tirx/cmake/modules/CUDA.cmake b/tirx/cmake/modules/CUDA.cmake
index 4e2cdd508c..29217b891c 100644
--- a/tirx/cmake/modules/CUDA.cmake
+++ b/tirx/cmake/modules/CUDA.cmake
@@ -16,7 +16,7 @@
# under the License.
# CUDA Module
-find_cuda(${USE_CUDA} ${USE_CUDNN})
+find_cuda(${USE_CUDA} OFF)
if(CUDA_FOUND)
# always set the includedir when CUDA is available
@@ -70,9 +70,8 @@ if(USE_CUDA)
endif()
tvm_file_glob(GLOB RUNTIME_CUDA_SRCS src/backend/cuda/runtime/*.cc)
- tvm_file_glob(GLOB VM_CUDA_BUILTIN_SRC_CC src/runtime/vm/cuda/*.cc)
- add_library(tvm_runtime_cuda_objs OBJECT ${RUNTIME_CUDA_SRCS} ${VM_CUDA_BUILTIN_SRC_CC})
+ add_library(tvm_runtime_cuda_objs OBJECT ${RUNTIME_CUDA_SRCS})
target_link_libraries(tvm_runtime_cuda_objs PUBLIC tvm_ffi_header)
# These sources compile into tvm_runtime_cuda.dll, so their TVM_RUNTIME_DLL /
# TVM_FFI_DLL symbols must be dllexport on MSVC (e.g. GetCudaDeviceCount in
@@ -92,71 +91,3 @@ if(USE_CUDA)
target_link_libraries(tvm_runtime_cuda PRIVATE ${CUDA_NVTX_LIBRARY})
endif()
endif(USE_CUDA)
-
-# Contrib sources gated by USE_CUDA go into libtvm_runtime_extra.
-# See the RuntimeExtra assembly block in CMakeLists.txt.
-
-if(USE_CUDA AND USE_CUDNN)
- message(STATUS "Build with cuDNN support")
- include_directories(SYSTEM ${CUDA_CUDNN_INCLUDE_DIRS})
- tvm_file_glob(GLOB CUDNN_RELAX_CONTRIB_SRC src/relax/backend/contrib/cudnn/*.cc)
- list(APPEND COMPILER_SRCS ${CUDNN_RELAX_CONTRIB_SRC})
- tvm_file_glob(GLOB CONTRIB_CUDNN_SRCS src/runtime/extra/contrib/cudnn/*.cc)
- add_library(tvm_cudnn_objs OBJECT ${CONTRIB_CUDNN_SRCS})
- target_link_libraries(tvm_cudnn_objs PRIVATE tvm_runtime_extra_defs)
- target_link_libraries(tvm_runtime_extra PRIVATE tvm_cudnn_objs ${CUDA_CUDNN_LIBRARY})
-endif(USE_CUDA AND USE_CUDNN)
-
-if(USE_CUDA AND USE_CUDNN_FRONTEND)
- message(STATUS "Build with cuDNN Frontend support")
- if(IS_DIRECTORY ${USE_CUDNN_FRONTEND})
- find_file(CUDNN_FRONTEND_HEADER cudnn_frontend.h HINTS ${USE_CUDNN_FRONTEND}/include)
- include_directories(SYSTEM ${USE_CUDNN_FRONTEND}/include)
- else()
- find_file(CUDNN_FRONTEND_HEADER cudnn_frontend.h)
- endif()
- if(NOT CUDNN_FRONTEND_HEADER)
- message(FATAL_ERROR "Cannot find cudnn_frontend.h, please set USE_CUDNN_FRONTEND to the path of the cuDNN frontend header")
- endif()
- tvm_file_glob(GLOB CONTRIB_CUDNN_FRONTEND_SRCS src/runtime/extra/contrib/cudnn/cudnn_frontend/*.cc)
- set_source_files_properties(${CONTRIB_CUDNN_SRCS} PROPERTIES COMPILE_DEFINITIONS TVM_USE_CUDNN_FRONTEND=1)
- add_library(tvm_cudnn_frontend_objs OBJECT ${CONTRIB_CUDNN_FRONTEND_SRCS})
- target_link_libraries(tvm_cudnn_frontend_objs PRIVATE tvm_runtime_extra_defs)
- target_link_libraries(tvm_runtime_extra PRIVATE tvm_cudnn_frontend_objs)
-endif(USE_CUDA AND USE_CUDNN_FRONTEND)
-
-if(USE_CUDA AND USE_CUBLAS)
- message(STATUS "Build with cuBLAS support")
- tvm_file_glob(GLOB CUBLAS_CONTRIB_SRC src/relax/backend/contrib/cublas/*.cc)
- list(APPEND COMPILER_SRCS ${CUBLAS_CONTRIB_SRC})
- tvm_file_glob(GLOB CONTRIB_CUBLAS_SRCS src/runtime/extra/contrib/cublas/*.cc)
- add_library(tvm_cublas_objs OBJECT ${CONTRIB_CUBLAS_SRCS})
- target_link_libraries(tvm_cublas_objs PRIVATE tvm_runtime_extra_defs)
- target_link_libraries(tvm_runtime_extra PRIVATE tvm_cublas_objs ${CUDA_CUBLAS_LIBRARY})
- if(NOT CUDA_CUBLASLT_LIBRARY STREQUAL "CUDA_CUBLASLT_LIBRARY-NOTFOUND")
- target_link_libraries(tvm_runtime_extra PRIVATE ${CUDA_CUBLASLT_LIBRARY})
- endif()
-endif(USE_CUDA AND USE_CUBLAS)
-
-if(USE_CUDA AND USE_THRUST)
- message(STATUS "Build with Thrust support")
- tvm_file_glob(GLOB CONTRIB_THRUST_SRC src/runtime/extra/contrib/thrust/*.cu)
- add_library(tvm_thrust_objs OBJECT ${CONTRIB_THRUST_SRC})
- target_link_libraries(tvm_thrust_objs PRIVATE tvm_runtime_extra_defs)
- target_compile_options(tvm_thrust_objs PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:--expt-extended-lambda>)
- if(NOT USE_THRUST MATCHES ${IS_TRUE_PATTERN})
- find_package(CCCL REQUIRED COMPONENTS Thrust)
- target_link_libraries(tvm_thrust_objs PRIVATE CCCL::Thrust)
- endif()
- target_link_libraries(tvm_runtime_extra PRIVATE tvm_thrust_objs)
-endif(USE_CUDA AND USE_THRUST)
-
-if(USE_CUDA AND USE_CURAND)
- message(STATUS "Build with cuRAND support")
- message(STATUS "${CUDA_CURAND_LIBRARY}")
- tvm_file_glob(GLOB CONTRIB_CURAND_SRC_CC src/runtime/extra/contrib/curand/*.cc)
- tvm_file_glob(GLOB CONTRIB_CURAND_SRC_CU src/runtime/extra/contrib/curand/*.cu)
- add_library(tvm_curand_objs OBJECT ${CONTRIB_CURAND_SRC_CC} ${CONTRIB_CURAND_SRC_CU})
- target_link_libraries(tvm_curand_objs PRIVATE tvm_runtime_extra_defs)
- target_link_libraries(tvm_runtime_extra PRIVATE tvm_curand_objs ${CUDA_CURAND_LIBRARY})
-endif(USE_CUDA AND USE_CURAND)
diff --git a/tirx/cmake/modules/ClangFlags.cmake b/tirx/cmake/modules/ClangFlags.cmake
index 61a88116c7..f296165726 100644
--- a/tirx/cmake/modules/ClangFlags.cmake
+++ b/tirx/cmake/modules/ClangFlags.cmake
@@ -82,7 +82,7 @@ if ("${CMAKE_CXX_COMPILER_ID}" STREQUAL "Clang")
-Wno-gnu-anonymous-struct
-Wno-nested-anon-types
)
- target_compile_options(tvm_objs PRIVATE $<$<COMPILE_LANGUAGE:CXX>: ${warning_opts}>)
+ target_compile_options(tvm_tirx_objs PRIVATE $<$<COMPILE_LANGUAGE:CXX>: ${warning_opts}>)
target_compile_options(tvm_runtime_objs PRIVATE $<$<COMPILE_LANGUAGE:CXX>: ${warning_opts}>)
diff --git a/tirx/cmake/modules/Logging.cmake b/tirx/cmake/modules/Logging.cmake
index 09bb094c29..434e284da3 100644
--- a/tirx/cmake/modules/Logging.cmake
+++ b/tirx/cmake/modules/Logging.cmake
@@ -21,9 +21,8 @@ include(FindPackageHandleStandardArgs)
if(USE_CUSTOM_LOGGING)
# Set and propagate TVM_LOG_CUSTOMIZE flag is custom logging has been requested
- target_compile_definitions(tvm_objs PUBLIC TVM_LOG_CUSTOMIZE=1)
+ target_compile_definitions(tvm_tirx_objs PUBLIC TVM_LOG_CUSTOMIZE=1)
target_compile_definitions(tvm_runtime_objs PUBLIC TVM_LOG_CUSTOMIZE=1)
- target_compile_definitions(tvm_libinfo_objs PUBLIC TVM_LOG_CUSTOMIZE=1)
- target_compile_definitions(tvm_compiler PUBLIC TVM_LOG_CUSTOMIZE=1)
+ target_compile_definitions(tvm_tirx PUBLIC TVM_LOG_CUSTOMIZE=1)
target_compile_definitions(tvm_runtime PUBLIC TVM_LOG_CUSTOMIZE=1)
endif()
diff --git a/tirx/cmake/modules/OpenCL.cmake b/tirx/cmake/modules/OpenCL.cmake
index 4a398ae79f..e24ec2a24f 100644
--- a/tirx/cmake/modules/OpenCL.cmake
+++ b/tirx/cmake/modules/OpenCL.cmake
@@ -25,7 +25,10 @@ if(USE_OPENCL)
"src/backend/opencl/runtime/opencl_wrapper/opencl_wrapper.cc"
)
list(APPEND RUNTIME_OPENCL_SRCS ${OPENCL_WRAPPER_SRCS})
- set(_opencl_include_dirs "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/OpenCL-Headers")
+ find_path(TVM_OPENCL_HEADERS_DIR CL/cl.h
+ HINTS "${TVM_OPENCL_HEADERS_ROOT}" ENV OPENCL_SDK
+ PATH_SUFFIXES include REQUIRED)
+ set(_opencl_include_dirs "${TVM_OPENCL_HEADERS_DIR}")
else()
find_opencl(${USE_OPENCL})
if(NOT OpenCL_FOUND)
@@ -61,11 +64,7 @@ if(USE_OPENCL)
"$<BUILD_INTERFACE:${_opencl_include_dirs}>")
tvm_configure_target_library(tvm_runtime_opencl RUNTIME_MODULE)
- if(USE_CLML_GRAPH_EXECUTOR)
- # CLML calls both OpenCL SDK and TVM OpenCL runtime APIs directly.
- target_link_libraries(tvm_clml_objs PRIVATE tvm_runtime_opencl)
- target_link_libraries(tvm_runtime_extra PRIVATE tvm_runtime_opencl ${_opencl_libs})
- endif()
+
if(USE_OPENCL_ENABLE_HOST_PTR)
target_compile_definitions(tvm_runtime_opencl_objs PRIVATE OPENCL_ENABLE_HOST_PTR)
diff --git a/tirx/cmake/modules/ROCM.cmake b/tirx/cmake/modules/ROCM.cmake
index 33abc75a2b..d910674dc2 100644
--- a/tirx/cmake/modules/ROCM.cmake
+++ b/tirx/cmake/modules/ROCM.cmake
@@ -50,34 +50,3 @@ if(USE_ROCM)
target_link_libraries(tvm_runtime_rocm PUBLIC tvm_runtime PRIVATE ${_rocm_libs})
tvm_configure_target_library(tvm_runtime_rocm RUNTIME_MODULE)
endif(USE_ROCM)
-
-# HIPBLAS contrib goes into libtvm_runtime_extra.
-if(USE_ROCM AND USE_HIPBLAS)
- message(STATUS "Build with HIPBLAS support")
- tvm_file_glob(GLOB HIPBLAS_CONTRIB_SRC src/relax/backend/contrib/hipblas/*.cc)
- list(APPEND COMPILER_SRCS ${HIPBLAS_CONTRIB_SRC})
- tvm_file_glob(GLOB HIPBLAS_CONTRIB_SRCS src/runtime/extra/contrib/hipblas/*.cc)
- add_library(tvm_hipblas_objs OBJECT ${HIPBLAS_CONTRIB_SRCS})
- target_link_libraries(tvm_hipblas_objs PRIVATE tvm_runtime_extra_defs)
- target_link_libraries(tvm_runtime_extra PRIVATE tvm_hipblas_objs ${ROCM_HIPBLAS_LIBRARY})
- if(NOT ROCM_HIPBLASLT_LIBRARY STREQUAL "ROCM_HIPBLASLT_LIBRARY-NOTFOUND")
- target_link_libraries(tvm_runtime_extra PRIVATE ${ROCM_HIPBLASLT_LIBRARY})
- endif()
-endif(USE_ROCM AND USE_HIPBLAS)
-
-if(USE_ROCM AND USE_THRUST)
- message(STATUS "Build with rocThrust support")
- # We need to override CXX to hipcc. This is required by rocthrust
- if(${CMAKE_CXX_COMPILER} MATCHES "hipcc$")
- message(STATUS "Using hipcc compiler to compile rocthrust code.")
- else()
- message(FATAL_ERROR "Set CXX=hipcc to compile rocthrust code.")
- endif()
-
- find_package(rocprim REQUIRED)
- find_package(rocthrust REQUIRED)
- set_source_files_properties(src/runtime/extra/contrib/thrust/thrust.cu PROPERTIES LANGUAGE CXX)
- add_library(tvm_rocthrust_objs OBJECT src/runtime/extra/contrib/thrust/thrust.cu)
- target_link_libraries(tvm_rocthrust_objs PRIVATE tvm_runtime_extra_defs)
- target_link_libraries(tvm_runtime_extra PRIVATE tvm_rocthrust_objs roc::rocthrust)
-endif(USE_ROCM AND USE_THRUST)
diff --git a/tirx/cmake/modules/contrib/Z3.cmake b/tirx/cmake/modules/contrib/Z3.cmake
index bf61ec1352..48dc4f8d9c 100644
--- a/tirx/cmake/modules/contrib/Z3.cmake
+++ b/tirx/cmake/modules/contrib/Z3.cmake
@@ -95,3 +95,15 @@ endif()
# Enable the real Z3 implementation inside the single src/sym/z3_prover.cc file.
add_compile_definitions(TVM_USE_Z3)
message(STATUS "Build with Z3 SMT solver support")
+
+if(TVM_BUILD_PYTHON_MODULE AND Z3_TARGET)
+ get_target_property(_z3_target_type ${Z3_TARGET} TYPE)
+ if(_z3_target_type STREQUAL "STATIC_LIBRARY" AND Python3_EXECUTABLE)
+ execute_process(COMMAND "${Python3_EXECUTABLE}" -c
+ "import importlib.util; from pathlib import Path; s=importlib.util.find_spec('mlc_z3_static'); print(Path(s.origin).parent / 'licenses/LICENSE-Z3.txt' if s else '')"
+ OUTPUT_VARIABLE _z3_license OUTPUT_STRIP_TRAILING_WHITESPACE)
+ if(EXISTS "${_z3_license}")
+ install(FILES "${_z3_license}" DESTINATION licenses)
+ endif()
+ endif()
+endif()
diff --git a/tirx/cmake/modules/Logging.cmake b/tirx/cmake/tvm_tirxConfig.cmake.in
similarity index 52%
copy from tirx/cmake/modules/Logging.cmake
copy to tirx/cmake/tvm_tirxConfig.cmake.in
index 09bb094c29..abce76e80a 100644
--- a/tirx/cmake/modules/Logging.cmake
+++ b/tirx/cmake/tvm_tirxConfig.cmake.in
@@ -15,15 +15,23 @@
# specific language governing permissions and limitations
# under the License.
-# This script configures the logging module and dependency on libbacktrace
-
-include(FindPackageHandleStandardArgs)
-
-if(USE_CUSTOM_LOGGING)
- # Set and propagate TVM_LOG_CUSTOMIZE flag is custom logging has been requested
- target_compile_definitions(tvm_objs PUBLIC TVM_LOG_CUSTOMIZE=1)
- target_compile_definitions(tvm_runtime_objs PUBLIC TVM_LOG_CUSTOMIZE=1)
- target_compile_definitions(tvm_libinfo_objs PUBLIC TVM_LOG_CUSTOMIZE=1)
- target_compile_definitions(tvm_compiler PUBLIC TVM_LOG_CUSTOMIZE=1)
- target_compile_definitions(tvm_runtime PUBLIC TVM_LOG_CUSTOMIZE=1)
+@PACKAGE_INIT@
+include(CMakeFindDependencyMacro)
+find_dependency(Python COMPONENTS Interpreter)
+if(NOT TARGET tvm_ffi::shared)
+ if(NOT tvm_ffi_DIR)
+ execute_process(COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir
+ OUTPUT_VARIABLE tvm_ffi_DIR OUTPUT_STRIP_TRAILING_WHITESPACE
+ COMMAND_ERROR_IS_FATAL ANY)
+ endif()
+ find_dependency(tvm_ffi CONFIG)
+endif()
+if("@USE_THREADS@")
+ find_dependency(Threads)
endif()
+include("${CMAKE_CURRENT_LIST_DIR}/tvm_tirxTargets.cmake")
+set(TVM_TIRX_VERSION "@TVM_TIRX_VERSION@")
+set(TVM_TIRX_CMAKE_DIR "${CMAKE_CURRENT_LIST_DIR}")
+set(TVM_TIRX_USE_CUDA "@USE_CUDA@")
+set(TVM_TIRX_USE_ROCM "@USE_ROCM@")
+set(TVM_TIRX_BUILD_FOR_ANDROID "@BUILD_FOR_ANDROID@")
diff --git a/tests/cpp/CMakeLists.txt b/tirx/tests/cpp/CMakeLists.txt
similarity index 78%
copy from tests/cpp/CMakeLists.txt
copy to tirx/tests/cpp/CMakeLists.txt
index 35c9fadb9e..255ea70ed0 100644
--- a/tests/cpp/CMakeLists.txt
+++ b/tirx/tests/cpp/CMakeLists.txt
@@ -16,7 +16,7 @@
# under the License.
# Enable ctest if gtest is available
-if(USE_GTEST)
+if(USE_GTEST AND NOT BUILD_TESTING STREQUAL "OFF")
# Check env var for backward compatibility. A better way to specify package
# locations is to use CMAKE_PREFIX_PATH or other standard CMake mechanism
# (see CMake documentation for `find_package`).
@@ -54,18 +54,18 @@ if(USE_GTEST)
endif()
endif()
-# Create the `cpptest` target when GTest is available.
+# Create the `tirx_cpptest` target when GTest is available.
if(GTEST_FOUND)
tvm_file_glob(GLOB_RECURSE TEST_SRCS "${CMAKE_CURRENT_SOURCE_DIR}/*.cc")
- add_executable(cpptest ${TEST_SRCS})
+ add_executable(tirx_cpptest ${TEST_SRCS})
# Preserve the executable location used before this became a subdirectory.
- set_target_properties(cpptest PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${PROJECT_BINARY_DIR}")
+ set_target_properties(tirx_cpptest PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}")
# include runtime files for unit testing
- target_link_libraries(cpptest PRIVATE tvm_compiler GTest::GTest GTest::Main GTest::gmock pthread dl)
+ target_link_libraries(tirx_cpptest PRIVATE tvm_tirx GTest::GTest GTest::Main GTest::gmock pthread dl)
if(DEFINED LLVM_LIBS)
# The TVM library is linked with LLVM libraries. If the LLVM libraries are
# static and the symbols are not hidden, then don't link them again into
- # cpptest since cpptest is itself linked against the TVM library. If static
+ # tirx_cpptest since tirx_cpptest is itself linked against the TVM library. If static
# LLVM libraries are linked in twice, it can cause issues with global
# variable initialization (cl::opt).
# If the LLVM libraries are dynamic, we have to link them again, since the
@@ -78,32 +78,32 @@ if(GTEST_FOUND)
endif()
endforeach()
if(DEFINED LLVM_SO OR HIDE_PRIVATE_SYMBOLS)
- target_link_libraries(cpptest PRIVATE ${LLVM_LIBS})
+ target_link_libraries(tirx_cpptest PRIVATE ${LLVM_LIBS})
endif()
endif()
- set_target_properties(cpptest PROPERTIES EXCLUDE_FROM_ALL 1)
- set_target_properties(cpptest PROPERTIES EXCLUDE_FROM_DEFAULT_BUILD 1)
- # cpptest is built into ${CMAKE_BINARY_DIR} but the TVM shared libs it
- # links against (libtvm_compiler.so / libtvm_runtime.so) now live in
+ set_target_properties(tirx_cpptest PROPERTIES EXCLUDE_FROM_ALL 1)
+ set_target_properties(tirx_cpptest PROPERTIES EXCLUDE_FROM_DEFAULT_BUILD 1)
+ # tirx_cpptest is built into ${CMAKE_BINARY_DIR} but the TVM shared libs it
+ # links against (libtvm_tirx.so / libtvm_runtime.so) now live in
# ${CMAKE_BINARY_DIR}/lib. CMake's default behaviour bakes the
# absolute path of build/lib/ into the binary's RUNPATH, which works in tree
# but is not relocatable. Set $ORIGIN/lib (or @loader_path/lib on macOS)
# explicitly so the test binary finds its TVM libs relative to its own
# location, mirroring the install-time RPATH setup.
if(APPLE)
- set_target_properties(cpptest PROPERTIES
+ set_target_properties(tirx_cpptest PROPERTIES
BUILD_RPATH "@loader_path/lib"
INSTALL_RPATH "@loader_path/lib"
)
elseif(UNIX)
- set_target_properties(cpptest PROPERTIES
+ set_target_properties(tirx_cpptest PROPERTIES
BUILD_RPATH "\$ORIGIN/lib"
INSTALL_RPATH "\$ORIGIN/lib"
)
endif()
- target_compile_definitions(cpptest PRIVATE "NDEBUG")
+ target_compile_definitions(tirx_cpptest PRIVATE "NDEBUG")
# For some reason, compile definitions are not propagated correctly, so we manually add them here
- target_compile_definitions(cpptest PUBLIC $<TARGET_PROPERTY:tvm_compiler,INTERFACE_COMPILE_DEFINITIONS>)
- gtest_discover_tests(cpptest WORKING_DIRECTORY "${PROJECT_BINARY_DIR}")
+ target_compile_definitions(tirx_cpptest PUBLIC $<TARGET_PROPERTY:tvm_tirx,INTERFACE_COMPILE_DEFINITIONS>)
+ gtest_discover_tests(tirx_cpptest WORKING_DIRECTORY "${CMAKE_BINARY_DIR}")
endif()
(tvm) 06/33: [REFACTOR][CI] Exercise both partitions and publish standalone core wheels
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 350ea906a3031261179fe25ba46d6697955c264e
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 01:36:39 2026 +0000
[REFACTOR][CI] Exercise both partitions and publish standalone core wheels
---
.github/actions/build-wheel-for-publish/action.yml | 12 +++---
.github/workflows/main.yml | 45 +++++++++++++++-------
.github/workflows/publish_wheel.yml | 10 ++---
.pre-commit-config.yaml | 4 +-
ci/jenkins/data.py | 8 +++-
ci/jenkins/generated/common.groovy | 1 +
.../package/manylinux_build_libtvm_runtime_cuda.sh | 5 ++-
.../package/windows_build_libtvm_runtime_cuda.bat | 4 +-
ci/wheel/manylinux_build_libtvm_runtime_cuda.sh | 5 ++-
ci/wheel/windows_build_libtvm_runtime_cuda.bat | 4 +-
docker/bash.sh | 2 +-
tests/scripts/release/test_release_package.sh | 2 +-
tests/scripts/setup-pytest-env.sh | 2 +-
tests/scripts/task_cpp_unittest.sh | 1 +
tests/scripts/task_java_unittest.sh | 2 +-
tests/scripts/task_python_docs.sh | 4 +-
tests/scripts/task_python_unittest.sh | 4 +-
tests/scripts/task_python_unittest_gpuonly.sh | 6 +--
tests/scripts/task_web_wasm.sh | 2 +-
.../tvm_tirx/backend/cuda/launch/generate.py | 2 +-
20 files changed, 76 insertions(+), 49 deletions(-)
diff --git a/.github/actions/build-wheel-for-publish/action.yml b/.github/actions/build-wheel-for-publish/action.yml
index d4a3ca14c2..9039669e48 100644
--- a/.github/actions/build-wheel-for-publish/action.yml
+++ b/.github/actions/build-wheel-for-publish/action.yml
@@ -15,9 +15,9 @@
# specific language governing permissions and limitations
# under the License.
-name: Build TVM Wheel
+name: Build TIRx Wheel
description: >
- Build and test the LLVM-enabled TVM wheel for a given OS/architecture
+ Build and test the standalone LLVM-enabled apache-tvm-tirx wheel for a given OS/architecture
combination using cibuildwheel.
inputs:
@@ -59,7 +59,7 @@ runs:
id: llvm-cache
with:
path: ${{ runner.os == 'Windows' && 'C:/opt/llvm' || '/opt/llvm' }}
- key: tvm-wheel-llvm-${{ env.LLVM_VERSION }}-${{ runner.os }}-${{ inputs.arch }}-v6
+ key: tirx-wheel-llvm-${{ env.LLVM_VERSION }}-${{ runner.os }}-${{ inputs.arch }}-v6
# ---- Install LLVM via conda (cache miss only) ----
- name: Setup conda
@@ -107,10 +107,10 @@ runs:
run: echo "TVM_CUDA_EXTRA_LIB=$(cygpath -m "$(pwd)")/build-wheel-cuda/lib/tvm_runtime_cuda.dll" >> "$GITHUB_ENV"
# ---- Build and test wheels ----
- - name: Build and test wheels
+ - name: Build and test TIRx core wheel
uses: pypa/cibuildwheel@294735312765b09d24a2fbec22660ce817587d55 # v4.1.0
with:
- package-dir: .
+ package-dir: tirx
output-dir: wheelhouse
env:
CIBW_BUILD: ${{ inputs.build }}
@@ -139,7 +139,7 @@ runs:
PATH="C:/opt/llvm/Library/bin;$PATH"
CMAKE_ARGS="-DUSE_LLVM='C:/opt/llvm/Library/bin/llvm-config.exe --link-static' -DZLIB_USE_STATIC_LIBS=ON -DCMAKE_PREFIX_PATH=C:/opt/llvm/Library ${{ inputs.cmake_defines }} -DTVM_PACKAGE_EXTRA_LIBS=${{ env.TVM_CUDA_EXTRA_LIB }}"
# Turns the wheel-specific assertions in
- # tests/python/all-platform-minimal-test/test_validate_runtime_library.py
+ # tirx/tests/python/all-platform-minimal-test/test_validate_runtime_library.py
# ON (they skip unless these are set). Every published wheel is LLVM-enabled,
# so TVM_WHEEL_EXPECT_LLVM is always "1"; the CUDA runtime is only bundled on
# the CUDA wheels, so TVM_WHEEL_EXPECT_CUDA_RUNTIME is "0" for CPU wheels.
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 6840582eed..31d8bb8e03 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -50,17 +50,22 @@ jobs:
- name: Detect environment variables
uses: ./.github/actions/detect-env-vars
id: env_vars
- - name: Build TVM wheel
+ - name: Build TIRx and TVM wheels
shell: bash -l {0}
env:
CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }}
run: |
+ set -euxo pipefail
pip install scikit-build-core
- export CMAKE_ARGS="-DUSE_LLVM=ON -DBUILD_TESTING=OFF"
+ export CMAKE_ARGS="-DUSE_LLVM=ON -DUSE_METAL=ON -DBUILD_TESTING=OFF"
+ pip wheel --no-deps -w dist ./tirx -v
+ export PIP_FIND_LINKS="$(pwd)/dist"
+ export CMAKE_ARGS="-DBUILD_TESTING=OFF"
pip wheel --no-deps -w dist . -v
- - name: Install TVM from wheel
+ - name: Install TIRx and TVM from wheels
shell: bash -l {0}
run: |
+ set -euxo pipefail
export TVM_FFI_VERSION="$(python -c 'from importlib.metadata import version; print(version("apache-tvm-ffi"))')"
pip install ml_dtypes
# --no-deps: this package requires apache-tvm-ffi>=0.1.14.post1, and the source-built core
@@ -74,25 +79,30 @@ jobs:
from importlib.metadata import requires, version
assert version("apache-tvm-ffi") == os.environ["TVM_FFI_VERSION"]
- assert "apache-tvm-ffi>=0.1.14.post1" in requires("apache-tvm")
- assert "apache-tvm-ffi-orcjit==0.1.2" in requires("apache-tvm")
+ assert version("apache-tvm-tirx") == "0.1.0a0"
+ assert "apache-tvm-tirx==0.1.0a0" in requires("apache-tvm")
+ assert "apache-tvm-ffi>=0.1.14.post1" in requires("apache-tvm-tirx")
+ assert "apache-tvm-ffi-orcjit==0.1.2" in requires("apache-tvm-tirx")
PY
- name: Test
shell: bash -l {0}
run: >-
python -m pytest -v tests/python/all-platform-minimal-test tests/python/script
+ tirx/tests/python/all-platform-minimal-test tirx/tests/python/script
- name: Minimal Metal Compile-Only
shell: bash -l {0}
run: |
- python -m pytest -v -s 'tests/python/codegen/gpu_common/test_allreduce.py::test_allreduce_sum_compile'
- python -m pytest -v -s 'tests/python/backend/metal/codegen/test_metal.py::test_func_with_trailing_pod_params'
+ set -euxo pipefail
+ python -m pytest -v -s 'tirx/tests/python/codegen/gpu_common/test_allreduce.py::test_allreduce_sum_compile'
+ python -m pytest -v -s 'tirx/tests/python/backend/metal/codegen/test_metal.py::test_func_with_trailing_pod_params'
- name: Minimal Metal Compile-and-Run
shell: bash -l {0}
run: |
- python -m pytest -v -s tests/python/backend/metal
- python -m pytest -v -s 'tests/python/codegen/gpu_common/test_gpu_intrinsics.py'
- python -m pytest -v -s 'tests/python/codegen/gpu_common/test_allreduce.py::test_allreduce_sum[dims0-metal]'
+ set -euxo pipefail
+ python -m pytest -v -s tirx/tests/python/backend/metal
+ python -m pytest -v -s 'tirx/tests/python/codegen/gpu_common/test_gpu_intrinsics.py'
+ python -m pytest -v -s 'tirx/tests/python/codegen/gpu_common/test_allreduce.py::test_allreduce_sum[dims0-metal]'
Windows:
if: ${{ github.repository == 'apache/tvm' }}
@@ -106,18 +116,24 @@ jobs:
- name: Detect environment variables
uses: ./.github/actions/detect-env-vars
id: env_vars
- - name: Install TVM
+ - name: Build and install TIRx and TVM
shell: cmd /C call {0}
env:
CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }}
run: |
- pip install scikit-build-core
+ pip install scikit-build-core || exit /b 1
set CMAKE_ARGS=-DUSE_LLVM=ON -DBUILD_TESTING=OFF
- pip install --no-deps . -v
+ pip wheel --no-deps -w dist ./tirx -v || exit /b 1
+ set PIP_FIND_LINKS=%CD%\dist
+ set CMAKE_ARGS=-DBUILD_TESTING=OFF
+ pip wheel --no-deps -w dist . -v || exit /b 1
+ for %%w in (dist\*.whl) do (
+ pip install --no-deps "%%w" || exit /b 1
+ )
- name: Install test dependencies
shell: cmd /C call {0}
run: |
- pip install psutil cloudpickle ml_dtypes numpy packaging scipy tornado typing_extensions
+ pip install psutil cloudpickle ml_dtypes numpy packaging scipy tornado typing_extensions || exit /b 1
rem --no-deps keeps the source-built apache-tvm-ffi (versioned below 0.1.14.post1 from the
rem shallow checkout) from being replaced by a PyPI build.
pip install --no-deps apache-tvm-ffi-orcjit==0.1.2
@@ -125,3 +141,4 @@ jobs:
shell: cmd /C call {0}
run: >-
python -m pytest -v tests/python/all-platform-minimal-test tests/python/script
+ tirx/tests/python/all-platform-minimal-test tirx/tests/python/script
diff --git a/.github/workflows/publish_wheel.yml b/.github/workflows/publish_wheel.yml
index f972d96339..45fcf9c598 100644
--- a/.github/workflows/publish_wheel.yml
+++ b/.github/workflows/publish_wheel.yml
@@ -110,7 +110,7 @@ jobs:
- name: Upload CUDA runtime sidecar
uses: actions/upload-artifact@v7.0.1
with:
- name: tvm-cuda-runtime-${{ matrix.arch }}
+ name: tirx-cuda-runtime-${{ matrix.arch }}
path: ${{ matrix.lib }}
if-no-files-found: error
@@ -166,10 +166,10 @@ jobs:
if: ${{ matrix.include_cuda_runtime == '1' }}
uses: actions/download-artifact@v8.0.1
with:
- name: tvm-cuda-runtime-${{ matrix.arch }}
+ name: tirx-cuda-runtime-${{ matrix.arch }}
path: build-wheel-cuda/lib
- - name: Build TVM wheel
+ - name: Build TIRx wheel
uses: ./.github/actions/build-wheel-for-publish
with:
arch: ${{ matrix.arch }}
@@ -180,7 +180,7 @@ jobs:
- name: Upload wheel artifact
uses: actions/upload-artifact@v7.0.1
with:
- name: tvm-wheel-${{ matrix.artifact_suffix }}
+ name: tirx-wheel-${{ matrix.artifact_suffix }}
path: wheelhouse/*.whl
if-no-files-found: error
@@ -196,7 +196,7 @@ jobs:
steps:
- uses: actions/download-artifact@v8.0.1
with:
- pattern: tvm-wheel-*
+ pattern: tirx-wheel-*
path: dist
merge-multiple: true
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index f38e1cb858..c76f80bf57 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -25,11 +25,11 @@ repos:
hooks:
- id: cuda-launch-generated
name: check CUDA launch generated files
- entry: python python/tvm/backend/cuda/launch/generate.py --check
+ entry: python tirx/python/tvm_tirx/backend/cuda/launch/generate.py --check
language: python
additional_dependencies: [ruff==0.12.3, clang-format==20.1.8]
pass_filenames: false
- files: ^(python/tvm/backend/cuda/launch/|src/backend/cuda/launch/|docs/tirx/api/cuda_launch.rst)
+ files: ^(tirx/python/tvm_tirx/backend/cuda/launch/|tirx/src/backend/cuda/launch/|docs/tirx/api/cuda_launch.rst)
- id: check-asf-header
name: check ASF Header
entry: python tests/lint/check_asf_header.py --check
diff --git a/ci/jenkins/data.py b/ci/jenkins/data.py
index fd75107e7d..74f7ccc96d 100644
--- a/ci/jenkins/data.py
+++ b/ci/jenkins/data.py
@@ -32,12 +32,18 @@ import sys
files_to_stash = {
# Executables and build files needed to run c++ tests
- "cpptest": ["build/cpptest", "build/build.ninja", "build/CMakeFiles/rules.ninja"],
+ "cpptest": [
+ "build/cpptest",
+ "build/tirx_cpptest",
+ "build/build.ninja",
+ "build/CMakeFiles/rules.ninja",
+ ],
# runtime files
"tvm_runtime": ["build/lib/libtvm_runtime.so", "build/config.cmake"],
# compiler files
"tvm_lib": [
"build/lib/libtvm_compiler.so",
+ "build/lib/libtvm_tirx.so",
"build/lib/libtvm_runtime.so",
"build/lib/libtvm_ffi.so",
"build/lib/libtvm_runtime_cuda.so",
diff --git a/ci/jenkins/generated/common.groovy b/ci/jenkins/generated/common.groovy
index cd8324b5f2..a22b9b218d 100644
--- a/ci/jenkins/generated/common.groovy
+++ b/ci/jenkins/generated/common.groovy
@@ -298,6 +298,7 @@ def make_cpp_tests(image, build_dir) {
--sccache-bucket tvm-sccache-prod \
--sccache-region us-west-2 \
--cmake-target cpptest \
+ --cmake-target tirx_cpptest \
--build-dir ${build_dir}
""",
label: 'Make C++ tests',
diff --git a/ci/scripts/package/manylinux_build_libtvm_runtime_cuda.sh b/ci/scripts/package/manylinux_build_libtvm_runtime_cuda.sh
index d0d847ae83..f1a43b27c2 100755
--- a/ci/scripts/package/manylinux_build_libtvm_runtime_cuda.sh
+++ b/ci/scripts/package/manylinux_build_libtvm_runtime_cuda.sh
@@ -39,6 +39,7 @@ dnf -y install cuda-driver-devel-13-1
# generator (no Ninja), so no build tools are installed here. Put the bundled
# CPython and CUDA toolchain on PATH for the CMake configure and nvcc.
export PATH="/opt/python/cp310-cp310/bin:/usr/local/cuda/bin:${PATH}"
+python -m pip install "apache-tvm-ffi>=0.1.14.post1"
nvcc --version
rm -rf "${build_dir}"
@@ -48,13 +49,13 @@ rm -rf "${build_dir}"
# intentionally not set: it would be a no-op here for the same reason (verified --
# the .so is byte-identical across arch values and carries no device code), and
# modern CMake fills in a default so configure does not fail without it.
-cmake -S "${repo_root}" -B "${build_dir}" \
+cmake -S "${repo_root}/tirx" -B "${build_dir}" \
-DCMAKE_BUILD_TYPE=Release \
-DBUILD_TESTING=OFF \
-DTVM_BUILD_PYTHON_MODULE=ON \
-DUSE_CUDA=/usr/local/cuda \
-DUSE_LLVM=OFF \
- -DUSE_CUBLAS=OFF -DUSE_CUDNN=OFF -DUSE_CUTLASS=OFF -DUSE_NCCL=OFF -DUSE_NVTX=OFF \
+ -DUSE_NVTX=OFF \
-DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc
cmake --build "${build_dir}" --target tvm_runtime tvm_runtime_cuda --parallel "${parallel}"
diff --git a/ci/scripts/package/windows_build_libtvm_runtime_cuda.bat b/ci/scripts/package/windows_build_libtvm_runtime_cuda.bat
index 4864c18a80..f541af389b 100644
--- a/ci/scripts/package/windows_build_libtvm_runtime_cuda.bat
+++ b/ci/scripts/package/windows_build_libtvm_runtime_cuda.bat
@@ -62,7 +62,7 @@ set "cuda_root=%CD%"
popd
set "CUDA_PATH=%cuda_root%"
-python -m pip install -U pip cmake ninja || exit /b 1
+python -m pip install -U pip cmake ninja "apache-tvm-ffi>=0.1.14.post1" || exit /b 1
"%nvcc_exe%" --version || exit /b 1
rem nvcc needs the MSVC host compiler (cl.exe), so locate VS via vswhere and run the
@@ -90,7 +90,7 @@ set "cuda_root_fwd=%cuda_root:\=/%"
set "nvcc_fwd=%nvcc_exe:\=/%"
call "%vcvars%" || exit /b 1
-cmake -S "%repo_root_fwd%" -B "%build_dir_fwd%" -G Ninja -DCMAKE_BUILD_TYPE=Release -DBUILD_TESTING=OFF -DTVM_BUILD_PYTHON_MODULE=ON -DUSE_CUDA="%cuda_root_fwd%" -DUSE_LLVM=OFF -DUSE_CUBLAS=OFF -DUSE_CUDNN=OFF -DUSE_CUTLASS=OFF -DUSE_NCCL=OFF -DUSE_NVTX=OFF -DCMAKE_CUDA_COMPILER="%nvcc_fwd%" -DCMAKE_CUDA_FLAGS="-allow-unsupported-compiler" || exit /b 1
+cmake -S "%repo_root_fwd%/tirx" -B "%build_dir_fwd%" -G Ninja -DCMAKE_BUILD_TYPE=Release -DBUILD_TESTING=OFF -DTVM_BUILD_PYTHON_MODULE=ON -DUSE_CUDA="%cuda_root_fwd%" -DUSE_LLVM=OFF -DUSE_NVTX=OFF -DCMAKE_CUDA_COMPILER="%nvcc_fwd%" -DCMAKE_CUDA_FLAGS="-allow-unsupported-compiler" || exit /b 1
cmake --build "%build_dir_fwd%" --target tvm_runtime tvm_runtime_cuda --config Release || exit /b 1
if not exist "%build_dir%\lib\tvm_runtime_cuda.dll" ( echo tvm_runtime_cuda.dll was not produced & exit /b 1 )
diff --git a/ci/wheel/manylinux_build_libtvm_runtime_cuda.sh b/ci/wheel/manylinux_build_libtvm_runtime_cuda.sh
index 4f793ccfd7..a53ee8b1d0 100755
--- a/ci/wheel/manylinux_build_libtvm_runtime_cuda.sh
+++ b/ci/wheel/manylinux_build_libtvm_runtime_cuda.sh
@@ -39,6 +39,7 @@ dnf -y install cuda-driver-devel-13-1
# generator (no Ninja), so no build tools are installed here. Put the bundled
# CPython and CUDA toolchain on PATH for the CMake configure and nvcc.
export PATH="/opt/python/cp310-cp310/bin:/usr/local/cuda/bin:${PATH}"
+python -m pip install "apache-tvm-ffi>=0.1.14.post1"
nvcc --version
rm -rf "${build_dir}"
@@ -48,13 +49,13 @@ rm -rf "${build_dir}"
# intentionally not set: it would be a no-op here for the same reason (verified --
# the .so is byte-identical across arch values and carries no device code), and
# modern CMake fills in a default so configure does not fail without it.
-cmake -S "${repo_root}" -B "${build_dir}" \
+cmake -S "${repo_root}/tirx" -B "${build_dir}" \
-DCMAKE_BUILD_TYPE=Release \
-DBUILD_TESTING=OFF \
-DTVM_BUILD_PYTHON_MODULE=ON \
-DUSE_CUDA=/usr/local/cuda \
-DUSE_LLVM=OFF \
- -DUSE_CUBLAS=OFF -DUSE_CUDNN=OFF -DUSE_CUTLASS=OFF -DUSE_NCCL=OFF -DUSE_NVTX=OFF \
+ -DUSE_NVTX=OFF \
-DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc
cmake --build "${build_dir}" --target tvm_runtime tvm_runtime_cuda --parallel "${parallel}"
diff --git a/ci/wheel/windows_build_libtvm_runtime_cuda.bat b/ci/wheel/windows_build_libtvm_runtime_cuda.bat
index 0c375b441a..e53c2b9281 100644
--- a/ci/wheel/windows_build_libtvm_runtime_cuda.bat
+++ b/ci/wheel/windows_build_libtvm_runtime_cuda.bat
@@ -62,7 +62,7 @@ set "cuda_root=%CD%"
popd
set "CUDA_PATH=%cuda_root%"
-python -m pip install -U pip cmake ninja || exit /b 1
+python -m pip install -U pip cmake ninja "apache-tvm-ffi>=0.1.14.post1" || exit /b 1
"%nvcc_exe%" --version || exit /b 1
rem nvcc needs the MSVC host compiler (cl.exe), so locate VS via vswhere and run the
@@ -90,7 +90,7 @@ set "cuda_root_fwd=%cuda_root:\=/%"
set "nvcc_fwd=%nvcc_exe:\=/%"
call "%vcvars%" || exit /b 1
-cmake -S "%repo_root_fwd%" -B "%build_dir_fwd%" -G Ninja -DCMAKE_BUILD_TYPE=Release -DBUILD_TESTING=OFF -DTVM_BUILD_PYTHON_MODULE=ON -DUSE_CUDA="%cuda_root_fwd%" -DUSE_LLVM=OFF -DUSE_CUBLAS=OFF -DUSE_CUDNN=OFF -DUSE_CUTLASS=OFF -DUSE_NCCL=OFF -DUSE_NVTX=OFF -DCMAKE_CUDA_COMPILER="%nvcc_fwd%" -DCMAKE_CUDA_FLAGS="-allow-unsupported-compiler" || exit /b 1
+cmake -S "%repo_root_fwd%/tirx" -B "%build_dir_fwd%" -G Ninja -DCMAKE_BUILD_TYPE=Release -DBUILD_TESTING=OFF -DTVM_BUILD_PYTHON_MODULE=ON -DUSE_CUDA="%cuda_root_fwd%" -DUSE_LLVM=OFF -DUSE_NVTX=OFF -DCMAKE_CUDA_COMPILER="%nvcc_fwd%" -DCMAKE_CUDA_FLAGS="-allow-unsupported-compiler" || exit /b 1
cmake --build "%build_dir_fwd%" --target tvm_runtime tvm_runtime_cuda --config Release || exit /b 1
if not exist "%build_dir%\lib\tvm_runtime_cuda.dll" ( echo tvm_runtime_cuda.dll was not produced & exit /b 1 )
diff --git a/docker/bash.sh b/docker/bash.sh
index 86ba2c33ea..fe0b1eeb2b 100755
--- a/docker/bash.sh
+++ b/docker/bash.sh
@@ -541,7 +541,7 @@ fi
# Set TVM import path inside the docker image
if [[ "${DOCKER_IMAGE_NAME}" == *"ci"* ]]; then
- DOCKER_ENV+=( --env PYTHONPATH="${REPO_MOUNT_POINT}"/python )
+ DOCKER_ENV+=( --env PYTHONPATH="${REPO_MOUNT_POINT}/tirx/python:${REPO_MOUNT_POINT}/python" )
fi
diff --git a/tests/scripts/release/test_release_package.sh b/tests/scripts/release/test_release_package.sh
index c16564d825..9effdd40ca 100644
--- a/tests/scripts/release/test_release_package.sh
+++ b/tests/scripts/release/test_release_package.sh
@@ -119,5 +119,5 @@ python3 -m pip install ${deps}
echo "[10/10] Import TVM and print path ..."
export TVM_HOME=$(pwd)
-export PYTHONPATH=$TVM_HOME/python:${PYTHONPATH:-}
+export PYTHONPATH=$TVM_HOME/tirx/python:$TVM_HOME/python:${PYTHONPATH:-}
python3 -c "import tvm; print(tvm.__path__)"
diff --git a/tests/scripts/setup-pytest-env.sh b/tests/scripts/setup-pytest-env.sh
index c62b770baf..e335b7903a 100755
--- a/tests/scripts/setup-pytest-env.sh
+++ b/tests/scripts/setup-pytest-env.sh
@@ -27,7 +27,7 @@ fi
set -ux
export TVM_PATH=`pwd`
-export PYTHONPATH="${TVM_PATH}/python"
+export PYTHONPATH="${TVM_PATH}/tirx/python:${TVM_PATH}/python"
export TVM_PYTEST_RESULT_DIR="${TVM_PATH}/build/pytest-results"
mkdir -p "${TVM_PYTEST_RESULT_DIR}"
diff --git a/tests/scripts/task_cpp_unittest.sh b/tests/scripts/task_cpp_unittest.sh
index bd526db185..424642673b 100755
--- a/tests/scripts/task_cpp_unittest.sh
+++ b/tests/scripts/task_cpp_unittest.sh
@@ -35,4 +35,5 @@ export OMP_NUM_THREADS=1
pushd "${BUILD_DIR}"
# run cpp test executable
./cpptest
+./tirx_cpptest
popd
diff --git a/tests/scripts/task_java_unittest.sh b/tests/scripts/task_java_unittest.sh
index 35621a5a13..29342c70c8 100755
--- a/tests/scripts/task_java_unittest.sh
+++ b/tests/scripts/task_java_unittest.sh
@@ -18,7 +18,7 @@
set -euxo pipefail
-export PYTHONPATH=python
+export PYTHONPATH=tirx/python:python
export LD_LIBRARY_PATH="lib:${LD_LIBRARY_PATH:-}"
# to avoid CI CPU thread throttling.
diff --git a/tests/scripts/task_python_docs.sh b/tests/scripts/task_python_docs.sh
index e45d439929..0a4b88df96 100755
--- a/tests/scripts/task_python_docs.sh
+++ b/tests/scripts/task_python_docs.sh
@@ -18,7 +18,7 @@
set -euxo pipefail
-export PYTHONPATH="$(pwd)/python"
+export PYTHONPATH="$(pwd)/tirx/python:$(pwd)/python"
# to avoid CI CPU thread throttling.
export TVM_BIND_THREADS=0
@@ -137,7 +137,7 @@ uv pip install --no-deps apache-tvm-ffi-orcjit==0.1.2
cd docs
-PYTHONPATH=$(pwd)/../python make htmldepoly SPHINXOPTS='-j auto' |& tee /tmp/$$.log.txt
+PYTHONPATH=$(pwd)/../tirx/python:$(pwd)/../python make htmldepoly SPHINXOPTS='-j auto' |& tee /tmp/$$.log.txt
if grep -E "failed to execute|Segmentation fault" < /tmp/$$.log.txt; then
echo "Some of sphinx-gallery item example failed to execute."
exit 1
diff --git a/tests/scripts/task_python_unittest.sh b/tests/scripts/task_python_unittest.sh
index 511b3ec2a5..c6c578cc5b 100755
--- a/tests/scripts/task_python_unittest.sh
+++ b/tests/scripts/task_python_unittest.sh
@@ -18,7 +18,7 @@
set -euxo pipefail
-export PYTHONPATH="$(pwd)/python"
+export PYTHONPATH="$(pwd)/tirx/python:$(pwd)/python"
export PYTEST_ADDOPTS="${CI_PYTEST_ADD_OPTIONS:-} ${PYTEST_ADDOPTS:-}"
# setup tvm-ffi into python folder
@@ -29,4 +29,4 @@ uv pip install -v --target=python ./3rdparty/tvm-ffi/
# --no-deps keeps the submodule-built tvm-ffi above from being replaced by a PyPI build.
uv pip install --no-deps apache-tvm-ffi-orcjit==0.1.2
-python3 -m pytest -vvs -n auto -m "${TVM_TEST_MARKER:-not gpu}" tests/python
+python3 -m pytest -vvs -n auto -m "${TVM_TEST_MARKER:-not gpu}" tests/python tirx/tests/python
diff --git a/tests/scripts/task_python_unittest_gpuonly.sh b/tests/scripts/task_python_unittest_gpuonly.sh
index ac9e112b84..c025bd6675 100755
--- a/tests/scripts/task_python_unittest_gpuonly.sh
+++ b/tests/scripts/task_python_unittest_gpuonly.sh
@@ -31,9 +31,9 @@ export TVM_TEST_TARGETS='cuda;metal;rocm;nvptx'
# only minimal functionality of Vulkan runtime.
export TVM_TEST_TARGETS='{"kind":"vulkan","from_device":0}'
-export PYTHONPATH="$(pwd)/python"
+export PYTHONPATH="$(pwd)/tirx/python:$(pwd)/python"
export PYTEST_ADDOPTS="${CI_PYTEST_ADD_OPTIONS:-} ${PYTEST_ADDOPTS:-}"
python3 -m pytest -vvs -n auto -m "${TVM_TEST_MARKER}" \
- tests/python/backend/vulkan/codegen/test_vulkan.py \
- tests/python/backend/vulkan/test_vulkan_module_export.py
+ tirx/tests/python/backend/vulkan/codegen/test_vulkan.py \
+ tirx/tests/python/backend/vulkan/test_vulkan_module_export.py
diff --git a/tests/scripts/task_web_wasm.sh b/tests/scripts/task_web_wasm.sh
index 5f4d8d309e..aed91a371b 100755
--- a/tests/scripts/task_web_wasm.sh
+++ b/tests/scripts/task_web_wasm.sh
@@ -18,7 +18,7 @@
set -euxo pipefail
-export PYTHONPATH=`pwd`/python
+export PYTHONPATH="$(pwd)/tirx/python:$(pwd)/python"
# setup tvm-ffi into python folder
uv pip install -v --target=python ./3rdparty/tvm-ffi/
diff --git a/tirx/python/tvm_tirx/backend/cuda/launch/generate.py b/tirx/python/tvm_tirx/backend/cuda/launch/generate.py
index 029e80a582..381ac5c913 100644
--- a/tirx/python/tvm_tirx/backend/cuda/launch/generate.py
+++ b/tirx/python/tvm_tirx/backend/cuda/launch/generate.py
@@ -318,7 +318,7 @@ def main():
DIRECTORY / "_generated.py": python_api(),
ROOT / "src/backend/cuda/launch/launch_fields.h": native_fields(),
ROOT / "src/backend/cuda/launch/launch_attributes.h": native_attributes(),
- ROOT / "docs/tirx/api/cuda_launch.rst": documentation(),
+ ROOT.parent / "docs/tirx/api/cuda_launch.rst": documentation(),
}
for path in list(outputs):
outputs[path] = format_output(path, outputs[path])
(tvm) 17/33: [FIX][Packaging] Require the FFI release with typed reflection attributes
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 8c3a0874fcbb87d38f3cf11686133654ece0f173
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 03:13:09 2026 +0000
[FIX][Packaging] Require the FFI release with typed reflection attributes
---
.github/workflows/main.yml | 2 +-
ci/scripts/package/manylinux_build_libtvm_runtime_cuda.sh | 2 +-
ci/scripts/package/windows_build_libtvm_runtime_cuda.bat | 2 +-
ci/wheel/manylinux_build_libtvm_runtime_cuda.sh | 2 +-
ci/wheel/windows_build_libtvm_runtime_cuda.bat | 2 +-
tirx/pyproject.toml | 4 ++--
6 files changed, 7 insertions(+), 7 deletions(-)
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 31d8bb8e03..8603fdc95f 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -81,7 +81,7 @@ jobs:
assert version("apache-tvm-ffi") == os.environ["TVM_FFI_VERSION"]
assert version("apache-tvm-tirx") == "0.1.0a0"
assert "apache-tvm-tirx==0.1.0a0" in requires("apache-tvm")
- assert "apache-tvm-ffi>=0.1.14.post1" in requires("apache-tvm-tirx")
+ assert "apache-tvm-ffi>=0.1.15rc0" in requires("apache-tvm-tirx")
assert "apache-tvm-ffi-orcjit==0.1.2" in requires("apache-tvm-tirx")
PY
diff --git a/ci/scripts/package/manylinux_build_libtvm_runtime_cuda.sh b/ci/scripts/package/manylinux_build_libtvm_runtime_cuda.sh
index f1a43b27c2..f1eb80ad77 100755
--- a/ci/scripts/package/manylinux_build_libtvm_runtime_cuda.sh
+++ b/ci/scripts/package/manylinux_build_libtvm_runtime_cuda.sh
@@ -39,7 +39,7 @@ dnf -y install cuda-driver-devel-13-1
# generator (no Ninja), so no build tools are installed here. Put the bundled
# CPython and CUDA toolchain on PATH for the CMake configure and nvcc.
export PATH="/opt/python/cp310-cp310/bin:/usr/local/cuda/bin:${PATH}"
-python -m pip install "apache-tvm-ffi>=0.1.14.post1"
+python -m pip install "apache-tvm-ffi>=0.1.15rc0"
nvcc --version
rm -rf "${build_dir}"
diff --git a/ci/scripts/package/windows_build_libtvm_runtime_cuda.bat b/ci/scripts/package/windows_build_libtvm_runtime_cuda.bat
index f541af389b..6e3b988b48 100644
--- a/ci/scripts/package/windows_build_libtvm_runtime_cuda.bat
+++ b/ci/scripts/package/windows_build_libtvm_runtime_cuda.bat
@@ -62,7 +62,7 @@ set "cuda_root=%CD%"
popd
set "CUDA_PATH=%cuda_root%"
-python -m pip install -U pip cmake ninja "apache-tvm-ffi>=0.1.14.post1" || exit /b 1
+python -m pip install -U pip cmake ninja "apache-tvm-ffi>=0.1.15rc0" || exit /b 1
"%nvcc_exe%" --version || exit /b 1
rem nvcc needs the MSVC host compiler (cl.exe), so locate VS via vswhere and run the
diff --git a/ci/wheel/manylinux_build_libtvm_runtime_cuda.sh b/ci/wheel/manylinux_build_libtvm_runtime_cuda.sh
index a53ee8b1d0..8a5f33bf46 100755
--- a/ci/wheel/manylinux_build_libtvm_runtime_cuda.sh
+++ b/ci/wheel/manylinux_build_libtvm_runtime_cuda.sh
@@ -39,7 +39,7 @@ dnf -y install cuda-driver-devel-13-1
# generator (no Ninja), so no build tools are installed here. Put the bundled
# CPython and CUDA toolchain on PATH for the CMake configure and nvcc.
export PATH="/opt/python/cp310-cp310/bin:/usr/local/cuda/bin:${PATH}"
-python -m pip install "apache-tvm-ffi>=0.1.14.post1"
+python -m pip install "apache-tvm-ffi>=0.1.15rc0"
nvcc --version
rm -rf "${build_dir}"
diff --git a/ci/wheel/windows_build_libtvm_runtime_cuda.bat b/ci/wheel/windows_build_libtvm_runtime_cuda.bat
index e53c2b9281..45e62fef10 100644
--- a/ci/wheel/windows_build_libtvm_runtime_cuda.bat
+++ b/ci/wheel/windows_build_libtvm_runtime_cuda.bat
@@ -62,7 +62,7 @@ set "cuda_root=%CD%"
popd
set "CUDA_PATH=%cuda_root%"
-python -m pip install -U pip cmake ninja "apache-tvm-ffi>=0.1.14.post1" || exit /b 1
+python -m pip install -U pip cmake ninja "apache-tvm-ffi>=0.1.15rc0" || exit /b 1
"%nvcc_exe%" --version || exit /b 1
rem nvcc needs the MSVC host compiler (cl.exe), so locate VS via vswhere and run the
diff --git a/tirx/pyproject.toml b/tirx/pyproject.toml
index ed16357b21..db8d5c071c 100644
--- a/tirx/pyproject.toml
+++ b/tirx/pyproject.toml
@@ -20,7 +20,7 @@
requires = [
"scikit-build-core>=0.11",
"mlc-z3-static>=4.16.0",
- "apache-tvm-ffi>=0.1.14.post1",
+ "apache-tvm-ffi>=0.1.15rc0",
]
build-backend = "scikit_build_core.build"
@@ -47,7 +47,7 @@ classifiers = [
"Topic :: Software Development :: Libraries :: Python Modules",
]
dependencies = [
- "apache-tvm-ffi>=0.1.14.post1",
+ "apache-tvm-ffi>=0.1.15rc0",
"apache-tvm-ffi-orcjit==0.1.2",
"ml_dtypes",
"numpy",
(tvm) 14/33: [FIX][Packaging] Limit source distributions to their owned Python packages
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit a0c78133878ef88ce5c5035d022a3d14657d03e4
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 02:10:47 2026 +0000
[FIX][Packaging] Limit source distributions to their owned Python packages
---
pyproject.toml | 2 +-
tirx/pyproject.toml | 2 +-
2 files changed, 2 insertions(+), 2 deletions(-)
diff --git a/pyproject.toml b/pyproject.toml
index 02badac363..cd3a9547bf 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -103,7 +103,7 @@ sdist.include = [
"/cmake/**",
"/src/**",
"/include/**",
- "/python/**",
+ "/python/tvm/**",
"/licenses/**",
"/LICENSE",
"/NOTICE",
diff --git a/tirx/pyproject.toml b/tirx/pyproject.toml
index c07ec984de..ed16357b21 100644
--- a/tirx/pyproject.toml
+++ b/tirx/pyproject.toml
@@ -103,7 +103,7 @@ sdist.include = [
"/cmake/**",
"/src/**",
"/include/**",
- "/python/**",
+ "/python/tvm_tirx/**",
"/licenses/**",
"/LICENSE",
"/NOTICE",
(tvm) 18/33: [FIX][Web] Locate WebAssembly assets in the parent package
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit bb9ec8c2e10325f9f8e2e31dc885389837c8122c
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 04:14:57 2026 +0000
[FIX][Web] Locate WebAssembly assets in the parent package
---
python/tvm/contrib/emcc.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/python/tvm/contrib/emcc.py b/python/tvm/contrib/emcc.py
index b1b50bee56..296f9ee9a2 100644
--- a/python/tvm/contrib/emcc.py
+++ b/python/tvm/contrib/emcc.py
@@ -21,7 +21,7 @@ import os
import subprocess
from pathlib import Path
-from tvm_tirx import libinfo
+from tvm import libinfo
def find_wasm_lib(name, optional=False):
(tvm) 24/33: [FIX][CMake] Keep private visibility on owned library targets
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit ee7de53ff088928ed3320f11d7cd1f4a42b9339c
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 05:26:24 2026 +0000
[FIX][CMake] Keep private visibility on owned library targets
---
CMakeLists.txt | 9 ++++++---
1 file changed, 6 insertions(+), 3 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 3e40018bb0..4bd18b92a0 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -129,9 +129,6 @@ if(MSVC)
endif()
else()
add_compile_options(-Wall)
- if(HIDE_PRIVATE_SYMBOLS)
- add_compile_options(-fvisibility=hidden)
- endif()
if(APPLE)
set(TVM_NO_UNDEFINED_SYMBOLS "-Wl,-undefined,error")
else()
@@ -178,6 +175,9 @@ add_library(tvm_runtime_extra_defs INTERFACE)
target_link_libraries(tvm_runtime_extra_defs INTERFACE
tvm_tirx::tvm_runtime tvm_tirx::private_headers tvm_ffi::header)
target_compile_definitions(tvm_runtime_extra_defs INTERFACE TVM_RUNTIME_EXTRA_EXPORTS TVM_FFI_EXPORTS)
+if(HIDE_PRIVATE_SYMBOLS AND NOT MSVC)
+ target_compile_options(tvm_runtime_extra_defs INTERFACE -fvisibility=hidden)
+endif()
target_link_libraries(tvm_runtime_extra PRIVATE tvm_runtime_extra_defs)
tvm_file_glob(GLOB _vm_sources src/runtime/vm/*.cc src/runtime/rpc/minrpc/*.cc)
target_sources(tvm_runtime_extra PRIVATE ${_vm_sources})
@@ -234,6 +234,9 @@ endif()
include(cmake/modules/RuntimeExtra.cmake)
target_link_libraries(tvm_runtime_extra PRIVATE ${TVM_RUNTIME_LINKER_LIBS} ${CMAKE_DL_LIBS})
add_library(tvm_compiler SHARED ${COMPILER_SRCS})
+if(HIDE_PRIVATE_SYMBOLS AND NOT MSVC)
+ target_compile_options(tvm_compiler PRIVATE -fvisibility=hidden)
+endif()
target_compile_definitions(tvm_compiler PRIVATE TVM_EXPORTS TVM_FFI_EXPORTS)
target_link_libraries(tvm_compiler PUBLIC tvm_tirx::tvm_tirx tvm_runtime_extra)
target_link_libraries(tvm_compiler PRIVATE tvm_tirx::private_headers ${TVM_LINKER_LIBS})
(tvm) 22/33: [FIX][CMake] Normalize the packaged Z3 license path on Windows
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 8fb70273cdafe7f301fd059a89028c2d4591bdf8
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 05:25:17 2026 +0000
[FIX][CMake] Normalize the packaged Z3 license path on Windows
---
tirx/cmake/modules/contrib/Z3.cmake | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/tirx/cmake/modules/contrib/Z3.cmake b/tirx/cmake/modules/contrib/Z3.cmake
index 48dc4f8d9c..6a33b41723 100644
--- a/tirx/cmake/modules/contrib/Z3.cmake
+++ b/tirx/cmake/modules/contrib/Z3.cmake
@@ -100,7 +100,7 @@ if(TVM_BUILD_PYTHON_MODULE AND Z3_TARGET)
get_target_property(_z3_target_type ${Z3_TARGET} TYPE)
if(_z3_target_type STREQUAL "STATIC_LIBRARY" AND Python3_EXECUTABLE)
execute_process(COMMAND "${Python3_EXECUTABLE}" -c
- "import importlib.util; from pathlib import Path; s=importlib.util.find_spec('mlc_z3_static'); print(Path(s.origin).parent / 'licenses/LICENSE-Z3.txt' if s else '')"
+ "import importlib.util; from pathlib import Path; s=importlib.util.find_spec('mlc_z3_static'); print((Path(s.origin).parent / 'licenses/LICENSE-Z3.txt').as_posix() if s else '')"
OUTPUT_VARIABLE _z3_license OUTPUT_STRIP_TRAILING_WHITESPACE)
if(EXISTS "${_z3_license}")
install(FILES "${_z3_license}" DESTINATION licenses)
(tvm) 20/33: [FIX][Python] Retain parent version and external-module headers
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 739ffe23c14a527585fedb4709d82a28c9582f6c
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 04:26:11 2026 +0000
[FIX][Python] Retain parent version and external-module headers
---
python/tvm/relax/frontend/nn/extern.py | 1 +
python/tvm/topi/__init__.py | 2 +-
2 files changed, 2 insertions(+), 1 deletion(-)
diff --git a/python/tvm/relax/frontend/nn/extern.py b/python/tvm/relax/frontend/nn/extern.py
index 3bde6c3b67..999315af36 100644
--- a/python/tvm/relax/frontend/nn/extern.py
+++ b/python/tvm/relax/frontend/nn/extern.py
@@ -316,6 +316,7 @@ class SourceModule(ExternModule): # pylint: disable=too-few-public-methods
The list of include paths.
"""
results = [
+ Path(tvm.libinfo.find_include_path()),
Path(libinfo.find_include_path()),
Path(tvm_ffi.libinfo.find_include_path()),
Path(tvm_ffi.libinfo.find_dlpack_include_path()),
diff --git a/python/tvm/topi/__init__.py b/python/tvm/topi/__init__.py
index 918db50510..f0db21581c 100644
--- a/python/tvm/topi/__init__.py
+++ b/python/tvm/topi/__init__.py
@@ -26,7 +26,7 @@ Some of the schedule function may have been specially optimized for a
specific workload.
"""
-from tvm_tirx.libinfo import __version__
+from tvm.libinfo import __version__
# Ensure C++ schedules get registered first, so python schedules can
# override them.
(tvm) 33/33: [FIX][Docs] Point wheel guidance to the standalone core configuration
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 58d620f53fe286bc7898abb91ce1ad561691852d
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 12:52:23 2026 +0000
[FIX][Docs] Point wheel guidance to the standalone core configuration
---
ci/wheel/README.md | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/ci/wheel/README.md b/ci/wheel/README.md
index 972c315309..3b677b593c 100644
--- a/ci/wheel/README.md
+++ b/ci/wheel/README.md
@@ -18,7 +18,7 @@
# TVM wheel packaging
The wheels are built by a standard `cibuildwheel` flow, configured in
-`.github/workflows/publish_wheel.yml` and `pyproject.toml` (`[tool.cibuildwheel]`
+`.github/workflows/publish_wheel.yml` and `tirx/pyproject.toml` (`[tool.cibuildwheel]`
and `[tool.scikit-build]`). This directory holds the few helper scripts that flow
invokes:
(tvm) 04/33: [REFACTOR][Tests] Preserve fixture and dialect ownership across the package split
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 3fdcc3f579de45d0c78caa9c873ce147861fbdbb
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 01:33:58 2026 +0000
[REFACTOR][Tests] Preserve fixture and dialect ownership across the package split
---
.../contrib/test_android/test_meta_schedule.py | 2 +-
tests/python/contrib/test_example_npu.py | 2 +-
.../python/contrib/test_tir_triton_integration.py | 4 +-
tests/python/disco/test_callback.py | 4 +-
tests/python/disco/test_ccl.py | 2 +-
tests/python/disco/test_loader.py | 2 +-
tests/python/disco/test_nvshmem.py | 4 +-
tests/python/disco/test_session.py | 4 +-
tests/python/driver/test_compile.py | 4 +-
tests/python/relax/backend/adreno/mod_utils.py | 4 +-
tests/python/relax/backend/adreno/test_clml_ops.py | 2 +-
.../relax/backend/adreno/test_texture_network.py | 2 +-
.../test_transform_fold_vdevice_scope_change.py | 2 +-
.../test_distributed_transform_lower_distir.py | 2 +-
...ributed_transform_lower_global_to_local_view.py | 2 +-
...est_distributed_transform_propagate_sharding.py | 2 +-
.../relax/script/test_relax_script_basic_usage.py | 12 +--
.../relax/script/test_relax_script_distributed.py | 8 +-
.../script/test_relax_script_dynamic_shape.py | 24 +++---
.../script/test_relax_script_error_handling.py | 8 +-
.../relax/script/test_relax_script_ir_builder.py | 4 +-
.../script/test_relax_script_meta_programming.py | 8 +-
.../test_relax_script_operators_arith_cmp.py | 4 +-
.../script/test_relax_script_operators_create.py | 4 +-
.../script/test_relax_script_operators_datatype.py | 4 +-
.../script/test_relax_script_operators_grad.py | 4 +-
.../script/test_relax_script_operators_image.py | 4 +-
.../script/test_relax_script_operators_index.py | 4 +-
.../test_relax_script_operators_linear_algebra.py | 4 +-
.../test_relax_script_operators_manipulate.py | 4 +-
.../relax/script/test_relax_script_operators_nn.py | 4 +-
.../script/test_relax_script_operators_search.py | 4 +-
.../script/test_relax_script_operators_set.py | 4 +-
.../script/test_relax_script_operators_sort.py | 4 +-
.../test_relax_script_operators_statistical.py | 4 +-
.../script/test_relax_script_operators_vision.py | 4 +-
.../relax/script/test_relax_script_printer.py | 60 ++++++-------
.../relax/script/test_relax_script_pyfunc.py | 4 +-
tests/python/relax/test_analysis.py | 4 +-
.../test_analysis_computable_at_compile_time.py | 2 +-
.../relax/test_analysis_contains_impure_call.py | 2 +-
.../python/relax/test_analysis_detect_recursion.py | 4 +-
.../relax/test_analysis_estimate_memory_usage.py | 4 +-
.../test_analysis_suggest_layout_transforms.py | 2 +-
tests/python/relax/test_analysis_type_analysis.py | 2 +-
tests/python/relax/test_analysis_well_formed.py | 4 +-
tests/python/relax/test_ast_printer.py | 4 +-
.../python/relax/test_backend_dispatch_sampling.py | 4 +-
.../relax/test_backend_dispatch_sort_scan.py | 2 +-
.../relax/test_backend_transform_shape_lower.py | 4 +-
tests/python/relax/test_base_py_module_printer.py | 4 +-
.../relax/test_base_py_module_symbolic_shape.py | 4 +-
tests/python/relax/test_bind_params.py | 2 +-
tests/python/relax/test_bind_symbolic_vars.py | 2 +-
tests/python/relax/test_binding_rewrite.py | 2 +-
tests/python/relax/test_blockbuilder_core.py | 2 +-
tests/python/relax/test_blockbuilder_emit_te.py | 2 +-
tests/python/relax/test_codegen_cublas.py | 2 +-
tests/python/relax/test_codegen_cudnn.py | 2 +-
tests/python/relax/test_codegen_cutlass.py | 4 +-
tests/python/relax/test_codegen_dnnl.py | 2 +-
tests/python/relax/test_codegen_hipblas.py | 2 +-
tests/python/relax/test_codegen_tensorrt.py | 2 +-
tests/python/relax/test_contrib_vllm.py | 2 +-
tests/python/relax/test_dataflow_inplace.py | 2 +-
tests/python/relax/test_dataflow_pattern.py | 4 +-
tests/python/relax/test_dataflow_rewriter.py | 4 +-
tests/python/relax/test_dlpack_integration.py | 4 +-
tests/python/relax/test_e2e_op_dynamic.py | 2 +-
tests/python/relax/test_expr.py | 2 +-
tests/python/relax/test_expr_functor.py | 2 +-
tests/python/relax/test_fast_math_transform.py | 2 +-
tests/python/relax/test_frontend_common.py | 2 +-
tests/python/relax/test_frontend_dynamo.py | 4 +-
.../relax/test_frontend_from_exported_program.py | 2 +-
tests/python/relax/test_frontend_from_fx.py | 2 +-
tests/python/relax/test_frontend_nn_exporter.py | 2 +-
.../python/relax/test_frontend_nn_extern_module.py | 2 +-
tests/python/relax/test_frontend_nn_modules.py | 2 +-
tests/python/relax/test_frontend_nn_op.py | 4 +-
tests/python/relax/test_frontend_nn_packing.py | 2 +-
tests/python/relax/test_frontend_nn_subroutines.py | 2 +-
tests/python/relax/test_frontend_nn_tensor.py | 2 +-
tests/python/relax/test_frontend_onnx.py | 4 +-
tests/python/relax/test_frontend_stablehlo.py | 2 +-
tests/python/relax/test_inline_functions.py | 2 +-
tests/python/relax/test_kill_after_last_use.py | 2 +-
tests/python/relax/test_lower_alloc_tensor.py | 2 +-
.../relax/test_meta_schedule_relax_integration.py | 2 +-
tests/python/relax/test_op_binary.py | 2 +-
tests/python/relax/test_op_ccl.py | 2 +-
tests/python/relax/test_op_create.py | 2 +-
tests/python/relax/test_op_datatype.py | 2 +-
tests/python/relax/test_op_grad.py | 2 +-
tests/python/relax/test_op_image.py | 2 +-
tests/python/relax/test_op_index.py | 4 +-
tests/python/relax/test_op_inspect.py | 2 +-
tests/python/relax/test_op_linear_algebra.py | 2 +-
tests/python/relax/test_op_manipulate.py | 2 +-
tests/python/relax/test_op_misc.py | 4 +-
tests/python/relax/test_op_nn.py | 2 +-
tests/python/relax/test_op_nn_convolution.py | 2 +-
tests/python/relax/test_op_nn_pooling.py | 2 +-
tests/python/relax/test_op_qdq.py | 2 +-
tests/python/relax/test_op_sampling.py | 2 +-
tests/python/relax/test_op_search.py | 2 +-
tests/python/relax/test_op_set.py | 2 +-
tests/python/relax/test_op_size.py | 2 +-
tests/python/relax/test_op_sort.py | 2 +-
tests/python/relax/test_op_statistical.py | 2 +-
tests/python/relax/test_op_take.py | 2 +-
tests/python/relax/test_op_ternary.py | 2 +-
tests/python/relax/test_op_unary.py | 2 +-
tests/python/relax/test_op_view.py | 2 +-
tests/python/relax/test_op_vision.py | 2 +-
.../python/relax/test_optimize_layout_transform.py | 4 +-
tests/python/relax/test_pipeline.py | 2 +-
tests/python/relax/test_pytorch_integration.py | 4 +-
tests/python/relax/test_relax_operators.py | 2 +-
.../python/relax/test_relax_to_pyfunc_converter.py | 4 +-
.../python/relax/test_remove_redundant_reshape.py | 2 +-
.../python/relax/test_runtime_builtin_rnn_state.py | 2 +-
tests/python/relax/test_testing_nn.py | 2 +-
tests/python/relax/test_tir_call_source_kernel.py | 4 +-
tests/python/relax/test_training_append_loss.py | 2 +-
tests/python/relax/test_training_loss.py | 2 +-
tests/python/relax/test_training_setup_trainer.py | 2 +-
.../python/relax/test_training_trainer_numeric.py | 2 +-
tests/python/relax/test_transform.py | 4 +-
.../relax/test_transform_allocate_workspace.py | 2 +-
tests/python/relax/test_transform_alter_op_impl.py | 4 +-
.../test_transform_annotate_tir_op_pattern.py | 2 +-
...st_transform_attach_attr_layout_free_tensors.py | 4 +-
.../relax/test_transform_attach_global_symbol.py | 4 +-
tests/python/relax/test_transform_bind_params.py | 4 +-
.../relax/test_transform_bind_symbolic_vars.py | 2 +-
.../relax/test_transform_bundle_model_params.py | 2 +-
.../relax/test_transform_canonicalize_bindings.py | 4 +-
tests/python/relax/test_transform_codegen_pass.py | 4 +-
.../test_transform_combine_parallel_matmul.py | 2 +-
.../relax/test_transform_compute_prim_value.py | 4 +-
.../relax/test_transform_convert_dataflow.py | 2 +-
tests/python/relax/test_transform_cse.py | 2 +-
.../relax/test_transform_dead_code_elimination.py | 2 +-
tests/python/relax/test_transform_decompose_ops.py | 2 +-
.../relax/test_transform_error_enrichment.py | 2 +-
.../relax/test_transform_expand_matmul_of_sum.py | 2 +-
.../relax/test_transform_expand_tuple_args.py | 2 +-
.../test_transform_fold_batch_norm_to_conv2d.py | 2 +-
tests/python/relax/test_transform_fold_constant.py | 4 +-
tests/python/relax/test_transform_fuse_ops.py | 4 +-
.../relax/test_transform_fuse_ops_by_pattern.py | 4 +-
tests/python/relax/test_transform_fuse_tir.py | 4 +-
.../relax/test_transform_fuse_transpose_matmul.py | 4 +-
tests/python/relax/test_transform_gradient.py | 2 +-
.../relax/test_transform_gradient_te_register.py | 2 +-
.../test_transform_inline_private_functions.py | 2 +-
.../relax/test_transform_ipc_allreduce_rewrite.py | 2 +-
tests/python/relax/test_transform_lambda_lift.py | 4 +-
.../relax/test_transform_lazy_transform_params.py | 4 +-
tests/python/relax/test_transform_legalize_ops.py | 4 +-
.../relax/test_transform_legalize_ops_binary.py | 4 +-
.../relax/test_transform_legalize_ops_ccl.py | 4 +-
.../test_transform_legalize_ops_create_datatype.py | 4 +-
.../test_transform_legalize_ops_distributed.py | 4 +-
.../relax/test_transform_legalize_ops_grad.py | 4 +-
.../relax/test_transform_legalize_ops_image.py | 4 +-
..._transform_legalize_ops_index_linear_algebra.py | 4 +-
.../test_transform_legalize_ops_manipulate.py | 4 +-
.../python/relax/test_transform_legalize_ops_nn.py | 4 +-
.../relax/test_transform_legalize_ops_qdq.py | 4 +-
...st_transform_legalize_ops_search_statistical.py | 4 +-
.../relax/test_transform_legalize_ops_unary.py | 2 +-
.../relax/test_transform_lift_transform_params.py | 4 +-
.../test_transform_lower_gpu_ipc_alloc_storage.py | 2 +-
.../test_transform_merge_composite_functions.py | 4 +-
.../test_transform_meta_schedule_apply_database.py | 2 +-
.../relax/test_transform_meta_schedule_tuning.py | 4 +-
tests/python/relax/test_transform_normalize.py | 2 +-
.../relax/test_transform_normalize_global_var.py | 4 +-
...st_transform_operator_specific_normalization.py | 2 +-
.../relax/test_transform_remove_unused_outputs.py | 2 +-
.../test_transform_remove_unused_parameters.py | 2 +-
..._transform_reorder_permute_dims_after_concat.py | 2 +-
.../test_transform_reorder_take_after_matmul.py | 2 +-
.../relax/test_transform_rewrite_cuda_graph.py | 4 +-
.../test_transform_rewrite_dataflow_reshape.py | 4 +-
...nsform_specialize_function_based_on_callsite.py | 2 +-
.../test_transform_split_layout_rewrite_preproc.py | 4 +-
.../test_transform_static_plan_block_memory.py | 4 +-
.../relax/test_transform_to_mixed_precision.py | 2 +-
.../relax/test_transform_topological_sort.py | 2 +-
.../relax/test_transform_update_param_type.py | 2 +-
.../relax/test_vm_alloc_storage_with_scope.py | 4 +-
tests/python/relax/test_vm_build.py | 4 +-
tests/python/relax/test_vm_builtin.py | 2 +-
tests/python/relax/test_vm_builtin_lower.py | 2 +-
tests/python/relax/test_vm_callback_function.py | 2 +-
tests/python/relax/test_vm_codegen_only.py | 4 +-
tests/python/relax/test_vm_cuda_graph.py | 4 +-
tests/python/relax/test_vm_execbuilder.py | 2 +-
tests/python/relax/texture/test_texture_nd.py | 2 +-
...st_s_tir_analysis_calculate_allocated_memory.py | 2 +-
.../test_s_tir_analysis_estimate_tir_flops.py | 2 +-
.../test_s_tir_analysis_identify_memcpy.py | 2 +-
.../test_s_tir_analysis_is_pure_function.py | 2 +-
.../s_tir/analysis/test_s_tir_analysis_oob.py | 2 +-
.../test_s_tir_analysis_verify_block_scope.py | 2 +-
.../s_tir/analysis/test_sblock_access_region.py | 2 +-
.../analysis/test_sblock_tensor_access_lca.py | 2 +-
.../python/s_tir/base/test_compilation_pipeline.py | 2 +-
.../python/s_tir/base/test_s_tir_texture_scope.py | 2 +-
.../s_tir/base/test_sblock_dependence_info.py | 2 +-
.../s_tir/base/test_tir_te_extern_function.py | 2 +-
tests/python/s_tir/dlight/test_benchmark.py | 4 +-
tests/python/s_tir/dlight/test_cpu_gemv.py | 2 +-
tests/python/s_tir/dlight/test_gpu_conv.py | 2 +-
tests/python/s_tir/dlight/test_gpu_fallback.py | 2 +-
tests/python/s_tir/dlight/test_gpu_gemv.py | 2 +-
.../s_tir/dlight/test_gpu_general_reduction.py | 2 +-
.../python/s_tir/dlight/test_gpu_low_batch_gemv.py | 2 +-
tests/python/s_tir/dlight/test_gpu_matmul.py | 2 +-
.../s_tir/dlight/test_gpu_matmul_tensorize.py | 2 +-
tests/python/s_tir/dlight/test_gpu_reduction.py | 2 +-
tests/python/s_tir/dlight/test_gpu_rmsnorm.py | 2 +-
tests/python/s_tir/dlight/test_gpu_transpose.py | 2 +-
tests/python/s_tir/dlight/test_primitives.py | 2 +-
.../meta_schedule/test_meta_schedule_arg_info.py | 2 +-
.../meta_schedule/test_meta_schedule_builder.py | 2 +-
.../meta_schedule/test_meta_schedule_cost_model.py | 2 +-
.../meta_schedule/test_meta_schedule_database.py | 2 +-
...schedule_feature_extractor_per_store_feature.py | 2 +-
.../test_meta_schedule_measure_callback.py | 2 +-
...eta_schedule_mutator_mutate_compute_location.py | 2 +-
.../test_meta_schedule_mutator_mutate_parallel.py | 2 +-
..._meta_schedule_mutator_mutate_thread_binding.py | 2 +-
.../test_meta_schedule_mutator_mutate_tile_size.py | 2 +-
.../test_meta_schedule_mutator_mutate_unroll.py | 2 +-
.../test_meta_schedule_post_order_apply.py | 2 +-
...ule_postproc_disallow_async_strided_mem_copy.py | 2 +-
...meta_schedule_postproc_disallow_dynamic_loop.py | 2 +-
..._schedule_postproc_rewrite_cooperative_fetch.py | 2 +-
.../test_meta_schedule_postproc_rewrite_layout.py | 2 +-
...e_postproc_rewrite_parallel_vectorize_unroll.py | 2 +-
...ta_schedule_postproc_rewrite_reduction_block.py | 2 +-
...est_meta_schedule_postproc_rewrite_tensorize.py | 2 +-
...meta_schedule_postproc_rewrite_unbound_block.py | 2 +-
.../test_meta_schedule_postproc_verify_gpu_code.py | 2 +-
...est_meta_schedule_postproc_verify_vtcm_limit.py | 2 +-
.../meta_schedule/test_meta_schedule_runner.py | 2 +-
...test_meta_schedule_schedule_rule_add_rfactor.py | 2 +-
...eta_schedule_schedule_rule_apply_custom_rule.py | 2 +-
.../test_meta_schedule_schedule_rule_auto_bind.py | 2 +-
...test_meta_schedule_schedule_rule_auto_inline.py | 2 +-
...chedule_schedule_rule_cross_thread_reduction.py | 2 +-
.../test_meta_schedule_schedule_rule_mlt.py | 2 +-
.../test_meta_schedule_schedule_rule_mlt_intrin.py | 2 +-
...dule_schedule_rule_parallel_vectorize_unroll.py | 2 +-
...hedule_schedule_rule_random_compute_location.py | 2 +-
.../test_meta_schedule_search_strategy.py | 2 +-
.../meta_schedule/test_meta_schedule_space_cpu.py | 2 +-
.../meta_schedule/test_meta_schedule_space_cuda.py | 2 +-
.../test_meta_schedule_space_cuda_async.py | 2 +-
.../test_meta_schedule_space_generator.py | 2 +-
.../test_meta_schedule_space_post_opt.py | 2 +-
.../test_meta_schedule_task_scheduler.py | 2 +-
.../test_meta_schedule_trace_apply.py | 2 +-
.../test_meta_schedule_tune_context.py | 2 +-
.../meta_schedule/test_meta_schedule_tune_tir.py | 2 +-
.../s_tir/schedule/test_tir_schedule_analysis.py | 2 +-
.../test_tir_schedule_annotate_tensor_access.py | 2 +-
.../schedule/test_tir_schedule_block_scope.py | 2 +-
.../s_tir/schedule/test_tir_schedule_blockize.py | 2 +-
.../schedule/test_tir_schedule_cache_index.py | 2 +-
.../schedule/test_tir_schedule_cache_read_write.py | 2 +-
.../s_tir/schedule/test_tir_schedule_compute_at.py | 2 +-
.../schedule/test_tir_schedule_compute_inline.py | 2 +-
.../test_tir_schedule_decompose_padding.py | 2 +-
.../s_tir/schedule/test_tir_schedule_error.py | 2 +-
.../s_tir/schedule/test_tir_schedule_for_kind.py | 2 +-
.../test_tir_schedule_fuse_reduction_epilogue.py | 2 +-
...ir_schedule_fuse_reduction_epilogue_clipping.py | 2 +-
...st_tir_schedule_fuse_reduction_epilogue_relu.py | 2 +-
.../s_tir/schedule/test_tir_schedule_merge.py | 2 +-
.../s_tir/schedule/test_tir_schedule_pad_einsum.py | 2 +-
.../s_tir/schedule/test_tir_schedule_partition.py | 2 +-
.../schedule/test_tir_schedule_read_write_at.py | 2 +-
.../s_tir/schedule/test_tir_schedule_reduction.py | 2 +-
.../s_tir/schedule/test_tir_schedule_reindex.py | 2 +-
.../s_tir/schedule/test_tir_schedule_reorder.py | 2 +-
.../test_tir_schedule_reorder_block_iter_var.py | 2 +-
.../s_tir/schedule/test_tir_schedule_rfactor.py | 2 +-
.../schedule/test_tir_schedule_rolling_buffer.py | 2 +-
.../s_tir/schedule/test_tir_schedule_sampling.py | 2 +-
.../s_tir/schedule/test_tir_schedule_set_dtype.py | 2 +-
.../s_tir/schedule/test_tir_schedule_set_scope.py | 2 +-
.../s_tir/schedule/test_tir_schedule_split_fuse.py | 2 +-
.../s_tir/schedule/test_tir_schedule_state.py | 10 +--
.../test_tir_schedule_state_cached_flags.py | 2 +-
.../schedule/test_tir_schedule_storage_align.py | 2 +-
.../s_tir/schedule/test_tir_schedule_tensorize.py | 2 +-
.../s_tir/schedule/test_tir_schedule_trace.py | 2 +-
.../s_tir/schedule/test_tir_schedule_transform.py | 2 +-
.../schedule/test_tir_schedule_transform_layout.py | 2 +-
.../s_tir/schedule/test_tir_schedule_utilities.py | 2 +-
.../schedule/test_tir_unsafe_hide_tensor_access.py | 2 +-
.../s_tir/script/test_s_tir_script_basic_usage.py | 16 ++--
.../s_tir/script/test_s_tir_script_blocks.py | 40 ++++-----
.../script/test_s_tir_script_dynamic_shape.py | 8 +-
.../script/test_s_tir_script_error_handling.py | 2 +-
.../s_tir/script/test_s_tir_script_ir_builder.py | 2 +-
.../script/test_s_tir_script_meta_programming.py | 4 +-
.../s_tir/script/test_s_tir_script_namespace.py | 22 ++---
.../s_tir/script/test_s_tir_script_printer.py | 82 +++++++++---------
.../script/test_s_tir_script_source_locations.py | 2 +-
tests/python/s_tir/test_arith_domain_touched.py | 2 +-
tests/python/s_tir/test_s_tir_renew_defs.py | 2 +-
...test_s_tir_transform_annotate_irregular_loop.py | 2 +-
.../test_s_tir_transform_canonicalize_loop.py | 2 +-
.../test_s_tir_transform_compact_tensor_region.py | 2 +-
...est_s_tir_transform_convert_blocks_to_opaque.py | 2 +-
.../test_s_tir_transform_default_gpu_schedule.py | 2 +-
.../test_s_tir_transform_inject_permuted_layout.py | 2 +-
.../test_s_tir_transform_inject_ptx_async_copy.py | 4 +-
...est_s_tir_transform_inject_software_pipeline.py | 2 +-
.../test_s_tir_transform_inject_virtual_thread.py | 2 +-
.../test_s_tir_transform_lift_thread_binding.py | 2 +-
.../test_s_tir_transform_loop_partition.py | 2 +-
...s_tir_transform_lower_cross_thread_reduction.py | 2 +-
.../test_s_tir_transform_lower_init_block.py | 2 +-
.../test_s_tir_transform_lower_match_tensor.py | 2 +-
.../test_s_tir_transform_lower_opaque_block.py | 2 +-
...test_s_tir_transform_lower_thread_all_reduce.py | 2 +-
.../test_s_tir_transform_lower_thread_binding.py | 2 +-
...transform_manifest_shared_memory_local_stage.py | 2 +-
...st_s_tir_transform_memhammer_lower_auto_copy.py | 2 +-
...form_merge_dynamic_shared_memory_allocations.py | 2 +-
.../test_s_tir_transform_mma_tensor_layout.py | 2 +-
...sform_plan_update_tensor_allocation_location.py | 2 +-
.../transform/test_s_tir_transform_remove_undef.py | 2 +-
...transform_remove_weight_layout_rewrite_block.py | 2 +-
.../test_s_tir_transform_rewrite_unsafe_select.py | 2 +-
.../transform/test_s_tir_transform_thread_sync.py | 2 +-
tests/python/script/test_constructor_contracts.py | 2 +-
tests/python/target/test_arm_target.py | 2 +-
tests/python/target/test_target_target.py | 2 +-
tests/python/te/test_te_create_function.py | 4 +-
.../codegen/conftest.py} | 34 +++-----
tests/python/tirx/codegen/test_tirx_codegen_fp8.py | 2 +-
.../conftest.py} | 28 +------
.../tvm_tirx/testing/_pytest.py} | 34 ++++++++
.../python/runtime/test_runtime_memory_scope.py | 2 +-
tirx/tests/python/tirx/codegen/conftest.py | 43 +++-------
tirx/tests/python/tirx/conftest.py | 97 +---------------------
.../python/tirx/script/test_tirx_script_tracing.py | 2 +-
tirx/tests/python/tirx/test_cuda_arch_marker.py | 2 +-
tirx/tests/python/tirx/test_cuda_backend_config.py | 2 +-
.../transform/test_tirx_transform_lower_intrin.py | 4 +-
358 files changed, 642 insertions(+), 750 deletions(-)
diff --git a/tests/python/contrib/test_android/test_meta_schedule.py b/tests/python/contrib/test_android/test_meta_schedule.py
index 4d28c6d786..1eab8b2b13 100644
--- a/tests/python/contrib/test_android/test_meta_schedule.py
+++ b/tests/python/contrib/test_android/test_meta_schedule.py
@@ -29,7 +29,7 @@ import tvm_tirx.testing, tvm
import tvm.topi.testing
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.builder import LocalBuilder
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from .infrastructure import get_android_gpu_target, get_rpc_runner
diff --git a/tests/python/contrib/test_example_npu.py b/tests/python/contrib/test_example_npu.py
index b0cee54f5a..9617e7a745 100644
--- a/tests/python/contrib/test_example_npu.py
+++ b/tests/python/contrib/test_example_npu.py
@@ -29,7 +29,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.backend.pattern_registry import get_patterns_with_prefix
from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions, RunCodegen
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
@tvm_tirx.script.ir_module
diff --git a/tests/python/contrib/test_tir_triton_integration.py b/tests/python/contrib/test_tir_triton_integration.py
index d4532f570a..ce5fb0122e 100644
--- a/tests/python/contrib/test_tir_triton_integration.py
+++ b/tests/python/contrib/test_tir_triton_integration.py
@@ -25,8 +25,8 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.frontend import nn
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/disco/test_callback.py b/tests/python/disco/test_callback.py
index 55d362017e..403cdc9c89 100644
--- a/tests/python/disco/test_callback.py
+++ b/tests/python/disco/test_callback.py
@@ -26,8 +26,8 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/disco/test_ccl.py b/tests/python/disco/test_ccl.py
index 56ef4b2791..64d091de3d 100644
--- a/tests/python/disco/test_ccl.py
+++ b/tests/python/disco/test_ccl.py
@@ -29,7 +29,7 @@ from tvm import relax as rx
from tvm.runtime import disco as di
from tvm.runtime.vm import VirtualMachine
from tvm.s_tir import dlight as dl
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
if di is None:
pytest.skip("disco runtime is not available", allow_module_level=True)
diff --git a/tests/python/disco/test_loader.py b/tests/python/disco/test_loader.py
index f4299df4b1..017239a1e4 100644
--- a/tests/python/disco/test_loader.py
+++ b/tests/python/disco/test_loader.py
@@ -32,7 +32,7 @@ from tvm.contrib import tvmjs
from tvm.runtime import disco as di
from tvm.s_tir import dlight as dl
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.target import Target
from tvm_tirx.testing import env
diff --git a/tests/python/disco/test_nvshmem.py b/tests/python/disco/test_nvshmem.py
index d0b1ca58c3..2b8b4a8909 100644
--- a/tests/python/disco/test_nvshmem.py
+++ b/tests/python/disco/test_nvshmem.py
@@ -34,8 +34,8 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir import TupleType as _TupleType
from tvm.runtime import disco as di
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/disco/test_session.py b/tests/python/disco/test_session.py
index 61c672804c..375c2d3d6b 100644
--- a/tests/python/disco/test_session.py
+++ b/tests/python/disco/test_session.py
@@ -35,8 +35,8 @@ import tvm_tirx.testing, tvm
from tvm.exec import disco_worker as _ # noqa: F401 # pylint: disable=unused-import
from tvm.runtime import disco as di
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
if di is None:
diff --git a/tests/python/driver/test_compile.py b/tests/python/driver/test_compile.py
index cba770948a..b71ba67b0a 100644
--- a/tests/python/driver/test_compile.py
+++ b/tests/python/driver/test_compile.py
@@ -25,8 +25,8 @@ import tvm_tirx.testing, tvm
from tvm import relax, te
from tvm_tirx.runtime import Executable
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/backend/adreno/mod_utils.py b/tests/python/relax/backend/adreno/mod_utils.py
index f239adf24e..fd3e2fe4a8 100644
--- a/tests/python/relax/backend/adreno/mod_utils.py
+++ b/tests/python/relax/backend/adreno/mod_utils.py
@@ -28,8 +28,8 @@ from tvm import relax, rpc
from tvm.relax.backend.adreno import clml
from tvm.relax.script import ir_builder as relax_builder
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/backend/adreno/test_clml_ops.py b/tests/python/relax/backend/adreno/test_clml_ops.py
index 45cb5e775f..08ade82ee7 100644
--- a/tests/python/relax/backend/adreno/test_clml_ops.py
+++ b/tests/python/relax/backend/adreno/test_clml_ops.py
@@ -53,7 +53,7 @@ from tvm.relax.backend.adreno import clml
from tvm.relax.backend.adreno.clml import OpenCLMLOffLoad, OpenCLMLOffLoadForLLM
from tvm.relax.script import ir_builder as relax_builder
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/backend/adreno/test_texture_network.py b/tests/python/relax/backend/adreno/test_texture_network.py
index aeeb39ac76..8039372deb 100644
--- a/tests/python/relax/backend/adreno/test_texture_network.py
+++ b/tests/python/relax/backend/adreno/test_texture_network.py
@@ -34,7 +34,7 @@ from tvm.relax.frontend.onnx import from_onnx
from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.transform.legalize_ops import adreno as legalize_adreno
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/backend/adreno/test_transform_fold_vdevice_scope_change.py b/tests/python/relax/backend/adreno/test_transform_fold_vdevice_scope_change.py
index bec2900153..1fa999d25d 100644
--- a/tests/python/relax/backend/adreno/test_transform_fold_vdevice_scope_change.py
+++ b/tests/python/relax/backend/adreno/test_transform_fold_vdevice_scope_change.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir.module import IRModule
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/distributed/test_distributed_transform_lower_distir.py b/tests/python/relax/distributed/test_distributed_transform_lower_distir.py
index fadbd12c3a..f5fe112c3f 100644
--- a/tests/python/relax/distributed/test_distributed_transform_lower_distir.py
+++ b/tests/python/relax/distributed/test_distributed_transform_lower_distir.py
@@ -24,7 +24,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.ir import assert_structural_equal
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py b/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
index a121224e07..891e201d7f 100644
--- a/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
+++ b/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
@@ -23,7 +23,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir import assert_structural_equal
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py b/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
index 4963526601..3ec43750ed 100644
--- a/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
+++ b/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
@@ -25,7 +25,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.ir import assert_structural_equal
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/script/test_relax_script_basic_usage.py b/tests/python/relax/script/test_relax_script_basic_usage.py
index b28d9f0240..8f631bd0ee 100644
--- a/tests/python/relax/script/test_relax_script_basic_usage.py
+++ b/tests/python/relax/script/test_relax_script_basic_usage.py
@@ -32,8 +32,8 @@ from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import DummyGlobalInfo, VDevice
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
@@ -46,9 +46,9 @@ def _check(
test,
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
@@ -1550,8 +1550,8 @@ def test_roundtrip_basic_usage(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
diff --git a/tests/python/relax/script/test_relax_script_distributed.py b/tests/python/relax/script/test_relax_script_distributed.py
index 85177dc543..009ebe93d9 100644
--- a/tests/python/relax/script/test_relax_script_distributed.py
+++ b/tests/python/relax/script/test_relax_script_distributed.py
@@ -27,8 +27,8 @@ from tvm_tirx.ir import Range
from tvm.relax import Call, SeqExpr, VarBinding
from tvm.relax.distributed import DeviceMesh
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
@@ -41,9 +41,9 @@ def _check(
test,
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
diff --git a/tests/python/relax/script/test_relax_script_dynamic_shape.py b/tests/python/relax/script/test_relax_script_dynamic_shape.py
index 72abb44241..4d1760f972 100644
--- a/tests/python/relax/script/test_relax_script_dynamic_shape.py
+++ b/tests/python/relax/script/test_relax_script_dynamic_shape.py
@@ -30,8 +30,8 @@ from tvm_tirx import IRModule, tirx
from tvm import relax
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
M = T.dynamic("M", "int64")
@@ -59,7 +59,7 @@ def test_type_vars_roundtrip():
def func[M: T.int64](x: R.Tensor((M, M * 2), "float32")):
return x
""",
- extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx},
+ extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx},
)
tvm_tirx.ir.assert_structural_equal(func, typed)
assert 'M = T.dynamic("M", dtype="int64")' not in script
@@ -76,13 +76,13 @@ def func[M: T.int64](x: R.Tensor((M, M * 2), "float32")):
tvm_tirx.ir.assert_structural_equal(
func,
tvm_tirx.script.from_source(
- script, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ script, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
),
)
tvm_tirx.ir.assert_structural_equal(
func,
tvm_tirx.script.from_source(
- repeated, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ repeated, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
),
)
@@ -105,7 +105,7 @@ def test_dynamic_module_symbol_identity():
restored = tvm_tirx.script.from_source(
source,
check_well_formed=False,
- extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx},
+ extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx},
)
first_n = restored["first"].params[0].ty.shape.values[0]
second_shape = restored["second"].params[0].ty.shape.values
@@ -124,9 +124,9 @@ def _check(
test,
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
@@ -250,9 +250,9 @@ def test_call_tir_with_tir_var():
repeated,
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
),
)
@@ -651,8 +651,8 @@ def test_roundtrip_dynamic_shape(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
diff --git a/tests/python/relax/script/test_relax_script_error_handling.py b/tests/python/relax/script/test_relax_script_error_handling.py
index da44aa11c3..e32592fe51 100644
--- a/tests/python/relax/script/test_relax_script_error_handling.py
+++ b/tests/python/relax/script/test_relax_script_error_handling.py
@@ -31,8 +31,8 @@ from tvm import relax
from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
@@ -45,9 +45,9 @@ def _check(
test,
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
diff --git a/tests/python/relax/script/test_relax_script_ir_builder.py b/tests/python/relax/script/test_relax_script_ir_builder.py
index 9121507813..4ec4557713 100644
--- a/tests/python/relax/script/test_relax_script_ir_builder.py
+++ b/tests/python/relax/script/test_relax_script_ir_builder.py
@@ -193,7 +193,7 @@ def test_function_subroutine_before_main():
"""The block builder can generate subroutines, and calls into subroutines"""
from tvm_tirx.script import ir as I
- from tvm_tirx.script import relax as R
+ from tvm.relax import script as R
# create with TVMScript
@I.ir_module
@@ -235,7 +235,7 @@ def test_function_subroutine_during_main():
"""Subroutines may be generated as needed, pausing the main function collection"""
from tvm_tirx.script import ir as I
- from tvm_tirx.script import relax as R
+ from tvm.relax import script as R
# create with TVMScript
@I.ir_module
diff --git a/tests/python/relax/script/test_relax_script_meta_programming.py b/tests/python/relax/script/test_relax_script_meta_programming.py
index 006d4435d9..557514d830 100644
--- a/tests/python/relax/script/test_relax_script_meta_programming.py
+++ b/tests/python/relax/script/test_relax_script_meta_programming.py
@@ -30,8 +30,8 @@ import tvm_tirx.testing, tvm
from tvm_tirx import IRModule, ir
from tvm import relax, topi
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
@@ -44,9 +44,9 @@ def _check(
test,
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
diff --git a/tests/python/relax/script/test_relax_script_operators_arith_cmp.py b/tests/python/relax/script/test_relax_script_operators_arith_cmp.py
index b0228601e4..3a14bbb24e 100644
--- a/tests/python/relax/script/test_relax_script_operators_arith_cmp.py
+++ b/tests/python/relax/script/test_relax_script_operators_arith_cmp.py
@@ -25,7 +25,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -34,7 +34,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_create.py b/tests/python/relax/script/test_relax_script_operators_create.py
index 89adacd68a..043ecf19e1 100644
--- a/tests/python/relax/script/test_relax_script_operators_create.py
+++ b/tests/python/relax/script/test_relax_script_operators_create.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_datatype.py b/tests/python/relax/script/test_relax_script_operators_datatype.py
index 2026ff5351..f394203bf4 100644
--- a/tests/python/relax/script/test_relax_script_operators_datatype.py
+++ b/tests/python/relax/script/test_relax_script_operators_datatype.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_grad.py b/tests/python/relax/script/test_relax_script_operators_grad.py
index 62fd09d7a4..bd43f81248 100644
--- a/tests/python/relax/script/test_relax_script_operators_grad.py
+++ b/tests/python/relax/script/test_relax_script_operators_grad.py
@@ -20,7 +20,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -29,7 +29,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_image.py b/tests/python/relax/script/test_relax_script_operators_image.py
index d59618d804..53cd473565 100644
--- a/tests/python/relax/script/test_relax_script_operators_image.py
+++ b/tests/python/relax/script/test_relax_script_operators_image.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_index.py b/tests/python/relax/script/test_relax_script_operators_index.py
index ea96339137..63bb8aa7c8 100644
--- a/tests/python/relax/script/test_relax_script_operators_index.py
+++ b/tests/python/relax/script/test_relax_script_operators_index.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_linear_algebra.py b/tests/python/relax/script/test_relax_script_operators_linear_algebra.py
index f5d33ebdff..0d7f37000d 100644
--- a/tests/python/relax/script/test_relax_script_operators_linear_algebra.py
+++ b/tests/python/relax/script/test_relax_script_operators_linear_algebra.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_manipulate.py b/tests/python/relax/script/test_relax_script_operators_manipulate.py
index 8fa3d5b380..5119f8f6cc 100644
--- a/tests/python/relax/script/test_relax_script_operators_manipulate.py
+++ b/tests/python/relax/script/test_relax_script_operators_manipulate.py
@@ -22,7 +22,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -31,7 +31,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_nn.py b/tests/python/relax/script/test_relax_script_operators_nn.py
index a81e09ae46..4268d3b7f3 100644
--- a/tests/python/relax/script/test_relax_script_operators_nn.py
+++ b/tests/python/relax/script/test_relax_script_operators_nn.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_search.py b/tests/python/relax/script/test_relax_script_operators_search.py
index ba99c0494a..9da26aee1d 100644
--- a/tests/python/relax/script/test_relax_script_operators_search.py
+++ b/tests/python/relax/script/test_relax_script_operators_search.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_set.py b/tests/python/relax/script/test_relax_script_operators_set.py
index 44f509873d..7781de2e9c 100644
--- a/tests/python/relax/script/test_relax_script_operators_set.py
+++ b/tests/python/relax/script/test_relax_script_operators_set.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_sort.py b/tests/python/relax/script/test_relax_script_operators_sort.py
index bd263acd70..15cf8f6f48 100644
--- a/tests/python/relax/script/test_relax_script_operators_sort.py
+++ b/tests/python/relax/script/test_relax_script_operators_sort.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_statistical.py b/tests/python/relax/script/test_relax_script_operators_statistical.py
index 12bdd5dcfd..569c230b06 100644
--- a/tests/python/relax/script/test_relax_script_operators_statistical.py
+++ b/tests/python/relax/script/test_relax_script_operators_statistical.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_operators_vision.py b/tests/python/relax/script/test_relax_script_operators_vision.py
index e9ae12fa83..666ee0a0de 100644
--- a/tests/python/relax/script/test_relax_script_operators_vision.py
+++ b/tests/python/relax/script/test_relax_script_operators_vision.py
@@ -21,7 +21,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check(
@@ -30,7 +30,7 @@ def _check(
):
test = parsed.script(show_meta=True)
roundtrip_mod = tvm_tirx.script.from_source(
- test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm_tirx.script.relax, "T": tvm_tirx.script.tirx}
+ test, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
)
tvm_tirx.ir.assert_structural_equal(parsed, roundtrip_mod)
if expect:
diff --git a/tests/python/relax/script/test_relax_script_printer.py b/tests/python/relax/script/test_relax_script_printer.py
index e690435747..d6abc532ed 100644
--- a/tests/python/relax/script/test_relax_script_printer.py
+++ b/tests/python/relax/script/test_relax_script_printer.py
@@ -32,8 +32,8 @@ from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax import TensorType
from tvm.relax.distributed import DeviceMesh, DTensorType, Placement
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.script.printer.scriptable import PrinterConfig, _script
@@ -166,7 +166,7 @@ def test_func():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@R.function
@@ -183,8 +183,8 @@ def test_module():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.relax import script as R
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -234,7 +234,7 @@ def test_function():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
@R.function
def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):
@@ -315,7 +315,7 @@ def test_lone_private_function():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
@R.function(private=True)
def main(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):
@@ -340,7 +340,7 @@ def test_extern_func():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -366,7 +366,7 @@ def test_extern_func_with_ty():
_assert_print_lines(
obj,
"""
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -389,9 +389,9 @@ def test_extern_func_with_ty_roundtrip():
mod.script(verbose_expr=True),
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
)
tvm_tirx.ir.assert_structural_equal(mod, roundtrip)
@@ -420,7 +420,7 @@ def test_nested_function(relax_prefix, comment_imports):
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -581,7 +581,7 @@ def test_prim_value():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@R.function
@@ -601,9 +601,9 @@ def func() -> T.int64:
float_script,
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
),
float_func,
@@ -627,9 +627,9 @@ def test_primitive_bindings_roundtrip_without_prim_value_marker():
source,
extra_vars={
"I": tvm_tirx.script.ir,
- "R": tvm_tirx.script.relax,
+ "R": tvm.relax.script,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
+ "Ts": tvm.s_tir.script,
},
),
func,
@@ -959,8 +959,8 @@ def test_module_cross_func_call():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.relax import script as R
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -983,8 +983,8 @@ class Module:
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.relax import script as R
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -1014,7 +1014,7 @@ def test_assert_op():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -1040,7 +1040,7 @@ def test_print():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -1066,7 +1066,7 @@ def test_private_function():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -1121,7 +1121,7 @@ def test_directly_construct_private_funcs():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -1168,7 +1168,7 @@ def test_reused_extern_func():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
@R.function
def func(x: R.Tensor((128, 128), dtype="float32")) -> R.Tensor((128, 128), dtype="float32"):
@@ -1198,7 +1198,7 @@ def test_inline_extern_func():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
@R.function
def func(x: R.Tensor((128, 128), dtype="float32")) -> R.Tensor((128, 128), dtype="float32"):
@@ -1248,7 +1248,7 @@ def test_hide_inferable_ty():
"""
from __future__ import annotations
-# from tvm_tirx.script import relax as R
+# from tvm.relax import script as R
@R.function
def func(A: R.Tensor((10, 20), dtype="float32"), B: R.Tensor(dtype="float32", ndim=2)) -> R.Tensor((10, 20), dtype="float32"):
@@ -1286,8 +1286,8 @@ def test_extern_func_roundtrip(show_all_ty):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
diff --git a/tests/python/relax/script/test_relax_script_pyfunc.py b/tests/python/relax/script/test_relax_script_pyfunc.py
index 8d298f30c0..38cfd0ed8f 100644
--- a/tests/python/relax/script/test_relax_script_pyfunc.py
+++ b/tests/python/relax/script/test_relax_script_pyfunc.py
@@ -23,8 +23,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax import BasePyModule
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
n = T.dynamic("n", "int32")
diff --git a/tests/python/relax/test_analysis.py b/tests/python/relax/test_analysis.py
index ac5fe5aece..b6f4edf585 100644
--- a/tests/python/relax/test_analysis.py
+++ b/tests/python/relax/test_analysis.py
@@ -34,8 +34,8 @@ from tvm.relax.analysis import (
used_vars,
)
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_computable_at_compile_time.py b/tests/python/relax/test_analysis_computable_at_compile_time.py
index 5a9f496aca..f44ac42d0a 100644
--- a/tests/python/relax/test_analysis_computable_at_compile_time.py
+++ b/tests/python/relax/test_analysis_computable_at_compile_time.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_contains_impure_call.py b/tests/python/relax/test_analysis_contains_impure_call.py
index 1f3d2158cd..a721f4f595 100644
--- a/tests/python/relax/test_analysis_contains_impure_call.py
+++ b/tests/python/relax/test_analysis_contains_impure_call.py
@@ -20,7 +20,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax as rx
from tvm.relax.analysis import contains_impure_call
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_simple_pure_case():
diff --git a/tests/python/relax/test_analysis_detect_recursion.py b/tests/python/relax/test_analysis_detect_recursion.py
index f90f0bfe1c..b2ca5889c6 100644
--- a/tests/python/relax/test_analysis_detect_recursion.py
+++ b/tests/python/relax/test_analysis_detect_recursion.py
@@ -19,8 +19,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax as rx
from tvm.relax.analysis import detect_recursion
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_estimate_memory_usage.py b/tests/python/relax/test_analysis_estimate_memory_usage.py
index 7396cc3bc1..615053ac15 100644
--- a/tests/python/relax/test_analysis_estimate_memory_usage.py
+++ b/tests/python/relax/test_analysis_estimate_memory_usage.py
@@ -19,8 +19,8 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax.analysis import estimate_memory_usage
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_suggest_layout_transforms.py b/tests/python/relax/test_analysis_suggest_layout_transforms.py
index cb873f83fc..3ead1c5b91 100644
--- a/tests/python/relax/test_analysis_suggest_layout_transforms.py
+++ b/tests/python/relax/test_analysis_suggest_layout_transforms.py
@@ -21,7 +21,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import relax, s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_type_analysis.py b/tests/python/relax/test_analysis_type_analysis.py
index 2dcf60f998..e1a3b9ea9e 100644
--- a/tests/python/relax/test_analysis_type_analysis.py
+++ b/tests/python/relax/test_analysis_type_analysis.py
@@ -26,7 +26,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import ir, tirx
from tvm import relax as rx
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_well_formed.py b/tests/python/relax/test_analysis_well_formed.py
index 4ca5ae5ff8..1f7ed61172 100644
--- a/tests/python/relax/test_analysis_well_formed.py
+++ b/tests/python/relax/test_analysis_well_formed.py
@@ -27,8 +27,8 @@ from tvm_tirx import tirx
from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
m = T.dynamic("m", "int64")
diff --git a/tests/python/relax/test_ast_printer.py b/tests/python/relax/test_ast_printer.py
index cb3be30e60..e784134a05 100644
--- a/tests/python/relax/test_ast_printer.py
+++ b/tests/python/relax/test_ast_printer.py
@@ -27,8 +27,8 @@ from tvm_tirx import tirx
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.testing import dump_ast
from tvm.relax.testing.ast_printer import ASTPrinter
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# Overload dump_ast to test both type and type annotations
diff --git a/tests/python/relax/test_backend_dispatch_sampling.py b/tests/python/relax/test_backend_dispatch_sampling.py
index df5424f331..53cce1dea3 100644
--- a/tests/python/relax/test_backend_dispatch_sampling.py
+++ b/tests/python/relax/test_backend_dispatch_sampling.py
@@ -23,8 +23,8 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.backend import DispatchSampling
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py b/tests/python/relax/test_backend_dispatch_sort_scan.py
index c8446535a3..bf8709e142 100644
--- a/tests/python/relax/test_backend_dispatch_sort_scan.py
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -31,7 +31,7 @@ from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.backend import DispatchSortScan
from tvm.s_tir import dlight
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_backend_transform_shape_lower.py b/tests/python/relax/test_backend_transform_shape_lower.py
index 6739031b33..eaa9bb6a4f 100644
--- a/tests/python/relax/test_backend_transform_shape_lower.py
+++ b/tests/python/relax/test_backend_transform_shape_lower.py
@@ -22,8 +22,8 @@ from tvm import relax
from tvm_tirx.ir import assert_structural_equal
from tvm.relax.testing.runtime_builtin import MakeShapeCode, MatchShapeCode
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# note: we expected RemovePurityChecking to be run first, so we force purity in most test cases
diff --git a/tests/python/relax/test_base_py_module_printer.py b/tests/python/relax/test_base_py_module_printer.py
index fa6aef20d7..71211c32ee 100644
--- a/tests/python/relax/test_base_py_module_printer.py
+++ b/tests/python/relax/test_base_py_module_printer.py
@@ -23,8 +23,8 @@ import tvm, tvm_tirx
from tvm_tirx.ir import StringImm
from tvm.relax.base_py_module import BasePyModule
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_base_py_module_symbolic_shape.py b/tests/python/relax/test_base_py_module_symbolic_shape.py
index 9301fe53d7..a5a5398cb7 100644
--- a/tests/python/relax/test_base_py_module_symbolic_shape.py
+++ b/tests/python/relax/test_base_py_module_symbolic_shape.py
@@ -24,8 +24,8 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import IRModule
from tvm.relax.base_py_module import BasePyModule
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_bind_params.py b/tests/python/relax/test_bind_params.py
index 4d467a7060..9d6a205e35 100644
--- a/tests/python/relax/test_bind_params.py
+++ b/tests/python/relax/test_bind_params.py
@@ -24,7 +24,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
param_specification = tvm_tirx.testing.parameter("by_string", "by_var")
param_shape = tvm_tirx.testing.parameter("static_shape", "dynamic_shape", "ndim", "arbitrary")
diff --git a/tests/python/relax/test_bind_symbolic_vars.py b/tests/python/relax/test_bind_symbolic_vars.py
index e5640636db..00d2790fa1 100644
--- a/tests/python/relax/test_bind_symbolic_vars.py
+++ b/tests/python/relax/test_bind_symbolic_vars.py
@@ -19,7 +19,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
replace_by_tir_var = tvm_tirx.testing.parameter(
diff --git a/tests/python/relax/test_binding_rewrite.py b/tests/python/relax/test_binding_rewrite.py
index b8771fff53..7b37c6b908 100644
--- a/tests/python/relax/test_binding_rewrite.py
+++ b/tests/python/relax/test_binding_rewrite.py
@@ -24,7 +24,7 @@ from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.analysis import name_to_binding
from tvm.relax.binding_rewrite import DataflowBlockRewrite
from tvm.relax.expr import DataflowVar, Var
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
@tvm_tirx.script.ir_module
diff --git a/tests/python/relax/test_blockbuilder_core.py b/tests/python/relax/test_blockbuilder_core.py
index 48a11a655a..6dee48dd1b 100644
--- a/tests/python/relax/test_blockbuilder_core.py
+++ b/tests/python/relax/test_blockbuilder_core.py
@@ -29,7 +29,7 @@ from tvm import te, topi
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax import ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx.function import Function
diff --git a/tests/python/relax/test_blockbuilder_emit_te.py b/tests/python/relax/test_blockbuilder_emit_te.py
index b94671b445..08c8c7bf4c 100644
--- a/tests/python/relax/test_blockbuilder_emit_te.py
+++ b/tests/python/relax/test_blockbuilder_emit_te.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
from tvm import relax as rx
from tvm import te
from tvm_tirx.ir.base import assert_structural_equal
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_codegen_cublas.py b/tests/python/relax/test_codegen_cublas.py
index 8c403286ed..f0ee137d19 100644
--- a/tests/python/relax/test_codegen_cublas.py
+++ b/tests/python/relax/test_codegen_cublas.py
@@ -29,7 +29,7 @@ from tvm import relax
from tvm.relax.backend.cuda.cublas import partition_for_cublas
from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.testing import get_relax_matmul_module
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script.ir_builder import IRBuilder
try:
diff --git a/tests/python/relax/test_codegen_cudnn.py b/tests/python/relax/test_codegen_cudnn.py
index 407a31086a..2bf2d540ab 100644
--- a/tests/python/relax/test_codegen_cudnn.py
+++ b/tests/python/relax/test_codegen_cudnn.py
@@ -30,7 +30,7 @@ from tvm.contrib.pickle_memoize import memoize
from tvm.relax.backend.cuda.cudnn import partition_for_cudnn
from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.testing import get_relax_stacked_attention_module
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/test_codegen_cutlass.py b/tests/python/relax/test_codegen_cutlass.py
index a87c6bab2e..f1f50ec6e7 100644
--- a/tests/python/relax/test_codegen_cutlass.py
+++ b/tests/python/relax/test_codegen_cutlass.py
@@ -38,8 +38,8 @@ from tvm.relax.testing import (
get_relax_stacked_attention_module,
)
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/test_codegen_dnnl.py b/tests/python/relax/test_codegen_dnnl.py
index 2d03350ee9..e3c7a38779 100644
--- a/tests/python/relax/test_codegen_dnnl.py
+++ b/tests/python/relax/test_codegen_dnnl.py
@@ -22,7 +22,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.contrib.pickle_memoize import memoize
from tvm.relax.dpl import make_fused_bias_activation_pattern
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
@tvm_tirx.script.ir_module
diff --git a/tests/python/relax/test_codegen_hipblas.py b/tests/python/relax/test_codegen_hipblas.py
index cba5382b32..26c47fa4a1 100644
--- a/tests/python/relax/test_codegen_hipblas.py
+++ b/tests/python/relax/test_codegen_hipblas.py
@@ -27,7 +27,7 @@ import tvm.topi.testing
from tvm import relax
from tvm.relax.backend.rocm.hipblas import partition_for_hipblas
from tvm.relax.testing import get_relax_matmul_module
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
try:
import ml_dtypes
diff --git a/tests/python/relax/test_codegen_tensorrt.py b/tests/python/relax/test_codegen_tensorrt.py
index 3f89e5437f..43a7d3b643 100644
--- a/tests/python/relax/test_codegen_tensorrt.py
+++ b/tests/python/relax/test_codegen_tensorrt.py
@@ -23,7 +23,7 @@ from tvm import relax
from tvm.contrib.pickle_memoize import memoize
from tvm.relax.backend.contrib.tensorrt import partition_for_tensorrt
from tvm.relax.dpl import is_op, make_fused_bias_activation_pattern, wildcard
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_contrib_vllm.py b/tests/python/relax/test_contrib_vllm.py
index b1891a3064..53ee0475c0 100644
--- a/tests/python/relax/test_contrib_vllm.py
+++ b/tests/python/relax/test_contrib_vllm.py
@@ -23,7 +23,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_dataflow_inplace.py b/tests/python/relax/test_dataflow_inplace.py
index e06f1680a7..0e71f2bb4a 100644
--- a/tests/python/relax/test_dataflow_inplace.py
+++ b/tests/python/relax/test_dataflow_inplace.py
@@ -31,7 +31,7 @@ from tvm.relax.testing.transform import (
dataflow_single_inplace_call,
)
from tvm.relax.transform import DataflowUseInplaceCalls
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_dataflow_pattern.py b/tests/python/relax/test_dataflow_pattern.py
index 629cb2bf59..d640bcbf65 100644
--- a/tests/python/relax/test_dataflow_pattern.py
+++ b/tests/python/relax/test_dataflow_pattern.py
@@ -31,8 +31,8 @@ from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.analysis import get_var2val
from tvm.relax.dpl import *
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_dataflow_rewriter.py b/tests/python/relax/test_dataflow_rewriter.py
index bce28e456c..5c60778706 100644
--- a/tests/python/relax/test_dataflow_rewriter.py
+++ b/tests/python/relax/test_dataflow_rewriter.py
@@ -20,8 +20,8 @@ import pytest
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_dlpack_integration.py b/tests/python/relax/test_dlpack_integration.py
index 8f4b7786df..45427c6aaa 100644
--- a/tests/python/relax/test_dlpack_integration.py
+++ b/tests/python/relax/test_dlpack_integration.py
@@ -30,8 +30,8 @@ import torch
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax import BasePyModule
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_e2e_op_dynamic.py b/tests/python/relax/test_e2e_op_dynamic.py
index 3f3d1ba6e4..fb111c89a6 100644
--- a/tests/python/relax/test_e2e_op_dynamic.py
+++ b/tests/python/relax/test_e2e_op_dynamic.py
@@ -27,7 +27,7 @@ import tvm.topi.testing
from tvm_tirx import tirx
from tvm import relax
from tvm.relax.transform import LegalizeOps
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
# TODO(tvm-team): `tirx.transform.DefaultGPUSchedule` does not work.
diff --git a/tests/python/relax/test_expr.py b/tests/python/relax/test_expr.py
index e054ffab1b..964c605e17 100644
--- a/tests/python/relax/test_expr.py
+++ b/tests/python/relax/test_expr.py
@@ -25,7 +25,7 @@ import tvm, tvm_tirx
from tvm import relax as rx
from tvm_tirx import tirx
from tvm.relax.expr import make_shape
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check_equal(x, y, map_free_vars=False):
diff --git a/tests/python/relax/test_expr_functor.py b/tests/python/relax/test_expr_functor.py
index cf8302db39..d8b563410a 100644
--- a/tests/python/relax/test_expr_functor.py
+++ b/tests/python/relax/test_expr_functor.py
@@ -41,7 +41,7 @@ from tvm.relax.expr import (
Var,
VarBinding,
)
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
m, n = tirx.Var("m", "int64"), tirx.Var("n", "int64")
x = relax.Var("x", R.Tensor([n], "float32"))
diff --git a/tests/python/relax/test_fast_math_transform.py b/tests/python/relax/test_fast_math_transform.py
index 8a2e6e4b88..51804983c7 100644
--- a/tests/python/relax/test_fast_math_transform.py
+++ b/tests/python/relax/test_fast_math_transform.py
@@ -24,7 +24,7 @@ from tvm import relax, topi
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.transform import FastMathTransform
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _run_pass_compare_output(Before, Expected):
diff --git a/tests/python/relax/test_frontend_common.py b/tests/python/relax/test_frontend_common.py
index 8420b0efdf..d49903c5b3 100644
--- a/tests/python/relax/test_frontend_common.py
+++ b/tests/python/relax/test_frontend_common.py
@@ -20,7 +20,7 @@ from tvm import relax
from tvm.relax.frontend import detach_params
from tvm.relax.frontend.common import autopad
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.script.parser import relax as R
diff --git a/tests/python/relax/test_frontend_dynamo.py b/tests/python/relax/test_frontend_dynamo.py
index 722bd6636c..b69c5b57fc 100644
--- a/tests/python/relax/test_frontend_dynamo.py
+++ b/tests/python/relax/test_frontend_dynamo.py
@@ -31,8 +31,8 @@ from tvm import relax
from tvm.relax.frontend.torch import relax_dynamo
from tvm.s_tir import meta_schedule as ms
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py
index 87978cb483..cbe277fb11 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -29,7 +29,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.frontend.torch import from_exported_program
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py
index 556ba02ceb..f84f7ee92c 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -29,7 +29,7 @@ from tvm import relax
from tvm.relax.frontend import detach_params
from tvm.relax.frontend.torch import from_fx
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_frontend_nn_exporter.py b/tests/python/relax/test_frontend_nn_exporter.py
index 67fc9362e6..86d34e84aa 100644
--- a/tests/python/relax/test_frontend_nn_exporter.py
+++ b/tests/python/relax/test_frontend_nn_exporter.py
@@ -24,7 +24,7 @@ from tvm import relax
from tvm_tirx.ir import assert_structural_equal
from tvm.relax.frontend import nn
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_frontend_nn_extern_module.py b/tests/python/relax/test_frontend_nn_extern_module.py
index a8c3c9c0f4..06cf8f9908 100644
--- a/tests/python/relax/test_frontend_nn_extern_module.py
+++ b/tests/python/relax/test_frontend_nn_extern_module.py
@@ -82,7 +82,7 @@ def _test_infer_sym(func, x, y, z): # pylint: disable=invalid-name
def _check_ir_equality(mod):
# pylint: disable=import-outside-toplevel
from tvm_tirx.script import ir as I
- from tvm_tirx.script import relax as R
+ from tvm.relax import script as R
from tvm_tirx.script import tirx as T
# pylint: enable=import-outside-toplevel
diff --git a/tests/python/relax/test_frontend_nn_modules.py b/tests/python/relax/test_frontend_nn_modules.py
index 118914384d..ab6d1d9b52 100644
--- a/tests/python/relax/test_frontend_nn_modules.py
+++ b/tests/python/relax/test_frontend_nn_modules.py
@@ -26,7 +26,7 @@ from tvm_tirx.ir import assert_structural_equal
from tvm.relax.frontend import nn
from tvm.relax.frontend.nn import core, modules, spec
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_frontend_nn_op.py b/tests/python/relax/test_frontend_nn_op.py
index 4b52ff7f41..676692730c 100644
--- a/tests/python/relax/test_frontend_nn_op.py
+++ b/tests/python/relax/test_frontend_nn_op.py
@@ -27,8 +27,8 @@ from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.frontend.nn import Module, Tensor, op, spec
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_frontend_nn_packing.py b/tests/python/relax/test_frontend_nn_packing.py
index 078cfdf66b..6342c17bd3 100644
--- a/tests/python/relax/test_frontend_nn_packing.py
+++ b/tests/python/relax/test_frontend_nn_packing.py
@@ -18,7 +18,7 @@
import tvm, tvm_tirx
from tvm.relax.frontend import nn
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _iter_binding_names(mod):
diff --git a/tests/python/relax/test_frontend_nn_subroutines.py b/tests/python/relax/test_frontend_nn_subroutines.py
index 0aa6a31532..11555869ca 100644
--- a/tests/python/relax/test_frontend_nn_subroutines.py
+++ b/tests/python/relax/test_frontend_nn_subroutines.py
@@ -21,7 +21,7 @@ from tvm import relax
from tvm_tirx.ir import assert_structural_equal
from tvm.relax.frontend import nn
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_linear():
diff --git a/tests/python/relax/test_frontend_nn_tensor.py b/tests/python/relax/test_frontend_nn_tensor.py
index 5e2414fdb9..01e22cea37 100644
--- a/tests/python/relax/test_frontend_nn_tensor.py
+++ b/tests/python/relax/test_frontend_nn_tensor.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.frontend.nn import Module, Tensor, spec
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_tensor_from_numpy():
diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py
index edafb8124e..2956cdc0b2 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -44,8 +44,8 @@ from tvm_tirx import tirx
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
bg = np.random.MT19937(0)
diff --git a/tests/python/relax/test_frontend_stablehlo.py b/tests/python/relax/test_frontend_stablehlo.py
index 81f95ce0d7..adec4238ef 100644
--- a/tests/python/relax/test_frontend_stablehlo.py
+++ b/tests/python/relax/test_frontend_stablehlo.py
@@ -33,7 +33,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.frontend.stablehlo import from_stablehlo
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_inline_functions.py b/tests/python/relax/test_inline_functions.py
index c8c6b860e8..3deb8799ce 100644
--- a/tests/python/relax/test_inline_functions.py
+++ b/tests/python/relax/test_inline_functions.py
@@ -20,7 +20,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_kill_after_last_use.py b/tests/python/relax/test_kill_after_last_use.py
index 0a37e7c3c3..8a21b44413 100644
--- a/tests/python/relax/test_kill_after_last_use.py
+++ b/tests/python/relax/test_kill_after_last_use.py
@@ -22,7 +22,7 @@ import tvm.relax
import tvm_tirx.testing, tvm
from tvm.relax.transform import KillAfterLastUse
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_lower_alloc_tensor.py b/tests/python/relax/test_lower_alloc_tensor.py
index 7cc6356906..928a317c4b 100644
--- a/tests/python/relax/test_lower_alloc_tensor.py
+++ b/tests/python/relax/test_lower_alloc_tensor.py
@@ -19,7 +19,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax.transform import LowerAllocTensor
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_basic():
diff --git a/tests/python/relax/test_meta_schedule_relax_integration.py b/tests/python/relax/test_meta_schedule_relax_integration.py
index 6cb286d1a0..8f8a35cb7e 100644
--- a/tests/python/relax/test_meta_schedule_relax_integration.py
+++ b/tests/python/relax/test_meta_schedule_relax_integration.py
@@ -30,7 +30,7 @@ from tvm_tirx.runtime import tensor as tvm_tensor
from tvm.runtime.vm import VirtualMachine
from tvm.s_tir import meta_schedule as ms
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_op_binary.py b/tests/python/relax/test_op_binary.py
index b1603a3fd4..b14f8a6d25 100644
--- a/tests/python/relax/test_op_binary.py
+++ b/tests/python/relax/test_op_binary.py
@@ -25,7 +25,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_ccl.py b/tests/python/relax/test_op_ccl.py
index f4533bbc74..3df7926737 100644
--- a/tests/python/relax/test_op_ccl.py
+++ b/tests/python/relax/test_op_ccl.py
@@ -22,7 +22,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_create.py b/tests/python/relax/test_op_create.py
index 559cb9b64b..1f18ea2e54 100644
--- a/tests/python/relax/test_op_create.py
+++ b/tests/python/relax/test_op_create.py
@@ -22,7 +22,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_op_datatype.py b/tests/python/relax/test_op_datatype.py
index bda623c0ee..bd1cb27811 100644
--- a/tests/python/relax/test_op_datatype.py
+++ b/tests/python/relax/test_op_datatype.py
@@ -22,7 +22,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_grad.py b/tests/python/relax/test_op_grad.py
index d34a1bf99b..ab7dd792ba 100644
--- a/tests/python/relax/test_op_grad.py
+++ b/tests/python/relax/test_op_grad.py
@@ -20,7 +20,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir import Op
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_image.py b/tests/python/relax/test_op_image.py
index 3cddf7d2ac..e90c16f02a 100644
--- a/tests/python/relax/test_op_image.py
+++ b/tests/python/relax/test_op_image.py
@@ -27,7 +27,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_index.py b/tests/python/relax/test_op_index.py
index 732b451686..0302ecbe5d 100644
--- a/tests/python/relax/test_op_index.py
+++ b/tests/python/relax/test_op_index.py
@@ -25,8 +25,8 @@ from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_op_inspect.py b/tests/python/relax/test_op_inspect.py
index 0de689e1e9..18df1a09bd 100644
--- a/tests/python/relax/test_op_inspect.py
+++ b/tests/python/relax/test_op_inspect.py
@@ -26,7 +26,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir import Op
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
# Parameterization for reading dtype of DLTensor. Chosen to have
diff --git a/tests/python/relax/test_op_linear_algebra.py b/tests/python/relax/test_op_linear_algebra.py
index 41ac265baa..66dd103926 100644
--- a/tests/python/relax/test_op_linear_algebra.py
+++ b/tests/python/relax/test_op_linear_algebra.py
@@ -23,7 +23,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_manipulate.py b/tests/python/relax/test_op_manipulate.py
index 253719574f..d6748bad1b 100644
--- a/tests/python/relax/test_op_manipulate.py
+++ b/tests/python/relax/test_op_manipulate.py
@@ -23,7 +23,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_op_misc.py b/tests/python/relax/test_op_misc.py
index 1684c065df..afb44600ef 100644
--- a/tests/python/relax/test_op_misc.py
+++ b/tests/python/relax/test_op_misc.py
@@ -20,8 +20,8 @@ import tvm_ffi
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax as rx
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_op_nn.py b/tests/python/relax/test_op_nn.py
index c2bcc276c0..1627282835 100644
--- a/tests/python/relax/test_op_nn.py
+++ b/tests/python/relax/test_op_nn.py
@@ -23,7 +23,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_nn_convolution.py b/tests/python/relax/test_op_nn_convolution.py
index dbf76545dc..f6810b4b9b 100644
--- a/tests/python/relax/test_op_nn_convolution.py
+++ b/tests/python/relax/test_op_nn_convolution.py
@@ -22,7 +22,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_conv1d_op_correctness():
diff --git a/tests/python/relax/test_op_nn_pooling.py b/tests/python/relax/test_op_nn_pooling.py
index 10de9d1f53..a69867ec78 100644
--- a/tests/python/relax/test_op_nn_pooling.py
+++ b/tests/python/relax/test_op_nn_pooling.py
@@ -23,7 +23,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_qdq.py b/tests/python/relax/test_op_qdq.py
index 5d3026464d..bb40d1e653 100644
--- a/tests/python/relax/test_op_qdq.py
+++ b/tests/python/relax/test_op_qdq.py
@@ -21,7 +21,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_sampling.py b/tests/python/relax/test_op_sampling.py
index e2076146c9..2600f73b5d 100644
--- a/tests/python/relax/test_op_sampling.py
+++ b/tests/python/relax/test_op_sampling.py
@@ -17,7 +17,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_ty: relax.Type):
diff --git a/tests/python/relax/test_op_search.py b/tests/python/relax/test_op_search.py
index 3e6703f509..2787e58d0c 100644
--- a/tests/python/relax/test_op_search.py
+++ b/tests/python/relax/test_op_search.py
@@ -24,7 +24,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_set.py b/tests/python/relax/test_op_set.py
index 0660600432..6eb5864266 100644
--- a/tests/python/relax/test_op_set.py
+++ b/tests/python/relax/test_op_set.py
@@ -22,7 +22,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_size.py b/tests/python/relax/test_op_size.py
index 222a6f07be..b58c16cb3e 100644
--- a/tests/python/relax/test_op_size.py
+++ b/tests/python/relax/test_op_size.py
@@ -21,7 +21,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_size():
diff --git a/tests/python/relax/test_op_sort.py b/tests/python/relax/test_op_sort.py
index c4f62f27b8..c58dd9b1d6 100644
--- a/tests/python/relax/test_op_sort.py
+++ b/tests/python/relax/test_op_sort.py
@@ -22,7 +22,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_statistical.py b/tests/python/relax/test_op_statistical.py
index 960ed1acda..b2b28d1ff7 100644
--- a/tests/python/relax/test_op_statistical.py
+++ b/tests/python/relax/test_op_statistical.py
@@ -24,7 +24,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_take.py b/tests/python/relax/test_op_take.py
index b703015d27..6690d1ebb0 100644
--- a/tests/python/relax/test_op_take.py
+++ b/tests/python/relax/test_op_take.py
@@ -21,7 +21,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
axis = tvm_tirx.testing.parameter(0, 1)
diff --git a/tests/python/relax/test_op_ternary.py b/tests/python/relax/test_op_ternary.py
index 8c94e68cd7..02a4991679 100644
--- a/tests/python/relax/test_op_ternary.py
+++ b/tests/python/relax/test_op_ternary.py
@@ -22,7 +22,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_unary.py b/tests/python/relax/test_op_unary.py
index 25166f7d90..6bce7b83f9 100644
--- a/tests/python/relax/test_op_unary.py
+++ b/tests/python/relax/test_op_unary.py
@@ -24,7 +24,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_op_correctness():
diff --git a/tests/python/relax/test_op_view.py b/tests/python/relax/test_op_view.py
index eb86257472..421348beda 100644
--- a/tests/python/relax/test_op_view.py
+++ b/tests/python/relax/test_op_view.py
@@ -21,7 +21,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_op_vision.py b/tests/python/relax/test_op_vision.py
index 65dc555478..a2ca7e67e0 100644
--- a/tests/python/relax/test_op_vision.py
+++ b/tests/python/relax/test_op_vision.py
@@ -29,7 +29,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import Op
from tvm.relax.transform import LegalizeOps
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_ty: relax.Type):
diff --git a/tests/python/relax/test_optimize_layout_transform.py b/tests/python/relax/test_optimize_layout_transform.py
index d3d4497059..adbe71785d 100644
--- a/tests/python/relax/test_optimize_layout_transform.py
+++ b/tests/python/relax/test_optimize_layout_transform.py
@@ -25,8 +25,8 @@ from tvm import relax
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.transform import DeadCodeElimination, FuseTIR, OptimizeLayoutTransform
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_pipeline.py b/tests/python/relax/test_pipeline.py
index 9891a75e1b..f8cfe28a90 100644
--- a/tests/python/relax/test_pipeline.py
+++ b/tests/python/relax/test_pipeline.py
@@ -21,7 +21,7 @@ import tvm_ffi
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_pytorch_integration.py b/tests/python/relax/test_pytorch_integration.py
index 0159865715..32d54036ae 100644
--- a/tests/python/relax/test_pytorch_integration.py
+++ b/tests/python/relax/test_pytorch_integration.py
@@ -32,8 +32,8 @@ import torch.nn.functional as F
import tvm, tvm_tirx
from tvm.relax import BasePyModule
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
n_matmul = T.dynamic("n", "int32")
diff --git a/tests/python/relax/test_relax_operators.py b/tests/python/relax/test_relax_operators.py
index 1be222ca9a..0733e9ec85 100644
--- a/tests/python/relax/test_relax_operators.py
+++ b/tests/python/relax/test_relax_operators.py
@@ -30,7 +30,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
m = T.dynamic("m", "int64")
diff --git a/tests/python/relax/test_relax_to_pyfunc_converter.py b/tests/python/relax/test_relax_to_pyfunc_converter.py
index 5555dd874e..8dd49fa9a8 100644
--- a/tests/python/relax/test_relax_to_pyfunc_converter.py
+++ b/tests/python/relax/test_relax_to_pyfunc_converter.py
@@ -30,8 +30,8 @@ import tvm, tvm_tirx
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.relax_to_pyfunc_converter import RelaxToPyFuncConverter
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
n_symbolic_add = T.dynamic("n", "int64")
diff --git a/tests/python/relax/test_remove_redundant_reshape.py b/tests/python/relax/test_remove_redundant_reshape.py
index cdceca3689..f67fadb78e 100644
--- a/tests/python/relax/test_remove_redundant_reshape.py
+++ b/tests/python/relax/test_remove_redundant_reshape.py
@@ -23,7 +23,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.transform import DeadCodeElimination, RemoveRedundantReshape
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _run_pass_compare_output(Before, Expected):
diff --git a/tests/python/relax/test_runtime_builtin_rnn_state.py b/tests/python/relax/test_runtime_builtin_rnn_state.py
index 495042a56a..e7606cbccd 100644
--- a/tests/python/relax/test_runtime_builtin_rnn_state.py
+++ b/tests/python/relax/test_runtime_builtin_rnn_state.py
@@ -25,7 +25,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir import dlight as dl
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_testing_nn.py b/tests/python/relax/test_testing_nn.py
index 512f8b6f55..42a81d6d59 100644
--- a/tests/python/relax/test_testing_nn.py
+++ b/tests/python/relax/test_testing_nn.py
@@ -20,7 +20,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.testing import nn
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_tir_call_source_kernel.py b/tests/python/relax/test_tir_call_source_kernel.py
index b312a412e4..900f5f1d5e 100644
--- a/tests/python/relax/test_tir_call_source_kernel.py
+++ b/tests/python/relax/test_tir_call_source_kernel.py
@@ -22,8 +22,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_training_append_loss.py b/tests/python/relax/test_training_append_loss.py
index 4048b378fc..03617986a5 100644
--- a/tests/python/relax/test_training_append_loss.py
+++ b/tests/python/relax/test_training_append_loss.py
@@ -21,7 +21,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.training import AppendLoss
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_simple():
diff --git a/tests/python/relax/test_training_loss.py b/tests/python/relax/test_training_loss.py
index 18680b6e3a..6c9c2af0c2 100644
--- a/tests/python/relax/test_training_loss.py
+++ b/tests/python/relax/test_training_loss.py
@@ -18,7 +18,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
@I.ir_module
diff --git a/tests/python/relax/test_training_setup_trainer.py b/tests/python/relax/test_training_setup_trainer.py
index 8537657570..c655c581fb 100644
--- a/tests/python/relax/test_training_setup_trainer.py
+++ b/tests/python/relax/test_training_setup_trainer.py
@@ -25,7 +25,7 @@ from tvm.relax.training import SetupTrainer
from tvm.relax.training.loss import MSELoss
from tvm.relax.training.optimizer import SGD, MomentumSGD
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_simple():
diff --git a/tests/python/relax/test_training_trainer_numeric.py b/tests/python/relax/test_training_trainer_numeric.py
index 80f6b64113..72c72b4c0f 100644
--- a/tests/python/relax/test_training_trainer_numeric.py
+++ b/tests/python/relax/test_training_trainer_numeric.py
@@ -24,7 +24,7 @@ from tvm.relax.training import SetupTrainer, Trainer
from tvm.relax.training.loss import MSELoss
from tvm.relax.training.optimizer import SGD, Adam
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def _get_backbone():
diff --git a/tests/python/relax/test_transform.py b/tests/python/relax/test_transform.py
index 9e7bf01625..c90559008e 100644
--- a/tests/python/relax/test_transform.py
+++ b/tests/python/relax/test_transform.py
@@ -25,8 +25,8 @@ from tvm import relax
from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_allocate_workspace.py b/tests/python/relax/test_transform_allocate_workspace.py
index 20da882271..07ca383065 100644
--- a/tests/python/relax/test_transform_allocate_workspace.py
+++ b/tests/python/relax/test_transform_allocate_workspace.py
@@ -18,7 +18,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
@I.ir_module
diff --git a/tests/python/relax/test_transform_alter_op_impl.py b/tests/python/relax/test_transform_alter_op_impl.py
index 784ac00c64..02139adf75 100644
--- a/tests/python/relax/test_transform_alter_op_impl.py
+++ b/tests/python/relax/test_transform_alter_op_impl.py
@@ -20,8 +20,8 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
kOperatorName = "operator_name"
diff --git a/tests/python/relax/test_transform_annotate_tir_op_pattern.py b/tests/python/relax/test_transform_annotate_tir_op_pattern.py
index 163b71bc95..497a555ad9 100644
--- a/tests/python/relax/test_transform_annotate_tir_op_pattern.py
+++ b/tests/python/relax/test_transform_annotate_tir_op_pattern.py
@@ -24,7 +24,7 @@ import tvm, tvm_tirx
import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm import relax
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_attach_attr_layout_free_tensors.py b/tests/python/relax/test_transform_attach_attr_layout_free_tensors.py
index 1fdcd642d2..625e187325 100644
--- a/tests/python/relax/test_transform_attach_attr_layout_free_tensors.py
+++ b/tests/python/relax/test_transform_attach_attr_layout_free_tensors.py
@@ -23,8 +23,8 @@ from tvm import relax
from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.transform import CombineParallelMatmul
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/test_transform_attach_global_symbol.py b/tests/python/relax/test_transform_attach_global_symbol.py
index 4b4ecbad2d..3542c98910 100644
--- a/tests/python/relax/test_transform_attach_global_symbol.py
+++ b/tests/python/relax/test_transform_attach_global_symbol.py
@@ -24,8 +24,8 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir import assert_structural_equal
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_bind_params.py b/tests/python/relax/test_transform_bind_params.py
index 6629e7b521..35af7e1218 100644
--- a/tests/python/relax/test_transform_bind_params.py
+++ b/tests/python/relax/test_transform_bind_params.py
@@ -23,8 +23,8 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
use_np_array = tvm_tirx.testing.parameter(False, True)
diff --git a/tests/python/relax/test_transform_bind_symbolic_vars.py b/tests/python/relax/test_transform_bind_symbolic_vars.py
index f2d7a42e9a..c4ba978eb9 100644
--- a/tests/python/relax/test_transform_bind_symbolic_vars.py
+++ b/tests/python/relax/test_transform_bind_symbolic_vars.py
@@ -23,7 +23,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_bundle_model_params.py b/tests/python/relax/test_transform_bundle_model_params.py
index 7ddbceefc7..b8e74bce66 100644
--- a/tests/python/relax/test_transform_bundle_model_params.py
+++ b/tests/python/relax/test_transform_bundle_model_params.py
@@ -27,7 +27,7 @@ import tvm_tirx.testing, tvm
import tvm.topi.testing
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_canonicalize_bindings.py b/tests/python/relax/test_transform_canonicalize_bindings.py
index 8215e90052..fe4100f461 100644
--- a/tests/python/relax/test_transform_canonicalize_bindings.py
+++ b/tests/python/relax/test_transform_canonicalize_bindings.py
@@ -26,8 +26,8 @@ from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.transform.transform import CanonicalizeBindings
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_codegen_pass.py b/tests/python/relax/test_transform_codegen_pass.py
index 237b4088f5..dcfc56e323 100644
--- a/tests/python/relax/test_transform_codegen_pass.py
+++ b/tests/python/relax/test_transform_codegen_pass.py
@@ -30,8 +30,8 @@ from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.dpl import is_op, wildcard
from tvm.relax.testing import transform
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.support import utils
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_transform_combine_parallel_matmul.py b/tests/python/relax/test_transform_combine_parallel_matmul.py
index 36e673ab00..1aa8767bce 100644
--- a/tests/python/relax/test_transform_combine_parallel_matmul.py
+++ b/tests/python/relax/test_transform_combine_parallel_matmul.py
@@ -22,7 +22,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.transform import CombineParallelMatmul
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/test_transform_compute_prim_value.py b/tests/python/relax/test_transform_compute_prim_value.py
index a8d5660e9e..9ec51abd9a 100644
--- a/tests/python/relax/test_transform_compute_prim_value.py
+++ b/tests/python/relax/test_transform_compute_prim_value.py
@@ -18,8 +18,8 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_convert_dataflow.py b/tests/python/relax/test_transform_convert_dataflow.py
index 91e264f714..39e8f1587d 100644
--- a/tests/python/relax/test_transform_convert_dataflow.py
+++ b/tests/python/relax/test_transform_convert_dataflow.py
@@ -19,7 +19,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
# functions that will not change
diff --git a/tests/python/relax/test_transform_cse.py b/tests/python/relax/test_transform_cse.py
index 9663268ee8..c776d11a60 100644
--- a/tests/python/relax/test_transform_cse.py
+++ b/tests/python/relax/test_transform_cse.py
@@ -22,7 +22,7 @@ import numpy as np
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax.transform import EliminateCommonSubexpr
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_transform_dead_code_elimination.py b/tests/python/relax/test_transform_dead_code_elimination.py
index d79517c944..1a1a899051 100644
--- a/tests/python/relax/test_transform_dead_code_elimination.py
+++ b/tests/python/relax/test_transform_dead_code_elimination.py
@@ -21,7 +21,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax.transform import DeadCodeElimination
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_transform_decompose_ops.py b/tests/python/relax/test_transform_decompose_ops.py
index 32494d9a84..74bf4d82b3 100644
--- a/tests/python/relax/test_transform_decompose_ops.py
+++ b/tests/python/relax/test_transform_decompose_ops.py
@@ -25,7 +25,7 @@ from tvm_tirx import IRModule
from tvm import relax
from tvm.relax import Function
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_error_enrichment.py b/tests/python/relax/test_transform_error_enrichment.py
index dbf441eb41..f469386f9f 100644
--- a/tests/python/relax/test_transform_error_enrichment.py
+++ b/tests/python/relax/test_transform_error_enrichment.py
@@ -59,7 +59,7 @@ def test_pass_error_renders_underlined_tvmscript():
"Location (TVMScript):\n"
"Access path: <root>.body.blocks[0].bindings[0].value\n\n"
"from __future__ import annotations\n\n"
- "# from tvm_tirx.script import relax as R\n\n"
+ "# from tvm.relax import script as R\n\n"
"@R.function\n"
'def main(x: R.Tensor((3, 4), dtype="float32"), '
'y: R.Tensor((5, 6), dtype="float32")) -> R.Tensor((3, 6), dtype="float32"):\n'
diff --git a/tests/python/relax/test_transform_expand_matmul_of_sum.py b/tests/python/relax/test_transform_expand_matmul_of_sum.py
index f56421935c..23c0f74008 100644
--- a/tests/python/relax/test_transform_expand_matmul_of_sum.py
+++ b/tests/python/relax/test_transform_expand_matmul_of_sum.py
@@ -22,7 +22,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
class Base:
diff --git a/tests/python/relax/test_transform_expand_tuple_args.py b/tests/python/relax/test_transform_expand_tuple_args.py
index 441dfb8144..a8dde501e4 100644
--- a/tests/python/relax/test_transform_expand_tuple_args.py
+++ b/tests/python/relax/test_transform_expand_tuple_args.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_fold_batch_norm_to_conv2d.py b/tests/python/relax/test_transform_fold_batch_norm_to_conv2d.py
index 51a79a94bb..6556e96d0a 100644
--- a/tests/python/relax/test_transform_fold_batch_norm_to_conv2d.py
+++ b/tests/python/relax/test_transform_fold_batch_norm_to_conv2d.py
@@ -26,7 +26,7 @@ from tvm_tirx.ir.module import IRModule
from tvm.relax.expr_functor import PyExprVisitor, visitor
from tvm.relax.script import ir_builder as relax_builder
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/test_transform_fold_constant.py b/tests/python/relax/test_transform_fold_constant.py
index 13f6e14089..7a5c60e822 100644
--- a/tests/python/relax/test_transform_fold_constant.py
+++ b/tests/python/relax/test_transform_fold_constant.py
@@ -23,8 +23,8 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_fuse_ops.py b/tests/python/relax/test_transform_fuse_ops.py
index 4b3cf4dfe2..7224458f3b 100644
--- a/tests/python/relax/test_transform_fuse_ops.py
+++ b/tests/python/relax/test_transform_fuse_ops.py
@@ -23,8 +23,8 @@ import tvm_tirx.testing, tvm
from tvm import relax, topi
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_fuse_ops_by_pattern.py b/tests/python/relax/test_transform_fuse_ops_by_pattern.py
index 3d9b31e3da..e98e2681ba 100644
--- a/tests/python/relax/test_transform_fuse_ops_by_pattern.py
+++ b/tests/python/relax/test_transform_fuse_ops_by_pattern.py
@@ -31,8 +31,8 @@ from tvm.relax.dpl.pattern import (
)
from tvm.relax.transform import PatternCheckContext
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_fuse_tir.py b/tests/python/relax/test_transform_fuse_tir.py
index a562773624..83477501d3 100644
--- a/tests/python/relax/test_transform_fuse_tir.py
+++ b/tests/python/relax/test_transform_fuse_tir.py
@@ -22,8 +22,8 @@ import tvm_tirx.testing, tvm
from tvm import relax, topi
from tvm.relax import ExternFunc as _ExternFunc
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_fuse_transpose_matmul.py b/tests/python/relax/test_transform_fuse_transpose_matmul.py
index 3924474906..031260d3aa 100644
--- a/tests/python/relax/test_transform_fuse_transpose_matmul.py
+++ b/tests/python/relax/test_transform_fuse_transpose_matmul.py
@@ -22,8 +22,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_gradient.py b/tests/python/relax/test_transform_gradient.py
index 5d35b52c74..960ac3020f 100644
--- a/tests/python/relax/test_transform_gradient.py
+++ b/tests/python/relax/test_transform_gradient.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir.base import assert_structural_equal
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_transform_gradient_te_register.py b/tests/python/relax/test_transform_gradient_te_register.py
index 29b429a368..c2966ba664 100644
--- a/tests/python/relax/test_transform_gradient_te_register.py
+++ b/tests/python/relax/test_transform_gradient_te_register.py
@@ -27,7 +27,7 @@ from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.training.utils import register_te_gradient
from tvm.relax.transform import Gradient
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_transform_inline_private_functions.py b/tests/python/relax/test_transform_inline_private_functions.py
index b59c40752c..78378d5eb4 100644
--- a/tests/python/relax/test_transform_inline_private_functions.py
+++ b/tests/python/relax/test_transform_inline_private_functions.py
@@ -23,7 +23,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_ipc_allreduce_rewrite.py b/tests/python/relax/test_transform_ipc_allreduce_rewrite.py
index 4db4f60f70..8370cc592a 100644
--- a/tests/python/relax/test_transform_ipc_allreduce_rewrite.py
+++ b/tests/python/relax/test_transform_ipc_allreduce_rewrite.py
@@ -19,7 +19,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_lambda_lift.py b/tests/python/relax/test_transform_lambda_lift.py
index a5ecfc352d..fb0cc2bc8f 100644
--- a/tests/python/relax/test_transform_lambda_lift.py
+++ b/tests/python/relax/test_transform_lambda_lift.py
@@ -25,8 +25,8 @@ from tvm import relax
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax import transform
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_lazy_transform_params.py b/tests/python/relax/test_transform_lazy_transform_params.py
index 3465230ef6..cde91a1c43 100644
--- a/tests/python/relax/test_transform_lazy_transform_params.py
+++ b/tests/python/relax/test_transform_lazy_transform_params.py
@@ -22,8 +22,8 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.transform import LazyTransformParams
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_legalize_ops.py b/tests/python/relax/test_transform_legalize_ops.py
index 73634c9be9..5303ba4eb5 100644
--- a/tests/python/relax/test_transform_legalize_ops.py
+++ b/tests/python/relax/test_transform_legalize_ops.py
@@ -26,8 +26,8 @@ from tvm import relax
from tvm.relax.transform import LegalizeOps
from tvm.relax.transform.legalize_ops.common import register_legalize
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_legalize_ops_binary.py b/tests/python/relax/test_transform_legalize_ops_binary.py
index 484afd300a..bef48180e7 100644
--- a/tests/python/relax/test_transform_legalize_ops_binary.py
+++ b/tests/python/relax/test_transform_legalize_ops_binary.py
@@ -20,8 +20,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
##################### Binary arithmetic #####################
diff --git a/tests/python/relax/test_transform_legalize_ops_ccl.py b/tests/python/relax/test_transform_legalize_ops_ccl.py
index 12beaa1b80..45aeefcf25 100644
--- a/tests/python/relax/test_transform_legalize_ops_ccl.py
+++ b/tests/python/relax/test_transform_legalize_ops_ccl.py
@@ -21,8 +21,8 @@ import tvm_tirx.testing, tvm
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_legalize_ops_create_datatype.py b/tests/python/relax/test_transform_legalize_ops_create_datatype.py
index 9f0662df14..664b116200 100644
--- a/tests/python/relax/test_transform_legalize_ops_create_datatype.py
+++ b/tests/python/relax/test_transform_legalize_ops_create_datatype.py
@@ -19,8 +19,8 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax.transform import LegalizeOps
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
##################### Creation #####################
diff --git a/tests/python/relax/test_transform_legalize_ops_distributed.py b/tests/python/relax/test_transform_legalize_ops_distributed.py
index 87d670e166..213a2f3cd0 100644
--- a/tests/python/relax/test_transform_legalize_ops_distributed.py
+++ b/tests/python/relax/test_transform_legalize_ops_distributed.py
@@ -21,8 +21,8 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_legalize_ops_grad.py b/tests/python/relax/test_transform_legalize_ops_grad.py
index 7b3ce4ea3f..4ba0b113e7 100644
--- a/tests/python/relax/test_transform_legalize_ops_grad.py
+++ b/tests/python/relax/test_transform_legalize_ops_grad.py
@@ -20,8 +20,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_legalize_ops_image.py b/tests/python/relax/test_transform_legalize_ops_image.py
index daf0e85cc2..72982fa3b6 100644
--- a/tests/python/relax/test_transform_legalize_ops_image.py
+++ b/tests/python/relax/test_transform_legalize_ops_image.py
@@ -20,8 +20,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.transform import LegalizeOps
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_legalize_ops_index_linear_algebra.py b/tests/python/relax/test_transform_legalize_ops_index_linear_algebra.py
index ba26612000..81b279d5f7 100644
--- a/tests/python/relax/test_transform_legalize_ops_index_linear_algebra.py
+++ b/tests/python/relax/test_transform_legalize_ops_index_linear_algebra.py
@@ -21,8 +21,8 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir import Op
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
##################### Indexing #####################
diff --git a/tests/python/relax/test_transform_legalize_ops_manipulate.py b/tests/python/relax/test_transform_legalize_ops_manipulate.py
index fb0888a5ca..639d9c6abd 100644
--- a/tests/python/relax/test_transform_legalize_ops_manipulate.py
+++ b/tests/python/relax/test_transform_legalize_ops_manipulate.py
@@ -23,8 +23,8 @@ from tvm import relax
from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
##################### Manipulation #####################
diff --git a/tests/python/relax/test_transform_legalize_ops_nn.py b/tests/python/relax/test_transform_legalize_ops_nn.py
index 27b9d6428a..4528e13513 100644
--- a/tests/python/relax/test_transform_legalize_ops_nn.py
+++ b/tests/python/relax/test_transform_legalize_ops_nn.py
@@ -23,8 +23,8 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
##################### Neural network #####################
diff --git a/tests/python/relax/test_transform_legalize_ops_qdq.py b/tests/python/relax/test_transform_legalize_ops_qdq.py
index 0275aa7de8..3ac002ad33 100644
--- a/tests/python/relax/test_transform_legalize_ops_qdq.py
+++ b/tests/python/relax/test_transform_legalize_ops_qdq.py
@@ -18,8 +18,8 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.relax.transform import LegalizeOps
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_legalize_ops_search_statistical.py b/tests/python/relax/test_transform_legalize_ops_search_statistical.py
index f21d4a873c..7d59c9446e 100644
--- a/tests/python/relax/test_transform_legalize_ops_search_statistical.py
+++ b/tests/python/relax/test_transform_legalize_ops_search_statistical.py
@@ -21,8 +21,8 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
##################### Search #####################
diff --git a/tests/python/relax/test_transform_legalize_ops_unary.py b/tests/python/relax/test_transform_legalize_ops_unary.py
index 8b5af413f9..ab88a99353 100644
--- a/tests/python/relax/test_transform_legalize_ops_unary.py
+++ b/tests/python/relax/test_transform_legalize_ops_unary.py
@@ -26,7 +26,7 @@ import tvm_tirx.testing, tvm
from tvm import topi
from tvm.relax.transform import LegalizeOps
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_lift_transform_params.py b/tests/python/relax/test_transform_lift_transform_params.py
index 4546b5d2e7..6f56df6511 100644
--- a/tests/python/relax/test_transform_lift_transform_params.py
+++ b/tests/python/relax/test_transform_lift_transform_params.py
@@ -29,8 +29,8 @@ pytest.importorskip("scipy") # tvm.topi.testing imports scipy
import tvm.topi.testing
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_lower_gpu_ipc_alloc_storage.py b/tests/python/relax/test_transform_lower_gpu_ipc_alloc_storage.py
index d1db285ecf..8d5800e09e 100644
--- a/tests/python/relax/test_transform_lower_gpu_ipc_alloc_storage.py
+++ b/tests/python/relax/test_transform_lower_gpu_ipc_alloc_storage.py
@@ -19,7 +19,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_merge_composite_functions.py b/tests/python/relax/test_transform_merge_composite_functions.py
index 05f9e322cb..7c4e9dd79d 100644
--- a/tests/python/relax/test_transform_merge_composite_functions.py
+++ b/tests/python/relax/test_transform_merge_composite_functions.py
@@ -19,8 +19,8 @@ import pytest
import tvm, tvm_tirx
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_meta_schedule_apply_database.py b/tests/python/relax/test_transform_meta_schedule_apply_database.py
index f6a397a8e9..8ad959ce21 100644
--- a/tests/python/relax/test_transform_meta_schedule_apply_database.py
+++ b/tests/python/relax/test_transform_meta_schedule_apply_database.py
@@ -22,7 +22,7 @@ from tvm_tirx import tirx
from tvm import relax
from tvm.s_tir import meta_schedule as ms
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
target = tvm_tirx.target.Target({"kind": "llvm", "num-cores": 16})
diff --git a/tests/python/relax/test_transform_meta_schedule_tuning.py b/tests/python/relax/test_transform_meta_schedule_tuning.py
index decf585208..9592acd404 100644
--- a/tests/python/relax/test_transform_meta_schedule_tuning.py
+++ b/tests/python/relax/test_transform_meta_schedule_tuning.py
@@ -43,8 +43,8 @@ from tvm import relax
from tvm_tirx.ir import transform
from tvm_tirx.ir.module import IRModule
from tvm_tirx.ir.transform import PassContext
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
target = tvm_tirx.target.Target({"kind": "llvm", "num-cores": 16})
diff --git a/tests/python/relax/test_transform_normalize.py b/tests/python/relax/test_transform_normalize.py
index 46c9428a3d..0513b3f18a 100644
--- a/tests/python/relax/test_transform_normalize.py
+++ b/tests/python/relax/test_transform_normalize.py
@@ -23,7 +23,7 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_normalize_global_var.py b/tests/python/relax/test_transform_normalize_global_var.py
index 29a53ce9cc..b8fa37e836 100644
--- a/tests/python/relax/test_transform_normalize_global_var.py
+++ b/tests/python/relax/test_transform_normalize_global_var.py
@@ -24,8 +24,8 @@ from tvm_tirx import tirx
from tvm import relax
from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_operator_specific_normalization.py b/tests/python/relax/test_transform_operator_specific_normalization.py
index e990d7eafb..55f598dc06 100644
--- a/tests/python/relax/test_transform_operator_specific_normalization.py
+++ b/tests/python/relax/test_transform_operator_specific_normalization.py
@@ -25,7 +25,7 @@ import tvm, tvm_tirx
import tvm.relax.testing.transform
import tvm_tirx.testing, tvm
from tvm import relax
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_transform_remove_unused_outputs.py b/tests/python/relax/test_transform_remove_unused_outputs.py
index 0fff54561e..19a700be16 100644
--- a/tests/python/relax/test_transform_remove_unused_outputs.py
+++ b/tests/python/relax/test_transform_remove_unused_outputs.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_remove_unused_parameters.py b/tests/python/relax/test_transform_remove_unused_parameters.py
index 16c2ffe5aa..870ce3f88d 100644
--- a/tests/python/relax/test_transform_remove_unused_parameters.py
+++ b/tests/python/relax/test_transform_remove_unused_parameters.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_reorder_permute_dims_after_concat.py b/tests/python/relax/test_transform_reorder_permute_dims_after_concat.py
index 23836513c0..257cd426ae 100644
--- a/tests/python/relax/test_transform_reorder_permute_dims_after_concat.py
+++ b/tests/python/relax/test_transform_reorder_permute_dims_after_concat.py
@@ -22,7 +22,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
class Base:
diff --git a/tests/python/relax/test_transform_reorder_take_after_matmul.py b/tests/python/relax/test_transform_reorder_take_after_matmul.py
index a7e1e6ba6a..3bc422cc82 100644
--- a/tests/python/relax/test_transform_reorder_take_after_matmul.py
+++ b/tests/python/relax/test_transform_reorder_take_after_matmul.py
@@ -22,7 +22,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_rewrite_cuda_graph.py b/tests/python/relax/test_transform_rewrite_cuda_graph.py
index f209b41be1..51266413d3 100644
--- a/tests/python/relax/test_transform_rewrite_cuda_graph.py
+++ b/tests/python/relax/test_transform_rewrite_cuda_graph.py
@@ -23,8 +23,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_rewrite_dataflow_reshape.py b/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
index d8872a280e..27caeed1b5 100644
--- a/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
+++ b/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
@@ -20,8 +20,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_specialize_function_based_on_callsite.py b/tests/python/relax/test_transform_specialize_function_based_on_callsite.py
index 13e368d580..15899251e8 100644
--- a/tests/python/relax/test_transform_specialize_function_based_on_callsite.py
+++ b/tests/python/relax/test_transform_specialize_function_based_on_callsite.py
@@ -25,7 +25,7 @@ from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.ir.module import IRModule
from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, visitor
from tvm.relax.transform.legalize_ops import adreno as legalize_adreno
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_transform_split_layout_rewrite_preproc.py b/tests/python/relax/test_transform_split_layout_rewrite_preproc.py
index 955f1d8524..28bf5da290 100644
--- a/tests/python/relax/test_transform_split_layout_rewrite_preproc.py
+++ b/tests/python/relax/test_transform_split_layout_rewrite_preproc.py
@@ -19,8 +19,8 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_static_plan_block_memory.py b/tests/python/relax/test_transform_static_plan_block_memory.py
index 6a56e74fd7..e30c4cf56f 100644
--- a/tests/python/relax/test_transform_static_plan_block_memory.py
+++ b/tests/python/relax/test_transform_static_plan_block_memory.py
@@ -22,8 +22,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_to_mixed_precision.py b/tests/python/relax/test_transform_to_mixed_precision.py
index 498f3f5741..b05dc19c68 100644
--- a/tests/python/relax/test_transform_to_mixed_precision.py
+++ b/tests/python/relax/test_transform_to_mixed_precision.py
@@ -21,7 +21,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm.relax.transform import ToMixedPrecision
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_transform_topological_sort.py b/tests/python/relax/test_transform_topological_sort.py
index cf5c631f5d..fb85944fd6 100644
--- a/tests/python/relax/test_transform_topological_sort.py
+++ b/tests/python/relax/test_transform_topological_sort.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_depth_first_from_inputs():
diff --git a/tests/python/relax/test_transform_update_param_type.py b/tests/python/relax/test_transform_update_param_type.py
index fca9d83b5a..d49cc29719 100644
--- a/tests/python/relax/test_transform_update_param_type.py
+++ b/tests/python/relax/test_transform_update_param_type.py
@@ -22,7 +22,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
class Base:
diff --git a/tests/python/relax/test_vm_alloc_storage_with_scope.py b/tests/python/relax/test_vm_alloc_storage_with_scope.py
index 203f54e93a..e2c4616832 100644
--- a/tests/python/relax/test_vm_alloc_storage_with_scope.py
+++ b/tests/python/relax/test_vm_alloc_storage_with_scope.py
@@ -24,8 +24,8 @@ import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.runtime import MemoryScope
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_vm_build.py b/tests/python/relax/test_vm_build.py
index 26bff991c0..13354415a4 100644
--- a/tests/python/relax/test_vm_build.py
+++ b/tests/python/relax/test_vm_build.py
@@ -36,8 +36,8 @@ from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.testing import nn
from tvm.relax.testing.vm import check_saved_func
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.support import cc, popen_pool, utils
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_vm_builtin.py b/tests/python/relax/test_vm_builtin.py
index c682e29be5..37489938ab 100644
--- a/tests/python/relax/test_vm_builtin.py
+++ b/tests/python/relax/test_vm_builtin.py
@@ -23,7 +23,7 @@ import tvm_tirx.script, tvm
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_multinomial_from_uniform():
diff --git a/tests/python/relax/test_vm_builtin_lower.py b/tests/python/relax/test_vm_builtin_lower.py
index f9c3a21295..dcb1a0bcb4 100644
--- a/tests/python/relax/test_vm_builtin_lower.py
+++ b/tests/python/relax/test_vm_builtin_lower.py
@@ -21,7 +21,7 @@ import tvm, tvm_tirx
import tvm_tirx.script, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_vm_callback_function.py b/tests/python/relax/test_vm_callback_function.py
index b1621f9ba9..f1fecb574a 100644
--- a/tests/python/relax/test_vm_callback_function.py
+++ b/tests/python/relax/test_vm_callback_function.py
@@ -21,7 +21,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
pytestmark = pytest.mark.skipif(not tvm_tirx.testing.device_enabled("llvm"), reason="llvm not enabled")
diff --git a/tests/python/relax/test_vm_codegen_only.py b/tests/python/relax/test_vm_codegen_only.py
index 866def31f1..950b672989 100644
--- a/tests/python/relax/test_vm_codegen_only.py
+++ b/tests/python/relax/test_vm_codegen_only.py
@@ -30,8 +30,8 @@ from tvm import relax
from tvm.relax.testing.runtime_builtin import MakeShapeCode, MatchShapeCode
from tvm.relax.testing.vm import check_saved_func
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_vm_cuda_graph.py b/tests/python/relax/test_vm_cuda_graph.py
index 65dfb73516..98379433a4 100644
--- a/tests/python/relax/test_vm_cuda_graph.py
+++ b/tests/python/relax/test_vm_cuda_graph.py
@@ -26,8 +26,8 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_vm_execbuilder.py b/tests/python/relax/test_vm_execbuilder.py
index 9f806f734f..af0a634290 100644
--- a/tests/python/relax/test_vm_execbuilder.py
+++ b/tests/python/relax/test_vm_execbuilder.py
@@ -23,7 +23,7 @@ import tvm_ffi
import tvm, tvm_tirx
from tvm import relax
from tvm.relax.testing.vm import check_saved_func
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
def test_vm_execute():
diff --git a/tests/python/relax/texture/test_texture_nd.py b/tests/python/relax/texture/test_texture_nd.py
index 0bce25b5c5..a1b0673215 100644
--- a/tests/python/relax/texture/test_texture_nd.py
+++ b/tests/python/relax/texture/test_texture_nd.py
@@ -29,7 +29,7 @@ from tvm import relax
from tvm.relax.transform.legalize_ops import adreno as legalize_adreno
from tvm.rpc import connect_tracker
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.support import ndk
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_calculate_allocated_memory.py b/tests/python/s_tir/analysis/test_s_tir_analysis_calculate_allocated_memory.py
index c4bc55fbcf..d6bade58b8 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_calculate_allocated_memory.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_calculate_allocated_memory.py
@@ -20,7 +20,7 @@ import pytest
import tvm, tvm_tirx
from tvm_tirx import tirx
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_estimate_tir_flops.py b/tests/python/s_tir/analysis/test_s_tir_analysis_estimate_tir_flops.py
index f6d28a45cf..ed4b6f71ff 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_estimate_tir_flops.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_estimate_tir_flops.py
@@ -24,7 +24,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir import IRModule
from tvm.s_tir.analysis import estimate_tir_flops
from tvm.s_tir.meta_schedule.testing.te_workload import create_te_workload
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py b/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py
index 47e86ce13f..2addeb29d4 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py
@@ -24,7 +24,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.ir import StringImm, TensorRegion
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
identify_memcpy = tvm.s_tir.analysis._ffi_api._identify_memcpy
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py b/tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py
index 7c7721277d..8ce4130f6f 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py
@@ -19,7 +19,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm.s_tir.analysis import assert_pure_function, is_pure_function
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_oob.py b/tests/python/s_tir/analysis/test_s_tir_analysis_oob.py
index bf2e9bd3bb..5a4c7d8735 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_oob.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_oob.py
@@ -17,7 +17,7 @@
import pytest
import tvm, tvm_tirx
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
index 7a13518400..f8077628d2 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
@@ -20,7 +20,7 @@ import pytest
import tvm, tvm_tirx
from tvm_tirx.ir import TensorRegion, TensorRegionType
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/analysis/test_sblock_access_region.py b/tests/python/s_tir/analysis/test_sblock_access_region.py
index 4410097d0b..bbcbffc9e5 100644
--- a/tests/python/s_tir/analysis/test_sblock_access_region.py
+++ b/tests/python/s_tir/analysis/test_sblock_access_region.py
@@ -20,7 +20,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import s_tir
from tvm_tirx.ir import Range, TensorRegion, TensorRegionType
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/analysis/test_sblock_tensor_access_lca.py b/tests/python/s_tir/analysis/test_sblock_tensor_access_lca.py
index 3e476ae7df..668d4f4b3b 100644
--- a/tests/python/s_tir/analysis/test_sblock_tensor_access_lca.py
+++ b/tests/python/s_tir/analysis/test_sblock_tensor_access_lca.py
@@ -16,7 +16,7 @@
# under the License.
import tvm, tvm_tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/base/test_compilation_pipeline.py b/tests/python/s_tir/base/test_compilation_pipeline.py
index 3682dd3159..a0e0e4411c 100644
--- a/tests/python/s_tir/base/test_compilation_pipeline.py
+++ b/tests/python/s_tir/base/test_compilation_pipeline.py
@@ -22,7 +22,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/base/test_s_tir_texture_scope.py b/tests/python/s_tir/base/test_s_tir_texture_scope.py
index 6425678045..40c911e596 100644
--- a/tests/python/s_tir/base/test_s_tir_texture_scope.py
+++ b/tests/python/s_tir/base/test_s_tir_texture_scope.py
@@ -23,7 +23,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm_tirx.ir.module import IRModule
from tvm.s_tir.backend.adreno import pipeline as adreno_pipeline
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx._compile import _split_host_device_modules
from tvm_tirx.tirx.compilation_pipeline import finalize_device_passes
diff --git a/tests/python/s_tir/base/test_sblock_dependence_info.py b/tests/python/s_tir/base/test_sblock_dependence_info.py
index 78df1f4a08..81c3f9fe5a 100644
--- a/tests/python/s_tir/base/test_sblock_dependence_info.py
+++ b/tests/python/s_tir/base/test_sblock_dependence_info.py
@@ -28,7 +28,7 @@ from tvm_tirx import tirx
from tvm_tirx.ir import IRModule
from tvm.s_tir import SBlockDependenceInfo
from tvm.s_tir.sblock_scope import DepKind
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx import Function
diff --git a/tests/python/s_tir/base/test_tir_te_extern_function.py b/tests/python/s_tir/base/test_tir_te_extern_function.py
index 57d4d77993..7cb1f41a80 100644
--- a/tests/python/s_tir/base/test_tir_te_extern_function.py
+++ b/tests/python/s_tir/base/test_tir_te_extern_function.py
@@ -24,7 +24,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import te
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# TODO(csullivan): Additional tests cases needed:
diff --git a/tests/python/s_tir/dlight/test_benchmark.py b/tests/python/s_tir/dlight/test_benchmark.py
index 4890e8b66f..5e563d0714 100644
--- a/tests/python/s_tir/dlight/test_benchmark.py
+++ b/tests/python/s_tir/dlight/test_benchmark.py
@@ -35,8 +35,8 @@ from tvm.s_tir.dlight.benchmark import (
)
from tvm.s_tir.meta_schedule.testing.local_rpc import LocalRPC
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
-from tvm_tirx.script import s_tir as Ts
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# The test function uses an undefined symbolic var in Relax.
diff --git a/tests/python/s_tir/dlight/test_cpu_gemv.py b/tests/python/s_tir/dlight/test_cpu_gemv.py
index 4b09008ef8..4503f5963a 100644
--- a/tests/python/s_tir/dlight/test_cpu_gemv.py
+++ b/tests/python/s_tir/dlight/test_cpu_gemv.py
@@ -22,7 +22,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm.s_tir import dlight as dl
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_conv.py b/tests/python/s_tir/dlight/test_gpu_conv.py
index 2ffa34ce44..e84f3f6577 100644
--- a/tests/python/s_tir/dlight/test_gpu_conv.py
+++ b/tests/python/s_tir/dlight/test_gpu_conv.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir import dlight as dl
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_fallback.py b/tests/python/s_tir/dlight/test_gpu_fallback.py
index 5941cdf198..fe0d60fea4 100644
--- a/tests/python/s_tir/dlight/test_gpu_fallback.py
+++ b/tests/python/s_tir/dlight/test_gpu_fallback.py
@@ -23,7 +23,7 @@ from tvm import s_tir
from tvm_tirx.ir import assert_structural_equal
from tvm.s_tir import dlight as dl
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_gemv.py b/tests/python/s_tir/dlight/test_gpu_gemv.py
index 27db2044ab..3818c3b863 100644
--- a/tests/python/s_tir/dlight/test_gpu_gemv.py
+++ b/tests/python/s_tir/dlight/test_gpu_gemv.py
@@ -22,7 +22,7 @@ from __future__ import annotations
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir import dlight as dl
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_general_reduction.py b/tests/python/s_tir/dlight/test_gpu_general_reduction.py
index bc0c23f826..7dd8fc42ce 100644
--- a/tests/python/s_tir/dlight/test_gpu_general_reduction.py
+++ b/tests/python/s_tir/dlight/test_gpu_general_reduction.py
@@ -22,7 +22,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir import IRModule, assert_structural_equal
from tvm.s_tir import dlight as dl
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py b/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
index 183e49d853..7fdb62c704 100644
--- a/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
+++ b/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
@@ -22,7 +22,7 @@ from unittest import mock
import tvm_tirx.testing, tvm
from tvm.s_tir import dlight as dl
from tvm.s_tir.dlight.gpu import low_batch_gemv
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_matmul.py b/tests/python/s_tir/dlight/test_gpu_matmul.py
index 7fa571a494..a60cf7eeca 100644
--- a/tests/python/s_tir/dlight/test_gpu_matmul.py
+++ b/tests/python/s_tir/dlight/test_gpu_matmul.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir import dlight as dl
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py b/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
index c0ed429e82..7c31b981f4 100644
--- a/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
+++ b/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
@@ -20,7 +20,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir import dlight as dl
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_reduction.py b/tests/python/s_tir/dlight/test_gpu_reduction.py
index f280015656..09392212ee 100644
--- a/tests/python/s_tir/dlight/test_gpu_reduction.py
+++ b/tests/python/s_tir/dlight/test_gpu_reduction.py
@@ -21,7 +21,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir import assert_structural_equal
from tvm.s_tir import dlight as dl
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_rmsnorm.py b/tests/python/s_tir/dlight/test_gpu_rmsnorm.py
index f389f9ad23..a5f8187057 100644
--- a/tests/python/s_tir/dlight/test_gpu_rmsnorm.py
+++ b/tests/python/s_tir/dlight/test_gpu_rmsnorm.py
@@ -20,7 +20,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx.ir import IRModule, assert_structural_equal
from tvm.s_tir import dlight as dl
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_gpu_transpose.py b/tests/python/s_tir/dlight/test_gpu_transpose.py
index 4ea67fd502..373384fd8f 100644
--- a/tests/python/s_tir/dlight/test_gpu_transpose.py
+++ b/tests/python/s_tir/dlight/test_gpu_transpose.py
@@ -20,7 +20,7 @@ import tvm, tvm_tirx
from tvm_tirx.ir import IRModule, assert_structural_equal
from tvm.s_tir import dlight as dl
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/dlight/test_primitives.py b/tests/python/s_tir/dlight/test_primitives.py
index 89728063ce..cb87931a53 100644
--- a/tests/python/s_tir/dlight/test_primitives.py
+++ b/tests/python/s_tir/dlight/test_primitives.py
@@ -21,7 +21,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_arg_info.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_arg_info.py
index d1fe6614b2..1adad84751 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_arg_info.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_arg_info.py
@@ -17,7 +17,7 @@
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
from tvm.s_tir.meta_schedule.arg_info import ArgInfo, TensorInfo
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_builder.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_builder.py
index 70fb6c8894..a6a31434a5 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_builder.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_builder.py
@@ -33,7 +33,7 @@ from tvm.s_tir.meta_schedule.builder import (
LocalBuilder,
PyBuilder,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_cost_model.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_cost_model.py
index 19f7d51098..d324cf2572 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_cost_model.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_cost_model.py
@@ -36,7 +36,7 @@ from tvm.s_tir.meta_schedule.search_strategy import MeasureCandidate
from tvm.s_tir.meta_schedule.tune_context import TuneContext
from tvm.s_tir.schedule.schedule import Schedule
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
requires_xgboost = pytest.mark.skipif(
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_database.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_database.py
index 2200f33e43..5ce1d9fa1a 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_database.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_database.py
@@ -34,7 +34,7 @@ from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.database import TuningRecord, Workload
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py
index 43268c1918..3c31bfea46 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py
@@ -27,7 +27,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir, te
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
N_FEATURES = 164
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_measure_callback.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_measure_callback.py
index 6256215174..5365a76cfd 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_measure_callback.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_measure_callback.py
@@ -23,7 +23,7 @@ import tvm, tvm_tirx
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.schedule import Schedule
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_compute_location.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_compute_location.py
index ed1217c528..003bb08134 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_compute_location.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_compute_location.py
@@ -17,7 +17,7 @@
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_parallel.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_parallel.py
index e2c371bf44..d11847c460 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_parallel.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_parallel.py
@@ -18,7 +18,7 @@
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_thread_binding.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_thread_binding.py
index cdfcaa20ad..2937fcd5ba 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_thread_binding.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_thread_binding.py
@@ -17,7 +17,7 @@
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_tile_size.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_tile_size.py
index 85a8a540a8..21a689fec1 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_tile_size.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_tile_size.py
@@ -20,7 +20,7 @@ from functools import reduce
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_unroll.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_unroll.py
index f771b2488d..52f3f14906 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_unroll.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_mutator_mutate_unroll.py
@@ -18,7 +18,7 @@
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_post_order_apply.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_post_order_apply.py
index e7bfaf3350..bccfa40a11 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_post_order_apply.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_post_order_apply.py
@@ -33,7 +33,7 @@ from tvm.s_tir.meta_schedule.schedule_rule import PyScheduleRule
from tvm.s_tir.meta_schedule.space_generator import PostOrderApply
from tvm.s_tir.schedule import SBlockRV, Schedule
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_disallow_async_strided_mem_copy.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_disallow_async_strided_mem_copy.py
index c0805a106e..032348671a 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_disallow_async_strided_mem_copy.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_disallow_async_strided_mem_copy.py
@@ -20,7 +20,7 @@
import tvm, tvm_tirx
from tvm_tirx import tirx
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_disallow_dynamic_loop.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_disallow_dynamic_loop.py
index b76daa3dc3..c8fdff1bcf 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_disallow_dynamic_loop.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_disallow_dynamic_loop.py
@@ -20,7 +20,7 @@
import tvm, tvm_tirx
from tvm_tirx import tirx
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_cooperative_fetch.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_cooperative_fetch.py
index 4350191143..88e6343c06 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_cooperative_fetch.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_cooperative_fetch.py
@@ -22,7 +22,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing import te_workload
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
from tvm.te import create_function
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_layout.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_layout.py
index 71362b1367..c8682407c2 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_layout.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_layout.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py
index 4e8ba7dbae..bcefe183b1 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py
@@ -21,7 +21,7 @@ import tvm_tirx.testing, tvm
from tvm.s_tir.meta_schedule.postproc import RewriteParallelVectorizeUnroll
from tvm.s_tir.schedule import Schedule
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,not-callable,misplaced-comparison-constant
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_reduction_block.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_reduction_block.py
index 2787419b4f..e5696d822a 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_reduction_block.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_reduction_block.py
@@ -20,7 +20,7 @@
import tvm, tvm_tirx
from tvm_tirx import tirx
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_tensorize.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_tensorize.py
index 4ee3612b85..fe1f5aee90 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_tensorize.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_tensorize.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.tensor_intrin import dot_product_common, x86
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_unbound_block.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_unbound_block.py
index d66098f084..7949c73878 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_unbound_block.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_unbound_block.py
@@ -20,7 +20,7 @@
import tvm, tvm_tirx
from tvm_tirx import tirx
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py
index 113c3f8ad6..d65bdbd718 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_vtcm_limit.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_vtcm_limit.py
index 30de265812..3f2bae97c8 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_vtcm_limit.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_vtcm_limit.py
@@ -20,7 +20,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_runner.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_runner.py
index 88e3827dcf..43316c34c1 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_runner.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_runner.py
@@ -59,7 +59,7 @@ from tvm.s_tir.meta_schedule.utils import (
get_global_func_with_default_on_worker,
)
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
from tvm_tirx.tirx import FloatImm
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_add_rfactor.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_add_rfactor.py
index 17d45e8bd6..2717ff6e86 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_add_rfactor.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_add_rfactor.py
@@ -21,7 +21,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
check_sketches,
generate_design_space,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
from tvm.te import create_function
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py
index 8585caec88..da02f86654 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py
@@ -22,7 +22,7 @@ import pytest
import tvm, tvm_tirx
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.schedule_rule import ApplyCustomRule
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_auto_bind.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_auto_bind.py
index 219b4d4838..334f5ae704 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_auto_bind.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_auto_bind.py
@@ -21,7 +21,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
check_sketches,
generate_design_space,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_auto_inline.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_auto_inline.py
index 403e150a2e..fd258b00d7 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_auto_inline.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_auto_inline.py
@@ -24,7 +24,7 @@ from tvm_tirx.ir.base import assert_structural_equal
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing.space_generation import generate_design_space
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_cross_thread_reduction.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_cross_thread_reduction.py
index c9fe87f93f..2eac2a1ead 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_cross_thread_reduction.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_cross_thread_reduction.py
@@ -23,7 +23,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
check_sketches,
generate_design_space,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
from tvm.te import create_function
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt.py
index d488e49b9a..4e0f9c3fda 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt.py
@@ -25,7 +25,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
generate_design_space,
print_sketches,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt_intrin.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt_intrin.py
index efb47fa766..72e67e126d 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt_intrin.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_mlt_intrin.py
@@ -27,7 +27,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
from tvm.s_tir.tensor_intrin.arm_cpu import DP4A_S8S8S32_INTRIN
from tvm.s_tir.tensor_intrin.x86 import AVX512_DOT_16x4_INTRIN as AVX512_INTRIN
from tvm.s_tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_parallel_vectorize_unroll.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_parallel_vectorize_unroll.py
index 00698287e3..766f8dc52d 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_parallel_vectorize_unroll.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_parallel_vectorize_unroll.py
@@ -22,7 +22,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
check_sketches,
generate_design_space,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_random_compute_location.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_random_compute_location.py
index 4248eb33ef..623a549982 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_random_compute_location.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_random_compute_location.py
@@ -21,7 +21,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
check_sketches,
generate_design_space,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_search_strategy.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_search_strategy.py
index ff0d22f977..282eea7b07 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_search_strategy.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_search_strategy.py
@@ -27,7 +27,7 @@ from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing.dummy_object import DummyMutator
from tvm.s_tir.schedule import Schedule, Trace
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
MATMUL_M = 32
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cpu.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cpu.py
index 63e8433038..f227a5efb1 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cpu.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cpu.py
@@ -24,7 +24,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
print_sketches,
)
from tvm.s_tir.meta_schedule.testing.te_workload import create_te_workload
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda.py
index 587a1fbf41..6de9d90fd4 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda.py
@@ -24,7 +24,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
print_sketches,
)
from tvm.s_tir.meta_schedule.testing.te_workload import create_te_workload
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda_async.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda_async.py
index 3057559c88..d96701fee3 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda_async.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda_async.py
@@ -24,7 +24,7 @@ from tvm.s_tir.meta_schedule.testing.space_generation import (
print_sketches,
)
from tvm.s_tir.meta_schedule.testing.te_workload import create_te_workload
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_generator.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_generator.py
index 301f24ef25..709bb2e130 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_generator.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_generator.py
@@ -32,7 +32,7 @@ from tvm.s_tir.meta_schedule.space_generator import (
from tvm.s_tir.meta_schedule.tune_context import TuneContext
from tvm.s_tir.schedule import Schedule
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_post_opt.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_post_opt.py
index c3097b329e..2cdf16e61b 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_post_opt.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_post_opt.py
@@ -26,7 +26,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.runner.config import EvaluatorConfig
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
from tvm_tirx.testing import env
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_task_scheduler.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_task_scheduler.py
index 688f9a56fa..77eab1efc9 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_task_scheduler.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_task_scheduler.py
@@ -28,7 +28,7 @@ from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing.dummy_object import DummyBuilder, DummyRunner
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,missing-docstring
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
index a637fbae0c..0c7f8ed0de 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
@@ -23,7 +23,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx.codegen import llvm_lookup_intrinsic_id
from tvm.s_tir import Schedule
from tvm.s_tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
from tvm_tirx.tirx import floordiv, floormod
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_context.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_context.py
index 53de2ba9ab..86c572aebc 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_context.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_context.py
@@ -25,7 +25,7 @@ import tvm_ffi
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir.meta_schedule import TuneContext
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_tir.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_tir.py
index b2f21327ee..9940825376 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_tir.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_tir.py
@@ -30,7 +30,7 @@ from tvm.s_tir.meta_schedule.testing.custom_builder_runner import run_module_via
from tvm.s_tir.meta_schedule.testing.local_rpc import LocalRPC
from tvm.s_tir.schedule import SBlockRV, Schedule
from tvm.s_tir.utils import derived_object
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
from tvm_tirx.testing import env
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
index eb65388f75..18940a1ce7 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
@@ -34,7 +34,7 @@ from tvm.s_tir.schedule.analysis import (
suggest_index_map,
)
from tvm.s_tir.tensor_intrin.x86 import dot_product_16x4_u8i8i32_desc
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm.te import create_function
from tvm_tirx.tirx import IndexMap, Var, decl_tensor, floordiv, floormod
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_annotate_tensor_access.py b/tests/python/s_tir/schedule/test_tir_schedule_annotate_tensor_access.py
index 7e6fbab2f5..5b020ad86f 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_annotate_tensor_access.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_annotate_tensor_access.py
@@ -23,7 +23,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_block_scope.py b/tests/python/s_tir/schedule/test_tir_schedule_block_scope.py
index b60150d72e..449ff60735 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_block_scope.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_block_scope.py
@@ -26,7 +26,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir
from tvm.s_tir.schedule import DepKind
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_blockize.py b/tests/python/s_tir/schedule/test_tir_schedule_blockize.py
index 00f89b494f..398baa0371 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_blockize.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_blockize.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir.schedule.testing import verify_trace_roundtrip
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_cache_index.py b/tests/python/s_tir/schedule/test_tir_schedule_cache_index.py
index 548b40326d..8ba7d03824 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_cache_index.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_cache_index.py
@@ -24,7 +24,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir.schedule.testing import verify_trace_roundtrip
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py b/tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py
index 8a5188da65..c84d2e9680 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py
@@ -29,7 +29,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_compute_at.py b/tests/python/s_tir/schedule/test_tir_schedule_compute_at.py
index f94c2b57a7..2d0f7af255 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_compute_at.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_compute_at.py
@@ -28,7 +28,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_compute_inline.py b/tests/python/s_tir/schedule/test_tir_schedule_compute_inline.py
index e74c2f0ce4..336c815430 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_compute_inline.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_compute_inline.py
@@ -26,7 +26,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_decompose_padding.py b/tests/python/s_tir/schedule/test_tir_schedule_decompose_padding.py
index 4ac5183456..5515f928dd 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_decompose_padding.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_decompose_padding.py
@@ -23,7 +23,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_error.py b/tests/python/s_tir/schedule/test_tir_schedule_error.py
index 115a46adfb..1e00c88cfb 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_error.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_error.py
@@ -24,7 +24,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_for_kind.py b/tests/python/s_tir/schedule/test_tir_schedule_for_kind.py
index f2113aafe8..7fc2be231e 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_for_kind.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_for_kind.py
@@ -27,7 +27,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue.py b/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue.py
index 890d28612f..b20c5f5e97 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue.py
@@ -26,7 +26,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_clipping.py b/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_clipping.py
index b96119e176..96a5f40c8f 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_clipping.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_clipping.py
@@ -26,7 +26,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_relu.py b/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_relu.py
index e40c333fe6..b98e66cb91 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_relu.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_relu.py
@@ -26,7 +26,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_merge.py b/tests/python/s_tir/schedule/test_tir_schedule_merge.py
index b14f29c1ea..0729d78402 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_merge.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_merge.py
@@ -25,7 +25,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_pad_einsum.py b/tests/python/s_tir/schedule/test_tir_schedule_pad_einsum.py
index c40f0ba642..28e3b836f3 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_pad_einsum.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_pad_einsum.py
@@ -25,7 +25,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_partition.py b/tests/python/s_tir/schedule/test_tir_schedule_partition.py
index 23547a22f8..dfc117d188 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_partition.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_partition.py
@@ -28,7 +28,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx.expr import IntImm
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py b/tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py
index 28b09b2c3d..cc9e5e7089 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py
@@ -43,7 +43,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_reduction.py b/tests/python/s_tir/schedule/test_tir_schedule_reduction.py
index c1da4be9f6..de645600cc 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_reduction.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_reduction.py
@@ -29,7 +29,7 @@ from tvm.s_tir.schedule.testing import (
verify_trace_roundtrip,
)
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_reindex.py b/tests/python/s_tir/schedule/test_tir_schedule_reindex.py
index a019bbf456..296cc6c3d1 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_reindex.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_reindex.py
@@ -26,7 +26,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_reorder.py b/tests/python/s_tir/schedule/test_tir_schedule_reorder.py
index 50a5ec3d27..cb5ecdbaee 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_reorder.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_reorder.py
@@ -27,7 +27,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_reorder_block_iter_var.py b/tests/python/s_tir/schedule/test_tir_schedule_reorder_block_iter_var.py
index 8de2331719..ba55e0396b 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_reorder_block_iter_var.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_reorder_block_iter_var.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir.schedule.testing import verify_trace_roundtrip
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_rfactor.py b/tests/python/s_tir/schedule/test_tir_schedule_rfactor.py
index 6f40f26697..c756d13fc4 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_rfactor.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_rfactor.py
@@ -26,7 +26,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_rolling_buffer.py b/tests/python/s_tir/schedule/test_tir_schedule_rolling_buffer.py
index d37b2ea507..ffbcd0c702 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_rolling_buffer.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_rolling_buffer.py
@@ -26,7 +26,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_sampling.py b/tests/python/s_tir/schedule/test_tir_schedule_sampling.py
index ddaa1a2c6b..12cff6aa2f 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_sampling.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_sampling.py
@@ -24,7 +24,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm.s_tir.schedule.testing import verify_trace_roundtrip
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_set_dtype.py b/tests/python/s_tir/schedule/test_tir_schedule_set_dtype.py
index 449304526f..c9f0fcc733 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_set_dtype.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_set_dtype.py
@@ -26,7 +26,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_set_scope.py b/tests/python/s_tir/schedule/test_tir_schedule_set_scope.py
index a316c44566..588b79412a 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_set_scope.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_set_scope.py
@@ -25,7 +25,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_split_fuse.py b/tests/python/s_tir/schedule/test_tir_schedule_split_fuse.py
index 76835bf875..a0a118b137 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_split_fuse.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_split_fuse.py
@@ -28,7 +28,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx.expr import IntImm
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_state.py b/tests/python/s_tir/schedule/test_tir_schedule_state.py
index d820b4f786..e46b152288 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_state.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_state.py
@@ -26,7 +26,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm_tirx.ir import IRModule
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
@@ -95,7 +95,7 @@ def block_in_opaque_block(
def replace_ir_builder(deep_copy=False, realize=False):
new_func = tvm_tirx.script.from_source(
elementwise.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
s = tvm.s_tir.ScheduleState(new_func, debug_mask="all")
target = tvm.s_tir.SBlock(
@@ -124,11 +124,11 @@ def replace_ir_builder(deep_copy=False, realize=False):
def replace_ir_builder_module(deep_copy=False, realize=False):
new_func = tvm_tirx.script.from_source(
elementwise.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
other_func = tvm_tirx.script.from_source(
elementwise.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
mod = IRModule(functions={"main": new_func, "other": other_func})
s = tvm.s_tir.ScheduleState(mod, debug_mask="all")
@@ -158,7 +158,7 @@ def replace_ir_builder_module(deep_copy=False, realize=False):
def replace_ir_builder_with_opaque():
func = tvm_tirx.script.from_source(
block_in_opaque_block.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
s = tvm.s_tir.ScheduleState(func, debug_mask="all")
gc.collect()
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py b/tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py
index d296c81333..5bf2b2ea29 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py
@@ -26,7 +26,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir
from tvm.s_tir.schedule.state import CachedFlags
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_storage_align.py b/tests/python/s_tir/schedule/test_tir_schedule_storage_align.py
index dc6328f530..7041896fa0 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_storage_align.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_storage_align.py
@@ -25,7 +25,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_tensorize.py b/tests/python/s_tir/schedule/test_tir_schedule_tensorize.py
index 05f7f433c6..db254cf127 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_tensorize.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_tensorize.py
@@ -38,7 +38,7 @@ from tvm.s_tir.tensor_intrin.arm_cpu import (
from tvm.s_tir.tensor_intrin.hexagon import VDMPY_i16i16i32_INTRIN, VRMPY_u8u8i32_INTRIN
from tvm.s_tir.tensor_intrin.rocm import AMDGPU_SDOT4_INTRIN
from tvm.s_tir.tensor_intrin.x86 import AVX512_DOT_16x4_INTRIN, VNNI_DOT_16x4_INTRIN
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_trace.py b/tests/python/s_tir/schedule/test_tir_schedule_trace.py
index cd2a0683f6..d3e6b25fe4 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_trace.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_trace.py
@@ -27,7 +27,7 @@ from tvm_tirx import tirx
from tvm import s_tir
from tvm.s_tir.schedule import Instruction, InstructionKind, LoopRV, SBlockRV, Trace
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_transform.py b/tests/python/s_tir/schedule/test_tir_schedule_transform.py
index a3d338ca68..d50cb42d7e 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_transform.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_transform.py
@@ -18,7 +18,7 @@ import tvm, tvm_tirx
from tvm.s_tir import Schedule
from tvm.s_tir.schedule.transform import tile_with_tensor_intrin
from tvm.s_tir.tensor_intrin.x86 import AVX512_DOT_16x4_INTRIN, VNNI_DOT_16x4_INTRIN
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py
index a4b24a4b9d..331e037a88 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py
@@ -29,7 +29,7 @@ from tvm.s_tir.schedule.testing import (
verify_trace_roundtrip,
)
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_utilities.py b/tests/python/s_tir/schedule/test_tir_schedule_utilities.py
index e6e3c52691..0921c276a6 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_utilities.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_utilities.py
@@ -28,7 +28,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable
diff --git a/tests/python/s_tir/schedule/test_tir_unsafe_hide_tensor_access.py b/tests/python/s_tir/schedule/test_tir_unsafe_hide_tensor_access.py
index 7dc932f6d5..53e823e736 100644
--- a/tests/python/s_tir/schedule/test_tir_unsafe_hide_tensor_access.py
+++ b/tests/python/s_tir/schedule/test_tir_unsafe_hide_tensor_access.py
@@ -25,7 +25,7 @@ from tvm.s_tir.schedule.testing import (
assert_structural_equal_ignore_global_symbol,
verify_trace_roundtrip,
)
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/script/test_s_tir_script_basic_usage.py b/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
index b3fe568b85..eb61d0d404 100644
--- a/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
+++ b/tests/python/s_tir/script/test_s_tir_script_basic_usage.py
@@ -26,7 +26,7 @@ import tvm_tirx.testing, tvm
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
from tvm_tirx.script import from_source
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
@@ -474,8 +474,8 @@ def test_for_thread_binding():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
@@ -761,8 +761,8 @@ def test_assert_stmt_roundtrip(error_type, message_parts):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, roundtrip, map_free_vars=True)
@@ -956,7 +956,7 @@ def func():
j{annotation} = i // 4
T.evaluate(j)
""",
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
binding = func.body[0].body.seq[0]
if mutable:
@@ -1076,8 +1076,8 @@ def test_roundtrip_basic_usage(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
diff --git a/tests/python/s_tir/script/test_s_tir_script_blocks.py b/tests/python/s_tir/script/test_s_tir_script_blocks.py
index 843d662f7f..8a99e5c18e 100644
--- a/tests/python/s_tir/script/test_s_tir_script_blocks.py
+++ b/tests/python/s_tir/script/test_s_tir_script_blocks.py
@@ -25,7 +25,7 @@ import tvm_tirx.testing, tvm
from tvm import s_tir
from tvm_tirx.ir import Range, TensorRegion, TensorRegionType
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# Expected completed types use the configured fixed index dtype, not shape dtype.
@@ -305,14 +305,14 @@ def expected_recursive_tensorslice_indices(
def test_complete_tensor_indices():
new_func = tvm_tirx.script.from_source(
func_with_tensorslice_indices.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
).with_attr("global_symbol", "main")
tvm_tirx.ir.assert_structural_equal(
new_func, expected_tensorslice_indices.with_attr("global_symbol", "main")
)
new_func = tvm_tirx.script.from_source(
func_with_recursive_tensorslice_indices.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
).with_attr("global_symbol", "main")
tvm_tirx.ir.assert_structural_equal(
new_func, expected_recursive_tensorslice_indices.with_attr("global_symbol", "main")
@@ -389,7 +389,7 @@ def expect_alloc_tensor_func(
def test_complete_alloc_tensor():
rt_func = tvm_tirx.script.from_source(
alloc_tensor_func.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
).with_attr("global_symbol", "main")
tvm_tirx.ir.assert_structural_equal(
rt_func, expect_alloc_tensor_func.with_attr("global_symbol", "main")
@@ -443,11 +443,11 @@ def test_alloc_zero_dim_tensor_round_trip():
func = alloc_zero_dim_tensor
func_with_block = alloc_zero_dim_tensor_block
rt_func = tvm_tirx.script.from_source(
- func.script(), extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir}
+ func.script(), extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
)
rt_func_with_block = tvm_tirx.script.from_source(
func_with_block.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
rt_mod = tvm.compile(rt_func, "llvm")
rt_mod_with_block = tvm.compile(rt_func_with_block, "llvm")
@@ -641,8 +641,8 @@ def test_matmul_original():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
@@ -663,8 +663,8 @@ def test_element_wise():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
@@ -687,8 +687,8 @@ def test_predicate():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
@@ -726,8 +726,8 @@ def test_match_tensor_region():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
@@ -776,8 +776,8 @@ def test_block_elements():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
@@ -817,8 +817,8 @@ def test_opaque_block():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
@@ -1025,8 +1025,8 @@ def test_roundtrip_blocks(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
diff --git a/tests/python/s_tir/script/test_s_tir_script_dynamic_shape.py b/tests/python/s_tir/script/test_s_tir_script_dynamic_shape.py
index 2d56718dac..9ed2c42e93 100644
--- a/tests/python/s_tir/script/test_s_tir_script_dynamic_shape.py
+++ b/tests/python/s_tir/script/test_s_tir_script_dynamic_shape.py
@@ -25,7 +25,7 @@ import pytest
import tvm_tirx.testing, tvm
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
from tvm_tirx.script import from_source
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
@@ -105,7 +105,7 @@ def gemm_dyn_shape(
def test_dynamic_shape_gemm():
gemm_dyn_shape_roundtrip = from_source(
gemm_dyn_shape.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
assert_structural_equal_ignore_global_symbol(gemm_dyn_shape, gemm_dyn_shape_roundtrip)
@@ -273,8 +273,8 @@ def test_roundtrip_dynamic_shape(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
diff --git a/tests/python/s_tir/script/test_s_tir_script_error_handling.py b/tests/python/s_tir/script/test_s_tir_script_error_handling.py
index b02c66f341..91e1dc83c1 100644
--- a/tests/python/s_tir/script/test_s_tir_script_error_handling.py
+++ b/tests/python/s_tir/script/test_s_tir_script_error_handling.py
@@ -30,7 +30,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import from_source
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/script/test_s_tir_script_ir_builder.py b/tests/python/s_tir/script/test_s_tir_script_ir_builder.py
index 59599cfa93..b2387b5d2e 100644
--- a/tests/python/s_tir/script/test_s_tir_script_ir_builder.py
+++ b/tests/python/s_tir/script/test_s_tir_script_ir_builder.py
@@ -25,7 +25,7 @@ from tvm import s_tir
from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.ir.location import SourceLoc, SourceName
from tvm.s_tir.script.ir_builder import function as build_function
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.tirx.script import ir_builder as T
diff --git a/tests/python/s_tir/script/test_s_tir_script_meta_programming.py b/tests/python/s_tir/script/test_s_tir_script_meta_programming.py
index fef353325f..596b18e8b8 100644
--- a/tests/python/s_tir/script/test_s_tir_script_meta_programming.py
+++ b/tests/python/s_tir/script/test_s_tir_script_meta_programming.py
@@ -26,7 +26,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
@@ -506,7 +506,7 @@ def matmul(
def test_multi_element_array_in_outmost_namespace():
func = matmul
rt_func = tvm_tirx.script.from_source(
- func.script(), extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir}
+ func.script(), extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
diff --git a/tests/python/s_tir/script/test_s_tir_script_namespace.py b/tests/python/s_tir/script/test_s_tir_script_namespace.py
index dcc05b300e..4dbd103c38 100644
--- a/tests/python/s_tir/script/test_s_tir_script_namespace.py
+++ b/tests/python/s_tir/script/test_s_tir_script_namespace.py
@@ -26,13 +26,13 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
@pytest.mark.skipif(sys.version_info < (3, 12), reason="PEP 695 requires Python 3.12")
def test_shared_operations_and_aliases():
- from tvm_tirx.script import s_tir as S
+ from tvm.s_tir import script as S
n = S.dynamic("n")
@@ -53,7 +53,7 @@ def test_shared_operations_and_aliases():
shared,
tvm_tirx.script.from_source(
shared.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
),
)
@@ -82,7 +82,7 @@ def test_mixed_module_roundtrip():
tvm_tirx.ir.assert_structural_equal(
Mixed,
tvm_tirx.script.from_source(
- script, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir}
+ script, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
),
)
direct_script = Mixed["direct"].script()
@@ -106,7 +106,7 @@ def test_tirx_construction_roundtrip_and_execution_are_independent(monkeypatch):
restored = tvm_tirx.script.from_source(
Direct.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
tvm_tirx.ir.assert_structural_equal(Direct, restored)
assert "s_tir" not in restored["main"].attrs
@@ -136,7 +136,7 @@ def test_legacy_mode_is_not_a_function_option(namespace, option):
tvm_tirx.script.from_source(
f"@{namespace}.function({option}=True, check_well_formed=False)\n"
"def main():\n T.evaluate(0)\n",
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
@@ -169,7 +169,7 @@ def test_tirx_rejects_s_tir_blocks():
tvm_tirx.script.from_source(
"@T.function(check_well_formed=False)\n"
"def main():\n with Ts.sblock('bad'):\n T.evaluate(0)\n",
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
@@ -179,7 +179,7 @@ def test_tirx_cannot_change_dialect_with_attribute(value):
tvm_tirx.script.from_source(
"@T.function(check_well_formed=False)\n"
f"def main():\n T.func_attr({{'s_tir': {value}}})\n T.evaluate(0)\n",
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
)
@@ -211,7 +211,7 @@ def test_s_tir_options_and_helpers():
scheduled,
tvm_tirx.script.from_source(
scheduled.script(),
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
),
)
@@ -223,7 +223,7 @@ def test_s_tir_options_and_helpers():
"tvm_tirx.tirx.script",
"tvm_tirx.script.parser",
"tvm.s_tir.script.ir_builder.frame",
- "tvm_tirx.script.ir_builder.s_tir.frame",
+ "tvm.s_tir.script.ir_builder",
],
)
def test_import_order(first):
@@ -233,7 +233,7 @@ def test_import_order(first):
"-c",
f"import {first}\n"
"import tvm.s_tir.script as direct\n"
- "from tvm_tirx.script import s_tir as Ts, tirx as T\n"
+ "from tvm.s_tir import script as Ts\nfrom tvm_tirx.script import tirx as T\n"
"from tvm_tirx.script.parser import s_tir as parser\n"
"assert Ts is direct is parser\n"
"assert Ts.Tensor is T.Tensor\n"
diff --git a/tests/python/s_tir/script/test_s_tir_script_printer.py b/tests/python/s_tir/script/test_s_tir_script_printer.py
index c85a6675cb..92bc63d305 100644
--- a/tests/python/s_tir/script/test_s_tir_script_printer.py
+++ b/tests/python/s_tir/script/test_s_tir_script_printer.py
@@ -30,7 +30,7 @@ from tvm_tirx.ir import Range, TensorRegion, TensorRegionType
from tvm.s_tir.script.ir_builder import function as build_function
from tvm_tirx.script import ir as I
from tvm_tirx.script import ir_builder as IB
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.tirx.script import ir_builder as TB
@@ -55,7 +55,7 @@ def test_function_symbolic_tensor_param_roundtrip():
assert source.index("def main[n](") < source.index("T.evaluate(n)")
tvm_tirx.ir.assert_structural_equal(
tvm_tirx.script.from_source(
- source, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir}
+ source, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
),
func,
)
@@ -76,7 +76,7 @@ def test_function_compound_tensor_shape_first_use_roundtrip():
assert source.index("def main[n](") < source.index("T.evaluate(n)")
tvm_tirx.ir.assert_structural_equal(
tvm_tirx.script.from_source(
- source, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir}
+ source, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
),
func,
)
@@ -109,7 +109,7 @@ def test_function_symbolic_alloc_tensor_roundtrip():
tvm_tirx.script.from_source(
source,
check_well_formed=False,
- extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir},
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
),
func,
)
@@ -132,7 +132,7 @@ def test_function():
expected="""
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -159,7 +159,7 @@ def test_function_tensor_data_use():
expected="""
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -188,7 +188,7 @@ def test_function_tensor_data_argument_is_scope_hint():
expected="""
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -286,7 +286,7 @@ def test_bind():
_assert_print(
obj,
"""
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function(private=True)
@@ -322,7 +322,7 @@ def test_remap():
v4, v5 = Ts.axis.remap("RS", [i4, i5])
expected_output = """
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -370,7 +370,7 @@ def test_root_block():
TB.evaluate(0)
expected_output = """
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -403,7 +403,7 @@ def test_private_function():
expected="""
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function(private=True)
@@ -424,7 +424,7 @@ def test_function_different_symbol():
expected_output = """
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -442,7 +442,7 @@ def test_return_statement():
return TB.int32(5)
expected_output = """
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -481,7 +481,7 @@ def test_custom_float_types(dtype):
TB.evaluate(getattr(TB, dtype)(0.0))
expected_output = f"""
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -523,7 +523,7 @@ def test_predicated_load_store():
expected_output = """
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -570,7 +570,7 @@ def test_predicated_tensor_load_store():
expected_output = """
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function(private=True)
@@ -607,7 +607,7 @@ def test_predicated_scalable_load_store():
expected_output = """
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -633,7 +633,7 @@ def test_masked_load_prevents_scalar_allocation_init_fusion():
assert "A[0] = T.masked_load" in source
tvm_tirx.ir.assert_structural_equal(
tvm_tirx.script.from_source(
- source, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir}
+ source, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
),
main,
)
@@ -649,7 +649,7 @@ def test_vload_with_explicit_scalable_data_type():
expected_output = """
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -672,7 +672,7 @@ def test_vectorize_llvm_pure_intrin():
expected_output = """
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -698,7 +698,7 @@ def test_func_with_loop_jumps():
expected_output = """
from __future__ import annotations
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -1272,8 +1272,8 @@ def test_same_name_var():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
@@ -1458,8 +1458,8 @@ def test_div_mod():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(func, rt_func, True)
@@ -1698,8 +1698,8 @@ def test_void_ptr_vs_handle():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
),
)
@@ -1718,8 +1718,8 @@ def test_void_ptr_vs_handle():
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
),
)
@@ -2072,7 +2072,7 @@ def test_annotation_multi_access_paths():
)
assert (
result
- == """# from tvm_tirx.script import s_tir as Ts
+ == """# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -2100,7 +2100,7 @@ def test_annotate_from_multi_obj():
)
assert (
result
- == """# from tvm_tirx.script import s_tir as Ts
+ == """# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -2133,7 +2133,7 @@ def test_disable_concise_scoping_when_scope_annotated():
)
assert (
result
- == """# from tvm_tirx.script import s_tir as Ts
+ == """# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@Ts.function
@@ -2163,7 +2163,7 @@ def test_ir_module():
_assert_module_representations(
mod,
"""
-# from tvm_tirx.script import s_tir as Ts
+# from tvm.s_tir import script as Ts
# from tvm_tirx.script import tirx as T
@T.ir_module
@@ -2223,8 +2223,8 @@ def test_roundtrip_lowered(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
@@ -2258,8 +2258,8 @@ def test_roundtrip_expressions(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
@@ -2291,8 +2291,8 @@ def test_roundtrip_tensors(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
@@ -2316,8 +2316,8 @@ def test_roundtrip_metadata(ir_generator):
extra_vars={
"I": tvm_tirx.script.ir,
"T": tvm_tirx.script.tirx,
- "Ts": tvm_tirx.script.s_tir,
- "R": tvm_tirx.script.relax,
+ "Ts": tvm.s_tir.script,
+ "R": tvm.relax.script,
},
)
tvm_tirx.ir.assert_structural_equal(original, after_roundtrip, map_free_vars=True)
diff --git a/tests/python/s_tir/script/test_s_tir_script_source_locations.py b/tests/python/s_tir/script/test_s_tir_script_source_locations.py
index 37f1f49c0f..3de5b52938 100644
--- a/tests/python/s_tir/script/test_s_tir_script_source_locations.py
+++ b/tests/python/s_tir/script/test_s_tir_script_source_locations.py
@@ -18,7 +18,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/test_arith_domain_touched.py b/tests/python/s_tir/test_arith_domain_touched.py
index 82571ea8b0..3a7394b77f 100644
--- a/tests/python/s_tir/test_arith_domain_touched.py
+++ b/tests/python/s_tir/test_arith_domain_touched.py
@@ -20,7 +20,7 @@ import pytest
import tvm_ffi
import tvm, tvm_tirx
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
m = T.dynamic("m", "int32")
diff --git a/tests/python/s_tir/test_s_tir_renew_defs.py b/tests/python/s_tir/test_s_tir_renew_defs.py
index 26105985e6..06210d940a 100644
--- a/tests/python/s_tir/test_s_tir_renew_defs.py
+++ b/tests/python/s_tir/test_s_tir_renew_defs.py
@@ -21,7 +21,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.ir import Var
from tvm.s_tir import SBlock
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx.function import Function
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_annotate_irregular_loop.py b/tests/python/s_tir/transform/test_s_tir_transform_annotate_irregular_loop.py
index fb17cb8a84..c1f6741770 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_annotate_irregular_loop.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_annotate_irregular_loop.py
@@ -23,7 +23,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_canonicalize_loop.py b/tests/python/s_tir/transform/test_s_tir_transform_canonicalize_loop.py
index 48c11d681b..0290952d7a 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_canonicalize_loop.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_canonicalize_loop.py
@@ -20,7 +20,7 @@ import pytest
import tvm, tvm_tirx
from tvm_tirx import tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_compact_tensor_region.py b/tests/python/s_tir/transform/test_s_tir_transform_compact_tensor_region.py
index 33bd10cac5..834c03b0ce 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_compact_tensor_region.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_compact_tensor_region.py
@@ -21,7 +21,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_convert_blocks_to_opaque.py b/tests/python/s_tir/transform/test_s_tir_transform_convert_blocks_to_opaque.py
index 03e51bec07..305934a423 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_convert_blocks_to_opaque.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_convert_blocks_to_opaque.py
@@ -22,7 +22,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir, te
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_default_gpu_schedule.py b/tests/python/s_tir/transform/test_s_tir_transform_default_gpu_schedule.py
index a93cf084d1..ab0a0d558b 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_default_gpu_schedule.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_default_gpu_schedule.py
@@ -19,7 +19,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm.s_tir.transform import DefaultGPUSchedule
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py
index 9a61db7898..66f8ed9d0d 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py
@@ -18,7 +18,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py
index 0d2c6b76c7..be44ac40d1 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py
@@ -25,7 +25,7 @@ import tvm.s_tir.testing
import tvm_tirx.testing, tvm
from tvm import s_tir
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
@@ -42,7 +42,7 @@ def test_cp_async_raw_dtype_round_trips():
T.s_tir.cp_async_raw(B.data, i * 16, A.data, i * 16, 16, ty="float16")
reparsed = tvm_tirx.script.from_source(
- f.script(), extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir}
+ f.script(), extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
)
tvm_tirx.ir.assert_structural_equal(f, reparsed)
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
index bff50ff2c0..4e0fb7c267 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
@@ -26,7 +26,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import te
from tvm.s_tir.meta_schedule.testing import te_workload
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
index d8d483822c..2f43c21c7a 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
@@ -20,7 +20,7 @@ import tvm_ffi
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lift_thread_binding.py b/tests/python/s_tir/transform/test_s_tir_transform_lift_thread_binding.py
index f052f00110..2937a64721 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lift_thread_binding.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lift_thread_binding.py
@@ -18,7 +18,7 @@
import tvm, tvm_tirx
from tvm_tirx import tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py b/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
index 1651d8c5da..c8c0e1ff61 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
@@ -22,7 +22,7 @@ import tvm_ffi
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.ir.module import IRModule
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py
index 1f8f04d210..eaece8d2cd 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_cross_thread_reduction.py
@@ -23,7 +23,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-member,invalid-name,unused-variable,unexpected-keyword-arg
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_init_block.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_init_block.py
index bd47d236e7..2ba8464acc 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_init_block.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_init_block.py
@@ -17,7 +17,7 @@
# ruff: noqa: F401, F821
import tvm, tvm_tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# pylint: disable=no-self-argument
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_match_tensor.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_match_tensor.py
index f44d283866..b236799b05 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_match_tensor.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_match_tensor.py
@@ -22,7 +22,7 @@ import pytest
import tvm, tvm_tirx
import tvm.s_tir
import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
index 546cc3b104..504c98eb64 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
@@ -19,7 +19,7 @@ from __future__ import annotations
import tvm, tvm_tirx
import tvm.s_tir
import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py
index e46c2f6ae2..bf9bfacd67 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py
@@ -22,7 +22,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import s_tir
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_binding.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_binding.py
index 1b8a390319..94930c2cf3 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_binding.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_thread_binding.py
@@ -22,7 +22,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_manifest_shared_memory_local_stage.py b/tests/python/s_tir/transform/test_s_tir_transform_manifest_shared_memory_local_stage.py
index 3a39f77407..3940f2a509 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_manifest_shared_memory_local_stage.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_manifest_shared_memory_local_stage.py
@@ -18,7 +18,7 @@
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
# fmt: off
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py b/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py
index 18981bfe84..725dbf929b 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py
@@ -23,7 +23,7 @@ import tvm_ffi
import tvm, tvm_tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py b/tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py
index 63af21b391..37283e1f55 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py
@@ -21,7 +21,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import s_tir
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm.topi.math import cast
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py b/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
index 957bc52a33..5b58bc1b1e 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
@@ -21,7 +21,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_plan_update_tensor_allocation_location.py b/tests/python/s_tir/transform/test_s_tir_transform_plan_update_tensor_allocation_location.py
index bd6d7e08d8..6ec0bf9a68 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_plan_update_tensor_allocation_location.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_plan_update_tensor_allocation_location.py
@@ -22,7 +22,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir
from tvm.s_tir.tensor_intrin.hexagon import VRMPY_u8u8i32_INTRIN
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_remove_undef.py b/tests/python/s_tir/transform/test_s_tir_transform_remove_undef.py
index dc2357864a..c366a9a8fe 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_remove_undef.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_remove_undef.py
@@ -20,7 +20,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_remove_weight_layout_rewrite_block.py b/tests/python/s_tir/transform/test_s_tir_transform_remove_weight_layout_rewrite_block.py
index a0ebd683b9..dfbaee591e 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_remove_weight_layout_rewrite_block.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_remove_weight_layout_rewrite_block.py
@@ -19,7 +19,7 @@ import sys
import tvm, tvm_tirx
from tvm_tirx.ir.module import IRModule
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx.function import Function
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_rewrite_unsafe_select.py b/tests/python/s_tir/transform/test_s_tir_transform_rewrite_unsafe_select.py
index cce7cd3826..7d43632576 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_rewrite_unsafe_select.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_rewrite_unsafe_select.py
@@ -18,7 +18,7 @@
import tvm, tvm_tirx
from tvm import s_tir
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py b/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
index cc205c6864..bbf76290c3 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
@@ -21,7 +21,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm import s_tir
from tvm_tirx.script import ir as I
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/script/test_constructor_contracts.py b/tests/python/script/test_constructor_contracts.py
index 6cb90862d4..a99d765641 100644
--- a/tests/python/script/test_constructor_contracts.py
+++ b/tests/python/script/test_constructor_contracts.py
@@ -26,7 +26,7 @@ import tvm_tirx.testing, tvm
from tvm_tirx import ir, tirx
from tvm import relax
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/target/test_arm_target.py b/tests/python/target/test_arm_target.py
index b0c3a6834b..acefce9a7f 100644
--- a/tests/python/target/test_arm_target.py
+++ b/tests/python/target/test_arm_target.py
@@ -25,7 +25,7 @@ import pytest
import tvm, tvm_tirx
from tvm_tirx import codegen
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/target/test_target_target.py b/tests/python/target/test_target_target.py
index 1ad749d77a..d9d9210f19 100644
--- a/tests/python/target/test_target_target.py
+++ b/tests/python/target/test_target_target.py
@@ -21,7 +21,7 @@ import tvm_ffi
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.target import Target
from tvm_tirx.testing import env
diff --git a/tests/python/te/test_te_create_function.py b/tests/python/te/test_te_create_function.py
index a4c90b9d25..8811415e40 100644
--- a/tests/python/te/test_te_create_function.py
+++ b/tests/python/te/test_te_create_function.py
@@ -25,7 +25,7 @@ import tvm, tvm_tirx
import tvm_tirx.testing, tvm
from tvm_tirx import tirx
from tvm import s_tir, te, topi
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
@@ -471,7 +471,7 @@ def test_tensor_attr():
)
func = te.create_function([A, B, C])
rt_func = tvm_tirx.script.from_source(
- func.script(), extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm_tirx.script.s_tir}
+ func.script(), extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
)
tvm_tirx.ir.assert_structural_equal(func, rt_func)
diff --git a/tests/python/s_tir/script/test_s_tir_script_source_locations.py b/tests/python/tirx/codegen/conftest.py
similarity index 60%
copy from tests/python/s_tir/script/test_s_tir_script_source_locations.py
copy to tests/python/tirx/codegen/conftest.py
index 37f1f49c0f..dabfb09775 100644
--- a/tests/python/s_tir/script/test_s_tir_script_source_locations.py
+++ b/tests/python/tirx/codegen/conftest.py
@@ -14,28 +14,20 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""S-TIR script source locations."""
+"""Hardware requirements for TIRx codegen tests."""
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
-from tvm_tirx.script import tirx as T
+from pathlib import Path
+from tvm_tirx.testing._pytest import (
+ _release_cuda_cache_between_tests as _release_cuda_cache_between_tests,
+)
+from tvm_tirx.testing._pytest import (
+ _source_codegen_without_cuda as _source_codegen_without_cuda,
+)
+from tvm_tirx.testing._pytest import (
+ codegen_collection_modifyitems,
+)
-def test_nesting_parsing():
- class dummy:
- pass
- for i in range(1):
-
- @tvm_tirx.script.ir_module
- class Module:
- @Ts.function
- def impl(
- A: T.Tensor((12, 196, 64), "float32"),
- ) -> None:
- T.evaluate(0)
-
-
-if __name__ == "__main__":
- tvm_tirx.testing.main()
+def pytest_collection_modifyitems(config, items):
+ codegen_collection_modifyitems(config, items, Path(__file__).resolve().parent)
diff --git a/tests/python/tirx/codegen/test_tirx_codegen_fp8.py b/tests/python/tirx/codegen/test_tirx_codegen_fp8.py
index 5838144065..1b5bfb3a49 100644
--- a/tests/python/tirx/codegen/test_tirx_codegen_fp8.py
+++ b/tests/python/tirx/codegen/test_tirx_codegen_fp8.py
@@ -19,7 +19,7 @@ import pytest
import tvm, tvm_tirx
import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/s_tir/script/test_s_tir_script_source_locations.py b/tests/python/tirx/conftest.py
similarity index 60%
copy from tests/python/s_tir/script/test_s_tir_script_source_locations.py
copy to tests/python/tirx/conftest.py
index 37f1f49c0f..8e18f0bd93 100644
--- a/tests/python/s_tir/script/test_s_tir_script_source_locations.py
+++ b/tests/python/tirx/conftest.py
@@ -14,28 +14,8 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""S-TIR script source locations."""
+"""Shared CUDA device and exact-architecture handling for TIRx tests."""
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx.script import s_tir as Ts
-from tvm_tirx.script import tirx as T
-
-
-def test_nesting_parsing():
- class dummy:
- pass
-
- for i in range(1):
-
- @tvm_tirx.script.ir_module
- class Module:
- @Ts.function
- def impl(
- A: T.Tensor((12, 196, 64), "float32"),
- ) -> None:
- T.evaluate(0)
-
-
-if __name__ == "__main__":
- tvm_tirx.testing.main()
+from tvm_tirx.testing import _pytest as _pytest
+from tvm_tirx.testing._pytest import pytest_configure as pytest_configure
+from tvm_tirx.testing._pytest import pytest_runtest_setup as pytest_runtest_setup
diff --git a/tirx/tests/python/tirx/conftest.py b/tirx/python/tvm_tirx/testing/_pytest.py
similarity index 79%
copy from tirx/tests/python/tirx/conftest.py
copy to tirx/python/tvm_tirx/testing/_pytest.py
index fa896876f5..89bd9497ff 100644
--- a/tirx/tests/python/tirx/conftest.py
+++ b/tirx/python/tvm_tirx/testing/_pytest.py
@@ -16,8 +16,10 @@
# under the License.
"""Shared CUDA device and exact-architecture handling for TIRx tests."""
+import gc
import os
import re
+from pathlib import Path
import pytest
@@ -110,3 +112,35 @@ def pytest_runtest_setup(item):
f"requires CUDA architecture {required} on {device_source}; "
f"actual architecture is {actual or 'unavailable'}"
)
+
+
+@pytest.fixture(autouse=True)
+def _source_codegen_without_cuda(monkeypatch):
+ # CPU-only source checks must not depend on a CUDA compiler being installed.
+ if not env.has_cuda():
+ monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+
+
+@pytest.fixture(autouse=True)
+def _release_cuda_cache_between_tests():
+ yield
+ gc.collect()
+ try:
+ import torch
+ except ImportError:
+ return
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+
+
+def codegen_collection_modifyitems(config, items, suite_root):
+ if env.has_cuda_compute(10):
+ return
+ suite_root = Path(suite_root).resolve()
+ skip = pytest.mark.skip(reason="requires a CUDA compute capability 10.0 device")
+ for item in items:
+ if (
+ Path(item.path).resolve().is_relative_to(suite_root)
+ and item.get_closest_marker("gpu") is not None
+ ):
+ item.add_marker(skip)
diff --git a/tirx/tests/python/runtime/test_runtime_memory_scope.py b/tirx/tests/python/runtime/test_runtime_memory_scope.py
index 6f0e0c7e55..23602b279d 100644
--- a/tirx/tests/python/runtime/test_runtime_memory_scope.py
+++ b/tirx/tests/python/runtime/test_runtime_memory_scope.py
@@ -126,7 +126,7 @@ def test_opencl_memory_scope_registration_load_order(load_order):
)
if not opencl_path.is_file():
pytest.skip("OpenCL runtime sidecar is not built")
- compiler = tvm_tirx.base._LOADED_LIBS.get("tvm_compiler")
+ compiler = tvm_tirx.base._LOADED_LIBS.get("tvm_tirx")
if compiler is None and load_order != "runtime_only":
pytest.skip("Compiler library is not available")
env = {
diff --git a/tirx/tests/python/tirx/codegen/conftest.py b/tirx/tests/python/tirx/codegen/conftest.py
index 769f579282..dabfb09775 100644
--- a/tirx/tests/python/tirx/codegen/conftest.py
+++ b/tirx/tests/python/tirx/codegen/conftest.py
@@ -16,41 +16,18 @@
# under the License.
"""Hardware requirements for TIRx codegen tests."""
-import gc
from pathlib import Path
-import pytest
-
-from tvm_tirx.testing import env
-
-
-@pytest.fixture(autouse=True)
-def _source_codegen_without_cuda(monkeypatch):
- # CPU-only source checks must not depend on a CUDA compiler being installed.
- if not env.has_cuda():
- monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
-
-
-@pytest.fixture(autouse=True)
-def _release_cuda_cache_between_tests():
- yield
- gc.collect()
- try:
- import torch
- except ImportError:
- return
- if torch.cuda.is_available():
- torch.cuda.empty_cache()
+from tvm_tirx.testing._pytest import (
+ _release_cuda_cache_between_tests as _release_cuda_cache_between_tests,
+)
+from tvm_tirx.testing._pytest import (
+ _source_codegen_without_cuda as _source_codegen_without_cuda,
+)
+from tvm_tirx.testing._pytest import (
+ codegen_collection_modifyitems,
+)
def pytest_collection_modifyitems(config, items):
- if env.has_cuda_compute(10):
- return
- suite_root = Path(__file__).resolve().parent
- skip = pytest.mark.skip(reason="requires a CUDA compute capability 10.0 device")
- for item in items:
- if (
- Path(item.path).resolve().is_relative_to(suite_root)
- and item.get_closest_marker("gpu") is not None
- ):
- item.add_marker(skip)
+ codegen_collection_modifyitems(config, items, Path(__file__).resolve().parent)
diff --git a/tirx/tests/python/tirx/conftest.py b/tirx/tests/python/tirx/conftest.py
index fa896876f5..8e18f0bd93 100644
--- a/tirx/tests/python/tirx/conftest.py
+++ b/tirx/tests/python/tirx/conftest.py
@@ -16,97 +16,6 @@
# under the License.
"""Shared CUDA device and exact-architecture handling for TIRx tests."""
-import os
-import re
-
-import pytest
-
-from tvm_tirx.testing import env
-
-_CUDA_ARCH_PATTERN = re.compile(r"sm_[1-9][0-9]*(?:a|f)?\Z")
-_XDIST_CUDA_DEVICE = None
-
-
-def _visible_cuda_archs():
- try:
- import torch
- except ImportError:
- return ()
-
- if not torch.cuda.is_available():
- return ()
- return tuple(env.cuda_arch(device_id) for device_id in range(torch.cuda.device_count()))
-
-
-def _require_homogeneous_cuda_archs():
- arches = _visible_cuda_archs()
- if not arches:
- return None
- if any(arch is None for arch in arches):
- raise pytest.UsageError(f"could not determine every visible CUDA architecture: {arches}")
- if len(set(arches)) != 1:
- profile = ", ".join(f"cuda:{device_id}={arch}" for device_id, arch in enumerate(arches))
- raise pytest.UsageError(
- "mixed-architecture CUDA pools are unsupported by the TIRx test suite: " + profile
- )
- return arches[0]
-
-
-def _set_cuda_device_for_xdist_worker():
- global _XDIST_CUDA_DEVICE
-
- try:
- import torch
- except ImportError:
- return None
-
- if not torch.cuda.is_available():
- return None
- if _XDIST_CUDA_DEVICE is None:
- worker = os.environ.get("PYTEST_XDIST_WORKER", "gw0")
- worker_index = int(worker[2:]) if worker.startswith("gw") and worker[2:].isdigit() else 0
- _XDIST_CUDA_DEVICE = worker_index % torch.cuda.device_count()
- torch.cuda.set_device(_XDIST_CUDA_DEVICE)
- return _XDIST_CUDA_DEVICE
-
-
-def pytest_configure(config):
- del config
- _require_homogeneous_cuda_archs()
- _set_cuda_device_for_xdist_worker()
-
-
-def pytest_runtest_setup(item):
- current_device = _set_cuda_device_for_xdist_worker()
- marker = item.get_closest_marker("cuda_arch")
- if marker is None:
- return
-
- unknown_kwargs = set(marker.kwargs) - {"device"}
- if unknown_kwargs:
- raise pytest.UsageError(
- f"cuda_arch marker has unsupported keyword(s): {sorted(unknown_kwargs)}"
- )
- arches = tuple(marker.args)
- if not arches or any(
- not isinstance(arch, str) or _CUDA_ARCH_PATTERN.fullmatch(arch) is None for arch in arches
- ):
- raise pytest.UsageError(
- "cuda_arch marker requires one or more canonical architectures such as sm_100a"
- )
-
- device_source = marker.kwargs.get("device", "cuda0")
- if device_source == "cuda0":
- device_id = 0
- elif device_source == "current":
- device_id = current_device if current_device is not None else 0
- else:
- raise pytest.UsageError("cuda_arch marker device must be 'cuda0' or 'current'")
-
- actual = env.cuda_arch(device_id)
- if actual not in arches:
- required = ", ".join(arches)
- pytest.skip(
- f"requires CUDA architecture {required} on {device_source}; "
- f"actual architecture is {actual or 'unavailable'}"
- )
+from tvm_tirx.testing import _pytest as _pytest
+from tvm_tirx.testing._pytest import pytest_configure as pytest_configure
+from tvm_tirx.testing._pytest import pytest_runtest_setup as pytest_runtest_setup
diff --git a/tirx/tests/python/tirx/script/test_tirx_script_tracing.py b/tirx/tests/python/tirx/script/test_tirx_script_tracing.py
index cde7d9e796..bea151d4ac 100644
--- a/tirx/tests/python/tirx/script/test_tirx_script_tracing.py
+++ b/tirx/tests/python/tirx/script/test_tirx_script_tracing.py
@@ -301,7 +301,7 @@ def test_shared_tensor_shapes_are_checked_by_compiled_host():
I.arg_("x", T.Tensor((n,), T.float32))
I.arg_("y", T.Tensor((n,), T.float32))
I.return_(n)
- compiled = tvm_tirx.compile(builder.get(), target="llvm", tir_pipeline="tirx")
+ compiled = tvm_tirx.tirx.compile(builder.get(), target="llvm", pipeline="tirx")
assert compiled(torch.zeros(3), torch.zeros(3)) == 3
for other, error, message in (
(torch.zeros(4), ValueError, "expected to match"),
diff --git a/tirx/tests/python/tirx/test_cuda_arch_marker.py b/tirx/tests/python/tirx/test_cuda_arch_marker.py
index 261657b424..c571924b07 100644
--- a/tirx/tests/python/tirx/test_cuda_arch_marker.py
+++ b/tirx/tests/python/tirx/test_cuda_arch_marker.py
@@ -27,7 +27,7 @@ def _tirx_conftest(pytestconfig):
for plugin in pytestconfig.pluginmanager.get_plugins():
plugin_path = getattr(plugin, "__file__", None)
if plugin_path is not None and Path(plugin_path).resolve() == expected:
- return plugin
+ return plugin._pytest
raise AssertionError(f"TIRx conftest plugin {expected} is not loaded")
diff --git a/tirx/tests/python/tirx/test_cuda_backend_config.py b/tirx/tests/python/tirx/test_cuda_backend_config.py
index df65a07b10..6c7ad26852 100644
--- a/tirx/tests/python/tirx/test_cuda_backend_config.py
+++ b/tirx/tests/python/tirx/test_cuda_backend_config.py
@@ -656,7 +656,7 @@ def main(A: T.Tensor((32,), "float32")):
target = TARGET
if target == "target_object":
target = tvm_tirx.target.Target("cuda", host="llvm")
-module = tvm_tirx.compile(main, target=target)
+module = tvm_tirx.tirx.compile(main, target=target)
metadata = module.mod.imports[0].inspect_source("backend_config")
from tvm_tirx.backend.config import parse_backend_config
assert parse_backend_config(metadata)["cuda"]["arch"] == "sm_100a"
diff --git a/tirx/tests/python/tirx/transform/test_tirx_transform_lower_intrin.py b/tirx/tests/python/tirx/transform/test_tirx_transform_lower_intrin.py
index 9fd84cd39d..f8c8b8f168 100644
--- a/tirx/tests/python/tirx/transform/test_tirx_transform_lower_intrin.py
+++ b/tirx/tests/python/tirx/transform/test_tirx_transform_lower_intrin.py
@@ -168,7 +168,7 @@ def test_lower_floordiv_and_floormod(dtype, kind, variant):
[tvm_tirx.tirx.transform.StmtSimplify(), tvm_tirx.tirx.transform.LowerIntrin()]
)(Before)
tvm_tirx.ir.assert_structural_equal(After, Expected)
- compiled = tvm_tirx.compile(After, "llvm")
+ compiled = tvm_tirx.tirx.compile(After, "llvm")
data = itertools.product(range(-10, 10), [y for y in range(-10, 10) if y])
for x, y in data:
if (variant == "pow2" and y != 8) or (variant == "offset" and y != 5):
@@ -226,7 +226,7 @@ def test_lower_floordiv_overflow_checks(dtype, clamp, data):
[tvm_tirx.tirx.transform.StmtSimplify(), tvm_tirx.tirx.transform.LowerIntrin()]
)(Before)
tvm_tirx.ir.assert_structural_equal(After, Expected)
- compiled = tvm_tirx.compile(After, "llvm")
+ compiled = tvm_tirx.tirx.compile(After, "llvm")
for value in data:
expected = min(max(value, -10), 32758) if clamp else value
assert compiled(value) == expected // 3
(tvm) 07/33: [REFACTOR][Native] Export core declarations and retain parent runtime consumers
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 9f9c9fe550f5d6d18c5b5527d0f74341b62de785
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 01:37:32 2026 +0000
[REFACTOR][Native] Export core declarations and retain parent runtime consumers
---
apps/android_rpc/app/src/main/jni/Android.mk | 2 +-
apps/android_rpc/app/src/main/jni/tvm_runtime.h | 20 ++++++++++----------
include/tvm/runtime/vm/executable.h | 2 +-
include/tvm/runtime/vm/tensor_cache_support.h | 8 ++++----
src/runtime/extra/contrib/clml/clml_runtime.h | 2 +-
src/runtime/extra/contrib/cublas/cublas.cc | 2 +-
src/runtime/extra/contrib/hipblas/hipblas.cc | 2 +-
src/runtime/extra/contrib/sort/sort.cc | 2 +-
src/runtime/rpc/rpc_socket_impl.cc | 4 ++--
src/runtime/vm/tensor_cache_support.cc | 4 ++--
tirx/include/tvm/codegen/codegen.h | 6 +++---
tirx/include/tvm/codegen/intrin_rule.h | 4 ++--
tirx/include/tvm/runtime/base.h | 19 +++++++++++++------
web/Makefile | 4 ++--
web/emcc/wasm_runtime.cc | 16 ++++++++--------
web/emcc/webgpu_runtime.cc | 2 +-
16 files changed, 53 insertions(+), 46 deletions(-)
diff --git a/apps/android_rpc/app/src/main/jni/Android.mk b/apps/android_rpc/app/src/main/jni/Android.mk
index b7540521c2..f38d54224a 100644
--- a/apps/android_rpc/app/src/main/jni/Android.mk
+++ b/apps/android_rpc/app/src/main/jni/Android.mk
@@ -36,7 +36,7 @@ include $(config)
LOCAL_SRC_FILES := org_apache_tvm_native_c_api.cc
LOCAL_LDFLAGS := -L$(SYSROOT)/usr/lib/ -llog
-LOCAL_C_INCLUDES := $(ROOT_PATH)/include \
+LOCAL_C_INCLUDES := $(ROOT_PATH)/include $(ROOT_PATH)/tirx/include \
$(ROOT_PATH)/3rdparty/tvm-ffi/include \
$(ROOT_PATH)/3rdparty/tvm-ffi/3rdparty/dlpack/include \
$(ROOT_PATH)/3rdparty/OpenCL-Headers
diff --git a/apps/android_rpc/app/src/main/jni/tvm_runtime.h b/apps/android_rpc/app/src/main/jni/tvm_runtime.h
index 40f1b76e56..fa5784e0d7 100644
--- a/apps/android_rpc/app/src/main/jni/tvm_runtime.h
+++ b/apps/android_rpc/app/src/main/jni/tvm_runtime.h
@@ -43,11 +43,11 @@
#include "../3rdparty/tvm-ffi/src/ffi/function.cc"
#include "../3rdparty/tvm-ffi/src/ffi/object.cc"
#include "../3rdparty/tvm-ffi/src/ffi/tensor.cc"
-#include "../src/runtime/cpu_device_api.cc"
-#include "../src/runtime/device_api.cc"
-#include "../src/runtime/logging.cc"
-#include "../src/runtime/memory_scope.cc"
-#include "../src/runtime/metadata.cc"
+#include "../tirx/src/runtime/cpu_device_api.cc"
+#include "../tirx/src/runtime/device_api.cc"
+#include "../tirx/src/runtime/logging.cc"
+#include "../tirx/src/runtime/memory_scope.cc"
+#include "../tirx/src/runtime/metadata.cc"
#include "../src/runtime/registry.cc"
#include "../src/runtime/rpc/rpc_channel.cc"
#include "../src/runtime/rpc/rpc_endpoint.cc"
@@ -57,12 +57,12 @@
#include "../src/runtime/rpc/rpc_server_env.cc"
#include "../src/runtime/rpc/rpc_session.cc"
#include "../src/runtime/rpc/rpc_socket_impl.cc"
-#include "../src/runtime/tensor.cc"
-#include "../src/runtime/thread_pool.cc"
-#include "../src/runtime/threading_backend.cc"
-#include "../src/runtime/timer.cc"
+#include "../tirx/src/runtime/tensor.cc"
+#include "../tirx/src/runtime/thread_pool.cc"
+#include "../tirx/src/runtime/threading_backend.cc"
+#include "../tirx/src/runtime/timer.cc"
#include "../src/runtime/vm/params.cc"
-#include "../src/runtime/workspace_pool.cc"
+#include "../tirx/src/runtime/workspace_pool.cc"
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
diff --git a/include/tvm/runtime/vm/executable.h b/include/tvm/runtime/vm/executable.h
index 5d4ff4aed3..6b1eecd1f2 100644
--- a/include/tvm/runtime/vm/executable.h
+++ b/include/tvm/runtime/vm/executable.h
@@ -85,7 +85,7 @@ struct VMFuncInfo {
* The executable contains information (e.g. data in different memory regions)
* to run in a virtual machine.
*/
-class TVM_RUNTIME_DLL VMExecutable : public ffi::ModuleObj {
+class TVM_RUNTIME_EXTRA_DLL VMExecutable : public ffi::ModuleObj {
public:
/*! \brief Get the property of the runtime module .*/
int GetPropertyMask() const final { return ffi::Module::kBinarySerializable; };
diff --git a/include/tvm/runtime/vm/tensor_cache_support.h b/include/tvm/runtime/vm/tensor_cache_support.h
index a02532124a..d5b1546e96 100644
--- a/include/tvm/runtime/vm/tensor_cache_support.h
+++ b/include/tvm/runtime/vm/tensor_cache_support.h
@@ -46,7 +46,7 @@ struct TensorCacheMetadata {
* \param staging_buffer The buffer to be used to avoid extra OpenCL copies. Pass in a nullptr
* in other cases
*/
- TVM_RUNTIME_DLL ffi::Tensor Load(Device device, const std::string* raw_data,
+ TVM_RUNTIME_EXTRA_DLL ffi::Tensor Load(Device device, const std::string* raw_data,
ffi::Optional<ffi::Tensor>* staging_buffer = nullptr) const;
/*! \brief Name of the parameter */
@@ -64,7 +64,7 @@ struct TensorCacheMetadata {
};
/*! \brief Load a FileRecord into memory */
- TVM_RUNTIME_DLL ffi::Array<ffi::Tensor> Load(
+ TVM_RUNTIME_EXTRA_DLL ffi::Array<ffi::Tensor> Load(
Device device, //
const std::string& path_prefix, //
std::string* raw_data_buffer, //
@@ -85,9 +85,9 @@ struct TensorCacheMetadata {
std::string path;
/*! \brief Load the metadata from a specific directory */
- TVM_RUNTIME_DLL static TensorCacheMetadata Load(const std::string& path);
+ TVM_RUNTIME_EXTRA_DLL static TensorCacheMetadata Load(const std::string& path);
/*! \brief Load the metadata from a given JSON string */
- TVM_RUNTIME_DLL static TensorCacheMetadata LoadFromStr(const std::string& json_str,
+ TVM_RUNTIME_EXTRA_DLL static TensorCacheMetadata LoadFromStr(const std::string& json_str,
const std::string& path);
};
diff --git a/src/runtime/extra/contrib/clml/clml_runtime.h b/src/runtime/extra/contrib/clml/clml_runtime.h
index df0b19b276..986088aaa0 100644
--- a/src/runtime/extra/contrib/clml/clml_runtime.h
+++ b/src/runtime/extra/contrib/clml/clml_runtime.h
@@ -50,7 +50,7 @@
#ifdef TVM_GRAPH_EXECUTOR_CLML
#include <CL/cl_qcom_ml_ops.h>
-#include "../../../../backend/opencl/runtime/opencl_common.h"
+#include "backend/opencl/runtime/opencl_common.h"
#define CAT_I(a, b) a##b
#define CAT(a, b) CAT_I(a, b)
diff --git a/src/runtime/extra/contrib/cublas/cublas.cc b/src/runtime/extra/contrib/cublas/cublas.cc
index 461bbee1f8..5c293fcd4b 100644
--- a/src/runtime/extra/contrib/cublas/cublas.cc
+++ b/src/runtime/extra/contrib/cublas/cublas.cc
@@ -27,7 +27,7 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
-#include "../../../../../3rdparty/compiler-rt/builtin_fp16.h"
+#include "builtin_fp16.h"
#include "../cblas/gemm_common.h"
#include "cublas_utils.h"
diff --git a/src/runtime/extra/contrib/hipblas/hipblas.cc b/src/runtime/extra/contrib/hipblas/hipblas.cc
index 18e136b0fd..f1e8f890ca 100644
--- a/src/runtime/extra/contrib/hipblas/hipblas.cc
+++ b/src/runtime/extra/contrib/hipblas/hipblas.cc
@@ -26,7 +26,7 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
-#include "../../../../../3rdparty/compiler-rt/builtin_fp16.h"
+#include "builtin_fp16.h"
#include "../cblas/gemm_common.h"
#include "hipblas_utils.h"
diff --git a/src/runtime/extra/contrib/sort/sort.cc b/src/runtime/extra/contrib/sort/sort.cc
index 6e3a99f935..8ef00c51f3 100644
--- a/src/runtime/extra/contrib/sort/sort.cc
+++ b/src/runtime/extra/contrib/sort/sort.cc
@@ -31,7 +31,7 @@
#include <algorithm>
#include <vector>
-#include "../../../../../3rdparty/compiler-rt/builtin_fp16.h"
+#include "builtin_fp16.h"
namespace tvm {
namespace contrib {
diff --git a/src/runtime/rpc/rpc_socket_impl.cc b/src/runtime/rpc/rpc_socket_impl.cc
index 7f7486d26c..cff92a17eb 100644
--- a/src/runtime/rpc/rpc_socket_impl.cc
+++ b/src/runtime/rpc/rpc_socket_impl.cc
@@ -115,8 +115,8 @@ ffi::Module RPCClientConnect(std::string url, int port, std::string key, bool en
return CreateRPCSessionModule(CreateClientSession(endpt));
}
-// TVM_RUNTIME_DLL needed for MSVC
-TVM_RUNTIME_DLL void RPCServerLoop(int sockfd) {
+// TVM_RUNTIME_EXTRA_DLL needed for MSVC
+TVM_RUNTIME_EXTRA_DLL void RPCServerLoop(int sockfd) {
support::TCPSocket sock(static_cast<support::TCPSocket::SockType>(sockfd));
RPCEndpoint::Create(std::make_unique<SockChannel>(sock), "SockServerLoop", "")->ServerLoop();
}
diff --git a/src/runtime/vm/tensor_cache_support.cc b/src/runtime/vm/tensor_cache_support.cc
index d28b68bf7a..0c7559bcd6 100644
--- a/src/runtime/vm/tensor_cache_support.cc
+++ b/src/runtime/vm/tensor_cache_support.cc
@@ -109,7 +109,7 @@ TensorCacheMetadata TensorCacheMetadata::LoadFromStr(const std::string& json_str
return result;
}
-TVM_RUNTIME_DLL TensorCacheMetadata TensorCacheMetadata::Load(const std::string& path) {
+TVM_RUNTIME_EXTRA_DLL TensorCacheMetadata TensorCacheMetadata::Load(const std::string& path) {
std::string json_str = support::ReadBytesFromFile(path + "/tensor-cache.json");
ffi::String err;
json::Value json_info = json::Parse(json_str, &err);
@@ -169,7 +169,7 @@ ffi::Tensor TensorCacheMetadata::FileRecord::ParamRecord::Load(
return arr;
}
-TVM_RUNTIME_DLL ffi::Array<ffi::Tensor> TensorCacheMetadata::FileRecord::Load(
+TVM_RUNTIME_EXTRA_DLL ffi::Array<ffi::Tensor> TensorCacheMetadata::FileRecord::Load(
Device device,
const std::string& path_prefix, //
std::string* raw_data_buffer, //
diff --git a/tirx/include/tvm/codegen/codegen.h b/tirx/include/tvm/codegen/codegen.h
index 2261b18290..6e2925a9cf 100644
--- a/tirx/include/tvm/codegen/codegen.h
+++ b/tirx/include/tvm/codegen/codegen.h
@@ -41,7 +41,7 @@ namespace codegen {
* \param target The target to be built.
* \return The generated module.
*/
-TVM_DLL ffi::Module Build(IRModule mod, Target target);
+TVM_TIRX_DLL ffi::Module Build(IRModule mod, Target target);
/*!
* \brief Pack imported device library to a C file.
@@ -54,7 +54,7 @@ TVM_DLL ffi::Module Build(IRModule mod, Target target);
* \param c_symbol_prefix Optional symbol prefix of the blob symbol.
* \return cstr The C string representation of the file.
*/
-TVM_DLL std::string PackImportsToC(const ffi::Module& m, bool system_lib,
+TVM_TIRX_DLL std::string PackImportsToC(const ffi::Module& m, bool system_lib,
const std::string& c_symbol_prefix = "");
/*!
@@ -69,7 +69,7 @@ TVM_DLL std::string PackImportsToC(const ffi::Module& m, bool system_lib,
* \param file_name Output path; its extension selects the object format.
* \param c_symbol_prefix Optional blob symbol prefix, valid only for system libraries.
*/
-TVM_DLL void PackImportsToObjectFile(const ffi::Module& mod, bool system_lib,
+TVM_TIRX_DLL void PackImportsToObjectFile(const ffi::Module& mod, bool system_lib,
const std::string& llvm_target_string,
const std::string& file_name,
const std::string& c_symbol_prefix = "");
diff --git a/tirx/include/tvm/codegen/intrin_rule.h b/tirx/include/tvm/codegen/intrin_rule.h
index d3d196f7d7..4d0f866d88 100644
--- a/tirx/include/tvm/codegen/intrin_rule.h
+++ b/tirx/include/tvm/codegen/intrin_rule.h
@@ -101,14 +101,14 @@ inline PrimExpr DispatchPureExtern(const PrimExpr& e) {
* The expansion clamps the input to [-4, 4] and evaluates numerator and
* denominator polynomials using the input element width; other types are rejected.
*/
-TVM_DLL PrimExpr DispatchFastErf(const PrimExpr& e);
+TVM_TIRX_DLL PrimExpr DispatchFastErf(const PrimExpr& e);
/*!
* \brief Lower a unary floating-point tanh call using sign-dependent exponential formulas.
* The expansion avoids the large-positive-input overflow of the symmetric formula;
* accuracy follows the floating-point arithmetic and exponential lowering in use.
*/
-TVM_DLL PrimExpr DispatchNumericallyStableTanh(const PrimExpr& e);
+TVM_TIRX_DLL PrimExpr DispatchNumericallyStableTanh(const PrimExpr& e);
} // namespace intrin
} // namespace codegen
diff --git a/tirx/include/tvm/runtime/base.h b/tirx/include/tvm/runtime/base.h
index 09ec0917f5..aa9c527d38 100644
--- a/tirx/include/tvm/runtime/base.h
+++ b/tirx/include/tvm/runtime/base.h
@@ -32,7 +32,7 @@
// passes the setuptools_scm-resolved version through CMake). The literal below is the
// fallback for a bare build with no override.
#ifndef TVM_VERSION
-#define TVM_VERSION "0.28.dev0"
+#define TVM_VERSION "0.1.0a0"
#endif
// The compiler, core runtime and runtime-extra libraries expose independent
@@ -77,12 +77,19 @@
#define TVM_DLL_EXPORT __attribute__((visibility("default")))
#endif
-// Core IR, TIRx and shared compiler infrastructure are currently linked into
-// libtvm_compiler. Keep their ownership explicit while inheriting the compiler
-// export policy, including downstream TVM_DLL overrides. Embedders may override
-// TVM_TIRX_DLL independently.
+// Core IR, TIRx and shared compiler infrastructure live in libtvm_tirx.
#ifndef TVM_TIRX_DLL
-#define TVM_TIRX_DLL TVM_DLL
+#if defined(__EMSCRIPTEN__)
+#define TVM_TIRX_DLL EMSCRIPTEN_KEEPALIVE
+#elif defined(_MSC_VER)
+#ifdef TVM_TIRX_EXPORTS
+#define TVM_TIRX_DLL __declspec(dllexport)
+#else
+#define TVM_TIRX_DLL __declspec(dllimport)
+#endif
+#else
+#define TVM_TIRX_DLL __attribute__((visibility("default")))
+#endif
#endif
// --- TVM_RUNTIME_DLL family (libtvm_runtime) ---
diff --git a/web/Makefile b/web/Makefile
index 7f802b5a21..12ab22f471 100644
--- a/web/Makefile
+++ b/web/Makefile
@@ -17,10 +17,10 @@
TVM_ROOT=$(realpath $(shell dirname $(firstword $(MAKEFILE_LIST))))/../
-INCLUDE_FLAGS = -I$(TVM_ROOT) -I$(TVM_ROOT)/include\
+INCLUDE_FLAGS = -I$(TVM_ROOT) -I$(TVM_ROOT)/include -I$(TVM_ROOT)/tirx/include\
-I$(TVM_ROOT)/3rdparty/tvm-ffi/include\
-I$(TVM_ROOT)/3rdparty/tvm-ffi/3rdparty/dlpack/include\
- -I$(TVM_ROOT)/3rdparty/compiler-rt
+ -I$(TVM_ROOT)/tirx/3rdparty/compiler-rt
.PHONY: clean all rmtypedep preparetest
diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc
index 07a52235ef..02948f145e 100644
--- a/web/emcc/wasm_runtime.cc
+++ b/web/emcc/wasm_runtime.cc
@@ -35,21 +35,21 @@
#include <cstring>
#include <limits>
-#include "src/runtime/cpu_device_api.cc"
-#include "src/runtime/device_api.cc"
+#include "tirx/src/runtime/cpu_device_api.cc"
+#include "tirx/src/runtime/device_api.cc"
#include "src/runtime/extra/contrib/sort/sort.cc"
-#include "src/runtime/logging.cc"
-#include "src/runtime/metadata.cc"
+#include "tirx/src/runtime/logging.cc"
+#include "tirx/src/runtime/metadata.cc"
#include "src/runtime/rpc/rpc_channel.cc"
#include "src/runtime/rpc/rpc_endpoint.cc"
#include "src/runtime/rpc/rpc_event_impl.cc"
#include "src/runtime/rpc/rpc_local_session.cc"
#include "src/runtime/rpc/rpc_module.cc"
#include "src/runtime/rpc/rpc_session.cc"
-#include "src/runtime/tensor.cc"
-#include "src/runtime/timer.cc"
+#include "tirx/src/runtime/tensor.cc"
+#include "tirx/src/runtime/timer.cc"
#include "src/runtime/vm/params.cc"
-#include "src/runtime/workspace_pool.cc"
+#include "tirx/src/runtime/workspace_pool.cc"
// relax setup
#include "3rdparty/tvm-ffi/src/ffi/backtrace.cc"
#include "3rdparty/tvm-ffi/src/ffi/container.cc"
@@ -69,7 +69,7 @@
#include "3rdparty/tvm-ffi/src/ffi/tensor.cc"
#include "3rdparty/tvm-ffi/src/ffi/testing/testing.cc"
// Enum registration requires the FFI type-attribute registry above.
-#include "src/runtime/memory_scope.cc"
+#include "tirx/src/runtime/memory_scope.cc"
#include "src/runtime/vm/attn_backend.cc"
#include "src/runtime/vm/builtin.cc"
#include "src/runtime/vm/bytecode.cc"
diff --git a/web/emcc/webgpu_runtime.cc b/web/emcc/webgpu_runtime.cc
index 780ac8824e..fb93db920c 100644
--- a/web/emcc/webgpu_runtime.cc
+++ b/web/emcc/webgpu_runtime.cc
@@ -38,7 +38,7 @@
#include <mutex>
#include <string>
-#include "../../src/runtime/metadata.h"
+#include "../../tirx/src/runtime/metadata.h"
namespace tvm {
namespace runtime {
(tvm) 10/33: [FIX][Tests] Keep mixed printer coverage with parent dialect registration
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 6c977552d07772d3c1bbd02f1198222182fd1d43
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 01:48:50 2026 +0000
[FIX][Tests] Keep mixed printer coverage with parent dialect registration
---
.../python/relax/script/test_relax_script_printer.py | 20 +++++++++++---------
.../s_tir/script/test_s_tir_script_namespace.py | 16 ++++++++++------
.../python/script/test_script_printer_entry.py | 1 +
3 files changed, 22 insertions(+), 15 deletions(-)
diff --git a/tests/python/relax/script/test_relax_script_printer.py b/tests/python/relax/script/test_relax_script_printer.py
index d6abc532ed..b86bdfed89 100644
--- a/tests/python/relax/script/test_relax_script_printer.py
+++ b/tests/python/relax/script/test_relax_script_printer.py
@@ -22,18 +22,19 @@ from __future__ import annotations
import pytest
from tvm_ffi.access_path import AccessPath
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule, tirx
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Range
-from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax import TensorType
-from tvm.relax.distributed import DeviceMesh, DTensorType, Placement
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm.relax.distributed import DeviceMesh, DTensorType, Placement
from tvm.s_tir import script as Ts
+from tvm_tirx import IRModule, tirx
+from tvm_tirx.ir import Range
+from tvm_tirx.ir import TupleType as _TupleType
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
from tvm_tirx.script.printer.scriptable import PrinterConfig, _script
@@ -434,11 +435,12 @@ class Module:
z: R.Tensor((), dtype="int32") = T.Call(nested, [x], ty=R.Tensor((), dtype="int32"))
return z
""".replace("R.", f"{relax_prefix}.")
- .replace("relax as R", f"relax as {relax_prefix}")
+ .replace("script as R", f"script as {relax_prefix}")
.replace("# from", "# from" if comment_imports else "from"),
)
tvm_tirx.ir.assert_structural_equal(
- NestedFunction, tvm_tirx.script.from_source(script, extra_vars={"I": I, "T": T, relax_prefix: R})
+ NestedFunction,
+ tvm_tirx.script.from_source(script, extra_vars={"I": I, "T": T, relax_prefix: R}),
)
diff --git a/tests/python/s_tir/script/test_s_tir_script_namespace.py b/tests/python/s_tir/script/test_s_tir_script_namespace.py
index 4dbd103c38..0321738ee2 100644
--- a/tests/python/s_tir/script/test_s_tir_script_namespace.py
+++ b/tests/python/s_tir/script/test_s_tir_script_namespace.py
@@ -23,10 +23,11 @@ import sys
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx.script import ir as I
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
@@ -82,7 +83,8 @@ def test_mixed_module_roundtrip():
tvm_tirx.ir.assert_structural_equal(
Mixed,
tvm_tirx.script.from_source(
- script, extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script}
+ script,
+ extra_vars={"I": tvm_tirx.script.ir, "T": tvm_tirx.script.tirx, "Ts": tvm.s_tir.script},
),
)
direct_script = Mixed["direct"].script()
@@ -223,15 +225,17 @@ def test_s_tir_options_and_helpers():
"tvm_tirx.tirx.script",
"tvm_tirx.script.parser",
"tvm.s_tir.script.ir_builder.frame",
- "tvm.s_tir.script.ir_builder",
+ "tvm_tirx.script.ir_builder.s_tir.frame",
],
)
def test_import_order(first):
+ # Parent initialization registers its deep builder alias with the shared core.
+ setup = "import tvm\n" if first == "tvm_tirx.script.ir_builder.s_tir.frame" else ""
subprocess.run(
[
sys.executable,
"-c",
- f"import {first}\n"
+ setup + f"import {first}\n"
"import tvm.s_tir.script as direct\n"
"from tvm.s_tir import script as Ts\nfrom tvm_tirx.script import tirx as T\n"
"from tvm_tirx.script.parser import s_tir as parser\n"
diff --git a/tirx/tests/python/script/test_script_printer_entry.py b/tests/python/script/test_script_printer_entry.py
similarity index 98%
rename from tirx/tests/python/script/test_script_printer_entry.py
rename to tests/python/script/test_script_printer_entry.py
index 81dc89d288..6d33e91767 100644
--- a/tirx/tests/python/script/test_script_printer_entry.py
+++ b/tests/python/script/test_script_printer_entry.py
@@ -22,6 +22,7 @@ import pytest
import tvm_ffi
from tvm_ffi.access_path import AccessPath
+import tvm # noqa: F401 # Register the parent Relax printer.
import tvm_tirx
from tvm_tirx.script.printer.scriptable import PrinterConfig, _script
(tvm) 29/33: [FIX][Docs] Resolve extracted TIRx API namespaces
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 78d2f04415d988a3f82da73dc1e4ee252d033371
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 09:05:08 2026 +0000
[FIX][Docs] Resolve extracted TIRx API namespaces
Point existing API references and tutorial imports at the extracted core package while retaining parent Relax and scheduling namespaces. Include the core Python path in documentation builds and update the existing alias resolution maps.
---
docs/conf.py | 15 +++---
docs/deep_dive/relax/tutorials/relax_creation.py | 8 +--
.../relax/tutorials/relax_transformation.py | 10 ++--
.../tensor_ir/tutorials/dlight_gpu_scheduling.py | 11 ++--
.../deep_dive/tensor_ir/tutorials/meta_schedule.py | 8 +--
docs/deep_dive/tensor_ir/tutorials/tir_creation.py | 19 +++----
.../tensor_ir/tutorials/tir_transformation.py | 19 +++----
docs/get_started/tutorials/ir_module.py | 17 +++---
docs/get_started/tutorials/quick_start.py | 9 ++--
docs/how_to/tutorials/bring_your_own_codegen.py | 55 ++++++++++----------
docs/how_to/tutorials/cross_compilation_and_rpc.py | 15 +++---
docs/how_to/tutorials/customize_opt.py | 26 ++++++----
docs/how_to/tutorials/e2e_opt_model.py | 9 ++--
.../how_to/tutorials/export_and_load_executable.py | 11 ++--
docs/how_to/tutorials/import_model.py | 7 +--
.../tutorials/mix_python_and_tvm_with_pymodule.py | 22 ++++----
docs/how_to/tutorials/optimize_llm.py | 31 ++++++-----
docs/reference/api/python/codegen.rst | 6 +--
docs/reference/api/python/error.rst | 6 +--
docs/reference/api/python/instrument.rst | 6 +--
docs/reference/api/python/ir.rst | 6 +--
docs/reference/api/python/runtime/runtime.rst | 6 +--
docs/reference/api/python/script/ir_builder.rst | 10 ++--
docs/reference/api/python/script/parser.rst | 44 ++++++++--------
docs/reference/api/python/script/printer.rst | 6 +--
docs/reference/api/python/script/script.rst | 40 +++++++--------
docs/reference/api/python/support.rst | 60 +++++++++++-----------
docs/reference/api/python/sym.rst | 6 +--
docs/reference/api/python/target.rst | 12 ++---
docs/reference/api/python/testing.rst | 6 +--
docs/reference/api/python/transform.rst | 6 +--
docs/tirx/api/analysis.rst | 2 +-
docs/tirx/api/backend.rst | 16 +++---
docs/tirx/api/compilation.rst | 12 ++---
docs/tirx/api/cuda.rst | 28 +++++-----
docs/tirx/api/cuda_compile.rst | 10 ++--
docs/tirx/api/execution.rst | 12 ++---
docs/tirx/api/layout.rst | 22 ++++----
docs/tirx/api/ptx.rst | 8 +--
docs/tirx/api/script.rst | 26 +++++-----
docs/tirx/api/tirx.rst | 2 +-
docs/tirx/api/trainium.rst | 12 ++---
docs/tirx/api/transform.rst | 2 +-
43 files changed, 344 insertions(+), 320 deletions(-)
diff --git a/docs/conf.py b/docs/conf.py
index 599e03ca38..361356a95e 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -51,6 +51,7 @@ else:
tvm_path = Path(os.pardir)
sys.path.insert(0, str(tvm_path.resolve() / "python"))
+sys.path.insert(0, str(tvm_path.resolve() / "tirx" / "python"))
sys.path.insert(0, str(tvm_path.resolve() / "docs"))
# -- General configuration ------------------------------------------------
@@ -315,7 +316,7 @@ language = None
# directories to ignore when looking for source files.
exclude_patterns = ["_build", "_staging"]
-# The TIRx API pages autodoc modules (tvm.tirx, tvm.backend.cuda) that re-export
+# The TIRx API pages autodoc modules (tvm_tirx.tirx, tvm_tirx.backend.cuda) that re-export
# common IR types (PrimExpr, Op, ...) from several modules, which makes a handful
# of autodoc'd cross references ambiguous. Silence that specific, benign category.
suppress_warnings = ["ref.python"]
@@ -549,8 +550,8 @@ autodoc_default_options = {
# Maps the original namespace to list of potential modules
# that we can import alias from.
tvm_alias_check_map = {
- "tvm.te": ["tvm.tirx"],
- "tvm.tirx": ["tvm.ir", "tvm.runtime"],
+ "tvm.te": ["tvm_tirx.tirx"],
+ "tvm_tirx.tirx": ["tvm_tirx.ir", "tvm_tirx.runtime"],
}
@@ -594,15 +595,15 @@ def update_alias_docstring(name, obj, lines):
tvm_class_name_rewrite_map = {
- "tvm.tirx": ["Var", "Call"],
+ "tvm_tirx.tirx": ["Var", "Call"],
"tvm.relax": ["Var", "Call", "StringImm"],
"tvm.relax.frontend.nn": ["Module"],
}
# When documenting modules under these prefixes, prefer types from the mapped module
-# to resolve ambiguous cross-references (e.g. Var exists in both tvm.tirx and tvm.relax).
+# to resolve ambiguous cross-references (e.g. Var exists in both tvm_tirx.tirx and tvm.relax).
tvm_module_type_preference = {
- "tvm.s_tir": "tvm.tirx",
+ "tvm.s_tir": "tvm_tirx.tirx",
}
@@ -715,7 +716,7 @@ def _patch_python_domain_find_obj():
return best_matches
# Check module type preference for cross-module resolution
- # (e.g. tvm.s_tir.analysis uses types from tvm.tirx).
+ # (e.g. tvm.s_tir.analysis uses types from tvm_tirx.tirx).
for prefix, preferred_mod in tvm_module_type_preference.items():
if modname.startswith(prefix):
preferred = [m for m in matches if m[0].startswith(preferred_mod + ".")]
diff --git a/docs/deep_dive/relax/tutorials/relax_creation.py b/docs/deep_dive/relax/tutorials/relax_creation.py
index 3d884e7942..166026596b 100644
--- a/docs/deep_dive/relax/tutorials/relax_creation.py
+++ b/docs/deep_dive/relax/tutorials/relax_creation.py
@@ -38,10 +38,10 @@ and relax NNModule API.
from tvm import relax, topi
from tvm.relax import ExternFunc as _ExternFunc
-from tvm.script import ir as I
-from tvm.script import relax as R
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
+from tvm.relax import script as R
+from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
+from tvm_tirx.script import tirx as T
n = T.dynamic("n", "int64")
diff --git a/docs/deep_dive/relax/tutorials/relax_transformation.py b/docs/deep_dive/relax/tutorials/relax_transformation.py
index 8abaa2d75e..49ff316d19 100644
--- a/docs/deep_dive/relax/tutorials/relax_transformation.py
+++ b/docs/deep_dive/relax/tutorials/relax_transformation.py
@@ -31,8 +31,10 @@ for optimizing and integrating with hardware backends.
# the :ref:`previous section <relax-creation>`.
import tvm
-from tvm import IRModule, relax
+import tvm_tirx
+from tvm import relax
from tvm.relax.frontend import nn
+from tvm_tirx import IRModule
class NNModule(nn.Module):
@@ -73,7 +75,7 @@ mod.show()
# in ML compilers. Note that in relax, fusion optimizations are done with the collaboration of
# a set of passes. We can apply them in a sequence.
-mod = tvm.ir.transform.Sequential(
+mod = tvm_tirx.ir.transform.Sequential(
[
tvm.relax.transform.AnnotateTIROpPattern(),
tvm.relax.transform.FuseOps(),
@@ -115,9 +117,9 @@ class ReluRewriter(PyExprMutator):
# Then we can write a pass to apply the mutator to the whole module.
-@tvm.ir.transform.module_pass(opt_level=0, name="ReluToGelu")
+@tvm_tirx.ir.transform.module_pass(opt_level=0, name="ReluToGelu")
class ReluToGelu: # pylint: disable=too-few-public-methods
- def transform_module(self, mod: IRModule, _ctx: tvm.ir.transform.PassContext) -> IRModule:
+ def transform_module(self, mod: IRModule, _ctx: tvm_tirx.ir.transform.PassContext) -> IRModule:
"""IRModule-level transformation"""
rewriter = ReluRewriter(mod)
for g_var, func in mod.functions_items():
diff --git a/docs/deep_dive/tensor_ir/tutorials/dlight_gpu_scheduling.py b/docs/deep_dive/tensor_ir/tutorials/dlight_gpu_scheduling.py
index 37c6f2d61c..1a2c2ea7cf 100644
--- a/docs/deep_dive/tensor_ir/tutorials/dlight_gpu_scheduling.py
+++ b/docs/deep_dive/tensor_ir/tutorials/dlight_gpu_scheduling.py
@@ -46,10 +46,11 @@ quality, and how to write custom rules.
# rules: ``Linear`` layers produce GEMM (matrix multiplication) kernels, ``LayerNorm``
# produces a general-reduction kernel, and ``ReLU`` is a simple elementwise op.
-import tvm
-from tvm import relax, tirx
+import tvm_tirx
+from tvm import relax
from tvm.relax.frontend import nn
from tvm.s_tir import dlight as dl
+from tvm_tirx import tirx
class DemoModel(nn.Module):
@@ -70,8 +71,8 @@ mod, params = DemoModel().export_tvm({"forward": {"x": nn.spec.Tensor((1, 768),
######################################################################
# Legalize Relax operators into TIR functions so that DLight has concrete kernels to schedule.
-device = tvm.cuda(0)
-target = tvm.target.Target.from_device(device)
+device = tvm_tirx.cuda(0)
+target = tvm_tirx.target.Target.from_device(device)
with target:
mod = relax.get_pipeline("zero")(mod)
@@ -263,7 +264,7 @@ from tvm.s_tir.dlight.base.schedule_rule import ScheduleRule
@ScheduleRule.from_callable("MyTileAndBind")
-def my_tile_and_bind(func: tirx.Function, target: tvm.target.Target, tunable: bool):
+def my_tile_and_bind(func: tirx.Function, target: tvm_tirx.target.Target, tunable: bool):
"""A minimal rule: for single-block injective functions, tile and bind to GPU threads."""
if not isinstance(func, tirx.Function):
return None
diff --git a/docs/deep_dive/tensor_ir/tutorials/meta_schedule.py b/docs/deep_dive/tensor_ir/tutorials/meta_schedule.py
index 584a05bcf2..030e96cbb3 100644
--- a/docs/deep_dive/tensor_ir/tutorials/meta_schedule.py
+++ b/docs/deep_dive/tensor_ir/tutorials/meta_schedule.py
@@ -78,7 +78,7 @@ the space of possible schedules.
import os
import tempfile
-import tvm
+import tvm_tirx
from tvm import relax
from tvm.relax.frontend import nn
@@ -100,8 +100,8 @@ class DemoModel(nn.Module):
input_shape = (1, 784)
mod, params = DemoModel().export_tvm({"forward": {"x": nn.spec.Tensor(input_shape, "float32")}})
-device = tvm.cuda(0)
-target = tvm.target.Target.from_device(device)
+device = tvm_tirx.cuda(0)
+target = tvm_tirx.target.Target.from_device(device)
######################################################################
# User-Facing Entry Points
@@ -129,7 +129,7 @@ target = tvm.target.Target.from_device(device)
if os.getenv("CI", "") != "true":
with target, tempfile.TemporaryDirectory() as tmp_dir:
- tuned_mod = tvm.ir.transform.Sequential(
+ tuned_mod = tvm_tirx.ir.transform.Sequential(
[
relax.get_pipeline("zero"),
relax.transform.MetaScheduleTuneTIR(
diff --git a/docs/deep_dive/tensor_ir/tutorials/tir_creation.py b/docs/deep_dive/tensor_ir/tutorials/tir_creation.py
index a96f0cb45f..15ab2f0a43 100644
--- a/docs/deep_dive/tensor_ir/tutorials/tir_creation.py
+++ b/docs/deep_dive/tensor_ir/tutorials/tir_creation.py
@@ -56,9 +56,10 @@ import numpy as np
import tvm_ffi
import tvm
-from tvm.script import ir as I
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
+import tvm_tirx
+from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
+from tvm_tirx.script import tirx as T
@I.ir_module
@@ -205,10 +206,10 @@ class DynamicShapeModule:
# Now let's check the runtime dynamic shape inference:
-def evaluate_dynamic_shape(lib: tvm.runtime.Module, m: int, n: int, k: int):
- A = tvm.runtime.tensor(np.random.uniform(size=(m, k)).astype("float32"))
- B = tvm.runtime.tensor(np.random.uniform(size=(k, n)).astype("float32"))
- C = tvm.runtime.tensor(np.zeros((m, n), dtype="float32"))
+def evaluate_dynamic_shape(lib: tvm_tirx.runtime.Module, m: int, n: int, k: int):
+ A = tvm_tirx.runtime.tensor(np.random.uniform(size=(m, k)).astype("float32"))
+ B = tvm_tirx.runtime.tensor(np.random.uniform(size=(k, n)).astype("float32"))
+ C = tvm_tirx.runtime.tensor(np.zeros((m, n), dtype="float32"))
lib(A, B, C)
return C.numpy()
@@ -265,7 +266,7 @@ C = te.compute((128, 128), lambda i, j: te.max(Y[i, j], 0), name="C")
# and one output parameter **C**.
te_func = te.create_function([A, B, C]).with_attr({"global_symbol": "mm_relu"})
-TEModule = tvm.IRModule({"mm_relu": te_func})
+TEModule = tvm_tirx.IRModule({"mm_relu": te_func})
TEModule.show()
######################################################################
@@ -283,5 +284,5 @@ Y = te.compute((M, N), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="Y")
C = te.compute((M, N), lambda i, j: te.max(Y[i, j], 0), name="C")
dyn_te_func = te.create_function([A, B, C]).with_attr({"global_symbol": "mm_relu"})
-DynamicTEModule = tvm.IRModule({"mm_relu": dyn_te_func})
+DynamicTEModule = tvm_tirx.IRModule({"mm_relu": dyn_te_func})
DynamicTEModule.show()
diff --git a/docs/deep_dive/tensor_ir/tutorials/tir_transformation.py b/docs/deep_dive/tensor_ir/tutorials/tir_transformation.py
index 6f53a1029c..7707419348 100644
--- a/docs/deep_dive/tensor_ir/tutorials/tir_transformation.py
+++ b/docs/deep_dive/tensor_ir/tutorials/tir_transformation.py
@@ -37,9 +37,10 @@ transformations of primitive tensor functions.
# First, let's take a look at the implementation of ``mm_relu`` in the previous section:
import tvm
-from tvm.script import ir as I
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
+import tvm_tirx
+from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
+from tvm_tirx.script import tirx as T
@I.ir_module
@@ -77,18 +78,18 @@ a_np = np.random.uniform(size=(128, 128)).astype("float32")
b_np = np.random.uniform(size=(128, 128)).astype("float32")
c_np = a_np @ b_np
-a_nd = tvm.runtime.tensor(a_np)
-b_nd = tvm.runtime.tensor(b_np)
-c_nd = tvm.runtime.tensor(np.zeros((128, 128), dtype="float32"))
+a_nd = tvm_tirx.runtime.tensor(a_np)
+b_nd = tvm_tirx.runtime.tensor(b_np)
+c_nd = tvm_tirx.runtime.tensor(np.zeros((128, 128), dtype="float32"))
-def evaluate(mod: tvm.IRModule):
- lib = tvm.tirx.compile(mod, target="llvm").mod
+def evaluate(mod: tvm_tirx.IRModule):
+ lib = tvm_tirx.tirx.compile(mod, target="llvm").mod
# check correctness
lib(a_nd, b_nd, c_nd)
np.testing.assert_allclose(c_nd.numpy(), c_np, rtol=1e-5)
# evaluate performance
- f_timer = lib.time_evaluator("main", tvm.cpu())
+ f_timer = lib.time_evaluator("main", tvm_tirx.cpu())
print(f_timer(a_nd, b_nd, c_nd))
diff --git a/docs/get_started/tutorials/ir_module.py b/docs/get_started/tutorials/ir_module.py
index beb5e7ed76..4219bb2ebb 100644
--- a/docs/get_started/tutorials/ir_module.py
+++ b/docs/get_started/tutorials/ir_module.py
@@ -42,6 +42,7 @@ from torch import nn
from torch.export import export
import tvm
+import tvm_tirx
from tvm import relax
from tvm.relax.frontend.torch import from_exported_program
@@ -118,8 +119,8 @@ mod_from_relax.show()
# directly output the IRModule in the TVMScript syntax, or alternatively,
# parse the TVMScript to obtain an IRModule.
-from tvm.script import ir as I
-from tvm.script import relax as R
+from tvm.relax import script as R
+from tvm_tirx.script import ir as I
@I.ir_module
@@ -234,11 +235,11 @@ mod.show()
# We can deploy the IRModule on CPU by specifying the target as ``llvm``.
exec = tvm.compile(mod, target="llvm")
-dev = tvm.cpu()
+dev = tvm_tirx.cpu()
vm = relax.VirtualMachine(exec, dev)
raw_data = np.random.rand(1, 784).astype("float32")
-data = tvm.runtime.tensor(raw_data, dev)
+data = tvm_tirx.runtime.tensor(raw_data, dev)
cpu_out = vm["main"](data, *params_from_torch["main"]).numpy()
print(cpu_out)
@@ -255,7 +256,7 @@ print(cpu_out)
from tvm.s_tir import dlight as dl
-with tvm.target.Target("cuda"):
+with tvm_tirx.target.Target("cuda"):
gpu_mod = dl.ApplyDefaultSchedule(
dl.gpu.Matmul(),
dl.gpu.Fallback(),
@@ -265,11 +266,11 @@ with tvm.target.Target("cuda"):
# Now we can compile the IRModule on GPU, the similar way as we did on CPU.
exec = tvm.compile(gpu_mod, target="cuda")
-dev = tvm.cuda(0)
+dev = tvm_tirx.cuda(0)
vm = relax.VirtualMachine(exec, dev)
# Need to allocate data and params on GPU device
-data = tvm.runtime.tensor(raw_data, dev)
-gpu_params = [tvm.runtime.tensor(p, dev) for p in params_from_torch["main"]]
+data = tvm_tirx.runtime.tensor(raw_data, dev)
+gpu_params = [tvm_tirx.runtime.tensor(p, dev) for p in params_from_torch["main"]]
gpu_out = vm["main"](data, *gpu_params).numpy()
print(gpu_out)
diff --git a/docs/get_started/tutorials/quick_start.py b/docs/get_started/tutorials/quick_start.py
index a1de5686a9..5f6ecd2919 100644
--- a/docs/get_started/tutorials/quick_start.py
+++ b/docs/get_started/tutorials/quick_start.py
@@ -80,6 +80,7 @@ to show how to use Apache TVM to compile a simple neural network.
#
import tvm
+import tvm_tirx
from tvm import relax
from tvm.relax.frontend import nn
@@ -137,14 +138,14 @@ mod = relax.get_pipeline("zero")(mod)
import numpy as np
-target = tvm.target.Target("llvm")
+target = tvm_tirx.target.Target("llvm")
ex = tvm.compile(mod, target)
-device = tvm.cpu()
+device = tvm_tirx.cpu()
vm = relax.VirtualMachine(ex, device)
data = np.random.rand(1, 784).astype("float32")
-tvm_data = tvm.runtime.tensor(data, device=device)
+tvm_data = tvm_tirx.runtime.tensor(data, device=device)
params = [np.random.rand(*param.shape).astype("float32") for _, param in param_spec]
-params = [tvm.runtime.tensor(param, device=device) for param in params]
+params = [tvm_tirx.runtime.tensor(param, device=device) for param in params]
print(vm["forward"](tvm_data, *params).numpy())
################################################################################
diff --git a/docs/how_to/tutorials/bring_your_own_codegen.py b/docs/how_to/tutorials/bring_your_own_codegen.py
index 96d94b3806..de1cebdbba 100644
--- a/docs/how_to/tutorials/bring_your_own_codegen.py
+++ b/docs/how_to/tutorials/bring_your_own_codegen.py
@@ -80,24 +80,24 @@ import tempfile
import numpy as np
-import tvm
-import tvm.relax.backend.contrib.example_npu
+import tvm.relax.backend.contrib.example_npu # noqa: F401
+import tvm_tirx
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.backend.contrib.tensorrt import partition_for_tensorrt
from tvm.relax.backend.pattern_registry import get_patterns_with_prefix
from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions, RunCodegen
-from tvm.script import relax as R
-has_example_npu_codegen = tvm.get_global_func("relax.ext.example_npu", True)
-has_example_npu_runtime = tvm.get_global_func("runtime.ExampleNPUJSONRuntimeCreate", True)
+has_example_npu_codegen = tvm_tirx.get_global_func("relax.ext.example_npu", True)
+has_example_npu_runtime = tvm_tirx.get_global_func("runtime.ExampleNPUJSONRuntimeCreate", True)
has_example_npu = has_example_npu_codegen and has_example_npu_runtime
-has_tensorrt_codegen = tvm.get_global_func("relax.ext.tensorrt", True) is not None
-_is_trt_runtime_enabled = tvm.get_global_func("relax.is_tensorrt_runtime_enabled", True)
+has_tensorrt_codegen = tvm_tirx.get_global_func("relax.ext.tensorrt", True) is not None
+_is_trt_runtime_enabled = tvm_tirx.get_global_func("relax.is_tensorrt_runtime_enabled", True)
has_tensorrt = (
has_tensorrt_codegen and _is_trt_runtime_enabled is not None and _is_trt_runtime_enabled()
)
-has_cuda = tvm.cuda(0).exist
+has_cuda = tvm_tirx.cuda(0).exist
######################################################################
# Step 2: Define the model
@@ -107,7 +107,7 @@ has_cuda = tvm.cuda(0).exist
# backends.
-@tvm.script.ir_module
+@tvm_tirx.script.ir_module
class ConvReLU:
@R.function
def main(
@@ -166,12 +166,13 @@ weight_np = np.random.randn(16, 3, 3, 3).astype("float32")
if has_example_npu:
npu_mod = RunCodegen()(npu_mod)
- with tvm.ir.transform.PassContext(opt_level=3):
- npu_exec = relax.build(npu_mod, tvm.target.Target("llvm"))
+ with tvm_tirx.ir.transform.PassContext(opt_level=3):
+ npu_exec = relax.build(npu_mod, tvm_tirx.target.Target("llvm"))
- npu_vm = relax.VirtualMachine(npu_exec, tvm.cpu())
+ npu_vm = relax.VirtualMachine(npu_exec, tvm_tirx.cpu())
npu_out = npu_vm["main"](
- tvm.runtime.tensor(data_np, tvm.cpu()), tvm.runtime.tensor(weight_np, tvm.cpu())
+ tvm_tirx.runtime.tensor(data_np, tvm_tirx.cpu()),
+ tvm_tirx.runtime.tensor(weight_np, tvm_tirx.cpu()),
)
assert npu_out.numpy().shape == (1, 16, 30, 30)
print("Example NPU run completed. Output shape:", npu_out.numpy().shape)
@@ -206,17 +207,19 @@ print(trt_mod)
# Build for CUDA, run on the GPU, and compare against the CPU reference.
if has_tensorrt and has_cuda:
- dev = tvm.cuda(0)
- with tvm.ir.transform.PassContext(opt_level=3):
+ dev = tvm_tirx.cuda(0)
+ with tvm_tirx.ir.transform.PassContext(opt_level=3):
trt_exec = relax.build(RunCodegen()(trt_mod), "cuda")
- trt_out = relax.VirtualMachine(trt_exec, dev)["main"](tvm.runtime.tensor(data_np, dev)).numpy()
+ trt_out = relax.VirtualMachine(trt_exec, dev)["main"](
+ tvm_tirx.runtime.tensor(data_np, dev)
+ ).numpy()
cpu_mod = relax.transform.LegalizeOps()(
relax.transform.BindParams("main", {"weight": weight_np})(ConvReLU)
)
cpu_exec = relax.build(cpu_mod, "llvm")
- cpu_out = relax.VirtualMachine(cpu_exec, tvm.cpu())["main"](
- tvm.runtime.tensor(data_np, tvm.cpu())
+ cpu_out = relax.VirtualMachine(cpu_exec, tvm_tirx.cpu())["main"](
+ tvm_tirx.runtime.tensor(data_np, tvm_tirx.cpu())
).numpy()
np.testing.assert_allclose(trt_out, cpu_out, rtol=1e-2, atol=1e-2)
@@ -235,12 +238,12 @@ if has_tensorrt and has_cuda:
fp16_mod = partition_for_tensorrt(
relax.transform.BindParams("main", {"weight": weight_np})(ConvReLU)
)
- with tvm.ir.transform.PassContext(
+ with tvm_tirx.ir.transform.PassContext(
opt_level=3, config={"relax.ext.tensorrt.options": {"use_fp16": True}}
):
fp16_exec = relax.build(RunCodegen()(fp16_mod), "cuda")
- fp16_out = relax.VirtualMachine(fp16_exec, tvm.cuda(0))["main"](
- tvm.runtime.tensor(data_np, tvm.cuda(0))
+ fp16_out = relax.VirtualMachine(fp16_exec, tvm_tirx.cuda(0))["main"](
+ tvm_tirx.runtime.tensor(data_np, tvm_tirx.cuda(0))
).numpy()
np.testing.assert_allclose(fp16_out, cpu_out, rtol=5e-2, atol=5e-2)
@@ -312,11 +315,11 @@ if has_torch and has_tensorrt and has_cuda:
print("After importing and partitioning the PyTorch model:")
print(torch_mod)
- torch_dev = tvm.cuda(0)
- with tvm.ir.transform.PassContext(opt_level=3):
+ torch_dev = tvm_tirx.cuda(0)
+ with tvm_tirx.ir.transform.PassContext(opt_level=3):
torch_exec = relax.build(RunCodegen()(torch_mod), "cuda")
deployed = relax.VirtualMachine(torch_exec, torch_dev)["main"](
- tvm.runtime.tensor(example_input.numpy(), torch_dev)
+ tvm_tirx.runtime.tensor(example_input.numpy(), torch_dev)
)[0].numpy()
np.testing.assert_allclose(deployed, torch_ref, rtol=1e-2, atol=1e-2)
@@ -331,9 +334,9 @@ if has_torch and has_tensorrt and has_cuda:
with tempfile.TemporaryDirectory() as tmpdir:
lib_path = os.path.join(tmpdir, "deployed_trt.so")
torch_exec.export_library(lib_path)
- loaded = tvm.runtime.load_module(lib_path)
+ loaded = tvm_tirx.runtime.load_module(lib_path)
reran = relax.VirtualMachine(loaded, torch_dev)["main"](
- tvm.runtime.tensor(example_input.numpy(), torch_dev)
+ tvm_tirx.runtime.tensor(example_input.numpy(), torch_dev)
)[0].numpy()
np.testing.assert_allclose(reran, torch_ref, rtol=1e-2, atol=1e-2)
print("Reloaded the exported library and reran; output", reran.shape, "still matches.")
diff --git a/docs/how_to/tutorials/cross_compilation_and_rpc.py b/docs/how_to/tutorials/cross_compilation_and_rpc.py
index 9fc425c7a9..f218582b82 100644
--- a/docs/how_to/tutorials/cross_compilation_and_rpc.py
+++ b/docs/how_to/tutorials/cross_compilation_and_rpc.py
@@ -98,13 +98,14 @@ import numpy as np
import tvm_ffi
import tvm
+import tvm_tirx
from tvm import rpc, te
-from tvm.support import utils
+from tvm_tirx.support import utils
-n = tvm.runtime.convert(1024)
+n = tvm_tirx.runtime.convert(1024)
A = te.placeholder((n,), name="A")
B = te.compute((n,), lambda i: A[i] + 1.0, name="B")
-mod = tvm.IRModule.from_expr(te.create_function([A, B]).with_attr("global_symbol", "add_one"))
+mod = tvm_tirx.IRModule.from_expr(te.create_function([A, B]).with_attr("global_symbol", "add_one"))
######################################################################
# Then we cross compile the kernel.
@@ -416,10 +417,10 @@ def run_opencl():
# them according to your environment.
opencl_device_host = "10.77.1.145"
opencl_device_port = 9090
- target = tvm.target.Target("opencl", host={"kind": "llvm", "mtriple": "aarch64-linux-gnu"})
+ target = tvm_tirx.target.Target("opencl", host={"kind": "llvm", "mtriple": "aarch64-linux-gnu"})
# create schedule for the above "add one" compute declaration
- mod = tvm.IRModule.from_expr(te.create_function([A, B]))
+ mod = tvm_tirx.IRModule.from_expr(te.create_function([A, B]))
sch = tvm.s_tir.Schedule(mod)
(x,) = sch.get_loops(block=sch.get_sblock("B"))
xo, xi = sch.split(x, [None, 32])
@@ -538,7 +539,7 @@ def run_pytorch_model_via_rpc():
if local_demo:
# For demonstration on local machine, use local target
- target = tvm.target.Target("llvm")
+ target = tvm_tirx.target.Target("llvm")
print("Using local target for demonstration")
else:
# Choose the appropriate target for your device:
@@ -560,7 +561,7 @@ def run_pytorch_model_via_rpc():
# - OpenCL: tvm.target.Target("opencl", host={"kind": "llvm", "mtriple": "aarch64-linux-gnu"})
#
# For this example, we use ARM 64-bit
- target = tvm.target.Target({"kind": "llvm", "mtriple": "aarch64-linux-gnu"})
+ target = tvm_tirx.target.Target({"kind": "llvm", "mtriple": "aarch64-linux-gnu"})
print(f"Cross-compiling for target: {target}")
# Apply optimization pipeline
diff --git a/docs/how_to/tutorials/customize_opt.py b/docs/how_to/tutorials/customize_opt.py
index e6e83c33c4..fc52128e6a 100644
--- a/docs/how_to/tutorials/customize_opt.py
+++ b/docs/how_to/tutorials/customize_opt.py
@@ -55,8 +55,10 @@ import tempfile
import numpy as np
import tvm
-from tvm import IRModule, relax
+import tvm_tirx
+from tvm import relax
from tvm.relax.frontend import nn
+from tvm_tirx import IRModule
######################################################################
# Composable IRModule Optimization
@@ -119,11 +121,11 @@ except ImportError as e:
# Define a new pass for CUBLAS dispatch
-@tvm.ir.transform.module_pass(opt_level=0, name="CublasDispatch")
+@tvm_tirx.ir.transform.module_pass(opt_level=0, name="CublasDispatch")
class CublasDispatch:
- def transform_module(self, mod: IRModule, _ctx: tvm.ir.transform.PassContext) -> IRModule:
+ def transform_module(self, mod: IRModule, _ctx: tvm_tirx.ir.transform.PassContext) -> IRModule:
# Check if CUBLAS is enabled
- if not tvm.get_global_func("relax.ext.cublas", True):
+ if not tvm_tirx.get_global_func("relax.ext.cublas", True):
raise Exception("CUBLAS is not enabled.")
# Get interested patterns
@@ -163,12 +165,12 @@ mod.show()
# To save CI time and avoid flakiness, we skip the tuning process in CI environment.
#
-device = tvm.cuda(0)
-target = tvm.target.Target.from_device(device)
+device = tvm_tirx.cuda(0)
+target = tvm_tirx.target.Target.from_device(device)
if os.getenv("CI", "") != "true":
trials = 2000
with target, tempfile.TemporaryDirectory() as tmp_dir:
- mod = tvm.ir.transform.Sequential(
+ mod = tvm_tirx.ir.transform.Sequential(
[
relax.get_pipeline("zero"),
relax.transform.MetaScheduleTuneTIR(work_dir=tmp_dir, max_trials_global=trials),
@@ -190,7 +192,7 @@ from tvm.s_tir import dlight as dl
# Apply DLight rules
with target:
- mod = tvm.ir.transform.Sequential(
+ mod = tvm_tirx.ir.transform.Sequential(
[
relax.get_pipeline("zero"),
dl.ApplyDefaultSchedule( # pylint: disable=not-callable
@@ -218,11 +220,13 @@ mod.show()
# We can build and deploy the optimized model to the TVM runtime.
ex = tvm.compile(mod, target="cuda")
-dev = tvm.cuda(0)
+dev = tvm_tirx.cuda(0)
vm = relax.VirtualMachine(ex, dev)
# Need to allocate data and params on GPU device
-data = tvm.runtime.tensor(np.random.rand(*input_shape).astype("float32"), dev)
-gpu_params = [tvm.runtime.tensor(np.random.rand(*p.shape).astype(p.dtype), dev) for _, p in params]
+data = tvm_tirx.runtime.tensor(np.random.rand(*input_shape).astype("float32"), dev)
+gpu_params = [
+ tvm_tirx.runtime.tensor(np.random.rand(*p.shape).astype(p.dtype), dev) for _, p in params
+]
gpu_out = vm["forward"](data, *gpu_params).numpy()
print(gpu_out)
diff --git a/docs/how_to/tutorials/e2e_opt_model.py b/docs/how_to/tutorials/e2e_opt_model.py
index 5de3d50485..8c35d1882e 100644
--- a/docs/how_to/tutorials/e2e_opt_model.py
+++ b/docs/how_to/tutorials/e2e_opt_model.py
@@ -68,6 +68,7 @@ torch_model = resnet18(weights=ResNet18_Weights.DEFAULT).eval()
# optimization.
import tvm
+import tvm_tirx
from tvm import relax
from tvm.relax.frontend.torch import from_exported_program
@@ -118,7 +119,7 @@ if not IS_IN_CI:
TOTAL_TRIALS = 512 # Change to 20000 for better performance if needed
MAX_TRIALS_PER_TASK = 16 # Change to more trials per task for better performance if needed
-target = tvm.target.Target("nvidia/geforce-rtx-3090-ti") # Change to your target device
+target = tvm_tirx.target.Target("nvidia/geforce-rtx-3090-ti") # Change to your target device
work_dir = "tuning_logs"
if not IS_IN_CI:
@@ -143,11 +144,11 @@ if not IS_IN_CI:
with target:
mod = tvm.s_tir.transform.DefaultGPUSchedule()(mod)
ex = tvm.compile(mod, target=target)
- dev = tvm.cuda(0)
+ dev = tvm_tirx.cuda(0)
vm = relax.VirtualMachine(ex, dev)
# Need to allocate data and params on GPU device
- gpu_data = tvm.runtime.tensor(np.random.rand(1, 3, 224, 224).astype("float32"), dev)
- gpu_params = [tvm.runtime.tensor(p, dev) for p in params["main"]]
+ gpu_data = tvm_tirx.runtime.tensor(np.random.rand(1, 3, 224, 224).astype("float32"), dev)
+ gpu_params = [tvm_tirx.runtime.tensor(p, dev) for p in params["main"]]
gpu_out = vm["main"](gpu_data, *gpu_params)[0].numpy()
print(gpu_out.shape)
diff --git a/docs/how_to/tutorials/export_and_load_executable.py b/docs/how_to/tutorials/export_and_load_executable.py
index 9972861c12..542266803a 100644
--- a/docs/how_to/tutorials/export_and_load_executable.py
+++ b/docs/how_to/tutorials/export_and_load_executable.py
@@ -65,6 +65,7 @@ except ImportError: # pragma: no cover
import tvm_ffi
import tvm
+import tvm_tirx
from tvm import relax
from tvm.relax.frontend.torch import from_exported_program
@@ -115,7 +116,7 @@ if RUN_EXAMPLE:
# executable. Passing ``workspace_dir`` keeps the intermediate packaging files,
# which is useful to inspect what was produced.
-TARGET = tvm.target.Target("llvm")
+TARGET = tvm_tirx.target.Target("llvm")
ARTIFACT_DIR = Path("relax_export_artifacts")
ARTIFACT_DIR.mkdir(exist_ok=True)
@@ -159,16 +160,16 @@ if RUN_EXAMPLE:
# runtime module directly.
if RUN_EXAMPLE:
- loaded_rt_mod = tvm.runtime.load_module(str(library_path))
- dev = tvm.cpu(0)
+ loaded_rt_mod = tvm_tirx.runtime.load_module(str(library_path))
+ dev = tvm_tirx.cpu(0)
vm = relax.VirtualMachine(loaded_rt_mod, dev)
# Prepare input data
input_tensor = torch.randn(1, 1, 28, 28, dtype=torch.float32)
- vm_input = tvm.runtime.tensor(input_tensor.numpy(), dev)
+ vm_input = tvm_tirx.runtime.tensor(input_tensor.numpy(), dev)
# Prepare parameters (allocate on target device)
- vm_params = [tvm.runtime.tensor(p, dev) for p in params["main"]]
+ vm_params = [tvm_tirx.runtime.tensor(p, dev) for p in params["main"]]
# Run inference: pass input data followed by all parameters
tvm_output = vm["main"](vm_input, *vm_params)
diff --git a/docs/how_to/tutorials/import_model.py b/docs/how_to/tutorials/import_model.py
index 888235a859..c8e31558fd 100644
--- a/docs/how_to/tutorials/import_model.py
+++ b/docs/how_to/tutorials/import_model.py
@@ -55,6 +55,7 @@ from torch import nn
from torch.export import export
import tvm
+import tvm_tirx
from tvm import relax
from tvm.relax.frontend.torch import from_exported_program
@@ -191,13 +192,13 @@ def convert_sigmoid(node: torch.fx.Node, importer: ExportedProgramImporter) -> r
mod_compiled = relax.get_pipeline("zero")(mod)
exec_module = tvm.compile(mod_compiled, target="llvm")
-dev = tvm.cpu()
+dev = tvm_tirx.cpu()
vm = relax.VirtualMachine(exec_module, dev)
# Run inference
input_data = np.random.rand(1, 3, 32, 32).astype("float32")
-tvm_input = tvm.runtime.tensor(input_data, dev)
-tvm_params = [tvm.runtime.tensor(p, dev) for p in params["main"]]
+tvm_input = tvm_tirx.runtime.tensor(input_data, dev)
+tvm_params = [tvm_tirx.runtime.tensor(p, dev) for p in params["main"]]
tvm_out = vm["main"](tvm_input, *tvm_params).numpy()
# Compare with PyTorch
diff --git a/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py b/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py
index 35e4110f9e..e195f431ce 100644
--- a/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py
+++ b/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py
@@ -56,13 +56,13 @@ try:
except ImportError:
torch = None
-import tvm
+import tvm_tirx
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.base_py_module import BasePyModule
-from tvm.script import ir as I
-from tvm.script import relax as R
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
+from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
+from tvm_tirx.script import tirx as T
IS_IN_CI = os.getenv("CI", "").lower() == "true"
HAS_TORCH = torch is not None
@@ -103,7 +103,7 @@ if RUN_EXAMPLE:
return self._convert_tvm_to_pytorch(result)
# TIR functions are JIT-compiled at instantiation
- mod = MyFirstModule(device=tvm.cpu(0))
+ mod = MyFirstModule(device=tvm_tirx.cpu(0))
x = torch.tensor([1.0, 2.0, 3.0, 4.0])
y = torch.tensor([10.0, 20.0, 30.0, 40.0])
@@ -171,7 +171,7 @@ if RUN_EXAMPLE:
print(f" [DEBUG] probs sum: {result.sum(dim=-1)}")
return result
- mod = DebugModule(device=tvm.cpu(0))
+ mod = DebugModule(device=tvm_tirx.cpu(0))
print("Running with debug prints:")
probs = mod.forward(torch.randn(2, 4), torch.randn(4, 3))
@@ -197,7 +197,7 @@ if RUN_EXAMPLE:
if RUN_EXAMPLE:
# Register a packed function (simulating an external library binding)
- @tvm.register_global_func("my_bias_add", override=True)
+ @tvm_tirx.register_global_func("my_bias_add", override=True)
def my_bias_add(x, bias, out):
"""Packed function: adds bias to each row of x."""
@@ -243,7 +243,7 @@ if RUN_EXAMPLE:
# 3. Python/PyTorch activation
return F.relu(h_biased)
- mod = PipelineModule(device=tvm.cpu(0))
+ mod = PipelineModule(device=tvm_tirx.cpu(0))
x = torch.randn(2, 4)
w = torch.randn(4, 3)
@@ -365,7 +365,7 @@ if RUN_EXAMPLE:
out = R.call_py_func("silu", (h,), ty_args=[R.Tensor((4, 8), "float32")])
return out
- mod = HybridVMModule(device=tvm.cpu(0))
+ mod = HybridVMModule(device=tvm_tirx.cpu(0))
x = torch.randn(4, 8)
# call_py_func is also callable from Python directly
@@ -412,7 +412,7 @@ if RUN_EXAMPLE:
) -> R.Tensor((n,), "float32"):
return R.add(x, y)
- mod = DynamicModule(device=tvm.cpu(0), target="llvm")
+ mod = DynamicModule(device=tvm_tirx.cpu(0), target="llvm")
# Inspect what the module contains
print("Functions:", mod.list_functions())
diff --git a/docs/how_to/tutorials/optimize_llm.py b/docs/how_to/tutorials/optimize_llm.py
index 7faa4f43bc..2c01922452 100644
--- a/docs/how_to/tutorials/optimize_llm.py
+++ b/docs/how_to/tutorials/optimize_llm.py
@@ -64,12 +64,14 @@ from pprint import pprint
from tvm_ffi import Shape
import tvm
-from tvm import relax, te, tirx
+import tvm_tirx
+from tvm import relax, te
from tvm.relax import register_pipeline
from tvm.relax.frontend import nn
from tvm.relax.frontend.nn import Tensor, op
from tvm.relax.frontend.nn.llm.kv_cache import PagedKVCache, TIRPagedKVCache
from tvm.s_tir import dlight
+from tvm_tirx import tirx
######################################################################
# First, we need to define the model configuration. The configuration includes the key parameters
@@ -92,8 +94,8 @@ class LlamaConfig:
head_dim: int = 64 # hidden_size // num_attention_heads
-dev = tvm.cuda(0)
-target = tvm.target.Target.from_device(dev)
+dev = tvm_tirx.cuda(0)
+target = tvm_tirx.target.Target.from_device(dev)
######################################################################
@@ -382,9 +384,11 @@ def _pipeline( # pylint: disable=too-many-arguments
):
ext_mods = ext_mods or []
- @tvm.ir.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext) -> tvm.ir.IRModule:
- seq = tvm.ir.transform.Sequential(
+ @tvm_tirx.ir.transform.module_pass(opt_level=0)
+ def _pipeline(
+ mod: tvm_tirx.ir.IRModule, _ctx: tvm_tirx.ir.transform.PassContext
+ ) -> tvm_tirx.ir.IRModule:
+ seq = tvm_tirx.ir.transform.Sequential(
[
# Phase 1. Passes on high-level operator graph
# We can enable cublas for further optimization
@@ -490,7 +494,8 @@ if not IS_IN_CI:
# Convert params into ndarray
params = [
- tvm.runtime.tensor(param_dict[k].astype("float16"), device=dev) for k in named_params.keys()
+ tvm_tirx.runtime.tensor(param_dict[k].astype("float16"), device=dev)
+ for k in named_params.keys()
]
@@ -524,7 +529,7 @@ if not IS_IN_CI:
input_len = len(prompt)
# Load prompt tokens into TVM ndarray on the target device
- tokens = tvm.runtime.tensor(np.array(prompt).astype("int32"), device=dev)
+ tokens = tvm_tirx.runtime.tensor(np.array(prompt).astype("int32"), device=dev)
######################################################################
# Create the KVCache
@@ -548,7 +553,7 @@ if not IS_IN_CI:
# The next step is to embed the tokens into the hidden states. We use the `embed` function
# compiled in the Relax IRModule to embed the tokens into the hidden states.
-nd_view_func = tvm.get_global_func("vm.builtin.reshape")
+nd_view_func = tvm_tirx.get_global_func("vm.builtin.reshape")
def embed(tokens, params):
@@ -563,9 +568,9 @@ def embed(tokens, params):
# ~~~~~~~
# Before running the forward pass, we first get some help functions for preparation.
-add_sequence_func = tvm.get_global_func("vm.builtin.kv_state_add_sequence")
-begin_forward_func = tvm.get_global_func("vm.builtin.kv_state_begin_forward")
-end_forward_func = tvm.get_global_func("vm.builtin.kv_state_end_forward")
+add_sequence_func = tvm_tirx.get_global_func("vm.builtin.kv_state_add_sequence")
+begin_forward_func = tvm_tirx.get_global_func("vm.builtin.kv_state_begin_forward")
+end_forward_func = tvm_tirx.get_global_func("vm.builtin.kv_state_end_forward")
######################################################################
# As we are creating a new sequence, we need to call `add_sequence_func` to initialize
@@ -610,7 +615,7 @@ if not IS_IN_CI:
print("The generated token:")
while last_token != tokenizer.eos_token_id:
- tokens = tvm.runtime.tensor(np.array([last_token]).astype("int32"), device=dev)
+ tokens = tvm_tirx.runtime.tensor(np.array([last_token]).astype("int32"), device=dev)
hidden_states = embed(tokens, params)
begin_forward_func(kv_cache, Shape([seq_id]), Shape([1]))
logits, kv_cache = vm["decode"](hidden_states, kv_cache, params)
diff --git a/docs/reference/api/python/codegen.rst b/docs/reference/api/python/codegen.rst
index 0fca6ed79f..c1ee15cc0f 100644
--- a/docs/reference/api/python/codegen.rst
+++ b/docs/reference/api/python/codegen.rst
@@ -15,8 +15,8 @@
specific language governing permissions and limitations
under the License.
-tvm.codegen
-===========
-.. automodule:: tvm.codegen
+tvm_tirx.codegen
+================
+.. automodule:: tvm_tirx.codegen
:members:
:show-inheritance:
diff --git a/docs/reference/api/python/error.rst b/docs/reference/api/python/error.rst
index a228f7b8bd..297e430093 100644
--- a/docs/reference/api/python/error.rst
+++ b/docs/reference/api/python/error.rst
@@ -15,9 +15,9 @@
specific language governing permissions and limitations
under the License.
-tvm.error
----------
-.. automodule:: tvm.error
+tvm_tirx.error
+--------------
+.. automodule:: tvm_tirx.error
:members:
:imported-members:
:autosummary:
diff --git a/docs/reference/api/python/instrument.rst b/docs/reference/api/python/instrument.rst
index 46d193ce15..10b4bee978 100644
--- a/docs/reference/api/python/instrument.rst
+++ b/docs/reference/api/python/instrument.rst
@@ -15,7 +15,7 @@
specific language governing permissions and limitations
under the License.
-tvm.ir.transform.instrument
-===========================
-.. automodule:: tvm.ir.transform.instrument
+tvm_tirx.ir.transform.instrument
+================================
+.. automodule:: tvm_tirx.ir.transform.instrument
:members:
diff --git a/docs/reference/api/python/ir.rst b/docs/reference/api/python/ir.rst
index 1747dee7be..60c4289b77 100644
--- a/docs/reference/api/python/ir.rst
+++ b/docs/reference/api/python/ir.rst
@@ -15,9 +15,9 @@
specific language governing permissions and limitations
under the License.
-tvm.ir
-------
-.. automodule:: tvm.ir
+tvm_tirx.ir
+-----------
+.. automodule:: tvm_tirx.ir
:members:
:imported-members:
:exclude-members: Expr
diff --git a/docs/reference/api/python/runtime/runtime.rst b/docs/reference/api/python/runtime/runtime.rst
index ae373080ae..04bcd64f2d 100644
--- a/docs/reference/api/python/runtime/runtime.rst
+++ b/docs/reference/api/python/runtime/runtime.rst
@@ -15,7 +15,7 @@
specific language governing permissions and limitations
under the License.
-tvm.runtime
------------
-.. automodule:: tvm.runtime
+tvm_tirx.runtime
+----------------
+.. automodule:: tvm_tirx.runtime
:members:
diff --git a/docs/reference/api/python/script/ir_builder.rst b/docs/reference/api/python/script/ir_builder.rst
index f26f70e1c1..564e9c1db6 100644
--- a/docs/reference/api/python/script/ir_builder.rst
+++ b/docs/reference/api/python/script/ir_builder.rst
@@ -15,12 +15,12 @@
specific language governing permissions and limitations
under the License.
-tvm.script.ir_builder
----------------------
+tvm_tirx.script.ir_builder
+--------------------------
-tvm.script.ir_builder
-*********************
-.. automodule:: tvm.script.ir_builder
+tvm_tirx.script.ir_builder
+**************************
+.. automodule:: tvm_tirx.script.ir_builder
:members:
:imported-members:
:exclude-members: Call, DataTypeImm, FuncType, GenericConst, MissingType, Op, PrimType, Range, StringImm, StringType, TensorRegion, TensorRegionType, Tuple, Type, Var
diff --git a/docs/reference/api/python/script/parser.rst b/docs/reference/api/python/script/parser.rst
index 31beb76afd..d4028219b8 100644
--- a/docs/reference/api/python/script/parser.rst
+++ b/docs/reference/api/python/script/parser.rst
@@ -15,44 +15,44 @@
specific language governing permissions and limitations
under the License.
-tvm.script.parser
------------------
+tvm_tirx.script.parser
+----------------------
-tvm.script.parser
-*****************
-.. automodule:: tvm.script.parser
+tvm_tirx.script.parser
+**********************
+.. automodule:: tvm_tirx.script.parser
:members:
:imported-members:
-tvm.script.parser.jit_support
-*****************************
+tvm_tirx.script.parser.jit_support
+**********************************
These helpers carry ``const_args``, the selected fixed argument map, into generated builder
execution. Parameter validation and caching remain with the JIT entry point.
-.. automodule:: tvm.script.parser.jit_support
+.. automodule:: tvm_tirx.script.parser.jit_support
:members: use_specialization, read_specialization_bindings
-tvm.script.parser.protocol_registry
-***********************************
+tvm_tirx.script.parser.protocol_registry
+****************************************
Language variants register source syntax policies independently of their
construction hooks. Namespace objects and lazy initialization callbacks are
-registered through :func:`tvm.script.parser.register_namespace` and
-:func:`tvm.script.parser.register_namespace_initializer`.
+registered through :func:`tvm_tirx.script.parser.register_namespace` and
+:func:`tvm_tirx.script.parser.register_namespace_initializer`.
-.. automodule:: tvm.script.parser.protocol_registry
+.. automodule:: tvm_tirx.script.parser.protocol_registry
:members: register_scalar_annotation, register_mutable_decl
The language variant aliases below share the public construction namespaces documented
in :doc:`script`. Parser entry points above use the canonical frontend.
-tvm.script.parser.ir
-********************
-.. automodule:: tvm.script.parser.ir
+tvm_tirx.script.parser.ir
+*************************
+.. automodule:: tvm_tirx.script.parser.ir
-tvm.script.parser.relax
-***********************
-.. automodule:: tvm.script.parser.relax
+tvm_tirx.script.parser.relax
+****************************
+.. automodule:: tvm_tirx.script.parser.relax
-tvm.script.parser.tirx
-**********************
-.. automodule:: tvm.script.parser.tirx
+tvm_tirx.script.parser.tirx
+***************************
+.. automodule:: tvm_tirx.script.parser.tirx
diff --git a/docs/reference/api/python/script/printer.rst b/docs/reference/api/python/script/printer.rst
index 098333de07..503ea719ad 100644
--- a/docs/reference/api/python/script/printer.rst
+++ b/docs/reference/api/python/script/printer.rst
@@ -15,12 +15,12 @@
specific language governing permissions and limitations
under the License.
-tvm.script.printer
+tvm_tirx.script.printer
-------------------------
-tvm.script.printer
+tvm_tirx.script.printer
*************************
-.. automodule:: tvm.script.printer
+.. automodule:: tvm_tirx.script.printer
:members:
:imported-members:
:exclude-members: Scriptable
diff --git a/docs/reference/api/python/script/script.rst b/docs/reference/api/python/script/script.rst
index dd287b4899..612b634c7a 100644
--- a/docs/reference/api/python/script/script.rst
+++ b/docs/reference/api/python/script/script.rst
@@ -15,22 +15,22 @@
specific language governing permissions and limitations
under the License.
-tvm.script
-----------
+tvm_tirx.script
+---------------
-tvm.script
-**********
-.. automodule:: tvm.script
+tvm_tirx.script
+***************
+.. automodule:: tvm_tirx.script
:members:
:imported-members:
:exclude-members: Any
-tvm.script.relax
+tvm.relax.script
****************
Operator attributes re-exported by this namespace are documented in :mod:`tvm.relax.op`.
-.. automodule:: tvm.script.relax
+.. automodule:: tvm.relax.script
:members:
:exclude-members: ExternFunc, ShapeExpr, TupleGetItem, Range, function, macro,
ArgmaxArgminAttrs, ArgsortAttrs, AstypeAttrs, CallInplacePackedAttrs, CallTIRInplaceAttrs,
@@ -42,28 +42,28 @@ Operator attributes re-exported by this namespace are documented in :mod:`tvm.re
SplitAttrs, SqueezeAttrs, StackAttrs, StatisticalAttrs, StridedSliceAttrs, TakeAttrs,
TileAttrs, TopKAttrs, ToVDeviceAttrs, TriluAttrs, WrapParamAttrs
-.. autofunction:: tvm.script.relax.function
+.. autofunction:: tvm.relax.script.function
-.. autofunction:: tvm.script.relax.macro
+.. autofunction:: tvm.relax.script.macro
-tvm.script.tirx
-***************
+tvm_tirx.script.tirx
+********************
-Operator attributes re-exported by this namespace are documented in :mod:`tvm.tirx`.
-Shared IR classes and builder infrastructure are documented in :mod:`tvm.ir` and
-:mod:`tvm.script.ir_builder`.
-The canonical ``MemoryScope`` class is documented in :mod:`tvm.runtime`.
+Operator attributes re-exported by this namespace are documented in :mod:`tvm_tirx.tirx`.
+Shared IR classes and builder infrastructure are documented in :mod:`tvm_tirx.ir` and
+:mod:`tvm_tirx.script.ir_builder`.
+The canonical ``MemoryScope`` class is documented in :mod:`tvm_tirx.runtime`.
-.. automodule:: tvm.script.tirx
+.. automodule:: tvm_tirx.script.tirx
:members:
:exclude-members: AlreadyEmitted, DataTypeImm, FuncType, GenericConst, IRBuilder,
IRModuleFrame, MemoryScope, MissingType, Op, PrimType, StringImm, StringType, TensorRegion,
TensorRegionType, Type, Range, meta_var, Var, Call, CommReducer, Reduce, FloatImm, IntImm, Cast, Add, Sub, Mul, Div, Mod, FloorDiv, FloorMod, LShift, RShift, BitwiseAnd, BitwiseOr, BitwiseXor, BitwiseNot, Min, Max, EQ, NE, LT, LE, GT, GE, And, Or, Not, Select, Ramp, Broadcast, Shuffle, CallEffectKind, IterVar, ComposeLayout, DtypeConstructor, ThreadScope, Iter, Layout, LetAnnotation, LocalVectorAnnotation, TileLayout, CallFFIKernelAttr, TensorMapEncodeTiledAttr, function, jit, inli [...]
-.. autofunction:: tvm.script.tirx.function
+.. autofunction:: tvm_tirx.script.tirx.function
-.. autofunction:: tvm.script.tirx.jit
+.. autofunction:: tvm_tirx.script.tirx.jit
-.. autofunction:: tvm.script.tirx.inline
+.. autofunction:: tvm_tirx.script.tirx.inline
-.. autofunction:: tvm.script.tirx.macro
+.. autofunction:: tvm_tirx.script.tirx.macro
diff --git a/docs/reference/api/python/support.rst b/docs/reference/api/python/support.rst
index 8042ce34c1..2c5991081c 100644
--- a/docs/reference/api/python/support.rst
+++ b/docs/reference/api/python/support.rst
@@ -15,41 +15,41 @@
specific language governing permissions and limitations
under the License.
-tvm.support
------------
-.. automodule:: tvm.support
+tvm_tirx.support
+----------------
+.. automodule:: tvm_tirx.support
:members:
:imported-members:
:autosummary:
-tvm.support.cc
-~~~~~~~~~~~~~~
-.. automodule:: tvm.support.cc
+tvm_tirx.support.cc
+~~~~~~~~~~~~~~~~~~~
+.. automodule:: tvm_tirx.support.cc
:members:
-tvm.support.nvcc
-~~~~~~~~~~~~~~~~
-.. automodule:: tvm.support.nvcc
+tvm_tirx.support.nvcc
+~~~~~~~~~~~~~~~~~~~~~
+.. automodule:: tvm_tirx.support.nvcc
:members:
-tvm.support.rocm
-~~~~~~~~~~~~~~~~
-.. automodule:: tvm.support.rocm
+tvm_tirx.support.rocm
+~~~~~~~~~~~~~~~~~~~~~
+.. automodule:: tvm_tirx.support.rocm
:members:
-tvm.support.ndk
-~~~~~~~~~~~~~~~
-.. automodule:: tvm.support.ndk
+tvm_tirx.support.ndk
+~~~~~~~~~~~~~~~~~~~~
+.. automodule:: tvm_tirx.support.ndk
:members:
-tvm.support.xcode
-~~~~~~~~~~~~~~~~~
-.. automodule:: tvm.support.xcode
+tvm_tirx.support.xcode
+~~~~~~~~~~~~~~~~~~~~~~
+.. automodule:: tvm_tirx.support.xcode
:members:
-tvm.support.clang
-~~~~~~~~~~~~~~~~~
-.. automodule:: tvm.support.clang
+tvm_tirx.support.clang
+~~~~~~~~~~~~~~~~~~~~~~
+.. automodule:: tvm_tirx.support.clang
:members:
tvm.contrib.emcc
@@ -57,17 +57,17 @@ tvm.contrib.emcc
.. automodule:: tvm.contrib.emcc
:members:
-tvm.support.popen_pool
-~~~~~~~~~~~~~~~~~~~~~~
-.. automodule:: tvm.support.popen_pool
+tvm_tirx.support.popen_pool
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+.. automodule:: tvm_tirx.support.popen_pool
:members:
-tvm.support.utils
-~~~~~~~~~~~~~~~~~
-.. automodule:: tvm.support.utils
+tvm_tirx.support.utils
+~~~~~~~~~~~~~~~~~~~~~~
+.. automodule:: tvm_tirx.support.utils
:members:
-tvm.support.tar
-~~~~~~~~~~~~~~~
-.. automodule:: tvm.support.tar
+tvm_tirx.support.tar
+~~~~~~~~~~~~~~~~~~~~
+.. automodule:: tvm_tirx.support.tar
:members:
diff --git a/docs/reference/api/python/sym.rst b/docs/reference/api/python/sym.rst
index 0dbb6ee72b..86d2fe971b 100644
--- a/docs/reference/api/python/sym.rst
+++ b/docs/reference/api/python/sym.rst
@@ -15,8 +15,8 @@
specific language governing permissions and limitations
under the License.
-tvm.sym
--------
-.. automodule:: tvm.sym
+tvm_tirx.sym
+------------
+.. automodule:: tvm_tirx.sym
:members:
:imported-members:
diff --git a/docs/reference/api/python/target.rst b/docs/reference/api/python/target.rst
index 9ccfb248d8..1e0b574914 100644
--- a/docs/reference/api/python/target.rst
+++ b/docs/reference/api/python/target.rst
@@ -15,15 +15,15 @@
specific language governing permissions and limitations
under the License.
-tvm.target
-----------
-.. automodule:: tvm.target
+tvm_tirx.target
+---------------
+.. automodule:: tvm_tirx.target
:members:
:imported-members:
:autosummary:
-tvm.target.tag
-~~~~~~~~~~~~~~
-.. automodule:: tvm.target.tag
+tvm_tirx.target.tag
+~~~~~~~~~~~~~~~~~~~
+.. automodule:: tvm_tirx.target.tag
:members:
:imported-members:
diff --git a/docs/reference/api/python/testing.rst b/docs/reference/api/python/testing.rst
index 5ef7181fd7..f16af323c9 100644
--- a/docs/reference/api/python/testing.rst
+++ b/docs/reference/api/python/testing.rst
@@ -15,8 +15,8 @@
specific language governing permissions and limitations
under the License.
-tvm.testing
------------
-.. automodule:: tvm.testing
+tvm_tirx.testing
+----------------
+.. automodule:: tvm_tirx.testing
:members:
:imported-members:
diff --git a/docs/reference/api/python/transform.rst b/docs/reference/api/python/transform.rst
index aebdbdbad5..3908c2c331 100644
--- a/docs/reference/api/python/transform.rst
+++ b/docs/reference/api/python/transform.rst
@@ -15,8 +15,8 @@
specific language governing permissions and limitations
under the License.
-tvm.ir.transform
-================
-.. automodule:: tvm.ir.transform
+tvm_tirx.ir.transform
+=====================
+.. automodule:: tvm_tirx.ir.transform
:members:
:imported-members:
diff --git a/docs/tirx/api/analysis.rst b/docs/tirx/api/analysis.rst
index 8efdf42349..82843eca8c 100644
--- a/docs/tirx/api/analysis.rst
+++ b/docs/tirx/api/analysis.rst
@@ -18,5 +18,5 @@
Analysis
========
-.. automodule:: tvm.tirx.analysis.analysis
+.. automodule:: tvm_tirx.tirx.analysis.analysis
:members:
diff --git a/docs/tirx/api/backend.rst b/docs/tirx/api/backend.rst
index 58f819ff2f..6730e79c1f 100644
--- a/docs/tirx/api/backend.rst
+++ b/docs/tirx/api/backend.rst
@@ -21,37 +21,37 @@ Backend Extension API
Backend loading
---------------
-``tvm.backend`` discovers and loads target-owned Python semantics. Depending
+``tvm_tirx.backend`` discovers and loads target-owned Python semantics. Depending
on the target, loading a backend registers its TVMScript namespaces,
tile-dispatch implementations, target tags, compilation-pipeline entry points,
and code-generation support.
-.. automodule:: tvm.backend
+.. automodule:: tvm_tirx.backend
:members:
:no-index:
CUDA registration
-----------------
-.. autofunction:: tvm.backend.cuda.register_backend
+.. autofunction:: tvm_tirx.backend.cuda.register_backend
:no-index:
-.. autofunction:: tvm.backend.cuda.script_namespace
+.. autofunction:: tvm_tirx.backend.cuda.script_namespace
:no-index:
-.. autofunction:: tvm.backend.cuda.script_namespaces
+.. autofunction:: tvm_tirx.backend.cuda.script_namespaces
:no-index:
Trainium registration
---------------------
-.. autofunction:: tvm.backend.trn.register_backend
+.. autofunction:: tvm_tirx.backend.trn.register_backend
:no-index:
-.. autofunction:: tvm.backend.trn.script_namespace
+.. autofunction:: tvm_tirx.backend.trn.script_namespace
:no-index:
-.. autofunction:: tvm.backend.trn.script_namespaces
+.. autofunction:: tvm_tirx.backend.trn.script_namespaces
:no-index:
These functions are backend integration points. Kernel-facing CUDA APIs are
diff --git a/docs/tirx/api/compilation.rst b/docs/tirx/api/compilation.rst
index 95e180a1c9..e7e04c6d37 100644
--- a/docs/tirx/api/compilation.rst
+++ b/docs/tirx/api/compilation.rst
@@ -18,15 +18,15 @@
Compilation
===========
-tvm.tirx.compile
-----------------
-.. currentmodule:: tvm.tirx
+tvm_tirx.tirx.compile
+---------------------
+.. currentmodule:: tvm_tirx.tirx
.. autofunction:: compile
:no-index:
-tvm.tirx.compilation_pipeline
------------------------------
-.. automodule:: tvm.tirx.compilation_pipeline
+tvm_tirx.tirx.compilation_pipeline
+----------------------------------
+.. automodule:: tvm_tirx.tirx.compilation_pipeline
:members:
:no-index:
diff --git a/docs/tirx/api/cuda.rst b/docs/tirx/api/cuda.rst
index 3762be0f6b..701263a222 100644
--- a/docs/tirx/api/cuda.rst
+++ b/docs/tirx/api/cuda.rst
@@ -23,7 +23,7 @@ CUDA helpers
The CUDA backend installs the ``Tx.cuda`` namespace when it is loaded::
- from tvm.script import tirx as Tx
+ from tvm_tirx.script import tirx as Tx
Tx.cuda.cta_sync()
leader = Tx.cuda.elect_sync()
@@ -91,56 +91,56 @@ provide ``.warp`` and ``.block`` variants.
CUDA language utilities
-----------------------
-``tvm.backend.cuda.lang`` provides reusable objects for hand-written CUDA TIRx
+``tvm_tirx.backend.cuda.lang`` provides reusable objects for hand-written CUDA TIRx
kernels, including schedulers, pipelines, barriers, and shared-memory pools.
-.. autoclass:: tvm.backend.cuda.lang.BaseTileScheduler
+.. autoclass:: tvm_tirx.backend.cuda.lang.BaseTileScheduler
:members:
-.. py:class:: tvm.backend.cuda.lang.ClusterPersistentScheduler2D(prefix, num_m_tiles, num_n_tiles, num_clusters, l2_group_size=8, cluster_m=1, cluster_n=1, serpentine=False)
+.. py:class:: tvm_tirx.backend.cuda.lang.ClusterPersistentScheduler2D(prefix, num_m_tiles, num_n_tiles, num_clusters, l2_group_size=8, cluster_m=1, cluster_n=1, serpentine=False)
Persistent two-dimensional tile scheduler with group-major or serpentine
traversal. ``init`` selects the first tile for a cluster, ``next_tile`` or
``next_tile_stride`` advances it, and ``valid`` reports whether the current
``m_idx`` / ``n_idx`` is in range.
-.. autoclass:: tvm.backend.cuda.lang.FlashAttentionLPTScheduler
+.. autoclass:: tvm_tirx.backend.cuda.lang.FlashAttentionLPTScheduler
:members:
-.. py:class:: tvm.backend.cuda.lang.FlashAttentionLinearScheduler(prefix, num_batches, num_heads, num_m_blocks, num_ctas)
+.. py:class:: tvm_tirx.backend.cuda.lang.FlashAttentionLinearScheduler(prefix, num_batches, num_heads, num_m_blocks, num_ctas)
Linear persistent scheduler over ``(batch, head, m_block)`` tasks.
``init`` starts from a CTA index, ``next_tile`` advances by ``num_ctas``,
and ``valid`` reports whether work remains.
-.. autoclass:: tvm.backend.cuda.lang.MBarrier
+.. autoclass:: tvm_tirx.backend.cuda.lang.MBarrier
:members:
-.. autoclass:: tvm.backend.cuda.lang.Pipeline
+.. autoclass:: tvm_tirx.backend.cuda.lang.Pipeline
:members:
-.. autoclass:: tvm.backend.cuda.lang.PipelineState
+.. autoclass:: tvm_tirx.backend.cuda.lang.PipelineState
:members:
-.. autoclass:: tvm.backend.cuda.lang.SMEMPool
+.. autoclass:: tvm_tirx.backend.cuda.lang.SMEMPool
:members:
-.. autoclass:: tvm.backend.cuda.lang.TCGen05Bar
+.. autoclass:: tvm_tirx.backend.cuda.lang.TCGen05Bar
:members:
-.. autoclass:: tvm.backend.cuda.lang.TMABar
+.. autoclass:: tvm_tirx.backend.cuda.lang.TMABar
:members:
IKET profiling
--------------
-.. automodule:: tvm.backend.cuda.iket
+.. automodule:: tvm_tirx.backend.cuda.iket
:members:
:no-index:
CUDA-specific transforms
------------------------
-.. automodule:: tvm.backend.cuda.transforms
+.. automodule:: tvm_tirx.backend.cuda.transforms
:members:
:no-index:
diff --git a/docs/tirx/api/cuda_compile.rst b/docs/tirx/api/cuda_compile.rst
index 2f5c1d8e47..c93558bc5a 100644
--- a/docs/tirx/api/cuda_compile.rst
+++ b/docs/tirx/api/cuda_compile.rst
@@ -56,26 +56,26 @@ compiler arguments are not modeled as separate Python fields.
.. code-block:: python
- from tvm.backend.cuda import BackendConfig
+ from tvm_tirx.backend.cuda import BackendConfig
cuda_config = BackendConfig(
compiler="nvrtc",
nvrtc=["--use_fast_math", "--ftz=false"],
)
- target = tvm.target.Target({
+ target = tvm_tirx.target.Target({
"kind": "cuda", "arch": "sm_100a",
"backend_config": {"cuda": cuda_config},
})
executable = tvm.compile(func, target=target)
-.. autoclass:: tvm.backend.cuda.BackendConfig
+.. autoclass:: tvm_tirx.backend.cuda.BackendConfig
:members:
:undoc-members:
-.. autoclass:: tvm.backend.cuda.compiler.CompilationResult
+.. autoclass:: tvm_tirx.backend.cuda.compiler.CompilationResult
:members:
-.. autofunction:: tvm.backend.cuda.compiler.compile_source
+.. autofunction:: tvm_tirx.backend.cuda.compiler.compile_source
See :doc:`../native_basics/cuda/compiling` for target-tag defaults, inheritance,
per-entry compilation, and artifact behavior. See :doc:`cuda_compile_options`
diff --git a/docs/tirx/api/execution.rst b/docs/tirx/api/execution.rst
index 6c22eebb30..d3acb99681 100644
--- a/docs/tirx/api/execution.rst
+++ b/docs/tirx/api/execution.rst
@@ -18,9 +18,9 @@
Thread Scope Enum Tags
======================
-tvm.tirx.thread_scope
----------------------
-.. automodule:: tvm.tirx.thread_scope
+tvm_tirx.tirx.thread_scope
+--------------------------
+.. automodule:: tvm_tirx.tirx.thread_scope
:members:
:no-index:
@@ -43,7 +43,7 @@ Same-backend hierarchy comparisons use explicit ranks; cross-backend comparisons
are rejected.
For example, CUDA TMA instructions accept ``thread`` and ``tcgen05.ld/st`` accept
``warpgroup``. CUDA index bindings and predicate analysis live in
-``tvm.backend.cuda.tile.exec_context``.
+``tvm_tirx.backend.cuda.tile.exec_context``.
Instruction APIs continue accepting ``scope="warp"`` and also accept qualified
names or ``ThreadScope`` values. Construction resolves the value before
@@ -52,10 +52,10 @@ Omitting scope uses the instruction backend's ``thread`` tag. Direct attribute
construction must provide the tag or its qualified name; shared attribute
schemas have no backend-independent default.
-.. automodule:: tvm.backend.cuda.thread_scope
+.. automodule:: tvm_tirx.backend.cuda.thread_scope
:members:
:no-index:
-.. automodule:: tvm.backend.trn.thread_scope
+.. automodule:: tvm_tirx.backend.trn.thread_scope
:members:
:no-index:
diff --git a/docs/tirx/api/layout.rst b/docs/tirx/api/layout.rst
index 158a0a2637..2b3b37172e 100644
--- a/docs/tirx/api/layout.rst
+++ b/docs/tirx/api/layout.rst
@@ -21,41 +21,41 @@ Layout IR
Layouts describe how logical tensor coordinates map to storage and execution
axes. See :doc:`../layout` for the programming model and worked examples.
-The primary classes are re-exported from ``tvm.tirx``, but are defined in this
+The primary classes are re-exported from ``tvm_tirx.tirx``, but are defined in this
module:
-.. autoclass:: tvm.tirx.layout.Layout
+.. autoclass:: tvm_tirx.tirx.layout.Layout
:members:
:no-index:
-.. autoclass:: tvm.tirx.layout.Axis
+.. autoclass:: tvm_tirx.tirx.layout.Axis
:members:
:no-index:
-.. autoclass:: tvm.tirx.layout.Iter
+.. autoclass:: tvm_tirx.tirx.layout.Iter
:members:
:no-index:
-.. autoclass:: tvm.tirx.layout.TileLayout
+.. autoclass:: tvm_tirx.tirx.layout.TileLayout
:members:
:no-index:
-.. autoclass:: tvm.tirx.layout.ComposeLayout
+.. autoclass:: tvm_tirx.tirx.layout.ComposeLayout
:members:
:no-index:
``S[...]`` and ``R[...]`` build shard and replica layout specifications,
respectively. Named axes such as ``laneid``, ``warpid``, ``tid_in_wg``,
``TLane``, and ``TCol`` belong to CUDA. Use ``T.cuda.Axis.laneid`` or
-import axes from ``tvm.backend.cuda.axis``. Trainium's ``P``, ``F``, and
+import axes from ``tvm_tirx.backend.cuda.axis``. Trainium's ``P``, ``F``, and
``Bank`` belong to ``T.trn.Axis``. The generic memory axis remains ``T.Axis.m``.
CUDA memory/instruction layout factories are defined in
-``tvm.backend.cuda.layout`` and exposed as ``T.cuda.layout``. Trainium
+``tvm_tirx.backend.cuda.layout`` and exposed as ``T.cuda.layout``. Trainium
layouts use explicit partition, free and bank axes. For example::
- from tvm.backend.trn.axis import P, F, Bank
- from tvm.tirx.layout import TileLayout, S
+ from tvm_tirx.backend.trn.axis import P, F, Bank
+ from tvm_tirx.tirx.layout import TileLayout, S
# A (128, 512) SBUF tile: partition p, free element f.
sbuf_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
@@ -74,7 +74,7 @@ the full tensor shape before slicing an operand region. The query does not
change the tensor type or its ``layout=None`` field. Noncompact global tensors
and other storage scopes require an explicit layout when the instruction uses one.
-.. automodule:: tvm.tirx.layout
+.. automodule:: tvm_tirx.tirx.layout
:members:
:no-index:
:exclude-members: Axis, ComposeLayout, Iter, Layout, TileLayout
diff --git a/docs/tirx/api/ptx.rst b/docs/tirx/api/ptx.rst
index c244389c77..f903f62097 100644
--- a/docs/tirx/api/ptx.rst
+++ b/docs/tirx/api/ptx.rst
@@ -22,11 +22,11 @@ The CUDA backend makes a table-driven PTX instruction namespace available
when authoring TIRx with TVMScript. This is an authoring namespace, not a
separate IR dialect::
- from tvm.script import tirx as Tx
+ from tvm_tirx.script import tirx as Tx
Tx.ptx.add.rn.f32(dst, lhs, rhs)
Each ``Tx.ptx`` call must match a form registered in
-:mod:`tvm.backend.cuda.ptx.table`. The table validates the instruction family,
+:mod:`tvm_tirx.backend.cuda.ptx.table`. The table validates the instruction family,
modifier combination, operand count, and operand types while constructing one
``tirx.ptx.*`` IR call. A new instruction is added by extending the table and,
when it needs a new representation, its code-generation support.
@@ -187,11 +187,11 @@ TVM CUDA codegen does not emit these operations; use ``Tx.ptx`` in new kernels.
Namespace API
-------------
-.. autoclass:: tvm.backend.cuda.ptx.PTXNamespace
+.. autoclass:: tvm_tirx.backend.cuda.ptx.PTXNamespace
:members:
:exclude-members: SINK
:special-members: __getitem__
-.. automodule:: tvm.backend.cuda.ptx
+.. automodule:: tvm_tirx.backend.cuda.ptx
:members:
:exclude-members: PTXNamespace
diff --git a/docs/tirx/api/script.rst b/docs/tirx/api/script.rst
index 196d407a9c..e4a85b764d 100644
--- a/docs/tirx/api/script.rst
+++ b/docs/tirx/api/script.rst
@@ -18,9 +18,9 @@
Core TVMScript
==============
-TIRx kernels use ``tvm.script.tirx`` for the parser and core IR builders::
+TIRx kernels use ``tvm_tirx.script.tirx`` for the parser and core IR builders::
- from tvm.script import tirx as Tx
+ from tvm_tirx.script import tirx as Tx
Tx.alloc_tensor(...)
@@ -30,7 +30,7 @@ authoring layers and TIRx IR, see :ref:`tirx-programming-model`.
Parser entry points
-------------------
-.. automodule:: tvm.script.tirx
+.. automodule:: tvm_tirx.script.tirx
:members: function, jit, inline, macro
:undoc-members:
:no-index:
@@ -39,18 +39,18 @@ Core IR builder
---------------
Common statement frames and construction helpers live in
-``tvm.script.ir_builder.frame`` and ``tvm.script.ir_builder.stmt``. The TIRx
+``tvm_tirx.script.ir_builder.frame`` and ``tvm_tirx.script.ir_builder.stmt``. The TIRx
namespace re-exports them alongside its function, tensor/layout, and execution
scope builders. Native shared statement builders live under
``include/tvm/script/ir_builder`` and ``src/script/ir_builder``; dialect-specific
thread-placement and tensor alias policies remain in the TIRx/S-TIR builders.
-.. automodule:: tvm.tirx.script.ir_builder.ir
+.. automodule:: tvm_tirx.tirx.script.ir_builder.ir
:members:
:no-index:
:exclude-members: LetAnnotation
-.. currentmodule:: tvm.tirx.script.ir_builder.ir
+.. currentmodule:: tvm_tirx.tirx.script.ir_builder.ir
.. py:class:: LetAnnotation(type_spec=None)
@@ -65,9 +65,9 @@ construct IR. ``IRBuilder.line_trace_scope()`` records file-backed Python source
lines. Script-facing void calls automatically emit into an active IRBuilder,
with or without line tracing::
- from tvm.script import tirx as Tx
- from tvm.script.ir_builder import IRBuilder
- from tvm.tirx.script import ir_builder as I
+ from tvm_tirx.script import tirx as Tx
+ from tvm_tirx.script.ir_builder import IRBuilder
+ from tvm_tirx.tirx.script import ir_builder as I
with IRBuilder() as builder:
with I.function_():
@@ -83,7 +83,7 @@ Auto emission applies to script-facing calls through ``Tx``, including PTX
and backend namespaces. A void call returns an ``AlreadyEmitted`` receipt;
passing that receipt to ``evaluate`` or ``emit_`` does not emit it again.
Calls producing values remain expressions. Raw IR constructors such as
-``tvm.ir.Call``, ``tvm.tirx.op`` and raw backend APIs remain pure even inside
+``tvm_tirx.ir.Call``, ``tvm_tirx.tirx.op`` and raw backend APIs remain pure even inside
an IRBuilder. Line tracing only supplies source locations; it neither enables
nor disables emission. The trace scope restores its previous trace and location
state on exit, including exceptional exits, and applies only to its owning builder.
@@ -121,9 +121,9 @@ A register fragment is an ordinary local tensor with an explicit layout.
For a two-CTA Layout-B accumulator, define the TMEM and register layouts
outside the kernel::
- from tvm.backend.cuda.axis import tid_in_wg
- from tvm.backend.cuda.layout import tmem_mma_operand_layout
- from tvm.tirx.layout import TileLayout, S
+ from tvm_tirx.backend.cuda.axis import tid_in_wg
+ from tvm_tirx.backend.cuda.layout import tmem_mma_operand_layout
+ from tvm_tirx.tirx.layout import TileLayout, S
layout_b = tmem_mma_operand_layout(
"D", (64, 128), "float32", M=128, cta_group=2)
diff --git a/docs/tirx/api/tirx.rst b/docs/tirx/api/tirx.rst
index f02bfe7720..c35e206a70 100644
--- a/docs/tirx/api/tirx.rst
+++ b/docs/tirx/api/tirx.rst
@@ -29,7 +29,7 @@ loads, and tensor-region constructors. Include ``tvm/tirx/type.h`` for
``TensorRegion`` expression from ``tvm/ir/expr.h`` and ``TensorRegionType``
from ``tvm/ir/type.h``.
-.. automodule:: tvm.tirx
+.. automodule:: tvm_tirx.tirx
:members:
:imported-members:
:exclude-members: Call, CommReducer, ComposeLayout, DispatchContext,
diff --git a/docs/tirx/api/trainium.rst b/docs/tirx/api/trainium.rst
index 6794934a3e..c909e9ba53 100644
--- a/docs/tirx/api/trainium.rst
+++ b/docs/tirx/api/trainium.rst
@@ -27,35 +27,35 @@ The Trainium backend installs ``Tx.nki``. Its current operations are
``tensorscalar``, ``tensorscalar_reduce``, ``scalar_tensor_tensor``,
``scalar_tensor_scalar``, ``memset``, ``identity``, and ``affine_select``.
-.. autoclass:: tvm.backend.trn.script.NKINamespace
+.. autoclass:: tvm_tirx.backend.trn.script.NKINamespace
:members:
:no-index:
Layout helpers
--------------
-.. automodule:: tvm.backend.trn.layout
+.. automodule:: tvm_tirx.backend.trn.layout
:members:
:no-index:
Compilation pipeline
--------------------
-.. automodule:: tvm.backend.trn.pipeline
+.. automodule:: tvm_tirx.backend.trn.pipeline
:members:
:no-index:
Transforms
----------
-.. automodule:: tvm.backend.trn.transform
+.. automodule:: tvm_tirx.backend.trn.transform
:members:
:no-index:
-.. autoclass:: tvm.backend.trn.transform.TrnNaiveAllocator
+.. autoclass:: tvm_tirx.backend.trn.transform.TrnNaiveAllocator
:members:
:no-index:
-.. autoclass:: tvm.backend.trn.transform.TrnPrivateTensorAlloc
+.. autoclass:: tvm_tirx.backend.trn.transform.TrnPrivateTensorAlloc
:members:
:no-index:
diff --git a/docs/tirx/api/transform.rst b/docs/tirx/api/transform.rst
index 258e225d31..1d7a924b06 100644
--- a/docs/tirx/api/transform.rst
+++ b/docs/tirx/api/transform.rst
@@ -19,7 +19,7 @@
Compiler Transforms
===================
-.. automodule:: tvm.tirx.transform
+.. automodule:: tvm_tirx.tirx.transform
:members:
:exclude-members: Attrs, PrimExpr, Op, Call
:imported-members:
(tvm) 05/33: [REFACTOR][Python] Load core and parent extensions from their owning packages
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 0bb1535c9c628d0ec28fcbc157152cecaf346eab
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 01:36:34 2026 +0000
[REFACTOR][Python] Load core and parent extensions from their owning packages
---
.../tvm_tirx/runtime => python/tvm}/__init__.py | 34 +++------
{tirx/python/tvm_tirx => python/tvm}/base.py | 41 +++--------
{tirx/python/tvm_tirx => python/tvm}/libinfo.py | 12 ++--
python/tvm/relax/backend/gpu_generic/cumsum.py | 2 +-
python/tvm/relax/backend/gpu_generic/sampling.py | 4 +-
python/tvm/relax/backend/patterns.py | 2 +-
.../tvm/relax/frontend/nn/llm/_decode_kernels.py | 2 +-
python/tvm/relax/frontend/nn/llm/_kernel_common.py | 6 +-
python/tvm/relax/frontend/nn/llm/_page_kernels.py | 2 +-
.../tvm/relax/frontend/nn/llm/_prefill_kernels.py | 4 +-
.../tvm/relax/frontend/nn/llm/_prefill_parallel.py | 2 +-
.../relax/frontend/nn/llm/position_embedding.py | 4 +-
python/tvm/relax/frontend/nn/llm/tree_attn.py | 4 +-
python/tvm/relax/frontend/nn/op.py | 7 +-
.../tvm/relax/frontend/tflite/tflite_frontend.py | 82 ++++++++++++++++------
python/tvm/relax/testing/attention.py | 4 +-
python/tvm/relax/testing/matmul.py | 4 +-
.../tvm/relax/transform/legalize_ops/inspect_op.py | 2 +-
python/tvm/rpc/minrpc.py | 30 +++++---
.../tvm_tirx => python/tvm}/runtime/__init__.py | 27 +------
python/tvm/s_tir/tensor_intrin/arm_cpu.py | 4 +-
.../tvm/s_tir/tensor_intrin/dot_product_common.py | 2 +-
python/tvm/s_tir/tensor_intrin/hexagon.py | 2 +-
python/tvm/s_tir/tensor_intrin/metal.py | 4 +-
python/tvm/s_tir/tensor_intrin/riscv_cpu.py | 2 +-
python/tvm/s_tir/tensor_intrin/rocm.py | 2 +-
python/tvm/s_tir/tensor_intrin/x86.py | 2 +-
tirx/python/tvm_tirx/__init__.py | 12 ----
tirx/python/tvm_tirx/backend/_autoload_backends.py | 21 +-----
tirx/python/tvm_tirx/backend/cuda/__init__.py | 2 +-
tirx/python/tvm_tirx/backend/metal/__init__.py | 2 +-
tirx/python/tvm_tirx/backend/opencl/__init__.py | 2 +-
tirx/python/tvm_tirx/backend/rocm/__init__.py | 2 +-
tirx/python/tvm_tirx/backend/vulkan/__init__.py | 2 +-
tirx/python/tvm_tirx/base.py | 15 ++--
tirx/python/tvm_tirx/libinfo.py | 27 ++++---
tirx/python/tvm_tirx/runtime/__init__.py | 7 --
tirx/python/tvm_tirx/script/parser/entry.py | 1 +
web/tests/python/prepare_test_libs.py | 4 +-
web/tests/python/relax_rpc_test.py | 5 +-
40 files changed, 178 insertions(+), 217 deletions(-)
diff --git a/tirx/python/tvm_tirx/runtime/__init__.py b/python/tvm/__init__.py
similarity index 50%
copy from tirx/python/tvm_tirx/runtime/__init__.py
copy to python/tvm/__init__.py
index b57d3af1b2..7320967ffb 100644
--- a/tirx/python/tvm_tirx/runtime/__init__.py
+++ b/python/tvm/__init__.py
@@ -15,28 +15,16 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""TVM runtime namespace."""
+# pylint: disable=redefined-builtin, wildcard-import
+"""TVM compiler dialects and runtime extensions built on TIRx."""
-from tvm_ffi import convert, Object
-from tvm_ffi._dtype import dtype as DataType, DataTypeCode
+from .libinfo import __version__
+from .base import _RUNTIME_ONLY
+from . import runtime
+from . import rpc
-# class exposures
-from .device import Device
-from .memory_scope import MemoryScope
-from ._tensor import Tensor, tensor, empty
-from .module import Module
-from .executable import Executable
-
-# function exposures
-from ._tensor import device, cpu, cuda, opencl, vulkan, metal
-from ._tensor import vpi, rocm, ext_dev, from_dlpack
-from .module import load_module, enabled, system_lib, load_static_library, num_threads
-
-try:
- from . import disco
-except (ImportError, ValueError):
- # disco C++ runtime is in libtvm_runtime_extra which may not be present.
- # Make the disco module optional.
- disco = None # type: ignore[assignment]
-
-from tvm_ffi import Shape as ShapeTuple
+if not _RUNTIME_ONLY:
+ from . import te
+ from . import s_tir
+ from .driver import compile
+ from . import relax
diff --git a/tirx/python/tvm_tirx/base.py b/python/tvm/base.py
similarity index 54%
copy from tirx/python/tvm_tirx/base.py
copy to python/tvm/base.py
index ca65863d5b..40469d954a 100644
--- a/tirx/python/tvm_tirx/base.py
+++ b/python/tvm/base.py
@@ -16,49 +16,30 @@
# under the License.
# coding: utf-8
# pylint: disable=invalid-name, import-outside-toplevel
-"""Base library for TVM."""
+"""Load parent TVM extensions after their standalone TIRx dependency."""
-import os
+from pathlib import Path
from tvm_ffi import libinfo as tvm_ffi_libinfo
from tvm_ffi.libinfo import load_lib_ctypes
-from . import libinfo
-
-# ----------------------------
-# library loading
-# ----------------------------
-
-# Whether only the runtime library is loaded (runtime-only wheel, or
-# ``TVM_USE_RUNTIME_LIB=1``). Set during library loading below.
-_RUNTIME_ONLY = os.environ.get("TVM_USE_RUNTIME_LIB") == "1"
-
-# Handles of the core libraries actually loaded, keyed by basename
-# (e.g. ``{"tvm_runtime": <CDLL>, "tvm_compiler": <CDLL>}``). Downstream /
-# autoloaded extensions can inspect this to skip duplicate libraries
-# (``"tvm_runtime" in _LOADED_LIBS``) and obtain the loaded handle.
-_LOADED_LIBS = {}
+from tvm_tirx.base import _RUNTIME_ONLY as _CORE_RUNTIME_ONLY
+from . import libinfo
-# runtime is loaded RTLD_GLOBAL to expose its symbols to subsequent loads;
-# compiler is loaded RTLD_LOCAL.
-_LOADED_LIBS["tvm_runtime"] = load_lib_ctypes(
- "tvm", "tvm_runtime", "RTLD_GLOBAL", extra_lib_paths=libinfo.package_lib_paths()
-)
-
+_RUNTIME_ONLY = _CORE_RUNTIME_ONLY
+_LOADED_LIBS = {
+ "tvm_runtime_extra": load_lib_ctypes(
+ "tvm", "tvm_runtime_extra", "RTLD_GLOBAL", extra_lib_paths=libinfo.package_lib_paths()
+ )
+}
if not _RUNTIME_ONLY:
- _lib_paths = libinfo.package_lib_paths()
+ _lib_paths = [Path(_LOADED_LIBS["tvm_runtime_extra"]._name).resolve().parent]
try:
tvm_ffi_libinfo._find_library_by_basename("tvm", "tvm_compiler", extra_lib_paths=_lib_paths)
except RuntimeError:
- # A runtime-only installation need not contain the compiler library.
_RUNTIME_ONLY = True
else:
_LOADED_LIBS["tvm_compiler"] = load_lib_ctypes(
"tvm", "tvm_compiler", "RTLD_LOCAL", extra_lib_paths=_lib_paths
)
-
-if _RUNTIME_ONLY:
- from tvm_ffi import registry as _tvm_ffi_registry
-
- _tvm_ffi_registry._SKIP_UNKNOWN_OBJECTS = True
diff --git a/tirx/python/tvm_tirx/libinfo.py b/python/tvm/libinfo.py
similarity index 89%
copy from tirx/python/tvm_tirx/libinfo.py
copy to python/tvm/libinfo.py
index 312d769f7e..18bc4c2daf 100644
--- a/tirx/python/tvm_tirx/libinfo.py
+++ b/python/tvm/libinfo.py
@@ -41,7 +41,7 @@ def package_lib_paths() -> list[Path]:
covers the wheel-install layout (``python/tvm/lib/``) and the in-tree dev
build layouts (``<worktree>/build/lib/`` and ``<worktree>/lib/``).
``TVM_LIBRARY_PATH`` is prepended when set so it takes priority. Callers
- pick the basenames they want (e.g. ``libtvm_runtime.so``) and the load
+ pick the basenames they want (e.g. ``libtvm_runtime_extra.so``) and the load
mode; this function only returns the search path.
"""
pkg = _rel_top_directory() # python/tvm/
@@ -56,8 +56,8 @@ def package_lib_paths() -> list[Path]:
return paths
-def find_libtvm_runtime() -> str:
- """Return the path to the ``libtvm_runtime`` shared library.
+def find_libtvm_runtime_extra() -> str:
+ """Return the path to the ``libtvm_runtime_extra`` shared library.
Returns
-------
@@ -70,11 +70,11 @@ def find_libtvm_runtime() -> str:
If the runtime library cannot be found.
"""
candidate = tvm_ffi_libinfo._find_library_by_basename(
- "tvm", "tvm_runtime", extra_lib_paths=package_lib_paths()
+ "tvm", "tvm_runtime_extra", extra_lib_paths=package_lib_paths()
)
if ret := tvm_ffi_libinfo._resolve_and_validate([candidate], cond=lambda _: True):
return ret
- raise RuntimeError("Cannot find libtvm_runtime")
+ raise RuntimeError("Cannot find libtvm_runtime_extra")
def find_include_path() -> str:
@@ -84,7 +84,7 @@ def find_include_path() -> str:
_rel_top_directory() / "include",
_dev_top_directory() / "include",
],
- cond=lambda p: (p / "tvm" / "runtime").is_dir(),
+ cond=lambda p: (p / "tvm" / "runtime" / "vm").is_dir(),
):
return ret
raise RuntimeError("Cannot find TVM include path.")
diff --git a/python/tvm/relax/backend/gpu_generic/cumsum.py b/python/tvm/relax/backend/gpu_generic/cumsum.py
index 2fb20bdb54..7610912f69 100644
--- a/python/tvm/relax/backend/gpu_generic/cumsum.py
+++ b/python/tvm/relax/backend/gpu_generic/cumsum.py
@@ -19,7 +19,7 @@
import math
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx import Function
diff --git a/python/tvm/relax/backend/gpu_generic/sampling.py b/python/tvm/relax/backend/gpu_generic/sampling.py
index 6e9d2fe499..0c28923e99 100644
--- a/python/tvm/relax/backend/gpu_generic/sampling.py
+++ b/python/tvm/relax/backend/gpu_generic/sampling.py
@@ -20,8 +20,8 @@
import math
from collections.abc import Callable
-import tvm, tvm_tirx
-from tvm_tirx.script import s_tir as Ts
+import tvm_tirx
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx import Function
diff --git a/python/tvm/relax/backend/patterns.py b/python/tvm/relax/backend/patterns.py
index 75de6ae575..f5cff44054 100644
--- a/python/tvm/relax/backend/patterns.py
+++ b/python/tvm/relax/backend/patterns.py
@@ -19,6 +19,7 @@
from collections.abc import Mapping
+from tvm.relax import script as R
from tvm.relax.dpl.pattern import (
DFPattern,
GlobalVarPattern,
@@ -28,7 +29,6 @@ from tvm.relax.dpl.pattern import (
is_tuple_get_item,
wildcard,
)
-from tvm_tirx.script import relax as R
from tvm_tirx.script import tirx as T
diff --git a/python/tvm/relax/frontend/nn/llm/_decode_kernels.py b/python/tvm/relax/frontend/nn/llm/_decode_kernels.py
index ec3ce3e670..a51b56fb5f 100644
--- a/python/tvm/relax/frontend/nn/llm/_decode_kernels.py
+++ b/python/tvm/relax/frontend/nn/llm/_decode_kernels.py
@@ -31,7 +31,7 @@ Contents:
import math
from typing import Any
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/python/tvm/relax/frontend/nn/llm/_kernel_common.py b/python/tvm/relax/frontend/nn/llm/_kernel_common.py
index e2f32620ae..d26010eefd 100644
--- a/python/tvm/relax/frontend/nn/llm/_kernel_common.py
+++ b/python/tvm/relax/frontend/nn/llm/_kernel_common.py
@@ -37,11 +37,11 @@ import enum
import math
from typing import Any
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm
from tvm import s_tir
+from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
from tvm_tirx.runtime import DataType
-from tvm_tirx.script import s_tir as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/python/tvm/relax/frontend/nn/llm/_page_kernels.py b/python/tvm/relax/frontend/nn/llm/_page_kernels.py
index 8c7f2a248d..95273278b6 100644
--- a/python/tvm/relax/frontend/nn/llm/_page_kernels.py
+++ b/python/tvm/relax/frontend/nn/llm/_page_kernels.py
@@ -32,7 +32,7 @@ This module contains:
from __future__ import annotations
# pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py b/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
index d0ee8e155a..36530e3dea 100644
--- a/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
+++ b/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
@@ -29,9 +29,9 @@ K/V loading path that is specific to its storage layout.
import math
from typing import Any, Literal
-import tvm, tvm_tirx
+import tvm
+from tvm.s_tir import script as Ts
from tvm_tirx import tirx
-from tvm_tirx.script import s_tir as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/python/tvm/relax/frontend/nn/llm/_prefill_parallel.py b/python/tvm/relax/frontend/nn/llm/_prefill_parallel.py
index eba66a746c..46253c9bf4 100644
--- a/python/tvm/relax/frontend/nn/llm/_prefill_parallel.py
+++ b/python/tvm/relax/frontend/nn/llm/_prefill_parallel.py
@@ -16,7 +16,7 @@
# under the License.
"""Parallel dot-product prefill for wide FP32 WebGPU attention heads."""
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from ._kernel_common import _get_kv_chunk_len, _get_seq_offset, _length_info_buffer, _rope
diff --git a/python/tvm/relax/frontend/nn/llm/position_embedding.py b/python/tvm/relax/frontend/nn/llm/position_embedding.py
index cddf2462ae..99ea24bedc 100644
--- a/python/tvm/relax/frontend/nn/llm/position_embedding.py
+++ b/python/tvm/relax/frontend/nn/llm/position_embedding.py
@@ -22,9 +22,9 @@ from collections.abc import Callable
from functools import partial
from typing import Any
-from tvm_tirx import tirx
from tvm.relax.frontend.nn import Tensor, op
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
from tvm_tirx.script import tirx as T
# pylint: disable=invalid-name
diff --git a/python/tvm/relax/frontend/nn/llm/tree_attn.py b/python/tvm/relax/frontend/nn/llm/tree_attn.py
index 1b52f319fd..777b1d638e 100644
--- a/python/tvm/relax/frontend/nn/llm/tree_attn.py
+++ b/python/tvm/relax/frontend/nn/llm/tree_attn.py
@@ -22,9 +22,9 @@
import math
from typing import Any
-from tvm_tirx import tirx
from tvm import s_tir
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/python/tvm/relax/frontend/nn/op.py b/python/tvm/relax/frontend/nn/op.py
index 10c925ed4f..e21e91f755 100644
--- a/python/tvm/relax/frontend/nn/op.py
+++ b/python/tvm/relax/frontend/nn/op.py
@@ -25,10 +25,10 @@ from typing import Any, TypeVar
import numpy as np
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
+from tvm.s_tir import script as Ts
from tvm_tirx import tirx as _tir
-from tvm_tirx.script import s_tir as Ts
from tvm_tirx.script import tirx as T
from ... import expr as rx
@@ -2181,7 +2181,8 @@ def extern(
out: OutType,
) -> OutType:
"""Invoke an extern function during runtime. The extern function must be registered with the "
- TVM runtime using `reflection::GlobalDef().def` (C++), or `tvm_tirx.register_global_func` (Python).
+ TVM runtime using `reflection::GlobalDef().def` (C++), or
+ `tvm_tirx.register_global_func` (Python).
Parameters
----------
diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py
index 0a66475452..27d5c21498 100644
--- a/python/tvm/relax/frontend/tflite/tflite_frontend.py
+++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py
@@ -27,13 +27,13 @@ import math
import numpy as np
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm_tirx
from tvm import relax
-from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax import op as _op
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
+from tvm_tirx.ir import TupleType as _TupleType
from .tflite_flexbuffer import FlexBufferDecoder
@@ -751,16 +751,22 @@ class OperatorConverter:
offsets = np.frombuffer(data, dtype="<i4", count=count + 1, offset=4).astype(np.int64)
if np.any(offsets < header_size) or np.any(offsets > len(data)):
- raise tvm_tirx.error.OpNotImplemented(f"{op_name} has out-of-bounds string tensor offsets")
+ raise tvm_tirx.error.OpNotImplemented(
+ f"{op_name} has out-of-bounds string tensor offsets"
+ )
if np.any(offsets[:-1] > offsets[1:]):
- raise tvm_tirx.error.OpNotImplemented(f"{op_name} has non-monotonic string tensor offsets")
+ raise tvm_tirx.error.OpNotImplemented(
+ f"{op_name} has non-monotonic string tensor offsets"
+ )
try:
values = [
data[int(offsets[i]) : int(offsets[i + 1])].decode("utf-8") for i in range(count)
]
except UnicodeDecodeError as e:
- raise tvm_tirx.error.OpNotImplemented(f"{op_name} has invalid UTF-8 string data: {e}") from e
+ raise tvm_tirx.error.OpNotImplemented(
+ f"{op_name} has invalid UTF-8 string data: {e}"
+ ) from e
shape = self._get_tensor_shape_tuple(tensor_wrapper)
expected_count = math.prod(shape) if shape else 1
if expected_count != count:
@@ -811,11 +817,15 @@ class OperatorConverter:
key_shape = self._get_tensor_shape_tuple(key_tensor)
value_shape = self._get_tensor_shape_tuple(value_tensor)
if key_shape != value_shape:
- raise tvm_tirx.error.OpNotImplemented("HASHTABLE_IMPORT requires keys and values same shape")
+ raise tvm_tirx.error.OpNotImplemented(
+ "HASHTABLE_IMPORT requires keys and values same shape"
+ )
if not self._has_tensor_buffer_data(key_tensor) or not self._has_tensor_buffer_data(
value_tensor
):
- raise tvm_tirx.error.OpNotImplemented("HASHTABLE_IMPORT requires constant keys and values")
+ raise tvm_tirx.error.OpNotImplemented(
+ "HASHTABLE_IMPORT requires constant keys and values"
+ )
if self._is_tflite_string_type(table_info["key_dtype"]):
keys = self._get_string_tensor_value(key_tensor, "HASHTABLE_IMPORT")
@@ -889,7 +899,9 @@ class OperatorConverter:
query_shape = self._get_tensor_shape_tuple(query_tensor)
output_shape = self._get_tensor_shape_tuple(output_tensor)
if output_shape != query_shape:
- raise tvm_tirx.error.OpNotImplemented("HASHTABLE_FIND output shape must match query shape")
+ raise tvm_tirx.error.OpNotImplemented(
+ "HASHTABLE_FIND output shape must match query shape"
+ )
query_values = self._get_string_tensor_value(query_tensor, "HASHTABLE_FIND")
default_values = self.get_tensor_value(default_tensor)
@@ -1554,7 +1566,9 @@ class OperatorConverter:
from tflite.LocalResponseNormalizationOptions import LocalResponseNormalizationOptions
if self.is_quantized(op):
- raise tvm_tirx.error.OpNotImplemented("TFlite quantized LRN operator is not supported yet.")
+ raise tvm_tirx.error.OpNotImplemented(
+ "TFlite quantized LRN operator is not supported yet."
+ )
input_tensors = self.get_input_tensors(op)
assert len(input_tensors) == 1, "input tensors length should be 1"
@@ -2096,7 +2110,9 @@ class OperatorConverter:
elif dtype.startswith(("int", "uint")):
op_fn = _op.bitwise_and
else:
- raise tvm_tirx.error.OpNotImplemented(f"STABLEHLO_AND with dtype {dtype} is not supported")
+ raise tvm_tirx.error.OpNotImplemented(
+ f"STABLEHLO_AND with dtype {dtype} is not supported"
+ )
return self.bb.normalize(op_fn(lhs, rhs))
def _convert_stablehlo_or(self, op):
@@ -2114,7 +2130,9 @@ class OperatorConverter:
elif dtype.startswith(("int", "uint")):
op_fn = _op.bitwise_or
else:
- raise tvm_tirx.error.OpNotImplemented(f"STABLEHLO_OR with dtype {dtype} is not supported")
+ raise tvm_tirx.error.OpNotImplemented(
+ f"STABLEHLO_OR with dtype {dtype} is not supported"
+ )
return self.bb.normalize(op_fn(lhs, rhs))
def _convert_stablehlo_ternary(self, op, relax_op):
@@ -2561,7 +2579,9 @@ class OperatorConverter:
elif direction == _DIR.STABLEHLO_COMPARISON_DIRECTION_GT:
descending = True
else:
- raise tvm_tirx.error.OpNotImplemented("STABLEHLO_SORT only supports LT or GT comparators")
+ raise tvm_tirx.error.OpNotImplemented(
+ "STABLEHLO_SORT only supports LT or GT comparators"
+ )
data = self.get_tensor_expr(input_tensors[0])
return self.bb.normalize(
@@ -2602,7 +2622,9 @@ class OperatorConverter:
return self.get_tensor_expr(input_tensors[0])
target = call_target_name or "<empty>"
- raise tvm_tirx.error.OpNotImplemented(f"STABLEHLO_CUSTOM_CALL target {target} is not supported")
+ raise tvm_tirx.error.OpNotImplemented(
+ f"STABLEHLO_CUSTOM_CALL target {target} is not supported"
+ )
def _convert_stablehlo_rng_bit_generator(self, op):
"""Convert STABLEHLO_RNG_BIT_GENERATOR to a bit-exact call_tir kernel."""
@@ -2638,7 +2660,9 @@ class OperatorConverter:
if algorithm == "threefry" and state_len != 2:
raise tvm_tirx.error.OpNotImplemented(f"{op_name} THREEFRY requires a u64[2] state")
if algorithm == "philox" and state_len not in (2, 3):
- raise tvm_tirx.error.OpNotImplemented(f"{op_name} PHILOX requires a u64[2] or u64[3] state")
+ raise tvm_tirx.error.OpNotImplemented(
+ f"{op_name} PHILOX requires a u64[2] or u64[3] state"
+ )
out_state_tensor, out_tensor = output_tensors
if self.get_tensor_type_str(out_state_tensor.tensor.Type()) != "uint64":
@@ -2650,7 +2674,9 @@ class OperatorConverter:
)
out_dtype = self.get_tensor_type_str(out_tensor.tensor.Type())
if out_dtype not in ("int32", "int64", "uint32", "uint64"):
- raise tvm_tirx.error.OpNotImplemented(f"{op_name} output dtype {out_dtype} is not supported")
+ raise tvm_tirx.error.OpNotImplemented(
+ f"{op_name} output dtype {out_dtype} is not supported"
+ )
out_shape = tuple(self._get_static_tensor_shape(out_tensor, op_name))
function = _build_stablehlo_rng_bit_generator_function(
@@ -2779,7 +2805,9 @@ class OperatorConverter:
def _check_tensor_metadata_match(self, actual, expected, op_name, tensor_role):
"""Validate that two TFLite tensors have matching static metadata."""
if self._get_tensor_metadata(actual) != self._get_tensor_metadata(expected):
- raise tvm_tirx.error.OpNotImplemented(f"{op_name} {tensor_role} tensor metadata mismatch")
+ raise tvm_tirx.error.OpNotImplemented(
+ f"{op_name} {tensor_role} tensor metadata mismatch"
+ )
def _check_subgraph_tensor_metadata(
self, subgraph, op_name, tensor_role, subgraph_indices, expected_tensors
@@ -3087,7 +3115,9 @@ class OperatorConverter:
input_tensors = self.get_input_tensors(op)
output_tensors = self.get_output_tensors(op)
if len(input_tensors) != 0 or len(output_tensors) != 0:
- raise tvm_tirx.error.OpNotImplemented("CALL_ONCE with inputs or outputs is not supported")
+ raise tvm_tirx.error.OpNotImplemented(
+ "CALL_ONCE with inputs or outputs is not supported"
+ )
init_subgraph = self._get_subgraph(init_subgraph_index, "CALL_ONCE")
if init_subgraph.InputsLength() != 0 or init_subgraph.OutputsLength() != 0:
@@ -3603,7 +3633,9 @@ class OperatorConverter:
data_shape = self._get_static_tensor_shape(input_tensors[0], "STABLEHLO_CONVOLUTION")
kernel_shape = self._get_static_tensor_shape(input_tensors[1], "STABLEHLO_CONVOLUTION")
if len(data_shape) != 4 or len(kernel_shape) != 4:
- raise tvm_tirx.error.OpNotImplemented("STABLEHLO_CONVOLUTION only supports 2D convolution")
+ raise tvm_tirx.error.OpNotImplemented(
+ "STABLEHLO_CONVOLUTION only supports 2D convolution"
+ )
if data_shape[3] != kernel_shape[2]:
raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_CONVOLUTION input channels must match kernel input channels"
@@ -3691,7 +3723,9 @@ class OperatorConverter:
"STABLEHLO_GATHER only supports collapsed_slice_dims matching the gather axis"
)
if len(slice_sizes) != len(data_shape):
- raise tvm_tirx.error.OpNotImplemented("STABLEHLO_GATHER slice_sizes must match operand rank")
+ raise tvm_tirx.error.OpNotImplemented(
+ "STABLEHLO_GATHER slice_sizes must match operand rank"
+ )
for i, (size, dim) in enumerate(zip(slice_sizes, data_shape)):
expected = 1 if i == axis else dim
if size != expected:
@@ -3703,7 +3737,9 @@ class OperatorConverter:
"STABLEHLO_GATHER only supports trailing index_vector_dim"
)
if not indices_shape or indices_shape[index_vector_dim] != 1:
- raise tvm_tirx.error.OpNotImplemented("STABLEHLO_GATHER only supports index vector size 1")
+ raise tvm_tirx.error.OpNotImplemented(
+ "STABLEHLO_GATHER only supports index vector size 1"
+ )
indices_batch_shape = indices_shape[:index_vector_dim]
expected_offset_dims = list(range(axis)) + list(
@@ -8515,7 +8551,7 @@ def _build_tflite_rfft2d_fft_function(input_shape, output_pair_shape):
# this kernel.
function_source = (
"from tvm_tirx.script.parser import tirx as T\n"
- "from tvm_tirx.script import s_tir as Ts\n"
+ "from tvm.s_tir import script as Ts\n"
"@Ts.function(private=True, check_well_formed=False)\n"
"def kernel(\n"
f" data: T.Tensor({tuple(int(x) for x in input_shape)}, 'float32'),\n"
diff --git a/python/tvm/relax/testing/attention.py b/python/tvm/relax/testing/attention.py
index 6aa1f4a946..67bdd2bc79 100644
--- a/python/tvm/relax/testing/attention.py
+++ b/python/tvm/relax/testing/attention.py
@@ -17,9 +17,9 @@
"""Relax script for attention module."""
-import tvm, tvm_tirx
+import tvm_tirx
+from tvm.relax import script as R
from tvm.relax.script import ir_builder as relax_builder
-from tvm_tirx.script import relax as R
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/python/tvm/relax/testing/matmul.py b/python/tvm/relax/testing/matmul.py
index 5ebbf3e60e..3251ccbbef 100644
--- a/python/tvm/relax/testing/matmul.py
+++ b/python/tvm/relax/testing/matmul.py
@@ -17,9 +17,9 @@
# ruff: noqa: RUF005
"""Utilities to construct matmul workloads."""
-import tvm, tvm_tirx
+import tvm_tirx
+from tvm.relax import script as R
from tvm.relax.script import ir_builder as relax_builder
-from tvm_tirx.script import relax as R
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/python/tvm/relax/transform/legalize_ops/inspect_op.py b/python/tvm/relax/transform/legalize_ops/inspect_op.py
index 068982e2ce..8bdf1931e8 100644
--- a/python/tvm/relax/transform/legalize_ops/inspect_op.py
+++ b/python/tvm/relax/transform/legalize_ops/inspect_op.py
@@ -19,8 +19,8 @@
import enum
+from tvm.s_tir import script as Ts
from tvm_tirx.ir import Call
-from tvm_tirx.script import s_tir as Ts
from tvm_tirx.script import tirx as T
from ... import op
diff --git a/python/tvm/rpc/minrpc.py b/python/tvm/rpc/minrpc.py
index 7e5c812e22..5885f75005 100644
--- a/python/tvm/rpc/minrpc.py
+++ b/python/tvm/rpc/minrpc.py
@@ -20,7 +20,8 @@ import os
import tvm_ffi
-from tvm_tirx import libinfo
+from tvm import libinfo
+from tvm_tirx import libinfo as core_libinfo
from tvm_tirx.support import cc
@@ -39,13 +40,15 @@ def find_minrpc_server_libpath(server="posix_popen_server"):
"""
curr_dir = os.path.dirname(os.path.realpath(os.path.expanduser(__file__)))
source_dir = os.path.abspath(os.path.join(curr_dir, "..", "..", ".."))
- minrpc_dir = os.path.join(source_dir, "src", "runtime", "rpc", "minrpc")
- path = os.path.join(minrpc_dir, server, f"{server}.cc")
-
- candidates = [path]
- if not os.path.isfile(path):
- raise RuntimeError(f"Cannot find minserver {server}, in candidates {candidates}")
- return minrpc_dir, path
+ roots = [os.path.join(curr_dir, ".."), source_dir]
+ candidates = []
+ for root in roots:
+ minrpc_dir = os.path.join(root, "src", "runtime", "rpc", "minrpc")
+ path = os.path.join(minrpc_dir, server, f"{server}.cc")
+ candidates.append(path)
+ if os.path.isfile(path):
+ return minrpc_dir, path
+ raise RuntimeError(f"Cannot find minserver {server}, in candidates {candidates}")
def with_minrpc(compile_func, server="posix_popen_server"):
@@ -65,7 +68,8 @@ def with_minrpc(compile_func, server="posix_popen_server"):
The return compilation.
"""
minrpc_dir, server_path = find_minrpc_server_libpath(server)
- runtime_path = libinfo.find_libtvm_runtime()
+ runtime_path = libinfo.find_libtvm_runtime_extra()
+ core_runtime_path = core_libinfo.find_libtvm_runtime()
tvm_ffi_path = tvm_ffi.libinfo.find_libtvm_ffi()
runtime_dir = os.path.abspath(os.path.dirname(runtime_path))
@@ -75,23 +79,27 @@ def with_minrpc(compile_func, server="posix_popen_server"):
# Note that however, this approach won't work on remote.
# Always recommend to link statically.
options += ["-Wl,-rpath=" + runtime_dir]
+ options += ["-Wl,-rpath=" + os.path.dirname(core_runtime_path)]
options += ["-Wl,-rpath=" + tvm_ffi_dir]
# Force the linker to retain the runtime shared lib dependency even if no
# symbol from it is directly referenced. The minrpc server binary only
# touches tvm_ffi APIs at link time but relies on global functions (e.g.
# ``rpc.CreateEventDrivenServer``) registered via static initializers
- # inside libtvm_runtime; with the default ``--as-needed`` those
+ # inside libtvm_runtime_extra; with the default ``--as-needed`` those
# registrations would never run.
options += ["-Wl,--no-as-needed"]
default_include_paths = [
libinfo.find_include_path(),
+ core_libinfo.find_include_path(),
tvm_ffi.libinfo.find_include_path(),
tvm_ffi.libinfo.find_dlpack_include_path(),
]
options += ["-I" + path for path in default_include_paths]
options += ["-I" + minrpc_dir]
fcompile = cc.cross_compiler(
- compile_func, options=options, add_files=[server_path, runtime_path, tvm_ffi_path]
+ compile_func,
+ options=options,
+ add_files=[server_path, runtime_path, core_runtime_path, tvm_ffi_path],
)
fcompile.__name__ = "with_minrpc"
fcompile.need_system_lib = True
diff --git a/tirx/python/tvm_tirx/runtime/__init__.py b/python/tvm/runtime/__init__.py
similarity index 50%
copy from tirx/python/tvm_tirx/runtime/__init__.py
copy to python/tvm/runtime/__init__.py
index b57d3af1b2..bfaec7febf 100644
--- a/tirx/python/tvm_tirx/runtime/__init__.py
+++ b/python/tvm/runtime/__init__.py
@@ -15,28 +15,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""TVM runtime namespace."""
+# pylint: disable=redefined-builtin, wildcard-import
+"""Parent runtime extensions: virtual machines and distributed execution."""
-from tvm_ffi import convert, Object
-from tvm_ffi._dtype import dtype as DataType, DataTypeCode
-
-# class exposures
-from .device import Device
-from .memory_scope import MemoryScope
-from ._tensor import Tensor, tensor, empty
-from .module import Module
-from .executable import Executable
-
-# function exposures
-from ._tensor import device, cpu, cuda, opencl, vulkan, metal
-from ._tensor import vpi, rocm, ext_dev, from_dlpack
-from .module import load_module, enabled, system_lib, load_static_library, num_threads
-
-try:
- from . import disco
-except (ImportError, ValueError):
- # disco C++ runtime is in libtvm_runtime_extra which may not be present.
- # Make the disco module optional.
- disco = None # type: ignore[assignment]
-
-from tvm_ffi import Shape as ShapeTuple
+from . import disco, vm
diff --git a/python/tvm/s_tir/tensor_intrin/arm_cpu.py b/python/tvm/s_tir/tensor_intrin/arm_cpu.py
index d566967fd3..8de05107fa 100644
--- a/python/tvm/s_tir/tensor_intrin/arm_cpu.py
+++ b/python/tvm/s_tir/tensor_intrin/arm_cpu.py
@@ -18,10 +18,10 @@
# ruff: noqa: E501, F401
"""Intrinsics for ARM tensorization."""
+from tvm.s_tir import script as Ts
+from tvm.s_tir.script.ir_builder import function as build_function
from tvm_tirx import tirx
from tvm_tirx.codegen import llvm_version_major
-from tvm.s_tir.script.ir_builder import function as build_function
-from tvm_tirx.script import s_tir as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/python/tvm/s_tir/tensor_intrin/dot_product_common.py b/python/tvm/s_tir/tensor_intrin/dot_product_common.py
index e3c8bc9ab2..7204ea6ce1 100644
--- a/python/tvm/s_tir/tensor_intrin/dot_product_common.py
+++ b/python/tvm/s_tir/tensor_intrin/dot_product_common.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name,missing-function-docstring
"""Dot product related intrinsics."""
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from .. import TensorIntrin
diff --git a/python/tvm/s_tir/tensor_intrin/hexagon.py b/python/tvm/s_tir/tensor_intrin/hexagon.py
index 34c2183330..74e4c9e9ce 100644
--- a/python/tvm/s_tir/tensor_intrin/hexagon.py
+++ b/python/tvm/s_tir/tensor_intrin/hexagon.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name,missing-function-docstring
"""Intrinsics for Hexagon tensorization."""
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from .. import TensorIntrin
diff --git a/python/tvm/s_tir/tensor_intrin/metal.py b/python/tvm/s_tir/tensor_intrin/metal.py
index bbee2ab12b..b0ca2411ff 100644
--- a/python/tvm/s_tir/tensor_intrin/metal.py
+++ b/python/tvm/s_tir/tensor_intrin/metal.py
@@ -19,9 +19,9 @@
from typing import Literal
-from tvm_tirx.runtime import MemoryScope
from tvm.s_tir import TensorIntrin
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
+from tvm_tirx.runtime import MemoryScope
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx import Expr, Function, Var
diff --git a/python/tvm/s_tir/tensor_intrin/riscv_cpu.py b/python/tvm/s_tir/tensor_intrin/riscv_cpu.py
index 4c970aa422..b31f411e78 100644
--- a/python/tvm/s_tir/tensor_intrin/riscv_cpu.py
+++ b/python/tvm/s_tir/tensor_intrin/riscv_cpu.py
@@ -21,9 +21,9 @@ import logging
import tvm_ffi
+from tvm.s_tir import script as Ts
from tvm_tirx.codegen import llvm_get_vector_width, target_has_features
from tvm_tirx.runtime import DataType
-from tvm_tirx.script import s_tir as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
diff --git a/python/tvm/s_tir/tensor_intrin/rocm.py b/python/tvm/s_tir/tensor_intrin/rocm.py
index 4d84fddbde..ccfaa9d075 100644
--- a/python/tvm/s_tir/tensor_intrin/rocm.py
+++ b/python/tvm/s_tir/tensor_intrin/rocm.py
@@ -26,8 +26,8 @@ lane dimension also lets CompactTensorAllocation remove the logical 64-lane
extent without a warp-memory lowering pass.
"""
+from tvm.s_tir import script as Ts
from tvm_tirx.runtime import convert
-from tvm_tirx.script import s_tir as Ts
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx.expr import Cast, IntImm
diff --git a/python/tvm/s_tir/tensor_intrin/x86.py b/python/tvm/s_tir/tensor_intrin/x86.py
index 959648b64a..e73370ab86 100644
--- a/python/tvm/s_tir/tensor_intrin/x86.py
+++ b/python/tvm/s_tir/tensor_intrin/x86.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name,missing-function-docstring
"""Intrinsics for x86 tensorization."""
-from tvm_tirx.script import s_tir as Ts
+from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
from .. import TensorIntrin
diff --git a/tirx/python/tvm_tirx/__init__.py b/tirx/python/tvm_tirx/__init__.py
index e6381cbe3b..efbdb05a5e 100644
--- a/tirx/python/tvm_tirx/__init__.py
+++ b/tirx/python/tvm_tirx/__init__.py
@@ -60,12 +60,6 @@ from .target import device_from_target
# Shared code generation helpers
from . import codegen
-# tvm.te
-from . import te
-
-# tvm.driver
-from .driver import compile
-
# others
from . import sym
@@ -75,12 +69,6 @@ from . import support
# Side-effect imports: register CUDA/ROCm FFI callbacks at TVM startup
from .support import rocm as _rocm, nvcc as _nvcc
-# Relax contain modules that are only available in compiler package
-# Do not import them if TVM is built with runtime only
-if not _RUNTIME_ONLY:
- # tvm.relax — registers itself via tvm.script.register_dialect in its __init__
- from . import relax
-
# NOTE: This file should be python2 compatible so we can
# raise proper error message when user run the package using
# an older version of the python
diff --git a/tirx/python/tvm_tirx/backend/_autoload_backends.py b/tirx/python/tvm_tirx/backend/_autoload_backends.py
index a348b7ec00..600075640b 100644
--- a/tirx/python/tvm_tirx/backend/_autoload_backends.py
+++ b/tirx/python/tvm_tirx/backend/_autoload_backends.py
@@ -21,12 +21,8 @@ from __future__ import annotations
import os
import warnings
from importlib.metadata import entry_points
-from pathlib import Path
-
-from tvm_ffi.libinfo import load_lib_ctypes
from tvm_tirx.backend.loader import load
-from tvm_tirx.base import _LOADED_LIBS
_BUILTIN_BACKENDS = (
"cuda",
@@ -48,17 +44,6 @@ def _load_builtin_backends() -> None:
for name in _BUILTIN_BACKENDS:
load(name)
- runtime_dir = Path(_LOADED_LIBS["tvm_runtime"]._name).resolve().parent
- try:
- # Load libtvm_runtime_extra if available for registration side effects.
- _LOADED_LIBS["tvm_runtime_extra"] = load_lib_ctypes(
- package="tvm",
- target_name="tvm_runtime_extra",
- extra_lib_paths=[runtime_dir],
- mode="RTLD_LOCAL",
- )
- except (OSError, FileNotFoundError, RuntimeError):
- pass
return None
@@ -78,10 +63,10 @@ def _autoload_backends() -> None:
_load_builtin_backends()
# Out-of-tree extensions opt into being loaded automatically at ``import tvm`` time
- # by declaring an entry point in the ``tvm.backends`` group:
- # [project.entry-points."tvm.backends"] tvm_foo = "tvm_foo:_autoload".
+ # by declaring an entry point in the ``tvm_tirx.backends`` group:
+ # [project.entry-points."tvm_tirx.backends"] tvm_foo = "tvm_foo:_autoload".
# Autoload can be disabled via ``TVM_DEVICE_BACKEND_AUTOLOAD=0``.
- for entry_pt in entry_points(group="tvm.backends"):
+ for entry_pt in entry_points(group="tvm_tirx.backends"):
try:
entry_pt.load()()
except Exception as e: # pylint: disable=broad-except
diff --git a/tirx/python/tvm_tirx/backend/cuda/__init__.py b/tirx/python/tvm_tirx/backend/cuda/__init__.py
index 01bd9316f6..ee50466435 100644
--- a/tirx/python/tvm_tirx/backend/cuda/__init__.py
+++ b/tirx/python/tvm_tirx/backend/cuda/__init__.py
@@ -72,7 +72,7 @@ def register_backend():
try:
# Runtime sidecars only need registration side effects; libtvm_runtime is global.
_LOADED_LIBS["tvm_runtime_cuda"] = load_lib_ctypes(
- package="tvm",
+ package="tvm_tirx",
target_name="tvm_runtime_cuda",
extra_lib_paths=[runtime_dir],
mode="RTLD_LOCAL",
diff --git a/tirx/python/tvm_tirx/backend/metal/__init__.py b/tirx/python/tvm_tirx/backend/metal/__init__.py
index 46afaa5f13..222d01c78c 100644
--- a/tirx/python/tvm_tirx/backend/metal/__init__.py
+++ b/tirx/python/tvm_tirx/backend/metal/__init__.py
@@ -50,7 +50,7 @@ def register_backend():
try:
# Runtime sidecars only need registration side effects; libtvm_runtime is global.
_LOADED_LIBS["tvm_runtime_metal"] = load_lib_ctypes(
- package="tvm",
+ package="tvm_tirx",
target_name="tvm_runtime_metal",
extra_lib_paths=[runtime_dir],
mode="RTLD_LOCAL",
diff --git a/tirx/python/tvm_tirx/backend/opencl/__init__.py b/tirx/python/tvm_tirx/backend/opencl/__init__.py
index 15d2277b67..633d540593 100644
--- a/tirx/python/tvm_tirx/backend/opencl/__init__.py
+++ b/tirx/python/tvm_tirx/backend/opencl/__init__.py
@@ -46,7 +46,7 @@ def register_backend():
try:
# Runtime sidecars only need registration side effects; libtvm_runtime is global.
_LOADED_LIBS["tvm_runtime_opencl"] = load_lib_ctypes(
- package="tvm",
+ package="tvm_tirx",
target_name="tvm_runtime_opencl",
extra_lib_paths=[runtime_dir],
mode="RTLD_LOCAL",
diff --git a/tirx/python/tvm_tirx/backend/rocm/__init__.py b/tirx/python/tvm_tirx/backend/rocm/__init__.py
index 20ff0f9ea1..ffd3ccfccd 100644
--- a/tirx/python/tvm_tirx/backend/rocm/__init__.py
+++ b/tirx/python/tvm_tirx/backend/rocm/__init__.py
@@ -45,7 +45,7 @@ def register_backend():
try:
# Runtime sidecars only need registration side effects; libtvm_runtime is global.
_LOADED_LIBS["tvm_runtime_rocm"] = load_lib_ctypes(
- package="tvm",
+ package="tvm_tirx",
target_name="tvm_runtime_rocm",
extra_lib_paths=[runtime_dir],
mode="RTLD_LOCAL",
diff --git a/tirx/python/tvm_tirx/backend/vulkan/__init__.py b/tirx/python/tvm_tirx/backend/vulkan/__init__.py
index 476fb716af..c13d17c574 100644
--- a/tirx/python/tvm_tirx/backend/vulkan/__init__.py
+++ b/tirx/python/tvm_tirx/backend/vulkan/__init__.py
@@ -53,7 +53,7 @@ def register_backend():
try:
# Runtime sidecars only need registration side effects; libtvm_runtime is global.
_LOADED_LIBS["tvm_runtime_vulkan"] = load_lib_ctypes(
- package="tvm",
+ package="tvm_tirx",
target_name="tvm_runtime_vulkan",
extra_lib_paths=[runtime_dir],
mode="RTLD_LOCAL",
diff --git a/tirx/python/tvm_tirx/base.py b/tirx/python/tvm_tirx/base.py
index ca65863d5b..8bc62431c2 100644
--- a/tirx/python/tvm_tirx/base.py
+++ b/tirx/python/tvm_tirx/base.py
@@ -19,6 +19,7 @@
"""Base library for TVM."""
import os
+from pathlib import Path
from tvm_ffi import libinfo as tvm_ffi_libinfo
from tvm_ffi.libinfo import load_lib_ctypes
@@ -34,7 +35,7 @@ from . import libinfo
_RUNTIME_ONLY = os.environ.get("TVM_USE_RUNTIME_LIB") == "1"
# Handles of the core libraries actually loaded, keyed by basename
-# (e.g. ``{"tvm_runtime": <CDLL>, "tvm_compiler": <CDLL>}``). Downstream /
+# (e.g. ``{"tvm_runtime": <CDLL>, "tvm_tirx": <CDLL>}``). Downstream /
# autoloaded extensions can inspect this to skip duplicate libraries
# (``"tvm_runtime" in _LOADED_LIBS``) and obtain the loaded handle.
_LOADED_LIBS = {}
@@ -43,19 +44,21 @@ _LOADED_LIBS = {}
# runtime is loaded RTLD_GLOBAL to expose its symbols to subsequent loads;
# compiler is loaded RTLD_LOCAL.
_LOADED_LIBS["tvm_runtime"] = load_lib_ctypes(
- "tvm", "tvm_runtime", "RTLD_GLOBAL", extra_lib_paths=libinfo.package_lib_paths()
+ "tvm_tirx", "tvm_runtime", "RTLD_GLOBAL", extra_lib_paths=libinfo.package_lib_paths()
)
if not _RUNTIME_ONLY:
- _lib_paths = libinfo.package_lib_paths()
+ _lib_paths = [Path(_LOADED_LIBS["tvm_runtime"]._name).resolve().parent]
try:
- tvm_ffi_libinfo._find_library_by_basename("tvm", "tvm_compiler", extra_lib_paths=_lib_paths)
+ tvm_ffi_libinfo._find_library_by_basename(
+ "tvm_tirx", "tvm_tirx", extra_lib_paths=_lib_paths
+ )
except RuntimeError:
# A runtime-only installation need not contain the compiler library.
_RUNTIME_ONLY = True
else:
- _LOADED_LIBS["tvm_compiler"] = load_lib_ctypes(
- "tvm", "tvm_compiler", "RTLD_LOCAL", extra_lib_paths=_lib_paths
+ _LOADED_LIBS["tvm_tirx"] = load_lib_ctypes(
+ "tvm_tirx", "tvm_tirx", "RTLD_LOCAL", extra_lib_paths=_lib_paths
)
if _RUNTIME_ONLY:
diff --git a/tirx/python/tvm_tirx/libinfo.py b/tirx/python/tvm_tirx/libinfo.py
index 312d769f7e..79345f706b 100644
--- a/tirx/python/tvm_tirx/libinfo.py
+++ b/tirx/python/tvm_tirx/libinfo.py
@@ -25,7 +25,7 @@ from tvm_ffi import libinfo as tvm_ffi_libinfo
def _rel_top_directory() -> Path:
- """Top directory in the installed (wheel) layout: ``python/tvm/``."""
+ """Top directory in the installed (wheel) layout: ``python/tvm_tirx/``."""
return Path(__file__).parent
@@ -37,22 +37,25 @@ def _dev_top_directory() -> Path:
def package_lib_paths() -> list[Path]:
"""Return search directories for TVM's shared libraries.
- Anchored on this file's location (``python/tvm/libinfo.py``), the list
- covers the wheel-install layout (``python/tvm/lib/``) and the in-tree dev
+ Anchored on this file's location (``python/tvm_tirx/libinfo.py``), the list
+ covers the wheel-install layout (``python/tvm_tirx/lib/``) and the in-tree dev
build layouts (``<worktree>/build/lib/`` and ``<worktree>/lib/``).
- ``TVM_LIBRARY_PATH`` is prepended when set so it takes priority. Callers
+ ``TVM_TIRX_LIBRARY_PATH`` is prepended when set so it takes priority. Callers
pick the basenames they want (e.g. ``libtvm_runtime.so``) and the load
mode; this function only returns the search path.
"""
- pkg = _rel_top_directory() # python/tvm/
+ pkg = _rel_top_directory() # python/tvm_tirx/
paths: list[Path] = []
- if os.environ.get("TVM_LIBRARY_PATH"):
- paths.append(Path(os.environ["TVM_LIBRARY_PATH"]))
+ if os.environ.get("TVM_TIRX_LIBRARY_PATH"):
+ paths.append(Path(os.environ["TVM_TIRX_LIBRARY_PATH"]))
paths += [
pkg / "lib", # wheel layout
_dev_top_directory() / "build" / "lib", # dev: <worktree>/build/lib
_dev_top_directory() / "lib", # dev: <worktree>/lib
]
+ dev = _dev_top_directory().resolve()
+ if dev.name == "tirx" and (dev.parent / "CMakeLists.txt").is_file():
+ paths.append(dev.parent / "build" / "lib")
return paths
@@ -70,7 +73,7 @@ def find_libtvm_runtime() -> str:
If the runtime library cannot be found.
"""
candidate = tvm_ffi_libinfo._find_library_by_basename(
- "tvm", "tvm_runtime", extra_lib_paths=package_lib_paths()
+ "tvm_tirx", "tvm_runtime", extra_lib_paths=package_lib_paths()
)
if ret := tvm_ffi_libinfo._resolve_and_validate([candidate], cond=lambda _: True):
return ret
@@ -90,10 +93,4 @@ def find_include_path() -> str:
raise RuntimeError("Cannot find TVM include path.")
-# The version is written by setuptools_scm into _version.py at build time
-# (see [tool.setuptools_scm] in pyproject.toml). The fallback keeps a source
-# checkout with no build run importable.
-try:
- from ._version import version as __version__
-except ImportError: # pragma: no cover - source tree without a build
- __version__ = "0.28.dev0"
+__version__ = "0.1.0a0"
diff --git a/tirx/python/tvm_tirx/runtime/__init__.py b/tirx/python/tvm_tirx/runtime/__init__.py
index b57d3af1b2..b9eebb95ff 100644
--- a/tirx/python/tvm_tirx/runtime/__init__.py
+++ b/tirx/python/tvm_tirx/runtime/__init__.py
@@ -32,11 +32,4 @@ from ._tensor import device, cpu, cuda, opencl, vulkan, metal
from ._tensor import vpi, rocm, ext_dev, from_dlpack
from .module import load_module, enabled, system_lib, load_static_library, num_threads
-try:
- from . import disco
-except (ImportError, ValueError):
- # disco C++ runtime is in libtvm_runtime_extra which may not be present.
- # Make the disco module optional.
- disco = None # type: ignore[assignment]
-
from tvm_ffi import Shape as ShapeTuple
diff --git a/tirx/python/tvm_tirx/script/parser/entry.py b/tirx/python/tvm_tirx/script/parser/entry.py
index 60302557e6..00d529b920 100644
--- a/tirx/python/tvm_tirx/script/parser/entry.py
+++ b/tirx/python/tvm_tirx/script/parser/entry.py
@@ -570,6 +570,7 @@ def _prepare_transpiler(
No builder frame or expression is created here.
"""
namespace = {
+ "tvm": sys.modules.get("tvm"),
"tvm_tirx": sys.modules.get("tvm_tirx"),
**_NAMESPACES,
**environment,
diff --git a/web/tests/python/prepare_test_libs.py b/web/tests/python/prepare_test_libs.py
index c6dc645b95..bdc70184fd 100644
--- a/web/tests/python/prepare_test_libs.py
+++ b/web/tests/python/prepare_test_libs.py
@@ -18,11 +18,11 @@
import os
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import relax, te
from tvm.contrib import emcc
+from tvm.relax import script as R
from tvm_tirx.script import ir as I
-from tvm_tirx.script import relax as R
def prepare_relax_lib(base_path):
diff --git a/web/tests/python/relax_rpc_test.py b/web/tests/python/relax_rpc_test.py
index 1444558ba6..e3d9cdbe48 100644
--- a/web/tests/python/relax_rpc_test.py
+++ b/web/tests/python/relax_rpc_test.py
@@ -18,10 +18,11 @@
import numpy as np
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import relax, rpc
from tvm.contrib import emcc
-from tvm_tirx.script import relax as R
+from tvm.relax import script as R
from tvm_tirx.support import utils
proxy_host = "127.0.0.1"
(tvm) 31/33: [FIX][Python] Resolve the incoming squeeze regression through the core namespace
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 0043f3ae2712e505f11ffa8055deff7009d39f8b
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 12:45:54 2026 +0000
[FIX][Python] Resolve the incoming squeeze regression through the core namespace
---
tests/python/relax/test_frontend_from_fx.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py
index ebff36a3ad..d3c6315c13 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -2586,7 +2586,7 @@ def test_squeeze_out_of_range_dim():
def forward(self, input):
return input.squeeze((0, 2))
- @tvm.script.ir_module
+ @tvm_tirx.script.ir_module
class Expected:
@R.function
def main(
(tvm) 02/33: [REFACTOR][Python] Rename extracted core imports to tvm_tirx
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 08fadc36f813a1063db102abaf77314a4865551b
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 00:56:49 2026 +0000
[REFACTOR][Python] Rename extracted core imports to tvm_tirx
---
apps/android_rpc/tests/android_rpc_test.py | 10 +-
apps/ios_rpc/tests/ios_rpc_test.py | 12 +-
jvm/core/src/test/scripts/prepare_test_libs.py | 22 +-
python/tvm/contrib/cblas.py | 6 +-
python/tvm/contrib/cublas.py | 6 +-
python/tvm/contrib/cublaslt.py | 4 +-
python/tvm/contrib/cudnn.py | 26 +-
python/tvm/contrib/cutlass/build.py | 27 +-
python/tvm/contrib/cutlass/gen_tensor_op.py | 10 +-
python/tvm/contrib/cutlass/library.py | 2 +-
python/tvm/contrib/dlpack.py | 4 +-
python/tvm/contrib/dnnl.py | 6 +-
python/tvm/contrib/emcc.py | 2 +-
python/tvm/contrib/hexagon/generate_take_op.py | 16 +-
python/tvm/contrib/hexagon/tools.py | 26 +-
python/tvm/contrib/hipblas.py | 6 +-
python/tvm/contrib/mkl.py | 8 +-
python/tvm/contrib/nnpack.py | 10 +-
python/tvm/contrib/random.py | 8 +-
python/tvm/contrib/tvmjs.py | 20 +-
python/tvm/driver/build_module.py | 20 +-
python/tvm/exec/disco_remote_socket_session.py | 4 +-
python/tvm/exec/disco_worker.py | 10 +-
python/tvm/exec/gpu_memory_bandwidth.py | 6 +-
python/tvm/relax/__init__.py | 12 +-
python/tvm/relax/analysis/analysis.py | 14 +-
python/tvm/relax/analysis/estimate_memory_usage.py | 12 +-
python/tvm/relax/backend/adreno/clml.py | 31 +-
python/tvm/relax/backend/adreno/pipeline.py | 20 +-
.../relax/backend/adreno/transform/transform.py | 12 +-
.../relax/backend/contrib/example_npu/patterns.py | 2 +-
python/tvm/relax/backend/contrib/nnapi.py | 6 +-
python/tvm/relax/backend/contrib/tensorrt.py | 18 +-
python/tvm/relax/backend/cpu_generic/pipeline.py | 18 +-
python/tvm/relax/backend/cuda/cublas.py | 24 +-
python/tvm/relax/backend/cuda/cudnn.py | 12 +-
python/tvm/relax/backend/cuda/cutlass.py | 16 +-
python/tvm/relax/backend/cuda/flashinfer.py | 20 +-
python/tvm/relax/backend/cuda/pipeline.py | 18 +-
python/tvm/relax/backend/dispatch_sampling.py | 6 +-
python/tvm/relax/backend/dispatch_sort_scan.py | 11 +-
python/tvm/relax/backend/gpu_generic/cumsum.py | 6 +-
python/tvm/relax/backend/gpu_generic/pipeline.py | 18 +-
python/tvm/relax/backend/gpu_generic/sampling.py | 10 +-
python/tvm/relax/backend/metal/coreml.py | 16 +-
python/tvm/relax/backend/patterns.py | 4 +-
python/tvm/relax/backend/rocm/hipblas.py | 14 +-
python/tvm/relax/backend/rocm/pipeline.py | 18 +-
python/tvm/relax/backend/utils.py | 2 +-
python/tvm/relax/base_py_module.py | 31 +-
python/tvm/relax/binding_rewrite.py | 10 +-
python/tvm/relax/block_builder.py | 24 +-
python/tvm/relax/distributed/global_info.py | 4 +-
.../tvm/relax/distributed/transform/transform.py | 18 +-
python/tvm/relax/distributed/type.py | 6 +-
python/tvm/relax/dpl/context.py | 4 +-
python/tvm/relax/dpl/pattern.py | 30 +-
python/tvm/relax/dpl/rewrite.py | 4 +-
python/tvm/relax/exec_builder.py | 6 +-
python/tvm/relax/expr.py | 84 +-
python/tvm/relax/expr_functor.py | 8 +-
python/tvm/relax/frontend/common.py | 16 +-
python/tvm/relax/frontend/nn/_tensor_op.py | 2 +-
python/tvm/relax/frontend/nn/core.py | 46 +-
python/tvm/relax/frontend/nn/exporter.py | 12 +-
python/tvm/relax/frontend/nn/extern.py | 14 +-
.../tvm/relax/frontend/nn/llm/_decode_kernels.py | 6 +-
python/tvm/relax/frontend/nn/llm/_kernel_common.py | 13 +-
python/tvm/relax/frontend/nn/llm/_page_kernels.py | 6 +-
.../tvm/relax/frontend/nn/llm/_prefill_kernels.py | 10 +-
.../tvm/relax/frontend/nn/llm/_prefill_parallel.py | 4 +-
python/tvm/relax/frontend/nn/llm/kv_cache.py | 56 +-
.../relax/frontend/nn/llm/position_embedding.py | 6 +-
python/tvm/relax/frontend/nn/llm/tree_attn.py | 17 +-
python/tvm/relax/frontend/nn/modules.py | 2 +-
python/tvm/relax/frontend/nn/op.py | 32 +-
python/tvm/relax/frontend/nn/subroutine.py | 3 +-
python/tvm/relax/frontend/nn/torch.py | 2 +-
python/tvm/relax/frontend/onnx/onnx_frontend.py | 259 ++---
.../frontend/stablehlo/stablehlo_translator.py | 15 +-
.../tvm/relax/frontend/tflite/tflite_frontend.py | 483 ++++-----
.../frontend/torch/base_fx_graph_translator.py | 33 +-
python/tvm/relax/frontend/torch/dynamo.py | 22 +-
.../frontend/torch/exported_program_translator.py | 72 +-
python/tvm/relax/frontend/torch/fx_translator.py | 22 +-
python/tvm/relax/global_info.py | 6 +-
python/tvm/relax/ir/instrument.py | 4 +-
python/tvm/relax/op/__init__.py | 6 +-
python/tvm/relax/op/_op_gradient.py | 4 +-
python/tvm/relax/op/base.py | 58 +-
python/tvm/relax/op/binary.py | 4 +-
python/tvm/relax/op/builtin/builtin.py | 8 +-
python/tvm/relax/op/ccl/ccl.py | 8 +-
python/tvm/relax/op/create.py | 10 +-
python/tvm/relax/op/datatype.py | 10 +-
python/tvm/relax/op/distributed/distributed.py | 8 +-
python/tvm/relax/op/grad/grad.py | 8 +-
python/tvm/relax/op/image/image.py | 12 +-
python/tvm/relax/op/index.py | 10 +-
python/tvm/relax/op/linear_algebra.py | 12 +-
python/tvm/relax/op/manipulate.py | 14 +-
python/tvm/relax/op/memory/memory.py | 8 +-
python/tvm/relax/op/memory/view.py | 6 +-
python/tvm/relax/op/nn/nn.py | 17 +-
python/tvm/relax/op/qdq.py | 10 +-
python/tvm/relax/op/sampling.py | 8 +-
python/tvm/relax/op/search.py | 8 +-
python/tvm/relax/op/set.py | 36 +-
python/tvm/relax/op/sorting.py | 8 +-
python/tvm/relax/op/statistical.py | 10 +-
python/tvm/relax/op/ternary.py | 4 +-
python/tvm/relax/op/unary.py | 4 +-
.../tvm/relax/op/vision/multibox_transform_loc.py | 10 +-
python/tvm/relax/op/vision/nms.py | 10 +-
python/tvm/relax/op/vision/roi_align.py | 10 +-
python/tvm/relax/op/vision/roi_pool.py | 10 +-
python/tvm/relax/op/vm/vm.py | 8 +-
python/tvm/relax/pipeline.py | 52 +-
python/tvm/relax/relax_to_pyfunc_converter.py | 21 +-
python/tvm/relax/script/__init__.py | 17 +-
python/tvm/relax/script/ir_builder/__init__.py | 4 +-
.../tvm/relax/script/ir_builder/distributed/ir.py | 12 +-
.../tvm/relax/script/ir_builder/distributed/op.py | 2 +-
python/tvm/relax/script/ir_builder/frame.py | 2 +-
python/tvm/relax/script/ir_builder/ir.py | 24 +-
python/tvm/relax/script/ir_builder/op.py | 38 +-
.../tvm/relax/script/ir_builder/parser_protocol.py | 60 +-
python/tvm/relax/script/printer.py | 6 +-
python/tvm/relax/testing/ast_printer.py | 30 +-
python/tvm/relax/testing/attention.py | 12 +-
python/tvm/relax/testing/lib_comparator.py | 16 +-
python/tvm/relax/testing/matmul.py | 8 +-
python/tvm/relax/testing/nn.py | 9 +-
python/tvm/relax/testing/transform.py | 20 +-
python/tvm/relax/testing/vm.py | 42 +-
python/tvm/relax/training/optimizer.py | 24 +-
python/tvm/relax/training/setup_trainer.py | 10 +-
python/tvm/relax/training/trainer.py | 38 +-
python/tvm/relax/training/utils.py | 6 +-
.../tvm/relax/transform/attach_external_modules.py | 4 +-
python/tvm/relax/transform/fast_math.py | 8 +-
.../fold_batch_norm_to_conv2d_for_inference.py | 7 +-
.../tvm/relax/transform/fuse_transpose_matmul.py | 9 +-
.../tvm/relax/transform/ipc_allreduce_rewrite.py | 16 +-
.../tvm/relax/transform/lazy_transform_params.py | 11 +-
python/tvm/relax/transform/legalize_ops/binary.py | 2 +-
python/tvm/relax/transform/legalize_ops/ccl.py | 5 +-
python/tvm/relax/transform/legalize_ops/common.py | 18 +-
python/tvm/relax/transform/legalize_ops/create.py | 9 +-
.../tvm/relax/transform/legalize_ops/datatype.py | 2 +-
.../relax/transform/legalize_ops/distributed.py | 5 +-
python/tvm/relax/transform/legalize_ops/grad.py | 9 +-
python/tvm/relax/transform/legalize_ops/image.py | 5 +-
python/tvm/relax/transform/legalize_ops/index.py | 9 +-
.../tvm/relax/transform/legalize_ops/inspect_op.py | 6 +-
.../relax/transform/legalize_ops/linear_algebra.py | 5 +-
.../tvm/relax/transform/legalize_ops/manipulate.py | 19 +-
python/tvm/relax/transform/legalize_ops/nn.py | 5 +-
python/tvm/relax/transform/legalize_ops/qdq.py | 15 +-
python/tvm/relax/transform/legalize_ops/search.py | 2 +-
.../relax/transform/legalize_ops/statistical.py | 5 +-
python/tvm/relax/transform/legalize_ops/unary.py | 2 +-
python/tvm/relax/transform/legalize_ops/vision.py | 9 +-
.../relax/transform/lower_gpu_ipc_alloc_storage.py | 14 +-
.../relax/transform/optimize_layout_transform.py | 4 +-
.../relax/transform/remove_redundant_reshape.py | 5 +-
python/tvm/relax/transform/transform.py | 262 ++---
python/tvm/relax/ty.py | 4 +-
python/tvm/relax/type.py | 6 +-
python/tvm/relax/utils.py | 44 +-
python/tvm/relax/vm_build.py | 56 +-
python/tvm/rpc/client.py | 8 +-
python/tvm/rpc/minrpc.py | 4 +-
python/tvm/rpc/proxy.py | 2 +-
python/tvm/rpc/remote.py | 28 +-
python/tvm/rpc/server.py | 16 +-
python/tvm/rpc/testing.py | 20 +-
python/tvm/rpc/tracker.py | 2 +-
python/tvm/runtime/disco/session.py | 6 +-
python/tvm/runtime/vm/param.py | 2 +-
python/tvm/runtime/vm/vm.py | 56 +-
python/tvm/s_tir/__init__.py | 12 +-
python/tvm/s_tir/_tensor_intrin.py | 4 +-
python/tvm/s_tir/analysis/__init__.py | 22 +-
python/tvm/s_tir/backend/adreno/pipeline.py | 2 +-
.../s_tir/backend/adreno/transform/transform.py | 4 +-
python/tvm/s_tir/data_layout.py | 2 +-
python/tvm/s_tir/dlight/adreno/base.py | 2 +-
python/tvm/s_tir/dlight/adreno/convolution.py | 5 +-
python/tvm/s_tir/dlight/adreno/fallback.py | 7 +-
python/tvm/s_tir/dlight/adreno/layout_transform.py | 5 +-
python/tvm/s_tir/dlight/adreno/pool.py | 5 +-
python/tvm/s_tir/dlight/adreno/utils.py | 2 +-
.../tvm/s_tir/dlight/analysis/common_analysis.py | 15 +-
python/tvm/s_tir/dlight/analysis/gemv.py | 7 +-
python/tvm/s_tir/dlight/base/schedule_rule.py | 5 +-
python/tvm/s_tir/dlight/base/transform.py | 9 +-
python/tvm/s_tir/dlight/base/utils.py | 13 +-
python/tvm/s_tir/dlight/benchmark/bench.py | 44 +-
python/tvm/s_tir/dlight/benchmark/extract.py | 88 +-
python/tvm/s_tir/dlight/benchmark/utils.py | 18 +-
python/tvm/s_tir/dlight/cpu/base.py | 2 +-
python/tvm/s_tir/dlight/cpu/gemv.py | 5 +-
python/tvm/s_tir/dlight/cpu/reduction.py | 7 +-
python/tvm/s_tir/dlight/gpu/base.py | 2 +-
python/tvm/s_tir/dlight/gpu/fallback.py | 21 +-
python/tvm/s_tir/dlight/gpu/gemv.py | 9 +-
python/tvm/s_tir/dlight/gpu/general_reduction.py | 7 +-
python/tvm/s_tir/dlight/gpu/low_batch_gemv.py | 11 +-
python/tvm/s_tir/dlight/gpu/matmul.py | 13 +-
python/tvm/s_tir/dlight/gpu/reduction.py | 9 +-
python/tvm/s_tir/dlight/gpu/rmsnorm.py | 12 +-
python/tvm/s_tir/dlight/gpu/transpose.py | 11 +-
python/tvm/s_tir/iter_var.py | 6 +-
python/tvm/s_tir/meta_schedule/arg_info.py | 6 +-
python/tvm/s_tir/meta_schedule/builder/builder.py | 6 +-
.../s_tir/meta_schedule/builder/local_builder.py | 12 +-
.../s_tir/meta_schedule/cost_model/cost_model.py | 2 +-
.../s_tir/meta_schedule/cost_model/mlp_model.py | 6 +-
.../s_tir/meta_schedule/cost_model/xgb_model.py | 4 +-
.../tvm/s_tir/meta_schedule/database/database.py | 8 +-
python/tvm/s_tir/meta_schedule/extracted_task.py | 6 +-
.../feature_extractor/feature_extractor.py | 4 +-
.../feature_extractor/random_feature_extractor.py | 6 +-
.../measure_callback/measure_callback.py | 2 +-
python/tvm/s_tir/meta_schedule/mutator/mutator.py | 2 +-
.../meta_schedule/post_optimization/post_opt.py | 2 +-
.../s_tir/meta_schedule/post_optimization/space.py | 2 +-
.../tvm/s_tir/meta_schedule/postproc/postproc.py | 2 +-
python/tvm/s_tir/meta_schedule/profiler.py | 2 +-
.../tvm/s_tir/meta_schedule/relax_integration.py | 32 +-
.../tvm/s_tir/meta_schedule/runner/local_runner.py | 12 +-
.../tvm/s_tir/meta_schedule/runner/rpc_runner.py | 4 +-
python/tvm/s_tir/meta_schedule/runner/runner.py | 2 +-
python/tvm/s_tir/meta_schedule/runner/utils.py | 8 +-
.../schedule/cuda/layout_transform.py | 8 +-
.../meta_schedule/schedule_rule/schedule_rule.py | 2 +-
.../search_strategy/search_strategy.py | 2 +-
.../space_generator/space_generator.py | 4 +-
.../meta_schedule/task_scheduler/task_scheduler.py | 2 +-
.../meta_schedule/testing/custom_builder_runner.py | 6 +-
.../testing/distributed_measure_candidates.py | 2 +-
.../meta_schedule/testing/space_generation.py | 4 +-
.../tvm/s_tir/meta_schedule/testing/te_workload.py | 5 +-
python/tvm/s_tir/meta_schedule/testing/tune_te.py | 6 +-
.../tvm/s_tir/meta_schedule/testing/tune_utils.py | 12 +-
.../meta_schedule/testing/validate_database.py | 58 +-
python/tvm/s_tir/meta_schedule/tir_integration.py | 6 +-
python/tvm/s_tir/meta_schedule/trace_apply.py | 4 +-
python/tvm/s_tir/meta_schedule/tune_context.py | 8 +-
python/tvm/s_tir/meta_schedule/utils.py | 4 +-
python/tvm/s_tir/op.py | 2 +-
python/tvm/s_tir/pipeline.py | 33 +-
python/tvm/s_tir/sblock_dependence_info.py | 6 +-
python/tvm/s_tir/sblock_scope.py | 4 +-
python/tvm/s_tir/schedule/analysis.py | 8 +-
python/tvm/s_tir/schedule/instruction.py | 4 +-
python/tvm/s_tir/schedule/schedule.py | 10 +-
python/tvm/s_tir/schedule/state.py | 6 +-
python/tvm/s_tir/schedule/testing.py | 10 +-
python/tvm/s_tir/schedule/trace.py | 12 +-
python/tvm/s_tir/script/__init__.py | 6 +-
python/tvm/s_tir/script/ir_builder/__init__.py | 4 +-
python/tvm/s_tir/script/ir_builder/frame.py | 4 +-
python/tvm/s_tir/script/ir_builder/ir.py | 26 +-
python/tvm/s_tir/script/ir_builder/op.py | 4 +-
.../tvm/s_tir/script/ir_builder/parser_protocol.py | 10 +-
python/tvm/s_tir/stmt.py | 6 +-
python/tvm/s_tir/tensor_intrin/__init__.py | 2 +-
python/tvm/s_tir/tensor_intrin/arm_cpu.py | 12 +-
.../tvm/s_tir/tensor_intrin/dot_product_common.py | 4 +-
python/tvm/s_tir/tensor_intrin/hexagon.py | 4 +-
python/tvm/s_tir/tensor_intrin/metal.py | 8 +-
python/tvm/s_tir/tensor_intrin/riscv_cpu.py | 10 +-
python/tvm/s_tir/tensor_intrin/rocm.py | 8 +-
python/tvm/s_tir/tensor_intrin/x86.py | 4 +-
python/tvm/s_tir/testing/runner.py | 14 +-
python/tvm/s_tir/testing/tir.py | 4 +-
python/tvm/s_tir/testing/utils.py | 2 +-
python/tvm/s_tir/transform/_ffi_api.py | 2 +-
python/tvm/s_tir/transform/transform.py | 60 +-
python/tvm/te/__init__.py | 14 +-
python/tvm/te/operation.py | 60 +-
python/tvm/te/reduction.py | 24 +-
python/tvm/te/tag.py | 2 +-
python/tvm/te/tensor.py | 8 +-
python/tvm/topi/__init__.py | 2 +-
python/tvm/topi/_te_tensor_overload.py | 6 +-
python/tvm/topi/gpu/scan.py | 48 +-
python/tvm/topi/gpu/scatter_elements.py | 11 +-
python/tvm/topi/gpu/scatter_nd.py | 11 +-
python/tvm/topi/gpu/sort.py | 74 +-
python/tvm/topi/image/grid_sample.py | 3 +-
python/tvm/topi/image/resize.py | 34 +-
python/tvm/topi/index_put.py | 7 +-
python/tvm/topi/math.py | 41 +-
python/tvm/topi/nn/batch_matmul.py | 8 +-
python/tvm/topi/nn/bitserial_conv2d.py | 14 +-
python/tvm/topi/nn/bitserial_dense.py | 8 +-
python/tvm/topi/nn/bitserial_util.py | 6 +-
python/tvm/topi/nn/bnn.py | 6 +-
python/tvm/topi/nn/conv2d.py | 16 +-
python/tvm/topi/nn/deformable_conv2d.py | 14 +-
python/tvm/topi/nn/dense.py | 16 +-
python/tvm/topi/nn/depth_to_space.py | 14 +-
python/tvm/topi/nn/depthwise_conv2d.py | 14 +-
python/tvm/topi/nn/dilate.py | 14 +-
python/tvm/topi/nn/elemwise.py | 18 +-
python/tvm/topi/nn/fifo_buffer.py | 22 +-
python/tvm/topi/nn/flatten.py | 6 +-
python/tvm/topi/nn/lstm.py | 3 +-
python/tvm/topi/nn/mapping.py | 2 +-
python/tvm/topi/nn/pad.py | 24 +-
python/tvm/topi/nn/pixel_shuffle.py | 10 +-
python/tvm/topi/nn/qnn.py | 9 +-
python/tvm/topi/nn/softmax.py | 2 +-
python/tvm/topi/nn/space_to_depth.py | 18 +-
python/tvm/topi/nn/utils.py | 6 +-
python/tvm/topi/scan.py | 8 +-
python/tvm/topi/scatter.py | 9 +-
python/tvm/topi/scatter_elements.py | 7 +-
python/tvm/topi/searchsorted.py | 4 +-
python/tvm/topi/signal.py | 7 +-
python/tvm/topi/sort.py | 34 +-
python/tvm/topi/sparse_reshape.py | 6 +-
python/tvm/topi/testing/poolnd_python.py | 2 +-
python/tvm/topi/transform.py | 40 +-
python/tvm/topi/unique.py | 11 +-
python/tvm/topi/utils.py | 84 +-
python/tvm/topi/vision/multibox_transform_loc.py | 22 +-
python/tvm/topi/vision/nms.py | 142 +--
python/tvm/topi/vision/nms_util.py | 26 +-
python/tvm/topi/vision/roi_align.py | 56 +-
python/tvm/topi/vision/roi_pool.py | 32 +-
tests/nightly/python/test_nnapi/infrastructure.py | 14 +-
.../test_nnapi/test_from_exported_to_cuda.py | 226 ++---
tests/nightly/python/test_nnapi/test_network.py | 14 +-
tests/nightly/python/test_nnapi/test_ops.py | 56 +-
.../test_minimal_target_codegen_llvm.py | 20 +-
tests/python/codegen/llvm/test_cross_compile.py | 26 +-
tests/python/contrib/pickle_memoize_script.py | 2 +-
.../python/contrib/test_android/infrastructure.py | 6 +-
.../contrib/test_android/test_meta_schedule.py | 8 +-
tests/python/contrib/test_cblas.py | 52 +-
tests/python/contrib/test_coreml_runtime.py | 8 +-
tests/python/contrib/test_cutlass_gemm.py | 68 +-
tests/python/contrib/test_dlpack.py | 16 +-
tests/python/contrib/test_example_npu.py | 28 +-
tests/python/contrib/test_hipblas.py | 40 +-
tests/python/contrib/test_memoize.py | 6 +-
tests/python/contrib/test_random.py | 56 +-
tests/python/contrib/test_rpc_proxy.py | 4 +-
tests/python/contrib/test_rpc_tracker.py | 4 +-
tests/python/contrib/test_sort.py | 28 +-
.../python/contrib/test_tir_triton_integration.py | 32 +-
tests/python/contrib/test_tvmjs.py | 18 +-
tests/python/disco/test_callback.py | 30 +-
tests/python/disco/test_ccl.py | 50 +-
tests/python/disco/test_custom_allreduce.py | 9 +-
tests/python/disco/test_loader.py | 20 +-
tests/python/disco/test_nvshmem.py | 38 +-
tests/python/disco/test_session.py | 24 +-
tests/python/driver/test_compile.py | 54 +-
tests/python/relax/backend/adreno/mod_utils.py | 44 +-
tests/python/relax/backend/adreno/test_clml_ops.py | 18 +-
.../relax/backend/adreno/test_texture_network.py | 20 +-
.../relax/backend/adreno/test_texture_ops.py | 68 +-
.../adreno/test_transform_annotate_custom_scope.py | 20 +-
.../test_transform_fold_vdevice_scope_change.py | 18 +-
tests/python/relax/backend/adreno/utils.py | 34 +-
tests/python/relax/conftest.py | 6 +-
.../distributed/test_distributed_dtensor_type.py | 26 +-
..._distributed_transform_legalize_redistribute.py | 12 +-
.../test_distributed_transform_lower_distir.py | 22 +-
...ributed_transform_lower_global_to_local_view.py | 18 +-
...est_distributed_transform_propagate_sharding.py | 18 +-
.../test_runtime_builtin_kv_cache_transfer.py | 66 +-
...est_runtime_builtin_kv_cache_transfer_kernel.py | 66 +-
.../relax/script/test_relax_script_basic_usage.py | 109 +-
.../relax/script/test_relax_script_distributed.py | 35 +-
.../script/test_relax_script_dynamic_shape.py | 101 +-
.../script/test_relax_script_error_handling.py | 45 +-
.../relax/script/test_relax_script_ir_builder.py | 29 +-
.../script/test_relax_script_meta_programming.py | 47 +-
.../test_relax_script_operators_arith_cmp.py | 21 +-
.../script/test_relax_script_operators_create.py | 21 +-
.../script/test_relax_script_operators_datatype.py | 21 +-
.../script/test_relax_script_operators_grad.py | 19 +-
.../script/test_relax_script_operators_image.py | 21 +-
.../script/test_relax_script_operators_index.py | 21 +-
.../test_relax_script_operators_linear_algebra.py | 21 +-
.../test_relax_script_operators_manipulate.py | 23 +-
.../relax/script/test_relax_script_operators_nn.py | 21 +-
.../script/test_relax_script_operators_search.py | 21 +-
.../script/test_relax_script_operators_set.py | 21 +-
.../script/test_relax_script_operators_sort.py | 21 +-
.../test_relax_script_operators_statistical.py | 21 +-
.../script/test_relax_script_operators_vision.py | 21 +-
.../relax/script/test_relax_script_printer.py | 169 ++--
.../relax/script/test_relax_script_pyfunc.py | 24 +-
tests/python/relax/test_analysis.py | 62 +-
.../test_analysis_computable_at_compile_time.py | 10 +-
.../relax/test_analysis_contains_impure_call.py | 16 +-
.../python/relax/test_analysis_detect_recursion.py | 44 +-
.../relax/test_analysis_estimate_memory_usage.py | 14 +-
.../test_analysis_suggest_layout_transforms.py | 39 +-
tests/python/relax/test_analysis_type_analysis.py | 100 +-
tests/python/relax/test_analysis_well_formed.py | 128 +--
tests/python/relax/test_ast_printer.py | 28 +-
.../python/relax/test_backend_dispatch_sampling.py | 22 +-
.../relax/test_backend_dispatch_sort_scan.py | 103 +-
.../relax/test_backend_transform_shape_lower.py | 44 +-
tests/python/relax/test_base_py_module_printer.py | 50 +-
.../relax/test_base_py_module_symbolic_shape.py | 75 +-
tests/python/relax/test_bind_params.py | 31 +-
tests/python/relax/test_bind_symbolic_vars.py | 28 +-
tests/python/relax/test_binding_rewrite.py | 50 +-
tests/python/relax/test_blockbuilder_core.py | 51 +-
tests/python/relax/test_blockbuilder_emit_te.py | 16 +-
tests/python/relax/test_codegen_coreml.py | 54 +-
tests/python/relax/test_codegen_cublas.py | 50 +-
tests/python/relax/test_codegen_cudnn.py | 34 +-
tests/python/relax/test_codegen_cutlass.py | 166 +--
tests/python/relax/test_codegen_dnnl.py | 22 +-
tests/python/relax/test_codegen_hipblas.py | 26 +-
tests/python/relax/test_codegen_tensorrt.py | 100 +-
tests/python/relax/test_contrib_vllm.py | 42 +-
tests/python/relax/test_dataflow_inplace.py | 51 +-
tests/python/relax/test_dataflow_pattern.py | 82 +-
tests/python/relax/test_dataflow_rewriter.py | 58 +-
tests/python/relax/test_dlpack_integration.py | 62 +-
tests/python/relax/test_e2e_op_dynamic.py | 43 +-
tests/python/relax/test_expr.py | 130 +--
tests/python/relax/test_expr_functor.py | 27 +-
tests/python/relax/test_fast_math_transform.py | 10 +-
tests/python/relax/test_frontend_common.py | 24 +-
tests/python/relax/test_frontend_dynamo.py | 47 +-
.../relax/test_frontend_from_exported_program.py | 580 +++++------
tests/python/relax/test_frontend_from_fx.py | 444 ++++----
tests/python/relax/test_frontend_nn_debug.py | 10 +-
tests/python/relax/test_frontend_nn_exporter.py | 17 +-
.../python/relax/test_frontend_nn_extern_module.py | 28 +-
tests/python/relax/test_frontend_nn_jit.py | 10 +-
.../relax/test_frontend_nn_llm_kernel_config.py | 48 +-
...test_frontend_nn_llm_sequence_prefill_masked.py | 64 +-
tests/python/relax/test_frontend_nn_modules.py | 16 +-
tests/python/relax/test_frontend_nn_mutator.py | 4 +-
tests/python/relax/test_frontend_nn_op.py | 119 +--
tests/python/relax/test_frontend_nn_packing.py | 10 +-
.../relax/test_frontend_nn_parameter_containers.py | 14 +-
tests/python/relax/test_frontend_nn_subroutines.py | 16 +-
tests/python/relax/test_frontend_nn_tensor.py | 16 +-
tests/python/relax/test_frontend_onnx.py | 385 +++----
tests/python/relax/test_frontend_onnx_backend.py | 8 +-
tests/python/relax/test_frontend_stablehlo.py | 40 +-
tests/python/relax/test_frontend_tflite.py | 438 ++++----
tests/python/relax/test_group_gemm_flashinfer.py | 40 +-
tests/python/relax/test_inline_functions.py | 34 +-
tests/python/relax/test_kill_after_last_use.py | 26 +-
tests/python/relax/test_lower_alloc_tensor.py | 12 +-
.../relax/test_meta_schedule_relax_integration.py | 18 +-
tests/python/relax/test_nn_llm_prefill_parallel.py | 28 +-
tests/python/relax/test_op_binary.py | 33 +-
tests/python/relax/test_op_ccl.py | 15 +-
tests/python/relax/test_op_create.py | 19 +-
tests/python/relax/test_op_datatype.py | 21 +-
tests/python/relax/test_op_distributed.py | 12 +-
tests/python/relax/test_op_grad.py | 12 +-
tests/python/relax/test_op_gradient_numeric.py | 124 +--
tests/python/relax/test_op_image.py | 23 +-
tests/python/relax/test_op_index.py | 51 +-
tests/python/relax/test_op_inspect.py | 66 +-
tests/python/relax/test_op_linear_algebra.py | 15 +-
tests/python/relax/test_op_manipulate.py | 59 +-
tests/python/relax/test_op_misc.py | 28 +-
tests/python/relax/test_op_nn.py | 17 +-
tests/python/relax/test_op_nn_convolution.py | 67 +-
tests/python/relax/test_op_nn_pooling.py | 113 +--
tests/python/relax/test_op_qdq.py | 15 +-
tests/python/relax/test_op_sampling.py | 10 +-
tests/python/relax/test_op_search.py | 15 +-
tests/python/relax/test_op_set.py | 15 +-
tests/python/relax/test_op_size.py | 26 +-
tests/python/relax/test_op_sort.py | 15 +-
tests/python/relax/test_op_statistical.py | 15 +-
tests/python/relax/test_op_take.py | 78 +-
tests/python/relax/test_op_ternary.py | 15 +-
tests/python/relax/test_op_unary.py | 15 +-
tests/python/relax/test_op_view.py | 110 +-
tests/python/relax/test_op_vision.py | 175 ++--
.../python/relax/test_optimize_layout_transform.py | 18 +-
tests/python/relax/test_pipeline.py | 62 +-
tests/python/relax/test_pytorch_integration.py | 14 +-
tests/python/relax/test_relax_operators.py | 226 ++---
.../python/relax/test_relax_to_pyfunc_converter.py | 18 +-
.../python/relax/test_remove_redundant_reshape.py | 10 +-
tests/python/relax/test_runtime_builtin.py | 92 +-
...runtime_builtin_paged_attention_kv_cache_cpu.py | 112 +--
..._builtin_paged_attention_kv_cache_flashinfer.py | 48 +-
...me_builtin_paged_attention_kv_cache_metadata.py | 94 +-
...ltin_paged_attention_kv_cache_mla_flashinfer.py | 74 +-
...ime_builtin_paged_attention_kv_cache_mla_tir.py | 88 +-
...runtime_builtin_paged_attention_kv_cache_tir.py | 84 +-
.../python/relax/test_runtime_builtin_rnn_state.py | 82 +-
.../relax/test_runtime_sampling_flashinfer.py | 22 +-
.../relax/test_tensorrt_engine_serialization.py | 59 +-
tests/python/relax/test_testing_nn.py | 18 +-
tests/python/relax/test_tir_call_source_kernel.py | 28 +-
tests/python/relax/test_training_append_loss.py | 10 +-
tests/python/relax/test_training_loss.py | 10 +-
tests/python/relax/test_training_optimizer.py | 10 +-
.../relax/test_training_optimizer_numeric.py | 23 +-
tests/python/relax/test_training_setup_trainer.py | 12 +-
.../python/relax/test_training_trainer_numeric.py | 32 +-
tests/python/relax/test_transform.py | 84 +-
.../relax/test_transform_allocate_workspace.py | 12 +-
tests/python/relax/test_transform_alter_op_impl.py | 16 +-
.../test_transform_annotate_tir_op_pattern.py | 42 +-
...st_transform_attach_attr_layout_free_tensors.py | 25 +-
.../relax/test_transform_attach_global_symbol.py | 25 +-
tests/python/relax/test_transform_bind_params.py | 72 +-
.../relax/test_transform_bind_symbolic_vars.py | 46 +-
.../relax/test_transform_bundle_model_params.py | 54 +-
.../relax/test_transform_canonicalize_bindings.py | 32 +-
tests/python/relax/test_transform_codegen_pass.py | 63 +-
.../test_transform_combine_parallel_matmul.py | 69 +-
.../relax/test_transform_compute_prim_value.py | 22 +-
.../relax/test_transform_convert_dataflow.py | 36 +-
.../python/relax/test_transform_convert_layout.py | 28 +-
tests/python/relax/test_transform_cse.py | 28 +-
.../relax/test_transform_dead_code_elimination.py | 46 +-
tests/python/relax/test_transform_decompose_ops.py | 27 +-
.../relax/test_transform_error_enrichment.py | 18 +-
.../relax/test_transform_expand_matmul_of_sum.py | 10 +-
.../relax/test_transform_expand_tuple_args.py | 16 +-
.../test_transform_fold_batch_norm_to_conv2d.py | 46 +-
tests/python/relax/test_transform_fold_constant.py | 104 +-
tests/python/relax/test_transform_fuse_ops.py | 42 +-
.../relax/test_transform_fuse_ops_by_pattern.py | 66 +-
tests/python/relax/test_transform_fuse_tir.py | 24 +-
.../relax/test_transform_fuse_transpose_matmul.py | 22 +-
tests/python/relax/test_transform_gradient.py | 16 +-
.../relax/test_transform_gradient_checkpoint.py | 12 +-
.../relax/test_transform_gradient_numeric.py | 36 +-
.../relax/test_transform_gradient_te_register.py | 21 +-
.../test_transform_inline_private_functions.py | 18 +-
.../relax/test_transform_ipc_allreduce_rewrite.py | 22 +-
tests/python/relax/test_transform_lambda_lift.py | 28 +-
.../relax/test_transform_lazy_transform_params.py | 64 +-
tests/python/relax/test_transform_legalize_ops.py | 72 +-
.../relax/test_transform_legalize_ops_binary.py | 320 +++---
.../relax/test_transform_legalize_ops_ccl.py | 30 +-
.../test_transform_legalize_ops_create_datatype.py | 186 ++--
.../test_transform_legalize_ops_distributed.py | 18 +-
.../relax/test_transform_legalize_ops_grad.py | 40 +-
.../relax/test_transform_legalize_ops_image.py | 34 +-
..._transform_legalize_ops_index_linear_algebra.py | 134 +--
.../test_transform_legalize_ops_manipulate.py | 212 ++--
.../python/relax/test_transform_legalize_ops_nn.py | 392 ++++----
.../relax/test_transform_legalize_ops_qdq.py | 78 +-
...st_transform_legalize_ops_search_statistical.py | 162 +--
.../relax/test_transform_legalize_ops_unary.py | 26 +-
.../relax/test_transform_lift_transform_params.py | 92 +-
.../test_transform_lower_gpu_ipc_alloc_storage.py | 14 +-
.../test_transform_merge_composite_functions.py | 48 +-
.../test_transform_meta_schedule_apply_database.py | 21 +-
.../relax/test_transform_meta_schedule_tuning.py | 34 +-
tests/python/relax/test_transform_normalize.py | 46 +-
.../relax/test_transform_normalize_global_var.py | 21 +-
...st_transform_operator_specific_normalization.py | 44 +-
.../python/relax/test_transform_realize_vdevice.py | 16 +-
.../relax/test_transform_remove_unused_outputs.py | 20 +-
.../test_transform_remove_unused_parameters.py | 22 +-
..._transform_reorder_permute_dims_after_concat.py | 10 +-
.../test_transform_reorder_take_after_matmul.py | 12 +-
.../relax/test_transform_rewrite_cuda_graph.py | 48 +-
.../test_transform_rewrite_dataflow_reshape.py | 64 +-
...nsform_specialize_function_based_on_callsite.py | 22 +-
.../test_transform_split_layout_rewrite_preproc.py | 20 +-
.../test_transform_static_plan_block_memory.py | 128 +--
.../relax/test_transform_to_mixed_precision.py | 48 +-
.../relax/test_transform_topological_sort.py | 32 +-
.../relax/test_transform_update_param_type.py | 10 +-
.../python/relax/test_transform_update_vdevice.py | 14 +-
tests/python/relax/test_type.py | 30 +-
tests/python/relax/test_utils.py | 22 +-
.../relax/test_vm_alloc_storage_with_scope.py | 26 +-
tests/python/relax/test_vm_build.py | 397 ++++----
tests/python/relax/test_vm_builtin.py | 28 +-
tests/python/relax/test_vm_builtin_lower.py | 18 +-
tests/python/relax/test_vm_callback_function.py | 26 +-
tests/python/relax/test_vm_codegen_only.py | 152 +--
tests/python/relax/test_vm_cuda_graph.py | 42 +-
tests/python/relax/test_vm_execbuilder.py | 84 +-
tests/python/relax/test_vm_instrument.py | 16 +-
tests/python/relax/test_vm_multi_device.py | 52 +-
tests/python/relax/texture/test_texture_nd.py | 42 +-
tests/python/runtime/rpc/test_rpc_base.py | 4 +-
tests/python/runtime/rpc/test_runtime_rpc.py | 60 +-
...st_s_tir_analysis_calculate_allocated_memory.py | 16 +-
.../test_s_tir_analysis_estimate_tir_flops.py | 10 +-
.../test_s_tir_analysis_identify_memcpy.py | 38 +-
.../test_s_tir_analysis_is_pure_function.py | 8 +-
.../s_tir/analysis/test_s_tir_analysis_oob.py | 18 +-
.../test_s_tir_analysis_verify_block_scope.py | 50 +-
.../s_tir/analysis/test_sblock_access_region.py | 100 +-
.../analysis/test_sblock_tensor_access_lca.py | 16 +-
.../python/s_tir/base/test_compilation_pipeline.py | 22 +-
.../python/s_tir/base/test_s_tir_texture_scope.py | 24 +-
.../s_tir/base/test_sblock_dependence_info.py | 16 +-
tests/python/s_tir/base/test_tir_data_layout.py | 8 +-
.../s_tir/base/test_tir_te_extern_function.py | 74 +-
tests/python/s_tir/dlight/test_benchmark.py | 18 +-
tests/python/s_tir/dlight/test_cpu_gemv.py | 50 +-
tests/python/s_tir/dlight/test_cpu_reduction.py | 19 +-
tests/python/s_tir/dlight/test_gpu_conv.py | 16 +-
tests/python/s_tir/dlight/test_gpu_fallback.py | 16 +-
tests/python/s_tir/dlight/test_gpu_gemv.py | 64 +-
.../s_tir/dlight/test_gpu_general_reduction.py | 20 +-
.../python/s_tir/dlight/test_gpu_low_batch_gemv.py | 34 +-
tests/python/s_tir/dlight/test_gpu_matmul.py | 44 +-
.../s_tir/dlight/test_gpu_matmul_tensorize.py | 24 +-
tests/python/s_tir/dlight/test_gpu_reduction.py | 14 +-
tests/python/s_tir/dlight/test_gpu_rmsnorm.py | 14 +-
tests/python/s_tir/dlight/test_gpu_transpose.py | 12 +-
tests/python/s_tir/dlight/test_primitives.py | 14 +-
.../meta_schedule/test_meta_schedule_arg_info.py | 4 +-
.../meta_schedule/test_meta_schedule_builder.py | 14 +-
.../meta_schedule/test_meta_schedule_cost_model.py | 14 +-
.../meta_schedule/test_meta_schedule_database.py | 56 +-
.../test_meta_schedule_feature_extractor.py | 4 +-
...schedule_feature_extractor_per_store_feature.py | 27 +-
.../test_meta_schedule_measure_callback.py | 10 +-
...eta_schedule_mutator_mutate_compute_location.py | 6 +-
.../test_meta_schedule_mutator_mutate_parallel.py | 6 +-
..._meta_schedule_mutator_mutate_thread_binding.py | 6 +-
.../test_meta_schedule_mutator_mutate_tile_size.py | 6 +-
.../test_meta_schedule_mutator_mutate_unroll.py | 6 +-
.../test_meta_schedule_post_order_apply.py | 22 +-
...ule_postproc_disallow_async_strided_mem_copy.py | 12 +-
...meta_schedule_postproc_disallow_dynamic_loop.py | 14 +-
..._schedule_postproc_rewrite_cooperative_fetch.py | 22 +-
.../test_meta_schedule_postproc_rewrite_layout.py | 38 +-
...e_postproc_rewrite_parallel_vectorize_unroll.py | 16 +-
...ta_schedule_postproc_rewrite_reduction_block.py | 20 +-
...est_meta_schedule_postproc_rewrite_tensorize.py | 22 +-
...meta_schedule_postproc_rewrite_unbound_block.py | 36 +-
.../test_meta_schedule_postproc_verify_gpu_code.py | 22 +-
...est_meta_schedule_postproc_verify_vtcm_limit.py | 18 +-
.../meta_schedule/test_meta_schedule_runner.py | 28 +-
...test_meta_schedule_schedule_rule_add_rfactor.py | 6 +-
...eta_schedule_schedule_rule_apply_custom_rule.py | 10 +-
.../test_meta_schedule_schedule_rule_auto_bind.py | 6 +-
...test_meta_schedule_schedule_rule_auto_inline.py | 50 +-
...chedule_schedule_rule_cross_thread_reduction.py | 10 +-
.../test_meta_schedule_schedule_rule_mlt.py | 10 +-
.../test_meta_schedule_schedule_rule_mlt_intrin.py | 8 +-
...dule_schedule_rule_parallel_vectorize_unroll.py | 14 +-
...hedule_schedule_rule_random_compute_location.py | 10 +-
.../test_meta_schedule_search_strategy.py | 24 +-
.../meta_schedule/test_meta_schedule_space_cpu.py | 6 +-
.../meta_schedule/test_meta_schedule_space_cuda.py | 6 +-
.../test_meta_schedule_space_cuda_async.py | 6 +-
...ule_space_cuda_async_multiple_initialization.py | 4 +-
.../test_meta_schedule_space_generator.py | 14 +-
.../test_meta_schedule_space_post_opt.py | 12 +-
.../test_meta_schedule_task_scheduler.py | 40 +-
.../test_meta_schedule_trace_apply.py | 54 +-
.../test_meta_schedule_tune_context.py | 14 +-
.../meta_schedule/test_meta_schedule_tune_tir.py | 16 +-
.../s_tir/schedule/test_tir_schedule_analysis.py | 22 +-
.../test_tir_schedule_annotate_tensor_access.py | 12 +-
.../schedule/test_tir_schedule_block_scope.py | 13 +-
.../s_tir/schedule/test_tir_schedule_blockize.py | 26 +-
.../schedule/test_tir_schedule_cache_index.py | 16 +-
.../schedule/test_tir_schedule_cache_read_write.py | 14 +-
.../s_tir/schedule/test_tir_schedule_compute_at.py | 23 +-
.../schedule/test_tir_schedule_compute_inline.py | 18 +-
.../test_tir_schedule_decompose_padding.py | 20 +-
.../s_tir/schedule/test_tir_schedule_error.py | 13 +-
.../s_tir/schedule/test_tir_schedule_for_kind.py | 12 +-
.../test_tir_schedule_fuse_reduction_epilogue.py | 26 +-
...ir_schedule_fuse_reduction_epilogue_clipping.py | 24 +-
...st_tir_schedule_fuse_reduction_epilogue_relu.py | 28 +-
.../schedule/test_tir_schedule_instruction.py | 4 +-
.../s_tir/schedule/test_tir_schedule_merge.py | 12 +-
.../s_tir/schedule/test_tir_schedule_pad_einsum.py | 10 +-
.../s_tir/schedule/test_tir_schedule_partition.py | 15 +-
.../schedule/test_tir_schedule_read_write_at.py | 12 +-
.../s_tir/schedule/test_tir_schedule_reduction.py | 24 +-
.../s_tir/schedule/test_tir_schedule_reindex.py | 16 +-
.../s_tir/schedule/test_tir_schedule_reorder.py | 12 +-
.../test_tir_schedule_reorder_block_iter_var.py | 14 +-
.../s_tir/schedule/test_tir_schedule_rfactor.py | 13 +-
.../schedule/test_tir_schedule_rolling_buffer.py | 21 +-
.../s_tir/schedule/test_tir_schedule_sampling.py | 12 +-
.../s_tir/schedule/test_tir_schedule_set_dtype.py | 14 +-
.../s_tir/schedule/test_tir_schedule_set_scope.py | 26 +-
.../s_tir/schedule/test_tir_schedule_split_fuse.py | 15 +-
.../s_tir/schedule/test_tir_schedule_state.py | 62 +-
.../test_tir_schedule_state_cached_flags.py | 13 +-
.../schedule/test_tir_schedule_storage_align.py | 14 +-
.../s_tir/schedule/test_tir_schedule_tensorize.py | 13 +-
.../test_tir_schedule_tensorize_mfma_numeric.py | 30 +-
.../s_tir/schedule/test_tir_schedule_trace.py | 15 +-
.../s_tir/schedule/test_tir_schedule_transform.py | 18 +-
.../schedule/test_tir_schedule_transform_layout.py | 58 +-
.../s_tir/schedule/test_tir_schedule_utilities.py | 26 +-
.../schedule/test_tir_unsafe_hide_tensor_access.py | 12 +-
.../s_tir/script/test_s_tir_script_basic_usage.py | 142 +--
.../s_tir/script/test_s_tir_script_blocks.py | 220 ++--
.../script/test_s_tir_script_dynamic_shape.py | 40 +-
.../script/test_s_tir_script_error_handling.py | 50 +-
.../s_tir/script/test_s_tir_script_ir_builder.py | 63 +-
.../script/test_s_tir_script_meta_programming.py | 72 +-
.../s_tir/script/test_s_tir_script_namespace.py | 78 +-
.../s_tir/script/test_s_tir_script_printer.py | 289 +++---
.../script/test_s_tir_script_source_locations.py | 12 +-
tests/python/s_tir/test_arith_domain_touched.py | 14 +-
tests/python/s_tir/test_s_tir_renew_defs.py | 50 +-
...test_s_tir_transform_annotate_irregular_loop.py | 33 +-
.../test_s_tir_transform_canonicalize_loop.py | 27 +-
.../test_s_tir_transform_compact_tensor_region.py | 33 +-
...est_s_tir_transform_convert_blocks_to_opaque.py | 19 +-
.../test_s_tir_transform_default_gpu_schedule.py | 108 +-
.../test_s_tir_transform_inject_permuted_layout.py | 14 +-
.../test_s_tir_transform_inject_ptx_async_copy.py | 86 +-
...est_s_tir_transform_inject_software_pipeline.py | 37 +-
.../test_s_tir_transform_inject_virtual_thread.py | 34 +-
.../test_s_tir_transform_lift_thread_binding.py | 13 +-
.../test_s_tir_transform_loop_partition.py | 106 +-
...s_tir_transform_lower_cross_thread_reduction.py | 26 +-
.../test_s_tir_transform_lower_init_block.py | 18 +-
.../test_s_tir_transform_lower_match_tensor.py | 18 +-
.../test_s_tir_transform_lower_opaque_block.py | 24 +-
...test_s_tir_transform_lower_thread_all_reduce.py | 82 +-
.../test_s_tir_transform_lower_thread_binding.py | 20 +-
...transform_manifest_shared_memory_local_stage.py | 16 +-
...st_s_tir_transform_memhammer_lower_auto_copy.py | 52 +-
...form_merge_dynamic_shared_memory_allocations.py | 18 +-
.../test_s_tir_transform_mma_tensor_layout.py | 25 +-
...sform_plan_update_tensor_allocation_location.py | 17 +-
.../transform/test_s_tir_transform_remove_undef.py | 20 +-
...transform_remove_weight_layout_rewrite_block.py | 12 +-
.../test_s_tir_transform_rewrite_unsafe_select.py | 14 +-
.../transform/test_s_tir_transform_thread_sync.py | 46 +-
tests/python/script/test_constructor_contracts.py | 15 +-
tests/python/sym/test_sym_canonical_simplify.py | 249 ++---
tests/python/sym/test_sym_modular_set.py | 100 +-
tests/python/target/test_arm_target.py | 40 +-
tests/python/target/test_target_target.py | 112 +--
tests/python/te/test_te_create_function.py | 85 +-
tests/python/te/test_te_tag.py | 4 +-
tests/python/te/test_te_tensor.py | 42 +-
tests/python/tirx/base/test_tirx_base.py | 48 +-
.../base/test_tirx_binding_type_propagation.py | 31 +-
tests/python/tirx/base/test_tirx_constructor.py | 298 +++---
tests/python/tirx/base/test_tirx_ops.py | 128 +--
tests/python/tirx/codegen/test_codegen_nvshmem.py | 22 +-
tests/python/tirx/codegen/test_tirx_codegen_fp8.py | 16 +-
.../tirx/script/test_tirx_script_ir_builder.py | 137 +--
.../python/tirx/script/test_tirx_script_printer.py | 87 +-
.../tirx/test_tirx_kernels_registry_correctness.py | 2 +-
tirx/python/tvm_tirx/backend/_autoload_backends.py | 6 +-
tirx/python/tvm_tirx/backend/adreno/__init__.py | 2 +-
tirx/python/tvm_tirx/backend/adreno/target_tags.py | 2 +-
tirx/python/tvm_tirx/backend/config.py | 4 +-
tirx/python/tvm_tirx/backend/cuda/__init__.py | 8 +-
tirx/python/tvm_tirx/backend/cuda/axis.py | 2 +-
.../python/tvm_tirx/backend/cuda/backend_config.py | 4 +-
.../tvm_tirx/backend/cuda/codegen/__init__.py | 4 +-
.../tvm_tirx/backend/cuda/codegen/registry.py | 2 +-
.../python/tvm_tirx/backend/cuda/codegen/schema.py | 2 +-
tirx/python/tvm_tirx/backend/cuda/compiler.py | 2 +-
tirx/python/tvm_tirx/backend/cuda/cpp/__init__.py | 4 +-
tirx/python/tvm_tirx/backend/cuda/cpp/asm.py | 8 +-
tirx/python/tvm_tirx/backend/cuda/cpp/builtins.py | 2 +-
.../python/tvm_tirx/backend/cuda/cpp/instrument.py | 8 +-
tirx/python/tvm_tirx/backend/cuda/host.py | 2 +-
tirx/python/tvm_tirx/backend/cuda/iket.py | 18 +-
.../tvm_tirx/backend/cuda/lang/alloc_pool.py | 14 +-
tirx/python/tvm_tirx/backend/cuda/lang/clc.py | 2 +-
tirx/python/tvm_tirx/backend/cuda/lang/pipeline.py | 8 +-
.../tvm_tirx/backend/cuda/lang/tile_scheduler.py | 6 +-
tirx/python/tvm_tirx/backend/cuda/launch/_impl.py | 6 +-
tirx/python/tvm_tirx/backend/cuda/layout.py | 12 +-
tirx/python/tvm_tirx/backend/cuda/op.py | 14 +-
tirx/python/tvm_tirx/backend/cuda/ptx/engine.py | 18 +-
.../tvm_tirx/backend/cuda/ptx/gen_coverage.py | 2 +-
.../tvm_tirx/backend/cuda/ptx/gen_helpers.py | 8 +-
tirx/python/tvm_tirx/backend/cuda/ptx/gen_stubs.py | 8 +-
tirx/python/tvm_tirx/backend/cuda/script.py | 6 +-
tirx/python/tvm_tirx/backend/cuda/target_tags.py | 2 +-
tirx/python/tvm_tirx/backend/cuda/thread_scope.py | 2 +-
.../tvm_tirx/backend/cuda/tile/exec_context.py | 2 +-
.../backend/cuda/tile/patterns/__init__.py | 6 +-
.../backend/cuda/tile/patterns/async_copy.py | 8 +-
.../tvm_tirx/backend/cuda/tile/patterns/common.py | 14 +-
.../backend/cuda/tile/patterns/descriptors.py | 10 +-
.../tvm_tirx/backend/cuda/tile/patterns/ewise.py | 8 +-
.../backend/cuda/tile/patterns/ldmatrix.py | 4 +-
.../tvm_tirx/backend/cuda/tile/patterns/memory.py | 2 +-
.../tvm_tirx/backend/cuda/tile/patterns/mma.py | 6 +-
.../backend/cuda/tile/patterns/register.py | 6 +-
.../backend/cuda/tile/patterns/tcgen05_mma.py | 16 +-
.../tvm_tirx/backend/cuda/tile/patterns/tma.py | 8 +-
.../tvm_tirx/backend/cuda/tile/patterns/tmem.py | 8 +-
.../backend/cuda/tile/patterns/tmem_copy.py | 6 +-
.../tvm_tirx/backend/cuda/transforms/__init__.py | 6 +-
tirx/python/tvm_tirx/backend/hexagon/linker.py | 6 +-
.../python/tvm_tirx/backend/hexagon/target_tags.py | 2 +-
tirx/python/tvm_tirx/backend/loader.py | 24 +-
tirx/python/tvm_tirx/backend/metal/__init__.py | 8 +-
tirx/python/tvm_tirx/backend/metal/op.py | 10 +-
tirx/python/tvm_tirx/backend/metal/script.py | 4 +-
tirx/python/tvm_tirx/backend/metal/target_tags.py | 4 +-
tirx/python/tvm_tirx/backend/opencl/__init__.py | 10 +-
tirx/python/tvm_tirx/backend/opencl/script.py | 4 +-
tirx/python/tvm_tirx/backend/rocm/__init__.py | 6 +-
tirx/python/tvm_tirx/backend/trn/__init__.py | 4 +-
tirx/python/tvm_tirx/backend/trn/axis.py | 2 +-
tirx/python/tvm_tirx/backend/trn/layout.py | 2 +-
tirx/python/tvm_tirx/backend/trn/op.py | 4 +-
tirx/python/tvm_tirx/backend/trn/pipeline.py | 16 +-
tirx/python/tvm_tirx/backend/trn/script.py | 4 +-
tirx/python/tvm_tirx/backend/trn/target_tags.py | 2 +-
tirx/python/tvm_tirx/backend/trn/thread_scope.py | 2 +-
.../tvm_tirx/backend/trn/tile/binary/default.py | 10 +-
.../tvm_tirx/backend/trn/tile/binary/utils.py | 12 +-
tirx/python/tvm_tirx/backend/trn/tile/common.py | 4 +-
.../backend/trn/tile/compose_op/binary_chain.py | 12 +-
.../backend/trn/tile/compose_op/binary_reduce.py | 12 +-
.../backend/trn/tile/compose_op/unary_reduce.py | 18 +-
.../tvm_tirx/backend/trn/tile/compose_op/utils.py | 2 +-
.../tvm_tirx/backend/trn/tile/copy/default.py | 12 +-
tirx/python/tvm_tirx/backend/trn/tile/dim_utils.py | 6 +-
.../tvm_tirx/backend/trn/tile/gemm/default.py | 18 +-
.../backend/trn/tile/instruction_generator.py | 12 +-
.../tvm_tirx/backend/trn/tile/instructions.py | 14 +-
.../tvm_tirx/backend/trn/tile/private_alloc.py | 18 +-
.../tvm_tirx/backend/trn/tile/reduction/utils.py | 16 +-
.../tvm_tirx/backend/trn/tile/select/default.py | 16 +-
.../tvm_tirx/backend/trn/tile/unary/default.py | 8 +-
.../tvm_tirx/backend/trn/tile/unary/utils.py | 16 +-
.../backend/trn/tile/unary/with_bias_scale.py | 10 +-
.../tvm_tirx/backend/trn/tile/workspace_utils.py | 4 +-
.../tvm_tirx/backend/trn/transform/__init__.py | 6 +-
.../backend/trn/transform/naive_allocator.py | 8 +-
.../backend/trn/transform/private_tensor_alloc.py | 12 +-
tirx/python/tvm_tirx/backend/vulkan/__init__.py | 6 +-
tirx/python/tvm_tirx/codegen.py | 2 +-
tirx/python/tvm_tirx/ir/_constant.py | 2 +-
tirx/python/tvm_tirx/ir/_ffi_api.py | 2 +-
tirx/python/tvm_tirx/ir/_overload_prim_expr.py | 6 +-
tirx/python/tvm_tirx/ir/attrs.py | 8 +-
tirx/python/tvm_tirx/ir/base.py | 2 +-
tirx/python/tvm_tirx/ir/expr.py | 54 +-
tirx/python/tvm_tirx/ir/function.py | 10 +-
tirx/python/tvm_tirx/ir/global_info.py | 2 +-
tirx/python/tvm_tirx/ir/location.py | 4 +-
tirx/python/tvm_tirx/ir/module.py | 4 +-
tirx/python/tvm_tirx/ir/op.py | 4 +-
tirx/python/tvm_tirx/ir/prim/expr.py | 4 +-
tirx/python/tvm_tirx/ir/prim/op.py | 6 +-
tirx/python/tvm_tirx/ir/stmt.py | 6 +-
tirx/python/tvm_tirx/ir/supply.py | 2 +-
.../tvm_tirx/ir/transform/_ffi_instrument_api.py | 2 +-
.../tvm_tirx/ir/transform/_ffi_transform_api.py | 2 +-
tirx/python/tvm_tirx/ir/transform/core.py | 36 +-
tirx/python/tvm_tirx/ir/transform/instrument.py | 2 +-
tirx/python/tvm_tirx/ir/type.py | 16 +-
tirx/python/tvm_tirx/runtime/_ffi_api.py | 2 +-
tirx/python/tvm_tirx/runtime/_tensor.py | 12 +-
tirx/python/tvm_tirx/runtime/executable.py | 8 +-
tirx/python/tvm_tirx/runtime/module.py | 24 +-
tirx/python/tvm_tirx/script/__init__.py | 46 +-
tirx/python/tvm_tirx/script/_ffi_api.py | 2 +-
tirx/python/tvm_tirx/script/ir_builder/__init__.py | 4 +-
tirx/python/tvm_tirx/script/ir_builder/_ffi_api.py | 2 +-
tirx/python/tvm_tirx/script/ir_builder/base.py | 12 +-
tirx/python/tvm_tirx/script/ir_builder/frame.py | 4 +-
tirx/python/tvm_tirx/script/ir_builder/ir.py | 20 +-
.../tvm_tirx/script/ir_builder/parser_protocol.py | 14 +-
tirx/python/tvm_tirx/script/ir_builder/stmt.py | 46 +-
tirx/python/tvm_tirx/script/ir_builder/tracing.py | 2 +-
tirx/python/tvm_tirx/script/parser/entry.py | 20 +-
.../tvm_tirx/script/parser/inspect_source.py | 2 +-
tirx/python/tvm_tirx/script/parser/ir.py | 2 +-
.../tvm_tirx/script/parser/protocol_registry.py | 2 +-
tirx/python/tvm_tirx/script/parser/transpile.py | 2 +-
tirx/python/tvm_tirx/script/printer/__init__.py | 2 +-
tirx/python/tvm_tirx/script/printer/_ffi_api.py | 2 +-
tirx/python/tvm_tirx/script/printer/doc.py | 4 +-
tirx/python/tvm_tirx/script/printer/doc_printer.py | 2 +-
tirx/python/tvm_tirx/script/printer/scriptable.py | 8 +-
tirx/python/tvm_tirx/script/tirx.pyi | 2 +-
tirx/python/tvm_tirx/support/__init__.py | 6 +-
tirx/python/tvm_tirx/support/cc.py | 2 +-
tirx/python/tvm_tirx/support/clang.py | 4 +-
tirx/python/tvm_tirx/support/libinfo.py | 2 +-
tirx/python/tvm_tirx/support/nvcc.py | 24 +-
tirx/python/tvm_tirx/support/popen_pool.py | 2 +-
tirx/python/tvm_tirx/support/popen_worker.py | 2 +-
tirx/python/tvm_tirx/support/rocm.py | 12 +-
tirx/python/tvm_tirx/support/tar.py | 2 +-
tirx/python/tvm_tirx/sym/_ffi_api.py | 2 +-
tirx/python/tvm_tirx/sym/analyzer.py | 16 +-
tirx/python/tvm_tirx/sym/bound.py | 6 +-
tirx/python/tvm_tirx/sym/int_set.py | 18 +-
tirx/python/tvm_tirx/sym/iter_affine_map.py | 24 +-
tirx/python/tvm_tirx/sym/pattern.py | 4 +-
tirx/python/tvm_tirx/sym/utils.py | 2 +-
tirx/python/tvm_tirx/target/__init__.py | 4 +-
tirx/python/tvm_tirx/target/_ffi_api.py | 2 +-
tirx/python/tvm_tirx/target/intrin.py | 2 +-
tirx/python/tvm_tirx/target/target.py | 4 +-
tirx/python/tvm_tirx/testing/_ffi_api.py | 2 +-
tirx/python/tvm_tirx/testing/env.py | 44 +-
tirx/python/tvm_tirx/testing/popen_pool.py | 4 +-
tirx/python/tvm_tirx/testing/utils.py | 32 +-
tirx/python/tvm_tirx/tirx/__init__.py | 22 +-
tirx/python/tvm_tirx/tirx/_compile.py | 56 +-
tirx/python/tvm_tirx/tirx/_ffi_api.py | 2 +-
tirx/python/tvm_tirx/tirx/_tensor_view.py | 70 +-
tirx/python/tvm_tirx/tirx/analysis/_ffi_api.py | 2 +-
tirx/python/tvm_tirx/tirx/analysis/analysis.py | 8 +-
tirx/python/tvm_tirx/tirx/bench.py | 22 +-
tirx/python/tvm_tirx/tirx/compilation_pipeline.py | 38 +-
tirx/python/tvm_tirx/tirx/expr.py | 18 +-
tirx/python/tvm_tirx/tirx/function.py | 46 +-
tirx/python/tvm_tirx/tirx/layout.py | 12 +-
tirx/python/tvm_tirx/tirx/op/__init__.py | 192 ++--
tirx/python/tvm_tirx/tirx/script/__init__.py | 6 +-
.../tvm_tirx/tirx/script/ir_builder/__init__.py | 6 +-
.../tirx/script/ir_builder/external_kernel.py | 14 +-
.../tvm_tirx/tirx/script/ir_builder/frame.py | 16 +-
tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py | 34 +-
tirx/python/tvm_tirx/tirx/script/ir_builder/op.py | 62 +-
.../tirx/script/ir_builder/parser_protocol.py | 96 +-
.../tvm_tirx/tirx/script/ir_builder/triton.py | 8 +-
.../tvm_tirx/tirx/script/ir_builder/utils.py | 4 +-
tirx/python/tvm_tirx/tirx/script/jit.py | 20 +-
tirx/python/tvm_tirx/tirx/tensor.py | 36 +-
tirx/python/tvm_tirx/tirx/thread_scope.py | 2 +-
tirx/python/tvm_tirx/tirx/tile/dispatch.py | 8 +-
tirx/python/tvm_tirx/tirx/tile/instruction.py | 22 +-
tirx/python/tvm_tirx/tirx/tile/views.py | 4 +-
tirx/python/tvm_tirx/tirx/transform/_ffi_api.py | 2 +-
tirx/python/tvm_tirx/tirx/transform/common.py | 6 +-
.../tvm_tirx/tirx/transform/function_pass.py | 12 +-
tirx/python/tvm_tirx/tirx/transform/transform.py | 68 +-
tirx/python/tvm_tirx/tirx/type.py | 6 +-
tirx/src/script/printer/printer.cc | 2 +-
.../test_validate_runtime_library.py | 6 +-
.../tests/python/backend/cuda/codegen/test_cuda.py | 382 +++----
.../python/backend/cuda/codegen/test_fastmath.py | 44 +-
tirx/tests/python/backend/cuda/codegen/test_fp4.py | 44 +-
tirx/tests/python/backend/cuda/codegen/test_fp8.py | 144 +--
.../python/backend/cuda/test_cuda_host_bundle.py | 84 +-
.../python/backend/cuda/test_cuda_module_export.py | 54 +-
.../python/backend/hexagon/codegen/test_hexagon.py | 26 +-
.../python/backend/metal/codegen/test_metal.py | 126 +--
.../backend/metal/test_metal_compile_callback.py | 28 +-
.../backend/metal/test_metal_module_export.py | 14 +-
.../python/backend/opencl/codegen/test_opencl.py | 68 +-
.../backend/opencl/test_opencl_module_export.py | 24 +-
.../tests/python/backend/rocm/codegen/test_rocm.py | 66 +-
.../python/backend/rocm/test_rocm_module_export.py | 24 +-
.../python/backend/vulkan/codegen/test_vulkan.py | 248 ++---
.../backend/vulkan/test_vulkan_module_export.py | 14 +-
.../python/backend/webgpu/codegen/test_webgpu.py | 36 +-
tirx/tests/python/codegen/c/test_c.py | 106 +-
.../python/codegen/gpu_common/test_allreduce.py | 54 +-
.../tests/python/codegen/gpu_common/test_device.py | 40 +-
.../codegen/gpu_common/test_gpu_intrinsics.py | 56 +-
tirx/tests/python/codegen/llvm/test_aarch64.py | 72 +-
tirx/tests/python/codegen/llvm/test_arm.py | 12 +-
tirx/tests/python/codegen/llvm/test_llvm.py | 288 +++---
tirx/tests/python/codegen/llvm/test_riscv.py | 32 +-
.../python/codegen/llvm/test_scalable_vector.py | 42 +-
tirx/tests/python/codegen/llvm/test_x86.py | 12 +-
.../python/codegen/test_argument_validation.py | 120 +--
tirx/tests/python/codegen/test_assert.py | 26 +-
tirx/tests/python/codegen/test_bool.py | 34 +-
tirx/tests/python/codegen/test_common.py | 48 +-
tirx/tests/python/codegen/test_extern.py | 38 +-
tirx/tests/python/codegen/test_static_init.py | 12 +-
tirx/tests/python/ir/test_ir_attrs.py | 60 +-
tirx/tests/python/ir/test_ir_json.py | 68 +-
tirx/tests/python/ir/test_ir_module.py | 12 +-
.../python/ir/test_ir_prim_expr_deep_equal.py | 22 +-
tirx/tests/python/ir/test_ir_type.py | 44 +-
tirx/tests/python/ir/test_nonnullable_ir.py | 24 +-
.../python/ir/test_roundtrip_runtime_module.py | 10 +-
.../python/runtime/test_runtime_device_api.py | 8 +-
.../python/runtime/test_runtime_memory_scope.py | 30 +-
tirx/tests/python/runtime/test_runtime_module.py | 54 +-
tirx/tests/python/runtime/test_runtime_tensor.py | 32 +-
tirx/tests/python/script/conftest.py | 8 +-
tirx/tests/python/script/minilang.py | 8 +-
tirx/tests/python/script/test_basic_usage.py | 28 +-
.../python/script/test_builder_void_emission.py | 12 +-
.../python/script/test_cross_function_call.py | 2 +-
tirx/tests/python/script/test_dynamic_shape.py | 2 +-
tirx/tests/python/script/test_error_handling.py | 8 +-
tirx/tests/python/script/test_ir_builder.py | 12 +-
.../tests/python/script/test_jit_specialization.py | 10 +-
tirx/tests/python/script/test_meta_programming.py | 6 +-
tirx/tests/python/script/test_parser_entry.py | 2 +-
.../tests/python/script/test_script_printer_doc.py | 8 +-
.../python/script/test_script_printer_entry.py | 16 +-
tirx/tests/python/script/test_source_locations.py | 8 +-
.../python/script/test_special_parser_protocol.py | 6 +-
tirx/tests/python/support/test_ccache.py | 6 +-
tirx/tests/python/support/test_nvcc.py | 6 +-
tirx/tests/python/support/test_popen_pool.py | 8 +-
tirx/tests/python/support/test_util.py | 2 +-
tirx/tests/python/sym/test_sym_analyzer_object.py | 116 +--
tirx/tests/python/sym/test_sym_const_int_bound.py | 146 +--
tirx/tests/python/sym/test_sym_deduce_bound.py | 202 ++--
.../tests/python/sym/test_sym_detect_clip_bound.py | 54 +-
.../python/sym/test_sym_detect_linear_equation.py | 52 +-
tirx/tests/python/sym/test_sym_intset.py | 256 ++---
tirx/tests/python/sym/test_sym_iter_affine_map.py | 714 ++++++-------
tirx/tests/python/sym/test_sym_rewrite_simplify.py | 1062 ++++++++++----------
tirx/tests/python/sym/test_sym_simplify.py | 62 +-
tirx/tests/python/sym/test_sym_z3.py | 64 +-
.../tests/python/target/test_llvm_features_info.py | 22 +-
tirx/tests/python/target/test_riscv_features.py | 8 +-
.../python/target/test_target_parser_mprofile.py | 4 +-
tirx/tests/python/target/test_x86_features.py | 18 +-
tirx/tests/python/testing/test_env.py | 76 +-
.../test_tirx_analysis_check_well_defined.py | 142 +--
.../analysis/test_tirx_analysis_undefined_vars.py | 48 +-
tirx/tests/python/tirx/base/test_tirx_host_func.py | 18 +-
.../tests/python/tirx/base/test_tirx_imm_values.py | 62 +-
tirx/tests/python/tirx/base/test_tirx_index_map.py | 94 +-
tirx/tests/python/tirx/base/test_tirx_intrin.py | 180 ++--
tirx/tests/python/tirx/base/test_tirx_nodes.py | 308 +++---
tirx/tests/python/tirx/base/test_tirx_op_types.py | 44 +-
.../python/tirx/base/test_tirx_ptx_cp_async.py | 20 +-
.../tirx/base/test_tirx_ptx_griddepcontrol.py | 20 +-
.../tests/python/tirx/base/test_tirx_ptx_mma_sp.py | 44 +-
.../tirx/base/test_tirx_ptx_scalar_f32_math.py | 30 +-
.../tirx/base/test_tirx_scalable_datatype.py | 14 +-
.../tests/python/tirx/base/test_tirx_specialize.py | 42 +-
.../python/tirx/base/test_tirx_tensor_type.py | 120 +--
.../python/tirx/base/test_tirx_texture_scope.py | 14 +-
tirx/tests/python/tirx/codegen/conftest.py | 2 +-
.../python/tirx/codegen/test_codegen_ampere.py | 30 +-
.../python/tirx/codegen/test_codegen_blackwell.py | 104 +-
.../tests/python/tirx/codegen/test_codegen_cuda.py | 150 +--
.../python/tirx/codegen/test_codegen_dsmem.py | 34 +-
.../python/tirx/codegen/test_codegen_hopper.py | 180 ++--
tirx/tests/python/tirx/codegen/test_codegen_nki.py | 14 +-
.../python/tirx/codegen/test_cuda_cta_reduce.py | 18 +-
.../python/tirx/codegen/test_cuda_wait_until.py | 10 +-
.../python/tirx/codegen/test_cuda_warp_reduce.py | 18 +-
tirx/tests/python/tirx/codegen/test_ptx_addr.py | 40 +-
tirx/tests/python/tirx/codegen/test_ptx_cvt.py | 4 +-
tirx/tests/python/tirx/codegen/test_ptx_dialect.py | 334 +++---
.../python/tirx/codegen/test_ptx_ld_st_ops.py | 28 +-
.../codegen/test_tcgen05_descriptor_encoders.py | 22 +-
tirx/tests/python/tirx/conftest.py | 2 +-
.../python/tirx/iket/iket_profile_workload.py | 14 +-
.../python/tirx/iket/test_iket_orchestration.py | 4 +-
tirx/tests/python/tirx/iket/test_iket_profiler.py | 88 +-
tirx/tests/python/tirx/layout_fixtures.py | 8 +-
.../tirx/script/test_tirx_script_basic_usage.py | 74 +-
.../tirx/script/test_tirx_script_dynamic_shape.py | 36 +-
.../tirx/script/test_tirx_script_error_handling.py | 8 +-
.../python/tirx/script/test_tirx_script_gpu.py | 16 +-
.../script/test_tirx_script_jit_specialization.py | 4 +-
.../script/test_tirx_script_meta_programming.py | 40 +-
.../script/test_tirx_script_printer_highlight.py | 8 +-
...test_tirx_script_printer_invisible_path_info.py | 16 +-
.../test_tirx_script_printer_structural_equal.py | 10 +-
.../script/test_tirx_script_printer_underlining.py | 14 +-
.../script/test_tirx_script_source_locations.py | 48 +-
.../python/tirx/script/test_tirx_script_tensors.py | 114 +--
.../python/tirx/script/test_tirx_script_tracing.py | 20 +-
tirx/tests/python/tirx/test_alloc_pool.py | 4 +-
tirx/tests/python/tirx/test_bench_utils.py | 8 +-
.../tests/python/tirx/test_compilation_pipeline.py | 16 +-
tirx/tests/python/tirx/test_control_flow.py | 18 +-
tirx/tests/python/tirx/test_cuda_backend_config.py | 164 +--
tirx/tests/python/tirx/test_cuda_launch.py | 60 +-
tirx/tests/python/tirx/test_exec_context.py | 2 +-
tirx/tests/python/tirx/test_inline.py | 4 +-
tirx/tests/python/tirx/test_jit.py | 20 +-
tirx/tests/python/tirx/test_layout.py | 38 +-
tirx/tests/python/tirx/test_op.py | 22 +-
.../tests/python/tirx/test_op_namespace_cleanup.py | 84 +-
.../tests/python/tirx/test_roundtrip_namespaces.py | 8 +-
.../python/tirx/test_tensor_data_reinfer_type.py | 6 +-
tirx/tests/python/tirx/test_thread_scope.py | 34 +-
tirx/tests/python/tirx/test_tile_namespace.py | 56 +-
tirx/tests/python/tirx/test_tirx_axis_enum.py | 14 +-
tirx/tests/python/tirx/test_verifier.py | 12 +-
tirx/tests/python/tirx/tile/cuda/conftest.py | 2 +-
.../python/tirx/tile/cuda/test_add_sub_mul_div.py | 166 +--
tirx/tests/python/tirx/tile/cuda/test_cp_async.py | 34 +-
.../python/tirx/tile/cuda/test_cp_async_bulk.py | 62 +-
.../tirx/tile/cuda/test_cp_async_bulk_tensor.py | 156 +--
tirx/tests/python/tirx/tile/cuda/test_cvt.py | 162 +--
tirx/tests/python/tirx/tile/cuda/test_ewise.py | 74 +-
.../tirx/tile/cuda/test_explicit_patterns.py | 70 +-
tirx/tests/python/tirx/tile/cuda/test_fma.py | 84 +-
.../tirx/tile/cuda/test_ld_st_cache_policy.py | 30 +-
.../python/tirx/tile/cuda/test_ld_st_gmem_smem.py | 68 +-
.../python/tirx/tile/cuda/test_ld_st_registers.py | 182 ++--
.../tirx/tile/cuda/test_ldmatrix_stmatrix.py | 52 +-
tirx/tests/python/tirx/tile/cuda/test_mma_sync.py | 116 +--
tirx/tests/python/tirx/tile/cuda/test_mov.py | 32 +-
.../python/tirx/tile/cuda/test_pattern_groups.py | 48 +-
.../python/tirx/tile/cuda/test_pattern_math.py | 24 +-
.../python/tirx/tile/cuda/test_pattern_mma.py | 66 +-
.../tirx/tile/cuda/test_pattern_transfers.py | 44 +-
.../tests/python/tirx/tile/cuda/test_tcgen05_cp.py | 74 +-
.../python/tirx/tile/cuda/test_tcgen05_ld_st.py | 202 ++--
.../python/tirx/tile/cuda/test_tcgen05_mma.py | 308 +++---
.../tirx/tile/cuda/test_tcgen05_mma_codegen.py | 20 +-
.../tirx/tile/cuda/test_tcgen05_mma_ws_split_n.py | 18 +-
tirx/tests/python/tirx/tile/test_dispatcher.py | 28 +-
tirx/tests/python/tirx/tile/trn/test_binary_trn.py | 52 +-
.../python/tirx/tile/trn/test_compose_op_trn.py | 118 +--
tirx/tests/python/tirx/tile/trn/test_copy_trn.py | 110 +-
tirx/tests/python/tirx/tile/trn/test_gemm_trn.py | 88 +-
.../python/tirx/tile/trn/test_private_alloc_trn.py | 38 +-
.../python/tirx/tile/trn/test_reduction_trn.py | 56 +-
tirx/tests/python/tirx/tile/trn/test_select_trn.py | 50 +-
tirx/tests/python/tirx/tile/trn/test_unary_trn.py | 54 +-
.../transform/test_tirx_transform_bf16_legalize.py | 100 +-
.../test_tirx_transform_common_subexpr_elim.py | 182 ++--
.../transform/test_tirx_transform_convert_ssa.py | 66 +-
.../test_tirx_transform_decorate_device_scope.py | 14 +-
.../test_tirx_transform_flatten_tensor.py | 60 +-
...est_tirx_transform_force_narrow_index_to_i32.py | 94 +-
.../transform/test_tirx_transform_fp8_legalize.py | 40 +-
.../transform/test_tirx_transform_function_pass.py | 26 +-
.../tirx/transform/test_tirx_transform_helpers.py | 64 +-
.../tirx/transform/test_tirx_transform_hoist_if.py | 86 +-
.../transform/test_tirx_transform_lower_intrin.py | 44 +-
.../transform/test_tirx_transform_lower_tirx.py | 170 ++--
.../test_tirx_transform_lower_tvm_builtin.py | 68 +-
.../test_tirx_transform_make_packed_api.py | 82 +-
.../test_tirx_transform_naive_allocator.py | 18 +-
.../test_tirx_transform_narrow_datatype.py | 48 +-
...st_tirx_transform_pointer_value_type_rewrite.py | 32 +-
.../transform/test_tirx_transform_remove_no_op.py | 64 +-
...est_tirx_transform_renormalize_split_pattern.py | 22 +-
.../tirx/transform/test_tirx_transform_simplify.py | 98 +-
.../test_tirx_transform_split_host_device.py | 124 +--
.../test_tirx_transform_storage_rewrite.py | 80 +-
.../transform/test_tirx_transform_unroll_loop.py | 52 +-
.../transform/test_tirx_transform_vectorize.py | 156 +--
web/tests/python/prepare_test_libs.py | 16 +-
web/tests/python/relax_rpc_test.py | 12 +-
web/tests/python/webgpu_rpc_test.py | 14 +-
1159 files changed, 20492 insertions(+), 20342 deletions(-)
diff --git a/apps/android_rpc/tests/android_rpc_test.py b/apps/android_rpc/tests/android_rpc_test.py
index e9997a14e5..7b41d41917 100644
--- a/apps/android_rpc/tests/android_rpc_test.py
+++ b/apps/android_rpc/tests/android_rpc_test.py
@@ -26,9 +26,9 @@ import os
import numpy as np
-import tvm
+import tvm, tvm_tirx
from tvm import rpc, te
-from tvm.support import ndk, utils
+from tvm_tirx.support import ndk, utils
# Set to be address of tvm proxy.
tracker_host = os.environ["TVM_TRACKER_HOST"]
@@ -48,7 +48,7 @@ test_vulkan = False
def test_rpc_module():
# graph
- n = tvm.runtime.convert(1024)
+ n = tvm_tirx.runtime.convert(1024)
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
a_np = np.random.uniform(size=1024).astype(A.dtype)
@@ -58,7 +58,7 @@ def test_rpc_module():
tracker = rpc.connect_tracker(tracker_host, tracker_port)
remote = tracker.request(key, priority=0, session_timeout=60)
- mod = tvm.IRModule.from_expr(te.create_function([A, B]).with_attr("global_symbol", "myadd"))
+ mod = tvm_tirx.IRModule.from_expr(te.create_function([A, B]).with_attr("global_symbol", "myadd"))
sch = tvm.s_tir.Schedule(mod)
(x,) = sch.get_loops(block=sch.get_sblock("B"))
xo, xi = sch.split(i, [None, 32])
@@ -66,7 +66,7 @@ def test_rpc_module():
sch.bind(xi, "threadIdx.x")
if test_opencl:
- f = tvm.compile(sch.mod, target=tvm.target.Target("opencl", host=target))
+ f = tvm.compile(sch.mod, target=tvm_tirx.target.Target("opencl", host=target))
path_dso_cl = temp.relpath("dev_lib_cl.so")
f.export_library(path_dso_cl, fcompile=ndk.create_shared)
diff --git a/apps/ios_rpc/tests/ios_rpc_test.py b/apps/ios_rpc/tests/ios_rpc_test.py
index e5865b937d..fdabecab12 100644
--- a/apps/ios_rpc/tests/ios_rpc_test.py
+++ b/apps/ios_rpc/tests/ios_rpc_test.py
@@ -24,9 +24,9 @@ import argparse
import numpy as np
-import tvm
+import tvm, tvm_tirx
from tvm import rpc, te
-from tvm.support import utils, xcode
+from tvm_tirx.support import utils, xcode
# Change target configuration, this is setting for iphone6s
arch = "arm64"
@@ -37,18 +37,18 @@ MODES = {"proxy": rpc.connect, "tracker": rpc.connect_tracker, "standalone": rpc
# override metal compiler to compile to iphone
-@tvm.register_global_func("tvm_callback_metal_compile")
+@tvm_tirx.register_global_func("tvm_callback_metal_compile")
def compile_metal(src, target):
return xcode.compile_metal(src, sdk=sdk)
def test_rpc_module(host, port, key, mode):
# graph
- n = tvm.runtime.convert(1024)
+ n = tvm_tirx.runtime.convert(1024)
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
temp = utils.tempdir()
- mod = tvm.IRModule.from_expr(te.create_function([A, B]).with_attr("global_symbol", "myadd"))
+ mod = tvm_tirx.IRModule.from_expr(te.create_function([A, B]).with_attr("global_symbol", "myadd"))
sch = tvm.s_tir.Schedule(mod)
(i,) = sch.get_loops(block=sch.get_sblock("B"))
i0, i1 = sch.split(i, [None, 32])
@@ -57,7 +57,7 @@ def test_rpc_module(host, port, key, mode):
# Build the dynamic lib.
# If we don't want to do metal and only use cpu, just set target to be target
- f = tvm.compile(sch.mod, target=tvm.target.Target("metal", host=target))
+ f = tvm.compile(sch.mod, target=tvm_tirx.target.Target("metal", host=target))
path_dso1 = temp.relpath("dev_lib.dylib")
f.export_library(path_dso1, fcompile=xcode.create_dylib, arch=arch, sdk=sdk)
diff --git a/jvm/core/src/test/scripts/prepare_test_libs.py b/jvm/core/src/test/scripts/prepare_test_libs.py
index 3b7f8206ca..67df2fcacc 100644
--- a/jvm/core/src/test/scripts/prepare_test_libs.py
+++ b/jvm/core/src/test/scripts/prepare_test_libs.py
@@ -19,24 +19,24 @@
import os
import sys
-import tvm
+import tvm, tvm_tirx
from tvm import relax, te
-from tvm.script import ir as I
-from tvm.script import relax as R
+from tvm_tirx.script import ir as I
+from tvm_tirx.script import relax as R
def prepare_relax_lib(base_path):
pipeline = relax.get_pipeline()
n = I.dynamic("n", "int64")
- @tvm.script.ir_module
+ @tvm_tirx.script.ir_module
class Mod:
@R.function
def main(x: R.Tensor([n], "float32"), y: R.Tensor([n], "float32")):
lv0 = R.add(x, y)
return lv0
- target = tvm.target.Target("llvm")
+ target = tvm_tirx.target.Target("llvm")
mod = pipeline(Mod)
ex = relax.build(mod, target)
@@ -46,34 +46,34 @@ def prepare_relax_lib(base_path):
def prepare_cpu_lib(base_path):
target = "llvm"
- if not tvm.runtime.enabled(target):
+ if not tvm_tirx.runtime.enabled(target):
raise RuntimeError(f"Target {target} is not enbaled")
n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.placeholder((n,), name="B")
C = te.compute(A.shape, lambda *i: A(*i) + B(*i), name="C")
- mod = tvm.IRModule.from_expr(te.create_function([A, B, C]).with_attr("global_symbol", "myadd"))
- fadd = tvm.tirx.compile(mod, target).mod
+ mod = tvm_tirx.IRModule.from_expr(te.create_function([A, B, C]).with_attr("global_symbol", "myadd"))
+ fadd = tvm_tirx.tirx.compile(mod, target).mod
lib_path = os.path.join(base_path, "add_cpu.so")
fadd.export_library(lib_path)
def prepare_gpu_lib(base_path):
- if not tvm.cuda().exist:
+ if not tvm_tirx.cuda().exist:
print("CUDA is not enabled, skip the generation")
return
n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.placeholder((n,), name="B")
C = te.compute(A.shape, lambda *i: A(*i) + B(*i), name="C")
- mod = tvm.IRModule.from_expr(te.create_function([A, B, C]).with_attr("global_symbol", "myadd"))
+ mod = tvm_tirx.IRModule.from_expr(te.create_function([A, B, C]).with_attr("global_symbol", "myadd"))
sch = tvm.s_tir.Schedule(mod)
sch.work_on("myadd")
(i,) = sch.get_loops(block=sch.get_sblock("C"))
i0, i1 = sch.split(i, [None, 32])
sch.bind(i0, "blockIdx.x")
sch.bind(i1, "threadIdx.x")
- fadd = tvm.tirx.compile(sch.mod, "cuda").mod
+ fadd = tvm_tirx.tirx.compile(sch.mod, "cuda").mod
lib_path = os.path.join(base_path, "add_cuda.so")
fadd.export_library(lib_path)
diff --git a/python/tvm/contrib/cblas.py b/python/tvm/contrib/cblas.py
index d2dd8012a8..9697f3acb0 100644
--- a/python/tvm/contrib/cblas.py
+++ b/python/tvm/contrib/cblas.py
@@ -16,7 +16,7 @@
# under the License.
"""External function interface to BLAS libraries."""
-import tvm
+import tvm, tvm_tirx
from tvm import te
@@ -45,7 +45,7 @@ def matmul(lhs, rhs, transa=False, transb=False, **kwargs):
return te.extern(
(n, m),
[lhs, rhs],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.cblas.matmul", ins[0], ins[1], outs[0], transa, transb
),
name="C",
@@ -79,7 +79,7 @@ def batch_matmul(lhs, rhs, transa=False, transb=False, iterative=False, **kwargs
return te.extern(
(b, n, m),
[lhs, rhs],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.cblas.batch_matmul"
if not iterative
else "tvm.contrib.cblas.batch_matmul_iterative",
diff --git a/python/tvm/contrib/cublas.py b/python/tvm/contrib/cublas.py
index a999d70f2a..c6840ab008 100644
--- a/python/tvm/contrib/cublas.py
+++ b/python/tvm/contrib/cublas.py
@@ -16,7 +16,7 @@
# under the License.
"""External function interface to cuBLAS libraries."""
-import tvm
+import tvm, tvm_tirx
from tvm import te
@@ -45,7 +45,7 @@ def matmul(lhs, rhs, transa=False, transb=False, dtype=None):
return te.extern(
(n, m),
[lhs, rhs],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.cublas.matmul", ins[0], ins[1], outs[0], transa, transb
),
dtype=dtype,
@@ -79,7 +79,7 @@ def batch_matmul(lhs, rhs, transa=False, transb=False, dtype=None):
return te.extern(
(b, n, m),
[lhs, rhs],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.cublas.batch_matmul", ins[0], ins[1], outs[0], transa, transb
),
dtype=dtype,
diff --git a/python/tvm/contrib/cublaslt.py b/python/tvm/contrib/cublaslt.py
index 44cb9c9f0b..893ffc5366 100644
--- a/python/tvm/contrib/cublaslt.py
+++ b/python/tvm/contrib/cublaslt.py
@@ -16,7 +16,7 @@
# under the License.
"""External function interface to cuBLASlt libraries."""
-import tvm
+import tvm, tvm_tirx
from tvm import te
@@ -47,7 +47,7 @@ def matmul(lhs, rhs, transa=False, transb=False, n=0, m=0, dtype=None):
return te.extern(
(n, m),
[lhs, rhs],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.cublaslt.matmul", ins[0], ins[1], outs[0], transa, transb
),
dtype=dtype,
diff --git a/python/tvm/contrib/cudnn.py b/python/tvm/contrib/cudnn.py
index e048ff5b28..718cac7a24 100644
--- a/python/tvm/contrib/cudnn.py
+++ b/python/tvm/contrib/cudnn.py
@@ -23,7 +23,7 @@ import ctypes
import numpy as np
import tvm_ffi
-import tvm
+import tvm, tvm_tirx
from tvm import te
# algos can be read from cudnn.h
@@ -51,7 +51,7 @@ def exists():
True if CuDNN support is enabled and a CuDNN-capable GPU
exists. Otherwise, False.
"""
- func = tvm.get_global_func("tvm.contrib.cudnn.exists", allow_missing=True)
+ func = tvm_tirx.get_global_func("tvm.contrib.cudnn.exists", allow_missing=True)
if func is None:
return False
@@ -199,7 +199,7 @@ def conv_output_shape(
c_output = w_shape[0]
x_chan = x_shape[1]
w_chan_input = w_shape[1]
- w_lanes = tvm.runtime.DataType(conv_dtype).lanes
+ w_lanes = tvm_tirx.runtime.DataType(conv_dtype).lanes
assert w_lanes == 1
x_shape = x_shape[2:]
w_shape = w_shape[2:]
@@ -207,7 +207,7 @@ def conv_output_shape(
else:
raise ValueError(f"Unknown CuDNN tensor format: '{tensor_format}'")
- x_lanes = tvm.runtime.DataType(data_dtype).lanes
+ x_lanes = tvm_tirx.runtime.DataType(data_dtype).lanes
assert x_chan * x_lanes == w_chan_input * groups, (
f"Mismatched dimensions, data has {x_chan // groups} channels/group "
f"(dimension {x_chan} with {x_lanes} lanes/value, {groups} groups), "
@@ -573,7 +573,7 @@ def conv_forward(
x_shape = list(x.shape)
- if isinstance(x.shape[0], tvm.tirx.expr.IntImm):
+ if isinstance(x.shape[0], tvm_tirx.tirx.expr.IntImm):
oshape = conv_output_shape(
tensor_format,
pad,
@@ -628,7 +628,7 @@ def conv_forward(
return te.extern(
oshape,
[x, w],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.cudnn.conv2d.forward",
conv_mode,
tensor_format,
@@ -651,7 +651,7 @@ def conv_forward(
return te.extern(
oshape,
[x, w],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.cudnn.conv3d.forward",
conv_mode,
tensor_format,
@@ -723,7 +723,7 @@ def conv_backward_data(
conv_dtype = dy.dtype if conv_dtype is None else conv_dtype
pad, stride, dilation, _, _ = _prepare_global_func_params(dims - 2, pad, stride, dilation)
- assert isinstance(dy.shape[0], tvm.tirx.expr.IntImm), (
+ assert isinstance(dy.shape[0], tvm_tirx.tirx.expr.IntImm), (
"Dynamic batch is not supported for cudnn conv2d backwad data yet."
)
@@ -752,7 +752,7 @@ def conv_backward_data(
return te.extern(
dx_shape,
[dy, w],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.cudnn.conv2d.backward_data",
conv_mode,
tensor_format,
@@ -817,7 +817,7 @@ def conv_backward_filter(
x_shape = list(x.shape)
- assert isinstance(x.shape[0], tvm.tirx.expr.IntImm), (
+ assert isinstance(x.shape[0], tvm_tirx.tirx.expr.IntImm), (
"Dynamic batch is not supported for cudnn conv2d backwad filter yet."
)
@@ -853,7 +853,7 @@ def conv_backward_filter(
return te.extern(
dw_shape,
[dy, x],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.cudnn.conv2d.backward_filter",
conv_mode,
tensor_format,
@@ -893,7 +893,7 @@ def softmax(x, axis=-1):
return te.extern(
x.shape,
[x],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.cudnn.softmax.forward", ins[0], outs[0], axis
),
name="y",
@@ -919,7 +919,7 @@ def log_softmax(x, axis=-1):
return te.extern(
x.shape,
[x],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.cudnn.log_softmax.forward", ins[0], outs[0], axis
),
name="y",
diff --git a/python/tvm/contrib/cutlass/build.py b/python/tvm/contrib/cutlass/build.py
index d9806d8166..aeea9e5b0a 100644
--- a/python/tvm/contrib/cutlass/build.py
+++ b/python/tvm/contrib/cutlass/build.py
@@ -28,9 +28,10 @@ from functools import reduce
from tvm_ffi import register_global_func
-import tvm
-from tvm import relax, runtime
-from tvm.support.nvcc import get_cuda_version
+import tvm, tvm_tirx
+from tvm_tirx import runtime
+from tvm import relax
+from tvm_tirx.support.nvcc import get_cuda_version
from tvm.topi.utils import get_const_tuple
from .gen_conv2d import CutlassConv2DProfiler
@@ -42,7 +43,7 @@ logger = logging.getLogger("cutlass")
def has_cutlass():
"""Returns true if the CUTLASS custom codegen is available"""
- return tvm.get_global_func("relax.ext.cutlass", True) is not None
+ return tvm_tirx.get_global_func("relax.ext.cutlass", True) is not None
def _get_cutlass_path():
@@ -108,7 +109,7 @@ def select_gemm_kernel(
):
"""Run CUTLASS profiler to select the best kernel, or return the default one for dynamic
workloads."""
- if any(isinstance(s, tvm.tirx.Any) for s in [MM, KK, NN]):
+ if any(isinstance(s, tvm_tirx.tirx.Any) for s in [MM, KK, NN]):
out = cutlass_profiler.get_default(
op_type, out_dtype, arg0_dtype, arg1_dtype, use_3xtf32, batched=batched
)
@@ -249,7 +250,7 @@ def handle_conv2d(
else:
conv_kind = ConvKind.Fprop
- if any(isinstance(s, tvm.tirx.Any) for s in d_shape):
+ if any(isinstance(s, tvm_tirx.tirx.Any) for s in d_shape):
out = cutlass_profiler.get_default(
op_type, out_dtype, data_dtype, weight_dtype, use_3xtf32, conv_kind, strides
)
@@ -399,14 +400,14 @@ def _extract_relax_function_signature(f):
def _extract_arg_idx(pattern_name, f):
- extract_func = tvm.get_global_func("relax.contrib.extract_arg_idx")
+ extract_func = tvm_tirx.get_global_func("relax.contrib.extract_arg_idx")
arg_indices = extract_func(pattern_name, f)
return {k: int(v) for k, v in arg_indices.items()}
def is_shape_valid_for_cutlass_matmul(
- lhs_shape: Sequence[tvm.ir.Expr],
- rhs_shape: Sequence[tvm.ir.Expr],
+ lhs_shape: Sequence[tvm_tirx.ir.Expr],
+ rhs_shape: Sequence[tvm_tirx.ir.Expr],
) -> bool:
"""
Check whether the shape of inputs can be handled by CUTLASS GEMM.
@@ -417,7 +418,7 @@ def is_shape_valid_for_cutlass_matmul(
as well as ND x 2D and 2D x ND. For example, it cannot handle matmul with shape
(2, 1, 4, 8) x (2, 3, 8, 16), because the batch stride of lhs is not constant.
"""
- if not isinstance(lhs_shape[-1], tvm.tirx.expr.IntImm | int):
+ if not isinstance(lhs_shape[-1], tvm_tirx.tirx.expr.IntImm | int):
# Reduction axis must be constant
return False
@@ -427,7 +428,7 @@ def is_shape_valid_for_cutlass_matmul(
# This could be regular matmul or batch matmul with shape ND x 2D or 2D x ND
return True
- analyzer = tvm.sym.Analyzer()
+ analyzer = tvm_tirx.sym.Analyzer()
# If one side has less dimensions, use 1 to fill the gap
batch_dim_pairs = list(
itertools.zip_longest(
@@ -836,7 +837,7 @@ def profile_relax_function(functions, options):
for f in functions:
annotator = CutlassRelaxFunctionAnnotator(
- tvm.IRModule.from_expr(f), conv2d_profiler, gemm_profiler, options
+ tvm_tirx.IRModule.from_expr(f), conv2d_profiler, gemm_profiler, options
)
annotated_functions.append(annotator.visit_expr(f))
@@ -876,7 +877,7 @@ def compile_cutlass_module(c_source_module, options):
c_source_module.export_library(lib_path, workspace_dir=tmp_dir, **compile_options)
# Recover static library
- return tvm.runtime.load_static_library(lib_path, function_names)
+ return tvm_tirx.runtime.load_static_library(lib_path, function_names)
def finalize_modules(lib, lib_path="compile.so", tmp_dir="./tmp"):
diff --git a/python/tvm/contrib/cutlass/gen_tensor_op.py b/python/tvm/contrib/cutlass/gen_tensor_op.py
index 477c1ee449..c5e1d2ec1e 100644
--- a/python/tvm/contrib/cutlass/gen_tensor_op.py
+++ b/python/tvm/contrib/cutlass/gen_tensor_op.py
@@ -28,8 +28,8 @@ import tempfile
import tvm_ffi
-from tvm.runtime import Object
-from tvm.tirx import IntImm
+from tvm_tirx.runtime import Object
+from tvm_tirx.tirx import IntImm
from . import _ffi_api as ffi
from .attention_operation import (
@@ -539,7 +539,7 @@ def instantiate_template(func_name, annotations, func_args):
attrs["M"] = annotations["M"]
attrs["group_size"] = annotations["group_size"]
- if not isinstance(attrs["M"], tvm.tirx.IntImm):
+ if not isinstance(attrs["M"], tvm_tirx.tirx.IntImm):
attrs["M"] = get_flattened_batch_dim(
func_args[lhs_arg_idx], int(annotations["batch_rank"])
)
@@ -888,7 +888,7 @@ def instantiate_template(func_name, annotations, func_args):
attrs = {"input": func_args[0], "gamma": func_args[1], "beta": func_args[2]}
attrs.update(dict(annotations))
- if not isinstance(attrs["M"], tvm.tirx.IntImm):
+ if not isinstance(attrs["M"], tvm_tirx.tirx.IntImm):
attrs["M"] = get_flattened_batch_dim(func_args[0], int(attrs["batch_rank"]))
code = instantiate_layer_norm_template(attrs)
@@ -899,7 +899,7 @@ def instantiate_template(func_name, annotations, func_args):
attrs = {"input": func_args[0], "weight": func_args[1]}
attrs.update(dict(annotations))
- if not isinstance(attrs["M"], tvm.tirx.IntImm):
+ if not isinstance(attrs["M"], tvm_tirx.tirx.IntImm):
attrs["M"] = get_flattened_batch_dim(func_args[0], int(attrs["batch_rank"]))
code = instantiate_rms_norm_template(attrs)
diff --git a/python/tvm/contrib/cutlass/library.py b/python/tvm/contrib/cutlass/library.py
index 82a9b6e71f..b3127a4ce9 100644
--- a/python/tvm/contrib/cutlass/library.py
+++ b/python/tvm/contrib/cutlass/library.py
@@ -22,7 +22,7 @@ import enum
import re
from enum import auto as enum_auto
-from tvm.tirx.expr import FloatImm, IntImm
+from tvm_tirx.tirx.expr import FloatImm, IntImm
class GeneratorTarget(enum.Enum):
diff --git a/python/tvm/contrib/dlpack.py b/python/tvm/contrib/dlpack.py
index 0eb2a87875..2b30e8385b 100644
--- a/python/tvm/contrib/dlpack.py
+++ b/python/tvm/contrib/dlpack.py
@@ -16,7 +16,7 @@
# under the License.
"""Wrapping functions to bridge frameworks with DLPack support to TVM"""
-import tvm.runtime
+import tvm_tirx.runtime, tvm
def convert_func(tvm_func, tensor_type, to_dlpack_func):
@@ -38,7 +38,7 @@ def convert_func(tvm_func, tensor_type, to_dlpack_func):
def _wrapper(*args):
args = tuple(
- tvm.runtime.from_dlpack(to_dlpack_func(arg)) if isinstance(arg, tensor_type) else arg
+ tvm_tirx.runtime.from_dlpack(to_dlpack_func(arg)) if isinstance(arg, tensor_type) else arg
for arg in args
)
return tvm_func(*args)
diff --git a/python/tvm/contrib/dnnl.py b/python/tvm/contrib/dnnl.py
index eebd7e1660..38afaad543 100644
--- a/python/tvm/contrib/dnnl.py
+++ b/python/tvm/contrib/dnnl.py
@@ -16,7 +16,7 @@
# under the License.
"""External function interface to BLAS libraries."""
-import tvm
+import tvm, tvm_tirx
from tvm import te
from ..topi.nn.utils import get_pad_tuple
@@ -47,7 +47,7 @@ def matmul(lhs, rhs, transa=False, transb=False, **kwargs):
return te.extern(
(n, m),
[lhs, rhs],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.dnnl.matmul", ins[0], ins[1], outs[0], transa, transb
),
name="C",
@@ -142,7 +142,7 @@ def dnnl_conv2d(
return te.extern(
out_shape,
[src, weights],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.dnnl.conv2d",
ins[0],
ins[1],
diff --git a/python/tvm/contrib/emcc.py b/python/tvm/contrib/emcc.py
index 296f9ee9a2..b1b50bee56 100644
--- a/python/tvm/contrib/emcc.py
+++ b/python/tvm/contrib/emcc.py
@@ -21,7 +21,7 @@ import os
import subprocess
from pathlib import Path
-from tvm import libinfo
+from tvm_tirx import libinfo
def find_wasm_lib(name, optional=False):
diff --git a/python/tvm/contrib/hexagon/generate_take_op.py b/python/tvm/contrib/hexagon/generate_take_op.py
index c77f08c81f..28ac81c5c8 100644
--- a/python/tvm/contrib/hexagon/generate_take_op.py
+++ b/python/tvm/contrib/hexagon/generate_take_op.py
@@ -16,8 +16,8 @@
# under the License.
# pylint: disable=missing-docstring, invalid-name, unnecessary-comprehension, unused-argument
-import tvm
-import tvm.testing
+import tvm, tvm_tirx
+import tvm_tirx.testing, tvm
from tvm import relax
from tvm.contrib.hexagon import hexagon_unary_ops
@@ -25,7 +25,7 @@ from tvm.contrib.hexagon import hexagon_unary_ops
def op_replace(call_node, func) -> bool:
if not isinstance(call_node, relax.Call):
return False
- call_tir_op = tvm.ir.Op.get("relax.call_tir")
+ call_tir_op = tvm_tirx.ir.Op.get("relax.call_tir")
if call_node.op != call_tir_op:
return False
ops = [
@@ -46,11 +46,11 @@ def op_replace(call_node, func) -> bool:
@relax.expr_functor.mutator
class Tanh2TakeReplace(tvm.relax.PyExprMutator):
- def __init__(self, mod: tvm.IRModule) -> None:
+ def __init__(self, mod: tvm_tirx.IRModule) -> None:
super().__init__(mod)
self.mod_ = mod
- def transform(self) -> tvm.IRModule:
+ def transform(self) -> tvm_tirx.IRModule:
# Iterate over all the nodes to check for the node replaceable
for global_var, func in self.mod_.functions.items():
# Skip non-relax functions
@@ -63,7 +63,7 @@ class Tanh2TakeReplace(tvm.relax.PyExprMutator):
return self.builder_.get()
def visit_call_(self, call_node: relax.Call) -> relax.Call:
- call_tir_op = tvm.ir.Op.get("relax.call_tir")
+ call_tir_op = tvm_tirx.ir.Op.get("relax.call_tir")
if call_node.op != call_tir_op:
return call_node
@@ -83,14 +83,14 @@ class Tanh2TakeReplace(tvm.relax.PyExprMutator):
take_func_gv = self.builder_.add_func(take_func, "take")
take_node = relax.call_tir(
take_func_gv,
- relax.expr.Tuple([call_node.args[1][0], relax.const(tvm.runtime.tensor(LUT))]),
+ relax.expr.Tuple([call_node.args[1][0], relax.const(tvm_tirx.runtime.tensor(LUT))]),
ty_args=[call_node.ty],
)
return take_node
return call_node
-@tvm.ir.transform.module_pass(opt_level=2, name="replace_tanh_take")
+@tvm_tirx.ir.transform.module_pass(opt_level=2, name="replace_tanh_take")
class PassReplaceWithTakeOpFunctions:
def transform_module(self, mod, ctx):
return Tanh2TakeReplace(mod).transform()
diff --git a/python/tvm/contrib/hexagon/tools.py b/python/tvm/contrib/hexagon/tools.py
index c72caede2a..218e5efcad 100644
--- a/python/tvm/contrib/hexagon/tools.py
+++ b/python/tvm/contrib/hexagon/tools.py
@@ -29,18 +29,18 @@ import tarfile
import numpy
from tvm_ffi import register_global_func
-import tvm
-import tvm.support.cc as cc
-from tvm.backend.hexagon.linker import (
+import tvm, tvm_tirx
+import tvm_tirx.support.cc as cc
+from tvm_tirx.backend.hexagon.linker import (
HEXAGON_LINK_MAIN as HEXAGON_LINK_MAIN,
)
-from tvm.backend.hexagon.linker import (
+from tvm_tirx.backend.hexagon.linker import (
hexagon_link as hexagon_link,
)
-from tvm.backend.hexagon.linker import (
+from tvm_tirx.backend.hexagon.linker import (
link_shared as link_shared,
)
-from tvm.backend.hexagon.linker import (
+from tvm_tirx.backend.hexagon.linker import (
register_linker as register_linker,
)
@@ -113,7 +113,7 @@ def link_shared_macos(so_name, objs, extra_args=None):
----------
so_name : str
Name of the shared library file.
- objs : list[str, tvm.ir.StringImm]
+ objs : list[str, tvm_tirx.ir.StringImm]
extra_args : dict (str->str) or Map<String,String>
Additional arguments:
'hex_arch' - Hexagon architecture, e.g. v68
@@ -127,7 +127,7 @@ def link_shared_macos(so_name, objs, extra_args=None):
# The list of object files can be passed as built-in Python strings,
# or as tvm.ir.StringImm's.
def to_str(s):
- if isinstance(s, tvm.ir.StringImm):
+ if isinstance(s, tvm_tirx.ir.StringImm):
return s.value
assert isinstance(s, str), 'argument "' + str(s) + '" should be a string or StrImm'
return s
@@ -227,7 +227,7 @@ def create_aot_shared(so_name: str | pathlib.Path, files, hexagon_arch: str, opt
def pack_imports(
- module: tvm.runtime.Module,
+ module: tvm_tirx.runtime.Module,
is_system_lib: bool, # pylint: disable=unused-argument
c_symbol_prefix: str,
workspace_dir: str,
@@ -238,7 +238,7 @@ def pack_imports(
Parameters
----------
- module: tvm.runtime.Module
+ module: tvm_tirx.runtime.Module
Module whose imported modules need to be serialized.
is_system_lib: bool
Flag whether the exported module will be used as a system library.
@@ -255,12 +255,12 @@ def pack_imports(
path_bin = os.path.join(workspace_dir, "imports.bin")
pack_to_bin_f_name = "runtime.ModulePackImportsToTensor"
- fpack_to_bin = tvm.get_global_func(pack_to_bin_f_name)
+ fpack_to_bin = tvm_tirx.get_global_func(pack_to_bin_f_name)
assert fpack_to_bin, f"Expecting {pack_to_bin_f_name} in registry"
fpack_to_bin(module).numpy().tofile(path_bin)
- mblob_symbol = c_symbol_prefix + tvm.get_global_func("runtime.ModuleImportsBlobName")()
+ mblob_symbol = c_symbol_prefix + tvm_tirx.get_global_func("runtime.ModuleImportsBlobName")()
binary_size = os.path.getsize(path_bin)
hexagon_toolchain = os.environ.get("HEXAGON_TOOLCHAIN")
@@ -344,7 +344,7 @@ def allocate_hexagon_array(dev, tensor_shape=None, dtype=None, data=None, mem_sc
physical_shape = [numpy.prod(tensor_shape)]
- arr = tvm.runtime.empty(physical_shape, dtype=dtype, device=dev, mem_scope=mem_scope)
+ arr = tvm_tirx.runtime.empty(physical_shape, dtype=dtype, device=dev, mem_scope=mem_scope)
if data is not None:
arr.copyfrom(data.reshape(physical_shape))
diff --git a/python/tvm/contrib/hipblas.py b/python/tvm/contrib/hipblas.py
index 078d6fef6e..e0ad63f91e 100644
--- a/python/tvm/contrib/hipblas.py
+++ b/python/tvm/contrib/hipblas.py
@@ -16,7 +16,7 @@
# under the License.
"""External function interface to hipBLAS libraries."""
-import tvm
+import tvm, tvm_tirx
from tvm import te
@@ -45,7 +45,7 @@ def matmul(lhs, rhs, transa=False, transb=False, dtype=None):
return te.extern(
(n, m),
[lhs, rhs],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.hipblas.matmul", ins[0], ins[1], outs[0], transa, transb
),
dtype=dtype,
@@ -79,7 +79,7 @@ def batch_matmul(lhs, rhs, transa=False, transb=False, dtype=None):
return te.extern(
(b, n, m),
[lhs, rhs],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.hipblas.batch_matmul", ins[0], ins[1], outs[0], transa, transb
),
dtype=dtype,
diff --git a/python/tvm/contrib/mkl.py b/python/tvm/contrib/mkl.py
index cc223b5f08..cb565c5563 100644
--- a/python/tvm/contrib/mkl.py
+++ b/python/tvm/contrib/mkl.py
@@ -16,7 +16,7 @@
# under the License.
"""External function interface to BLAS libraries."""
-import tvm
+import tvm, tvm_tirx
from tvm import te
@@ -45,7 +45,7 @@ def matmul(lhs, rhs, transa=False, transb=False, **kwargs):
return te.extern(
(n, m),
[lhs, rhs],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.mkl.matmul", ins[0], ins[1], outs[0], transa, transb
),
name="C",
@@ -78,7 +78,7 @@ def matmul_u8s8s32(lhs, rhs, transa=False, transb=False, **kwargs):
return te.extern(
(n, m),
[lhs, rhs],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.mkl.matmul_u8s8s32", ins[0], ins[1], outs[0], transa, transb
),
name="C",
@@ -112,7 +112,7 @@ def batch_matmul(lhs, rhs, transa=False, transb=False, iterative=False, **kwargs
return te.extern(
(b, n, m),
[lhs, rhs],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.mkl.batch_matmul"
if not iterative
else "tvm.contrib.mkl.batch_matmul_iterative",
diff --git a/python/tvm/contrib/nnpack.py b/python/tvm/contrib/nnpack.py
index e6068db7c9..9b6cc30246 100644
--- a/python/tvm/contrib/nnpack.py
+++ b/python/tvm/contrib/nnpack.py
@@ -19,7 +19,7 @@
import tvm_ffi
-import tvm
+import tvm, tvm_tirx
from tvm import te
@@ -50,7 +50,7 @@ def fully_connected_inference(lhs, rhs, nthreads=1):
return te.extern(
(m,),
[lhs, rhs],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.nnpack.fully_connected_inference", ins[0], ins[1], outs[0], nthreads
),
name="C",
@@ -114,7 +114,7 @@ def convolution_inference(
return te.extern(
(batch, output_channels, output_height, output_width),
[data, kernel, bias] if bias is not None else [data, kernel],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.nnpack.convolution_inference",
ins[0],
ins[1],
@@ -177,7 +177,7 @@ def convolution_inference_without_weight_transform(
return te.extern(
(batch, output_channels, output_height, output_width),
[data, transformed_kernel, bias] if bias is not None else [data, transformed_kernel],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.nnpack.convolution_inference_without_weight_transform",
ins[0],
ins[1],
@@ -223,7 +223,7 @@ def convolution_inference_weight_transform(
return te.extern(
(output_channels, input_channels, transform_tile_size, transform_tile_size),
[kernel],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.nnpack.convolution_inference_weight_transform",
ins[0],
outs[0],
diff --git a/python/tvm/contrib/random.py b/python/tvm/contrib/random.py
index 0fcfc567a4..76462c6ecd 100644
--- a/python/tvm/contrib/random.py
+++ b/python/tvm/contrib/random.py
@@ -18,7 +18,7 @@
import tvm_ffi
-import tvm
+import tvm, tvm_tirx
from tvm import te
@@ -43,7 +43,7 @@ def randint(low, high, size, dtype="int32"):
return te.extern(
size,
[],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.random.randint", int(low), int(high), outs[0]
),
dtype=dtype,
@@ -77,7 +77,7 @@ def uniform(low, high, size):
return te.extern(
size,
[],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.random.uniform", float(low), float(high), outs[0]
),
dtype="float32",
@@ -107,7 +107,7 @@ def normal(loc, scale, size):
return te.extern(
size,
[],
- lambda ins, outs: tvm.tirx.call_packed(
+ lambda ins, outs: tvm_tirx.tirx.call_packed(
"tvm.contrib.random.normal", float(loc), float(scale), outs[0]
),
dtype="float32",
diff --git a/python/tvm/contrib/tvmjs.py b/python/tvm/contrib/tvmjs.py
index 3c64a3917d..b616d2b757 100644
--- a/python/tvm/contrib/tvmjs.py
+++ b/python/tvm/contrib/tvmjs.py
@@ -36,8 +36,8 @@ try:
except ImportError:
ml_dtypes = None
-import tvm
-from tvm.runtime import DataType
+import tvm, tvm_tirx
+from tvm_tirx.runtime import DataType
from .emcc import create_wasm as create_tvmjs_wasm
from .emcc import find_wasm_lib
@@ -202,8 +202,8 @@ class TensorCacheShardingManager:
def dump_tensor_cache(
- params: Mapping[str, np.ndarray | tvm.runtime.Tensor]
- | Iterator[tuple[str, np.ndarray | tvm.runtime.Tensor]],
+ params: Mapping[str, np.ndarray | tvm_tirx.runtime.Tensor]
+ | Iterator[tuple[str, np.ndarray | tvm_tirx.runtime.Tensor]],
cache_dir: str,
encode_format: str | Mapping[str, str] = "f32-to-bf16",
meta_data=None,
@@ -216,8 +216,8 @@ def dump_tensor_cache(
Parameters
----------
params: Union[
- Mapping[str, Union[np.ndarray, tvm.runtime.Tensor]],
- Iterator[Tuple[str, Union[np.ndarray, tvm.runtime.Tensor]]],
+ Mapping[str, Union[np.ndarray, tvm_tirx.runtime.Tensor]],
+ Iterator[Tuple[str, Union[np.ndarray, tvm_tirx.runtime.Tensor]]],
]
The parameter dictionary or generator
@@ -295,7 +295,7 @@ def dump_tensor_cache(
v = v.numpy()
# prefer to preserve original dtype, especially if the format was bfloat16
- dtype = origin_v.dtype if isinstance(origin_v, tvm.runtime.Tensor) else v.dtype
+ dtype = origin_v.dtype if isinstance(origin_v, tvm_tirx.runtime.Tensor) else v.dtype
if dtype in DataType._NUMPY_DTYPE_TO_STR:
dtype = DataType._NUMPY_DTYPE_TO_STR[dtype]
@@ -356,7 +356,7 @@ def dump_tensor_cache(
os.remove(b16_nd_cache_json)
-def load_tensor_cache(cachepath: str, device: tvm.runtime.Device):
+def load_tensor_cache(cachepath: str, device: tvm_tirx.runtime.Device):
"""Load the tensor cache from the directory or json.
@@ -365,7 +365,7 @@ def load_tensor_cache(cachepath: str, device: tvm.runtime.Device):
cachepath: str
Path to the location or json file.
- device: tvm.runtime.Device
+ device: tvm_tirx.runtime.Device
The device we would like to load the data from.
"""
if not cachepath.endswith(".json"):
@@ -390,7 +390,7 @@ def load_tensor_cache(cachepath: str, device: tvm.runtime.Device):
offset = rec["byteOffset"]
nbytes = rec["nbytes"]
- arr = tvm.runtime.empty(shape, dtype, device=device)
+ arr = tvm_tirx.runtime.empty(shape, dtype, device=device)
assert offset + nbytes <= len(raw_data)
buffer_source = raw_data[offset : offset + nbytes]
if dtype == "float8_e4m3fn":
diff --git a/python/tvm/driver/build_module.py b/python/tvm/driver/build_module.py
index 15cc37c4e9..1052c615aa 100644
--- a/python/tvm/driver/build_module.py
+++ b/python/tvm/driver/build_module.py
@@ -20,11 +20,11 @@
from collections.abc import Callable
-import tvm
-from tvm.ir.module import IRModule
-from tvm.runtime import Executable
-from tvm.target import Target
-from tvm.tirx import Function
+import tvm, tvm_tirx
+from tvm_tirx.ir.module import IRModule
+from tvm_tirx.runtime import Executable
+from tvm_tirx.target import Target
+from tvm_tirx.tirx import Function
def _contains_relax(mod: Function | IRModule) -> bool:
@@ -40,8 +40,8 @@ def compile( # pylint: disable=redefined-builtin
mod: Function | IRModule,
target: Target | None = None,
*,
- relax_pipeline: tvm.ir.transform.Pass | Callable | str | None = "default",
- tir_pipeline: tvm.ir.transform.Pass | Callable | str | None = "default",
+ relax_pipeline: tvm_tirx.ir.transform.Pass | Callable | str | None = "default",
+ tir_pipeline: tvm_tirx.ir.transform.Pass | Callable | str | None = "default",
) -> Executable:
"""
Compile an IRModule to a runtime executable.
@@ -56,10 +56,10 @@ def compile( # pylint: disable=redefined-builtin
TIR or Relax functions.
target : Optional[Target]
The target platform and its backend compilation configuration.
- relax_pipeline : Optional[Union[tvm.ir.transform.Pass, Callable, str]]
+ relax_pipeline : Optional[Union[tvm_tirx.ir.transform.Pass, Callable, str]]
The compilation pipeline to use for Relax functions.
Only used if the module contains Relax functions.
- tir_pipeline : Optional[Union[tvm.ir.transform.Pass, Callable, str]]
+ tir_pipeline : Optional[Union[tvm_tirx.ir.transform.Pass, Callable, str]]
The compilation pipeline to use for TIR functions.
Returns
@@ -75,4 +75,4 @@ def compile( # pylint: disable=redefined-builtin
relax_pipeline=relax_pipeline,
tir_pipeline=tir_pipeline,
)
- return tvm.tirx.compile(mod, target, pipeline=tir_pipeline)
+ return tvm_tirx.tirx.compile(mod, target, pipeline=tir_pipeline)
diff --git a/python/tvm/exec/disco_remote_socket_session.py b/python/tvm/exec/disco_remote_socket_session.py
index b79acf4342..3141157ffa 100644
--- a/python/tvm/exec/disco_remote_socket_session.py
+++ b/python/tvm/exec/disco_remote_socket_session.py
@@ -20,7 +20,7 @@
import sys
-import tvm
+import tvm, tvm_tirx
from . import disco_worker as _ # pylint: disable=unused-import
@@ -32,5 +32,5 @@ if __name__ == "__main__":
server_host = sys.argv[1]
server_port = int(sys.argv[2])
num_workers = int(sys.argv[3])
- func = tvm.get_global_func("runtime.disco.RemoteSocketSession")
+ func = tvm_tirx.get_global_func("runtime.disco.RemoteSocketSession")
func(server_host, server_port, num_workers)
diff --git a/python/tvm/exec/disco_worker.py b/python/tvm/exec/disco_worker.py
index 2052ae7b04..11f97d298e 100644
--- a/python/tvm/exec/disco_worker.py
+++ b/python/tvm/exec/disco_worker.py
@@ -25,8 +25,8 @@ from collections.abc import Callable
from tvm_ffi import Shape, get_global_func, register_global_func
from tvm_ffi.core import String
-import tvm
-from tvm.runtime import Tensor, tensor
+import tvm, tvm_tirx
+from tvm_tirx.runtime import Tensor, tensor
@register_global_func("tests.disco.add_one", override=True)
@@ -62,7 +62,7 @@ def _shape_tuple_func(x: Shape):
@register_global_func("tests.disco.test_callback", override=True)
-def _make_callback(device: tvm.runtime.Device) -> Callable[[str, int], Tensor]:
+def _make_callback(device: tvm_tirx.runtime.Device) -> Callable[[str, int], Tensor]:
"""For use in tests/python/disco/test_callback.py
This function simulates a callback to be used for lazy parameter
@@ -70,7 +70,7 @@ def _make_callback(device: tvm.runtime.Device) -> Callable[[str, int], Tensor]:
Parameters
----------
- device: tvm.runtime.Device
+ device: tvm_tirx.runtime.Device
The device on which parameters should be located, when
returned by the callback function.
@@ -94,7 +94,7 @@ def _make_callback(device: tvm.runtime.Device) -> Callable[[str, int], Tensor]:
arr = np.arange(4).reshape([2, 2]).astype("float32")
else:
raise ValueError(f"Unexpected index {param_index}")
- return tvm.runtime.tensor(arr, device=device)
+ return tvm_tirx.runtime.tensor(arr, device=device)
return fget_item
diff --git a/python/tvm/exec/gpu_memory_bandwidth.py b/python/tvm/exec/gpu_memory_bandwidth.py
index 2b7b87e3e9..ffec0d618a 100644
--- a/python/tvm/exec/gpu_memory_bandwidth.py
+++ b/python/tvm/exec/gpu_memory_bandwidth.py
@@ -22,7 +22,7 @@ import itertools
import numpy as np
-import tvm
+import tvm, tvm_tirx
from tvm import te
from tvm.s_tir.meta_schedule.runner import EvaluatorConfig, RPCConfig
from tvm.s_tir.testing.runner import local_run, rpc_run
@@ -179,9 +179,9 @@ def main(): # pylint: disable=too-many-locals
"""Entry point"""
args = _parse_args()
# pylint: disable=invalid-name
- target = tvm.target.Target(args.target)
+ target = tvm_tirx.target.Target(args.target)
if args.target_host is not None:
- target = tvm.target.Target(args.target, host=args.target_host)
+ target = tvm_tirx.target.Target(args.target, host=args.target_host)
dtype = args.dtype
rpcConfig = None
if args.rpc_host is not None and args.rpc_port is not None and args.rpc_key is not None:
diff --git a/python/tvm/relax/__init__.py b/python/tvm/relax/__init__.py
index 110e347fd1..c4f910a70b 100644
--- a/python/tvm/relax/__init__.py
+++ b/python/tvm/relax/__init__.py
@@ -20,7 +20,7 @@
from tvm.runtime import vm
from tvm.runtime.vm import VirtualMachine, VMInstrumentReturnKind
-from tvm.ir import Call
+from tvm_tirx.ir import Call
# Global information
from .global_info import DummyGlobalInfo, VDevice
@@ -111,19 +111,19 @@ from .vm_build import build, VMExecutable
from .binding_rewrite import DataflowBlockRewrite
-import tvm.script
+import tvm_tirx.script, tvm
-tvm.script.register_dialect("relax", "tvm.relax.script", builder_path="tvm.relax.script.ir_builder")
+tvm_tirx.script.register_dialect("relax", "tvm.relax.script", builder_path="tvm.relax.script.ir_builder")
-def _check_script_module(module: "tvm.ir.IRModule") -> None:
+def _check_script_module(module: "tvm_tirx.ir.IRModule") -> None:
# Delay builder imports until validation, keeping dialect/runtime bootstrap safe.
from tvm.relax.script.ir_builder.parser_protocol import _check_module_well_formed
_check_module_well_formed(module)
-tvm.script.register_module_validator(_check_script_module, prepend=True)
+tvm_tirx.script.register_module_validator(_check_script_module, prepend=True)
def _initialize_script_namespace() -> None:
@@ -132,7 +132,7 @@ def _initialize_script_namespace() -> None:
script._initialize()
-from tvm.script.parser import register_namespace_initializer as _register_namespace_initializer
+from tvm_tirx.script.parser import register_namespace_initializer as _register_namespace_initializer
_register_namespace_initializer(_initialize_script_namespace, aliases=("relax",))
diff --git a/python/tvm/relax/analysis/analysis.py b/python/tvm/relax/analysis/analysis.py
index b57c96b2dc..b35b215068 100644
--- a/python/tvm/relax/analysis/analysis.py
+++ b/python/tvm/relax/analysis/analysis.py
@@ -24,13 +24,13 @@ configuring the passes and scripting them in Python.
from collections.abc import Callable
from enum import IntEnum
-import tvm
-from tvm import IRModule, tirx
-from tvm.ir import Call, Type
+import tvm, tvm_tirx
+from tvm_tirx import IRModule, tirx
+from tvm_tirx.ir import Call, Type
from tvm.relax.expr import Binding, DataflowBlock, Expr, Function, GlobalVar
from tvm.relax.type import FuncType
from tvm.s_tir import SBlock
-from tvm.tirx import IndexMap, Var
+from tvm_tirx.tirx import IndexMap, Var
from . import _ffi_api
@@ -489,7 +489,7 @@ def well_formed(obj: IRModule | Function, check_ty: bool = True) -> None:
Parameters
----------
- obj : Union[tvm.IRModule, Function]
+ obj : Union[tvm_tirx.IRModule, Function]
The input IRModule or relax.Function.
check_ty : bool
@@ -512,7 +512,7 @@ def check_well_formed(obj: IRModule | Function, check_ty: bool = True) -> bool:
Parameters
----------
- obj : Union[tvm.IRModule, Function]
+ obj : Union[tvm_tirx.IRModule, Function]
The input IRModule or relax.Function.
check_ty : bool
@@ -566,7 +566,7 @@ def suggest_layout_transforms(
return _ffi_api.suggest_layout_transforms(func, write_tensor_index_maps) # type: ignore
-def detect_recursion(mod: tvm.IRModule) -> list[list[GlobalVar]]:
+def detect_recursion(mod: tvm_tirx.IRModule) -> list[list[GlobalVar]]:
"""
Find all sets of recursive or mutually recursive functions in the module.
diff --git a/python/tvm/relax/analysis/estimate_memory_usage.py b/python/tvm/relax/analysis/estimate_memory_usage.py
index 78159d9c53..59cf45add1 100644
--- a/python/tvm/relax/analysis/estimate_memory_usage.py
+++ b/python/tvm/relax/analysis/estimate_memory_usage.py
@@ -17,9 +17,9 @@
# pylint: disable=abstract-method,unused-argument
# pylint: disable=missing-function-docstring,missing-module-docstring
-import tvm
-from tvm.ir import Call, Op
-from tvm.ir.module import IRModule
+import tvm, tvm_tirx
+from tvm_tirx.ir import Call, Op
+from tvm_tirx.ir.module import IRModule
from ..expr import Expr, Function, ShapeExpr
from ..expr_functor import PyExprVisitor, visitor
@@ -124,11 +124,11 @@ def estimate_memory_usage(mod: IRModule | Function) -> str:
)
size: int = 1
for dim_len in shape.values:
- if not isinstance(dim_len, tvm.tirx.IntImm):
+ if not isinstance(dim_len, tvm_tirx.tirx.IntImm):
self.total_dyn_size_tensor_num += 1
return -1
size *= dim_len.value
- dtype = tvm.DataType(dtype_str)
+ dtype = tvm_tirx.DataType(dtype_str)
return size * ((dtype.bits + 7) // 8) * dtype.lanes
def accumulate_builtin_tensor_alloc(self, shape: Expr, dtype_str: str) -> None:
@@ -177,6 +177,6 @@ def estimate_memory_usage(mod: IRModule | Function) -> str:
return "- Function " + func_name + ":\n" + est
if isinstance(mod, Function):
- mod = tvm.IRModule({tvm.ir.GlobalVar("foo"): mod})
+ mod = tvm_tirx.IRModule({tvm_tirx.ir.GlobalVar("foo"): mod})
return MemoryEstimator().estimate(mod)
diff --git a/python/tvm/relax/backend/adreno/clml.py b/python/tvm/relax/backend/adreno/clml.py
index f7f0884bdf..69cd06e867 100644
--- a/python/tvm/relax/backend/adreno/clml.py
+++ b/python/tvm/relax/backend/adreno/clml.py
@@ -17,9 +17,10 @@
# pylint: disable=invalid-name, unused-argument, pointless-exception-statement
"""Pattern table for CLML backend"""
-import tvm
-from tvm import IRModule, relax, tirx
-from tvm.ir.transform import PassContext, module_pass
+import tvm, tvm_tirx
+from tvm_tirx import IRModule, tirx
+from tvm import relax
+from tvm_tirx.ir.transform import PassContext, module_pass
from tvm.relax import transform
from tvm.relax.dpl.pattern import (
GlobalVarPattern,
@@ -37,7 +38,7 @@ from ..pattern_registry import register_patterns
def _dtype_str(dtype):
- return str(dtype.dtype) if isinstance(dtype, tvm.ir.PrimType) else str(dtype)
+ return str(dtype.dtype) if isinstance(dtype, tvm_tirx.ir.PrimType) else str(dtype)
@mutator
@@ -62,7 +63,7 @@ class AppendReshapeToBNRewriter(PyExprMutator):
def visit_tuple_getitem_(self, op: TupleGetItem):
tuple_value = op.tuple_value
- reshape_op = tvm.ir.Op.get("relax.reshape")
+ reshape_op = tvm_tirx.ir.Op.get("relax.reshape")
if isinstance(tuple_value, relax.Var) and tuple_value in self.bn_vars:
bn_call = self.bn_vars[tuple_value]
@@ -82,7 +83,7 @@ class AppendReshapeToBNRewriter(PyExprMutator):
@transform.function_pass(opt_level=0, name="AppendReshapeToBN")
class AppendReshapeToBNRewriterPass:
def transform_function(
- self, func: relax.Function, mod: IRModule, _ctx: tvm.ir.transform.PassContext
+ self, func: relax.Function, mod: IRModule, _ctx: tvm_tirx.ir.transform.PassContext
) -> relax.Function:
updated_func = AppendReshapeToBNRewriter(mod).visit_expr(func)
updated_func = relax.analysis.remove_all_unused(updated_func)
@@ -98,7 +99,7 @@ def clml_sdk_version():
# Registry: "relax.get_openclml_version" — returns the CLML SDK version
# that TVM was built against; registered unconditionally in codegen.cc.
# Grep hint: grep -rn 'relax.get_openclml_version' src/
- get_version = tvm.get_global_func("relax.get_openclml_version", allow_missing=True)
+ get_version = tvm_tirx.get_global_func("relax.get_openclml_version", allow_missing=True)
if get_version is None:
return 2
return int(get_version())
@@ -112,7 +113,7 @@ def is_clml_runtime_enabled():
ret: bool
True if present, False if not.
"""
- check_enabled = tvm.get_global_func("relax.op.is_openclml_runtime_enabled", True)
+ check_enabled = tvm_tirx.get_global_func("relax.op.is_openclml_runtime_enabled", True)
if check_enabled:
return check_enabled()
return False
@@ -450,7 +451,7 @@ def clml_pattern_table():
}
for param in params.values():
- if not isinstance(param, tvm.ir.GenericConst):
+ if not isinstance(param, tvm_tirx.ir.GenericConst):
return False
base_shape = None
@@ -612,7 +613,7 @@ class OpenCLMLOffLoad:
"relax.nn.conv2d": ["NCHW", "OIHW"],
"relax.nn.conv2d_transpose": ["NCHW", "OIHW"],
}
- seq = tvm.ir.transform.Sequential(
+ seq = tvm_tirx.ir.transform.Sequential(
[
transform.ConvertLayout(clml_layouts),
transform.Normalize(),
@@ -695,15 +696,15 @@ clml_llm_patterns = [
register_patterns(clml_llm_patterns)
-@tvm.ir.transform.module_pass(opt_level=0, name="OpenCLMLOffLoadForLLM")
+@tvm_tirx.ir.transform.module_pass(opt_level=0, name="OpenCLMLOffLoadForLLM")
class OpenCLMLOffLoadForLLM:
"""A compiler pass that partition the graph with dequant Matmul to CLML backend offload."""
- def __init__(self, target: tvm.target.Target) -> None:
+ def __init__(self, target: tvm_tirx.target.Target) -> None:
"""Initializer.
Parameters
----------
- target : tvm.target.Target
+ target : tvm_tirx.target.Target
Target device.
"""
self.target = target
@@ -711,12 +712,12 @@ class OpenCLMLOffLoadForLLM:
def transform_module(
self,
mod: IRModule,
- _ctx: tvm.ir.transform.PassContext,
+ _ctx: tvm_tirx.ir.transform.PassContext,
) -> IRModule:
"""Apply required passed to transform"""
if "adreno" in self.target.keys and (clml_sdk_version() >= 5):
- mod = tvm.ir.transform.Sequential(
+ mod = tvm_tirx.ir.transform.Sequential(
[
transform.Normalize(),
transform.FuseOpsByPattern(clml_llm_patterns, annotate_codegen=True),
diff --git a/python/tvm/relax/backend/adreno/pipeline.py b/python/tvm/relax/backend/adreno/pipeline.py
index b99bbdd473..819e57a217 100644
--- a/python/tvm/relax/backend/adreno/pipeline.py
+++ b/python/tvm/relax/backend/adreno/pipeline.py
@@ -16,12 +16,12 @@
# under the License.
"""The Relax Adreno GPU backend compilation pipeline and other passes."""
-import tvm
+import tvm, tvm_tirx
from tvm import relax
from tvm.relax.transform.legalize_ops import adreno as legalize_adreno
-def library_dispatch_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def library_dispatch_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default library dispatch passes for Adreno GPU backend."""
if "clml" in target.keys:
return [
@@ -32,7 +32,7 @@ def library_dispatch_passes(target: tvm.target.Target): # pylint: disable=unuse
return []
-def legalize_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def legalize_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default legalization passes for Adreno GPU backend."""
desired_layouts = {"relax.nn.conv2d": ["NCHW4c", "OIHW4o", "NCHW4c"]}
skip_ops = [
@@ -44,7 +44,7 @@ def legalize_passes(target: tvm.target.Target): # pylint: disable=unused-argume
pass_list.extend(
[
- tvm.tirx.transform.BindTarget(tvm.target.Target.current(allow_none=False)),
+ tvm_tirx.tirx.transform.BindTarget(tvm_tirx.target.Target.current(allow_none=False)),
relax.transform.DecomposeOpsForInference(),
]
)
@@ -110,23 +110,23 @@ def legalize_passes(target: tvm.target.Target): # pylint: disable=unused-argume
return pass_list
-def dataflow_lower_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def dataflow_lower_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default dataflow lowering passes for Adreno GPU backend."""
return relax.backend.gpu_generic.dataflow_lower_passes(target)
-def finalize_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def finalize_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default finalization passes for Adreno GPU backend."""
return relax.backend.gpu_generic.finalize_passes(target)
-def get_default_pipeline(target: tvm.target.Target):
+def get_default_pipeline(target: tvm_tirx.target.Target):
"""Return the default compilation pipeline for Adreno GPU."""
- @tvm.ir.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext):
+ @tvm_tirx.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm_tirx.ir.IRModule, _ctx: tvm_tirx.ir.transform.PassContext):
with target:
- seq = tvm.ir.transform.Sequential(
+ seq = tvm_tirx.ir.transform.Sequential(
library_dispatch_passes(target)
+ legalize_passes(target)
+ dataflow_lower_passes(target)
diff --git a/python/tvm/relax/backend/adreno/transform/transform.py b/python/tvm/relax/backend/adreno/transform/transform.py
index 5b591bbf46..5db8cea36b 100644
--- a/python/tvm/relax/backend/adreno/transform/transform.py
+++ b/python/tvm/relax/backend/adreno/transform/transform.py
@@ -17,33 +17,33 @@
# pylint: disable=invalid-name
"""Adreno Relax transformation passes."""
-import tvm.ir
-from tvm.target import Target
+import tvm_tirx.ir, tvm
+from tvm_tirx.target import Target
from . import _ffi_api
-def AnnotateCustomMemoryScope(target: Target | None = None) -> tvm.ir.transform.Pass:
+def AnnotateCustomMemoryScope(target: Target | None = None) -> tvm_tirx.ir.transform.Pass:
"""Allocate the memory scope information. This is Adreno specific pass to annotate
The memory scope information and realize the same with RealizeVDevice pass followed by
updating the Functiontion var_buffer mapping using SpecializeFunctionBasedOnCallSite.
Returns
-------
- ret: tvm.ir.transform.Pass
+ ret: tvm_tirx.ir.transform.Pass
The registered pass for allocating workspace.
"""
return _ffi_api.AnnotateCustomMemoryScope(target) # type: ignore
-def FoldVDeviceScopeChange() -> tvm.ir.transform.Pass:
+def FoldVDeviceScopeChange() -> tvm_tirx.ir.transform.Pass:
"""This pass is a texture specific pass that can optimize unnecessary to_device copies.
Like texture_scope -> ToVDevice -> global scope. In this case the producer can directly
store into global scope avoiding unnecessary device copy.
Returns
-------
- ret: tvm.ir.transform.Pass
+ ret: tvm_tirx.ir.transform.Pass
The registered pass for allocating workspace.
"""
return _ffi_api.FoldVDeviceScopeChange() # type: ignore
diff --git a/python/tvm/relax/backend/contrib/example_npu/patterns.py b/python/tvm/relax/backend/contrib/example_npu/patterns.py
index 41ae185357..37fb27a988 100644
--- a/python/tvm/relax/backend/contrib/example_npu/patterns.py
+++ b/python/tvm/relax/backend/contrib/example_npu/patterns.py
@@ -24,7 +24,7 @@ tiling, and fusion strategies.
from typing import ClassVar
-from tvm.ir import Op
+from tvm_tirx.ir import Op
from tvm.relax.dpl.pattern import is_op, wildcard
from tvm.relax.transform import PatternCheckContext
diff --git a/python/tvm/relax/backend/contrib/nnapi.py b/python/tvm/relax/backend/contrib/nnapi.py
index b8f9125cd5..3288239030 100644
--- a/python/tvm/relax/backend/contrib/nnapi.py
+++ b/python/tvm/relax/backend/contrib/nnapi.py
@@ -19,7 +19,7 @@
from collections.abc import Mapping
-from tvm.ir import IRModule
+from tvm_tirx.ir import IRModule
from tvm.relax.dpl.pattern import (
DFPattern,
is_op,
@@ -303,14 +303,14 @@ def partition_for_nnapi(mod: IRModule, feature_level: int | None = None) -> IRMo
Parameters
----------
- mod : tvm.ir.IRModule
+ mod : tvm_tirx.ir.IRModule
The module to run passes on.
feature_level : Optional[int]
The maximum NNAPI feature level.
Returns
-------
- mod : tvm.ir.IRModule
+ mod : tvm_tirx.ir.IRModule
Annotated and partitioned module.
"""
patterns = get_patterns_with_prefix("nnapi")
diff --git a/python/tvm/relax/backend/contrib/tensorrt.py b/python/tvm/relax/backend/contrib/tensorrt.py
index b3a3c2b9cd..9f2b170c7a 100644
--- a/python/tvm/relax/backend/contrib/tensorrt.py
+++ b/python/tvm/relax/backend/contrib/tensorrt.py
@@ -23,8 +23,8 @@ converter registered under the same name (the converters are keyed by
out of the module and annotates them for the ``tensorrt`` codegen.
"""
-from tvm import tirx
-from tvm.ir import IRModule
+from tvm_tirx import tirx
+from tvm_tirx.ir import IRModule
from tvm.relax.dpl.pattern import is_op, wildcard
from tvm.relax.expr import ShapeExpr
from tvm.relax.transform import (
@@ -193,14 +193,14 @@ def partition_for_tensorrt(mod: IRModule) -> IRModule:
Parameters
----------
- mod : tvm.ir.IRModule
+ mod : tvm_tirx.ir.IRModule
The module to partition. Bind model parameters (e.g. via
``relax.transform.BindParams``) before calling this so that weights are
available to TensorRT as constants.
Returns
-------
- mod : tvm.ir.IRModule
+ mod : tvm_tirx.ir.IRModule
The module with TensorRT-supported subgraphs grouped into composite
functions annotated for the ``tensorrt`` codegen.
@@ -229,12 +229,12 @@ def partition_for_tensorrt(mod: IRModule) -> IRModule:
After binding model weights as constants, partition and generate code on the
compilation GPU before exporting::
- import tvm
+ import tvm, tvm_tirx
from tvm import relax
from tvm.relax.backend.contrib.tensorrt import partition_for_tensorrt
partitioned = partition_for_tensorrt(mod)
- with tvm.ir.transform.PassContext(
+ with tvm_tirx.ir.transform.PassContext(
config={
"relax.ext.tensorrt.options": {
"build_at_compile_time": True,
@@ -247,11 +247,11 @@ def partition_for_tensorrt(mod: IRModule) -> IRModule:
Load the library and run inference in a new process on the deployment GPU::
- import tvm
+ import tvm, tvm_tirx
from tvm import relax
- executable = tvm.runtime.load_module("model.so")
- vm = relax.VirtualMachine(executable, tvm.cuda(0))
+ executable = tvm_tirx.runtime.load_module("model.so")
+ vm = relax.VirtualMachine(executable, tvm_tirx.cuda(0))
result = vm["main"](*inputs)
"""
patterns = get_patterns_with_prefix("tensorrt")
diff --git a/python/tvm/relax/backend/cpu_generic/pipeline.py b/python/tvm/relax/backend/cpu_generic/pipeline.py
index 36a1798f67..dec0d99303 100644
--- a/python/tvm/relax/backend/cpu_generic/pipeline.py
+++ b/python/tvm/relax/backend/cpu_generic/pipeline.py
@@ -16,11 +16,11 @@
# under the License.
"""The Relax CPU backend compilation pipeline and other passes."""
-import tvm
+import tvm, tvm_tirx
from tvm import relax
-def library_dispatch_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def library_dispatch_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default library dispatch passes for CPU backend."""
return [
relax.backend.DispatchSampling(),
@@ -28,7 +28,7 @@ def library_dispatch_passes(target: tvm.target.Target): # pylint: disable=unuse
]
-def legalize_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def legalize_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default legalization passes for CPU backend."""
return [
tvm.relax.transform.LegalizeOps(),
@@ -39,7 +39,7 @@ def legalize_passes(target: tvm.target.Target): # pylint: disable=unused-argume
]
-def dataflow_lower_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def dataflow_lower_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default dataflow lowering passes for CPU backend."""
return [
relax.transform.RewriteDataflowReshape(),
@@ -49,7 +49,7 @@ def dataflow_lower_passes(target: tvm.target.Target): # pylint: disable=unused-
]
-def finalize_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def finalize_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default finalization passes for CPU backend."""
return [
relax.transform.StaticPlanBlockMemory(),
@@ -62,13 +62,13 @@ def finalize_passes(target: tvm.target.Target): # pylint: disable=unused-argume
]
-def get_default_pipeline(target: tvm.target.Target):
+def get_default_pipeline(target: tvm_tirx.target.Target):
"""Return the default compilation pipeline for CPU."""
- @tvm.ir.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext):
+ @tvm_tirx.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm_tirx.ir.IRModule, _ctx: tvm_tirx.ir.transform.PassContext):
with target:
- seq = tvm.ir.transform.Sequential(
+ seq = tvm_tirx.ir.transform.Sequential(
library_dispatch_passes(target)
+ legalize_passes(target)
+ dataflow_lower_passes(target)
diff --git a/python/tvm/relax/backend/cuda/cublas.py b/python/tvm/relax/backend/cuda/cublas.py
index f3d5ce708f..7d525b9a1c 100644
--- a/python/tvm/relax/backend/cuda/cublas.py
+++ b/python/tvm/relax/backend/cuda/cublas.py
@@ -20,11 +20,11 @@
import operator
from functools import reduce
-import tvm
-from tvm import DataType
+import tvm, tvm_tirx
+from tvm_tirx import DataType
from tvm.relax import transform
from tvm.relax.transform import PatternCheckContext
-from tvm.sym import Analyzer
+from tvm_tirx.sym import Analyzer
from ..pattern_registry import get_patterns_with_prefix, register_patterns
from ..patterns import (
@@ -74,7 +74,7 @@ def _check_matmul(context: PatternCheckContext) -> bool:
lhs_shape = lhs.ty.shape.values
rhs_shape = rhs.ty.shape.values
- if not isinstance(lhs_shape[-1], tvm.tirx.expr.IntImm | int):
+ if not isinstance(lhs_shape[-1], tvm_tirx.tirx.expr.IntImm | int):
# Reduction axis must be constant
return False
@@ -82,7 +82,7 @@ def _check_matmul(context: PatternCheckContext) -> bool:
if lhs_shape[-1] % 4 != 0:
# Reduction axis must be multiples of 4 for IGEMM
return False
- if not isinstance(rhs_shape[-1], tvm.tirx.expr.IntImm | int) or rhs_shape[-1] % 4 != 0:
+ if not isinstance(rhs_shape[-1], tvm_tirx.tirx.expr.IntImm | int) or rhs_shape[-1] % 4 != 0:
# Rows number must be multiples of 4 for IGEMM
return False
elif lhs_dtype == "float8_e4m3fn" and rhs_dtype == "float8_e4m3fn":
@@ -102,12 +102,12 @@ def _check_matmul(context: PatternCheckContext) -> bool:
# cuBLAS FP8 operations require all tensors being aligned to 16 bytes.
if (
- not isinstance(rhs_shape[-1], tvm.tirx.expr.IntImm | int)
+ not isinstance(rhs_shape[-1], tvm_tirx.tirx.expr.IntImm | int)
or rhs_shape[-1] % (16 // DataType(lhs_dtype).itemsize) != 0
):
return False
if (
- not isinstance(rhs_shape[-2], tvm.tirx.expr.IntImm | int)
+ not isinstance(rhs_shape[-2], tvm_tirx.tirx.expr.IntImm | int)
or rhs_shape[-2] % (16 // DataType(out_dtype).itemsize) != 0
):
return False
@@ -122,7 +122,7 @@ def _check_matmul(context: PatternCheckContext) -> bool:
bias = context.annotated_expr["bias"]
bias_shape = bias.ty.shape.values
bias_batches = reduce(operator.mul, bias_shape[:-1], 1)
- if not isinstance(bias_batches, tvm.tirx.expr.IntImm | int) or int(bias_batches) > 1:
+ if not isinstance(bias_batches, tvm_tirx.tirx.expr.IntImm | int) or int(bias_batches) > 1:
# cuBLAS only supports bias vector
return False
@@ -133,8 +133,8 @@ def _check_matmul(context: PatternCheckContext) -> bool:
# must be equal. If lhs is batched but rhs is not, we can use the regular GEMM by
# flattening all batch axes into the M axis.
return (
- isinstance(lhs_batches, tvm.ir.Var)
- or isinstance(rhs_batches, tvm.ir.Var)
+ isinstance(lhs_batches, tvm_tirx.ir.Var)
+ or isinstance(rhs_batches, tvm_tirx.ir.Var)
or (analyzer.can_prove_equal(lhs_batches, rhs_batches))
or (analyzer.can_prove(lhs_batches >= 1) and analyzer.can_prove(rhs_batches == 1))
)
@@ -226,7 +226,7 @@ def partition_for_cublas(mod, bind_constants=False):
Parameters
----------
- mod: tvm.IRModule
+ mod: tvm_tirx.IRModule
The IRModule to be partitioned.
bind_constants : bool
@@ -234,7 +234,7 @@ def partition_for_cublas(mod, bind_constants=False):
Returns
-------
- mod: tvm.IRModule
+ mod: tvm_tirx.IRModule
The resulting IRModule, containing partitioned subgraphs to be
offloaded to the cuBLAS backend.
"""
diff --git a/python/tvm/relax/backend/cuda/cudnn.py b/python/tvm/relax/backend/cuda/cudnn.py
index 523a651592..c9ebaf630c 100644
--- a/python/tvm/relax/backend/cuda/cudnn.py
+++ b/python/tvm/relax/backend/cuda/cudnn.py
@@ -20,7 +20,7 @@
import operator
from functools import partial, reduce
-import tvm
+import tvm, tvm_tirx
from tvm import relax
from tvm.relax import PyExprMutator, expr_functor, transform
from tvm.relax.transform import PatternCheckContext
@@ -142,18 +142,18 @@ def partition_for_cudnn(mod):
Parameters
----------
- mod: tvm.IRModule
+ mod: tvm_tirx.IRModule
The IRModule to be partitioned.
Returns
-------
- mod: tvm.IRModule
+ mod: tvm_tirx.IRModule
The resulting IRModule, containing partitioned subgraphs to be
offloaded to the cuDNN backend.
"""
patterns = get_patterns_with_prefix("cudnn")
- return tvm.ir.transform.Sequential(
+ return tvm_tirx.ir.transform.Sequential(
[
transform.FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True),
annotate_workspace,
@@ -191,7 +191,7 @@ class WorkspaceAnnotator(PyExprMutator):
out_size_1d = _shape_1d(f.ret_ty.shape)
# This needs to be in sync with the actual value that the kernel expects.
workspace_size_bytes = out_size_1d * {"float16": 2, "float32": 4}[out_dtype]
- if not isinstance(workspace_size_bytes, int | tvm.tirx.expr.IntImm):
+ if not isinstance(workspace_size_bytes, int | tvm_tirx.tirx.expr.IntImm):
# Tempororay workaround for dynamic shape workload. Will be removed when
# workspace for dynamic shape workload is implemented.
workspace_size_bytes = 8
@@ -200,7 +200,7 @@ class WorkspaceAnnotator(PyExprMutator):
return f
-@tvm.ir.transform.module_pass(opt_level=0)
+@tvm_tirx.ir.transform.module_pass(opt_level=0)
def annotate_workspace(mod, _):
"""Pass to annotate a workspace requirement for each cuDNN-offloaded function."""
annotator = WorkspaceAnnotator(mod)
diff --git a/python/tvm/relax/backend/cuda/cutlass.py b/python/tvm/relax/backend/cuda/cutlass.py
index 7ef24ad5b5..c8ca45b190 100644
--- a/python/tvm/relax/backend/cuda/cutlass.py
+++ b/python/tvm/relax/backend/cuda/cutlass.py
@@ -22,7 +22,7 @@ import operator
from collections.abc import Mapping, Sequence
from functools import reduce
-import tvm
+import tvm, tvm_tirx
from tvm.contrib.cutlass.build import is_shape_valid_for_cutlass_matmul
from tvm.relax import (
Call,
@@ -83,7 +83,7 @@ def _has_dependency(from_var: Var, to_var: Var, var_usages: Mapping[Var, Sequenc
def _is_same_shape(shape1, shape2):
- analyzer = tvm.sym.Analyzer()
+ analyzer = tvm_tirx.sym.Analyzer()
return all([analyzer.can_prove_equal(s1, s2) for s1, s2 in zip(shape1, shape2)])
@@ -136,7 +136,7 @@ def _check_conv2d(context: PatternCheckContext) -> bool:
# Check if any dimensions are symbolic.
for dim in data.ty.shape.values:
- if isinstance(dim, tvm.ir.Var):
+ if isinstance(dim, tvm_tirx.ir.Var):
return False
# pylint: disable=invalid-name
@@ -551,7 +551,7 @@ class WorkspaceAnnotator(PyExprMutator):
out_size_1d = _shape_1d(f.ret_ty.shape)
# This needs to be in sync with the actual value that the kernel expects.
workspace_size_bytes = out_size_1d * {"float16": 2, "float32": 4}[out_dtype]
- if not isinstance(workspace_size_bytes, int | tvm.tirx.expr.IntImm):
+ if not isinstance(workspace_size_bytes, int | tvm_tirx.tirx.expr.IntImm):
# Tempororay workaround for dynamic shape workload. Will be removed when
# workspace for dynamic shape workload is implemented.
workspace_size_bytes = 8
@@ -560,7 +560,7 @@ class WorkspaceAnnotator(PyExprMutator):
return f
-@tvm.ir.transform.module_pass(opt_level=0)
+@tvm_tirx.ir.transform.module_pass(opt_level=0)
def annotate_workspace(mod, _):
"""Pass to annotate a workspace requirement for each CUTLASS-offloaded function."""
annotator = WorkspaceAnnotator(mod)
@@ -577,7 +577,7 @@ def partition_for_cutlass(mod, annotate_codegen=True, use_flash_mqa=True):
Parameters
----------
- mod: tvm.IRModule
+ mod: tvm_tirx.IRModule
The IRModule to be partitioned.
annotate_codegen: bool
@@ -591,7 +591,7 @@ def partition_for_cutlass(mod, annotate_codegen=True, use_flash_mqa=True):
Returns
-------
- mod: tvm.IRModule
+ mod: tvm_tirx.IRModule
The resulting IRModule, containing partitioned subgraphs to be
compiled by the CUTLASS backend.
"""
@@ -609,7 +609,7 @@ def partition_for_cutlass(mod, annotate_codegen=True, use_flash_mqa=True):
mod[func_name] = func
patterns = get_patterns_with_prefix("cutlass")
- return tvm.ir.transform.Sequential(
+ return tvm_tirx.ir.transform.Sequential(
[
transform.FuseOpsByPattern(
patterns, bind_constants=False, annotate_codegen=annotate_codegen
diff --git a/python/tvm/relax/backend/cuda/flashinfer.py b/python/tvm/relax/backend/cuda/flashinfer.py
index b3ff3bfc21..f92510a09a 100644
--- a/python/tvm/relax/backend/cuda/flashinfer.py
+++ b/python/tvm/relax/backend/cuda/flashinfer.py
@@ -20,8 +20,8 @@
import re
from pathlib import Path
-import tvm
-from tvm.target import Target
+import tvm, tvm_tirx
+from tvm_tirx.target import Target
def _rename_exported_func_names(source_paths: list[Path], prefix: str):
@@ -48,9 +48,9 @@ def _rename_exported_func_names(source_paths: list[Path], prefix: str):
source_path.write_text("".join(lines), encoding="utf-8")
-def _load_flashinfer_modules(object_files: list[Path]) -> list[tvm.runtime.Module]:
+def _load_flashinfer_modules(object_files: list[Path]) -> list[tvm_tirx.runtime.Module]:
return [
- tvm.runtime.load_static_library(str(obj_path.absolute()), func_names=[])
+ tvm_tirx.runtime.load_static_library(str(obj_path.absolute()), func_names=[])
for obj_path in object_files
]
@@ -63,7 +63,7 @@ def gen_flashinfer_prefill_module(
v_head_dim: int,
enable_inline_rope: bool,
return_static_libs: bool = False,
-) -> list[tvm.runtime.Module]:
+) -> list[tvm_tirx.runtime.Module]:
"""Generate a FlashInfer module for prefill.
Parameters
@@ -158,7 +158,7 @@ def gen_flashinfer_decode_module(
v_head_dim: int,
enable_inline_rope: bool,
return_static_libs: bool = False,
-) -> list[tvm.runtime.Module]:
+) -> list[tvm_tirx.runtime.Module]:
"""Generate a FlashInfer module for decode.
Parameters
@@ -236,7 +236,7 @@ def gen_flashinfer_mla_module(
head_dim_ckv: int,
head_dim_kpe: int,
return_static_libs: bool = False,
-) -> list[tvm.runtime.Module]:
+) -> list[tvm_tirx.runtime.Module]:
"""Generate a FlashInfer module for MLA.
Parameters
@@ -300,7 +300,7 @@ def gen_flashinfer_mla_module(
def gen_grouped_gemm_module(
target: Target, return_static_libs: bool = False
-) -> list[tvm.runtime.Module]:
+) -> list[tvm_tirx.runtime.Module]:
"""Generate a FlashInfer module for FP8 grouped GEMM.
Parameters
@@ -314,7 +314,7 @@ def gen_grouped_gemm_module(
Returns
-------
- List[tvm.runtime.Module]
+ List[tvm_tirx.runtime.Module]
A list of compiled static library modules for FlashInfer FP8 grouped GEMM kernels.
Note
@@ -334,7 +334,7 @@ def gen_grouped_gemm_module(
"in https://docs.flashinfer.ai to install FlashInfer."
)
- compute_version = "".join(tvm.support.nvcc.get_target_compute_version(target).split("."))
+ compute_version = "".join(tvm_tirx.support.nvcc.get_target_compute_version(target).split("."))
if compute_version == "100":
jit_spec = gen_gemm_sm100_module()
else:
diff --git a/python/tvm/relax/backend/cuda/pipeline.py b/python/tvm/relax/backend/cuda/pipeline.py
index 8073cf80f7..f8f818d369 100644
--- a/python/tvm/relax/backend/cuda/pipeline.py
+++ b/python/tvm/relax/backend/cuda/pipeline.py
@@ -16,11 +16,11 @@
# under the License.
"""The Relax CUDA backend compilation pipeline and other passes."""
-import tvm
+import tvm, tvm_tirx
from tvm import relax
-def library_dispatch_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def library_dispatch_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default library dispatch passes for CUDA backend."""
return [
relax.backend.DispatchSampling(),
@@ -28,7 +28,7 @@ def library_dispatch_passes(target: tvm.target.Target): # pylint: disable=unuse
]
-def legalize_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def legalize_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default legalization passes for CUDA backend."""
from tvm.s_tir import dlight as dl # pylint: disable=import-outside-toplevel
@@ -48,7 +48,7 @@ def legalize_passes(target: tvm.target.Target): # pylint: disable=unused-argume
]
-def dataflow_lower_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def dataflow_lower_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default dataflow lowering passes for CUDA backend."""
return [
relax.transform.RewriteDataflowReshape(),
@@ -58,7 +58,7 @@ def dataflow_lower_passes(target: tvm.target.Target): # pylint: disable=unused-
]
-def finalize_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def finalize_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default finalization passes for CUDA backend."""
return [
relax.transform.StaticPlanBlockMemory(),
@@ -72,13 +72,13 @@ def finalize_passes(target: tvm.target.Target): # pylint: disable=unused-argume
]
-def get_default_pipeline(target: tvm.target.Target):
+def get_default_pipeline(target: tvm_tirx.target.Target):
"""Return the default compilation pipeline for CUDA."""
- @tvm.ir.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext):
+ @tvm_tirx.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm_tirx.ir.IRModule, _ctx: tvm_tirx.ir.transform.PassContext):
with target:
- seq = tvm.ir.transform.Sequential(
+ seq = tvm_tirx.ir.transform.Sequential(
library_dispatch_passes(target)
+ legalize_passes(target)
+ dataflow_lower_passes(target)
diff --git a/python/tvm/relax/backend/dispatch_sampling.py b/python/tvm/relax/backend/dispatch_sampling.py
index bf9b079dc5..398c1bf0e6 100644
--- a/python/tvm/relax/backend/dispatch_sampling.py
+++ b/python/tvm/relax/backend/dispatch_sampling.py
@@ -18,9 +18,9 @@
"""Dispatch sampling operators to platform dependent implementation."""
from tvm import relax
-from tvm.ir import Op
-from tvm.ir.module import IRModule
-from tvm.ir.transform import PassContext, module_pass
+from tvm_tirx.ir import Op
+from tvm_tirx.ir.module import IRModule
+from tvm_tirx.ir.transform import PassContext, module_pass
from tvm.relax import expr_functor
from .utils import BackendDispatcher
diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py b/python/tvm/relax/backend/dispatch_sort_scan.py
index f91728daa0..47828fda29 100644
--- a/python/tvm/relax/backend/dispatch_sort_scan.py
+++ b/python/tvm/relax/backend/dispatch_sort_scan.py
@@ -20,13 +20,14 @@
from functools import reduce
from operator import mul
-from tvm import DataType, relax, topi
+from tvm_tirx import DataType
+from tvm import relax, topi
from tvm.contrib.thrust import can_use_thrust
-from tvm.ir import GlobalVar, Op
-from tvm.ir.module import IRModule
-from tvm.ir.transform import PassContext, module_pass
+from tvm_tirx.ir import GlobalVar, Op
+from tvm_tirx.ir.module import IRModule
+from tvm_tirx.ir.transform import PassContext, module_pass
from tvm.relax import expr_functor
-from tvm.target import Target
+from tvm_tirx.target import Target
from .utils import BackendDispatcher
diff --git a/python/tvm/relax/backend/gpu_generic/cumsum.py b/python/tvm/relax/backend/gpu_generic/cumsum.py
index a0622e741a..2fb20bdb54 100644
--- a/python/tvm/relax/backend/gpu_generic/cumsum.py
+++ b/python/tvm/relax/backend/gpu_generic/cumsum.py
@@ -19,9 +19,9 @@
import math
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-from tvm.tirx import Function
+from tvm_tirx.script import s_tir as Ts
+from tvm_tirx.script import tirx as T
+from tvm_tirx.tirx import Function
def _is_power_of_two(n: int):
diff --git a/python/tvm/relax/backend/gpu_generic/pipeline.py b/python/tvm/relax/backend/gpu_generic/pipeline.py
index 8ded70313d..7b3bf7a5b0 100644
--- a/python/tvm/relax/backend/gpu_generic/pipeline.py
+++ b/python/tvm/relax/backend/gpu_generic/pipeline.py
@@ -16,11 +16,11 @@
# under the License.
"""The Relax generic GPU backend compilation pipeline and other passes."""
-import tvm
+import tvm, tvm_tirx
from tvm import relax
-def library_dispatch_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def library_dispatch_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default library dispatch passes for generic GPU backend."""
return [
relax.backend.DispatchSampling(),
@@ -28,7 +28,7 @@ def library_dispatch_passes(target: tvm.target.Target): # pylint: disable=unuse
]
-def legalize_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def legalize_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default legalization passes for generic GPU backend."""
from tvm.s_tir import dlight as dl # pylint: disable=import-outside-toplevel
@@ -48,7 +48,7 @@ def legalize_passes(target: tvm.target.Target): # pylint: disable=unused-argume
]
-def dataflow_lower_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def dataflow_lower_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default dataflow lowering passes for generic GPU backend."""
return [
relax.transform.RewriteDataflowReshape(),
@@ -58,7 +58,7 @@ def dataflow_lower_passes(target: tvm.target.Target): # pylint: disable=unused-
]
-def finalize_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def finalize_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default finalization passes for generic GPU backend."""
return [
relax.transform.StaticPlanBlockMemory(),
@@ -71,13 +71,13 @@ def finalize_passes(target: tvm.target.Target): # pylint: disable=unused-argume
]
-def get_default_pipeline(target: tvm.target.Target):
+def get_default_pipeline(target: tvm_tirx.target.Target):
"""Return the default compilation pipeline for generic GPU."""
- @tvm.ir.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext):
+ @tvm_tirx.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm_tirx.ir.IRModule, _ctx: tvm_tirx.ir.transform.PassContext):
with target:
- seq = tvm.ir.transform.Sequential(
+ seq = tvm_tirx.ir.transform.Sequential(
library_dispatch_passes(target)
+ legalize_passes(target)
+ dataflow_lower_passes(target)
diff --git a/python/tvm/relax/backend/gpu_generic/sampling.py b/python/tvm/relax/backend/gpu_generic/sampling.py
index 9f6d054a3f..6e9d2fe499 100644
--- a/python/tvm/relax/backend/gpu_generic/sampling.py
+++ b/python/tvm/relax/backend/gpu_generic/sampling.py
@@ -20,10 +20,10 @@
import math
from collections.abc import Callable
-import tvm
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-from tvm.tirx import Function
+import tvm, tvm_tirx
+from tvm_tirx.script import s_tir as Ts
+from tvm_tirx.script import tirx as T
+from tvm_tirx.tirx import Function
def _is_power_of_two(n: int):
@@ -72,7 +72,7 @@ def gpu_multinomial_from_uniform(
The generated function
"""
- target = tvm.target.Target.current()
+ target = tvm_tirx.target.Target.current()
target_dtype = "int32" if "webgpu" in str(target) else "int64"
TX = T.int64(tx_len) # threadIdx.x
diff --git a/python/tvm/relax/backend/metal/coreml.py b/python/tvm/relax/backend/metal/coreml.py
index fa25e1a900..185ad7e26e 100644
--- a/python/tvm/relax/backend/metal/coreml.py
+++ b/python/tvm/relax/backend/metal/coreml.py
@@ -22,15 +22,15 @@ import shutil
import tvm_ffi
-import tvm
+import tvm, tvm_tirx
from tvm.contrib import coreml_runtime
-from tvm.ir import Call, GenericConst, PrimType
+from tvm_tirx.ir import Call, GenericConst, PrimType
from tvm.relax import transform
from tvm.relax.dpl.pattern import is_op, wildcard
from tvm.relax.expr import BindingBlock, Function, SeqExpr, Var, VarBinding
from tvm.relax.transform import PatternCheckContext
from tvm.relax.type import TensorType
-from tvm.support.xcode import compile_coreml
+from tvm_tirx.support.xcode import compile_coreml
from ...expr_functor import PyExprVisitor, visitor
from ..pattern_registry import get_patterns_with_prefix, register_patterns
@@ -150,12 +150,12 @@ def partition_for_coreml(mod):
Parameters
----------
- mod: tvm.IRModule
+ mod: tvm_tirx.IRModule
The IRModule to be partitioned.
Returns
-------
- mod: tvm.IRModule
+ mod: tvm_tirx.IRModule
The resulting IRModule, containing partitioned subgraphs to be
offloaded to the coreml backend.
"""
@@ -286,7 +286,7 @@ class CallNodeInfoCollector(PyExprVisitor):
def visit_call_(self, call: Call) -> None:
self.attrs.append(call.attrs)
for arg in call.args:
- if tvm.ir.is_prim_expr(arg):
+ if tvm_tirx.ir.is_prim_expr(arg):
self.primvals.append(arg)
if isinstance(arg, GenericConst):
self.consts.append(arg)
@@ -319,7 +319,7 @@ class CodegenCoreML(PyExprVisitor):
self.model_inputs_ = []
self.buf_idx_ = 0
- getter = tvm.get_global_func("relax.analysis.get_var2val")
+ getter = tvm_tirx.get_global_func("relax.analysis.get_var2val")
assert getter, "Cannot find `relax.analysis.get_var2val` function."
self.var2val = getter(function)
@@ -484,6 +484,6 @@ def coreml_compiler(funcs, options, constant_names):
shutil.rmtree(mlmodelc_path)
builder.compile(model_dir)
- dev = tvm.cpu(0)
+ dev = tvm_tirx.cpu(0)
compiled_funcs.append(coreml_runtime.create(name, mlmodelc_path, dev).module)
return compiled_funcs
diff --git a/python/tvm/relax/backend/patterns.py b/python/tvm/relax/backend/patterns.py
index 740ca48d2d..75de6ae575 100644
--- a/python/tvm/relax/backend/patterns.py
+++ b/python/tvm/relax/backend/patterns.py
@@ -28,8 +28,8 @@ from tvm.relax.dpl.pattern import (
is_tuple_get_item,
wildcard,
)
-from tvm.script import relax as R
-from tvm.script import tirx as T
+from tvm_tirx.script import relax as R
+from tvm_tirx.script import tirx as T
def _with_bias_activation_pattern(
diff --git a/python/tvm/relax/backend/rocm/hipblas.py b/python/tvm/relax/backend/rocm/hipblas.py
index 75c47321ee..942f399e3f 100644
--- a/python/tvm/relax/backend/rocm/hipblas.py
+++ b/python/tvm/relax/backend/rocm/hipblas.py
@@ -20,7 +20,7 @@
import operator
from functools import reduce
-import tvm
+import tvm, tvm_tirx
from tvm.relax import transform
from tvm.relax.transform import PatternCheckContext
@@ -56,7 +56,7 @@ def _check_matmul(context: PatternCheckContext) -> bool:
lhs_shape = lhs.ty.shape.values
rhs_shape = rhs.ty.shape.values
- if not isinstance(lhs_shape[-1], tvm.tirx.expr.IntImm | int):
+ if not isinstance(lhs_shape[-1], tvm_tirx.tirx.expr.IntImm | int):
# Reduction axis must be constant
return False
@@ -75,7 +75,7 @@ def _check_matmul(context: PatternCheckContext) -> bool:
bias = context.annotated_expr["bias"]
bias_shape = bias.ty.shape.values
bias_batches = reduce(operator.mul, bias_shape[:-1], 1)
- if not isinstance(bias_batches, tvm.tirx.expr.IntImm | int) or int(bias_batches) > 1:
+ if not isinstance(bias_batches, tvm_tirx.tirx.expr.IntImm | int) or int(bias_batches) > 1:
# hipblas only supports bias vector
return False
@@ -84,8 +84,8 @@ def _check_matmul(context: PatternCheckContext) -> bool:
# must be equal. If lhs is batched but rhs is not, we can use the regular GEMM by
# flattening all batch axes into the M axis.
return (
- isinstance(lhs_batches, tvm.ir.Var)
- or isinstance(rhs_batches, tvm.ir.Var)
+ isinstance(lhs_batches, tvm_tirx.ir.Var)
+ or isinstance(rhs_batches, tvm_tirx.ir.Var)
or (int(lhs_batches) == int(rhs_batches))
or (lhs_batches >= 1 and rhs_batches == 1)
)
@@ -167,12 +167,12 @@ def partition_for_hipblas(mod):
Parameters
----------
- mod: tvm.IRModule
+ mod: tvm_tirx.IRModule
The IRModule to be partitioned.
Returns
-------
- mod: tvm.IRModule
+ mod: tvm_tirx.IRModule
The resulting IRModule, containing partitioned subgraphs to be
offloaded to the hipblas backend.
"""
diff --git a/python/tvm/relax/backend/rocm/pipeline.py b/python/tvm/relax/backend/rocm/pipeline.py
index f031f617da..03bd853c5a 100644
--- a/python/tvm/relax/backend/rocm/pipeline.py
+++ b/python/tvm/relax/backend/rocm/pipeline.py
@@ -16,11 +16,11 @@
# under the License.
"""The Relax ROCm backend compilation pipeline and other passes."""
-import tvm
+import tvm, tvm_tirx
from tvm import relax
-def library_dispatch_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def library_dispatch_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default library dispatch passes for ROCm backend."""
return [
relax.backend.DispatchSampling(),
@@ -28,7 +28,7 @@ def library_dispatch_passes(target: tvm.target.Target): # pylint: disable=unuse
]
-def legalize_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def legalize_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default legalization passes for ROCm backend."""
from tvm.s_tir import dlight as dl # pylint: disable=import-outside-toplevel
@@ -48,7 +48,7 @@ def legalize_passes(target: tvm.target.Target): # pylint: disable=unused-argume
]
-def dataflow_lower_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def dataflow_lower_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default dataflow lowering passes for ROCm backend."""
return [
relax.transform.RewriteDataflowReshape(),
@@ -58,7 +58,7 @@ def dataflow_lower_passes(target: tvm.target.Target): # pylint: disable=unused-
]
-def finalize_passes(target: tvm.target.Target): # pylint: disable=unused-argument
+def finalize_passes(target: tvm_tirx.target.Target): # pylint: disable=unused-argument
"""The default finalization passes for ROCm backend."""
return [
relax.transform.StaticPlanBlockMemory(),
@@ -71,13 +71,13 @@ def finalize_passes(target: tvm.target.Target): # pylint: disable=unused-argume
]
-def get_default_pipeline(target: tvm.target.Target):
+def get_default_pipeline(target: tvm_tirx.target.Target):
"""Return the default compilation pipeline for ROCm."""
- @tvm.ir.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext):
+ @tvm_tirx.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm_tirx.ir.IRModule, _ctx: tvm_tirx.ir.transform.PassContext):
with target:
- seq = tvm.ir.transform.Sequential(
+ seq = tvm_tirx.ir.transform.Sequential(
library_dispatch_passes(target)
+ legalize_passes(target)
+ dataflow_lower_passes(target)
diff --git a/python/tvm/relax/backend/utils.py b/python/tvm/relax/backend/utils.py
index 47c01ac431..bcb6ea50d7 100644
--- a/python/tvm/relax/backend/utils.py
+++ b/python/tvm/relax/backend/utils.py
@@ -20,7 +20,7 @@
from tvm import relax
from tvm.relax import DataflowVar, PyExprMutator
from tvm.relax.transform import PatternCheckContext
-from tvm.target import Target
+from tvm_tirx.target import Target
class BackendDispatcher(PyExprMutator):
diff --git a/python/tvm/relax/base_py_module.py b/python/tvm/relax/base_py_module.py
index cf5f56824f..192710b3ad 100644
--- a/python/tvm/relax/base_py_module.py
+++ b/python/tvm/relax/base_py_module.py
@@ -26,11 +26,12 @@ from typing import Any, Optional, Union
import numpy as np
from tvm_ffi import Function
-import tvm
-from tvm import relax, tirx
-from tvm.ir import IRModule
-from tvm.runtime import Device, Tensor
-from tvm.target import Target
+import tvm, tvm_tirx
+from tvm_tirx import tirx
+from tvm import relax
+from tvm_tirx.ir import IRModule
+from tvm_tirx.runtime import Device, Tensor
+from tvm_tirx.target import Target
try:
from torch.utils.dlpack import to_dlpack as to_dlpack_legacy
@@ -82,7 +83,7 @@ class BasePyModule:
# globals this needs may already be cleared.
if not owned or sys.is_finalizing():
return
- unregister_py_func = tvm.get_global_func("vm.builtin.unregister_py_func")
+ unregister_py_func = tvm_tirx.get_global_func("vm.builtin.unregister_py_func")
for func_name in owned:
unregister_py_func(func_name)
except Exception: # pylint: disable=broad-except
@@ -158,7 +159,7 @@ class BasePyModule:
def _compile_functions(self):
"""Compile TIR and Relax functions using JIT compilation."""
# Compile TIR functions first
- tir_mod = tvm.IRModule(
+ tir_mod = tvm_tirx.IRModule(
{
gv: func
for gv, func in self.ir_mod.functions_items()
@@ -222,7 +223,7 @@ class BasePyModule:
return
try:
- register_py_func = tvm.get_global_func("vm.builtin.register_py_func")
+ register_py_func = tvm_tirx.get_global_func("vm.builtin.register_py_func")
except ValueError:
return
@@ -302,7 +303,7 @@ class BasePyModule:
if func_name not in self.extern_funcs:
try:
- self.extern_funcs[func_name] = tvm.get_global_func(func_name)
+ self.extern_funcs[func_name] = tvm_tirx.get_global_func(func_name)
except ValueError as error:
raise ValueError(
f"Function '{func_name}' not found as a global function. "
@@ -422,14 +423,14 @@ class BasePyModule:
if _FASTER_DLPACK_EXTENSION is not None:
try:
dlpack = torch.to_dlpack(tensor)
- return tvm.runtime.from_dlpack(dlpack)
+ return tvm_tirx.runtime.from_dlpack(dlpack)
except (AttributeError, ValueError):
pass # Fall through to the next method
# 2. Try modern `torch.to_dlpack` (preferred for PyTorch >= 1.7)
try:
dlpack = torch.to_dlpack(tensor)
- return tvm.runtime.from_dlpack(dlpack)
+ return tvm_tirx.runtime.from_dlpack(dlpack)
except (AttributeError, ValueError):
pass # Fall through to the next method
@@ -437,7 +438,7 @@ class BasePyModule:
if to_dlpack_legacy:
try:
dlpack = to_dlpack_legacy(tensor)
- return tvm.runtime.from_dlpack(dlpack)
+ return tvm_tirx.runtime.from_dlpack(dlpack)
except (AttributeError, ValueError) as error_legacy:
print(
f"Warning: Legacy DLPack conversion failed ({error_legacy}), "
@@ -446,12 +447,12 @@ class BasePyModule:
# 4. If all DLPack methods fail, use numpy fallback
numpy_array = tensor.detach().cpu().numpy()
- return tvm.runtime.tensor(numpy_array, device=self.device)
+ return tvm_tirx.runtime.tensor(numpy_array, device=self.device)
# For other types (like scalars, lists), convert to numpy first
try:
numpy_array = np.array(tensor, dtype=np.float32)
- return tvm.runtime.tensor(numpy_array, device=self.device)
+ return tvm_tirx.runtime.tensor(numpy_array, device=self.device)
except (TypeError, ValueError) as error:
raise TypeError(
f"Unsupported type for conversion to TVM Tensor: {type(tensor)}"
@@ -638,7 +639,7 @@ class BasePyModule:
This method extends the standard IRModule show() method to handle
Python functions stored in the IRModule's pyfuncs attribute.
"""
- from tvm.script.printer.highlight import cprint # pylint: disable=import-outside-toplevel
+ from tvm_tirx.script.printer.highlight import cprint # pylint: disable=import-outside-toplevel
if black_format is None:
env = os.environ.get("TVM_BLACK_FORMAT")
diff --git a/python/tvm/relax/binding_rewrite.py b/python/tvm/relax/binding_rewrite.py
index bd46029243..b62451afe6 100644
--- a/python/tvm/relax/binding_rewrite.py
+++ b/python/tvm/relax/binding_rewrite.py
@@ -19,8 +19,8 @@
import tvm_ffi
-import tvm
-from tvm.runtime import Object
+import tvm, tvm_tirx
+from tvm_tirx.runtime import Object
from . import Binding, DataflowBlock, Expr, Function, Var, _ffi_api
@@ -138,18 +138,18 @@ class DataflowBlockRewrite(Object):
ret.__name__ = self.func_name
return ret
- def mutate_irmodule(self, irmodule: tvm.IRModule) -> tvm.IRModule:
+ def mutate_irmodule(self, irmodule: tvm_tirx.IRModule) -> tvm_tirx.IRModule:
"""
Return an updated IRModule by replacing the old function with the mutated root function.
Parameters
----------
- irmodule : tvm.IRModule
+ irmodule : tvm_tirx.IRModule
The base IRModule to update.
Returns
-------
- tvm.IRModule
+ tvm_tirx.IRModule
The updated IRModule.
"""
ret = _ffi_api.dfb_rewrite_mutate_irmodule(self, irmodule) # type: ignore
diff --git a/python/tvm/relax/block_builder.py b/python/tvm/relax/block_builder.py
index a4d82777fd..14677bb40e 100644
--- a/python/tvm/relax/block_builder.py
+++ b/python/tvm/relax/block_builder.py
@@ -23,11 +23,11 @@ from typing import Any, Optional
import tvm_ffi
-import tvm
+import tvm, tvm_tirx
from tvm import relax as rx
-from tvm import tirx
-from tvm.ir.module import IRModule
-from tvm.runtime import Object
+from tvm_tirx import tirx
+from tvm_tirx.ir.module import IRModule
+from tvm_tirx.runtime import Object
from . import _ffi_api
from .expr import BaseFunc, Binding, BindingBlock, Expr, GlobalVar, Tuple, Var
@@ -87,7 +87,7 @@ class TestingScope:
self._bb = block_builder
shape_vars = []
for var in def_vars:
- if isinstance(var, tvm.ir.Var):
+ if isinstance(var, tvm_tirx.ir.Var):
shape_vars.append(var)
else:
raise ValueError("def_vars can only contain Vars")
@@ -479,7 +479,7 @@ class BlockBuilder(Object):
m = T.dynamic("m", "int64")
n = T.dynamic("n", "int64")
- @tvm.script.ir_module
+ @tvm_tirx.script.ir_module
class Module:
@Ts.function
def te_func(
@@ -530,7 +530,7 @@ class BlockBuilder(Object):
m = T.dynamic("m", "int64")
n = T.dynamic("n", "int64")
- @tvm.script.ir_module
+ @tvm_tirx.script.ir_module
class Module:
@Ts.function
def te_func(
@@ -621,7 +621,7 @@ class BlockBuilder(Object):
Returns
-------
- gvar: tvm.ir.GlobalVar
+ gvar: tvm_tirx.ir.GlobalVar
A GlobalVar representing the function
"""
@@ -685,18 +685,18 @@ class BlockBuilder(Object):
"""
return _ffi_api.BlockBuilderNormalize(self, expr) # type: ignore
- def get(self) -> tvm.IRModule:
+ def get(self) -> tvm_tirx.IRModule:
"""Return intermediate IRModule. For the situation where the IRModule is needed in the
middle of a building process.
Returns
-------
- ret : tvm.IRModule
+ ret : tvm_tirx.IRModule
An IRModule with Relax and TIR functions being built.
"""
return _ffi_api.BlockBuilderGetContextIRModule(self) # type: ignore
- def finalize(self) -> tvm.IRModule:
+ def finalize(self) -> tvm_tirx.IRModule:
"""Finalize the building process and return the result IRModule.
Possibly rename GlobalVars in the IRModule to ensure name uniqueness and the invariant:
@@ -708,7 +708,7 @@ class BlockBuilder(Object):
Returns
-------
- ret : tvm.IRModule
+ ret : tvm_tirx.IRModule
An IRModule with Relax and TIR functions being built.
"""
return _ffi_api.BlockBuilderFinalize(self) # type: ignore
diff --git a/python/tvm/relax/distributed/global_info.py b/python/tvm/relax/distributed/global_info.py
index 125e6652f5..eaf388b432 100644
--- a/python/tvm/relax/distributed/global_info.py
+++ b/python/tvm/relax/distributed/global_info.py
@@ -20,8 +20,8 @@
import tvm_ffi
from tvm_ffi import Shape
-from tvm.ir import Range
-from tvm.ir.global_info import GlobalInfo
+from tvm_tirx.ir import Range
+from tvm_tirx.ir.global_info import GlobalInfo
from . import _ffi_api as ffi
diff --git a/python/tvm/relax/distributed/transform/transform.py b/python/tvm/relax/distributed/transform/transform.py
index 87721e0b11..f2da847e80 100644
--- a/python/tvm/relax/distributed/transform/transform.py
+++ b/python/tvm/relax/distributed/transform/transform.py
@@ -17,52 +17,52 @@
# pylint: disable=invalid-name
"""Relax distributed-related transformation passes."""
-import tvm.ir
+import tvm_tirx.ir, tvm
from . import _ffi_api
-def PropagateSharding() -> tvm.ir.transform.Pass:
+def PropagateSharding() -> tvm_tirx.ir.transform.Pass:
"""Propagate sharding information.
Returns
-------
- ret : tvm.ir.transform.Pass
+ ret : tvm_tirx.ir.transform.Pass
The registered pass
"""
return _ffi_api.PropagateSharding() # type: ignore
-def LowerGlobalViewToLocalView() -> tvm.ir.transform.Pass:
+def LowerGlobalViewToLocalView() -> tvm_tirx.ir.transform.Pass:
"""Lower global view TIR to local view
Returns
-------
- ret : tvm.ir.transform.Pass
+ ret : tvm_tirx.ir.transform.Pass
The registered pass
"""
return _ffi_api.LowerGlobalViewToLocalView() # type: ignore
-def LegalizeRedistribute() -> tvm.ir.transform.Pass:
+def LegalizeRedistribute() -> tvm_tirx.ir.transform.Pass:
"""Legalize redistribute op to ccl op.
S->R: R.ccl.allgather
R->S: R.dist.redistribute_replica_to_shard
Returns
-------
- ret : tvm.ir.transform.Pass
+ ret : tvm_tirx.ir.transform.Pass
The registered pass
"""
return _ffi_api.LegalizeRedistribute() # type: ignore
-def LowerDistIR() -> tvm.ir.transform.Pass:
+def LowerDistIR() -> tvm_tirx.ir.transform.Pass:
"""Lower DistIR to Relax
Returns
-------
- ret : tvm.ir.transform.Pass
+ ret : tvm_tirx.ir.transform.Pass
The registered pass
"""
return _ffi_api.LowerDistIR() # type: ignore
diff --git a/python/tvm/relax/distributed/type.py b/python/tvm/relax/distributed/type.py
index 2d44c47e8b..3873b3df53 100644
--- a/python/tvm/relax/distributed/type.py
+++ b/python/tvm/relax/distributed/type.py
@@ -21,10 +21,10 @@ import enum
import tvm_ffi
-from tvm.ir import Location
-from tvm.ir.location import UNKNOWN_LOC
+from tvm_tirx.ir import Location
+from tvm_tirx.ir.location import UNKNOWN_LOC
from tvm.relax.type import TensorType, Type
-from tvm.runtime import Object
+from tvm_tirx.runtime import Object
from . import _ffi_api
from .global_info import DeviceMesh
diff --git a/python/tvm/relax/dpl/context.py b/python/tvm/relax/dpl/context.py
index adc955ebbc..232fd97523 100644
--- a/python/tvm/relax/dpl/context.py
+++ b/python/tvm/relax/dpl/context.py
@@ -17,14 +17,14 @@
"""The Graph Matching Context Manager for Dataflow Pattern Language."""
-import tvm
+import tvm, tvm_tirx
from ..expr import DataflowBlock, Var
from . import _ffi as ffi
from .pattern import DFPattern
-class PatternContext(tvm.runtime.Object):
+class PatternContext(tvm_tirx.runtime.Object):
"""A context object for doing graph (topogical) pattern matching."""
def __init__(self, incremental=False):
diff --git a/python/tvm/relax/dpl/pattern.py b/python/tvm/relax/dpl/pattern.py
index caa04bb8f0..517eb16cee 100644
--- a/python/tvm/relax/dpl/pattern.py
+++ b/python/tvm/relax/dpl/pattern.py
@@ -25,12 +25,12 @@ from typing import Union
import tvm_ffi
from tvm_ffi import Array
-import tvm
-from tvm.ir.op import Op
+import tvm, tvm_tirx
+from tvm_tirx.ir.op import Op
-from ...ir import make_node
-from ...ir.base import Node
-from ...runtime import Object
+from tvm_tirx.ir import make_node
+from tvm_tirx.ir.base import Node
+from tvm_tirx.runtime import Object
from ..expr import Expr, Var
from . import _ffi as ffi
@@ -160,7 +160,7 @@ class DFPattern(Node):
has_shape assumes that the matched relax.Expr only has one
output tensor. Use is_tuple for those with multiple outputs.
"""
- if not isinstance(shape, list | tuple) and not tvm.ir.is_prim_expr(shape):
+ if not isinstance(shape, list | tuple) and not tvm_tirx.ir.is_prim_expr(shape):
raise ValueError("has_shape takes a list or tuple as input.")
return ShapePattern(pattern=self, shape=shape)
@@ -613,11 +613,11 @@ class ShapePattern(DFPattern):
pattern: tvm.relax.dpl.DFPattern
The input pattern that needs type annotation.
- shape: List[tvm.ir.Expr]
+ shape: List[tvm_tirx.ir.Expr]
The shape to match.
"""
- def __init__(self, pattern: "DFPattern", shape: list[tvm.ir.Expr]):
+ def __init__(self, pattern: "DFPattern", shape: list[tvm_tirx.ir.Expr]):
self.__init_handle_by_constructor__(ffi.ShapePattern, pattern, shape) # type: ignore
@@ -642,11 +642,11 @@ class PrimArrPattern(DFPattern):
Parameters
----------
- shape : List[tvm.ir.Expr]
+ shape : List[tvm_tirx.ir.Expr]
The shape to match.
"""
- def __init__(self, shape: list[tvm.ir.Expr]):
+ def __init__(self, shape: list[tvm_tirx.ir.Expr]):
self.__init_handle_by_constructor__(ffi.PrimArrPattern, shape) # type: ignore
def __getitem__(self, index: int):
@@ -668,11 +668,11 @@ class AttrPattern(DFPattern):
pattern: tvm.relax.dpl.DFPattern
The input pattern.
- attrs: tvm.ir.attrs.Attrs
+ attrs: tvm_tirx.ir.attrs.Attrs
The attributes to match.
"""
- def __init__(self, pattern: "DFPattern", attrs: tvm.ir.attrs.Attrs):
+ def __init__(self, pattern: "DFPattern", attrs: tvm_tirx.ir.attrs.Attrs):
self.__init_handle_by_constructor__(ffi.AttrPattern, pattern, attrs) # type: ignore
@@ -744,7 +744,7 @@ def is_op(op_name: str) -> ExprPattern:
Parameters
----------
op_name: String
- The name of the tvm.ir.op.Op object
+ The name of the tvm_tirx.ir.op.Op object
Returns
-------
@@ -830,13 +830,13 @@ def has_dtype(dtype: str, pattern: DFPattern = None) -> DataTypePattern:
return DataTypePattern(pattern, dtype)
-def is_shape(shape: list[tvm.ir.Expr]) -> "PrimArrPattern":
+def is_shape(shape: list[tvm_tirx.ir.Expr]) -> "PrimArrPattern":
"""
Directly matches a shape which is an array of Expr
Parameters
----------
- shape : List[tvm.ir.Expr]
+ shape : List[tvm_tirx.ir.Expr]
The expected shape
Returns
diff --git a/python/tvm/relax/dpl/rewrite.py b/python/tvm/relax/dpl/rewrite.py
index a0ac3665a3..c2a848c3ca 100644
--- a/python/tvm/relax/dpl/rewrite.py
+++ b/python/tvm/relax/dpl/rewrite.py
@@ -20,8 +20,8 @@ from collections.abc import Callable
from tvm_ffi import register_object
-from tvm.ir import IRModule
-from tvm.runtime import Object
+from tvm_tirx.ir import IRModule
+from tvm_tirx.runtime import Object
from ..expr import Expr, Function, Var
from . import _ffi as ffi
diff --git a/python/tvm/relax/exec_builder.py b/python/tvm/relax/exec_builder.py
index e6cdd8ac81..3768bc7786 100644
--- a/python/tvm/relax/exec_builder.py
+++ b/python/tvm/relax/exec_builder.py
@@ -23,7 +23,7 @@ from enum import IntEnum
import tvm_ffi
from tvm_ffi import Shape
-import tvm
+import tvm, tvm_tirx
from . import _ffi_api
from .vm_build import VMExecutable
@@ -110,7 +110,7 @@ class ExecBuilder(tvm_ffi.core.Object):
def emit_call(
self,
name: str,
- args: list[tvm.runtime.Tensor | tvm.DataType] | None = None,
+ args: list[tvm_tirx.runtime.Tensor | tvm_tirx.DataType] | None = None,
dst: int | None = None,
) -> None:
"""emit a call instruction which calls a packed function."""
@@ -124,7 +124,7 @@ class ExecBuilder(tvm_ffi.core.Object):
shape_tuple = Shape(arg)
new_arg = self.convert_constant(shape_tuple)
args_.append(new_arg)
- elif isinstance(arg, tvm.runtime.Tensor | tvm.DataType | Shape):
+ elif isinstance(arg, tvm_tirx.runtime.Tensor | tvm_tirx.DataType | Shape):
new_arg = self.convert_constant(arg)
args_.append(new_arg)
else:
diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py
index dbcf98b485..442243fe2a 100644
--- a/python/tvm/relax/expr.py
+++ b/python/tvm/relax/expr.py
@@ -26,23 +26,23 @@ import numpy as _np # type: ignore
import tvm_ffi
from tvm_ffi.core import String
-import tvm.ir
+import tvm_tirx.ir, tvm
import tvm.relax
-import tvm.runtime
-from tvm import DataType
-from tvm.ir import Scriptable
-from tvm.ir.location import UNKNOWN_LOC, UnknownLoc
+import tvm_tirx.runtime, tvm
+from tvm_tirx import DataType
+from tvm_tirx.ir import Scriptable
+from tvm_tirx.ir.location import UNKNOWN_LOC, UnknownLoc
-from ..ir import BaseFunc, Location, Node
-from ..ir.expr import _CallableExprWithOp
+from tvm_tirx.ir import BaseFunc, Location, Node
+from tvm_tirx.ir.expr import _CallableExprWithOp
from . import _ffi_api
# It is a workaround for mypy: https://github.com/python/mypy/issues/7866#issuecomment-549454370
# This feature is not supported until python 3.10:
# https://docs.python.org/3.10/whatsnew/3.10.html#pep-613-typealias
-Expr = tvm.ir.Expr
-Type = tvm.ir.Type # pylint: disable=invalid-name
-GlobalVar = tvm.ir.GlobalVar
+Expr = tvm_tirx.ir.Expr
+Type = tvm_tirx.ir.Type # pylint: disable=invalid-name
+GlobalVar = tvm_tirx.ir.GlobalVar
def prim_value(value: Expr | int | float, dtype: str | None = None) -> Expr:
@@ -71,16 +71,16 @@ def prim_value(value: Expr | int | float, dtype: str | None = None) -> Expr:
it uses bool, int64, or float64. The script alias ``R.prim_value`` retains
these programmatic conversion semantics.
"""
- if tvm.ir.is_prim_expr(value):
+ if tvm_tirx.ir.is_prim_expr(value):
return value
if isinstance(value, bool | _np.bool_):
- return tvm.tirx.IntImm(dtype or "bool", int(value))
+ return tvm_tirx.tirx.IntImm(dtype or "bool", int(value))
if isinstance(value, Integral):
- return tvm.tirx.IntImm(dtype or "int64", int(value))
+ return tvm_tirx.tirx.IntImm(dtype or "int64", int(value))
if isinstance(value, Real):
- return tvm.tirx.FloatImm(dtype or "float64", float(value))
+ return tvm_tirx.tirx.FloatImm(dtype or "float64", float(value))
tvm_value = tvm_ffi.convert(value)
- if tvm.ir.is_prim_expr(tvm_value):
+ if tvm_tirx.ir.is_prim_expr(tvm_value):
return tvm_value
raise TypeError(f"Cannot convert {value} with type {type(value)} to `Expr`")
@@ -98,7 +98,7 @@ Type.is_base_of = _relax_type_is_base_of # type: ignore[attr-defined]
def _is_tensor_or_missing_type(ty: Type) -> bool:
- return isinstance(ty, tvm.relax.TensorType | tvm.ir.MissingType)
+ return isinstance(ty, tvm.relax.TensorType | tvm_tirx.ir.MissingType)
def _binary_op_helper(lhs: Expr, rhs: Expr, op: Callable):
@@ -160,8 +160,8 @@ class IfExpr(_CallableExprWithOp):
# Compatibility aliases. Tuple expressions are owned by the common IR.
-Tuple = tvm.ir.Tuple
-TupleGetItem = tvm.ir.TupleGetItem
+Tuple = tvm_tirx.ir.Tuple
+TupleGetItem = tvm_tirx.ir.TupleGetItem
@tvm_ffi.register_object("relax.expr.ShapeExpr")
@@ -209,7 +209,7 @@ def make_shape(shape: list[Any] | tuple[Any, ...]) -> ShapeExpr:
# Ordinary Relax bindings use the canonical IR Var directly.
-Var = tvm.ir.Var
+Var = tvm_tirx.ir.Var
@tvm_ffi.register_object("relax.expr.DataflowVar")
@@ -250,7 +250,7 @@ class DataflowVar(Var):
raise TypeError("name must be a str")
if ty is not None:
- ty = tvm.runtime.convert(ty)
+ ty = tvm_tirx.runtime.convert(ty)
if not isinstance(ty, Type):
raise TypeError(
"ty needs to be an instance of Type. "
@@ -372,7 +372,7 @@ class Function(BaseFunc, Scriptable):
body: Expr
ret_ty: Type
is_pure: bool
- attrs: tvm.ir.DictAttrs
+ attrs: tvm_tirx.ir.DictAttrs
loc: Location
def __init__(
@@ -381,11 +381,11 @@ class Function(BaseFunc, Scriptable):
body: Expr,
ret_ty: Type | None = None,
is_pure: bool | None = True,
- attrs: tvm.ir.DictAttrs | None = None,
+ attrs: tvm_tirx.ir.DictAttrs | None = None,
loc: Location = UNKNOWN_LOC,
) -> None:
if attrs is None:
- attrs = tvm.ir.DictAttrs({})
+ attrs = tvm_tirx.ir.DictAttrs({})
self.__init_handle_by_constructor__(
_ffi_api.Function,
params,
@@ -401,12 +401,12 @@ class Function(BaseFunc, Scriptable):
params: list[Var],
ret_ty: Type,
is_pure: bool | None = True,
- attrs: tvm.ir.DictAttrs | None = None,
+ attrs: tvm_tirx.ir.DictAttrs | None = None,
loc: Location = UNKNOWN_LOC,
):
"""Construct a relax.Function but without body"""
if attrs is None:
- attrs = tvm.ir.DictAttrs({})
+ attrs = tvm_tirx.ir.DictAttrs({})
return _ffi_api.FunctionCreateEmpty(params, ret_ty, is_pure, attrs, loc) # type: ignore
def __call__(self, *args):
@@ -417,14 +417,14 @@ class Function(BaseFunc, Scriptable):
args: List[relax.Expr]
Arguments.
"""
- return tvm.ir.Call(self, args, None, None)
+ return tvm_tirx.ir.Call(self, args, None, None)
- def bind_symbolic_vars(self, binding_map: Mapping[str | tvm.tirx.Var, Expr]) -> "Function":
+ def bind_symbolic_vars(self, binding_map: Mapping[str | tvm_tirx.tirx.Var, Expr]) -> "Function":
"""Return a new function with updated symbolic variable
Parameters
----------
- binding_map: Mapping[str | tvm.tirx.Var, Expr]
+ binding_map: Mapping[str | tvm_tirx.tirx.Var, Expr]
The mapping of values to be replaced. Keys may be either
a `tirx.Var` or a string name of the variable. If the
@@ -441,7 +441,7 @@ class Function(BaseFunc, Scriptable):
# Relax uses int64 for symbolic variables, but the FFI
# converts python integers into int32.
binding_map = {
- key: tvm.tirx.const(value, "int64") if isinstance(value, int) else value
+ key: tvm_tirx.tirx.const(value, "int64") if isinstance(value, int) else value
for key, value in binding_map.items()
}
@@ -451,7 +451,7 @@ class Function(BaseFunc, Scriptable):
self,
binding_map: Mapping[
str | Var,
- int | float | Expr | tvm.runtime.Tensor | _np.ndarray,
+ int | float | Expr | tvm_tirx.runtime.Tensor | _np.ndarray,
],
) -> "Function":
"""Return a new function with updated symbolic variable
@@ -460,7 +460,7 @@ class Function(BaseFunc, Scriptable):
----------
binding_map: Mapping[
str | Var,
- int | float | Expr | tvm.runtime.Tensor | _np.ndarray,
+ int | float | Expr | tvm_tirx.runtime.Tensor | _np.ndarray,
]
The mapping of values to be replaced.
@@ -487,8 +487,8 @@ class Function(BaseFunc, Scriptable):
if isinstance(value, int):
# Relax uses int64 for symbolic variables, but the FFI
# converts python integers into int32.
- return tvm.tirx.const(value, "int64")
- elif isinstance(value, _np.ndarray | tvm.runtime.Tensor):
+ return tvm_tirx.tirx.const(value, "int64")
+ elif isinstance(value, _np.ndarray | tvm_tirx.runtime.Tensor):
return tvm.relax.const(value)
else:
return value
@@ -498,7 +498,7 @@ class Function(BaseFunc, Scriptable):
return _ffi_api.FunctionBindParams(self, binding_map) # type: ignore
def inline_functions(
- self, function_map: Mapping[str | tvm.ir.GlobalVar, "Function"]
+ self, function_map: Mapping[str | tvm_tirx.ir.GlobalVar, "Function"]
) -> "Function":
return _ffi_api.FunctionInlineFunctions(self, function_map) # type: ignore
@@ -530,13 +530,13 @@ def extern(name: str, ty: Type | None = None, loc: Location = UNKNOWN_LOC):
def const(
- value: bool | int | float | _np.ndarray | tvm.runtime.Tensor, dtype: str | None = None
-) -> tvm.ir.GenericConst:
+ value: bool | int | float | _np.ndarray | tvm_tirx.runtime.Tensor, dtype: str | None = None
+) -> tvm_tirx.ir.GenericConst:
"""Create a constant value.
Parameters
----------
- value: bool | int | float | numpy.ndarray | tvm.runtime.Tensor
+ value: bool | int | float | numpy.ndarray | tvm_tirx.runtime.Tensor
The constant value.
dtype: Optional[str]
@@ -554,7 +554,7 @@ def const(
# Needed for bf16 and fp8 support (does not come with numpy)
import ml_dtypes # pylint: disable=unused-import,import-outside-toplevel
- if isinstance(dtype, tvm.ir.PrimType):
+ if isinstance(dtype, tvm_tirx.ir.PrimType):
dtype = dtype.dtype
if isinstance(value, Number | (bool | list)):
@@ -573,9 +573,9 @@ def const(
if isinstance(value, _np.ndarray | _np.generic):
if dtype is not None:
value = value.astype(dtype)
- value = tvm.runtime.tensor(value)
+ value = tvm_tirx.runtime.tensor(value)
- if not isinstance(value, tvm.runtime.Tensor):
+ if not isinstance(value, tvm_tirx.runtime.Tensor):
raise ValueError("value has to be scalar or Tensor")
return _ffi_api.MakeTensorConst(value, None, UNKNOWN_LOC)
@@ -588,7 +588,7 @@ class TEPlaceholderOp(tvm.te.tensor.Operation):
def te_tensor(
value: Expr,
- tir_var_map: dict[tvm.ir.Var, tvm.tirx.Expr],
+ tir_var_map: dict[tvm_tirx.ir.Var, tvm_tirx.tirx.Expr],
name: str = "rxplaceholder",
):
"""Create a TE tensor from a Relax expression, with primitive variables
@@ -599,7 +599,7 @@ def te_tensor(
value : Expr
The relax expression, which is required to have TensorType.
- tir_var_map : Dict[tvm.ir.Var, tvm.tirx.Expr]
+ tir_var_map : Dict[tvm_tirx.ir.Var, tvm_tirx.tirx.Expr]
The exact-identity mapping used to refresh canonical primitive Vars
that appear in the input shape.
diff --git a/python/tvm/relax/expr_functor.py b/python/tvm/relax/expr_functor.py
index bfd77b176e..ba38df59fe 100644
--- a/python/tvm/relax/expr_functor.py
+++ b/python/tvm/relax/expr_functor.py
@@ -23,12 +23,12 @@ from collections.abc import Callable
import tvm_ffi
from tvm import te as _te
-from tvm import tirx as _tirx
-from tvm.ir import Call, DataTypeImm, GenericConst, Op, StringImm, TensorLoad, is_prim_expr
-from tvm.runtime import Object
+from tvm_tirx import tirx as _tirx
+from tvm_tirx.ir import Call, DataTypeImm, GenericConst, Op, StringImm, TensorLoad, is_prim_expr
+from tvm_tirx.runtime import Object
from tvm.s_tir.utils import derived_object
-from ..ir.module import IRModule
+from tvm_tirx.ir.module import IRModule
from . import _ffi_api
from .block_builder import BlockBuilder
from .expr import (
diff --git a/python/tvm/relax/frontend/common.py b/python/tvm/relax/frontend/common.py
index b4aad8f778..fc20e93ccd 100644
--- a/python/tvm/relax/frontend/common.py
+++ b/python/tvm/relax/frontend/common.py
@@ -19,34 +19,34 @@
import numpy as _np
-import tvm
+import tvm, tvm_tirx
from tvm import topi
-def detach_params(mod: tvm.IRModule) -> tuple[tvm.IRModule, dict[str, list[tvm.runtime.Tensor]]]:
+def detach_params(mod: tvm_tirx.IRModule) -> tuple[tvm_tirx.IRModule, dict[str, list[tvm_tirx.runtime.Tensor]]]:
"""Detach the attribute "params" in the functions of the input IRModule as
separate dictionary of params.
Parameters
----------
- mod : tvm.IRModule
+ mod : tvm_tirx.IRModule
The IRModule whose functions' "param" attribute is going to be detached.
Returns
-------
- detached_mod : tvm.IRModule
+ detached_mod : tvm_tirx.IRModule
The IRModule after the detachment.
- params_dict : Dict[str, List[tvm.runtime.Tensor]]
+ params_dict : Dict[str, List[tvm_tirx.runtime.Tensor]]
The detached params. The dict keys corresponds to the names of the
functions in the input IRModule that have attribute "params".
"""
- detached_mod = tvm.IRModule()
+ detached_mod = tvm_tirx.IRModule()
params_dict = dict()
for gv, func in mod.functions_items():
if "params" in func.attrs:
params = list(func.attrs["params"])
- if not all([isinstance(param, tvm.runtime.Tensor) for param in params]):
+ if not all([isinstance(param, tvm_tirx.runtime.Tensor) for param in params]):
raise ValueError('The value "params" attribute is expected to be a list of Tensor.')
params_dict[gv.name_hint] = params
detached_mod[gv] = func.without_attr("params")
@@ -112,7 +112,7 @@ def autopad(
pad = _np.concatenate([_np.zeros([2, 2], dtype="int64"), pad], axis=0)
if pad_type not in ["constant", "edge", "reflect"]:
- raise tvm.error.OpAttributeInvalid(
+ raise tvm_tirx.error.OpAttributeInvalid(
"Value " + pad_type + ' in attribute "mode" is invalid for operator Pad.'
)
diff --git a/python/tvm/relax/frontend/nn/_tensor_op.py b/python/tvm/relax/frontend/nn/_tensor_op.py
index 6285f8008f..3c8abfb18d 100644
--- a/python/tvm/relax/frontend/nn/_tensor_op.py
+++ b/python/tvm/relax/frontend/nn/_tensor_op.py
@@ -17,7 +17,7 @@
# ruff: noqa: F821
"""Adding member operators to nn.Tensor."""
-from tvm import tirx
+from tvm_tirx import tirx
def _op():
diff --git a/python/tvm/relax/frontend/nn/core.py b/python/tvm/relax/frontend/nn/core.py
index e93152403f..476d721b7b 100644
--- a/python/tvm/relax/frontend/nn/core.py
+++ b/python/tvm/relax/frontend/nn/core.py
@@ -34,14 +34,14 @@ from typing import (
import numpy as np # type: ignore
-import tvm.runtime
-from tvm import tirx
-from tvm.ir import IRModule
-from tvm.ir.transform import Pass
-from tvm.runtime import Device
-from tvm.runtime import device as as_device
+import tvm_tirx.runtime, tvm
+from tvm_tirx import tirx
+from tvm_tirx.ir import IRModule
+from tvm_tirx.ir.transform import Pass
+from tvm_tirx.runtime import Device
+from tvm_tirx.runtime import device as as_device
from tvm.runtime.vm import VirtualMachine
-from tvm.target import Target
+from tvm_tirx.target import Target
from .... import relax as rx
from ...block_builder import BlockBuilder
@@ -159,7 +159,7 @@ class Tensor(_TensorOp):
never be created directly by users in usual cases, and the only exception is to indicate
the shape/dtype of return values of an external function.
- If shape is a string `name`, we create a symbolic shape `tvm.tirx.Var(name, "int64")`.
+ If shape is a string `name`, we create a symbolic shape `tvm_tirx.tirx.Var(name, "int64")`.
"""
new_shape = []
for expr in shape:
@@ -172,9 +172,9 @@ class Tensor(_TensorOp):
expr = tirx.Var(expr, "int64")
new_shape.append(expr)
continue
- if not tvm.ir.is_prim_expr(expr):
+ if not tvm_tirx.ir.is_prim_expr(expr):
raise TypeError(f"Invalid shape: {shape}")
- assert expr.ty == tvm.ir.PrimType("int64")
+ assert expr.ty == tvm_tirx.ir.PrimType("int64")
new_shape.append(expr)
return Tensor(
_expr=rx.Var(
@@ -190,8 +190,8 @@ class Tensor(_TensorOp):
def shape(self) -> list[int | tirx.Expr]:
"""Returns the shape of the tensor as a list of integers.
- An integer can be a python int or tvm.tirx.Expr, depending on whether the shape is
- fully static, for example, [1, 2, tvm.tirx.Var("n", "int64")] is a valid shape where
+ An integer can be a python int or tvm_tirx.tirx.Expr, depending on whether the shape is
+ fully static, for example, [1, 2, tvm_tirx.tirx.Var("n", "int64")] is a valid shape where
the last dimension is dynamic while the first two dimensions are always static constants.
Returns
@@ -253,7 +253,7 @@ class Parameter(Tensor):
----------
shape : Sequence[Union[int, str, tirx.Expr]]
The shape of the parameter. If it is a string `name`, we create a symbolic shape
- `tvm.tirx.Var(name, "int64")`.
+ `tvm_tirx.tirx.Var(name, "int64")`.
dtype : Optional[str]
The data type of the parameter. If not specified, the default dtype will be used.
"""
@@ -266,14 +266,14 @@ class Parameter(Tensor):
@property
def data(self) -> Tensor | None:
"""Returns the concrete value of the parameter if it is bound to a concrete value,
- otherwise returns None. The returned value is a tvm.runtime.Tensor."""
+ otherwise returns None. The returned value is a tvm_tirx.runtime.Tensor."""
return self._data
@data.setter
- def data(self, data: Union[None, tvm.runtime.Tensor, np.ndarray, "torch.Tensor"]) -> None:
+ def data(self, data: Union[None, tvm_tirx.runtime.Tensor, np.ndarray, "torch.Tensor"]) -> None:
"""Set the concrete value of the parameter. The data should be one of the following:
- None: unbind the parameter to concrete values
- - tvm.runtime.Tensor
+ - tvm_tirx.runtime.Tensor
- numpy.ndarray
- torch.Tensor and any other DLPack-compliant tensors
"""
@@ -281,10 +281,10 @@ class Parameter(Tensor):
self._data = data
return
# Try to do zero-copy if possible
- if isinstance(data, tvm.runtime.Tensor):
+ if isinstance(data, tvm_tirx.runtime.Tensor):
pass
elif isinstance(data, np.ndarray):
- data = tvm.runtime.tensor(data)
+ data = tvm_tirx.runtime.tensor(data)
elif hasattr(data, "__dlpack__"):
data = _from_dlpack(data)
else:
@@ -479,7 +479,7 @@ class Module(SubroutineMixin):
Returns
-------
- irmodule : tvm.ir.IRModule
+ irmodule : tvm_tirx.ir.IRModule
The converted tvm IR representation of the model.
params : List[Tuple[str, Parameter]]
A list of Parameters corresponding to the weights of the model.
@@ -847,15 +847,15 @@ def _attribute_finder(root: Module, prefix: str, condition_yield: Callable[[Any]
)
-def _from_dlpack(tensor) -> tvm.runtime.Tensor:
+def _from_dlpack(tensor) -> tvm_tirx.runtime.Tensor:
try:
- return tvm.runtime.from_dlpack(tensor)
+ return tvm_tirx.runtime.from_dlpack(tensor)
except RuntimeError:
pass
# special logic for PyTorch
device_type = tensor.device.type
device_id = tensor.device.index or 0
- return tvm.runtime.tensor(
+ return tvm_tirx.runtime.tensor(
tensor.numpy(),
device=Device(
Device._DEVICE_NAME_TO_TYPE[device_type],
@@ -866,7 +866,7 @@ def _from_dlpack(tensor) -> tvm.runtime.Tensor:
def _param_to_tensor(
params: list[tuple[str, Parameter]], device: Device
-) -> list[tvm.runtime.Tensor]:
+) -> list[tvm_tirx.runtime.Tensor]:
results = []
missing = []
for name, param in params:
diff --git a/python/tvm/relax/frontend/nn/exporter.py b/python/tvm/relax/frontend/nn/exporter.py
index a271ad9fea..e4fbbc947d 100644
--- a/python/tvm/relax/frontend/nn/exporter.py
+++ b/python/tvm/relax/frontend/nn/exporter.py
@@ -21,9 +21,9 @@ import operator
import threading
import typing
-import tvm
-from tvm import tirx
-from tvm.ir import IRModule
+import tvm, tvm_tirx
+from tvm_tirx import tirx
+from tvm_tirx.ir import IRModule
from .... import relax as rx
from ...block_builder import BlockBuilder
@@ -178,7 +178,7 @@ def _emit_method( # pylint: disable=too-many-locals,too-many-branches,too-many-
raise TypeError(f"Unsupported return type: {type(expr)}")
def _convert_input(arg):
- if isinstance(arg, tvm.ir.Var):
+ if isinstance(arg, tvm_tirx.ir.Var):
return rx.Var(arg.name, ty=ShapeType(values=[arg]))
if isinstance(arg, core.Tensor | core.Object):
return arg._expr # pylint: disable=protected-access
@@ -203,7 +203,7 @@ def _emit_method( # pylint: disable=too-many-locals,too-many-branches,too-many-
for name, param in params:
# Make sure the a symbolic shape is not re-registered (same as _method_spec_to_inputs)
# e.g. we do not see `vocab_size` for `lm_head` and `vocab_size_1` for `embed_tokens`
- new_shape = [_get_var(x) if isinstance(x, tvm.ir.Var) else x for x in param.shape]
+ new_shape = [_get_var(x) if isinstance(x, tvm_tirx.ir.Var) else x for x in param.shape]
var = core.Tensor.placeholder(new_shape, param.dtype, name)._expr
inputs.append(var)
param._expr = var
@@ -264,7 +264,7 @@ def _emit_method( # pylint: disable=too-many-locals,too-many-branches,too-many-
return type(arg.elements)(ret)
if isinstance(arg, core.Tensor):
return core.Tensor(_expr=var)
- if isinstance(arg, tvm.ir.Var):
+ if isinstance(arg, tvm_tirx.ir.Var):
return arg
raise TypeError(f"Unsupported input type: {type(arg)}")
diff --git a/python/tvm/relax/frontend/nn/extern.py b/python/tvm/relax/frontend/nn/extern.py
index 9fe90e6f5b..c88b140501 100644
--- a/python/tvm/relax/frontend/nn/extern.py
+++ b/python/tvm/relax/frontend/nn/extern.py
@@ -26,10 +26,10 @@ from pathlib import Path
import tvm_ffi
-import tvm
-from tvm import libinfo, tirx
-from tvm.runtime import Module, load_static_library
-from tvm.support import cc as _cc
+import tvm, tvm_tirx
+from tvm_tirx import libinfo, tirx
+from tvm_tirx.runtime import Module, load_static_library
+from tvm_tirx.support import cc as _cc
from ...op import call_dps_packed
from . import core
@@ -60,8 +60,8 @@ class ExternModule:
if isinstance(arg, float):
return rx.prim_value(tirx.FloatImm("float64", arg))
if isinstance(arg, str):
- return tvm.ir.StringImm(arg)
- if tvm.ir.is_prim_expr(arg):
+ return tvm_tirx.ir.StringImm(arg)
+ if tvm_tirx.ir.is_prim_expr(arg):
return rx.prim_value(arg)
if isinstance(arg, tuple | list):
return rx.Tuple([_convert(e, f"{name}_{i}") for i, e in enumerate(arg)])
@@ -277,7 +277,7 @@ class SourceModule(ExternModule): # pylint: disable=too-few-public-methods
f"Using environment variable `TVM_HOME`, but it is not a directory: {tvm_path!s}"
)
else:
- import tvm # pylint: disable=import-outside-toplevel
+ import tvm, tvm_tirx # pylint: disable=import-outside-toplevel
tvm_path = Path(tvm.__file__).parent
assert tvm_path.is_dir()
diff --git a/python/tvm/relax/frontend/nn/llm/_decode_kernels.py b/python/tvm/relax/frontend/nn/llm/_decode_kernels.py
index c8471e22fc..ec3ce3e670 100644
--- a/python/tvm/relax/frontend/nn/llm/_decode_kernels.py
+++ b/python/tvm/relax/frontend/nn/llm/_decode_kernels.py
@@ -31,9 +31,9 @@ Contents:
import math
from typing import Any
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-from tvm.target import Target
+from tvm_tirx.script import s_tir as Ts
+from tvm_tirx.script import tirx as T
+from tvm_tirx.target import Target
from ._kernel_common import (
_get_kv_chunk_len,
diff --git a/python/tvm/relax/frontend/nn/llm/_kernel_common.py b/python/tvm/relax/frontend/nn/llm/_kernel_common.py
index ec60ede494..e2f32620ae 100644
--- a/python/tvm/relax/frontend/nn/llm/_kernel_common.py
+++ b/python/tvm/relax/frontend/nn/llm/_kernel_common.py
@@ -37,12 +37,13 @@ import enum
import math
from typing import Any
-import tvm
-from tvm import s_tir, tirx
-from tvm.runtime import DataType
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-from tvm.target import Target
+import tvm, tvm_tirx
+from tvm_tirx import tirx
+from tvm import s_tir
+from tvm_tirx.runtime import DataType
+from tvm_tirx.script import s_tir as Ts
+from tvm_tirx.script import tirx as T
+from tvm_tirx.target import Target
from .position_embedding import switch_rope_freq_func
diff --git a/python/tvm/relax/frontend/nn/llm/_page_kernels.py b/python/tvm/relax/frontend/nn/llm/_page_kernels.py
index f913f339cf..8c7f2a248d 100644
--- a/python/tvm/relax/frontend/nn/llm/_page_kernels.py
+++ b/python/tvm/relax/frontend/nn/llm/_page_kernels.py
@@ -32,9 +32,9 @@ This module contains:
from __future__ import annotations
# pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-from tvm.target import Target
+from tvm_tirx.script import s_tir as Ts
+from tvm_tirx.script import tirx as T
+from tvm_tirx.target import Target
from ._kernel_common import get_max_num_threads_per_block
diff --git a/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py b/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
index 4d31ee05c8..d0ee8e155a 100644
--- a/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
+++ b/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
@@ -29,11 +29,11 @@ K/V loading path that is specific to its storage layout.
import math
from typing import Any, Literal
-import tvm
-from tvm import tirx
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-from tvm.target import Target
+import tvm, tvm_tirx
+from tvm_tirx import tirx
+from tvm_tirx.script import s_tir as Ts
+from tvm_tirx.script import tirx as T
+from tvm_tirx.target import Target
from ._kernel_common import (
_alloc_mha_qkvo_buffers,
diff --git a/python/tvm/relax/frontend/nn/llm/_prefill_parallel.py b/python/tvm/relax/frontend/nn/llm/_prefill_parallel.py
index 70dfdd86ae..eba66a746c 100644
--- a/python/tvm/relax/frontend/nn/llm/_prefill_parallel.py
+++ b/python/tvm/relax/frontend/nn/llm/_prefill_parallel.py
@@ -16,8 +16,8 @@
# under the License.
"""Parallel dot-product prefill for wide FP32 WebGPU attention heads."""
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
+from tvm_tirx.script import s_tir as Ts
+from tvm_tirx.script import tirx as T
from ._kernel_common import _get_kv_chunk_len, _get_seq_offset, _length_info_buffer, _rope
diff --git a/python/tvm/relax/frontend/nn/llm/kv_cache.py b/python/tvm/relax/frontend/nn/llm/kv_cache.py
index cf46018522..d040ff03ed 100644
--- a/python/tvm/relax/frontend/nn/llm/kv_cache.py
+++ b/python/tvm/relax/frontend/nn/llm/kv_cache.py
@@ -37,13 +37,13 @@ test suite can continue to import them via ``tvm.relax.frontend.nn.llm.kv_cache`
import math
from typing import Any, Literal
-import tvm
+import tvm, tvm_tirx
from tvm import relax as rx
-from tvm import tirx
-from tvm.ir import TupleType as _TupleType
+from tvm_tirx import tirx
+from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.frontend.nn import Object, Tensor
-from tvm.target import Target
+from tvm_tirx.target import Target
# Re-export enums + kernel factories so existing ``from kv_cache import ...``
# users (test suite, tree_attn.py, mlc-llm, etc.) continue to work after the
@@ -121,7 +121,7 @@ __all__ = [
class PagedKVCache(Object): # pylint: disable=too-few-public-methods
"""The Paged KV Cache used in LLM batching for efficient attention computation."""
- extern_mods: list[tvm.runtime.Module] = []
+ extern_mods: list[tvm_tirx.runtime.Module] = []
def attention_with_fused_qkv(
self,
@@ -485,18 +485,18 @@ class FlashInferPagedKVCache(PagedKVCache): # pylint: disable=too-few-public-me
bb = rx.BlockBuilder.current()
mha_functions = (
[
- rx.Tuple([tvm.ir.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_paged_run"), rx.ExternFunc("batch_prefill_plan")]),
- rx.Tuple([tvm.ir.StringImm("flashinfer"), rx.ExternFunc("batch_decode_run"), rx.ExternFunc("batch_decode_plan")]),
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target, sliding_window_size=layer_sliding_window_size), "tir_attention_prefill_sliding_window")]),
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_decode_sliding_window")]),
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache")]),
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(tree_attn(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_paged_run"), rx.ExternFunc("batch_prefill_plan")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("flashinfer"), rx.ExternFunc("batch_decode_run"), rx.ExternFunc("batch_decode_plan")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target, sliding_window_size=layer_sliding_window_size), "tir_attention_prefill_sliding_window")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_decode_sliding_window")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(tree_attn(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask")]),
]
if attn_kind_single == "mha"
else [rx.Tuple([]) for _ in range(6)]
)
- ragged_prefill_function = rx.Tuple([tvm.ir.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_ragged_run"), rx.ExternFunc("batch_prefill_plan")]) if attn_kind_single == "mha" else rx.Tuple([tvm.ir.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_ragged_run"), rx.ExternFunc("batch_prefill_plan"), rx.prim_value(mla_original_qk_head_dim), rx.prim_value(mla_original_v_head_dim)])
- mla_function = rx.Tuple([tvm.ir.StringImm("flashinfer"), rx.ExternFunc("batch_mla_run"), rx.ExternFunc("batch_mla_plan")] if attn_kind_single == "mla" else [])
+ ragged_prefill_function = rx.Tuple([tvm_tirx.ir.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_ragged_run"), rx.ExternFunc("batch_prefill_plan")]) if attn_kind_single == "mha" else rx.Tuple([tvm_tirx.ir.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_ragged_run"), rx.ExternFunc("batch_prefill_plan"), rx.prim_value(mla_original_qk_head_dim), rx.prim_value(mla_original_v_head_dim)])
+ mla_function = rx.Tuple([tvm_tirx.ir.StringImm("flashinfer"), rx.ExternFunc("batch_mla_run"), rx.ExternFunc("batch_mla_plan")] if attn_kind_single == "mla" else [])
attn_merge_functions = [
bb.add_func(_merge_state_inplace(num_attention_heads, v_head_dim, dtype, target, "tir_attention_merge_state"), "tir_attention_merge_state"),
]
@@ -673,13 +673,13 @@ class TIRPagedKVCache(PagedKVCache): # pylint: disable=too-few-public-methods
raise ValueError("MLA is not supported in TIR kernels for now.")
args.extend(
[
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(_attention_prefill_ragged_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, v_head_dim, dtype, rope_scaling), "tir_attention_prefill_ragged_cpu")]),
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(_attention_prefill_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling), "tir_attention_prefill_cpu")]),
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(_attention_decode_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling), "tir_attention_decode_cpu")]),
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(_attention_prefill_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, sliding_window_size=layer_sliding_window_size), "tir_attention_prefill_cpu_sliding_window")]),
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(_attention_decode_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling), "tir_attention_decode_cpu_sliding_window")]),
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(tree_attn_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling), "tir_attention_prefill_with_tree_mask_cpu")]),
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache_cpu")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(_attention_prefill_ragged_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, v_head_dim, dtype, rope_scaling), "tir_attention_prefill_ragged_cpu")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(_attention_prefill_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling), "tir_attention_prefill_cpu")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(_attention_decode_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling), "tir_attention_decode_cpu")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(_attention_prefill_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, sliding_window_size=layer_sliding_window_size), "tir_attention_prefill_cpu_sliding_window")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(_attention_decode_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling), "tir_attention_decode_cpu_sliding_window")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(tree_attn_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling), "tir_attention_prefill_with_tree_mask_cpu")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache_cpu")]),
rx.Tuple([]), # f_mla_prefill
rx.Tuple([bb.add_func(_merge_state_inplace_cpu(dtype), "tir_attention_merge_state_cpu")]),
bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale, qk_head_dim, num_attention_heads, num_key_value_heads, dtype, rope_scaling, rotary_dim), "tir_split_rotary"),
@@ -691,20 +691,20 @@ class TIRPagedKVCache(PagedKVCache): # pylint: disable=too-few-public-methods
else:
ragged_qk_head_dim = qk_head_dim if attn_kind_single == "mha" else mla_original_qk_head_dim
ragged_v_head_dim = v_head_dim if attn_kind_single == "mha" else mla_original_v_head_dim
- args.append(rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(_attention_prefill_ragged(num_key_value_heads if attn_kind_single == "mha" else num_attention_heads, num_attention_heads, ragged_qk_head_dim, ragged_v_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_ragged")]))
+ args.append(rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(_attention_prefill_ragged(num_key_value_heads if attn_kind_single == "mha" else num_attention_heads, num_attention_heads, ragged_qk_head_dim, ragged_v_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_ragged")]))
mha_functions = (
[
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling, target), "tir_attention_prefill")]),
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling, target), "tir_attention_decode")]),
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target, sliding_window_size=layer_sliding_window_size), "tir_attention_prefill_sliding_window")]),
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_decode_sliding_window")]),
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache")]),
- rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(tree_attn(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling, target), "tir_attention_prefill")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling, target), "tir_attention_decode")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target, sliding_window_size=layer_sliding_window_size), "tir_attention_prefill_sliding_window")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_decode_sliding_window")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache")]),
+ rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(tree_attn(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask")]),
]
if attn_kind_single == "mha"
else [rx.Tuple([]) for _ in range(6)]
)
- mla_function = rx.Tuple([tvm.ir.StringImm("tirx"), bb.add_func(_attention_prefill_mla(num_attention_heads, v_head_dim, qk_head_dim - v_head_dim, dtype, False, target), "tir_attention_prefill_mla")] if attn_kind_single == "mla" else [])
+ mla_function = rx.Tuple([tvm_tirx.ir.StringImm("tirx"), bb.add_func(_attention_prefill_mla(num_attention_heads, v_head_dim, qk_head_dim - v_head_dim, dtype, False, target), "tir_attention_prefill_mla")] if attn_kind_single == "mla" else [])
attn_merge_functions = [
bb.add_func(_merge_state_inplace(num_attention_heads, v_head_dim, dtype, target, "tir_attention_merge_state"), "tir_attention_merge_state"),
]
diff --git a/python/tvm/relax/frontend/nn/llm/position_embedding.py b/python/tvm/relax/frontend/nn/llm/position_embedding.py
index 009d2da28d..cddf2462ae 100644
--- a/python/tvm/relax/frontend/nn/llm/position_embedding.py
+++ b/python/tvm/relax/frontend/nn/llm/position_embedding.py
@@ -22,10 +22,10 @@ from collections.abc import Callable
from functools import partial
from typing import Any
-from tvm import tirx
+from tvm_tirx import tirx
from tvm.relax.frontend.nn import Tensor, op
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
+from tvm_tirx.script import s_tir as Ts
+from tvm_tirx.script import tirx as T
# pylint: disable=invalid-name
diff --git a/python/tvm/relax/frontend/nn/llm/tree_attn.py b/python/tvm/relax/frontend/nn/llm/tree_attn.py
index 95ea6eb784..1b52f319fd 100644
--- a/python/tvm/relax/frontend/nn/llm/tree_attn.py
+++ b/python/tvm/relax/frontend/nn/llm/tree_attn.py
@@ -22,10 +22,11 @@
import math
from typing import Any
-from tvm import s_tir, tirx
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-from tvm.target import Target
+from tvm_tirx import tirx
+from tvm import s_tir
+from tvm_tirx.script import s_tir as Ts
+from tvm_tirx.script import tirx as T
+from tvm_tirx.target import Target
# The helpers below are shared with the main KV-cache kernels. They live in
# ``_kernel_common`` so that ``kv_cache.py``, this file, and the split kernel
@@ -83,7 +84,7 @@ def tree_attn_cpu(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any]):
Returns
-------
- mod : tvm.IRModule
+ mod : tvm_tirx.IRModule
The generated IR module.
"""
group_size = h_q // h_kv
@@ -262,7 +263,7 @@ def tree_attn(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any], target: Target)
Returns
-------
- mod : tvm.IRModule
+ mod : tvm_tirx.IRModule
The generated IR module.
"""
# pylint: disable=invalid-name,line-too-long
@@ -573,7 +574,7 @@ def tree_attn_with_paged_kv_cache_cpu(h_kv, h_q, d, dtype, rope_scaling: dict[st
Returns
-------
- mod : tvm.IRModule
+ mod : tvm_tirx.IRModule
The generated IR module.
"""
global_symbol = "tree_attn_paged_kv_cpu"
@@ -745,7 +746,7 @@ def tree_attn_with_paged_kv_cache(
Returns
-------
- mod : tvm.IRModule
+ mod : tvm_tirx.IRModule
The generated IR module.
"""
# pylint: disable=invalid-name, line-too-long
diff --git a/python/tvm/relax/frontend/nn/modules.py b/python/tvm/relax/frontend/nn/modules.py
index 0842c51b5e..5016ffbe43 100644
--- a/python/tvm/relax/frontend/nn/modules.py
+++ b/python/tvm/relax/frontend/nn/modules.py
@@ -21,7 +21,7 @@
from collections.abc import Sequence
from tvm import relax as rx
-from tvm import tirx
+from tvm_tirx import tirx
from . import op
from .core import Effect, Module, ModuleList, Parameter, Tensor, get_default_dtype
diff --git a/python/tvm/relax/frontend/nn/op.py b/python/tvm/relax/frontend/nn/op.py
index af7565f6d0..10c925ed4f 100644
--- a/python/tvm/relax/frontend/nn/op.py
+++ b/python/tvm/relax/frontend/nn/op.py
@@ -25,11 +25,11 @@ from typing import Any, TypeVar
import numpy as np
-import tvm
+import tvm, tvm_tirx
from tvm import te
-from tvm import tirx as _tir
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
+from tvm_tirx import tirx as _tir
+from tvm_tirx.script import s_tir as Ts
+from tvm_tirx.script import tirx as T
from ... import expr as rx
from ... import op as _op
@@ -2078,7 +2078,7 @@ def tensor_ir_op(
call_tir_args.append(arg._expr)
elif isinstance(arg, rx.ShapeExpr):
call_tir_args.extend(arg.values)
- elif tvm.ir.is_prim_expr(arg):
+ elif tvm_tirx.ir.is_prim_expr(arg):
call_tir_args.append(arg)
else:
raise TypeError(
@@ -2091,7 +2091,7 @@ def tensor_ir_op(
else:
out_ty = [x._expr.ty for x in out]
- out_ty = out_ty[0] if len(out_ty) == 1 else tvm.ir.TupleType(out_ty)
+ out_ty = out_ty[0] if len(out_ty) == 1 else tvm_tirx.ir.TupleType(out_ty)
bb = BlockBuilder.current()
global_var = bb.add_func(func, name_hint)
@@ -2148,7 +2148,7 @@ def tensor_ir_inplace_op(
call_tir_args.append(arg._expr)
elif isinstance(arg, rx.ShapeExpr):
call_tir_args.extend(arg.values)
- elif tvm.ir.is_prim_expr(arg):
+ elif tvm_tirx.ir.is_prim_expr(arg):
call_tir_args.append(arg)
else:
raise TypeError(
@@ -2161,7 +2161,7 @@ def tensor_ir_inplace_op(
else:
out_ty = [x._expr.ty for x in out]
- out_ty = out_ty[0] if len(out_ty) == 1 else tvm.ir.TupleType(out_ty)
+ out_ty = out_ty[0] if len(out_ty) == 1 else tvm_tirx.ir.TupleType(out_ty)
bb = BlockBuilder.current()
global_var = bb.add_func(func, name_hint)
@@ -2181,7 +2181,7 @@ def extern(
out: OutType,
) -> OutType:
"""Invoke an extern function during runtime. The extern function must be registered with the "
- TVM runtime using `reflection::GlobalDef().def` (C++), or `tvm.register_global_func` (Python).
+ TVM runtime using `reflection::GlobalDef().def` (C++), or `tvm_tirx.register_global_func` (Python).
Parameters
----------
@@ -2209,8 +2209,8 @@ def extern(
if isinstance(arg, float):
return rx.prim_value(_tir.FloatImm("float64", arg))
if isinstance(arg, str):
- return tvm.ir.StringImm(arg)
- if tvm.ir.is_prim_expr(arg):
+ return tvm_tirx.ir.StringImm(arg)
+ if tvm_tirx.ir.is_prim_expr(arg):
return rx.prim_value(arg)
if isinstance(arg, tuple | list):
return rx.Tuple([_convert(e, f"{name}_{i}") for i, e in enumerate(arg)])
@@ -2238,7 +2238,7 @@ def debug_func(
.. code-block:: python
- @tvm.register_global_func(name_of_debug_func)
+ @tvm_tirx.register_global_func(name_of_debug_func)
def debug_func(lineno: str, arg_0, arg_1, ...) -> None:
...
@@ -2274,10 +2274,10 @@ def debug_func(
converted_args.append(rx.prim_value(_tir.IntImm("int64", arg)))
elif isinstance(arg, float):
converted_args.append(rx.prim_value(_tir.FloatImm("float32", arg)))
- elif tvm.ir.is_prim_expr(arg):
+ elif tvm_tirx.ir.is_prim_expr(arg):
converted_args.append(rx.prim_value(arg))
elif isinstance(arg, str):
- converted_args.append(tvm.ir.StringImm(arg))
+ converted_args.append(tvm_tirx.ir.StringImm(arg))
else:
raise TypeError(f"Unsupported type {type(arg)}")
@@ -2285,8 +2285,8 @@ def debug_func(
rx.call_pure_packed(
"vm.builtin.invoke_debug_func",
io.effect,
- tvm.ir.StringImm(name),
- tvm.ir.StringImm(_line_info),
+ tvm_tirx.ir.StringImm(name),
+ tvm_tirx.ir.StringImm(_line_info),
*converted_args,
ty_args=[rx.AnyType()],
),
diff --git a/python/tvm/relax/frontend/nn/subroutine.py b/python/tvm/relax/frontend/nn/subroutine.py
index b3bb6138ab..368baea32b 100644
--- a/python/tvm/relax/frontend/nn/subroutine.py
+++ b/python/tvm/relax/frontend/nn/subroutine.py
@@ -26,7 +26,8 @@ import typing
import tvm_ffi
-from tvm import ir, relax
+from tvm_tirx import ir
+from tvm import relax
from tvm.relax.frontend import nn
diff --git a/python/tvm/relax/frontend/nn/torch.py b/python/tvm/relax/frontend/nn/torch.py
index 0ef8baa001..6f287ea16e 100644
--- a/python/tvm/relax/frontend/nn/torch.py
+++ b/python/tvm/relax/frontend/nn/torch.py
@@ -23,7 +23,7 @@ from typing import Any
import torch
from tvm_ffi import Array, Shape
-from tvm.runtime import Tensor, _tensor
+from tvm_tirx.runtime import Tensor, _tensor
from tvm.runtime.vm import VirtualMachine
from . import core
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 4eec3df11f..484aead264 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -56,12 +56,13 @@ except ImportError as err:
import tvm_ffi
-import tvm
-from tvm import relax, tirx, topi
-from tvm.ir import Expr as TvmExpr
-from tvm.ir import IRModule
-from tvm.ir.supply import UniqueNameSupply
-from tvm.runtime import DataType, DataTypeCode
+import tvm, tvm_tirx
+from tvm_tirx import tirx
+from tvm import relax, topi
+from tvm_tirx.ir import Expr as TvmExpr
+from tvm_tirx.ir import IRModule
+from tvm_tirx.ir.supply import UniqueNameSupply
+from tvm_tirx.runtime import DataType, DataTypeCode
from tvm.topi.utils import get_const_tuple
from ..common import autopad
@@ -98,15 +99,15 @@ def get_type(elem_type: str | int) -> str:
def get_constant(
- var: tvm.ir.GenericConst | relax.Var,
+ var: tvm_tirx.ir.GenericConst | relax.Var,
params: list[dict[str, relax.Var]],
-) -> tvm.ir.GenericConst | relax.Var:
+) -> tvm_tirx.ir.GenericConst | relax.Var:
"""Attempt to convert a variable to a constant if possible.
This is the primary function meant to interact with params.
Parameters
----------
- var: Union[tvm.ir.GenericConst, relax.Var]
+ var: Union[tvm_tirx.ir.GenericConst, relax.Var]
The input value to try to convert to a constant.
params: List[Dict[str, relax.Var]]
The parameters for the graph. Contains both the global registry of nodes
@@ -115,7 +116,7 @@ def get_constant(
Returns
-------
- var : Union[tvm.ir.GenericConst, relax.Var]
+ var : Union[tvm_tirx.ir.GenericConst, relax.Var]
The input value converted to a constant if possible. If the value
isn't found in params, the input variable is returned unmodified.
"""
@@ -133,7 +134,7 @@ def get_constant(
return var
-def get_value(token, value_dict: dict[str, tvm.tirx.Var]) -> int | tvm.tirx.Var:
+def get_value(token, value_dict: dict[str, tvm_tirx.tirx.Var]) -> int | tvm_tirx.tirx.Var:
"""Convert a token to an integer constant or symbolic variable.
Parameters
@@ -146,7 +147,7 @@ def get_value(token, value_dict: dict[str, tvm.tirx.Var]) -> int | tvm.tirx.Var:
Returns
-------
- Union[int, tvm.tirx.Var]
+ Union[int, tvm_tirx.tirx.Var]
The decoded token
"""
@@ -154,12 +155,12 @@ def get_value(token, value_dict: dict[str, tvm.tirx.Var]) -> int | tvm.tirx.Var:
return int(token)
except ValueError:
if token not in value_dict or token == "?":
- value_dict[token] = tvm.tirx.Var(token, "int64")
+ value_dict[token] = tvm_tirx.tirx.Var(token, "int64")
value = value_dict[token]
return value
-def parse_shape_name(name: str, value_dict: dict[str, tvm.tirx.Var]) -> tirx.Expr | tvm.tirx.Var:
+def parse_shape_name(name: str, value_dict: dict[str, tvm_tirx.tirx.Var]) -> tirx.Expr | tvm_tirx.tirx.Var:
"""Converts expressions in the shape dimension name to prim expressions.
Parameters
@@ -172,7 +173,7 @@ def parse_shape_name(name: str, value_dict: dict[str, tvm.tirx.Var]) -> tirx.Exp
Returns
-------
- Union[tirx.Expr, tvm.tirx.Var]
+ Union[tirx.Expr, tvm_tirx.tirx.Var]
The expression of the shape dimension.
"""
@@ -209,7 +210,7 @@ def parse_shape_name(name: str, value_dict: dict[str, tvm.tirx.Var]) -> tirx.Exp
def get_info(
- info_proto: onnx.onnx_ml_pb2.ValueInfoProto, value_dict: dict[str, tvm.tirx.Var]
+ info_proto: onnx.onnx_ml_pb2.ValueInfoProto, value_dict: dict[str, tvm_tirx.tirx.Var]
) -> tuple[str, list, str, list, dict]:
"""Extract the shape from a ValueInfoProto.
@@ -268,13 +269,13 @@ def get_numpy(tensor_proto: onnx.onnx_ml_pb2.TensorProto) -> _np.ndarray:
def get_prim_expr_list(
- inputs: tvm.ir.GenericConst | relax.ShapeExpr,
+ inputs: tvm_tirx.ir.GenericConst | relax.ShapeExpr,
) -> list[int | tirx.Expr]:
"""Attempt to convert a variable to list of Expr if possible.
Parameters
----------
- inputs : Union[tvm.ir.GenericConst, relax.ShapeExpr, tvm.tirx.Expr]
+ inputs : Union[tvm_tirx.ir.GenericConst, relax.ShapeExpr, tvm_tirx.tirx.Expr]
The input value to try to convert to a list of Expr.
Returns
@@ -282,14 +283,14 @@ def get_prim_expr_list(
ret : List[Union[int, tirx.Expr]]
The input value converted to a list of Expr if possible.
"""
- if isinstance(inputs, tvm.ir.GenericConst):
+ if isinstance(inputs, tvm_tirx.ir.GenericConst):
np_value = inputs.value.numpy()
if np_value.ndim != 1:
raise ValueError(f"Cannot cast {type(inputs)} to list of Expr")
return np_value.tolist()
elif isinstance(inputs, relax.ShapeExpr):
return inputs.values
- elif tvm.ir.is_prim_expr(inputs):
+ elif tvm_tirx.ir.is_prim_expr(inputs):
return [inputs]
else:
raise ValueError(f"Cannot cast {type(inputs)} to list of Expr")
@@ -563,7 +564,7 @@ class MatMulInteger16(OnnxOpConverter):
def _to_numpy(x):
- if tvm.ir.is_prim_expr(x):
+ if tvm_tirx.ir.is_prim_expr(x):
if isinstance(x, tirx.IntImm | tirx.FloatImm):
return _np.array(x.value)
return x
@@ -595,12 +596,12 @@ class BinaryBase(OnnxOpConverter):
"""Base implementation for binary operations."""
if cls.numpy_op is None or cls.relax_op is None:
raise ValueError("Numpy and Relax operators must be defined for BinaryBase.")
- is_prim_var = [tvm.ir.is_prim_var(inp) for inp in inputs]
+ is_prim_var = [tvm_tirx.ir.is_prim_var(inp) for inp in inputs]
if all(
- is_var or not isinstance(inp, tvm.ir.Call | tvm.ir.Var)
+ is_var or not isinstance(inp, tvm_tirx.ir.Call | tvm_tirx.ir.Var)
for inp, is_var in zip(inputs, is_prim_var)
):
- has_prim_expr = any(tvm.ir.is_prim_expr(inp) for inp in inputs)
+ has_prim_expr = any(tvm_tirx.ir.is_prim_expr(inp) for inp in inputs)
x = _to_numpy(inputs[0])
y = _to_numpy(inputs[1])
output = cls.numpy_op(x, y) # pylint: disable=not-callable
@@ -611,7 +612,7 @@ class BinaryBase(OnnxOpConverter):
if x.dtype == y.dtype and not _np.issubdtype(output.dtype, _np.bool_):
# no numpy precision widening
output = output.astype(x.dtype)
- if all([isinstance(inp, tvm.ir.GenericConst) for inp in inputs]):
+ if all([isinstance(inp, tvm_tirx.ir.GenericConst) for inp in inputs]):
return relax.const(output, output.dtype) # pylint: disable=not-callable
return cls.relax_op(inputs[0], inputs[1]) # pylint: disable=not-callable
@@ -658,9 +659,9 @@ class Div(BinaryBase):
@staticmethod
def _as_scalar_prim_expr(expr, dtype):
- if tvm.ir.is_prim_expr(expr):
+ if tvm_tirx.ir.is_prim_expr(expr):
return expr
- if isinstance(expr, tvm.ir.GenericConst):
+ if isinstance(expr, tvm_tirx.ir.GenericConst):
data = expr.value.numpy()
if data.size == 1:
return tirx.const(data.item(), dtype)
@@ -668,7 +669,7 @@ class Div(BinaryBase):
@staticmethod
def _is_zero(expr):
- if isinstance(expr, tvm.ir.GenericConst):
+ if isinstance(expr, tvm_tirx.ir.GenericConst):
return bool(_np.any(expr.value.numpy() == 0))
if isinstance(expr, tirx.IntImm):
return int(expr.value) == 0
@@ -688,12 +689,12 @@ class Div(BinaryBase):
try:
lhs_dtype = (
str(getattr(inputs[0], "dtype", None) or inputs[0].ty)
- if tvm.ir.is_prim_expr(inputs[0])
+ if tvm_tirx.ir.is_prim_expr(inputs[0])
else inputs[0].ty.dtype.dtype
)
rhs_dtype = (
str(getattr(inputs[1], "dtype", None) or inputs[1].ty)
- if tvm.ir.is_prim_expr(inputs[1])
+ if tvm_tirx.ir.is_prim_expr(inputs[1])
else inputs[1].ty.dtype.dtype
)
lhs_code = DataType(lhs_dtype).type_code
@@ -709,13 +710,13 @@ class Div(BinaryBase):
if cls._is_zero(inputs[1]):
raise ValueError("ONNX Div with integer inputs encountered divisor value 0.")
- if all(isinstance(inp, tvm.ir.GenericConst) for inp in inputs):
+ if all(isinstance(inp, tvm_tirx.ir.GenericConst) for inp in inputs):
lhs = inputs[0].value.numpy()
rhs = inputs[1].value.numpy()
output = cls._numpy_integer_divide(lhs, rhs, lhs_code == DataTypeCode.INT)
return relax.const(output, lhs_dtype)
- has_prim_expr = any(tvm.ir.is_prim_expr(inp) for inp in inputs)
+ has_prim_expr = any(tvm_tirx.ir.is_prim_expr(inp) for inp in inputs)
lhs = cls._as_scalar_prim_expr(inputs[0], lhs_dtype)
rhs = cls._as_scalar_prim_expr(inputs[1], rhs_dtype)
if has_prim_expr and lhs is not None and rhs is not None:
@@ -834,10 +835,10 @@ class Equal(OnnxOpConverter):
@classmethod
def _impl_v13(cls, bb, inputs, attr, params):
- if all([isinstance(inp, tvm.ir.GenericConst) for inp in inputs]):
+ if all([isinstance(inp, tvm_tirx.ir.GenericConst) for inp in inputs]):
output = inputs[0].value.numpy() == inputs[1].value.numpy()
return relax.const(output, output.dtype)
- elif all([isinstance(inp, tvm.ir.GenericConst | relax.ShapeExpr) for inp in inputs]):
+ elif all([isinstance(inp, tvm_tirx.ir.GenericConst | relax.ShapeExpr) for inp in inputs]):
lhs = get_prim_expr_list(inputs[0])
rhs = get_prim_expr_list(inputs[1])
if len(lhs) != len(rhs):
@@ -901,7 +902,7 @@ class BitwiseNot(OnnxOpConverter):
@classmethod
def _impl_v18(cls, bb, inputs, attr, params):
- if isinstance(inputs[0], tvm.ir.GenericConst):
+ if isinstance(inputs[0], tvm_tirx.ir.GenericConst):
return relax.const(_np.bitwise_not(inputs[0].value.numpy()), inputs[0].ty.dtype)
return relax.op.bitwise_not(inputs[0])
@@ -1089,8 +1090,8 @@ class Hardmax(OnnxOpConverter):
iota = relax.op.reshape(relax.op.arange(0, axis_extent, dtype="int64"), iota_shape)
return relax.op.astype(relax.op.equal(iota, argmax), dtype)
argmax = relax.op.argmax(data, axis=normalized_axis)
- on_value = relax.prim_value(tvm.tirx.const(1.0, dtype))
- off_value = relax.prim_value(tvm.tirx.const(0.0, dtype))
+ on_value = relax.prim_value(tvm_tirx.tirx.const(1.0, dtype))
+ off_value = relax.prim_value(tvm_tirx.tirx.const(0.0, dtype))
return relax.op.one_hot(argmax, on_value, off_value, axis_extent, normalized_axis)
@classmethod
@@ -1133,7 +1134,7 @@ class Transpose(OnnxOpConverter):
elif hasattr(data.ty, "shape") and data.ty.shape:
input_ndim = len(data.ty.shape)
else:
- if isinstance(data, tvm.ir.GenericConst):
+ if isinstance(data, tvm_tirx.ir.GenericConst):
input_ndim = data.value.numpy().ndim
else:
input_ndim = None
@@ -1148,7 +1149,7 @@ class Transpose(OnnxOpConverter):
f"must equal the number of input tensor dimensions ({input_ndim})"
)
- if isinstance(data, tvm.ir.GenericConst):
+ if isinstance(data, tvm_tirx.ir.GenericConst):
output = _np.transpose(data.value.numpy(), axes)
return relax.const(output, output.dtype)
@@ -1170,14 +1171,14 @@ class Unsqueeze(OnnxOpConverter):
axes = get_constant(inputs[1], params)
data_ndim = _get_known_tensor_rank(data)
- if tvm.ir.is_prim_expr(data) and isinstance(axes, tvm.ir.GenericConst):
+ if tvm_tirx.ir.is_prim_expr(data) and isinstance(axes, tvm_tirx.ir.GenericConst):
constant_axes = _normalize_constant_axes(
list(map(int, axes.value.numpy().tolist())), 1, "Unsqueeze"
)
if constant_axes == [0]:
return relax.ShapeExpr([data])
raise NotImplementedError("Unsqueeze with symbolic scalar inputs only supports axis 0.")
- if isinstance(data, tvm.ir.GenericConst) and isinstance(axes, tvm.ir.GenericConst):
+ if isinstance(data, tvm_tirx.ir.GenericConst) and isinstance(axes, tvm_tirx.ir.GenericConst):
constant_axes = _normalize_constant_axes(
list(map(int, axes.value.numpy().tolist())),
data.value.numpy().ndim + axes.value.numpy().size,
@@ -1196,7 +1197,7 @@ class Unsqueeze(OnnxOpConverter):
expanded = expanded.reshape(new_shape)
return relax.const(expanded, data.ty.dtype)
- if isinstance(axes, tvm.ir.GenericConst):
+ if isinstance(axes, tvm_tirx.ir.GenericConst):
if data_ndim is None:
raise ValueError("Unsqueeze requires a statically known input rank.")
constant_axes = _normalize_constant_axes(
@@ -1234,7 +1235,7 @@ class Concat(OnnxOpConverter):
def is_shape_like(x: Any) -> bool:
if isinstance(x, relax.ShapeExpr):
return True
- elif isinstance(x, tvm.ir.GenericConst):
+ elif isinstance(x, tvm_tirx.ir.GenericConst):
return x.ty.ndim == 1 and x.ty.dtype == "int64"
else:
return False
@@ -1257,14 +1258,14 @@ class Concat(OnnxOpConverter):
for inp in resolved:
if isinstance(inp, relax.ShapeExpr):
const_inputs.extend(inp.values)
- elif isinstance(inp, tvm.ir.GenericConst):
+ elif isinstance(inp, tvm_tirx.ir.GenericConst):
const_inputs.extend(inp.value.numpy().tolist())
else:
raise NotImplementedError(f"Unsupported input type: {type(inp)}")
return relax.ShapeExpr(const_inputs)
# If all inputs are constant, perform computation directly.
- if all([isinstance(inp, tvm.ir.GenericConst) for inp in inputs]):
+ if all([isinstance(inp, tvm_tirx.ir.GenericConst) for inp in inputs]):
const_inputs = []
for inp in inputs:
const_inputs.append(inp.value.numpy())
@@ -1287,12 +1288,12 @@ class Cast(OnnxOpConverter):
shape = [int(x) for x in shape]
return relax.const(shape, to_type)
inputs = [bb.normalize(relax.op.shape_to_tensor(shape))]
- if isinstance(inputs[0], tvm.ir.GenericConst):
+ if isinstance(inputs[0], tvm_tirx.ir.GenericConst):
output = inputs[0].value.numpy().astype(to_type)
return relax.const(output, to_type)
- if tvm.ir.is_prim_expr(inputs[0]):
+ if tvm_tirx.ir.is_prim_expr(inputs[0]):
if isinstance(inputs[0], tirx.IntImm | tirx.FloatImm):
- return tvm.tirx.const(inputs[0].value, to_type)
+ return tvm_tirx.tirx.const(inputs[0].value, to_type)
return inputs[0].astype(to_type)
try:
@@ -1425,7 +1426,7 @@ class Gather(OnnxOpConverter):
axis = attr.get("axis", 0)
# If all inputs are constant, we can compute directly.
- if all([isinstance(inp, tvm.ir.GenericConst) for inp in [data, indices]]):
+ if all([isinstance(inp, tvm_tirx.ir.GenericConst) for inp in [data, indices]]):
output = _np.take(data.value.numpy(), indices.value.numpy(), axis=axis)
return relax.const(output, output.dtype)
@@ -1439,7 +1440,7 @@ class Gather(OnnxOpConverter):
# non-scalar index such as (1,) must keep its rank, so only a true
# 0-D index collapses to a scalar.
if isinstance(data, relax.ShapeExpr):
- if isinstance(indices, tvm.ir.GenericConst) and indices.value.numpy().ndim == 0:
+ if isinstance(indices, tvm_tirx.ir.GenericConst) and indices.value.numpy().ndim == 0:
np_index = int(indices.value.numpy().item())
shape_val = data[np_index]
return relax.prim_value(shape_val)
@@ -1733,7 +1734,7 @@ class Reshape(OnnxOpConverter):
if isinstance(data, relax.ShapeExpr):
# Preserve identity flatten for shape values to keep shape-specialized
# handling in downstream shape-construction patterns.
- if isinstance(new_shape, tvm.ir.GenericConst):
+ if isinstance(new_shape, tvm_tirx.ir.GenericConst):
new_shape_values = new_shape.value.numpy().tolist()
if new_shape_values == [-1]:
return data
@@ -1741,7 +1742,7 @@ class Reshape(OnnxOpConverter):
# Other reshape targets follow regular int64 tensor reshape semantics.
data = bb.normalize(relax.op.shape_to_tensor(data))
- if isinstance(data, tvm.ir.GenericConst) and isinstance(new_shape, tvm.ir.GenericConst):
+ if isinstance(data, tvm_tirx.ir.GenericConst) and isinstance(new_shape, tvm_tirx.ir.GenericConst):
data_array = data.value.numpy()
new_shape_values = new_shape.value.numpy().tolist()
if not allowzero:
@@ -1751,7 +1752,7 @@ class Reshape(OnnxOpConverter):
]
out = _np.reshape(data_array, new_shape_values)
return relax.const(out, out.dtype)
- if isinstance(new_shape, tvm.ir.GenericConst):
+ if isinstance(new_shape, tvm_tirx.ir.GenericConst):
new_shape_values = new_shape.value.numpy().tolist()
if allowzero and 0 in new_shape_values:
new_shape = _tensor_to_shape_expr(
@@ -1768,12 +1769,12 @@ class Where(OnnxOpConverter):
@classmethod
def _impl_v16(cls, bb, inputs, attr, params):
- if all([isinstance(inp, tvm.ir.GenericConst) for inp in inputs]):
+ if all([isinstance(inp, tvm_tirx.ir.GenericConst) for inp in inputs]):
np_inputs = [inp.value.numpy() for inp in inputs]
output = _np.where(*np_inputs)
return relax.const(output, output.dtype)
- if all([isinstance(inp, tvm.ir.GenericConst | relax.ShapeExpr) for inp in inputs]):
+ if all([isinstance(inp, tvm_tirx.ir.GenericConst | relax.ShapeExpr) for inp in inputs]):
try:
condition, x, y = [get_prim_expr_list(inp) for inp in inputs]
except ValueError:
@@ -1906,7 +1907,7 @@ class Trilu(OnnxOpConverter):
k = 0
else:
k = get_constant(k, params)
- if isinstance(k, tvm.ir.GenericConst):
+ if isinstance(k, tvm_tirx.ir.GenericConst):
k = int(k.value.numpy().item())
if isinstance(k, tirx.IntImm):
k = int(k)
@@ -1921,7 +1922,7 @@ class Trilu(OnnxOpConverter):
row_idx = relax.op.reshape(relax.op.arange(0, m, dtype="int64"), (m, 1))
col_idx = relax.op.reshape(relax.op.arange(0, n, dtype="int64"), (1, n))
diff = relax.op.subtract(col_idx, row_idx)
- if tvm.ir.is_prim_expr(k):
+ if tvm_tirx.ir.is_prim_expr(k):
shape_value = k if str(k.ty) == "int64" else k.astype("int64")
k_int64 = bb.normalize(relax.op.shape_to_tensor(relax.ShapeExpr([shape_value])))
k_int64 = bb.normalize(relax.op.squeeze(k_int64, axis=[0]))
@@ -2198,7 +2199,7 @@ class Conv(OnnxOpConverter):
msg = (
f'Value {attr["auto_pad"]} in attribute "auto_pad" of operator Conv is invalid.'
)
- raise tvm.error.OpAttributeInvalid(msg)
+ raise tvm_tirx.error.OpAttributeInvalid(msg)
attr.pop("auto_pad")
conv_out = bb.normalize(
@@ -2288,7 +2289,7 @@ class ConvTranspose(OnnxOpConverter):
elif auto_pad == "NOTSET":
pass
else:
- raise tvm.error.OpAttributeInvalid(
+ raise tvm_tirx.error.OpAttributeInvalid(
f'Value {auto_pad} in attribute "auto_pad" of operator '
"ConvTranspose is invalid."
)
@@ -2330,7 +2331,7 @@ class CumSum(OnnxOpConverter):
axis_input = get_constant(inputs[1], params)
exclusive = attr.get("exclusive", 0) != 0
- if isinstance(axis_input, tvm.ir.GenericConst):
+ if isinstance(axis_input, tvm_tirx.ir.GenericConst):
axis_data = axis_input.value.numpy()
if axis_data.ndim == 0:
axis = int(axis_data.item())
@@ -2376,7 +2377,7 @@ class Squeeze(OnnxOpConverter):
def _impl_v13(cls, bb, inputs, attr, params):
data = inputs[0]
axis = get_constant(inputs[1], params)
- if isinstance(axis, tvm.ir.GenericConst):
+ if isinstance(axis, tvm_tirx.ir.GenericConst):
axis = tuple([int(x) for x in axis.value.numpy()])
return cls._squeeze(bb, data, axis)
@@ -2405,7 +2406,7 @@ class Squeeze(OnnxOpConverter):
@classmethod
def _squeeze(cls, bb, data, axis):
# If data is constant, perform computation directly.
- if isinstance(data, tvm.ir.GenericConst):
+ if isinstance(data, tvm_tirx.ir.GenericConst):
if isinstance(axis, tuple | type(None)):
out_data = _np.squeeze(data.value.numpy(), axis)
else:
@@ -2483,7 +2484,7 @@ class ConstantOfShape(OnnxOpConverter):
else:
dtype = "float32"
# If shape is a constant, treat it as a ShapeExpr.
- if isinstance(shape, tvm.ir.GenericConst):
+ if isinstance(shape, tvm_tirx.ir.GenericConst):
shape = relax.ShapeExpr(list(shape.value.numpy()))
# Special case where requested shape are constant
@@ -2603,10 +2604,10 @@ class Neg(OnnxOpConverter):
@classmethod
def _impl_v13(cls, bb, inputs, attr, params):
- if isinstance(inputs[0], tvm.ir.GenericConst):
+ if isinstance(inputs[0], tvm_tirx.ir.GenericConst):
data_np = inputs[0].value.numpy()
return relax.const(_np.negative(data_np), inputs[0].ty.dtype)
- if tvm.ir.is_prim_expr(inputs[0]):
+ if tvm_tirx.ir.is_prim_expr(inputs[0]):
return -inputs[0]
return relax.op.negative(inputs[0])
@@ -2616,7 +2617,7 @@ class Abs(OnnxOpConverter):
@classmethod
def _impl_v13(cls, bb, inputs, attr, params):
- if isinstance(inputs[0], tvm.ir.GenericConst):
+ if isinstance(inputs[0], tvm_tirx.ir.GenericConst):
output = _np.abs(inputs[0].value.numpy())
return relax.const(output, output.dtype)
return relax.op.abs(inputs[0])
@@ -2781,7 +2782,7 @@ class MultiInputBase(OnnxOpConverter):
def _impl_v1(cls, bb, inputs, attr, params):
if cls.numpy_op is None or cls.relax_op is None:
raise NotImplementedError("numpy_op and relax_op must be defined for MultiInputBase")
- if all([isinstance(inp, tvm.ir.GenericConst) for inp in inputs]):
+ if all([isinstance(inp, tvm_tirx.ir.GenericConst) for inp in inputs]):
np_inputs = [inp.value.numpy() for inp in inputs]
# numpy_op (np.mean/np.sum/np.min/np.max) reduces its first arg,
# treating any further positional args as `axis`, so calling it as
@@ -2856,7 +2857,7 @@ class Log(OnnxOpConverter):
@classmethod
def _impl_v13(cls, bb, inputs, attr, params):
- if isinstance(inputs[0], tvm.ir.GenericConst):
+ if isinstance(inputs[0], tvm_tirx.ir.GenericConst):
return relax.const(_np.log(inputs[0].value.numpy()), inputs[0].ty.dtype)
return relax.op.log(inputs[0])
@@ -2938,7 +2939,7 @@ class Split(OnnxOpConverter):
if splits is not None:
splits_rank = splits.ty.ndim
if splits is not None and splits_rank > 0:
- if isinstance(splits, tvm.ir.GenericConst):
+ if isinstance(splits, tvm_tirx.ir.GenericConst):
splits = splits.value.numpy()
indices = []
index = 0
@@ -2956,7 +2957,7 @@ class Split(OnnxOpConverter):
def get_prim_value_list(values):
new_values = []
for v in list(values):
- if tvm.ir.is_prim_expr(v):
+ if tvm_tirx.ir.is_prim_expr(v):
new_values.append(relax.prim_value(v))
else:
new_values.append(v)
@@ -2966,11 +2967,11 @@ def get_prim_value_list(values):
def _get_known_tensor_rank(expr: relax.Expr) -> int | None:
"""Return the statically known rank of an expression when available."""
- if isinstance(expr, tvm.ir.GenericConst):
+ if isinstance(expr, tvm_tirx.ir.GenericConst):
return len(expr.value.numpy().shape)
if isinstance(expr, relax.ShapeExpr):
return 1
- if tvm.ir.is_prim_expr(expr):
+ if tvm_tirx.ir.is_prim_expr(expr):
return 0
ty = expr.ty
if isinstance(ty, relax.TensorType):
@@ -2983,14 +2984,14 @@ def _get_known_tensor_length(expr: relax.Expr | None) -> int | None:
if expr is None:
return None
- if isinstance(expr, tvm.ir.GenericConst):
+ if isinstance(expr, tvm_tirx.ir.GenericConst):
np_value = expr.value.numpy()
if np_value.ndim != 1:
raise ValueError(f"Expected a 1-D tensor, but got ndim={np_value.ndim}.")
return int(np_value.shape[0])
if isinstance(expr, relax.ShapeExpr):
return len(expr.values)
- if tvm.ir.is_prim_expr(expr):
+ if tvm_tirx.ir.is_prim_expr(expr):
return 1
ty = expr.ty
if not isinstance(ty, relax.TensorType):
@@ -3029,9 +3030,9 @@ def _as_int64_tensor(bb: relax.BlockBuilder, expr: relax.Expr) -> relax.Expr:
if isinstance(expr, relax.ShapeExpr):
return bb.normalize(relax.op.shape_to_tensor(expr))
- if tvm.ir.is_prim_expr(expr):
+ if tvm_tirx.ir.is_prim_expr(expr):
return bb.normalize(relax.op.full((1,), expr, dtype="int64"))
- if isinstance(expr, tvm.ir.GenericConst):
+ if isinstance(expr, tvm_tirx.ir.GenericConst):
if expr.ty.dtype == "int64":
return expr
return bb.normalize(relax.op.astype(expr, "int64"))
@@ -3135,8 +3136,8 @@ class Slice(OnnxOpConverter):
axes = get_constant(inputs[3], params)
steps = get_constant(inputs[4], params)
all_constant_params = all(
- isinstance(param, tvm.ir.GenericConst | relax.ShapeExpr)
- or tvm.ir.is_prim_expr(param)
+ isinstance(param, tvm_tirx.ir.GenericConst | relax.ShapeExpr)
+ or tvm_tirx.ir.is_prim_expr(param)
or param is None
for param in [starts, ends, axes, steps]
)
@@ -3245,7 +3246,7 @@ class Slice(OnnxOpConverter):
f"Slice expects steps and starts to have the same length, but got "
f"{steps_len} and {axes_len}."
)
- if isinstance(steps_tensor, tvm.ir.GenericConst) and _np.any(
+ if isinstance(steps_tensor, tvm_tirx.ir.GenericConst) and _np.any(
steps_tensor.value.numpy() == 0
):
raise ValueError("Slice step values must be non-zero.")
@@ -3283,7 +3284,7 @@ class Pad(OnnxOpConverter):
if constant_value is None:
constant_value = 0.0
- if isinstance(pads, tvm.ir.GenericConst):
+ if isinstance(pads, tvm_tirx.ir.GenericConst):
pad_before, pad_after = _np.split(pads.value.numpy(), 2)
pad_before = _np.ndarray.tolist(pad_before)
pad_after = _np.ndarray.tolist(pad_after)
@@ -3292,7 +3293,7 @@ class Pad(OnnxOpConverter):
pad_mode = attr.get("mode", b"constant").decode("utf-8")
if pad_mode not in ["constant", "edge", "reflect"]:
- raise tvm.error.OpAttributeInvalid(
+ raise tvm_tirx.error.OpAttributeInvalid(
"Value " + pad_mode + ' in attribute "mode" is invalid for operator Pad.'
)
@@ -3318,7 +3319,7 @@ class Pad(OnnxOpConverter):
else:
constant_value = 0.0
- if not isinstance(pads, tvm.ir.GenericConst):
+ if not isinstance(pads, tvm_tirx.ir.GenericConst):
raise ValueError("Dynamic pads are not supported yet.")
pad_before, pad_after = _np.split(pads.value.numpy(), 2)
@@ -3328,7 +3329,7 @@ class Pad(OnnxOpConverter):
axes_input = inputs[3] if len(inputs) > 3 else None
if axes_input is not None:
axes_const = get_constant(axes_input, params)
- if not isinstance(axes_const, tvm.ir.GenericConst):
+ if not isinstance(axes_const, tvm_tirx.ir.GenericConst):
raise ValueError("Dynamic axes are not supported for Pad yet.")
axes = axes_const.value.numpy().tolist()
@@ -3357,7 +3358,7 @@ class Pad(OnnxOpConverter):
pad_before, pad_after, constant_value = cls._parse_pads_and_axes(inputs, params)
pad_mode = attr.get("mode", b"constant").decode("utf-8")
if pad_mode not in allowed_modes:
- raise tvm.error.OpAttributeInvalid(
+ raise tvm_tirx.error.OpAttributeInvalid(
"Value " + pad_mode + ' in attribute "mode" is invalid for operator Pad.'
)
@@ -3404,7 +3405,7 @@ class Tile(OnnxOpConverter):
@classmethod
def _impl_v13(cls, bb, inputs, attr, params):
reps = get_constant(inputs[1], params)
- if isinstance(reps, tvm.ir.GenericConst):
+ if isinstance(reps, tvm_tirx.ir.GenericConst):
reps = reps.value.numpy().tolist()
return bb.emit_te(topi.tile, inputs[0], reps)
@@ -3464,11 +3465,11 @@ class Expand(OnnxOpConverter):
assert len(data_shape) == len(target_shape)
# Apply ONNX v13 Expand broadcasting rules
for i, s in enumerate(target_shape):
- if isinstance(s, tvm.tirx.IntImm):
+ if isinstance(s, tvm_tirx.tirx.IntImm):
if s.value == -1:
# -1 means preserve the input dimension
target_shape[i] = data_shape[i]
- elif isinstance(data_shape[i], tvm.tirx.IntImm) and data_shape[i].value == 1:
+ elif isinstance(data_shape[i], tvm_tirx.tirx.IntImm) and data_shape[i].value == 1:
# Input dimension is 1, can broadcast to any target dimension >= 1
if s.value < 1:
raise ValueError(
@@ -3476,7 +3477,7 @@ class Expand(OnnxOpConverter):
f"at possition {i}. Target dimensions must be >= 1."
)
elif (
- isinstance(data_shape[i], tvm.tirx.IntImm)
+ isinstance(data_shape[i], tvm_tirx.tirx.IntImm)
and s.value == data_shape[i].value
):
# Dimensions match, no change needed
@@ -3486,7 +3487,7 @@ class Expand(OnnxOpConverter):
# This would "squeeze" the dimension - preserve input for safety
target_shape[i] = data_shape[i]
else:
- if isinstance(data_shape[i], tvm.tirx.IntImm):
+ if isinstance(data_shape[i], tvm_tirx.tirx.IntImm):
raise ValueError(
f"ONNX Expand: Cannot broadcast input shape {original_data_shape} "
f"to target shape {original_target_shape}. "
@@ -3501,7 +3502,7 @@ class Expand(OnnxOpConverter):
return relax.op.broadcast_to(data, relax.ShapeExpr(target_shape))
# If possible, directly expand to constant shape.
- if isinstance(shape, tvm.ir.GenericConst):
+ if isinstance(shape, tvm_tirx.ir.GenericConst):
new_shape = shape.value.numpy().tolist()
# ONNX Expand operator requires preserving target rank and broadcasting
# according to standard rules. Dimensions are right-aligned.
@@ -3557,7 +3558,7 @@ class Expand(OnnxOpConverter):
shape_vars = []
for i in range(shape_ndim):
- shape_vars.append(tvm.tirx.Var(f"x_{i}", "int64"))
+ shape_vars.append(tvm_tirx.tirx.Var(f"x_{i}", "int64"))
bb.match_cast(shape_dataflow_var, relax.ShapeType(shape_vars))
# Applying broadcasting rules for dynamic shapes
@@ -3722,7 +3723,7 @@ class Dropout(OnnxOpConverter):
ratio = 0.5
if len(inputs) >= 2 and inputs[1] is not None:
const = get_constant(inputs[1], params)
- if isinstance(const, tvm.ir.GenericConst):
+ if isinstance(const, tvm_tirx.ir.GenericConst):
ratio = float(const.value.numpy())
return relax.op.nn.dropout(inputs[0], ratio)
@@ -3862,7 +3863,7 @@ class Resize(OnnxOpConverter):
roi_static: list[float] | None = None
roi_dynamic_vec: relax.Expr | None = None
if roi is not None:
- if isinstance(roi, tvm.ir.GenericConst):
+ if isinstance(roi, tvm_tirx.ir.GenericConst):
roi_np = roi.value.numpy().tolist()
if len(roi_np) == 2 * ndims:
roi_static = roi_np[2:ndims] + roi_np[ndims + 2 : 2 * ndims]
@@ -3882,7 +3883,7 @@ class Resize(OnnxOpConverter):
# Convert scales to sizes if needed.
if scales is not None:
- if isinstance(scales, tvm.ir.GenericConst):
+ if isinstance(scales, tvm_tirx.ir.GenericConst):
scales = scales.value.numpy()
elif isinstance(scales, relax.expr.ShapeExpr):
scales = [int(val.value) for val in scales.values]
@@ -3894,7 +3895,7 @@ class Resize(OnnxOpConverter):
sizes.append((scales[i] * dim).astype("int64"))
sizes = sizes[2:]
else:
- if isinstance(sizes, tvm.ir.GenericConst):
+ if isinstance(sizes, tvm_tirx.ir.GenericConst):
sizes = sizes.value.numpy().astype("int64").tolist()[2:]
elif isinstance(sizes, relax.expr.ShapeExpr):
sizes = [int(val.value) for val in sizes.values][2:]
@@ -3969,7 +3970,7 @@ class AffineGrid(OnnxOpConverter):
align_corners = bool(attr.get("align_corners", 0))
# Extract size values
- if isinstance(size, tvm.ir.GenericConst):
+ if isinstance(size, tvm_tirx.ir.GenericConst):
size_vals = size.value.numpy().astype("int64").tolist()
elif isinstance(size, relax.expr.ShapeExpr):
size_vals = [int(v.value) for v in size.values]
@@ -4092,14 +4093,14 @@ class Range(OnnxOpConverter):
delta = get_constant(inputs[2], params)
def get_scalar_dtype(x):
- if tvm.ir.is_prim_expr(x):
+ if tvm_tirx.ir.is_prim_expr(x):
return str(getattr(x, "dtype", None) or x.ty)
return str(x.ty.dtype)
out_dtype = get_scalar_dtype(start)
def get_scalar_value(x):
- if isinstance(x, tvm.ir.GenericConst):
+ if isinstance(x, tvm_tirx.ir.GenericConst):
value = x.value.numpy()
if value.size != 1:
raise ValueError("Range scalar input must have exactly one element.")
@@ -4111,7 +4112,7 @@ class Range(OnnxOpConverter):
delta = get_scalar_value(delta)
def is_dynamic_scalar(x):
- return tvm.ir.is_prim_expr(x) or isinstance(x, relax.Expr)
+ return tvm_tirx.ir.is_prim_expr(x) or isinstance(x, relax.Expr)
out_dtype_is_float = _relax_dtype_is_floating_point(out_dtype)
@@ -4129,7 +4130,7 @@ class Range(OnnxOpConverter):
count_dtype = "float64" if out_dtype_is_float else "int64"
def scalar_expr(x, dtype):
- if tvm.ir.is_prim_expr(x):
+ if tvm_tirx.ir.is_prim_expr(x):
expr_dtype = str(getattr(x, "dtype", None) or x.ty)
if expr_dtype != "int64":
x = tirx.Cast(x, "int64")
@@ -4577,7 +4578,7 @@ class Flatten(OnnxOpConverter):
if axis == 0:
new_shape = (1, -1)
else:
- shape_flags = [isinstance(x, tvm.script.tirx.IntImm) for x in data_shape[0:axis]]
+ shape_flags = [isinstance(x, tvm_tirx.script.tirx.IntImm) for x in data_shape[0:axis]]
if all(shape_flags):
data_shape = [x.value for x in data_shape[0:axis]]
@@ -4813,7 +4814,7 @@ class ReduceMax(OnnxOpConverter):
axes = None
if len(inputs) > 1 and inputs[1] is not None:
axes_const = get_constant(inputs[1], params)
- assert isinstance(axes_const, tvm.ir.GenericConst), (
+ assert isinstance(axes_const, tvm_tirx.ir.GenericConst), (
"Only constant axes currently supported"
)
axes = axes_const.value.numpy().tolist()
@@ -4849,7 +4850,7 @@ class ReduceMin(OnnxOpConverter):
axes = None
if len(inputs) > 1 and inputs[1] is not None:
axes_const = get_constant(inputs[1], params)
- assert isinstance(axes_const, tvm.ir.GenericConst), (
+ assert isinstance(axes_const, tvm_tirx.ir.GenericConst), (
"Only constant axes currently supported"
)
axes = axes_const.value.numpy().tolist()
@@ -4885,7 +4886,7 @@ class ReduceSum(OnnxOpConverter):
axes = None
if len(inputs) > 1 and inputs[1] is not None:
axes_const = get_constant(inputs[1], params)
- assert isinstance(axes_const, tvm.ir.GenericConst), (
+ assert isinstance(axes_const, tvm_tirx.ir.GenericConst), (
"Only constant axes currently supported"
)
axes = axes_const.value.numpy().tolist()
@@ -4921,7 +4922,7 @@ class ReduceMean(OnnxOpConverter):
axes = None
if len(inputs) > 1 and inputs[1] is not None:
axes_const = get_constant(inputs[1], params)
- assert isinstance(axes_const, tvm.ir.GenericConst), (
+ assert isinstance(axes_const, tvm_tirx.ir.GenericConst), (
"Only constant axes currently supported"
)
axes = axes_const.value.numpy().tolist()
@@ -4957,7 +4958,7 @@ class ReduceProd(OnnxOpConverter):
axes = None
if len(inputs) > 1 and inputs[1] is not None:
axes_const = get_constant(inputs[1], params)
- assert isinstance(axes_const, tvm.ir.GenericConst), (
+ assert isinstance(axes_const, tvm_tirx.ir.GenericConst), (
"Only constant axes currently supported"
)
axes = axes_const.value.numpy().tolist()
@@ -4999,7 +5000,7 @@ class ReduceLogSumExp(OnnxOpConverter):
axes = None
if len(inputs) > 1 and inputs[1] is not None:
axes_const = get_constant(inputs[1], params)
- assert isinstance(axes_const, tvm.ir.GenericConst), (
+ assert isinstance(axes_const, tvm_tirx.ir.GenericConst), (
"Only constant axes currently supported"
)
axes = axes_const.value.numpy().tolist()
@@ -5044,7 +5045,7 @@ class ReduceLogSum(OnnxOpConverter):
axes = None
if len(inputs) > 1 and inputs[1] is not None:
axes_const = get_constant(inputs[1], params)
- assert isinstance(axes_const, tvm.ir.GenericConst), (
+ assert isinstance(axes_const, tvm_tirx.ir.GenericConst), (
"Only constant axes currently supported"
)
axes = axes_const.value.numpy().tolist()
@@ -5080,7 +5081,7 @@ class ReduceSumSquare(OnnxOpConverter):
axes = None
if len(inputs) > 1 and inputs[1] is not None:
axes_const = get_constant(inputs[1], params)
- assert isinstance(axes_const, tvm.ir.GenericConst), (
+ assert isinstance(axes_const, tvm_tirx.ir.GenericConst), (
"Only constant axes currently supported"
)
axes = axes_const.value.numpy().tolist()
@@ -5116,7 +5117,7 @@ class ReduceL1(OnnxOpConverter):
axes = None
if len(inputs) > 1 and inputs[1] is not None:
axes_const = get_constant(inputs[1], params)
- assert isinstance(axes_const, tvm.ir.GenericConst), (
+ assert isinstance(axes_const, tvm_tirx.ir.GenericConst), (
"Only constant axes currently supported"
)
axes = axes_const.value.numpy().tolist()
@@ -5152,7 +5153,7 @@ class ReduceL2(OnnxOpConverter):
axes = None
if len(inputs) > 1 and inputs[1] is not None:
axes_const = get_constant(inputs[1], params)
- assert isinstance(axes_const, tvm.ir.GenericConst), (
+ assert isinstance(axes_const, tvm_tirx.ir.GenericConst), (
"Only constant axes currently supported"
)
axes = axes_const.value.numpy().tolist()
@@ -5269,7 +5270,7 @@ class TopK(OnnxOpConverter):
def _impl_v11(cls, bb, inputs, attr, params):
data = inputs[0]
k = get_constant(inputs[1], params)
- if not isinstance(k, tvm.ir.GenericConst):
+ if not isinstance(k, tvm_tirx.ir.GenericConst):
raise ValueError("TopK k must be a constant")
k = int(k.value.numpy().item())
axis = attr.get("axis", -1)
@@ -5368,9 +5369,9 @@ class OneHot(OnnxOpConverter):
depth = get_constant(inputs[1], params)
values = get_constant(inputs[2], params)
axis = attr.get("axis", -1)
- assert isinstance(depth, tvm.ir.GenericConst), "Only constant depth currently supported."
+ assert isinstance(depth, tvm_tirx.ir.GenericConst), "Only constant depth currently supported."
depth = depth.value.numpy().tolist()
- assert isinstance(values, tvm.ir.GenericConst), "Only constant values currently supported."
+ assert isinstance(values, tvm_tirx.ir.GenericConst), "Only constant values currently supported."
values = values.value.numpy().tolist()
off_value, on_value = values
off_value, on_value = (
@@ -5648,7 +5649,7 @@ class SequenceErase(OnnxOpConverter):
if len(inputs) == 2:
position = inputs[1]
# Non constant position is not supported.
- if isinstance(position, tvm.ir.GenericConst):
+ if isinstance(position, tvm_tirx.ir.GenericConst):
position = int(position.value.numpy())
else:
raise NotImplementedError("Position must be a constant.")
@@ -5680,7 +5681,7 @@ class SequenceInsert(OnnxOpConverter):
if len(inputs) == 3:
position = inputs[2]
# Non constant position is not supported.
- if isinstance(position, tvm.ir.GenericConst):
+ if isinstance(position, tvm_tirx.ir.GenericConst):
position = int(position.value.numpy())
else:
raise NotImplementedError("Position must be a constant.")
@@ -5742,7 +5743,7 @@ class SplitToSequence(OnnxOpConverter):
split = _np.array(1)
else:
split = inputs[1]
- if not isinstance(split, tvm.ir.GenericConst):
+ if not isinstance(split, tvm_tirx.ir.GenericConst):
raise ValueError("Only constant split supported for SplitToSequence")
split = split.value.numpy()
@@ -5784,7 +5785,7 @@ class SequenceAt(OnnxOpConverter):
def _impl_v11(cls, bb, inputs, attr, params):
input_sequence = inputs[0]
position = inputs[1]
- assert isinstance(position, tvm.ir.GenericConst), (
+ assert isinstance(position, tvm_tirx.ir.GenericConst), (
"Only constant position supported for SequenceAt"
)
position = int(position.value.numpy())
@@ -5818,7 +5819,7 @@ class NonMaxSuppression(OnnxOpConverter):
center_point_box = attr.get("center_point_box", 0)
if max_output_boxes_per_class is not None and isinstance(
- max_output_boxes_per_class, tvm.ir.GenericConst
+ max_output_boxes_per_class, tvm_tirx.ir.GenericConst
):
max_output_boxes_per_class = int(max_output_boxes_per_class.value.numpy().item())
elif max_output_boxes_per_class is not None and isinstance(
@@ -5833,7 +5834,7 @@ class NonMaxSuppression(OnnxOpConverter):
else:
max_output_boxes_per_class = 0 # Default value
- if iou_threshold is not None and isinstance(iou_threshold, tvm.ir.GenericConst):
+ if iou_threshold is not None and isinstance(iou_threshold, tvm_tirx.ir.GenericConst):
iou_threshold = float(iou_threshold.value.numpy().item())
elif iou_threshold is not None and isinstance(iou_threshold, relax.Var):
var_name = iou_threshold.name
@@ -5845,7 +5846,7 @@ class NonMaxSuppression(OnnxOpConverter):
else:
iou_threshold = 0.5 # Default value
- if score_threshold is not None and isinstance(score_threshold, tvm.ir.GenericConst):
+ if score_threshold is not None and isinstance(score_threshold, tvm_tirx.ir.GenericConst):
score_threshold = float(score_threshold.value.numpy().item())
elif score_threshold is not None and isinstance(score_threshold, relax.Var):
var_name = score_threshold.name
@@ -5914,7 +5915,7 @@ class AllClassNMS(OnnxOpConverter):
center_point_box = attr.get("center_point_box", 0)
if max_output_boxes_per_class is not None and isinstance(
- max_output_boxes_per_class, tvm.ir.GenericConst
+ max_output_boxes_per_class, tvm_tirx.ir.GenericConst
):
max_output_boxes_per_class = int(max_output_boxes_per_class.value.numpy().item())
elif max_output_boxes_per_class is not None and isinstance(
@@ -5929,7 +5930,7 @@ class AllClassNMS(OnnxOpConverter):
else:
max_output_boxes_per_class = 0 # Default value
- if iou_threshold is not None and isinstance(iou_threshold, tvm.ir.GenericConst):
+ if iou_threshold is not None and isinstance(iou_threshold, tvm_tirx.ir.GenericConst):
iou_threshold = float(iou_threshold.value.numpy().item())
elif iou_threshold is not None and isinstance(iou_threshold, relax.Var):
var_name = iou_threshold.name
@@ -5941,7 +5942,7 @@ class AllClassNMS(OnnxOpConverter):
else:
iou_threshold = 0.5 # Default value
- if score_threshold is not None and isinstance(score_threshold, tvm.ir.GenericConst):
+ if score_threshold is not None and isinstance(score_threshold, tvm_tirx.ir.GenericConst):
score_threshold = float(score_threshold.value.numpy().item())
elif score_threshold is not None and isinstance(score_threshold, relax.Var):
var_name = score_threshold.name
@@ -6320,7 +6321,7 @@ class ONNXGraphImporter:
opset : opset version
Returns
-------
- mod : tvm.IRModule
+ mod : tvm_tirx.IRModule
The returned relax module
"""
has_if = any(node.op_type == "If" for node in graph.node)
@@ -6474,7 +6475,7 @@ class ONNXGraphImporter:
if unsupported_ops:
msg = "The following operators are not supported for frontend ONNX: "
msg += ", ".join(unsupported_ops)
- raise tvm.error.OpNotImplemented(msg)
+ raise tvm_tirx.error.OpNotImplemented(msg)
def _construct_nodes(self, graph: onnx.onnx_ml_pb2.GraphProto):
"""Nodes are stored as directed acyclic graph."""
@@ -6582,9 +6583,9 @@ class ONNXGraphImporter:
name = value_proto
return name
- def _parse_array(self, tensor_proto: onnx.onnx_ml_pb2.TensorProto) -> tvm.runtime.tensor:
+ def _parse_array(self, tensor_proto: onnx.onnx_ml_pb2.TensorProto) -> tvm_tirx.runtime.tensor:
np_array = get_numpy(tensor_proto).reshape(tuple(tensor_proto.dims))
- return tvm.runtime.tensor(np_array)
+ return tvm_tirx.runtime.tensor(np_array)
def _parse_attr(self, attr_proto: onnx.onnx_ml_pb2.AttributeProto) -> dict[str, Any]:
"""Convert a list of AttributeProto to a dict, with names as keys."""
@@ -6749,7 +6750,7 @@ def from_onnx(
Returns
-------
- mod : tvm.IRModule
+ mod : tvm_tirx.IRModule
The relax module for compilation
"""
# Error if the model version is below 1.1.0
diff --git a/python/tvm/relax/frontend/stablehlo/stablehlo_translator.py b/python/tvm/relax/frontend/stablehlo/stablehlo_translator.py
index d6330b4058..2e85f7608c 100644
--- a/python/tvm/relax/frontend/stablehlo/stablehlo_translator.py
+++ b/python/tvm/relax/frontend/stablehlo/stablehlo_translator.py
@@ -21,13 +21,14 @@
from collections.abc import Callable
from typing import Any
-import tvm
-from tvm import relax, tirx
+import tvm, tvm_tirx
+from tvm_tirx import tirx
+from tvm import relax
def _is_tensor_var(value) -> bool:
"""Whether ``value`` belongs to Relax's tensor domain."""
- return isinstance(value, tvm.ir.Var) and isinstance(value.ty, relax.TensorType)
+ return isinstance(value, tvm_tirx.ir.Var) and isinstance(value.ty, relax.TensorType)
class StableHLOImporter:
@@ -357,7 +358,7 @@ class StableHLOImporter:
"stablehlo.return": self._return,
}
- def from_stablehlo(self, model, input_info: list[tuple[tuple[int], str]]) -> tvm.IRModule:
+ def from_stablehlo(self, model, input_info: list[tuple[tuple[int], str]]) -> tvm_tirx.IRModule:
"""Convert a StableHLO Module to a Relax program.
Parameters
@@ -370,7 +371,7 @@ class StableHLOImporter:
Returns
-------
- output : tvm.IRModule
+ output : tvm_tirx.IRModule
The result IRModule with entry function "main"
"""
from jaxlib import mlir
@@ -423,7 +424,7 @@ class StableHLOImporter:
def from_stablehlo(
stablehlo_module,
input_info: list[tuple[tuple[int], str]] | None = None,
-) -> tvm.IRModule:
+) -> tvm_tirx.IRModule:
"""Convert a StableHLO Module to a Relax program
Parameters
@@ -436,7 +437,7 @@ def from_stablehlo(
Returns
-------
- output : tvm.IRModule
+ output : tvm_tirx.IRModule
The result IRModule with entry function "main"
"""
from jax._src.interpreters import mlir as jax_mlir
diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py
index 7cd4838c6e..0a66475452 100644
--- a/python/tvm/relax/frontend/tflite/tflite_frontend.py
+++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py
@@ -27,12 +27,13 @@ import math
import numpy as np
-import tvm
-from tvm import relax, tirx
-from tvm.ir import TupleType as _TupleType
+import tvm, tvm_tirx
+from tvm_tirx import tirx
+from tvm import relax
+from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax import op as _op
-from tvm.script import s_tir as Ts
+from tvm_tirx.script import s_tir as Ts
from .tflite_flexbuffer import FlexBufferDecoder
@@ -502,11 +503,11 @@ class OperatorConverter:
)
if len(raise_msg) > 0:
- raise tvm.error.OpNotImplemented(raise_msg)
+ raise tvm_tirx.error.OpNotImplemented(raise_msg)
def convert_operator_marker(self, op, op_name):
"""Reject TFLite marker builtins with an explicit diagnostic."""
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"TFLite operator marker {op_name} is not a Relax tensor operator"
)
@@ -604,12 +605,12 @@ class OperatorConverter:
return (container, shared_name)
if fallback_tensor is not None:
return ("", get_tensor_name(self.subgraph, fallback_tensor.tensor_idx))
- raise tvm.error.OpNotImplemented("VAR_HANDLE requires VarHandleOptions")
+ raise tvm_tirx.error.OpNotImplemented("VAR_HANDLE requires VarHandleOptions")
def _get_resource_key_for_handle(self, tensor, op_name):
tensor_name = get_tensor_name(self.subgraph, tensor.tensor_idx)
if tensor_name not in self.resource_handles:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"{op_name} requires a VAR_HANDLE in the same TFLite subgraph"
)
return self.resource_handles[tensor_name]
@@ -619,7 +620,7 @@ class OperatorConverter:
input_tensors = self.get_input_tensors(op)
output_tensors = self.get_output_tensors(op)
if len(input_tensors) != 0 or len(output_tensors) != 1:
- raise tvm.error.OpNotImplemented("VAR_HANDLE expects no inputs and one output")
+ raise tvm_tirx.error.OpNotImplemented("VAR_HANDLE expects no inputs and one output")
resource_key = self._get_var_handle_resource_key(op, output_tensors[0])
resource_tensor_name = get_tensor_name(self.subgraph, output_tensors[0].tensor_idx)
@@ -629,7 +630,7 @@ class OperatorConverter:
def convert_assign_variable(self, op):
"""Convert the CALL_ONCE initialization subset of ASSIGN_VARIABLE."""
if not self.conversion_state["in_call_once_init"]:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"ASSIGN_VARIABLE outside CALL_ONCE initialization is not supported by the "
"Relax TFLite frontend yet because it requires mutable resource state modeling."
)
@@ -637,7 +638,7 @@ class OperatorConverter:
input_tensors = self.get_input_tensors(op)
output_tensors = self.get_output_tensors(op)
if len(input_tensors) != 2 or len(output_tensors) != 0:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"ASSIGN_VARIABLE expects a resource handle and value input with no outputs"
)
@@ -652,12 +653,12 @@ class OperatorConverter:
input_tensors = self.get_input_tensors(op)
output_tensors = self.get_output_tensors(op)
if len(input_tensors) != 1 or len(output_tensors) != 1:
- raise tvm.error.OpNotImplemented("READ_VARIABLE expects one input and one output")
+ raise tvm_tirx.error.OpNotImplemented("READ_VARIABLE expects one input and one output")
resource_key = self._get_resource_key_for_handle(input_tensors[0], "READ_VARIABLE")
resource_values = self.conversion_state["resource_values"]
if resource_key not in resource_values:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"READ_VARIABLE requires a resource initialized by a supported CALL_ONCE subgraph"
)
return resource_values[resource_key]
@@ -691,9 +692,9 @@ class OperatorConverter:
pass
if key_dtype is None or value_dtype is None:
- raise tvm.error.OpNotImplemented("HASHTABLE requires HashtableOptions")
+ raise tvm_tirx.error.OpNotImplemented("HASHTABLE requires HashtableOptions")
if not self._is_supported_hashtable_type_pair(key_dtype, value_dtype):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"TFLite HASHTABLE only supports int64/string or string/int64 tables"
)
@@ -705,12 +706,12 @@ class OperatorConverter:
key_dtype,
value_dtype,
)
- raise tvm.error.OpNotImplemented("HASHTABLE requires HashtableOptions")
+ raise tvm_tirx.error.OpNotImplemented("HASHTABLE requires HashtableOptions")
def _get_hashtable_info_for_handle(self, tensor, op_name):
tensor_name = get_tensor_name(self.subgraph, tensor.tensor_idx)
if tensor_name not in self.hashtable_handles:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"{op_name} requires a HASHTABLE in the same TFLite subgraph"
)
return self.hashtable_handles[tensor_name]
@@ -732,38 +733,38 @@ class OperatorConverter:
def _get_string_tensor_value(self, tensor_wrapper, op_name):
"""Decode a constant TFLite string tensor buffer."""
if not self._is_tflite_string_type(tensor_wrapper.tensor.Type()):
- raise tvm.error.OpNotImplemented(f"{op_name} requires a TensorType.STRING tensor")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} requires a TensorType.STRING tensor")
if not self._has_tensor_buffer_data(tensor_wrapper):
- raise tvm.error.OpNotImplemented(f"{op_name} requires a constant string tensor")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} requires a constant string tensor")
data = bytes(tensor_wrapper.buffer.DataAsNumpy())
if len(data) < 4:
- raise tvm.error.OpNotImplemented(f"{op_name} has an invalid string tensor buffer")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} has an invalid string tensor buffer")
count = int(np.frombuffer(data, dtype="<i4", count=1)[0])
if count < 0:
- raise tvm.error.OpNotImplemented(f"{op_name} has an invalid string tensor count")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} has an invalid string tensor count")
header_size = 4 * (count + 2)
if len(data) < header_size:
- raise tvm.error.OpNotImplemented(f"{op_name} has an invalid string tensor offsets")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} has an invalid string tensor offsets")
offsets = np.frombuffer(data, dtype="<i4", count=count + 1, offset=4).astype(np.int64)
if np.any(offsets < header_size) or np.any(offsets > len(data)):
- raise tvm.error.OpNotImplemented(f"{op_name} has out-of-bounds string tensor offsets")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} has out-of-bounds string tensor offsets")
if np.any(offsets[:-1] > offsets[1:]):
- raise tvm.error.OpNotImplemented(f"{op_name} has non-monotonic string tensor offsets")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} has non-monotonic string tensor offsets")
try:
values = [
data[int(offsets[i]) : int(offsets[i + 1])].decode("utf-8") for i in range(count)
]
except UnicodeDecodeError as e:
- raise tvm.error.OpNotImplemented(f"{op_name} has invalid UTF-8 string data: {e}") from e
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} has invalid UTF-8 string data: {e}") from e
shape = self._get_tensor_shape_tuple(tensor_wrapper)
expected_count = math.prod(shape) if shape else 1
if expected_count != count:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"{op_name} string tensor buffer count does not match its shape"
)
return np.array(values, dtype=object).reshape(shape)
@@ -773,7 +774,7 @@ class OperatorConverter:
input_tensors = self.get_input_tensors(op)
output_tensors = self.get_output_tensors(op)
if len(input_tensors) != 0 or len(output_tensors) != 1:
- raise tvm.error.OpNotImplemented("HASHTABLE expects no inputs and one output")
+ raise tvm_tirx.error.OpNotImplemented("HASHTABLE expects no inputs and one output")
table_key, key_dtype, value_dtype = self._get_hashtable_key(op, output_tensors[0])
table_tensor_name = get_tensor_name(self.subgraph, output_tensors[0].tensor_idx)
@@ -787,7 +788,7 @@ class OperatorConverter:
def convert_hashtable_import(self, op):
"""Convert static metadata for the CALL_ONCE HASHTABLE_IMPORT subset."""
if not self.conversion_state["in_call_once_init"]:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_IMPORT outside CALL_ONCE initialization is not supported by the "
"Relax TFLite frontend yet because it requires mutable resource state modeling."
)
@@ -795,7 +796,7 @@ class OperatorConverter:
input_tensors = self.get_input_tensors(op)
output_tensors = self.get_output_tensors(op)
if len(input_tensors) != 3 or len(output_tensors) != 0:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_IMPORT expects table, keys, and values inputs with no outputs"
)
@@ -806,15 +807,15 @@ class OperatorConverter:
key_tensor.tensor.Type() != table_info["key_dtype"]
or value_tensor.tensor.Type() != table_info["value_dtype"]
):
- raise tvm.error.OpNotImplemented("HASHTABLE_IMPORT key/value dtypes mismatch")
+ raise tvm_tirx.error.OpNotImplemented("HASHTABLE_IMPORT key/value dtypes mismatch")
key_shape = self._get_tensor_shape_tuple(key_tensor)
value_shape = self._get_tensor_shape_tuple(value_tensor)
if key_shape != value_shape:
- raise tvm.error.OpNotImplemented("HASHTABLE_IMPORT requires keys and values same shape")
+ raise tvm_tirx.error.OpNotImplemented("HASHTABLE_IMPORT requires keys and values same shape")
if not self._has_tensor_buffer_data(key_tensor) or not self._has_tensor_buffer_data(
value_tensor
):
- raise tvm.error.OpNotImplemented("HASHTABLE_IMPORT requires constant keys and values")
+ raise tvm_tirx.error.OpNotImplemented("HASHTABLE_IMPORT requires constant keys and values")
if self._is_tflite_string_type(table_info["key_dtype"]):
keys = self._get_string_tensor_value(key_tensor, "HASHTABLE_IMPORT")
@@ -826,7 +827,7 @@ class OperatorConverter:
values = self.get_tensor_value(value_tensor)
if np.unique(keys).size != keys.size:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_IMPORT with duplicate keys is not supported"
)
@@ -849,7 +850,7 @@ class OperatorConverter:
input_tensors = self.get_input_tensors(op)
output_tensors = self.get_output_tensors(op)
if len(input_tensors) != 3 or len(output_tensors) != 1:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_FIND expects table, query, and default inputs with one output"
)
@@ -859,7 +860,7 @@ class OperatorConverter:
table_key = table_info["table_key"]
hashtable_values = self.conversion_state["hashtable_values"]
if table_key not in hashtable_values:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_FIND requires a table initialized by a supported CALL_ONCE subgraph"
)
table_values = hashtable_values[table_key]
@@ -869,26 +870,26 @@ class OperatorConverter:
or default_tensor.tensor.Type() != table_values["value_dtype"]
or output_tensor.tensor.Type() != table_values["value_dtype"]
):
- raise tvm.error.OpNotImplemented("HASHTABLE_FIND key/value dtypes mismatch")
+ raise tvm_tirx.error.OpNotImplemented("HASHTABLE_FIND key/value dtypes mismatch")
if not (
self._is_tflite_string_type(table_values["key_dtype"])
and table_values["value_dtype"] == TensorType.INT64
):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_FIND only supports constant string -> int64 tables"
)
if not self._has_tensor_buffer_data(query_tensor):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_FIND with runtime string queries is not supported"
)
if not self._has_tensor_buffer_data(default_tensor):
- raise tvm.error.OpNotImplemented("HASHTABLE_FIND requires constant default values")
+ raise tvm_tirx.error.OpNotImplemented("HASHTABLE_FIND requires constant default values")
query_shape = self._get_tensor_shape_tuple(query_tensor)
output_shape = self._get_tensor_shape_tuple(output_tensor)
if output_shape != query_shape:
- raise tvm.error.OpNotImplemented("HASHTABLE_FIND output shape must match query shape")
+ raise tvm_tirx.error.OpNotImplemented("HASHTABLE_FIND output shape must match query shape")
query_values = self._get_string_tensor_value(query_tensor, "HASHTABLE_FIND")
default_values = self.get_tensor_value(default_tensor)
@@ -898,7 +899,7 @@ class OperatorConverter:
elif default_shape == query_shape:
result = default_values.astype(np.int64).copy()
else:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_FIND default value must be scalar or match query shape"
)
@@ -920,7 +921,7 @@ class OperatorConverter:
input_tensors = self.get_input_tensors(op)
output_tensors = self.get_output_tensors(op)
if len(input_tensors) != 3 or len(output_tensors) != 2:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_LOOKUP expects lookup, key, and value inputs with two outputs"
)
@@ -931,19 +932,19 @@ class OperatorConverter:
lookup_tensor.tensor.Type() != TensorType.INT32
or key_tensor.tensor.Type() != TensorType.INT32
):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_LOOKUP requires int32 lookup and key tensors"
)
if self._is_tflite_string_type(value_tensor.tensor.Type()):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_LOOKUP with TensorType.STRING values is not supported"
)
if value_tensor.tensor.Type() != output_tensor.tensor.Type():
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_LOOKUP output dtype must match the value tensor dtype"
)
if hits_tensor.tensor.Type() != TensorType.UINT8:
- raise tvm.error.OpNotImplemented("HASHTABLE_LOOKUP hits output must be uint8")
+ raise tvm_tirx.error.OpNotImplemented("HASHTABLE_LOOKUP hits output must be uint8")
lookup_shape = to_int_list(self.get_tensor_shape(lookup_tensor))
key_shape = to_int_list(self.get_tensor_shape(key_tensor))
@@ -952,23 +953,23 @@ class OperatorConverter:
hits_shape = to_int_list(self.get_tensor_shape(hits_tensor))
if len(lookup_shape) != 1 or len(key_shape) != 1 or len(value_shape) < 1:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_LOOKUP requires rank-1 lookup/key and rank>=1 value tensors"
)
if key_shape[0] != value_shape[0]:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_LOOKUP requires key and value tensors to agree on row count"
)
if key_shape[0] == 0:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_LOOKUP requires a non-empty key/value table"
)
if output_shape != [lookup_shape[0]] + value_shape[1:]:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_LOOKUP output shape must match lookup count and value tail shape"
)
if hits_shape != [lookup_shape[0]]:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_LOOKUP hits output shape must match lookup count"
)
@@ -1000,17 +1001,17 @@ class OperatorConverter:
input_tensors = self.get_input_tensors(op)
output_tensors = self.get_output_tensors(op)
if len(input_tensors) != 1 or len(output_tensors) != 1:
- raise tvm.error.OpNotImplemented("HASHTABLE_SIZE expects one input and one output")
+ raise tvm_tirx.error.OpNotImplemented("HASHTABLE_SIZE expects one input and one output")
from tflite.TensorType import TensorType
if output_tensors[0].tensor.Type() != TensorType.INT64:
- raise tvm.error.OpNotImplemented("HASHTABLE_SIZE output must be int64")
+ raise tvm_tirx.error.OpNotImplemented("HASHTABLE_SIZE output must be int64")
table_info = self._get_hashtable_info_for_handle(input_tensors[0], "HASHTABLE_SIZE")
table_key = table_info["table_key"]
hashtable_values = self.conversion_state["hashtable_values"]
if table_key not in hashtable_values:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"HASHTABLE_SIZE requires a table initialized by a supported CALL_ONCE subgraph"
)
return relax.const(np.array([hashtable_values[table_key]["size"]], dtype=np.int64), "int64")
@@ -1112,7 +1113,7 @@ class OperatorConverter:
# Ensure that all zero points are zeros
zero_point = tflite_zero_point
if not np.all(zero_point == 0):
- raise tvm.error.OpAttributeInvalid(
+ raise tvm_tirx.error.OpAttributeInvalid(
"TFLite per-axis quantization restricts all zero points to be"
+ " 0, but a non-zero value is observed"
)
@@ -1332,8 +1333,8 @@ class OperatorConverter:
# Get min/max of the output dtype. This will be used to ensure that clip a_min/a_max are not
# beyond the dtype range.
- qmin = float(tvm.tirx.min_value(dtype).value)
- qmax = float(tvm.tirx.max_value(dtype).value)
+ qmin = float(tvm_tirx.tirx.min_value(dtype).value)
+ qmax = float(tvm_tirx.tirx.max_value(dtype).value)
# The input expr is a quantized tensor with its scale and zero point. We calculate the
# suitable clip off points based on these scale and zero point.
@@ -1347,7 +1348,7 @@ class OperatorConverter:
return relax.op.clip(expr, min=max(qmin, quantize(0.0)), max=qmax)
fused_activation_fn_str = self.activation_fn_type[fused_activation_fn]
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"Quantized activation {fused_activation_fn_str} is not supported yet."
)
@@ -1374,7 +1375,7 @@ class OperatorConverter:
target_value, success = try_infer_value(
target_expr,
parameters={
- k: tvm.runtime.tensor(np.array(v)) for k, v in self.exp_tab.params.items()
+ k: tvm_tirx.runtime.tensor(np.array(v)) for k, v in self.exp_tab.params.items()
},
)
"""
@@ -1525,7 +1526,7 @@ class OperatorConverter:
input_tensor_rank = len(input_tensor.tensor.ShapeAsNumpy())
if self.is_quantized(op):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"TFLite quantized L2_NORMALIZATION operator is not supported yet."
)
@@ -1539,7 +1540,7 @@ class OperatorConverter:
# if we have fused activation fn
if output_tensor.qnn_params:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"TFLite quantized L2_NORMALIZATION operator is not supported yet."
)
out = self.convert_fused_activation_function(out, fused_activation_fn)
@@ -1553,7 +1554,7 @@ class OperatorConverter:
from tflite.LocalResponseNormalizationOptions import LocalResponseNormalizationOptions
if self.is_quantized(op):
- raise tvm.error.OpNotImplemented("TFlite quantized LRN operator is not supported yet.")
+ raise tvm_tirx.error.OpNotImplemented("TFlite quantized LRN operator is not supported yet.")
input_tensors = self.get_input_tensors(op)
assert len(input_tensors) == 1, "input tensors length should be 1"
@@ -1693,7 +1694,7 @@ class OperatorConverter:
def is_dynamic(tensor):
return self.has_expr(tensor.tensor_idx) and not isinstance(
- self.get_expr(tensor.tensor_idx), tvm.ir.GenericConst
+ self.get_expr(tensor.tensor_idx), tvm_tirx.ir.GenericConst
)
def static_scalar(tensor):
@@ -2095,7 +2096,7 @@ class OperatorConverter:
elif dtype.startswith(("int", "uint")):
op_fn = _op.bitwise_and
else:
- raise tvm.error.OpNotImplemented(f"STABLEHLO_AND with dtype {dtype} is not supported")
+ raise tvm_tirx.error.OpNotImplemented(f"STABLEHLO_AND with dtype {dtype} is not supported")
return self.bb.normalize(op_fn(lhs, rhs))
def _convert_stablehlo_or(self, op):
@@ -2113,7 +2114,7 @@ class OperatorConverter:
elif dtype.startswith(("int", "uint")):
op_fn = _op.bitwise_or
else:
- raise tvm.error.OpNotImplemented(f"STABLEHLO_OR with dtype {dtype} is not supported")
+ raise tvm_tirx.error.OpNotImplemented(f"STABLEHLO_OR with dtype {dtype} is not supported")
return self.bb.normalize(op_fn(lhs, rhs))
def _convert_stablehlo_ternary(self, op, relax_op):
@@ -2145,7 +2146,7 @@ class OperatorConverter:
# A malformed flatbuffer may declare a BuiltinOptions2 type without
# carrying the actual options table. Fail cleanly instead of raising
# an opaque AttributeError when accessing the missing payload.
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"{options_cls.__name__} is required but missing from the operator"
)
# Look up the expected BuiltinOptions2 enum value by matching the class
@@ -2165,7 +2166,7 @@ class OperatorConverter:
try:
return [int(dim) for dim in self.get_tensor_shape(tensor)]
except (TypeError, ValueError) as err:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"{op_name} requires statically-known tensor shapes"
) from err
@@ -2179,7 +2180,7 @@ class OperatorConverter:
"""Return expr dtype if the StableHLO subset supports it."""
dtype = expr.ty.dtype.dtype
if not dtype.startswith("float"):
- raise tvm.error.OpNotImplemented(f"{op_name} with dtype {dtype} is not supported")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} with dtype {dtype} is not supported")
return dtype
def _convert_stablehlo_cbrt(self, op):
@@ -2219,7 +2220,7 @@ class OperatorConverter:
def _get_stablehlo_simple_body_op(self, body_subgraph_index, parent_op_name, input_count):
"""Return the single operator from a simple StableHLO body subgraph."""
if body_subgraph_index <= 0 or body_subgraph_index >= self.model.SubgraphsLength():
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"{parent_op_name} requires a valid non-main body subgraph"
)
@@ -2229,7 +2230,7 @@ class OperatorConverter:
or body_subgraph.OutputsLength() != 1
or body_subgraph.OperatorsLength() != 1
):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"{parent_op_name} only supports single-op body subgraphs"
)
@@ -2240,13 +2241,13 @@ class OperatorConverter:
):
"""Validate that the StableHLO reduce init value matches the Relax identity."""
if self.has_expr(init_tensor.tensor_idx):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"{parent_op_name} with dynamic init values is not supported"
)
init_value = np.asarray(self.get_tensor_value(init_tensor))
if init_value.shape not in [(), (1,)]:
- raise tvm.error.OpNotImplemented(f"{parent_op_name} requires scalar init values")
+ raise tvm_tirx.error.OpNotImplemented(f"{parent_op_name} requires scalar init values")
dtype = init_value.dtype
scalar = init_value.item()
@@ -2269,12 +2270,12 @@ class OperatorConverter:
else:
is_identity = False
else:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"{parent_op_name} reducer {reducer_name} is not supported"
)
if not is_identity:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"{parent_op_name} init value must match the reducer identity"
)
@@ -2300,7 +2301,7 @@ class OperatorConverter:
"STABLEHLO_MULTIPLY": relax.op.prod,
}
if reducer_name not in reducers:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"STABLEHLO_REDUCE reducer {reducer_name} is not supported"
)
@@ -2322,14 +2323,14 @@ class OperatorConverter:
)
reducer_name = self.get_op_code_str(body_op)
if reducer_name != "STABLEHLO_MAXIMUM":
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_REDUCE_WINDOW only supports MAXIMUM reducer windows"
)
self._check_stablehlo_reduce_init(input_tensors[1], reducer_name, "STABLEHLO_REDUCE_WINDOW")
data_shape = self._get_static_tensor_shape(input_tensors[0], "STABLEHLO_REDUCE_WINDOW")
if len(data_shape) != 4:
- raise tvm.error.OpNotImplemented("STABLEHLO_REDUCE_WINDOW only supports 4D input")
+ raise tvm_tirx.error.OpNotImplemented("STABLEHLO_REDUCE_WINDOW only supports 4D input")
window_dimensions = self._get_stablehlo_i64_vector(opts.WindowDimensionsAsNumpy(), [])
window_strides = self._get_stablehlo_i64_vector(
@@ -2352,23 +2353,23 @@ class OperatorConverter:
or len(window_dilations) != 4
or len(padding) != 8
):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_REDUCE_WINDOW only supports rank-4 window attributes"
)
if window_dimensions[0] != 1 or window_dimensions[3] != 1:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_REDUCE_WINDOW only supports pooling over spatial dimensions"
)
if window_strides[0] != 1 or window_strides[3] != 1:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_REDUCE_WINDOW only supports unit batch/channel strides"
)
if base_dilations != [1, 1, 1, 1]:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_REDUCE_WINDOW with base dilation is not supported"
)
if padding[0] != 0 or padding[1] != 0 or padding[6] != 0 or padding[7] != 0:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_REDUCE_WINDOW only supports spatial padding"
)
@@ -2408,23 +2409,23 @@ class OperatorConverter:
index_vector_dim = int(opts.IndexVectorDim())
if indices_rank == 0 or index_vector_dim != indices_rank - 1:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_SCATTER only supports trailing index-vector dimensions"
)
if update_window_dims:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_SCATTER only supports point updates without update windows"
)
if inserted_window_dims != list(range(operand_rank)):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_SCATTER only supports point updates for every operand dimension"
)
if scatter_dims_to_operand_dims != list(range(operand_rank)):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_SCATTER only supports canonical scatter-to-operand dimensions"
)
if indices_shape[-1] != operand_rank or updates_shape != indices_shape[:-1]:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_SCATTER requires point update shapes to match scatter indices"
)
@@ -2439,7 +2440,7 @@ class OperatorConverter:
"STABLEHLO_MULTIPLY": "mul",
}
if reducer_name not in reductions:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"STABLEHLO_SCATTER reducer {reducer_name} is not supported"
)
@@ -2457,7 +2458,7 @@ class OperatorConverter:
input_tensors = self.get_input_tensors(op)
output_tensors = self.get_output_tensors(op)
if len(output_tensors) != 1:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_COMPOSITE only supports single-output decompositions"
)
@@ -2467,7 +2468,7 @@ class OperatorConverter:
composite_name.decode("utf-8") if composite_name is not None else "<unnamed>"
)
if opts.CompositeAttributesLength() != 0:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"STABLEHLO_COMPOSITE {composite_name} with composite attributes is not supported"
)
@@ -2476,16 +2477,16 @@ class OperatorConverter:
decomposition_subgraph_index <= 0
or decomposition_subgraph_index >= self.model.SubgraphsLength()
):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"STABLEHLO_COMPOSITE {composite_name} requires a valid decomposition subgraph"
)
decomposition_subgraph = self.model.Subgraphs(decomposition_subgraph_index)
if decomposition_subgraph.InputsLength() != len(input_tensors):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"STABLEHLO_COMPOSITE {composite_name} decomposition input count mismatch"
)
if decomposition_subgraph.OutputsLength() != 1:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"STABLEHLO_COMPOSITE {composite_name} only supports single-output decompositions"
)
@@ -2527,20 +2528,20 @@ class OperatorConverter:
input_tensors = self.get_input_tensors(op)
output_tensors = self.get_output_tensors(op)
if len(input_tensors) != 1 or len(output_tensors) != 1:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_SORT only supports single-input single-output sort"
)
opts = self._get_stablehlo_options(op, StablehloSortOptions)
if opts.IsStable():
- raise tvm.error.OpNotImplemented("STABLEHLO_SORT stable sort is not supported")
+ raise tvm_tirx.error.OpNotImplemented("STABLEHLO_SORT stable sort is not supported")
body_op = self._get_stablehlo_simple_body_op(
int(opts.ComparatorSubgraphIndex()), "STABLEHLO_SORT", 2
)
comparator_name = self.get_op_code_str(body_op)
if comparator_name != "STABLEHLO_COMPARE":
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"STABLEHLO_SORT comparator {comparator_name} is not supported"
)
@@ -2549,7 +2550,7 @@ class OperatorConverter:
compare_opts.CompareType()
== StablehloComparisonType.STABLEHLO_COMPARISON_TYPE_FLOAT_TOTAL_ORDER
):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_SORT with TOTALORDER comparator is not supported"
)
@@ -2560,7 +2561,7 @@ class OperatorConverter:
elif direction == _DIR.STABLEHLO_COMPARISON_DIRECTION_GT:
descending = True
else:
- raise tvm.error.OpNotImplemented("STABLEHLO_SORT only supports LT or GT comparators")
+ raise tvm_tirx.error.OpNotImplemented("STABLEHLO_SORT only supports LT or GT comparators")
data = self.get_tensor_expr(input_tensors[0])
return self.bb.normalize(
@@ -2584,15 +2585,15 @@ class OperatorConverter:
# dropped for single-device import. TFLite has no runtime kernel
# for general STABLEHLO_CUSTOM_CALL targets.
if opts.HasSideEffect():
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_CUSTOM_CALL Sharding with side effects is not supported"
)
if opts.CalledComputationsLength() != 0:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_CUSTOM_CALL Sharding with called computations is not supported"
)
if len(input_tensors) != 1 or len(output_tensors) != 1:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_CUSTOM_CALL Sharding requires one input and one output"
)
self._check_tensor_metadata_match(
@@ -2601,7 +2602,7 @@ class OperatorConverter:
return self.get_tensor_expr(input_tensors[0])
target = call_target_name or "<empty>"
- raise tvm.error.OpNotImplemented(f"STABLEHLO_CUSTOM_CALL target {target} is not supported")
+ raise tvm_tirx.error.OpNotImplemented(f"STABLEHLO_CUSTOM_CALL target {target} is not supported")
def _convert_stablehlo_rng_bit_generator(self, op):
"""Convert STABLEHLO_RNG_BIT_GENERATOR to a bit-exact call_tir kernel."""
@@ -2612,7 +2613,7 @@ class OperatorConverter:
input_tensors = self.get_input_tensors(op)
output_tensors = self.get_output_tensors(op)
if len(input_tensors) != 1 or len(output_tensors) != 2:
- raise tvm.error.OpNotImplemented(f"{op_name} expects one input and two outputs")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} expects one input and two outputs")
opts = self._get_stablehlo_options(op, StablehloRngBitGeneratorOptions)
algorithm_enum = opts.Algorithm()
@@ -2622,34 +2623,34 @@ class OperatorConverter:
elif algorithm_enum in (RngAlgorithm.PHILOX, RngAlgorithm.DEFAULT):
algorithm = "philox"
else:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"{op_name} algorithm {algorithm_enum} is not supported"
)
state_tensor = input_tensors[0]
if self.get_tensor_type_str(state_tensor.tensor.Type()) != "uint64":
- raise tvm.error.OpNotImplemented(f"{op_name} requires a uint64 initial state")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} requires a uint64 initial state")
state_shape = self._get_static_tensor_shape(state_tensor, op_name)
if len(state_shape) != 1:
- raise tvm.error.OpNotImplemented(f"{op_name} requires a 1-D initial state")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} requires a 1-D initial state")
state_len = int(state_shape[0])
# State-length constraints mirror the TFLite runtime kernel.
if algorithm == "threefry" and state_len != 2:
- raise tvm.error.OpNotImplemented(f"{op_name} THREEFRY requires a u64[2] state")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} THREEFRY requires a u64[2] state")
if algorithm == "philox" and state_len not in (2, 3):
- raise tvm.error.OpNotImplemented(f"{op_name} PHILOX requires a u64[2] or u64[3] state")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} PHILOX requires a u64[2] or u64[3] state")
out_state_tensor, out_tensor = output_tensors
if self.get_tensor_type_str(out_state_tensor.tensor.Type()) != "uint64":
- raise tvm.error.OpNotImplemented(f"{op_name} output state must be uint64")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} output state must be uint64")
out_state_shape = self._get_static_tensor_shape(out_state_tensor, op_name)
if list(out_state_shape) != list(state_shape):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"{op_name} output state shape must match the initial state"
)
out_dtype = self.get_tensor_type_str(out_tensor.tensor.Type())
if out_dtype not in ("int32", "int64", "uint32", "uint64"):
- raise tvm.error.OpNotImplemented(f"{op_name} output dtype {out_dtype} is not supported")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} output dtype {out_dtype} is not supported")
out_shape = tuple(self._get_static_tensor_shape(out_tensor, op_name))
function = _build_stablehlo_rng_bit_generator_function(
@@ -2692,11 +2693,11 @@ class OperatorConverter:
op_options = op.BuiltinOptions()
if op_options is None:
- raise tvm.error.OpNotImplemented(f"{options_cls.__name__} is required")
+ raise tvm_tirx.error.OpNotImplemented(f"{options_cls.__name__} is required")
options_type = getattr(BuiltinOptions, options_cls.__name__, None)
if options_type is not None and op.BuiltinOptionsType() != options_type:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"Unexpected BuiltinOptions type: expected "
f"{options_cls.__name__}, got {op.BuiltinOptionsType()}"
)
@@ -2707,9 +2708,9 @@ class OperatorConverter:
def _get_subgraph(self, subgraph_index, op_name, allow_main=False):
"""Return a validated TFLite subgraph by index."""
if subgraph_index < 0 or subgraph_index >= self.model.SubgraphsLength():
- raise tvm.error.OpNotImplemented(f"{op_name} requires a valid subgraph index")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} requires a valid subgraph index")
if not allow_main and subgraph_index == 0:
- raise tvm.error.OpNotImplemented(f"{op_name} cannot target the main subgraph")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} cannot target the main subgraph")
return self.model.Subgraphs(subgraph_index)
def _make_tuple_or_single(self, exprs):
@@ -2726,9 +2727,9 @@ class OperatorConverter:
"""Validate a referenced subgraph's input and output counts."""
subgraph = self._get_subgraph(subgraph_index, op_name)
if input_count is not None and subgraph.InputsLength() != input_count:
- raise tvm.error.OpNotImplemented(f"{op_name} subgraph input count mismatch")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} subgraph input count mismatch")
if output_count is not None and subgraph.OutputsLength() != output_count:
- raise tvm.error.OpNotImplemented(f"{op_name} subgraph output count mismatch")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} subgraph output count mismatch")
return subgraph
def _check_subgraph_interface(
@@ -2778,7 +2779,7 @@ class OperatorConverter:
def _check_tensor_metadata_match(self, actual, expected, op_name, tensor_role):
"""Validate that two TFLite tensors have matching static metadata."""
if self._get_tensor_metadata(actual) != self._get_tensor_metadata(expected):
- raise tvm.error.OpNotImplemented(f"{op_name} {tensor_role} tensor metadata mismatch")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} {tensor_role} tensor metadata mismatch")
def _check_subgraph_tensor_metadata(
self, subgraph, op_name, tensor_role, subgraph_indices, expected_tensors
@@ -2800,7 +2801,7 @@ class OperatorConverter:
tensor = tensor.tensor
dtype = self.get_tensor_type_str(tensor.Type())
if dtype != "bool" or tensor.ShapeLength() != 0:
- raise tvm.error.OpNotImplemented(f"{op_name} requires a scalar bool condition")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} requires a scalar bool condition")
def _get_subgraph_params(self, subgraph):
"""Create Relax parameters for a TFLite subgraph."""
@@ -2831,7 +2832,7 @@ class OperatorConverter:
lowering_stack = self.conversion_state["lowering_stack"]
if subgraph_index in lowering_stack:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"Recursive TFLite {op_name} subgraphs are not supported"
)
@@ -2982,7 +2983,7 @@ class OperatorConverter:
input_tensors = self.get_input_tensors(op)
output_tensors = self.get_output_tensors(op)
if len(input_tensors) < 1:
- raise tvm.error.OpNotImplemented("IF requires a condition input")
+ raise tvm_tirx.error.OpNotImplemented("IF requires a condition input")
self._require_scalar_bool_tensor(input_tensors[0], "IF")
branch_input_count = len(input_tensors) - 1
@@ -3019,9 +3020,9 @@ class OperatorConverter:
output_tensors = self.get_output_tensors(op)
loop_var_count = len(input_tensors)
if loop_var_count == 0:
- raise tvm.error.OpNotImplemented(f"{op_name} requires loop-carried inputs")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} requires loop-carried inputs")
if len(output_tensors) != loop_var_count:
- raise tvm.error.OpNotImplemented(f"{op_name} output count must match input count")
+ raise tvm_tirx.error.OpNotImplemented(f"{op_name} output count must match input count")
cond_subgraph = self._check_subgraph_interface(
cond_subgraph_index,
@@ -3086,11 +3087,11 @@ class OperatorConverter:
input_tensors = self.get_input_tensors(op)
output_tensors = self.get_output_tensors(op)
if len(input_tensors) != 0 or len(output_tensors) != 0:
- raise tvm.error.OpNotImplemented("CALL_ONCE with inputs or outputs is not supported")
+ raise tvm_tirx.error.OpNotImplemented("CALL_ONCE with inputs or outputs is not supported")
init_subgraph = self._get_subgraph(init_subgraph_index, "CALL_ONCE")
if init_subgraph.InputsLength() != 0 or init_subgraph.OutputsLength() != 0:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"CALL_ONCE with non-empty init subgraph I/O is not supported"
)
if init_subgraph.OperatorsLength() != 0:
@@ -3103,7 +3104,7 @@ class OperatorConverter:
for op_idx in range(init_subgraph.OperatorsLength()):
op_name = self.get_op_code_str(init_subgraph.Operators(op_idx))
if op_name not in supported_init_ops:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"CALL_ONCE init subgraph operator {op_name} is not supported"
)
@@ -3310,7 +3311,7 @@ class OperatorConverter:
# TOTALORDER compare is not expressible via Relax comparison ops.
if compare_type == StablehloComparisonType.STABLEHLO_COMPARISON_TYPE_FLOAT_TOTAL_ORDER:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_COMPARE with TOTALORDER comparison type is not supported"
)
@@ -3325,7 +3326,7 @@ class OperatorConverter:
}
relax_fn = direction_map.get(direction)
if relax_fn is None:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"Unsupported StableHLO comparison direction: {direction}"
)
@@ -3352,11 +3353,11 @@ class OperatorConverter:
interior = [int(d) for d in opts.InteriorPaddingAsNumpy()]
if any(d != 0 for d in interior):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_PAD with interior (dilation) padding is not supported"
)
if any(d < 0 for d in edge_low) or any(d < 0 for d in edge_high):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_PAD with negative edge padding (crop) is not supported"
)
@@ -3367,7 +3368,7 @@ class OperatorConverter:
if not self.has_expr(pad_value_tensor.tensor_idx):
pad_val = float(self.get_tensor_value(pad_value_tensor))
else:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_PAD with dynamic padding value is not supported"
)
@@ -3404,14 +3405,14 @@ class OperatorConverter:
# (assumes start values are constant in the flatbuffer)
# TODO: support dynamic start indices via Relax arithmetic
if any(self.has_expr(t.tensor_idx) for t in input_tensors[1:]):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_DYNAMIC_SLICE with dynamic start indices is not supported"
)
start_vals = [int(self.get_tensor_value(t)) for t in input_tensors[1:]]
operand_shape = [int(d) for d in self.get_tensor_shape(input_tensors[0])]
for start, size, dim in zip(start_vals, slice_sizes, operand_shape):
if start < 0 or start + size > dim:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_DYNAMIC_SLICE with out-of-bounds start indices is not supported"
)
end_vals = [s + sz for s, sz in zip(start_vals, slice_sizes)]
@@ -3449,13 +3450,13 @@ class OperatorConverter:
update_shape = self._get_static_tensor_shape(update_tensor, op_name)
rank = len(operand_shape)
if len(update_shape) != rank or len(start_tensors) != rank:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_DYNAMIC_UPDATE_SLICE requires operand, update, "
"and start-index ranks to match"
)
for dim, size in zip(operand_shape, update_shape):
if size > dim:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_DYNAMIC_UPDATE_SLICE update shape must be smaller than "
"or equal to operand shape for all dimensions"
)
@@ -3472,7 +3473,7 @@ class OperatorConverter:
start_vals = [int(np.asarray(self.get_tensor_value(t)).item()) for t in start_tensors]
for start, size, dim in zip(start_vals, update_shape, operand_shape):
if start < 0 or start + size > dim:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_DYNAMIC_UPDATE_SLICE with out-of-bounds update "
"indices is not supported"
)
@@ -3534,13 +3535,13 @@ class OperatorConverter:
lhs_shape = self._get_static_tensor_shape(input_tensors[0], "STABLEHLO_DOT_GENERAL")
rhs_shape = self._get_static_tensor_shape(input_tensors[1], "STABLEHLO_DOT_GENERAL")
if len(lhs_shape) != 2 or len(rhs_shape) != 2:
- raise tvm.error.OpNotImplemented("STABLEHLO_DOT_GENERAL only supports 2D matmul")
+ raise tvm_tirx.error.OpNotImplemented("STABLEHLO_DOT_GENERAL only supports 2D matmul")
if lhs_batch_dims or rhs_batch_dims:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_DOT_GENERAL with batching dimensions is not supported"
)
if lhs_contract_dims != [1] or rhs_contract_dims != [0]:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_DOT_GENERAL only supports canonical contracting dimensions"
)
@@ -3567,15 +3568,15 @@ class OperatorConverter:
opts.OutputSpatialDimensionsAsNumpy(), []
)
if input_spatial_dims != [1, 2]:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_CONVOLUTION only supports NHWC input layout"
)
if kernel_spatial_dims != [0, 1]:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_CONVOLUTION only supports HWIO kernel layout"
)
if output_spatial_dims != [1, 2]:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_CONVOLUTION only supports NHWC output layout"
)
@@ -3587,24 +3588,24 @@ class OperatorConverter:
or int(opts.OutputBatchDimension()) != 0
or int(opts.OutputFeatureDimension()) != 3
):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_CONVOLUTION only supports canonical NHWC/HWIO dimension numbers"
)
if int(opts.BatchGroupCount()) != 1:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_CONVOLUTION with batch_group_count > 1 is not supported"
)
if int(opts.FeatureGroupCount()) != 1:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_CONVOLUTION with feature_group_count > 1 is not supported"
)
data_shape = self._get_static_tensor_shape(input_tensors[0], "STABLEHLO_CONVOLUTION")
kernel_shape = self._get_static_tensor_shape(input_tensors[1], "STABLEHLO_CONVOLUTION")
if len(data_shape) != 4 or len(kernel_shape) != 4:
- raise tvm.error.OpNotImplemented("STABLEHLO_CONVOLUTION only supports 2D convolution")
+ raise tvm_tirx.error.OpNotImplemented("STABLEHLO_CONVOLUTION only supports 2D convolution")
if data_shape[3] != kernel_shape[2]:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_CONVOLUTION input channels must match kernel input channels"
)
@@ -3618,19 +3619,19 @@ class OperatorConverter:
)
if len(window_strides) != 2 or len(rhs_dilation) != 2:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_CONVOLUTION only supports two spatial dimensions"
)
if lhs_dilation != [1, 1]:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_CONVOLUTION with lhs dilation is not supported"
)
if any(window_reversal):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_CONVOLUTION with window reversal is not supported"
)
if len(padding) != 4:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_CONVOLUTION only supports 2D low/high padding"
)
@@ -3679,43 +3680,43 @@ class OperatorConverter:
output_shape = [int(d) for d in self.get_tensor_shape(output_tensors[0])]
if len(start_index_map) != 1:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_GATHER only supports one start_index_map entry"
)
axis = start_index_map[0]
if axis < 0 or axis >= len(data_shape):
- raise tvm.error.OpNotImplemented(f"Unsupported STABLEHLO_GATHER axis: {axis}")
+ raise tvm_tirx.error.OpNotImplemented(f"Unsupported STABLEHLO_GATHER axis: {axis}")
if collapsed_slice_dims != [axis]:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_GATHER only supports collapsed_slice_dims matching the gather axis"
)
if len(slice_sizes) != len(data_shape):
- raise tvm.error.OpNotImplemented("STABLEHLO_GATHER slice_sizes must match operand rank")
+ raise tvm_tirx.error.OpNotImplemented("STABLEHLO_GATHER slice_sizes must match operand rank")
for i, (size, dim) in enumerate(zip(slice_sizes, data_shape)):
expected = 1 if i == axis else dim
if size != expected:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_GATHER only supports take-equivalent slice_sizes"
)
if index_vector_dim != len(indices_shape) - 1:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_GATHER only supports trailing index_vector_dim"
)
if not indices_shape or indices_shape[index_vector_dim] != 1:
- raise tvm.error.OpNotImplemented("STABLEHLO_GATHER only supports index vector size 1")
+ raise tvm_tirx.error.OpNotImplemented("STABLEHLO_GATHER only supports index vector size 1")
indices_batch_shape = indices_shape[:index_vector_dim]
expected_offset_dims = list(range(axis)) + list(
range(axis + len(indices_batch_shape), len(data_shape) + len(indices_batch_shape) - 1)
)
if offset_dims != expected_offset_dims:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_GATHER offset_dims do not match Relax take output layout"
)
expected_output_shape = data_shape[:axis] + indices_batch_shape + data_shape[axis + 1 :]
if output_shape != expected_output_shape:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"STABLEHLO_GATHER output shape does not match Relax take semantics"
)
@@ -3741,7 +3742,7 @@ class OperatorConverter:
def convert_gelu(self, op):
"""Convert TFLite GELU"""
if self.is_quantized(op):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"The TFLite to Relax converter does not support quantized GELU operator yet."
)
@@ -3824,7 +3825,7 @@ class OperatorConverter:
if input_is_quantized and not comparison_op:
if not output_tensor.qnn_params:
- raise tvm.error.OpAttributeInvalid(
+ raise tvm_tirx.error.OpAttributeInvalid(
"Quantized TFLite elemwise operator output must have quantization parameters"
)
out = self.quantize(out, output_tensor)
@@ -3847,7 +3848,7 @@ class OperatorConverter:
def convert_cumsum(self, op):
"""Convert TFLite CUMSUM"""
if self.is_quantized(op):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"The TFLite to Relax converter does not support quantized CUMSUM operator yet."
)
@@ -3860,7 +3861,7 @@ class OperatorConverter:
input_expr = self.get_tensor_expr(input_tensors[0])
if self.has_expr(input_tensors[1].tensor_idx):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"The TFLite to Relax converter does not support dynamic axis for CUMSUM yet."
)
axis = self.get_tensor_value(input_tensors[1])
@@ -3874,7 +3875,7 @@ class OperatorConverter:
cumsum_options.Init(op_options.Bytes, op_options.Pos)
exclusive = cumsum_options.Exclusive()
if cumsum_options.Reverse():
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"The TFLite to Relax converter does not support reverse CUMSUM operator yet."
)
@@ -4223,7 +4224,7 @@ class OperatorConverter:
def _check_random_output_dtype(self, op_name, output_dtype, supported_dtypes):
if output_dtype not in supported_dtypes:
supported = ", ".join(supported_dtypes)
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"The TFLite {op_name} converter currently supports output dtype(s) "
f"{supported} only, but got {output_dtype}."
)
@@ -4483,7 +4484,7 @@ class OperatorConverter:
in_f32 = self.dequantize(in_expr, input_tensor)
weight_axis = weight_tensor.qnn_params["axis"]
if weight_axis != 0:
- raise tvm.error.OpAttributeInvalid(
+ raise tvm_tirx.error.OpAttributeInvalid(
f"FC weight QuantizedDimension() must be 0 (output-channel "
f"axis in [OC,IC] layout), got {weight_axis}"
)
@@ -4601,7 +4602,7 @@ class OperatorConverter:
if fused_activation_fn == ActivationFunctionType.TANH:
return relax.op.tanh(in_expr)
fused_activation_fn_str = self.activation_fn_type[fused_activation_fn]
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"Fused activation {fused_activation_fn_str} is not supported yet."
)
@@ -4641,7 +4642,7 @@ class OperatorConverter:
conv_options.Init(op_options.Bytes, op_options.Pos)
depth_multiplier = conv_options.DepthMultiplier()
else:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"Operator {conv_type} is not supported for frontend TFLite."
)
@@ -4758,7 +4759,7 @@ class OperatorConverter:
params["padding"] = [pad_top, pad_left, pad_bottom, pad_right]
else:
- raise tvm.error.OpAttributeUnImplemented(
+ raise tvm_tirx.error.OpAttributeUnImplemented(
f"Padding format {padding} is not supported for operator Conv."
)
@@ -4773,13 +4774,13 @@ class OperatorConverter:
if is_depthwise_conv:
if per_channel_depthwise:
if weight_axis != 3:
- raise tvm.error.OpAttributeInvalid(
+ raise tvm_tirx.error.OpAttributeInvalid(
"Per-channel DepthwiseConv2D weight QuantizedDimension() must be 3 "
f"(the flattened output-channel axis), got {weight_axis}"
)
scale_count = weight_tensor.tensor.Quantization().ScaleLength()
if scale_count != in_channels:
- raise tvm.error.OpAttributeInvalid(
+ raise tvm_tirx.error.OpAttributeInvalid(
"Per-channel DepthwiseConv2D weight scale count must match "
f"input_channels * depth_multiplier ({in_channels}), got {scale_count}"
)
@@ -4795,7 +4796,7 @@ class OperatorConverter:
w_f32 = relax.op.reshape(w_f32, depthwise_weight_shape)
else:
if weight_axis != 0:
- raise tvm.error.OpAttributeInvalid(
+ raise tvm_tirx.error.OpAttributeInvalid(
"Per-tensor DepthwiseConv2D weight QuantizedDimension() must be 0, "
f"got {weight_axis}"
)
@@ -4807,7 +4808,7 @@ class OperatorConverter:
)
else:
if weight_axis != 0:
- raise tvm.error.OpAttributeInvalid(
+ raise tvm_tirx.error.OpAttributeInvalid(
f"Conv2D weight QuantizedDimension() must be 0 (output-channel "
f"axis in [OC,KH,KW,IC] layout), got {weight_axis}"
)
@@ -4979,12 +4980,12 @@ class OperatorConverter:
if do_pad:
params["padding"] = [pad_front, pad_top, pad_left, pad_back, pad_bottom, pad_right]
else:
- raise tvm.error.OpAttributeUnImplemented(
+ raise tvm_tirx.error.OpAttributeUnImplemented(
f"Padding format {padding} is not supported for operator Conv3D."
)
if input_tensor.qnn_params:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"Quantized Conv3D is not yet supported in the Relax frontend."
)
@@ -5010,7 +5011,7 @@ class OperatorConverter:
# Handle fused activation.
if output_tensor.qnn_params:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"Quantized Conv3D is not yet supported in the Relax frontend."
)
@@ -5125,12 +5126,12 @@ class OperatorConverter:
output_pad = max(0, dim_stride - dim_kernel)
params["output_padding"][idx] = output_pad
else:
- raise tvm.error.OpAttributeUnImplemented(
+ raise tvm_tirx.error.OpAttributeUnImplemented(
f"Padding format {padding} is not supported for operator Conv3DTranspose."
)
if input_tensor.qnn_params:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"Quantized Conv3DTranspose is not yet supported in the Relax frontend."
)
@@ -5156,7 +5157,7 @@ class OperatorConverter:
# Handle fused activation.
if output_tensor.qnn_params:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"Quantized Conv3DTranspose is not yet supported in the Relax frontend."
)
@@ -5357,7 +5358,7 @@ class OperatorConverter:
elif np.issubdtype(np_dtype, np.integer):
identity = np.iinfo(np_dtype).min
else:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"UNSORTED_SEGMENT_MAX does not support output dtype {output_dtype}."
)
return relax.op.full(output_shape, relax.const(identity, output_dtype), output_dtype)
@@ -5368,7 +5369,7 @@ class OperatorConverter:
elif np.issubdtype(np_dtype, np.integer):
identity = np.iinfo(np_dtype).max
else:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"UNSORTED_SEGMENT_MIN does not support output dtype {output_dtype}."
)
return relax.op.full(output_shape, relax.const(identity, output_dtype), output_dtype)
@@ -5379,20 +5380,20 @@ class OperatorConverter:
if op_name == "SEGMENT_SUM":
segment_ids_tensor = input_tensors[1]
if self.has_expr(segment_ids_tensor.tensor_idx):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"TFLite SEGMENT_SUM with runtime segment_ids is not supported, "
"because TFLite does not encode a reliable output segment count."
)
segment_ids = self.get_tensor_value(segment_ids_tensor)
if np.any(segment_ids < 0):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"TFLite SEGMENT_SUM with negative segment ids is not supported."
)
return int(np.max(segment_ids)) + 1 if segment_ids.size else 0
num_segments_tensor = input_tensors[2]
if self.has_expr(num_segments_tensor.tensor_idx):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"TFLite {op_name} with runtime num_segments is not supported."
)
num_segments_value = self.get_tensor_value(num_segments_tensor)
@@ -5424,7 +5425,7 @@ class OperatorConverter:
if not self.has_expr(segment_ids_tensor.tensor_idx):
segment_ids_value = self.get_tensor_value(segment_ids_tensor)
if np.any(segment_ids_value < 0):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"TFLite {op_name} with negative segment ids is not supported."
)
@@ -5487,7 +5488,7 @@ class OperatorConverter:
from tflite.ReverseSequenceOptions import ReverseSequenceOptions
if self.is_quantized(op):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"TFLite does not support quantized REVERSE_SEQUENCE operator yet."
)
@@ -5535,54 +5536,54 @@ class OperatorConverter:
input_tensors = self.get_input_tensors(op)
output_tensors = self.get_output_tensors(op)
if len(input_tensors) != 2 or len(output_tensors) != 1:
- raise tvm.error.OpNotImplemented("RFFT2D expects two inputs and one output")
+ raise tvm_tirx.error.OpNotImplemented("RFFT2D expects two inputs and one output")
data_tensor, fft_length_tensor = input_tensors
output_tensor = output_tensors[0]
if data_tensor.tensor.Type() != TensorType.FLOAT32:
- raise tvm.error.OpNotImplemented("RFFT2D input must be float32")
+ raise tvm_tirx.error.OpNotImplemented("RFFT2D input must be float32")
if not self._is_tflite_complex64_type(output_tensor.tensor.Type()):
- raise tvm.error.OpNotImplemented("RFFT2D output must be COMPLEX64")
+ raise tvm_tirx.error.OpNotImplemented("RFFT2D output must be COMPLEX64")
if (
data_tensor.tensor.Sparsity() is not None
or fft_length_tensor.tensor.Sparsity() is not None
):
- raise tvm.error.OpNotImplemented("RFFT2D does not support sparse inputs")
+ raise tvm_tirx.error.OpNotImplemented("RFFT2D does not support sparse inputs")
input_shape = tuple(to_int_list(self.get_tensor_shape(data_tensor)))
tflite_output_shape = tuple(to_int_list(self.get_tensor_shape(output_tensor)))
if len(input_shape) < 2:
- raise tvm.error.OpNotImplemented("RFFT2D input rank must be at least 2")
+ raise tvm_tirx.error.OpNotImplemented("RFFT2D input rank must be at least 2")
try:
fft_length_value = self.get_tensor_value_or_prefetched(fft_length_tensor)
except (ValueError, TypeError):
- raise tvm.error.OpNotImplemented("RFFT2D requires a constant fft_length") from None
+ raise tvm_tirx.error.OpNotImplemented("RFFT2D requires a constant fft_length") from None
if fft_length_value is None:
- raise tvm.error.OpNotImplemented("RFFT2D requires a constant fft_length")
+ raise tvm_tirx.error.OpNotImplemented("RFFT2D requires a constant fft_length")
# Reject non-integer fft_length tensors before astype("int64") can silently
# truncate (e.g. float32 2.7 -> int64 2).
if not np.issubdtype(fft_length_value.dtype, np.integer):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"RFFT2D fft_length must be an integer tensor, got dtype {fft_length_value.dtype!r}"
)
fft_length = fft_length_value.astype("int64")
if tuple(fft_length.shape) != (2,):
- raise tvm.error.OpNotImplemented("RFFT2D fft_length must be a length-2 tensor")
+ raise tvm_tirx.error.OpNotImplemented("RFFT2D fft_length must be a length-2 tensor")
height = int(fft_length[0])
width = int(fft_length[1])
if height <= 0 or width <= 0:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"RFFT2D fft_length must be positive, got ({height}, {width})"
)
if height != int(input_shape[-2]) or width != int(input_shape[-1]):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"RFFT2D currently supports fft_length matching the input spatial shape"
)
expected_tflite_output_shape = input_shape[:-2] + (height, width // 2 + 1)
if tflite_output_shape != expected_tflite_output_shape:
- raise tvm.error.OpNotImplemented("RFFT2D output shape does not match fft_length")
+ raise tvm_tirx.error.OpNotImplemented("RFFT2D output shape does not match fft_length")
relax_output_shape = self._get_relax_tensor_shape(output_tensor)
# Dispatch: power-of-2 sizes use the O(N^2 log N) Cooley-Tukey FFT kernel;
@@ -5737,7 +5738,7 @@ class OperatorConverter:
pad_left, pad_right = get_pad_value(input_w, filter_w, stride_w)
params["padding"] = [pad_top, pad_left, pad_bottom, pad_right]
else:
- raise tvm.error.OpAttributeUnImplemented(
+ raise tvm_tirx.error.OpAttributeUnImplemented(
f"Padding format {padding} for operator Pool2D is not supported."
)
@@ -5770,7 +5771,7 @@ class OperatorConverter:
avg_pool_exp = relax.op.nn.avg_pool2d(square_exp, **params)
out = relax.op.sqrt(avg_pool_exp)
else:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"Operator {pool_type} pool is not supported for frontend TFLite."
)
@@ -5874,7 +5875,7 @@ class OperatorConverter:
mode = "REFLECT" if mode_byte == 0 else "SYMMETRIC"
if mode == "SYMMETRIC":
- raise tvm.error.OpAttributeUnImplemented(
+ raise tvm_tirx.error.OpAttributeUnImplemented(
"MIRROR_PAD with SYMMETRIC mode is not yet supported."
)
# Flatten tuple-of-tuples to a list for relax.op.nn.pad
@@ -5999,7 +6000,7 @@ class OperatorConverter:
from tflite.LSTMOptions import LSTMOptions
if self.is_quantized(op):
- raise tvm.error.OpNotImplemented("TFLite quantized LSTM is not supported yet.")
+ raise tvm_tirx.error.OpNotImplemented("TFLite quantized LSTM is not supported yet.")
input_tensors = self.get_input_tensors(op)
assert len(input_tensors) == 24, (
@@ -6022,7 +6023,7 @@ class OperatorConverter:
# Only coupled input-forget gate is supported.
if input_tensors[1].tensor_idx != -1 or input_tensors[5].tensor_idx != -1:
- raise tvm.error.OpNotImplemented("Only coupled input-forget LSTM is supported.")
+ raise tvm_tirx.error.OpNotImplemented("Only coupled input-forget LSTM is supported.")
# Peephole, projection, and layer norm are not modeled yet.
if (
@@ -6030,7 +6031,7 @@ class OperatorConverter:
or any(t.tensor_idx != -1 for t in input_tensors[16:18])
or any(t.tensor_idx != -1 for t in input_tensors[20:24])
):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"Peephole, projection, and layer norm LSTM are not supported yet."
)
@@ -6145,7 +6146,7 @@ class OperatorConverter:
from tflite.SVDFOptions import SVDFOptions
if self.is_quantized(op):
- raise tvm.error.OpNotImplemented("TFLite quantized SVDF is not supported yet.")
+ raise tvm_tirx.error.OpNotImplemented("TFLite quantized SVDF is not supported yet.")
input_tensors = self.get_input_tensors(op)
assert len(input_tensors) == 5, (
@@ -6172,7 +6173,7 @@ class OperatorConverter:
batch_size = int(batch_size)
num_filters = to_int_list(self.get_tensor_shape(input_tensors[1]))[0]
if num_filters % rank != 0:
- raise tvm.error.OpNotImplemented("SVDF num_filters must be divisible by rank.")
+ raise tvm_tirx.error.OpNotImplemented("SVDF num_filters must be divisible by rank.")
num_units = num_filters // rank
memory_size = to_int_list(self.get_tensor_shape(input_tensors[2]))[1]
@@ -6236,7 +6237,7 @@ class OperatorConverter:
from tflite.SequenceRNNOptions import SequenceRNNOptions
if self.is_quantized(op):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"TFLite quantized UNIDIRECTIONAL_SEQUENCE_RNN is not supported yet."
)
@@ -6351,7 +6352,7 @@ class OperatorConverter:
from tflite.UnidirectionalSequenceLSTMOptions import UnidirectionalSequenceLSTMOptions
if self.is_quantized(op):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"TFLite quantized UNIDIRECTIONAL_SEQUENCE_LSTM is not supported yet."
)
@@ -6374,13 +6375,13 @@ class OperatorConverter:
# Only coupled input-forget gate is supported.
if input_tensors[1].tensor_idx != -1 or input_tensors[5].tensor_idx != -1:
- raise tvm.error.OpNotImplemented("Only coupled input-forget LSTM is supported.")
+ raise tvm_tirx.error.OpNotImplemented("Only coupled input-forget LSTM is supported.")
if any(input_tensors[idx].tensor_idx != -1 for idx in [9, 10, 11]):
- raise tvm.error.OpNotImplemented("TFLite peephole LSTM is not supported yet.")
+ raise tvm_tirx.error.OpNotImplemented("TFLite peephole LSTM is not supported yet.")
if any(input_tensors[idx].tensor_idx != -1 for idx in [16, 17]):
- raise tvm.error.OpNotImplemented("TFLite projection LSTM is not supported yet.")
+ raise tvm_tirx.error.OpNotImplemented("TFLite projection LSTM is not supported yet.")
if any(input_tensors[idx].tensor_idx != -1 for idx in [20, 21, 22, 23]):
- raise tvm.error.OpNotImplemented("TFLite layer-norm LSTM is not supported yet.")
+ raise tvm_tirx.error.OpNotImplemented("TFLite layer-norm LSTM is not supported yet.")
# Weights (transposed once outside the loop).
w_f_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[2]))
@@ -6500,7 +6501,7 @@ class OperatorConverter:
from tflite.BuiltinOptions import BuiltinOptions
if self.is_quantized(op):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"TFLite quantized BIDIRECTIONAL_SEQUENCE_RNN is not supported yet."
)
@@ -6520,7 +6521,7 @@ class OperatorConverter:
fused_activation_fn = rnn_opts.FusedActivationFunction()
merge_outputs = rnn_opts.MergeOutputs()
if any(input_tensors[idx].tensor_idx != -1 for idx in [9, 10, 11]):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"TFLite BIDIRECTIONAL_SEQUENCE_RNN aux input is not supported yet."
)
@@ -6663,7 +6664,7 @@ class OperatorConverter:
from tflite.BuiltinOptions import BuiltinOptions
if self.is_quantized(op):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"TFLite quantized BIDIRECTIONAL_SEQUENCE_LSTM is not supported yet."
)
@@ -6687,11 +6688,11 @@ class OperatorConverter:
# ── Forward LSTM weights (transposed once outside the loop) ──
if input_tensors[1].tensor_idx != -1 or input_tensors[5].tensor_idx != -1:
- raise tvm.error.OpNotImplemented("Only coupled input-forget LSTM is supported.")
+ raise tvm_tirx.error.OpNotImplemented("Only coupled input-forget LSTM is supported.")
if any(input_tensors[idx].tensor_idx != -1 for idx in [9, 10, 11]):
- raise tvm.error.OpNotImplemented("TFLite peephole LSTM is not supported yet.")
+ raise tvm_tirx.error.OpNotImplemented("TFLite peephole LSTM is not supported yet.")
if any(input_tensors[idx].tensor_idx != -1 for idx in [16, 17]):
- raise tvm.error.OpNotImplemented("TFLite projection LSTM is not supported yet.")
+ raise tvm_tirx.error.OpNotImplemented("TFLite projection LSTM is not supported yet.")
fw_w_f_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[2]))
fw_w_c_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[3]))
@@ -6705,13 +6706,13 @@ class OperatorConverter:
# ── Backward LSTM weights (transposed once outside the loop) ──
if input_tensors[18].tensor_idx != -1 or input_tensors[22].tensor_idx != -1:
- raise tvm.error.OpNotImplemented("Only coupled input-forget LSTM is supported.")
+ raise tvm_tirx.error.OpNotImplemented("Only coupled input-forget LSTM is supported.")
if any(input_tensors[idx].tensor_idx != -1 for idx in [26, 27, 28]):
- raise tvm.error.OpNotImplemented("TFLite peephole LSTM is not supported yet.")
+ raise tvm_tirx.error.OpNotImplemented("TFLite peephole LSTM is not supported yet.")
if any(input_tensors[idx].tensor_idx != -1 for idx in [33, 34]):
- raise tvm.error.OpNotImplemented("TFLite projection LSTM is not supported yet.")
+ raise tvm_tirx.error.OpNotImplemented("TFLite projection LSTM is not supported yet.")
if any(input_tensors[idx].tensor_idx != -1 for idx in range(39, 48)):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"TFLite BIDIRECTIONAL_SEQUENCE_LSTM aux input is not supported yet."
)
@@ -6962,7 +6963,7 @@ class OperatorConverter:
)
if self.has_expr(dense_shape_tensor.tensor_idx):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"TFLite EMBEDDING_LOOKUP_SPARSE with runtime dense_shape is not supported."
)
@@ -6975,7 +6976,7 @@ class OperatorConverter:
if not self.has_expr(ids_tensor.tensor_idx):
ids_value = self.get_tensor_value(ids_tensor)
if np.any(ids_value < 0):
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"TFLite EMBEDDING_LOOKUP_SPARSE with negative ids is not supported."
)
@@ -7016,7 +7017,7 @@ class OperatorConverter:
elif combiner == CombinerType.SQRTN:
denominator_updates = relax.op.multiply(weights, weights)
else:
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
f"Unsupported TFLite EMBEDDING_LOOKUP_SPARSE combiner value {combiner}"
)
@@ -7338,7 +7339,7 @@ class OperatorConverter:
in_f32 = self.dequantize(in_expr, input_tensor)
weight_axis = weights_tensor.qnn_params["axis"]
if weight_axis != 0:
- raise tvm.error.OpAttributeInvalid(
+ raise tvm_tirx.error.OpAttributeInvalid(
f"TransposeConv weight QuantizedDimension() must be 0 "
f"(output-channel axis in OHWI layout), got {weight_axis}"
)
@@ -7982,11 +7983,11 @@ class OperatorConverter:
on_val = self.get_tensor_value(on_value).item()
off_val = self.get_tensor_value(off_value).item()
if "float" in dtype:
- on_prim = relax.prim_value(tvm.tirx.FloatImm(dtype, float(on_val)))
- off_prim = relax.prim_value(tvm.tirx.FloatImm(dtype, float(off_val)))
+ on_prim = relax.prim_value(tvm_tirx.tirx.FloatImm(dtype, float(on_val)))
+ off_prim = relax.prim_value(tvm_tirx.tirx.FloatImm(dtype, float(off_val)))
else:
- on_prim = relax.prim_value(tvm.tirx.IntImm(dtype, int(on_val)))
- off_prim = relax.prim_value(tvm.tirx.IntImm(dtype, int(off_val)))
+ on_prim = relax.prim_value(tvm_tirx.tirx.IntImm(dtype, int(on_val)))
+ off_prim = relax.prim_value(tvm_tirx.tirx.IntImm(dtype, int(off_val)))
out = relax.op.one_hot(indices_expr, on_prim, off_prim, depth, axis)
@@ -8324,7 +8325,7 @@ def _build_tflite_rfft2d_function(input_shape, output_pair_shape):
(batch, out_y, out_x) iteration is structured as S-TIR spatial axes so a
downstream ``tvm.tir.schedule`` pass can parallelize it.
"""
- from tvm.script.parser import tirx as T
+ from tvm_tirx.script.parser import tirx as T
batch = 1
for dim in input_shape[:-2]:
@@ -8398,7 +8399,7 @@ def _build_tflite_rfft2d_fft_function(input_shape, output_pair_shape):
Complexity is ``O(B * H * W * (log2(H) + log2(W)))``, vs the DFT
reference kernel's ``O(B * H * W * H * W)``.
"""
- from tvm.script.parser import tirx as T
+ from tvm_tirx.script.parser import tirx as T
batch = 1
for dim in input_shape[:-2]:
@@ -8513,8 +8514,8 @@ def _build_tflite_rfft2d_fft_function(input_shape, output_pair_shape):
# pipeline (which expects an SBlockRealize at the function body) accepts
# this kernel.
function_source = (
- "from tvm.script.parser import tirx as T\n"
- "from tvm.script import s_tir as Ts\n"
+ "from tvm_tirx.script.parser import tirx as T\n"
+ "from tvm_tirx.script import s_tir as Ts\n"
"@Ts.function(private=True, check_well_formed=False)\n"
"def kernel(\n"
f" data: T.Tensor({tuple(int(x) for x in input_shape)}, 'float32'),\n"
@@ -8593,7 +8594,7 @@ def _build_stablehlo_rng_bit_generator_function(algorithm, state_len, out_dtype,
tensor. The updated state keeps the key unchanged and only advances the counter,
which is the only behaviour the runtime relies on.
"""
- from tvm.script.parser import tirx as T
+ from tvm_tirx.script.parser import tirx as T
total = 1
for dim in out_shape:
@@ -8774,7 +8775,7 @@ def prepare_dense_matrix_from_sparse(sparse_tensor, sparse_tensor_value, sparse_
elif VectorType(v_type) == VectorType.Uint8:
return N.Uint8Flags
else:
- raise tvm.error.OpNotImplemented(f"The provided type {v_type} is not supported")
+ raise tvm_tirx.error.OpNotImplemented(f"The provided type {v_type} is not supported")
def _get_flattened_index(indices, shape):
index = 0
@@ -8877,7 +8878,7 @@ def prepare_dense_matrix_from_sparse(sparse_tensor, sparse_tensor_value, sparse_
def get_scalar_from_constant(expr):
"""Returns scalar value from Relax constant scalar."""
- assert isinstance(expr, tvm.ir.GenericConst) and not expr.value.shape, (
+ assert isinstance(expr, tvm_tirx.ir.GenericConst) and not expr.value.shape, (
"Expr is not a constant scalar."
)
value = expr.value.numpy()
@@ -9029,7 +9030,7 @@ def from_tflite(
Returns
-------
- mod : tvm.IRModule
+ mod : tvm_tirx.IRModule
The relax module for compilation.
params : dict of str to tvm.nd.NDArray
@@ -9150,7 +9151,7 @@ def from_tflite(
# inputs cannot be represented. This also covers runtime string queries
# for ops like HASHTABLE_FIND, whose constant-foldable subset is handled
# in the op converter.
- raise tvm.error.OpNotImplemented(
+ raise tvm_tirx.error.OpNotImplemented(
"Relax TFLite frontend does not support TensorType.STRING graph inputs "
"(e.g. runtime string queries)"
)
@@ -9182,7 +9183,7 @@ def from_tflite(
_, param_value_list = map(list, zip(*exp_tab.params.values()))
func_attrs = {}
func_attrs["num_input"] = len(input_list)
- func_attrs["params"] = [tvm.runtime.tensor(arr) for arr in param_value_list]
+ func_attrs["params"] = [tvm_tirx.runtime.tensor(arr) for arr in param_value_list]
relax_mod["main"] = relax_mod["main"].with_attrs(func_attrs)
return relax_mod
diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
index b6fc656e98..6b8a6a2e85 100644
--- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
+++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
@@ -28,10 +28,11 @@ from functools import reduce
import tvm_ffi
-import tvm
-from tvm import relax, tirx
-from tvm.ir import PrimType
-from tvm.runtime import DataTypeCode
+import tvm, tvm_tirx
+from tvm_tirx import tirx
+from tvm import relax
+from tvm_tirx.ir import PrimType
+from tvm_tirx.runtime import DataTypeCode
def _diagonal_einsum_subscripts(ndim: int, dim1: int, dim2: int) -> str:
@@ -720,13 +721,13 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
lhs, rhs = self.retrieve_args(node)
- if (isinstance(lhs, tvm.ir.Var) and isinstance(lhs.ty, relax.TensorType)) or (
- isinstance(rhs, tvm.ir.Var) and isinstance(rhs.ty, relax.TensorType)
+ if (isinstance(lhs, tvm_tirx.ir.Var) and isinstance(lhs.ty, relax.TensorType)) or (
+ isinstance(rhs, tvm_tirx.ir.Var) and isinstance(rhs.ty, relax.TensorType)
):
return call_binary_op(relax_op, lhs, rhs)
- elif isinstance(lhs, tvm.ir.GenericConst) and not isinstance(rhs, tvm.ir.GenericConst):
+ elif isinstance(lhs, tvm_tirx.ir.GenericConst) and not isinstance(rhs, tvm_tirx.ir.GenericConst):
return call_binary_op(relax_op, lhs, relax.const(rhs, dtype=lhs.ty.dtype))
- elif isinstance(rhs, tvm.ir.GenericConst) and not isinstance(lhs, tvm.ir.GenericConst):
+ elif isinstance(rhs, tvm_tirx.ir.GenericConst) and not isinstance(lhs, tvm_tirx.ir.GenericConst):
return call_binary_op(relax_op, relax.const(lhs, dtype=rhs.ty.dtype), rhs)
return intrinsic_op(lhs, rhs)
@@ -2434,18 +2435,18 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
return input_shape[axis]
return val
- if tvm.ir.is_prim_expr(bound):
+ if tvm_tirx.ir.is_prim_expr(bound):
value = _adjust(bound)
return relax.prim_value(value)
bound = _adjust(bound)
- if not tvm.ir.is_prim_expr(bound):
+ if not tvm_tirx.ir.is_prim_expr(bound):
bound = relax.prim_value(bound)
return bound
start = _normalize_bound(start)
end = _normalize_bound(end)
- if not tvm.ir.is_prim_expr(step):
+ if not tvm_tirx.ir.is_prim_expr(step):
step = relax.prim_value(step)
return self.block_builder.emit(
@@ -2954,9 +2955,9 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
if tensor_meta is not None and len(tensor_meta.shape) == 1:
num_selected = tensor_meta.shape[0]
if not isinstance(num_selected, int):
- num_selected = tvm.ir.Var(str(num_selected), "int64")
+ num_selected = tvm_tirx.ir.Var(str(num_selected), "int64")
else:
- num_selected = tvm.ir.Var(f"{node.name}_num_selected", "int64")
+ num_selected = tvm_tirx.ir.Var(f"{node.name}_num_selected", "int64")
indices = self.block_builder.match_cast(
indices, relax.TensorType([1, num_selected], "int64")
)
@@ -3045,7 +3046,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
x = self.env[node.args[0]]
if isinstance(x, list | tuple | relax.ShapeExpr | relax.Tuple):
return x[node.args[1]]
- elif isinstance(x, tvm.ir.Var) and isinstance(x.ty, relax.TupleType | relax.TensorType):
+ elif isinstance(x, tvm_tirx.ir.Var) and isinstance(x.ty, relax.TupleType | relax.TensorType):
if isinstance(x.ty, relax.TupleType):
return self.block_builder.emit(relax.TupleGetItem(x, node.args[1]))
@@ -3120,7 +3121,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
for i in expand_dim:
sliced_shape.insert(i, 1)
return self.block_builder.emit(relax.op.reshape(sliced, sliced_shape))
- elif isinstance(x, tvm.ir.GenericConst):
+ elif isinstance(x, tvm_tirx.ir.GenericConst):
dtype = x.ty.dtype
return relax.const(x.value.numpy()[node.args[1]], dtype)
else:
@@ -3130,7 +3131,7 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
x = self.env[node.args[0]]
shape = self.shape_of(x)
dtype = x.ty.dtype
- analyzer = tvm.sym.Analyzer()
+ analyzer = tvm_tirx.sym.Analyzer()
has_single_element = shape is not None and all(
analyzer.can_prove_equal(dim, 1) for dim in shape
)
diff --git a/python/tvm/relax/frontend/torch/dynamo.py b/python/tvm/relax/frontend/torch/dynamo.py
index e39dabd34e..e946bbb4a7 100644
--- a/python/tvm/relax/frontend/torch/dynamo.py
+++ b/python/tvm/relax/frontend/torch/dynamo.py
@@ -24,7 +24,7 @@ import functools
import tvm_ffi
-import tvm
+import tvm, tvm_tirx
from tvm.relax import build as relax_build
from .fx_translator import from_fx
@@ -37,12 +37,12 @@ def device_from_inputs(example_inputs):
return None
-def relax_dynamo(pipeline: tvm.ir.transform.Pass | None = None):
+def relax_dynamo(pipeline: tvm_tirx.ir.transform.Pass | None = None):
"""A helper function to create a relax backend.
Parameters
----------
- pipeline : Optional[tvm.ir.transform.Pass]
+ pipeline : Optional[tvm_tirx.ir.transform.Pass]
The pipeline to be applied to the relax module before sent to build.
Returns
@@ -63,7 +63,7 @@ def relax_dynamo(pipeline: tvm.ir.transform.Pass | None = None):
# ffi::Tensor items returned in containers back to torch.Tensor when
# the call site passed torch.Tensor inputs.
return nd_tensor
- if isinstance(nd_tensor, tvm.runtime.Tensor):
+ if isinstance(nd_tensor, tvm_tirx.runtime.Tensor):
return torch.from_numpy(nd_tensor.numpy())
elif isinstance(nd_tensor, tvm_ffi.Array):
return tuple(to_torch_tensor(x) for x in nd_tensor)
@@ -99,7 +99,7 @@ def relax_dynamo(pipeline: tvm.ir.transform.Pass | None = None):
for s in tensor.shape:
if isinstance(s, torch.SymInt):
if str(s) not in shape_vars:
- shape_vars[str(s)] = tvm.tirx.Var(str(s), "int64")
+ shape_vars[str(s)] = tvm_tirx.tirx.Var(str(s), "int64")
shape.append(shape_vars[str(s)])
else:
shape.append(s)
@@ -108,11 +108,11 @@ def relax_dynamo(pipeline: tvm.ir.transform.Pass | None = None):
mod = from_fx(graph_module, input_info)
if device.type == "cuda":
- dev = tvm.cuda(device.index)
- target = tvm.target.Target("cuda")
+ dev = tvm_tirx.cuda(device.index)
+ target = tvm_tirx.target.Target("cuda")
else:
- dev = tvm.cpu(0)
- target = tvm.target.Target(llvm_target())
+ dev = tvm_tirx.cpu(0)
+ target = tvm_tirx.target.Target(llvm_target())
# invoke optimization pipeline.
if pipeline is None:
@@ -149,7 +149,7 @@ def relax_dynamo(pipeline: tvm.ir.transform.Pass | None = None):
return _relax_backend
-def dynamo_capture_subgraphs(model, *params, **kwargs) -> tvm.IRModule:
+def dynamo_capture_subgraphs(model, *params, **kwargs) -> tvm_tirx.IRModule:
"""Capture subgraphs of the PyTorch model using torch.compile into an IRModule.
Parameters
@@ -177,7 +177,7 @@ def dynamo_capture_subgraphs(model, *params, **kwargs) -> tvm.IRModule:
keep_params_as_input = "keep_params_as_input" in kwargs and kwargs["keep_params_as_input"]
kwargs.pop("keep_params_as_input", None)
- mod = tvm.IRModule()
+ mod = tvm_tirx.IRModule()
def _capture(graph_module: fx.GraphModule, example_inputs):
assert isinstance(graph_module, torch.fx.GraphModule)
diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py
index 61ae87c723..5735e80a35 100644
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -33,7 +33,7 @@ except ImportError as err:
"torch is required by the PyTorch frontend. Install it with: pip install torch"
) from err
-import tvm
+import tvm, tvm_tirx
from tvm import relax
from .base_fx_graph_translator import BaseFXGraphImporter
@@ -43,7 +43,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
"""An importer from ExportedProgram to Relax."""
@staticmethod
- def _convert_pytorch_tensor_to_tvm(tensor_value: torch.Tensor) -> tvm.runtime.Tensor:
+ def _convert_pytorch_tensor_to_tvm(tensor_value: torch.Tensor) -> tvm_tirx.runtime.Tensor:
"""Convert a PyTorch tensor to TVM tensor, handling sparse tensors.
Parameters
@@ -53,7 +53,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
Returns
-------
- tvm.runtime.Tensor
+ tvm_tirx.runtime.Tensor
The converted TVM tensor.
"""
# PyTorch sparse tensors (layout != torch.strided) must be converted to dense.
@@ -65,12 +65,12 @@ class ExportedProgramImporter(BaseFXGraphImporter):
# Try DLPack conversion first (faster)
try:
- return tvm.runtime.from_dlpack(tensor_detached)
+ return tvm_tirx.runtime.from_dlpack(tensor_detached)
except (RuntimeError, BufferError):
# Fallback: convert to numpy and then to TVM tensor
# This handles cases where DLPack conversion fails
tensor_cpu = tensor_detached.cpu().contiguous()
- return tvm.runtime.tensor(tensor_cpu.numpy())
+ return tvm_tirx.runtime.tensor(tensor_cpu.numpy())
########## Unary Ops ##########
@@ -485,9 +485,9 @@ class ExportedProgramImporter(BaseFXGraphImporter):
else:
seq_len, batch_size, input_size = input_shape
- seq_len = int(seq_len) if isinstance(seq_len, tvm.tirx.IntImm) else seq_len
- batch_size = int(batch_size) if isinstance(batch_size, tvm.tirx.IntImm) else batch_size
- input_size = int(input_size) if isinstance(input_size, tvm.tirx.IntImm) else input_size
+ seq_len = int(seq_len) if isinstance(seq_len, tvm_tirx.tirx.IntImm) else seq_len
+ batch_size = int(batch_size) if isinstance(batch_size, tvm_tirx.tirx.IntImm) else batch_size
+ input_size = int(input_size) if isinstance(input_size, tvm_tirx.tirx.IntImm) else input_size
# Extract hidden size from the LSTM parameters
# The parameters are: [weight_ih, weight_hh, bias_ih, bias_hh]
# weight_ih shape: (4 * hidden_size, input_size)
@@ -789,9 +789,9 @@ class ExportedProgramImporter(BaseFXGraphImporter):
else:
seq_len, batch_size, input_size = input_shape
- seq_len = int(seq_len) if isinstance(seq_len, tvm.tirx.IntImm) else seq_len
- batch_size = int(batch_size) if isinstance(batch_size, tvm.tirx.IntImm) else batch_size
- input_size = int(input_size) if isinstance(input_size, tvm.tirx.IntImm) else input_size
+ seq_len = int(seq_len) if isinstance(seq_len, tvm_tirx.tirx.IntImm) else seq_len
+ batch_size = int(batch_size) if isinstance(batch_size, tvm_tirx.tirx.IntImm) else batch_size
+ input_size = int(input_size) if isinstance(input_size, tvm_tirx.tirx.IntImm) else input_size
# Extract hidden size from parameters
# For bidirectional: params has weights for both directions
@@ -1167,11 +1167,11 @@ class ExportedProgramImporter(BaseFXGraphImporter):
# tensor's own dimension size (common with dynamic shapes).
if isinstance(start, int) and start == 0 and isinstance(step, int) and step == 1:
in_shape = self.shape_of(x)
- if in_shape is not None and tvm.ir.is_prim_expr(end_val):
+ if in_shape is not None and tvm_tirx.ir.is_prim_expr(end_val):
actual_dim = dim if dim >= 0 else len(in_shape) + dim
dim_expr = in_shape[actual_dim]
- if tvm.ir.is_prim_expr(dim_expr):
- if tvm.ir.prim.expr_deep_equal(end_val, dim_expr):
+ if tvm_tirx.ir.is_prim_expr(dim_expr):
+ if tvm_tirx.ir.prim.expr_deep_equal(end_val, dim_expr):
return x
axes = [dim]
@@ -1468,8 +1468,8 @@ class ExportedProgramImporter(BaseFXGraphImporter):
assert storage_offset == 0, "as_strided with non-zero storage_offset is not supported yet"
# Only handle view-like cases where the provided strides align with a contiguous layout.
- can_check = all(isinstance(dim, int | tvm.tirx.IntImm) for dim in size) and all(
- isinstance(st, int | tvm.tirx.IntImm) for st in stride
+ can_check = all(isinstance(dim, int | tvm_tirx.tirx.IntImm) for dim in size) and all(
+ isinstance(st, int | tvm_tirx.tirx.IntImm) for st in stride
)
if can_check:
expected_stride = []
@@ -1588,7 +1588,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
graph_module, # torch.fx.GraphModule
operands: list[relax.Expr],
name_hint: str,
- ) -> tvm.ir.GlobalVar:
+ ) -> tvm_tirx.ir.GlobalVar:
"""Translate a branch subgraph (GraphModule) into a Relax function in the IRModule.
Parameters
@@ -1602,13 +1602,13 @@ class ExportedProgramImporter(BaseFXGraphImporter):
Returns
-------
- gv : tvm.ir.GlobalVar
+ gv : tvm_tirx.ir.GlobalVar
The GlobalVar referring to the newly added Relax function.
"""
# Cache: avoid importing the same subgraph twice.
cache_key = id(graph_module)
if not hasattr(self, "_branch_cache"):
- self._branch_cache: dict[int, tvm.ir.GlobalVar] = {}
+ self._branch_cache: dict[int, tvm_tirx.ir.GlobalVar] = {}
if not hasattr(self, "_branch_counter"):
self._branch_counter: int = 0
if cache_key in self._branch_cache:
@@ -1632,7 +1632,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
# Create fresh symbolic variables to avoid sharing with the caller function.
if orig_si.shape is not None:
new_shape = [
- tvm.ir.Var(s.name, s.ty) if tvm.ir.is_prim_var(s) else s
+ tvm_tirx.ir.Var(s.name, s.ty) if tvm_tirx.ir.is_prim_var(s) else s
for s in orig_si.shape
]
si = relax.TensorType(new_shape, orig_si.dtype)
@@ -1735,7 +1735,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
else node.kwargs.get("assert_msg", "Assertion Failed")
)
return self.block_builder.emit(
- relax.op.assert_op(condition, "{}", tvm.ir.StringImm(message))
+ relax.op.assert_op(condition, "{}", tvm_tirx.ir.StringImm(message))
)
def create_convert_map(
@@ -2097,8 +2097,8 @@ class ExportedProgramImporter(BaseFXGraphImporter):
}
def _process_derived_symbol(
- self, symbol, torch_symbol_to_relax_var: dict[str, tvm.tirx.Var]
- ) -> tuple[str, tvm.tirx.Expr | None]:
+ self, symbol, torch_symbol_to_relax_var: dict[str, tvm_tirx.tirx.Var]
+ ) -> tuple[str, tvm_tirx.tirx.Expr | None]:
"""Process a sympy symbol to generate a descriptive name and TIR expression."""
import sympy
@@ -2111,10 +2111,10 @@ class ExportedProgramImporter(BaseFXGraphImporter):
tir_expr = None
for arg in symbol.args:
if isinstance(arg, sympy.Integer):
- term = tvm.tirx.IntImm("int64", int(arg))
+ term = tvm_tirx.tirx.IntImm("int64", int(arg))
elif isinstance(arg, sympy.Symbol):
term = torch_symbol_to_relax_var.setdefault(
- str(arg), tvm.tirx.Var(str(arg), "int64")
+ str(arg), tvm_tirx.tirx.Var(str(arg), "int64")
)
else:
_, term = self._process_derived_symbol(arg, torch_symbol_to_relax_var)
@@ -2129,14 +2129,14 @@ class ExportedProgramImporter(BaseFXGraphImporter):
elif isinstance(symbol, sympy.Add):
tir_expr = tir_expr + term
- if isinstance(tir_expr, tvm.tirx.Add):
+ if isinstance(tir_expr, tvm_tirx.tirx.Add):
for const, var in [(tir_expr.a, tir_expr.b), (tir_expr.b, tir_expr.a)]:
- if isinstance(const, tvm.tirx.IntImm) and tvm.ir.is_prim_var(var):
+ if isinstance(const, tvm_tirx.tirx.IntImm) and tvm_tirx.ir.is_prim_var(var):
return f"{var.name}___{const.value}", tir_expr
- if isinstance(tir_expr, tvm.tirx.Mul):
+ if isinstance(tir_expr, tvm_tirx.tirx.Mul):
for const, var in [(tir_expr.a, tir_expr.b), (tir_expr.b, tir_expr.a)]:
- if isinstance(const, tvm.tirx.IntImm) and tvm.ir.is_prim_var(var):
+ if isinstance(const, tvm_tirx.tirx.IntImm) and tvm_tirx.ir.is_prim_var(var):
return f"{var.name}_{const.value}", tir_expr
return str(symbol), tir_expr
@@ -2147,7 +2147,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
"""Create relax input vars."""
parameters_buffers_constants = OrderedDict()
user_inputs = OrderedDict()
- torch_symbol_to_relax_var: dict[str, tvm.tirx.Var] = {}
+ torch_symbol_to_relax_var: dict[str, tvm_tirx.tirx.Var] = {}
range_constraints = {}
if hasattr(exported_program, "range_constraints"):
@@ -2208,7 +2208,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
relax_shape.append(derived_expr)
else:
shape_var = torch_symbol_to_relax_var.setdefault(
- symbol_name, tvm.tirx.Var(symbol_name, "int64")
+ symbol_name, tvm_tirx.tirx.Var(symbol_name, "int64")
)
relax_shape.append(shape_var)
else:
@@ -2230,7 +2230,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
unwrap_unit_return_tuple: bool,
no_bind_return_tuple: bool,
custom_convert_map: dict[str, Callable[[fx.Node, BaseFXGraphImporter], relax.Var]] | None,
- ) -> tvm.IRModule:
+ ) -> tvm_tirx.IRModule:
"""Convert a PyTorch ExportedProgram to a Relax program."""
# Update the conversion map with custom ops if provided.
@@ -2359,7 +2359,7 @@ def from_exported_program(
dict[str, Callable[[fx.Node, BaseFXGraphImporter], relax.Var]] | None
) = None,
run_ep_decomposition: bool = True,
-) -> tvm.IRModule:
+) -> tvm_tirx.IRModule:
"""Convert a PyTorch ExportedProgram to a Relax program
Parameters
@@ -2388,7 +2388,7 @@ def from_exported_program(
Returns
-------
- output : tvm.IRModule
+ output : tvm_tirx.IRModule
The import result IRModule, with the function "main" containing the
translated logic.
@@ -2401,7 +2401,7 @@ def from_exported_program(
.. code-block:: python
# Import the importer.
- import tvm
+ import tvm, tvm_tirx
from tvm.relax.frontend.torch import from_exported_program
import torch
from torch.export import export
@@ -2423,7 +2423,7 @@ def from_exported_program(
exported_program = export(torch_model, args=example_args)
# Use the importer to import the ExportedProgram to Relax.
- mod: tvm.IRModule = from_exported_program(exported_program)
+ mod: tvm_tirx.IRModule = from_exported_program(exported_program)
"""
def _is_sparse_tensor(value: object) -> bool:
diff --git a/python/tvm/relax/frontend/torch/fx_translator.py b/python/tvm/relax/frontend/torch/fx_translator.py
index 2e8300e096..e13f0dd403 100644
--- a/python/tvm/relax/frontend/torch/fx_translator.py
+++ b/python/tvm/relax/frontend/torch/fx_translator.py
@@ -22,7 +22,7 @@
from collections.abc import Callable
from functools import partial, reduce
-import tvm
+import tvm, tvm_tirx
from tvm import relax
from .base_fx_graph_translator import BaseFXGraphImporter
@@ -178,19 +178,19 @@ class TorchFXImporter(BaseFXGraphImporter):
lhs, rhs = self.retrieve_args(node)
- if (isinstance(lhs, tvm.ir.Var) and isinstance(lhs.ty, relax.TensorType)) or (
- isinstance(rhs, tvm.ir.Var) and isinstance(rhs.ty, relax.TensorType)
+ if (isinstance(lhs, tvm_tirx.ir.Var) and isinstance(lhs.ty, relax.TensorType)) or (
+ isinstance(rhs, tvm_tirx.ir.Var) and isinstance(rhs.ty, relax.TensorType)
):
output = call_binary_op(relax_op, lhs, rhs)
self.env[node.args[0]] = output
return output
- elif isinstance(lhs, tvm.ir.GenericConst):
+ elif isinstance(lhs, tvm_tirx.ir.GenericConst):
output = call_binary_op(relax_op, lhs, relax.const(rhs, dtype=lhs.ty.dtype))
self.env[node.args[0]] = output
return output
- elif isinstance(rhs, tvm.ir.GenericConst):
+ elif isinstance(rhs, tvm_tirx.ir.GenericConst):
output = call_binary_op(relax_op, relax.const(lhs, dtype=rhs.ty.dtype), rhs)
self.env[node.args[0]] = output
return output
@@ -1073,7 +1073,7 @@ class TorchFXImporter(BaseFXGraphImporter):
unwrap_unit_return_tuple: bool,
no_bind_return_tuple: bool,
custom_convert_map: dict | None = None,
- ) -> tvm.IRModule:
+ ) -> tvm_tirx.IRModule:
"""Convert a PyTorch FX GraphModule to a Relax program."""
from torch import fx
@@ -1100,7 +1100,7 @@ class TorchFXImporter(BaseFXGraphImporter):
dtype = self._convert_data_type(str(param.data.dtype))
inputs.append(relax.Var(name, relax.TensorType(shape, dtype)))
self.params[param] = inputs[-1]
- params.append(tvm.runtime.tensor(param.data.cpu().numpy()))
+ params.append(tvm_tirx.runtime.tensor(param.data.cpu().numpy()))
else:
func_attrs = None
@@ -1108,7 +1108,7 @@ class TorchFXImporter(BaseFXGraphImporter):
self._check_unsupported_func_type(graph.nodes)
sym_vars = {
- v.name: v for shape, _ in input_info for v in shape if isinstance(v, tvm.ir.Var)
+ v.name: v for shape, _ in input_info for v in shape if isinstance(v, tvm_tirx.ir.Var)
}
with self.block_builder.function(name=func_name, params=inputs.copy(), attrs=func_attrs):
@@ -1190,7 +1190,7 @@ def from_fx(
no_bind_return_tuple: bool = False,
custom_convert_map: dict | None = None,
default_image_layout: str = "NCHW",
-) -> tvm.IRModule:
+) -> tvm_tirx.IRModule:
"""Convert a PyTorch FX GraphModule to a Relax program
Parameters
@@ -1221,7 +1221,7 @@ def from_fx(
Returns
-------
- output : tvm.IRModule
+ output : tvm_tirx.IRModule
The import result IRModule, with the function "main" containing the
translated logic.
If `keep_params_as_input` is true, the "main" function have an attribute
@@ -1269,7 +1269,7 @@ def from_fx(
raise RuntimeError("Failed to export the PyTorch model to FX.")
# Use the importer to import the PyTorch model to Relax.
- mod: tvm.IRModule = from_fx(graph_module, input_info)
+ mod: tvm_tirx.IRModule = from_fx(graph_module, input_info)
# Print out the imported model.
print(mod.script())
diff --git a/python/tvm/relax/global_info.py b/python/tvm/relax/global_info.py
index 868acaae9a..24af98dcc1 100644
--- a/python/tvm/relax/global_info.py
+++ b/python/tvm/relax/global_info.py
@@ -18,9 +18,9 @@
import tvm_ffi
-from tvm.ir import GlobalInfo
-from tvm.runtime import Device, MemoryScope, convert
-from tvm.target import Target
+from tvm_tirx.ir import GlobalInfo
+from tvm_tirx.runtime import Device, MemoryScope, convert
+from tvm_tirx.target import Target
from . import _ffi_api
diff --git a/python/tvm/relax/ir/instrument.py b/python/tvm/relax/ir/instrument.py
index 9f61bc6f08..f92d0525c0 100644
--- a/python/tvm/relax/ir/instrument.py
+++ b/python/tvm/relax/ir/instrument.py
@@ -16,11 +16,11 @@
# under the License.
"""Common relax pass instrumentation across IR variants."""
-import tvm
+import tvm, tvm_tirx
from tvm import relax
-@tvm.ir.transform.pass_instrument
+@tvm_tirx.ir.transform.pass_instrument
class WellFormedInstrument:
"""An instrument that checks the input/output IRModule of the Pass
is well formed. It will skip specific passes, like Normalize.
diff --git a/python/tvm/relax/op/__init__.py b/python/tvm/relax/op/__init__.py
index 72302aa00b..1fef6ef86a 100644
--- a/python/tvm/relax/op/__init__.py
+++ b/python/tvm/relax/op/__init__.py
@@ -254,7 +254,7 @@ from . import _op_gradient
def _register_op_make():
# pylint: disable=import-outside-toplevel
from .. import expr
- from tvm.ir import _tensor_expr_overload
+ from tvm_tirx.ir import _tensor_expr_overload
from . import _ffi_api
def _add(lhs, rhs):
@@ -274,10 +274,10 @@ def _register_op_make():
def _call(func, *args, attrs=None):
if not isinstance(
- func.ty, expr.tvm.ir.FuncType | expr.tvm.relax.FuncType | expr.tvm.ir.MissingType
+ func.ty, expr.tvm_tirx.ir.FuncType | expr.tvm.relax.FuncType | expr.tvm_tirx.ir.MissingType
):
return NotImplemented
- return expr.tvm.ir.Call(func, args, attrs=attrs)
+ return expr.tvm_tirx.ir.Call(func, args, attrs=attrs)
operators = {}
operators["astype"] = (
diff --git a/python/tvm/relax/op/_op_gradient.py b/python/tvm/relax/op/_op_gradient.py
index dd1baf9081..0ccdab518b 100644
--- a/python/tvm/relax/op/_op_gradient.py
+++ b/python/tvm/relax/op/_op_gradient.py
@@ -21,9 +21,9 @@ import functools
import operator
from tvm import relax
-from tvm.ir import Call, PrimType
+from tvm_tirx.ir import Call, PrimType
from tvm.relax.type import ShapeType
-from tvm.sym import Analyzer
+from tvm_tirx.sym import Analyzer
from ..block_builder import BlockBuilder
from ..expr import Expr, ShapeExpr, Var
diff --git a/python/tvm/relax/op/base.py b/python/tvm/relax/op/base.py
index 3297e5442b..40ac3e33da 100644
--- a/python/tvm/relax/op/base.py
+++ b/python/tvm/relax/op/base.py
@@ -21,13 +21,13 @@ from collections.abc import Callable
import tvm_ffi
-import tvm
-import tvm.runtime
-from tvm.ir import Attrs, Call, Op
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC, Location
-from tvm.ir.op import _make_op_api
-from tvm.runtime import DataTypeCode, Object
+import tvm, tvm_tirx
+import tvm_tirx.runtime, tvm
+from tvm_tirx.ir import Attrs, Call, Op
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.ir.op import _make_op_api
+from tvm_tirx.runtime import DataTypeCode, Object
from ..expr import Expr, ExternFunc, GlobalVar, Var
from ..type import Type
@@ -61,7 +61,7 @@ def register_gradient(
result : callable
The registered gradient, or a decorator if fgradient is not supplied.
"""
- return tvm.ir.register_op_attr(op_name, "FPrimalGradient", fgradient, override)
+ return tvm_tirx.ir.register_op_attr(op_name, "FPrimalGradient", fgradient, override)
def null_value(*, ty=None, loc: Location = UNKNOWN_LOC) -> Call:
@@ -86,7 +86,7 @@ def _wrap_inline_arg_tuple(args) -> Expr:
"""
if isinstance(args, tuple | list):
- return tvm.ir.Tuple([convert_to_expr(arg) for arg in args])
+ return tvm_tirx.ir.Tuple([convert_to_expr(arg) for arg in args])
elif (
isinstance(args, Expr)
and not isinstance(args, tvm.relax.Tuple)
@@ -102,19 +102,19 @@ def _wrap_tir_arg_tuple(func, args) -> Expr:
if not isinstance(args, tuple | list):
return _wrap_inline_arg_tuple(args)
signature = getattr(func, "ty", None)
- arg_types = signature.arg_types if isinstance(signature, tvm.ir.FuncType) else ()
+ arg_types = signature.arg_types if isinstance(signature, tvm_tirx.ir.FuncType) else ()
values = []
for i, value in enumerate(args):
ty = arg_types[i] if i < len(arg_types) else None
- if isinstance(ty, tvm.ir.PrimType) and ty.dtype.lanes == 1:
+ if isinstance(ty, tvm_tirx.ir.PrimType) and ty.dtype.lanes == 1:
if type(value) in (int, bool) and ty.matches_code(
DataTypeCode.INT, DataTypeCode.UINT, DataTypeCode.BOOL
):
- value = tvm.tirx.IntImm(ty, value)
+ value = tvm_tirx.tirx.IntImm(ty, value)
elif type(value) is float and ty.matches_code(DataTypeCode.FLOAT, DataTypeCode.BFLOAT):
- value = tvm.tirx.FloatImm(ty, value)
+ value = tvm_tirx.tirx.FloatImm(ty, value)
values.append(convert_to_expr(value))
- return tvm.ir.Tuple(values)
+ return tvm_tirx.ir.Tuple(values)
_call_tir = _make_op_api(Op.get("relax.call_tir"), __name__)
@@ -160,7 +160,7 @@ def call_tir_packed(gvar: GlobalVar, args: Expr, *, ty=None, loc: Location = UNK
----------
gvar : GlobalVar
The GlobalVar referring to a TIRx Function with its native
- ``tvm.ir.FuncType`` signature.
+ ``tvm_tirx.ir.FuncType`` signature.
args : Expr
The ordered arguments, supplied as an inline Relax tuple, a Python
@@ -206,7 +206,7 @@ def call_tir_packed(gvar: GlobalVar, args: Expr, *, ty=None, loc: Location = UNK
A caller that knows the scalar function has no effects may assert purity::
result = relax.call_pure_packed(
- tvm.ir.Op.get("relax.call_tir_packed"), add_scalar, (a, b)
+ tvm_tirx.ir.Op.get("relax.call_tir_packed"), add_scalar, (a, b)
)
A native copy function with a void return writes to a caller-owned tensor::
@@ -426,13 +426,13 @@ def invoke_closure(
) # type: ignore
-def render_object(val: tvm.Object) -> str:
+def render_object(val: tvm_tirx.Object) -> str:
"""
Given a TVM Object, renders it in string form. Used for Relax printing and assertions.
Parameters
----------
- val: tvm.Object
+ val: tvm_tirx.Object
An object to render
Returns
@@ -440,7 +440,7 @@ def render_object(val: tvm.Object) -> str:
ret: str
A string representing the value, ideally human-readable
"""
- if isinstance(val, tvm.runtime.Tensor):
+ if isinstance(val, tvm_tirx.runtime.Tensor):
return str(val)
if isinstance(val, tvm_ffi.Array):
fields = ", ".join([render_object(val[i]) for i in range(len(val))])
@@ -448,8 +448,8 @@ def render_object(val: tvm.Object) -> str:
return str(val)
-@tvm.register_global_func("relax.run.shape_to_tensor")
-def relax_shape_to_tensor(shape_tuple: tvm_ffi.Shape) -> tvm.runtime.Tensor:
+@tvm_tirx.register_global_func("relax.run.shape_to_tensor")
+def relax_shape_to_tensor(shape_tuple: tvm_ffi.Shape) -> tvm_tirx.runtime.Tensor:
"""
Takes a Shape and convert it to Tensor.
@@ -458,11 +458,11 @@ def relax_shape_to_tensor(shape_tuple: tvm_ffi.Shape) -> tvm.runtime.Tensor:
shape_tuple: tvm_ffi.Shape
Shape tuple that we want to convert to Tensor at runtime
"""
- return tvm.runtime.tensor([int(v) for v in shape_tuple])
+ return tvm_tirx.runtime.tensor([int(v) for v in shape_tuple])
-@tvm.register_global_func("relax.run.print")
-def relax_print(format_str: str, *format_args: tvm.Object) -> None:
+@tvm_tirx.register_global_func("relax.run.print")
+def relax_print(format_str: str, *format_args: tvm_tirx.Object) -> None:
"""
Takes a list of values to print, formats with the given format string.
If the format string is empty, simply prints.
@@ -492,8 +492,8 @@ def print(format: str | Expr, *values: Expr, ty=None, loc: Location = UNKNOWN_LO
return Call("relax.print", [format, *values], ty=ty, loc=loc)
-@tvm.register_global_func("relax.run.assert_op")
-def relax_assert_op(condition: tvm.Object, format_str: str, *format_args: tvm.Object) -> None:
+@tvm_tirx.register_global_func("relax.run.assert_op")
+def relax_assert_op(condition: tvm_tirx.Object, format_str: str, *format_args: tvm_tirx.Object) -> None:
"""
A variadic function. The first value serves as the assertion condition:
If the condition is true, then the operator does nothing.
@@ -507,13 +507,13 @@ def relax_assert_op(condition: tvm.Object, format_str: str, *format_args: tvm.Ob
Parameters
----------
- condition: tvm.Object
+ condition: tvm_tirx.Object
The assertion condition. Must be a boolean scalar.
format_str: str
A Python-style format string for the error message.
- format_args: List[tvm.Object]
+ format_args: List[tvm_tirx.Object]
Values used for formatting the string.
"""
if not isinstance(format_str, str):
@@ -523,7 +523,7 @@ def relax_assert_op(condition: tvm.Object, format_str: str, *format_args: tvm.Ob
if isinstance(condition, bool | int):
val = condition
- elif isinstance(condition, tvm.runtime.Tensor):
+ elif isinstance(condition, tvm_tirx.runtime.Tensor):
# may happen if the original program had unknown shape or dtype for the tensor's type
dtype = condition.dtype
if dtype != "bool":
diff --git a/python/tvm/relax/op/binary.py b/python/tvm/relax/op/binary.py
index 1348a21073..765f8bf15d 100644
--- a/python/tvm/relax/op/binary.py
+++ b/python/tvm/relax/op/binary.py
@@ -17,8 +17,8 @@
# pylint: disable=redefined-builtin, invalid-name
"""Relax binary arithmetic and comparison operators."""
-from tvm.ir import Call as _Call
-from tvm.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from ..expr import Expr
diff --git a/python/tvm/relax/op/builtin/builtin.py b/python/tvm/relax/op/builtin/builtin.py
index f0a2921643..6b1bd6989d 100644
--- a/python/tvm/relax/op/builtin/builtin.py
+++ b/python/tvm/relax/op/builtin/builtin.py
@@ -15,10 +15,10 @@
# specific language governing permissions and limitations
"""The builtin Relax operators."""
-from tvm.ir import Call, DataTypeImm, GenericConst, StringImm
-from tvm.ir import Call as _Call
-from tvm.ir.location import UNKNOWN_LOC, Location
-from tvm.runtime import MemoryScope
+from tvm_tirx.ir import Call, DataTypeImm, GenericConst, StringImm
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.runtime import MemoryScope
from ...expr import Expr, prim_value
from ...utils import convert_to_expr
diff --git a/python/tvm/relax/op/ccl/ccl.py b/python/tvm/relax/op/ccl/ccl.py
index 0496411661..9b4e22f98e 100644
--- a/python/tvm/relax/op/ccl/ccl.py
+++ b/python/tvm/relax/op/ccl/ccl.py
@@ -18,10 +18,10 @@
import tvm_ffi
-from tvm.ir import Attrs
-from tvm.ir import Call as _Call
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.ir import Attrs
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from ...expr import Expr
diff --git a/python/tvm/relax/op/create.py b/python/tvm/relax/op/create.py
index 2fcb937d27..df2c500f12 100644
--- a/python/tvm/relax/op/create.py
+++ b/python/tvm/relax/op/create.py
@@ -18,11 +18,11 @@
import tvm_ffi
-from tvm import DataType, DataTypeCode
-from tvm.ir import Attrs, PrimType, is_prim_expr
-from tvm.ir import Call as _Call
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx import DataType, DataTypeCode
+from tvm_tirx.ir import Attrs, PrimType, is_prim_expr
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from ..expr import Expr, ShapeExpr, prim_value
diff --git a/python/tvm/relax/op/datatype.py b/python/tvm/relax/op/datatype.py
index 712e6cc361..0384d10c1a 100644
--- a/python/tvm/relax/op/datatype.py
+++ b/python/tvm/relax/op/datatype.py
@@ -18,11 +18,11 @@
import tvm_ffi
-from tvm import DataType
-from tvm.ir import Attrs, PrimType
-from tvm.ir import Call as _Call
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx import DataType
+from tvm_tirx.ir import Attrs, PrimType
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from ..expr import Expr
diff --git a/python/tvm/relax/op/distributed/distributed.py b/python/tvm/relax/op/distributed/distributed.py
index 79d621d3ff..8cb65f985e 100644
--- a/python/tvm/relax/op/distributed/distributed.py
+++ b/python/tvm/relax/op/distributed/distributed.py
@@ -17,10 +17,10 @@
# pylint: disable=redefined-builtin
"""Operators for distributed Relax."""
-from tvm.ir import Call, Op
-from tvm.ir.attrs import make_node
-from tvm.ir.location import UNKNOWN_LOC, Location
-from tvm.ir.op import _make_op_api
+from tvm_tirx.ir import Call, Op
+from tvm_tirx.ir.attrs import make_node
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.ir.op import _make_op_api
from tvm.relax.distributed import DeviceMesh, Placement
from ...expr import Expr
diff --git a/python/tvm/relax/op/grad/grad.py b/python/tvm/relax/op/grad/grad.py
index da992c8e26..8b2bc5f1d3 100644
--- a/python/tvm/relax/op/grad/grad.py
+++ b/python/tvm/relax/op/grad/grad.py
@@ -22,10 +22,10 @@ gradients for forward operators. The ty inference for grad operators just return
ty of the input.
"""
-from tvm.ir import Call as _Call
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC as _UNKNOWN_LOC
-from tvm.ir.location import Location as _Location
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC as _UNKNOWN_LOC
+from tvm_tirx.ir.location import Location as _Location
from ...expr import Expr
diff --git a/python/tvm/relax/op/image/image.py b/python/tvm/relax/op/image/image.py
index d085871bd6..773ee5b82f 100644
--- a/python/tvm/relax/op/image/image.py
+++ b/python/tvm/relax/op/image/image.py
@@ -20,12 +20,12 @@ from typing import cast
import tvm_ffi
-from tvm import DataType
-from tvm.ir import Attrs, is_prim_expr
-from tvm.ir import Call as _Call
-from tvm.ir import PrimType as _PrimType
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx import DataType
+from tvm_tirx.ir import Attrs, is_prim_expr
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir import PrimType as _PrimType
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from ...expr import Expr, ShapeExpr
diff --git a/python/tvm/relax/op/index.py b/python/tvm/relax/op/index.py
index 0e09dac18e..51c0d6d370 100644
--- a/python/tvm/relax/op/index.py
+++ b/python/tvm/relax/op/index.py
@@ -18,11 +18,11 @@
import tvm_ffi
-from tvm.error import InternalError as _InternalError
-from tvm.ir import Attrs
-from tvm.ir import Call as _Call
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.error import InternalError as _InternalError
+from tvm_tirx.ir import Attrs
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from tvm.relax.expr import Tuple as _Tuple
from ..expr import Expr
diff --git a/python/tvm/relax/op/linear_algebra.py b/python/tvm/relax/op/linear_algebra.py
index f10ce64b18..5482b728cf 100644
--- a/python/tvm/relax/op/linear_algebra.py
+++ b/python/tvm/relax/op/linear_algebra.py
@@ -19,12 +19,12 @@
import tvm_ffi
-from tvm import DataType
-from tvm.ir import Attrs
-from tvm.ir import Call as _Call
-from tvm.ir import PrimType as _PrimType
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx import DataType
+from tvm_tirx.ir import Attrs
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir import PrimType as _PrimType
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from ..expr import Expr
from ..expr import Tuple as RxTuple
diff --git a/python/tvm/relax/op/manipulate.py b/python/tvm/relax/op/manipulate.py
index b08fa2c30d..4cc764145c 100644
--- a/python/tvm/relax/op/manipulate.py
+++ b/python/tvm/relax/op/manipulate.py
@@ -20,13 +20,13 @@ from collections.abc import Callable
import tvm_ffi
-from tvm.error import InternalError as _InternalError
-from tvm.ir import Attrs, is_prim_expr
-from tvm.ir import Call as _Call
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC, Location
-from tvm.runtime import DataTypeCode
-from tvm.tirx import FloatImm, IndexMap, IntImm
+from tvm_tirx.error import InternalError as _InternalError
+from tvm_tirx.ir import Attrs, is_prim_expr
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.runtime import DataTypeCode
+from tvm_tirx.tirx import FloatImm, IndexMap, IntImm
from ..expr import Expr, ShapeExpr, prim_value
from ..expr import Tuple as RxTuple
diff --git a/python/tvm/relax/op/memory/memory.py b/python/tvm/relax/op/memory/memory.py
index 089f761b2c..b8a0d375e7 100644
--- a/python/tvm/relax/op/memory/memory.py
+++ b/python/tvm/relax/op/memory/memory.py
@@ -15,10 +15,10 @@
# specific language governing permissions and limitations
"""Relax memory primitives."""
-from tvm.ir import Call, DataTypeImm, GenericConst, StringImm
-from tvm.ir import Call as _Call
-from tvm.ir.location import UNKNOWN_LOC, Location
-from tvm.runtime import MemoryScope
+from tvm_tirx.ir import Call, DataTypeImm, GenericConst, StringImm
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.runtime import MemoryScope
from ...expr import Expr, prim_value
from ...utils import convert_to_expr
diff --git a/python/tvm/relax/op/memory/view.py b/python/tvm/relax/op/memory/view.py
index e0adadda2b..0ffc4fe4d2 100644
--- a/python/tvm/relax/op/memory/view.py
+++ b/python/tvm/relax/op/memory/view.py
@@ -28,9 +28,9 @@ while keeping the same underlying data.
from collections.abc import Sequence
-from tvm.ir import Call as _Call
-from tvm.ir import DataTypeImm, GenericConst
-from tvm.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir import DataTypeImm, GenericConst
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from tvm.relax import Expr, ShapeExpr, Tuple
from tvm.relax.expr import prim_value
diff --git a/python/tvm/relax/op/nn/nn.py b/python/tvm/relax/op/nn/nn.py
index e99f11135e..e1085a155e 100644
--- a/python/tvm/relax/op/nn/nn.py
+++ b/python/tvm/relax/op/nn/nn.py
@@ -18,14 +18,15 @@
import tvm_ffi
-from tvm import DataType, relax
-from tvm.error import InternalError as _InternalError
-from tvm.ir import Attrs
-from tvm.ir import Call as _Call
-from tvm.ir import PrimType as _PrimType
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC, Location
-from tvm.tirx import FloatImm
+from tvm_tirx import DataType
+from tvm import relax
+from tvm_tirx.error import InternalError as _InternalError
+from tvm_tirx.ir import Attrs
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir import PrimType as _PrimType
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.tirx import FloatImm
from ...expr import Expr
diff --git a/python/tvm/relax/op/qdq.py b/python/tvm/relax/op/qdq.py
index c3a1c9c60d..0a3173be37 100644
--- a/python/tvm/relax/op/qdq.py
+++ b/python/tvm/relax/op/qdq.py
@@ -18,11 +18,11 @@
import tvm_ffi
-from tvm.ir import Attrs
-from tvm.ir import Call as _Call
-from tvm.ir import PrimType as _PrimType
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.ir import Attrs
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir import PrimType as _PrimType
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from ..expr import Expr
diff --git a/python/tvm/relax/op/sampling.py b/python/tvm/relax/op/sampling.py
index 0ded7aea43..65175f3574 100644
--- a/python/tvm/relax/op/sampling.py
+++ b/python/tvm/relax/op/sampling.py
@@ -18,10 +18,10 @@
import tvm_ffi
-from tvm.ir import Attrs
-from tvm.ir import Call as _Call
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.ir import Attrs
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from ..expr import Expr
diff --git a/python/tvm/relax/op/search.py b/python/tvm/relax/op/search.py
index c19b1ac20c..08bbaa340a 100644
--- a/python/tvm/relax/op/search.py
+++ b/python/tvm/relax/op/search.py
@@ -19,10 +19,10 @@
import tvm_ffi
-from tvm.ir import Attrs
-from tvm.ir import Call as _Call
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.ir import Attrs
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from ..expr import Expr
diff --git a/python/tvm/relax/op/set.py b/python/tvm/relax/op/set.py
index 32eb00a447..72438badee 100644
--- a/python/tvm/relax/op/set.py
+++ b/python/tvm/relax/op/set.py
@@ -19,9 +19,9 @@
import numpy as np # type: ignore
-import tvm
-from tvm.ir import Call as _Call
-from tvm.ir.location import UNKNOWN_LOC, Location
+import tvm, tvm_tirx
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from ..expr import Expr, prim_value
@@ -91,15 +91,15 @@ def unique(
)
-@tvm.register_global_func("relax.run.unique")
+@tvm_tirx.register_global_func("relax.run.unique")
def numpy_unique(
- x: tvm.runtime.tensor,
+ x: tvm_tirx.runtime.tensor,
sorted: int,
return_index: int,
return_inverse: int,
return_counts: int,
axis: int | None = None,
-) -> tvm.runtime.tensor:
+) -> tvm_tirx.runtime.tensor:
"""Returns the unique elements of the input tensor.
Uses numpy.unique to compute unique elements.
@@ -123,7 +123,7 @@ def numpy_unique(
if not sorted:
indices = np.unique(x_numpy, return_index=True, axis=axis)[1]
unique_values = np.take(x_numpy, builtins.sorted(indices), axis=axis)
- return tvm.runtime.tensor(unique_values)
+ return tvm_tirx.runtime.tensor(unique_values)
# Otherwise, numpy returns a tuple
unique_values = result[0]
@@ -139,23 +139,23 @@ def numpy_unique(
sort_order = np.argsort(indices)
unique_values = np.take(unique_values, sort_order, axis=axis)
indices = np.sort(indices)
- output_list.append(tvm.runtime.tensor(unique_values))
- output_list.append(tvm.runtime.tensor(indices))
+ output_list.append(tvm_tirx.runtime.tensor(unique_values))
+ output_list.append(tvm_tirx.runtime.tensor(indices))
elif not sorted:
# Need to get indices to reorder
_, indices = np.unique(x_numpy, return_index=True, axis=axis)
sort_order = np.argsort(indices)
unique_values = np.take(unique_values, sort_order, axis=axis)
- output_list.append(tvm.runtime.tensor(unique_values))
+ output_list.append(tvm_tirx.runtime.tensor(unique_values))
if bool(return_index):
indices_from_result = result[result_idx]
result_idx += 1
- output_list.append(tvm.runtime.tensor(np.sort(indices_from_result)))
+ output_list.append(tvm_tirx.runtime.tensor(np.sort(indices_from_result)))
else:
# Sorted case
- output_list.append(tvm.runtime.tensor(unique_values))
+ output_list.append(tvm_tirx.runtime.tensor(unique_values))
if bool(return_index):
- output_list.append(tvm.runtime.tensor(result[result_idx]))
+ output_list.append(tvm_tirx.runtime.tensor(result[result_idx]))
result_idx += 1
if bool(return_inverse):
@@ -171,7 +171,7 @@ def numpy_unique(
# When axis is None, it has length X.size (flattened)
# When axis is specified, it has length X.shape[axis]
# numpy.unique already returns 1D inverse_indices, so no reshaping needed
- output_list.append(tvm.runtime.tensor(inverse_indices))
+ output_list.append(tvm_tirx.runtime.tensor(inverse_indices))
result_idx += 1
if bool(return_counts):
@@ -181,7 +181,7 @@ def numpy_unique(
_, orig_indices = np.unique(x_numpy, return_index=True, axis=axis)
sort_order = np.argsort(orig_indices)
counts = counts[sort_order]
- output_list.append(tvm.runtime.tensor(counts))
+ output_list.append(tvm_tirx.runtime.tensor(counts))
return tuple(output_list)
@@ -217,7 +217,7 @@ def nonzero(x: Expr, *, ty=None, loc: Location = UNKNOWN_LOC) -> Expr:
return _Call("relax.nonzero", [x], ty=ty, loc=loc) # type: ignore
-@tvm.register_global_func("relax.run.nonzero")
-def numpy_nonzero(x: tvm.runtime.tensor) -> tvm.runtime.tensor:
+@tvm_tirx.register_global_func("relax.run.nonzero")
+def numpy_nonzero(x: tvm_tirx.runtime.tensor) -> tvm_tirx.runtime.tensor:
np_result = np.atleast_1d(x.numpy()).nonzero()
- return tvm.runtime.tensor(np.stack(np_result, axis=0))
+ return tvm_tirx.runtime.tensor(np.stack(np_result, axis=0))
diff --git a/python/tvm/relax/op/sorting.py b/python/tvm/relax/op/sorting.py
index 8b3a0ad0aa..fa2ed7ac5d 100644
--- a/python/tvm/relax/op/sorting.py
+++ b/python/tvm/relax/op/sorting.py
@@ -18,10 +18,10 @@
import tvm_ffi
-from tvm.ir import Attrs, GenericConst
-from tvm.ir import Call as _Call
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.ir import Attrs, GenericConst
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from ..expr import Expr
diff --git a/python/tvm/relax/op/statistical.py b/python/tvm/relax/op/statistical.py
index bf1b6a16d8..3751bfa9d9 100644
--- a/python/tvm/relax/op/statistical.py
+++ b/python/tvm/relax/op/statistical.py
@@ -19,11 +19,11 @@
import tvm_ffi
-from tvm import DataType
-from tvm.ir import Attrs, PrimType
-from tvm.ir import Call as _Call
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx import DataType
+from tvm_tirx.ir import Attrs, PrimType
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from ..expr import Expr
diff --git a/python/tvm/relax/op/ternary.py b/python/tvm/relax/op/ternary.py
index d348269cdc..511448d3de 100644
--- a/python/tvm/relax/op/ternary.py
+++ b/python/tvm/relax/op/ternary.py
@@ -17,8 +17,8 @@
# pylint: disable=redefined-builtin, invalid-name
"""Relax ternary arithmetic operators."""
-from tvm.ir import Call as _Call
-from tvm.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from ..expr import Expr
diff --git a/python/tvm/relax/op/unary.py b/python/tvm/relax/op/unary.py
index 95a882c2ac..35a403d707 100644
--- a/python/tvm/relax/op/unary.py
+++ b/python/tvm/relax/op/unary.py
@@ -17,8 +17,8 @@
# pylint: disable=redefined-builtin, invalid-name
"""Relax unary arithmetic operators."""
-from tvm.ir import Call as _Call
-from tvm.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from ..expr import Expr
from ..utils import convert_to_expr
diff --git a/python/tvm/relax/op/vision/multibox_transform_loc.py b/python/tvm/relax/op/vision/multibox_transform_loc.py
index b43bd474df..bdc36bea9b 100644
--- a/python/tvm/relax/op/vision/multibox_transform_loc.py
+++ b/python/tvm/relax/op/vision/multibox_transform_loc.py
@@ -18,11 +18,11 @@
import tvm_ffi
-from tvm.ir import Attrs
-from tvm.ir import Call as _Call
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC as _UNKNOWN_LOC
-from tvm.ir.location import Location as _Location
+from tvm_tirx.ir import Attrs
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC as _UNKNOWN_LOC
+from tvm_tirx.ir.location import Location as _Location
@tvm_ffi.register_object("relax.attrs.MultiboxTransformLocAttrs")
diff --git a/python/tvm/relax/op/vision/nms.py b/python/tvm/relax/op/vision/nms.py
index 876b40830c..b570c466f9 100644
--- a/python/tvm/relax/op/vision/nms.py
+++ b/python/tvm/relax/op/vision/nms.py
@@ -18,11 +18,11 @@
import tvm_ffi
-from tvm.ir import Attrs
-from tvm.ir import Call as _Call
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC as _UNKNOWN_LOC
-from tvm.ir.location import Location as _Location
+from tvm_tirx.ir import Attrs
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC as _UNKNOWN_LOC
+from tvm_tirx.ir.location import Location as _Location
@tvm_ffi.register_object("relax.attrs.AllClassNonMaximumSuppressionAttrs")
diff --git a/python/tvm/relax/op/vision/roi_align.py b/python/tvm/relax/op/vision/roi_align.py
index ea11d1a5c3..31f21eff5f 100644
--- a/python/tvm/relax/op/vision/roi_align.py
+++ b/python/tvm/relax/op/vision/roi_align.py
@@ -18,11 +18,11 @@
import tvm_ffi
-from tvm.ir import Attrs
-from tvm.ir import Call as _Call
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC as _UNKNOWN_LOC
-from tvm.ir.location import Location as _Location
+from tvm_tirx.ir import Attrs
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC as _UNKNOWN_LOC
+from tvm_tirx.ir.location import Location as _Location
from ..base import Expr
diff --git a/python/tvm/relax/op/vision/roi_pool.py b/python/tvm/relax/op/vision/roi_pool.py
index 0946d70785..20dffa253e 100644
--- a/python/tvm/relax/op/vision/roi_pool.py
+++ b/python/tvm/relax/op/vision/roi_pool.py
@@ -18,11 +18,11 @@
import tvm_ffi
-from tvm.ir import Attrs
-from tvm.ir import Call as _Call
-from tvm.ir.attrs import make_node as _make_attrs
-from tvm.ir.location import UNKNOWN_LOC as _UNKNOWN_LOC
-from tvm.ir.location import Location as _Location
+from tvm_tirx.ir import Attrs
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.attrs import make_node as _make_attrs
+from tvm_tirx.ir.location import UNKNOWN_LOC as _UNKNOWN_LOC
+from tvm_tirx.ir.location import Location as _Location
from ..base import Expr
diff --git a/python/tvm/relax/op/vm/vm.py b/python/tvm/relax/op/vm/vm.py
index 3a9e91b6f6..d398906cf1 100644
--- a/python/tvm/relax/op/vm/vm.py
+++ b/python/tvm/relax/op/vm/vm.py
@@ -15,10 +15,10 @@
# specific language governing permissions and limitations
"""Relax vm primitives."""
-from tvm.ir import Call, DataTypeImm, GenericConst, StringImm
-from tvm.ir import Call as _Call
-from tvm.ir.location import UNKNOWN_LOC, Location
-from tvm.runtime import MemoryScope
+from tvm_tirx.ir import Call, DataTypeImm, GenericConst, StringImm
+from tvm_tirx.ir import Call as _Call
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.runtime import MemoryScope
from ...expr import Expr, Tuple, prim_value
from ...utils import convert_to_expr
diff --git a/python/tvm/relax/pipeline.py b/python/tvm/relax/pipeline.py
index 87226a3ce5..81bbfb1284 100644
--- a/python/tvm/relax/pipeline.py
+++ b/python/tvm/relax/pipeline.py
@@ -23,7 +23,7 @@ as it is or serves as a basis to do further composition.
# pylint: disable=unused-argument
-import tvm
+import tvm, tvm_tirx
from tvm.s_tir import meta_schedule as ms
from . import backend, transform
@@ -43,24 +43,24 @@ def zero_pipeline(*, enable_warning: bool = False):
the database. By default we don't print warning.
"""
- @tvm.ir.transform.module_pass(opt_level=0)
- def f_zero_pipeline(mod: tvm.ir.IRModule, ctx: tvm.ir.transform.PassContext) -> tvm.ir.IRModule:
+ @tvm_tirx.ir.transform.module_pass(opt_level=0)
+ def f_zero_pipeline(mod: tvm_tirx.ir.IRModule, ctx: tvm_tirx.ir.transform.PassContext) -> tvm_tirx.ir.IRModule:
"""Pipeline that applies pre-tuned logs.
Parameters
----------
- mod : tvm.ir.IRModule
+ mod : tvm_tirx.ir.IRModule
Input IRModule.
- ctx : tvm.ir.transform.PassContext
+ ctx : tvm_tirx.ir.transform.PassContext
The pass context
Returns
-------
- mod: tvm.ir.IRModule
+ mod: tvm_tirx.ir.IRModule
The result transformed module.
"""
- seq = tvm.ir.transform.Sequential(
+ seq = tvm_tirx.ir.transform.Sequential(
[
transform.LegalizeOps(enable_warning=enable_warning),
transform.AnnotateTIROpPattern(),
@@ -80,9 +80,9 @@ def zero_pipeline(*, enable_warning: bool = False):
def default_build_pipeline():
"""The default compilation pipeline used in tvm.compile"""
- @tvm.ir.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext) -> tvm.ir.IRModule:
- seq = tvm.ir.transform.Sequential(
+ @tvm_tirx.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm_tirx.ir.IRModule, _ctx: tvm_tirx.ir.transform.PassContext) -> tvm_tirx.ir.IRModule:
+ seq = tvm_tirx.ir.transform.Sequential(
[
backend.DispatchSampling(),
backend.DispatchSortScan(),
@@ -109,7 +109,7 @@ def default_build_pipeline():
def static_shape_tuning_pipeline(
total_trials: int,
- target: str | tvm.target.Target,
+ target: str | tvm_tirx.target.Target,
work_dir: str = "tuning_logs",
cpu_weight_prepack: bool = False,
max_trials_per_task: int | None = None,
@@ -121,7 +121,7 @@ def static_shape_tuning_pipeline(
total_trials : int
Total number of trials to run.
- target : Union[str, tvm.target.Target]
+ target : Union[str, tvm_tirx.target.Target]
The target device to tune the model.
work_dir : str
@@ -158,18 +158,18 @@ def static_shape_tuning_pipeline(
)(mod)
ex = tvm.compile(mod, target=target)
- vm = relax.VirtualMachine(ex, device=tvm.cpu())
+ vm = relax.VirtualMachine(ex, device=tvm_tirx.cpu())
# Transform the params using the vm function
# the name should be f"{func_name}_transform_params"
params = vm["main_transform_params"](params["main"])
- input_data = tvm.runtime.tensor(np.random.randn(1, 3, 224, 224).astype("float32"))
+ input_data = tvm_tirx.runtime.tensor(np.random.randn(1, 3, 224, 224).astype("float32"))
out = vm["main"](input_data, *params).numpy()
"""
- @tvm.ir.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext) -> tvm.ir.IRModule:
+ @tvm_tirx.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm_tirx.ir.IRModule, _ctx: tvm_tirx.ir.transform.PassContext) -> tvm_tirx.ir.IRModule:
if cpu_weight_prepack:
pre_tuning_layout_rewrite = [transform.AttachAttrLayoutFreeTensors()]
post_tuning_layout_rewrite = [
@@ -181,8 +181,8 @@ def static_shape_tuning_pipeline(
pre_tuning_layout_rewrite = []
post_tuning_layout_rewrite = []
- with tvm.target.Target(target):
- mod = tvm.ir.transform.Sequential(
+ with tvm_tirx.target.Target(target):
+ mod = tvm_tirx.ir.transform.Sequential(
[
transform.DecomposeOpsForInference(),
transform.CanonicalizeBindings(),
@@ -197,7 +197,7 @@ def static_shape_tuning_pipeline(
max_trials_per_task=max_trials_per_task,
)
if total_trials > 0
- else tvm.ir.transform.Sequential([])
+ else tvm_tirx.ir.transform.Sequential([])
),
transform.MetaScheduleApplyDatabase(work_dir),
*post_tuning_layout_rewrite,
@@ -218,7 +218,7 @@ PIPELINE_MAP = {
}
-def get_pipeline(name: str = "zero", **kwargs) -> tvm.ir.transform.Pass:
+def get_pipeline(name: str = "zero", **kwargs) -> tvm_tirx.ir.transform.Pass:
"""Get pre-build pipeline by name
Parameters
@@ -230,7 +230,7 @@ def get_pipeline(name: str = "zero", **kwargs) -> tvm.ir.transform.Pass:
Returns
-------
- pipeline: tvm.ir.transform.Pass
+ pipeline: tvm_tirx.ir.transform.Pass
The transformation pipeline.
"""
@@ -259,7 +259,7 @@ def register_pipeline(name: str):
return _register
-def library_dispatch_passes(target: tvm.target.Target):
+def library_dispatch_passes(target: tvm_tirx.target.Target):
"""Get the default library dispatch passes for the given target."""
if target.kind.name == "cuda":
return backend.cuda.library_dispatch_passes(target)
@@ -276,7 +276,7 @@ def library_dispatch_passes(target: tvm.target.Target):
raise ValueError(f"Target {target} is not yet supported by library dispatch passes.")
-def legalize_passes(target: tvm.target.Target):
+def legalize_passes(target: tvm_tirx.target.Target):
"""Get the default legalization passes for the given target."""
if target.kind.name == "cuda":
return backend.cuda.legalize_passes(target)
@@ -293,7 +293,7 @@ def legalize_passes(target: tvm.target.Target):
raise ValueError(f"Target {target} is not yet supported by legalize passes.")
-def dataflow_lower_passes(target: tvm.target.Target):
+def dataflow_lower_passes(target: tvm_tirx.target.Target):
"""Get the default legalization passes for the given target."""
if target.kind.name == "cuda":
return backend.cuda.dataflow_lower_passes(target)
@@ -310,7 +310,7 @@ def dataflow_lower_passes(target: tvm.target.Target):
raise ValueError(f"Target {target} is not yet supported by dataflow lowering passes.")
-def finalize_passes(target: tvm.target.Target):
+def finalize_passes(target: tvm_tirx.target.Target):
"""Get the default legalization passes for the given target."""
if target.kind.name == "cuda":
return backend.cuda.finalize_passes(target)
@@ -327,7 +327,7 @@ def finalize_passes(target: tvm.target.Target):
raise ValueError(f"Target {target} is not yet supported by finalization passes.")
-def get_default_pipeline(target: tvm.target.Target):
+def get_default_pipeline(target: tvm_tirx.target.Target):
"""Get the default Relax compilation pipeline for the given target."""
if target.kind.name == "cuda":
return backend.cuda.get_default_pipeline(target)
diff --git a/python/tvm/relax/relax_to_pyfunc_converter.py b/python/tvm/relax/relax_to_pyfunc_converter.py
index 925caf79cb..a31c8cd018 100644
--- a/python/tvm/relax/relax_to_pyfunc_converter.py
+++ b/python/tvm/relax/relax_to_pyfunc_converter.py
@@ -28,9 +28,10 @@ import numpy # pylint: disable=unused-import
import torch
import torch.nn.functional as F
-import tvm
-from tvm import relax, runtime
-from tvm.ir import IRModule, Op
+import tvm, tvm_tirx
+from tvm_tirx import runtime
+from tvm import relax
+from tvm_tirx.ir import IRModule, Op
class RelaxToPyFuncConverter:
@@ -395,7 +396,7 @@ class RelaxExpressionConverter:
return self._convert_var(expr, args)
elif isinstance(expr, relax.Call):
return self._convert_call(expr, args)
- elif isinstance(expr, tvm.ir.GenericConst):
+ elif isinstance(expr, tvm_tirx.ir.GenericConst):
return self._convert_constant(expr)
elif isinstance(expr, relax.SeqExpr):
return self._convert_seq_expr(expr, args)
@@ -621,8 +622,8 @@ class RelaxExpressionConverter:
if global_var.name_hint == func_name and hasattr(func, "body"):
try:
# Compile the TIR function
- target = tvm.target.Target("llvm")
- with tvm.target.Target(target):
+ target = tvm_tirx.target.Target("llvm")
+ with tvm_tirx.target.Target(target):
tir_function = tvm.compile(func, target=target)
break
except (RuntimeError, ValueError, TypeError) as compile_e:
@@ -633,7 +634,7 @@ class RelaxExpressionConverter:
# If not found in current module, try global registry
if tir_function is None:
- tir_function = tvm.get_global_func(func_name)
+ tir_function = tvm_tirx.get_global_func(func_name)
if tir_function is None:
if len(converted_args) >= 2:
@@ -743,7 +744,7 @@ class RelaxExpressionConverter:
try:
# Get the packed function from TVM
- packed_function = tvm.get_global_func(func_name)
+ packed_function = tvm_tirx.get_global_func(func_name)
if packed_function is None:
return converted_args[0] if converted_args else torch.tensor([])
@@ -778,7 +779,7 @@ class RelaxExpressionConverter:
# Fallback: return the first argument
return converted_args[0] if converted_args else torch.tensor([])
- def _convert_constant(self, const: tvm.ir.GenericConst) -> Any:
+ def _convert_constant(self, const: tvm_tirx.ir.GenericConst) -> Any:
"""Convert a Relax constant to Python equivalent."""
if hasattr(const, "value"):
data = const.value
@@ -1043,7 +1044,7 @@ class RelaxExpressionConverter:
)
else:
shape = (int(shape_arg),)
- elif isinstance(shape_arg, tvm.ir.GenericConst):
+ elif isinstance(shape_arg, tvm_tirx.ir.GenericConst):
# Constant tensor case
shape_data = shape_arg.value.numpy()
shape = tuple(int(v) for v in shape_data)
diff --git a/python/tvm/relax/script/__init__.py b/python/tvm/relax/script/__init__.py
index b882ea22e7..23f88ea29e 100644
--- a/python/tvm/relax/script/__init__.py
+++ b/python/tvm/relax/script/__init__.py
@@ -27,10 +27,10 @@ from typing import TYPE_CHECKING
from typing import Any as _Any
if TYPE_CHECKING:
- from tvm.ir import IRModule
+ from tvm_tirx.ir import IRModule
from tvm.relax.base_py_module import BasePyModule
- from tvm.runtime import Device
- from tvm.target import Target
+ from tvm_tirx.runtime import Device
+ from tvm_tirx.target import Target
_initialized = False
@@ -43,7 +43,7 @@ def _initialize() -> None:
return
_initializing = True
try:
- from tvm.script.parser import entry, register_namespace
+ from tvm_tirx.script.parser import entry, register_namespace
from . import ir_builder as builder
@@ -83,7 +83,7 @@ class _PyModuleFactory:
self.__name__ = original_class.__name__
def __call__(self, device: Device | None = None, target: Target | None = None) -> BasePyModule:
- from tvm import cpu, ir
+ from tvm_tirx import cpu, ir
from tvm.relax.base_py_module import BasePyModule
source = self.ir_module
@@ -157,10 +157,11 @@ def py_module(
"""
def apply(source_class: type) -> IRModule | _PyModuleFactory:
- from tvm import ir, relax
+ from tvm_tirx import ir
+ from tvm import relax
from tvm.relax.base_py_module import BasePyModule
- from tvm.script.parser.entry import parse
- from tvm.script.parser.inspect_source import acquire_source, capture_definition_scope
+ from tvm_tirx.script.parser.entry import parse
+ from tvm_tirx.script.parser.inspect_source import acquire_source, capture_definition_scope
if not _inspect.isclass(source_class):
raise TypeError(f"Expect a class, but got: {source_class}")
diff --git a/python/tvm/relax/script/ir_builder/__init__.py b/python/tvm/relax/script/ir_builder/__init__.py
index 48e712073f..c2460bed58 100644
--- a/python/tvm/relax/script/ir_builder/__init__.py
+++ b/python/tvm/relax/script/ir_builder/__init__.py
@@ -17,8 +17,8 @@
# pylint: disable=redefined-builtin, wrong-import-order, no-member, invalid-name
"""Public exports for the Relax builder."""
-from tvm.ir.type import _init_type_api
-from tvm.script.ir_builder.ir import constexpr
+from tvm_tirx.ir.type import _init_type_api
+from tvm_tirx.script.ir_builder.ir import constexpr
from . import distributed as dist
from . import ir as _native
diff --git a/python/tvm/relax/script/ir_builder/distributed/ir.py b/python/tvm/relax/script/ir_builder/distributed/ir.py
index 35e93903d9..72ad38a204 100644
--- a/python/tvm/relax/script/ir_builder/distributed/ir.py
+++ b/python/tvm/relax/script/ir_builder/distributed/ir.py
@@ -22,21 +22,21 @@ from numbers import Number
import numpy as _np # type: ignore
-import tvm
-from tvm.ir import GenericConst
+import tvm, tvm_tirx
+from tvm_tirx.ir import GenericConst
from tvm.relax.distributed import DeviceMesh, DTensorType
-from tvm.runtime import _tensor
+from tvm_tirx.runtime import _tensor
def const(
- value: bool | int | float | _np.ndarray | tvm.runtime.Tensor,
+ value: bool | int | float | _np.ndarray | tvm_tirx.runtime.Tensor,
ty: DTensorType,
) -> GenericConst:
"""Create a distributed constant value with the specified tensor type.
Parameters
----------
- value : bool, int, float, numpy.ndarray or tvm.runtime.Tensor
+ value : bool, int, float, numpy.ndarray or tvm_tirx.runtime.Tensor
The constant value.
ty : DTensorType
@@ -53,7 +53,7 @@ def const(
Python scalars and NumPy values are converted to the dtype specified by
``ty.tensor_ty``. The dtype is not inferred from the Python value.
"""
- ty = tvm.runtime.convert(ty)
+ ty = tvm_tirx.runtime.convert(ty)
if not isinstance(ty, DTensorType):
raise TypeError("ty needs to be an instance of DTensorType. ")
dtype = str(ty.tensor_ty.dtype)
diff --git a/python/tvm/relax/script/ir_builder/distributed/op.py b/python/tvm/relax/script/ir_builder/distributed/op.py
index ec0fac1a98..7c73e3a6dc 100644
--- a/python/tvm/relax/script/ir_builder/distributed/op.py
+++ b/python/tvm/relax/script/ir_builder/distributed/op.py
@@ -17,7 +17,7 @@
# pylint: disable=redefined-builtin, wrong-import-order, no-member, invalid-name
"""Distributed Relax expression operators."""
-from tvm.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from tvm.relax.distributed import DeviceMesh, Placement
from tvm.relax.expr import Expr
from tvm.relax.op import call_tir as call_tir
diff --git a/python/tvm/relax/script/ir_builder/frame.py b/python/tvm/relax/script/ir_builder/frame.py
index 3ffc0bc2aa..355bc64669 100644
--- a/python/tvm/relax/script/ir_builder/frame.py
+++ b/python/tvm/relax/script/ir_builder/frame.py
@@ -18,7 +18,7 @@
from tvm_ffi import register_object as _register_object
-from tvm.script.ir_builder.base import IRBuilderFrame
+from tvm_tirx.script.ir_builder.base import IRBuilderFrame
@_register_object("script.ir_builder.relax.RelaxFrame")
diff --git a/python/tvm/relax/script/ir_builder/ir.py b/python/tvm/relax/script/ir_builder/ir.py
index bb2f5662ff..2357d635e1 100644
--- a/python/tvm/relax/script/ir_builder/ir.py
+++ b/python/tvm/relax/script/ir_builder/ir.py
@@ -24,12 +24,12 @@ import inspect
import numbers as _numbers
import re as _re
-from tvm import ir as _ir
+from tvm_tirx import ir as _ir
from tvm import relax
from tvm import relax as _relax
-from tvm.ir import IRModule
-from tvm.ir.location import UNKNOWN_LOC, Location
-from tvm.ir.type import _init_type_api, _normalize_type
+from tvm_tirx.ir import IRModule
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.ir.type import _init_type_api, _normalize_type
from tvm.relax import Expr, ExternFunc, ShapeExpr, TupleGetItem, const
from tvm.relax.distributed import DeviceMesh as _DeviceMesh
from tvm.relax.distributed import DTensorType as _DTensorType
@@ -37,7 +37,7 @@ from tvm.relax.distributed import Placement as _Placement
from tvm.relax.distributed import device_mesh
from tvm.relax.dpl import PatternMatchingRewriter
from tvm.relax.global_info import DummyGlobalInfo, VDevice
-from tvm.runtime._tensor import (
+from tvm_tirx.runtime._tensor import (
cpu,
cuda,
device,
@@ -50,11 +50,11 @@ from tvm.runtime._tensor import (
vulkan,
webgpu,
)
-from tvm.script.ir_builder.base import LocationEntry as _LocationEntry
-from tvm.script.ir_builder.base import resolve_global_info_args as _resolve_global_info_args
-from tvm.script.ir_builder.ir import _global_infos
-from tvm.script.ir_builder.ir import dtype as dtype
-from tvm.script.ir_builder.parser_protocol import resolve_global_info_ as _resolve_global_info
+from tvm_tirx.script.ir_builder.base import LocationEntry as _LocationEntry
+from tvm_tirx.script.ir_builder.base import resolve_global_info_args as _resolve_global_info_args
+from tvm_tirx.script.ir_builder.ir import _global_infos
+from tvm_tirx.script.ir_builder.ir import dtype as dtype
+from tvm_tirx.script.ir_builder.parser_protocol import resolve_global_info_ as _resolve_global_info
from . import _ffi_api
@@ -207,8 +207,8 @@ def rewriter(rewriter_mod: IRModule | type) -> PatternMatchingRewriter:
"""
if not isinstance(rewriter_mod, IRModule):
- from tvm.script.parser.entry import parse
- from tvm.script.parser.inspect_source import capture_definition_scope
+ from tvm_tirx.script.parser.entry import parse
+ from tvm_tirx.script.parser.inspect_source import capture_definition_scope
if not inspect.isclass(rewriter_mod):
raise TypeError(f"Expect a class, but got: {rewriter_mod}")
diff --git a/python/tvm/relax/script/ir_builder/op.py b/python/tvm/relax/script/ir_builder/op.py
index 6c3b94f326..8af2b60e99 100644
--- a/python/tvm/relax/script/ir_builder/op.py
+++ b/python/tvm/relax/script/ir_builder/op.py
@@ -24,13 +24,13 @@ import numbers as _numbers
from types import SimpleNamespace as _Namespace
from typing import Any
-import tvm
-from tvm import ir as _ir
+import tvm, tvm_tirx
+from tvm_tirx import ir as _ir
from tvm import relax as _relax
-from tvm import tirx as _tir
-from tvm.ir.location import UNKNOWN_LOC, Location
-from tvm.ir.op import _init_op_api
-from tvm.ir.prim import _ffi_api as _prim_ffi
+from tvm_tirx import tirx as _tir
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.ir.op import _init_op_api
+from tvm_tirx.ir.prim import _ffi_api as _prim_ffi
from tvm.relax import Call, Expr, ExternFunc
from tvm.relax.global_info import VDevice
from tvm.relax.op import (
@@ -234,8 +234,8 @@ from tvm.relax.op import to_vdevice as _to_vdevice
from tvm.relax.op.builtin import stop_lift_params
from tvm.relax.type import Type
from tvm.relax.utils import convert_to_expr
-from tvm.script.ir_builder import base as _base
-from tvm.script.ir_builder.base import auto_emit_void_call as _auto_emit_void_call
+from tvm_tirx.script.ir_builder import base as _base
+from tvm_tirx.script.ir_builder.base import auto_emit_void_call as _auto_emit_void_call
from . import distributed as dist # noqa: F401
from .ir import _value, lookup_vdevice
@@ -321,7 +321,7 @@ def call_packed(
is_default = True
attrs = None
if kwargs or not is_default:
- attrs = tvm.ir.attrs.make_node(attrs_type_key, **kwargs)
+ attrs = tvm_tirx.ir.attrs.make_node(attrs_type_key, **kwargs)
return Call(op, args, attrs=attrs, ty_args=ty_args)
@@ -475,7 +475,7 @@ def select(condition, true_value, false_value):
def if_then_else_(condition: Any, true_value: Any, false_value: Any) -> Any:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.if_then_else_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.if_then_else_`."""
if hasattr(condition, "__tvm_ffi_object__"):
condition = condition.__tvm_ffi_object__()
if not isinstance(condition, _ir.Expr):
@@ -499,22 +499,22 @@ def if_then_else_(condition: Any, true_value: Any, false_value: Any) -> Any:
def and_(*values: Any) -> Any:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.and_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.and_`."""
return logical_and(*values)
def or_(*values: Any) -> Any:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.or_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.or_`."""
return logical_or(*values)
def not_(value: Any) -> Any:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.not_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.not_`."""
return logical_not(value)
def lt_(lhs: Any, rhs: Any, *, loc: _Loc = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.lt_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.lt_`."""
if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for value in (lhs, rhs)):
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
@@ -524,7 +524,7 @@ def lt_(lhs: Any, rhs: Any, *, loc: _Loc = None) -> _ir.Expr:
def le_(lhs: Any, rhs: Any, *, loc: _Loc = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.le_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.le_`."""
if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for value in (lhs, rhs)):
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
@@ -534,7 +534,7 @@ def le_(lhs: Any, rhs: Any, *, loc: _Loc = None) -> _ir.Expr:
def gt_(lhs: Any, rhs: Any, *, loc: _Loc = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.gt_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.gt_`."""
if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for value in (lhs, rhs)):
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
@@ -544,7 +544,7 @@ def gt_(lhs: Any, rhs: Any, *, loc: _Loc = None) -> _ir.Expr:
def ge_(lhs: Any, rhs: Any, *, loc: _Loc = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.ge_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.ge_`."""
if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for value in (lhs, rhs)):
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
@@ -554,7 +554,7 @@ def ge_(lhs: Any, rhs: Any, *, loc: _Loc = None) -> _ir.Expr:
def eq_(lhs: Any, rhs: Any, *, loc: _Loc = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.eq_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.eq_`."""
if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for value in (lhs, rhs)):
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
@@ -564,7 +564,7 @@ def eq_(lhs: Any, rhs: Any, *, loc: _Loc = None) -> _ir.Expr:
def ne_(lhs: Any, rhs: Any, *, loc: _Loc = None) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.ne_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.ne_`."""
if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for value in (lhs, rhs)):
lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
diff --git a/python/tvm/relax/script/ir_builder/parser_protocol.py b/python/tvm/relax/script/ir_builder/parser_protocol.py
index 639624a13a..90b6c5ec28 100644
--- a/python/tvm/relax/script/ir_builder/parser_protocol.py
+++ b/python/tvm/relax/script/ir_builder/parser_protocol.py
@@ -19,7 +19,7 @@
Hooks normalize syntax operands and construct native builder frames and IR directly.
For example, generated ``X.if_(condition)`` creates the native conditional frame;
``X.then_()`` and ``X.else_()`` enter its branches. See the corresponding shared
-``tvm.script.ir_builder.parser_protocol`` hooks for operand and loc contracts.
+``tvm_tirx.script.ir_builder.parser_protocol`` hooks for operand and loc contracts.
"""
from __future__ import annotations
@@ -30,16 +30,16 @@ from typing import Any, NoReturn
import tvm_ffi as _ffi
-from tvm import ir as _ir
+from tvm_tirx import ir as _ir
from tvm import relax as _relax
-from tvm.ir.location import UnknownLoc
+from tvm_tirx.ir.location import UnknownLoc
from tvm.relax import Call, Expr, Var, VarBinding
from tvm.relax.type import Type
from tvm.relax.utils import gen_call_tir_inputs
-from tvm.runtime import Object as tvm_Object
-from tvm.script.ir_builder import base as _base
-from tvm.script.ir_builder.base import IRBuilder as _IRBuilder
-from tvm.script.ir_builder.parser_protocol import decl_function
+from tvm_tirx.runtime import Object as tvm_Object
+from tvm_tirx.script.ir_builder import base as _base
+from tvm_tirx.script.ir_builder.base import IRBuilder as _IRBuilder
+from tvm_tirx.script.ir_builder.parser_protocol import decl_function
from . import _ffi_api
from . import frame as _frame
@@ -83,7 +83,7 @@ def function_(
reference: _ir.Var | None = None,
loc: _Loc = None,
) -> _frame.FunctionFrame:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.function_`.
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.function_`.
Public pure/private options pass to native purity and visibility controls.
Local function bodies require their previously declared reference.
@@ -106,7 +106,7 @@ def arg_(name: str, ty: Any, *, loc: _Loc = None) -> _ir.Var:
Source parameter name.
ty : Type, Var, or callable
Parameter annotation; corresponds to ``annotation`` in the shared
- :func:`tvm.script.ir_builder.parser_protocol.arg_` contract. Primitive
+ :func:`tvm_tirx.script.ir_builder.parser_protocol.arg_` contract. Primitive
annotations create a fresh parameter; an existing variable retains
its identity.
loc : LocationEntry, Location or None, optional
@@ -125,7 +125,7 @@ def arg_(name: str, ty: Any, *, loc: _Loc = None) -> _ir.Var:
def func_name_(name: str) -> None:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.func_name_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.func_name_`."""
return _ffi_api.FuncName(name)
@@ -141,7 +141,7 @@ def func_attr(attrs: dict[str, tvm_Object]) -> None:
def func_ret_type_(annotation: Any, *, loc: _Loc = None) -> None:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.func_ret_type_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.func_ret_type_`."""
return _ffi_api.FuncRetType(_native._type(_base._return_annotation(annotation)))
@@ -160,7 +160,7 @@ def check_well_formed_(function: _relax.Function) -> None:
Notes
-----
- See :func:`tvm.script.ir_builder.parser_protocol.check_well_formed_`
+ See :func:`tvm_tirx.script.ir_builder.parser_protocol.check_well_formed_`
for the shared whole-module validation coordinator.
Validate the completed Relax function. Whole-module cross-function
@@ -195,7 +195,7 @@ def _check_module_well_formed(module: _ir.IRModule) -> None:
def call_global_var_(function: _ir.GlobalVar, args: Sequence[Any]) -> _ir.Expr:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.call_global_var_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.call_global_var_`."""
return _relax.Call(function, [_relax.utils.convert_to_expr(value) for value in args])
@@ -247,7 +247,7 @@ def bind_(
name_loc: _Loc = None,
frame_value: bool = False,
) -> Any:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.bind_`.
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.bind_`.
Relax values emit a binding or MatchCast; value_loc belongs to the actual
RHS and loc to the binding target. Metadata passes through unchanged.
@@ -311,7 +311,7 @@ def decl_mutable_cell_(
loc: _Loc = None,
name_loc: _Loc = None,
) -> NoReturn:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.decl_mutable_cell_`.
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.decl_mutable_cell_`.
Relax has immutable bindings and rejects mutable storage declarations.
"""
@@ -319,7 +319,7 @@ def decl_mutable_cell_(
def set_mutable_cell_(target: Any, value: Any, *, loc: _Loc = None) -> NoReturn:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.set_mutable_cell_`.
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.set_mutable_cell_`.
Relax has immutable bindings and rejects mutable storage updates.
"""
@@ -327,7 +327,7 @@ def set_mutable_cell_(target: Any, value: Any, *, loc: _Loc = None) -> NoReturn:
def unpack_(value: Any) -> Any:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.unpack_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.unpack_`."""
if isinstance(value, _relax.Tuple):
return _python.tuple(value.fields)
if isinstance(value, _relax.Expr) and isinstance(value.ty, _ir.TupleType):
@@ -420,7 +420,7 @@ def emit_var_binding(value: VarBinding) -> Var:
def emit_(value: Any, *, loc: _Loc = None) -> None:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.emit_`.
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.emit_`.
Only void expressions may be discarded. Receipts and None emit nothing;
non-void expressions raise ValueError and unsupported host values raise TypeError.
@@ -443,12 +443,12 @@ def emit_(value: Any, *, loc: _Loc = None) -> None:
def setitem_(target: Any, key: Any, value: Any, *, loc: _Loc = None) -> NoReturn:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.setitem_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.setitem_`."""
raise TypeError("Relax does not support indexed assignment")
def setattr_(target: Any, name: str, value: Any, *, loc: _Loc = None) -> NoReturn:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.setattr_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.setattr_`."""
raise TypeError("Relax does not support attribute assignment")
@@ -466,24 +466,24 @@ def setattr_(target: Any, name: str, value: Any, *, loc: _Loc = None) -> NoRetur
def if_(condition: Any, *, loc: _Loc = None) -> _frame.IfFrame:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.if_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.if_`."""
if not isinstance(condition, _relax.Expr):
condition = _relax.prim_value(condition)
return _base.at_(loc, _ffi_api.If(condition))
def then_(*, loc: _Loc = None) -> _frame.ThenFrame:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.then_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.then_`."""
return _base.at_(loc, _ffi_api.Then())
def else_(*, loc: _Loc = None) -> _frame.ElseFrame:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.else_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.else_`."""
return _base.at_(loc, _ffi_api.Else())
def for_(iterable: Any, *, names: str | Sequence[str] | None = None, loc: _Loc = None) -> NoReturn:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.for_`.
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.for_`.
Relax rejects source loops; use supported functional control flow.
"""
@@ -491,17 +491,17 @@ def for_(iterable: Any, *, names: str | Sequence[str] | None = None, loc: _Loc =
def while_(condition: Any, *, loc: _Loc = None) -> NoReturn:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.while_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.while_`."""
raise TypeError("Relax does not support imperative while loops")
def range_(*args: Any, annotations: dict[str, Any] | None = None) -> NoReturn:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.range_`."""
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.range_`."""
raise TypeError("Relax does not support imperative for loops")
def break_(*, loc: _Loc = None) -> NoReturn:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.break_`.
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.break_`.
Relax rejects loop-control statements.
"""
@@ -509,7 +509,7 @@ def break_(*, loc: _Loc = None) -> NoReturn:
def continue_(*, loc: _Loc = None) -> NoReturn:
- """Implements :func:`tvm.script.ir_builder.parser_protocol.continue_`.
+ """Implements :func:`tvm_tirx.script.ir_builder.parser_protocol.continue_`.
Relax rejects loop-control statements.
"""
@@ -528,7 +528,7 @@ def func_ret_value(value: Expr) -> None:
... 109591 lines suppressed ...
(tvm) 21/33: [FIX][CUDA] Bound embedded launch source literals for MSVC
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 144ecab59d852bb3217389a53acf2125cb3cb69f
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 05:25:17 2026 +0000
[FIX][CUDA] Bound embedded launch source literals for MSVC
---
tirx/python/tvm_tirx/backend/cuda/launch/generate.py | 12 +++++++++---
tirx/src/backend/cuda/launch/launch_source.h | 18 ++++++++++++------
2 files changed, 21 insertions(+), 9 deletions(-)
diff --git a/tirx/python/tvm_tirx/backend/cuda/launch/generate.py b/tirx/python/tvm_tirx/backend/cuda/launch/generate.py
index 381ac5c913..5e546a650b 100644
--- a/tirx/python/tvm_tirx/backend/cuda/launch/generate.py
+++ b/tirx/python/tvm_tirx/backend/cuda/launch/generate.py
@@ -333,14 +333,20 @@ def main():
.replace('#include "launch_attributes.h"', ""),
]
)
+ # MSVC limits individual string literals to 16380 bytes before concatenation.
+ chunks = [""]
+ for line in ("\n" + support + "\n").splitlines(keepends=True):
+ if len((chunks[-1] + line).encode("utf-8")) > 8000:
+ chunks.append("")
+ chunks[-1] += line
outputs[ROOT / "src/backend/cuda/launch/launch_source.h"] = (
CPP_LICENSE
+ "\n// Generated by launch/generate.py; edit the registry or launch_config.h.\n"
"#ifndef TVM_BACKEND_CUDA_LAUNCH_LAUNCH_SOURCE_H_\n"
"#define TVM_BACKEND_CUDA_LAUNCH_LAUNCH_SOURCE_H_\n"
- 'namespace tvm_cuda_launch {\ninline constexpr const char* kSupportSource = R"TVM_CUDA_LAUNCH(\n'
- + support
- + '\n)TVM_CUDA_LAUNCH";\n}\n#endif\n'
+ "namespace tvm_cuda_launch {\ninline constexpr const char* kSupportSource =\n"
+ + "\n".join(f'R"TVM_CUDA_LAUNCH({chunk})TVM_CUDA_LAUNCH"' for chunk in chunks)
+ + ";\n}\n#endif\n"
)
source_path = ROOT / "src/backend/cuda/launch/launch_source.h"
outputs[source_path] = format_output(source_path, outputs[source_path])
diff --git a/tirx/src/backend/cuda/launch/launch_source.h b/tirx/src/backend/cuda/launch/launch_source.h
index 1f8dbd9563..df5125f8a1 100644
--- a/tirx/src/backend/cuda/launch/launch_source.h
+++ b/tirx/src/backend/cuda/launch/launch_source.h
@@ -21,7 +21,8 @@
#ifndef TVM_BACKEND_CUDA_LAUNCH_LAUNCH_SOURCE_H_
#define TVM_BACKEND_CUDA_LAUNCH_LAUNCH_SOURCE_H_
namespace tvm_cuda_launch {
-inline constexpr const char* kSupportSource = R"TVM_CUDA_LAUNCH(
+inline constexpr const char* kSupportSource =
+ R"TVM_CUDA_LAUNCH(
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
@@ -223,7 +224,8 @@ inline void NormalizeLegacyPresence(LaunchValues* values) {
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
+)TVM_CUDA_LAUNCH"
+ R"TVM_CUDA_LAUNCH( * distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
@@ -400,7 +402,8 @@ inline unsigned EncodeDriverAttributes(const LaunchValues& config,
#if CUDA_VERSION >= 13020
auto& attr = (*attrs)[count++];
attr.id = CU_LAUNCH_ATTRIBUTE_PORTABLE_CLUSTER_SIZE_MODE;
- attr.value.portableClusterSizeMode = static_cast<decltype(attr.value.portableClusterSizeMode)>(
+)TVM_CUDA_LAUNCH"
+ R"TVM_CUDA_LAUNCH( attr.value.portableClusterSizeMode = static_cast<decltype(attr.value.portableClusterSizeMode)>(
config.Get(Field::kPortableClusterSizeMode).integer);
#else
throw std::invalid_argument(
@@ -574,7 +577,8 @@ inline unsigned EncodeRuntimeAttributes(const LaunchValues& config,
attr.val.accessPolicyWindow.missProp =
static_cast<decltype(attr.val.accessPolicyWindow.missProp)>(
config.Get(Field::kAccessPolicyWindowMissProp).integer);
-#else
+)TVM_CUDA_LAUNCH"
+ R"TVM_CUDA_LAUNCH(#else
throw std::invalid_argument(
"CUDA launch field access_policy_window requires toolkit 12.0 or newer");
#endif
@@ -759,7 +763,8 @@ inline LaunchValues Normalize(LaunchValues config, const KernelRequirements& ker
}
}
for (Field flag :
- {Field::kCooperative, Field::kProgrammaticStreamSerialization,
+)TVM_CUDA_LAUNCH"
+ R"TVM_CUDA_LAUNCH( {Field::kCooperative, Field::kProgrammaticStreamSerialization,
Field::kNvlinkUtilCentricScheduling, Field::kProgrammaticEventTriggerAtBlockStart}) {
Require(config.Int(flag) == 0 || config.Int(flag) == 1,
std::string(kFields[static_cast<size_t>(flag)].name) + " must be a boolean");
@@ -933,7 +938,8 @@ inline void EncodeRuntime(const LaunchValues& values, const KernelRequirements&
*block[axis] = values.Int(Axis(Field::kBlockX, axis));
if (kernel.required_block_size) {
#if CUDART_VERSION >= 13000
- *grid[axis] /= values.Int(Axis(Field::kClusterX, axis));
+)TVM_CUDA_LAUNCH"
+ R"TVM_CUDA_LAUNCH( *grid[axis] /= values.Int(Axis(Field::kClusterX, axis));
*block[axis] = 1;
#else
throw std::invalid_argument("CUDA required block size requires CUDA 13 or newer");
(tvm) 16/33: [STYLE] Format relocated include paths and printer imports
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit e7fff21d102e6794f80286b3b22db95aa96fadc4
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 02:10:48 2026 +0000
[STYLE] Format relocated include paths and printer imports
---
apps/android_rpc/app/src/main/jni/tvm_runtime.h | 12 ++++++------
include/tvm/runtime/vm/tensor_cache_support.h | 7 ++++---
src/runtime/extra/contrib/cublas/cublas.cc | 2 +-
src/runtime/extra/contrib/hipblas/hipblas.cc | 2 +-
tirx/include/tvm/codegen/codegen.h | 8 ++++----
tirx/src/script/printer/printer.cc | 4 ++--
web/emcc/wasm_runtime.cc | 12 ++++++------
7 files changed, 24 insertions(+), 23 deletions(-)
diff --git a/apps/android_rpc/app/src/main/jni/tvm_runtime.h b/apps/android_rpc/app/src/main/jni/tvm_runtime.h
index fa5784e0d7..60638071c6 100644
--- a/apps/android_rpc/app/src/main/jni/tvm_runtime.h
+++ b/apps/android_rpc/app/src/main/jni/tvm_runtime.h
@@ -43,11 +43,6 @@
#include "../3rdparty/tvm-ffi/src/ffi/function.cc"
#include "../3rdparty/tvm-ffi/src/ffi/object.cc"
#include "../3rdparty/tvm-ffi/src/ffi/tensor.cc"
-#include "../tirx/src/runtime/cpu_device_api.cc"
-#include "../tirx/src/runtime/device_api.cc"
-#include "../tirx/src/runtime/logging.cc"
-#include "../tirx/src/runtime/memory_scope.cc"
-#include "../tirx/src/runtime/metadata.cc"
#include "../src/runtime/registry.cc"
#include "../src/runtime/rpc/rpc_channel.cc"
#include "../src/runtime/rpc/rpc_endpoint.cc"
@@ -57,11 +52,16 @@
#include "../src/runtime/rpc/rpc_server_env.cc"
#include "../src/runtime/rpc/rpc_session.cc"
#include "../src/runtime/rpc/rpc_socket_impl.cc"
+#include "../src/runtime/vm/params.cc"
+#include "../tirx/src/runtime/cpu_device_api.cc"
+#include "../tirx/src/runtime/device_api.cc"
+#include "../tirx/src/runtime/logging.cc"
+#include "../tirx/src/runtime/memory_scope.cc"
+#include "../tirx/src/runtime/metadata.cc"
#include "../tirx/src/runtime/tensor.cc"
#include "../tirx/src/runtime/thread_pool.cc"
#include "../tirx/src/runtime/threading_backend.cc"
#include "../tirx/src/runtime/timer.cc"
-#include "../src/runtime/vm/params.cc"
#include "../tirx/src/runtime/workspace_pool.cc"
#ifdef TVM_OPENCL_RUNTIME
diff --git a/include/tvm/runtime/vm/tensor_cache_support.h b/include/tvm/runtime/vm/tensor_cache_support.h
index d5b1546e96..0293810deb 100644
--- a/include/tvm/runtime/vm/tensor_cache_support.h
+++ b/include/tvm/runtime/vm/tensor_cache_support.h
@@ -46,8 +46,9 @@ struct TensorCacheMetadata {
* \param staging_buffer The buffer to be used to avoid extra OpenCL copies. Pass in a nullptr
* in other cases
*/
- TVM_RUNTIME_EXTRA_DLL ffi::Tensor Load(Device device, const std::string* raw_data,
- ffi::Optional<ffi::Tensor>* staging_buffer = nullptr) const;
+ TVM_RUNTIME_EXTRA_DLL ffi::Tensor Load(
+ Device device, const std::string* raw_data,
+ ffi::Optional<ffi::Tensor>* staging_buffer = nullptr) const;
/*! \brief Name of the parameter */
std::string name;
@@ -88,7 +89,7 @@ struct TensorCacheMetadata {
TVM_RUNTIME_EXTRA_DLL static TensorCacheMetadata Load(const std::string& path);
/*! \brief Load the metadata from a given JSON string */
TVM_RUNTIME_EXTRA_DLL static TensorCacheMetadata LoadFromStr(const std::string& json_str,
- const std::string& path);
+ const std::string& path);
};
} // namespace vm
diff --git a/src/runtime/extra/contrib/cublas/cublas.cc b/src/runtime/extra/contrib/cublas/cublas.cc
index 5c293fcd4b..a7cd28b1b1 100644
--- a/src/runtime/extra/contrib/cublas/cublas.cc
+++ b/src/runtime/extra/contrib/cublas/cublas.cc
@@ -27,8 +27,8 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
-#include "builtin_fp16.h"
#include "../cblas/gemm_common.h"
+#include "builtin_fp16.h"
#include "cublas_utils.h"
namespace tvm {
diff --git a/src/runtime/extra/contrib/hipblas/hipblas.cc b/src/runtime/extra/contrib/hipblas/hipblas.cc
index f1e8f890ca..c2120192c9 100644
--- a/src/runtime/extra/contrib/hipblas/hipblas.cc
+++ b/src/runtime/extra/contrib/hipblas/hipblas.cc
@@ -26,8 +26,8 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
-#include "builtin_fp16.h"
#include "../cblas/gemm_common.h"
+#include "builtin_fp16.h"
#include "hipblas_utils.h"
namespace tvm {
diff --git a/tirx/include/tvm/codegen/codegen.h b/tirx/include/tvm/codegen/codegen.h
index 6e2925a9cf..8fdb4a1b0d 100644
--- a/tirx/include/tvm/codegen/codegen.h
+++ b/tirx/include/tvm/codegen/codegen.h
@@ -55,7 +55,7 @@ TVM_TIRX_DLL ffi::Module Build(IRModule mod, Target target);
* \return cstr The C string representation of the file.
*/
TVM_TIRX_DLL std::string PackImportsToC(const ffi::Module& m, bool system_lib,
- const std::string& c_symbol_prefix = "");
+ const std::string& c_symbol_prefix = "");
/*!
* \brief Pack imported device libraries into an object file using LLVM.
@@ -70,9 +70,9 @@ TVM_TIRX_DLL std::string PackImportsToC(const ffi::Module& m, bool system_lib,
* \param c_symbol_prefix Optional blob symbol prefix, valid only for system libraries.
*/
TVM_TIRX_DLL void PackImportsToObjectFile(const ffi::Module& mod, bool system_lib,
- const std::string& llvm_target_string,
- const std::string& file_name,
- const std::string& c_symbol_prefix = "");
+ const std::string& llvm_target_string,
+ const std::string& file_name,
+ const std::string& c_symbol_prefix = "");
namespace attr {
constexpr const char* kParallelStridePattern = "parallel_stride_pattern";
diff --git a/tirx/src/script/printer/printer.cc b/tirx/src/script/printer/printer.cc
index a8ffd598e2..feeacc432c 100644
--- a/tirx/src/script/printer/printer.cc
+++ b/tirx/src/script/printer/printer.cc
@@ -372,8 +372,8 @@ ffi::String RenderScript(const ffi::ObjectRef& obj, const PrinterConfig& config)
->Attr("load_json")
->Call({LiteralDoc::Str(json, std::nullopt)}),
std::nullopt));
- header.push_back(
- ffi::String(loader == "tvm_tirx" ? "import tvm_tirx\n" : "import tvm_tirx as " + loader + "\n"));
+ header.push_back(ffi::String(loader == "tvm_tirx" ? "import tvm_tirx\n"
+ : "import tvm_tirx as " + loader + "\n"));
} else {
block.value()->stmts.push_back(
CommentDoc("Metadata omitted. Use show_meta=True in script() method to show it."));
diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc
index 02948f145e..a9f07eb0c0 100644
--- a/web/emcc/wasm_runtime.cc
+++ b/web/emcc/wasm_runtime.cc
@@ -35,20 +35,20 @@
#include <cstring>
#include <limits>
-#include "tirx/src/runtime/cpu_device_api.cc"
-#include "tirx/src/runtime/device_api.cc"
#include "src/runtime/extra/contrib/sort/sort.cc"
-#include "tirx/src/runtime/logging.cc"
-#include "tirx/src/runtime/metadata.cc"
#include "src/runtime/rpc/rpc_channel.cc"
#include "src/runtime/rpc/rpc_endpoint.cc"
#include "src/runtime/rpc/rpc_event_impl.cc"
#include "src/runtime/rpc/rpc_local_session.cc"
#include "src/runtime/rpc/rpc_module.cc"
#include "src/runtime/rpc/rpc_session.cc"
+#include "src/runtime/vm/params.cc"
+#include "tirx/src/runtime/cpu_device_api.cc"
+#include "tirx/src/runtime/device_api.cc"
+#include "tirx/src/runtime/logging.cc"
+#include "tirx/src/runtime/metadata.cc"
#include "tirx/src/runtime/tensor.cc"
#include "tirx/src/runtime/timer.cc"
-#include "src/runtime/vm/params.cc"
#include "tirx/src/runtime/workspace_pool.cc"
// relax setup
#include "3rdparty/tvm-ffi/src/ffi/backtrace.cc"
@@ -69,7 +69,6 @@
#include "3rdparty/tvm-ffi/src/ffi/tensor.cc"
#include "3rdparty/tvm-ffi/src/ffi/testing/testing.cc"
// Enum registration requires the FFI type-attribute registry above.
-#include "tirx/src/runtime/memory_scope.cc"
#include "src/runtime/vm/attn_backend.cc"
#include "src/runtime/vm/builtin.cc"
#include "src/runtime/vm/bytecode.cc"
@@ -80,6 +79,7 @@
#include "src/runtime/vm/rnn_state.cc"
#include "src/runtime/vm/tensor_cache_support.cc"
#include "src/runtime/vm/vm.cc"
+#include "tirx/src/runtime/memory_scope.cc"
// --- Implementations of backend and wasm runtime API. ---
(tvm) 08/33: [REFACTOR][Script] Emit imports from each dialect owning package
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 72ac640cf4748b57e2cfa621e981031de918d63c
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 01:37:32 2026 +0000
[REFACTOR][Script] Emit imports from each dialect owning package
---
tirx/src/script/printer/printer.cc | 11 +++++++----
tirx/src/tirx/ir/index_map.cc | 2 +-
2 files changed, 8 insertions(+), 5 deletions(-)
diff --git a/tirx/src/script/printer/printer.cc b/tirx/src/script/printer/printer.cc
index 271260ab8e..a8ffd598e2 100644
--- a/tirx/src/script/printer/printer.cc
+++ b/tirx/src/script/printer/printer.cc
@@ -360,9 +360,9 @@ ffi::String RenderScript(const ffi::ObjectRef& obj, const PrinterConfig& config)
if (auto id = child.as<IdDoc>()) used.insert(id.value()->name);
},
&visited);
- std::string loader = "tvm";
+ std::string loader = "tvm_tirx";
for (size_t suffix = 1; used.count(loader); ++suffix)
- loader = "tvm_" + std::to_string(suffix);
+ loader = "tvm_tirx_" + std::to_string(suffix);
ffi::String json = ffi::json::Stringify(
ffi::ToJSONGraph(metadata, ffi::json::Object{{"tvm_version", TVM_VERSION}}), 2);
block.value()->stmts.insert(block.value()->stmts.begin(),
@@ -373,7 +373,7 @@ ffi::String RenderScript(const ffi::ObjectRef& obj, const PrinterConfig& config)
->Call({LiteralDoc::Str(json, std::nullopt)}),
std::nullopt));
header.push_back(
- ffi::String(loader == "tvm" ? "import tvm\n" : "import tvm as " + loader + "\n"));
+ ffi::String(loader == "tvm_tirx" ? "import tvm_tirx\n" : "import tvm_tirx as " + loader + "\n"));
} else {
block.value()->stmts.push_back(
CommentDoc("Metadata omitted. Use show_meta=True in script() method to show it."));
@@ -411,7 +411,10 @@ ffi::String RenderScript(const ffi::ObjectRef& obj, const PrinterConfig& config)
canonical.resize(canonical.size() - std::string(".prefix").size());
if (!namespaces.count(canonical)) continue;
ffi::String alias = config->GetExtraConfig<ffi::String>(key, fallback);
- ffi::String import = "from tvm_tirx.script import " + canonical + " as " + std::string(alias);
+ ffi::String import =
+ canonical == "s_tir" || canonical == "relax"
+ ? "from tvm." + canonical + " import script as " + std::string(alias)
+ : "from tvm_tirx.script import " + canonical + " as " + std::string(alias);
if (config->GetExtraConfig<bool>("ir.comment_imports", false) && !config->show_meta) {
header.push_back(CommentDoc(import));
} else {
diff --git a/tirx/src/tirx/ir/index_map.cc b/tirx/src/tirx/ir/index_map.cc
index 547e385035..28fb38565c 100644
--- a/tirx/src/tirx/ir/index_map.cc
+++ b/tirx/src/tirx/ir/index_map.cc
@@ -473,7 +473,7 @@ ffi::String IndexMapNode::ToPythonString(
std::string inverse_lambda_expr =
IndexMap2PythonLambdaExpr(inverse->initial_indices, inverse->final_indices);
std::ostringstream oss;
- oss << "tvm.tirx.IndexMap.from_func(" << lambda_expr
+ oss << "tvm_tirx.tirx.IndexMap.from_func(" << lambda_expr
<< ", inverse_index_map=" << inverse_lambda_expr << ")";
return ffi::String(oss.str());
}
(tvm) 27/33: [FIX][IR] Keep the name-supply template constructor local to consumers
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit c5f78007c509661f4f00a289443e90acf01988f9
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 07:54:10 2026 +0000
[FIX][IR] Keep the name-supply template constructor local to consumers
---
tirx/include/tvm/ir/unique_name_supply.h | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/tirx/include/tvm/ir/unique_name_supply.h b/tirx/include/tvm/ir/unique_name_supply.h
index b0cadde29c..cc5a8f3745 100644
--- a/tirx/include/tvm/ir/unique_name_supply.h
+++ b/tirx/include/tvm/ir/unique_name_supply.h
@@ -128,7 +128,7 @@ class UniqueNameSupply : public ffi::ObjectRef {
* The functor should return the name of the dereferenced object.
*/
template <typename Iter, typename Lambda>
- TVM_TIRX_DLL UniqueNameSupply(Iter begin, Iter end, Lambda f) : UniqueNameSupply("") {
+ UniqueNameSupply(Iter begin, Iter end, Lambda f) : UniqueNameSupply("") {
for (auto it = begin; it != end; ++it) {
this->operator->()->ReserveName(f(*it), false);
}
(tvm) 15/33: [FIX][Lint] Apply repository lint ownership to the extracted tree
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 64ac09abdacdb28ff12a6560e37abe377103113a
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 02:10:47 2026 +0000
[FIX][Lint] Apply repository lint ownership to the extracted tree
---
.pre-commit-config.yaml | 8 ++++----
tests/lint/check_asf_header.py | 1 +
tests/lint/check_file_type.py | 4 ++--
3 files changed, 7 insertions(+), 6 deletions(-)
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index c76f80bf57..ec8c900505 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -53,13 +53,13 @@ repos:
- id: check-merge-conflict
- id: check-symlinks
- id: end-of-file-fixer
- exclude: ^3rdparty/
+ exclude: ^(tirx/)?3rdparty/
- id: mixed-line-ending
- exclude: ^3rdparty/
+ exclude: ^(tirx/)?3rdparty/
- id: requirements-txt-fixer
exclude: ^docker/python/bootstrap
- id: trailing-whitespace
- exclude: ^3rdparty/
+ exclude: ^(tirx/)?3rdparty/
- id: check-yaml
- id: check-toml
@@ -95,7 +95,7 @@ repos:
hooks:
- id: clang-format
types_or: [c, c++, cuda, objective-c++]
- exclude: ^3rdparty/
+ exclude: ^(tirx/)?3rdparty/
# ---------- Cython ----------
- repo: https://github.com/MarcoGorelli/cython-lint
diff --git a/tests/lint/check_asf_header.py b/tests/lint/check_asf_header.py
index d6237c600a..9cc1ebf73a 100644
--- a/tests/lint/check_asf_header.py
+++ b/tests/lint/check_asf_header.py
@@ -183,6 +183,7 @@ FMT_MAP = {
# Files and patterns to skip (3rdparty, generated files, etc.).
SKIP_LIST: list[str] = [
"3rdparty/*",
+ "tirx/3rdparty/*",
"ffi/3rdparty/*",
".github/*",
".txdev/*",
diff --git a/tests/lint/check_file_type.py b/tests/lint/check_file_type.py
index 9318cf8b6c..9d301bec18 100644
--- a/tests/lint/check_file_type.py
+++ b/tests/lint/check_file_type.py
@@ -93,7 +93,7 @@ ALLOW_FILE_NAME = {
}
# List of specific files allowed in relpath to <proj_root>
-ALLOW_SPECIFIC_FILE = {"LICENSE", "NOTICE", "KEYS"}
+ALLOW_SPECIFIC_FILE = {"LICENSE", "NOTICE", "KEYS", "tirx/LICENSE", "tirx/NOTICE"}
def filename_allowed(name: str) -> bool:
@@ -121,7 +121,7 @@ def filename_allowed(name: str) -> bool:
if basename == "Dockerfile" or basename.startswith("Dockerfile."):
return True
- if name.startswith("3rdparty"):
+ if name.startswith(("3rdparty", "tirx/3rdparty")):
return True
if name.startswith(".txdev") or name.startswith(".claude"):
(tvm) 23/33: [FIX][Metal] Restore the compiler callback after its owning test
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 1d8feb685786d01f1114c2b9825c23f439fb6637
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 05:25:18 2026 +0000
[FIX][Metal] Restore the compiler callback after its owning test
---
.../python/backend/metal/codegen/test_metal.py | 25 ++++++++++++++++------
1 file changed, 18 insertions(+), 7 deletions(-)
diff --git a/tirx/tests/python/backend/metal/codegen/test_metal.py b/tirx/tests/python/backend/metal/codegen/test_metal.py
index 1fb69ddc2a..d582ef1cc0 100644
--- a/tirx/tests/python/backend/metal/codegen/test_metal.py
+++ b/tirx/tests/python/backend/metal/codegen/test_metal.py
@@ -16,6 +16,7 @@
# under the License.
import numpy as np
import pytest
+import tvm_ffi
import tvm_tirx
import tvm_tirx.testing
@@ -202,16 +203,26 @@ def test_func_with_trailing_pod_params():
with T.launch_thread("threadIdx.x", 16) as i:
B[i] = A[i] + x
- @tvm_tirx.register_global_func("tvm_callback_metal_compile")
+ previous_callback = tvm_tirx.get_global_func("tvm_callback_metal_compile", allow_missing=True)
+
+ @tvm_tirx.register_global_func("tvm_callback_metal_compile", override=True)
def compile_metal(src, target):
return xcode.compile_metal(src)
- mod = tvm_tirx.IRModule({"main": func})
-
- f = tvm_tirx.tirx.compile(mod, target="metal").mod
- src: str = f.imports[0].inspect_source()
- occurrences = src.count("struct func_kernel_args_t")
- assert occurrences == 1, occurrences
+ try:
+ mod = tvm_tirx.IRModule({"main": func})
+
+ f = tvm_tirx.tirx.compile(mod, target="metal").mod
+ src: str = f.imports[0].inspect_source()
+ occurrences = src.count("struct func_kernel_args_t")
+ assert occurrences == 1, occurrences
+ finally:
+ if previous_callback is None:
+ tvm_ffi.registry.remove_global_func("tvm_callback_metal_compile")
+ else:
+ tvm_tirx.register_global_func(
+ "tvm_callback_metal_compile", previous_callback, override=True
+ )
def test_codegen_simdgroup_buffer_data():
(tvm) 28/33: [FIX][Testing] Probe the loaded NCCL runtime capability
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 04f31ef66962d3c4587f618f5b246931eba906a7
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 07:54:21 2026 +0000
[FIX][Testing] Probe the loaded NCCL runtime capability
---
tirx/python/tvm_tirx/testing/env.py | 7 ++++++-
1 file changed, 6 insertions(+), 1 deletion(-)
diff --git a/tirx/python/tvm_tirx/testing/env.py b/tirx/python/tvm_tirx/testing/env.py
index 0fbf896493..7e58d5c269 100644
--- a/tirx/python/tvm_tirx/testing/env.py
+++ b/tirx/python/tvm_tirx/testing/env.py
@@ -284,7 +284,12 @@ def has_cublas() -> bool:
def has_nccl() -> bool:
"""True if NCCL was built in and a CUDA device is present."""
- return has_cuda() and build_flag_enabled("USE_NCCL")
+ try:
+ return has_cuda() and (
+ tvm_tirx.get_global_func("runtime.disco.nccl.init_ccl", allow_missing=True) is not None
+ )
+ except Exception: # pylint: disable=broad-except
+ return False
def has_hipblas() -> bool:
(tvm) 32/33: [STYLE][CI] Format the native build dependency command
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 674b3ae17d10c177384003f479d36355a63c51a2
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 12:46:21 2026 +0000
[STYLE][CI] Format the native build dependency command
---
tests/scripts/task_build.py | 12 ++++++++++--
1 file changed, 10 insertions(+), 2 deletions(-)
diff --git a/tests/scripts/task_build.py b/tests/scripts/task_build.py
index 3160870aaf..f7bfd558e6 100755
--- a/tests/scripts/task_build.py
+++ b/tests/scripts/task_build.py
@@ -93,8 +93,16 @@ if __name__ == "__main__":
if not ffi_marker.is_file() or ffi_marker.read_text().strip() != ffi_identity:
subprocess.run(
[
- "uv", "pip", "install", "--python", sys.executable,
- "-v", "--no-deps", "--reinstall", "--target=python", "./3rdparty/tvm-ffi/",
+ "uv",
+ "pip",
+ "install",
+ "--python",
+ sys.executable,
+ "-v",
+ "--no-deps",
+ "--reinstall",
+ "--target=python",
+ "./3rdparty/tvm-ffi/",
],
cwd=REPO_ROOT,
env=env,
(tvm) 12/33: [FIX][Script] Synchronize the generated TIRx namespace reference
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit f393d43497b2ad7dd502419fa2a903d6029fff5c
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 01:48:54 2026 +0000
[FIX][Script] Synchronize the generated TIRx namespace reference
---
tirx/python/tvm_tirx/script/tirx.pyi | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/tirx/python/tvm_tirx/script/tirx.pyi b/tirx/python/tvm_tirx/script/tirx.pyi
index e7ee8136fd..3acfadbce1 100644
--- a/tirx/python/tvm_tirx/script/tirx.pyi
+++ b/tirx/python/tvm_tirx/script/tirx.pyi
@@ -2634,5 +2634,5 @@ class _PTX:
ptx: _PTX
-# Every other tvm.script.tirx member stays dynamically typed, as before.
+# Every other tvm_tirx.script.tirx member stays dynamically typed, as before.
def __getattr__(name: str) -> Any: ...
(tvm) 11/33: [REFACTOR][Python] Format imports after the core namespace split
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 1614b98573ebcfacb72e3c36dd6b4d9795ad8392
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 01:48:54 2026 +0000
[REFACTOR][Python] Format imports after the core namespace split
---
apps/android_rpc/tests/android_rpc_test.py | 7 +-
apps/ios_rpc/tests/ios_rpc_test.py | 7 +-
jvm/core/src/test/scripts/prepare_test_libs.py | 11 +-
python/tvm/contrib/cblas.py | 2 +-
python/tvm/contrib/cublas.py | 2 +-
python/tvm/contrib/cublaslt.py | 2 +-
python/tvm/contrib/cudnn.py | 2 +-
python/tvm/contrib/cutlass/build.py | 6 +-
python/tvm/contrib/dlpack.py | 6 +-
python/tvm/contrib/dnnl.py | 2 +-
python/tvm/contrib/hexagon/generate_take_op.py | 9 +-
python/tvm/contrib/hexagon/tools.py | 2 +-
python/tvm/contrib/hipblas.py | 2 +-
python/tvm/contrib/mkl.py | 2 +-
python/tvm/contrib/nnpack.py | 2 +-
python/tvm/contrib/random.py | 2 +-
python/tvm/contrib/tvmjs.py | 3 +-
python/tvm/driver/build_module.py | 3 +-
python/tvm/exec/disco_remote_socket_session.py | 3 +-
python/tvm/exec/disco_worker.py | 2 +-
python/tvm/exec/gpu_memory_bandwidth.py | 3 +-
python/tvm/relax/__init__.py | 7 +-
python/tvm/relax/analysis/analysis.py | 7 +-
python/tvm/relax/analysis/estimate_memory_usage.py | 2 +-
python/tvm/relax/backend/adreno/clml.py | 6 +-
python/tvm/relax/backend/adreno/pipeline.py | 3 +-
.../relax/backend/adreno/transform/transform.py | 2 +-
.../relax/backend/contrib/example_npu/patterns.py | 2 +-
python/tvm/relax/backend/contrib/nnapi.py | 2 +-
python/tvm/relax/backend/contrib/tensorrt.py | 4 +-
python/tvm/relax/backend/cpu_generic/pipeline.py | 3 +-
python/tvm/relax/backend/cuda/cublas.py | 5 +-
python/tvm/relax/backend/cuda/cudnn.py | 2 +-
python/tvm/relax/backend/cuda/cutlass.py | 2 +-
python/tvm/relax/backend/cuda/flashinfer.py | 2 +-
python/tvm/relax/backend/cuda/pipeline.py | 3 +-
python/tvm/relax/backend/dispatch_sampling.py | 2 +-
python/tvm/relax/backend/dispatch_sort_scan.py | 4 +-
python/tvm/relax/backend/gpu_generic/pipeline.py | 3 +-
python/tvm/relax/backend/metal/coreml.py | 5 +-
python/tvm/relax/backend/rocm/hipblas.py | 2 +-
python/tvm/relax/backend/rocm/pipeline.py | 3 +-
python/tvm/relax/base_py_module.py | 9 +-
python/tvm/relax/binding_rewrite.py | 2 +-
python/tvm/relax/block_builder.py | 2 +-
.../tvm/relax/distributed/transform/transform.py | 2 +-
python/tvm/relax/distributed/type.py | 2 +-
python/tvm/relax/dpl/context.py | 2 +-
python/tvm/relax/dpl/pattern.py | 6 +-
python/tvm/relax/exec_builder.py | 2 +-
python/tvm/relax/expr.py | 10 +-
python/tvm/relax/expr_functor.py | 4 +-
python/tvm/relax/frontend/common.py | 6 +-
python/tvm/relax/frontend/nn/core.py | 4 +-
python/tvm/relax/frontend/nn/exporter.py | 2 +-
python/tvm/relax/frontend/nn/extern.py | 5 +-
python/tvm/relax/frontend/nn/llm/kv_cache.py | 6 +-
python/tvm/relax/frontend/nn/subroutine.py | 2 +-
python/tvm/relax/frontend/nn/torch.py | 2 +-
python/tvm/relax/frontend/onnx/onnx_frontend.py | 30 +-
.../frontend/stablehlo/stablehlo_translator.py | 4 +-
.../frontend/torch/base_fx_graph_translator.py | 16 +-
python/tvm/relax/frontend/torch/dynamo.py | 3 +-
.../frontend/torch/exported_program_translator.py | 2 +-
python/tvm/relax/frontend/torch/fx_translator.py | 2 +-
python/tvm/relax/ir/instrument.py | 2 +-
python/tvm/relax/op/__init__.py | 3 +-
python/tvm/relax/op/_op_gradient.py | 2 +-
python/tvm/relax/op/base.py | 9 +-
python/tvm/relax/op/distributed/distributed.py | 2 +-
python/tvm/relax/op/index.py | 2 +-
python/tvm/relax/op/memory/view.py | 4 +-
python/tvm/relax/op/nn/nn.py | 2 +-
python/tvm/relax/op/set.py | 2 +-
python/tvm/relax/pipeline.py | 14 +-
python/tvm/relax/relax_to_pyfunc_converter.py | 5 +-
python/tvm/relax/script/__init__.py | 6 +-
.../tvm/relax/script/ir_builder/distributed/ir.py | 4 +-
.../tvm/relax/script/ir_builder/distributed/op.py | 2 +-
python/tvm/relax/script/ir_builder/ir.py | 8 +-
python/tvm/relax/script/ir_builder/op.py | 12 +-
.../tvm/relax/script/ir_builder/parser_protocol.py | 4 +-
python/tvm/relax/testing/ast_printer.py | 4 +-
python/tvm/relax/testing/lib_comparator.py | 4 +-
python/tvm/relax/testing/nn.py | 5 +-
python/tvm/relax/testing/transform.py | 12 +-
python/tvm/relax/testing/vm.py | 2 +-
python/tvm/relax/training/optimizer.py | 8 +-
python/tvm/relax/training/setup_trainer.py | 2 +-
python/tvm/relax/training/trainer.py | 23 +-
python/tvm/relax/training/utils.py | 4 +-
python/tvm/relax/transform/fast_math.py | 5 +-
.../fold_batch_norm_to_conv2d_for_inference.py | 6 +-
.../tvm/relax/transform/fuse_transpose_matmul.py | 4 +-
.../tvm/relax/transform/ipc_allreduce_rewrite.py | 4 +-
.../tvm/relax/transform/lazy_transform_params.py | 4 +-
python/tvm/relax/transform/legalize_ops/ccl.py | 4 +-
python/tvm/relax/transform/legalize_ops/common.py | 2 +-
python/tvm/relax/transform/legalize_ops/create.py | 4 +-
.../relax/transform/legalize_ops/distributed.py | 2 +-
python/tvm/relax/transform/legalize_ops/grad.py | 2 +-
python/tvm/relax/transform/legalize_ops/image.py | 2 +-
python/tvm/relax/transform/legalize_ops/index.py | 4 +-
.../relax/transform/legalize_ops/linear_algebra.py | 2 +-
.../tvm/relax/transform/legalize_ops/manipulate.py | 7 +-
python/tvm/relax/transform/legalize_ops/nn.py | 2 +-
python/tvm/relax/transform/legalize_ops/qdq.py | 4 +-
.../relax/transform/legalize_ops/statistical.py | 2 +-
python/tvm/relax/transform/legalize_ops/vision.py | 4 +-
.../relax/transform/lower_gpu_ipc_alloc_storage.py | 4 +-
.../relax/transform/optimize_layout_transform.py | 4 +-
.../relax/transform/remove_redundant_reshape.py | 4 +-
python/tvm/relax/transform/transform.py | 7 +-
python/tvm/relax/type.py | 2 +-
python/tvm/relax/utils.py | 12 +-
python/tvm/relax/vm_build.py | 2 +-
python/tvm/rpc/client.py | 3 +-
python/tvm/rpc/remote.py | 2 +-
python/tvm/rpc/testing.py | 2 +-
python/tvm/runtime/disco/session.py | 1 +
python/tvm/runtime/vm/vm.py | 2 +-
python/tvm/s_tir/__init__.py | 7 +-
python/tvm/s_tir/analysis/__init__.py | 3 +-
python/tvm/s_tir/dlight/adreno/convolution.py | 2 +-
python/tvm/s_tir/dlight/adreno/fallback.py | 2 +-
python/tvm/s_tir/dlight/adreno/layout_transform.py | 2 +-
python/tvm/s_tir/dlight/adreno/pool.py | 2 +-
.../tvm/s_tir/dlight/analysis/common_analysis.py | 8 +-
python/tvm/s_tir/dlight/analysis/gemv.py | 4 +-
python/tvm/s_tir/dlight/base/schedule_rule.py | 2 +-
python/tvm/s_tir/dlight/base/transform.py | 2 +-
python/tvm/s_tir/dlight/base/utils.py | 4 +-
python/tvm/s_tir/dlight/benchmark/bench.py | 12 +-
python/tvm/s_tir/dlight/benchmark/extract.py | 7 +-
python/tvm/s_tir/dlight/benchmark/utils.py | 2 +-
python/tvm/s_tir/dlight/cpu/gemv.py | 2 +-
python/tvm/s_tir/dlight/cpu/reduction.py | 2 +-
python/tvm/s_tir/dlight/gpu/fallback.py | 4 +-
python/tvm/s_tir/dlight/gpu/gemv.py | 4 +-
python/tvm/s_tir/dlight/gpu/general_reduction.py | 2 +-
python/tvm/s_tir/dlight/gpu/low_batch_gemv.py | 4 +-
python/tvm/s_tir/dlight/gpu/matmul.py | 4 +-
python/tvm/s_tir/dlight/gpu/reduction.py | 4 +-
python/tvm/s_tir/dlight/gpu/rmsnorm.py | 5 +-
python/tvm/s_tir/dlight/gpu/transpose.py | 6 +-
.../s_tir/meta_schedule/builder/local_builder.py | 4 +-
.../s_tir/meta_schedule/cost_model/mlp_model.py | 4 +-
.../s_tir/meta_schedule/cost_model/xgb_model.py | 2 +-
.../tvm/s_tir/meta_schedule/database/database.py | 2 +-
.../feature_extractor/random_feature_extractor.py | 2 +-
python/tvm/s_tir/meta_schedule/mutator/mutator.py | 2 +-
.../tvm/s_tir/meta_schedule/postproc/postproc.py | 2 +-
.../tvm/s_tir/meta_schedule/relax_integration.py | 6 +-
.../tvm/s_tir/meta_schedule/runner/local_runner.py | 4 +-
.../tvm/s_tir/meta_schedule/runner/rpc_runner.py | 2 +-
python/tvm/s_tir/meta_schedule/runner/utils.py | 4 +-
.../schedule/cuda/layout_transform.py | 2 +-
.../meta_schedule/schedule_rule/schedule_rule.py | 2 +-
.../search_strategy/search_strategy.py | 2 +-
.../space_generator/space_generator.py | 2 +-
.../meta_schedule/testing/custom_builder_runner.py | 2 +-
.../meta_schedule/testing/space_generation.py | 2 +-
.../tvm/s_tir/meta_schedule/testing/te_workload.py | 2 +-
python/tvm/s_tir/meta_schedule/testing/tune_te.py | 2 +-
.../tvm/s_tir/meta_schedule/testing/tune_utils.py | 2 +-
.../meta_schedule/testing/validate_database.py | 8 +-
python/tvm/s_tir/meta_schedule/tir_integration.py | 2 +-
python/tvm/s_tir/meta_schedule/trace_apply.py | 2 +-
python/tvm/s_tir/meta_schedule/tune_context.py | 2 +-
python/tvm/s_tir/meta_schedule/utils.py | 2 +-
python/tvm/s_tir/pipeline.py | 4 +-
python/tvm/s_tir/sblock_dependence_info.py | 2 +-
python/tvm/s_tir/sblock_scope.py | 2 +-
python/tvm/s_tir/schedule/schedule.py | 2 +-
python/tvm/s_tir/schedule/state.py | 2 +-
python/tvm/s_tir/schedule/testing.py | 4 +-
python/tvm/s_tir/schedule/trace.py | 2 +-
python/tvm/s_tir/script/ir_builder/ir.py | 6 +-
python/tvm/s_tir/stmt.py | 2 +-
python/tvm/s_tir/testing/runner.py | 4 +-
python/tvm/s_tir/testing/tir.py | 3 +-
python/tvm/te/operation.py | 11 +-
python/tvm/te/reduction.py | 2 +-
python/tvm/te/tag.py | 4 +-
python/tvm/topi/_te_tensor_overload.py | 2 +-
python/tvm/topi/gpu/scan.py | 9 +-
python/tvm/topi/gpu/scatter_elements.py | 4 +-
python/tvm/topi/gpu/scatter_nd.py | 4 +-
python/tvm/topi/gpu/sort.py | 35 ++-
python/tvm/topi/image/grid_sample.py | 2 +-
python/tvm/topi/image/resize.py | 3 +-
python/tvm/topi/index_put.py | 2 +-
python/tvm/topi/math.py | 5 +-
python/tvm/topi/nn/batch_matmul.py | 2 +-
python/tvm/topi/nn/bitserial_conv2d.py | 2 +-
python/tvm/topi/nn/bitserial_dense.py | 6 +-
python/tvm/topi/nn/bitserial_util.py | 2 +-
python/tvm/topi/nn/bnn.py | 2 +-
python/tvm/topi/nn/conv2d.py | 2 +-
python/tvm/topi/nn/deformable_conv2d.py | 2 +-
python/tvm/topi/nn/dense.py | 9 +-
python/tvm/topi/nn/depth_to_space.py | 2 +-
python/tvm/topi/nn/depthwise_conv2d.py | 2 +-
python/tvm/topi/nn/dilate.py | 2 +-
python/tvm/topi/nn/elemwise.py | 3 +-
python/tvm/topi/nn/fifo_buffer.py | 3 +-
python/tvm/topi/nn/flatten.py | 3 +-
python/tvm/topi/nn/lstm.py | 2 +-
python/tvm/topi/nn/mapping.py | 2 +-
python/tvm/topi/nn/pad.py | 11 +-
python/tvm/topi/nn/pixel_shuffle.py | 3 +-
python/tvm/topi/nn/qnn.py | 5 +-
python/tvm/topi/nn/softmax.py | 2 +-
python/tvm/topi/nn/space_to_depth.py | 2 +-
python/tvm/topi/nn/utils.py | 3 +-
python/tvm/topi/scan.py | 4 +-
python/tvm/topi/scatter.py | 2 +-
python/tvm/topi/scatter_elements.py | 2 +-
python/tvm/topi/signal.py | 2 +-
python/tvm/topi/sort.py | 19 +-
python/tvm/topi/sparse_reshape.py | 2 +-
python/tvm/topi/testing/poolnd_python.py | 2 +-
python/tvm/topi/transform.py | 11 +-
python/tvm/topi/unique.py | 2 +-
python/tvm/topi/utils.py | 7 +-
python/tvm/topi/vision/multibox_transform_loc.py | 2 +-
python/tvm/topi/vision/nms.py | 11 +-
python/tvm/topi/vision/nms_util.py | 3 +-
python/tvm/topi/vision/roi_align.py | 2 +-
python/tvm/topi/vision/roi_pool.py | 10 +-
tests/nightly/python/test_nnapi/infrastructure.py | 3 +-
.../test_nnapi/test_from_exported_to_cuda.py | 9 +-
tests/nightly/python/test_nnapi/test_network.py | 7 +-
tests/nightly/python/test_nnapi/test_ops.py | 5 +-
.../test_minimal_target_codegen_llvm.py | 5 +-
tests/python/codegen/llvm/test_cross_compile.py | 5 +-
tests/python/contrib/pickle_memoize_script.py | 2 +-
.../python/contrib/test_android/infrastructure.py | 3 +-
.../contrib/test_android/test_meta_schedule.py | 5 +-
tests/python/contrib/test_cblas.py | 5 +-
tests/python/contrib/test_coreml_runtime.py | 6 +-
tests/python/contrib/test_cutlass_gemm.py | 4 +-
tests/python/contrib/test_dlpack.py | 5 +-
tests/python/contrib/test_example_npu.py | 7 +-
tests/python/contrib/test_hipblas.py | 5 +-
tests/python/contrib/test_memoize.py | 2 +-
tests/python/contrib/test_random.py | 5 +-
tests/python/contrib/test_rpc_proxy.py | 2 +-
tests/python/contrib/test_rpc_tracker.py | 2 +-
tests/python/contrib/test_sort.py | 5 +-
.../python/contrib/test_tir_triton_integration.py | 9 +-
tests/python/contrib/test_tvmjs.py | 2 +-
tests/python/disco/test_callback.py | 7 +-
tests/python/disco/test_ccl.py | 9 +-
tests/python/disco/test_custom_allreduce.py | 4 +-
tests/python/disco/test_loader.py | 7 +-
tests/python/disco/test_nvshmem.py | 16 +-
tests/python/disco/test_session.py | 9 +-
tests/python/driver/test_compile.py | 9 +-
tests/python/relax/backend/adreno/mod_utils.py | 9 +-
tests/python/relax/backend/adreno/test_clml_ops.py | 7 +-
.../relax/backend/adreno/test_texture_network.py | 7 +-
.../relax/backend/adreno/test_texture_ops.py | 5 +-
.../adreno/test_transform_annotate_custom_scope.py | 11 +-
.../test_transform_fold_vdevice_scope_change.py | 7 +-
tests/python/relax/backend/adreno/utils.py | 5 +-
tests/python/relax/conftest.py | 2 +-
.../distributed/test_distributed_dtensor_type.py | 4 +-
..._distributed_transform_legalize_redistribute.py | 4 +-
.../test_distributed_transform_lower_distir.py | 7 +-
...ributed_transform_lower_global_to_local_view.py | 7 +-
...est_distributed_transform_propagate_sharding.py | 6 +-
.../test_runtime_builtin_kv_cache_transfer.py | 5 +-
...est_runtime_builtin_kv_cache_transfer_kernel.py | 6 +-
.../relax/script/test_relax_script_basic_usage.py | 13 +-
.../relax/script/test_relax_script_distributed.py | 15 +-
.../script/test_relax_script_dynamic_shape.py | 17 +-
.../script/test_relax_script_error_handling.py | 13 +-
.../relax/script/test_relax_script_ir_builder.py | 10 +-
.../script/test_relax_script_meta_programming.py | 11 +-
.../test_relax_script_operators_arith_cmp.py | 9 +-
.../script/test_relax_script_operators_create.py | 9 +-
.../script/test_relax_script_operators_datatype.py | 9 +-
.../script/test_relax_script_operators_grad.py | 7 +-
.../script/test_relax_script_operators_image.py | 9 +-
.../script/test_relax_script_operators_index.py | 9 +-
.../test_relax_script_operators_linear_algebra.py | 9 +-
.../test_relax_script_operators_manipulate.py | 11 +-
.../relax/script/test_relax_script_operators_nn.py | 9 +-
.../script/test_relax_script_operators_search.py | 9 +-
.../script/test_relax_script_operators_set.py | 9 +-
.../script/test_relax_script_operators_sort.py | 9 +-
.../test_relax_script_operators_statistical.py | 9 +-
.../script/test_relax_script_operators_vision.py | 9 +-
.../relax/script/test_relax_script_pyfunc.py | 6 +-
tests/python/relax/test_analysis.py | 8 +-
.../test_analysis_computable_at_compile_time.py | 5 +-
.../relax/test_analysis_contains_impure_call.py | 6 +-
.../python/relax/test_analysis_detect_recursion.py | 6 +-
.../relax/test_analysis_estimate_memory_usage.py | 6 +-
.../test_analysis_suggest_layout_transforms.py | 5 +-
tests/python/relax/test_analysis_type_analysis.py | 17 +-
tests/python/relax/test_analysis_well_formed.py | 11 +-
tests/python/relax/test_ast_printer.py | 8 +-
.../python/relax/test_backend_dispatch_sampling.py | 12 +-
.../relax/test_backend_dispatch_sort_scan.py | 13 +-
.../relax/test_backend_transform_shape_lower.py | 10 +-
tests/python/relax/test_base_py_module_printer.py | 9 +-
.../relax/test_base_py_module_symbolic_shape.py | 8 +-
tests/python/relax/test_bind_params.py | 8 +-
tests/python/relax/test_bind_symbolic_vars.py | 4 +-
tests/python/relax/test_binding_rewrite.py | 6 +-
tests/python/relax/test_blockbuilder_core.py | 11 +-
tests/python/relax/test_blockbuilder_emit_te.py | 5 +-
tests/python/relax/test_codegen_coreml.py | 7 +-
tests/python/relax/test_codegen_cublas.py | 7 +-
tests/python/relax/test_codegen_cudnn.py | 7 +-
tests/python/relax/test_codegen_cutlass.py | 11 +-
tests/python/relax/test_codegen_dnnl.py | 7 +-
tests/python/relax/test_codegen_hipblas.py | 7 +-
tests/python/relax/test_codegen_tensorrt.py | 7 +-
tests/python/relax/test_contrib_vllm.py | 11 +-
tests/python/relax/test_dataflow_inplace.py | 7 +-
tests/python/relax/test_dataflow_pattern.py | 9 +-
tests/python/relax/test_dataflow_rewriter.py | 4 +-
tests/python/relax/test_dlpack_integration.py | 4 +-
tests/python/relax/test_e2e_op_dynamic.py | 9 +-
tests/python/relax/test_expr.py | 19 +-
tests/python/relax/test_expr_functor.py | 13 +-
tests/python/relax/test_fast_math_transform.py | 7 +-
tests/python/relax/test_frontend_common.py | 6 +-
tests/python/relax/test_frontend_dynamo.py | 11 +-
.../relax/test_frontend_from_exported_program.py | 7 +-
tests/python/relax/test_frontend_from_fx.py | 6 +-
tests/python/relax/test_frontend_nn_debug.py | 6 +-
tests/python/relax/test_frontend_nn_exporter.py | 10 +-
.../python/relax/test_frontend_nn_extern_module.py | 9 +-
tests/python/relax/test_frontend_nn_jit.py | 6 +-
.../relax/test_frontend_nn_llm_kernel_config.py | 5 +-
...test_frontend_nn_llm_sequence_prefill_masked.py | 4 +-
tests/python/relax/test_frontend_nn_modules.py | 9 +-
tests/python/relax/test_frontend_nn_mutator.py | 2 +-
tests/python/relax/test_frontend_nn_op.py | 17 +-
tests/python/relax/test_frontend_nn_packing.py | 4 +-
.../relax/test_frontend_nn_parameter_containers.py | 12 +-
tests/python/relax/test_frontend_nn_subroutines.py | 8 +-
tests/python/relax/test_frontend_nn_tensor.py | 7 +-
tests/python/relax/test_frontend_onnx.py | 15 +-
tests/python/relax/test_frontend_onnx_backend.py | 3 +-
tests/python/relax/test_frontend_stablehlo.py | 7 +-
tests/python/relax/test_frontend_tflite.py | 31 ++-
tests/python/relax/test_group_gemm_flashinfer.py | 9 +-
tests/python/relax/test_inline_functions.py | 6 +-
tests/python/relax/test_kill_after_last_use.py | 7 +-
tests/python/relax/test_lower_alloc_tensor.py | 6 +-
.../relax/test_meta_schedule_relax_integration.py | 11 +-
tests/python/relax/test_nn_llm_prefill_parallel.py | 5 +-
tests/python/relax/test_op_binary.py | 8 +-
tests/python/relax/test_op_ccl.py | 9 +-
tests/python/relax/test_op_create.py | 8 +-
tests/python/relax/test_op_datatype.py | 8 +-
tests/python/relax/test_op_distributed.py | 8 +-
tests/python/relax/test_op_grad.py | 6 +-
tests/python/relax/test_op_gradient_numeric.py | 9 +-
tests/python/relax/test_op_image.py | 9 +-
tests/python/relax/test_op_index.py | 11 +-
tests/python/relax/test_op_inspect.py | 5 +-
tests/python/relax/test_op_linear_algebra.py | 8 +-
tests/python/relax/test_op_manipulate.py | 8 +-
tests/python/relax/test_op_misc.py | 4 +-
tests/python/relax/test_op_nn.py | 8 +-
tests/python/relax/test_op_nn_convolution.py | 15 +-
tests/python/relax/test_op_nn_pooling.py | 24 +-
tests/python/relax/test_op_qdq.py | 8 +-
tests/python/relax/test_op_sampling.py | 4 +-
tests/python/relax/test_op_search.py | 8 +-
tests/python/relax/test_op_set.py | 8 +-
tests/python/relax/test_op_size.py | 6 +-
tests/python/relax/test_op_sort.py | 8 +-
tests/python/relax/test_op_statistical.py | 8 +-
tests/python/relax/test_op_take.py | 7 +-
tests/python/relax/test_op_ternary.py | 8 +-
tests/python/relax/test_op_unary.py | 8 +-
tests/python/relax/test_op_view.py | 7 +-
tests/python/relax/test_op_vision.py | 11 +-
.../python/relax/test_optimize_layout_transform.py | 9 +-
tests/python/relax/test_pipeline.py | 13 +-
tests/python/relax/test_pytorch_integration.py | 4 +-
tests/python/relax/test_relax_operators.py | 35 ++-
.../python/relax/test_relax_to_pyfunc_converter.py | 7 +-
.../python/relax/test_remove_redundant_reshape.py | 5 +-
tests/python/relax/test_runtime_builtin.py | 7 +-
...runtime_builtin_paged_attention_kv_cache_cpu.py | 16 +-
..._builtin_paged_attention_kv_cache_flashinfer.py | 4 +-
...me_builtin_paged_attention_kv_cache_metadata.py | 14 +-
...ltin_paged_attention_kv_cache_mla_flashinfer.py | 12 +-
...ime_builtin_paged_attention_kv_cache_mla_tir.py | 12 +-
...runtime_builtin_paged_attention_kv_cache_tir.py | 4 +-
.../python/relax/test_runtime_builtin_rnn_state.py | 6 +-
.../relax/test_runtime_sampling_flashinfer.py | 8 +-
.../relax/test_tensorrt_engine_serialization.py | 7 +-
tests/python/relax/test_testing_nn.py | 6 +-
tests/python/relax/test_tir_call_source_kernel.py | 7 +-
tests/python/relax/test_training_append_loss.py | 6 +-
tests/python/relax/test_training_loss.py | 4 +-
tests/python/relax/test_training_optimizer.py | 6 +-
.../relax/test_training_optimizer_numeric.py | 7 +-
tests/python/relax/test_training_setup_trainer.py | 8 +-
.../python/relax/test_training_trainer_numeric.py | 7 +-
tests/python/relax/test_transform.py | 14 +-
.../relax/test_transform_allocate_workspace.py | 6 +-
tests/python/relax/test_transform_alter_op_impl.py | 6 +-
.../test_transform_annotate_tir_op_pattern.py | 6 +-
...st_transform_attach_attr_layout_free_tensors.py | 9 +-
.../relax/test_transform_attach_global_symbol.py | 11 +-
tests/python/relax/test_transform_bind_params.py | 11 +-
.../relax/test_transform_bind_symbolic_vars.py | 8 +-
.../relax/test_transform_bundle_model_params.py | 7 +-
.../relax/test_transform_canonicalize_bindings.py | 12 +-
tests/python/relax/test_transform_codegen_pass.py | 11 +-
.../test_transform_combine_parallel_matmul.py | 7 +-
.../relax/test_transform_compute_prim_value.py | 7 +-
.../relax/test_transform_convert_dataflow.py | 6 +-
.../python/relax/test_transform_convert_layout.py | 5 +-
tests/python/relax/test_transform_cse.py | 5 +-
.../relax/test_transform_dead_code_elimination.py | 5 +-
tests/python/relax/test_transform_decompose_ops.py | 11 +-
.../relax/test_transform_error_enrichment.py | 4 +-
.../relax/test_transform_expand_matmul_of_sum.py | 4 +-
.../relax/test_transform_expand_tuple_args.py | 7 +-
.../test_transform_fold_batch_norm_to_conv2d.py | 9 +-
tests/python/relax/test_transform_fold_constant.py | 11 +-
tests/python/relax/test_transform_fuse_ops.py | 11 +-
.../relax/test_transform_fuse_ops_by_pattern.py | 7 +-
tests/python/relax/test_transform_fuse_tir.py | 7 +-
.../relax/test_transform_fuse_transpose_matmul.py | 6 +-
tests/python/relax/test_transform_gradient.py | 6 +-
.../relax/test_transform_gradient_checkpoint.py | 6 +-
.../relax/test_transform_gradient_numeric.py | 5 +-
.../relax/test_transform_gradient_te_register.py | 11 +-
.../test_transform_inline_private_functions.py | 9 +-
.../relax/test_transform_ipc_allreduce_rewrite.py | 14 +-
tests/python/relax/test_transform_lambda_lift.py | 12 +-
.../relax/test_transform_lazy_transform_params.py | 9 +-
tests/python/relax/test_transform_legalize_ops.py | 9 +-
.../relax/test_transform_legalize_ops_binary.py | 8 +-
.../relax/test_transform_legalize_ops_ccl.py | 8 +-
.../test_transform_legalize_ops_create_datatype.py | 6 +-
.../test_transform_legalize_ops_distributed.py | 9 +-
.../relax/test_transform_legalize_ops_grad.py | 8 +-
.../relax/test_transform_legalize_ops_image.py | 6 +-
..._transform_legalize_ops_index_linear_algebra.py | 10 +-
.../test_transform_legalize_ops_manipulate.py | 10 +-
.../python/relax/test_transform_legalize_ops_nn.py | 10 +-
.../relax/test_transform_legalize_ops_qdq.py | 6 +-
...st_transform_legalize_ops_search_statistical.py | 10 +-
.../relax/test_transform_legalize_ops_unary.py | 9 +-
.../relax/test_transform_lift_transform_params.py | 7 +-
.../test_transform_lower_gpu_ipc_alloc_storage.py | 6 +-
.../test_transform_merge_composite_functions.py | 4 +-
.../test_transform_meta_schedule_apply_database.py | 9 +-
.../relax/test_transform_meta_schedule_tuning.py | 9 +-
tests/python/relax/test_transform_normalize.py | 8 +-
.../relax/test_transform_normalize_global_var.py | 13 +-
...st_transform_operator_specific_normalization.py | 5 +-
.../python/relax/test_transform_realize_vdevice.py | 5 +-
.../relax/test_transform_remove_unused_outputs.py | 7 +-
.../test_transform_remove_unused_parameters.py | 7 +-
..._transform_reorder_permute_dims_after_concat.py | 4 +-
.../test_transform_reorder_take_after_matmul.py | 4 +-
.../relax/test_transform_rewrite_cuda_graph.py | 6 +-
.../test_transform_rewrite_dataflow_reshape.py | 6 +-
...nsform_specialize_function_based_on_callsite.py | 9 +-
.../test_transform_split_layout_rewrite_preproc.py | 6 +-
.../test_transform_static_plan_block_memory.py | 6 +-
.../relax/test_transform_to_mixed_precision.py | 4 +-
.../relax/test_transform_topological_sort.py | 7 +-
.../relax/test_transform_update_param_type.py | 4 +-
.../python/relax/test_transform_update_vdevice.py | 5 +-
tests/python/relax/test_type.py | 4 +-
tests/python/relax/test_utils.py | 2 +-
.../relax/test_vm_alloc_storage_with_scope.py | 9 +-
tests/python/relax/test_vm_build.py | 31 ++-
tests/python/relax/test_vm_builtin.py | 9 +-
tests/python/relax/test_vm_builtin_lower.py | 6 +-
tests/python/relax/test_vm_callback_function.py | 9 +-
tests/python/relax/test_vm_codegen_only.py | 8 +-
tests/python/relax/test_vm_cuda_graph.py | 11 +-
tests/python/relax/test_vm_execbuilder.py | 5 +-
tests/python/relax/test_vm_instrument.py | 5 +-
tests/python/relax/test_vm_multi_device.py | 5 +-
tests/python/relax/texture/test_texture_nd.py | 9 +-
tests/python/runtime/rpc/test_rpc_base.py | 2 +-
tests/python/runtime/rpc/test_runtime_rpc.py | 9 +-
...st_s_tir_analysis_calculate_allocated_memory.py | 5 +-
.../test_s_tir_analysis_estimate_tir_flops.py | 7 +-
.../test_s_tir_analysis_identify_memcpy.py | 7 +-
.../test_s_tir_analysis_is_pure_function.py | 4 +-
.../s_tir/analysis/test_s_tir_analysis_oob.py | 3 +-
.../test_s_tir_analysis_verify_block_scope.py | 9 +-
.../s_tir/analysis/test_sblock_access_region.py | 7 +-
.../analysis/test_sblock_tensor_access_lca.py | 2 +-
.../python/s_tir/base/test_compilation_pipeline.py | 4 +-
.../python/s_tir/base/test_s_tir_texture_scope.py | 13 +-
.../s_tir/base/test_sblock_dependence_info.py | 11 +-
tests/python/s_tir/base/test_tir_data_layout.py | 7 +-
.../s_tir/base/test_tir_te_extern_function.py | 9 +-
tests/python/s_tir/dlight/test_benchmark.py | 6 +-
tests/python/s_tir/dlight/test_cpu_gemv.py | 3 +-
tests/python/s_tir/dlight/test_cpu_reduction.py | 6 +-
tests/python/s_tir/dlight/test_gpu_conv.py | 4 +-
tests/python/s_tir/dlight/test_gpu_fallback.py | 6 +-
tests/python/s_tir/dlight/test_gpu_gemv.py | 5 +-
.../s_tir/dlight/test_gpu_general_reduction.py | 8 +-
.../python/s_tir/dlight/test_gpu_low_batch_gemv.py | 4 +-
tests/python/s_tir/dlight/test_gpu_matmul.py | 4 +-
.../s_tir/dlight/test_gpu_matmul_tensorize.py | 4 +-
tests/python/s_tir/dlight/test_gpu_reduction.py | 6 +-
tests/python/s_tir/dlight/test_gpu_rmsnorm.py | 6 +-
tests/python/s_tir/dlight/test_gpu_transpose.py | 7 +-
tests/python/s_tir/dlight/test_primitives.py | 5 +-
.../meta_schedule/test_meta_schedule_arg_info.py | 2 +-
.../meta_schedule/test_meta_schedule_builder.py | 9 +-
.../meta_schedule/test_meta_schedule_cost_model.py | 7 +-
.../meta_schedule/test_meta_schedule_database.py | 11 +-
.../test_meta_schedule_feature_extractor.py | 2 +-
...schedule_feature_extractor_per_store_feature.py | 7 +-
.../test_meta_schedule_measure_callback.py | 4 +-
.../test_meta_schedule_post_order_apply.py | 9 +-
...ule_postproc_disallow_async_strided_mem_copy.py | 5 +-
...meta_schedule_postproc_disallow_dynamic_loop.py | 5 +-
..._schedule_postproc_rewrite_cooperative_fetch.py | 11 +-
.../test_meta_schedule_postproc_rewrite_layout.py | 7 +-
...e_postproc_rewrite_parallel_vectorize_unroll.py | 7 +-
...ta_schedule_postproc_rewrite_reduction_block.py | 5 +-
...est_meta_schedule_postproc_rewrite_tensorize.py | 5 +-
...meta_schedule_postproc_rewrite_unbound_block.py | 9 +-
.../test_meta_schedule_postproc_verify_gpu_code.py | 7 +-
...est_meta_schedule_postproc_verify_vtcm_limit.py | 7 +-
.../meta_schedule/test_meta_schedule_runner.py | 9 +-
...test_meta_schedule_schedule_rule_add_rfactor.py | 4 +-
...eta_schedule_schedule_rule_apply_custom_rule.py | 5 +-
.../test_meta_schedule_schedule_rule_auto_bind.py | 2 +-
...test_meta_schedule_schedule_rule_auto_inline.py | 9 +-
...chedule_schedule_rule_cross_thread_reduction.py | 6 +-
.../test_meta_schedule_schedule_rule_mlt.py | 5 +-
.../test_meta_schedule_schedule_rule_mlt_intrin.py | 4 +-
...dule_schedule_rule_parallel_vectorize_unroll.py | 4 +-
...hedule_schedule_rule_random_compute_location.py | 4 +-
.../test_meta_schedule_search_strategy.py | 6 +-
.../meta_schedule/test_meta_schedule_space_cpu.py | 2 +-
.../meta_schedule/test_meta_schedule_space_cuda.py | 2 +-
.../test_meta_schedule_space_cuda_async.py | 2 +-
...ule_space_cuda_async_multiple_initialization.py | 2 +-
.../test_meta_schedule_space_generator.py | 7 +-
.../test_meta_schedule_space_post_opt.py | 7 +-
.../test_meta_schedule_task_scheduler.py | 6 +-
.../test_meta_schedule_trace_apply.py | 9 +-
.../test_meta_schedule_tune_context.py | 7 +-
.../meta_schedule/test_meta_schedule_tune_tir.py | 7 +-
.../s_tir/schedule/test_tir_schedule_analysis.py | 13 +-
.../test_tir_schedule_annotate_tensor_access.py | 9 +-
.../schedule/test_tir_schedule_block_scope.py | 9 +-
.../s_tir/schedule/test_tir_schedule_blockize.py | 9 +-
.../schedule/test_tir_schedule_cache_index.py | 9 +-
.../schedule/test_tir_schedule_cache_read_write.py | 9 +-
.../s_tir/schedule/test_tir_schedule_compute_at.py | 9 +-
.../schedule/test_tir_schedule_compute_inline.py | 9 +-
.../test_tir_schedule_decompose_padding.py | 9 +-
.../s_tir/schedule/test_tir_schedule_error.py | 7 +-
.../s_tir/schedule/test_tir_schedule_for_kind.py | 9 +-
.../test_tir_schedule_fuse_reduction_epilogue.py | 19 +-
...ir_schedule_fuse_reduction_epilogue_clipping.py | 9 +-
...st_tir_schedule_fuse_reduction_epilogue_relu.py | 9 +-
.../schedule/test_tir_schedule_instruction.py | 3 +-
.../s_tir/schedule/test_tir_schedule_merge.py | 9 +-
.../s_tir/schedule/test_tir_schedule_pad_einsum.py | 9 +-
.../s_tir/schedule/test_tir_schedule_partition.py | 9 +-
.../schedule/test_tir_schedule_read_write_at.py | 9 +-
.../s_tir/schedule/test_tir_schedule_reduction.py | 9 +-
.../s_tir/schedule/test_tir_schedule_reindex.py | 9 +-
.../s_tir/schedule/test_tir_schedule_reorder.py | 9 +-
.../test_tir_schedule_reorder_block_iter_var.py | 9 +-
.../s_tir/schedule/test_tir_schedule_rfactor.py | 9 +-
.../schedule/test_tir_schedule_rolling_buffer.py | 9 +-
.../s_tir/schedule/test_tir_schedule_sampling.py | 7 +-
.../s_tir/schedule/test_tir_schedule_set_dtype.py | 9 +-
.../s_tir/schedule/test_tir_schedule_set_scope.py | 9 +-
.../s_tir/schedule/test_tir_schedule_split_fuse.py | 9 +-
.../s_tir/schedule/test_tir_schedule_state.py | 7 +-
.../test_tir_schedule_state_cached_flags.py | 9 +-
.../schedule/test_tir_schedule_storage_align.py | 9 +-
.../s_tir/schedule/test_tir_schedule_tensorize.py | 9 +-
.../test_tir_schedule_tensorize_mfma_numeric.py | 5 +-
.../s_tir/schedule/test_tir_schedule_trace.py | 14 +-
.../s_tir/schedule/test_tir_schedule_transform.py | 4 +-
.../schedule/test_tir_schedule_transform_layout.py | 12 +-
.../s_tir/schedule/test_tir_schedule_utilities.py | 11 +-
.../schedule/test_tir_unsafe_hide_tensor_access.py | 9 +-
.../s_tir/script/test_s_tir_script_basic_usage.py | 7 +-
.../s_tir/script/test_s_tir_script_blocks.py | 12 +-
.../script/test_s_tir_script_dynamic_shape.py | 5 +-
.../script/test_s_tir_script_error_handling.py | 7 +-
.../s_tir/script/test_s_tir_script_ir_builder.py | 12 +-
.../script/test_s_tir_script_meta_programming.py | 14 +-
.../s_tir/script/test_s_tir_script_printer.py | 24 +-
.../script/test_s_tir_script_source_locations.py | 4 +-
tests/python/s_tir/test_arith_domain_touched.py | 3 +-
tests/python/s_tir/test_s_tir_renew_defs.py | 6 +-
...test_s_tir_transform_annotate_irregular_loop.py | 13 +-
.../test_s_tir_transform_canonicalize_loop.py | 5 +-
.../test_s_tir_transform_compact_tensor_region.py | 7 +-
...est_s_tir_transform_convert_blocks_to_opaque.py | 9 +-
.../test_s_tir_transform_default_gpu_schedule.py | 6 +-
.../test_s_tir_transform_inject_permuted_layout.py | 9 +-
.../test_s_tir_transform_inject_ptx_async_copy.py | 10 +-
...est_s_tir_transform_inject_software_pipeline.py | 9 +-
.../test_s_tir_transform_inject_virtual_thread.py | 11 +-
.../test_s_tir_transform_lift_thread_binding.py | 5 +-
.../test_s_tir_transform_loop_partition.py | 15 +-
...s_tir_transform_lower_cross_thread_reduction.py | 5 +-
.../test_s_tir_transform_lower_init_block.py | 3 +-
.../test_s_tir_transform_lower_match_tensor.py | 5 +-
.../test_s_tir_transform_lower_opaque_block.py | 5 +-
...test_s_tir_transform_lower_thread_all_reduce.py | 7 +-
.../test_s_tir_transform_lower_thread_binding.py | 5 +-
...transform_manifest_shared_memory_local_stage.py | 5 +-
...st_s_tir_transform_memhammer_lower_auto_copy.py | 9 +-
...form_merge_dynamic_shared_memory_allocations.py | 9 +-
.../test_s_tir_transform_mma_tensor_layout.py | 10 +-
...sform_plan_update_tensor_allocation_location.py | 9 +-
.../transform/test_s_tir_transform_remove_undef.py | 7 +-
...transform_remove_weight_layout_rewrite_block.py | 5 +-
.../test_s_tir_transform_rewrite_unsafe_select.py | 5 +-
.../transform/test_s_tir_transform_thread_sync.py | 7 +-
tests/python/script/test_constructor_contracts.py | 8 +-
tests/python/sym/test_sym_canonical_simplify.py | 48 +++-
tests/python/sym/test_sym_modular_set.py | 4 +-
tests/python/target/test_arm_target.py | 5 +-
tests/python/target/test_target_target.py | 17 +-
tests/python/te/test_te_create_function.py | 14 +-
tests/python/te/test_te_tag.py | 3 +-
tests/python/te/test_te_tensor.py | 5 +-
tests/python/tirx/base/test_tirx_base.py | 3 +-
.../base/test_tirx_binding_type_propagation.py | 16 +-
tests/python/tirx/base/test_tirx_constructor.py | 3 +-
tests/python/tirx/base/test_tirx_ops.py | 5 +-
tests/python/tirx/codegen/test_codegen_nvshmem.py | 6 +-
tests/python/tirx/codegen/test_tirx_codegen_fp8.py | 5 +-
.../tirx/script/test_tirx_script_ir_builder.py | 12 +-
.../python/tirx/script/test_tirx_script_printer.py | 24 +-
.../tvm_tirx/backend/cuda/tile/patterns/ewise.py | 12 +-
tirx/python/tvm_tirx/backend/trn/pipeline.py | 4 +-
tirx/python/tvm_tirx/ir/type.py | 4 +-
tirx/python/tvm_tirx/script/__init__.py | 5 +-
tirx/python/tvm_tirx/script/printer/scriptable.py | 4 +-
tirx/python/tvm_tirx/tirx/__init__.py | 8 +-
tirx/python/tvm_tirx/tirx/_compile.py | 4 +-
tirx/python/tvm_tirx/tirx/_tensor_view.py | 11 +-
tirx/python/tvm_tirx/tirx/compilation_pipeline.py | 4 +-
tirx/python/tvm_tirx/tirx/op/__init__.py | 6 +-
tirx/python/tvm_tirx/tirx/tensor.py | 4 +-
.../tvm_tirx/tirx/transform/function_pass.py | 3 +-
.../tests/python/backend/cuda/codegen/test_cuda.py | 48 +++-
.../python/backend/cuda/test_cuda_host_bundle.py | 6 +-
.../python/backend/cuda/test_cuda_module_export.py | 4 +-
.../python/backend/metal/codegen/test_metal.py | 3 +-
.../tests/python/backend/rocm/codegen/test_rocm.py | 4 +-
.../python/backend/webgpu/codegen/test_webgpu.py | 3 +-
tirx/tests/python/codegen/c/test_c.py | 4 +-
.../python/codegen/gpu_common/test_allreduce.py | 4 +-
tirx/tests/python/codegen/llvm/test_llvm.py | 4 +-
tirx/tests/python/ir/test_ir_json.py | 12 +-
tirx/tests/python/ir/test_nonnullable_ir.py | 7 +-
.../python/runtime/test_runtime_memory_scope.py | 4 +-
tirx/tests/python/runtime/test_runtime_module.py | 4 +-
tirx/tests/python/script/test_basic_usage.py | 4 +-
tirx/tests/python/sym/test_sym_analyzer_object.py | 4 +-
tirx/tests/python/sym/test_sym_deduce_bound.py | 16 +-
.../tests/python/sym/test_sym_detect_clip_bound.py | 8 +-
tirx/tests/python/sym/test_sym_intset.py | 26 +-
tirx/tests/python/sym/test_sym_iter_affine_map.py | 60 +++-
tirx/tests/python/sym/test_sym_rewrite_simplify.py | 303 ++++++++++++++++-----
tirx/tests/python/sym/test_sym_simplify.py | 6 +-
.../tests/python/target/test_llvm_features_info.py | 11 +-
tirx/tests/python/target/test_x86_features.py | 3 +-
tirx/tests/python/testing/test_env.py | 4 +-
.../test_tirx_analysis_check_well_defined.py | 6 +-
tirx/tests/python/tirx/base/test_tirx_intrin.py | 31 ++-
tirx/tests/python/tirx/base/test_tirx_nodes.py | 10 +-
.../python/tirx/base/test_tirx_tensor_type.py | 13 +-
.../tests/python/tirx/codegen/test_codegen_cuda.py | 16 +-
.../python/tirx/codegen/test_cuda_wait_until.py | 4 +-
tirx/tests/python/tirx/codegen/test_ptx_addr.py | 12 +-
tirx/tests/python/tirx/codegen/test_ptx_dialect.py | 16 +-
.../python/tirx/codegen/test_ptx_ld_st_ops.py | 24 +-
.../codegen/test_tcgen05_descriptor_encoders.py | 4 +-
tirx/tests/python/tirx/iket/test_iket_profiler.py | 8 +-
.../tirx/script/test_tirx_script_basic_usage.py | 8 +-
.../tirx/script/test_tirx_script_dynamic_shape.py | 22 +-
.../python/tirx/script/test_tirx_script_gpu.py | 4 +-
.../script/test_tirx_script_meta_programming.py | 8 +-
...test_tirx_script_printer_invisible_path_info.py | 4 +-
.../script/test_tirx_script_source_locations.py | 8 +-
.../python/tirx/script/test_tirx_script_tensors.py | 8 +-
tirx/tests/python/tirx/test_cuda_backend_config.py | 4 +-
.../tests/python/tirx/test_op_namespace_cleanup.py | 8 +-
.../tests/python/tirx/test_roundtrip_namespaces.py | 4 +-
tirx/tests/python/tirx/test_tile_namespace.py | 7 +-
tirx/tests/python/tirx/test_tirx_axis_enum.py | 6 +-
.../python/tirx/tile/cuda/test_add_sub_mul_div.py | 12 +-
tirx/tests/python/tirx/tile/cuda/test_cp_async.py | 4 +-
.../tirx/tile/cuda/test_cp_async_bulk_tensor.py | 16 +-
.../python/tirx/tile/cuda/test_ld_st_registers.py | 8 +-
tirx/tests/python/tirx/tile/cuda/test_mma_sync.py | 16 +-
.../python/tirx/tile/cuda/test_pattern_mma.py | 12 +-
.../tests/python/tirx/tile/cuda/test_tcgen05_cp.py | 8 +-
.../python/tirx/tile/cuda/test_tcgen05_ld_st.py | 16 +-
.../python/tirx/tile/cuda/test_tcgen05_mma.py | 8 +-
.../tirx/tile/cuda/test_tcgen05_mma_codegen.py | 4 +-
.../tirx/tile/cuda/test_tcgen05_mma_ws_split_n.py | 6 +-
tirx/tests/python/tirx/tile/trn/test_copy_trn.py | 5 +-
.../transform/test_tirx_transform_bf16_legalize.py | 4 +-
.../transform/test_tirx_transform_lower_tirx.py | 9 +-
.../test_tirx_transform_make_packed_api.py | 11 +-
.../test_tirx_transform_narrow_datatype.py | 4 +-
.../test_tirx_transform_split_host_device.py | 5 +-
web/tests/python/webgpu_rpc_test.py | 7 +-
728 files changed, 3238 insertions(+), 2051 deletions(-)
diff --git a/apps/android_rpc/tests/android_rpc_test.py b/apps/android_rpc/tests/android_rpc_test.py
index 7b41d41917..af450e2693 100644
--- a/apps/android_rpc/tests/android_rpc_test.py
+++ b/apps/android_rpc/tests/android_rpc_test.py
@@ -26,7 +26,8 @@ import os
import numpy as np
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import rpc, te
from tvm_tirx.support import ndk, utils
@@ -58,7 +59,9 @@ def test_rpc_module():
tracker = rpc.connect_tracker(tracker_host, tracker_port)
remote = tracker.request(key, priority=0, session_timeout=60)
- mod = tvm_tirx.IRModule.from_expr(te.create_function([A, B]).with_attr("global_symbol", "myadd"))
+ mod = tvm_tirx.IRModule.from_expr(
+ te.create_function([A, B]).with_attr("global_symbol", "myadd")
+ )
sch = tvm.s_tir.Schedule(mod)
(x,) = sch.get_loops(block=sch.get_sblock("B"))
xo, xi = sch.split(i, [None, 32])
diff --git a/apps/ios_rpc/tests/ios_rpc_test.py b/apps/ios_rpc/tests/ios_rpc_test.py
index fdabecab12..858dd4b0d5 100644
--- a/apps/ios_rpc/tests/ios_rpc_test.py
+++ b/apps/ios_rpc/tests/ios_rpc_test.py
@@ -24,7 +24,8 @@ import argparse
import numpy as np
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import rpc, te
from tvm_tirx.support import utils, xcode
@@ -48,7 +49,9 @@ def test_rpc_module(host, port, key, mode):
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
temp = utils.tempdir()
- mod = tvm_tirx.IRModule.from_expr(te.create_function([A, B]).with_attr("global_symbol", "myadd"))
+ mod = tvm_tirx.IRModule.from_expr(
+ te.create_function([A, B]).with_attr("global_symbol", "myadd")
+ )
sch = tvm.s_tir.Schedule(mod)
(i,) = sch.get_loops(block=sch.get_sblock("B"))
i0, i1 = sch.split(i, [None, 32])
diff --git a/jvm/core/src/test/scripts/prepare_test_libs.py b/jvm/core/src/test/scripts/prepare_test_libs.py
index 67df2fcacc..995c7904a6 100644
--- a/jvm/core/src/test/scripts/prepare_test_libs.py
+++ b/jvm/core/src/test/scripts/prepare_test_libs.py
@@ -19,7 +19,8 @@
import os
import sys
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import relax, te
from tvm_tirx.script import ir as I
from tvm_tirx.script import relax as R
@@ -52,7 +53,9 @@ def prepare_cpu_lib(base_path):
A = te.placeholder((n,), name="A")
B = te.placeholder((n,), name="B")
C = te.compute(A.shape, lambda *i: A(*i) + B(*i), name="C")
- mod = tvm_tirx.IRModule.from_expr(te.create_function([A, B, C]).with_attr("global_symbol", "myadd"))
+ mod = tvm_tirx.IRModule.from_expr(
+ te.create_function([A, B, C]).with_attr("global_symbol", "myadd")
+ )
fadd = tvm_tirx.tirx.compile(mod, target).mod
lib_path = os.path.join(base_path, "add_cpu.so")
fadd.export_library(lib_path)
@@ -66,7 +69,9 @@ def prepare_gpu_lib(base_path):
A = te.placeholder((n,), name="A")
B = te.placeholder((n,), name="B")
C = te.compute(A.shape, lambda *i: A(*i) + B(*i), name="C")
- mod = tvm_tirx.IRModule.from_expr(te.create_function([A, B, C]).with_attr("global_symbol", "myadd"))
+ mod = tvm_tirx.IRModule.from_expr(
+ te.create_function([A, B, C]).with_attr("global_symbol", "myadd")
+ )
sch = tvm.s_tir.Schedule(mod)
sch.work_on("myadd")
(i,) = sch.get_loops(block=sch.get_sblock("C"))
diff --git a/python/tvm/contrib/cblas.py b/python/tvm/contrib/cblas.py
index 9697f3acb0..2097ed7e81 100644
--- a/python/tvm/contrib/cblas.py
+++ b/python/tvm/contrib/cblas.py
@@ -16,7 +16,7 @@
# under the License.
"""External function interface to BLAS libraries."""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
diff --git a/python/tvm/contrib/cublas.py b/python/tvm/contrib/cublas.py
index c6840ab008..6ca052d2cf 100644
--- a/python/tvm/contrib/cublas.py
+++ b/python/tvm/contrib/cublas.py
@@ -16,7 +16,7 @@
# under the License.
"""External function interface to cuBLAS libraries."""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
diff --git a/python/tvm/contrib/cublaslt.py b/python/tvm/contrib/cublaslt.py
index 893ffc5366..cb40d79cb2 100644
--- a/python/tvm/contrib/cublaslt.py
+++ b/python/tvm/contrib/cublaslt.py
@@ -16,7 +16,7 @@
# under the License.
"""External function interface to cuBLASlt libraries."""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
diff --git a/python/tvm/contrib/cudnn.py b/python/tvm/contrib/cudnn.py
index 718cac7a24..d0fc7a1438 100644
--- a/python/tvm/contrib/cudnn.py
+++ b/python/tvm/contrib/cudnn.py
@@ -23,7 +23,7 @@ import ctypes
import numpy as np
import tvm_ffi
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
# algos can be read from cudnn.h
diff --git a/python/tvm/contrib/cutlass/build.py b/python/tvm/contrib/cutlass/build.py
index aeea9e5b0a..16abdb40eb 100644
--- a/python/tvm/contrib/cutlass/build.py
+++ b/python/tvm/contrib/cutlass/build.py
@@ -28,11 +28,11 @@ from functools import reduce
from tvm_ffi import register_global_func
-import tvm, tvm_tirx
-from tvm_tirx import runtime
+import tvm_tirx
from tvm import relax
-from tvm_tirx.support.nvcc import get_cuda_version
from tvm.topi.utils import get_const_tuple
+from tvm_tirx import runtime
+from tvm_tirx.support.nvcc import get_cuda_version
from .gen_conv2d import CutlassConv2DProfiler
from .gen_gemm import CutlassGemmProfiler
diff --git a/python/tvm/contrib/dlpack.py b/python/tvm/contrib/dlpack.py
index 2b30e8385b..c38621966e 100644
--- a/python/tvm/contrib/dlpack.py
+++ b/python/tvm/contrib/dlpack.py
@@ -16,7 +16,7 @@
# under the License.
"""Wrapping functions to bridge frameworks with DLPack support to TVM"""
-import tvm_tirx.runtime, tvm
+import tvm_tirx.runtime
def convert_func(tvm_func, tensor_type, to_dlpack_func):
@@ -38,7 +38,9 @@ def convert_func(tvm_func, tensor_type, to_dlpack_func):
def _wrapper(*args):
args = tuple(
- tvm_tirx.runtime.from_dlpack(to_dlpack_func(arg)) if isinstance(arg, tensor_type) else arg
+ tvm_tirx.runtime.from_dlpack(to_dlpack_func(arg))
+ if isinstance(arg, tensor_type)
+ else arg
for arg in args
)
return tvm_func(*args)
diff --git a/python/tvm/contrib/dnnl.py b/python/tvm/contrib/dnnl.py
index 38afaad543..ceb4fcfaac 100644
--- a/python/tvm/contrib/dnnl.py
+++ b/python/tvm/contrib/dnnl.py
@@ -16,7 +16,7 @@
# under the License.
"""External function interface to BLAS libraries."""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
from ..topi.nn.utils import get_pad_tuple
diff --git a/python/tvm/contrib/hexagon/generate_take_op.py b/python/tvm/contrib/hexagon/generate_take_op.py
index 28ac81c5c8..96e7d38e4e 100644
--- a/python/tvm/contrib/hexagon/generate_take_op.py
+++ b/python/tvm/contrib/hexagon/generate_take_op.py
@@ -16,8 +16,9 @@
# under the License.
# pylint: disable=missing-docstring, invalid-name, unnecessary-comprehension, unused-argument
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm.contrib.hexagon import hexagon_unary_ops
@@ -83,7 +84,9 @@ class Tanh2TakeReplace(tvm.relax.PyExprMutator):
take_func_gv = self.builder_.add_func(take_func, "take")
take_node = relax.call_tir(
take_func_gv,
- relax.expr.Tuple([call_node.args[1][0], relax.const(tvm_tirx.runtime.tensor(LUT))]),
+ relax.expr.Tuple(
+ [call_node.args[1][0], relax.const(tvm_tirx.runtime.tensor(LUT))]
+ ),
ty_args=[call_node.ty],
)
return take_node
diff --git a/python/tvm/contrib/hexagon/tools.py b/python/tvm/contrib/hexagon/tools.py
index 218e5efcad..97c92b3ceb 100644
--- a/python/tvm/contrib/hexagon/tools.py
+++ b/python/tvm/contrib/hexagon/tools.py
@@ -29,7 +29,7 @@ import tarfile
import numpy
from tvm_ffi import register_global_func
-import tvm, tvm_tirx
+import tvm_tirx
import tvm_tirx.support.cc as cc
from tvm_tirx.backend.hexagon.linker import (
HEXAGON_LINK_MAIN as HEXAGON_LINK_MAIN,
diff --git a/python/tvm/contrib/hipblas.py b/python/tvm/contrib/hipblas.py
index e0ad63f91e..f9529eadcf 100644
--- a/python/tvm/contrib/hipblas.py
+++ b/python/tvm/contrib/hipblas.py
@@ -16,7 +16,7 @@
# under the License.
"""External function interface to hipBLAS libraries."""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
diff --git a/python/tvm/contrib/mkl.py b/python/tvm/contrib/mkl.py
index cb565c5563..fc10051144 100644
--- a/python/tvm/contrib/mkl.py
+++ b/python/tvm/contrib/mkl.py
@@ -16,7 +16,7 @@
# under the License.
"""External function interface to BLAS libraries."""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
diff --git a/python/tvm/contrib/nnpack.py b/python/tvm/contrib/nnpack.py
index 9b6cc30246..7b27ccd46e 100644
--- a/python/tvm/contrib/nnpack.py
+++ b/python/tvm/contrib/nnpack.py
@@ -19,7 +19,7 @@
import tvm_ffi
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
diff --git a/python/tvm/contrib/random.py b/python/tvm/contrib/random.py
index 76462c6ecd..8b592ad65f 100644
--- a/python/tvm/contrib/random.py
+++ b/python/tvm/contrib/random.py
@@ -18,7 +18,7 @@
import tvm_ffi
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
diff --git a/python/tvm/contrib/tvmjs.py b/python/tvm/contrib/tvmjs.py
index b616d2b757..8a20a6bfad 100644
--- a/python/tvm/contrib/tvmjs.py
+++ b/python/tvm/contrib/tvmjs.py
@@ -36,7 +36,8 @@ try:
except ImportError:
ml_dtypes = None
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm_tirx.runtime import DataType
from .emcc import create_wasm as create_tvmjs_wasm
diff --git a/python/tvm/driver/build_module.py b/python/tvm/driver/build_module.py
index 1052c615aa..01bbf50ceb 100644
--- a/python/tvm/driver/build_module.py
+++ b/python/tvm/driver/build_module.py
@@ -20,7 +20,8 @@
from collections.abc import Callable
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm_tirx.ir.module import IRModule
from tvm_tirx.runtime import Executable
from tvm_tirx.target import Target
diff --git a/python/tvm/exec/disco_remote_socket_session.py b/python/tvm/exec/disco_remote_socket_session.py
index 3141157ffa..23f326cb21 100644
--- a/python/tvm/exec/disco_remote_socket_session.py
+++ b/python/tvm/exec/disco_remote_socket_session.py
@@ -20,7 +20,8 @@
import sys
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from . import disco_worker as _ # pylint: disable=unused-import
diff --git a/python/tvm/exec/disco_worker.py b/python/tvm/exec/disco_worker.py
index 11f97d298e..205691a2b6 100644
--- a/python/tvm/exec/disco_worker.py
+++ b/python/tvm/exec/disco_worker.py
@@ -25,7 +25,7 @@ from collections.abc import Callable
from tvm_ffi import Shape, get_global_func, register_global_func
from tvm_ffi.core import String
-import tvm, tvm_tirx
+import tvm_tirx
from tvm_tirx.runtime import Tensor, tensor
diff --git a/python/tvm/exec/gpu_memory_bandwidth.py b/python/tvm/exec/gpu_memory_bandwidth.py
index ffec0d618a..30bb1cb1a8 100644
--- a/python/tvm/exec/gpu_memory_bandwidth.py
+++ b/python/tvm/exec/gpu_memory_bandwidth.py
@@ -22,7 +22,8 @@ import itertools
import numpy as np
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import te
from tvm.s_tir.meta_schedule.runner import EvaluatorConfig, RPCConfig
from tvm.s_tir.testing.runner import local_run, rpc_run
diff --git a/python/tvm/relax/__init__.py b/python/tvm/relax/__init__.py
index c4f910a70b..123f2b887f 100644
--- a/python/tvm/relax/__init__.py
+++ b/python/tvm/relax/__init__.py
@@ -111,9 +111,12 @@ from .vm_build import build, VMExecutable
from .binding_rewrite import DataflowBlockRewrite
-import tvm_tirx.script, tvm
+import tvm_tirx.script
+import tvm
-tvm_tirx.script.register_dialect("relax", "tvm.relax.script", builder_path="tvm.relax.script.ir_builder")
+tvm_tirx.script.register_dialect(
+ "relax", "tvm.relax.script", builder_path="tvm.relax.script.ir_builder"
+)
def _check_script_module(module: "tvm_tirx.ir.IRModule") -> None:
diff --git a/python/tvm/relax/analysis/analysis.py b/python/tvm/relax/analysis/analysis.py
index b35b215068..5781dc7f2c 100644
--- a/python/tvm/relax/analysis/analysis.py
+++ b/python/tvm/relax/analysis/analysis.py
@@ -24,12 +24,13 @@ configuring the passes and scripting them in Python.
from collections.abc import Callable
from enum import IntEnum
-import tvm, tvm_tirx
-from tvm_tirx import IRModule, tirx
-from tvm_tirx.ir import Call, Type
+import tvm
+import tvm_tirx
from tvm.relax.expr import Binding, DataflowBlock, Expr, Function, GlobalVar
from tvm.relax.type import FuncType
from tvm.s_tir import SBlock
+from tvm_tirx import IRModule, tirx
+from tvm_tirx.ir import Call, Type
from tvm_tirx.tirx import IndexMap, Var
from . import _ffi_api
diff --git a/python/tvm/relax/analysis/estimate_memory_usage.py b/python/tvm/relax/analysis/estimate_memory_usage.py
index 59cf45add1..dd23be876e 100644
--- a/python/tvm/relax/analysis/estimate_memory_usage.py
+++ b/python/tvm/relax/analysis/estimate_memory_usage.py
@@ -17,7 +17,7 @@
# pylint: disable=abstract-method,unused-argument
# pylint: disable=missing-function-docstring,missing-module-docstring
-import tvm, tvm_tirx
+import tvm_tirx
from tvm_tirx.ir import Call, Op
from tvm_tirx.ir.module import IRModule
diff --git a/python/tvm/relax/backend/adreno/clml.py b/python/tvm/relax/backend/adreno/clml.py
index 69cd06e867..aca8a0b6ef 100644
--- a/python/tvm/relax/backend/adreno/clml.py
+++ b/python/tvm/relax/backend/adreno/clml.py
@@ -17,10 +17,8 @@
# pylint: disable=invalid-name, unused-argument, pointless-exception-statement
"""Pattern table for CLML backend"""
-import tvm, tvm_tirx
-from tvm_tirx import IRModule, tirx
+import tvm_tirx
from tvm import relax
-from tvm_tirx.ir.transform import PassContext, module_pass
from tvm.relax import transform
from tvm.relax.dpl.pattern import (
GlobalVarPattern,
@@ -33,6 +31,8 @@ from tvm.relax.dpl.pattern import (
from tvm.relax.expr import TupleGetItem, VarBinding
from tvm.relax.expr_functor import PyExprMutator, mutator
from tvm.relax.transform import PatternCheckContext
+from tvm_tirx import IRModule, tirx
+from tvm_tirx.ir.transform import PassContext, module_pass
from ..pattern_registry import register_patterns
diff --git a/python/tvm/relax/backend/adreno/pipeline.py b/python/tvm/relax/backend/adreno/pipeline.py
index 819e57a217..3061ea2bd9 100644
--- a/python/tvm/relax/backend/adreno/pipeline.py
+++ b/python/tvm/relax/backend/adreno/pipeline.py
@@ -16,7 +16,8 @@
# under the License.
"""The Relax Adreno GPU backend compilation pipeline and other passes."""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import relax
from tvm.relax.transform.legalize_ops import adreno as legalize_adreno
diff --git a/python/tvm/relax/backend/adreno/transform/transform.py b/python/tvm/relax/backend/adreno/transform/transform.py
index 5db8cea36b..220246b8d5 100644
--- a/python/tvm/relax/backend/adreno/transform/transform.py
+++ b/python/tvm/relax/backend/adreno/transform/transform.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name
"""Adreno Relax transformation passes."""
-import tvm_tirx.ir, tvm
+import tvm_tirx.ir
from tvm_tirx.target import Target
from . import _ffi_api
diff --git a/python/tvm/relax/backend/contrib/example_npu/patterns.py b/python/tvm/relax/backend/contrib/example_npu/patterns.py
index 37fb27a988..72545cc1a0 100644
--- a/python/tvm/relax/backend/contrib/example_npu/patterns.py
+++ b/python/tvm/relax/backend/contrib/example_npu/patterns.py
@@ -24,9 +24,9 @@ tiling, and fusion strategies.
from typing import ClassVar
-from tvm_tirx.ir import Op
from tvm.relax.dpl.pattern import is_op, wildcard
from tvm.relax.transform import PatternCheckContext
+from tvm_tirx.ir import Op
from ...pattern_registry import register_patterns
diff --git a/python/tvm/relax/backend/contrib/nnapi.py b/python/tvm/relax/backend/contrib/nnapi.py
index 3288239030..e3ccfc5bc2 100644
--- a/python/tvm/relax/backend/contrib/nnapi.py
+++ b/python/tvm/relax/backend/contrib/nnapi.py
@@ -19,13 +19,13 @@
from collections.abc import Mapping
-from tvm_tirx.ir import IRModule
from tvm.relax.dpl.pattern import (
DFPattern,
is_op,
wildcard,
)
from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions
+from tvm_tirx.ir import IRModule
from ..pattern_registry import get_patterns_with_prefix, register_patterns
diff --git a/python/tvm/relax/backend/contrib/tensorrt.py b/python/tvm/relax/backend/contrib/tensorrt.py
index 9f2b170c7a..20c5a76d98 100644
--- a/python/tvm/relax/backend/contrib/tensorrt.py
+++ b/python/tvm/relax/backend/contrib/tensorrt.py
@@ -23,8 +23,6 @@ converter registered under the same name (the converters are keyed by
out of the module and annotates them for the ``tensorrt`` codegen.
"""
-from tvm_tirx import tirx
-from tvm_tirx.ir import IRModule
from tvm.relax.dpl.pattern import is_op, wildcard
from tvm.relax.expr import ShapeExpr
from tvm.relax.transform import (
@@ -32,6 +30,8 @@ from tvm.relax.transform import (
MergeCompositeFunctions,
PatternCheckContext,
)
+from tvm_tirx import tirx
+from tvm_tirx.ir import IRModule
from ..pattern_registry import Pattern, get_patterns_with_prefix, register_patterns
diff --git a/python/tvm/relax/backend/cpu_generic/pipeline.py b/python/tvm/relax/backend/cpu_generic/pipeline.py
index dec0d99303..1d28cf7e07 100644
--- a/python/tvm/relax/backend/cpu_generic/pipeline.py
+++ b/python/tvm/relax/backend/cpu_generic/pipeline.py
@@ -16,7 +16,8 @@
# under the License.
"""The Relax CPU backend compilation pipeline and other passes."""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import relax
diff --git a/python/tvm/relax/backend/cuda/cublas.py b/python/tvm/relax/backend/cuda/cublas.py
index 7d525b9a1c..88ee94917c 100644
--- a/python/tvm/relax/backend/cuda/cublas.py
+++ b/python/tvm/relax/backend/cuda/cublas.py
@@ -20,10 +20,11 @@
import operator
from functools import reduce
-import tvm, tvm_tirx
-from tvm_tirx import DataType
+import tvm
+import tvm_tirx
from tvm.relax import transform
from tvm.relax.transform import PatternCheckContext
+from tvm_tirx import DataType
from tvm_tirx.sym import Analyzer
from ..pattern_registry import get_patterns_with_prefix, register_patterns
diff --git a/python/tvm/relax/backend/cuda/cudnn.py b/python/tvm/relax/backend/cuda/cudnn.py
index c9ebaf630c..9c6ae08433 100644
--- a/python/tvm/relax/backend/cuda/cudnn.py
+++ b/python/tvm/relax/backend/cuda/cudnn.py
@@ -20,7 +20,7 @@
import operator
from functools import partial, reduce
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import relax
from tvm.relax import PyExprMutator, expr_functor, transform
from tvm.relax.transform import PatternCheckContext
diff --git a/python/tvm/relax/backend/cuda/cutlass.py b/python/tvm/relax/backend/cuda/cutlass.py
index c8ca45b190..4db8c862d1 100644
--- a/python/tvm/relax/backend/cuda/cutlass.py
+++ b/python/tvm/relax/backend/cuda/cutlass.py
@@ -22,7 +22,7 @@ import operator
from collections.abc import Mapping, Sequence
from functools import reduce
-import tvm, tvm_tirx
+import tvm_tirx
from tvm.contrib.cutlass.build import is_shape_valid_for_cutlass_matmul
from tvm.relax import (
Call,
diff --git a/python/tvm/relax/backend/cuda/flashinfer.py b/python/tvm/relax/backend/cuda/flashinfer.py
index f92510a09a..65f4bbf859 100644
--- a/python/tvm/relax/backend/cuda/flashinfer.py
+++ b/python/tvm/relax/backend/cuda/flashinfer.py
@@ -20,7 +20,7 @@
import re
from pathlib import Path
-import tvm, tvm_tirx
+import tvm_tirx
from tvm_tirx.target import Target
diff --git a/python/tvm/relax/backend/cuda/pipeline.py b/python/tvm/relax/backend/cuda/pipeline.py
index f8f818d369..3ef90cef99 100644
--- a/python/tvm/relax/backend/cuda/pipeline.py
+++ b/python/tvm/relax/backend/cuda/pipeline.py
@@ -16,7 +16,8 @@
# under the License.
"""The Relax CUDA backend compilation pipeline and other passes."""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import relax
diff --git a/python/tvm/relax/backend/dispatch_sampling.py b/python/tvm/relax/backend/dispatch_sampling.py
index 398c1bf0e6..77e6cc2076 100644
--- a/python/tvm/relax/backend/dispatch_sampling.py
+++ b/python/tvm/relax/backend/dispatch_sampling.py
@@ -18,10 +18,10 @@
"""Dispatch sampling operators to platform dependent implementation."""
from tvm import relax
+from tvm.relax import expr_functor
from tvm_tirx.ir import Op
from tvm_tirx.ir.module import IRModule
from tvm_tirx.ir.transform import PassContext, module_pass
-from tvm.relax import expr_functor
from .utils import BackendDispatcher
diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py b/python/tvm/relax/backend/dispatch_sort_scan.py
index 47828fda29..bd57cafe69 100644
--- a/python/tvm/relax/backend/dispatch_sort_scan.py
+++ b/python/tvm/relax/backend/dispatch_sort_scan.py
@@ -20,13 +20,13 @@
from functools import reduce
from operator import mul
-from tvm_tirx import DataType
from tvm import relax, topi
from tvm.contrib.thrust import can_use_thrust
+from tvm.relax import expr_functor
+from tvm_tirx import DataType
from tvm_tirx.ir import GlobalVar, Op
from tvm_tirx.ir.module import IRModule
from tvm_tirx.ir.transform import PassContext, module_pass
-from tvm.relax import expr_functor
from tvm_tirx.target import Target
from .utils import BackendDispatcher
diff --git a/python/tvm/relax/backend/gpu_generic/pipeline.py b/python/tvm/relax/backend/gpu_generic/pipeline.py
index 7b3bf7a5b0..016564d3c4 100644
--- a/python/tvm/relax/backend/gpu_generic/pipeline.py
+++ b/python/tvm/relax/backend/gpu_generic/pipeline.py
@@ -16,7 +16,8 @@
# under the License.
"""The Relax generic GPU backend compilation pipeline and other passes."""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import relax
diff --git a/python/tvm/relax/backend/metal/coreml.py b/python/tvm/relax/backend/metal/coreml.py
index 185ad7e26e..64f3c21caf 100644
--- a/python/tvm/relax/backend/metal/coreml.py
+++ b/python/tvm/relax/backend/metal/coreml.py
@@ -22,14 +22,15 @@ import shutil
import tvm_ffi
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm.contrib import coreml_runtime
-from tvm_tirx.ir import Call, GenericConst, PrimType
from tvm.relax import transform
from tvm.relax.dpl.pattern import is_op, wildcard
from tvm.relax.expr import BindingBlock, Function, SeqExpr, Var, VarBinding
from tvm.relax.transform import PatternCheckContext
from tvm.relax.type import TensorType
+from tvm_tirx.ir import Call, GenericConst, PrimType
from tvm_tirx.support.xcode import compile_coreml
from ...expr_functor import PyExprVisitor, visitor
diff --git a/python/tvm/relax/backend/rocm/hipblas.py b/python/tvm/relax/backend/rocm/hipblas.py
index 942f399e3f..9d750a5194 100644
--- a/python/tvm/relax/backend/rocm/hipblas.py
+++ b/python/tvm/relax/backend/rocm/hipblas.py
@@ -20,7 +20,7 @@
import operator
from functools import reduce
-import tvm, tvm_tirx
+import tvm_tirx
from tvm.relax import transform
from tvm.relax.transform import PatternCheckContext
diff --git a/python/tvm/relax/backend/rocm/pipeline.py b/python/tvm/relax/backend/rocm/pipeline.py
index 03bd853c5a..df4f07f8ba 100644
--- a/python/tvm/relax/backend/rocm/pipeline.py
+++ b/python/tvm/relax/backend/rocm/pipeline.py
@@ -16,7 +16,8 @@
# under the License.
"""The Relax ROCm backend compilation pipeline and other passes."""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import relax
diff --git a/python/tvm/relax/base_py_module.py b/python/tvm/relax/base_py_module.py
index 192710b3ad..82542770b4 100644
--- a/python/tvm/relax/base_py_module.py
+++ b/python/tvm/relax/base_py_module.py
@@ -26,9 +26,10 @@ from typing import Any, Optional, Union
import numpy as np
from tvm_ffi import Function
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx
from tvm import relax
+from tvm_tirx import tirx
from tvm_tirx.ir import IRModule
from tvm_tirx.runtime import Device, Tensor
from tvm_tirx.target import Target
@@ -639,7 +640,9 @@ class BasePyModule:
This method extends the standard IRModule show() method to handle
Python functions stored in the IRModule's pyfuncs attribute.
"""
- from tvm_tirx.script.printer.highlight import cprint # pylint: disable=import-outside-toplevel
+ from tvm_tirx.script.printer.highlight import (
+ cprint, # pylint: disable=import-outside-toplevel
+ )
if black_format is None:
env = os.environ.get("TVM_BLACK_FORMAT")
diff --git a/python/tvm/relax/binding_rewrite.py b/python/tvm/relax/binding_rewrite.py
index b62451afe6..f67481961b 100644
--- a/python/tvm/relax/binding_rewrite.py
+++ b/python/tvm/relax/binding_rewrite.py
@@ -19,7 +19,7 @@
import tvm_ffi
-import tvm, tvm_tirx
+import tvm_tirx
from tvm_tirx.runtime import Object
from . import Binding, DataflowBlock, Expr, Function, Var, _ffi_api
diff --git a/python/tvm/relax/block_builder.py b/python/tvm/relax/block_builder.py
index 14677bb40e..8bfcc8b4c3 100644
--- a/python/tvm/relax/block_builder.py
+++ b/python/tvm/relax/block_builder.py
@@ -23,7 +23,7 @@ from typing import Any, Optional
import tvm_ffi
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import relax as rx
from tvm_tirx import tirx
from tvm_tirx.ir.module import IRModule
diff --git a/python/tvm/relax/distributed/transform/transform.py b/python/tvm/relax/distributed/transform/transform.py
index f2da847e80..07a68676fe 100644
--- a/python/tvm/relax/distributed/transform/transform.py
+++ b/python/tvm/relax/distributed/transform/transform.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name
"""Relax distributed-related transformation passes."""
-import tvm_tirx.ir, tvm
+import tvm_tirx.ir
from . import _ffi_api
diff --git a/python/tvm/relax/distributed/type.py b/python/tvm/relax/distributed/type.py
index 3873b3df53..f7aa670951 100644
--- a/python/tvm/relax/distributed/type.py
+++ b/python/tvm/relax/distributed/type.py
@@ -21,9 +21,9 @@ import enum
import tvm_ffi
+from tvm.relax.type import TensorType, Type
from tvm_tirx.ir import Location
from tvm_tirx.ir.location import UNKNOWN_LOC
-from tvm.relax.type import TensorType, Type
from tvm_tirx.runtime import Object
from . import _ffi_api
diff --git a/python/tvm/relax/dpl/context.py b/python/tvm/relax/dpl/context.py
index 232fd97523..748ae71545 100644
--- a/python/tvm/relax/dpl/context.py
+++ b/python/tvm/relax/dpl/context.py
@@ -17,7 +17,7 @@
"""The Graph Matching Context Manager for Dataflow Pattern Language."""
-import tvm, tvm_tirx
+import tvm_tirx
from ..expr import DataflowBlock, Var
from . import _ffi as ffi
diff --git a/python/tvm/relax/dpl/pattern.py b/python/tvm/relax/dpl/pattern.py
index 517eb16cee..38ce8e4636 100644
--- a/python/tvm/relax/dpl/pattern.py
+++ b/python/tvm/relax/dpl/pattern.py
@@ -25,12 +25,12 @@ from typing import Union
import tvm_ffi
from tvm_ffi import Array
-import tvm, tvm_tirx
-from tvm_tirx.ir.op import Op
-
+import tvm_tirx
from tvm_tirx.ir import make_node
from tvm_tirx.ir.base import Node
+from tvm_tirx.ir.op import Op
from tvm_tirx.runtime import Object
+
from ..expr import Expr, Var
from . import _ffi as ffi
diff --git a/python/tvm/relax/exec_builder.py b/python/tvm/relax/exec_builder.py
index 3768bc7786..04ce0099fd 100644
--- a/python/tvm/relax/exec_builder.py
+++ b/python/tvm/relax/exec_builder.py
@@ -23,7 +23,7 @@ from enum import IntEnum
import tvm_ffi
from tvm_ffi import Shape
-import tvm, tvm_tirx
+import tvm_tirx
from . import _ffi_api
from .vm_build import VMExecutable
diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py
index 442243fe2a..4d357fe8bc 100644
--- a/python/tvm/relax/expr.py
+++ b/python/tvm/relax/expr.py
@@ -26,15 +26,15 @@ import numpy as _np # type: ignore
import tvm_ffi
from tvm_ffi.core import String
-import tvm_tirx.ir, tvm
+import tvm
import tvm.relax
-import tvm_tirx.runtime, tvm
+import tvm_tirx.ir
+import tvm_tirx.runtime
from tvm_tirx import DataType
-from tvm_tirx.ir import Scriptable
+from tvm_tirx.ir import BaseFunc, Location, Node, Scriptable
+from tvm_tirx.ir.expr import _CallableExprWithOp
from tvm_tirx.ir.location import UNKNOWN_LOC, UnknownLoc
-from tvm_tirx.ir import BaseFunc, Location, Node
-from tvm_tirx.ir.expr import _CallableExprWithOp
from . import _ffi_api
# It is a workaround for mypy: https://github.com/python/mypy/issues/7866#issuecomment-549454370
diff --git a/python/tvm/relax/expr_functor.py b/python/tvm/relax/expr_functor.py
index ba38df59fe..c4744fbb68 100644
--- a/python/tvm/relax/expr_functor.py
+++ b/python/tvm/relax/expr_functor.py
@@ -23,12 +23,12 @@ from collections.abc import Callable
import tvm_ffi
from tvm import te as _te
+from tvm.s_tir.utils import derived_object
from tvm_tirx import tirx as _tirx
from tvm_tirx.ir import Call, DataTypeImm, GenericConst, Op, StringImm, TensorLoad, is_prim_expr
+from tvm_tirx.ir.module import IRModule
from tvm_tirx.runtime import Object
-from tvm.s_tir.utils import derived_object
-from tvm_tirx.ir.module import IRModule
from . import _ffi_api
from .block_builder import BlockBuilder
from .expr import (
diff --git a/python/tvm/relax/frontend/common.py b/python/tvm/relax/frontend/common.py
index fc20e93ccd..147baacbfd 100644
--- a/python/tvm/relax/frontend/common.py
+++ b/python/tvm/relax/frontend/common.py
@@ -19,11 +19,13 @@
import numpy as _np
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import topi
-def detach_params(mod: tvm_tirx.IRModule) -> tuple[tvm_tirx.IRModule, dict[str, list[tvm_tirx.runtime.Tensor]]]:
+def detach_params(
+ mod: tvm_tirx.IRModule,
+) -> tuple[tvm_tirx.IRModule, dict[str, list[tvm_tirx.runtime.Tensor]]]:
"""Detach the attribute "params" in the functions of the input IRModule as
separate dictionary of params.
diff --git a/python/tvm/relax/frontend/nn/core.py b/python/tvm/relax/frontend/nn/core.py
index 476d721b7b..ecdfc5a8bf 100644
--- a/python/tvm/relax/frontend/nn/core.py
+++ b/python/tvm/relax/frontend/nn/core.py
@@ -34,13 +34,13 @@ from typing import (
import numpy as np # type: ignore
-import tvm_tirx.runtime, tvm
+import tvm_tirx.runtime
+from tvm.runtime.vm import VirtualMachine
from tvm_tirx import tirx
from tvm_tirx.ir import IRModule
from tvm_tirx.ir.transform import Pass
from tvm_tirx.runtime import Device
from tvm_tirx.runtime import device as as_device
-from tvm.runtime.vm import VirtualMachine
from tvm_tirx.target import Target
from .... import relax as rx
diff --git a/python/tvm/relax/frontend/nn/exporter.py b/python/tvm/relax/frontend/nn/exporter.py
index e4fbbc947d..dd908a209e 100644
--- a/python/tvm/relax/frontend/nn/exporter.py
+++ b/python/tvm/relax/frontend/nn/exporter.py
@@ -21,7 +21,7 @@ import operator
import threading
import typing
-import tvm, tvm_tirx
+import tvm_tirx
from tvm_tirx import tirx
from tvm_tirx.ir import IRModule
diff --git a/python/tvm/relax/frontend/nn/extern.py b/python/tvm/relax/frontend/nn/extern.py
index c88b140501..3bde6c3b67 100644
--- a/python/tvm/relax/frontend/nn/extern.py
+++ b/python/tvm/relax/frontend/nn/extern.py
@@ -26,7 +26,8 @@ from pathlib import Path
import tvm_ffi
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm_tirx import libinfo, tirx
from tvm_tirx.runtime import Module, load_static_library
from tvm_tirx.support import cc as _cc
@@ -277,7 +278,7 @@ class SourceModule(ExternModule): # pylint: disable=too-few-public-methods
f"Using environment variable `TVM_HOME`, but it is not a directory: {tvm_path!s}"
)
else:
- import tvm, tvm_tirx # pylint: disable=import-outside-toplevel
+ import tvm
tvm_path = Path(tvm.__file__).parent
assert tvm_path.is_dir()
diff --git a/python/tvm/relax/frontend/nn/llm/kv_cache.py b/python/tvm/relax/frontend/nn/llm/kv_cache.py
index d040ff03ed..d86acfd734 100644
--- a/python/tvm/relax/frontend/nn/llm/kv_cache.py
+++ b/python/tvm/relax/frontend/nn/llm/kv_cache.py
@@ -37,12 +37,12 @@ test suite can continue to import them via ``tvm.relax.frontend.nn.llm.kv_cache`
import math
from typing import Any, Literal
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import relax as rx
-from tvm_tirx import tirx
-from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.frontend.nn import Object, Tensor
+from tvm_tirx import tirx
+from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.target import Target
# Re-export enums + kernel factories so existing ``from kv_cache import ...``
diff --git a/python/tvm/relax/frontend/nn/subroutine.py b/python/tvm/relax/frontend/nn/subroutine.py
index 368baea32b..067e959ffb 100644
--- a/python/tvm/relax/frontend/nn/subroutine.py
+++ b/python/tvm/relax/frontend/nn/subroutine.py
@@ -26,9 +26,9 @@ import typing
import tvm_ffi
-from tvm_tirx import ir
from tvm import relax
from tvm.relax.frontend import nn
+from tvm_tirx import ir
def _camel_to_snake(name):
diff --git a/python/tvm/relax/frontend/nn/torch.py b/python/tvm/relax/frontend/nn/torch.py
index 6f287ea16e..edc568762d 100644
--- a/python/tvm/relax/frontend/nn/torch.py
+++ b/python/tvm/relax/frontend/nn/torch.py
@@ -23,8 +23,8 @@ from typing import Any
import torch
from tvm_ffi import Array, Shape
-from tvm_tirx.runtime import Tensor, _tensor
from tvm.runtime.vm import VirtualMachine
+from tvm_tirx.runtime import Tensor, _tensor
from . import core
from . import spec as _spec
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 484aead264..601f6bf19e 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -56,14 +56,14 @@ except ImportError as err:
import tvm_ffi
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm_tirx
from tvm import relax, topi
+from tvm.topi.utils import get_const_tuple
+from tvm_tirx import tirx
from tvm_tirx.ir import Expr as TvmExpr
from tvm_tirx.ir import IRModule
from tvm_tirx.ir.supply import UniqueNameSupply
from tvm_tirx.runtime import DataType, DataTypeCode
-from tvm.topi.utils import get_const_tuple
from ..common import autopad
@@ -160,7 +160,9 @@ def get_value(token, value_dict: dict[str, tvm_tirx.tirx.Var]) -> int | tvm_tirx
return value
-def parse_shape_name(name: str, value_dict: dict[str, tvm_tirx.tirx.Var]) -> tirx.Expr | tvm_tirx.tirx.Var:
+def parse_shape_name(
+ name: str, value_dict: dict[str, tvm_tirx.tirx.Var]
+) -> tirx.Expr | tvm_tirx.tirx.Var:
"""Converts expressions in the shape dimension name to prim expressions.
Parameters
@@ -1178,7 +1180,9 @@ class Unsqueeze(OnnxOpConverter):
if constant_axes == [0]:
return relax.ShapeExpr([data])
raise NotImplementedError("Unsqueeze with symbolic scalar inputs only supports axis 0.")
- if isinstance(data, tvm_tirx.ir.GenericConst) and isinstance(axes, tvm_tirx.ir.GenericConst):
+ if isinstance(data, tvm_tirx.ir.GenericConst) and isinstance(
+ axes, tvm_tirx.ir.GenericConst
+ ):
constant_axes = _normalize_constant_axes(
list(map(int, axes.value.numpy().tolist())),
data.value.numpy().ndim + axes.value.numpy().size,
@@ -1742,7 +1746,9 @@ class Reshape(OnnxOpConverter):
# Other reshape targets follow regular int64 tensor reshape semantics.
data = bb.normalize(relax.op.shape_to_tensor(data))
- if isinstance(data, tvm_tirx.ir.GenericConst) and isinstance(new_shape, tvm_tirx.ir.GenericConst):
+ if isinstance(data, tvm_tirx.ir.GenericConst) and isinstance(
+ new_shape, tvm_tirx.ir.GenericConst
+ ):
data_array = data.value.numpy()
new_shape_values = new_shape.value.numpy().tolist()
if not allowzero:
@@ -3469,7 +3475,9 @@ class Expand(OnnxOpConverter):
if s.value == -1:
# -1 means preserve the input dimension
target_shape[i] = data_shape[i]
- elif isinstance(data_shape[i], tvm_tirx.tirx.IntImm) and data_shape[i].value == 1:
+ elif (
+ isinstance(data_shape[i], tvm_tirx.tirx.IntImm) and data_shape[i].value == 1
+ ):
# Input dimension is 1, can broadcast to any target dimension >= 1
if s.value < 1:
raise ValueError(
@@ -5369,9 +5377,13 @@ class OneHot(OnnxOpConverter):
depth = get_constant(inputs[1], params)
values = get_constant(inputs[2], params)
axis = attr.get("axis", -1)
- assert isinstance(depth, tvm_tirx.ir.GenericConst), "Only constant depth currently supported."
+ assert isinstance(depth, tvm_tirx.ir.GenericConst), (
+ "Only constant depth currently supported."
+ )
depth = depth.value.numpy().tolist()
- assert isinstance(values, tvm_tirx.ir.GenericConst), "Only constant values currently supported."
+ assert isinstance(values, tvm_tirx.ir.GenericConst), (
+ "Only constant values currently supported."
+ )
values = values.value.numpy().tolist()
off_value, on_value = values
off_value, on_value = (
diff --git a/python/tvm/relax/frontend/stablehlo/stablehlo_translator.py b/python/tvm/relax/frontend/stablehlo/stablehlo_translator.py
index 2e85f7608c..fc3ed3de06 100644
--- a/python/tvm/relax/frontend/stablehlo/stablehlo_translator.py
+++ b/python/tvm/relax/frontend/stablehlo/stablehlo_translator.py
@@ -21,9 +21,9 @@
from collections.abc import Callable
from typing import Any
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm_tirx
from tvm import relax
+from tvm_tirx import tirx
def _is_tensor_var(value) -> bool:
diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
index 6b8a6a2e85..9e8fbd5d1c 100644
--- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
+++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
@@ -28,9 +28,9 @@ from functools import reduce
import tvm_ffi
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm_tirx
from tvm import relax
+from tvm_tirx import tirx
from tvm_tirx.ir import PrimType
from tvm_tirx.runtime import DataTypeCode
@@ -725,9 +725,13 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
isinstance(rhs, tvm_tirx.ir.Var) and isinstance(rhs.ty, relax.TensorType)
):
return call_binary_op(relax_op, lhs, rhs)
- elif isinstance(lhs, tvm_tirx.ir.GenericConst) and not isinstance(rhs, tvm_tirx.ir.GenericConst):
+ elif isinstance(lhs, tvm_tirx.ir.GenericConst) and not isinstance(
+ rhs, tvm_tirx.ir.GenericConst
+ ):
return call_binary_op(relax_op, lhs, relax.const(rhs, dtype=lhs.ty.dtype))
- elif isinstance(rhs, tvm_tirx.ir.GenericConst) and not isinstance(lhs, tvm_tirx.ir.GenericConst):
+ elif isinstance(rhs, tvm_tirx.ir.GenericConst) and not isinstance(
+ lhs, tvm_tirx.ir.GenericConst
+ ):
return call_binary_op(relax_op, relax.const(lhs, dtype=rhs.ty.dtype), rhs)
return intrinsic_op(lhs, rhs)
@@ -3046,7 +3050,9 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
x = self.env[node.args[0]]
if isinstance(x, list | tuple | relax.ShapeExpr | relax.Tuple):
return x[node.args[1]]
- elif isinstance(x, tvm_tirx.ir.Var) and isinstance(x.ty, relax.TupleType | relax.TensorType):
+ elif isinstance(x, tvm_tirx.ir.Var) and isinstance(
+ x.ty, relax.TupleType | relax.TensorType
+ ):
if isinstance(x.ty, relax.TupleType):
return self.block_builder.emit(relax.TupleGetItem(x, node.args[1]))
diff --git a/python/tvm/relax/frontend/torch/dynamo.py b/python/tvm/relax/frontend/torch/dynamo.py
index e946bbb4a7..41b5cd774d 100644
--- a/python/tvm/relax/frontend/torch/dynamo.py
+++ b/python/tvm/relax/frontend/torch/dynamo.py
@@ -24,7 +24,8 @@ import functools
import tvm_ffi
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm.relax import build as relax_build
from .fx_translator import from_fx
diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py
index 5735e80a35..947bbec223 100644
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -33,7 +33,7 @@ except ImportError as err:
"torch is required by the PyTorch frontend. Install it with: pip install torch"
) from err
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import relax
from .base_fx_graph_translator import BaseFXGraphImporter
diff --git a/python/tvm/relax/frontend/torch/fx_translator.py b/python/tvm/relax/frontend/torch/fx_translator.py
index e13f0dd403..7d31fa3468 100644
--- a/python/tvm/relax/frontend/torch/fx_translator.py
+++ b/python/tvm/relax/frontend/torch/fx_translator.py
@@ -22,7 +22,7 @@
from collections.abc import Callable
from functools import partial, reduce
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import relax
from .base_fx_graph_translator import BaseFXGraphImporter
diff --git a/python/tvm/relax/ir/instrument.py b/python/tvm/relax/ir/instrument.py
index f92d0525c0..dc47a09b65 100644
--- a/python/tvm/relax/ir/instrument.py
+++ b/python/tvm/relax/ir/instrument.py
@@ -16,7 +16,7 @@
# under the License.
"""Common relax pass instrumentation across IR variants."""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import relax
diff --git a/python/tvm/relax/op/__init__.py b/python/tvm/relax/op/__init__.py
index 1fef6ef86a..4ef0a7ac90 100644
--- a/python/tvm/relax/op/__init__.py
+++ b/python/tvm/relax/op/__init__.py
@@ -274,7 +274,8 @@ def _register_op_make():
def _call(func, *args, attrs=None):
if not isinstance(
- func.ty, expr.tvm_tirx.ir.FuncType | expr.tvm.relax.FuncType | expr.tvm_tirx.ir.MissingType
+ func.ty,
+ expr.tvm_tirx.ir.FuncType | expr.tvm.relax.FuncType | expr.tvm_tirx.ir.MissingType,
):
return NotImplemented
return expr.tvm_tirx.ir.Call(func, args, attrs=attrs)
diff --git a/python/tvm/relax/op/_op_gradient.py b/python/tvm/relax/op/_op_gradient.py
index 0ccdab518b..43fbda67c6 100644
--- a/python/tvm/relax/op/_op_gradient.py
+++ b/python/tvm/relax/op/_op_gradient.py
@@ -21,8 +21,8 @@ import functools
import operator
from tvm import relax
-from tvm_tirx.ir import Call, PrimType
from tvm.relax.type import ShapeType
+from tvm_tirx.ir import Call, PrimType
from tvm_tirx.sym import Analyzer
from ..block_builder import BlockBuilder
diff --git a/python/tvm/relax/op/base.py b/python/tvm/relax/op/base.py
index 40ac3e33da..7761ce0a70 100644
--- a/python/tvm/relax/op/base.py
+++ b/python/tvm/relax/op/base.py
@@ -21,8 +21,9 @@ from collections.abc import Callable
import tvm_ffi
-import tvm, tvm_tirx
-import tvm_tirx.runtime, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.runtime
from tvm_tirx.ir import Attrs, Call, Op
from tvm_tirx.ir.attrs import make_node as _make_attrs
from tvm_tirx.ir.location import UNKNOWN_LOC, Location
@@ -493,7 +494,9 @@ def print(format: str | Expr, *values: Expr, ty=None, loc: Location = UNKNOWN_LO
@tvm_tirx.register_global_func("relax.run.assert_op")
-def relax_assert_op(condition: tvm_tirx.Object, format_str: str, *format_args: tvm_tirx.Object) -> None:
+def relax_assert_op(
+ condition: tvm_tirx.Object, format_str: str, *format_args: tvm_tirx.Object
+) -> None:
"""
A variadic function. The first value serves as the assertion condition:
If the condition is true, then the operator does nothing.
diff --git a/python/tvm/relax/op/distributed/distributed.py b/python/tvm/relax/op/distributed/distributed.py
index 8cb65f985e..41bb4ceea0 100644
--- a/python/tvm/relax/op/distributed/distributed.py
+++ b/python/tvm/relax/op/distributed/distributed.py
@@ -17,11 +17,11 @@
# pylint: disable=redefined-builtin
"""Operators for distributed Relax."""
+from tvm.relax.distributed import DeviceMesh, Placement
from tvm_tirx.ir import Call, Op
from tvm_tirx.ir.attrs import make_node
from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from tvm_tirx.ir.op import _make_op_api
-from tvm.relax.distributed import DeviceMesh, Placement
from ...expr import Expr
diff --git a/python/tvm/relax/op/index.py b/python/tvm/relax/op/index.py
index 51c0d6d370..a32bd46dcd 100644
--- a/python/tvm/relax/op/index.py
+++ b/python/tvm/relax/op/index.py
@@ -18,12 +18,12 @@
import tvm_ffi
+from tvm.relax.expr import Tuple as _Tuple
from tvm_tirx.error import InternalError as _InternalError
from tvm_tirx.ir import Attrs
from tvm_tirx.ir import Call as _Call
from tvm_tirx.ir.attrs import make_node as _make_attrs
from tvm_tirx.ir.location import UNKNOWN_LOC, Location
-from tvm.relax.expr import Tuple as _Tuple
from ..expr import Expr
from ..utils import convert_to_expr
diff --git a/python/tvm/relax/op/memory/view.py b/python/tvm/relax/op/memory/view.py
index 0ffc4fe4d2..5ab65a6154 100644
--- a/python/tvm/relax/op/memory/view.py
+++ b/python/tvm/relax/op/memory/view.py
@@ -28,11 +28,11 @@ while keeping the same underlying data.
from collections.abc import Sequence
+from tvm.relax import Expr, ShapeExpr, Tuple
+from tvm.relax.expr import prim_value
from tvm_tirx.ir import Call as _Call
from tvm_tirx.ir import DataTypeImm, GenericConst
from tvm_tirx.ir.location import UNKNOWN_LOC, Location
-from tvm.relax import Expr, ShapeExpr, Tuple
-from tvm.relax.expr import prim_value
from ..base import null_value
diff --git a/python/tvm/relax/op/nn/nn.py b/python/tvm/relax/op/nn/nn.py
index e1085a155e..262cf8c011 100644
--- a/python/tvm/relax/op/nn/nn.py
+++ b/python/tvm/relax/op/nn/nn.py
@@ -18,8 +18,8 @@
import tvm_ffi
-from tvm_tirx import DataType
from tvm import relax
+from tvm_tirx import DataType
from tvm_tirx.error import InternalError as _InternalError
from tvm_tirx.ir import Attrs
from tvm_tirx.ir import Call as _Call
diff --git a/python/tvm/relax/op/set.py b/python/tvm/relax/op/set.py
index 72438badee..b7323cea20 100644
--- a/python/tvm/relax/op/set.py
+++ b/python/tvm/relax/op/set.py
@@ -19,7 +19,7 @@
import numpy as np # type: ignore
-import tvm, tvm_tirx
+import tvm_tirx
from tvm_tirx.ir import Call as _Call
from tvm_tirx.ir.location import UNKNOWN_LOC, Location
diff --git a/python/tvm/relax/pipeline.py b/python/tvm/relax/pipeline.py
index 81bbfb1284..53ef83f5ed 100644
--- a/python/tvm/relax/pipeline.py
+++ b/python/tvm/relax/pipeline.py
@@ -23,7 +23,7 @@ as it is or serves as a basis to do further composition.
# pylint: disable=unused-argument
-import tvm, tvm_tirx
+import tvm_tirx
from tvm.s_tir import meta_schedule as ms
from . import backend, transform
@@ -44,7 +44,9 @@ def zero_pipeline(*, enable_warning: bool = False):
"""
@tvm_tirx.ir.transform.module_pass(opt_level=0)
- def f_zero_pipeline(mod: tvm_tirx.ir.IRModule, ctx: tvm_tirx.ir.transform.PassContext) -> tvm_tirx.ir.IRModule:
+ def f_zero_pipeline(
+ mod: tvm_tirx.ir.IRModule, ctx: tvm_tirx.ir.transform.PassContext
+ ) -> tvm_tirx.ir.IRModule:
"""Pipeline that applies pre-tuned logs.
Parameters
@@ -81,7 +83,9 @@ def default_build_pipeline():
"""The default compilation pipeline used in tvm.compile"""
@tvm_tirx.ir.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm_tirx.ir.IRModule, _ctx: tvm_tirx.ir.transform.PassContext) -> tvm_tirx.ir.IRModule:
+ def _pipeline(
+ mod: tvm_tirx.ir.IRModule, _ctx: tvm_tirx.ir.transform.PassContext
+ ) -> tvm_tirx.ir.IRModule:
seq = tvm_tirx.ir.transform.Sequential(
[
backend.DispatchSampling(),
@@ -169,7 +173,9 @@ def static_shape_tuning_pipeline(
"""
@tvm_tirx.ir.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm_tirx.ir.IRModule, _ctx: tvm_tirx.ir.transform.PassContext) -> tvm_tirx.ir.IRModule:
+ def _pipeline(
+ mod: tvm_tirx.ir.IRModule, _ctx: tvm_tirx.ir.transform.PassContext
+ ) -> tvm_tirx.ir.IRModule:
if cpu_weight_prepack:
pre_tuning_layout_rewrite = [transform.AttachAttrLayoutFreeTensors()]
post_tuning_layout_rewrite = [
diff --git a/python/tvm/relax/relax_to_pyfunc_converter.py b/python/tvm/relax/relax_to_pyfunc_converter.py
index a31c8cd018..b1643f51d4 100644
--- a/python/tvm/relax/relax_to_pyfunc_converter.py
+++ b/python/tvm/relax/relax_to_pyfunc_converter.py
@@ -28,9 +28,10 @@ import numpy # pylint: disable=unused-import
import torch
import torch.nn.functional as F
-import tvm, tvm_tirx
-from tvm_tirx import runtime
+import tvm
+import tvm_tirx
from tvm import relax
+from tvm_tirx import runtime
from tvm_tirx.ir import IRModule, Op
diff --git a/python/tvm/relax/script/__init__.py b/python/tvm/relax/script/__init__.py
index 23f88ea29e..a075886235 100644
--- a/python/tvm/relax/script/__init__.py
+++ b/python/tvm/relax/script/__init__.py
@@ -27,8 +27,8 @@ from typing import TYPE_CHECKING
from typing import Any as _Any
if TYPE_CHECKING:
- from tvm_tirx.ir import IRModule
from tvm.relax.base_py_module import BasePyModule
+ from tvm_tirx.ir import IRModule
from tvm_tirx.runtime import Device
from tvm_tirx.target import Target
@@ -83,8 +83,8 @@ class _PyModuleFactory:
self.__name__ = original_class.__name__
def __call__(self, device: Device | None = None, target: Target | None = None) -> BasePyModule:
- from tvm_tirx import cpu, ir
from tvm.relax.base_py_module import BasePyModule
+ from tvm_tirx import cpu, ir
source = self.ir_module
instance_module = ir.IRModule(
@@ -157,9 +157,9 @@ def py_module(
"""
def apply(source_class: type) -> IRModule | _PyModuleFactory:
- from tvm_tirx import ir
from tvm import relax
from tvm.relax.base_py_module import BasePyModule
+ from tvm_tirx import ir
from tvm_tirx.script.parser.entry import parse
from tvm_tirx.script.parser.inspect_source import acquire_source, capture_definition_scope
diff --git a/python/tvm/relax/script/ir_builder/distributed/ir.py b/python/tvm/relax/script/ir_builder/distributed/ir.py
index 72ad38a204..0d54c0934e 100644
--- a/python/tvm/relax/script/ir_builder/distributed/ir.py
+++ b/python/tvm/relax/script/ir_builder/distributed/ir.py
@@ -22,9 +22,9 @@ from numbers import Number
import numpy as _np # type: ignore
-import tvm, tvm_tirx
-from tvm_tirx.ir import GenericConst
+import tvm_tirx
from tvm.relax.distributed import DeviceMesh, DTensorType
+from tvm_tirx.ir import GenericConst
from tvm_tirx.runtime import _tensor
diff --git a/python/tvm/relax/script/ir_builder/distributed/op.py b/python/tvm/relax/script/ir_builder/distributed/op.py
index 7c73e3a6dc..efa41f88b0 100644
--- a/python/tvm/relax/script/ir_builder/distributed/op.py
+++ b/python/tvm/relax/script/ir_builder/distributed/op.py
@@ -17,13 +17,13 @@
# pylint: disable=redefined-builtin, wrong-import-order, no-member, invalid-name
"""Distributed Relax expression operators."""
-from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from tvm.relax.distributed import DeviceMesh, Placement
from tvm.relax.expr import Expr
from tvm.relax.op import call_tir as call_tir
from tvm.relax.op.distributed import annotate_sharding as _annotate_sharding
from tvm.relax.op.distributed import call_tir_local_view, redistribute_replica_to_shard
from tvm.relax.op.distributed import redistribute as _redistribute
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from .ir import _lookup_device_mesh
diff --git a/python/tvm/relax/script/ir_builder/ir.py b/python/tvm/relax/script/ir_builder/ir.py
index 2357d635e1..3cc977e31e 100644
--- a/python/tvm/relax/script/ir_builder/ir.py
+++ b/python/tvm/relax/script/ir_builder/ir.py
@@ -24,12 +24,8 @@ import inspect
import numbers as _numbers
import re as _re
-from tvm_tirx import ir as _ir
from tvm import relax
from tvm import relax as _relax
-from tvm_tirx.ir import IRModule
-from tvm_tirx.ir.location import UNKNOWN_LOC, Location
-from tvm_tirx.ir.type import _init_type_api, _normalize_type
from tvm.relax import Expr, ExternFunc, ShapeExpr, TupleGetItem, const
from tvm.relax.distributed import DeviceMesh as _DeviceMesh
from tvm.relax.distributed import DTensorType as _DTensorType
@@ -37,6 +33,10 @@ from tvm.relax.distributed import Placement as _Placement
from tvm.relax.distributed import device_mesh
from tvm.relax.dpl import PatternMatchingRewriter
from tvm.relax.global_info import DummyGlobalInfo, VDevice
+from tvm_tirx import ir as _ir
+from tvm_tirx.ir import IRModule
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.ir.type import _init_type_api, _normalize_type
from tvm_tirx.runtime._tensor import (
cpu,
cuda,
diff --git a/python/tvm/relax/script/ir_builder/op.py b/python/tvm/relax/script/ir_builder/op.py
index 8af2b60e99..43e6eb5ec6 100644
--- a/python/tvm/relax/script/ir_builder/op.py
+++ b/python/tvm/relax/script/ir_builder/op.py
@@ -24,13 +24,8 @@ import numbers as _numbers
from types import SimpleNamespace as _Namespace
from typing import Any
-import tvm, tvm_tirx
-from tvm_tirx import ir as _ir
+import tvm_tirx
from tvm import relax as _relax
-from tvm_tirx import tirx as _tir
-from tvm_tirx.ir.location import UNKNOWN_LOC, Location
-from tvm_tirx.ir.op import _init_op_api
-from tvm_tirx.ir.prim import _ffi_api as _prim_ffi
from tvm.relax import Call, Expr, ExternFunc
from tvm.relax.global_info import VDevice
from tvm.relax.op import (
@@ -234,6 +229,11 @@ from tvm.relax.op import to_vdevice as _to_vdevice
from tvm.relax.op.builtin import stop_lift_params
from tvm.relax.type import Type
from tvm.relax.utils import convert_to_expr
+from tvm_tirx import ir as _ir
+from tvm_tirx import tirx as _tir
+from tvm_tirx.ir.location import UNKNOWN_LOC, Location
+from tvm_tirx.ir.op import _init_op_api
+from tvm_tirx.ir.prim import _ffi_api as _prim_ffi
from tvm_tirx.script.ir_builder import base as _base
from tvm_tirx.script.ir_builder.base import auto_emit_void_call as _auto_emit_void_call
diff --git a/python/tvm/relax/script/ir_builder/parser_protocol.py b/python/tvm/relax/script/ir_builder/parser_protocol.py
index 90b6c5ec28..0fe229930e 100644
--- a/python/tvm/relax/script/ir_builder/parser_protocol.py
+++ b/python/tvm/relax/script/ir_builder/parser_protocol.py
@@ -30,12 +30,12 @@ from typing import Any, NoReturn
import tvm_ffi as _ffi
-from tvm_tirx import ir as _ir
from tvm import relax as _relax
-from tvm_tirx.ir.location import UnknownLoc
from tvm.relax import Call, Expr, Var, VarBinding
from tvm.relax.type import Type
from tvm.relax.utils import gen_call_tir_inputs
+from tvm_tirx import ir as _ir
+from tvm_tirx.ir.location import UnknownLoc
from tvm_tirx.runtime import Object as tvm_Object
from tvm_tirx.script.ir_builder import base as _base
from tvm_tirx.script.ir_builder.base import IRBuilder as _IRBuilder
diff --git a/python/tvm/relax/testing/ast_printer.py b/python/tvm/relax/testing/ast_printer.py
index 3aa3946086..bfcaae5198 100644
--- a/python/tvm/relax/testing/ast_printer.py
+++ b/python/tvm/relax/testing/ast_printer.py
@@ -24,10 +24,10 @@ but it can be useful for tutorials and debugging.
from collections.abc import Iterable
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import relax
-from tvm_tirx.ir.expr import Expr
from tvm.relax import ExprFunctor
+from tvm_tirx.ir.expr import Expr
def wrap_quotes(text: str) -> str:
diff --git a/python/tvm/relax/testing/lib_comparator.py b/python/tvm/relax/testing/lib_comparator.py
index de26dc94a1..2a25ceb5da 100644
--- a/python/tvm/relax/testing/lib_comparator.py
+++ b/python/tvm/relax/testing/lib_comparator.py
@@ -19,8 +19,8 @@
from collections.abc import Iterable
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
class LibCompareVMInstrument:
diff --git a/python/tvm/relax/testing/nn.py b/python/tvm/relax/testing/nn.py
index c94d574062..2ecb14e6f3 100644
--- a/python/tvm/relax/testing/nn.py
+++ b/python/tvm/relax/testing/nn.py
@@ -23,10 +23,11 @@ from typing import Any
import numpy as np # type: ignore
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx
from tvm import relax, topi
from tvm.relax.op.grad.grad import end_checkpoint, start_checkpoint
+from tvm_tirx import tirx
def emit(expr: relax.Expr, name_hint: str = "") -> relax.Var:
diff --git a/python/tvm/relax/testing/transform.py b/python/tvm/relax/testing/transform.py
index b939c9b2a4..a46ef9be19 100644
--- a/python/tvm/relax/testing/transform.py
+++ b/python/tvm/relax/testing/transform.py
@@ -22,10 +22,10 @@ import os
import tvm_ffi
-import tvm, tvm_tirx
+import tvm_tirx
+from tvm.relax.expr import DataflowBlock, Var
from tvm_tirx.ir import Call
from tvm_tirx.ir.module import IRModule
-from tvm.relax.expr import DataflowBlock, Var
from tvm_tirx.runtime import Object
@@ -42,7 +42,9 @@ def dataflow_liveness_analysis(block: DataflowBlock) -> dict[Var, tuple[int, int
if "PYTEST_CURRENT_TEST" not in os.environ:
logging.warning("The function dataflow_liveness_analysis is exposed for testing only.")
- live_ranges = tvm_tirx.get_global_func("relax.testing.transform.DataflowLivenessAnalysis")(block) # type: ignore
+ live_ranges = tvm_tirx.get_global_func("relax.testing.transform.DataflowLivenessAnalysis")(
+ block
+ ) # type: ignore
ret = {}
for var, live_range in live_ranges.items():
ret[var] = tuple(live_range)
@@ -58,7 +60,9 @@ def dataflow_alias_analysis(
if "PYTEST_CURRENT_TEST" not in os.environ:
logging.warning("The function dataflow_alias_analysis is exposed for testing only.")
- alias_sets, tuple_map = tvm_tirx.get_global_func("relax.testing.transform.DataflowAliasAnalysis")(
+ alias_sets, tuple_map = tvm_tirx.get_global_func(
+ "relax.testing.transform.DataflowAliasAnalysis"
+ )(
block,
inputs,
) # type: ignore
diff --git a/python/tvm/relax/testing/vm.py b/python/tvm/relax/testing/vm.py
index 168b5c376d..f41d368b51 100644
--- a/python/tvm/relax/testing/vm.py
+++ b/python/tvm/relax/testing/vm.py
@@ -21,7 +21,7 @@ from typing import Any
import numpy as np # type: ignore
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import relax
from tvm_tirx.runtime import Object
diff --git a/python/tvm/relax/training/optimizer.py b/python/tvm/relax/training/optimizer.py
index 924285487d..f704dbad41 100644
--- a/python/tvm/relax/training/optimizer.py
+++ b/python/tvm/relax/training/optimizer.py
@@ -23,7 +23,8 @@ from typing import Optional, Union
import numpy as np # type: ignore
import tvm_ffi
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from ..block_builder import BlockBuilder
from ..expr import Function, TupleGetItem, Var, const
@@ -150,7 +151,10 @@ class Optimizer:
f"type {x.ty}"
)
data_type = tvm_tirx.DataType(x.ty.dtype.dtype)
- if data_type.type_code not in (tvm_tirx.DataTypeCode.BFLOAT, tvm_tirx.DataTypeCode.FLOAT):
+ if data_type.type_code not in (
+ tvm_tirx.DataTypeCode.BFLOAT,
+ tvm_tirx.DataTypeCode.FLOAT,
+ ):
raise ValueError(
f"Optimizers only support Tensor parameters of floating point dtype, but dtype "
f"of {x.name} is {x.ty.dtype}"
diff --git a/python/tvm/relax/training/setup_trainer.py b/python/tvm/relax/training/setup_trainer.py
index 9824345a71..2dd448125f 100644
--- a/python/tvm/relax/training/setup_trainer.py
+++ b/python/tvm/relax/training/setup_trainer.py
@@ -17,7 +17,7 @@
# pylint: disable=not-callable, unused-argument
"""Setup Trainer Pass."""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm_tirx.ir.module import IRModule
from tvm_tirx.tirx.expr import IntImm
diff --git a/python/tvm/relax/training/trainer.py b/python/tvm/relax/training/trainer.py
index 34f8895550..dacf61fa35 100644
--- a/python/tvm/relax/training/trainer.py
+++ b/python/tvm/relax/training/trainer.py
@@ -19,7 +19,7 @@
import numpy as np # type: ignore
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import relax
from tvm_tirx.ir.module import IRModule
from tvm_tirx.runtime._tensor import Tensor
@@ -140,14 +140,18 @@ class Trainer:
def zero_init_params(self):
"""Zero initialize all parameters. Requires all parameters have static shapes."""
self._params = [
- tvm_tirx.runtime.tensor(np.zeros(self._get_shape_list(p), p.ty.dtype.dtype), self.device)
+ tvm_tirx.runtime.tensor(
+ np.zeros(self._get_shape_list(p), p.ty.dtype.dtype), self.device
+ )
for p in self._param_vars
]
def zero_init_states(self):
"""Zero initialize all states. Requires all states have static shapes."""
self._states = [
- tvm_tirx.runtime.tensor(np.zeros(self._get_shape_list(s), s.ty.dtype.dtype), self.device)
+ tvm_tirx.runtime.tensor(
+ np.zeros(self._get_shape_list(s), s.ty.dtype.dtype), self.device
+ )
for s in self._state_vars
]
@@ -155,8 +159,9 @@ class Trainer:
self,
params: list[np.ndarray | Tensor] | dict[str, np.ndarray | Tensor],
):
- """Load parameters from a dict or a list. Will convert parameters into tvm_tirx.runtime.Tensor
- in self.device.
+ """Load parameters from a dict or a list.
+
+ Will convert parameters into tvm_tirx.runtime.Tensor in self.device.
Parameters
----------
@@ -181,7 +186,9 @@ class Trainer:
for key, val in params.items():
if key not in self._param_name_to_pos:
raise ValueError(f"Parameter {key} is not found in the model")
- self._params[self._param_name_to_pos[key]] = tvm_tirx.runtime.tensor(val, self.device)
+ self._params[self._param_name_to_pos[key]] = tvm_tirx.runtime.tensor(
+ val, self.device
+ )
else:
raise ValueError("The type of extern_params should be either list or dict")
@@ -215,7 +222,9 @@ class Trainer:
for key, val in states.items():
if key not in self._param_name_to_pos:
raise ValueError(f"Parameter {key} is not found in the model")
- self._states[self._param_name_to_pos[key]] = tvm_tirx.runtime.tensor(val, self.device)
+ self._states[self._param_name_to_pos[key]] = tvm_tirx.runtime.tensor(
+ val, self.device
+ )
else:
raise ValueError("The type of extern_states should be either list or dict")
diff --git a/python/tvm/relax/training/utils.py b/python/tvm/relax/training/utils.py
index 98b09b9ab5..ca291d22d9 100644
--- a/python/tvm/relax/training/utils.py
+++ b/python/tvm/relax/training/utils.py
@@ -21,10 +21,10 @@ from collections.abc import Callable
from tvm_ffi import register_global_func
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import relax
-from tvm_tirx.ir import Call
from tvm.relax.block_builder import BlockBuilder
+from tvm_tirx.ir import Call
from ..expr import Function, Var
from . import _ffi_api
diff --git a/python/tvm/relax/transform/fast_math.py b/python/tvm/relax/transform/fast_math.py
index ac271d1ee2..eace9821e6 100644
--- a/python/tvm/relax/transform/fast_math.py
+++ b/python/tvm/relax/transform/fast_math.py
@@ -17,10 +17,11 @@
# pylint: disable=invalid-name, unused-argument, redefined-argument-from-local
"""Relax Use Fast Math pass."""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import topi
-from tvm_tirx.ir.module import IRModule
from tvm.relax import Call, Expr, PyExprMutator, expr_functor
+from tvm_tirx.ir.module import IRModule
@expr_functor.mutator
diff --git a/python/tvm/relax/transform/fold_batch_norm_to_conv2d_for_inference.py b/python/tvm/relax/transform/fold_batch_norm_to_conv2d_for_inference.py
index af436deef8..7fe0df50ae 100644
--- a/python/tvm/relax/transform/fold_batch_norm_to_conv2d_for_inference.py
+++ b/python/tvm/relax/transform/fold_batch_norm_to_conv2d_for_inference.py
@@ -17,12 +17,12 @@
# pylint: disable=invalid-name, unused-argument, redefined-argument-from-local
"""Relax Fold Batchnorm into Conv2D."""
-from tvm_tirx import tirx
from tvm import relax
-from tvm_tirx.ir.module import IRModule
-from tvm_tirx.ir.transform import PassContext
from tvm.relax import Expr
from tvm.relax.dpl import TupleGetItemPattern, is_const, is_op, rewrite_call, wildcard
+from tvm_tirx import tirx
+from tvm_tirx.ir.module import IRModule
+from tvm_tirx.ir.transform import PassContext
from . import function_pass
diff --git a/python/tvm/relax/transform/fuse_transpose_matmul.py b/python/tvm/relax/transform/fuse_transpose_matmul.py
index 9cbcf1c45a..48a88362d8 100644
--- a/python/tvm/relax/transform/fuse_transpose_matmul.py
+++ b/python/tvm/relax/transform/fuse_transpose_matmul.py
@@ -23,11 +23,11 @@ Note that
ops into library (e.g. cuBLAS) calls, please run dispatch pass before this pass.
"""
-import tvm, tvm_tirx
-from tvm_tirx import IRModule, tirx
+import tvm_tirx
from tvm import relax, te
from tvm.relax.dpl.pattern import is_op, wildcard
from tvm.relax.expr_functor import PyExprMutator, mutator
+from tvm_tirx import IRModule, tirx
@tvm_tirx.ir.transform.module_pass(opt_level=0, name="FuseTransposeMatmul")
diff --git a/python/tvm/relax/transform/ipc_allreduce_rewrite.py b/python/tvm/relax/transform/ipc_allreduce_rewrite.py
index da1f994f0c..530fcccf86 100644
--- a/python/tvm/relax/transform/ipc_allreduce_rewrite.py
+++ b/python/tvm/relax/transform/ipc_allreduce_rewrite.py
@@ -18,11 +18,11 @@
The pass is written in Python for experiment, fast development.
"""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import relax
-from tvm_tirx.ir.module import IRModule
from tvm.relax.expr import Expr, Var
from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, visitor
+from tvm_tirx.ir.module import IRModule
@tvm_tirx.ir.transform.module_pass(opt_level=0, name="IPCAllReduceRewrite")
diff --git a/python/tvm/relax/transform/lazy_transform_params.py b/python/tvm/relax/transform/lazy_transform_params.py
index 62a57b930f..f2099ab0e1 100644
--- a/python/tvm/relax/transform/lazy_transform_params.py
+++ b/python/tvm/relax/transform/lazy_transform_params.py
@@ -18,10 +18,10 @@
# ruff: noqa: RUF005
"""Relax LazyTransformParams pass."""
-import tvm, tvm_tirx
-from tvm_tirx import IRModule
+import tvm_tirx
from tvm import relax
from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, visitor
+from tvm_tirx import IRModule
@visitor
diff --git a/python/tvm/relax/transform/legalize_ops/ccl.py b/python/tvm/relax/transform/legalize_ops/ccl.py
index b7478ed593..ddc9647749 100644
--- a/python/tvm/relax/transform/legalize_ops/ccl.py
+++ b/python/tvm/relax/transform/legalize_ops/ccl.py
@@ -18,10 +18,10 @@
# ruff: noqa: RUF005
"""Default legalization function for ccl operators."""
-from tvm_tirx import sym, tirx
from tvm import topi
-from tvm_tirx.ir import Call
from tvm.relax import ExternFunc as _ExternFunc
+from tvm_tirx import sym, tirx
+from tvm_tirx.ir import Call
from ...block_builder import BlockBuilder
from ...expr import Expr, ShapeExpr
diff --git a/python/tvm/relax/transform/legalize_ops/common.py b/python/tvm/relax/transform/legalize_ops/common.py
index 37a3948174..875d98d300 100644
--- a/python/tvm/relax/transform/legalize_ops/common.py
+++ b/python/tvm/relax/transform/legalize_ops/common.py
@@ -18,7 +18,7 @@
from collections.abc import Callable
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
from tvm_tirx.ir import Call, GenericConst
from tvm_tirx.runtime import DataTypeCode
diff --git a/python/tvm/relax/transform/legalize_ops/create.py b/python/tvm/relax/transform/legalize_ops/create.py
index a6e5664a54..b0dc36d3b0 100644
--- a/python/tvm/relax/transform/legalize_ops/create.py
+++ b/python/tvm/relax/transform/legalize_ops/create.py
@@ -20,9 +20,9 @@
import numpy as np
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm_tirx
from tvm import te, topi
+from tvm_tirx import tirx
from tvm_tirx.ir import Call
from ...block_builder import BlockBuilder
diff --git a/python/tvm/relax/transform/legalize_ops/distributed.py b/python/tvm/relax/transform/legalize_ops/distributed.py
index 0515774bc3..0e641eaa87 100644
--- a/python/tvm/relax/transform/legalize_ops/distributed.py
+++ b/python/tvm/relax/transform/legalize_ops/distributed.py
@@ -17,8 +17,8 @@
# pylint: disable=invalid-name
"""Default legalization function for distir-related operators."""
-from tvm_tirx import tirx
from tvm import relax
+from tvm_tirx import tirx
from tvm_tirx.ir import Call
from ...block_builder import BlockBuilder
diff --git a/python/tvm/relax/transform/legalize_ops/grad.py b/python/tvm/relax/transform/legalize_ops/grad.py
index e1fb9e66a1..8b63a39f55 100644
--- a/python/tvm/relax/transform/legalize_ops/grad.py
+++ b/python/tvm/relax/transform/legalize_ops/grad.py
@@ -19,8 +19,8 @@
import logging
-from tvm_tirx import tirx
from tvm import te, topi
+from tvm_tirx import tirx
from tvm_tirx.ir import Call
from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.tirx.script import ir_builder as T
diff --git a/python/tvm/relax/transform/legalize_ops/image.py b/python/tvm/relax/transform/legalize_ops/image.py
index fdf96190e3..d446b716ce 100644
--- a/python/tvm/relax/transform/legalize_ops/image.py
+++ b/python/tvm/relax/transform/legalize_ops/image.py
@@ -17,8 +17,8 @@
# pylint: disable=invalid-name
"""Default legalization function for image operators."""
-from tvm_tirx import tirx
from tvm import topi
+from tvm_tirx import tirx
from tvm_tirx.ir import Call
from ...block_builder import BlockBuilder
diff --git a/python/tvm/relax/transform/legalize_ops/index.py b/python/tvm/relax/transform/legalize_ops/index.py
index 0d012fe782..e7350d69fb 100644
--- a/python/tvm/relax/transform/legalize_ops/index.py
+++ b/python/tvm/relax/transform/legalize_ops/index.py
@@ -17,9 +17,9 @@
# pylint: disable=invalid-name
"""Default legalization function for index operators."""
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm_tirx
from tvm import te, topi
+from tvm_tirx import tirx
from tvm_tirx.ir import Call, PrimType
from ...block_builder import BlockBuilder
diff --git a/python/tvm/relax/transform/legalize_ops/linear_algebra.py b/python/tvm/relax/transform/legalize_ops/linear_algebra.py
index 7bb47538d1..cf916a0a99 100644
--- a/python/tvm/relax/transform/legalize_ops/linear_algebra.py
+++ b/python/tvm/relax/transform/legalize_ops/linear_algebra.py
@@ -17,8 +17,8 @@
# pylint: disable=invalid-name
"""Default legalization function for linear algebra operators."""
-from tvm_tirx import DataTypeCode, tirx
from tvm import relax, te, topi
+from tvm_tirx import DataTypeCode, tirx
from tvm_tirx.ir import Call
from ...block_builder import BlockBuilder
diff --git a/python/tvm/relax/transform/legalize_ops/manipulate.py b/python/tvm/relax/transform/legalize_ops/manipulate.py
index 45750c1c74..fe89cbef63 100644
--- a/python/tvm/relax/transform/legalize_ops/manipulate.py
+++ b/python/tvm/relax/transform/legalize_ops/manipulate.py
@@ -18,13 +18,14 @@
# ruff: noqa: RUF005
"""Default legalization function for manipulate operators."""
-import tvm, tvm_tirx
-from tvm_tirx import DataTypeCode, tirx
+import tvm
+import tvm_tirx
from tvm import relax, te, topi
-from tvm_tirx.ir import Call
from tvm.relax.op.base import call_tir
from tvm.relax.type import TensorType
from tvm.relax.utils import gen_call_tir_inputs
+from tvm_tirx import DataTypeCode, tirx
+from tvm_tirx.ir import Call
from tvm_tirx.tirx.expr import IntImm
from ...block_builder import BlockBuilder
diff --git a/python/tvm/relax/transform/legalize_ops/nn.py b/python/tvm/relax/transform/legalize_ops/nn.py
index 79bf74faba..c082a67725 100644
--- a/python/tvm/relax/transform/legalize_ops/nn.py
+++ b/python/tvm/relax/transform/legalize_ops/nn.py
@@ -20,8 +20,8 @@
import logging
import math
-from tvm_tirx import tirx
from tvm import s_tir, te, topi
+from tvm_tirx import tirx
from tvm_tirx.ir import Call
from ...block_builder import BlockBuilder
diff --git a/python/tvm/relax/transform/legalize_ops/qdq.py b/python/tvm/relax/transform/legalize_ops/qdq.py
index e24d6675ce..6596d0c784 100644
--- a/python/tvm/relax/transform/legalize_ops/qdq.py
+++ b/python/tvm/relax/transform/legalize_ops/qdq.py
@@ -17,9 +17,9 @@
# pylint: disable=invalid-name
"""Default legalization function for quantize/dequantize operators."""
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm_tirx
from tvm import te
+from tvm_tirx import tirx
from tvm_tirx.ir import Call
from tvm_tirx.runtime import DataTypeCode
diff --git a/python/tvm/relax/transform/legalize_ops/statistical.py b/python/tvm/relax/transform/legalize_ops/statistical.py
index 70a39c52f1..20c5eaba85 100644
--- a/python/tvm/relax/transform/legalize_ops/statistical.py
+++ b/python/tvm/relax/transform/legalize_ops/statistical.py
@@ -19,8 +19,8 @@
from collections.abc import Callable
-from tvm_tirx import tirx
from tvm import te, topi
+from tvm_tirx import tirx
from tvm_tirx.ir import Call
from ...block_builder import BlockBuilder
diff --git a/python/tvm/relax/transform/legalize_ops/vision.py b/python/tvm/relax/transform/legalize_ops/vision.py
index fd61c98ab5..35840e8d03 100644
--- a/python/tvm/relax/transform/legalize_ops/vision.py
+++ b/python/tvm/relax/transform/legalize_ops/vision.py
@@ -16,9 +16,9 @@
# under the License.
"""Default legalization function for vision network related operators."""
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm_tirx
from tvm import relax, te, topi
+from tvm_tirx import tirx
from tvm_tirx.ir import Call
from ...block_builder import BlockBuilder
diff --git a/python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py b/python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py
index 0438202c9a..fc328e17ae 100644
--- a/python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py
+++ b/python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py
@@ -18,11 +18,11 @@
The pass is written in Python for experiment, fast development.
"""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import relax
-from tvm_tirx.ir.module import IRModule
from tvm.relax.expr import Expr
from tvm.relax.expr_functor import PyExprMutator, mutator
+from tvm_tirx.ir.module import IRModule
@tvm_tirx.ir.transform.module_pass(opt_level=0, name="LowerGPUIPCAllocStorage")
diff --git a/python/tvm/relax/transform/optimize_layout_transform.py b/python/tvm/relax/transform/optimize_layout_transform.py
index 62c7b4c291..5d6bab3ec4 100644
--- a/python/tvm/relax/transform/optimize_layout_transform.py
+++ b/python/tvm/relax/transform/optimize_layout_transform.py
@@ -19,10 +19,10 @@
import tvm_ffi
-from tvm_tirx.ir.module import IRModule
-from tvm_tirx.ir.transform import PassContext
from tvm.relax import Expr
from tvm.relax.dpl import TuplePattern, is_op, rewrite_call, wildcard
+from tvm_tirx.ir.module import IRModule
+from tvm_tirx.ir.transform import PassContext
from . import function_pass
diff --git a/python/tvm/relax/transform/remove_redundant_reshape.py b/python/tvm/relax/transform/remove_redundant_reshape.py
index 512a30c252..bcfc9d6910 100644
--- a/python/tvm/relax/transform/remove_redundant_reshape.py
+++ b/python/tvm/relax/transform/remove_redundant_reshape.py
@@ -19,11 +19,11 @@
import tvm_ffi
-from tvm_tirx import IRModule
from tvm import relax
-from tvm_tirx.ir.transform import PassContext
from tvm.relax import Expr
from tvm.relax.dpl import is_op, rewrite_call, wildcard
+from tvm_tirx import IRModule
+from tvm_tirx.ir.transform import PassContext
from . import function_pass
diff --git a/python/tvm/relax/transform/transform.py b/python/tvm/relax/transform/transform.py
index ada9d309a8..2f634621a5 100644
--- a/python/tvm/relax/transform/transform.py
+++ b/python/tvm/relax/transform/transform.py
@@ -28,7 +28,8 @@ import numpy as np # type: ignore
import tvm_ffi
from tvm_ffi import Array
-import tvm_tirx.ir, tvm
+import tvm
+import tvm_tirx.ir
from tvm.relax import Expr, Type, Var
from tvm.relax.dpl import DFPattern
from tvm_tirx.runtime import Object, Tensor
@@ -1404,7 +1405,9 @@ def ToMixedPrecision(
return _ffi_api.ToMixedPrecision(out_dtype, fp16_input_names) # type: ignore
-def SplitCallTIRByPattern(patterns: list[Function], fcodegen: Callable) -> tvm_tirx.ir.transform.Pass:
+def SplitCallTIRByPattern(
+ patterns: list[Function], fcodegen: Callable
+) -> tvm_tirx.ir.transform.Pass:
"""Split a Function into 2 parts: the first part is a TIR Function which is
matched with some pattern, and the second part is the rest of the original
Function. It will call fcodegen to generate the code for the matched pattern
diff --git a/python/tvm/relax/type.py b/python/tvm/relax/type.py
index 44533ced87..a2078eea78 100644
--- a/python/tvm/relax/type.py
+++ b/python/tvm/relax/type.py
@@ -21,10 +21,10 @@
import tvm_ffi
from tvm_ffi import Array
+from tvm.relax.global_info import VDevice
from tvm_tirx.ir import AnyType, EnvFunc, Location, PrimType, TupleType
from tvm_tirx.ir.location import UNKNOWN_LOC, UnknownLoc
from tvm_tirx.ir.type import _normalize_type
-from tvm.relax.global_info import VDevice
from . import _ffi_api
from .expr import Expr, ShapeExpr, Type
diff --git a/python/tvm/relax/utils.py b/python/tvm/relax/utils.py
index 5254974a98..93683ff0a8 100644
--- a/python/tvm/relax/utils.py
+++ b/python/tvm/relax/utils.py
@@ -28,12 +28,12 @@ from typing import Any, Optional
import tvm_ffi
from tvm_ffi import Array, Map
-import tvm, tvm_tirx
-from tvm_tirx.ir import StringImm
+import tvm
+import tvm_tirx
from tvm.relax.global_info import VDevice
-
from tvm_tirx import tirx
-from tvm_tirx.ir import Attrs, Type
+from tvm_tirx.ir import Attrs, StringImm, Type
+
from ..te import Tensor as te_Tensor
from ..te import create_function
from . import _ffi_api
@@ -328,7 +328,9 @@ def gen_call_tir_inputs(
return ShapeExpr(
[
tvm_ffi.structural_map(
- value, (tvm_tirx.ir.Var, lambda var: tir_var_inverse_map.get(var, var)), order="post"
+ value,
+ (tvm_tirx.ir.Var, lambda var: tir_var_inverse_map.get(var, var)),
+ order="post",
)
for value in shape_values
]
diff --git a/python/tvm/relax/vm_build.py b/python/tvm/relax/vm_build.py
index bea225f026..e5ee01200e 100644
--- a/python/tvm/relax/vm_build.py
+++ b/python/tvm/relax/vm_build.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name, no-member
"""VM build logics"""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import relax
from tvm_tirx.ir.module import IRModule
from tvm_tirx.runtime import Executable
diff --git a/python/tvm/rpc/client.py b/python/tvm/rpc/client.py
index fcf679fa90..555f715efa 100644
--- a/python/tvm/rpc/client.py
+++ b/python/tvm/rpc/client.py
@@ -27,7 +27,8 @@ import time
import tvm_ffi
from tvm_ffi import DLDeviceType
-import tvm_tirx.runtime, tvm
+import tvm
+import tvm_tirx.runtime
from tvm_tirx.support import utils
from . import _ffi_api, base, server
diff --git a/python/tvm/rpc/remote.py b/python/tvm/rpc/remote.py
index 07b7292f0f..4e9bfc8b22 100644
--- a/python/tvm/rpc/remote.py
+++ b/python/tvm/rpc/remote.py
@@ -23,7 +23,7 @@ A remote tensor is not a local tensor or a DLPack provider. Use ``numpy`` or
import numpy as np
import tvm_ffi
-import tvm_tirx.runtime, tvm
+import tvm_tirx.runtime
from . import _ffi_api
diff --git a/python/tvm/rpc/testing.py b/python/tvm/rpc/testing.py
index dfa0085d7f..00ca910aa4 100644
--- a/python/tvm/rpc/testing.py
+++ b/python/tvm/rpc/testing.py
@@ -20,7 +20,7 @@
import numpy as np
-import tvm, tvm_tirx
+import tvm_tirx
# RPC test functions to be registered for unit-tests purposes
diff --git a/python/tvm/runtime/disco/session.py b/python/tvm/runtime/disco/session.py
index d8e7ca198e..6188600ca7 100644
--- a/python/tvm/runtime/disco/session.py
+++ b/python/tvm/runtime/disco/session.py
@@ -31,6 +31,7 @@ from tvm_ffi import Object, Shape, get_global_func, register_global_func, regist
from tvm_tirx.runtime._tensor import Tensor
from tvm_tirx.runtime._tensor import tensor as _as_Tensor
from tvm_tirx.runtime.device import Device
+
from . import _ffi_api, process_pool # pylint: disable=unused-import
diff --git a/python/tvm/runtime/vm/vm.py b/python/tvm/runtime/vm/vm.py
index f9f66cc7b7..2857971aca 100644
--- a/python/tvm/runtime/vm/vm.py
+++ b/python/tvm/runtime/vm/vm.py
@@ -26,7 +26,7 @@ from typing import Any
import numpy as np # type: ignore
from tvm_ffi import Function, register_global_func
-import tvm, tvm_tirx
+import tvm_tirx
from tvm_tirx.runtime import Device, Object
diff --git a/python/tvm/s_tir/__init__.py b/python/tvm/s_tir/__init__.py
index dcbbea3368..fb3218011e 100644
--- a/python/tvm/s_tir/__init__.py
+++ b/python/tvm/s_tir/__init__.py
@@ -18,9 +18,12 @@
# pylint: disable=invalid-name
"""S-TIR namespace for scheduable TensorIR"""
-import tvm_tirx.script, tvm
+import tvm_tirx.script
+import tvm
-tvm_tirx.script.register_dialect("s_tir", "tvm.s_tir.script", builder_path="tvm.s_tir.script.ir_builder")
+tvm_tirx.script.register_dialect(
+ "s_tir", "tvm.s_tir.script", builder_path="tvm.s_tir.script.ir_builder"
+)
from .iter_var import IterVar
from ._tensor_intrin import TensorIntrin
diff --git a/python/tvm/s_tir/analysis/__init__.py b/python/tvm/s_tir/analysis/__init__.py
index 927b0f058a..20a635d234 100644
--- a/python/tvm/s_tir/analysis/__init__.py
+++ b/python/tvm/s_tir/analysis/__init__.py
@@ -20,7 +20,8 @@
# pylint: disable=invalid-name
from typing import Optional, Union
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm_tirx.ir import IRModule, TensorRegion
from tvm_tirx.tirx.expr import Var
from tvm.s_tir import SBlock
diff --git a/python/tvm/s_tir/dlight/adreno/convolution.py b/python/tvm/s_tir/dlight/adreno/convolution.py
index ebdded601a..d751948fef 100644
--- a/python/tvm/s_tir/dlight/adreno/convolution.py
+++ b/python/tvm/s_tir/dlight/adreno/convolution.py
@@ -17,8 +17,8 @@
# pylint: disable=missing-docstring, invalid-name
"""A Conv2d schedule rule for Adreno GPU operators."""
-from tvm_tirx import tirx
from tvm import s_tir
+from tvm_tirx import tirx
from tvm_tirx.target import Target
from .. import analysis
diff --git a/python/tvm/s_tir/dlight/adreno/fallback.py b/python/tvm/s_tir/dlight/adreno/fallback.py
index f1617fbf24..3b81a5e6f4 100644
--- a/python/tvm/s_tir/dlight/adreno/fallback.py
+++ b/python/tvm/s_tir/dlight/adreno/fallback.py
@@ -35,8 +35,8 @@
from functools import cache
-from tvm_tirx import tirx
from tvm import s_tir
+from tvm_tirx import tirx
from tvm_tirx.runtime import MemoryScope
from tvm_tirx.target import Target
diff --git a/python/tvm/s_tir/dlight/adreno/layout_transform.py b/python/tvm/s_tir/dlight/adreno/layout_transform.py
index 5f07497684..805b4f8938 100644
--- a/python/tvm/s_tir/dlight/adreno/layout_transform.py
+++ b/python/tvm/s_tir/dlight/adreno/layout_transform.py
@@ -37,8 +37,8 @@
"Schedules for Texture Based Layout Transforms"
-from tvm_tirx import tirx
from tvm import s_tir
+from tvm_tirx import tirx
from tvm_tirx.target import Target
from .. import analysis
diff --git a/python/tvm/s_tir/dlight/adreno/pool.py b/python/tvm/s_tir/dlight/adreno/pool.py
index 444e19f737..872ca347e6 100644
--- a/python/tvm/s_tir/dlight/adreno/pool.py
+++ b/python/tvm/s_tir/dlight/adreno/pool.py
@@ -18,8 +18,8 @@
# ruff: noqa: F841
"""Pool schedule rule for Adreno operators."""
-from tvm_tirx import tirx
from tvm import s_tir
+from tvm_tirx import tirx
from tvm_tirx.target import Target
from .. import analysis
diff --git a/python/tvm/s_tir/dlight/analysis/common_analysis.py b/python/tvm/s_tir/dlight/analysis/common_analysis.py
index 076b457237..2488a16d22 100644
--- a/python/tvm/s_tir/dlight/analysis/common_analysis.py
+++ b/python/tvm/s_tir/dlight/analysis/common_analysis.py
@@ -26,13 +26,13 @@ from typing import Literal
import tvm_ffi
from tvm_ffi import get_global_func
-import tvm, tvm_tirx
-from tvm_tirx import ir, tirx
+import tvm_tirx
from tvm import s_tir
-from tvm_tirx.ir import TensorRegion
-from tvm_tirx.runtime import MemoryScope
from tvm.s_tir import Schedule
from tvm.s_tir.schedule import SBlockRV
+from tvm_tirx import ir, tirx
+from tvm_tirx.ir import TensorRegion
+from tvm_tirx.runtime import MemoryScope
from tvm_tirx.target.target import Target
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
diff --git a/python/tvm/s_tir/dlight/analysis/gemv.py b/python/tvm/s_tir/dlight/analysis/gemv.py
index de7ebd9876..82da265d90 100644
--- a/python/tvm/s_tir/dlight/analysis/gemv.py
+++ b/python/tvm/s_tir/dlight/analysis/gemv.py
@@ -18,9 +18,9 @@
import tvm_ffi
-import tvm, tvm_tirx
-from tvm_tirx import sym, tirx
+import tvm_tirx
from tvm import s_tir
+from tvm_tirx import sym, tirx
from .common_analysis import (
SBlockInfo,
diff --git a/python/tvm/s_tir/dlight/base/schedule_rule.py b/python/tvm/s_tir/dlight/base/schedule_rule.py
index 59ed260c17..ebc981ca87 100644
--- a/python/tvm/s_tir/dlight/base/schedule_rule.py
+++ b/python/tvm/s_tir/dlight/base/schedule_rule.py
@@ -18,8 +18,8 @@
from collections.abc import Callable
-from tvm_tirx import tirx
from tvm import s_tir
+from tvm_tirx import tirx
from tvm_tirx.target import Target
diff --git a/python/tvm/s_tir/dlight/base/transform.py b/python/tvm/s_tir/dlight/base/transform.py
index 869ce9a387..065c79470e 100644
--- a/python/tvm/s_tir/dlight/base/transform.py
+++ b/python/tvm/s_tir/dlight/base/transform.py
@@ -19,8 +19,8 @@ Apply ScheduleRules onto an IRModule to generate default schedules without tunin
or a space for MetaSchedule tuning
"""
-from tvm_tirx import tirx
from tvm import s_tir
+from tvm_tirx import tirx
from tvm_tirx.ir import IRModule
from tvm_tirx.ir.transform import PassContext, module_pass
from tvm_tirx.target import Target
diff --git a/python/tvm/s_tir/dlight/base/utils.py b/python/tvm/s_tir/dlight/base/utils.py
index 680401af10..3b0aa612f5 100644
--- a/python/tvm/s_tir/dlight/base/utils.py
+++ b/python/tvm/s_tir/dlight/base/utils.py
@@ -17,9 +17,9 @@
# pylint: disable=missing-docstring
"""Utility methods for generic GPU."""
-import tvm, tvm_tirx
-from tvm_tirx import DataType, tirx
+import tvm_tirx
from tvm import s_tir
+from tvm_tirx import DataType, tirx
from tvm_tirx.ir import PrimType
from tvm_tirx.target import Target
diff --git a/python/tvm/s_tir/dlight/benchmark/bench.py b/python/tvm/s_tir/dlight/benchmark/bench.py
index 47f29d5f11..ade079e74d 100644
--- a/python/tvm/s_tir/dlight/benchmark/bench.py
+++ b/python/tvm/s_tir/dlight/benchmark/bench.py
@@ -19,11 +19,11 @@
from collections.abc import Callable
from typing import TYPE_CHECKING, Optional
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import relax
-from tvm_tirx.ir import IRModule
from tvm.s_tir.meta_schedule.runner import EvaluatorConfig
from tvm.s_tir.meta_schedule.testing.tune_utils import generate_input_data
+from tvm_tirx.ir import IRModule
from tvm_tirx.tirx import Function
from .extract import extract_all_func_info_from_relax, extract_func_info_from_function
@@ -116,7 +116,9 @@ def benchmark(
else:
# normal case like [1, n, 128], generate random tensor
input_tensors.append(
- tvm_tirx.runtime.tensor(generate_input_data(list(input_shape), input_dtype), device=dev)
+ tvm_tirx.runtime.tensor(
+ generate_input_data(list(input_shape), input_dtype), device=dev
+ )
)
# append scalar input tensors for rotary embedding
input_tensors.extend(scalar_input_tensors)
@@ -144,7 +146,9 @@ def benchmark(
_, profile_result = rpc_run(
rt_mod,
device_type=dev._DEVICE_TYPE_TO_NAME[dev.dlpack_device_type()],
- args=[w.numpy() if isinstance(w, tvm_tirx.runtime.Tensor) else w for w in input_tensors],
+ args=[
+ w.numpy() if isinstance(w, tvm_tirx.runtime.Tensor) else w for w in input_tensors
+ ],
rpc_config=rpc_config,
evaluator_config=evaluator_config,
)
diff --git a/python/tvm/s_tir/dlight/benchmark/extract.py b/python/tvm/s_tir/dlight/benchmark/extract.py
index 86b3cdd9d1..638cabb392 100644
--- a/python/tvm/s_tir/dlight/benchmark/extract.py
+++ b/python/tvm/s_tir/dlight/benchmark/extract.py
@@ -21,7 +21,8 @@ from pathlib import Path
import cloudpickle
import tvm_ffi
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import relax
from .utils import default_dym_var_sample_func, get_func_name_from_gv
@@ -217,7 +218,9 @@ def extract_all_func_info_from_relax(
]
The function input information and dynamic shape variable dictionary.
"""
- relax_func_dict: dict[tvm_tirx.ir.GlobalVar, dict[tvm_tirx.ir.GlobalVar, list[tuple[list, int]]]] = {}
+ relax_func_dict: dict[
+ tvm_tirx.ir.GlobalVar, dict[tvm_tirx.ir.GlobalVar, list[tuple[list, int]]]
+ ] = {}
for gv, func in mod.functions_items(): # pylint: disable=invalid-name,too-many-nested-blocks
if isinstance(func, tvm.relax.Function):
for block in func.body.blocks:
diff --git a/python/tvm/s_tir/dlight/benchmark/utils.py b/python/tvm/s_tir/dlight/benchmark/utils.py
index 819b3ca047..a427723426 100644
--- a/python/tvm/s_tir/dlight/benchmark/utils.py
+++ b/python/tvm/s_tir/dlight/benchmark/utils.py
@@ -18,7 +18,7 @@
from typing import Any
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import relax
INPUT_SHAPE_TYPE = list[tuple[tuple[int, ...], str]] # pylint: disable=invalid-name
diff --git a/python/tvm/s_tir/dlight/cpu/gemv.py b/python/tvm/s_tir/dlight/cpu/gemv.py
index 15f021f036..05b06b1c6d 100644
--- a/python/tvm/s_tir/dlight/cpu/gemv.py
+++ b/python/tvm/s_tir/dlight/cpu/gemv.py
@@ -16,8 +16,8 @@
# under the License.
"""A rule for GEMV and DecodeGEMV."""
-from tvm_tirx import tirx
from tvm import s_tir
+from tvm_tirx import tirx
from tvm_tirx.target import Target
from ..analysis import SBlockInfo, normalize_function
diff --git a/python/tvm/s_tir/dlight/cpu/reduction.py b/python/tvm/s_tir/dlight/cpu/reduction.py
index affaf0aa94..17b24dde80 100644
--- a/python/tvm/s_tir/dlight/cpu/reduction.py
+++ b/python/tvm/s_tir/dlight/cpu/reduction.py
@@ -16,8 +16,8 @@
# under the License.
"""CPU reduction rule for operators including softmax, layer norm, RMS norm, etc."""
-from tvm_tirx import tirx
from tvm import s_tir
+from tvm_tirx import tirx
from tvm_tirx.codegen import llvm_get_vector_width
from tvm_tirx.target import Target
diff --git a/python/tvm/s_tir/dlight/gpu/fallback.py b/python/tvm/s_tir/dlight/gpu/fallback.py
index 4b1c4ab3ba..d01da74e82 100644
--- a/python/tvm/s_tir/dlight/gpu/fallback.py
+++ b/python/tvm/s_tir/dlight/gpu/fallback.py
@@ -19,9 +19,9 @@
import tvm_ffi
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm_tirx
from tvm import s_tir
+from tvm_tirx import tirx
from tvm_tirx.target import Target
from .. import base
diff --git a/python/tvm/s_tir/dlight/gpu/gemv.py b/python/tvm/s_tir/dlight/gpu/gemv.py
index 820ac3d8fc..316ab89bf0 100644
--- a/python/tvm/s_tir/dlight/gpu/gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/gemv.py
@@ -19,9 +19,9 @@
from functools import reduce
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm_tirx
from tvm import s_tir
+from tvm_tirx import tirx
from tvm_tirx.target import Target
from ..analysis import (
diff --git a/python/tvm/s_tir/dlight/gpu/general_reduction.py b/python/tvm/s_tir/dlight/gpu/general_reduction.py
index c576979a4e..97de071111 100644
--- a/python/tvm/s_tir/dlight/gpu/general_reduction.py
+++ b/python/tvm/s_tir/dlight/gpu/general_reduction.py
@@ -19,8 +19,8 @@
import tvm_ffi
-from tvm_tirx import sym, tirx
from tvm import s_tir
+from tvm_tirx import sym, tirx
from tvm_tirx.runtime import MemoryScope
from tvm_tirx.target import Target
diff --git a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
index fbf6f6c8fb..7c45d2c5f5 100644
--- a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
@@ -22,9 +22,9 @@ from typing import Literal
import tvm_ffi
-import tvm, tvm_tirx
-from tvm_tirx import sym, tirx
+import tvm_tirx
from tvm import s_tir
+from tvm_tirx import sym, tirx
from tvm_tirx.target import Target
from ..analysis import (
diff --git a/python/tvm/s_tir/dlight/gpu/matmul.py b/python/tvm/s_tir/dlight/gpu/matmul.py
index ca7a8c6aae..32055d4c5b 100644
--- a/python/tvm/s_tir/dlight/gpu/matmul.py
+++ b/python/tvm/s_tir/dlight/gpu/matmul.py
@@ -21,11 +21,11 @@
from dataclasses import dataclass
from enum import Enum
-from tvm_tirx import tirx
from tvm import s_tir
-from tvm_tirx.ir import Range
from tvm.s_tir import IterVar
from tvm.s_tir.schedule.schedule import SBlockRV
+from tvm_tirx import tirx
+from tvm_tirx.ir import Range
from tvm_tirx.script import tirx as T
from tvm_tirx.target import Target
from tvm_tirx.tirx import Expr, Var
diff --git a/python/tvm/s_tir/dlight/gpu/reduction.py b/python/tvm/s_tir/dlight/gpu/reduction.py
index f50146ddc2..aae34c7f6a 100644
--- a/python/tvm/s_tir/dlight/gpu/reduction.py
+++ b/python/tvm/s_tir/dlight/gpu/reduction.py
@@ -21,9 +21,9 @@ from collections.abc import Mapping
import tvm_ffi
-import tvm, tvm_tirx
-from tvm_tirx import sym, tirx
+import tvm_tirx
from tvm import s_tir
+from tvm_tirx import sym, tirx
from tvm_tirx.target import Target
from ..analysis import (
diff --git a/python/tvm/s_tir/dlight/gpu/rmsnorm.py b/python/tvm/s_tir/dlight/gpu/rmsnorm.py
index 7fadb890c0..01ec6e9609 100644
--- a/python/tvm/s_tir/dlight/gpu/rmsnorm.py
+++ b/python/tvm/s_tir/dlight/gpu/rmsnorm.py
@@ -17,10 +17,11 @@
# pylint: disable=missing-docstring
"""A RMS norm schedule rule for GPU operators."""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
+from tvm.s_tir import SBlock
from tvm_tirx import tirx
from tvm_tirx.ir import Call, TensorLoad, TensorStore
-from tvm.s_tir import SBlock
from tvm_tirx.target import Target
from tvm_tirx.tirx.expr import Cast
diff --git a/python/tvm/s_tir/dlight/gpu/transpose.py b/python/tvm/s_tir/dlight/gpu/transpose.py
index e0daa44b80..7606fc83b5 100644
--- a/python/tvm/s_tir/dlight/gpu/transpose.py
+++ b/python/tvm/s_tir/dlight/gpu/transpose.py
@@ -16,12 +16,12 @@
# under the License.
"""Reduction rule for operators including softmax, layer norm, RMS norm, etc"""
-import tvm, tvm_tirx
-from tvm_tirx import sym, tirx
+import tvm_tirx
from tvm import s_tir
-from tvm_tirx.ir import TensorLoad
from tvm.s_tir import Schedule
from tvm.s_tir.schedule import SBlockRV
+from tvm_tirx import sym, tirx
+from tvm_tirx.ir import TensorLoad
from tvm_tirx.target import Target
from ..analysis import detect_dominant_read, normalize_function
diff --git a/python/tvm/s_tir/meta_schedule/builder/local_builder.py b/python/tvm/s_tir/meta_schedule/builder/local_builder.py
index 95f2757a36..3fc3aaa5b0 100644
--- a/python/tvm/s_tir/meta_schedule/builder/local_builder.py
+++ b/python/tvm/s_tir/meta_schedule/builder/local_builder.py
@@ -24,10 +24,10 @@ from typing import Optional, Union
from tvm_ffi import register_global_func
-from tvm_tirx.ir import IRModule
-from tvm_tirx.runtime import Module, Tensor
from tvm.runtime.vm import load_param_dict, save_param_dict
from tvm.s_tir.utils import derived_object
+from tvm_tirx.ir import IRModule
+from tvm_tirx.runtime import Module, Tensor
from tvm_tirx.support.popen_pool import MapResult, PopenPoolExecutor, StatusKind
from tvm_tirx.target import Target
diff --git a/python/tvm/s_tir/meta_schedule/cost_model/mlp_model.py b/python/tvm/s_tir/meta_schedule/cost_model/mlp_model.py
index 3e2ca52fb1..18528ef6ae 100644
--- a/python/tvm/s_tir/meta_schedule/cost_model/mlp_model.py
+++ b/python/tvm/s_tir/meta_schedule/cost_model/mlp_model.py
@@ -32,10 +32,10 @@ import numpy as np # type: ignore
import torch # type: ignore
from tvm.s_tir.utils import derived_object
-from tvm_tirx.support.tar import tar, untar
-
from tvm_tirx.runtime import Tensor
+from tvm_tirx.support.tar import tar, untar
from tvm_tirx.target import Target
+
from ..cost_model import PyCostModel
from ..database import JSONDatabase
from ..feature_extractor import FeatureExtractor, PerStoreFeature
diff --git a/python/tvm/s_tir/meta_schedule/cost_model/xgb_model.py b/python/tvm/s_tir/meta_schedule/cost_model/xgb_model.py
index 1f8daafaee..04094213fb 100644
--- a/python/tvm/s_tir/meta_schedule/cost_model/xgb_model.py
+++ b/python/tvm/s_tir/meta_schedule/cost_model/xgb_model.py
@@ -26,9 +26,9 @@ from typing import TYPE_CHECKING, Any, Literal, NamedTuple, Optional
import numpy as np # type: ignore
from tvm.s_tir.utils import derived_object
+from tvm_tirx.runtime import Tensor
from tvm_tirx.support.tar import tar, untar
-from tvm_tirx.runtime import Tensor
from ..cost_model import PyCostModel
from ..feature_extractor import FeatureExtractor
from ..logging import get_logger
diff --git a/python/tvm/s_tir/meta_schedule/database/database.py b/python/tvm/s_tir/meta_schedule/database/database.py
index 1c59a65baa..d1420f22a6 100644
--- a/python/tvm/s_tir/meta_schedule/database/database.py
+++ b/python/tvm/s_tir/meta_schedule/database/database.py
@@ -27,9 +27,9 @@ from typing import Literal
from tvm_ffi import register_object
+from tvm.s_tir.schedule import Schedule, Trace
from tvm_tirx.ir.module import IRModule
from tvm_tirx.runtime import Object
-from tvm.s_tir.schedule import Schedule, Trace
from tvm_tirx.target import Target
from .. import _ffi_api
diff --git a/python/tvm/s_tir/meta_schedule/feature_extractor/random_feature_extractor.py b/python/tvm/s_tir/meta_schedule/feature_extractor/random_feature_extractor.py
index 63059869d5..a23bb88860 100644
--- a/python/tvm/s_tir/meta_schedule/feature_extractor/random_feature_extractor.py
+++ b/python/tvm/s_tir/meta_schedule/feature_extractor/random_feature_extractor.py
@@ -18,7 +18,7 @@
import numpy as np # type: ignore
-import tvm_tirx.runtime, tvm
+import tvm_tirx.runtime
from tvm.s_tir.utils import derived_object
from ..feature_extractor import PyFeatureExtractor
diff --git a/python/tvm/s_tir/meta_schedule/mutator/mutator.py b/python/tvm/s_tir/meta_schedule/mutator/mutator.py
index fb0061eab8..e506522d2d 100644
--- a/python/tvm/s_tir/meta_schedule/mutator/mutator.py
+++ b/python/tvm/s_tir/meta_schedule/mutator/mutator.py
@@ -27,8 +27,8 @@ from typing import Literal
from tvm_ffi import register_object
-from tvm_tirx.runtime import Object
from tvm.s_tir.schedule import Trace
+from tvm_tirx.runtime import Object
from .. import _ffi_api
diff --git a/python/tvm/s_tir/meta_schedule/postproc/postproc.py b/python/tvm/s_tir/meta_schedule/postproc/postproc.py
index c93890b365..f388f262b1 100644
--- a/python/tvm/s_tir/meta_schedule/postproc/postproc.py
+++ b/python/tvm/s_tir/meta_schedule/postproc/postproc.py
@@ -27,8 +27,8 @@ from typing import Literal
from tvm_ffi import register_object
-from tvm_tirx.runtime import Object
from tvm.s_tir.schedule import Schedule
+from tvm_tirx.runtime import Object
from .. import _ffi_api
diff --git a/python/tvm/s_tir/meta_schedule/relax_integration.py b/python/tvm/s_tir/meta_schedule/relax_integration.py
index 8a7f1cc39f..7a9174bb0c 100644
--- a/python/tvm/s_tir/meta_schedule/relax_integration.py
+++ b/python/tvm/s_tir/meta_schedule/relax_integration.py
@@ -19,7 +19,7 @@
import warnings
from typing import TYPE_CHECKING, Union
-import tvm, tvm_tirx
+import tvm_tirx
# isort: off
from typing import Literal
@@ -466,7 +466,9 @@ def compile_relax(
is_gpu_target = relax_pipeline_mod.BackendDispatcher.is_gpu_target(target)
@tvm_tirx.ir.transform.module_pass(opt_level=3)
- def _ms_pipeline(mod: tvm_tirx.ir.IRModule, _ctx: tvm_tirx.ir.transform.PassContext) -> tvm_tirx.ir.IRModule:
+ def _ms_pipeline(
+ mod: tvm_tirx.ir.IRModule, _ctx: tvm_tirx.ir.transform.PassContext
+ ) -> tvm_tirx.ir.IRModule:
fuse_seq = [
*dispatch_passes,
relax.transform.LegalizeOps(enable_warning=enable_warning),
diff --git a/python/tvm/s_tir/meta_schedule/runner/local_runner.py b/python/tvm/s_tir/meta_schedule/runner/local_runner.py
index 945eb8dd9a..39f8621436 100644
--- a/python/tvm/s_tir/meta_schedule/runner/local_runner.py
+++ b/python/tvm/s_tir/meta_schedule/runner/local_runner.py
@@ -21,11 +21,11 @@ import subprocess
from collections.abc import Callable
from contextlib import contextmanager
-import tvm, tvm_tirx
+import tvm_tirx
from tvm.s_tir.utils import derived_object
+from tvm_tirx.runtime import Device, Module
from tvm_tirx.support.popen_pool import PopenPoolExecutor
-from tvm_tirx.runtime import Device, Module
from ..logging import get_logger
from ..profiler import Profiler
from ..utils import get_global_func_with_default_on_worker
diff --git a/python/tvm/s_tir/meta_schedule/runner/rpc_runner.py b/python/tvm/s_tir/meta_schedule/runner/rpc_runner.py
index fb01df5805..ba27115d9c 100644
--- a/python/tvm/s_tir/meta_schedule/runner/rpc_runner.py
+++ b/python/tvm/s_tir/meta_schedule/runner/rpc_runner.py
@@ -22,8 +22,8 @@ from collections.abc import Callable
from contextlib import contextmanager
from tvm.rpc import RPCSession
-from tvm_tirx.runtime import Device, Module
from tvm.s_tir.utils import derived_object
+from tvm_tirx.runtime import Device, Module
from tvm_tirx.support.popen_pool import PopenPoolExecutor
from ..logging import get_logger
diff --git a/python/tvm/s_tir/meta_schedule/runner/utils.py b/python/tvm/s_tir/meta_schedule/runner/utils.py
index 9c86183cd3..ce135346f8 100644
--- a/python/tvm/s_tir/meta_schedule/runner/utils.py
+++ b/python/tvm/s_tir/meta_schedule/runner/utils.py
@@ -20,9 +20,9 @@ import itertools
from collections.abc import Callable
from typing import Any
-import tvm_tirx.runtime, tvm
-
+import tvm_tirx.runtime
from tvm_tirx.runtime import Device, Module
+
from .config import EvaluatorConfig
T_ARG_INFO_JSON_OBJ = list[Any] # pylint: disable=invalid-name
diff --git a/python/tvm/s_tir/meta_schedule/schedule/cuda/layout_transform.py b/python/tvm/s_tir/meta_schedule/schedule/cuda/layout_transform.py
index 717bb3ee13..64aa0acffe 100644
--- a/python/tvm/s_tir/meta_schedule/schedule/cuda/layout_transform.py
+++ b/python/tvm/s_tir/meta_schedule/schedule/cuda/layout_transform.py
@@ -20,7 +20,7 @@
import math
from collections import deque
-import tvm, tvm_tirx
+import tvm_tirx
from tvm.s_tir import meta_schedule
from tvm.s_tir.schedule import ExprRV, LoopRV, SBlockRV, Schedule
diff --git a/python/tvm/s_tir/meta_schedule/schedule_rule/schedule_rule.py b/python/tvm/s_tir/meta_schedule/schedule_rule/schedule_rule.py
index ca2f2b72b1..f6543cacb8 100644
--- a/python/tvm/s_tir/meta_schedule/schedule_rule/schedule_rule.py
+++ b/python/tvm/s_tir/meta_schedule/schedule_rule/schedule_rule.py
@@ -30,8 +30,8 @@ from typing import Literal
from tvm_ffi import register_object
-from tvm_tirx.runtime import Object
from tvm.s_tir.schedule import SBlockRV, Schedule
+from tvm_tirx.runtime import Object
from .. import _ffi_api
diff --git a/python/tvm/s_tir/meta_schedule/search_strategy/search_strategy.py b/python/tvm/s_tir/meta_schedule/search_strategy/search_strategy.py
index e171348dbb..9541d21201 100644
--- a/python/tvm/s_tir/meta_schedule/search_strategy/search_strategy.py
+++ b/python/tvm/s_tir/meta_schedule/search_strategy/search_strategy.py
@@ -29,8 +29,8 @@ from typing import Literal
# isort: on
from tvm_ffi import register_object
-from tvm_tirx.runtime import Object
from tvm.s_tir.schedule import Schedule
+from tvm_tirx.runtime import Object
from .. import _ffi_api
from ..arg_info import ArgInfo
diff --git a/python/tvm/s_tir/meta_schedule/space_generator/space_generator.py b/python/tvm/s_tir/meta_schedule/space_generator/space_generator.py
index ea4039aba6..20afd791a2 100644
--- a/python/tvm/s_tir/meta_schedule/space_generator/space_generator.py
+++ b/python/tvm/s_tir/meta_schedule/space_generator/space_generator.py
@@ -29,9 +29,9 @@ from typing import Literal
# isort: on
from tvm_ffi import register_object
+from tvm.s_tir.schedule import Schedule
from tvm_tirx.ir import IRModule
from tvm_tirx.runtime import Object
-from tvm.s_tir.schedule import Schedule
from .. import _ffi_api
diff --git a/python/tvm/s_tir/meta_schedule/testing/custom_builder_runner.py b/python/tvm/s_tir/meta_schedule/testing/custom_builder_runner.py
index fcc7345c53..1d03b20b63 100644
--- a/python/tvm/s_tir/meta_schedule/testing/custom_builder_runner.py
+++ b/python/tvm/s_tir/meta_schedule/testing/custom_builder_runner.py
@@ -21,8 +21,8 @@ from collections.abc import Callable
import numpy as np # type: ignore
-from tvm_tirx.runtime import Executable, Module
from tvm.s_tir.meta_schedule.runner import RPCConfig
+from tvm_tirx.runtime import Executable, Module
def run_module_via_rpc(
diff --git a/python/tvm/s_tir/meta_schedule/testing/space_generation.py b/python/tvm/s_tir/meta_schedule/testing/space_generation.py
index f1ec0b2988..9a7dbeda54 100644
--- a/python/tvm/s_tir/meta_schedule/testing/space_generation.py
+++ b/python/tvm/s_tir/meta_schedule/testing/space_generation.py
@@ -23,11 +23,11 @@ from typing import Literal
import tvm_ffi
-from tvm_tirx.ir import IRModule
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.schedule import Trace
from tvm.s_tir.schedule.testing import verify_trace_roundtrip
+from tvm_tirx.ir import IRModule
from tvm_tirx.target import Target
diff --git a/python/tvm/s_tir/meta_schedule/testing/te_workload.py b/python/tvm/s_tir/meta_schedule/testing/te_workload.py
index a5d7ccad81..a1c391997f 100644
--- a/python/tvm/s_tir/meta_schedule/testing/te_workload.py
+++ b/python/tvm/s_tir/meta_schedule/testing/te_workload.py
@@ -19,8 +19,8 @@
# pylint: disable=missing-docstring
-from tvm_tirx import tirx
from tvm import te, topi
+from tvm_tirx import tirx
from tvm_tirx.target import Target
diff --git a/python/tvm/s_tir/meta_schedule/testing/tune_te.py b/python/tvm/s_tir/meta_schedule/testing/tune_te.py
index 0510ea5dd1..95b375f1fd 100644
--- a/python/tvm/s_tir/meta_schedule/testing/tune_te.py
+++ b/python/tvm/s_tir/meta_schedule/testing/tune_te.py
@@ -19,7 +19,7 @@
import argparse
import logging
-import tvm, tvm_tirx
+import tvm_tirx
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing.te_workload import create_te_workload
from tvm.s_tir.testing import strtobool
diff --git a/python/tvm/s_tir/meta_schedule/testing/tune_utils.py b/python/tvm/s_tir/meta_schedule/testing/tune_utils.py
index d0610c0aec..7112da246d 100644
--- a/python/tvm/s_tir/meta_schedule/testing/tune_utils.py
+++ b/python/tvm/s_tir/meta_schedule/testing/tune_utils.py
@@ -20,7 +20,7 @@ from collections.abc import Callable
import numpy as np # type: ignore
-import tvm, tvm_tirx
+import tvm_tirx
from tvm_tirx.runtime import Tensor
diff --git a/python/tvm/s_tir/meta_schedule/testing/validate_database.py b/python/tvm/s_tir/meta_schedule/testing/validate_database.py
index 04d2d37f4f..e812fe72ac 100644
--- a/python/tvm/s_tir/meta_schedule/testing/validate_database.py
+++ b/python/tvm/s_tir/meta_schedule/testing/validate_database.py
@@ -29,8 +29,8 @@ import numpy as np # type: ignore
import tvm_ffi
from tvm_ffi import get_global_func, register_global_func
-import tvm, tvm_tirx
-from tvm_tirx.ir import IRModule
+import tvm
+import tvm_tirx
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing.tune_utils import generate_input_data
@@ -38,6 +38,7 @@ from tvm.s_tir.meta_schedule.utils import remove_build_dir
from tvm.s_tir.schedule import Trace
from tvm.s_tir.tensor_intrin import * # type: ignore # pylint: disable=wildcard-import,unused-wildcard-import
from tvm.s_tir.testing import strtobool
+from tvm_tirx.ir import IRModule
from tvm_tirx.support import describe
from tvm_tirx.target import Target
@@ -443,7 +444,8 @@ def make_alloc_arg_and_check(
The list of argument lists.
"""
return [
- [tvm_tirx.runtime.tensor(arg, device=device) for arg in inputs] for _ in range(alloc_repeat)
+ [tvm_tirx.runtime.tensor(arg, device=device) for arg in inputs]
+ for _ in range(alloc_repeat)
]
def f_with_args_run_evaluator_common(
diff --git a/python/tvm/s_tir/meta_schedule/tir_integration.py b/python/tvm/s_tir/meta_schedule/tir_integration.py
index 051bedbd54..3a51927607 100644
--- a/python/tvm/s_tir/meta_schedule/tir_integration.py
+++ b/python/tvm/s_tir/meta_schedule/tir_integration.py
@@ -23,8 +23,8 @@ from typing import Literal
from tvm_ffi import register_global_func
# isort: on
-from tvm_tirx import ir, tirx
from tvm.s_tir.schedule import Schedule as _Schedule
+from tvm_tirx import ir, tirx
from tvm_tirx.target import Target
from tvm_tirx.tirx.expr import IntImm
diff --git a/python/tvm/s_tir/meta_schedule/trace_apply.py b/python/tvm/s_tir/meta_schedule/trace_apply.py
index 04ce73a05d..aad7327406 100644
--- a/python/tvm/s_tir/meta_schedule/trace_apply.py
+++ b/python/tvm/s_tir/meta_schedule/trace_apply.py
@@ -17,8 +17,8 @@
"""Specialized applications of trace"""
from tvm.s_tir import Schedule, Trace
-
from tvm_tirx.target import Target
+
from . import _ffi_api
diff --git a/python/tvm/s_tir/meta_schedule/tune_context.py b/python/tvm/s_tir/meta_schedule/tune_context.py
index 825f79a35c..2c1ad072ff 100644
--- a/python/tvm/s_tir/meta_schedule/tune_context.py
+++ b/python/tvm/s_tir/meta_schedule/tune_context.py
@@ -26,9 +26,9 @@ from tvm_ffi import register_object, register_global_func
# isort: on
+from tvm.s_tir import Schedule
from tvm_tirx import IRModule
from tvm_tirx.runtime import Object
-from tvm.s_tir import Schedule
from tvm_tirx.target import Target
from tvm_tirx.tirx import Function
diff --git a/python/tvm/s_tir/meta_schedule/utils.py b/python/tvm/s_tir/meta_schedule/utils.py
index d43fd9dd3e..2982f5e65e 100644
--- a/python/tvm/s_tir/meta_schedule/utils.py
+++ b/python/tvm/s_tir/meta_schedule/utils.py
@@ -24,8 +24,8 @@ from typing import Any
import numpy as np # type: ignore
from tvm_ffi import Array, Function, Map, get_global_func, register_global_func
-from tvm_tirx.ir import IRModule
from tvm.rpc import RPCSession
+from tvm_tirx.ir import IRModule
from tvm_tirx.tirx import FloatImm, IntImm
diff --git a/python/tvm/s_tir/pipeline.py b/python/tvm/s_tir/pipeline.py
index 5f34c2dbad..84473e3f99 100644
--- a/python/tvm/s_tir/pipeline.py
+++ b/python/tvm/s_tir/pipeline.py
@@ -20,9 +20,9 @@
import tvm_ffi
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm_tirx
from tvm import s_tir
+from tvm_tirx import tirx
from tvm_tirx.tirx import compilation_pipeline as tir_pipeline
tir = tirx # alias for backward compat
diff --git a/python/tvm/s_tir/sblock_dependence_info.py b/python/tvm/s_tir/sblock_dependence_info.py
index 94d542d905..21af539cb1 100644
--- a/python/tvm/s_tir/sblock_dependence_info.py
+++ b/python/tvm/s_tir/sblock_dependence_info.py
@@ -19,9 +19,9 @@ to store the block level dependences"""
from tvm_ffi import register_object
+from tvm.s_tir import SBlock
from tvm_tirx.ir.module import IRModule
from tvm_tirx.runtime import Object
-from tvm.s_tir import SBlock
from tvm_tirx.tirx import Function
from . import _ffi_api
diff --git a/python/tvm/s_tir/sblock_scope.py b/python/tvm/s_tir/sblock_scope.py
index a4919c57dd..13cc9de92f 100644
--- a/python/tvm/s_tir/sblock_scope.py
+++ b/python/tvm/s_tir/sblock_scope.py
@@ -21,9 +21,9 @@ from typing import Optional
from tvm_ffi import register_object
+from tvm.s_tir import SBlock
from tvm_tirx.ir import For
from tvm_tirx.runtime import Object
-from tvm.s_tir import SBlock
from . import _ffi_api
diff --git a/python/tvm/s_tir/schedule/schedule.py b/python/tvm/s_tir/schedule/schedule.py
index 6655e76ca8..5f58c4b65d 100644
--- a/python/tvm/s_tir/schedule/schedule.py
+++ b/python/tvm/s_tir/schedule/schedule.py
@@ -22,10 +22,10 @@ from typing import Literal
from tvm_ffi import register_object as _register_object
+from tvm.s_tir import SBlock
from tvm_tirx.error import register_error
from tvm_tirx.ir import Expr, For, GlobalVar, IRModule, is_prim_expr
from tvm_tirx.runtime import DataTypeCode, MemoryScope, Object
-from tvm.s_tir import SBlock
from tvm_tirx.tirx import FloatImm, Function, IntImm, Var, is_tensor_var
from tvm_tirx.tirx.function import IndexMap
diff --git a/python/tvm/s_tir/schedule/state.py b/python/tvm/s_tir/schedule/state.py
index f00e8b2157..7279573047 100644
--- a/python/tvm/s_tir/schedule/state.py
+++ b/python/tvm/s_tir/schedule/state.py
@@ -22,9 +22,9 @@ from enum import IntEnum
from tvm_ffi import register_object
+from tvm.s_tir import SBlock, SBlockRealize
from tvm_tirx.ir import For, IRModule
from tvm_tirx.runtime import Object
-from tvm.s_tir import SBlock, SBlockRealize
from tvm_tirx.tirx import Function
from ..sblock_scope import SBlockScope, StmtSRef
diff --git a/python/tvm/s_tir/schedule/testing.py b/python/tvm/s_tir/schedule/testing.py
index 50650f407c..c413754031 100644
--- a/python/tvm/s_tir/schedule/testing.py
+++ b/python/tvm/s_tir/schedule/testing.py
@@ -20,9 +20,9 @@
from collections.abc import Sequence
from typing import Any
-import tvm, tvm_tirx
-from tvm_tirx.ir import IRModule, assert_structural_equal
+import tvm_tirx
from tvm.s_tir.schedule import Schedule, Trace
+from tvm_tirx.ir import IRModule, assert_structural_equal
from tvm_tirx.tirx import Function
diff --git a/python/tvm/s_tir/schedule/trace.py b/python/tvm/s_tir/schedule/trace.py
index 5d7c677c50..44c6342df8 100644
--- a/python/tvm/s_tir/schedule/trace.py
+++ b/python/tvm/s_tir/schedule/trace.py
@@ -23,11 +23,11 @@ from typing import TYPE_CHECKING, Any
from tvm_ffi import Array, Map
from tvm_ffi import register_object as _register_object
+from tvm_tirx.ir import save_json
from tvm_tirx.runtime import Object
from tvm_tirx.tirx.expr import FloatImm, IntImm
from tvm_tirx.tirx.function import IndexMap
-from tvm_tirx.ir import save_json
from . import _ffi_api
from .instruction import ATTR_TYPE, INPUT_RV_TYPE, Instruction
diff --git a/python/tvm/s_tir/script/ir_builder/ir.py b/python/tvm/s_tir/script/ir_builder/ir.py
index 78981c4e48..6bd52f8b50 100644
--- a/python/tvm/s_tir/script/ir_builder/ir.py
+++ b/python/tvm/s_tir/script/ir_builder/ir.py
@@ -21,17 +21,17 @@ import threading
from numbers import Integral
from typing import Any
+from tvm.s_tir.iter_var import IterVar
+from tvm.s_tir.op import comm_reducer
+from tvm.te import CommReducer, Reduce
from tvm_tirx import ir
from tvm_tirx import tirx as tir
from tvm_tirx.ir import TensorLoad, TensorRegion, is_prim_expr
from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from tvm_tirx.ir.type import _normalize_type
from tvm_tirx.runtime import MemoryScope
-from tvm.s_tir.iter_var import IterVar
-from tvm.s_tir.op import comm_reducer
from tvm_tirx.script.ir_builder.base import MISSING, auto_emit_void_call
from tvm_tirx.script.parser.protocol_registry import register_mutable_decl as _register_mutable_decl
-from tvm.te import CommReducer, Reduce
from tvm_tirx.tirx import Expr, IntImm, Var
from tvm_tirx.tirx.layout import Layout
from tvm_tirx.tirx.script.ir_builder.ir import _get_layout, _record_meta_resource
diff --git a/python/tvm/s_tir/stmt.py b/python/tvm/s_tir/stmt.py
index f43f56d258..05bed5bcb2 100644
--- a/python/tvm/s_tir/stmt.py
+++ b/python/tvm/s_tir/stmt.py
@@ -20,10 +20,10 @@ from collections.abc import Mapping, Sequence
import tvm_ffi
+from tvm.s_tir import IterVar
from tvm_tirx.ir import Expr, Location, Scriptable, SeqStmt, Stmt, TensorRegion, Var, const
from tvm_tirx.ir.location import UNKNOWN_LOC
from tvm_tirx.runtime import Object
-from tvm.s_tir import IterVar
from . import _ffi_api
diff --git a/python/tvm/s_tir/testing/runner.py b/python/tvm/s_tir/testing/runner.py
index f0e995a70c..1b56bc35d3 100644
--- a/python/tvm/s_tir/testing/runner.py
+++ b/python/tvm/s_tir/testing/runner.py
@@ -23,8 +23,8 @@ from typing import TYPE_CHECKING, Literal, Optional, Union
if TYPE_CHECKING:
import numpy as np
- from tvm_tirx.runtime import Device, Module, Tensor
from tvm.s_tir.meta_schedule.runner import EvaluatorConfig, RPCConfig
+ from tvm_tirx.runtime import Device, Module, Tensor
# pylint: disable=import-outside-toplevel,protected-access
@@ -123,8 +123,8 @@ def local_run( # pylint: disable=too-many-arguments,too-many-locals
import os.path as osp
import tempfile
- from tvm_tirx.runtime import device, load_module
from tvm.s_tir.meta_schedule.runner import EvaluatorConfig
+ from tvm_tirx.runtime import device, load_module
evaluator_config = EvaluatorConfig._normalized(evaluator_config)
export_func, output_format = _normalize_export_func(export_func, output_format)
diff --git a/python/tvm/s_tir/testing/tir.py b/python/tvm/s_tir/testing/tir.py
index f3d8c58ccc..e7c45933af 100644
--- a/python/tvm/s_tir/testing/tir.py
+++ b/python/tvm/s_tir/testing/tir.py
@@ -37,7 +37,8 @@ def mfma_schedule(
shared_scope="shared",
):
"""Create a tensorized schedule for GEMM with MFMA intrinsics."""
- import tvm, tvm_tirx
+ import tvm
+ import tvm_tirx
ir_module = tvm_tirx.IRModule({"main": workload})
sch = tvm.s_tir.Schedule(ir_module)
diff --git a/python/tvm/te/operation.py b/python/tvm/te/operation.py
index bf29094864..53d6b90d77 100644
--- a/python/tvm/te/operation.py
+++ b/python/tvm/te/operation.py
@@ -23,9 +23,10 @@ from numbers import Integral as _Integral
from tvm_ffi import Array
-import tvm_tirx.sym._ffi_api, tvm
-import tvm_tirx.tirx, tvm
-import tvm_tirx.tirx._ffi_api, tvm
+import tvm
+import tvm_tirx.sym._ffi_api
+import tvm_tirx.tirx
+import tvm_tirx.tirx._ffi_api
from tvm_tirx.ir import is_prim_expr
from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from tvm_tirx.runtime import convert
@@ -355,7 +356,9 @@ def extern(
return res[0] if len(res) == 1 else res
-def extern_function(input_tensors: list[_tensor.Tensor], function: tvm_tirx.tirx.Function, **kwargs):
+def extern_function(
+ input_tensors: list[_tensor.Tensor], function: tvm_tirx.tirx.Function, **kwargs
+):
"""Compute tensors via a schedulable TIR Function
Parameters
diff --git a/python/tvm/te/reduction.py b/python/tvm/te/reduction.py
index f1d7bd8d04..fa56018df9 100644
--- a/python/tvm/te/reduction.py
+++ b/python/tvm/te/reduction.py
@@ -22,7 +22,7 @@ from typing import TYPE_CHECKING
import tvm_ffi
from tvm_ffi import Array
-import tvm, tvm_tirx
+import tvm_tirx
from tvm_tirx import ir
from tvm_tirx import tirx as tir
from tvm_tirx.ir import Expr, Scriptable, Var, const
diff --git a/python/tvm/te/tag.py b/python/tvm/te/tag.py
index bdf2103f16..29c4754c56 100644
--- a/python/tvm/te/tag.py
+++ b/python/tvm/te/tag.py
@@ -91,6 +91,8 @@ def tag_scope(tag):
# or use tag_scope as decorator
@tvm.te.tag_scope(tag="conv")
def compute_relu(data):
- return te.compute(data.shape, lambda *i: tvm_tirx.tirx.Select(data(*i) < 0, 0.0, data(*i)))
+ return te.compute(
+ data.shape, lambda *i: tvm_tirx.tirx.Select(data(*i) < 0, 0.0, data(*i))
+ )
"""
return TagScope(tag)
diff --git a/python/tvm/topi/_te_tensor_overload.py b/python/tvm/topi/_te_tensor_overload.py
index c9ba8111a6..2d0ce42c98 100644
--- a/python/tvm/topi/_te_tensor_overload.py
+++ b/python/tvm/topi/_te_tensor_overload.py
@@ -17,10 +17,10 @@
"""Register TOPI implementations for TE tensor overload hooks."""
from tvm import te
+from tvm.te import _te_tensor_overload as _overload
from tvm_tirx.ir import _overload_prim_expr as _expr
from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from tvm_tirx.runtime import DataTypeCode
-from tvm.te import _te_tensor_overload as _overload
from . import broadcast as _broadcast
from . import math as _math
diff --git a/python/tvm/topi/gpu/scan.py b/python/tvm/topi/gpu/scan.py
index ab72e2ab9a..937ae36b74 100644
--- a/python/tvm/topi/gpu/scan.py
+++ b/python/tvm/topi/gpu/scan.py
@@ -20,7 +20,8 @@
import operator
from collections.abc import Callable
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import te
from tvm.contrib.thrust import can_use_rocthrust, can_use_thrust
from tvm_tirx.script.ir_builder import IRBuilder
@@ -158,7 +159,8 @@ def exclusive_scan_ir(data, output, reduction=None, binop=operator.add, identity
T.tensor_store(
middle,
T.tensor_indices(middle, 0),
- start[T.tensor_indices(start, 0)] + tvm_tirx.tirx.indexdiv(width, 2),
+ start[T.tensor_indices(start, 0)]
+ + tvm_tirx.tirx.indexdiv(width, 2),
)
T.tensor_store(
end,
@@ -257,7 +259,8 @@ def exclusive_scan_ir(data, output, reduction=None, binop=operator.add, identity
T.tensor_store(
middle,
T.tensor_indices(middle, 0),
- start[T.tensor_indices(start, 0)] + tvm_tirx.tirx.indexdiv(width, 2),
+ start[T.tensor_indices(start, 0)]
+ + tvm_tirx.tirx.indexdiv(width, 2),
)
T.tensor_store(
end,
diff --git a/python/tvm/topi/gpu/scatter_elements.py b/python/tvm/topi/gpu/scatter_elements.py
index 155fd07ba7..7fb95ad036 100644
--- a/python/tvm/topi/gpu/scatter_elements.py
+++ b/python/tvm/topi/gpu/scatter_elements.py
@@ -17,9 +17,9 @@
# pylint: disable=invalid-name
"""scatter_elements related operators"""
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm_tirx
from tvm import te
+from tvm_tirx import tirx
from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.tirx.script import ir_builder as T
diff --git a/python/tvm/topi/gpu/scatter_nd.py b/python/tvm/topi/gpu/scatter_nd.py
index 8abc5e6580..67f780979c 100644
--- a/python/tvm/topi/gpu/scatter_nd.py
+++ b/python/tvm/topi/gpu/scatter_nd.py
@@ -18,9 +18,9 @@
# ruff: noqa: E741
"""scatter_nd related operators"""
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm_tirx
from tvm import te # hide redefinition of min and max
+from tvm_tirx import tirx
from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.tirx.script import ir_builder as T
diff --git a/python/tvm/topi/gpu/sort.py b/python/tvm/topi/gpu/sort.py
index 31384b4afd..cfe9b313a9 100644
--- a/python/tvm/topi/gpu/sort.py
+++ b/python/tvm/topi/gpu/sort.py
@@ -18,7 +18,8 @@
# ruff: noqa: RUF005
"""Sort related operators"""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import te
from tvm.s_tir.script import ir_builder as Ts
from tvm_tirx.script.ir_builder import IRBuilder
@@ -148,7 +149,9 @@ def _odd_even_sort(
],
)
- T.evaluate(tvm_tirx.ir.Call("tirx.gpu_storage_sync", [tvm_tirx.ir.StringImm("shared")], ty="void"))
+ T.evaluate(
+ tvm_tirx.ir.Call("tirx.gpu_storage_sync", [tvm_tirx.ir.StringImm("shared")], ty="void")
+ )
idxm = tvm_tirx.tirx.indexmod
# OddEvenTransposeSort
@@ -177,7 +180,9 @@ def _odd_even_sort(
)
T.tensor_store(tmp_values_swap, [tid + n + 1], temp_values[0])
T.evaluate(
- tvm_tirx.ir.Call("tirx.gpu_storage_sync", [tvm_tirx.ir.StringImm("shared")], ty="void")
+ tvm_tirx.ir.Call(
+ "tirx.gpu_storage_sync", [tvm_tirx.ir.StringImm("shared")], ty="void"
+ )
)
## Copy sorted data to output
@@ -607,7 +612,9 @@ def _sort_common(
# calculate the start, mid, and end points of this section
start_pos = width * bz
- middle = cast(tvm.te.min(start_pos + tvm_tirx.tirx.indexdiv(width, 2), size), target_dtype)
+ middle = cast(
+ tvm.te.min(start_pos + tvm_tirx.tirx.indexdiv(width, 2), size), target_dtype
+ )
end = cast(tvm.te.min(start_pos + width, size), target_dtype)
with T.if_(start_pos < size):
@@ -652,7 +659,9 @@ def _sort_common(
## if the final sorted data ended up in the swap, copy it to the real output
nthread_bx = ceil_div(size, nthread_tx)
with T.if_(
- tvm_tirx.tirx.all(upper_lim > lower_lim, tvm_tirx.tirx.indexmod(upper_lim - lower_lim, 2) == 1)
+ tvm_tirx.tirx.all(
+ upper_lim > lower_lim, tvm_tirx.tirx.indexmod(upper_lim - lower_lim, 2) == 1
+ )
):
with T.then_():
with T.frame_scope(_thread_launches(nthread_tx, nthread_bx, nthread_by)) as (
@@ -834,7 +843,9 @@ def sort_thrust(data, axis=-1, is_ascend=1, workspace=None):
value_buf = tvm_tirx.tirx.decl_tensor(
data.shape, data.dtype, "value_buf", data_alignment=8, layout=None
)
- indices_buf = tvm_tirx.tirx.decl_tensor(data.shape, dtype, "out_buf", data_alignment=8, layout=None)
+ indices_buf = tvm_tirx.tirx.decl_tensor(
+ data.shape, dtype, "out_buf", data_alignment=8, layout=None
+ )
def f_compute(ins, outs):
args = ["tvm.contrib.thrust.sort", ins[0], outs[0], outs[1], is_ascend]
@@ -901,7 +912,9 @@ def argsort(data, axis=-1, is_ascend=1, dtype="float32", ret_type="indices"):
value_swap_buf = tvm_tirx.tirx.decl_tensor(
data.shape, data.dtype, "value_swap_buf", data_alignment=8, layout=None
)
- indices_buf = tvm_tirx.tirx.decl_tensor(data.shape, dtype, "out_buf", data_alignment=8, layout=None)
+ indices_buf = tvm_tirx.tirx.decl_tensor(
+ data.shape, dtype, "out_buf", data_alignment=8, layout=None
+ )
indices_swap_buf = tvm_tirx.tirx.decl_tensor(
data.shape, dtype, "out_swap_buf", data_alignment=8, layout=None
)
@@ -1130,7 +1143,9 @@ def topk_thrust(
else:
workspace_buf = None
out_bufs = [
- tvm_tirx.tirx.decl_tensor(data.shape, data.dtype, "value_buf", data_alignment=8, layout=None),
+ tvm_tirx.tirx.decl_tensor(
+ data.shape, data.dtype, "value_buf", data_alignment=8, layout=None
+ ),
tvm_tirx.tirx.decl_tensor(data.shape, dtype, "indices_buf", data_alignment=8, layout=None),
]
@@ -1222,7 +1237,9 @@ def searchsorted(sorted_sequence, values, right=False, out_dtype="int64"):
values_ptr = values_buf
indices_ptr = indices_buf
- max_threads = int(tvm_tirx.target.Target.current(allow_none=False).attrs["max_num_threads"])
+ max_threads = int(
+ tvm_tirx.target.Target.current(allow_none=False).attrs["max_num_threads"]
+ )
nthread_tx = max_threads
nthread_bx = ceil_div(num_search, nthread_tx)
with T.frame_scope(
diff --git a/python/tvm/topi/image/grid_sample.py b/python/tvm/topi/image/grid_sample.py
index c49c44358a..50b395badd 100644
--- a/python/tvm/topi/image/grid_sample.py
+++ b/python/tvm/topi/image/grid_sample.py
@@ -17,8 +17,8 @@
# pylint: disable=invalid-name
"""affine_grid and grid_sample operator"""
-from tvm_tirx import tirx
from tvm import te
+from tvm_tirx import tirx
def affine_grid(data, target_shape, align_corners=True):
diff --git a/python/tvm/topi/image/resize.py b/python/tvm/topi/image/resize.py
index 1c252db001..d0c5b07b48 100644
--- a/python/tvm/topi/image/resize.py
+++ b/python/tvm/topi/image/resize.py
@@ -18,7 +18,8 @@
# ruff: noqa: E741, F821
"""TVM operator input resize compute."""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import te
from tvm.topi.utils import nchw_pack_layout, nchw_xc_layout
diff --git a/python/tvm/topi/index_put.py b/python/tvm/topi/index_put.py
index 22d13eb3cf..03948f9870 100644
--- a/python/tvm/topi/index_put.py
+++ b/python/tvm/topi/index_put.py
@@ -17,8 +17,8 @@
# pylint: disable=invalid-name
"""IndexPut operator"""
-from tvm_tirx import tirx
from tvm import te
+from tvm_tirx import tirx
from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.tirx.script import ir_builder as T
diff --git a/python/tvm/topi/math.py b/python/tvm/topi/math.py
index e25c3b3acd..2a002975f0 100644
--- a/python/tvm/topi/math.py
+++ b/python/tvm/topi/math.py
@@ -17,9 +17,10 @@
"""Elementwise operators"""
# pylint: disable=redefined-builtin,unused-argument
-import tvm, tvm_tirx
-from tvm_tirx import DataTypeCode
+import tvm
+import tvm_tirx
from tvm import te
+from tvm_tirx import DataTypeCode
from tvm_tirx.ir.location import UNKNOWN_LOC, Location
from . import cpp, tag
diff --git a/python/tvm/topi/nn/batch_matmul.py b/python/tvm/topi/nn/batch_matmul.py
index d27194a65f..c709cf09d7 100644
--- a/python/tvm/topi/nn/batch_matmul.py
+++ b/python/tvm/topi/nn/batch_matmul.py
@@ -20,7 +20,7 @@
# pylint: disable=invalid-name
import logging
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
from ..utils import get_const_tuple
diff --git a/python/tvm/topi/nn/bitserial_conv2d.py b/python/tvm/topi/nn/bitserial_conv2d.py
index 50be9a49c8..0dbafcd15a 100644
--- a/python/tvm/topi/nn/bitserial_conv2d.py
+++ b/python/tvm/topi/nn/bitserial_conv2d.py
@@ -18,7 +18,7 @@
# pylint: disable=unused-argument, redefined-builtin
"""Bitserial Conv2D operators"""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
from ..utils import get_const_tuple
diff --git a/python/tvm/topi/nn/bitserial_dense.py b/python/tvm/topi/nn/bitserial_dense.py
index 5a5ec27d90..bb0a510b2e 100644
--- a/python/tvm/topi/nn/bitserial_dense.py
+++ b/python/tvm/topi/nn/bitserial_dense.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name, too-many-locals, too-many-arguments
"""Bitserial Dense operator."""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
from tvm.topi.utils import get_const_tuple
@@ -70,7 +70,9 @@ def bitserial_dense(
matmul = te.compute(
oshape,
lambda i, j: te.sum(
- tvm_tirx.tirx.popcount(weight_packed[j, wb, k] & data_packed[i, db, k]).astype(out_dtype)
+ tvm_tirx.tirx.popcount(weight_packed[j, wb, k] & data_packed[i, db, k]).astype(
+ out_dtype
+ )
<< (db + wb).astype(out_dtype),
axis=[wb, db, k],
),
diff --git a/python/tvm/topi/nn/bitserial_util.py b/python/tvm/topi/nn/bitserial_util.py
index bdf64c40e5..6cfd4967af 100644
--- a/python/tvm/topi/nn/bitserial_util.py
+++ b/python/tvm/topi/nn/bitserial_util.py
@@ -19,7 +19,7 @@
import numpy as np
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
from tvm.topi.transform import concatenate
diff --git a/python/tvm/topi/nn/bnn.py b/python/tvm/topi/nn/bnn.py
index 1a57b3ae33..c951b2a1e4 100644
--- a/python/tvm/topi/nn/bnn.py
+++ b/python/tvm/topi/nn/bnn.py
@@ -16,7 +16,7 @@
# under the License.
"""Binary Neural Network (BNN) Operators"""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
from .. import tag
diff --git a/python/tvm/topi/nn/conv2d.py b/python/tvm/topi/nn/conv2d.py
index 7579b486cd..873fb83ee9 100644
--- a/python/tvm/topi/nn/conv2d.py
+++ b/python/tvm/topi/nn/conv2d.py
@@ -25,7 +25,7 @@ from collections.abc import Sequence
import numpy as np
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
from ..utils import get_const_int, get_const_tuple, simplify, tag
diff --git a/python/tvm/topi/nn/deformable_conv2d.py b/python/tvm/topi/nn/deformable_conv2d.py
index 0212647e48..9acc522250 100644
--- a/python/tvm/topi/nn/deformable_conv2d.py
+++ b/python/tvm/topi/nn/deformable_conv2d.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name, too-many-locals, too-many-arguments
"""Deformable Conv2D operators"""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
from ..cpp.utils import bilinear_sample_nchw, bilinear_sample_nhwc
diff --git a/python/tvm/topi/nn/dense.py b/python/tvm/topi/nn/dense.py
index 820ba534f5..862bac6bdd 100644
--- a/python/tvm/topi/nn/dense.py
+++ b/python/tvm/topi/nn/dense.py
@@ -18,7 +18,7 @@
# ruff: noqa: E741, F821
"""TVM operator fully connected compute."""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
from .. import add, tag
@@ -105,7 +105,12 @@ def matmul(
batch_dims_b = [1] * (result_ndim - len(batch_dims_b)) + batch_dims_b
for idx, (l, r) in enumerate(zip(batch_dims_a, batch_dims_b)):
- if not tvm_tirx.ir.is_prim_var(l) and not tvm_tirx.ir.is_prim_var(r) and int(l) != 1 and int(r) != 1:
+ if (
+ not tvm_tirx.ir.is_prim_var(l)
+ and not tvm_tirx.ir.is_prim_var(r)
+ and int(l) != 1
+ and int(r) != 1
+ ):
assert int(l) == int(r), (
"Batch dimensions of dense do not match: "
f"{tensor_a.shape[:-2]} vs {tensor_b.shape[:-2]}."
diff --git a/python/tvm/topi/nn/depth_to_space.py b/python/tvm/topi/nn/depth_to_space.py
index d4ef288c66..e851722340 100644
--- a/python/tvm/topi/nn/depth_to_space.py
+++ b/python/tvm/topi/nn/depth_to_space.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name
"""TVM operator depth_to_space compute."""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
from .. import tag
diff --git a/python/tvm/topi/nn/depthwise_conv2d.py b/python/tvm/topi/nn/depthwise_conv2d.py
index f3cb8162e8..dce41304c7 100644
--- a/python/tvm/topi/nn/depthwise_conv2d.py
+++ b/python/tvm/topi/nn/depthwise_conv2d.py
@@ -22,7 +22,7 @@ from collections import namedtuple
import numpy as np
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
from ..utils import get_const_tuple, simplify
diff --git a/python/tvm/topi/nn/dilate.py b/python/tvm/topi/nn/dilate.py
index 5278c20a23..578167a1e2 100644
--- a/python/tvm/topi/nn/dilate.py
+++ b/python/tvm/topi/nn/dilate.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name
"""Dilation operators"""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
from .. import tag, utils
diff --git a/python/tvm/topi/nn/elemwise.py b/python/tvm/topi/nn/elemwise.py
index 176b0a58c9..fcb5d587da 100644
--- a/python/tvm/topi/nn/elemwise.py
+++ b/python/tvm/topi/nn/elemwise.py
@@ -16,7 +16,8 @@
# under the License.
"""Elementwise operators"""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import te
from .. import tag
diff --git a/python/tvm/topi/nn/fifo_buffer.py b/python/tvm/topi/nn/fifo_buffer.py
index 184a31fafd..977270b43e 100644
--- a/python/tvm/topi/nn/fifo_buffer.py
+++ b/python/tvm/topi/nn/fifo_buffer.py
@@ -18,7 +18,8 @@
"""FIFO buffer op"""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import te
from .. import tag
diff --git a/python/tvm/topi/nn/flatten.py b/python/tvm/topi/nn/flatten.py
index a264a6c0da..35e1956697 100644
--- a/python/tvm/topi/nn/flatten.py
+++ b/python/tvm/topi/nn/flatten.py
@@ -16,7 +16,8 @@
# under the License.
"""TVM operator flatten compute."""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import te
from .. import tag
diff --git a/python/tvm/topi/nn/lstm.py b/python/tvm/topi/nn/lstm.py
index fe7ab070b9..10f5185672 100644
--- a/python/tvm/topi/nn/lstm.py
+++ b/python/tvm/topi/nn/lstm.py
@@ -18,9 +18,9 @@
# ruff: noqa: E731
"""General LSTM implementation using TE scan."""
-from tvm_tirx import tirx
from tvm import te
from tvm.topi import tag
+from tvm_tirx import tirx
def lstm(
diff --git a/python/tvm/topi/nn/mapping.py b/python/tvm/topi/nn/mapping.py
index c9ad384c7d..dbc3ab56d9 100644
--- a/python/tvm/topi/nn/mapping.py
+++ b/python/tvm/topi/nn/mapping.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name, line-too-long
"""Operators of one-to-one-mapping on the first input"""
-import tvm, tvm_tirx
+import tvm
from tvm import te
from .. import tag
diff --git a/python/tvm/topi/nn/pad.py b/python/tvm/topi/nn/pad.py
index c84fca124d..b4ad136f8a 100644
--- a/python/tvm/topi/nn/pad.py
+++ b/python/tvm/topi/nn/pad.py
@@ -16,7 +16,8 @@
# under the License.
"""Pad the data by constant value"""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import te
from .. import tag
@@ -92,7 +93,9 @@ def pad(data, pad_before, pad_after=None, pad_value=0.0, name="PadInput", attrs=
dshape.append(dim)
out_shape = tuple(ana.simplify(dshape[i] + pad_before[i] + pad_after[i]) for i in range(n))
pad_value = (
- pad_value if tvm_tirx.ir.is_prim_expr(pad_value) else tvm_tirx.tirx.const(pad_value, data.dtype)
+ pad_value
+ if tvm_tirx.ir.is_prim_expr(pad_value)
+ else tvm_tirx.tirx.const(pad_value, data.dtype)
)
def _pad(*indices):
@@ -259,7 +262,9 @@ def replicate_pad(data, pad_before, pad_after=None, name="ReplicatePadInput"):
# Branchless edge clamp. This is bit-identical to the nested
# if_then_else form (0 below 0, size-1 at or above size, orig_idx
# otherwise) but lowers to min/max instead of per-element branches.
- clamped_idx = tvm_tirx.tirx.max(tvm_tirx.tirx.const(0, "int32"), tvm_tirx.tirx.min(size - 1, orig_idx))
+ clamped_idx = tvm_tirx.tirx.max(
+ tvm_tirx.tirx.const(0, "int32"), tvm_tirx.tirx.min(size - 1, orig_idx)
+ )
index_tuple.append(clamped_idx)
return data(*index_tuple)
diff --git a/python/tvm/topi/nn/pixel_shuffle.py b/python/tvm/topi/nn/pixel_shuffle.py
index b650a49c89..21dfd33e5f 100644
--- a/python/tvm/topi/nn/pixel_shuffle.py
+++ b/python/tvm/topi/nn/pixel_shuffle.py
@@ -17,7 +17,8 @@
# ruff: noqa: RUF005
"""TVM operator pixel shuffle compute."""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
def pixel_shuffle(data, upscale_factor, name="PixelShuffle"):
diff --git a/python/tvm/topi/nn/qnn.py b/python/tvm/topi/nn/qnn.py
index 7e261f5954..18c0a7c05e 100644
--- a/python/tvm/topi/nn/qnn.py
+++ b/python/tvm/topi/nn/qnn.py
@@ -16,9 +16,10 @@
# under the License.
"""Quantized Neural Network (QNN) Operators"""
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx
from tvm import te, topi
+from tvm_tirx import tirx
SQNN_DISABLE = 0
SQNN_INT8 = 1
diff --git a/python/tvm/topi/nn/softmax.py b/python/tvm/topi/nn/softmax.py
index 8eac0de89e..acdd3a99dc 100644
--- a/python/tvm/topi/nn/softmax.py
+++ b/python/tvm/topi/nn/softmax.py
@@ -18,7 +18,7 @@
# ruff: noqa: RUF005
"""TVM operator for softmax and log_softmax compute."""
-import tvm, tvm_tirx
+import tvm
from tvm import te, topi
diff --git a/python/tvm/topi/nn/space_to_depth.py b/python/tvm/topi/nn/space_to_depth.py
index 0f98f8148a..a168fbe90a 100644
--- a/python/tvm/topi/nn/space_to_depth.py
+++ b/python/tvm/topi/nn/space_to_depth.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name
"""TVM operator space_to_depth compute."""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
from .. import tag
diff --git a/python/tvm/topi/nn/utils.py b/python/tvm/topi/nn/utils.py
index 6710faa1e3..69e9e9055c 100644
--- a/python/tvm/topi/nn/utils.py
+++ b/python/tvm/topi/nn/utils.py
@@ -17,7 +17,8 @@
# pylint: disable=invalid-name, unused-variable
"""NN operator common utilities"""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from ..utils import get_const_int
diff --git a/python/tvm/topi/scan.py b/python/tvm/topi/scan.py
index 88d9d7f7f6..3f737c4a20 100644
--- a/python/tvm/topi/scan.py
+++ b/python/tvm/topi/scan.py
@@ -20,12 +20,12 @@
import operator
from collections.abc import Callable
-import tvm, tvm_tirx
+import tvm
from tvm_tirx.script.ir_builder import IRBuilder
+from tvm_tirx.tirx import decl_tensor
from tvm_tirx.tirx.script import ir_builder as T
from ..te import extern
-from tvm_tirx.tirx import decl_tensor
from . import utils
from .math import cast
diff --git a/python/tvm/topi/scatter.py b/python/tvm/topi/scatter.py
index d534243b32..91b9bda8b2 100644
--- a/python/tvm/topi/scatter.py
+++ b/python/tvm/topi/scatter.py
@@ -18,8 +18,8 @@
# ruff: noqa: E741
"""ScatterND operator"""
-from tvm_tirx import DataTypeCode, ir, tirx
from tvm import te # hide redefinition of min and max
+from tvm_tirx import DataTypeCode, ir, tirx
from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.sym.analyzer import Analyzer
from tvm_tirx.tirx.script import ir_builder as T
diff --git a/python/tvm/topi/scatter_elements.py b/python/tvm/topi/scatter_elements.py
index 395e50accc..4dcfbac76d 100644
--- a/python/tvm/topi/scatter_elements.py
+++ b/python/tvm/topi/scatter_elements.py
@@ -16,8 +16,8 @@
# under the License.
"""ScatterElements operator"""
-from tvm_tirx import tirx
from tvm import te
+from tvm_tirx import tirx
from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.tirx.script import ir_builder as T
diff --git a/python/tvm/topi/signal.py b/python/tvm/topi/signal.py
index d783a11dbf..36a75da38b 100644
--- a/python/tvm/topi/signal.py
+++ b/python/tvm/topi/signal.py
@@ -19,8 +19,8 @@
from math import pi
-from tvm_tirx import ir, tirx
from tvm import te
+from tvm_tirx import ir, tirx
from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.tirx.script import ir_builder as T
diff --git a/python/tvm/topi/sort.py b/python/tvm/topi/sort.py
index 0b2886de95..af99d445cc 100644
--- a/python/tvm/topi/sort.py
+++ b/python/tvm/topi/sort.py
@@ -17,7 +17,8 @@
# pylint: disable=too-many-arguments
"""Argsort operator"""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import te
from .utils import get_const_tuple
@@ -51,7 +52,9 @@ def sort(data, axis=-1, is_ascend=1):
data_buf = tvm_tirx.tirx.decl_tensor(
data.shape, data.dtype, "data_buf", data_alignment=8, layout=None
)
- out_buf = tvm_tirx.tirx.decl_tensor(data.shape, data.dtype, "out_buf", data_alignment=8, layout=None)
+ out_buf = tvm_tirx.tirx.decl_tensor(
+ data.shape, data.dtype, "out_buf", data_alignment=8, layout=None
+ )
out = te.extern(
data.shape,
[data],
@@ -136,7 +139,9 @@ def argsort(data, valid_count=None, axis=-1, is_ascend=1, dtype="float32"):
tag="argsort_nms_cpu",
)
else:
- out_buf = tvm_tirx.tirx.decl_tensor(data.shape, dtype, "out_buf", data_alignment=8, layout=None)
+ out_buf = tvm_tirx.tirx.decl_tensor(
+ data.shape, dtype, "out_buf", data_alignment=8, layout=None
+ )
out = te.extern(
data.shape,
[data],
@@ -196,11 +201,15 @@ def topk(data, k=1, axis=-1, ret_type="both", is_ascend=False, dtype="int64"):
out_bufs = []
if ret_type in ["both", "values"]:
out_bufs.append(
- tvm_tirx.tirx.decl_tensor(out_shape, data.dtype, "value_buf", data_alignment=8, layout=None)
+ tvm_tirx.tirx.decl_tensor(
+ out_shape, data.dtype, "value_buf", data_alignment=8, layout=None
+ )
)
if ret_type in ["both", "indices"]:
out_bufs.append(
- tvm_tirx.tirx.decl_tensor(out_shape, dtype, "indices_buf", data_alignment=8, layout=None)
+ tvm_tirx.tirx.decl_tensor(
+ out_shape, dtype, "indices_buf", data_alignment=8, layout=None
+ )
)
out_shapes = [out_shape] * len(out_bufs)
diff --git a/python/tvm/topi/sparse_reshape.py b/python/tvm/topi/sparse_reshape.py
index 178f08e047..eece8aa2df 100644
--- a/python/tvm/topi/sparse_reshape.py
+++ b/python/tvm/topi/sparse_reshape.py
@@ -17,8 +17,8 @@
# pylint: disable=invalid-name, too-many-arguments, too-many-nested-blocks
"""Sparse_Reshape operator"""
-from tvm_tirx.script.ir_builder import IRBuilder
from tvm.te import div, extern, floordiv, floormod
+from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.tirx import Cast, decl_tensor
from tvm_tirx.tirx.script import ir_builder as T
diff --git a/python/tvm/topi/testing/poolnd_python.py b/python/tvm/topi/testing/poolnd_python.py
index 0aa57d1de0..5ab39f97af 100644
--- a/python/tvm/topi/testing/poolnd_python.py
+++ b/python/tvm/topi/testing/poolnd_python.py
@@ -23,7 +23,7 @@ import math
import numpy as np
-import tvm, tvm_tirx
+import tvm
def _get_supported_layout(dims: int):
diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py
index 3aebc33ec0..be83f667c0 100644
--- a/python/tvm/topi/transform.py
+++ b/python/tvm/topi/transform.py
@@ -21,7 +21,8 @@ from math import pi
import numpy as np
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import te, topi
from . import cpp, tag
@@ -320,7 +321,9 @@ def strided_set(a, v, begin, end, strides=None):
strides = [tvm_tirx.tirx.const(1, "int32")] * n
else:
strides = [
- tvm_tirx.tirx.if_then_else(strides.shape[0] > i, strides[i], tvm_tirx.tirx.const(1, "int32"))
+ tvm_tirx.tirx.if_then_else(
+ strides.shape[0] > i, strides[i], tvm_tirx.tirx.const(1, "int32")
+ )
for i in range(n)
]
@@ -352,7 +355,9 @@ def strided_set(a, v, begin, end, strides=None):
for i in range(n):
from_val.append(within_index(begin[i], end[i], strides[i], indices[i]))
index_tuple.append(make_idx(begin[i], end[i], strides[i], a.shape[i], indices[i]))
- return tvm_tirx.tirx.if_then_else(tvm_tirx.tirx.all(*from_val), v(*index_tuple), a(*indices))
+ return tvm_tirx.tirx.if_then_else(
+ tvm_tirx.tirx.all(*from_val), v(*index_tuple), a(*indices)
+ )
return te.compute(a.shape, _select, name="strided_set")
diff --git a/python/tvm/topi/unique.py b/python/tvm/topi/unique.py
index c15848247b..a150fa2846 100644
--- a/python/tvm/topi/unique.py
+++ b/python/tvm/topi/unique.py
@@ -17,8 +17,8 @@
# pylint: disable=invalid-name
"""Unique operator"""
-from tvm_tirx import tirx
from tvm import te
+from tvm_tirx import tirx
from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.tirx.script import ir_builder as T
diff --git a/python/tvm/topi/utils.py b/python/tvm/topi/utils.py
index ff3898102a..0f6eaa24a8 100644
--- a/python/tvm/topi/utils.py
+++ b/python/tvm/topi/utils.py
@@ -22,7 +22,8 @@ from numbers import Integral
import numpy as np
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import te
from tvm.s_tir import sbijective_layout, slayout
from tvm_tirx.sym.utils import get_const_int
@@ -411,7 +412,9 @@ def get_shape(src_shape, src_layout, dst_layout):
assert len(src_layout) == len(dst_layout), f"Incompatible layout {src_layout} vs {dst_layout}"
layout_mapping = sbijective_layout(src_layout, dst_layout)
- dst_indices = layout_mapping.forward_index(tvm_tirx.runtime.convert(list(range(len(src_layout)))))
+ dst_indices = layout_mapping.forward_index(
+ tvm_tirx.runtime.convert(list(range(len(src_layout))))
+ )
return get_const_tuple(tuple([src_shape[i.value] for i in dst_indices]))
diff --git a/python/tvm/topi/vision/multibox_transform_loc.py b/python/tvm/topi/vision/multibox_transform_loc.py
index 87d907c2d0..96f9e140bd 100644
--- a/python/tvm/topi/vision/multibox_transform_loc.py
+++ b/python/tvm/topi/vision/multibox_transform_loc.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name
"""Multibox location transform (SSD / TFLite DetectionPostProcess decode)."""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te, topi
diff --git a/python/tvm/topi/vision/nms.py b/python/tvm/topi/vision/nms.py
index 17fa589fb2..969832c949 100644
--- a/python/tvm/topi/vision/nms.py
+++ b/python/tvm/topi/vision/nms.py
@@ -17,7 +17,8 @@
# pylint: disable=import-error, invalid-name, no-member, too-many-locals, too-many-arguments, undefined-variable, too-many-nested-blocks, too-many-branches, too-many-statements, too-many-function-args
"""Non-maximum suppression operator"""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import te
from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.tirx import if_then_else
@@ -53,7 +54,9 @@ def _get_valid_counts_ir(
if id_index < 0:
is_valid = score > score_threshold
else:
- is_valid = tvm_tirx.tirx.all(score > score_threshold, data[(i, j, id_index)] >= 0)
+ is_valid = tvm_tirx.tirx.all(
+ score > score_threshold, data[(i, j, id_index)] >= 0
+ )
with T.if_(is_valid):
with T.then_():
@@ -952,7 +955,9 @@ def _nms_loop(
on_new_invalidated_box_func(i, k)
with T.serial(0, batch_size) as i:
- nkeep = if_then_else(tvm_tirx.tirx.all(top_k > 0, top_k < valid_count[i]), top_k, valid_count[i])
+ nkeep = if_then_else(
+ tvm_tirx.tirx.all(top_k > 0, top_k < valid_count[i]), top_k, valid_count[i]
+ )
with T.if_(tvm_tirx.tirx.all(iou_threshold > te.const(0), valid_count[i] > te.const(0))):
with T.then_():
diff --git a/python/tvm/topi/vision/nms_util.py b/python/tvm/topi/vision/nms_util.py
index 108ded1a6c..6b4fd8ad48 100644
--- a/python/tvm/topi/vision/nms_util.py
+++ b/python/tvm/topi/vision/nms_util.py
@@ -18,7 +18,8 @@
# ruff: noqa: E741
"""Common utilities used in Non-maximum suppression operators"""
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import te
from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.tirx.script import ir_builder as T
diff --git a/python/tvm/topi/vision/roi_align.py b/python/tvm/topi/vision/roi_align.py
index 7314f23670..f6f3792822 100644
--- a/python/tvm/topi/vision/roi_align.py
+++ b/python/tvm/topi/vision/roi_align.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name
"""ROI Align operator"""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
from ..cpp.utils import bilinear_sample_nchw, bilinear_sample_nhwc
diff --git a/python/tvm/topi/vision/roi_pool.py b/python/tvm/topi/vision/roi_pool.py
index 09e28ab01b..457120f0ca 100644
--- a/python/tvm/topi/vision/roi_pool.py
+++ b/python/tvm/topi/vision/roi_pool.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name
"""ROI Pool operator"""
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import te
@@ -53,8 +53,12 @@ def roi_pool_nchw(data, rois, pooled_size, spatial_scale):
roi_h = te.max(roi_end_h - roi_start_h + 1, tvm_tirx.tirx.const(1, "int32"))
roi_w = te.max(roi_end_w - roi_start_w + 1, tvm_tirx.tirx.const(1, "int32"))
- bin_h = tvm_tirx.tirx.Cast(roi_h, roi_dtype) / tvm_tirx.tirx.const(float(pooled_size_h), roi_dtype)
- bin_w = tvm_tirx.tirx.Cast(roi_w, roi_dtype) / tvm_tirx.tirx.const(float(pooled_size_w), roi_dtype)
+ bin_h = tvm_tirx.tirx.Cast(roi_h, roi_dtype) / tvm_tirx.tirx.const(
+ float(pooled_size_h), roi_dtype
+ )
+ bin_w = tvm_tirx.tirx.Cast(roi_w, roi_dtype) / tvm_tirx.tirx.const(
+ float(pooled_size_w), roi_dtype
+ )
hstart = te.floor(tvm_tirx.tirx.Cast(ph, roi_dtype) * bin_h).astype("int32")
wstart = te.floor(tvm_tirx.tirx.Cast(pw, roi_dtype) * bin_w).astype("int32")
diff --git a/tests/nightly/python/test_nnapi/infrastructure.py b/tests/nightly/python/test_nnapi/infrastructure.py
index 76547f1feb..cd054010a7 100644
--- a/tests/nightly/python/test_nnapi/infrastructure.py
+++ b/tests/nightly/python/test_nnapi/infrastructure.py
@@ -18,7 +18,8 @@
import numpy as np
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
import tvm_tirx.script.relax as R
from tvm.relax.backend.contrib.nnapi import partition_for_nnapi
from tvm_tirx.support import ndk, utils
diff --git a/tests/nightly/python/test_nnapi/test_from_exported_to_cuda.py b/tests/nightly/python/test_nnapi/test_from_exported_to_cuda.py
index ab9ec30fa8..30942cdbf3 100644
--- a/tests/nightly/python/test_nnapi/test_from_exported_to_cuda.py
+++ b/tests/nightly/python/test_nnapi/test_from_exported_to_cuda.py
@@ -24,8 +24,9 @@ from torch.export import export
from torch.nn import Softmax, Upsample
from torch.nn import functional as F
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm.relax.frontend.torch import from_exported_program
@@ -60,7 +61,9 @@ def assert_torch_output_vs_tvm_from_exported_to_cuda(raw_data, torch_module, tar
for i in range(len(pytorch_out)):
actual = gpu_out[i].numpy()
desired = pytorch_out[i].detach().numpy()
- tvm_tirx.testing.assert_allclose(actual=actual, desired=desired, rtol=1e-5, atol=1e-5)
+ tvm_tirx.testing.assert_allclose(
+ actual=actual, desired=desired, rtol=1e-5, atol=1e-5
+ )
else:
actual = gpu_out[0].numpy()
desired = pytorch_out.detach().numpy()
diff --git a/tests/nightly/python/test_nnapi/test_network.py b/tests/nightly/python/test_nnapi/test_network.py
index d74b2c4029..6493bbdfda 100644
--- a/tests/nightly/python/test_nnapi/test_network.py
+++ b/tests/nightly/python/test_nnapi/test_network.py
@@ -23,7 +23,8 @@ pytest.importorskip("onnx")
import onnx
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from test_nnapi.conftest import remote
from test_nnapi.infrastructure import build_and_run # , build_and_run_vm
from tvm.contrib.download import download_testdata
@@ -127,7 +128,9 @@ def test_network(name, dtype):
for _name, (shape, _dtype) in inputs.items():
input_data[_name] = np.random.uniform(-1.0, 1.0, shape).astype(_dtype)
- inputs_tvm: list[tvm_tirx.runtime.Tensor] = [tvm_tirx.runtime.tensor(v) for k, v in input_data.items()]
+ inputs_tvm: list[tvm_tirx.runtime.Tensor] = [
+ tvm_tirx.runtime.tensor(v) for k, v in input_data.items()
+ ]
outputs = _build_and_run_network(remote_obj, tracker, mod, inputs_tvm)
nnapi_out = outputs[0]
expected_out = outputs[1]
diff --git a/tests/nightly/python/test_nnapi/test_ops.py b/tests/nightly/python/test_nnapi/test_ops.py
index cd528873dc..adbe44a4b1 100644
--- a/tests/nightly/python/test_nnapi/test_ops.py
+++ b/tests/nightly/python/test_nnapi/test_ops.py
@@ -19,8 +19,9 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.script
import tvm_tirx.script.relax as R
from test_nnapi.conftest import remote
from test_nnapi.infrastructure import build_and_run
diff --git a/tests/python/all-platform-minimal-test/test_minimal_target_codegen_llvm.py b/tests/python/all-platform-minimal-test/test_minimal_target_codegen_llvm.py
index 82d2758972..dd297f414f 100644
--- a/tests/python/all-platform-minimal-test/test_minimal_target_codegen_llvm.py
+++ b/tests/python/all-platform-minimal-test/test_minimal_target_codegen_llvm.py
@@ -24,8 +24,9 @@ import re
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import te, topi
from tvm_tirx.support import utils
from tvm_tirx.testing import env
diff --git a/tests/python/codegen/llvm/test_cross_compile.py b/tests/python/codegen/llvm/test_cross_compile.py
index 106ecfab95..c1135013e8 100644
--- a/tests/python/codegen/llvm/test_cross_compile.py
+++ b/tests/python/codegen/llvm/test_cross_compile.py
@@ -23,8 +23,9 @@ import struct
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import rpc
from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/contrib/pickle_memoize_script.py b/tests/python/contrib/pickle_memoize_script.py
index 0f7309704e..f0d73e3910 100755
--- a/tests/python/contrib/pickle_memoize_script.py
+++ b/tests/python/contrib/pickle_memoize_script.py
@@ -19,7 +19,7 @@
import sys
-import tvm, tvm_tirx
+import tvm
@tvm.contrib.pickle_memoize.memoize("test_memoize_save_data", save_at_exit=True)
diff --git a/tests/python/contrib/test_android/infrastructure.py b/tests/python/contrib/test_android/infrastructure.py
index 44fbc8b4c8..174ffe4523 100644
--- a/tests/python/contrib/test_android/infrastructure.py
+++ b/tests/python/contrib/test_android/infrastructure.py
@@ -20,7 +20,8 @@
import os
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm.s_tir.meta_schedule.runner import EvaluatorConfig, RPCConfig, RPCRunner
diff --git a/tests/python/contrib/test_android/test_meta_schedule.py b/tests/python/contrib/test_android/test_meta_schedule.py
index 1eab8b2b13..cd2e6e396a 100644
--- a/tests/python/contrib/test_android/test_meta_schedule.py
+++ b/tests/python/contrib/test_android/test_meta_schedule.py
@@ -25,11 +25,12 @@ import pytest
pytest.importorskip("scipy") # tvm.topi.testing imports scipy
-import tvm_tirx.testing, tvm
+import tvm
import tvm.topi.testing
+import tvm_tirx.testing
from tvm.s_tir import meta_schedule as ms
-from tvm.s_tir.meta_schedule.builder import LocalBuilder
from tvm.s_tir import script as Ts
+from tvm.s_tir.meta_schedule.builder import LocalBuilder
from tvm_tirx.script import tirx as T
from .infrastructure import get_android_gpu_target, get_rpc_runner
diff --git a/tests/python/contrib/test_cblas.py b/tests/python/contrib/test_cblas.py
index 7bd1f4030f..867f20db28 100644
--- a/tests/python/contrib/test_cblas.py
+++ b/tests/python/contrib/test_cblas.py
@@ -21,9 +21,10 @@ import pytest
pytest.importorskip("scipy") # tvm.topi.testing imports scipy
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
import tvm.topi.testing
+import tvm_tirx
+import tvm_tirx.testing
from tvm import te
from tvm.contrib import cblas, dnnl, mkl
diff --git a/tests/python/contrib/test_coreml_runtime.py b/tests/python/contrib/test_coreml_runtime.py
index 0880d2c0d3..ae07c13334 100644
--- a/tests/python/contrib/test_coreml_runtime.py
+++ b/tests/python/contrib/test_coreml_runtime.py
@@ -20,7 +20,8 @@ import os
import numpy as np
import pytest
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import rpc, te
from tvm.contrib import coreml_runtime
from tvm_tirx.support import utils, xcode
@@ -76,7 +77,8 @@ def test_coreml_runtime():
runtime = coreml_runtime.create("main", model_path, dev, session=session)
for name in inputs:
runtime.set_input(
- name, dev.tensor(inputs[name]) if session else tvm_tirx.runtime.tensor(inputs[name], dev)
+ name,
+ dev.tensor(inputs[name]) if session else tvm_tirx.runtime.tensor(inputs[name], dev),
)
runtime.invoke()
tvm_outputs = [runtime.get_output(i).numpy() for i in range(runtime.get_num_outputs())]
diff --git a/tests/python/contrib/test_cutlass_gemm.py b/tests/python/contrib/test_cutlass_gemm.py
index 2e4dde1339..db780ab7fc 100644
--- a/tests/python/contrib/test_cutlass_gemm.py
+++ b/tests/python/contrib/test_cutlass_gemm.py
@@ -19,8 +19,8 @@ import ml_dtypes
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.contrib.pickle_memoize import memoize
from tvm_tirx.testing import env
diff --git a/tests/python/contrib/test_dlpack.py b/tests/python/contrib/test_dlpack.py
index b9625f0377..2486be0331 100644
--- a/tests/python/contrib/test_dlpack.py
+++ b/tests/python/contrib/test_dlpack.py
@@ -16,8 +16,9 @@
# under the License.
import numpy as np
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import te
from tvm.contrib.dlpack import to_pytorch_func
diff --git a/tests/python/contrib/test_example_npu.py b/tests/python/contrib/test_example_npu.py
index 9617e7a745..fd0efe7bd7 100644
--- a/tests/python/contrib/test_example_npu.py
+++ b/tests/python/contrib/test_example_npu.py
@@ -24,12 +24,12 @@ implementation using TVM's testing infrastructure.
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.backend.pattern_registry import get_patterns_with_prefix
from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions, RunCodegen
-from tvm.relax import script as R
@tvm_tirx.script.ir_module
@@ -224,7 +224,6 @@ def test_example_npu_codegen():
@example_npu_enabled
def test_example_npu_runtime_execution():
"""Test end-to-end execution with the example NPU runtime"""
- import tvm.relax.backend.contrib.example_npu
# Create simple test inputs
np.random.seed(42)
diff --git a/tests/python/contrib/test_hipblas.py b/tests/python/contrib/test_hipblas.py
index c2338cd9ce..9e424bbdd3 100644
--- a/tests/python/contrib/test_hipblas.py
+++ b/tests/python/contrib/test_hipblas.py
@@ -18,8 +18,9 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import te
from tvm.contrib import hipblas
from tvm_tirx.testing import env
diff --git a/tests/python/contrib/test_memoize.py b/tests/python/contrib/test_memoize.py
index bdd3f05720..3740b11c28 100644
--- a/tests/python/contrib/test_memoize.py
+++ b/tests/python/contrib/test_memoize.py
@@ -23,7 +23,7 @@ import subprocess
import sys
import tempfile
-import tvm_tirx.testing, tvm
+import tvm_tirx.testing
TEST_SCRIPT_FILE = pathlib.Path(__file__).with_name("pickle_memoize_script.py").resolve()
diff --git a/tests/python/contrib/test_random.py b/tests/python/contrib/test_random.py
index 005984c466..8eb7a9e594 100644
--- a/tests/python/contrib/test_random.py
+++ b/tests/python/contrib/test_random.py
@@ -23,8 +23,9 @@ import threading
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import rpc, te
from tvm.contrib import random
diff --git a/tests/python/contrib/test_rpc_proxy.py b/tests/python/contrib/test_rpc_proxy.py
index 17907f5856..479f90e245 100644
--- a/tests/python/contrib/test_rpc_proxy.py
+++ b/tests/python/contrib/test_rpc_proxy.py
@@ -21,7 +21,7 @@ import logging
import multiprocessing
import time
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import rpc
diff --git a/tests/python/contrib/test_rpc_tracker.py b/tests/python/contrib/test_rpc_tracker.py
index 809dd9da0f..95a97fe776 100644
--- a/tests/python/contrib/test_rpc_tracker.py
+++ b/tests/python/contrib/test_rpc_tracker.py
@@ -20,7 +20,7 @@
import logging
import time
-import tvm, tvm_tirx
+import tvm_tirx
from tvm import rpc
diff --git a/tests/python/contrib/test_sort.py b/tests/python/contrib/test_sort.py
index 424231de42..8f5364aef4 100644
--- a/tests/python/contrib/test_sort.py
+++ b/tests/python/contrib/test_sort.py
@@ -20,8 +20,9 @@
# pylint: disable=invalid-name
import numpy as np
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import te
diff --git a/tests/python/contrib/test_tir_triton_integration.py b/tests/python/contrib/test_tir_triton_integration.py
index ce5fb0122e..eadf3c7a36 100644
--- a/tests/python/contrib/test_tir_triton_integration.py
+++ b/tests/python/contrib/test_tir_triton_integration.py
@@ -20,13 +20,14 @@ import sys
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm.relax.frontend import nn
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm.relax.frontend import nn
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/contrib/test_tvmjs.py b/tests/python/contrib/test_tvmjs.py
index 2ef538f682..cbe966dec3 100644
--- a/tests/python/contrib/test_tvmjs.py
+++ b/tests/python/contrib/test_tvmjs.py
@@ -24,7 +24,7 @@ import tempfile
import numpy as np
import pytest
-import tvm_tirx.testing, tvm
+import tvm_tirx.testing
from tvm.contrib import tvmjs
dtype = tvm_tirx.testing.parameter(
diff --git a/tests/python/disco/test_callback.py b/tests/python/disco/test_callback.py
index 403cdc9c89..967dc3f071 100644
--- a/tests/python/disco/test_callback.py
+++ b/tests/python/disco/test_callback.py
@@ -23,11 +23,12 @@ import tempfile
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx.script import ir as I
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/disco/test_ccl.py b/tests/python/disco/test_ccl.py
index 64d091de3d..b772b45ee5 100644
--- a/tests/python/disco/test_ccl.py
+++ b/tests/python/disco/test_ccl.py
@@ -22,14 +22,15 @@ import tempfile
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import get_global_func
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax as rx
+from tvm.relax import script as R
from tvm.runtime import disco as di
from tvm.runtime.vm import VirtualMachine
from tvm.s_tir import dlight as dl
-from tvm.relax import script as R
+from tvm_tirx import get_global_func
if di is None:
pytest.skip("disco runtime is not available", allow_module_level=True)
diff --git a/tests/python/disco/test_custom_allreduce.py b/tests/python/disco/test_custom_allreduce.py
index fe8edd632f..e0c8f275e0 100644
--- a/tests/python/disco/test_custom_allreduce.py
+++ b/tests/python/disco/test_custom_allreduce.py
@@ -23,8 +23,8 @@ import numpy as np
import pytest
from tvm_ffi import Shape
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.runtime import disco
from tvm_tirx.runtime import DataType
diff --git a/tests/python/disco/test_loader.py b/tests/python/disco/test_loader.py
index 017239a1e4..9ef42b65ff 100644
--- a/tests/python/disco/test_loader.py
+++ b/tests/python/disco/test_loader.py
@@ -25,14 +25,15 @@ import numpy as np
import pytest
from tvm_ffi import Shape, register_global_func
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax as rx
from tvm.contrib import tvmjs
+from tvm.relax import script as R
from tvm.runtime import disco as di
from tvm.s_tir import dlight as dl
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm_tirx.target import Target
from tvm_tirx.testing import env
diff --git a/tests/python/disco/test_nvshmem.py b/tests/python/disco/test_nvshmem.py
index 2b8b4a8909..f9cb5be833 100644
--- a/tests/python/disco/test_nvshmem.py
+++ b/tests/python/disco/test_nvshmem.py
@@ -29,13 +29,14 @@ import numpy as np
import pytest
from tvm_ffi import Shape
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx.ir import TupleType as _TupleType
-from tvm.runtime import disco as di
-from tvm_tirx.script import ir as I
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax import script as R
+from tvm.runtime import disco as di
from tvm.s_tir import script as Ts
+from tvm_tirx.ir import TupleType as _TupleType
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
@@ -340,7 +341,10 @@ def _kernel_compile(compile_mode):
target = tvm_tirx.target.Target("cuda")
tvm.compile(
NvshmemQueryModule,
- target={**tvm_tirx.target.Target(target).export(), "backend_config": backend_config},
+ target={
+ **tvm_tirx.target.Target(target).export(),
+ "backend_config": backend_config,
+ },
).export_library(path)
mod = sess.load_vm_module(path)
result = mod["main"]()
diff --git a/tests/python/disco/test_session.py b/tests/python/disco/test_session.py
index 375c2d3d6b..ae184afe07 100644
--- a/tests/python/disco/test_session.py
+++ b/tests/python/disco/test_session.py
@@ -28,15 +28,16 @@ import pytest
from tvm_ffi import Shape
from tvm_ffi.core import String
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
# Imported for the side effect of registering the tests.disco.* worker functions.
from tvm.exec import disco_worker as _ # noqa: F401 # pylint: disable=unused-import
-from tvm.runtime import disco as di
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm.runtime import disco as di
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
if di is None:
diff --git a/tests/python/driver/test_compile.py b/tests/python/driver/test_compile.py
index b71ba67b0a..cb0bab40f5 100644
--- a/tests/python/driver/test_compile.py
+++ b/tests/python/driver/test_compile.py
@@ -20,13 +20,14 @@ import platform
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax, te
-from tvm_tirx.runtime import Executable
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.runtime import Executable
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/backend/adreno/mod_utils.py b/tests/python/relax/backend/adreno/mod_utils.py
index fd3e2fe4a8..a0448305eb 100644
--- a/tests/python/relax/backend/adreno/mod_utils.py
+++ b/tests/python/relax/backend/adreno/mod_utils.py
@@ -22,14 +22,15 @@ import json
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax, rpc
+from tvm.relax import script as R
from tvm.relax.backend.adreno import clml
from tvm.relax.script import ir_builder as relax_builder
-from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/backend/adreno/test_clml_ops.py b/tests/python/relax/backend/adreno/test_clml_ops.py
index 08ade82ee7..8fa7dcb061 100644
--- a/tests/python/relax/backend/adreno/test_clml_ops.py
+++ b/tests/python/relax/backend/adreno/test_clml_ops.py
@@ -46,14 +46,15 @@ from mod_utils import (
)
from utils import skip_unless_adreno_clml, verify_results
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax, rpc
+from tvm.relax import script as R
from tvm.relax.backend.adreno import clml
from tvm.relax.backend.adreno.clml import OpenCLMLOffLoad, OpenCLMLOffLoadForLLM
from tvm.relax.script import ir_builder as relax_builder
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/backend/adreno/test_texture_network.py b/tests/python/relax/backend/adreno/test_texture_network.py
index 8039372deb..23d487be6a 100644
--- a/tests/python/relax/backend/adreno/test_texture_network.py
+++ b/tests/python/relax/backend/adreno/test_texture_network.py
@@ -27,14 +27,15 @@ pytest.importorskip("onnx")
import onnx
from utils import verify_results
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.frontend.onnx import from_onnx
from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.transform.legalize_ops import adreno as legalize_adreno
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/backend/adreno/test_texture_ops.py b/tests/python/relax/backend/adreno/test_texture_ops.py
index de127a10c0..8d86edc449 100644
--- a/tests/python/relax/backend/adreno/test_texture_ops.py
+++ b/tests/python/relax/backend/adreno/test_texture_ops.py
@@ -18,8 +18,9 @@
import pytest
from utils import skip_unless_adreno_opencl_vulkan, verify_results
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm # noqa: F401 # Register the parent Relax script dialect.
+import tvm_tirx
+import tvm_tirx.testing
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
diff --git a/tests/python/relax/backend/adreno/test_transform_annotate_custom_scope.py b/tests/python/relax/backend/adreno/test_transform_annotate_custom_scope.py
index febfc4bce9..a40ba3e30b 100644
--- a/tests/python/relax/backend/adreno/test_transform_annotate_custom_scope.py
+++ b/tests/python/relax/backend/adreno/test_transform_annotate_custom_scope.py
@@ -16,12 +16,13 @@
# under the License.
# ruff: noqa: F401, F841
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir.module import IRModule
from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, visitor
from tvm.relax.transform.legalize_ops import adreno as legalize_adreno
+from tvm_tirx.ir.module import IRModule
from tvm_tirx.runtime import MemoryScope
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
@@ -89,7 +90,9 @@ def verify(mod, expected):
# "relax.nn.layer_norm",
]
with tgt:
- mod = tvm_tirx.tirx.transform.BindTarget(tvm_tirx.target.Target.current(allow_none=False))(mod)
+ mod = tvm_tirx.tirx.transform.BindTarget(tvm_tirx.target.Target.current(allow_none=False))(
+ mod
+ )
mod = tvm.relax.transform.DecomposeOpsForInference()(mod)
mod = tvm.relax.transform.FoldConstant()(mod)
desired_layouts = {"relax.nn.conv2d": ["NCHW4c", "OIHW4o", "NCHW4c"]}
diff --git a/tests/python/relax/backend/adreno/test_transform_fold_vdevice_scope_change.py b/tests/python/relax/backend/adreno/test_transform_fold_vdevice_scope_change.py
index 1fa999d25d..94c38f74b2 100644
--- a/tests/python/relax/backend/adreno/test_transform_fold_vdevice_scope_change.py
+++ b/tests/python/relax/backend/adreno/test_transform_fold_vdevice_scope_change.py
@@ -18,11 +18,12 @@
from __future__ import annotations
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir.module import IRModule
from tvm.s_tir import script as Ts
+from tvm_tirx.ir.module import IRModule
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/backend/adreno/utils.py b/tests/python/relax/backend/adreno/utils.py
index 9eb74444ef..ae0091a0ae 100644
--- a/tests/python/relax/backend/adreno/utils.py
+++ b/tests/python/relax/backend/adreno/utils.py
@@ -21,8 +21,9 @@ import tempfile
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm_tirx.support import ndk
diff --git a/tests/python/relax/conftest.py b/tests/python/relax/conftest.py
index e12d4ed001..cbda4a86de 100644
--- a/tests/python/relax/conftest.py
+++ b/tests/python/relax/conftest.py
@@ -16,7 +16,7 @@
import pytest
-import tvm, tvm_tirx
+import tvm_tirx
from tvm.relax.ir.instrument import WellFormedInstrument
diff --git a/tests/python/relax/distributed/test_distributed_dtensor_type.py b/tests/python/relax/distributed/test_distributed_dtensor_type.py
index d6683b8950..29fa26cf00 100644
--- a/tests/python/relax/distributed/test_distributed_dtensor_type.py
+++ b/tests/python/relax/distributed/test_distributed_dtensor_type.py
@@ -18,8 +18,8 @@
import pytest
import tvm_ffi
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax as rx
from tvm_tirx import tirx
from tvm_tirx.ir import Range
diff --git a/tests/python/relax/distributed/test_distributed_transform_legalize_redistribute.py b/tests/python/relax/distributed/test_distributed_transform_legalize_redistribute.py
index 89d1957989..583c0ce018 100644
--- a/tests/python/relax/distributed/test_distributed_transform_legalize_redistribute.py
+++ b/tests/python/relax/distributed/test_distributed_transform_legalize_redistribute.py
@@ -18,8 +18,8 @@
# type: ignore
from __future__ import annotations
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
diff --git a/tests/python/relax/distributed/test_distributed_transform_lower_distir.py b/tests/python/relax/distributed/test_distributed_transform_lower_distir.py
index f5fe112c3f..a4f986e907 100644
--- a/tests/python/relax/distributed/test_distributed_transform_lower_distir.py
+++ b/tests/python/relax/distributed/test_distributed_transform_lower_distir.py
@@ -19,12 +19,13 @@
# type: ignore
from __future__ import annotations
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
+from tvm.s_tir import script as Ts
from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.ir import assert_structural_equal
-from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py b/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
index 891e201d7f..ad7ba265b5 100644
--- a/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
+++ b/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
@@ -19,11 +19,12 @@
# type: ignore
from __future__ import annotations
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import assert_structural_equal
from tvm.s_tir import script as Ts
+from tvm_tirx.ir import assert_structural_equal
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py b/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
index 3ec43750ed..e8aaab7978 100644
--- a/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
+++ b/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
@@ -20,12 +20,12 @@
from __future__ import annotations
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
+from tvm.s_tir import script as Ts
from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.ir import assert_structural_equal
-from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py b/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py
index f9c4be4d93..51db3bed34 100644
--- a/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py
+++ b/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py
@@ -30,8 +30,9 @@ import torch
import tvm_ffi
from tvm_ffi import Shape
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax.frontend.nn.llm.kv_cache import (
AttnKind,
RopeMode,
diff --git a/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer_kernel.py b/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer_kernel.py
index 8ffaf107f4..dd73a6415f 100644
--- a/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer_kernel.py
+++ b/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer_kernel.py
@@ -20,10 +20,10 @@ import numpy as np
import pytest
from tvm_ffi import Shape
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx.runtime import Device
+import tvm_tirx
+import tvm_tirx.testing
from tvm.runtime import disco as di
+from tvm_tirx.runtime import Device
page_size = 4
num_layers = 4
diff --git a/tests/python/relax/script/test_relax_script_basic_usage.py b/tests/python/relax/script/test_relax_script_basic_usage.py
index 8f631bd0ee..64bdeb4dbc 100644
--- a/tests/python/relax/script/test_relax_script_basic_usage.py
+++ b/tests/python/relax/script/test_relax_script_basic_usage.py
@@ -23,17 +23,18 @@ import sys
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule, tirx
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax, topi
-from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import DummyGlobalInfo, VDevice
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx import IRModule, tirx
+from tvm_tirx.ir import TupleType as _TupleType
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/script/test_relax_script_distributed.py b/tests/python/relax/script/test_relax_script_distributed.py
index 009ebe93d9..5d5e4eaf06 100644
--- a/tests/python/relax/script/test_relax_script_distributed.py
+++ b/tests/python/relax/script/test_relax_script_distributed.py
@@ -18,17 +18,18 @@
from __future__ import annotations
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Range
from tvm.relax import Call, SeqExpr, VarBinding
-from tvm.relax.distributed import DeviceMesh
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm.relax.distributed import DeviceMesh
from tvm.s_tir import script as Ts
+from tvm_tirx import IRModule
+from tvm_tirx.ir import Range
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/script/test_relax_script_dynamic_shape.py b/tests/python/relax/script/test_relax_script_dynamic_shape.py
index 4d1760f972..31359fef11 100644
--- a/tests/python/relax/script/test_relax_script_dynamic_shape.py
+++ b/tests/python/relax/script/test_relax_script_dynamic_shape.py
@@ -23,15 +23,16 @@ import sys
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule, tirx
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx import IRModule, tirx
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
M = T.dynamic("M", "int64")
@@ -76,13 +77,15 @@ def func[M: T.int64](x: R.Tensor((M, M * 2), "float32")):
tvm_tirx.ir.assert_structural_equal(
func,
tvm_tirx.script.from_source(
- script, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
+ script,
+ extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx},
),
)
tvm_tirx.ir.assert_structural_equal(
func,
tvm_tirx.script.from_source(
- repeated, extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx}
+ repeated,
+ extra_vars={"I": tvm_tirx.script.ir, "R": tvm.relax.script, "T": tvm_tirx.script.tirx},
),
)
diff --git a/tests/python/relax/script/test_relax_script_error_handling.py b/tests/python/relax/script/test_relax_script_error_handling.py
index e32592fe51..95a8c2d055 100644
--- a/tests/python/relax/script/test_relax_script_error_handling.py
+++ b/tests/python/relax/script/test_relax_script_error_handling.py
@@ -23,16 +23,17 @@ import inspect
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx import IRModule
+from tvm_tirx.ir import TupleType as _TupleType
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/script/test_relax_script_ir_builder.py b/tests/python/relax/script/test_relax_script_ir_builder.py
index 4ec4557713..6cd867e67c 100644
--- a/tests/python/relax/script/test_relax_script_ir_builder.py
+++ b/tests/python/relax/script/test_relax_script_ir_builder.py
@@ -16,12 +16,12 @@
# under the License.
"""Relax script ir builder."""
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.script import ir_builder as R
+from tvm_tirx import tirx
from tvm_tirx.script.ir_builder.base import IRBuilder
@@ -192,8 +192,8 @@ def test_regression_py_print():
def test_function_subroutine_before_main():
"""The block builder can generate subroutines, and calls into subroutines"""
- from tvm_tirx.script import ir as I
from tvm.relax import script as R
+ from tvm_tirx.script import ir as I
# create with TVMScript
@I.ir_module
@@ -234,8 +234,8 @@ def test_function_subroutine_before_main():
def test_function_subroutine_during_main():
"""Subroutines may be generated as needed, pausing the main function collection"""
- from tvm_tirx.script import ir as I
from tvm.relax import script as R
+ from tvm_tirx.script import ir as I
# create with TVMScript
@I.ir_module
diff --git a/tests/python/relax/script/test_relax_script_meta_programming.py b/tests/python/relax/script/test_relax_script_meta_programming.py
index 557514d830..7462828f1c 100644
--- a/tests/python/relax/script/test_relax_script_meta_programming.py
+++ b/tests/python/relax/script/test_relax_script_meta_programming.py
@@ -24,14 +24,15 @@ from typing import TypeVar
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule, ir
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax, topi
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx import IRModule, ir
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/script/test_relax_script_operators_arith_cmp.py b/tests/python/relax/script/test_relax_script_operators_arith_cmp.py
index 3a14bbb24e..4face373ab 100644
--- a/tests/python/relax/script/test_relax_script_operators_arith_cmp.py
+++ b/tests/python/relax/script/test_relax_script_operators_arith_cmp.py
@@ -20,12 +20,13 @@ from collections.abc import Callable
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import script as R
+from tvm_tirx import IRModule
def _check(
diff --git a/tests/python/relax/script/test_relax_script_operators_create.py b/tests/python/relax/script/test_relax_script_operators_create.py
index 043ecf19e1..ef17453384 100644
--- a/tests/python/relax/script/test_relax_script_operators_create.py
+++ b/tests/python/relax/script/test_relax_script_operators_create.py
@@ -16,12 +16,13 @@
# under the License.
"""Relax script operators create."""
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import script as R
+from tvm_tirx import IRModule
def _check(
diff --git a/tests/python/relax/script/test_relax_script_operators_datatype.py b/tests/python/relax/script/test_relax_script_operators_datatype.py
index f394203bf4..2324468636 100644
--- a/tests/python/relax/script/test_relax_script_operators_datatype.py
+++ b/tests/python/relax/script/test_relax_script_operators_datatype.py
@@ -16,12 +16,13 @@
# under the License.
"""Relax script operators datatype."""
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import script as R
+from tvm_tirx import IRModule
def _check(
diff --git a/tests/python/relax/script/test_relax_script_operators_grad.py b/tests/python/relax/script/test_relax_script_operators_grad.py
index bd43f81248..0c852f6af5 100644
--- a/tests/python/relax/script/test_relax_script_operators_grad.py
+++ b/tests/python/relax/script/test_relax_script_operators_grad.py
@@ -16,11 +16,12 @@
# under the License.
"""Relax script operators grad."""
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import script as R
+from tvm_tirx import IRModule
def _check(
diff --git a/tests/python/relax/script/test_relax_script_operators_image.py b/tests/python/relax/script/test_relax_script_operators_image.py
index 53cd473565..036a5873a8 100644
--- a/tests/python/relax/script/test_relax_script_operators_image.py
+++ b/tests/python/relax/script/test_relax_script_operators_image.py
@@ -16,12 +16,13 @@
# under the License.
"""Relax script operators image."""
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import script as R
+from tvm_tirx import IRModule
def _check(
diff --git a/tests/python/relax/script/test_relax_script_operators_index.py b/tests/python/relax/script/test_relax_script_operators_index.py
index 63bb8aa7c8..ebcf4d30da 100644
--- a/tests/python/relax/script/test_relax_script_operators_index.py
+++ b/tests/python/relax/script/test_relax_script_operators_index.py
@@ -16,12 +16,13 @@
# under the License.
"""Relax script operators index."""
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import script as R
+from tvm_tirx import IRModule
def _check(
diff --git a/tests/python/relax/script/test_relax_script_operators_linear_algebra.py b/tests/python/relax/script/test_relax_script_operators_linear_algebra.py
index 0d7f37000d..3928b2b585 100644
--- a/tests/python/relax/script/test_relax_script_operators_linear_algebra.py
+++ b/tests/python/relax/script/test_relax_script_operators_linear_algebra.py
@@ -16,12 +16,13 @@
# under the License.
"""Relax script operators linear algebra."""
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import script as R
+from tvm_tirx import IRModule
def _check(
diff --git a/tests/python/relax/script/test_relax_script_operators_manipulate.py b/tests/python/relax/script/test_relax_script_operators_manipulate.py
index 5119f8f6cc..f50b5ca035 100644
--- a/tests/python/relax/script/test_relax_script_operators_manipulate.py
+++ b/tests/python/relax/script/test_relax_script_operators_manipulate.py
@@ -16,13 +16,14 @@
# under the License.
"""Relax script operators manipulate."""
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm_tirx import IRModule
+from tvm_tirx.script import ir as I
def _check(
diff --git a/tests/python/relax/script/test_relax_script_operators_nn.py b/tests/python/relax/script/test_relax_script_operators_nn.py
index 4268d3b7f3..dff9a36ef8 100644
--- a/tests/python/relax/script/test_relax_script_operators_nn.py
+++ b/tests/python/relax/script/test_relax_script_operators_nn.py
@@ -16,12 +16,13 @@
# under the License.
"""Relax script operators nn."""
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import script as R
+from tvm_tirx import IRModule
def _check(
diff --git a/tests/python/relax/script/test_relax_script_operators_search.py b/tests/python/relax/script/test_relax_script_operators_search.py
index 9da26aee1d..e3c65b7f3c 100644
--- a/tests/python/relax/script/test_relax_script_operators_search.py
+++ b/tests/python/relax/script/test_relax_script_operators_search.py
@@ -16,12 +16,13 @@
# under the License.
"""Relax script operators search."""
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import script as R
+from tvm_tirx import IRModule
def _check(
diff --git a/tests/python/relax/script/test_relax_script_operators_set.py b/tests/python/relax/script/test_relax_script_operators_set.py
index 7781de2e9c..552c14c6f5 100644
--- a/tests/python/relax/script/test_relax_script_operators_set.py
+++ b/tests/python/relax/script/test_relax_script_operators_set.py
@@ -16,12 +16,13 @@
# under the License.
"""Relax script operators set."""
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import script as R
+from tvm_tirx import IRModule
def _check(
diff --git a/tests/python/relax/script/test_relax_script_operators_sort.py b/tests/python/relax/script/test_relax_script_operators_sort.py
index 15cf8f6f48..32ba2e360e 100644
--- a/tests/python/relax/script/test_relax_script_operators_sort.py
+++ b/tests/python/relax/script/test_relax_script_operators_sort.py
@@ -16,12 +16,13 @@
# under the License.
"""Relax script operators sort."""
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import script as R
+from tvm_tirx import IRModule
def _check(
diff --git a/tests/python/relax/script/test_relax_script_operators_statistical.py b/tests/python/relax/script/test_relax_script_operators_statistical.py
index 569c230b06..0682598157 100644
--- a/tests/python/relax/script/test_relax_script_operators_statistical.py
+++ b/tests/python/relax/script/test_relax_script_operators_statistical.py
@@ -16,12 +16,13 @@
# under the License.
"""Relax script operators statistical."""
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import script as R
+from tvm_tirx import IRModule
def _check(
diff --git a/tests/python/relax/script/test_relax_script_operators_vision.py b/tests/python/relax/script/test_relax_script_operators_vision.py
index 666ee0a0de..2368f788f3 100644
--- a/tests/python/relax/script/test_relax_script_operators_vision.py
+++ b/tests/python/relax/script/test_relax_script_operators_vision.py
@@ -16,12 +16,13 @@
# under the License.
"""Relax script operators vision."""
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import script as R
+from tvm_tirx import IRModule
def _check(
diff --git a/tests/python/relax/script/test_relax_script_pyfunc.py b/tests/python/relax/script/test_relax_script_pyfunc.py
index 38cfd0ed8f..5ed40a8372 100644
--- a/tests/python/relax/script/test_relax_script_pyfunc.py
+++ b/tests/python/relax/script/test_relax_script_pyfunc.py
@@ -19,12 +19,12 @@
import pytest
import torch
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax import BasePyModule
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
n = T.dynamic("n", "int32")
diff --git a/tests/python/relax/test_analysis.py b/tests/python/relax/test_analysis.py
index b6f4edf585..d9ad7d0de5 100644
--- a/tests/python/relax/test_analysis.py
+++ b/tests/python/relax/test_analysis.py
@@ -18,10 +18,11 @@
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax as rx
from tvm.relax import ExternFunc as _ExternFunc
+from tvm.relax import script as R
from tvm.relax.analysis import (
all_global_vars,
all_vars,
@@ -33,9 +34,8 @@ from tvm.relax.analysis import (
udchain,
used_vars,
)
-from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_computable_at_compile_time.py b/tests/python/relax/test_analysis_computable_at_compile_time.py
index f44ac42d0a..2a256e348d 100644
--- a/tests/python/relax/test_analysis_computable_at_compile_time.py
+++ b/tests/python/relax/test_analysis_computable_at_compile_time.py
@@ -17,8 +17,9 @@
# ruff: noqa: F841
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_contains_impure_call.py b/tests/python/relax/test_analysis_contains_impure_call.py
index a721f4f595..234fbfd70b 100644
--- a/tests/python/relax/test_analysis_contains_impure_call.py
+++ b/tests/python/relax/test_analysis_contains_impure_call.py
@@ -16,11 +16,11 @@
# under the License.
# ruff: noqa: F841
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax as rx
-from tvm.relax.analysis import contains_impure_call
from tvm.relax import script as R
+from tvm.relax.analysis import contains_impure_call
def test_simple_pure_case():
diff --git a/tests/python/relax/test_analysis_detect_recursion.py b/tests/python/relax/test_analysis_detect_recursion.py
index b2ca5889c6..85aa4888d5 100644
--- a/tests/python/relax/test_analysis_detect_recursion.py
+++ b/tests/python/relax/test_analysis_detect_recursion.py
@@ -15,11 +15,11 @@
# specific language governing permissions and limitations
# under the License.
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax as rx
-from tvm.relax.analysis import detect_recursion
from tvm.relax import script as R
+from tvm.relax.analysis import detect_recursion
from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_estimate_memory_usage.py b/tests/python/relax/test_analysis_estimate_memory_usage.py
index 615053ac15..dc692021da 100644
--- a/tests/python/relax/test_analysis_estimate_memory_usage.py
+++ b/tests/python/relax/test_analysis_estimate_memory_usage.py
@@ -16,10 +16,10 @@
# under the License.
# ruff: noqa: E501
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm.relax.analysis import estimate_memory_usage
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax import script as R
+from tvm.relax.analysis import estimate_memory_usage
from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_suggest_layout_transforms.py b/tests/python/relax/test_analysis_suggest_layout_transforms.py
index 3ead1c5b91..3fcd3e93cb 100644
--- a/tests/python/relax/test_analysis_suggest_layout_transforms.py
+++ b/tests/python/relax/test_analysis_suggest_layout_transforms.py
@@ -18,10 +18,11 @@
import pytest
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx.testing
from tvm import relax, s_tir
from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_analysis_type_analysis.py b/tests/python/relax/test_analysis_type_analysis.py
index e1a3b9ea9e..d14e863ca6 100644
--- a/tests/python/relax/test_analysis_type_analysis.py
+++ b/tests/python/relax/test_analysis_type_analysis.py
@@ -21,12 +21,13 @@
import pytest
import tvm_ffi
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import ir, tirx
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax as rx
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax import script as R
+from tvm_tirx import ir, tirx
from tvm_tirx.script import tirx as T
@@ -37,7 +38,9 @@ def test_get_static_type_basic():
# prim
s1 = tvm_tirx.ir.PrimType("float32")
- tvm_tirx.ir.assert_structural_equal(rx.analysis.get_static_type(s1), tvm_tirx.ir.PrimType("float32"))
+ tvm_tirx.ir.assert_structural_equal(
+ rx.analysis.get_static_type(s1), tvm_tirx.ir.PrimType("float32")
+ )
def test_get_static_type_shape():
@@ -635,7 +638,11 @@ def _generate_prim_test_cases():
for dtype in dtypes:
# LCA of a PrimType with itself yields itself
- yield (tvm_tirx.ir.PrimType(dtype), tvm_tirx.ir.PrimType(dtype), tvm_tirx.ir.PrimType(dtype))
+ yield (
+ tvm_tirx.ir.PrimType(dtype),
+ tvm_tirx.ir.PrimType(dtype),
+ tvm_tirx.ir.PrimType(dtype),
+ )
for dtype_a in dtypes:
for dtype_b in dtypes:
diff --git a/tests/python/relax/test_analysis_well_formed.py b/tests/python/relax/test_analysis_well_formed.py
index 1f7ed61172..30591e37ee 100644
--- a/tests/python/relax/test_analysis_well_formed.py
+++ b/tests/python/relax/test_analysis_well_formed.py
@@ -20,15 +20,16 @@ from __future__ import annotations
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax as rx
-from tvm_tirx import tirx
-from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
+from tvm_tirx.ir import TupleType as _TupleType
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
m = T.dynamic("m", "int64")
diff --git a/tests/python/relax/test_ast_printer.py b/tests/python/relax/test_ast_printer.py
index e784134a05..7092f0671e 100644
--- a/tests/python/relax/test_ast_printer.py
+++ b/tests/python/relax/test_ast_printer.py
@@ -20,15 +20,15 @@ from functools import partial
import numpy as np
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax as rx
-from tvm_tirx import tirx
from tvm.relax import ExternFunc as _ExternFunc
+from tvm.relax import script as R
from tvm.relax.testing import dump_ast
from tvm.relax.testing.ast_printer import ASTPrinter
-from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
from tvm_tirx.script import tirx as T
# Overload dump_ast to test both type and type annotations
diff --git a/tests/python/relax/test_backend_dispatch_sampling.py b/tests/python/relax/test_backend_dispatch_sampling.py
index 53cce1dea3..f9b2fb6bb1 100644
--- a/tests/python/relax/test_backend_dispatch_sampling.py
+++ b/tests/python/relax/test_backend_dispatch_sampling.py
@@ -17,14 +17,14 @@
# pylint: disable=missing-docstring
# ruff: noqa: E501
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx.ir.base import assert_structural_equal
-from tvm.relax.backend import DispatchSampling
-from tvm_tirx.script import ir as I
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm.relax import script as R
+from tvm.relax.backend import DispatchSampling
from tvm.s_tir import script as Ts
+from tvm_tirx.ir.base import assert_structural_equal
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py b/tests/python/relax/test_backend_dispatch_sort_scan.py
index bf8709e142..2cfe5f1073 100644
--- a/tests/python/relax/test_backend_dispatch_sort_scan.py
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -21,17 +21,18 @@ import numpy as np
import pytest
import tvm_ffi
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax, topi
from tvm.contrib.thrust import can_use_thrust
-from tvm_tirx.ir.base import assert_structural_equal
+from tvm.relax import script as R
from tvm.relax.backend import DispatchSortScan
from tvm.s_tir import dlight
+from tvm_tirx import tirx
+from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_backend_transform_shape_lower.py b/tests/python/relax/test_backend_transform_shape_lower.py
index eaa9bb6a4f..028cc3a199 100644
--- a/tests/python/relax/test_backend_transform_shape_lower.py
+++ b/tests/python/relax/test_backend_transform_shape_lower.py
@@ -16,14 +16,14 @@
# under the License.
# ruff: noqa: F841
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import assert_structural_equal
-from tvm.relax.testing.runtime_builtin import MakeShapeCode, MatchShapeCode
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm.relax.testing.runtime_builtin import MakeShapeCode, MatchShapeCode
from tvm.s_tir import script as Ts
+from tvm_tirx.ir import assert_structural_equal
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
# note: we expected RemovePurityChecking to be run first, so we force purity in most test cases
diff --git a/tests/python/relax/test_base_py_module_printer.py b/tests/python/relax/test_base_py_module_printer.py
index 71211c32ee..640971db30 100644
--- a/tests/python/relax/test_base_py_module_printer.py
+++ b/tests/python/relax/test_base_py_module_printer.py
@@ -19,12 +19,13 @@
import pytest
-import tvm, tvm_tirx
-from tvm_tirx.ir import StringImm
-from tvm.relax.base_py_module import BasePyModule
-from tvm_tirx.script import ir as I
+import tvm
+import tvm_tirx
from tvm.relax import script as R
+from tvm.relax.base_py_module import BasePyModule
from tvm.s_tir import script as Ts
+from tvm_tirx.ir import StringImm
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_base_py_module_symbolic_shape.py b/tests/python/relax/test_base_py_module_symbolic_shape.py
index a5a5398cb7..f73d48ff63 100644
--- a/tests/python/relax/test_base_py_module_symbolic_shape.py
+++ b/tests/python/relax/test_base_py_module_symbolic_shape.py
@@ -19,13 +19,13 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-from tvm_tirx import tirx
+import tvm_tirx
from tvm import relax
-from tvm_tirx.ir import IRModule
-from tvm.relax.base_py_module import BasePyModule
from tvm.relax import script as R
+from tvm.relax.base_py_module import BasePyModule
from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
+from tvm_tirx.ir import IRModule
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_bind_params.py b/tests/python/relax/test_bind_params.py
index 9d6a205e35..52adacdbc5 100644
--- a/tests/python/relax/test_bind_params.py
+++ b/tests/python/relax/test_bind_params.py
@@ -19,12 +19,12 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import script as R
+from tvm_tirx import tirx
param_specification = tvm_tirx.testing.parameter("by_string", "by_var")
param_shape = tvm_tirx.testing.parameter("static_shape", "dynamic_shape", "ndim", "arbitrary")
diff --git a/tests/python/relax/test_bind_symbolic_vars.py b/tests/python/relax/test_bind_symbolic_vars.py
index 00d2790fa1..14e846196e 100644
--- a/tests/python/relax/test_bind_symbolic_vars.py
+++ b/tests/python/relax/test_bind_symbolic_vars.py
@@ -16,8 +16,8 @@
# under the License.
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_binding_rewrite.py b/tests/python/relax/test_binding_rewrite.py
index 7b37c6b908..871fe109cb 100644
--- a/tests/python/relax/test_binding_rewrite.py
+++ b/tests/python/relax/test_binding_rewrite.py
@@ -18,13 +18,13 @@
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax import ExternFunc as _ExternFunc
+from tvm.relax import script as R
from tvm.relax.analysis import name_to_binding
from tvm.relax.binding_rewrite import DataflowBlockRewrite
from tvm.relax.expr import DataflowVar, Var
-from tvm.relax import script as R
@tvm_tirx.script.ir_module
diff --git a/tests/python/relax/test_blockbuilder_core.py b/tests/python/relax/test_blockbuilder_core.py
index 6dee48dd1b..9027266b55 100644
--- a/tests/python/relax/test_blockbuilder_core.py
+++ b/tests/python/relax/test_blockbuilder_core.py
@@ -20,16 +20,17 @@
# The test here do not depend on tvmscript to cover most basic features
import pytest
-import tvm, tvm_tirx
+import tvm
import tvm.contrib.cblas
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax as rx
-from tvm_tirx import tirx
from tvm import te, topi
-from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax import ExternFunc
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir.base import assert_structural_equal
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx.function import Function
diff --git a/tests/python/relax/test_blockbuilder_emit_te.py b/tests/python/relax/test_blockbuilder_emit_te.py
index 08c8c7bf4c..2b3e960921 100644
--- a/tests/python/relax/test_blockbuilder_emit_te.py
+++ b/tests/python/relax/test_blockbuilder_emit_te.py
@@ -18,11 +18,12 @@
# The tests here depend on tvmscript
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import relax as rx
from tvm import te
-from tvm_tirx.ir.base import assert_structural_equal
from tvm.s_tir import script as Ts
+from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_codegen_coreml.py b/tests/python/relax/test_codegen_coreml.py
index ea96426af4..fb813fccd8 100644
--- a/tests/python/relax/test_codegen_coreml.py
+++ b/tests/python/relax/test_codegen_coreml.py
@@ -19,8 +19,9 @@ import importlib.util
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
target, dev = "llvm", tvm_tirx.cpu()
@@ -28,7 +29,7 @@ target, dev = "llvm", tvm_tirx.cpu()
def _has_xcode():
try:
- import tvm_tirx.support.xcode, tvm
+ import tvm_tirx.support.xcode
tvm_tirx.support.xcode.xcrun([])
return True
diff --git a/tests/python/relax/test_codegen_cublas.py b/tests/python/relax/test_codegen_cublas.py
index f0ee137d19..c268c3d74d 100644
--- a/tests/python/relax/test_codegen_cublas.py
+++ b/tests/python/relax/test_codegen_cublas.py
@@ -18,18 +18,19 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm_tirx.testing import env
pytest.importorskip("scipy") # tvm.topi.testing imports scipy
import tvm.topi.testing
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.backend.cuda.cublas import partition_for_cublas
from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.testing import get_relax_matmul_module
-from tvm.relax import script as R
from tvm_tirx.script.ir_builder import IRBuilder
try:
diff --git a/tests/python/relax/test_codegen_cudnn.py b/tests/python/relax/test_codegen_cudnn.py
index 2bf2d540ab..d3d3c5d354 100644
--- a/tests/python/relax/test_codegen_cudnn.py
+++ b/tests/python/relax/test_codegen_cudnn.py
@@ -18,8 +18,9 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm_tirx.testing import env
pytest.importorskip("scipy") # tvm.topi.testing imports scipy
@@ -27,10 +28,10 @@ pytest.importorskip("scipy") # tvm.topi.testing imports scipy
import tvm.topi.testing
from tvm import relax
from tvm.contrib.pickle_memoize import memoize
+from tvm.relax import script as R
from tvm.relax.backend.cuda.cudnn import partition_for_cudnn
from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.testing import get_relax_stacked_attention_module
-from tvm.relax import script as R
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/test_codegen_cutlass.py b/tests/python/relax/test_codegen_cutlass.py
index f1f50ec6e7..ca42273b2a 100644
--- a/tests/python/relax/test_codegen_cutlass.py
+++ b/tests/python/relax/test_codegen_cutlass.py
@@ -18,10 +18,11 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx.ir import TupleType as _TupleType
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax import ExternFunc as _ExternFunc
+from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.testing import env
pytest.importorskip("scipy") # tvm.topi.testing imports scipy
@@ -30,6 +31,7 @@ import tvm.topi.testing
from tvm import relax
from tvm.contrib.cutlass.build import is_shape_valid_for_cutlass_matmul
from tvm.contrib.pickle_memoize import memoize
+from tvm.relax import script as R
from tvm.relax.backend.cuda.cutlass import partition_for_cutlass
from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.testing import (
@@ -37,9 +39,8 @@ from tvm.relax.testing import (
get_relax_matmul_module,
get_relax_stacked_attention_module,
)
-from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/test_codegen_dnnl.py b/tests/python/relax/test_codegen_dnnl.py
index e3c7a38779..c2d694eac7 100644
--- a/tests/python/relax/test_codegen_dnnl.py
+++ b/tests/python/relax/test_codegen_dnnl.py
@@ -17,12 +17,13 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm.contrib.pickle_memoize import memoize
-from tvm.relax.dpl import make_fused_bias_activation_pattern
from tvm.relax import script as R
+from tvm.relax.dpl import make_fused_bias_activation_pattern
@tvm_tirx.script.ir_module
diff --git a/tests/python/relax/test_codegen_hipblas.py b/tests/python/relax/test_codegen_hipblas.py
index 26c47fa4a1..85e76a4f3e 100644
--- a/tests/python/relax/test_codegen_hipblas.py
+++ b/tests/python/relax/test_codegen_hipblas.py
@@ -17,17 +17,18 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm_tirx.testing import env
pytest.importorskip("scipy") # tvm.topi.testing imports scipy
import tvm.topi.testing
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.backend.rocm.hipblas import partition_for_hipblas
from tvm.relax.testing import get_relax_matmul_module
-from tvm.relax import script as R
try:
import ml_dtypes
diff --git a/tests/python/relax/test_codegen_tensorrt.py b/tests/python/relax/test_codegen_tensorrt.py
index 43a7d3b643..0821283d05 100644
--- a/tests/python/relax/test_codegen_tensorrt.py
+++ b/tests/python/relax/test_codegen_tensorrt.py
@@ -17,13 +17,14 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm.contrib.pickle_memoize import memoize
+from tvm.relax import script as R
from tvm.relax.backend.contrib.tensorrt import partition_for_tensorrt
from tvm.relax.dpl import is_op, make_fused_bias_activation_pattern, wildcard
-from tvm.relax import script as R
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_contrib_vllm.py b/tests/python/relax/test_contrib_vllm.py
index 53ee0475c0..1833b4f679 100644
--- a/tests/python/relax/test_contrib_vllm.py
+++ b/tests/python/relax/test_contrib_vllm.py
@@ -18,12 +18,13 @@ import numpy as np
import pytest
import tvm_ffi
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
@@ -767,7 +768,9 @@ def test_reconstruct_from_cache():
num_blocks = 1
reshape_and_cache_func = tvm_tirx.get_global_func("tvm.contrib.vllm.reshape_and_cache")
- reconstruct_from_cache_func = tvm_tirx.get_global_func("tvm.contrib.vllm.reconstruct_from_cache")
+ reconstruct_from_cache_func = tvm_tirx.get_global_func(
+ "tvm.contrib.vllm.reconstruct_from_cache"
+ )
def run_and_check():
dev = tvm_tirx.cuda(0)
diff --git a/tests/python/relax/test_dataflow_inplace.py b/tests/python/relax/test_dataflow_inplace.py
index 0e71f2bb4a..2da251af0c 100644
--- a/tests/python/relax/test_dataflow_inplace.py
+++ b/tests/python/relax/test_dataflow_inplace.py
@@ -20,10 +20,9 @@ import numpy as np
import pytest
import torch
-import tvm, tvm_tirx
-from tvm_tirx import testing
+import tvm
+import tvm_tirx
from tvm import relax
-from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax.testing.transform import (
dataflow_alias_analysis,
dataflow_inplace_analysis,
@@ -32,6 +31,8 @@ from tvm.relax.testing.transform import (
)
from tvm.relax.transform import DataflowUseInplaceCalls
from tvm.s_tir import script as Ts
+from tvm_tirx import testing
+from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_dataflow_pattern.py b/tests/python/relax/test_dataflow_pattern.py
index d640bcbf65..f27e88ec0a 100644
--- a/tests/python/relax/test_dataflow_pattern.py
+++ b/tests/python/relax/test_dataflow_pattern.py
@@ -24,15 +24,16 @@ import math
import pytest
import tvm_ffi
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx.testing
from tvm import relax as rx
-from tvm_tirx import tirx
-from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
+from tvm.relax import script as R
from tvm.relax.analysis import get_var2val
from tvm.relax.dpl import *
-from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
+from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_dataflow_rewriter.py b/tests/python/relax/test_dataflow_rewriter.py
index 5c60778706..88126b9547 100644
--- a/tests/python/relax/test_dataflow_rewriter.py
+++ b/tests/python/relax/test_dataflow_rewriter.py
@@ -18,10 +18,10 @@
import pytest
-import tvm_tirx.testing, tvm
-from tvm_tirx.script import ir as I
+import tvm_tirx.testing
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_dlpack_integration.py b/tests/python/relax/test_dlpack_integration.py
index 45427c6aaa..740fe02519 100644
--- a/tests/python/relax/test_dlpack_integration.py
+++ b/tests/python/relax/test_dlpack_integration.py
@@ -27,8 +27,8 @@ This test verifies:
import pytest
import torch
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax import BasePyModule
from tvm.relax import script as R
from tvm.s_tir import script as Ts
diff --git a/tests/python/relax/test_e2e_op_dynamic.py b/tests/python/relax/test_e2e_op_dynamic.py
index fb111c89a6..982886f302 100644
--- a/tests/python/relax/test_e2e_op_dynamic.py
+++ b/tests/python/relax/test_e2e_op_dynamic.py
@@ -18,16 +18,17 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
pytest.importorskip("scipy") # tvm.topi.testing imports scipy
import tvm.topi.testing
-from tvm_tirx import tirx
from tvm import relax
-from tvm.relax.transform import LegalizeOps
from tvm.relax import script as R
+from tvm.relax.transform import LegalizeOps
+from tvm_tirx import tirx
from tvm_tirx.script import tirx as T
# TODO(tvm-team): `tirx.transform.DefaultGPUSchedule` does not work.
diff --git a/tests/python/relax/test_expr.py b/tests/python/relax/test_expr.py
index 964c605e17..57dc468c29 100644
--- a/tests/python/relax/test_expr.py
+++ b/tests/python/relax/test_expr.py
@@ -21,11 +21,12 @@ import numpy as np
import pytest
import tvm_ffi
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import relax as rx
-from tvm_tirx import tirx
-from tvm.relax.expr import make_shape
from tvm.relax import script as R
+from tvm.relax.expr import make_shape
+from tvm_tirx import tirx
def _check_equal(x, y, map_free_vars=False):
@@ -73,7 +74,9 @@ def test_var_name_keyword_compatibility() -> None:
def test_tensor_type_empty_dtype_is_unknown() -> None:
- tvm_tirx.ir.assert_structural_equal(rx.TensorType([1, 2], dtype=""), rx.TensorType([1, 2], None))
+ tvm_tirx.ir.assert_structural_equal(
+ rx.TensorType([1, 2], dtype=""), rx.TensorType([1, 2], None)
+ )
def test_relax_expr_ty_running_example() -> None:
@@ -414,7 +417,9 @@ def test_call_reinfer_type_from_current_inputs():
func = rx.Var("func", rx.FuncType([original_ty], original_ty))
_check_type_missing(tvm_tirx.ir.reinfer_type(rx.Call(func, [x])))
- invalid = rx.Call("relax.abs", [rx.Var("scalar", tvm_tirx.ir.PrimType("int32"))], ty=original_ty)
+ invalid = rx.Call(
+ "relax.abs", [rx.Var("scalar", tvm_tirx.ir.PrimType("int32"))], ty=original_ty
+ )
with pytest.raises(TypeError):
tvm_tirx.ir.reinfer_type(invalid)
@@ -427,7 +432,9 @@ def test_call_reinfer_type_from_current_inputs():
)
_check_equal(tvm_tirx.ir.reinfer_type(changed_explicit), changed_ty)
- _check_equal(tvm_tirx.ir.reinfer_type(rx.op.astype(x, "float16")), rx.TensorType([2], "float16"))
+ _check_equal(
+ tvm_tirx.ir.reinfer_type(rx.op.astype(x, "float16")), rx.TensorType([2], "float16")
+ )
_check_equal(tvm_tirx.ir.reinfer_type(rx.op.astype(x, "int32")), rx.TensorType([2], "int32"))
diff --git a/tests/python/relax/test_expr_functor.py b/tests/python/relax/test_expr_functor.py
index d8b563410a..a9fb88783e 100644
--- a/tests/python/relax/test_expr_functor.py
+++ b/tests/python/relax/test_expr_functor.py
@@ -17,13 +17,12 @@
# ruff: noqa: E501, F841
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax, te
-from tvm_tirx.ir import Call, GenericConst, Op, StringImm
-from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax import PyExprMutator, PyExprVisitor
+from tvm.relax import script as R
from tvm.relax.expr import (
BindingBlock,
DataflowBlock,
@@ -41,7 +40,9 @@ from tvm.relax.expr import (
Var,
VarBinding,
)
-from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Call, GenericConst, Op, StringImm
+from tvm_tirx.ir.base import assert_structural_equal
m, n = tirx.Var("m", "int64"), tirx.Var("n", "int64")
x = relax.Var("x", R.Tensor([n], "float32"))
diff --git a/tests/python/relax/test_fast_math_transform.py b/tests/python/relax/test_fast_math_transform.py
index 51804983c7..321c2216fa 100644
--- a/tests/python/relax/test_fast_math_transform.py
+++ b/tests/python/relax/test_fast_math_transform.py
@@ -19,12 +19,13 @@
import pytest
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx.testing
from tvm import relax, topi
-from tvm_tirx.ir.base import assert_structural_equal
+from tvm.relax import script as R
from tvm.relax.transform import FastMathTransform
+from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
def _run_pass_compare_output(Before, Expected):
diff --git a/tests/python/relax/test_frontend_common.py b/tests/python/relax/test_frontend_common.py
index d49903c5b3..3734216e54 100644
--- a/tests/python/relax/test_frontend_common.py
+++ b/tests/python/relax/test_frontend_common.py
@@ -14,13 +14,13 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm.relax.frontend import detach_params
from tvm.relax.frontend.common import autopad
-from tvm_tirx.script import ir as I
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
from tvm_tirx.script.parser import relax as R
diff --git a/tests/python/relax/test_frontend_dynamo.py b/tests/python/relax/test_frontend_dynamo.py
index b69c5b57fc..71a8bd8a46 100644
--- a/tests/python/relax/test_frontend_dynamo.py
+++ b/tests/python/relax/test_frontend_dynamo.py
@@ -24,15 +24,16 @@ import torch
import torch._dynamo as dynamo
from packaging import version
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.frontend.torch import relax_dynamo
from tvm.s_tir import meta_schedule as ms
-from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py
index cbe277fb11..ded394285e 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -24,12 +24,13 @@ from torch import nn
from torch.export import export
from torch.nn import Module
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.frontend.torch import from_exported_program
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py
index f84f7ee92c..ebff36a3ad 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -23,13 +23,13 @@ import torch.nn.functional as F
from torch import fx
from torch.nn import Module
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.frontend import detach_params
from tvm.relax.frontend.torch import from_fx
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_frontend_nn_debug.py b/tests/python/relax/test_frontend_nn_debug.py
index 827a4409ba..7d08b90711 100644
--- a/tests/python/relax/test_frontend_nn_debug.py
+++ b/tests/python/relax/test_frontend_nn_debug.py
@@ -17,11 +17,11 @@
# pylint: disable=missing-docstring
import torch
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax.frontend import nn
from tvm.relax.frontend.nn import op, spec
+from tvm_tirx import tirx
from tvm_tirx.runtime import Tensor
diff --git a/tests/python/relax/test_frontend_nn_exporter.py b/tests/python/relax/test_frontend_nn_exporter.py
index 86d34e84aa..94c87e8a2e 100644
--- a/tests/python/relax/test_frontend_nn_exporter.py
+++ b/tests/python/relax/test_frontend_nn_exporter.py
@@ -17,14 +17,14 @@
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import assert_structural_equal
+from tvm.relax import script as R
from tvm.relax.frontend import nn
+from tvm_tirx import tirx
+from tvm_tirx.ir import assert_structural_equal
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_frontend_nn_extern_module.py b/tests/python/relax/test_frontend_nn_extern_module.py
index 06cf8f9908..991c020ecd 100644
--- a/tests/python/relax/test_frontend_nn_extern_module.py
+++ b/tests/python/relax/test_frontend_nn_extern_module.py
@@ -23,9 +23,10 @@ from pathlib import Path
import numpy as np
import tvm_ffi
-import tvm, tvm_tirx
-import tvm_tirx.libinfo, tvm
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.libinfo
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax.frontend import nn
@@ -81,8 +82,8 @@ def _test_infer_sym(func, x, y, z): # pylint: disable=invalid-name
def _check_ir_equality(mod):
# pylint: disable=import-outside-toplevel
- from tvm_tirx.script import ir as I
from tvm.relax import script as R
+ from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
# pylint: enable=import-outside-toplevel
diff --git a/tests/python/relax/test_frontend_nn_jit.py b/tests/python/relax/test_frontend_nn_jit.py
index 7f9db65c26..4be15bb43d 100644
--- a/tests/python/relax/test_frontend_nn_jit.py
+++ b/tests/python/relax/test_frontend_nn_jit.py
@@ -18,11 +18,11 @@
import pytest
import torch
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax.frontend import nn
from tvm.relax.frontend.nn import spec
+from tvm_tirx import tirx
@pytest.mark.parametrize("debug", [True, False])
diff --git a/tests/python/relax/test_frontend_nn_llm_kernel_config.py b/tests/python/relax/test_frontend_nn_llm_kernel_config.py
index 6c04f1a2e2..61cad58b4c 100644
--- a/tests/python/relax/test_frontend_nn_llm_kernel_config.py
+++ b/tests/python/relax/test_frontend_nn_llm_kernel_config.py
@@ -17,8 +17,9 @@
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax.frontend.nn.llm._decode_kernels import _merge_state_inplace
from tvm.relax.frontend.nn.llm._kernel_common import (
_get_prefill_kernel_config,
diff --git a/tests/python/relax/test_frontend_nn_llm_sequence_prefill_masked.py b/tests/python/relax/test_frontend_nn_llm_sequence_prefill_masked.py
index c49420f25c..acd37a1e74 100644
--- a/tests/python/relax/test_frontend_nn_llm_sequence_prefill_masked.py
+++ b/tests/python/relax/test_frontend_nn_llm_sequence_prefill_masked.py
@@ -45,8 +45,8 @@ import math
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax.frontend.nn.llm.kv_cache import _attention_sequence_prefill_with_mask
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_frontend_nn_modules.py b/tests/python/relax/test_frontend_nn_modules.py
index ab6d1d9b52..f76fc60047 100644
--- a/tests/python/relax/test_frontend_nn_modules.py
+++ b/tests/python/relax/test_frontend_nn_modules.py
@@ -19,14 +19,15 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import assert_structural_equal
+from tvm.relax import script as R
from tvm.relax.frontend import nn
from tvm.relax.frontend.nn import core, modules, spec
+from tvm_tirx.ir import assert_structural_equal
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_frontend_nn_mutator.py b/tests/python/relax/test_frontend_nn_mutator.py
index 6b47fafbc7..3677423e94 100644
--- a/tests/python/relax/test_frontend_nn_mutator.py
+++ b/tests/python/relax/test_frontend_nn_mutator.py
@@ -17,7 +17,7 @@
from typing import Any
-import tvm, tvm_tirx
+import tvm_tirx
from tvm.relax.frontend import nn
diff --git a/tests/python/relax/test_frontend_nn_op.py b/tests/python/relax/test_frontend_nn_op.py
index 676692730c..1639009a17 100644
--- a/tests/python/relax/test_frontend_nn_op.py
+++ b/tests/python/relax/test_frontend_nn_op.py
@@ -19,16 +19,17 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax, s_tir
-from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
-from tvm.relax.frontend.nn import Module, Tensor, op, spec
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm.relax.frontend.nn import Module, Tensor, op, spec
from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
+from tvm_tirx.ir import TupleType as _TupleType
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
@@ -988,7 +989,9 @@ def test_multinomial_from_uniform():
np_prob = np_rand / np_rand.sum(axis=1, keepdims=True)
nd_prob = tvm_tirx.runtime.tensor(np_prob, dev)
# special sample to get deterministic results
- nd_sample = tvm_tirx.runtime.tensor(np.array([[1], [0], [1], [1], [0], [1]]).astype(np.float32), dev)
+ nd_sample = tvm_tirx.runtime.tensor(
+ np.array([[1], [0], [1], [1], [0], [1]]).astype(np.float32), dev
+ )
nd_sample_indices = tvm_tirx.runtime.tensor(
np.array([[0], [1], [1], [2], [2], [2]]).astype(np.int64), dev
)
diff --git a/tests/python/relax/test_frontend_nn_packing.py b/tests/python/relax/test_frontend_nn_packing.py
index 6342c17bd3..63e18c9d2d 100644
--- a/tests/python/relax/test_frontend_nn_packing.py
+++ b/tests/python/relax/test_frontend_nn_packing.py
@@ -15,10 +15,10 @@
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name,missing-docstring
-import tvm, tvm_tirx
+import tvm_tirx
+from tvm.relax import script as R
from tvm.relax.frontend import nn
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
def _iter_binding_names(mod):
diff --git a/tests/python/relax/test_frontend_nn_parameter_containers.py b/tests/python/relax/test_frontend_nn_parameter_containers.py
index 1e5d4a6f88..9c25464e38 100644
--- a/tests/python/relax/test_frontend_nn_parameter_containers.py
+++ b/tests/python/relax/test_frontend_nn_parameter_containers.py
@@ -20,8 +20,8 @@ from typing import Any
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax.frontend import nn
@@ -171,8 +171,12 @@ def test_load_state_dict():
assert unexpected_keys == []
tvm_tirx.testing.assert_allclose(m.list_params[0].data.numpy(), np.full((4,), 1.0, "float32"))
tvm_tirx.testing.assert_allclose(m.list_params[1].data.numpy(), np.full((4,), 2.0, "float32"))
- tvm_tirx.testing.assert_allclose(m.dict_params["foo"].data.numpy(), np.full((4,), 3.0, "float32"))
- tvm_tirx.testing.assert_allclose(m.dict_params["bar"].data.numpy(), np.full((4,), 4.0, "float32"))
+ tvm_tirx.testing.assert_allclose(
+ m.dict_params["foo"].data.numpy(), np.full((4,), 3.0, "float32")
+ )
+ tvm_tirx.testing.assert_allclose(
+ m.dict_params["bar"].data.numpy(), np.full((4,), 4.0, "float32")
+ )
def test_export_tvm_parameter_names():
diff --git a/tests/python/relax/test_frontend_nn_subroutines.py b/tests/python/relax/test_frontend_nn_subroutines.py
index 11555869ca..d4d0c202c8 100644
--- a/tests/python/relax/test_frontend_nn_subroutines.py
+++ b/tests/python/relax/test_frontend_nn_subroutines.py
@@ -15,13 +15,13 @@
# specific language governing permissions and limitations
# under the License.
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import assert_structural_equal
+from tvm.relax import script as R
from tvm.relax.frontend import nn
+from tvm_tirx.ir import assert_structural_equal
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
def test_linear():
diff --git a/tests/python/relax/test_frontend_nn_tensor.py b/tests/python/relax/test_frontend_nn_tensor.py
index 01e22cea37..a8c4b54875 100644
--- a/tests/python/relax/test_frontend_nn_tensor.py
+++ b/tests/python/relax/test_frontend_nn_tensor.py
@@ -18,11 +18,12 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm.relax.frontend.nn import Module, Tensor, spec
from tvm.relax import script as R
+from tvm.relax.frontend.nn import Module, Tensor, spec
def test_tensor_from_numpy():
diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py
index 2956cdc0b2..4035a8f845 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -38,14 +38,15 @@ import onnxruntime
import tvm_ffi
from onnx import ModelProto, TensorProto, helper, numpy_helper
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm.relax.frontend.onnx import from_onnx
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm.relax.frontend.onnx import from_onnx
from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
bg = np.random.MT19937(0)
@@ -11765,7 +11766,9 @@ def test_symbolic_shape_deduction():
)
model = helper.make_model(graph, producer_name="test_shape_deduction")
tvm_model = from_onnx(model, keep_params_in_input=True)
- tvm_tirx.ir.assert_structural_equal(tvm_model["main"].without_attr("params"), expected["main"])
+ tvm_tirx.ir.assert_structural_equal(
+ tvm_model["main"].without_attr("params"), expected["main"]
+ )
batch = T.dynamic("batch", "int64")
seq = T.dynamic("seq", "int64")
diff --git a/tests/python/relax/test_frontend_onnx_backend.py b/tests/python/relax/test_frontend_onnx_backend.py
index 2159266648..6fa706c150 100644
--- a/tests/python/relax/test_frontend_onnx_backend.py
+++ b/tests/python/relax/test_frontend_onnx_backend.py
@@ -40,7 +40,8 @@ import onnx
import onnx.backend.test
from onnx.backend.base import Backend, BackendRep
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx
diff --git a/tests/python/relax/test_frontend_stablehlo.py b/tests/python/relax/test_frontend_stablehlo.py
index adec4238ef..336acd6d05 100644
--- a/tests/python/relax/test_frontend_stablehlo.py
+++ b/tests/python/relax/test_frontend_stablehlo.py
@@ -28,12 +28,13 @@ import functools
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.frontend.stablehlo import from_stablehlo
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_frontend_tflite.py b/tests/python/relax/test_frontend_tflite.py
index 0016b32b75..ecf3e32494 100644
--- a/tests/python/relax/test_frontend_tflite.py
+++ b/tests/python/relax/test_frontend_tflite.py
@@ -34,8 +34,9 @@ import tensorflow as tf
import tflite.Model
from tensorflow.keras import applications as keras_app
-import tvm, tvm_tirx
+import tvm
import tvm.relax.frontend.tflite.tflite_frontend as tflite_frontend
+import tvm_tirx
from tvm import relax
from tvm.relax.frontend.tflite import from_tflite
from tvm_tirx.script.parser import ir as I
@@ -5177,7 +5178,9 @@ def _build_tflite_operator_marker_model(builtin_name):
@pytest.mark.parametrize("builtin_name", ["DELEGATE", "PLACEHOLDER_FOR_GREATER_OP_CODES"])
def test_operator_marker_unsupported(builtin_name):
"""TFLite marker builtins report explicit unsupported diagnostics."""
- with pytest.raises(tvm_tirx.error.OpNotImplemented, match=f"TFLite operator marker {builtin_name}"):
+ with pytest.raises(
+ tvm_tirx.error.OpNotImplemented, match=f"TFLite operator marker {builtin_name}"
+ ):
_load_model_from_buffer(_build_tflite_operator_marker_model(builtin_name))
@@ -5945,7 +5948,9 @@ def test_call_subgraph_nested_call():
def test_call_subgraph_invalid_index_unsupported():
"""Test CALL rejects invalid subgraph indices before lowering."""
- with pytest.raises(tvm_tirx.error.OpNotImplemented, match="CALL requires a valid subgraph index"):
+ with pytest.raises(
+ tvm_tirx.error.OpNotImplemented, match="CALL requires a valid subgraph index"
+ ):
_load_model_from_buffer(_build_tflite_call_model(call_subgraph_index=2))
@@ -6287,7 +6292,9 @@ def test_if_subgraphs_multi_output():
def test_if_subgraphs_non_bool_condition_unsupported():
"""Test IF rejects non-bool condition tensors."""
- with pytest.raises(tvm_tirx.error.OpNotImplemented, match="IF requires a scalar bool condition"):
+ with pytest.raises(
+ tvm_tirx.error.OpNotImplemented, match="IF requires a scalar bool condition"
+ ):
_load_model_from_buffer(_build_tflite_if_model(condition_type=_tfl_tensor_type.INT32))
@@ -6757,13 +6764,17 @@ def test_while_subgraphs_two_loop_vars():
def test_while_subgraphs_non_bool_condition_unsupported():
"""Test WHILE rejects cond subgraphs that do not return scalar bool."""
- with pytest.raises(tvm_tirx.error.OpNotImplemented, match="WHILE requires a scalar bool condition"):
+ with pytest.raises(
+ tvm_tirx.error.OpNotImplemented, match="WHILE requires a scalar bool condition"
+ ):
_load_model_from_buffer(_build_tflite_while_model(cond_output_type=_tfl_tensor_type.INT32))
def test_while_subgraphs_invalid_index_unsupported():
"""Test WHILE rejects invalid cond/body subgraph indices before lowering."""
- with pytest.raises(tvm_tirx.error.OpNotImplemented, match="WHILE requires a valid subgraph index"):
+ with pytest.raises(
+ tvm_tirx.error.OpNotImplemented, match="WHILE requires a valid subgraph index"
+ ):
_load_model_from_buffer(_build_tflite_while_model(cond_subgraph_index=3))
@@ -6785,7 +6796,9 @@ def test_while_subgraphs_loop_state_metadata_mismatch_unsupported():
def test_while_subgraphs_output_count_mismatch_unsupported():
"""Test WHILE rejects body subgraphs whose output arity does not match loop vars."""
- with pytest.raises(tvm_tirx.error.OpNotImplemented, match="WHILE subgraph output count mismatch"):
+ with pytest.raises(
+ tvm_tirx.error.OpNotImplemented, match="WHILE subgraph output count mismatch"
+ ):
_load_model_from_buffer(_build_tflite_while_model(body_outputs=[]))
@@ -9150,7 +9163,9 @@ def test_stablehlo_rng_bit_generator_constant_state():
def test_stablehlo_rng_bit_generator_unsupported_output_dtype():
"""TFLite STABLEHLO_RNG_BIT_GENERATOR rejects non-integer output dtypes."""
buf = _build_stablehlo_rng_model(_tfl_rng_algorithm.PHILOX, 3, [2, 3], _tfl_tensor_type.FLOAT32)
- with pytest.raises(tvm_tirx.error.OpNotImplemented, match="output dtype float32 is not supported"):
+ with pytest.raises(
+ tvm_tirx.error.OpNotImplemented, match="output dtype float32 is not supported"
+ ):
_load_model_from_buffer(buf)
diff --git a/tests/python/relax/test_group_gemm_flashinfer.py b/tests/python/relax/test_group_gemm_flashinfer.py
index 64f11e1230..43fbd942c6 100644
--- a/tests/python/relax/test_group_gemm_flashinfer.py
+++ b/tests/python/relax/test_group_gemm_flashinfer.py
@@ -24,8 +24,9 @@ import numpy as np
import pytest
import torch
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
DEFAULT_WORKSPACE_SIZE = 32 * 1024 * 1024
@@ -424,7 +425,9 @@ def test_grouped_gemm_correctness(
else:
output = tvm_tirx.runtime.empty(output_shape, dtype="float32", device=device)
- int_workspace = tvm_tirx.runtime.empty((DEFAULT_WORKSPACE_SIZE,), dtype="int32", device=device)
+ int_workspace = tvm_tirx.runtime.empty(
+ (DEFAULT_WORKSPACE_SIZE,), dtype="int32", device=device
+ )
float_workspace = tvm_tirx.runtime.empty(
(DEFAULT_WORKSPACE_SIZE,), dtype="float32", device=device
)
diff --git a/tests/python/relax/test_inline_functions.py b/tests/python/relax/test_inline_functions.py
index 3deb8799ce..9dc3e283f3 100644
--- a/tests/python/relax/test_inline_functions.py
+++ b/tests/python/relax/test_inline_functions.py
@@ -17,10 +17,10 @@
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx.script import ir as I
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax import script as R
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_kill_after_last_use.py b/tests/python/relax/test_kill_after_last_use.py
index 8a21b44413..f947e855c0 100644
--- a/tests/python/relax/test_kill_after_last_use.py
+++ b/tests/python/relax/test_kill_after_last_use.py
@@ -17,12 +17,13 @@
import pytest
-import tvm, tvm_tirx
+import tvm
import tvm.relax
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
+from tvm.relax import script as R
from tvm.relax.transform import KillAfterLastUse
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_lower_alloc_tensor.py b/tests/python/relax/test_lower_alloc_tensor.py
index 928a317c4b..1d7cd3a1cd 100644
--- a/tests/python/relax/test_lower_alloc_tensor.py
+++ b/tests/python/relax/test_lower_alloc_tensor.py
@@ -15,11 +15,11 @@
# specific language governing permissions and limitations
# under the License.
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
+from tvm.relax import script as R
from tvm.relax.transform import LowerAllocTensor
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
def test_basic():
diff --git a/tests/python/relax/test_meta_schedule_relax_integration.py b/tests/python/relax/test_meta_schedule_relax_integration.py
index 8f8a35cb7e..75e706430f 100644
--- a/tests/python/relax/test_meta_schedule_relax_integration.py
+++ b/tests/python/relax/test_meta_schedule_relax_integration.py
@@ -22,15 +22,16 @@ import tempfile
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.runtime import cpu as tvm_cpu
-from tvm_tirx.runtime import tensor as tvm_tensor
+from tvm.relax import script as R
from tvm.runtime.vm import VirtualMachine
from tvm.s_tir import meta_schedule as ms
+from tvm_tirx.runtime import cpu as tvm_cpu
+from tvm_tirx.runtime import tensor as tvm_tensor
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_nn_llm_prefill_parallel.py b/tests/python/relax/test_nn_llm_prefill_parallel.py
index 172e023804..7e437e1568 100644
--- a/tests/python/relax/test_nn_llm_prefill_parallel.py
+++ b/tests/python/relax/test_nn_llm_prefill_parallel.py
@@ -17,8 +17,9 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax.frontend.nn.llm._prefill_kernels import (
_attention_prefill,
_attention_prefill_ragged,
diff --git a/tests/python/relax/test_op_binary.py b/tests/python/relax/test_op_binary.py
index b14f8a6d25..163db0c972 100644
--- a/tests/python/relax/test_op_binary.py
+++ b/tests/python/relax/test_op_binary.py
@@ -19,13 +19,13 @@ from collections.abc import Callable
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import VDevice
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
def test_op_correctness():
diff --git a/tests/python/relax/test_op_ccl.py b/tests/python/relax/test_op_ccl.py
index 3df7926737..cbc8f251ed 100644
--- a/tests/python/relax/test_op_ccl.py
+++ b/tests/python/relax/test_op_ccl.py
@@ -17,12 +17,13 @@
# ruff: noqa: F401
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
def test_op_correctness():
diff --git a/tests/python/relax/test_op_create.py b/tests/python/relax/test_op_create.py
index 1f18ea2e54..1f377283fe 100644
--- a/tests/python/relax/test_op_create.py
+++ b/tests/python/relax/test_op_create.py
@@ -16,13 +16,13 @@
# under the License.
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import VDevice
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_op_datatype.py b/tests/python/relax/test_op_datatype.py
index bd1cb27811..43e6756a32 100644
--- a/tests/python/relax/test_op_datatype.py
+++ b/tests/python/relax/test_op_datatype.py
@@ -17,12 +17,12 @@
import numpy as np # type: ignore
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
def test_op_correctness():
diff --git a/tests/python/relax/test_op_distributed.py b/tests/python/relax/test_op_distributed.py
index d5fbd72bdb..76a0aa42fd 100644
--- a/tests/python/relax/test_op_distributed.py
+++ b/tests/python/relax/test_op_distributed.py
@@ -16,8 +16,8 @@
# under the License.
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm_tirx.script.parser import relax as R
@@ -34,7 +34,9 @@ def test_redistribute_R_to_S():
# Distributed placement requires the contextual builder inference rule.
distributed_abs = relax.op.abs(x)
- tvm_tirx.ir.assert_structural_equal(tvm_tirx.ir.reinfer_type(distributed_abs), tvm_tirx.ir.Type.missing())
+ tvm_tirx.ir.assert_structural_equal(
+ tvm_tirx.ir.reinfer_type(distributed_abs), tvm_tirx.ir.Type.missing()
+ )
_check_inference(bb, distributed_abs, x.ty)
_check_inference(
diff --git a/tests/python/relax/test_op_grad.py b/tests/python/relax/test_op_grad.py
index ab7dd792ba..067a6ee759 100644
--- a/tests/python/relax/test_op_grad.py
+++ b/tests/python/relax/test_op_grad.py
@@ -16,11 +16,11 @@
# under the License.
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import script as R
+from tvm_tirx.ir import Op
def test_op_correctness():
diff --git a/tests/python/relax/test_op_gradient_numeric.py b/tests/python/relax/test_op_gradient_numeric.py
index 1a89790dcd..07bdeb0bd5 100644
--- a/tests/python/relax/test_op_gradient_numeric.py
+++ b/tests/python/relax/test_op_gradient_numeric.py
@@ -22,14 +22,15 @@ import numpy as np
import pytest
import tvm_ffi
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Call
-from tvm_tirx.ir.op import Op
from tvm.relax.testing.gradient import check_numerical_grads
from tvm.relax.transform import LegalizeOps
from tvm.relax.type import TensorType, TupleType
+from tvm_tirx.ir import Call
+from tvm_tirx.ir.op import Op
def relax_check_gradients(
diff --git a/tests/python/relax/test_op_image.py b/tests/python/relax/test_op_image.py
index e90c16f02a..cbd738109c 100644
--- a/tests/python/relax/test_op_image.py
+++ b/tests/python/relax/test_op_image.py
@@ -17,17 +17,18 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
pytest.importorskip("scipy") # tvm.topi.testing imports scipy
import tvm.topi.testing
-from tvm_tirx import tirx
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import VDevice
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
def test_op_correctness():
diff --git a/tests/python/relax/test_op_index.py b/tests/python/relax/test_op_index.py
index 0302ecbe5d..e99f53a82b 100644
--- a/tests/python/relax/test_op_index.py
+++ b/tests/python/relax/test_op_index.py
@@ -18,15 +18,16 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import VDevice
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_op_inspect.py b/tests/python/relax/test_op_inspect.py
index 18df1a09bd..fc2967eab7 100644
--- a/tests/python/relax/test_op_inspect.py
+++ b/tests/python/relax/test_op_inspect.py
@@ -22,11 +22,12 @@ import numpy as np
import pytest
import tvm_ffi
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx.testing
from tvm import relax
+from tvm.relax import script as R
from tvm_tirx.ir import Op
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm_tirx.script import tirx as T
# Parameterization for reading dtype of DLTensor. Chosen to have
diff --git a/tests/python/relax/test_op_linear_algebra.py b/tests/python/relax/test_op_linear_algebra.py
index 66dd103926..8f44178b67 100644
--- a/tests/python/relax/test_op_linear_algebra.py
+++ b/tests/python/relax/test_op_linear_algebra.py
@@ -17,13 +17,13 @@
# ruff: noqa: F841
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import VDevice
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
def test_op_correctness():
diff --git a/tests/python/relax/test_op_manipulate.py b/tests/python/relax/test_op_manipulate.py
index d6748bad1b..b1f068acc5 100644
--- a/tests/python/relax/test_op_manipulate.py
+++ b/tests/python/relax/test_op_manipulate.py
@@ -17,13 +17,13 @@
# ruff: noqa: E731, F841
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import VDevice
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_op_misc.py b/tests/python/relax/test_op_misc.py
index afb44600ef..7006085344 100644
--- a/tests/python/relax/test_op_misc.py
+++ b/tests/python/relax/test_op_misc.py
@@ -17,8 +17,8 @@
# ruff: noqa: F841
import tvm_ffi
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax as rx
from tvm.relax import script as R
from tvm.s_tir import script as Ts
diff --git a/tests/python/relax/test_op_nn.py b/tests/python/relax/test_op_nn.py
index 1627282835..7c1d0ac1ef 100644
--- a/tests/python/relax/test_op_nn.py
+++ b/tests/python/relax/test_op_nn.py
@@ -17,13 +17,13 @@
# ruff: noqa: F841
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import VDevice
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
def test_op_correctness():
diff --git a/tests/python/relax/test_op_nn_convolution.py b/tests/python/relax/test_op_nn_convolution.py
index f6810b4b9b..c720701f88 100644
--- a/tests/python/relax/test_op_nn_convolution.py
+++ b/tests/python/relax/test_op_nn_convolution.py
@@ -16,13 +16,13 @@
# under the License.
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import VDevice
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
def test_conv1d_op_correctness():
@@ -917,7 +917,12 @@ def test_conv2d_infer_ty_shape_symbolic():
bb,
relax.op.nn.conv2d(x0, w0, strides=(2, 2), padding=(1, 1), dilation=(2, 2)),
relax.TensorType(
- (n, ko, tvm_tirx.tirx.floordiv(ih + 3, 2) + 1 - kh, tvm_tirx.tirx.floordiv(iw + 3, 2) + 1 - kw),
+ (
+ n,
+ ko,
+ tvm_tirx.tirx.floordiv(ih + 3, 2) + 1 - kh,
+ tvm_tirx.tirx.floordiv(iw + 3, 2) + 1 - kw,
+ ),
"float32",
),
)
diff --git a/tests/python/relax/test_op_nn_pooling.py b/tests/python/relax/test_op_nn_pooling.py
index a69867ec78..da1fbe2898 100644
--- a/tests/python/relax/test_op_nn_pooling.py
+++ b/tests/python/relax/test_op_nn_pooling.py
@@ -17,13 +17,13 @@
# ruff: noqa: E741, F841
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import VDevice
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
def test_op_correctness():
@@ -381,7 +381,9 @@ def test_max_pool2d_infer_ty_ceil_mode_symbolic():
relax.op.nn.max_pool2d(
x, pool_size=(3, 3), strides=(2, 2), padding=(1, 1), dilation=(2, 2), ceil_mode=True
),
- relax.TensorType((n, c, tvm_tirx.tirx.floordiv(ih, 2), tvm_tirx.tirx.floordiv(iw, 2)), "float32"),
+ relax.TensorType(
+ (n, c, tvm_tirx.tirx.floordiv(ih, 2), tvm_tirx.tirx.floordiv(iw, 2)), "float32"
+ ),
)
@@ -613,7 +615,13 @@ def test_max_pool3d_infer_ty_ceil_mode_symbolic():
ceil_mode=True,
),
relax.TensorType(
- (n, c, tvm_tirx.tirx.floordiv(id_, 2), tvm_tirx.tirx.floordiv(ih, 2), tvm_tirx.tirx.floordiv(iw, 2)),
+ (
+ n,
+ c,
+ tvm_tirx.tirx.floordiv(id_, 2),
+ tvm_tirx.tirx.floordiv(ih, 2),
+ tvm_tirx.tirx.floordiv(iw, 2),
+ ),
"float32",
),
)
@@ -1040,7 +1048,9 @@ def test_avg_pool2d_infer_ty_ceil_mode_symbolic():
relax.op.nn.avg_pool2d(
x, pool_size=(3, 3), strides=(2, 2), padding=(1, 1), dilation=(2, 2), ceil_mode=True
),
- relax.TensorType((n, c, tvm_tirx.tirx.floordiv(ih, 2), tvm_tirx.tirx.floordiv(iw, 2)), "float32"),
+ relax.TensorType(
+ (n, c, tvm_tirx.tirx.floordiv(ih, 2), tvm_tirx.tirx.floordiv(iw, 2)), "float32"
+ ),
)
diff --git a/tests/python/relax/test_op_qdq.py b/tests/python/relax/test_op_qdq.py
index bb40d1e653..b975c43b53 100644
--- a/tests/python/relax/test_op_qdq.py
+++ b/tests/python/relax/test_op_qdq.py
@@ -16,12 +16,12 @@
# under the License.
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
def test_op_correctness():
diff --git a/tests/python/relax/test_op_sampling.py b/tests/python/relax/test_op_sampling.py
index 2600f73b5d..8787364721 100644
--- a/tests/python/relax/test_op_sampling.py
+++ b/tests/python/relax/test_op_sampling.py
@@ -14,8 +14,8 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import script as R
diff --git a/tests/python/relax/test_op_search.py b/tests/python/relax/test_op_search.py
index 2787e58d0c..ccd2eaa9b3 100644
--- a/tests/python/relax/test_op_search.py
+++ b/tests/python/relax/test_op_search.py
@@ -18,13 +18,13 @@ from collections.abc import Callable
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import VDevice
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
def test_op_correctness():
diff --git a/tests/python/relax/test_op_set.py b/tests/python/relax/test_op_set.py
index 6eb5864266..03469aaccd 100644
--- a/tests/python/relax/test_op_set.py
+++ b/tests/python/relax/test_op_set.py
@@ -16,13 +16,13 @@
# under the License.
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import VDevice
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
def test_op_correctness():
diff --git a/tests/python/relax/test_op_size.py b/tests/python/relax/test_op_size.py
index b58c16cb3e..97b658df53 100644
--- a/tests/python/relax/test_op_size.py
+++ b/tests/python/relax/test_op_size.py
@@ -17,11 +17,11 @@
import numpy as np
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm_tirx.script import ir as I
def test_op_size():
diff --git a/tests/python/relax/test_op_sort.py b/tests/python/relax/test_op_sort.py
index c58dd9b1d6..ec8af7961f 100644
--- a/tests/python/relax/test_op_sort.py
+++ b/tests/python/relax/test_op_sort.py
@@ -16,13 +16,13 @@
# under the License.
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import VDevice
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
def test_op_correctness():
diff --git a/tests/python/relax/test_op_statistical.py b/tests/python/relax/test_op_statistical.py
index b2b28d1ff7..f20ed589fd 100644
--- a/tests/python/relax/test_op_statistical.py
+++ b/tests/python/relax/test_op_statistical.py
@@ -18,13 +18,13 @@ from collections.abc import Callable
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import VDevice
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
def test_op_correctness():
diff --git a/tests/python/relax/test_op_take.py b/tests/python/relax/test_op_take.py
index 6690d1ebb0..4da2c1c8f0 100644
--- a/tests/python/relax/test_op_take.py
+++ b/tests/python/relax/test_op_take.py
@@ -18,10 +18,11 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx.script import ir as I
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax import script as R
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
axis = tvm_tirx.testing.parameter(0, 1)
diff --git a/tests/python/relax/test_op_ternary.py b/tests/python/relax/test_op_ternary.py
index 02a4991679..a24ba4eea7 100644
--- a/tests/python/relax/test_op_ternary.py
+++ b/tests/python/relax/test_op_ternary.py
@@ -16,13 +16,13 @@
# under the License.
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import VDevice
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
def test_op_correctness():
diff --git a/tests/python/relax/test_op_unary.py b/tests/python/relax/test_op_unary.py
index 6bce7b83f9..f369d07df9 100644
--- a/tests/python/relax/test_op_unary.py
+++ b/tests/python/relax/test_op_unary.py
@@ -18,13 +18,13 @@ from collections.abc import Callable
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import Op
from tvm.relax import VDevice
from tvm.relax import script as R
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
def test_op_correctness():
diff --git a/tests/python/relax/test_op_view.py b/tests/python/relax/test_op_view.py
index 421348beda..359828177f 100644
--- a/tests/python/relax/test_op_view.py
+++ b/tests/python/relax/test_op_view.py
@@ -18,10 +18,11 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx.script import ir as I
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax import script as R
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_op_vision.py b/tests/python/relax/test_op_vision.py
index a2ca7e67e0..48ac48d10c 100644
--- a/tests/python/relax/test_op_vision.py
+++ b/tests/python/relax/test_op_vision.py
@@ -18,18 +18,19 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm_tirx.testing import env
pytest.importorskip("scipy") # tvm.topi.testing imports scipy
import tvm.topi.testing
-from tvm_tirx import tirx
from tvm import relax
-from tvm_tirx.ir import Op
-from tvm.relax.transform import LegalizeOps
from tvm.relax import script as R
+from tvm.relax.transform import LegalizeOps
+from tvm_tirx import tirx
+from tvm_tirx.ir import Op
def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_ty: relax.Type):
diff --git a/tests/python/relax/test_optimize_layout_transform.py b/tests/python/relax/test_optimize_layout_transform.py
index adbe71785d..735129f551 100644
--- a/tests/python/relax/test_optimize_layout_transform.py
+++ b/tests/python/relax/test_optimize_layout_transform.py
@@ -20,13 +20,14 @@
import numpy as np
import pytest
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir.base import assert_structural_equal
-from tvm.relax.transform import DeadCodeElimination, FuseTIR, OptimizeLayoutTransform
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm.relax.transform import DeadCodeElimination, FuseTIR, OptimizeLayoutTransform
from tvm.s_tir import script as Ts
+from tvm_tirx.ir.base import assert_structural_equal
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_pipeline.py b/tests/python/relax/test_pipeline.py
index f8cfe28a90..f54576a1b7 100644
--- a/tests/python/relax/test_pipeline.py
+++ b/tests/python/relax/test_pipeline.py
@@ -18,8 +18,9 @@ import numpy as np
import pytest
import tvm_ffi
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import script as R
from tvm_tirx.script import tirx as T
@@ -123,7 +124,9 @@ def test_pipeline_with_kv_cache():
kv, kv_cache = vm["main"](x, y, tvm_ffi.Shape(np_shape), kv_cache)
cache_np[i, :] = x_np + y_np
- tvm_tirx.testing.assert_allclose(kv.numpy(), cache_np[: np_shape[0], :], rtol=1e-7, atol=1e-7)
+ tvm_tirx.testing.assert_allclose(
+ kv.numpy(), cache_np[: np_shape[0], :], rtol=1e-7, atol=1e-7
+ )
@pytest.mark.parametrize("target_name", ["vulkan", "webgpu"])
@@ -212,7 +215,9 @@ def test_default_cuda_pipeline_schedules_power():
# Verify: At least one compute kernel has thread binding.
# (Using any() avoids false-fail from host-side shape helper kernels
# that may be generated by finalize_passes but don't need GPU binding)
- functions = [func for _, func in mod.functions.items() if isinstance(func, tvm_tirx.tirx.Function)]
+ functions = [
+ func for _, func in mod.functions.items() if isinstance(func, tvm_tirx.tirx.Function)
+ ]
assert functions, "expected at least one TIR Function after pipeline"
has_bound_kernel = any(_has_thread_binding(f) for f in functions)
diff --git a/tests/python/relax/test_pytorch_integration.py b/tests/python/relax/test_pytorch_integration.py
index 32d54036ae..64b65070c0 100644
--- a/tests/python/relax/test_pytorch_integration.py
+++ b/tests/python/relax/test_pytorch_integration.py
@@ -29,11 +29,11 @@ import pytest
import torch
import torch.nn.functional as F
-import tvm, tvm_tirx
+import tvm_tirx
from tvm.relax import BasePyModule
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
n_matmul = T.dynamic("n", "int32")
diff --git a/tests/python/relax/test_relax_operators.py b/tests/python/relax/test_relax_operators.py
index 0733e9ec85..9781a98a08 100644
--- a/tests/python/relax/test_relax_operators.py
+++ b/tests/python/relax/test_relax_operators.py
@@ -26,11 +26,12 @@ import numpy as np
import pytest
import tvm_ffi
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
m = T.dynamic("m", "int64")
@@ -424,7 +425,9 @@ def test_call_tir_canonical_roundtrip(name, tuple_result, result_kind, namespace
tvm_tirx.tirx.Var(f"p{i}", T.Tensor((3,), "float32"))
for i in range(1 + output_count - int(inplace))
]
- func = tvm_tirx.tirx.Function(params, [tvm_tirx.ir.Evaluate(0)], ret_type=tvm_tirx.ir.TupleType([]))
+ func = tvm_tirx.tirx.Function(
+ params, [tvm_tirx.ir.Evaluate(0)], ret_type=tvm_tirx.ir.TupleType([])
+ )
callee = relax.BlockBuilder().add_func(func, "callee")
x = relax.Var("x", tensor_type)
kwargs = {}
@@ -440,14 +443,18 @@ def test_call_tir_canonical_roundtrip(name, tuple_result, result_kind, namespace
assert call.args[0].same_as(callee)
assert call.args[1].fields[0].same_as(x)
tvm_tirx.ir.assert_structural_equal(call.ty_args[0], output_type)
- tvm_tirx.ir.assert_structural_equal(call.ty, output_type if result_type is None else result_type)
+ tvm_tirx.ir.assert_structural_equal(
+ call.ty, output_type if result_type is None else result_type
+ )
# An attrs object and semantic keywords construct the same canonical Call.
with_attrs = builder(
callee, call.args[1], ty_args=[output_type], attrs=call.attrs, ty=result_type, loc=loc
)
tvm_tirx.ir.assert_structural_equal(call, with_attrs)
# Roundtrip printing requires a type-consistent Call after constructor checks.
- call = tvm_tirx.ir.Call(call.op, call.args, attrs=call.attrs, ty_args=call.ty_args, loc=call.loc)
+ call = tvm_tirx.ir.Call(
+ call.op, call.args, attrs=call.attrs, ty_args=call.ty_args, loc=call.loc
+ )
printed = call.script()
assert printed.startswith(f"R.{name}(")
assert "ty_args=" in printed and "out_ty=" not in printed
@@ -493,9 +500,13 @@ def test_call_dps_packed_canonical_roundtrip(tuple_size, callee_kind, result_kin
assert call.args[0].same_as(callee)
else:
assert isinstance(call.args[0], relax.ExternFunc)
- tvm_tirx.ir.assert_structural_equal(call.ty, output_type if result_type is None else result_type)
+ tvm_tirx.ir.assert_structural_equal(
+ call.ty, output_type if result_type is None else result_type
+ )
# Roundtrip printing requires a type-consistent Call after constructor checks.
- call = tvm_tirx.ir.Call(call.op, call.args, attrs=call.attrs, ty_args=call.ty_args, loc=call.loc)
+ call = tvm_tirx.ir.Call(
+ call.op, call.args, attrs=call.attrs, ty_args=call.ty_args, loc=call.loc
+ )
printed = call.script()
assert printed.startswith("R.call_dps_packed(")
if callee_kind in ("string", "extern"):
@@ -555,9 +566,13 @@ def test_call_py_func_canonical_roundtrip(tuple_size, result_kind, builder):
assert call.args[0].value == "callback"
assert len(call.args[1].fields) == 1 and call.args[1].fields[0].same_as(x)
tvm_tirx.ir.assert_structural_equal(call.ty_args[0], output_type)
- tvm_tirx.ir.assert_structural_equal(call.ty, output_type if result_type is None else result_type)
+ tvm_tirx.ir.assert_structural_equal(
+ call.ty, output_type if result_type is None else result_type
+ )
# Roundtrip printing requires a type-consistent Call after constructor checks.
- call = tvm_tirx.ir.Call(call.op, call.args, attrs=call.attrs, ty_args=call.ty_args, loc=call.loc)
+ call = tvm_tirx.ir.Call(
+ call.op, call.args, attrs=call.attrs, ty_args=call.ty_args, loc=call.loc
+ )
printed = call.script()
assert printed.startswith("R.call_py_func(")
assert "ty_args=" in printed and "out_ty=" not in printed
diff --git a/tests/python/relax/test_relax_to_pyfunc_converter.py b/tests/python/relax/test_relax_to_pyfunc_converter.py
index 8dd49fa9a8..bc8c08d69e 100644
--- a/tests/python/relax/test_relax_to_pyfunc_converter.py
+++ b/tests/python/relax/test_relax_to_pyfunc_converter.py
@@ -26,12 +26,13 @@ import pytest
import torch
import torch.nn.functional as F
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm.relax import ExternFunc as _ExternFunc
-from tvm.relax.relax_to_pyfunc_converter import RelaxToPyFuncConverter
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm.relax.relax_to_pyfunc_converter import RelaxToPyFuncConverter
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
n_symbolic_add = T.dynamic("n", "int64")
diff --git a/tests/python/relax/test_remove_redundant_reshape.py b/tests/python/relax/test_remove_redundant_reshape.py
index f67fadb78e..086ce4fbb1 100644
--- a/tests/python/relax/test_remove_redundant_reshape.py
+++ b/tests/python/relax/test_remove_redundant_reshape.py
@@ -19,11 +19,12 @@
Test relax transform - Eliminate redundant reshape operations
"""
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx.testing
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.transform import DeadCodeElimination, RemoveRedundantReshape
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
def _run_pass_compare_output(Before, Expected):
diff --git a/tests/python/relax/test_runtime_builtin.py b/tests/python/relax/test_runtime_builtin.py
index cf1f4a3cb5..2f58b27741 100644
--- a/tests/python/relax/test_runtime_builtin.py
+++ b/tests/python/relax/test_runtime_builtin.py
@@ -19,11 +19,12 @@ import numpy as np
import pytest
import tvm_ffi
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.contrib import tvmjs
-from tvm_tirx.ir import assert_structural_equal
from tvm.relax.testing.runtime_builtin import MakeShapeCode, MatchShapeCode
+from tvm_tirx.ir import assert_structural_equal
from tvm_tirx.support import utils
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
index 8615f36982..b37b2c7a1d 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
@@ -28,8 +28,8 @@ import scipy.special
import tvm_ffi
from tvm_ffi import Shape
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax.frontend.nn.llm.kv_cache import (
AttnKind,
_attention_decode_cpu,
@@ -141,7 +141,9 @@ def set_global_func(head_dim, dtype, layer_sliding_window_size=1024):
fprepare_import = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_prepare_import")
fimport_page_group = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_import_page_group")
ffinish_import = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_finish_import")
- fget_sequence_length = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_get_sequence_length")
+ fget_sequence_length = tvm_tirx.get_global_func(
+ "vm.builtin.attention_kv_cache_get_sequence_length"
+ )
target = tvm_tirx.target.Target.from_device(device)
cache_key = (
@@ -522,7 +524,9 @@ def apply_attention(
queries_np = global_new_q[layer_id]
keys_np = global_new_k[layer_id]
values_np = global_new_v[layer_id]
- qkv = tvm_tirx.runtime.tensor(np.concatenate([queries_np, keys_np, values_np], axis=1), device)
+ qkv = tvm_tirx.runtime.tensor(
+ np.concatenate([queries_np, keys_np, values_np], axis=1), device
+ )
outputs = tvm_tirx.runtime.empty(queries_np.shape, dtype, device=device)
fattention_with_fuse_qkv(kv_cache, layer_id, sm_scale, qkv, outputs)
@@ -755,7 +759,9 @@ def test_attention_with_shared_kv(source_attn_kind, layer_window_size):
shared_q = rng.standard_normal((append_length, num_qo_heads, head_dim)).astype(dtype)
fbegin_forward(kv_cache, Shape([0]), Shape([append_length]), None)
- qkv = tvm_tirx.runtime.tensor(np.concatenate([source_q, current_k, current_v], axis=1), device)
+ qkv = tvm_tirx.runtime.tensor(
+ np.concatenate([source_q, current_k, current_v], axis=1), device
+ )
source_output = tvm_tirx.runtime.empty(source_q.shape, dtype, device=device)
fattention_with_fuse_qkv(kv_cache, 0, sm_scale, qkv, source_output)
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py
index 9430c45259..23f33e5095 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py
@@ -22,8 +22,8 @@ import torch
import tvm_ffi
from tvm_ffi import Shape
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm.relax.frontend.nn.llm.kv_cache import (
AttnKind,
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_metadata.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_metadata.py
index 0b6a84733a..0c4d4d625a 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_metadata.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_metadata.py
@@ -20,10 +20,10 @@ import pytest
import tvm_ffi
from tvm_ffi import Shape
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx.error import InternalError
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax.frontend.nn.llm.kv_cache import AttnKind, RopeMode
+from tvm_tirx.error import InternalError
reserved_nseq = 4
maximum_total_seq_length = 128
@@ -117,7 +117,9 @@ def test_checkpoint_metadata_reports_layout_pages_and_groups():
fprepare_import = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_prepare_import")
fimport_page_group = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_import_page_group")
ffinish_import = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_finish_import")
- fget_sequence_length = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_get_sequence_length")
+ fget_sequence_length = tvm_tirx.get_global_func(
+ "vm.builtin.attention_kv_cache_get_sequence_length"
+ )
kv_cache = create_kv_cache()
append_tokens(kv_cache)
@@ -217,7 +219,9 @@ def test_checkpoint_metadata_rejects_unsupported_layouts_and_sequence_ids():
rope_ext_cache = create_kv_cache(
rope_ext_factors=tvm_tirx.runtime.empty((head_dim // 2,), "float32", device=device)
)
- dst = tvm_tirx.runtime.empty((1, 1, 2, num_kv_heads, page_size, head_dim), "float16", device=device)
+ dst = tvm_tirx.runtime.empty(
+ (1, 1, 2, num_kv_heads, page_size, head_dim), "float16", device=device
+ )
with pytest.raises(InternalError, match="sliding-window"):
fget_layout_hash(sliding_cache)
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py
index 9c373cb96b..ccf530ad57 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py
@@ -23,8 +23,8 @@ import torch
import tvm_ffi
from tvm_ffi import Shape
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm.relax.frontend.nn.llm.kv_cache import (
AttnKind,
@@ -367,11 +367,15 @@ def apply_attention(
outputs1 = tvm_tirx.runtime.empty(
(total_seq_length, num_attention_heads, v_head_dim), dtype, device=device
)
- lse1 = tvm_tirx.runtime.empty((total_seq_length, num_attention_heads), "float32", device=device)
+ lse1 = tvm_tirx.runtime.empty(
+ (total_seq_length, num_attention_heads), "float32", device=device
+ )
outputs2 = tvm_tirx.runtime.empty(
(total_seq_length, num_attention_heads, kv_lora_rank), dtype, device=device
)
- lse2 = tvm_tirx.runtime.empty((total_seq_length, num_attention_heads), "float32", device=device)
+ lse2 = tvm_tirx.runtime.empty(
+ (total_seq_length, num_attention_heads), "float32", device=device
+ )
fappend_mla_kv(kv_cache, layer_id, key_value)
if not is_decode_request:
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py
index 64164b0cf5..ad21aa3ee4 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py
@@ -22,8 +22,8 @@ import torch
import tvm_ffi
from tvm_ffi import Shape
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax.frontend.nn.llm.kv_cache import (
AttnKind,
RopeMode,
@@ -297,11 +297,15 @@ def apply_attention(
outputs1 = tvm_tirx.runtime.empty(
(total_seq_length, num_attention_heads, v_head_dim), dtype, device=device
)
- lse1 = tvm_tirx.runtime.empty((total_seq_length, num_attention_heads), "float32", device=device)
+ lse1 = tvm_tirx.runtime.empty(
+ (total_seq_length, num_attention_heads), "float32", device=device
+ )
outputs2 = tvm_tirx.runtime.empty(
(total_seq_length, num_attention_heads, kv_lora_rank), dtype, device=device
)
- lse2 = tvm_tirx.runtime.empty((total_seq_length, num_attention_heads), "float32", device=device)
+ lse2 = tvm_tirx.runtime.empty(
+ (total_seq_length, num_attention_heads), "float32", device=device
+ )
fappend_mla_kv(kv_cache, layer_id, key_value)
if not is_decode_request:
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
index 6f9159103a..f898e074f8 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
@@ -26,8 +26,8 @@ import torch
import tvm_ffi
from tvm_ffi import Shape
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax.frontend.nn.llm.kv_cache import (
AttnKind,
RopeMode,
diff --git a/tests/python/relax/test_runtime_builtin_rnn_state.py b/tests/python/relax/test_runtime_builtin_rnn_state.py
index e7606cbccd..094aa00559 100644
--- a/tests/python/relax/test_runtime_builtin_rnn_state.py
+++ b/tests/python/relax/test_runtime_builtin_rnn_state.py
@@ -21,11 +21,11 @@ import numpy as np
import pytest
from tvm_ffi import Shape
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm_tirx
+import tvm_tirx.testing
from tvm.s_tir import dlight as dl
from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_runtime_sampling_flashinfer.py b/tests/python/relax/test_runtime_sampling_flashinfer.py
index 2ff7d55d3e..93fe857792 100644
--- a/tests/python/relax/test_runtime_sampling_flashinfer.py
+++ b/tests/python/relax/test_runtime_sampling_flashinfer.py
@@ -19,8 +19,8 @@ import random
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm_tirx.support import utils
@@ -49,7 +49,9 @@ def test_sampling():
# Probability tensor (each row sums to 1)
probs_np = np.array([[0.1, 0.2, 0.3, 0.2, 0.2] for _ in range(batch_size)], dtype="float32")
- target = tvm_tirx.testing.run_with_gpu_lock(lambda: tvm_tirx.target.Target.from_device(tvm_tirx.cuda()))
+ target = tvm_tirx.testing.run_with_gpu_lock(
+ lambda: tvm_tirx.target.Target.from_device(tvm_tirx.cuda())
+ )
sampling_mod = load_module(
"flashinfer_sampling",
relax.backend.cuda.flashinfer.gen_sampling_module(
diff --git a/tests/python/relax/test_tensorrt_engine_serialization.py b/tests/python/relax/test_tensorrt_engine_serialization.py
index ea22890740..dea08cefc6 100644
--- a/tests/python/relax/test_tensorrt_engine_serialization.py
+++ b/tests/python/relax/test_tensorrt_engine_serialization.py
@@ -26,12 +26,13 @@ import textwrap
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm.relax.backend.contrib.tensorrt import partition_for_tensorrt
from tvm.relax.dpl import is_op, wildcard
+from tvm_tirx import tirx
from tvm_tirx.support import cc
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_testing_nn.py b/tests/python/relax/test_testing_nn.py
index 42a81d6d59..f724deebab 100644
--- a/tests/python/relax/test_testing_nn.py
+++ b/tests/python/relax/test_testing_nn.py
@@ -15,12 +15,12 @@
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: RUF005
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.testing import nn
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_tir_call_source_kernel.py b/tests/python/relax/test_tir_call_source_kernel.py
index 900f5f1d5e..022f715cb0 100644
--- a/tests/python/relax/test_tir_call_source_kernel.py
+++ b/tests/python/relax/test_tir_call_source_kernel.py
@@ -18,12 +18,13 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_training_append_loss.py b/tests/python/relax/test_training_append_loss.py
index 03617986a5..5832111315 100644
--- a/tests/python/relax/test_training_append_loss.py
+++ b/tests/python/relax/test_training_append_loss.py
@@ -17,11 +17,11 @@
# ruff: noqa: E501, F841
import pytest
-import tvm_tirx.testing, tvm
-from tvm_tirx.ir.base import assert_structural_equal
+import tvm_tirx.testing
+from tvm.relax import script as R
from tvm.relax.training import AppendLoss
+from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
def test_simple():
diff --git a/tests/python/relax/test_training_loss.py b/tests/python/relax/test_training_loss.py
index 6c9c2af0c2..8030e3070d 100644
--- a/tests/python/relax/test_training_loss.py
+++ b/tests/python/relax/test_training_loss.py
@@ -14,11 +14,11 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-import tvm_tirx.testing, tvm
+import tvm_tirx.testing
from tvm import relax
+from tvm.relax import script as R
from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
@I.ir_module
diff --git a/tests/python/relax/test_training_optimizer.py b/tests/python/relax/test_training_optimizer.py
index 040d15ed25..b875b00611 100644
--- a/tests/python/relax/test_training_optimizer.py
+++ b/tests/python/relax/test_training_optimizer.py
@@ -18,11 +18,11 @@
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.training.optimizer import SGD, Adam, MomentumSGD
+from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.script.parser import relax as R
diff --git a/tests/python/relax/test_training_optimizer_numeric.py b/tests/python/relax/test_training_optimizer_numeric.py
index 2eeb7a39a4..b798be978f 100644
--- a/tests/python/relax/test_training_optimizer_numeric.py
+++ b/tests/python/relax/test_training_optimizer_numeric.py
@@ -22,12 +22,13 @@ import numpy as np
import pytest
import tvm_ffi
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm.relax.training.optimizer import SGD, Adam, MomentumSGD
from tvm.runtime.vm import VirtualMachine
+from tvm_tirx import IRModule
from tvm_tirx.script.parser import relax as R
from tvm_tirx.testing import assert_allclose
diff --git a/tests/python/relax/test_training_setup_trainer.py b/tests/python/relax/test_training_setup_trainer.py
index c655c581fb..312af9fd6c 100644
--- a/tests/python/relax/test_training_setup_trainer.py
+++ b/tests/python/relax/test_training_setup_trainer.py
@@ -17,15 +17,15 @@
# ruff: noqa: E501
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir.base import assert_structural_equal
+from tvm.relax import script as R
from tvm.relax.training import SetupTrainer
from tvm.relax.training.loss import MSELoss
from tvm.relax.training.optimizer import SGD, MomentumSGD
+from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
def test_simple():
diff --git a/tests/python/relax/test_training_trainer_numeric.py b/tests/python/relax/test_training_trainer_numeric.py
index 72c72b4c0f..b62d8d20aa 100644
--- a/tests/python/relax/test_training_trainer_numeric.py
+++ b/tests/python/relax/test_training_trainer_numeric.py
@@ -17,14 +17,15 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.training import SetupTrainer, Trainer
from tvm.relax.training.loss import MSELoss
from tvm.relax.training.optimizer import SGD, Adam
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
def _get_backbone():
diff --git a/tests/python/relax/test_transform.py b/tests/python/relax/test_transform.py
index c90559008e..66c8ec4738 100644
--- a/tests/python/relax/test_transform.py
+++ b/tests/python/relax/test_transform.py
@@ -18,15 +18,15 @@
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import TupleType as _TupleType
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.ir import TupleType as _TupleType
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
@@ -382,7 +382,9 @@ def test_call_dps_packed_rewrite_nested_tuple_output():
relax.TensorType((4,), "int32"),
relax.TensorType((5, 6), "float16"),
]
- output_ty = tvm_tirx.ir.TupleType([flat_output_types[0], tvm_tirx.ir.TupleType(flat_output_types[1:])])
+ output_ty = tvm_tirx.ir.TupleType(
+ [flat_output_types[0], tvm_tirx.ir.TupleType(flat_output_types[1:])]
+ )
x = relax.Var("x", input_ty)
call = relax.Call(
diff --git a/tests/python/relax/test_transform_allocate_workspace.py b/tests/python/relax/test_transform_allocate_workspace.py
index 07ca383065..61c513f9fe 100644
--- a/tests/python/relax/test_transform_allocate_workspace.py
+++ b/tests/python/relax/test_transform_allocate_workspace.py
@@ -14,11 +14,11 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm_tirx.script import ir as I
@I.ir_module
diff --git a/tests/python/relax/test_transform_alter_op_impl.py b/tests/python/relax/test_transform_alter_op_impl.py
index 02139adf75..e79e5dd55d 100644
--- a/tests/python/relax/test_transform_alter_op_impl.py
+++ b/tests/python/relax/test_transform_alter_op_impl.py
@@ -16,12 +16,12 @@
# under the License.
# ruff: noqa: E501, E731, F841
-import tvm_tirx.testing, tvm
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import TupleType as _TupleType
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.ir import TupleType as _TupleType
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
kOperatorName = "operator_name"
diff --git a/tests/python/relax/test_transform_annotate_tir_op_pattern.py b/tests/python/relax/test_transform_annotate_tir_op_pattern.py
index 497a555ad9..84b868724e 100644
--- a/tests/python/relax/test_transform_annotate_tir_op_pattern.py
+++ b/tests/python/relax/test_transform_annotate_tir_op_pattern.py
@@ -20,9 +20,9 @@ import enum
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_attach_attr_layout_free_tensors.py b/tests/python/relax/test_transform_attach_attr_layout_free_tensors.py
index 625e187325..4dedf16c74 100644
--- a/tests/python/relax/test_transform_attach_attr_layout_free_tensors.py
+++ b/tests/python/relax/test_transform_attach_attr_layout_free_tensors.py
@@ -17,14 +17,15 @@
# ruff: noqa: F401
import numpy as np
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx.testing
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.transform import CombineParallelMatmul
-from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/test_transform_attach_global_symbol.py b/tests/python/relax/test_transform_attach_global_symbol.py
index 3542c98910..a4cb7eecf9 100644
--- a/tests/python/relax/test_transform_attach_global_symbol.py
+++ b/tests/python/relax/test_transform_attach_global_symbol.py
@@ -18,14 +18,15 @@
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx
+import tvm_tirx.script
from tvm import relax
-from tvm_tirx.ir import assert_structural_equal
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
+from tvm_tirx.ir import assert_structural_equal
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_bind_params.py b/tests/python/relax/test_transform_bind_params.py
index 35af7e1218..fa6cebf6d3 100644
--- a/tests/python/relax/test_transform_bind_params.py
+++ b/tests/python/relax/test_transform_bind_params.py
@@ -18,9 +18,10 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import ExternFunc as _ExternFunc
from tvm.relax import script as R
@@ -133,7 +134,9 @@ def test_bind_params_consistent_explicit_and_inferred_symbolic_var():
n = relax.Var("n", prim_type)
data = relax.Var("data", relax.TensorType([n], "float32"))
before = relax.Function([n, data], n, prim_type)
- explicit_value = tvm_tirx.tirx.Add(tvm_tirx.tirx.const(2, "int64"), tvm_tirx.tirx.const(3, "int64"))
+ explicit_value = tvm_tirx.tirx.Add(
+ tvm_tirx.tirx.const(2, "int64"), tvm_tirx.tirx.const(3, "int64")
+ )
after = before.bind_params({n: explicit_value, data: np.zeros([5], "float32")})
diff --git a/tests/python/relax/test_transform_bind_symbolic_vars.py b/tests/python/relax/test_transform_bind_symbolic_vars.py
index c4ba978eb9..8f50915428 100644
--- a/tests/python/relax/test_transform_bind_symbolic_vars.py
+++ b/tests/python/relax/test_transform_bind_symbolic_vars.py
@@ -17,13 +17,13 @@
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_bundle_model_params.py b/tests/python/relax/test_transform_bundle_model_params.py
index b8e74bce66..6fee509a4c 100644
--- a/tests/python/relax/test_transform_bundle_model_params.py
+++ b/tests/python/relax/test_transform_bundle_model_params.py
@@ -22,12 +22,13 @@ import pytest
pytest.importorskip("scipy") # tvm.topi.testing imports scipy
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
import tvm.topi.testing
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_canonicalize_bindings.py b/tests/python/relax/test_transform_canonicalize_bindings.py
index fe4100f461..ea8883d5e8 100644
--- a/tests/python/relax/test_transform_canonicalize_bindings.py
+++ b/tests/python/relax/test_transform_canonicalize_bindings.py
@@ -18,16 +18,16 @@
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
+from tvm.relax import script as R
+from tvm.relax.transform.transform import CanonicalizeBindings
+from tvm.s_tir import script as Ts
from tvm_tirx.ir import TupleType as _TupleType
from tvm_tirx.ir.base import assert_structural_equal
-from tvm.relax.transform.transform import CanonicalizeBindings
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
-from tvm.s_tir import script as Ts
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_codegen_pass.py b/tests/python/relax/test_transform_codegen_pass.py
index dcfc56e323..0a033139f8 100644
--- a/tests/python/relax/test_transform_codegen_pass.py
+++ b/tests/python/relax/test_transform_codegen_pass.py
@@ -22,16 +22,17 @@ import tempfile
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax, s_tir
from tvm.relax import ExternFunc as _ExternFunc
+from tvm.relax import script as R
from tvm.relax.dpl import is_op, wildcard
from tvm.relax.testing import transform
-from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
from tvm_tirx.support import utils
from tvm_tirx.testing import env
diff --git a/tests/python/relax/test_transform_combine_parallel_matmul.py b/tests/python/relax/test_transform_combine_parallel_matmul.py
index 1aa8767bce..4f5a22a009 100644
--- a/tests/python/relax/test_transform_combine_parallel_matmul.py
+++ b/tests/python/relax/test_transform_combine_parallel_matmul.py
@@ -17,12 +17,13 @@
# ruff: noqa: E731, F401
import pytest
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx.testing
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.script import ir_builder as relax_builder
from tvm.relax.transform import CombineParallelMatmul
-from tvm.relax import script as R
+from tvm_tirx import tirx
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/test_transform_compute_prim_value.py b/tests/python/relax/test_transform_compute_prim_value.py
index 9ec51abd9a..33ddf91feb 100644
--- a/tests/python/relax/test_transform_compute_prim_value.py
+++ b/tests/python/relax/test_transform_compute_prim_value.py
@@ -15,11 +15,12 @@
# specific language governing permissions and limitations
# under the License.
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx.script import ir as I
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_convert_dataflow.py b/tests/python/relax/test_transform_convert_dataflow.py
index 39e8f1587d..585d0dc93e 100644
--- a/tests/python/relax/test_transform_convert_dataflow.py
+++ b/tests/python/relax/test_transform_convert_dataflow.py
@@ -15,11 +15,11 @@
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: F841
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm_tirx.script import ir as I
# functions that will not change
diff --git a/tests/python/relax/test_transform_convert_layout.py b/tests/python/relax/test_transform_convert_layout.py
index ff321dfb1c..5aa5cfc40a 100644
--- a/tests/python/relax/test_transform_convert_layout.py
+++ b/tests/python/relax/test_transform_convert_layout.py
@@ -16,8 +16,9 @@
# under the License.
# ruff: noqa: E501, F841
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax.transform import ConvertLayout, Normalize
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
diff --git a/tests/python/relax/test_transform_cse.py b/tests/python/relax/test_transform_cse.py
index c776d11a60..54a8ae5cc6 100644
--- a/tests/python/relax/test_transform_cse.py
+++ b/tests/python/relax/test_transform_cse.py
@@ -19,8 +19,9 @@
import numpy as np
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax.transform import EliminateCommonSubexpr
from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
diff --git a/tests/python/relax/test_transform_dead_code_elimination.py b/tests/python/relax/test_transform_dead_code_elimination.py
index 1a1a899051..983a235755 100644
--- a/tests/python/relax/test_transform_dead_code_elimination.py
+++ b/tests/python/relax/test_transform_dead_code_elimination.py
@@ -18,8 +18,9 @@
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax.transform import DeadCodeElimination
from tvm.s_tir import script as Ts
from tvm_tirx.script.parser import ir as I
diff --git a/tests/python/relax/test_transform_decompose_ops.py b/tests/python/relax/test_transform_decompose_ops.py
index 74bf4d82b3..9d872554ad 100644
--- a/tests/python/relax/test_transform_decompose_ops.py
+++ b/tests/python/relax/test_transform_decompose_ops.py
@@ -18,14 +18,15 @@
from typing import Union
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
-from tvm_tirx import IRModule
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
from tvm.relax import Function
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm_tirx import IRModule
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_error_enrichment.py b/tests/python/relax/test_transform_error_enrichment.py
index f469386f9f..3606823b34 100644
--- a/tests/python/relax/test_transform_error_enrichment.py
+++ b/tests/python/relax/test_transform_error_enrichment.py
@@ -23,8 +23,8 @@ pass name plus the offending location rendered as underlined TVMScript.
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm_tirx.ir import IRModule
diff --git a/tests/python/relax/test_transform_expand_matmul_of_sum.py b/tests/python/relax/test_transform_expand_matmul_of_sum.py
index 23c0f74008..69a1239270 100644
--- a/tests/python/relax/test_transform_expand_matmul_of_sum.py
+++ b/tests/python/relax/test_transform_expand_matmul_of_sum.py
@@ -19,10 +19,10 @@ import inspect
import pytest
-import tvm_tirx.testing, tvm
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm_tirx.script import ir as I
class Base:
diff --git a/tests/python/relax/test_transform_expand_tuple_args.py b/tests/python/relax/test_transform_expand_tuple_args.py
index a8dde501e4..a3f3cb5ac5 100644
--- a/tests/python/relax/test_transform_expand_tuple_args.py
+++ b/tests/python/relax/test_transform_expand_tuple_args.py
@@ -16,10 +16,11 @@
# under the License.
# ruff: noqa: F401
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx.script import ir as I
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax import script as R
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_fold_batch_norm_to_conv2d.py b/tests/python/relax/test_transform_fold_batch_norm_to_conv2d.py
index 6556e96d0a..e958b38470 100644
--- a/tests/python/relax/test_transform_fold_batch_norm_to_conv2d.py
+++ b/tests/python/relax/test_transform_fold_batch_norm_to_conv2d.py
@@ -19,14 +19,15 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir.module import IRModule
+from tvm.relax import script as R
from tvm.relax.expr_functor import PyExprVisitor, visitor
from tvm.relax.script import ir_builder as relax_builder
+from tvm_tirx.ir.module import IRModule
from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm_tirx.script.ir_builder import IRBuilder
diff --git a/tests/python/relax/test_transform_fold_constant.py b/tests/python/relax/test_transform_fold_constant.py
index 7a5c60e822..e706c07300 100644
--- a/tests/python/relax/test_transform_fold_constant.py
+++ b/tests/python/relax/test_transform_fold_constant.py
@@ -17,14 +17,15 @@
# ruff: noqa: F841
import numpy as np
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import TupleType as _TupleType
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.ir import TupleType as _TupleType
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_fuse_ops.py b/tests/python/relax/test_transform_fuse_ops.py
index 7224458f3b..b5d4518849 100644
--- a/tests/python/relax/test_transform_fuse_ops.py
+++ b/tests/python/relax/test_transform_fuse_ops.py
@@ -18,13 +18,14 @@
from __future__ import annotations
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax, topi
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
@@ -1717,7 +1718,9 @@ def test_primitive_call_arg_not_inlined():
def main(x: R.Tensor((4,), dtype="int64")):
cls = Before
with R.dataflow():
- value: T.int64 = R.call_pure_packed("get_scalar", ty_args=tvm_tirx.ir.PrimType("int64"))
+ value: T.int64 = R.call_pure_packed(
+ "get_scalar", ty_args=tvm_tirx.ir.PrimType("int64")
+ )
lv = R.call_tir(
cls.add_scalar,
(x, value),
diff --git a/tests/python/relax/test_transform_fuse_ops_by_pattern.py b/tests/python/relax/test_transform_fuse_ops_by_pattern.py
index e98e2681ba..de993811fa 100644
--- a/tests/python/relax/test_transform_fuse_ops_by_pattern.py
+++ b/tests/python/relax/test_transform_fuse_ops_by_pattern.py
@@ -18,8 +18,10 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
+import tvm
+import tvm_tirx
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.backend.cuda.cublas import partition_for_cublas
from tvm.relax.backend.cuda.cutlass import partition_for_cutlass
from tvm.relax.dpl.pattern import (
@@ -30,9 +32,8 @@ from tvm.relax.dpl.pattern import (
wildcard,
)
from tvm.relax.transform import PatternCheckContext
-from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_fuse_tir.py b/tests/python/relax/test_transform_fuse_tir.py
index 83477501d3..3efc7f9efa 100644
--- a/tests/python/relax/test_transform_fuse_tir.py
+++ b/tests/python/relax/test_transform_fuse_tir.py
@@ -17,13 +17,14 @@
from __future__ import annotations
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax, topi
from tvm.relax import ExternFunc as _ExternFunc
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_fuse_transpose_matmul.py b/tests/python/relax/test_transform_fuse_transpose_matmul.py
index 031260d3aa..1995edec66 100644
--- a/tests/python/relax/test_transform_fuse_transpose_matmul.py
+++ b/tests/python/relax/test_transform_fuse_transpose_matmul.py
@@ -18,12 +18,12 @@
import numpy as np
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_gradient.py b/tests/python/relax/test_transform_gradient.py
index 960ac3020f..7cdbcd86df 100644
--- a/tests/python/relax/test_transform_gradient.py
+++ b/tests/python/relax/test_transform_gradient.py
@@ -18,11 +18,11 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir.base import assert_structural_equal
from tvm.s_tir import script as Ts
+from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_transform_gradient_checkpoint.py b/tests/python/relax/test_transform_gradient_checkpoint.py
index b62c1eb35c..9ef48d0593 100644
--- a/tests/python/relax/test_transform_gradient_checkpoint.py
+++ b/tests/python/relax/test_transform_gradient_checkpoint.py
@@ -17,12 +17,12 @@
# ruff: noqa: E501, F811
"""Unit tests for gradient with checkpointing."""
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.block_builder import BlockBuilder
from tvm.relax.testing import nn
+from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
diff --git a/tests/python/relax/test_transform_gradient_numeric.py b/tests/python/relax/test_transform_gradient_numeric.py
index 9ce0613259..3127e6e08a 100644
--- a/tests/python/relax/test_transform_gradient_numeric.py
+++ b/tests/python/relax/test_transform_gradient_numeric.py
@@ -18,8 +18,9 @@
import numpy as np
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
from tvm.relax.testing.gradient import check_numerical_grads
from tvm_tirx.script.parser import ir as I
diff --git a/tests/python/relax/test_transform_gradient_te_register.py b/tests/python/relax/test_transform_gradient_te_register.py
index c2966ba664..7c74a5f577 100644
--- a/tests/python/relax/test_transform_gradient_te_register.py
+++ b/tests/python/relax/test_transform_gradient_te_register.py
@@ -19,15 +19,16 @@
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm_tirx import tirx
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir import TupleType as _TupleType
-from tvm_tirx.ir.base import assert_structural_equal
from tvm.relax.training.utils import register_te_gradient
from tvm.relax.transform import Gradient
from tvm.s_tir import script as Ts
+from tvm_tirx import tirx
+from tvm_tirx.ir import TupleType as _TupleType
+from tvm_tirx.ir.base import assert_structural_equal
from tvm_tirx.script.parser import ir as I
from tvm_tirx.script.parser import relax as R
from tvm_tirx.script.parser import tirx as T
diff --git a/tests/python/relax/test_transform_inline_private_functions.py b/tests/python/relax/test_transform_inline_private_functions.py
index 78378d5eb4..1c825b4e27 100644
--- a/tests/python/relax/test_transform_inline_private_functions.py
+++ b/tests/python/relax/test_transform_inline_private_functions.py
@@ -18,12 +18,13 @@
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_ipc_allreduce_rewrite.py b/tests/python/relax/test_transform_ipc_allreduce_rewrite.py
index 8370cc592a..4a071b2ee9 100644
--- a/tests/python/relax/test_transform_ipc_allreduce_rewrite.py
+++ b/tests/python/relax/test_transform_ipc_allreduce_rewrite.py
@@ -15,11 +15,11 @@
# specific language governing permissions and limitations
# under the License.
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
@@ -68,7 +68,9 @@ def test_ipc_allreduce_rewrite():
mod,
(
Expected
- if tvm_tirx.get_global_func("runtime.disco.cuda_ipc.custom_allreduce", allow_missing=True)
+ if tvm_tirx.get_global_func(
+ "runtime.disco.cuda_ipc.custom_allreduce", allow_missing=True
+ )
is not None
else Module
),
@@ -124,7 +126,9 @@ def test_ipc_allreduce_spread_along_reshape():
mod,
(
Expected
- if tvm_tirx.get_global_func("runtime.disco.cuda_ipc.custom_allreduce", allow_missing=True)
+ if tvm_tirx.get_global_func(
+ "runtime.disco.cuda_ipc.custom_allreduce", allow_missing=True
+ )
is not None
else Module
),
diff --git a/tests/python/relax/test_transform_lambda_lift.py b/tests/python/relax/test_transform_lambda_lift.py
index fb0cc2bc8f..5e5538786f 100644
--- a/tests/python/relax/test_transform_lambda_lift.py
+++ b/tests/python/relax/test_transform_lambda_lift.py
@@ -18,15 +18,15 @@
import tvm_ffi
-import tvm, tvm_tirx
-import tvm_tirx.script, tvm
-import tvm_tirx.testing, tvm
+import tvm_tirx
+import tvm_tirx.script
+import tvm_tirx.testing
from tvm import relax
-from tvm_tirx.ir.base import assert_structural_equal
-from tvm.relax import transform
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm.relax import transform
from tvm.s_tir import script as Ts
+from tvm_tirx.ir.base import assert_structural_equal
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_lazy_transform_params.py b/tests/python/relax/test_transform_lazy_transform_params.py
index cde91a1c43..6f2ebde888 100644
--- a/tests/python/relax/test_transform_lazy_transform_params.py
+++ b/tests/python/relax/test_transform_lazy_transform_params.py
@@ -17,13 +17,14 @@
# ruff: noqa: F841
import numpy as np
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
-from tvm.relax.transform import LazyTransformParams
-from tvm_tirx.script import ir as I
from tvm.relax import script as R
+from tvm.relax.transform import LazyTransformParams
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_legalize_ops.py b/tests/python/relax/test_transform_legalize_ops.py
index 5303ba4eb5..12e0a50c65 100644
--- a/tests/python/relax/test_transform_legalize_ops.py
+++ b/tests/python/relax/test_transform_legalize_ops.py
@@ -20,14 +20,15 @@ from __future__ import annotations
import pytest
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
+import tvm
+import tvm_tirx
+import tvm_tirx.testing
from tvm import relax
+from tvm.relax import script as R
from tvm.relax.transform import LegalizeOps
from tvm.relax.transform.legalize_ops.common import register_legalize
-from tvm_tirx.script import ir as I
-from tvm.relax import script as R
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
diff --git a/tests/python/relax/test_transform_legalize_ops_binary.py b/tests/python/relax/test_transform_legalize_ops_binary.py
index bef48180e7..845af56647 100644
--- a/tests/python/relax/test_transform_legalize_ops_binary.py
+++ b/tests/python/relax/test_transform_legalize_ops_binary.py
@@ -16,12 +16,12 @@
# under the License.
# ruff: noqa: E501
-import tvm, tvm_tirx
-import tvm_tirx.testing, tvm
-from tvm.relax.transform import LegalizeOps
-from tvm_tirx.script import ir as I
+import tvm_tirx
+import tvm_tirx.testing
from tvm.relax import script as R
+from tvm.relax.transform import LegalizeOps
from tvm.s_tir import script as Ts
+from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
##################### Binary arithmetic #####################
diff --git a/tests/python/relax/test_transform_legalize_ops_ccl.py b/tests/python/relax/test_transform_legalize_ops_ccl.py
index 45aeefcf25..ff38e8116f 100644
--- a/tests/python/relax/test_transform_legalize_ops_ccl.py
+++ b/tests/python/relax/test_transform_legalize_ops_ccl.py
@@ -16,13 +16,13 @@
# under the License.
# ruff: noqa: E501, F841
-import tvm, tvm_tirx
... 7379 lines suppressed ...
(tvm) 13/33: [FIX][CI] Build both native partitions with the pinned FFI dependency
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit fc53ebd4b7cadb983d7ee52776bf4a8635fdfe04
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 02:10:47 2026 +0000
[FIX][CI] Build both native partitions with the pinned FFI dependency
---
ci/jenkins/data.py | 4 ++-
tests/python/ci/test_jenkins_helpers.py | 12 ++++++++-
tests/scripts/task_build.py | 44 +++++++++++++++++++++++++++++----
tests/scripts/task_cpp_unittest.sh | 6 +++++
tests/scripts/task_python_unittest.sh | 4 +--
5 files changed, 61 insertions(+), 9 deletions(-)
diff --git a/ci/jenkins/data.py b/ci/jenkins/data.py
index 74f7ccc96d..d6ddd98990 100644
--- a/ci/jenkins/data.py
+++ b/ci/jenkins/data.py
@@ -29,6 +29,7 @@ data-driven artifact lists.
"""
import sys
+from glob import glob
files_to_stash = {
# Executables and build files needed to run c++ tests
@@ -45,7 +46,8 @@ files_to_stash = {
"build/lib/libtvm_compiler.so",
"build/lib/libtvm_tirx.so",
"build/lib/libtvm_runtime.so",
- "build/lib/libtvm_ffi.so",
+ "python/tvm_ffi",
+ *glob("python/apache_tvm_ffi-*.dist-info"),
"build/lib/libtvm_runtime_cuda.so",
"build/lib/libtvm_runtime_vulkan.so",
"build/lib/libtvm_runtime_opencl.so",
diff --git a/tests/python/ci/test_jenkins_helpers.py b/tests/python/ci/test_jenkins_helpers.py
index cfc0c69ced..ef3a0a3445 100644
--- a/tests/python/ci/test_jenkins_helpers.py
+++ b/tests/python/ci/test_jenkins_helpers.py
@@ -25,7 +25,7 @@ import sys
import pytest
-from .test_utils import JENKINS_SCRIPT_ROOT, REPO_ROOT
+from .test_utils import JENKINS_SCRIPT_ROOT, REPO_ROOT, TempGit
def load_s3(script_root=JENKINS_SCRIPT_ROOT):
@@ -93,6 +93,16 @@ def test_build_environment_and_target(tmp_path):
script.parent.mkdir(parents=True)
shutil.copy(REPO_ROOT / "tests/scripts/task_build.py", script)
(tmp_path / "build").mkdir()
+ ffi_path = tmp_path / "3rdparty/tvm-ffi"
+ ffi_path.mkdir(parents=True)
+ ffi_repo = TempGit(ffi_path)
+ ffi_repo.run("config", "user.name", "CI Test")
+ ffi_repo.run("config", "user.email", "ci-test@example.com")
+ ffi_repo.run("commit", "--allow-empty", "-m", "FFI fixture")
+ ffi_commit = ffi_repo.run("rev-parse", "HEAD", stdout=subprocess.PIPE).stdout.strip()
+ ffi_marker = tmp_path / "python/tvm_ffi/.source-commit"
+ ffi_marker.parent.mkdir(parents=True)
+ ffi_marker.write_text(f"{ffi_commit} {sys.implementation.cache_tag}\n")
executable(
tmp_path / "cmake",
"""
diff --git a/tests/scripts/task_build.py b/tests/scripts/task_build.py
index c240ec51c0..3160870aaf 100755
--- a/tests/scripts/task_build.py
+++ b/tests/scripts/task_build.py
@@ -21,6 +21,7 @@ import multiprocessing
import os
import shutil
import subprocess
+import sys
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parent.parent.parent
@@ -34,12 +35,19 @@ if __name__ == "__main__":
parser.add_argument("--sccache-bucket", required=False, help="sccache bucket name")
parser.add_argument("--sccache-region", required=False, help="sccache region")
parser.add_argument("--build-dir", default="build", help="build folder")
- parser.add_argument("--cmake-target", help="optional build target")
+ parser.add_argument(
+ "--cmake-target", action="append", default=[], help="build target (repeatable)"
+ )
+ parser.add_argument(
+ "--ffi-only", action="store_true", help="prepare the source-pinned FFI package"
+ )
parser.add_argument("--debug", required=False, action="store_true", help="build in debug mode")
args = parser.parse_args()
sccache_exe = shutil.which("sccache")
- if args.cmake_target in ["standalone_crt", "crttest"]:
+ if args.cmake_target and all(
+ target in {"standalone_crt", "crttest"} for target in args.cmake_target
+ ):
logging.info("Skipping standalone_crt build")
exit(0)
@@ -47,6 +55,13 @@ if __name__ == "__main__":
build_dir = Path(os.getcwd()) / args.build_dir
build_dir = build_dir.relative_to(REPO_ROOT)
env = os.environ.copy()
+ env["PYTHONPATH"] = os.pathsep.join(
+ [
+ str(REPO_ROOT / "tirx" / "python"),
+ str(REPO_ROOT / "python"),
+ os.environ.get("PYTHONPATH", ""),
+ ]
+ )
if use_sccache:
if args.sccache_bucket and "AWS_ACCESS_KEY_ID" in os.environ:
@@ -68,6 +83,27 @@ if __name__ == "__main__":
reason = "<unknown>"
logging.info(f"Not using sccache, reason: {reason}")
+ # Build from this checkout's pinned submodule once, then transfer the exact
+ # Python/native package together with the TVM libraries to test workers.
+ ffi_commit = subprocess.check_output(
+ ["git", "-C", str(REPO_ROOT / "3rdparty" / "tvm-ffi"), "rev-parse", "HEAD"], text=True
+ ).strip()
+ ffi_identity = f"{ffi_commit} {sys.implementation.cache_tag}"
+ ffi_marker = REPO_ROOT / "python" / "tvm_ffi" / ".source-commit"
+ if not ffi_marker.is_file() or ffi_marker.read_text().strip() != ffi_identity:
+ subprocess.run(
+ [
+ "uv", "pip", "install", "--python", sys.executable,
+ "-v", "--no-deps", "--reinstall", "--target=python", "./3rdparty/tvm-ffi/",
+ ],
+ cwd=REPO_ROOT,
+ env=env,
+ check=True,
+ )
+ ffi_marker.write_text(ffi_identity + "\n")
+ if args.ffi_only:
+ sys.exit(0)
+
if use_sccache:
subprocess.run(["sccache", "--start-server"], env=env, check=False)
logging.info("===== sccache stats =====")
@@ -89,9 +125,7 @@ if __name__ == "__main__":
check=True,
)
- ninja_args = [f"-j{num_cpus}"]
- if args.cmake_target:
- ninja_args.insert(0, args.cmake_target)
+ ninja_args = [*args.cmake_target, f"-j{num_cpus}"]
if os.environ.get("VERBOSE", "true").lower() in {"1", "true", "yes"}:
ninja_args.append("-v")
subprocess.run(["cmake", "--build", ".", "--", *ninja_args], cwd=build_dir, env=env, check=True)
diff --git a/tests/scripts/task_cpp_unittest.sh b/tests/scripts/task_cpp_unittest.sh
index 424642673b..efe6417733 100755
--- a/tests/scripts/task_cpp_unittest.sh
+++ b/tests/scripts/task_cpp_unittest.sh
@@ -32,6 +32,12 @@ fi
export TVM_BIND_THREADS=0
export OMP_NUM_THREADS=1
+# Test workers receive the exact FFI package in their own checkout. Native
+# binaries must find that copy even when their build workspace had another path.
+ffi_lib_dir="$(pwd)/python/tvm_ffi/lib"
+export LD_LIBRARY_PATH="${ffi_lib_dir}:${LD_LIBRARY_PATH:-}"
+export DYLD_LIBRARY_PATH="${ffi_lib_dir}:${DYLD_LIBRARY_PATH:-}"
+
pushd "${BUILD_DIR}"
# run cpp test executable
./cpptest
diff --git a/tests/scripts/task_python_unittest.sh b/tests/scripts/task_python_unittest.sh
index c6c578cc5b..85b5d692f0 100755
--- a/tests/scripts/task_python_unittest.sh
+++ b/tests/scripts/task_python_unittest.sh
@@ -21,8 +21,8 @@ set -euxo pipefail
export PYTHONPATH="$(pwd)/tirx/python:$(pwd)/python"
export PYTEST_ADDOPTS="${CI_PYTEST_ADD_OPTIONS:-} ${PYTEST_ADDOPTS:-}"
-# setup tvm-ffi into python folder
-uv pip install -v --target=python ./3rdparty/tvm-ffi/
+# Reuse the exact source-pinned FFI package transferred by the native build.
+python3 tests/scripts/task_build.py --ffi-only
# LLVMModule JIT execution. Installed here as well as in the CI image because a docker/
# change does not reach the test containers in the same run (they use the pinned tag).
(tvm) 19/33: [FIX][CUTLASS] Propagate owning runtime headers to the vendor binding
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit ea1c75476defc94a57c1cbd9e17bb692ca569b3c
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 04:26:10 2026 +0000
[FIX][CUTLASS] Propagate owning runtime headers to the vendor binding
---
cmake/modules/contrib/CUTLASS.cmake | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/cmake/modules/contrib/CUTLASS.cmake b/cmake/modules/contrib/CUTLASS.cmake
index 6d81ba9238..ac0958eb96 100644
--- a/cmake/modules/contrib/CUTLASS.cmake
+++ b/cmake/modules/contrib/CUTLASS.cmake
@@ -30,7 +30,7 @@ if(USE_CUDA AND USE_CUTLASS)
${PROJECT_SOURCE_DIR}/3rdparty/cutlass_fpA_intB_gemm
${PROJECT_SOURCE_DIR}/3rdparty/cutlass_fpA_intB_gemm/cutlass/include
)
- target_link_libraries(fpA_intB_gemm_tvm PRIVATE tvm_ffi_header)
+ target_link_libraries(fpA_intB_gemm_tvm PRIVATE tvm_runtime_extra_defs)
set(CUTLASS_FPA_INTB_RUNTIME_SRCS "")
list(APPEND CUTLASS_FPA_INTB_RUNTIME_SRCS src/runtime/extra/contrib/cutlass/weight_preprocess.cc)
(tvm) 25/33: [CI] Dump Python stacks when a test stops making progress
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx/standalone-package-cleanup
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 5cd26dc7db68153821bae59dbe80836d3842f612
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 15 05:26:24 2026 +0000
[CI] Dump Python stacks when a test stops making progress
---
tests/scripts/task_python_unittest.sh | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/tests/scripts/task_python_unittest.sh b/tests/scripts/task_python_unittest.sh
index 85b5d692f0..9b319b7a67 100755
--- a/tests/scripts/task_python_unittest.sh
+++ b/tests/scripts/task_python_unittest.sh
@@ -29,4 +29,4 @@ python3 tests/scripts/task_build.py --ffi-only
# --no-deps keeps the submodule-built tvm-ffi above from being replaced by a PyPI build.
uv pip install --no-deps apache-tvm-ffi-orcjit==0.1.2
-python3 -m pytest -vvs -n auto -m "${TVM_TEST_MARKER:-not gpu}" tests/python tirx/tests/python
+python3 -m pytest -vvs -n auto -o faulthandler_timeout=120 -m "${TVM_TEST_MARKER:-not gpu}" tests/python tirx/tests/python