chore: import upstream snapshot with attribution

python/tvm/relax/testing/__init__.py

@@ -0,0 +1,24 @@
+# isort: skip_file
+# 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.
+# pylint: disable=wildcard-import, redefined-builtin
+"""The Relax testing namespace containing nn and translator."""
+
+from .nn import *
+from .ast_printer import dump_ast
+from .matmul import *
+from .attention import *

python/tvm/relax/testing/ast_printer.py

@@ -0,0 +1,377 @@
+# 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.
+# pylint: disable=redefined-builtin, abstract-method, arguments-differ
+"""
+Utility script for printing Relax modules as AST diagrams,
+only intended to show how the AST is put together.
+It is not a pretty-printer and, in fact, is more of an ugly-printer,
+but it can be useful for tutorials and debugging.
+"""
+
+from collections.abc import Iterable
+
+import tvm
+from tvm import relax
+from tvm.ir.expr import Expr
+from tvm.relax import ExprFunctor
+
+
+def wrap_quotes(text: str) -> str:
+    """
+    Wraps the text in quotes.
+    """
+    return f'"{text}"'
+
+
+class ASTPrinter(ExprFunctor):
+    """
+    Class for recursing down ASTs and printing them in a very simple format,
+    mainly for instructive purposes and, perhaps, debugging.
+    """
+
+    def __init__(
+        self,
+        indent_str="    ",
+        include_ty_annotations=True,
+        include_call_attrs=True,
+    ):
+        self.indent_str = indent_str
+        self.include_ty_annotations = include_ty_annotations
+        self.include_call_attrs = include_call_attrs
+
+    def visit_expr(self, expr: relax.Expr) -> str:
+        # extend so we also dispatch to bindings and binding blocks,
+        # a little silly but IRFunctor hasn't been ported to Python
+        if isinstance(expr, relax.DataflowBlock):
+            return self.visit_dataflow_block_(expr)
+        if isinstance(expr, relax.BindingBlock):
+            return self.visit_binding_block_(expr)
+        if isinstance(expr, relax.Binding):
+            return self.visit_binding_(expr)
+        return super().visit_expr(expr)
+
+    def indent(self, text: str) -> str:
+        """
+        Indent all lines of the input.
+        """
+        if text == "":
+            return ""
+        lines = text.split("\n")
+        return self.indent_str + f"\n{self.indent_str}".join(lines)
+
+    def build_ast_node(self, nodename: str, force_newline=False, **kwargs: str) -> str:
+        """
+        Returns 'nodename(..., fields[i][0]=fields[i][1], ...)'
+        with appropriate indentation
+        """
+        return self.build_list(
+            map(lambda field: f"{field[0]}={field[1]}", kwargs.items()),
+            open_tok=f"{nodename}(",
+            close_tok=")",
+            force_newline=force_newline,
+        )
+
+    def build_expr(self, node: relax.Expr, nodename: str, force_newline=False, **kwargs: str):
+        """
+        Renders a Relax expression as a string using `build_ast_node`.
+        Handles whether to include the ty fields.
+        """
+        fields = kwargs.copy()
+        if not node.ty.is_missing() and self.include_ty_annotations:
+            fields["ty"] = self.visit_ty_(node.ty)
+        return self.build_ast_node(nodename, force_newline=force_newline, **fields)
+
+    def build_list(
+        self, members: Iterable[str], open_tok="[", close_tok="]", force_newline=False
+    ) -> str:
+        """
+        Builds a list of the members given, appropriately indented,
+        with each field on a line.
+        (special case: if there is only one field, then we do not put it on a new line
+        unless that field contains a newline or `force_newline` is set to true).
+        `open_tok` and `close_tok` are used to open and close the list, respectively.
+        """
+        mem_list = list(members)
+        if not mem_list:
+            return f"{open_tok}{close_tok}"
+        if len(mem_list) == 1 and not force_newline and "\n" not in mem_list[0]:
+            return f"{open_tok}{mem_list[0]}{close_tok}"
+        member_lines = ",\n".join(map(self.indent, mem_list))
+        return f"{open_tok}\n{member_lines}\n{close_tok}"
+
+    def visit_constant_(self, op: relax.Constant) -> str:
+        # simple rule of thumb: keep scalars inline, but anything larger goes on a new one
+        force_newline = len(op.data.shape) > 0
+        return self.build_expr(op, "Constant", force_newline=force_newline, data=str(op.data))
+
+    def visit_tuple_(self, op: relax.Tuple) -> str:
+        return self.build_expr(op, "Tuple", fields=self.build_list(map(self.visit_expr, op.fields)))
+
+    def visit_dataflow_var_(self, op: relax.DataflowVar) -> str:
+        return self.build_expr(op, "DataflowVar", name_hint=wrap_quotes(op.name_hint))
+
+    def visit_var_(self, op: relax.Var) -> str:
+        return self.build_expr(op, "Var", name_hint=wrap_quotes(op.name_hint))
+
+    def visit_shape_expr_(self, op: relax.ShapeExpr) -> str:
+        return self.build_expr(
+            op, "ShapeExpr", values=self.build_list(map(self.visit_prim_expr_field_, op.values))
+        )
+
+    def visit_extern_func_(self, op: relax.ExternFunc) -> str:
+        # ExternFunc does not inherit from relax.Expr either,
+        # so it doesn't have ty fields and we don't use build_expr
+        return self.build_ast_node("ExternFunc", global_symbol=wrap_quotes(op.global_symbol))
+
+    def visit_global_var_(self, op: relax.GlobalVar) -> str:
+        return self.build_expr(op, "GlobalVar", name_hint=wrap_quotes(op.name_hint))
+
+    def visit_function_(self, op: relax.Function) -> str:
+        fields = {
+            "params": self.build_list(map(self.visit_expr, op.params)),
+            "body": self.visit_expr(op.body),
+            "ret_ty": self.visit_ty_(op.ret_ty),
+            "is_pure": op.is_pure,
+        }
+        if op.attrs:
+            fields["attrs"] = self.build_list(
+                map(
+                    lambda kv: f"{wrap_quotes(str(kv[0]))}: {wrap_quotes(str(kv[1]))}",
+                    op.attrs.items(),
+                ),
+                open_tok="{",
+                close_tok="}",
+            )
+        return self.build_expr(op, "Function", **fields)
+
+    def visit_call_(self, op: relax.Call) -> str:
+        fields = {
+            "op": self.visit_expr(op.op),
+            "args": self.build_list(map(self.visit_expr, op.args)),
+        }
+        if op.ty_args:
+            fields["ty_args"] = self.build_list(map(self.visit_ty_, op.ty_args))
+        if op.attrs and self.include_call_attrs:
+
+            def display_attrs(attr_key):
+                attr_val = op.attrs[attr_key]
+
+                if isinstance(attr_val, str):
+                    # attrs can be strings but also other types;
+                    # we want to wrap strings in quotes
+                    # (__repr__ would work but it uses single quotes)
+                    attr_val = wrap_quotes(attr_val)
+                elif isinstance(attr_val, tvm.tirx.IntImm):
+                    if attr_val.dtype == "bool":
+                        attr_val = bool(attr_val.value)
+                    else:
+                        attr_val = int(attr_val.value)
+
+                return f"{wrap_quotes(attr_key)}: {attr_val}"
+
+            fields["attrs"] = self.build_list(
+                map(display_attrs, op.attrs.keys()),
+                open_tok="{",
+                close_tok="}",
+            )
+        return self.build_expr(op, "Call", **fields)
+
+    def visit_seq_expr_(self, op: relax.SeqExpr) -> str:
+        return self.build_expr(
+            op,
+            "SeqExpr",
+            blocks=self.build_list(map(self.visit_binding_block_, op.blocks)),
+            body=self.visit_expr(op.body),
+        )
+
+    def visit_if_(self, op: relax.If) -> str:
+        return self.build_expr(
+            op,
+            "If",
+            cond=self.visit_expr(op.cond),
+            true_branch=self.visit_expr(op.true_branch),
+            false_branch=self.visit_expr(op.false_branch),
+        )
+
+    def visit_string_imm_(self, op: relax.StringImm) -> str:
+        return self.build_expr(op, "StringImm", value=wrap_quotes(op.value))
+
+    def visit_data_type_imm_(self, op: relax.DataTypeImm) -> str:
+        return self.build_expr(op, "DataTypeImm", value=op.value)
+
+    def visit_op_(self, op: tvm.ir.Op) -> str:
+        # TODO: List other attributes?
+        # op is not actually a Relax expr and does not have
+        # ty fields, so we don't use build_expr here
+        return self.build_ast_node("Op", name=wrap_quotes(op.name))
+
+    def visit_prim_expr_field_(self, prim_expr: Expr) -> str:
+        # TODO: We may want to print Expr ASTs, but this is a simplification for now
+        return self.build_ast_node("Expr", value=f"`{prim_expr!s}`")
+
+    def visit_expr_fallback_(self, op: Expr) -> str:
+        if not tvm.ir.is_prim_expr(op):
+            raise ValueError(f"Invalid Relax expression {op} ({type(op)})")
+        return self.visit_prim_expr_field_(op)
+
+    def visit_tuple_getitem_(self, op: relax.TupleGetItem) -> str:
+        return self.build_expr(
+            op,
+            "TupleGetItem",
+            tuple_value=self.visit_expr(op.tuple_value),
+            index=str(op.index),
+        )
+
+    def visit_type_(self, type_node: relax.Type) -> str:
+        """
+        Recurse down types and print their ASTs too
+        """
+        if isinstance(type_node, relax.ShapeType):
+            return self.build_ast_node("ShapeType", ndim=str(type_node.ndim))
+        if isinstance(type_node, relax.AnyType):
+            return self.build_ast_node("AnyType")
+        if isinstance(type_node, relax.PackedFuncType):
+            return self.build_ast_node("PackedFuncType")
+        if isinstance(type_node, tvm.ir.PrimType):
+            return self.build_ast_node("PrimType", dtype=type_node.dtype)
+        if isinstance(type_node, relax.TensorType):
+            fields = {}
+            if type_node.ndim is not None:
+                fields["ndim"] = str(type_node.ndim)
+            if type_node.dtype != "":
+                fields["dtype"] = type_node.dtype
+            return self.build_ast_node("TensorType", **fields)
+        if isinstance(type_node, relax.TupleType):
+            return self.build_ast_node(
+                "TupleType", fields=self.build_list(map(self.visit_type_, type_node.fields))
+            )
+        if isinstance(type_node, relax.FuncType):
+            fields = {}
+            if type_node.params is not None:
+                fields["params"] = self.build_list(map(self.visit_type_, type_node.params))
+            fields["ret"] = self.visit_type_(type_node.ret)
+            fields["purity"] = bool(type_node.purity)
+            return self.build_ast_node(
+                "FuncType",
+                **fields,
+            )
+        raise ValueError(f"Invalid Relax Type {type_node} ({type(type_node)})")
+
+    def visit_ty_(self, ty_node: relax.Type) -> str:
+        """
+        Recurse down type and print their ASTs too
+        """
+        if isinstance(ty_node, relax.ShapeType):
+            fields = {}
+            fields["ndim"] = str(ty_node.ndim)
+            if ty_node.values is not None:
+                fields["values"] = self.build_list(map(self.visit_prim_expr_field_, ty_node.values))
+            return self.build_ast_node("ShapeType", **fields)
+        elif isinstance(ty_node, relax.AnyType):
+            return self.build_ast_node("AnyType")
+        elif isinstance(ty_node, tvm.ir.PrimType):
+            return self.build_ast_node("PrimType", dtype=ty_node.dtype)
+        elif isinstance(ty_node, relax.TensorType):
+            fields = {}
+            fields["dtype"] = ty_node.dtype
+            if ty_node.shape:
+                fields["shape"] = self.visit_expr(ty_node.shape)
+            else:
+                fields["ndim"] = str(ty_node.ndim)
+            return self.build_ast_node("TensorType", **fields)
+        elif isinstance(ty_node, relax.TupleType):
+            return self.build_ast_node(
+                "TupleType",
+                fields=self.build_list(map(self.visit_ty_, ty_node.fields)),
+            )
+        elif isinstance(ty_node, relax.FuncType):
+            fields = {}
+            if ty_node.params is not None:
+                fields["params"] = self.build_list(map(self.visit_ty_, ty_node.params))
+            fields["ret"] = self.visit_ty_(ty_node.ret)
+            fields["purity"] = bool(ty_node.purity)
+            return self.build_ast_node("FuncType", **fields)
+        else:
+            raise ValueError(f"Invalid Relax Type {ty_node} ({type(ty_node)})")
+
+    def visit_binding_block_(self, block: relax.BindingBlock) -> str:
+        """
+        Recurse down binding blocks
+        """
+        return self.build_ast_node(
+            "BindingBlock",
+            bindings=self.build_list(map(self.visit_binding_, block.bindings), force_newline=True),
+        )
+
+    def visit_dataflow_block_(self, block: relax.DataflowBlock) -> str:
+        """
+        Recurse down a dataflow block
+        """
+        return self.build_ast_node(
+            "DataflowBlock",
+            bindings=self.build_list(map(self.visit_binding_, block.bindings), force_newline=True),
+        )
+
+    def visit_binding_(self, binding: relax.Binding) -> str:
+        """
+        Distinguish between binding types
+        """
+        if isinstance(binding, relax.MatchCast):
+            return self.visit_match_cast_(binding)
+        if isinstance(binding, relax.VarBinding):
+            return self.visit_var_binding_(binding)
+        raise ValueError(f"Invalid binding type in {binding}: {type(binding)}")
+
+    def visit_match_cast_(self, match_cast: relax.MatchCast) -> str:
+        """
+        Handle match shape
+        """
+        fields = {
+            "var": self.visit_expr(match_cast.var),
+            "value": self.visit_expr(match_cast.value),
+            "ty": self.visit_ty_(match_cast.ty),
+        }
+        return self.build_ast_node("MatchCast", **fields)
+
+    def visit_var_binding_(self, var_binding: relax.VarBinding) -> str:
+        """
+        Handle ordinary var bindings
+        """
+        return self.build_ast_node(
+            "VarBinding",
+            var=self.visit_expr(var_binding.var),
+            value=self.visit_expr(var_binding.value),
+        )
+
+
+def dump_ast(
+    exp: relax.Expr,
+    indent_str="    ",
+    include_ty_annotations=True,
+    include_call_attrs=True,
+) -> str:
+    """
+    Dump an AST in a text format.
+    Can vary the indentation string and choose whether to include
+    type and shape annotations or call attributes.
+    """
+    printer = ASTPrinter(
+        indent_str=indent_str,
+        include_ty_annotations=include_ty_annotations,
+        include_call_attrs=include_call_attrs,
+    )
+    return printer.visit_expr(exp)

