chore: import upstream snapshot with attribution

docs/deep_dive/relax/abstraction.rst

@@ -0,0 +1,73 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+.. _relax-abstraction:
+
+Graph Abstraction for ML Models
+-------------------------------
+Graph abstraction is a key technique used in machine learning (ML) compilers
+to represent and reason about the structure and data flow of ML models. By
+abstracting the model into a graph representation, the compiler can perform
+various optimizations to improve performance and efficiency. This tutorial will
+cover the basics of graph abstraction, its key elements of Relax IR, and how it enables optimization in ML compilers.
+
+What is Graph Abstraction?
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+Graph abstraction is the process of representing an ML model as a directed graph,
+where the nodes represent computational operations (e.g., matrix multiplication,
+convolution) and the edges represent the flow of data between these operations.
+This abstraction allows the compiler to analyze the dependencies and
+relationships between different parts of the model.
+
+.. code:: python
+
+    from tvm.script import relax as R
+
+    @R.function
+    def main(
+        x: R.Tensor((1, 784), dtype="float32"),
+        weight: R.Tensor((784, 256), dtype="float32"),
+        bias: R.Tensor((256,), dtype="float32"),
+    ) -> R.Tensor((1, 256), dtype="float32"):
+        with R.dataflow():
+            lv0 = R.matmul(x, weight)
+            lv1 = R.add(lv0, bias)
+            gv = R.nn.relu(lv1)
+            R.output(gv)
+        return gv
+
+Key Features of Relax
+~~~~~~~~~~~~~~~~~~~~~
+Relax, the graph representation utilized in Apache TVM,
+facilitates end-to-end optimization of ML models through several crucial
+features:
+
+- **First-class symbolic shape**: Relax employs symbolic shapes to represent
+  tensor dimensions, enabling global tracking of dynamic shape relationships
+  across tensor operators and function calls.
+
+- **Multi-level abstractions**: Relax supports cross-level abstractions, from
+  high-level neural network layers to low-level tensor operations, enabling
+  optimizations that span different hierarchies within the model.
+
+- **Composable transformations**: Relax offers a framework for composable
+  transformations that can be selectively applied to different model components.
+  This includes capabilities such as partial lowering and partial specialization,
+  providing flexible customization and optimization options.
+
+These features collectively empower Relax to offer a powerful and adaptable approach
+to ML model optimization within the Apache TVM ecosystem.

docs/deep_dive/relax/dpl.rst

@@ -0,0 +1,557 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+.. _relax-dpl:
+
+Dataflow Pattern Language (DPL)
+===============================
+The Dataflow Pattern Language (DPL) is Relax's built-in facility for
+**pattern matching and rewriting** on computation graphs. It lets you describe a
+sub-graph structure you are looking for, search for it inside a Relax function,
+and optionally replace it with a new structure -- all without hand-writing a
+full IR visitor.
+
+DPL is used throughout the TVM stack:
+
+- **Operator fusion** -- ``FuseOpsByPattern`` groups matched operators into a
+  single fused function.
+- **Backend dispatch** -- CUTLASS, cuBLAS, cuDNN and other backends register
+  patterns so the compiler can route sub-graphs to optimized library kernels.
+- **Custom graph transforms** -- users write their own patterns and rewriters
+  to perform project-specific optimizations.
+
+The typical workflow has three steps:
+
+1. **Build a pattern** that describes the sub-graph shape (e.g. ``matmul`` followed
+   by ``add``).
+2. **Match** the pattern against Relax IR to locate all occurrences.
+3. **Rewrite** each match into a replacement expression.
+
+The public API lives in ``tvm.relax.dpl`` (source: ``python/tvm/relax/dpl/``).
+
+
+Building Patterns
+-----------------
+A *pattern* is a lightweight description of what an expression should look like.
+Patterns are built by combining small building blocks.
+
+Basic Patterns
+~~~~~~~~~~~~~~
+The most common leaf patterns are:
+
+- ``wildcard()`` -- matches any expression.
+- ``is_op("relax.add")`` -- matches a specific Relax operator.
+- ``is_const()`` -- matches any constant value.
+- ``is_var(name)`` -- matches a ``Var`` node (optionally with a given name).
+- ``is_dfv(name)`` -- matches a ``DataflowVar`` node.
+- ``is_gv(name)`` -- matches a ``GlobalVar``.
+
+.. code:: python
+
+    from tvm.relax.dpl import wildcard, is_op, is_const
+
+    # Match any relax.add call, regardless of arguments
+    add_pattern = is_op("relax.add")(wildcard(), wildcard())
+
+Call Patterns
+~~~~~~~~~~~~~
+Calling a pattern as a function produces a ``CallPattern``.  The callee is the
+pattern itself, and the positional arguments are patterns for each operand:
+
+.. code:: python
+
+    x = wildcard()
+    w = wildcard()
+
+    # Match: relax.matmul(x, w)
+    matmul = is_op("relax.matmul")(x, w)
+
+For operators with variadic arguments, pass ``varg_default_wildcard=True`` so
+that extra arguments are matched by implicit wildcards:
+
+.. code:: python
+
+    # Match relax.concat with any number of inputs
+    concat = is_op("relax.concat")(wildcard(), varg_default_wildcard=True)
+
+DPL also provides specialized helpers for common call patterns:
+
+- ``is_call_tir(func_name, args)`` -- matches ``R.call_tir(func_name, (args...,))``.
+- ``is_call_dps_packed(func_name, args)`` -- matches ``R.call_dps_packed``.
+- ``is_call_packed(func_name, args)`` -- matches ``R.call_packed``.
+
+.. code:: python
+
+    from tvm.relax.dpl import is_call_tir, wildcard
+
+    # Match a call_tir that calls the function "decode"
+    decode = is_call_tir("decode", args=[wildcard(), wildcard()])
+
+Tuple Patterns
+~~~~~~~~~~~~~~
+``TuplePattern`` matches a Relax tuple with a fixed number of fields.
+It supports indexing with ``[]`` to create ``TupleGetItemPattern``:
+
+.. code:: python
+
+    from tvm.relax.dpl import is_tuple, wildcard
+
+    a, b = wildcard(), wildcard()
+    tup = is_tuple([a, b])
+
+    # Match: getting the first element from the tuple
+    first = tup[0]
+
+Constraints
+~~~~~~~~~~~
+Any pattern can be further narrowed by attaching constraints:
+
+- ``.has_dtype(dtype)`` -- the matched expression must have the given data type.
+- ``.has_shape(shape)`` -- the matched expression must have the given shape.
+- ``.has_attr(attrs)`` -- the matched call must carry the given attributes.
+- ``.has_ty(ty)`` -- the matched expression must have the given type.
+
+.. code:: python
+
+    # Match a float16 matmul
+    fp16_matmul = is_op("relax.matmul")(wildcard(), wildcard()).has_dtype("float16")
+
+Logical Combinators
+~~~~~~~~~~~~~~~~~~~
+Patterns can be combined with logical operators:
+
+- ``pat_a | pat_b`` -- match if **either** pattern matches (``OrPattern``).
+- ``pat_a & pat_b`` -- match if **both** patterns match (``AndPattern``).
+- ``~pat`` -- match anything **except** ``pat`` (``NotPattern``).
+
+.. code:: python
+
+    # Match either relu or gelu activation
+    activation = is_op("relax.nn.relu")(wildcard()) | is_op("relax.nn.gelu")(wildcard())
+
+Sequence Patterns
+~~~~~~~~~~~~~~~~~
+When a pattern spans multiple bindings inside a ``DataflowBlock``, use
+*sequence operators* to express producer-consumer relationships:
+
+- ``a ^ b`` (``used_by``) -- ``a`` is used by ``b`` (``a`` may also be used
+  elsewhere).
+- ``a >> b`` (``only_used_by``) -- ``a`` is **only** used by ``b`` (no other
+  consumers).
+
+These return a ``PatternSeq`` that can be chained:
+
+.. code:: python
+
+    x = wildcard()
+    matmul = is_op("relax.matmul")(x, wildcard())
+    add = is_op("relax.add")(matmul, wildcard())
+
+    # matmul result is exclusively consumed by the add
+    seq = matmul >> add
+
+High-level Helpers
+~~~~~~~~~~~~~~~~~~
+``make_fused_bias_activation_pattern`` builds a common
+``op -> optional bias -> optional activation`` chain in one call:
+
+.. code:: python
+
+    from tvm.relax.dpl import make_fused_bias_activation_pattern
+
+    # conv2d + bias + relu
+    pattern = make_fused_bias_activation_pattern(
+        "relax.nn.conv2d",
+        with_bias=True,
+        activation="relax.nn.relu",
+    )
+
+
+Matching Without Rewriting
+--------------------------
+Sometimes you only need to **detect** a structure without replacing it.
+Every ``DFPattern`` exposes two matching methods:
+
+- ``pattern.match(expr)`` -- returns ``True`` if the pattern matches.
+- ``pattern.extract_matched_expr(expr)`` -- returns a
+  ``dict[DFPattern, Expr]`` mapping each sub-pattern to the concrete
+  expression it matched, or ``None`` on failure.
+
+.. code:: python
+
+    from tvm.relax.dpl import wildcard, is_op
+
+    x = wildcard()
+    y = wildcard()
+    add_pat = is_op("relax.add")(x, y)
+
+    # Assume `expr` is a Relax expression: R.add(a, b)
+    if add_pat.match(expr):
+        matched = add_pat.extract_matched_expr(expr)
+        # matched[x] -> the expression that matched `x`
+        # matched[y] -> the expression that matched `y`
+
+When matching across variable bindings (e.g., ``lv0 = ...; lv1 = f(lv0)``),
+the matcher needs a ``var2val`` map so it can see through binding
+boundaries.  Use ``tvm.relax.analysis.get_var2val(func)`` to build one:
+
+.. code:: python
+
+    from tvm.relax.analysis import get_var2val
+
+    var2val = get_var2val(func)
+    matched = pattern.extract_matched_expr(expr, var2val=var2val)
+
+
+Rewriting Matched Patterns
+--------------------------
+
+``rewrite_call``
+~~~~~~~~~~~~~~~~
+``rewrite_call`` is the simplest rewrite API.  It walks every expression in a
+function, and when the pattern matches, it calls your callback to produce a
+replacement.
+
+.. code:: python
+
+    rewrite_call(pattern, rewriter, func) -> Function
+
+The callback signature is:
+
+.. code:: python
+
+    def rewriter(
+        matched_expr: Expr,
+        matchings: dict[DFPattern, Expr],
+    ) -> Expr:
+        ...
+
+**Example -- replace** ``reshape(reshape(x, s1), s2)`` **with**
+``reshape(x, s2)``:
+
+.. code:: python
+
+    from tvm import relax
+    from tvm.relax.dpl import wildcard, is_op, rewrite_call
+
+    inp = wildcard()
+    shape1, shape2 = wildcard(), wildcard()
+    inner = is_op("relax.reshape")(inp, shape1)
+    outer = is_op("relax.reshape")(inner, shape2)
+
+    def rewriter(expr, matchings):
+        # Keep the original input but use the outermost target shape
+        return relax.op.reshape(matchings[inp], matchings[outer].args[1])
+
+    new_func = rewrite_call(outer, rewriter, func)
+
+``rewrite_call`` is best for **local, single-expression** rewrites.
+
+
+``rewrite_bindings`` with ``PatternContext``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+When a rewrite involves **multiple bindings** across a ``DataflowBlock``
+(e.g., merging three separate matmuls into one), use ``rewrite_bindings``
+together with ``PatternContext``.
+
+``PatternContext`` enables topological (graph-level) matching on an entire
+dataflow block rather than on individual expressions.
+
+.. code:: python
+
+    rewrite_bindings(ctx, rewriter, func) -> Function
+
+The callback receives *variables* rather than expressions:
+
+.. code:: python
+
+    def rewriter(
+        matchings: dict[DFPattern, Var],
+        bindings: dict[Var, Expr],
+    ) -> dict[Var, Expr]:
+        ...
+
+- ``matchings[pat]`` returns the **bound variable** (``Var``) whose right-hand
+  side matched ``pat``.  The ``Var`` itself carries ``ty`` and can be
+  used directly in new expressions.
+- ``bindings`` maps each ``Var`` to its bound ``Expr`` (the right-hand side),
+  useful when you need to inspect the original expression.
+
+**Example -- merge three parallel matmuls into one**:
+
+.. code:: python
+
+    from tvm.script import relax as R
+    from tvm.relax.dpl import wildcard, is_op, rewrite_bindings, PatternContext
+
+    with PatternContext() as ctx:
+        inp_pat = wildcard()
+        w1, w2, w3 = wildcard(), wildcard(), wildcard()
+
+        matmul1 = is_op("relax.matmul")(inp_pat, w1)
+        matmul2 = is_op("relax.matmul")(inp_pat, w2)
+        matmul3 = is_op("relax.matmul")(inp_pat, w3)
+
+        def rewriter(matchings, _bindings):
+            inp = matchings[inp_pat]
+            W1 = matchings[w1]
+            W2 = matchings[w2]
+            W3 = matchings[w3]
+            width = W1.ty.shape[1]
+
+            concat_w = R.concat([W1, W2, W3], axis=1)
+            merged = R.matmul(inp, concat_w)
+
+            return {
+                matchings[matmul1]: R.strided_slice(
+                    merged, axes=[2], begin=[0], end=[width],
+                ),
+                matchings[matmul2]: R.strided_slice(
+                    merged, axes=[2], begin=[width], end=[width * 2],
+                ),
+                matchings[matmul3]: R.strided_slice(
+                    merged, axes=[2], begin=[width * 2], end=[width * 3],
+                ),
+            }
+
+        new_func = rewrite_bindings(ctx, rewriter, func)
+
+
+Declarative Rewriting with ``@R.rewriter``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+For straightforward one-to-one replacements you can declare the pattern and
+its replacement as two Relax functions in a single ``IRModule``.  The
+``@R.rewriter`` decorator turns the module into a ``PatternMatchingRewriter``
+object that can be applied directly.
+
+.. code:: python
+
+    from tvm.script import relax as R
+
+    @R.rewriter
+    class RewriteAddToPackedCall:
+        @R.function
+        def pattern(
+            A: R.Tensor([16], "float32"),
+            B: R.Tensor([16], "float32"),
+        ):
+            C = R.add(A, B)
+            return C
+
+        @R.function
+        def replacement(
+            A: R.Tensor([16], "float32"),
+            B: R.Tensor([16], "float32"),
+        ):
+            C = R.call_pure_packed(
+                "my_fast_add",
+                A,
+                B,
+                ty_args=R.Tensor([16], "float32"),
+            )
+            return C
+
+    # Apply to an IRModule or a single function
+    rewritten_mod = RewriteAddToPackedCall(mod)
+
+Composing Rewriters
+~~~~~~~~~~~~~~~~~~~
+Multiple ``PatternMatchingRewriter`` objects can be combined with the ``|``
+operator so they run as a single pass:
+
+.. code:: python
+
+    combined = rewriter_a | rewriter_b
+    result = combined(mod)
+
+The left-hand rewriter is tried first; the right-hand rewriter only applies to
+bindings that were **not** already modified by the left.
+
+
+Using DPL in Compiler Passes
+-----------------------------
+The most common way DPL appears in the TVM codebase is through the
+``FuseOpsByPattern`` pass, which uses ``FusionPattern`` objects to drive
+operator fusion.
+
+``FusionPattern``
+~~~~~~~~~~~~~~~~~
+A ``FusionPattern`` bundles four pieces of information:
+
+- ``name`` -- a string label (e.g., ``"cutlass.matmul"``).
+- ``pattern`` -- a ``DFPattern`` that describes the sub-graph to match.
+- ``annotation_patterns`` -- a ``dict[str, DFPattern]`` that names interesting
+  sub-patterns so the check function can inspect them.
+- ``check`` -- an optional ``Callable[[PatternCheckContext], bool]`` that
+  performs additional validation after a structural match succeeds.
+
+.. code:: python
+
+    from tvm.relax.dpl import wildcard, is_op
+    from tvm.relax.transform import FusionPattern
+
+    x = wildcard()
+    w = wildcard()
+    matmul = is_op("relax.matmul")(x, w)
+    bias = wildcard()
+    add = is_op("relax.add")(matmul, bias)
+
+    pattern = FusionPattern(
+        name="my_backend.matmul_bias",
+        pattern=add,
+        annotation_patterns={"matmul": matmul, "bias": bias, "lhs": x, "rhs": w},
+        check=my_check_fn,
+    )
+
+``PatternCheckContext``
+~~~~~~~~~~~~~~~~~~~~~~~
+When ``FuseOpsByPattern`` finds a structural match, it calls the ``check``
+function with a ``PatternCheckContext`` that provides:
+
+- ``matched_expr`` -- the root expression of the match.
+- ``annotated_expr`` -- a ``dict[str, Expr]`` resolved from the
+  ``annotation_patterns``.
+- ``matched_bindings`` -- a ``dict[Var, Expr]`` of bindings being fused.
+- ``var_usages`` -- a ``dict[Var, Sequence[Var]]`` of variable use chains.
+- ``value_to_bound_var`` -- a ``dict[Expr, Var]`` mapping values back to
+  their bound variables.
+
+Use the check function to enforce constraints that cannot be expressed
+structurally (dtype restrictions, shape compatibility, attribute values, etc.):
+
+.. code:: python
+
+    from tvm.relax.transform import PatternCheckContext
+
+    def my_check_fn(ctx: PatternCheckContext) -> bool:
+        matmul_expr = ctx.annotated_expr["matmul"]
+        # Only accept float16 output
+        if matmul_expr.ty.dtype != "float16":
+            return False
+        return True
+
+``FuseOpsByPattern``
+~~~~~~~~~~~~~~~~~~~~
+``FuseOpsByPattern`` is a module-level pass that takes a list of
+``FusionPattern`` (or equivalent tuples) and groups every match into a fused
+sub-function.
+
+.. code:: python
+
+    from tvm.relax.dpl import wildcard, is_op
+    from tvm.relax.transform import FuseOpsByPattern
+
+    # 1. Define the pattern
+    w = wildcard()
+    x = wildcard()
+    wT = is_op("relax.permute_dims")(w)
+    o = is_op("relax.matmul")(x, wT)
+    annotations = {"o": o, "w": w, "x": x, "wT": wT}
+
+    def check(ctx):
+        transpose_call = ctx.annotated_expr["wT"]
+        ndim = transpose_call.args[0].ty.ndim
+        if ndim == -1:
+            return False
+        if ndim == 2 and transpose_call.attrs.axes is None:
+            return True
+        axes = list(range(ndim))
+        axes[-1], axes[-2] = axes[-2], axes[-1]
+        return list(transpose_call.attrs.axes) == axes
+
+    # 2. Run the pass
+    mod = FuseOpsByPattern(
+        [("transpose_matmul_fuse", o, annotations, check)],
+        bind_constants=False,
+    )(mod)
+
+When ``annotate_codegen=True``, each fused function is additionally wrapped
+with ``Codegen`` and ``global_symbol`` attributes, which is how backends like
+CUTLASS and cuBLAS register themselves for external code generation.
+
+
+Quick Reference
+---------------
+
+**Pattern construction**
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 65
+
+   * - API
+     - Description
+   * - ``wildcard()``
+     - Match any expression
+   * - ``is_op(op_name)``
+     - Match a Relax operator by name
+   * - ``is_const()``
+     - Match any constant
+   * - ``is_var(name)`` / ``is_dfv(name)`` / ``is_gv(name)``
+     - Match ``Var`` / ``DataflowVar`` / ``GlobalVar``
+   * - ``is_tuple(fields)``
+     - Match a tuple with given field patterns
+   * - ``is_call_tir(name, args)``
+     - Match ``R.call_tir``
+   * - ``is_call_dps_packed(name, args)``
+     - Match ``R.call_dps_packed``
+   * - ``is_call_packed(name, args)``
+     - Match ``R.call_packed``
+   * - ``make_fused_bias_activation_pattern(...)``
+     - Build ``op + bias + activation`` chain
+   * - ``.has_dtype()`` / ``.has_shape()`` / ``.has_attr()`` / ``.has_ty()``
+     - Attach constraints
+   * - ``|`` / ``&`` / ``~``
+     - Or / And / Not combinators
+   * - ``^`` / ``>>``
+     - used_by / only_used_by (sequence)
+
+**Matching and rewriting**
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 65
+
+   * - API
+     - Description
+   * - ``pattern.match(expr)``
+     - Returns ``True`` if pattern matches
+   * - ``pattern.extract_matched_expr(expr)``
+     - Returns ``dict[DFPattern, Expr]`` or ``None``
+   * - ``rewrite_call(pattern, rewriter, func)``
+     - Rewrite individual expressions
+   * - ``rewrite_bindings(ctx, rewriter, func)``
+     - Rewrite across bindings in a ``DataflowBlock``
+   * - ``PatternMatchingRewriter.from_module(mod)``
+     - Declarative rewriter from ``IRModule``
+   * - ``@R.rewriter``
+     - Decorator shorthand for ``from_module``
+
+**Pass integration**
+
+.. list-table::
+   :header-rows: 1
+   :widths: 35 65
+
+   * - API
+     - Description
+   * - ``FusionPattern(name, pattern, annotations, check)``
+     - Bundle pattern with metadata for ``FuseOpsByPattern``
+   * - ``PatternCheckContext``
+     - Runtime context passed to check functions
+   * - ``FuseOpsByPattern(patterns, ...)``
+     - Module pass that fuses matched sub-graphs

