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