Add printer for Layout/BijectiveLayout (#3582)

diff --git a/python/tvm/tensor.py b/python/tvm/tensor.py
index db8fb272a551eeb939ce743b8cc6cc3e778bd0b0..1cadf0621823e7db7479a1ef2d7cb3013444e2ab 100644 (file)
--- 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
index 9f3d0fe494a2c1ed800ca1ec40506666c31fa197..fa823c80aa47c33b0483bf2b0256113b863655b1 100644 (file)
--- 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