[MLIR] Fix bug in simplify affine map with operands

Fix bug in simplify affine map with operands utility; the wrong LHS and
RHS were being used in some cases post simplification. While on this,
also handle a corner case of undefined expressions.

Differential Revision: https://reviews.llvm.org/D138584

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

@@ -679,8 +679,8 @@ static void simplifyExprAndOperands(AffineExpr &expr,
     return;
 
   // Simplify the child expressions first.
-  auto lhs = binExpr.getLHS();
-  auto rhs = binExpr.getRHS();
+  AffineExpr lhs = binExpr.getLHS();
+  AffineExpr rhs = binExpr.getRHS();
   simplifyExprAndOperands(lhs, operands);
   simplifyExprAndOperands(rhs, operands);
   expr = getAffineBinaryOpExpr(binExpr.getKind(), lhs, rhs);
@@ -691,11 +691,18 @@ static void simplifyExprAndOperands(AffineExpr &expr,
     return;
   }
 
+  // The `lhs` and `rhs` may be different post construction of simplified expr.
+  lhs = binExpr.getLHS();
+  rhs = binExpr.getRHS();
   auto rhsConst = rhs.dyn_cast<AffineConstantExpr>();
   if (!rhsConst)
     return;
 
   int64_t rhsConstVal = rhsConst.getValue();
+  // Undefined exprsessions aren't touched; IR can still be valid with them.
+  if (rhsConstVal == 0)
+    return;
+
   AffineExpr quotientTimesDiv, rem;
   int64_t divisor;
 

mlir/test/Dialect/Affine/canonicalize.mlir

@@ -1155,6 +1155,8 @@ module {
 
 // Simplification of maps exploiting operand info.
 
+// CHECK: #[[$MAP_SIMPLER:.*]] = affine_map<(d0, d1) -> (((d0 + d1) mod 458313) floordiv 227)>
+
 // CHECK-LABEL: func @simplify_with_operands
 func.func @simplify_with_operands(%N: index, %A: memref<?x32xf32>) {
   // CHECK-NEXT: affine.for %[[I:.*]] = 0 to %{{.*}}
@@ -1186,5 +1188,15 @@ func.func @simplify_with_operands(%N: index, %A: memref<?x32xf32>) {
     "test.foo"(%x) : (f32) -> ()
   }
 
+  affine.for %arg0 = 0 to %N step 128 {
+    affine.for %arg4 = 0 to 32 step 32 {
+      affine.for %arg5 = 0 to 128 {
+        // CHECK: affine.apply #[[$MAP_SIMPLER]]
+        %x = affine.apply affine_map<(d0, d1, d2) -> (((d0 + d2) mod 458313) floordiv 227 + d1 floordiv 256)>(%arg0, %arg4, %arg5)
+        "test.foo"(%x) : (index) -> ()
+      }
+    }
+  }
+
   return
 }