docs/deep_dive/relax/index.rst

@@ -0,0 +1,35 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+.. _relax-deep-dive:
+
+Relax
+=====
+Relax is a high-level abstraction for graph optimization and transformation in Apache TVM stack.
+Additionally, Apache TVM combines Relax and TensorIR together for cross-level
+optimization. Hence, Relax is usually working closely with TensorIR for representing and optimizing
+the whole IRModule
+
+
+.. toctree::
+    :maxdepth: 2
+
+    abstraction
+    learning
+    dpl
+    tutorials/relax_creation
+    tutorials/relax_transformation

docs/deep_dive/relax/learning.rst

@@ -0,0 +1,276 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+.. _relax-learning:
+
+Understand Relax Abstraction
+============================
+Relax is a graph abstraction used in Apache TVM, which
+helps to end-to-end optimize ML models. The principal objective of Relax
+is to depict the structure and data flow of ML models, including the
+dependencies and relationships between different parts of the model, as
+well as how to execute the model on hardware.
+
+End to End Model Execution
+--------------------------
+
+In this chapter, we will use the following model as an example. This is
+a two-layer neural network that consists of two linear operations with
+relu activation.
+
+.. image:: /_static/img/e2e_fashionmnist_mlp_model.png
+   :width: 85%
+   :align: center
+
+
+High-Level Operations Representation
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Let us begin by reviewing a Numpy implementation of the model.
+
+.. code:: python
+
+    def numpy_mlp(data, w0, b0, w1, b1):
+        lv0 = data @ w0 + b0
+        lv1 = np.maximum(lv0, 0)
+        lv2 = lv1 @ w1 + b1
+        return lv2
+
+The above example code shows the high-level array operations to perform the end-to-end model
+execution. Of course, we can rewrite the above code using Relax as follows:
+
+.. code:: python
+
+    from tvm.script import relax as R
+
+    @R.function
+    def relax_mlp(
+        data: R.Tensor(("n", 784), dtype="float32"),
+        w0: R.Tensor((784, 128), dtype="float32"),
+        b0: R.Tensor((128,), dtype="float32"),
+        w1: R.Tensor((128, 10), dtype="float32"),
+        b1: R.Tensor((10,), dtype="float32"),
+    ) -> R.Tensor(("n", 10), dtype="float32"):
+        with R.dataflow():
+            lv0 = R.matmul(data, w0) + b0
+            lv1 = R.nn.relu(lv0)
+            lv2 = R.matmul(lv1, w1) + b1
+            R.output(lv2)
+        return lv2
+
+Low-Level Integration
+~~~~~~~~~~~~~~~~~~~~~
+
+However, again from the pov of machine learning compilation (MLC), we would like to see
+through the details under the hood of these array computations.
+
+For the purpose of illustrating details under the hood, we will again write examples in low-level numpy:
+
+We will use a loop instead of array functions when necessary to demonstrate the possible loop computations.
+When possible, we always explicitly allocate arrays via numpy.empty and pass them around.
+The code block below shows a low-level numpy implementation of the same model.
+
+.. code:: python
+
+    def lnumpy_linear(X: np.ndarray, W: np.ndarray, B: np.ndarray, Z: np.ndarray):
+        n, m, K = X.shape[0], W.shape[1], X.shape[1]
+        Y = np.empty((n, m), dtype="float32")
+        for i in range(n):
+            for j in range(m):
+                for k in range(K):
+                    if k == 0:
+                        Y[i, j] = 0
+                    Y[i, j] = Y[i, j] + X[i, k] * W[k, j]
+
+        for i in range(n):
+            for j in range(m):
+                Z[i, j] = Y[i, j] + B[j]
+
+
+    def lnumpy_relu0(X: np.ndarray, Y: np.ndarray):
+        n, m = X.shape
+        for i in range(n):
+            for j in range(m):
+                Y[i, j] = np.maximum(X[i, j], 0)
+
+    def lnumpy_mlp(data, w0, b0, w1, b1):
+        n = data.shape[0]
+        lv0 = np.empty((n, 128), dtype="float32")
+        lnumpy_linear(data, w0, b0, lv0)
+
+        lv1 = np.empty((n, 128), dtype="float32")
+        lnumpy_relu0(lv0, lv1)
+
+        out = np.empty((n, 10), dtype="float32")
+        lnumpy_linear(lv1, w1, b1, out)
+        return out
+
+With the low-level NumPy example in mind, now we are ready to introduce an Relax abstraction
+for the end-to-end model execution. The code block below shows a TVMScript implementation of the model.
+
+.. code:: python
+
+    from tvm.script import ir as I
+    from tvm.script import tirx as T
+    from tvm.script import relax as R
+
+    @I.ir_module
+    class Module:
+        @T.prim_func(private=True)
+        def linear(x: T.handle, w: T.handle, b: T.handle, z: T.handle):
+            M, N, K = T.int64(), T.int64(), T.int64()
+            X = T.match_buffer(x, (M, K), "float32")
+            W = T.match_buffer(w, (K, N), "float32")
+            B = T.match_buffer(b, (N,), "float32")
+            Z = T.match_buffer(z, (M, N), "float32")
+            Y = T.alloc_buffer((M, N), "float32")
+            for i, j, k in T.grid(M, N, K):
+                with T.sblock("Y"):
+                    v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
+                    with T.init():
+                        Y[v_i, v_j] = T.float32(0.0)
+                    Y[v_i, v_j] = Y[v_i, v_j] + X[v_i, v_k] * W[v_k, v_j]
+            for i, j in T.grid(M, N):
+                with T.sblock("Z"):
+                    v_i, v_j = T.axis.remap("SS", [i, j])
+                    Z[v_i, v_j] = Y[v_i, v_j] + B[v_j]
+
+        @T.prim_func(private=True)
+        def relu(x: T.handle, y: T.handle):
+            M, N = T.int64(), T.int64()
+            X = T.match_buffer(x, (M, N), "float32")
+            Y = T.match_buffer(y, (M, N), "float32")
+            for i, j in T.grid(M, N):
+                with T.sblock("Y"):
+                    v_i, v_j = T.axis.remap("SS", [i, j])
+                    Y[v_i, v_j] = T.max(X[v_i, v_j], T.float32(0.0))
+
+        @R.function
+        def main(
+            x: R.Tensor(("n", 784), dtype="float32"),
+            w0: R.Tensor((784, 256), dtype="float32"),
+            b0: R.Tensor((256,), dtype="float32"),
+            w1: R.Tensor((256, 10), dtype="float32"),
+            b1: R.Tensor((10,), dtype="float32")
+        ) -> R.Tensor(("n", 10), dtype="float32"):
+            cls = Module
+            n = T.int64()
+            with R.dataflow():
+                lv = R.call_tir(cls.linear, (x, w0, b0), out_ty=R.Tensor((n, 256), dtype="float32"))
+                lv1 = R.call_tir(cls.relu, (lv,), out_ty=R.Tensor((n, 256), dtype="float32"))
+                lv2 = R.call_tir(cls.linear, (lv1, w1, b1), out_ty=R.Tensor((n, 10), dtype="float32"))
+                R.output(lv2)
+            return lv2
+
+The above code contains kinds of functions: the primitive tensor functions (``T.prim_func``) and a
+``R.function`` (relax function). Relax function is a new type of abstraction representing
+high-level neural network executions.
+
+Note that the above relax module natively supports symbolic shapes, see the ``"n"`` in the
+tensor shapes in ``main`` function and ``M``, ``N``, ``K`` in the ``linear`` function. This is
+a key feature of Relax abstraction, which enables the compiler to track dynamic shape relations
+globally across tensor operators and function calls.
+
+Again it is helpful to see the TVMScript code and low-level numpy code side-by-side and check the
+corresponding elements, and we are going to walk through each of them in detail. Since we already
+learned about primitive tensor functions, we are going to focus on the high-level execution part.
+
+Key Elements of Relax
+---------------------
+This section will introduce the key elements of Relax abstraction and how it enables optimization
+in ML compilers.
+
+Type
+~~~~
+Type is the Relax representation of expression type information. It can
+be ``TensorType``, ``TupleType``, etc. In the above example, we use ``TensorType``
+(short in ``R.Tensor`` in TVMScript) to represent the shape and dtype of the tensor of the inputs,
+outputs, and intermediate results.
+
+R.call_tir
+~~~~~~~~~~
+The ``R.call_tir`` function is a new abstraction in Relax that allows calling primitive tensor
+functions in the same IRModule. This is a key feature of Relax that enables cross-level
+abstractions, from high-level neural network layers to low-level tensor operations.
+Taking one line from the above code as an example:
+
+.. code:: python
+
+    lv = R.call_tir(cls.linear, (x, w0, b0), out_ty=R.Tensor((n, 256), dtype="float32"))
+
+To explain what does ``R.call_tir`` work, let us review an equivalent low-level numpy
+implementation of the operation, as follows:
+
+.. code:: python
+
+    lv0 = np.empty((n, 256), dtype="float32")
+    lnumpy_linear(x, w0, b0, lv0)
+
+Specifically, ``call_tir`` allocates an output tensor res, then pass the inputs and the output
+to the prim_func. After executing prim_func the result is populated in res, then we can return
+the result.
+
+This convention is called **destination passing**, The idea is that input and output are explicitly
+allocated outside and passed to the low-level primitive function. This style is commonly used
+in low-level library designs, so higher-level frameworks can handle that memory allocation
+decision. Note that not all tensor operations can be presented in this style (specifically,
+there are operations whose output shape depends on the input). Nevertheless, in common practice,
+it is usually helpful to write the low-level function in this style when possible.
+
+Dataflow Block
+~~~~~~~~~~~~~~
+Another important element in a relax function is the R.dataflow() scope annotation.
+
+.. code:: python
+
+    with R.dataflow():
+        lv = R.call_tir(cls.linear, (x, w0, b0), out_ty=R.Tensor((n, 256), dtype="float32"))
+        lv1 = R.call_tir(cls.relu, (lv,), out_ty=R.Tensor((n, 256), dtype="float32"))
+        lv2 = R.call_tir(cls.linear, (lv1, w1, b1), out_ty=R.Tensor((n, 10), dtype="float32"))
+        R.output(lv2)
+
+Before we talk about the dataflow block, let us first introduce the concept of **pure** and
+**side-effect**. A function is **pure** or **side-effect free** if:
+
+- it only reads from its inputs and returns the result via its output
+- it will not change other parts of the program (such as incrementing a global counter).
+
+For example, all ``R.call_tir`` functions are pure functions, as they only read from their inputs
+and write the output to another new allocated tensor. However, the **inplace operations** are not
+pure functions, in other words, they are side-effect functions, because they will change the existing
+intermediate or input tensors.
+
+A dataflow block is a way for us to mark the computational graph regions of the program.
+Specifically, within a dataflow block, all the operations need to be **side-effect free**.
+Outside a dataflow block, the operations can contain side-effect.
+
+.. note::
+
+    A common question that arises is why we need to manually mark dataflow blocks instead of
+    automatically inferring them. There are two main reasons for this approach:
+
+    - Automatic inference of dataflow blocks can be challenging and imprecise, particularly
+      when dealing with calls to packed functions (such as cuBLAS integrations). By manually
+      marking dataflow blocks, we enable the compiler to accurately understand and optimize
+      the program's dataflow.
+    - Many optimizations can only be applied within dataflow blocks. For instance, fusion
+      optimization is limited to operations within a single dataflow block. If the compiler
+      were to incorrectly infer dataflow boundaries, it might miss crucial optimization
+      opportunities, potentially impacting the program's performance.
+
+By allowing manual marking of dataflow blocks, we ensure that the compiler has the most
+accurate information to work with, leading to more effective optimizations.