python/tvm/relax/testing/attention.py

@@ -0,0 +1,150 @@
+# 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 script for attention module."""
+
+import tvm
+from tvm.script import relax as R
+from tvm.script import tirx as T
+from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder import relax as relax_builder
+
+
+def get_relax_attention_module(
+    q_shape,
+    k_shape,
+    v_shape,
+    *,
+    dtype,
+    bias_shape=None,
+    qk_scale=None,
+    causal_mask=None,
+    window_size=None,
+):  # pylint: disable=too-many-arguments, too-many-locals, invalid-name
+    """Get a relax module for attention."""
+
+    if qk_scale is not None:
+        qk_scale = T.FloatImm("float32", qk_scale)
+
+    if window_size is not None:
+        window_size = T.IntImm("int32", window_size)
+
+    with IRBuilder() as builder:
+        with relax_builder.function():
+            R.func_name("main")
+            q = R.arg("q", R.Tensor(q_shape, dtype))
+            k = R.arg("k", R.Tensor(k_shape, dtype))
+            v = R.arg("v", R.Tensor(v_shape, dtype))
+            bias = None
+            if bias_shape is not None and bias_shape != "none":
+                bias = R.arg("bias", R.Tensor(bias_shape, dtype))
+
+            with R.dataflow() as frame:
+                result = R.emit(R.nn.attention(q, k, v, bias, qk_scale, causal_mask, window_size))
+                R.output(result)
+
+            R.func_ret_value(frame.output_vars[0])
+
+    func = builder.get()
+    return tvm.IRModule({"main": func})
+
+
+def get_relax_stacked_attention_module(
+    qkv,
+    b,
+    s,
+    n,
+    h,
+    h_v,
+    op,
+    bias=None,
+    qk_scale=None,
+    single_shape=False,
+    layout="BS3NH",
+):  # pylint: disable=too-many-arguments, too-many-locals, too-many-branches, invalid-name
+    # pylint: disable=too-many-statements
+    """Get a relax module for stacked attention."""
+    dtype = str(qkv.dtype)
+    assert layout in ["BS3NH", "SBN3H"]
+
+    if qk_scale is not None:
+        qk_scale = T.FloatImm("float32", qk_scale)
+
+    if single_shape:
+        if layout == "BS3NH":
+            qk_shape = R.shape([b, s, n, h])
+        elif layout == "SBN3H":
+            qk_shape = R.shape([b, s, n, h])
+        v_shape = qk_shape
+    else:
+        if layout == "BS3NH":
+            qk_shape = [b, s, n, h]
+            v_shape = [b, s, n, h_v]
+        elif layout == "SBN3H":
+            qk_shape = [s, b, n, h]
+            v_shape = [s, b, n, h_v]
+
+    if layout == "BS3NH":
+        split_axis = 2
+        split_sections = [n * h, n * h * 2]
+    elif layout == "SBN3H":
+        split_axis = 3
+        split_sections = [h, h * 2]
+
+    with IRBuilder() as builder:
+        with relax_builder.function():
+            R.func_name("main")
+            qkv = R.arg("qkv", R.Tensor(qkv.shape, dtype))
+            if bias is not None:
+                bias = R.arg("bias", R.Tensor(bias.shape, dtype))
+            with R.dataflow() as frame:
+                if op == "split":
+                    qkv_tuple = R.split(qkv, split_sections, axis=split_axis)
+                    q = qkv_tuple[0]
+                    k = qkv_tuple[1]
+                    v = qkv_tuple[2]
+                elif op == "strided_slice":
+                    q = R.strided_slice(qkv, [split_axis], [0], [split_sections[0]], [1])
+                    k = R.strided_slice(
+                        qkv, [split_axis], [split_sections[0]], [split_sections[1]], [1]
+                    )
+                    v = R.strided_slice(
+                        qkv,
+                        [split_axis],
+                        [split_sections[1]],
+                        [int(qkv.ty.shape[split_axis])],
+                        [1],
+                    )
+                else:
+                    raise NotImplementedError()
+                if layout == "BS3NH":
+                    q = R.reshape(q, qk_shape)
+                    k = R.reshape(k, qk_shape)
+                    v = R.reshape(v, v_shape)
+                elif layout == "SBN3H":
+                    q = R.permute_dims(q, [1, 0, 2, 3])
+                    k = R.permute_dims(k, [1, 0, 2, 3])
+                    v = R.permute_dims(v, [1, 0, 2, 3])
+                result = R.emit(R.nn.attention(q, k, v, bias, qk_scale))
+                if layout == "SBN3H":
+                    result = R.emit(R.permute_dims(result, [1, 0, 2, 3]))
+                R.output(result)
+
+            R.func_ret_value(frame.output_vars[0])
+
+    func = builder.get()
+    return tvm.IRModule({"main": func})

