[MLIR] Simplify Semi-affine expressions by rule based matching and replacing "expr - q * (expr floordiv q)" with "expr mod q" expression.

[arnab-oss]

Summary

Add rule based matching for detecting and transforming "expr - q * (expr floordiv q)"
to "expr mod q", where q is a symbolic exxpression, in simplifyAdd function.

[bondhugula]

mlir/lib/IR/AffineExpr.cpp

620

You don't need the extra llrhsBinOpExpr on the second disjunct.

mlir/test/Dialect/Affine/simplify-affine-structures.mlir

487

Sorted order here -- move this up.

492

CHECK-LABEL here while at this.

[bondhugula]

mlir/lib/IR/AffineExpr.cpp

608

-> Check if lrhsBin... is of the form ...

mlir/test/Dialect/Affine/simplify-affine-structures.mlir

493

Nit: use underscore between semi and affine.

[arnab-oss]

Added test cases.

[arnab-oss]

Addressed all comments by @bondhugula

[arnab-oss]

ping @dcaballe and @springerm , can you please take a look and review this? This is a continuation of semi affine expression simplification PR.

[dcaballe]

LGTM, thanks!

mlir/lib/IR/AffineExpr.cpp

static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs) {
   // When doing successive additions, bring constant to the right: turn (d0 + 2)
   // + d1 into (d0 + d1) + 2.
   if (lBin && lBin.getKind() == AffineExprKind::Add) {
     if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) {
       return lBin.getLHS() + rhs + lrhs;
     }
   }
 
-  // Detect and transform "expr - c * (expr floordiv c)" to "expr mod c". This
-  // leads to a much more efficient form when 'c' is a power of two, and in
-  // general a more compact and readable form.
+  // Detect and transform "expr - q * (expr floordiv q)" to "expr mod q", where
+  // q may be a constant or symbolic expression. This leads to a much more
+  // efficient form when 'c' is a power of two, and in general a more compact
+  // and readable form.
 
   // Process '(expr floordiv c) * (-c)'.
   if (!rBinOpExpr)
     return nullptr;
 
   auto lrhs = rBinOpExpr.getLHS();
   auto rrhs = rBinOpExpr.getRHS();
 
+  AffineExpr llrhs, rlrhs;
+
+  // Check if lrhsBinOpExpr is of the form (expr floordiv q) * q, where q is a
+  // symbolic expression.
+  auto lrhsBinOpExpr = lrhs.dyn_cast<AffineBinaryOpExpr>();
+  // Check rrhsConstOpExpr = -1.
+  auto rrhsConstOpExpr = rrhs.dyn_cast<AffineConstantExpr>();
+  if (rrhsConstOpExpr && rrhsConstOpExpr.getValue() == -1 && lrhsBinOpExpr &&
+      lrhsBinOpExpr.getKind() == AffineExprKind::Mul) {
+    // Check llrhs = expr floordiv q.
+    llrhs = lrhsBinOpExpr.getLHS();
+    // Check rlrhs = q.
+    rlrhs = lrhsBinOpExpr.getRHS();
+    auto llrhsBinOpExpr = llrhs.dyn_cast<AffineBinaryOpExpr>();
+    if (!llrhsBinOpExpr || llrhsBinOpExpr.getKind() != AffineExprKind::FloorDiv)
+      return nullptr;
+    if (llrhsBinOpExpr.getRHS() == rlrhs && lhs == llrhsBinOpExpr.getLHS())
+      return lhs % rlrhs;
+  }
+
   // Process lrhs, which is 'expr floordiv c'.
   AffineBinaryOpExpr lrBinOpExpr = lrhs.dyn_cast<AffineBinaryOpExpr>();
   if (!lrBinOpExpr || lrBinOpExpr.getKind() != AffineExprKind::FloorDiv)
     return nullptr;
 
