[Refactor] Bring up tirx namespace

CMakeLists.txt

@@ -271,7 +271,7 @@ tvm_file_glob(GLOB_RECURSE COMPILER_SRCS
     src/ir/*.cc
     src/arith/*.cc
     src/te/*.cc
-    src/tir/*.cc
+    src/tirx/*.cc
     src/s_tir/*.cc
     src/topi/*.cc
     src/driver/*.cc

docs/README.md

@@ -15,6 +15,22 @@
 <!--- specific language governing permissions and limitations -->
 <!--- under the License. -->
 
+<!--- or more contributor license agreements.  See the NOTICE file -->
+<!--- distributed with this work for additional information -->
+<!--- regarding copyright ownership.  The ASF licenses this file -->
+<!--- to you under the Apache License, Version 2.0 (the -->
+<!--- "License"); you may not use this file except in compliance -->
+<!--- with the License.  You may obtain a copy of the License at -->
+
+<!---   http://www.apache.org/licenses/LICENSE-2.0 -->
+
+<!--- Unless required by applicable law or agreed to in writing, -->
+<!--- software distributed under the License is distributed on an -->
+<!--- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -->
+<!--- KIND, either express or implied.  See the License for the -->
+<!--- specific language governing permissions and limitations -->
+<!--- under the License. -->
+
 # TVM Documentation
 
 This folder contains the source of TVM's documentation, hosted at https://tvm.apache.org/docs

docs/arch/index.rst

@@ -64,10 +64,10 @@ contains a collection of functions. Currently, we support two primary variants o
 - **relax::Function** is a high-level functional program representation. A relax.Function represents high-level graph structure,
   usually corresponds to an end-to-end model or a sub-graph of the overall model. You can view a relax.Function as a computational
   graph with additional support for control-flow, and complex data structures.
-- **tir::PrimFunc** is a low-level program representation that contains elements including loop-nest choices, multi-dimensional load/store,
+- **tirx::PrimFunc** is a low-level program representation that contains elements including loop-nest choices, multi-dimensional load/store,
   threading, and vector/tensor instructions. It is usually used to represent an operator program that executes a (possibly-fused) layer in a model.
 
-During the compilation and transformation, all relax operators are lowered to ``tir::PrimFunc`` or ``TVM PackedFunc``, which can be executed directly
+During the compilation and transformation, all relax operators are lowered to ``tirx::PrimFunc`` or ``TVM PackedFunc``, which can be executed directly
 on the target device, while the calls to relax operators are lowered to calls to low-level functions (e.g. ``R.call_tir`` or ``R.call_dps``).
 
 Transformations
@@ -83,15 +83,14 @@ relax transformations
 relax transformations contain a collection of passes that apply to relax functions. The optimizations include common graph-level
 optimizations such as constant folding and dead-code elimination for operators, and backend-specific optimizations such as library dispatch.
 
-tir transformations
-^^^^^^^^^^^^^^^^^^^
-tir transformations contain a collection of passes that apply to tir functions. There are two major types of transformations:
+tirx transformations
+^^^^^^^^^^^^^^^^^^^^
 
 - **TensorIR schedule**: TensorIR schedules are designed to optimize the TensorIR functions for a specific target, with user-guided instructions and control how the target code is generated.
-  For CPU targets, TIR PrimFunc can generate valid code and execute on the target device without schedule but with very-low performance. However, for GPU targets, the schedule is essential
-  for generating valid code with thread bindings. For more details, please refer to the :ref:`TensorIR Transformation <tir-transform>` section. Additionally, we provides ``MetaSchedule`` to
+  For CPU targets, tirx PrimFunc can generate valid code and execute on the target device without schedule but with very-low performance. However, for GPU targets, the schedule is essential
+  for generating valid code with thread bindings. For more details, please refer to the :ref:`TensorIR Transformation <tirx-transform>` section. Additionally, we provides ``MetaSchedule`` to
   automate the search of TensorIR schedule.
-- **Lowering Passes**: These passes usually perform after the schedule is applied, transforming a TIR PrimFunc into another functionally equivalent PrimFunc, but closer to the
+- **Lowering Passes**: These passes usually perform after the schedule is applied, transforming a tirx PrimFunc into another functionally equivalent PrimFunc, but closer to the
   target-specific representation. For example, there are passes to flatten multi-dimensional access to one-dimensional pointer access, to expand the intrinsics into target-specific ones,
   and to decorate the function entry to meet the runtime calling convention.
 
@@ -102,12 +101,12 @@ focus on optimizations that are not covered by them.
 
 cross-level transformations
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
-Apache TVM enables cross-level optimization of end-to-end models. As the IRModule includes both relax and tir functions, the cross-level transformations are designed to mutate
+Apache TVM enables cross-level optimization of end-to-end models. As the IRModule includes both relax and tirx functions, the cross-level transformations are designed to mutate
 the IRModule by applying different transformations to these two types of functions.
 
-For example, ``relax.LegalizeOps`` pass mutates the IRModule by lowering relax operators, adding corresponding TIR PrimFunc into the IRModule, and replacing the relax operators
-with calls to the lowered TIR PrimFunc. Another example is operator fusion pipeline in relax (including ``relax.FuseOps`` and ``relax.FuseTIR``), which fuses multiple consecutive tensor operations
-into one. Different from the previous implementations, relax fusion pipeline analyzes the pattern of TIR functions and detects the best fusion rules automatically rather
+For example, ``relax.LegalizeOps`` pass mutates the IRModule by lowering relax operators, adding corresponding tirx PrimFunc into the IRModule, and replacing the relax operators
+with calls to the lowered tirx PrimFunc. Another example is operator fusion pipeline in relax (including ``relax.FuseOps`` and ``relax.FuseTIR``), which fuses multiple consecutive tensor operations
+into one. Different from the previous implementations, relax fusion pipeline analyzes the pattern of tirx functions and detects the best fusion rules automatically rather
 than human-defined operator fusion patterns.
 
 Target Translation
@@ -172,12 +171,12 @@ Summary and Discussions
 
 In summary, the key data structures in the compilation flows are:
 
-- IRModule: contains relax.Function and tir.PrimFunc
+- IRModule: contains relax.Function and tirx.PrimFunc
 - runtime.Module: contains runtime.PackedFunc
 
 Most parts of the compilation are transformations among the key data structures.
 
-- relax/transform and tir/transform are deterministic rule-based transformations
+- relax/transform and tirx/transform are deterministic rule-based transformations
 - meta-schedule contains the search-based transformations
 
 Finally, the compilation flow example is only a typical use-case of the TVM stack.
@@ -237,9 +236,9 @@ Thanks to the node module, we can directly access any field of the TVM's IRNode
 
 .. code-block:: python
 
-    x = tvm.tir.Var("x", "int32")
-    y = tvm.tir.Add(x, x)
-    # a and b are fields of a tir.Add node
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Add(x, x)
+    # a and b are fields of a tirx.Add node
     # we can directly use the field name to access the IR structures
     assert y.a == x
 
@@ -249,14 +248,14 @@ The ability to save/store, and inspect an IR node provides a foundation for maki
 tvm/ir
 ------
 The `tvm/ir` folder contains the unified data structure and interfaces across all IR function variants.
-The components in `tvm/ir` are shared by `tvm/relax` and `tvm/tir`, notable ones include
+The components in `tvm/ir` are shared by `tvm/relax` and `tvm/tirx`, notable ones include
 
 - IRModule
 - Type
 - PassContext and Pass
 - Op
 
-Different variants of functions(e.g. relax.Function and tir.PrimFunc) can co-exist in an IRModule.
+Different variants of functions(e.g. relax.Function and tirx.PrimFunc) can co-exist in an IRModule.
 While these variants may not have the same content representation, they use the same data structure to represent types.
 As a consequence, we use the same data structure to represent function (type) signatures of these variants.
 The unified type system allows one function variant to call another function
@@ -267,8 +266,8 @@ The following code snippet gives an example of PassContext configuration.
 
 .. code-block:: python
 
-    # configure the behavior of the tir.UnrollLoop pass
-    with tvm.transform.PassContext(config={"tir.UnrollLoop": { "auto_max_step": 10 }}):
+    # configure the behavior of the tirx.UnrollLoop pass
+    with tvm.transform.PassContext(config={"tirx.UnrollLoop": { "auto_max_step": 10 }}):
         # code affected by the pass context
 
 
@@ -304,34 +303,34 @@ Relax is the high-level IR used to represent the computational graph of a model.
 Note that Relax usually works closely with the TensorIR IRModule, most of the transformations are applied on both Relax and TensorIR functions
 in the IRModule. Please refer to the :ref:`Relax Deep Dive <relax-deep-dive>` for more details.
 
-tvm/tir
--------
+tvm/tirx
+--------
 
-TIR contains the definition of the low-level program representations. We use `tir::PrimFunc` to represent functions that can be transformed by TIR passes.
-Besides the IR data structures, the tir module also includes:
+tirx contains the definition of the low-level program representations. We use `tirx::PrimFunc` to represent functions that can be transformed by tirx passes.
+Besides the IR data structures, the tirx module also includes:
 
-- A set of schedule primitives to control the generated code in ``tir/schedule``.
-- A set of builtin intrinsics in ``tir/tensor_intrin``.
-- A set of analysis passes to analyze the TIR functions in ``tir/analysis``.
-- A set of transformation passes to lower or optimize the TIR functions in ``tir/transform``.
+- A set of schedule primitives to control the generated code in ``tirx/schedule``.
+- A set of builtin intrinsics in ``tirx/tensor_intrin``.
+- A set of analysis passes to analyze the tirx functions in ``tirx/analysis``.
+- A set of transformation passes to lower or optimize the tirx functions in ``tirx/transform``.
 
 Please refer to the :ref:`TensorIR Deep Dive <tensor-ir-deep-dive>` for more details.
 
 tvm/arith
 ---------
 
-This module is closely tied to the TIR. One of the key problems in the low-level code generation is the analysis of the indices'
+This module is closely tied to tirx. One of the key problems in the low-level code generation is the analysis of the indices'
 arithmetic properties — the positiveness, variable bound, and the integer set that describes the iterator space. arith module provides
-a collection of tools that do (primarily integer) analysis. A TIR pass can use these analyses to simplify and optimize the code.
+a collection of tools that do (primarily integer) analysis. A tirx pass can use these analyses to simplify and optimize the code.
 
 tvm/te and tvm/topi
 -------------------
 
 TE stands for Tensor Expression. TE is a domain-specific language (DSL) for describing tensor computations. Importantly, a tensor expression
-itself is not a self-contained function that can be stored into IRModule. We can use ``te.create_prim_func`` to convert a tensor expression to a ``tir::PrimFunc``
+itself is not a self-contained function that can be stored into IRModule. We can use ``te.create_prim_func`` to convert a tensor expression to a ``tirx::PrimFunc``
 and then integrate it into the IRModule.
 
-While possible to construct operators directly via TIR or tensor expressions (TE) for each use case, it is tedious to do so.
+While possible to construct operators directly via tirx or tensor expressions (TE) for each use case, it is tedious to do so.
 `topi` (Tensor operator inventory) provides a set of pre-defined operators defined by numpy and found in common deep learning workloads.
 
 tvm/s_tir/meta_schedule
@@ -343,7 +342,7 @@ and can be used to optimize TensorIR schedules. Note that MetaSchedule only work
 tvm/dlight
 ----------
 
-DLight is a set of pre-defined, easy-to-use, and performant TIR schedules. DLight aims:
+DLight is a set of pre-defined, easy-to-use, and performant tirx schedules. DLight aims:
 
 - Fully support **dynamic shape workloads**.
 - **Light weight**. DLight schedules provides tuning-free schedule with reasonable performance.

docs/arch/pass_infra.rst

@@ -31,7 +31,7 @@ transformation using the analysis result collected during and/or before traversa
 However, as TVM evolves quickly, the need for a more systematic and efficient
 way to manage these passes is becoming apparent. In addition, a generic
 framework that manages the passes across different layers of the TVM stack (e.g.
-Relax and tir) paves the way for developers to quickly prototype and plug the
+Relax and tirx) paves the way for developers to quickly prototype and plug the
 implemented passes into the system.
 
 This doc describes the design of such an infra that takes the advantage of the
@@ -166,7 +166,7 @@ Pass Constructs
 ^^^^^^^^^^^^^^^
 
 The pass infra is designed in a hierarchical manner, and it could work at
-different granularities of Relax/tir programs. A pure virtual class ``PassNode`` is
+different granularities of Relax/tirx programs. A pure virtual class ``PassNode`` is
 introduced to serve as the base of the different optimization passes. This class
 contains several virtual methods that must be implemented by the
 subclasses at the level of modules, functions, or sequences of passes.
@@ -222,13 +222,13 @@ Function-Level Passes
 ^^^^^^^^^^^^^^^^^^^^^
 
 Function-level passes are used to implement various intra-function level
-optimizations for a given Relax/tir module. It fetches one function at a time from
+optimizations for a given Relax/tirx module. It fetches one function at a time from
 the function list of a module for optimization and yields a rewritten Relax
-``Function`` or tir ``PrimFunc``. Most of passes can be classified into this category, such as
+``Function`` or tirx ``PrimFunc``. Most of passes can be classified into this category, such as
 common subexpression elimination and inference simplification in Relax as well as vectorization
-and flattening storage in tir, etc.
+and flattening storage in tirx, etc.
 
-Note that the scope of passes at this level is either a Relax function or a tir primitive function.
+Note that the scope of passes at this level is either a Relax function or a tirx primitive function.
 Therefore, we cannot add or delete a function through these passes as they are not aware of
 the global information.
 
@@ -571,9 +571,9 @@ loop unrolling pass
 
 .. code:: c++
 
-    TVM_REGISTER_PASS_CONFIG_OPTION("tir.UnrollLoop", UnrollLoopConfig);
+    TVM_REGISTER_PASS_CONFIG_OPTION("tirx.UnrollLoop", UnrollLoopConfig);
 
-Please refer to `src/tir/transform/unroll_loop.cc`_ for more details.
+Please refer to `src/tirx/transform/unroll_loop.cc`_ for more details.
 
 .. _pass_instrument_py_frontend:
 
@@ -656,7 +656,7 @@ new ``PassInstrument`` are called.
 
 .. _python/tvm/ir/instrument.py: https://github.com/apache/tvm/blob/main/python/tvm/ir/instrument.py
 
-.. _src/tir/transform/unroll_loop.cc: https://github.com/apache/tvm/blob/main/src/tir/transform/unroll_loop.cc
+.. _src/tirx/transform/unroll_loop.cc: https://github.com/apache/tvm/blob/main/src/tirx/transform/unroll_loop.cc
 
 .. _use pass infra: https://github.com/apache/tvm/blob/main/docs/how_to/tutorials/customize_opt.py
 

docs/arch/runtime.rst

@@ -241,7 +241,7 @@ For example,  we can access the value field of the IntImmNode.
 
     import tvm
 
-    x = tvm.tir.IntImm("int32", 1)
+    x = tvm.tirx.IntImm("int32", 1)
     # access the value field of IntImmNode
     print(x.value)
 

docs/arch/runtimes/vulkan.rst

@@ -254,6 +254,6 @@ string are all false boolean flags.
   validated with `spvValidate`_.
 
 * ``TVM_VULKAN_DEBUG_SHADER_SAVEPATH`` - A path to a directory.  If
-  set to a non-empty string, the Vulkan codegen will save tir, binary
+  set to a non-empty string, the Vulkan codegen will save tirx, binary
   SPIR-V, and disassembled SPIR-V shaders to this directory, to be
   used for debugging purposes.

docs/conf.py

@@ -479,8 +479,8 @@ def force_gc(gallery_conf, fname):
 # Maps the original namespace to list of potential modules
 # that we can import alias from.
 tvm_alias_check_map = {
-    "tvm.te": ["tvm.tir"],
-    "tvm.tir": ["tvm.ir", "tvm.runtime"],
+    "tvm.te": ["tvm.tirx"],
+    "tvm.tirx": ["tvm.ir", "tvm.runtime"],
 }
 
 ## Setup header and other configs
@@ -606,7 +606,7 @@ def update_alias_docstring(name, obj, lines):
 
 
 tvm_class_name_rewrite_map = {
-    "tvm.tir": ["Var", "Call"],
+    "tvm.tirx": ["Var", "Call"],
     "tvm.relax": ["Var", "Call"],
     "tvm.relax.frontend.nn": ["Module"],
 }
@@ -616,7 +616,7 @@ def distinguish_class_name(name: str, lines: list[str]):
     """Distinguish the docstring of type annotations.
 
     In the whole TVM, there are many classes with the same name but in different modules,
-    e.g. ``tir.Var``, ``relax.Var``. This function is used to distinguish them in the docstring,
+    e.g. ``tirx.Var``, ``relax.Var``. This function is used to distinguish them in the docstring,
     by adding the module name as prefix.
 
     To be specific, this function will check the current object name, and if it in the specific

docs/contribute/pull_request.rst

@@ -207,7 +207,7 @@ each time (e.g. you can test a change in CPU and GPU while retaining incremental
     python tests/scripts/ci.py cpu --unittest
 
     # quickly iterate by running a specific test and skipping the rebuild each time
-    python tests/scripts/ci.py cpu --skip-build --tests tests/python/tir-transform/test_tir_transform_inject_rolling_buffer.py::test_upscale
+    python tests/scripts/ci.py cpu --skip-build --tests tests/python/tirx-transform/test_tir_transform_inject_rolling_buffer.py::test_upscale
 
     # run the CPU build and drop into a shell in the container
     python tests/scripts/ci.py cpu --interactive
@@ -261,7 +261,7 @@ If you want to run a single test:
   export PYTHONPATH=python
   rm -rf python/tvm/*.pyc python/tvm/*/*.pyc python/tvm/*/*/*.pyc
 
-  python -m pytest -v tests/python/tir-transform/test_tir_transform_storage_rewrite.py
+  python -m pytest -v tests/python/tirx-transform/test_tir_transform_storage_rewrite.py
 
   # Additionally if you want to run a single test, for example test_all_elemwise inside a file.
   python -m pytest -v -k "test_all_elemwise" tests/python/frontend/tflite/test_forward.py

docs/deep_dive/relax/tutorials/relax_creation.py

@@ -40,7 +40,7 @@
 from tvm import relax, topi
 from tvm.script import ir as I
 from tvm.script import relax as R
-from tvm.script import tir as T
+from tvm.script import tirx as T
 
 
 @I.ir_module

docs/deep_dive/tensor_ir/abstraction.rst

@@ -29,7 +29,7 @@ the compute statements themselves.
 
 .. code:: python
 
-    from tvm.script import tir as T
+    from tvm.script import tirx as T
 
     @T.prim_func
     def main(

docs/deep_dive/tensor_ir/learning.rst

@@ -15,7 +15,7 @@
     specific language governing permissions and limitations
     under the License.
 
-.. _tir-learning:
+.. _tirx-learning:
 
 Understand TensorIR Abstraction
 ===============================

docs/deep_dive/tensor_ir/tutorials/tir_creation.py

@@ -23,7 +23,7 @@
 -----------------
 In this section, we will introduce the methods to write a TensorIR function
 in Apache TVM. This tutorial presumes familiarity with the fundamental concepts of TensorIR.
-If not already acquainted, please refer to :ref:`tir-learning` initially.
+If not already acquainted, please refer to :ref:`tirx-learning` initially.
 
 .. note::
 
@@ -49,14 +49,14 @@
 #
 # Standard Format
 # ***************
-# Let's take an example of ``mm_relu`` from :ref:`tir-learning`. Here is the complete
+# Let's take an example of ``mm_relu`` from :ref:`tirx-learning`. Here is the complete
 # format of the ir_module and in TVMScript:
 
 import numpy as np
 
 import tvm
 from tvm.script import ir as I
-from tvm.script import tir as T
+from tvm.script import tirx as T
 
 
 @I.ir_module