chore: import upstream snapshot with attribution

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()