Fix test expectation for affine.apply mod 1 (always returns 0)

Author: sherylll

mlir/lib/Dialect/Affine/Utils/Utils.cpp

@@ -79,19 +79,19 @@ class AffineApplyExpander
   ///             negative = a < 0 in
   ///         select negative, remainder + b, remainder.
   ///
-  /// Special case for power of 2: use bitwise AND (x & (n-1)) for non-negative
+  /// Special case for power-of-2 RHS: use bitwise AND (x & (n-1)) for non-negative
   /// x.
   Value visitModExpr(AffineBinaryOpExpr expr) {
     if (auto rhsConst = dyn_cast<AffineConstantExpr>(expr.getRHS())) {
-      if (rhsConst.getValue() <= 0) {
+      int64_t rhsValue = rhsConst.getValue();
+      if (rhsValue <= 0) {
         emitError(loc, "modulo by non-positive value is not supported");
         return nullptr;
       }
 
       // Special case: x mod n where n is a power of 2 can be optimized to x &
-      // (n-1)
-      int64_t rhsValue = rhsConst.getValue();
-      if (rhsValue > 0 && (rhsValue & (rhsValue - 1)) == 0) {
+      // (n-1).
+      if ((rhsValue & (rhsValue - 1)) == 0) {
         auto lhs = visit(expr.getLHS());
         assert(lhs && "unexpected affine expr lowering failure");
 

mlir/test/Conversion/AffineToStandard/lower-affine.mlir

@@ -928,11 +928,18 @@ func.func @affine_parallel_with_reductions_i64(%arg0: memref<3x3xi64>, %arg1: me
 // CHECK:    }
 // CHECK:  }
 
-#map_mod_8 = affine_map<(i) -> (i mod 8)>
 // CHECK-LABEL: func @affine_apply_mod_8
 func.func @affine_apply_mod_8(%arg0 : index) -> (index) {
   // CHECK-NEXT: %[[c7:.*]] = arith.constant 7 : index
   // CHECK-NEXT: %[[v0:.*]] = arith.andi %arg0, %[[c7]] : index
-  %0 = affine.apply #map_mod_8 (%arg0)
+  %0 = affine.apply affine_map<(i) -> (i mod 8)> (%arg0)
+  return %0 : index
+}
+
+// CHECK-LABEL: func @affine_apply_mod_1
+func.func @affine_apply_mod_1(%arg0 : index) -> (index) {
+  // CHECK-NEXT: %[[c0:.*]] = arith.constant 0 : index
+  // CHECK-NEXT: return %[[c0]] : index
+  %0 = affine.apply affine_map<(i) -> (i mod 1)> (%arg0)
   return %0 : index
 }