python/tvm/relax/testing/lib_comparator.py

@@ -0,0 +1,133 @@
+# 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.
+# pylint: disable=unused-argument
+"""Tools to compare libraries."""
+
+from collections.abc import Iterable
+
+import tvm
+import tvm.testing
+
+
+class LibCompareVMInstrument:
+    """Instrument class to compare libs.
+
+    This class build an instrument function that
+    pair tests an existing compiled relax vm implementation
+    and an extra module, which can sits in another backend
+    but offers a same subset of compiled TIR functions.
+
+    The instrumentation enables us to automatically
+    check and compare each ops being called in the pipeline
+    by looking up the same name in the provided mod and run testing.
+
+    Parameters
+    ----------
+    mod: runtime.Module
+        The module of interest to be validated.
+
+    device: runtime.Device
+        The device to run the target module on.
+
+    verbose: bool
+        Whether print out messages.
+
+    rtol: float
+        rtol used in validation
+
+    atol: float
+        atol used in validation
+    """
+
+    def __init__(self, mod, device, verbose=True, rtol=1e-5, atol=1e-5):
+        self.mod = mod
+        self.device = device
+        self.verbose = verbose
+        self.counter = 0
+        self.rtol = rtol
+        self.atol = atol
+
+    def compare(
+        self,
+        name: str,
+        ref_args: list[tvm.runtime.Tensor] | tuple[tvm.runtime.Tensor, ...],
+        new_args: list[tvm.runtime.Tensor] | tuple[tvm.runtime.Tensor, ...],
+        ret_indices: Iterable[int],
+    ):
+        """Comparison function, can be overloaded.
+
+        Parameters
+        ----------
+        name: str
+            Name of the function.
+
+        ref_args:
+            The reference arguments.
+
+        new_args:
+            The args to be passed to the comparison function.
+
+        ret_indices:
+            List of indices to validate return values.
+        """
+        my_func = self.mod.get_function(name, query_imports=True)
+        if self.verbose:
+            print(f"[{self.counter}] Validating {name} ...")
+        my_func(*new_args)
+        for rindex in ret_indices:
+            tvm.testing.assert_allclose(
+                new_args[rindex].numpy(), ref_args[rindex].numpy(), atol=self.atol, rtol=self.rtol
+            )
+        if self.verbose:
+            print(f"[{self.counter}] Validating {name}, passed.")
+        self.counter += 1
+
+    def skip_instrument(self, func, name, before_run, ret_val, *args):
+        return False
+
+    def __call__(self, func, name, before_run, ret_val, *args):
+        if before_run:
+            return
+        if name.startswith("vm.builtin."):
+            return
+        if any(not isinstance(x, tvm.runtime.Tensor) for x in args):
+            return
+        try:
+            self.mod.get_function(name, query_imports=True)
+        except AttributeError:
+            if self.verbose:
+                print(f"Cannot find {name}, skip...")
+            return
+
+        if self.skip_instrument(func, name, before_run, ret_val, *args):
+            return
+
+        new_args = []
+        # not always true, true for most ops.
+        ret_indices = (len(args) - 1,)
+        temp_args = []
+        for i, arg in enumerate(args):
+            arr = tvm.runtime.empty(arg.shape, arg.dtype, device=self.device)
+            # copy from cpu since we look at different device
+            if i not in ret_indices:
+                temp_cpu = arg.copyto(tvm.cpu())
+                temp_args.append(temp_cpu)
+                arr.copyfrom(temp_cpu)
+            new_args.append(arr)
+        # wait until all copy complete before we release temp_cpu
+        self.device.sync()
+        self.compare(name, args, new_args, ret_indices)

