Cleanup spurious printing bits in EDSCs

This CL also makes ScopedEDSCContexts to reset the Bindable numbering when
creating a new context.
This is useful to write minimal tests that don't use FileCheck pattern
captures for now.

PiperOrigin-RevId: 230079997

diff --git a/mlir/include/mlir/EDSC/Types.h b/mlir/include/mlir/EDSC/Types.h
index 90d0ba680d808d7c940e653433dbc0f230b701ba..1ac4081694280085bcd28c84384582598dfa99fb 100644 (file)
--- a/mlir/include/mlir/EDSC/Types.h
+++ b/mlir/include/mlir/EDSC/Types.h
@@ -227,7 +227,10 @@ struct Bindable : public Expr {
     assert(!ptr || isa<Bindable>() && "expected Bindable");
   }
 
+  friend struct ScopedEDSCContext;
+
 private:
+  static void resetIds() { newId() = 0; }
   static unsigned &newId();
 };
 

diff --git a/mlir/lib/EDSC/Types.cpp b/mlir/lib/EDSC/Types.cpp
index f3368cd5dbd23b7a6dec2aca2d30b0029f1e8ee7..97d21434934b3b190be46ba9315c83fcc1ca2470 100644 (file)
--- a/mlir/lib/EDSC/Types.cpp
+++ b/mlir/lib/EDSC/Types.cpp
@@ -78,7 +78,10 @@ struct StmtStorage {
 
 } // namespace detail
 
-ScopedEDSCContext::ScopedEDSCContext() { Expr::globalAllocator() = &allocator; }
+ScopedEDSCContext::ScopedEDSCContext() {
+  Expr::globalAllocator() = &allocator;
+  Bindable::resetIds();
+}
 
 ScopedEDSCContext::~ScopedEDSCContext() { Expr::globalAllocator() = nullptr; }
 
@@ -241,7 +244,7 @@ void Expr::print(raw_ostream &os) const {
   } else if (auto un = this->dyn_cast<UnaryExpr>()) {
     os << "unknown_unary";
   } else if (auto bin = this->dyn_cast<BinaryExpr>()) {
-    os << bin.getLHS();
+    os << "(" << bin.getLHS();
     switch (bin.getKind()) {
     case ExprKind::Add:
       os << " + ";
@@ -271,7 +274,8 @@ void Expr::print(raw_ostream &os) const {
       os << "unknown_binary";
     }
     }
-    os << bin.getRHS();
+    os << bin.getRHS() << ")";
+    return;
   } else if (auto ter = this->dyn_cast<TernaryExpr>()) {
     switch (ter.getKind()) {
     case ExprKind::Select:
@@ -451,8 +455,7 @@ void Stmt::print(raw_ostream &os, Twine indent) const {
   if (auto stmtExpr = rhs.dyn_cast<StmtBlockLikeExpr>()) {
     switch (stmtExpr.getKind()) {
     case ExprKind::For:
-      os << indent << "for(idx(" << lhs << ")=" << stmtExpr << ") {";
-      os << " // @" << storage;
+      os << indent << "for(" << lhs << " = " << stmtExpr << ") {";
       os << "\n";
       for (const auto &s : getEnclosedStmts()) {
         if (!s.getRHS().isa<StmtBlockLikeExpr>()) {
@@ -477,7 +480,7 @@ void Stmt::print(raw_ostream &os, Twine indent) const {
     }
     }
   } else {
-    os << "lhs(" << lhs << ") = " << rhs;
+    os << lhs << " = " << rhs;
   }
 }
 

diff --git a/mlir/test/EDSC/lower-edsc.mlir b/mlir/test/EDSC/lower-edsc.mlir
index 5d05b5932cd438552497edff62deb9160a81edc9..8786e9238ac2b7a6e6008c125b6752ba9b155048 100644 (file)
--- a/mlir/test/EDSC/lower-edsc.mlir
+++ b/mlir/test/EDSC/lower-edsc.mlir
@@ -10,20 +10,20 @@ func @fn() {
 }
 
 // CHECK: block {
-// CHECK:   for(idx({{.*}})=[[zero1:.*]] to {{.*}} step [[step1:.*]]) {
-// CHECK:     for(idx({{.*}})=[[zero1]] to {{.*}} step [[step1]]) {
-// CHECK:       for(idx({{.*}})=[[zero1]] to {{.*}} step [[step1]]) {
-// CHECK:         for(idx({{.*}})=[[zero1]] to {{.*}} step [[step1]]) {
-// CHECK:           lhs({{.*}}) = store( ... );
+// CHECK:   for({{.*}}=[[zero1:.*]] to {{.*}} step [[step1:.*]]) {
+// CHECK:     for({{.*}}=[[zero1]] to {{.*}} step [[step1]]) {
+// CHECK:       for({{.*}}=[[zero1]] to {{.*}} step [[step1]]) {
+// CHECK:         for({{.*}}=[[zero1]] to {{.*}} step [[step1]]) {
+// CHECK:           {{.*}} = store( ... );
 // CHECK:         };
 // CHECK:       };
 // CHECK:     };
 // CHECK:   }
 // CHECK: }
 // CHECK: block {
-// CHECK:   for(idx({{.*}})=[[zero2:.*]] to {{.*}} step [[step2:.*]]) {
-// CHECK:     for(idx({{.*}})=[[zero2]] to {{.*}} step [[step2]]) {
-// CHECK:       lhs({{.*}}) = store( ... );
+// CHECK:   for({{.*}}=[[zero2:.*]] to {{.*}} step [[step2:.*]]) {
+// CHECK:     for({{.*}}=[[zero2]] to {{.*}} step [[step2]]) {
+// CHECK:       {{.*}} = store( ... );
 // CHECK:     };
 // CHECK:   }
 // CHECK: }