Relax pretty printer (apache#8)

* Relax pretty printer initial prototype

* call into TVMScriptPrinter for PrimFuncs

* most round-trip tests pass

* address comments

* fix typo

Author: altanh

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