python/tvm/relax/testing/matmul.py

@@ -0,0 +1,65 @@
+# 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: RUF005
+"""Utilities to construct matmul workloads."""
+
+import tvm
+from tvm.script import relax as R
+from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder import relax as relax_builder
+
+
+def get_relax_matmul_module(
+    x_shape,
+    y_shape,
+    in_dtype,
+    out_dtype=None,
+    transposed_y=False,
+    bias_shape=None,
+    activation=None,
+    residual_bin_op=None,
+    residual_activation=None,
+):
+    """Create a matmul op followd by epilogue operations."""
+    out_dtype = out_dtype if out_dtype is not None else in_dtype
+    with IRBuilder() as builder:
+        with relax_builder.function():
+            R.func_name("main")
+            x = R.arg("x", R.Tensor(x_shape, in_dtype))
+            y = R.arg("y", R.Tensor(y_shape, in_dtype))
+            if bias_shape is not None:
+                bias = R.arg("bias", R.Tensor(bias_shape, out_dtype))
+
+            with R.dataflow() as frame:
+                if transposed_y:
+                    axes = list(range(len(y_shape) - 2)) + [-1, -2]
+                    y = R.emit(R.permute_dims(y, axes=axes))
+                result = R.emit(R.matmul(x, y, out_dtype=out_dtype))
+                if bias_shape is not None:
+                    result = R.emit(result + bias)
+                if activation is not None:
+                    result = R.emit(activation(result))
+                if residual_bin_op is not None:
+                    result = R.emit(residual_bin_op(result, x))
+                    if residual_activation is not None:
+                        result = R.emit(residual_activation(result))
+                R.output(result)
+
+            R.func_ret_value(frame.output_vars[0])
+
+    func = builder.get()
+    return tvm.IRModule({"main": func})

