chore: import upstream snapshot with attribution

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)