Relax pretty printer (apache#8)

* Relax pretty printer initial prototype

* call into TVMScriptPrinter for PrimFuncs

* most round-trip tests pass

* address comments

* fix typo

include/tvm/ir/type_functor.h

@@ -27,6 +27,7 @@
 #include <tvm/node/functor.h>
 #include <tvm/relay/adt.h>
 #include <tvm/relay/expr.h>
+#include <tvm/relax/type.h>
 
 #include <string>
 #include <utility>
@@ -89,6 +90,9 @@ class TypeFunctor<R(const Type& n, Args...)> {
  virtual R VisitType_(const TypeDataNode* op, Args... args) TYPE_FUNCTOR_DEFAULT;
  virtual R VisitType_(const PrimTypeNode* op, Args... args) TYPE_FUNCTOR_DEFAULT;
  virtual R VisitType_(const PointerTypeNode* op, Args... args) TYPE_FUNCTOR_DEFAULT;
+ virtual R VisitType_(const relax::ShapeTypeNode* op, Args... args) TYPE_FUNCTOR_DEFAULT;
+ virtual R VisitType_(const relax::DynTensorTypeNode* op, Args... args) TYPE_FUNCTOR_DEFAULT;
+ virtual R VisitType_(const relax::DimTypeNode* op, Args... args) TYPE_FUNCTOR_DEFAULT;
  virtual R VisitTypeDefault_(const Object* op, Args...) {
  LOG(FATAL) << "Do not have a default for " << op->GetTypeKey();
  throw; // unreachable, written to stop compiler warning
@@ -112,6 +116,9 @@ class TypeFunctor<R(const Type& n, Args...)> {
  TVM_TYPE_FUNCTOR_DISPATCH(TypeDataNode);
  TVM_TYPE_FUNCTOR_DISPATCH(PrimTypeNode);
  TVM_TYPE_FUNCTOR_DISPATCH(PointerTypeNode);
+ TVM_TYPE_FUNCTOR_DISPATCH(relax::ShapeTypeNode);
+ TVM_TYPE_FUNCTOR_DISPATCH(relax::DynTensorTypeNode);
+ TVM_TYPE_FUNCTOR_DISPATCH(relax::DimTypeNode);
  return vtable;
  }
 };

include/tvm/relax/expr.h

@@ -96,6 +96,7 @@ class VarNode : public ExprNode {
  }
 
  bool SEqualReduce(const VarNode* other, SEqualReducer equal) const {
+ equal->MarkGraphNode();
  return equal(vid, other->vid) && equal(type_annotation, other->type_annotation) &&
  // Do we use the analysis information in equality?
  equal(checked_type_, other->checked_type_) && equal(shape_, other->shape_);
@@ -140,6 +141,7 @@ class DataflowVarNode : public VarNode {
  }
 
  bool SEqualReduce(const DataflowVarNode* other, SEqualReducer equal) const {
+ equal->MarkGraphNode();
  return equal(vid, other->vid) && equal(type_annotation, other->type_annotation) &&
  equal(shape_, other->shape_) && equal(checked_type_, other->checked_type_);
  }

include/tvm/relax/ir_functor.h

@@ -0,0 +1,120 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/relax/ir_functor.h
+ * \brief A generic visitor for traversing Relax IR nodes.
+ */
+#ifndef TVM_RELAX_IR_FUNCTOR_H_
+#define TVM_RELAX_IR_FUNCTOR_H_
+
+#include <tvm/node/functor.h>
+#include <tvm/node/node.h>
+#include <tvm/relax/expr.h>
+#include <tvm/relay/expr.h>
+
+namespace tvm {
+namespace relax {
+
+template <typename FType>
+class IRFunctor;
+
+#define IR_FUNCTOR_DEFAULT \
+ { return VisitNodeDefault_(op, std::forward<Args>(args)...); }
+
+#define RELAX_IR_FUNCTOR_DISPATCH(OP) \
+ vtable.template set_dispatch<OP>([](const ObjectRef& n, TSelf* self, Args... args) { \
+ return self->VisitNode_(static_cast<const OP*>(n.get()), std::forward<Args>(args)...); \
+ });
+
+template <typename R, typename... Args>
+class IRFunctor<R(const ObjectRef& n, Args...)> {
+ private:
+ using TSelf = IRFunctor<R(const ObjectRef& n, Args...)>;
+ using FType = NodeFunctor<R(const ObjectRef& n, TSelf* self, Args...)>;
+
+ public:
+ using result_type = R;
+ virtual ~IRFunctor() {}
+
+ R operator()(const ObjectRef& n, Args... args) {
+ return VisitNode(n, std::forward<Args>(args)...);
+ }
+
+ virtual R VisitNode(const ObjectRef& n, Args... args) {
+ ICHECK(n.defined()) << "Found null pointer node while traversing AST. The previous pass may "
+ "have generated invalid data.";
+ static FType vtable = InitVTable();
+ return vtable(n, this, std::forward<Args>(args)...);
+ }
+
+ // IR nodes inherited from Relay
+ virtual R VisitNode_(const relay::ConstantNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+ virtual R VisitNode_(const relay::TupleNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+ virtual R VisitNode_(const relay::GlobalVarNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+ virtual R VisitNode_(const relay::CallNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+ virtual R VisitNode_(const relay::IfNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+ virtual R VisitNode_(const OpNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+ virtual R VisitNode_(const relay::TupleGetItemNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+
+ // IR nodes introduced by Relax
+ virtual R VisitNode_(const relax::VarNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+ virtual R VisitNode_(const relax::DataflowVarNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+ virtual R VisitNode_(const relax::ShapeExprNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+ virtual R VisitNode_(const relax::MatchShapeNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+ virtual R VisitNode_(const relax::VarBindingNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+ virtual R VisitNode_(const relax::BindingBlockNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+ virtual R VisitNode_(const relax::DataflowBlockNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+ virtual R VisitNode_(const relax::SeqExprNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+ virtual R VisitNode_(const relax::FunctionNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+ virtual R VisitNode_(const relax::ExternFuncNode* op, Args... args) IR_FUNCTOR_DEFAULT;
+
+ virtual R VisitNodeDefault_(const Object* op, Args...) {
+ LOG(FATAL) << "no default visitor implemented for " << op->GetTypeKey();
+ throw;
+ }
+
+ private:
+ static FType InitVTable() {
+ FType vtable;
+ RELAX_IR_FUNCTOR_DISPATCH(relay::ConstantNode);
+ RELAX_IR_FUNCTOR_DISPATCH(relay::TupleNode);
+ RELAX_IR_FUNCTOR_DISPATCH(relay::GlobalVarNode);
+ RELAX_IR_FUNCTOR_DISPATCH(relay::CallNode);
+ RELAX_IR_FUNCTOR_DISPATCH(relay::IfNode);
+ RELAX_IR_FUNCTOR_DISPATCH(OpNode);
+ RELAX_IR_FUNCTOR_DISPATCH(relay::TupleGetItemNode);
+ RELAX_IR_FUNCTOR_DISPATCH(relax::VarNode);
+ RELAX_IR_FUNCTOR_DISPATCH(relax::DataflowVarNode);
+ RELAX_IR_FUNCTOR_DISPATCH(relax::ShapeExprNode);
+ RELAX_IR_FUNCTOR_DISPATCH(relax::MatchShapeNode);
+ RELAX_IR_FUNCTOR_DISPATCH(relax::VarBindingNode);
+ RELAX_IR_FUNCTOR_DISPATCH(relax::BindingBlockNode);
+ RELAX_IR_FUNCTOR_DISPATCH(relax::DataflowBlockNode);
+ RELAX_IR_FUNCTOR_DISPATCH(relax::SeqExprNode);
+ RELAX_IR_FUNCTOR_DISPATCH(relax::FunctionNode);
+ RELAX_IR_FUNCTOR_DISPATCH(relax::ExternFuncNode);
+ return vtable;
+ }
+};
+
+} // namespace relax
+} // namespace tvm
+
+#endif // TVM_RELAX_IR_FUNCTOR_H_

python/tvm/relax/expr.py

@@ -129,7 +129,8 @@ def __init__(self, bindings: List[Binding], span: Span = None) -> None:
 
 @tvm._ffi.register_object("relax.expr.DataflowBlock")
 class DataflowBlock(BindingBlock):
- pass
+ def __init__(self, bindings: List[Binding], span: Span = None) -> None:
+ self.__init_handle_by_constructor__(_ffi_api.DataflowBlock, bindings, span)
 
 
 @tvm._ffi.register_object("relax.expr.SeqExpr")