python/tvm/relax/testing/nn.py

@@ -0,0 +1,354 @@
+# 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.
+# pylint: disable=redefined-builtin, invalid-name
+# ruff: noqa: RUF005
+"""PyTorch-like nn.Module API for constructing workloads."""
+
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np  # type: ignore
+
+import tvm
+from tvm import relax, tirx, topi
+from tvm.relax.op.grad.grad import end_checkpoint, start_checkpoint
+
+
+def emit(expr: relax.Expr, name_hint: str = "") -> relax.Var:
+    return relax.BlockBuilder.current().emit(expr, name_hint=name_hint)
+
+
+def emit_te(func: Callable, *args: Any, **kwargs: Any) -> relax.Var:
+    return relax.BlockBuilder.current().emit_te(func, *args, **kwargs)
+
+
+def checkpoint(
+    func: Callable, *args: Any, **kwargs: Any
+) -> relax.Var | list[relax.Var] | list[Any]:
+    """Mark function(*args, **kwargs) should be computed in a checkpointed manner during backward.
+
+    To be specific, args and kwargs will be checkpointed, and func(*args, **kwargs) will be
+    recomputed in the backward stage.
+    """
+    args = [start_checkpoint(v) if isinstance(v, relax.Expr) else v for v in args]
+    kwargs = {k: start_checkpoint(v) if isinstance(v, relax.Expr) else v for k, v in kwargs.items()}
+    result = func(*args, **kwargs)
+    if isinstance(result, list | tuple):
+        result = [end_checkpoint(v) if isinstance(v, relax.Expr) else v for v in result]
+    else:
+        assert isinstance(result, relax.Expr)
+        result = end_checkpoint(result)
+    return result
+
+
+def emit_checkpoint(
+    func: Callable, *args: Any, **kwargs: Any
+) -> relax.Var | list[relax.Var] | list[Any]:
+    """Mark function(*args, **kwargs) should be computed in a checkpointed manner during backward.
+
+    To be specific, args and kwargs will be checkpointed, and func(*args, **kwargs) will be
+    recomputed in the backward stage.
+
+    This interface will additionally emit the exprs marked with start_checkpoint() and
+    end_checkpoint() with suffix "_scp" and "_ecp" respectively, for easily understanding the
+    result tvmscript.
+    """
+    bb = relax.BlockBuilder.current()
+    args = [
+        bb.emit(start_checkpoint(v), v.name_hint + "_scp") if isinstance(v, relax.Var) else v
+        for v in args
+    ]
+    kwargs = {
+        k: bb.emit(start_checkpoint(v), v.name_hint + "_scp") if isinstance(v, relax.Var) else v
+        for k, v in kwargs.items()
+    }
+    result = func(*args, **kwargs)
+    if isinstance(result, list | tuple):
+        result = list(result)
+        for i, v in enumerate(result):
+            if isinstance(v, relax.Expr):
+                if not isinstance(v, relax.Var):
+                    v = bb.emit(v)
+                result[i] = bb.emit(end_checkpoint(v), v.name_hint + "_ecp")
+    else:
+        assert isinstance(result, relax.Expr)
+        result_emit = bb.emit(result)
+        result = bb.emit(end_checkpoint(result_emit), result_emit.name_hint + "_ecp")
+
+    return result
+
+
+def emit_checkpoint_sequential(
+    functions: list[Callable],
+    segments: int | list[int],
+    input: relax.Var,
+    checkpoint_last: bool = False,
+) -> relax.Var | list[relax.Var] | list[Any]:
+    """A helper function for checkpointing sequential models. This interface has similar purpose
+    as torch.utils.checkpoint.checkpoint_sequential.
+
+    Sequential models execute a list of modules/functions in order (sequentially). Therefore, we
+    can divide such a model in various segments and checkpoint each segment. By default, we will
+    checkpoint all segments except the last, meaning their inputs will be saved from the forward
+    stage and they will be recomputed in the backward stage.
+
+    Parameters
+    ----------
+    functions : List[Callable]
+        The list of functions to be executed sequentially.
+
+    segments : Union[int, List[int]]
+        The segments. If segments is int `n`, functions will be evenly divided into `n` segments;
+        if segments is a list of ints, it marks the start of every segment.
+
+    input : relax.Var
+        The input of the first function.
+
+    checkpoint_last : bool
+        Whether the last segment will be checkpointed. Default: False
+
+    Returns
+    -------
+    output : Union[relax.Var, List[relax.Var], List[Any]]
+        The emited output of the last function.
+    """
+    bb = relax.BlockBuilder.current()
+
+    def run_function(start, end, functions):
+        def forward(input):
+            for j in range(start, end):
+                input = functions[j](input)
+            return input
+
+        return forward
+
+    n = len(functions)
+    if not isinstance(segments, list):
+        segments = list(range(0, n, n // segments)) + [n]
+    if segments[-1] != n:
+        segments = segments + [n]
+
+    assert len(segments) >= 2
+
+    for i in range(len(segments) - 1):
+        if i == len(segments) - 2 and not checkpoint_last:
+            input = run_function(segments[i], segments[i + 1], functions)(input)
+        else:
+            input = emit_checkpoint(run_function(segments[i], segments[i + 1], functions), input)
+
+    assert isinstance(input, relax.Expr)
+    if not isinstance(input, relax.Var):
+        input = bb.emit(input)
+    return input
+
+
+def _try_unique_name(name: str):
+    """Attempt to uniquify the name
+
+    If a `relax.BlockBuilder` is active, use it to return a unique
+    name.  Otherwise, return the name itself.
+
+    Two distinct variables in Relax may have identical names.
+    However, for user readability, it is convenient to have all names
+    be unique within a Relax function.  If a Placeholder or Parameter
+    is defined within an active `relax.BlockBuilder`, that context may
+    be used to provide a unique name.  Otherwise, allow the duplicate
+    names.
+
+    Parameters
+    ----------
+    name: str
+
+        The variable name
+
+    Returns
+    -------
+    updated_name: str
+
+        The updated variable name
+
+
+    """
+    block_builder = relax.BlockBuilder.current()
+    if block_builder is None:
+        return name
+    else:
+        return block_builder.get_unique_name(name)
+
+
+class Placeholder(relax.Var):
+    """A placeholder variable that can represent model input."""
+
+    def __init__(self, shape: list[Any] | tuple[Any, ...], dtype="float32", name="data"):
+        if not isinstance(shape, list | tuple):
+            raise TypeError("the shape of Placeholder is expected to be a list or a tuple")
+        super().__init__(_try_unique_name(name), relax.TensorType(shape, dtype))
+
+
+class Parameter(relax.Var):
+    """A special kind of relax Var that represents model parameter(weight)."""
+
+    def __init__(self, shape: list[Any] | tuple[Any, ...], dtype="float32", name="param"):
+        if not isinstance(shape, list | tuple):
+            raise TypeError("the shape of Parameter is expected to be a list or a tuple")
+
+        super().__init__(_try_unique_name(name), relax.TensorType(shape, dtype))
+
+
+class Module(tvm.relax.frontend.nn.SubroutineMixin):
+    """Base class for all model modules.
+
+    A neural network or a layer can subclass this class.
+
+    By default, calls into this module will generate the `relax.Expr`
+    representing the output within the current function body.  Setting
+    the variable "define_subrouine" to True; either at the
+    `nn.Module`, subclass, or instance level; will instead produce a
+    subroutine within the same module, which is then called within the
+    current function body.
+
+    Example
+    -------
+    .. code-block:: python
+
+        # Define a linear layer
+        class Linear(Module)
+            def __init__(self, in_features, out_features, bias=True):
+                self.in_features = in_features
+                self.out_features = out_features
+                self.weight = Parameter((in_features, out_features), name="linear_weight")
+                if bias:
+                    self.bias = Parameter((out_features,), name="linear_bias")
+                else:
+                    self.bias = None
+
+            # All submodules should implement forward.
+            # Defines the forward computation performed at every call.
+            def forward(self, input: relax.Expr) -> relax.Var:
+                y = emit_te(topi.matmul, input, self.weight)
+                if self.bias is not None:
+                    y = emit_te(topi.add, y, self.bias)
+                return y
+    """
+
+    define_subroutine: bool = False
+
+    def parameters(self) -> list[Parameter]:
+        """Return the list of parameters in the module."""
+        return _unpack_params(self.__dict__)
+
+    def forward(self, input: relax.Expr):
+        """Define the computation performed at every call."""
+        raise NotImplementedError()
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
+
+
+def _unpack_params(value: object) -> list[relax.Var]:
+    if isinstance(value, Parameter):
+        return [value]
+    if isinstance(value, Module):
+        return value.parameters()
+    if isinstance(value, dict):
+        params = []
+        for v in value.values():
+            params += _unpack_params(v)
+        return params
+    if isinstance(value, list | tuple):
+        params = []
+        for v in value:
+            params += _unpack_params(v)
+        return params
+    return []
+
+
+def init_params(mod: tvm.IRModule) -> list[tvm.runtime.Tensor]:
+    """Utility function to initialize model's parameters."""
+    shape_dict = {v.name_hint: v.ty.shape for v in mod["main"].params}
+    params = []
+    for k, v in shape_dict.items():
+        if k.startswith("data"):
+            continue
+        if isinstance(v, relax.ShapeExpr):
+            shape = []
+            for i in v:
+                if isinstance(i, tirx.IntImm):
+                    shape.append(int(i))
+                else:
+                    raise TypeError("cannot initialize for unknown-shape parameters.")
+            params.append(tvm.runtime.tensor(np.zeros(shape).astype(np.float32)))
+        else:
+            raise TypeError("cannot initialize for unknown-shape parameters.")
+    return params
+
+
+class Sequential(Module):
+    """A sequential container that concatenates modules in it.
+
+    Example
+    -------
+    .. code-block:: python
+
+        model = nn.Sequential(
+                    nn.Conv2d(1, 20, 5),
+                    nn.ReLU(),
+                    nn.Conv2d(20, 64, 5),
+                    nn.ReLU()
+                )
+    """
+
+    def __init__(self, *modules: Module):
+        self.modules = modules
+
+    def forward(self, input: relax.Expr) -> relax.Var:
+        for module in self.modules:
+            input = module(input)
+        return input
+
+
+class ReLU(Module):
+    """Applies the rectified linear unit activation function on the input."""
+
+    def forward(self, input: relax.Expr) -> relax.Var:
+        return emit_te(topi.nn.relu, input)
+
+
+class LogSoftmax(Module):
+    """Applies log softmax activation function on the input."""
+
+    def forward(self, input: relax.Expr) -> relax.Var:
+        return emit_te(topi.nn.log_softmax, input)
+
+
+class Linear(Module):
+    """Applies a linear transformation to the input data: :math:`y = xA + b`."""
+
+    def __init__(self, in_features, out_features, bias=True):
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = Parameter((in_features, out_features), name="linear_weight")
+        if bias:
+            self.bias = Parameter((out_features,), name="linear_bias")
+        else:
+            self.bias = None
+
+    def forward(self, input: relax.Expr) -> relax.Var:
+        y = emit_te(topi.matmul, input, self.weight)
+        if self.bias is not None:
+            y = emit_te(topi.add, y, self.bias)
+        return y

python/tvm/relax/testing/runtime_builtin.py

@@ -0,0 +1,37 @@
+# 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.
+
+# pylint: disable=invalid-name
+"""Testing utilities for runtime builtin functions."""
+
+from enum import IntEnum
+
+
+class MatchShapeCode(IntEnum):
+    """Code passed to match shape builtin"""
+
+    ASSERT_EQUAL_TO_IMM = 0
+    STORE_TO_HEAP = 1
+    NO_OP = 2
+    ASSERT_EQUAL_TO_LOAD = 3
+
+
+class MakeShapeCode(IntEnum):
+    """Code passed to match shape builtin"""
+
+    USE_IMM = 0
+    LOAD_SHAPE = 1

python/tvm/relax/testing/transform.py

@@ -0,0 +1,125 @@
+# 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.
+# pylint: disable=unused-argument, invalid-name, no-else-return, abstract-method, arguments-differ
+"""Relax transformation passes for testing"""
+
+import logging
+import os
+
+import tvm_ffi
+
+import tvm
+from tvm.ir import Call
+from tvm.ir.module import IRModule
+from tvm.relax.expr import DataflowBlock, Var
+from tvm.runtime import Object
+
+
+def ApplyEmptyCppMutator() -> tvm.ir.transform.Pass:
+    """Create empty cpp mutator"""
+    packed_func = tvm.get_global_func("relax.testing.transform.ApplyEmptyCppMutator")
+    return packed_func()
+
+
+def dataflow_liveness_analysis(block: DataflowBlock) -> dict[Var, tuple[int, int]]:
+    """
+    Inner function for the dataflow inplace transformation exposed for testing.
+    """
+    if "PYTEST_CURRENT_TEST" not in os.environ:
+        logging.warning("The function dataflow_liveness_analysis is exposed for testing only.")
+
+    live_ranges = tvm.get_global_func("relax.testing.transform.DataflowLivenessAnalysis")(block)  # type: ignore
+    ret = {}
+    for var, live_range in live_ranges.items():
+        ret[var] = tuple(live_range)
+    return ret  # type: ignore
+
+
+def dataflow_alias_analysis(
+    block: DataflowBlock, inputs: list[Var]
+) -> tuple[dict[Var, set[int]], dict[int, list[set[int]]]]:
+    """
+    Inner function for the dataflow inplace transformation exposed for testing.
+    """
+    if "PYTEST_CURRENT_TEST" not in os.environ:
+        logging.warning("The function dataflow_alias_analysis is exposed for testing only.")
+
+    alias_sets, tuple_map = tvm.get_global_func("relax.testing.transform.DataflowAliasAnalysis")(
+        block,
+        inputs,
+    )  # type: ignore
+    res_alias_sets = {}
+    res_tuple_map = {}
+    for var, alias_set in alias_sets.items():
+        res_alias_sets[var] = set(alias_set)
+    for idx, elem_alias_sets in tuple_map.items():
+        res_tuple_map[idx] = [set(alias_set) for alias_set in elem_alias_sets]
+    return res_alias_sets, res_tuple_map  # type: ignore
+
+
+@tvm_ffi.register_object("relax.transform.InplaceOpportunity")
+class InplaceOpportunity(Object):
+    """
+    Represents an opportunity to make a binding in-place. Exposed only for testing;
+    the constructor is not exposed.
+
+    Parameters:
+    -----------
+    binding_idx: int
+        Index of the binding within its block
+
+    arg_idxs: List[int]
+        Indices of arguments that are eligible to be used as in-place targets.
+    """
+
+    def __init__(self, _binding_idx, _arg_idxs):
+        raise NotImplementedError("Constructor for InplaceOpportunity not exposed!")
+
+
+def dataflow_inplace_analysis(
+    block: DataflowBlock, inputs: list[Var], mod: IRModule
+) -> tuple[list[tuple[int, set[int]]], list[tuple[int, set[int]]]]:
+    """
+    Inner function for the dataflow inplace transformation exposed for testing.
+    """
+    if "PYTEST_CURRENT_TEST" not in os.environ:
+        logging.warning("The function dataflow_inplace_analysis is exposed for testing only.")
+    index_lists = tvm.get_global_func("relax.testing.transform.DataflowInplaceAnalysis")(
+        block, inputs, mod
+    )  # type: ignore
+
+    def convert(opp_list):
+        return list(map(lambda opp: (int(opp.binding_idx), set(map(int, opp.arg_idxs))), opp_list))
+
+    return (convert(index_lists[0]), convert(index_lists[1]))  # type: ignore
+
+
+def dataflow_single_inplace_call(
+    mod: IRModule, call: Call, inplace_indices: list[int]
+) -> tuple[Call, IRModule]:
+    """
+    Inner function for the dataflow inplace transformation exposed for testing.
+    """
+    if "PYTEST_CURRENT_TEST" not in os.environ:
+        logging.warning("The function dataflow_single_inplace_call is exposed for testing only.")
+
+    ret = tvm.get_global_func("relax.testing.transform.SingleInplaceCall")(
+        mod,
+        call,
+        inplace_indices,
+    )  # type: ignore
+    return (ret[0], ret[1])  # type: ignore

python/tvm/relax/testing/vm.py

@@ -0,0 +1,92 @@
+# 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.
+# pylint: disable=invalid-name
+"""Testing utilities for relax VM"""
+
+from typing import Any
+
+import numpy as np  # type: ignore
+
+import tvm
+from tvm import relax
+from tvm.runtime import Object
+
+
+@tvm.register_global_func("test.vm.move")
+def move(src):
+    return src
+
+
+@tvm.register_global_func("test.vm.add")
+def add(a, b):
+    ret = a.numpy() + b.numpy()
+    return tvm.runtime.tensor(ret)
+
+
+@tvm.register_global_func("test.vm.mul")
+def mul(a, b):
+    ret = a.numpy() * b.numpy()
+    return tvm.runtime.tensor(ret)
+
+
+@tvm.register_global_func("test.vm.equal_zero")
+def equal_zero(a):
+    ret = np.all(a.numpy() == 0)
+    return tvm.runtime.tensor(ret)
+
+
+@tvm.register_global_func("test.vm.subtract_one")
+def subtract_one(a):
+    ret = np.subtract(a.numpy(), 1)
+    return tvm.runtime.tensor(ret)
+
+
+@tvm.register_global_func("test.vm.identity")
+def identity_packed(a, b):
+    b[:] = tvm.runtime.tensor(a.numpy())
+
+
+@tvm.register_global_func("test.vm.tile")
+def tile_packed(a, b):
+    b[:] = tvm.runtime.tensor(np.tile(a.numpy(), (1, 2)))
+
+
+@tvm.register_global_func("test.vm.add_scalar")
+def add_scalar(a, b):
+    return a + b
+
+
+@tvm.register_global_func("test.vm.get_device_id")
+def get_device_id(device):
+    return device.index
+
+
+def check_saved_func(vm: relax.VirtualMachine, func_name: str, *inputs: list[Any]) -> Object:
+    # uses save_function to create a closure with the given inputs
+    # and ensure the result is the same
+    # (assumes the functions return tensors and that they're idempotent)
+    saved_name = f"{func_name}_saved"
+    vm.save_function(func_name, saved_name, *inputs)
+    res1 = vm[func_name](*inputs)
+    res2 = vm[saved_name]()
+    tvm.testing.assert_allclose(res1.numpy(), res2.numpy(), rtol=1e-7, atol=1e-7)
+    return res1
+
+
+@tvm.register_global_func("test.vm.check_if_defined")
+def check_if_defined(obj: tvm.Object) -> tvm.tirx.IntImm:
+    return tvm.runtime.convert(obj is not None)