[REFACTOR] Consume standalone TIRx as a pinned dependency (#20803)

Consume Apache TIRx as a pinned submodule at `3rdparty/tvm-tirx`,
replacing the owned TIRx source tree and retaining TVM's compiler and
runtime extensions.

Update native and Python builds, packaging, CI, development scripts, and
documentation to use the standalone component and its recursive FFI
dependency. Keep the parent distribution focused on TVM extensions while
consuming TIRx through its installed interface.

Author: tqchen

.agents/skills/tir-test/SKILL.md

@@ -3,8 +3,8 @@ name: tir-test
 description: Run TVM's in-tree TIRx tests and tirx-kernels import and correctness checks.
 ---
 
-Run the compiler suite in `tirx/tests/python/tirx/`, using the `tvm_tirx` package
-in `tirx/python/`, together with the downstream correctness test in
+Run the compiler suite in `3rdparty/tvm-tirx/tests/python/tirx/`, using the `tvm_tirx` package
+in `3rdparty/tvm-tirx/python/`, together with the downstream correctness test in
 `tests/python/tirx/test_tirx_kernels_registry_correctness.py`.
 
 Both tirx-kernels import gates and its correctness coverage are enabled by
@@ -25,9 +25,9 @@ Run all commands below in the same Bash shell from the TVM repository root.
    export TIR_TEST_XDIST_WORKERS="${TIR_TEST_XDIST_WORKERS:-16}"
    export TIR_TEST_WITH_KERNELS="${TIR_TEST_WITH_KERNELS:-1}"
 
-   [[ -d "$TIR_TEST_TVM_ROOT/tirx/python/tvm_tirx" &&
-      -d "$TIR_TEST_TVM_ROOT/tirx/tests/python/tirx" ]] || {
-     echo "run from the TVM repository root containing tirx/" >&2
+   [[ -d "$TIR_TEST_TVM_ROOT/3rdparty/tvm-tirx/python/tvm_tirx" &&
+      -d "$TIR_TEST_TVM_ROOT/3rdparty/tvm-tirx/tests/python/tirx" ]] || {
+     echo "run from the TVM repository root containing 3rdparty/tvm-tirx/" >&2
      exit 2
    }
    [[ "$TIR_TEST_WITH_KERNELS" == 0 || "$TIR_TEST_WITH_KERNELS" == 1 ]] || {
@@ -89,7 +89,7 @@ Run all commands below in the same Bash shell from the TVM repository root.
    export CUDA_VISIBLE_DEVICES
    CUDA_VISIBLE_DEVICES="$(IFS=,; printf '%s' "${TIR_TEST_SELECTED_GPUS[*]}")"
    # Preserve reference packages and other paths supplied by branch-shell.
-   export PYTHONPATH="${TIR_TEST_TVM_ROOT}/.local/python:${TIR_TEST_TVM_ROOT}/tirx/python:${TIR_TEST_TVM_ROOT}/python:${PYTHONPATH:-}"
+   export PYTHONPATH="${TIR_TEST_TVM_ROOT}/.local/python:${TIR_TEST_TVM_ROOT}/3rdparty/tvm-tirx/python:${TIR_TEST_TVM_ROOT}/python:${PYTHONPATH:-}"
    export TVM_LIBRARY_PATH="${TIR_TEST_TVM_ROOT}/build/lib"
    export TVM_TIRX_LIBRARY_PATH="${TIR_TEST_TVM_ROOT}/build/lib"
 
@@ -111,7 +111,7 @@ Run all commands below in the same Bash shell from the TVM repository root.
    from tvm_tirx.base import _LOADED_LIBS, _RUNTIME_ONLY
 
    root = Path(os.environ["TIR_TEST_TVM_ROOT"])
-   assert Path(tvm_tirx.__file__).resolve().is_relative_to(root / "tirx/python"), (
+   assert Path(tvm_tirx.__file__).resolve().is_relative_to(root / "3rdparty/tvm-tirx/python"), (
        f"imported another TIRx worktree: {tvm_tirx.__file__}"
    )
    assert not _RUNTIME_ONLY, "build the TIRx compiler before running its tests"
@@ -147,14 +147,14 @@ Run all commands below in the same Bash shell from the TVM repository root.
    `TIR_TEST_NUM_GPUS=4` to cap the count or `TIR_TEST_GPUS=3,4,5,6` to override
    selection with exact physical IDs.
 
-   Use dependencies compatible with `tirx/pyproject.toml`. A non-editable FFI
+   Use dependencies compatible with `3rdparty/tvm-tirx/pyproject.toml`. A non-editable FFI
    installation in `.local/python/` takes priority over the inherited environment.
    If native libraries are missing or stale, build this worktree with the same
    `PYTHONPATH` before testing.
 
 2. Run downstream import gates and select the full test coverage:
    ```bash
-   TIR_TEST_PATHS=(tirx/tests/python/tirx/)
+   TIR_TEST_PATHS=(3rdparty/tvm-tirx/tests/python/tirx/)
    if [[ "$TIR_TEST_WITH_KERNELS" == 1 ]]; then
      TIR_TEST_KERNELS_ROOT="${TIR_TEST_KERNELS_ROOT:-$(dirname "$TIR_TEST_TVM_ROOT")/tirx-kernels}"
      [[ -d "$TIR_TEST_KERNELS_ROOT/tirx_kernels" ]] || {
@@ -178,7 +178,7 @@ Run all commands below in the same Bash shell from the TVM repository root.
    "$PYTHON_BIN" -m pytest "${TIR_TEST_PATHS[@]}" -n "$TIR_TEST_XDIST_WORKERS" 2>&1 | tee "$TIR_TEST_LOG"
    ```
 
-   `tirx/tests/python/tirx/conftest.py` uses `tvm_tirx.testing._pytest` to assign
+   `3rdparty/tvm-tirx/tests/python/tirx/conftest.py` uses `tvm_tirx.testing._pytest` to assign
    each xdist worker's PyTorch current device round-robin across the visible
    GPUs. Tests that explicitly use `tvm_tirx.cuda(0)` still run on the first
    selected physical GPU.

.agents/skills/tirx-ptx-dialect/SKILL.md

@@ -178,7 +178,7 @@ Run from the TVM repo root with the workspace `PYTHONPATH` (see `tir-test`).
      codegen test, since certification does not enumerate those helpers;
    - only when the instruction belongs to the MegaMoE extracted-intrinsics
      surface → its kernel and mnemonic list in
-     `tirx/tests/python/tirx/codegen/test_cuda.py::test_megamoe_extracted_intrinsics_codegen`
+     `3rdparty/tvm-tirx/tests/python/tirx/codegen/test_cuda.py::test_megamoe_extracted_intrinsics_codegen`
      (that test covers one workload, not the dialect).
 3. **Regenerate the stub** — any change to a family's tokens or shapes makes
    `test_ptx_stub_up_to_date` fail until you do:
@@ -277,7 +277,7 @@ Do this as one commit series, in this order:
    made the first time this table changed versions.
 3. **Retarget the citations.** Update the module docstring's version
    statement and URL, then run the renumbering over `table.py`,
-   `engine.py`, `render.py`, `tirx/tests/python/tirx/codegen/test_ptx_*.py` and
+   `engine.py`, `render.py`, `3rdparty/tvm-tirx/tests/python/tirx/codegen/test_ptx_*.py` and
    `test_cuda.py`, skipping any block that already carries the new
    numbering (`_PTX_94_ENTRIES` was such a block for 9.3→9.4). Hand-fix the
    forms a regex cannot: ranges that span an insertion (`9.7.x.a-b`), brace

.github/actions/build-wheel-for-publish/action.yml

@@ -110,7 +110,7 @@ runs:
     - name: Build and test TIRx core wheel
       uses: pypa/cibuildwheel@294735312765b09d24a2fbec22660ce817587d55  # v4.1.0
       with:
-        package-dir: tirx
+        package-dir: 3rdparty/tvm-tirx
         output-dir: wheelhouse
       env:
         CIBW_BUILD: ${{ inputs.build }}
@@ -139,7 +139,7 @@ runs:
           PATH="C:/opt/llvm/Library/bin;$PATH"
           CMAKE_ARGS="-DTVM_TIRX_USE_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_TIRX_PACKAGE_EXTRA_LIBS=${{ env.TVM_CUDA_EXTRA_LIB }}"
         # Turns the wheel-specific assertions in
-        # tirx/tests/python/runtime/test_validate_wheel_library.py
+        # 3rdparty/tvm-tirx/tests/python/runtime/test_validate_wheel_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.

.github/actions/setup/action.yml

@@ -44,4 +44,4 @@ runs:
   - name: Install tvm-ffi pip package
     shell: bash -l {0}
     run: |
-      pip install -v ./tirx/3rdparty/tvm-ffi
+      pip install -v ./3rdparty/tvm-tirx/3rdparty/tvm-ffi

.github/workflows/main.yml

@@ -58,51 +58,40 @@ jobs:
           set -euxo pipefail
           pip install "scikit-build-core>=0.11" "setuptools-scm>=8"
           export CMAKE_ARGS="-DTVM_TIRX_USE_LLVM=ON -DTVM_TIRX_USE_METAL=ON -DTVM_TIRX_BUILD_TESTS=OFF"
-          pip wheel --no-deps -w dist ./tirx -v
+          pip wheel --no-deps -w dist ./3rdparty/tvm-tirx -v
           pip install --no-deps --force-reinstall dist/apache_tvm_tirx-*.whl
           export CMAKE_ARGS="-DBUILD_TESTING=OFF"
           pip wheel --no-deps --no-build-isolation -w dist . -v
       - 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
           # versions as 0.1.dev1+g<sha> from the shallow checkout, which is < 0.1.0 under PEP 440.
           # Without it pip replaces the source-matched core with a PyPI build.
           pip install --no-deps apache-tvm-ffi-orcjit==0.1.2
           # Keep the source-matched tvm-ffi installed by the setup action.
           pip install --no-deps dist/*.whl
-          python - <<'PY'
-          import os
-          from importlib.metadata import requires, version
-
-          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.15rc0" 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 tirx/tests/python/codegen/llvm/test_minimal.py
-          tirx/tests/python/runtime/test_validate_wheel_library.py tirx/tests/python/script
+          python -m pytest -v 3rdparty/tvm-tirx/tests/python/codegen/llvm/test_minimal.py
+          3rdparty/tvm-tirx/tests/python/runtime/test_validate_wheel_library.py 3rdparty/tvm-tirx/tests/python/script
       - name: Minimal Metal Compile-Only
         shell: bash -l {0}
         run: |
           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_basic.py::test_func_with_trailing_pod_params'
+          python -m pytest -v -s '3rdparty/tvm-tirx/tests/python/codegen/gpu_common/test_allreduce.py::test_allreduce_sum_compile'
+          python -m pytest -v -s '3rdparty/tvm-tirx/tests/python/backend/metal/codegen/test_basic.py::test_func_with_trailing_pod_params'
       - name: Minimal Metal Compile-and-Run
         shell: bash -l {0}
         run: |
           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_intrinsics.py'
-          python -m pytest -v -s 'tirx/tests/python/codegen/gpu_common/test_allreduce.py::test_allreduce_sum[dims0-metal]'
+          python -m pytest -v -s 3rdparty/tvm-tirx/tests/python/backend/metal
+          python -m pytest -v -s '3rdparty/tvm-tirx/tests/python/codegen/gpu_common/test_intrinsics.py'
+          python -m pytest -v -s '3rdparty/tvm-tirx/tests/python/codegen/gpu_common/test_allreduce.py::test_allreduce_sum[dims0-metal]'
 
   Windows:
     if: ${{ github.repository == 'apache/tvm' }}
@@ -123,7 +112,7 @@ jobs:
         run: |
           pip install "scikit-build-core>=0.11" "setuptools-scm>=8" || exit /b 1
           set CMAKE_ARGS=-DTVM_TIRX_USE_LLVM=ON -DTVM_TIRX_BUILD_TESTS=OFF
-          pip wheel --no-deps -w dist ./tirx -v || exit /b 1
+          pip wheel --no-deps -w dist ./3rdparty/tvm-tirx -v || exit /b 1
           for %%w in (dist\apache_tvm_tirx-*.whl) do (
             pip install --no-deps --force-reinstall "%%w" || exit /b 1
           )
@@ -142,5 +131,5 @@ jobs:
       - name: Test
         shell: cmd /C call {0}
         run: >-
-          python -m pytest -v tirx/tests/python/codegen/llvm/test_minimal.py
-          tirx/tests/python/runtime/test_validate_wheel_library.py tirx/tests/python/script
+          python -m pytest -v 3rdparty/tvm-tirx/tests/python/codegen/llvm/test_minimal.py
+          3rdparty/tvm-tirx/tests/python/runtime/test_validate_wheel_library.py 3rdparty/tvm-tirx/tests/python/script

.gitmodules

@@ -7,6 +7,6 @@
 [submodule "3rdparty/libflash_attn"]
 	path = 3rdparty/libflash_attn
 	url = https://github.com/tlc-pack/libflash_attn
-[submodule "3rdparty/tvm-ffi"]
-	path = tirx/3rdparty/tvm-ffi
-	url = https://github.com/apache/tvm-ffi
+[submodule "3rdparty/tvm-tirx"]
+	path = 3rdparty/tvm-tirx
+	url = https://github.com/apache/tvm-tirx.git

.pre-commit-config.yaml

@@ -23,13 +23,6 @@ repos:
   # ---------- Local hooks ----------
   - repo: local
     hooks:
-      - id: cuda-launch-generated
-        name: check CUDA launch generated files
-        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: ^(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
@@ -53,13 +46,13 @@ repos:
       - id: check-merge-conflict
       - id: check-symlinks
       - id: end-of-file-fixer
-        exclude: ^(tirx/)?3rdparty/
+        exclude: ^3rdparty/
       - id: mixed-line-ending
-        exclude: ^(tirx/)?3rdparty/
+        exclude: ^3rdparty/
       - id: requirements-txt-fixer
         exclude: ^docker/python/bootstrap
       - id: trailing-whitespace
-        exclude: ^(tirx/)?3rdparty/
+        exclude: ^3rdparty/
       - id: check-yaml
       - id: check-toml
 
@@ -95,7 +88,7 @@ repos:
     hooks:
       - id: clang-format
         types_or: [c, c++, cuda, objective-c++]
-        exclude: ^(tirx/)?3rdparty/
+        exclude: ^3rdparty/
 
   # ---------- Cython ----------
   - repo: https://github.com/MarcoGorelli/cython-lint

3rdparty/tvm-tirx

@@ -62,7 +62,8 @@ cmake --build build --parallel
 Development should use `PYTHONPATH`, not editable installs:
 
 ```bash
-export PYTHONPATH="$(pwd)/python:$(pwd)/.local/python"
+export PYTHONPATH="$(pwd)/python:$(pwd)/3rdparty/tvm-tirx/python:$(pwd)/.local/python"
+export TVM_TIRX_LIBRARY_PATH="$(pwd)/build/lib"
 ```
 
 Do not use `pip install -e` for TVM or `tvm-ffi`; editable installs can make
@@ -73,8 +74,8 @@ one worktree silently import another worktree's code.
 Run the smallest relevant test first, then broaden as needed. Common examples:
 
 ```bash
-python -m pytest tirx/tests/python/codegen/llvm/test_minimal.py -xvs
-python -m pytest tirx/tests/python/tirx/test_base.py -xvs
+python -m pytest 3rdparty/tvm-tirx/tests/python/codegen/llvm/test_minimal.py -xvs
+python -m pytest 3rdparty/tvm-tirx/tests/python/tirx/test_base.py -xvs
 ./build/cpptest
 ```
 

AGENTS.md

@@ -85,7 +85,7 @@ set_property(CACHE TVM_USE_GTEST PROPERTY STRINGS AUTO ON OFF)
 # ######### Core dependencies ##########
 
 set(_installed_tirx ON)
-if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tirx/CMakeLists.txt")
+if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm-tirx/CMakeLists.txt")
   set(_installed_tirx OFF)
 endif()
 option(TVM_USE_INSTALLED_TIRX "Build against an installed apache-tvm-tirx package" ${_installed_tirx})
@@ -149,8 +149,8 @@ if(TVM_USE_INSTALLED_TIRX)
         "TVM_OPENCL_DLL=__attribute__((visibility(\"default\")))")
     endif()
   endif()
-  if(NOT TVM_TIRX_VERSION STREQUAL "0.1.0a0")
-    message(FATAL_ERROR "TVM requires apache-tvm-tirx 0.1.0a0")
+  if(NOT TVM_TIRX_VERSION STREQUAL "0.1.0a2")
+    message(FATAL_ERROR "TVM requires apache-tvm-tirx 0.1.0a2")
   endif()
 else()
   # Keep source-built core libraries beside the parent libraries.
@@ -161,7 +161,7 @@ else()
       list(APPEND _tirx_output_defaults ${_kind})
     endif()
   endforeach()
-  add_subdirectory(tirx)
+  add_subdirectory(3rdparty/tvm-tirx)
   if(TARGET tirx_cpptest)
     set_target_properties(tirx_cpptest PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}")
     foreach(_config IN LISTS CMAKE_CONFIGURATION_TYPES)
@@ -234,7 +234,7 @@ add_library(tvm_runtime_extra_defs INTERFACE)
 target_link_libraries(tvm_runtime_extra_defs INTERFACE
   tvm_runtime::shared tvm_ffi::shared)
 target_include_directories(tvm_runtime_extra_defs SYSTEM INTERFACE
-  "${CMAKE_CURRENT_SOURCE_DIR}/tirx/3rdparty/compiler-rt")
+  "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm-tirx/3rdparty/compiler-rt")
 target_compile_definitions(tvm_runtime_extra_defs INTERFACE TVM_RUNTIME_EXTRA_EXPORTS TVM_FFI_EXPORTS)
 target_compile_features(tvm_runtime_extra_defs INTERFACE cxx_std_20)
 target_link_libraries(tvm_runtime_extra PRIVATE tvm_runtime_extra_defs)
@@ -304,7 +304,7 @@ endif()
 if(TARGET tvm_clml_objs)
   # CLML shares OpenCL implementation types; these are source inputs, not an
   # installed core interface. The parent sdist carries their bounded closure.
-  target_include_directories(tvm_clml_objs PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/tirx/src")
+  target_include_directories(tvm_clml_objs PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm-tirx/src")
 endif()
 # Retain the upstream runtime-extra assembly module and its target composition.
 include(cmake/modules/RuntimeExtra.cmake)
@@ -313,7 +313,7 @@ target_sources(tvm_compiler PRIVATE ${COMPILER_SRCS})
 target_compile_definitions(tvm_compiler PRIVATE TVM_EXPORTS TVM_FFI_EXPORTS)
 target_link_libraries(tvm_compiler PUBLIC tvm_tirx_compiler::shared tvm_ffi::shared tvm_runtime_extra)
 target_include_directories(tvm_compiler SYSTEM PRIVATE
-  "${CMAKE_CURRENT_SOURCE_DIR}/tirx/3rdparty/compiler-rt")
+  "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm-tirx/3rdparty/compiler-rt")
 foreach(_target tvm_compiler tvm_runtime_extra)
   target_include_directories(${_target} PUBLIC
     "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>"