[MLIR] Fix bug in simplify affine map with operands

[bondhugula]

Summary

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.

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

 static void simplifyExprAndOperands(AffineExpr &expr,
                                     ArrayRef<Value> operands) {
   // We do this only for certain floordiv/mod expressions.
   auto binExpr = expr.dyn_cast<AffineBinaryOpExpr>();
   if (!binExpr)
     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);
 
   binExpr = expr.dyn_cast<AffineBinaryOpExpr>();
   if (!binExpr || (binExpr.getKind() != AffineExprKind::FloorDiv &&
                    binExpr.getKind() != AffineExprKind::Mod)) {
     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;
 
   // Simplify expressions of the form e = (e_1 + e_2) floordiv c or (e_1 + e_2)
   // mod c, where e_1 is a multiple of `k` and 0 <= e_2 < k. In such cases, if
   // `c` % `k` == 0, (e_1 + e_2) floordiv c can be simplified to e_1 floordiv c.
   // And when k % c == 0, (e_1 + e_2) mod c can be simplified to e_2 mod c.
   if (isQTimesDPlusR(lhs, operands, divisor, quotientTimesDiv, rem)) {

mlir/test/Dialect/Affine/canonicalize.mlir

func.func @fold_const_splat_global() -> memref<32x64xf32> {
     return %s: memref<32x64xf32>
   }
 }
 
 // -----
 
 // 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 %{{.*}}
   affine.for %i = 0 to %N step 32 {
     // CHECK-NEXT: affine.for %[[II:.*]] = 0 to 32
     affine.for %ii = 0 to 32 {
       // %ii is less than 32 and %i divides 32.
       // CHECK: affine.load %{{.*}}[0, 0]
func.func @simplify_with_operands(%N: index, %A: memref<?x32xf32>) {
 
   // Should not simplify.
   affine.for %i = -1 to 32 {
     // CHECK: affine.load %{{.*}}[%{{.*}} floordiv {{.*}}, %{{.*}} mod {{.*}}] :
     %x = affine.load %A[%i floordiv 32, %i mod 32] : 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
 }