-  auto llrhs = lrBinOpExpr.getLHS();
-  auto rlrhs = lrBinOpExpr.getRHS();
+  llrhs = lrBinOpExpr.getLHS();
+  rlrhs = lrBinOpExpr.getRHS();
 
   if (lhs == llrhs && rlrhs == -rrhs) {
     return lhs % rlrhs;
   }
   return nullptr;
 }
 
 AffineExpr AffineExpr::operator+(int64_t v) const {

mlir/test/Dialect/Affine/simplify-affine-structures.mlir

func @test_not_trivially_true_or_false_returning_three_results() -> (index, index, index) {
   }
   return %res#0, %res#1, %res#2 : index, index, index
 }
 
 // -----
 
 // Test simplification of mod expressions.
 // CHECK-DAG:   #[[$MOD:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s3 + s4 * s1 + (s0 - s1) mod s2)>
 // CHECK-DAG:   #[[$SIMPLIFIED_MOD_RHS:.*]] = affine_map<()[s0, s1, s2, s3] -> (s3 mod (s2 - s0 * s1))>
 // CHECK-DAG:   #[[$MODULO_AND_PRODUCT:.*]] = affine_map<()[s0, s1, s2, s3] -> (s0 * s1 + s3 - (-s0 + s3) mod s2)>
 // CHECK-LABEL: func @semiaffine_simplification_mod
 // CHECK-SAME:  (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index)
 func @semiaffine_simplification_mod(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) -> (index, index, index) {
   %a = affine.apply affine_map<(d0, d1)[s0, s1, s2, s3] -> ((-(d1 * s0 - (s0 - s1) mod s2) + s3) + (d0 * s1 + d1 * s0))>(%arg0, %arg1)[%arg2, %arg3, %arg4, %arg5]
   %b = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> (d0 mod (s0 - s1 * s2 + s3 - s0))>(%arg0)[%arg0, %arg1, %arg2, %arg3]
   %c = affine.apply affine_map<(d0)[s0, s1, s2] -> (d0 + (d0 + s0) mod s2 + s0 * s1 - (d0 + s0) mod s2 - (d0 - s0) mod s2)>(%arg0)[%arg1, %arg2, %arg3]
   return %a, %b, %c : index, index, index
 }
 // CHECK-NEXT: %[[RESULT0:.*]] = affine.apply #[[$MOD]]()[%[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG5]], %[[ARG0]]]
 // CHECK-NEXT: %[[RESULT1:.*]] = affine.apply #[[$SIMPLIFIED_MOD_RHS]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG0]]]
 // CHECK-NEXT: %[[RESULT2:.*]] = affine.apply #[[$MODULO_AND_PRODUCT]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG0]]]
 // CHECK-NEXT: return %[[RESULT0]], %[[RESULT1]], %[[RESULT2]]
 

 func @semiaffine_simplification_product(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) -> (index, index) {
   %a = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> ((-(s0 - (s0 - s1) * s2) + s3) + (d0 + s0))>(%arg0)[%arg1, %arg2, %arg3, %arg4]
   %b = affine.apply affine_map<(d0, d1, d2)[s0, s1] -> (d0 + d1 * s1 + d1 + d0 * s0 + d1 * s0 + d2 * s1 + d2)>(%arg0, %arg1, %arg2)[%arg3, %arg4]
   return %a, %b : index, index
 }
 // CHECK-NEXT: %[[RESULT0:.*]] = affine.apply #[[$PRODUCT]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG0]]]
 // CHECK-NEXT: %[[RESULT1:.*]] = affine.apply #[[$SUM_OF_PRODUCTS]]()[%[[ARG3]], %[[ARG4]], %[[ARG0]], %[[ARG1]], %[[ARG2]]]
 // CHECK-NEXT: return %[[RESULT0]], %[[RESULT1]]
+
+// -----
+
+// CHECK-DAG: #[[$SIMPLIFIED_MAP:.*]] = affine_map<()[s0, s1, s2, s3] -> ((-s0 + s2 + s3) mod (s0 + s1))>
+// CHECK-LABEL: func @semi_affine_simplification_euclidean_lemma
+// CHECK-SAME: (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index)
+func @semi_affine_simplification_euclidean_lemma(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) -> (index, index) {
+  %a = affine.apply affine_map<(d0, d1)[s0, s1] -> ((d0 + d1) - ((d0 + d1) floordiv (s0 - s1)) * (s0 - s1) - (d0 + d1) mod (s0 - s1))>(%arg0, %arg1)[%arg2, %arg3]
+  %b = affine.apply affine_map<(d0, d1)[s0, s1] -> ((d0 + d1 - s0) - ((d0 + d1 - s0) floordiv (s0 + s1)) * (s0 + s1))>(%arg0, %arg1)[%arg2, %arg3]
+  return %a, %b : index, index
+}
+// CHECK-NEXT: %[[ZERO:.*]] = arith.constant 0 : index
+// CHECK-NEXT: %[[RESULT:.*]] = affine.apply #[[$SIMPLIFIED_MAP]]()[%[[ARG2]], %[[ARG3]], %[[ARG0]], %[[ARG1]]]
+// CHECK-NEXT: return %[[ZERO]], %[[RESULT]]