Add printer for Layout/BijectiveLayout (apache#3582)

neo-ai · Sep 6, 2019 · 4b7d9f1 · 4b7d9f1
1 parent 168f1bb
commit 4b7d9f1
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/python/tvm/tensor.py b/python/tvm/tensor.py
@@ -225,12 +225,6 @@ class Layout(NodeBase):
     --------
     layout : Declare a layout
     """
-    def __str__(self):
-        return self.name
-
-    def __repr__(self):
-        return "Layout(" + self.name + ")"
-
     def __len__(self):
         return _api_internal._LayoutNdim(self)
 

diff --git a/src/lang/data_layout.cc b/src/lang/data_layout.cc
@@ -195,6 +195,11 @@ int32_t Layout::FactorOf(const LayoutAxis& axis) const {
   return -1;
 }
 
+TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable)
+.set_dispatch<LayoutNode>([](const LayoutNode* l, IRPrinter* p) {
+    p->stream << "Layout(" << l->name << ")";
+  });
+
 inline bool GetStoreRule(Array<Expr>* rule,
                          const Layout& src_layout,
                          const Layout& dst_layout) {
@@ -346,4 +351,10 @@ BijectiveLayout BijectiveLayoutNode::make(const Layout& src_layout,
   return BijectiveLayout(n);
 }
 
+TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable)
+.set_dispatch<BijectiveLayoutNode>([](const BijectiveLayoutNode* b, IRPrinter* p) {
+    p->stream << "BijectiveLayout(" << b->src_layout.name()
+              << "->" << b->dst_layout.name() << ")";
+  });
+
 }  // namespace tvm