chore: import upstream snapshot with attribution

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.