Add printer for Layout/BijectiveLayout (#3582)

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)
 

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