From 119af6712e2a1f375aa6a25d38a812dcf29f1503 Mon Sep 17 00:00:00 2001
From: Nicolas Vasilache <ntv@google.com>
Date: Sat, 19 Jan 2019 13:20:28 -0800
Subject: [PATCH] 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
---
 mlir/include/mlir/EDSC/Types.h |  3 +++
 mlir/lib/EDSC/Types.cpp        | 15 +++++++++------
 mlir/test/EDSC/lower-edsc.mlir | 16 ++++++++--------
 3 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/mlir/include/mlir/EDSC/Types.h b/mlir/include/mlir/EDSC/Types.h
index 90d0ba680d80..1ac408169428 100644
--- 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 f3368cd5dbd2..97d21434934b 100644
--- 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 5d05b5932cd4..8786e9238ac2 100644
--- 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: }