This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch unity-staging in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 869826ecb5fff6fdcfc2ee2153db1e31e4bc47f4 Author: Junru Shao <[email protected]> AuthorDate: Fri Feb 10 07:41:11 2023 -0800 [Unity] Relax TVMScript Printer (#13944) This PR introduces Relax as a dialect supported by the TVMScript Printer. Some caveats: - Needs to rebase to mainline before merging. - Some tests are skiped because some operators are not upstreamed to the unity branch yet. Co-authored-by: Tianqi Chen <[email protected]> Co-authored-by: Yuchen Jin <[email protected]> Co-authored-by: Steven S. Lyubomirsky <[email protected]> Co-authored-by: Yong Wu <[email protected]> Co-authored-by: Prakalp Srivastava <[email protected]> Co-authored-by: Sunghyun Park <[email protected]> Co-authored-by: Ruihang Lai <[email protected]> Co-authored-by: Hongyi Jin <[email protected]> Co-authored-by: Bohan Hou <[email protected]> Co-authored-by: Siyuan Feng <[email protected]> --- python/tvm/relax/expr.py | 53 +-- python/tvm/script/parser/core/entry.py | 1 - src/relax/ir/expr.cc | 45 -- src/relax/ir/struct_info.cc | 43 -- src/script/printer/relax/binding.cc | 87 ++++ src/script/printer/relax/call.cc | 212 +++++++++ src/script/printer/relax/expr.cc | 136 ++++++ src/script/printer/relax/function.cc | 78 ++++ src/script/printer/relax/region.cc | 100 +++++ src/script/printer/relax/struct_info.cc | 149 +++++++ src/script/printer/relax/tir.cc | 89 ++++ src/script/printer/relax/type.cc | 89 ++++ src/script/printer/relax/utils.h | 101 +++++ tests/python/relax/test_tvmscript_printer_relax.py | 488 +++++++++++++++++++++ 14 files changed, 1541 insertions(+), 130 deletions(-) diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py index 138724ed06..f1cf815d8e 100644 --- a/python/tvm/relax/expr.py +++ b/python/tvm/relax/expr.py @@ -22,16 +22,18 @@ from numbers import Number from typing import Any, Callable, Dict, List, Optional, Union import numpy as _np # type: ignore + import tvm import tvm._ffi -import tvm.relax import tvm.ir +import tvm.relax from tvm import DataType from tvm._ffi import base as _base -from tvm.runtime import ndarray as _nd, Object +from tvm.runtime import Object +from tvm.runtime import ndarray as _nd from ..ir import BaseFunc, Node, SourceName, Span -from ..runtime import String +from ..runtime import Scriptable, String from ..tir import PrimExpr from . import _ffi_api @@ -55,7 +57,7 @@ class Id(Object): # NOTE: place base struct info in expr to avoid cyclic dep # from expr to struct info. -class StructInfo(Node): +class StructInfo(Node, Scriptable): """The base class of all StructInfo. StructInfo contains both the static type @@ -110,7 +112,7 @@ def _binary_rhs_helper(rhs: "ExprWithOp") -> "ExprWithOp": raise TypeError(f"type {type(rhs)} not supported") -class ExprWithOp(Expr): +class ExprWithOp(Expr, Scriptable): """Basetype of all relax expressions that defines op overloading.""" def astype(self, dtype: Union[str, DataType]) -> "ExprWithOp": @@ -436,7 +438,7 @@ class DataflowVar(Var): @tvm._ffi.register_object("relax.expr.PrimValue") -class PrimValue(Expr): +class PrimValue(Expr, Scriptable): """The prim expr representing the value.""" value: PrimExpr @@ -448,7 +450,7 @@ class PrimValue(Expr): @tvm._ffi.register_object("relax.expr.StringImm") -class StringImm(Expr): +class StringImm(Expr, Scriptable): """Represent a string literal constant.""" value: str @@ -458,7 +460,7 @@ class StringImm(Expr): @tvm._ffi.register_object("relax.expr.DataTypeImm") -class DataTypeImm(Expr): +class DataTypeImm(Expr, Scriptable): """Represent a data type constant.""" value: DataType @@ -468,11 +470,9 @@ class DataTypeImm(Expr): @tvm._ffi.register_object("relax.expr.Binding") -class Binding(Node): +class Binding(Node, Scriptable): """The base class of a binding in Relax.""" - ... - @tvm._ffi.register_object("relax.expr.MatchCast") class MatchCast(Binding): @@ -548,7 +548,7 @@ class SeqExpr(ExprWithOp): @tvm._ffi.register_object("relax.expr.Function") -class Function(BaseFunc): +class Function(BaseFunc, Scriptable): """A Relax function.""" params: List[Var] @@ -588,35 +588,6 @@ class Function(BaseFunc): """ return Call(self, args, None, None) - def script(self, show_meta: bool = False) -> str: - """Print relax.Function into TVMScript - - Parameters - ---------- - show_meta : bool - Whether to show meta information - - Returns - ------- - script : str - The TVM Script of the relax.Function - """ - return tvm._ffi.get_global_func("script.AsRelaxScript")(self, show_meta) # type: ignore - - def show(self, style: str = "light") -> None: - """ - A sugar for print highlighted TVM script. - - Parameters - ---------- - style : str, optional - Pygments styles extended by "light" (default) and "dark", by default "light" - """ - from tvm.script.highlight import cprint # pylint: disable=import-outside-toplevel - - # Use deferred import to avoid circular import while keeping cprint under tvm/script - cprint(self, style=style) - @tvm._ffi.register_object("relax.expr.ExternFunc") class ExternFunc(BaseFunc): diff --git a/python/tvm/script/parser/core/entry.py b/python/tvm/script/parser/core/entry.py index 3c01b54a9f..d8a11f5b46 100644 --- a/python/tvm/script/parser/core/entry.py +++ b/python/tvm/script/parser/core/entry.py @@ -54,7 +54,6 @@ def parse(program: Union[doc.AST, Any, str], extra_vars: Dict[str, Any] = None) "tir": tir, "relax": relax, "R": relax, - "tvm": tvm, } source = Source(program) diff --git a/src/relax/ir/expr.cc b/src/relax/ir/expr.cc index 45868a488a..a0aaea886d 100644 --- a/src/relax/ir/expr.cc +++ b/src/relax/ir/expr.cc @@ -94,13 +94,6 @@ TVM_REGISTER_GLOBAL("relax.Call") .set_body_typed([](Expr op, Array<Expr> args, Attrs attrs, Array<StructInfo> sinfo_args, Span span) { return Call(op, args, attrs, sinfo_args, span); }); -TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) - .set_dispatch<CallNode>([](const ObjectRef& ref, ReprPrinter* p) { - auto* node = static_cast<const CallNode*>(ref.get()); - p->stream << "CallNode(" << node->op << ", " << node->args << ", " << node->attrs << ", " - << node->sinfo_args << ")"; - }); - If::If(Expr cond, Expr true_branch, Expr false_branch, Span span) { ObjectPtr<IfNode> n = make_object<IfNode>(); n->cond = std::move(cond); @@ -137,13 +130,6 @@ TVM_REGISTER_GLOBAL("relax.If") return If(cond, true_branch, false_branch, span); }); -TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) - .set_dispatch<IfNode>([](const ObjectRef& ref, ReprPrinter* p) { - auto* node = static_cast<const IfNode*>(ref.get()); - p->stream << "IfNode(" << node->cond << ", " << node->true_branch << ", " - << node->false_branch << ")"; - }); - Tuple::Tuple(tvm::Array<relay::Expr> fields, Span span) { ObjectPtr<TupleNode> n = make_object<TupleNode>(); n->fields = std::move(fields); @@ -179,12 +165,6 @@ Tuple WithFields(Tuple tuple, Optional<Array<Expr>> opt_fields, Optional<Span> o return tuple; } -TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) - .set_dispatch<TupleNode>([](const ObjectRef& ref, ReprPrinter* p) { - auto* node = static_cast<const TupleNode*>(ref.get()); - p->stream << "Tuple(" << node->fields << ")"; - }); - TupleGetItem::TupleGetItem(Expr tuple, int index, Span span) { ObjectPtr<TupleGetItemNode> n = make_object<TupleGetItemNode>(); n->tuple = std::move(tuple); @@ -216,12 +196,6 @@ TVM_REGISTER_GLOBAL("relax.TupleGetItem").set_body_typed([](Expr tuple, int inde return TupleGetItem(tuple, index); }); -TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) - .set_dispatch<TupleGetItemNode>([](const ObjectRef& ref, ReprPrinter* p) { - auto* node = static_cast<const TupleGetItemNode*>(ref.get()); - p->stream << "TupleGetItemNode(" << node->tuple << ", " << node->index << ")"; - }); - TVM_REGISTER_NODE_TYPE(ShapeExprNode); ShapeExpr::ShapeExpr(Array<PrimExpr> values, Span span) { @@ -245,19 +219,6 @@ TVM_REGISTER_GLOBAL("relax.ShapeExpr").set_body_typed([](Array<PrimExpr> values, return ShapeExpr(values, span); }); -TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) - .set_dispatch<ShapeExprNode>([](const ObjectRef& ref, ReprPrinter* p) { - const ShapeExprNode* node = static_cast<const ShapeExprNode*>(ref.get()); - p->stream << "ShapeExpr("; - for (auto it = node->values.begin(); it != node->values.end(); it++) { - if (it != node->values.begin()) { - p->stream << ", "; - } - p->stream << *it; - } - p->stream << ")"; - }); - TVM_REGISTER_NODE_TYPE(VarNode); Var::Var(Id vid, Optional<StructInfo> struct_info_annotation, Span span) { @@ -572,12 +533,6 @@ TVM_REGISTER_GLOBAL("relax.ExternFunc").set_body_typed([](String global_symbol, return ExternFunc(global_symbol, span); }); -TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) - .set_dispatch<ExternFuncNode>([](const ObjectRef& ref, ReprPrinter* p) { - const auto* node = static_cast<const ExternFuncNode*>(ref.get()); - p->stream << "ExternFunc(\"" << node->global_symbol << "\")"; - }); - Expr GetShapeOf(const Expr& expr) { // default case, to be normalized. ICHECK(expr->struct_info_.defined()) << "GetShapeOf can only be applied to normalized expr"; diff --git a/src/relax/ir/struct_info.cc b/src/relax/ir/struct_info.cc index 9db7cea672..4004ad28d5 100644 --- a/src/relax/ir/struct_info.cc +++ b/src/relax/ir/struct_info.cc @@ -41,11 +41,6 @@ TVM_REGISTER_GLOBAL("relax.ObjectStructInfo").set_body_typed([](Span span) { return ObjectStructInfo(span); }); -TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) - .set_dispatch<ObjectStructInfoNode>([](const ObjectRef& ref, ReprPrinter* p) { - p->stream << "ObjectStructInfo()"; - }); - // Prim PrimStructInfo::PrimStructInfo(DataType dtype, Span span) { ObjectPtr<PrimStructInfoNode> n = make_object<PrimStructInfoNode>(); @@ -60,12 +55,6 @@ TVM_REGISTER_GLOBAL("relax.PrimStructInfo").set_body_typed([](DataType dtype, Sp return PrimStructInfo(dtype, span); }); -TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) - .set_dispatch<PrimStructInfoNode>([](const ObjectRef& ref, ReprPrinter* p) { - const auto* node = static_cast<const PrimStructInfoNode*>(ref.get()); - p->stream << "PrimStructInfo(" << node->dtype << ")"; - }); - // Shape ShapeStructInfo::ShapeStructInfo(Array<PrimExpr> values, Span span) { ObjectPtr<ShapeStructInfoNode> n = make_object<ShapeStructInfoNode>(); @@ -102,16 +91,6 @@ TVM_REGISTER_GLOBAL("relax.ShapeStructInfo") } }); -TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) - .set_dispatch<ShapeStructInfoNode>([](const ObjectRef& ref, ReprPrinter* p) { - const auto* node = static_cast<const ShapeStructInfoNode*>(ref.get()); - if (node->values.defined()) { - p->stream << "ShapeStructInfo(" << node->values.value() << ")"; - } else { - p->stream << "ShapeStructInfo(ndim=" << node->ndim << ")"; - } - }); - // Tensor TensorStructInfo::TensorStructInfo(Expr shape, DataType dtype, Span span) { ObjectPtr<TensorStructInfoNode> n = make_object<TensorStructInfoNode>(); @@ -150,16 +129,6 @@ TVM_REGISTER_GLOBAL("relax.TensorStructInfo") } }); -TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) - .set_dispatch<TensorStructInfoNode>([](const ObjectRef& ref, ReprPrinter* p) { - const auto* node = static_cast<const TensorStructInfoNode*>(ref.get()); - if (node->shape.defined()) { - p->stream << "TensorStructInfo(" << node->shape.value() << ", " << node->dtype << ")"; - } else { - p->stream << "TensorStructInfo(" << node->dtype << ", ndim=" << node->ndim << ")"; - } - }); - // Tuple TupleStructInfo::TupleStructInfo(Array<StructInfo> fields, Span span) { ObjectPtr<TupleStructInfoNode> n = make_object<TupleStructInfoNode>(); @@ -175,12 +144,6 @@ TVM_REGISTER_GLOBAL("relax.TupleStructInfo") return TupleStructInfo(fields, span); }); -TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) - .set_dispatch<TupleStructInfoNode>([](const ObjectRef& ref, ReprPrinter* p) { - const auto* node = static_cast<const TupleStructInfoNode*>(ref.get()); - p->stream << "TupleStructInfo(" << node->fields << ")"; - }); - // Func FuncStructInfo::FuncStructInfo(Array<StructInfo> params, StructInfo ret, Span span) { ObjectPtr<FuncStructInfoNode> n = make_object<FuncStructInfoNode>(); @@ -223,12 +186,6 @@ TVM_REGISTER_GLOBAL("relax.FuncStructInfoOpaqueFunc") } }); -TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) - .set_dispatch<FuncStructInfoNode>([](const ObjectRef& ref, ReprPrinter* p) { - const auto* node = static_cast<const FuncStructInfoNode*>(ref.get()); - p->stream << "FuncStructInfo(" << node->params << ", " << node->ret << ")"; - }); - // Helper functions void UpdateStructInfo(Expr expr, StructInfo struct_info) { ICHECK(!expr->struct_info_.defined()) diff --git a/src/script/printer/relax/binding.cc b/src/script/printer/relax/binding.cc new file mode 100644 index 0000000000..8a50fe9698 --- /dev/null +++ b/src/script/printer/relax/binding.cc @@ -0,0 +1,87 @@ +/* + * 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. + */ +#include "./utils.h" + +namespace tvm { +namespace script { +namespace printer { + +IfDoc PrintIfExpr(const relax::If& n, const ObjectPath& n_p, const IRDocsifier& d, // + const Optional<ExprDoc>& var, const Optional<ExprDoc>& ann) { + using relax::SeqExpr; + ExprDoc cond = d->AsDoc<ExprDoc>(n->cond, n_p->Attr("cond")); + std::vector<Array<StmtDoc>> branches{ + PrintSeqExpr(Downcast<SeqExpr>(n->true_branch), n_p->Attr("true_branch"), d, false), + PrintSeqExpr(Downcast<SeqExpr>(n->false_branch), n_p->Attr("false_branch"), d, false), + }; + if (var.defined()) { + for (Array<StmtDoc>& stmts : branches) { + ExprDoc ret = Downcast<ExprStmtDoc>(stmts.back())->expr; + stmts.Set(stmts.size() - 1, AssignDoc(var.value(), ret, ann)); + } + } + return IfDoc(cond, branches[0], branches[1]); +} + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::MatchCast>( + "", [](relax::MatchCast n, ObjectPath n_p, IRDocsifier d) -> Doc { + using relax::StructInfo; + using relax::MatchStructInfo; + Optional<ExprDoc> ann = StructInfoAsAnn(n->var, n_p->Attr("var"), d, n->value); + ExprDoc rhs = Relax(d, "match_cast") + ->Call({d->AsDoc<ExprDoc>(n->value, n_p->Attr("value")), + d->AsDoc<ExprDoc>(n->struct_info, n_p->Attr("struct_info_"))}); + ExprDoc lhs = DefineVar(n->var, d->frames.back(), d); + return AssignDoc(lhs, rhs, ann); + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::VarBinding>( // + "", [](relax::VarBinding n, ObjectPath n_p, IRDocsifier d) -> Doc { + if (const auto if_ = n->value.as<relax::IfNode>()) { + Optional<ExprDoc> ann = StructInfoAsAnn(n->var, n_p->Attr("var"), d, n->value); + ExprDoc lhs = DefineVar(n->var, d->frames.back(), d); + return PrintIfExpr(GetRef<relax::If>(if_), n_p->Attr("value"), d, lhs, ann); + } else if (n->value->IsInstance<tvm::BaseFuncNode>()) { + IdDoc lhs = DefineVar(n->var, d->frames.back(), d); + d->cfg->binding_names.push_back(lhs->name); + Doc ret = d->AsDoc(n->value, n_p->Attr("value")); + d->cfg->binding_names.pop_back(); + return ret; + } else { + ExprDoc rhs = d->AsDoc<ExprDoc>(n->value, n_p->Attr("value")); + Optional<ExprDoc> ann = StructInfoAsAnn(n->var, n_p->Attr("var"), d, n->value); + ExprDoc lhs = DefineVar(n->var, d->frames.back(), d); + return AssignDoc(lhs, rhs, ann); + } + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::If>("", [](relax::If n, ObjectPath n_p, IRDocsifier d) -> Doc { + return PrintIfExpr(n, n_p, d, NullOpt, NullOpt); + }); + +TVM_SCRIPT_REPR(relax::MatchCastNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::VarBindingNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::IfNode, ReprPrintRelax); + +} // namespace printer +} // namespace script +} // namespace tvm diff --git a/src/script/printer/relax/call.cc b/src/script/printer/relax/call.cc new file mode 100644 index 0000000000..2feb2082c5 --- /dev/null +++ b/src/script/printer/relax/call.cc @@ -0,0 +1,212 @@ +/* + * 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. + */ +#include "./utils.h" + +namespace tvm { +namespace script { +namespace printer { + +class AttrPrinter : public tvm::AttrVisitor { + public: + explicit AttrPrinter(const ObjectPath& p, const IRDocsifier& d, Array<String>* keys, + Array<ExprDoc>* values) + : p(p), d(d), keys(keys), values(values) {} + + void Visit(const char* key, double* value) final { + keys->push_back(key); + values->push_back(LiteralDoc::Float(*value, p->Attr(key))); + } + + void Visit(const char* key, int64_t* value) final { + keys->push_back(key); + values->push_back(LiteralDoc::Int(*value, p->Attr(key))); + } + + void Visit(const char* key, uint64_t* value) final { + keys->push_back(key); + values->push_back(LiteralDoc::Int(*value, p->Attr(key))); + } + + void Visit(const char* key, int* value) final { + keys->push_back(key); + values->push_back(LiteralDoc::Int(*value, p->Attr(key))); + } + + void Visit(const char* key, bool* value) final { + keys->push_back(key); + values->push_back(LiteralDoc::Boolean(*value, p->Attr(key))); + } + + void Visit(const char* key, std::string* value) final { + keys->push_back(key); + values->push_back(LiteralDoc::Str(*value, p->Attr(key))); + } + + void Visit(const char* key, DataType* value) final { + keys->push_back(key); + values->push_back(LiteralDoc::DataType(*value, p->Attr(key))); + } + + void Visit(const char* key, runtime::ObjectRef* value) final { + keys->push_back(key); + values->push_back(d->AsDoc<ExprDoc>(*value, p->Attr(key))); + } + + void Visit(const char* key, void** value) final { + LOG(FATAL) << "TypeError: void is not allowed in Attrs"; + } + + void Visit(const char* key, runtime::NDArray* value) final { + LOG(FATAL) << "TypeError: NDArray is not allowed in Attrs"; + } + + const ObjectPath& p; + const IRDocsifier& d; + Array<String>* keys; + Array<ExprDoc>* values; +}; + +ExprDoc PrintCallee(const relax::Expr& n, const ObjectPath& n_p, const IRDocsifier& d) { + // TODO(@junrushao): handle callee better + if (const auto* ext = n.as<relax::ExternFuncNode>()) { + return LiteralDoc::Str(ext->global_symbol, n_p); + } else if (const auto* gv = n.as<tvm::GlobalVarNode>()) { + IdDoc callee(gv->name_hint); + callee->source_paths.push_back(n_p); + return callee; + } else { + return d->AsDoc<ExprDoc>(n, n_p); + } +} + +Optional<ExprDoc> PrintCallTIR(const relax::Call& n, const ObjectPath& n_p, const IRDocsifier& d) { + static const Op& call_tir_op = Op::Get("relax.call_tir"); + if (!n->op.same_as(call_tir_op)) { + return NullOpt; + } + ICHECK(n->args.size() == 2 || n->args.size() == 3); + ICHECK(n->sinfo_args.size() == 1); + Array<ExprDoc> args; + Array<String> kwargs_keys; + Array<ExprDoc> kwargs_values; + // Step 1. Print n->args[0], the callee + args.push_back(PrintCallee(n->args[0], n_p->Attr("args")->ArrayIndex(0), d)); + // Step 2. Print n->args[1], the input arguments + args.push_back(d->AsDoc<ExprDoc>(n->args[1], n_p->Attr("args")->ArrayIndex(1))); + // Step 3. Print n->sinfo_args, the output struct info + relax::StructInfo o_sinfo = n->sinfo_args[0]; + ObjectPath o_sinfo_p = n_p->Attr("sinfo_args")->ArrayIndex(0); + kwargs_keys.push_back("out_sinfo"); + if (const auto* o = o_sinfo.as<relax::TupleStructInfoNode>()) { + Array<ExprDoc> fields; + ObjectPath fields_p = o_sinfo_p->Attr("fields"); + for (int i = 0, l = o->fields.size(); i < l; ++i) { + fields.push_back(d->AsDoc<ExprDoc>(o->fields[i], fields_p->ArrayIndex(i))); + } + kwargs_values.push_back(ListDoc(fields)); + } else { + kwargs_values.push_back(d->AsDoc<ExprDoc>(o_sinfo, o_sinfo_p)); + } + // Step 4. Print n->args[2], the tir variables + if (n->args.size() == 3) { + kwargs_keys.push_back("tir_vars"); + kwargs_values.push_back(d->AsDoc<ExprDoc>(n->args[2], n_p->Attr("args")->ArrayIndex(2))); + } + return Relax(d, "call_tir")->Call(args, kwargs_keys, kwargs_values); +} + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::Call>( // + "", [](relax::Call n, ObjectPath n_p, IRDocsifier d) -> Doc { + // Special case: call_tir + if (Optional<ExprDoc> doc = PrintCallTIR(n, n_p, d)) { + return doc.value(); + } + ExprDoc prefix{nullptr}; + Array<ExprDoc> args; + Array<String> kwargs_keys; + Array<ExprDoc> kwargs_values; + // Step 1. Print op + if (const auto* op = n->op.as<relax::ExternFuncNode>()) { + prefix = Relax(d, "call_packed"); + args.push_back(LiteralDoc::Str(op->global_symbol, n_p->Attr("op"))); + } else if (const auto* op = n->op.as<tvm::GlobalVarNode>()) { + prefix = IdDoc(op->name_hint); + prefix->source_paths.push_back(n_p->Attr("op")); + } else if (const auto* op = n->op.as<tvm::OpNode>()) { + std::string name = op->name; + if (name.rfind("relax.", 0) == 0) { + prefix = Relax(d, name.substr(6)); + } else { + prefix = IdDoc(name); + } + prefix->source_paths.push_back(n_p->Attr("op")); + } else if (n->op->IsInstance<relax::VarNode>()) { + prefix = d->AsDoc<ExprDoc>(n->op, n_p->Attr("op")); + } else { + LOG(FATAL) << "TypeError: Unsupported op: " << n->op->GetTypeKey(); + } + // Step 2. Print args + if (!n->args.empty()) { + args.push_back(PrintCallee(n->args[0], n_p->Attr("args")->ArrayIndex(0), d)); + } + for (int i = 1, l = n->args.size(); i < l; ++i) { + args.push_back(d->AsDoc<ExprDoc>(n->args[i], n_p->Attr("args")->ArrayIndex(i))); + } + // Step 3. Print attrs + if (n->attrs.defined()) { + if (n->op->IsInstance<relax::ExternFuncNode>()) { + kwargs_keys.push_back("attrs_type_key"); + kwargs_values.push_back(LiteralDoc::Str(n->attrs->GetTypeKey(), n_p->Attr("attrs"))); + } + if (const auto* attrs = n->attrs.as<tvm::DictAttrsNode>()) { + std::vector<std::pair<String, ObjectRef>> sorted; + for (const auto& kv : attrs->dict) { + sorted.push_back(kv); + } + std::sort(sorted.begin(), sorted.end()); + for (const auto& kv : sorted) { + kwargs_keys.push_back(kv.first); + kwargs_values.push_back( + d->AsDoc<ExprDoc>(kv.second, n_p->Attr("attrs")->Attr(kv.first))); + } + } else { + AttrPrinter printer(n_p->Attr("attrs"), d, &kwargs_keys, &kwargs_values); + const_cast<BaseAttrsNode*>(n->attrs.get())->VisitAttrs(&printer); + } + } + // Step 4. Print type_args + if (n->sinfo_args.size() > 0) { + ObjectPath sinfo_args_p = n_p->Attr("sinfo_args"); + Array<ExprDoc> sinfo_args; + for (int i = 0, l = n->sinfo_args.size(); i < l; ++i) { + sinfo_args.push_back( + d->AsDoc<ExprDoc>(n->sinfo_args[i], sinfo_args_p->ArrayIndex(i))); + } + kwargs_keys.push_back("sinfo_args"); + kwargs_values.push_back(TupleDoc(sinfo_args)); + } + return prefix->Call(args, kwargs_keys, kwargs_values); + }); + +TVM_SCRIPT_REPR(relax::CallNode, ReprPrintRelax); + +} // namespace printer +} // namespace script +} // namespace tvm diff --git a/src/script/printer/relax/expr.cc b/src/script/printer/relax/expr.cc new file mode 100644 index 0000000000..a786932fc3 --- /dev/null +++ b/src/script/printer/relax/expr.cc @@ -0,0 +1,136 @@ +/* + * 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. + */ +#include "./utils.h" + +namespace tvm { +namespace script { +namespace printer { + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::PrimValue>( // + "", [](relax::PrimValue n, ObjectPath n_p, IRDocsifier d) -> Doc { + // TODO(@junrushao): float numbers + return Relax(d, "prim_value")->Call({d->AsDoc<ExprDoc>(n->value, n_p->Attr("value"))}); + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::StringImm>( // + "", [](relax::StringImm n, ObjectPath n_p, IRDocsifier d) -> Doc { + return Relax(d, "str")->Call({LiteralDoc::Str(n->value, n_p->Attr("value"))}); + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::DataTypeImm>( // + "", [](relax::DataTypeImm n, ObjectPath n_p, IRDocsifier d) -> Doc { + return Relax(d, "dtype")->Call({LiteralDoc::DataType(n->value, n_p->Attr("value"))}); + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::Tuple>( // + "", [](relax::Tuple n, ObjectPath n_p, IRDocsifier d) -> Doc { + // TODO(@junrushao): revisit tuple printing + if (n->fields.empty()) { + return Relax(d, "tuple")->Call({}); + } + Array<ExprDoc> fields_doc; + ObjectPath fields_p = n_p->Attr("fields"); + for (int i = 0, l = n->fields.size(); i < l; ++i) { + fields_doc.push_back(d->AsDoc<ExprDoc>(n->fields[i], fields_p->ArrayIndex(i))); + } + return TupleDoc(fields_doc); + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::TupleGetItem>( // + "", [](relax::TupleGetItem n, ObjectPath n_p, IRDocsifier d) -> Doc { + ExprDoc idx = LiteralDoc::Int(n->index, n_p->Attr("index")); + return d->AsDoc<ExprDoc>(n->tuple, n_p->Attr("tuple"))[{idx}]; + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::ShapeExpr>( // + "", [](relax::ShapeExpr n, ObjectPath n_p, IRDocsifier d) -> Doc { + Array<ExprDoc> values_doc; + ObjectPath values_p = n_p->Attr("values"); + for (int i = 0, l = n->values.size(); i < l; ++i) { + values_doc.push_back(PrintShapeVar(n->values[i], values_p->ArrayIndex(i), d)); + } + return TupleDoc(values_doc); + }); + +Optional<ExprDoc> SpecialScalar(const runtime::NDArray& n, const ObjectPath& p) { + DataType dtype = n.DataType(); + const void* data = n->data; + if (n->ndim != 0 || n->device.device_type != kDLCPU) { + return NullOpt; + } + if (dtype == DataType::Int(32)) { + return LiteralDoc::Int(*reinterpret_cast<const int32_t*>(data), p); + } else if (dtype == DataType::Int(64)) { + return LiteralDoc::Int(*reinterpret_cast<const int64_t*>(data), p); + } else if (dtype == DataType::Float(32)) { + return LiteralDoc::Float(*reinterpret_cast<const float*>(data), p); + } else if (dtype == DataType::Float(64)) { + return LiteralDoc::Float(*reinterpret_cast<const double*>(data), p); + } else if (dtype == DataType::Bool()) { + return LiteralDoc::Boolean(*reinterpret_cast<const uint8_t*>(data), p); + } else { + return NullOpt; + } +} + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::Constant>( // + "", [](relax::Constant n, ObjectPath n_p, IRDocsifier d) -> Doc { + if (Optional<ExprDoc> s = SpecialScalar(n->data, n_p->Attr("data"))) { + return Relax(d, "const") + ->Call({ + s.value(), + LiteralDoc::DataType(n->data.DataType(), n_p->Attr("data")->Attr("dtype")), + }); + } + return d->AddMetadata(n); + }); + +Doc PrintRelaxVar(relax::Var n, ObjectPath p, IRDocsifier d) { + if (!d->IsVarDefined(n)) { + ExprDoc ann = d->AsDoc<ExprDoc>(n->struct_info_, p->Attr("struct_info_")); + Frame f = d->frames.back(); + ExprDoc var = DefineVar(n, f, d); + f->stmts.push_back(AssignDoc(var, NullOpt, ann)); + } + return d->GetVarDoc(n).value(); +} + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable).set_dispatch<relax::Var>("", PrintRelaxVar); +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable).set_dispatch<relax::DataflowVar>("", PrintRelaxVar); + +TVM_SCRIPT_REPR(relax::PrimValueNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::StringImmNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::DataTypeImmNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::TupleNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::TupleGetItemNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::ShapeExprNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::VarNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::DataflowVarNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::ConstantNode, ReprPrintRelax); + +} // namespace printer +} // namespace script +} // namespace tvm diff --git a/src/script/printer/relax/function.cc b/src/script/printer/relax/function.cc new file mode 100644 index 0000000000..fa085fcad4 --- /dev/null +++ b/src/script/printer/relax/function.cc @@ -0,0 +1,78 @@ +/* + * 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. + */ +#include "./utils.h" + +namespace tvm { +namespace script { +namespace printer { + +TVM_REGISTER_NODE_TYPE(RelaxFrameNode); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::Function>("", [](relax::Function n, ObjectPath n_p, IRDocsifier d) -> Doc { + std::unordered_set<const tir::VarNode*> func_vars; + With<RelaxFrame> f(d); + (*f)->AddDispatchToken(d, "relax"); + (*f)->is_func = true; + (*f)->func_vars = &func_vars; + // Step 1. Print the return type + Optional<ExprDoc> ret_type = NullOpt; + if (const auto& func_sinfo = relax::MatchStructInfo<relax::FuncStructInfo>(n)) { + ret_type = d->AsDoc<ExprDoc>(func_sinfo.value()->ret, // + n_p->Attr("struct_info_")->Attr("ret")); + } + // Step 2. Print params + Array<AssignDoc> params; + { + ObjectPath params_p = n_p->Attr("params"); + for (int i = 0, l = n->params.size(); i < l; ++i) { + params.push_back(AssignDoc( + /*lhs=*/DefineVar(n->params[i], *f, d), + /*rhs=*/NullOpt, StructInfoAsAnn(n->params[i], params_p->ArrayIndex(i), d, NullOpt))); + } + } + // Step 3. Clean up func variables + (*f)->func_vars = nullptr; + // Step 4. Print attributes + if (n->attrs.defined() && !n->attrs->dict.empty()) { + (*f)->stmts.push_back( + ExprStmtDoc(Relax(d, "func_attr") // + ->Call({d->AsDoc<ExprDoc>(n->attrs, n_p->Attr("attrs"))}))); + } + // Step 5. Print body + Array<StmtDoc> body = + PrintSeqExpr(Downcast<relax::SeqExpr>(n->body), n_p->Attr("body"), d, /*use_ret=*/true); + (*f)->stmts.insert((*f)->stmts.end(), body.begin(), body.end()); + return HeaderWrapper(d, FunctionDoc(IdDoc(FindFunctionName(d, n).value_or("main")), params, + {Relax(d, "function")}, ret_type, (*f)->stmts)); + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::ExternFunc>( // + "", [](relax::ExternFunc n, ObjectPath n_p, IRDocsifier d) -> Doc { + // TODO(@junrushao): print more information out of extern function. + return ExprStmtDoc(LiteralDoc::Str(n->global_symbol, n_p)); + }); + +TVM_SCRIPT_REPR(relax::FunctionNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::ExternFuncNode, ReprPrintRelax); + +} // namespace printer +} // namespace script +} // namespace tvm diff --git a/src/script/printer/relax/region.cc b/src/script/printer/relax/region.cc new file mode 100644 index 0000000000..1ac0b5ba14 --- /dev/null +++ b/src/script/printer/relax/region.cc @@ -0,0 +1,100 @@ +/* + * 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. + */ +#include "./utils.h" + +namespace tvm { +namespace script { +namespace printer { + +Array<StmtDoc> PrintSeqExpr(const relax::SeqExpr& n, const ObjectPath& n_p, const IRDocsifier& d, + bool use_ret) { + With<RelaxFrame> f(d); + const Array<relax::BindingBlock>& blocks = n->blocks; + ObjectPath blocks_p = n_p->Attr("blocks"); + Array<StmtDoc>* stmts = &(*f)->stmts; + for (int i = 0, l = blocks.size(); i < l; ++i) { + Doc block = d->AsDoc(blocks[i], blocks_p->ArrayIndex(i)); + if (const auto* stmt_block = block.as<StmtBlockDocNode>()) { + stmts->insert(stmts->end(), stmt_block->stmts.begin(), stmt_block->stmts.end()); + } else if (const auto* stmt = block.as<StmtDocNode>()) { + stmts->push_back(GetRef<StmtDoc>(stmt)); + } else { + LOG(FATAL) << "TypeError: Unknown type: " << block->GetTypeKey(); + } + } + ExprDoc ret = d->AsDoc<ExprDoc>(n->body, n_p->Attr("body")); + if (use_ret) { + stmts->push_back(ReturnDoc(ret)); + } else { + stmts->push_back(ExprStmtDoc(ret)); + } + return *stmts; +} + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::SeqExpr>("", [](relax::SeqExpr n, ObjectPath n_p, IRDocsifier d) -> Doc { + return StmtBlockDoc(PrintSeqExpr(n, n_p, d, false)); + }); + +Array<StmtDoc> PrintBindingBlock(const relax::BindingBlock& n, const ObjectPath& n_p, + const IRDocsifier& d, Array<ExprDoc>* non_dataflow_vars) { + const Array<relax::Binding>& bindings = n->bindings; + ObjectPath bindings_p = n_p->Attr("bindings"); + Array<StmtDoc> stmts; + for (int i = 0, l = bindings.size(); i < l; ++i) { + const relax::Binding& binding = bindings[i]; + ObjectPath binding_p = bindings_p->ArrayIndex(i); + ICHECK(binding->var.defined()); + Doc binding_doc = d->AsDoc(binding, binding_p); + if (const auto* stmt = binding_doc.as<StmtDocNode>()) { + stmts.push_back(GetRef<StmtDoc>(stmt)); + } else if (const auto* stmt_block = binding_doc.as<StmtBlockDocNode>()) { + stmts.insert(stmts.end(), stmt_block->stmts.begin(), stmt_block->stmts.end()); + } else { + LOG(FATAL) << "TypeError: Unknown type: " << binding_doc->GetTypeKey(); + } + if (non_dataflow_vars != nullptr && !binding->var->IsInstance<relax::DataflowVarNode>()) { + non_dataflow_vars->push_back(d->AsDoc<ExprDoc>(binding->var, binding_p->Attr("var"))); + } + } + return stmts; +} + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::BindingBlock>( // + "", [](relax::BindingBlock n, ObjectPath n_p, IRDocsifier d) -> Doc { + return StmtBlockDoc(PrintBindingBlock(n, n_p, d, nullptr)); + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::DataflowBlock>( // + "", [](relax::DataflowBlock n, ObjectPath n_p, IRDocsifier d) -> Doc { + Array<ExprDoc> non_dataflow_vars; + Array<StmtDoc> stmts = PrintBindingBlock(n, n_p, d, &non_dataflow_vars); + stmts.push_back(ExprStmtDoc(Relax(d, "output")->Call(non_dataflow_vars))); + return ScopeDoc(NullOpt, Relax(d, "dataflow")->Call({}), stmts); + }); + +TVM_SCRIPT_REPR(relax::SeqExprNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::BindingBlockNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::DataflowBlockNode, ReprPrintRelax); + +} // namespace printer +} // namespace script +} // namespace tvm diff --git a/src/script/printer/relax/struct_info.cc b/src/script/printer/relax/struct_info.cc new file mode 100644 index 0000000000..6f4a66c991 --- /dev/null +++ b/src/script/printer/relax/struct_info.cc @@ -0,0 +1,149 @@ +/* + * 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. + */ +#include <tvm/tir/stmt_functor.h> + +#include "./utils.h" + +namespace tvm { +namespace script { +namespace printer { + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::ObjectStructInfo>( // + "", [](relax::ObjectStructInfo n, ObjectPath n_p, IRDocsifier d) -> Doc { + return Relax(d, "Object"); + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::PrimStructInfo>( + "", [](relax::PrimStructInfo n, ObjectPath n_p, IRDocsifier d) -> Doc { + return Relax(d, "Prim")->Call({LiteralDoc::DataType(n->dtype, n_p->Attr("dtype"))}); + }); + +ExprDoc PrintShapeVar(const PrimExpr& e, const ObjectPath& e_p, const IRDocsifier& d) { + ExprDoc expr_doc = d->AsDoc<ExprDoc>(e, e_p); + // Step 1. Find if `func_vars` are being collected + const RelaxFrameNode* f = nullptr; + for (const Frame& frame : d->frames) { + if (const auto* relax_frame = frame.as<RelaxFrameNode>()) { + if (relax_frame->func_vars) { + f = relax_frame; + break; + } + } + } + // Step 2. Figure out if the PrimExpr contains at least a func var + bool func_var_mode = false; + if (f != nullptr) { + tir::PostOrderVisit(e, [f, &func_var_mode](const ObjectRef& obj) -> void { + if (const auto* var = obj.as<tir::VarNode>()) { + if (f->func_vars->count(var)) { + func_var_mode = true; + } + } + }); + } + // Step 3. Stringify the PrimExpr if func var exists + if (func_var_mode) { + return LiteralDoc::Str(DocToPythonScript(expr_doc, d->cfg), e_p); + } + return expr_doc; +} + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::ShapeStructInfo>( + "", [](relax::ShapeStructInfo n, ObjectPath n_p, IRDocsifier d) -> Doc { + if (n->values.defined()) { + Array<PrimExpr> shape = n->values.value(); + ObjectPath shape_p = n_p->Attr("values"); + Array<ExprDoc> shape_docs; + for (int i = 0, ndim = shape.size(); i < ndim; ++i) { + shape_docs.push_back(PrintShapeVar(shape[i], shape_p->ArrayIndex(i), d)); + } + return Relax(d, "Shape")->Call({ListDoc(shape_docs)}); + } + return Relax(d, "Shape") + ->Call({}, {"ndim"}, {LiteralDoc::Int(n->ndim, n_p->Attr("ndim"))}); + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::TensorStructInfo>( // + "", [](relax::TensorStructInfo n, ObjectPath n_p, IRDocsifier d) -> Doc { + Array<ExprDoc> args; + Array<String> kwargs_keys; + Array<ExprDoc> kwargs_values; + if (n->shape.defined()) { + args.push_back(d->AsDoc<ExprDoc>(n->shape.value(), n_p->Attr("shape"))); + } + if (!n->IsUnknownDtype()) { + kwargs_keys.push_back("dtype"); + kwargs_values.push_back(LiteralDoc::DataType(n->dtype, n_p->Attr("dtype"))); + } + if (!n->shape.defined() && !n->IsUnknownNdim()) { + kwargs_keys.push_back("ndim"); + kwargs_values.push_back(LiteralDoc::Int(n->ndim, n_p->Attr("ndim"))); + } + if (args.empty() && kwargs_keys.empty()) { + return Relax(d, "Tensor"); + } + return Relax(d, "Tensor")->Call(args, kwargs_keys, kwargs_values); + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::TupleStructInfo>( // + "", [](relax::TupleStructInfo n, ObjectPath n_p, IRDocsifier d) -> Doc { + if (n->fields.empty()) { + return Relax(d, "Tuple"); + } + Array<ExprDoc> fields_doc; + ObjectPath fields_p = n_p->Attr("fields"); + for (int i = 0, l = n->fields.size(); i < l; ++i) { + fields_doc.push_back(d->AsDoc<ExprDoc>(n->fields[i], fields_p->ArrayIndex(i))); + } + return Relax(d, "Tuple")->Call(fields_doc); + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::FuncStructInfo>( // + "", [](relax::FuncStructInfo n, ObjectPath n_p, IRDocsifier d) -> Doc { + if (n->IsOpaque()) { + return Relax(d, "Callable"); + } + // TODO(@junrushao): track symbolic shape relation + Array<ExprDoc> params_doc; + Array<relax::StructInfo> params = n->params.value(); + ObjectPath params_p = n_p->Attr("params"); + for (int i = 0, n_params = params.size(); i < n_params; ++i) { + params_doc.push_back(d->AsDoc<ExprDoc>(params[i], params_p->ArrayIndex(i))); + } + return Relax(d, "Callable") + ->Call({TupleDoc(params_doc), // + d->AsDoc<ExprDoc>(n->ret, n_p->Attr("ret"))}); + }); + +TVM_SCRIPT_REPR(relax::ObjectStructInfoNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::PrimStructInfoNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::ShapeStructInfoNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::TensorStructInfoNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::TupleStructInfoNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::FuncStructInfoNode, ReprPrintRelax); + +} // namespace printer +} // namespace script +} // namespace tvm diff --git a/src/script/printer/relax/tir.cc b/src/script/printer/relax/tir.cc new file mode 100644 index 0000000000..2c8bb0f1da --- /dev/null +++ b/src/script/printer/relax/tir.cc @@ -0,0 +1,89 @@ +/* + * 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. + */ +#include <tvm/ir/expr.h> + +#include "./utils.h" + +namespace tvm { +namespace script { +namespace printer { + +Doc PrintTIRVar(tir::Var n, ObjectPath n_p, IRDocsifier d) { + ICHECK(n->dtype.is_int() && n->dtype.is_scalar()) << "TypeError: Relax only uses " + "scalar integer TIR variables, but gets: " + << n; + if (!d->IsVarDefined(n)) { + // Find the outmost Relax function frame. If not exist, the outmost Relax frame. + RelaxFrameNode* f = nullptr; + for (const Frame& frame : d->frames) { + if (const auto* relax_frame = frame.as<RelaxFrameNode>()) { + if (relax_frame->is_func) { + f = const_cast<RelaxFrameNode*>(relax_frame); + break; + } else if (f == nullptr) { + f = const_cast<RelaxFrameNode*>(relax_frame); + } + } + } + // There should be at least one Relax frame + if (f == nullptr) { + LOG(FATAL) << "IndexError: No relax environment is found when printing a TIR var under " + "relax's dispatch token"; + } + // If the Relax function frame is collecting func vars + if (f->func_vars) { + ICHECK(f->is_func); + f->func_vars->insert(n.get()); + } + IdDoc var = d->Define(n, GetRef<Frame>(f), n->name_hint.empty() ? "v" : n->name_hint); + var->source_paths.push_back(n_p); + f->stmts.push_back(AssignDoc(var, + TIR(d, "Var")->Call({ + LiteralDoc::Str(var->name, n_p->Attr("name_hint")), + LiteralDoc::DataType(n->dtype, n_p->Attr("dtype")), + }), + NullOpt)); + } + if (Optional<ExprDoc> doc = d->GetVarDoc(n)) { + return doc.value(); + } + LOG(FATAL) << "IndexError: Variable is not defined in the environment: " << n; +} + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable).set_dispatch<tir::Var>("relax", PrintTIRVar); +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable).set_dispatch<tir::SizeVar>("relax", PrintTIRVar); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<tvm::IntImm>( // + "relax", [](tvm::IntImm n, ObjectPath n_p, IRDocsifier d) -> Doc { // + // TODO(@junrushao): support non-int64 cases + return LiteralDoc::Int(n->value, n_p); + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<tvm::GlobalVar>( // + "relax", [](tvm::GlobalVar n, ObjectPath n_p, IRDocsifier d) -> Doc { // + IdDoc ret(n->name_hint); + ret->source_paths.push_back(n_p); + return ret; + }); + +} // namespace printer +} // namespace script +} // namespace tvm diff --git a/src/script/printer/relax/type.cc b/src/script/printer/relax/type.cc new file mode 100644 index 0000000000..d13d90b1d5 --- /dev/null +++ b/src/script/printer/relax/type.cc @@ -0,0 +1,89 @@ +/* + * 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. + */ +#include "./utils.h" + +namespace tvm { +namespace script { +namespace printer { + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::ShapeType>( // + "", [](relax::ShapeType n, ObjectPath n_p, IRDocsifier d) -> Doc { + return Relax(d, "Shape") + ->Call({}, {"ndim"}, {LiteralDoc::Int(n->ndim, n_p->Attr("ndim"))}); + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::ObjectType>( // + "", [](relax::ObjectType n, ObjectPath n_p, IRDocsifier d) -> Doc { + return Relax(d, "Object"); + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::DynTensorType>( + "", [](relax::DynTensorType n, ObjectPath n_p, IRDocsifier d) -> Doc { + return Relax(d, "Tensor") + ->Call({}, {"ndim", "dtype"}, + {LiteralDoc::Int(n->ndim, n_p->Attr("ndim")), + LiteralDoc::DataType(n->dtype, n_p->Attr("dtype"))}); + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<relax::PackedFuncType>( + "", [](relax::PackedFuncType n, ObjectPath n_p, IRDocsifier d) -> Doc { + return Relax(d, "PackedFunc"); // TODO(@junrushao): verify if this is correct + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<tvm::TupleType>( // + "relax", [](tvm::TupleType n, ObjectPath n_p, IRDocsifier d) -> Doc { + if (n->fields.empty()) { + return Relax(d, "Tuple"); + } + Array<ExprDoc> fields_doc; + ObjectPath fields_p = n_p->Attr("fields"); + for (int i = 0, l = n->fields.size(); i < l; ++i) { + fields_doc.push_back(d->AsDoc<ExprDoc>(n->fields[i], fields_p->ArrayIndex(i))); + } + return Relax(d, "Tuple")->Call(fields_doc); + }); + +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<tvm::FuncType>( + "relax", [](tvm::FuncType n, ObjectPath n_p, IRDocsifier d) -> Doc { + Array<ExprDoc> arg_types_doc; + Array<Type> arg_types = n->arg_types; + ObjectPath arg_types_p = n_p->Attr("arg_types"); + for (int i = 0, n_params = arg_types.size(); i < n_params; ++i) { + arg_types_doc.push_back(d->AsDoc<ExprDoc>(arg_types[i], arg_types_p->ArrayIndex(i))); + } + return Relax(d, "Callable") + ->Call({TupleDoc(arg_types_doc), // + d->AsDoc<ExprDoc>(n->ret_type, n_p->Attr("ret_type"))}); + }); + +TVM_SCRIPT_REPR(relax::ShapeTypeNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::ObjectTypeNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::DynTensorTypeNode, ReprPrintRelax); +TVM_SCRIPT_REPR(relax::PackedFuncTypeNode, ReprPrintRelax); +TVM_REGISTER_GLOBAL("script.printer.ReprPrintRelax").set_body_typed(ReprPrintRelax); + +} // namespace printer +} // namespace script +} // namespace tvm diff --git a/src/script/printer/relax/utils.h b/src/script/printer/relax/utils.h new file mode 100644 index 0000000000..7702f7b22d --- /dev/null +++ b/src/script/printer/relax/utils.h @@ -0,0 +1,101 @@ +/* + * 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. + */ +#ifndef TVM_SCRIPT_PRINTER_RELAX_UTILS_H_ +#define TVM_SCRIPT_PRINTER_RELAX_UTILS_H_ + +#include <tvm/relax/struct_info.h> +#include <tvm/relax/type.h> +#include <tvm/script/printer/ir_docsifier.h> + +#include <string> +#include <unordered_set> +#include <utility> +#include <vector> + +#include "../utils.h" + +namespace tvm { +namespace script { +namespace printer { + +class RelaxFrameNode : public FrameNode { + public: + bool is_func = false; + std::unordered_set<const tir::VarNode*>* func_vars = nullptr; + + void VisitAttrs(AttrVisitor* v) { + FrameNode::VisitAttrs(v); + v->Visit("is_global_func", &is_func); + // `func_var_to_define` is not visited + } + + static constexpr const char* _type_key = "script.printer.RelaxFrame"; + TVM_DECLARE_FINAL_OBJECT_INFO(RelaxFrameNode, FrameNode); +}; + +class RelaxFrame : public Frame { + public: + explicit RelaxFrame(const IRDocsifier& d) { + ObjectPtr<RelaxFrameNode> n = make_object<RelaxFrameNode>(); + n->stmts.clear(); + n->d = d.get(); + n->is_func = false; + n->func_vars = nullptr; + data_ = std::move(n); + } + + TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(RelaxFrame, Frame, RelaxFrameNode); +}; + +/*! \brief Redirected method for the ReprPrinter */ +inline std::string ReprPrintRelax(const ObjectRef& obj, const PrinterConfig& cfg) { + IRDocsifier d(cfg); + With<RelaxFrame> f(d); + (*f)->AddDispatchToken(d, "relax"); + return Docsify(obj, d, *f, cfg); +} + +inline IdDoc DefineVar(const relax::Var& var, const Frame& frame, const IRDocsifier& d) { + return d->Define(var, frame, var->name_hint().empty() ? "v" : var->name_hint()); +} + +inline Optional<ExprDoc> StructInfoAsAnn(const relax::Var& v, const ObjectPath& v_p, + const IRDocsifier& d, const Optional<relax::Expr>& rhs) { + if (!v->struct_info_.defined()) { + return NullOpt; + } + if (const auto* call = rhs.as<relax::CallNode>()) { + static const Op& call_tir_op = Op::Get("relax.call_tir"); + if (call->op.same_as(call_tir_op)) { + return NullOpt; + } + } + return d->AsDoc<ExprDoc>(v->struct_info_, v_p->Attr("struct_info_")); +} + +Array<StmtDoc> PrintSeqExpr(const relax::SeqExpr& n, const ObjectPath& n_p, const IRDocsifier& d, + bool use_ret); + +ExprDoc PrintShapeVar(const PrimExpr& e, const ObjectPath& e_p, const IRDocsifier& d); + +} // namespace printer +} // namespace script +} // namespace tvm + +#endif // TVM_SCRIPT_PRINTER_RELAX_UTILS_H_ diff --git a/tests/python/relax/test_tvmscript_printer_relax.py b/tests/python/relax/test_tvmscript_printer_relax.py new file mode 100644 index 0000000000..e2cb8bc5fc --- /dev/null +++ b/tests/python/relax/test_tvmscript_printer_relax.py @@ -0,0 +1,488 @@ +# 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=missing-docstring +import pytest +from tvm import IRModule, relax, tir +from tvm.script import relax as R + + +def _assert_print(obj, expected): + if not isinstance(obj, str): + obj = obj.script(verbose_expr=True) + obj = obj.strip() + assert obj == expected.strip(), "\n" + obj + + +def test_function(): + @R.function + def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)): # type: ignore + return a + + _assert_print( + func, + """ +# from tvm.script import relax as R + [email protected] +def main(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)): + return a""", + ) + + +def test_extern_func(): + @R.function + def relax_func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)): # type: ignore + return a + + obj = IRModule( + { + "func": relax_func, + "my_ext": relax.ExternFunc("my_ext"), + } + ) + _assert_print( + obj, + """ +# from tvm.script import ir as I +# from tvm.script import relax as R + [email protected]_module +class Module: + "my_ext" + @R.function + def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)): + return a +""", + ) + + +def test_object_struct_info(): + obj = relax.ObjectStructInfo() + _assert_print( + obj, + "R.Object", + ) + + +def test_prim_struct_info(): + obj = relax.PrimStructInfo("float32") + _assert_print(obj, 'R.Prim("float32")') + + +def test_shape_struct_info_0(): + obj = relax.ShapeStructInfo(ndim=-1) + _assert_print(obj, "R.Shape(ndim=-1)") + + +def test_shape_struct_info_1(): + obj = relax.ShapeStructInfo([1, 2, 3]) + _assert_print(obj, "R.Shape([1, 2, 3])") + + +def test_shape_struct_info_2(): + obj = relax.ShapeStructInfo([1, tir.Var("a", "int64"), 3]) + _assert_print( + obj, + """ +a = T.Var("a", "int64") +R.Shape([1, a, 3])""", + ) + + +def test_tensor_struct_info(): + obj = relax.TensorStructInfo( + shape=relax.ShapeExpr([1, tir.Var("a", "int64"), 3]), + dtype="float32", + ) + _assert_print( + obj, + """ +a = T.Var("a", "int64") +R.Tensor((1, a, 3), dtype="float32") +""", + ) + + +def test_tuple_struct_info_empty(): + obj = relax.TupleStructInfo([]) + _assert_print(obj, "R.Tuple") + + +def test_tuple_struct_info(): + obj = relax.TupleStructInfo( + [ + relax.PrimStructInfo("float32"), + relax.ObjectStructInfo(), + relax.ShapeStructInfo([1, tir.Var("a", "int64"), 3]), + ] + ) + _assert_print( + obj, + """ +a = T.Var("a", "int64") +R.Tuple(R.Prim("float32"), R.Object, R.Shape([1, a, 3])) +""", + ) + + +def test_func_struct_info(): + obj = relax.FuncStructInfo( + params=[ + relax.PrimStructInfo("float32"), + relax.ObjectStructInfo(), + relax.ShapeStructInfo([1, tir.Var("a", "int64"), 3]), + ], + ret=relax.TensorStructInfo( + shape=relax.ShapeExpr([1, 2, 3]), + dtype="float32", + ), + ) + _assert_print( + obj, + """ +a = T.Var("a", "int64") +R.Callable((R.Prim("float32"), R.Object, R.Shape([1, a, 3])), R.Tensor((1, 2, 3), dtype="float32")) +""", + ) + + +def test_shape_type(): + obj = relax.ShapeType(ndim=3) + _assert_print(obj, "R.Shape(ndim=3)") + + +def test_object_type(): + obj = relax.ObjectType() + _assert_print(obj, "R.Object") + + +def test_dyn_tensor_type(): + obj = relax.DynTensorType() + _assert_print(obj, 'R.Tensor(ndim=-1, dtype="float32")') + + +def test_packed_func_type(): + obj = relax.PackedFuncType() + _assert_print(obj, "R.PackedFunc") + + +def test_tuple_type(): + obj = relax.TupleType([relax.ShapeType(ndim=3), relax.ObjectType()]) + _assert_print( + obj._relax_script(), # pylint: disable=protected-access + "R.Tuple(R.Shape(ndim=3), R.Object)", + ) + + +def test_func_type(): + obj = relax.FuncType( + arg_types=[ + relax.ObjectType(), + relax.ShapeType(ndim=3), + ], + ret_type=relax.DynTensorType( + ndim=3, + dtype="float32", + ), + ) + _assert_print( + obj._relax_script(), # pylint: disable=protected-access + 'R.Callable((R.Object, R.Shape(ndim=3)), R.Tensor(ndim=3, dtype="float32"))', + ) + + +def test_prim_value(): + obj = relax.PrimValue(1) + _assert_print(obj, "R.prim_value(1)") + + +def test_string_imm(): + obj = relax.StringImm("hello") + _assert_print(obj, 'R.str("hello")') + + +def test_data_type_imm(): + obj = relax.DataTypeImm("float32") + _assert_print(obj, 'R.dtype("float32")') + + +def test_var(): + obj = relax.Var("a", relax.TensorStructInfo([1, tir.Var("x", "int64"), 3], "float32")) + _assert_print( + obj, + """ +x = T.Var("x", "int64") +a: R.Tensor((1, x, 3), dtype="float32") +a""", + ) + + +def test_dataflow_var(): + obj = relax.DataflowVar("a", relax.TensorStructInfo([1, tir.Var("x", "int64"), 3], "float32")) + _assert_print( + obj, + """ +x = T.Var("x", "int64") +a: R.Tensor((1, x, 3), dtype="float32") +a""", + ) + + +def test_tuple(): + obj = relax.Tuple( + [ + relax.Var("a", relax.TensorStructInfo([1, tir.Var("x", "int64"), 3], "float32")), + relax.Var("b", relax.TensorStructInfo([1, tir.Var("y", "int64"), 3], "float32")), + relax.Var("c", relax.TensorStructInfo([1, tir.Var("z", "int64"), 3], "float32")), + ] + ) + _assert_print( + obj, + """ +x = T.Var("x", "int64") +a: R.Tensor((1, x, 3), dtype="float32") +y = T.Var("y", "int64") +b: R.Tensor((1, y, 3), dtype="float32") +z = T.Var("z", "int64") +c: R.Tensor((1, z, 3), dtype="float32") +(a, b, c) +""", + ) + + +def test_tuple_get_item(): + obj = relax.TupleGetItem( + relax.Tuple( + [ + relax.Var("a", relax.TensorStructInfo([1, tir.Var("x", "int64"), 3], "float32")), + relax.Var("b", relax.TensorStructInfo([1, tir.Var("y", "int64"), 3], "float32")), + relax.Var("c", relax.TensorStructInfo([1, tir.Var("z", "int64"), 3], "float32")), + ] + ), + 0, + ) + _assert_print( + obj, + """ +x = T.Var("x", "int64") +a: R.Tensor((1, x, 3), dtype="float32") +y = T.Var("y", "int64") +b: R.Tensor((1, y, 3), dtype="float32") +z = T.Var("z", "int64") +c: R.Tensor((1, z, 3), dtype="float32") +(a, b, c)[0] +""", + ) + + +def test_shape_expr(): + obj = relax.ShapeExpr([1, 2, 3]) + _assert_print(obj, "(1, 2, 3)") + + +def test_call(): + x = tir.Var("x", "int64") + a = relax.Var("a", relax.TensorStructInfo([1, x, 3], "float32")) + obj = relax.call_tir("my_func", args=a, out_sinfo=a.struct_info, tir_vars=[x]) + _assert_print( + obj, + """ +x = T.Var("x", "int64") +a: R.Tensor((1, x, 3), dtype="float32") +R.call_tir("my_func", (a,), out_sinfo=R.Tensor((1, x, 3), dtype="float32"), tir_vars=(x,)) +""", + ) + + [email protected](reason="`relax.op.sin` is not upstreamed yet") +def test_seq_expr(): + x = tir.Var("x", "int64") + a = relax.Var("a", relax.TensorStructInfo([1, x, 3], "float32")) + b = relax.DataflowVar("b", relax.TensorStructInfo([1, x, 3], "float32")) + c = relax.Var("c", relax.TensorStructInfo([1, x, 3], "float32")) + + obj = relax.SeqExpr( + blocks=[ + relax.DataflowBlock( + bindings=[ + relax.VarBinding(b, relax.op.sin(a)), + relax.VarBinding(c, relax.op.sin(b)), + ] + ), + ], + body=c, + ) + _assert_print( + obj, + """ +x = T.Var("x", "int64") +a: R.Tensor((1, x, 3), dtype="float32") +with R.dataflow(): + b: R.Tensor((1, x, 3), dtype="float32") = R.sin(a) + c: R.Tensor((1, x, 3), dtype="float32") = R.sin(b) + R.output(c) +c +""", + ) + + [email protected](reason="`relax.op.sin` is not upstreamed yet") +def test_binding_block(): + x = tir.Var("x", "int64") + a = relax.Var("a", relax.TensorStructInfo([1, x, 3], "float32")) + b = relax.Var("b", relax.TensorStructInfo([1, x, 3], "float32")) + c = relax.Var("c", relax.TensorStructInfo([1, x, 3], "float32")) + obj = relax.BindingBlock( + bindings=[ + relax.VarBinding(b, relax.op.sin(a)), + relax.VarBinding(c, relax.op.sin(b)), + ] + ) + _assert_print( + obj, + """ +x = T.Var("x", "int64") +a: R.Tensor((1, x, 3), dtype="float32") +b: R.Tensor((1, x, 3), dtype="float32") = R.sin(a) +c: R.Tensor((1, x, 3), dtype="float32") = R.sin(b) +""", + ) + + [email protected](reason="`relax.op.sin` is not upstreamed yet") +def test_dataflow_block(): + x = tir.Var("x", "int64") + a = relax.Var("a", relax.TensorStructInfo([1, x, 3], "float32")) + b = relax.DataflowVar("b", relax.TensorStructInfo([1, x, 3], "float32")) + c = relax.Var("c", relax.TensorStructInfo([1, x, 3], "float32")) + obj = relax.DataflowBlock( + bindings=[ + relax.VarBinding(b, relax.op.sin(a)), + relax.VarBinding(c, relax.op.sin(b)), + ] + ) + _assert_print( + obj, + """ +x = T.Var("x", "int64") +a: R.Tensor((1, x, 3), dtype="float32") +with R.dataflow(): + b: R.Tensor((1, x, 3), dtype="float32") = R.sin(a) + c: R.Tensor((1, x, 3), dtype="float32") = R.sin(b) + R.output(c) +""", + ) + + +def test_match_cast(): + x = tir.Var("x", "int64") + a = relax.Var("a", relax.TensorStructInfo([1, x, 3])) + b = relax.Var("b", relax.TensorStructInfo([1, 5, 3])) + obj = relax.MatchCast( + var=b, + value=a, + struct_info=b.struct_info, + ) + _assert_print( + obj, + """ +x = T.Var("x", "int64") +a: R.Tensor((1, x, 3), dtype="float32") +b: R.Tensor((1, 5, 3), dtype="float32") = R.match_cast(a, R.Tensor((1, 5, 3), dtype="float32")) +""", + ) + + [email protected](reason="`relax.op.sin` is not upstreamed yet") +def test_var_binding(): + x = tir.Var("x", "int64") + a = relax.Var("a", relax.TensorStructInfo([1, x, 3], "float32")) + b = relax.Var("b", relax.TensorStructInfo([1, x, 3], "float32")) + obj = relax.VarBinding(b, relax.op.sin(a)) + _assert_print( + obj, + """ +x = T.Var("x", "int64") +a: R.Tensor((1, x, 3), dtype="float32") +b: R.Tensor((1, x, 3), dtype="float32") = R.sin(a) +""", + ) + + +def test_if(): + a = relax.Var("a", relax.TensorStructInfo([], "bool")) + b = relax.Var("b", relax.TensorStructInfo([1, 2, 3], "float32")) + c = relax.Var("c", relax.TensorStructInfo([1, 2, 3], "float32")) + obj = relax.If( + a, + relax.SeqExpr([], b), + relax.SeqExpr([], c), + ) + _assert_print( + obj, + """ +a: R.Tensor((), dtype="bool") +if a: + b: R.Tensor((1, 2, 3), dtype="float32") + b +else: + c: R.Tensor((1, 2, 3), dtype="float32") + c +""", + ) + + +if __name__ == "__main__": + test_function() + test_extern_func() + + test_object_struct_info() + test_prim_struct_info() + test_shape_struct_info_0() + test_shape_struct_info_1() + test_shape_struct_info_2() + test_tensor_struct_info() + test_tuple_struct_info_empty() + test_tuple_struct_info() + test_func_struct_info() + + test_shape_type() + test_object_type() + test_dyn_tensor_type() + test_packed_func_type() + test_tuple_type() + test_func_type() + + test_prim_value() + test_string_imm() + test_data_type_imm() + + test_var() + test_dataflow_var() + # + test_tuple() + test_tuple_get_item() + test_shape_expr() + test_call() + + test_seq_expr() + test_binding_block() + test_dataflow_block() + + test_match_cast() + test_var_binding() + test_if()