docs/deep_dive/relax/tutorials/README.txt

@@ -0,0 +1,2 @@
+Deep Dive: Relax
+----------------

docs/deep_dive/relax/tutorials/relax_creation.py

@@ -0,0 +1,285 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E402
+
+"""
+.. _relax-creation:
+
+Relax Creation
+==============
+This tutorial demonstrates how to create Relax functions and programs.
+We'll cover various ways to define Relax functions, including using TVMScript,
+and relax NNModule API.
+"""
+
+
+######################################################################
+# Create Relax programs using TVMScript
+# -------------------------------------
+# TVMScript is a domain-specific language for representing Apache TVM's
+# intermediate representation (IR). It is a Python dialect that can be used
+# to define an IRModule, which contains both TensorIR and Relax functions.
+#
+# In this section, we will show how to define a simple MLP model with only
+# high-level Relax operators using TVMScript.
+
+from tvm import relax, topi
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tirx as T
+
+
+@I.ir_module
+class RelaxModule:
+    @R.function
+    def forward(
+        data: R.Tensor(("n", 784), dtype="float32"),
+        w0: R.Tensor((128, 784), dtype="float32"),
+        b0: R.Tensor((128,), dtype="float32"),
+        w1: R.Tensor((10, 128), dtype="float32"),
+        b1: R.Tensor((10,), dtype="float32"),
+    ) -> R.Tensor(("n", 10), dtype="float32"):
+        with R.dataflow():
+            lv0 = R.matmul(data, R.permute_dims(w0)) + b0
+            lv1 = R.nn.relu(lv0)
+            lv2 = R.matmul(lv1, R.permute_dims(w1)) + b1
+            R.output(lv2)
+        return lv2
+
+
+RelaxModule.show()
+
+######################################################################
+# Relax is not only a graph-level IR, but also supports cross-level
+# representation and transformation. To be specific, we can directly call
+# TensorIR functions in Relax function.
+
+
+@I.ir_module
+class RelaxModuleWithTIR:
+    @T.prim_func(s_tir=True)
+    def relu(x: T.handle, y: T.handle):
+        n = T.int64()
+        m = T.int64()
+        X = T.match_buffer(x, (n, m), "float32")
+        Y = T.match_buffer(y, (n, m), "float32")
+        for i, j in T.grid(n, m):
+            with T.sblock("relu"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                Y[vi, vj] = T.max(X[vi, vj], T.float32(0))
+
+    @R.function
+    def forward(
+        data: R.Tensor(("n", 784), dtype="float32"),
+        w0: R.Tensor((128, 784), dtype="float32"),
+        b0: R.Tensor((128,), dtype="float32"),
+        w1: R.Tensor((10, 128), dtype="float32"),
+        b1: R.Tensor((10,), dtype="float32"),
+    ) -> R.Tensor(("n", 10), dtype="float32"):
+        n = T.int64()
+        cls = RelaxModuleWithTIR
+        with R.dataflow():
+            lv0 = R.matmul(data, R.permute_dims(w0)) + b0
+            lv1 = R.call_tir(cls.relu, lv0, R.Tensor((n, 128), dtype="float32"))
+            lv2 = R.matmul(lv1, R.permute_dims(w1)) + b1
+            R.output(lv2)
+        return lv2
+
+
+RelaxModuleWithTIR.show()
+
+######################################################################
+# .. note::
+#
+#   You may notice that the printed output is different from the written
+#   TVMScript code. This is because we print the IRModule in a standard
+#   format, while we support syntax sugar for the input
+#
+#   For example, we can combine multiple operators into a single line, as
+#
+#   .. code-block:: python
+#
+#     lv0 = R.matmul(data, R.permute_dims(w0)) + b0
+#
+#   However, the normalized expression requires only one operation in one
+#   binding. So the printed output is different from the written TVMScript code,
+#   as
+#
+#   .. code-block:: python
+#
+#     lv: R.Tensor((784, 128), dtype="float32") = R.permute_dims(w0, axes=None)
+#     lv1: R.Tensor((n, 128), dtype="float32") = R.matmul(data, lv)
+#     lv0: R.Tensor((n, 128), dtype="float32") = R.add(lv1, b0)
+#
+
+######################################################################
+# Create Relax programs using NNModule API
+# ----------------------------------------
+# Besides TVMScript, we also provide a PyTorch-like API for defining neural networks.
+# It is designed to be more intuitive and easier to use than TVMScript.
+#
+# In this section, we will show how to define the same MLP model using
+# Relax NNModule API.
+
+from tvm.relax.frontend import nn
+
+
+class NNModule(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(784, 128)
+        self.relu1 = nn.ReLU()
+        self.fc2 = nn.Linear(128, 10)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.relu1(x)
+        x = self.fc2(x)
+        return x
+
+
+######################################################################
+# After we define the NNModule, we can export it to TVM IRModule via
+# ``export_tvm``.
+
+mod, params = NNModule().export_tvm({"forward": {"x": nn.spec.Tensor(("n", 784), "float32")}})
+mod.show()
+
+######################################################################
+# We can also insert customized function calls into the NNModule, such as
+# Tensor Expression(TE), TensorIR functions or other TVM packed functions.
+
+
+@T.prim_func(s_tir=True)
+def tir_linear(x: T.handle, w: T.handle, b: T.handle, z: T.handle):
+    M = T.int64()
+    N = T.int64()
+    K = T.int64()
+    X = T.match_buffer(x, (M, K), "float32")
+    W = T.match_buffer(w, (N, K), "float32")
+    B = T.match_buffer(b, (N,), "float32")
+    Z = T.match_buffer(z, (M, N), "float32")
+    for i, j, k in T.grid(M, N, K):
+        with T.sblock("linear"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                Z[vi, vj] = 0
+            Z[vi, vj] = Z[vi, vj] + X[vi, vk] * W[vj, vk]
+    for i, j in T.grid(M, N):
+        with T.sblock("add"):
+            vi, vj = T.axis.remap("SS", [i, j])
+            Z[vi, vj] = Z[vi, vj] + B[vj]
+
+
+class NNModuleWithTIR(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(784, 128)
+        self.fc2 = nn.Linear(128, 10)
+
+    def forward(self, x):
+        n = x.shape[0]
+        # We can call external functions using nn.extern
+        x = nn.extern(
+            "env.linear",
+            [x, self.fc1.weight, self.fc1.bias],
+            out=nn.Tensor.placeholder((n, 128), "float32"),
+        )
+        # We can also call TensorIR via Tensor Expression API in TOPI
+        x = nn.tensor_expr_op(topi.nn.relu, "relu", [x])
+        # We can also call other TVM packed functions
+        x = nn.tensor_ir_op(
+            tir_linear,
+            "tir_linear",
+            [x, self.fc2.weight, self.fc2.bias],
+            out=nn.Tensor.placeholder((n, 10), "float32"),
+        )
+        return x
+
+
+mod, params = NNModuleWithTIR().export_tvm(
+    {"forward": {"x": nn.spec.Tensor(("n", 784), "float32")}}
+)
+mod.show()
+
+
+######################################################################
+# Create Relax programs using Block Builder API
+# ---------------------------------------------
+# In addition to the above APIs, we also provide a Block Builder API for
+# creating Relax programs. It is a IR builder API, which is more
+# low-level and widely used in TVM's internal logic, e.g writing a
+# customized pass.
+
+bb = relax.BlockBuilder()
+n = T.int64()
+x = relax.Var("x", R.Tensor((n, 784), "float32"))
+fc1_weight = relax.Var("fc1_weight", R.Tensor((128, 784), "float32"))
+fc1_bias = relax.Var("fc1_bias", R.Tensor((128,), "float32"))
+fc2_weight = relax.Var("fc2_weight", R.Tensor((10, 128), "float32"))
+fc2_bias = relax.Var("fc2_bias", R.Tensor((10,), "float32"))
+with bb.function("forward", [x, fc1_weight, fc1_bias, fc2_weight, fc2_bias]):
+    with bb.dataflow():
+        lv0 = bb.emit(relax.op.matmul(x, relax.op.permute_dims(fc1_weight)) + fc1_bias)
+        lv1 = bb.emit(relax.op.nn.relu(lv0))
+        gv = bb.emit(relax.op.matmul(lv1, relax.op.permute_dims(fc2_weight)) + fc2_bias)
+        bb.emit_output(gv)
+    bb.emit_func_output(gv)
+
+mod = bb.get()
+mod.show()
+
+######################################################################
+# Also, Block Builder API supports building cross-level IRModule with both
+# Relax functions, TensorIR functions and other TVM packed functions.
+
+bb = relax.BlockBuilder()
+with bb.function("forward", [x, fc1_weight, fc1_bias, fc2_weight, fc2_bias]):
+    with bb.dataflow():
+        lv0 = bb.emit(
+            relax.call_dps_packed(
+                "env.linear",
+                [x, fc1_weight, fc1_bias],
+                out_ty=relax.TensorType((n, 128), "float32"),
+            )
+        )
+        lv1 = bb.emit_te(topi.nn.relu, lv0)
+        tir_gv = bb.add_func(tir_linear, "tir_linear")
+        gv = bb.emit(
+            relax.call_tir(
+                tir_gv,
+                [lv1, fc2_weight, fc2_bias],
+                out_ty=relax.TensorType((n, 10), "float32"),
+            )
+        )
+        bb.emit_output(gv)
+    bb.emit_func_output(gv)
+mod = bb.get()
+mod.show()
+
+######################################################################
+# Note that the Block Builder API is not as user-friendly as the above APIs,
+# but it is lowest-level API and works closely with the IR definition. We
+# recommend using the above APIs for users who only want to define and
+# transform a ML model. But for those who want to build more complex
+# transformations, the Block Builder API is a more flexible choice.
+
+######################################################################
+# Summary
+# -------
+# This tutorial demonstrates how to create Relax programs using TVMScript,
+# NNModule API, Block Builder API and PackedFunc API for different use cases.

docs/deep_dive/relax/tutorials/relax_transformation.py

@@ -0,0 +1,142 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E402
+
+"""
+.. _relax-transform:
+
+Transformation
+--------------
+In this section, we will dive into the transformation of Relax programs.
+Transformations is one of the key ingredients of the compilation flows
+for optimizing and integrating with hardware backends.
+"""
+
+######################################################################
+# Let's first create a simple Relax program as what we have done in
+# the :ref:`previous section <relax-creation>`.
+
+import tvm
+from tvm import IRModule, relax
+from tvm.relax.frontend import nn
+
+
+class NNModule(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(784, 128)
+        self.relu1 = nn.ReLU()
+        self.fc2 = nn.Linear(128, 10)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.relu1(x)
+        x = self.fc2(x)
+        return x
+
+
+origin_mod, params = NNModule().export_tvm(
+    {"forward": {"x": nn.spec.Tensor(("n", 784), "float32")}}
+)
+origin_mod.show()
+
+######################################################################
+# Apply transformations
+# ~~~~~~~~~~~~~~~~~~~~~
+# Passes are the main way to apply transformations to the program.
+# We can apply passes to the program. As first step, let's apply
+# a built-in pass ``LegalizeOps`` to lower the high-level operators
+# into low-level operators.
+
+mod = tvm.relax.transform.LegalizeOps()(origin_mod)
+mod.show()
+
+######################################################################
+# As we can see from the output, the high-level operators (aka ``relax.op``) in the program
+# are replaced by their corresponding low-level operators (aka ``relax.call_tir``).
+#
+# Then let's trying to apply the operator fusion, which is a wide-used optimization technique
+# in ML compilers. Note that in relax, fusion optimizations are done with the collaboration of
+# a set of passes. We can apply them in a sequence.
+
+mod = tvm.ir.transform.Sequential(
+    [
+        tvm.relax.transform.AnnotateTIROpPattern(),
+        tvm.relax.transform.FuseOps(),
+        tvm.relax.transform.FuseTIR(),
+    ]
+)(mod)
+mod.show()
+
+######################################################################
+# As result, we can see that the ``matmul``, ``add`` and ``relu`` operators are fused
+# into one kernel (aka one ``call_tir``).
+#
+# For all built-in passes, please refer to :py:class:`relax.transform`.
+#
+# Custom Passes
+# ~~~~~~~~~~~~~
+# We can also define our own passes. Let's take an example of rewriting the ``relu``
+# operator to ``gelu`` operator.
+#
+# First, we need to write a Relax IR Mutator to do the rewriting.
+
+from tvm.relax.expr_functor import PyExprMutator, mutator
+
+
+@mutator
+class ReluRewriter(PyExprMutator):
+    def __init__(self, mod):
+        super().__init__(mod)
+
+    def visit_call_(self, call: relax.Call) -> relax.Expr:
+        # visit the relax.Call expr, and only handle the case when op is relax.nn.relu
+        if call.op.name == "relax.nn.relu":
+            return relax.op.nn.gelu(call.args[0])
+
+        return super().visit_call_(call)
+
+
+######################################################################
+# Then we can write a pass to apply the mutator to the whole module.
+
+
+@tvm.transform.module_pass(opt_level=0, name="ReluToGelu")
+class ReluToGelu:  # pylint: disable=too-few-public-methods
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        rewriter = ReluRewriter(mod)
+        for g_var, func in mod.functions_items():
+            if isinstance(func, relax.Function):
+                func = rewriter.visit_expr(func)
+                rewriter.builder_.update_func(g_var, func)
+        return rewriter.builder_.get()
+
+
+mod = ReluToGelu()(origin_mod)
+mod.show()
+
+######################################################################
+# The printed output shows that the ``relax.nn.relu`` operator is
+# rewritten to ``relax.nn.gelu`` operator.
+#
+# For the details of the mutator, please refer to :py:class:`relax.expr_functor.PyExprMutator`.
+#
+# Summary
+# ~~~~~~~
+# In this section, we have shown how to apply transformations to the Relax program.
+# We have also shown how to define and apply custom transformations.

docs/deep_dive/tensor_ir/abstraction.rst

@@ -0,0 +1,72 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+.. _tir-abstraction:
+
+Tensor Program Abstraction
+--------------------------
+Before we dive into the details of TensorIR, let's first introduce what is a primitive tensor
+function. Primitive tensor functions are functions that correspond to a single "unit" of
+computational operation. For example, a convolution operation can be a primitive tensor function,
+and a fused convolution + relu operation can also be a primitive tensor function.
+Usually, a typical abstraction for primitive tensor function implementation contains the following
+elements: multi-dimensional buffers, loop nests that drive the tensor computations, and finally,
+the compute statements themselves.
+
+.. code:: python
+
+    from tvm.script import tirx as T
+
+    @T.prim_func
+    def main(
+        A: T.Buffer((128,), "float32"),
+        B: T.Buffer((128,), "float32"),
+        C: T.Buffer((128,), "float32"),
+    ) -> None:
+        for i in range(128):
+            with T.sblock("C"):
+                vi = T.axis.spatial(128, i)
+                C[vi] = A[vi] + B[vi]
+
+Key Elements of Tensor Programs
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+The demonstrated primitive tensor function calculates the element-wise sum of two vectors.
+The function:
+
+- Accepts three **multi-dimensional buffers** as parameters, and generates one **multi-dimensional
+  buffer** as output.
+- Incorporates a solitary **loop nest** ``i`` that facilitates the computation.
+- Features a singular **compute statement** that calculates the element-wise sum of the two
+  vectors.
+
+Extra Structure in TensorIR
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Crucially, we are unable to execute arbitrary transformations on the program, as certain
+computations rely on the loop's sequence. Fortunately, the majority of primitive tensor
+functions we focus on possess favorable properties, such as independence among loop iterations.
+For instance, the aforementioned program includes block and iteration annotations:
+
+- The **block annotation** ``with T.sblock("C")`` signifies that the block is the fundamental
+  computation unit designated for scheduling. A block may encompass a single computation
+  statement, multiple computation statements with loops, or opaque intrinsics such as Tensor
+  Core instructions.
+- The **iteration annotation** ``T.axis.spatial``, indicating that variable ``vi`` is mapped
+  to ``i``, and all iterations are independent.
+
+While this information isn't crucial for *executing* the specific program, it proves useful when
+transforming the program. Consequently, we can confidently parallelize or reorder loops associated
+with ``vi``, provided we traverse all the index elements from 0 to 128.

docs/deep_dive/tensor_ir/index.rst

@@ -0,0 +1,43 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+.. _tensor-ir-deep-dive:
+
+TensorIR
+========
+TensorIR is one of the core abstractions in the Apache TVM stack, used to
+represent and optimize primitive tensor functions.
+
+The TensorIR codebase consists of two modules (split from the former ``tir``):
+
+- **tirx** — Core IR definitions and lowering (PrimFunc, Buffer, SBlock,
+  expressions, statements, lowering passes).
+- **s_tir** (Schedulable TIR) — Schedule primitives, MetaSchedule, DLight,
+  and tensor intrinsics.
+
+In TVMScript, both modules are accessed via
+``from tvm.script import tirx as T``.
+
+.. toctree::
+    :maxdepth: 2
+
+    abstraction
+    learning
+    tutorials/tir_creation
+    tutorials/tir_transformation
+    tutorials/dlight_gpu_scheduling
+    tutorials/meta_schedule

docs/deep_dive/tensor_ir/learning.rst

@@ -0,0 +1,255 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+.. _tirx-learning:
+
+Understand TensorIR Abstraction
+===============================
+TensorIR is the tensor program abstraction in Apache TVM, which is one of the standard
+machine learning compilation frameworks. The principal objective of tensor program abstraction
+is to depict loops and associated hardware acceleration options, including threading, the
+application of specialized hardware instructions, and memory access.
+
+To help our explanations, let us use the following sequence of tensor computations as
+a motivating example. Specifically, for two :math:`128 \times 128` matrices ``A`` and ``B``, let us perform the
+following two steps of tensor computations.
+
+.. math::
+
+  Y_{i, j} &= \sum_k A_{i, k} \times B_{k, j} \\
+  C_{i, j} &= \mathbb{relu}(Y_{i, j}) = \mathbb{max}(Y_{i, j}, 0)
+
+
+The above computations resemble a typical primitive tensor function commonly seen in neural networks,
+a linear layer with relu activation. We use TensorIR to depict the above computations as follows.
+
+Before we invoke TensorIR, let's use native Python codes with NumPy to show the computation:
+
+.. code:: python
+
+    def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):
+        Y = np.empty((128, 128), dtype="float32")
+        for i in range(128):
+            for j in range(128):
+                for k in range(128):
+                    if k == 0:
+                        Y[i, j] = 0
+                    Y[i, j] = Y[i, j] + A[i, k] * B[k, j]
+        for i in range(128):
+            for j in range(128):
+                C[i, j] = max(Y[i, j], 0)
+
+With the low-level NumPy example in mind, now we are ready to introduce TensorIR. The code block
+below shows a TensorIR implementation of ``mm_relu``. The particular code is implemented in a
+language called TVMScript, which is a domain-specific dialect embedded in python AST.
+
+.. code:: python
+
+    from tvm.script import tirx as T
+
+    @tvm.script.ir_module
+    class MyModule:
+        @T.prim_func
+        def mm_relu(A: T.Buffer((128, 128), "float32"),
+                    B: T.Buffer((128, 128), "float32"),
+                    C: T.Buffer((128, 128), "float32")):
+            Y = T.alloc_buffer((128, 128), dtype="float32")
+            for i, j, k in T.grid(128, 128, 128):
+                with T.sblock("Y"):
+                    vi = T.axis.spatial(128, i)
+                    vj = T.axis.spatial(128, j)
+                    vk = T.axis.reduce(128, k)
+                    with T.init():
+                        Y[vi, vj] = T.float32(0)
+                    Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
+            for i, j in T.grid(128, 128):
+                with T.sblock("C"):
+                    vi = T.axis.spatial(128, i)
+                    vj = T.axis.spatial(128, j)
+                    C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
+
+
+Next, let's invest the elements in the above TensorIR program.
+
+Function Parameters and Buffers
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+**The function parameters correspond to the same set of parameters on the numpy function.**
+
+.. code:: python
+
+    # TensorIR
+    def mm_relu(A: T.Buffer((128, 128), "float32"),
+                B: T.Buffer((128, 128), "float32"),
+                C: T.Buffer((128, 128), "float32")):
+        ...
+    # NumPy
+    def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):
+        ...
+
+Here ``A``, ``B``, and ``C`` takes a type named ``T.Buffer``, which with shape
+argument ``(128, 128)`` and data type ``float32``. This additional information
+helps possible MLC process to generate code that specializes in the shape and data
+type.
+
+**Similarly, TensorIR also uses a buffer type in intermediate result allocation.**
+
+.. code:: python
+
+    # TensorIR
+    Y = T.alloc_buffer((128, 128), dtype="float32")
+    # NumPy
+    Y = np.empty((128, 128), dtype="float32")
+
+Loop Iterations
+~~~~~~~~~~~~~~~
+**There are also direct correspondence of loop iterations.**
+
+``T.grid`` is a syntactic sugar in TensorIR for us to write multiple nested iterators.
+
+.. code:: python
+
+    # TensorIR with `T.grid`
+    for i, j, k in T.grid(128, 128, 128):
+        ...
+    # TensorIR with `range`
+    for i in range(128):
+        for j in range(128):
+            for k in range(128):
+                ...
+    # NumPy
+    for i in range(128):
+        for j in range(128):
+            for k in range(128):
+                ...
+
+Computational Block
+~~~~~~~~~~~~~~~~~~~
+A significant distinction lies in computational statements:
+**TensorIR incorporates an additional construct termed** ``T.sblock``.
+
+.. code:: python
+
+    # TensorIR
+    with T.sblock("Y"):
+        vi = T.axis.spatial(128, i)
+        vj = T.axis.spatial(128, j)
+        vk = T.axis.reduce(128, k)
+        with T.init():
+            Y[vi, vj] = T.float32(0)
+        Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
+    # NumPy
+    vi, vj, vk = i, j, k
+    if vk == 0:
+        Y[vi, vj] = 0
+    Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
+
+A **block** represents a fundamental computation unit within TensorIR. Importantly,
+a block encompasses more information than standard NumPy code. It comprises a set of block axes
+``(vi, vj, vk)`` and the computations delineated around them.
+
+.. code:: python
+
+    vi = T.axis.spatial(128, i)
+    vj = T.axis.spatial(128, j)
+    vk = T.axis.reduce(128, k)
+
+The above three lines declare the **key properties** about block axes in the following syntax.
+
+.. code:: python
+
+    [block_axis] = T.axis.[axis_type]([axis_range], [mapped_value])
+
+These three lines convey the following details:
+
+- They specify the binding of ``vi``, ``vj``, ``vk`` (in this instance, to ``i``, ``j``, ``k``).
+- They declare the original range intended for ``vi``, ``vj``, ``vk``
+  (the 128 in ``T.axis.spatial(128, i)``).
+- They announce the properties of the iterators (spatial, reduce).
+
+Block Axis Properties
+~~~~~~~~~~~~~~~~~~~~~
+Let's delve deeper into the properties of the block axis. These properties signify the axis's
+relationship to the computation in progress. The block comprises three axes ``vi``, ``vj``, and
+``vk``, meanwhile the block reads the buffer ``A[vi, vk]``, ``B[vk, vj]`` and writes the buffer
+``Y[vi, vj]``. Strictly speaking, the block performs (reduction) updates to Y, which we label
+as write for the time being, as we don't require the value of Y from another block.
+
+Significantly, for a fixed value of ``vi`` and ``vj``, the computation block yields a point
+value at a spatial location of ``Y`` (``Y[vi, vj]``) that is independent of other locations in ``Y``
+(with different ``vi``, ``vj`` values). We can refer to ``vi``, ``vj`` as **spatial axes** since
+they directly correspond to the start of a spatial region of buffers that the block writes to.
+The axes involved in reduction (``vk``) are designated as **reduce axes**.
+
+Why Extra Information in Block
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+One crucial observation is that the additional information (block axis range and their properties)
+makes the block to be **self-contained** when it comes to the iterations that it is supposed to
+carry out independent from the external loop-nest ``i, j, k``.
+
+The block axis information also provides additional properties that help us to validate the correctness of the
+external loops that are used to carry out the computation. For example, the above code block will result in an
+error because the loop expects an iterator of size 128, but we only bound it to a for loop of size 127.
+
+.. code:: python
+
+    # wrong program due to loop and block iteration mismatch
+    for i in range(127):
+        with T.sblock("C"):
+            vi = T.axis.spatial(128, i)
+            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+            error here due to iterator size mismatch
+            ...
+
+Sugars for Block Axes Binding
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+In situations where each of the block axes is directly mapped to an outer loop iterator,
+we can use ``T.axis.remap`` to declare the block axis in a single line.
+
+.. code:: python
+
+    # SSR means the properties of each axes are "spatial", "spatial", "reduce"
+    vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+
+which is equivalent to
+
+.. code:: python
+
+    vi = T.axis.spatial(range_of_i, i)
+    vj = T.axis.spatial(range_of_j, j)
+    vk = T.axis.reduce (range_of_k, k)
+
+So we can also write the programs as follows.
+
+.. code:: python
+
+    @tvm.script.ir_module
+    class MyModuleWithAxisRemapSugar:
+        @T.prim_func
+        def mm_relu(A: T.Buffer((128, 128), "float32"),
+                    B: T.Buffer((128, 128), "float32"),
+                    C: T.Buffer((128, 128), "float32")):
+            Y = T.alloc_buffer((128, 128), dtype="float32")
+            for i, j, k in T.grid(128, 128, 128):
+                with T.sblock("Y"):
+                    vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                    with T.init():
+                        Y[vi, vj] = T.float32(0)
+                    Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
+            for i, j in T.grid(128, 128):
+                with T.sblock("C"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    C[vi, vj] = T.max(Y[vi, vj], T.float32(0))

docs/deep_dive/tensor_ir/tutorials/README.txt

@@ -0,0 +1,2 @@
+Deep Dive: TensorIR
+-------------------

docs/deep_dive/tensor_ir/tutorials/dlight_gpu_scheduling.py

@@ -0,0 +1,316 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E402, E501
+
+"""
+.. _dlight_gpu_scheduling:
+
+DLight: Rule-Based GPU Scheduling
+==================================
+TIR functions produced by Relax legalization need GPU-specific scheduling — thread binding,
+loop tiling, shared memory usage — before they can run efficiently on a GPU. There are two
+main approaches in TVM:
+
+- **MetaSchedule**: explores a search space to find the best schedule. High quality, but
+  compilation takes minutes to hours.
+- **DLight**: applies pre-defined scheduling rules deterministically. No tuning required,
+  compilation completes in seconds. Performance is excellent for well-known patterns
+  (e.g., GEMM, GEMV in LLM workloads) and fair for the rest.
+
+This tutorial covers how DLight works, what rules are available, how to diagnose scheduling
+quality, and how to write custom rules.
+
+.. contents:: Table of Contents
+    :local:
+    :depth: 1
+"""
+
+######################################################################
+# Prepare a Model
+# ---------------
+# We build a small model with ``nn.Module`` that is rich enough to trigger multiple DLight
+# rules: ``Linear`` layers produce GEMM (matrix multiplication) kernels, ``LayerNorm``
+# produces a general-reduction kernel, and ``ReLU`` is a simple elementwise op.
+
+import tvm
+from tvm import relax, tirx
+from tvm.relax.frontend import nn
+from tvm.s_tir import dlight as dl
+
+
+class DemoModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(768, 768)
+        self.relu = nn.ReLU()
+        self.norm = nn.LayerNorm(768)
+        self.fc2 = nn.Linear(768, 256)
+
+    def forward(self, x):
+        x = self.norm(self.relu(self.fc1(x)))
+        return self.fc2(x)
+
+
+mod, params = DemoModel().export_tvm({"forward": {"x": nn.spec.Tensor((1, 768), "float32")}})
+
+######################################################################
+# Legalize Relax operators into TIR functions so that DLight has concrete kernels to schedule.
+
+device = tvm.cuda(0)
+target = tvm.target.Target.from_device(device)
+with target:
+    mod = relax.get_pipeline("zero")(mod)
+
+######################################################################
+# At this point every TIR function in ``mod`` is **unscheduled** — it has no thread bindings
+# and would not run efficiently on a GPU. Let's see what functions we have:
+for gv, func in mod.functions_items():
+    if isinstance(func, tirx.PrimFunc):
+        print(f"  {gv.name_hint}")
+
+######################################################################
+# Basic Usage: ApplyDefaultSchedule
+# ---------------------------------
+# ``ApplyDefaultSchedule`` is an ``IRModule`` pass. It iterates over every TIR function in the
+# module and tries the given rules **in order**. For each function the first rule whose
+# ``apply()`` returns a non-``None`` schedule wins; subsequent rules are skipped.
+# After scheduling, the function is marked with ``tirx.is_scheduled`` so it won't be
+# scheduled again by a later ``ApplyDefaultSchedule`` call.
+
+######################################################################
+# Here we use a common subset of rules. The full catalog (including ``LowBatchGEMV``,
+# ``Transpose``, ``RMSNorm``) is listed in the next section.
+
+with target:
+    scheduled_mod = dl.ApplyDefaultSchedule(
+        dl.gpu.Matmul(),  # GEMM: dense matrix multiplication
+        dl.gpu.GEMV(),  # matrix-vector products
+        dl.gpu.Reduction(),  # simple reductions (sum, max, ...)
+        dl.gpu.GeneralReduction(),  # compound reductions (softmax, layer norm, ...)
+        dl.gpu.Fallback(),  # catch-all for anything unmatched above
+    )(mod)
+
+scheduled_mod.show()
+
+######################################################################
+# Compared with the unscheduled IR, you can now see thread bindings
+# (``blockIdx.x``, ``threadIdx.x``, ...) and loop transformations in each TIR function.
+
+######################################################################
+# Rule Catalog
+# ------------
+# DLight ships a set of GPU scheduling rules. Each rule is a subclass of
+# ``ScheduleRule`` and implements an ``apply(func, target, tunable)`` method that returns
+# a ``Schedule`` if the rule matches, or ``None`` to pass.
+#
+# The built-in GPU rules, roughly from most specific to most general:
+#
+# .. list-table::
+#    :header-rows: 1
+#    :widths: 20 40 40
+#
+#    * - Rule
+#      - Pattern
+#      - Typical operators
+#    * - ``Matmul``
+#      - GEMM index pattern ``C[S,I,J] += A[S,I,K] * B[S,J,K]``
+#      - ``nn.Linear``, batched matmul
+#    * - ``GEMV``
+#      - Matrix-vector multiply (one dimension is 1)
+#      - single-batch decode in attention
+#    * - ``LowBatchGEMV``
+#      - Low-batch GEMM scheduled with a GEMV strategy
+#      - small-batch decode
+#    * - ``Reduction``
+#      - Simple accumulation ``X[...] += Y[...]``
+#      - sum, max, argmax
+#    * - ``GeneralReduction``
+#      - Spatial dims followed by reduction dims (``S* R*``)
+#      - softmax, layer norm, RMS norm
+#    * - ``Transpose``
+#      - Read/write indices are permutations of each other
+#      - 2-D transpose
+#    * - ``RMSNorm``
+#      - Contains an ``rsqrt`` operation
+#      - RMS normalization
+#    * - ``Fallback``
+#      - Any function (always matches)
+#      - generic catch-all
+#
+# **Rule order matters.** ``ApplyDefaultSchedule`` stops at the first match, so:
+#
+# - Put **specialized** rules first (``Matmul``, ``GEMV``) — they have strict matching
+#   conditions but produce high-quality schedules.
+# - Put **general** rules later (``GeneralReduction``, ``Fallback``) — they match broadly
+#   but with less optimal schedules.
+# - If you put ``Fallback`` first, it would "steal" every function and no specialized
+#   rule would ever run.
+
+######################################################################
+# Diagnosing Schedule Quality
+# ---------------------------
+# A common question is: *which rule scheduled which function?* ``ApplyDefaultSchedule``
+# does not log this directly, but you can figure it out by applying rules one at a time.
+#
+# **Step 1**: Apply each rule individually and record which functions it claims.
+
+from collections import OrderedDict
+
+rules = OrderedDict(
+    [
+        ("Matmul", dl.gpu.Matmul()),
+        ("GEMV", dl.gpu.GEMV()),
+        ("LowBatchGEMV", dl.gpu.LowBatchGEMV()),
+        ("Reduction", dl.gpu.Reduction()),
+        ("GeneralReduction", dl.gpu.GeneralReduction()),
+        ("Transpose", dl.gpu.Transpose()),
+        ("RMSNorm", dl.gpu.RMSNorm()),
+    ]
+)
+
+rule_assignment = {}
+for rule_name, rule in rules.items():
+    with target:
+        test_mod = dl.ApplyDefaultSchedule(rule)(mod)
+    for gv, func in test_mod.functions_items():
+        if isinstance(func, tirx.PrimFunc) and gv.name_hint not in rule_assignment:
+            if "tirx.is_scheduled" in func.attrs and func.attrs["tirx.is_scheduled"] == 1:
+                rule_assignment[gv.name_hint] = rule_name
+
+######################################################################
+# **Step 2**: Functions not claimed by any specialized rule will fall through to ``Fallback``.
+
+all_tir_funcs = [
+    gv.name_hint for gv, func in mod.functions_items() if isinstance(func, tirx.PrimFunc)
+]
+fallback_funcs = [name for name in all_tir_funcs if name not in rule_assignment]
+
+print("Rule assignments:")
+for name, rule_name in sorted(rule_assignment.items()):
+    print(f"  {name:40s} -> {rule_name}")
+if fallback_funcs:
+    print("Handled by Fallback (may have suboptimal performance):")
+    for name in sorted(fallback_funcs):
+        print(f"  {name}")
+
+######################################################################
+# If an important kernel lands in the Fallback bucket, you have three options:
+#
+# 1. Write a **custom DLight rule** for it (see below).
+# 2. Use **MetaSchedule** to auto-tune that specific function.
+# 3. Manually schedule it with the ``tvm.s_tir.Schedule`` API.
+
+######################################################################
+# DLight vs MetaSchedule
+# ----------------------
+# The two systems are complementary, not competing:
+#
+# .. list-table::
+#    :header-rows: 1
+#    :widths: 20 40 40
+#
+#    * -
+#      - DLight
+#      - MetaSchedule
+#    * - Mechanism
+#      - Deterministic rule matching
+#      - Search-space exploration
+#    * - Compile time
+#      - Seconds
+#      - Minutes to hours
+#    * - Performance
+#      - Excellent on known patterns, fair otherwise
+#      - Near-optimal with sufficient search budget
+#    * - Best for
+#      - Default path, rapid iteration, CI
+#      - Hot-spot tuning in production
+#
+# A practical workflow:
+#
+# 1. Run ``ApplyDefaultSchedule`` with the full rule set to cover all functions.
+# 2. Profile the compiled model to identify hot-spot kernels.
+# 3. Use ``MetaScheduleTuneTIR`` to auto-tune only those kernels.
+#
+# Note that ``MetaScheduleTuneTIR`` does **not** automatically skip functions already
+# scheduled by DLight — it processes every ``PrimFunc`` in the module. In practice this
+# is harmless (tuning an already-scheduled function simply re-explores its space), but if
+# you want to avoid the extra search cost, filter the module or use ``MetaScheduleTuneIRMod``
+# with ``op_names`` to target specific functions.
+
+######################################################################
+# Writing a Custom Rule
+# ---------------------
+# You can extend DLight by writing your own ``ScheduleRule``. The simplest way is
+# ``ScheduleRule.from_callable``, which wraps a plain function into a rule **instance**.
+
+from tvm import s_tir
+from tvm.s_tir.dlight.analysis import normalize_prim_func
+from tvm.s_tir.dlight.base.schedule_rule import ScheduleRule
+
+
+@ScheduleRule.from_callable("MyTileAndBind")
+def my_tile_and_bind(func: tirx.PrimFunc, target: tvm.target.Target, tunable: bool):
+    """A minimal rule: for single-block injective functions, tile and bind to GPU threads."""
+    if not isinstance(func, tirx.PrimFunc):
+        return None
+    sch = s_tir.Schedule(func)
+    # Use normalize_prim_func to get block info with correct spatial/reduction classification.
+    # This is the same analysis used by built-in DLight rules.
+    block_infos = normalize_prim_func(sch)
+    if block_infos is None or len(block_infos) != 1:
+        return None  # only handle single-block functions
+    info = block_infos[0]
+    if not info.is_injective():
+        return None  # skip reductions — dom_kind() uses iter_type, not loop kind
+    loops = sch.get_loops(info.block_rv)
+    if len(loops) == 0:
+        return None
+    fused = sch.fuse(*loops)
+    bx, tx = sch.split(fused, factors=[None, 256])
+    sch.bind(bx, "blockIdx.x")
+    sch.bind(tx, "threadIdx.x")
+    return sch
+
+
+######################################################################
+# Insert the custom rule into the rule chain. Note that ``from_callable`` returns an
+# **instance**, so pass it directly — do not call ``my_tile_and_bind()`` again.
+
+with target:
+    custom_mod = dl.ApplyDefaultSchedule(
+        dl.gpu.Matmul(),
+        dl.gpu.GeneralReduction(),
+        my_tile_and_bind,  # our custom rule, tried before Fallback
+        dl.gpu.Fallback(),
+    )(mod)
+
+custom_mod.show()
+
+######################################################################
+# To build a production-quality rule, subclass ``ScheduleRule`` directly and implement
+# ``apply()`` with full analysis logic (see ``tvm.s_tir.dlight.gpu.Matmul`` for an example).
+
+######################################################################
+# Summary
+# -------
+# - **DLight** provides fast, deterministic GPU scheduling via rule matching.
+# - Rules are tried in order; the first match wins. Put specialized rules before general ones.
+# - Use the **single-rule probing** technique to diagnose which rule handles each function.
+# - Combine DLight with MetaSchedule: DLight for baseline coverage, MetaSchedule for hot-spot tuning.
+# - Extend DLight by writing custom ``ScheduleRule`` implementations.
+#
+# For DLight's role in the broader optimization pipeline, see :ref:`customize_opt`.

docs/deep_dive/tensor_ir/tutorials/meta_schedule.py

@@ -0,0 +1,307 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E402
+
+"""
+.. _meta_schedule_deep_dive:
+
+MetaSchedule: Search-Based Auto-Tuning
+=======================================
+MetaSchedule is TVM's search-based auto-tuning framework, located in
+``python/tvm/s_tir/meta_schedule/``. It explores different TIR schedules
+(loop tiling, vectorization, thread binding, etc.) and measures them on real
+hardware to find the fastest implementation for each operator.
+
+While **DLight** (see :ref:`dlight_gpu_scheduling`) provides rule-based scheduling with zero
+search time, MetaSchedule trades compilation time for better performance by searching over
+the space of possible schedules.
+
+.. contents:: Table of Contents
+    :local:
+    :depth: 1
+"""
+
+######################################################################
+# Architecture Overview
+# ---------------------
+# A MetaSchedule tuning session involves the following components:
+#
+# - **ExtractedTask**: A unique TIR workload extracted from a Relax IRModule,
+#   with a ``task_name`` and ``weight`` (call frequency in the graph).
+# - **TuneContext**: Container holding all resources for a single tuning task
+#   (module, target, space generator, search strategy, etc.).
+# - **SpaceGenerator** (default: ``PostOrderApply``): Generates the design space
+#   of possible schedules by applying ``ScheduleRule`` instances to each block.
+# - **SearchStrategy** (default: ``EvolutionarySearch``): Explores the design
+#   space using an evolutionary algorithm guided by a cost model.
+# - **CostModel** (default: ``XGBModel``): Predicts schedule performance using
+#   XGBoost, reducing the number of actual hardware measurements needed.
+#   Alternatives include ``MLPModel`` (neural network) and ``RandomModel``
+#   (baseline).
+# - **Builder** / **Runner**: Compile and execute candidates on real hardware to
+#   obtain measured run times.
+# - **Database** (default: ``JSONDatabase``): Persistently stores tuning records
+#   (schedule traces + measured run times) for later retrieval.
+# - **TaskScheduler** (default: ``GradientBasedScheduler``): Allocates tuning
+#   budget across multiple tasks based on their weights and estimated improvement
+#   potential.
+#
+# The tuning loop works as follows:
+#
+# 1. The **TaskScheduler** picks a task to tune.
+# 2. The **SpaceGenerator** produces candidate schedules from the design space.
+# 3. The **SearchStrategy** selects candidates (guided by the **CostModel**),
+#    sends them to the **Builder** and **Runner** for measurement.
+# 4. Measured results are committed to the **Database** and used to update the
+#    **CostModel** for the next iteration.
+# 5. Repeat until the trial budget is exhausted.
+
+######################################################################
+# Prepare a Model
+# ---------------
+# We reuse a simple model to demonstrate MetaSchedule APIs.
+
+import os
+import tempfile
+
+import tvm
+from tvm import relax
+from tvm.relax.frontend import nn
+
+
+class DemoModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(784, 256)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(256, 10, bias=False)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        return x
+
+
+input_shape = (1, 784)
+mod, params = DemoModel().export_tvm({"forward": {"x": nn.spec.Tensor(input_shape, "float32")}})
+
+device = tvm.cuda(0)
+target = tvm.target.Target.from_device(device)
+
+######################################################################
+# User-Facing Entry Points
+# ------------------------
+# MetaSchedule provides several levels of API, from high-level transforms to
+# low-level tuning functions.
+#
+# Transform-Based API (Recommended)
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+# These are Relax passes that can be composed into a ``Sequential`` pipeline:
+#
+# - **MetaScheduleTuneIRMod**: Tunes an entire IRModule. Supports ``op_names``
+#   for selective operator tuning.
+# - **MetaScheduleTuneTIR**: Tunes all TIR functions individually (no
+#   ``op_names`` filtering).
+# - **MetaScheduleApplyDatabase**: Applies the best schedules from the tuning
+#   database. Only replaces functions that have records; the rest are left
+#   unchanged.
+#
+# Here is a typical tune-and-apply pipeline:
+#
+# .. note::
+#
+#    To save CI time and avoid flakiness, we skip the tuning process in CI.
+
+if os.getenv("CI", "") != "true":
+    with target, tempfile.TemporaryDirectory() as tmp_dir:
+        tuned_mod = tvm.ir.transform.Sequential(
+            [
+                relax.get_pipeline("zero"),
+                relax.transform.MetaScheduleTuneTIR(
+                    work_dir=tmp_dir,
+                    max_trials_global=300,
+                ),
+                relax.transform.MetaScheduleApplyDatabase(work_dir=tmp_dir),
+            ]
+        )(mod)
+
+    tuned_mod.show()
+
+######################################################################
+# Inspecting Tunable Tasks
+# ------------------------
+# Before tuning, use ``extract_tasks`` to see what MetaSchedule will tune:
+
+from tvm.s_tir.meta_schedule.relax_integration import extract_tasks
+
+with target:
+    legalized_mod = relax.get_pipeline("zero")(mod)
+
+tasks = extract_tasks(legalized_mod, target)
+for i, task in enumerate(tasks):
+    print(f"Task {i}: {task.task_name}  (weight={task.weight})")
+
+######################################################################
+# Each ``ExtractedTask`` has:
+#
+# - ``task_name``: Derived from the PrimFunc name (e.g., ``"fused_matmul_add_relu"``).
+# - ``weight``: How many ``call_tir`` sites invoke this workload. The task
+#   scheduler uses weights to allocate more budget to frequently-called operators.
+# - ``dispatched``: List of candidate TIR modules for this workload.
+
+######################################################################
+# Selective Operator Tuning
+# -------------------------
+# ``MetaScheduleTuneIRMod`` accepts an ``op_names`` parameter to tune only
+# operators whose task name contains any of the given strings:
+#
+# .. code-block:: python
+#
+#     with target:
+#         mod = tvm.ir.transform.Sequential([
+#             relax.transform.MetaScheduleTuneIRMod(
+#                 params={},
+#                 work_dir="./tuning_logs",
+#                 max_trials_global=300,
+#                 op_names=["matmul"],  # Only tune matmul-related operators
+#             ),
+#             relax.transform.MetaScheduleApplyDatabase(work_dir="./tuning_logs"),
+#         ])(mod)
+#
+# Operators without tuning records are left unscheduled -- you can apply DLight or
+# other rule-based schedules to cover them afterward.
+#
+# .. note::
+#
+#    ``MetaScheduleTuneTIR`` does not support ``op_names`` filtering. Use
+#    ``MetaScheduleTuneIRMod`` when you need selective tuning.
+
+######################################################################
+# Database
+# --------
+# When using a fixed ``work_dir``, tuning results are persisted in two
+# newline-delimited JSON files:
+#
+# - ``database_workload.json``: One line per unique workload (structural hash +
+#   serialized IRModule).
+# - ``database_tuning_record.json``: One line per tuning record (workload index +
+#   schedule trace + measured run times).
+#
+# Records are appended incrementally as tuning progresses.
+#
+# Resumption Semantics
+# ~~~~~~~~~~~~~~~~~~~~
+# When you re-run tuning with the same ``work_dir``, existing records are loaded
+# and used as warm-start seeds for the evolutionary search. The tuner does
+# **not** skip already-seen workloads entirely -- it starts from a better initial
+# population, so re-runs are faster than starting from scratch but still consume
+# trials.
+#
+# Once tuning is done, subsequent compilations only need
+# ``MetaScheduleApplyDatabase``:
+#
+# .. code-block:: python
+#
+#     with target:
+#         mod = relax.transform.MetaScheduleApplyDatabase(
+#             work_dir="./tuning_logs"
+#         )(mod)
+#
+# Database Implementations
+# ~~~~~~~~~~~~~~~~~~~~~~~~
+# MetaSchedule ships several database backends:
+#
+# - **JSONDatabase**: Persistent file-based storage (default). Created
+#   automatically when you pass ``work_dir``.
+# - **MemoryDatabase**: In-memory, non-persistent. Useful for testing.
+# - **UnionDatabase**: Queries all sub-databases and returns the globally best
+#   record.
+# - **OrderedUnionDatabase**: Queries sub-databases in order; returns from the
+#   first one that has a match.
+# - **ScheduleFnDatabase**: Wraps a user-provided scheduling function.
+
+######################################################################
+# Cross-Model Database Reuse
+# --------------------------
+# MetaSchedule identifies workloads by their structural hash. If two models
+# contain operators with the same shape, dtype, and computation, they share the
+# same hash and can reuse tuning records.
+#
+# module_equality Options
+# ~~~~~~~~~~~~~~~~~~~~~~~
+# - ``"structural"`` (default): Exact structural match. Safe but strict.
+# - ``"anchor-block"``: Match based on the dominant compute block, ignoring
+#   surrounding context. More permissive -- enables sharing across fused operators
+#   that have the same core computation but different fusion boundaries.
+#
+# ``OrderedUnionDatabase`` enables a layered lookup strategy: check a local
+# database first, then fall back to a shared team database:
+#
+# .. code-block:: python
+#
+#     from tvm.s_tir.meta_schedule.database import JSONDatabase, OrderedUnionDatabase
+#
+#     local_db = JSONDatabase(work_dir="./my_tuning_logs")
+#     shared_db = JSONDatabase(work_dir="/shared/tuning_db")
+#     combined_db = OrderedUnionDatabase(local_db, shared_db)
+#
+#     with target, combined_db:
+#         mod = relax.transform.MetaScheduleApplyDatabase()(mod)
+
+######################################################################
+# Key Parameters Reference
+# ------------------------
+#
+# .. list-table::
+#    :header-rows: 1
+#    :widths: 25 75
+#
+#    * - Parameter
+#      - Description
+#    * - ``max_trials_global``
+#      - Total trial budget shared across all tasks. Set proportional to the
+#        number of tasks (e.g., 200-500 trials per task for good results).
+#    * - ``max_trials_per_task``
+#      - Per-task trial cap. Defaults to ``max_trials_global`` if not set.
+#    * - ``op_names``
+#      - List of strings to filter tasks by name (substring match).
+#        ``MetaScheduleTuneIRMod`` only.
+#    * - ``work_dir``
+#      - Directory for database files and logs. Use a fixed path to enable
+#        persistence and resumption.
+#    * - ``cost_model``
+#      - ``"xgb"`` (XGBoost, default), ``"mlp"`` (neural network), or
+#        ``"random"`` (baseline). Only available via ``tune_relax``.
+#    * - ``runner``
+#      - ``"local"`` (default) or an ``RPCRunner`` instance for remote devices.
+#        Only available via ``tune_relax``.
+#    * - ``module_equality``
+#      - ``"structural"`` (default) or ``"anchor-block"`` for more permissive
+#        cross-model matching. Only available via ``tune_relax``.
+
+######################################################################
+# Summary
+# -------
+# - **MetaSchedule** finds high-quality TIR schedules by searching over the
+#   design space and measuring on real hardware.
+# - Use ``MetaScheduleTuneTIR`` for full-module tuning, or
+#   ``MetaScheduleTuneIRMod`` with ``op_names`` for selective tuning.
+# - Tuning records persist in ``work_dir`` and can be reused across runs and
+#   models with the same operator shapes.
+# - Combine with DLight: use DLight for fast baseline coverage, then MetaSchedule
+#   for hot-spot tuning (see :ref:`dlight_gpu_scheduling`).

docs/deep_dive/tensor_ir/tutorials/tir_creation.py

@@ -0,0 +1,289 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E402
+
+"""
+.. _tir-creation:
+
+TensorIR Creation
+-----------------
+In this section, we will introduce the methods to write a TensorIR function
+in Apache TVM. This tutorial presumes familiarity with the fundamental concepts of TensorIR.
+If not already acquainted, please refer to :ref:`tirx-learning` initially.
+
+.. note::
+
+    This tutorial concentrates on the construction of **standalone** TensorIR functions. The
+    techniques presented here are not requisite for end users to compile Relax models.
+
+"""
+
+######################################################################
+# Create TensorIR using TVMScript
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+# The most straightforward way to create a TensorIR function via TVMScript.
+# TVMScript is a TVM Python dialect that represents TensorIR in TVM.
+#
+# .. important::
+#
+#     While TVMScript employs Python syntax and AST, ensuring full compatibility
+#     with Python tools like auto-completion and linting, it is not a native Python
+#     language and cannot be executed by a Python interpreter.
+#
+#     More precisely, the decorator **@tvm.script** extracts the Python AST from
+#     the decorated function, subsequently parsing it into TensorIR.
+#
+# Standard Format
+# ***************
+# Let's take an example of ``mm_relu`` from :ref:`tirx-learning`. Here is the complete
+# format of the ir_module and in TVMScript:
+
+import numpy as np
+import tvm_ffi
+
+import tvm
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+@I.ir_module
+class MyModule:
+    @T.prim_func(s_tir=True)
+    def mm_relu(
+        A: T.Buffer((128, 128), "float32"),
+        B: T.Buffer((128, 128), "float32"),
+        C: T.Buffer((128, 128), "float32"),
+    ):
+        Y = T.alloc_buffer((128, 128), dtype="float32")
+        for i in range(128):
+            for j in range(128):
+                for k in range(128):
+                    with T.sblock("Y"):
+                        vi = T.axis.spatial(128, i)
+                        vj = T.axis.spatial(128, j)
+                        vk = T.axis.reduce(128, k)
+                        T.reads(A[vi, vk], B[vk, vj])
+                        T.writes(Y[vi, vj])
+                        with T.init():
+                            Y[vi, vj] = T.float32(0)
+                        Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
+        for i in range(128):
+            for j in range(128):
+                with T.sblock("C"):
+                    vi = T.axis.spatial(128, i)
+                    vj = T.axis.spatial(128, j)
+                    T.reads(Y[vi, vj])
+                    T.writes(C[vi, vj])
+                    C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
+
+
+######################################################################
+# Concise with Syntactic Sugar
+# ****************************
+# For ease of writing, we can employ the following syntactic sugar to
+# streamline the code:
+#
+# - Utilize ``T.grid`` to condense nested loops;
+# - Employ ``T.axis.remap`` to abbreviate block iterator annotations;
+# - Exclude ``T.reads`` and ``T.writes`` for blocks whose content can
+#   be inferred from the block body;
+
+
+@I.ir_module
+class ConciseModule:
+    @T.prim_func(s_tir=True)
+    def mm_relu(
+        A: T.Buffer((128, 128), "float32"),
+        B: T.Buffer((128, 128), "float32"),
+        C: T.Buffer((128, 128), "float32"),
+    ):
+        Y = T.alloc_buffer((128, 128), dtype="float32")
+        for i, j, k in T.grid(128, 128, 128):
+            with T.sblock("Y"):
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                with T.init():
+                    Y[vi, vj] = T.float32(0)
+                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
+        for i, j in T.grid(128, 128):
+            with T.sblock("C"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
+
+
+######################################################################
+# We can use the following code to verify that the two modules are equivalent:
+
+print(tvm_ffi.structural_equal(MyModule, ConciseModule))
+
+######################################################################
+# Interactive with Python Variables
+# *********************************
+# Despite TVMScript not being executed by a Python interpreter, limited
+# interaction with Python is feasible. For instance, Python variables can
+# be used to ascertain the shape and data type of a TensorIR.
+
+# Python variables
+M = N = K = 128
+dtype = "float32"
+
+
+# IRModule in TVMScript
+@I.ir_module
+class ConciseModuleFromPython:
+    @T.prim_func(s_tir=True)
+    def mm_relu(
+        A: T.Buffer((M, K), dtype),
+        B: T.Buffer((K, N), dtype),
+        C: T.Buffer((M, N), dtype),
+    ):
+        Y = T.alloc_buffer((M, N), dtype)
+        for i, j, k in T.grid(M, N, K):
+            with T.sblock("Y"):
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                with T.init():
+                    Y[vi, vj] = T.cast(T.float32(0), dtype)
+                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
+        for i, j in T.grid(M, N):
+            with T.sblock("C"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                C[vi, vj] = T.max(Y[vi, vj], T.cast(T.float32(0), dtype))
+
+
+######################################################################
+# Check the equivalence:
+
+print(tvm_ffi.structural_equal(ConciseModule, ConciseModuleFromPython))
+
+
+######################################################################
+# TensorIR Function with Dynamic Shapes
+# *************************************
+# Despite TVMScript not being executed by a Python interpreter, limited
+# interaction with Python is feasible. For instance, Python variables can
+# be used to ascertain the shape and data type of a TensorIR.
+
+
+@I.ir_module
+class DynamicShapeModule:
+    @T.prim_func(s_tir=True)
+    def mm_relu(a: T.handle, b: T.handle, c: T.handle):
+        # Dynamic shape definition
+        M = T.int32()
+        N = T.int32()
+        K = T.int32()
+
+        # Bind the input buffers with the dynamic shapes
+        A = T.match_buffer(a, [M, K], dtype)
+        B = T.match_buffer(b, [K, N], dtype)
+        C = T.match_buffer(c, [M, N], dtype)
+        Y = T.alloc_buffer((M, N), dtype)
+        for i, j, k in T.grid(M, N, K):
+            with T.sblock("Y"):
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                with T.init():
+                    Y[vi, vj] = T.cast(T.float32(0), dtype)
+                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
+        for i, j in T.grid(M, N):
+            with T.sblock("C"):
+                vi, vj = T.axis.remap("SS", [i, j])
+                C[vi, vj] = T.max(Y[vi, vj], T.cast(T.float32(0), dtype))
+
+
+######################################################################
+# Now let's check the runtime dynamic shape inference:
+
+
+def evaluate_dynamic_shape(lib: tvm.runtime.Module, m: int, n: int, k: int):
+    A = tvm.runtime.tensor(np.random.uniform(size=(m, k)).astype("float32"))
+    B = tvm.runtime.tensor(np.random.uniform(size=(k, n)).astype("float32"))
+    C = tvm.runtime.tensor(np.zeros((m, n), dtype="float32"))
+    lib(A, B, C)
+    return C.numpy()
+
+
+# Compile lib only once
+dyn_shape_lib = tvm.compile(DynamicShapeModule, target="llvm")
+# Able to handle different shapes
+print(evaluate_dynamic_shape(dyn_shape_lib, m=4, n=4, k=4))
+print(evaluate_dynamic_shape(dyn_shape_lib, m=64, n=64, k=128))
+
+######################################################################
+# Create TensorIR using Tensor Expression
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+# Often, the specifics of TensorIR are disregarded in favor of expressing the computation more
+# succinctly, leading to the pragmatic generation of TensorIR. This is where Tensor Expression
+# (TE) becomes relevant.
+#
+# Tensor Expression (TE) serves as a domain-specific language delineating a sequence of
+# computations through an expression-like API.
+#
+# .. note::
+#
+#   Tensor Expression comprises two components within the TVM stack: the expression and the
+#   schedule. The expression is the domain-specific language embodying the computation pattern,
+#   precisely what we're addressing in this section. Conversely, the TE schedule is the legacy
+#   scheduling method, has been superseded by the TensorIR schedule in the current TVM stack.
+#
+# Create Static-Shape Functions
+# *****************************
+# We use the same example of ``mm_relu`` from the last subsection to demonstrate the
+# TE creation method.
+
+from tvm import te
+
+A = te.placeholder((128, 128), "float32", name="A")
+B = te.placeholder((128, 128), "float32", name="B")
+k = te.reduce_axis((0, 128), "k")
+Y = te.compute((128, 128), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="Y")
+C = te.compute((128, 128), lambda i, j: te.max(Y[i, j], 0), name="C")
+
+######################################################################
+# Here ``te.compute`` takes the signature ``te.compute(output_shape, fcompute)``.
+# And the fcompute function describes how we want to compute the value of each
+# element ``Y[i, j]`` for a given index:
+#
+# .. code:: python
+#
+#   lambda i, j: te.sum(A[i, k] * B[k, j], axis=k)
+#
+# The aforementioned lambda expression encapsulates the computation:
+# :math:`Y_{i, j} = \sum_k A_{i, k} \times B_{k, j}`. Upon defining the computation,
+# we can formulate a TensorIR function by incorporating the pertinent parameters of interest.
+# In this specific instance, we aim to construct a function with two input parameters **A, B**
+# and one output parameter **C**.
+
+te_func = te.create_prim_func([A, B, C]).with_attr({"global_symbol": "mm_relu"})
+TEModule = tvm.IRModule({"mm_relu": te_func})
+TEModule.show()
+
+######################################################################
+# Create Dynamic-Shape Functions
+# ******************************
+# We can also create a dynamic-shape function using Tensor Expression. The only difference
+# is that we need to specify the shape of the input tensors as symbolic variables.
+
+# Declare symbolic variables
+M, N, K = te.var("m"), te.var("n"), te.var("k")
+A = te.placeholder((M, N), "float32", name="A")
+B = te.placeholder((K, N), "float32", name="B")
+k = te.reduce_axis((0, K), "k")
+Y = te.compute((M, N), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="Y")
+C = te.compute((M, N), lambda i, j: te.max(Y[i, j], 0), name="C")
+
+dyn_te_func = te.create_prim_func([A, B, C]).with_attr({"global_symbol": "mm_relu"})
+DynamicTEModule = tvm.IRModule({"mm_relu": dyn_te_func})
+DynamicTEModule.show()

docs/deep_dive/tensor_ir/tutorials/tir_transformation.py

@@ -0,0 +1,177 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E402
+
+"""
+.. _tirx-transform:
+
+Transformation
+--------------
+In this section, we will get to the main ingredients of the compilation flows -
+transformations of primitive tensor functions.
+"""
+
+######################################################################
+# In the :ref:`previous section <tirx-learning>`, we have given an example of how to write
+# ``mm_relu`` using TensorIR. In practice, there can be multiple ways to implement
+# the same functionality, and each implementation can result in different performance.
+#
+# .. note::
+#   This tutorial primarily illustrates the application of TensorIR Transformation,
+#   rather than delving into optimization techniques.
+#
+# First, let's take a look at the implementation of ``mm_relu`` in the previous section:
+
+import tvm
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+@I.ir_module
+class MyModule:
+    @T.prim_func(s_tir=True)
+    def main(
+        A: T.Buffer((128, 128), "float32"),
+        B: T.Buffer((128, 128), "float32"),
+        C: T.Buffer((128, 128), "float32"),
+    ):
+        T.func_attr({"tirx.noalias": True})
+        with T.sblock("root"):
+            T.reads()
+            T.writes()
+            Y = T.sblock_alloc_buffer((128, 128))
+            for i, j, k in T.grid(128, 128, 128):
+                with T.sblock("Y"):
+                    vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                    with T.init():
+                        Y[vi, vj] = T.float32(0)
+                    Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
+            for i, j in T.grid(128, 128):
+                with T.sblock("C"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
+
+
+######################################################################
+# Before we transform the function, let's first evaluate the performance of the
+# original implementation.
+
+import numpy as np
+
+a_np = np.random.uniform(size=(128, 128)).astype("float32")
+b_np = np.random.uniform(size=(128, 128)).astype("float32")
+c_np = a_np @ b_np
+
+a_nd = tvm.runtime.tensor(a_np)
+b_nd = tvm.runtime.tensor(b_np)
+c_nd = tvm.runtime.tensor(np.zeros((128, 128), dtype="float32"))
+
+
+def evaluate(mod: tvm.IRModule):
+    lib = tvm.tirx.build(mod, target="llvm")
+    # check correctness
+    lib(a_nd, b_nd, c_nd)
+    np.testing.assert_allclose(c_nd.numpy(), c_np, rtol=1e-5)
+    # evaluate performance
+    f_timer = lib.time_evaluator("main", tvm.cpu())
+    print(f_timer(a_nd, b_nd, c_nd))
+
+
+evaluate(MyModule)
+
+######################################################################
+# Initialization Schedule
+# ***********************
+# We initiate the process of code transformation by establishing a Schedule helper class,
+# utilizing the provided **MyModule** as input.
+
+sch = tvm.s_tir.Schedule(MyModule)
+
+######################################################################
+# Loop Tiling
+# ***********
+# Subsequently, we execute the requisite operations to acquire a reference to
+# block **Y** and its associated loops.
+
+block_Y = sch.get_sblock("Y")
+i, j, k = sch.get_loops(block_Y)
+
+######################################################################
+# We now proceed to execute the transformations. The initial modification involves
+# splitting loop ``j`` into two separate loops, with the inner loop possessing a
+# length of 8. It is crucial to understand that the transformation process is procedural;
+# thus, inadvertent execution of the block twice will yield an error stating the
+# non-existence of variable ``j``.
+
+j0, j1 = sch.split(j, factors=[None, 8])
+
+######################################################################
+# The outcome of the transformation can be examined, as it is retained within ``sch.mod``.
+
+sch.mod.show()
+
+######################################################################
+# Following the initial transformation phase, two supplementary loops, ``j_0`` and ``j_1``,
+# have been generated with respective ranges of 16 and 8. The subsequent
+# action involves reordering these two loops.
+
+sch.reorder(j0, k, j1)
+sch.mod.show()
+evaluate(sch.mod)
+
+######################################################################
+# Leverage Localities
+# *******************
+# Subsequently, we will execute two additional transformation steps to achieve a different
+# variant. First, we employ a primitive known as **reverse_compute_at** to relocate block
+# **C** to an inner loop of **Y**.
+
+block_C = sch.get_sblock("C")
+sch.reverse_compute_at(block_C, j0)
+sch.mod.show()
+
+######################################################################
+# Rewrite Reduction
+# *****************
+# Until now, the reduction initialization and update step have been maintained together
+# within a single block body. This amalgamated form facilitates loop transformations,
+# as the outer loops ``i``, ``j`` of initialization and updates generally need to remain
+# synchronized.
+#
+# Following the loop transformations, we can segregate the initialization of Y's elements
+# from the reduction update via the **decompose_reduction** primitive.
+
+sch.decompose_reduction(block_Y, k)
+sch.mod.show()
+evaluate(sch.mod)
+
+######################################################################
+# Trace the Transformation
+# ************************
+# TensorIR schedule is a procedural language, and the transformation is executed in a
+# step-by-step manner. We can trace the transformation by printing the schedule or the
+# history of the schedule.
+#
+# We've already see the schedule by printing ``sch.mod``. We can also print the history
+# of the schedule by ``sch.trace``.
+
+sch.trace.show()
+
+######################################################################
+# Alternatively, we can output the IRModule in conjunction with the historical trace.
+
+sch.show()