From f9e50199e99a3b2ad6ca893d24adb3123a2c78a6 Mon Sep 17 00:00:00 2001
From: Nicolas Vasilache <ntv@google.com>
Date: Sat, 22 Sep 2018 10:35:44 -0700
Subject: [PATCH] [MLIR] Fix AsmPrinter.cpp for single ssa-id AffineMap

The AsmPrinter wrongly assumes that all single ssa-id AffineMap
are the identity map for the purpose of printing.
This CL adds the missing level of indirection as well as a test.

This bug was originally shaken off by the experimental TC->MLIR path.

Before this CL, the test would print:
```
mlfunc @mlfuncsimplemap(%arg0 : affineint, %arg1 : affineint, %arg2 : affineint) {
  for %i0 = 0 to %arg0 {
    for %i1 = 0 to %i0 {
                   ~~~   should be %arg1
      %c42_i32 = constant 42 : i32
    }
  }
  return
}
```

PiperOrigin-RevId: 214120817
---
 mlir/lib/IR/AsmPrinter.cpp | 11 ++++++++---
 mlir/test/IR/parser.mlir   | 12 ++++++++++++
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 61e0ed15d13d..131b7f8475f2 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -1431,9 +1431,14 @@ void MLFunctionPrinter::printBound(AffineBound bound, const char *prefix) {
       return;
     }
 
-    // Print bound that consists of a single SSA id.
-    if (isa<AffineDimExpr>(expr) || isa<AffineSymbolExpr>(expr)) {
-      printOperand(bound.getOperand(0));
+    // Print bound that consists of a single SSA id, we need an indirection
+    // to achieve this.
+    if (auto *dimExpr = dyn_cast<AffineDimExpr>(expr)) {
+      printOperand(bound.getOperand(dimExpr->getPosition()));
+      return;
+    } else if (auto *symExpr = dyn_cast<AffineSymbolExpr>(expr)) {
+      printOperand(
+          bound.getOperand(map->getNumDims() + symExpr->getPosition()));
       return;
     }
   } else {
diff --git a/mlir/test/IR/parser.mlir b/mlir/test/IR/parser.mlir
index 659f6459dee4..421b450b2ccd 100644
--- a/mlir/test/IR/parser.mlir
+++ b/mlir/test/IR/parser.mlir
@@ -454,3 +454,15 @@ mlfunc @mlfuncattrempty() -> ()
   attributes {} {
   return
 }
+
+// CHECK-label mlfunc @mlfuncsimplemap
+#mapsimple0 = ()[s0, s1, s2] -> (s0)
+#mapsimple1 = (d0)[s0, s1, s2] -> (s1)
+mlfunc @mlfuncsimplemap(%arg0 : affineint, %arg1 : affineint, %arg2 : affineint) -> () {
+  for %i0 = 0 to #mapsimple0()[%arg0, %arg1, %arg2] { // CHECK: for %i0 = 0 to %arg0 {
+    for %i1 = 0 to #mapsimple1(%i0)[%arg0, %arg1, %arg2] { // CHECK: for %i1 = 0 to %arg1 {
+      %c42_i32 = constant 42 : i32
+    }
+  }
+  return
+}