Fix asm printer for affine expr

- fixes tensorflow/mlir#201 Signed-off-by: Uday Bondhugula <uday@polymagelabs.com> Closes tensorflow/mlir#204 COPYBARA_INTEGRATE_REVIEW=tensorflow/mlir#204 from bondhugula:printfix 3f8a5b65391f45598258b2735fecaa409fbde848 PiperOrigin-RevId: 279115720
llvm · Nov 7, 2019 · eb47d5e · eb47d5e
1 parent 8f00b44
commit eb47d5e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
@@ -1182,9 +1182,12 @@ void ModulePrinter::printAffineExprInternal(
 
     // Pretty print multiplication with -1.
     auto rhsConst = rhsExpr.dyn_cast<AffineConstantExpr>();
-    if (rhsConst && rhsConst.getValue() == -1) {
+    if (rhsConst && binOp.getKind() == AffineExprKind::Mul &&
+        rhsConst.getValue() == -1) {
       os << "-";
       printAffineExprInternal(lhsExpr, BindingStrength::Strong, printValueName);
+      if (enclosingTightness == BindingStrength::Strong)
+        os << ')';
       return;
     }
 

diff --git a/mlir/test/IR/affine-map.mlir b/mlir/test/IR/affine-map.mlir
@@ -174,6 +174,9 @@
 // CHECK: #map{{[0-9]+}} = (d0) -> (d0 - (d0 + 1))
 #map53 = (d0) -> (d0 - (d0 + 1))
 
+// CHECK: #map{{[0-9]+}} = (d0)[s0] -> ((-s0) floordiv 4, d0 floordiv -1)
+#map54 = (d0)[s0] -> (-s0 floordiv 4, d0 floordiv -1)
+
 // Single identity maps are removed.
 // CHECK: func @f0(memref<2x4xi8, 1>)
 func @f0(memref<2x4xi8, #map0, 1>)
@@ -343,3 +346,6 @@ func @f52(memref<1xi8, #map52>)
 
 // CHECK: func @f53(memref<1xi8, #map{{[0-9]+}}>)
 func @f53(memref<1xi8, #map53>)
+
+// CHECK: func @f54(memref<10xi32, #map{{[0-9]+}}>)
+func @f54(memref<10xi32, #map54>)