[Mlir-commits] [mlir] Lower affine modulo by powers of two using bitwise AND (PR #146311)

[Yuxi Sun]

https://github.com/sherylll updated https://github.com/llvm/llvm-project/pull/146311

>From db4de47d6af94d48d90db9c5812a8dcb32aa6cc2 Mon Sep 17 00:00:00 2001
From: sherylll <sherylsun14 at gmail.com>
Date: Mon, 30 Jun 2025 05:48:26 +0000
Subject: [PATCH 1/2] Lower affine modulo by powers of two using bitwise AND

---
 mlir/lib/Dialect/Affine/Utils/Utils.cpp           | 15 +++++++++++++++
 .../Conversion/AffineToStandard/lower-affine.mlir |  9 +++++++++
 2 files changed, 24 insertions(+)

diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 66b3f2a4f93a5..0cffe52dd0615 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -80,12 +80,27 @@ class AffineApplyExpander
   ///         let remainder = srem a, b;
   ///             negative = a < 0 in
   ///         select negative, remainder + b, remainder.
+  ///
+  /// Special case for power of 2: 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) {
         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) {
+        auto lhs = visit(expr.getLHS());
+        assert(lhs && "unexpected affine expr lowering failure");
+
+        Value maskCst =
+            builder.create<arith::ConstantIndexOp>(loc, rhsValue - 1);
+        return builder.create<arith::AndIOp>(loc, lhs, maskCst);
+      }
     }
 
     auto lhs = visit(expr.getLHS());
diff --git a/mlir/test/Conversion/AffineToStandard/lower-affine.mlir b/mlir/test/Conversion/AffineToStandard/lower-affine.mlir
index 550ea71882e14..07f7c64fe6ea5 100644
--- a/mlir/test/Conversion/AffineToStandard/lower-affine.mlir
+++ b/mlir/test/Conversion/AffineToStandard/lower-affine.mlir
@@ -927,3 +927,12 @@ func.func @affine_parallel_with_reductions_i64(%arg0: memref<3x3xi64>, %arg1: me
 // CHECK:      scf.reduce.return %[[RES]] : i64
 // 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)
+  return %0 : index
+}

>From 1624c8407f80b7914cb9443aa3e32b089cdc9123 Mon Sep 17 00:00:00 2001
From: sherylll <sherylsun14 at gmail.com>
Date: Sat, 1 Nov 2025 06:19:30 +0000
Subject: [PATCH 2/2] Fix test expectation for affine.apply mod 1 (always
 returns 0)

---
 mlir/lib/Dialect/Affine/Utils/Utils.cpp               | 10 +++++-----
 .../Conversion/AffineToStandard/lower-affine.mlir     | 11 +++++++++--
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index a81ec13d5707a..234c004d23c23 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/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");
 
diff --git a/mlir/test/Conversion/AffineToStandard/lower-affine.mlir b/mlir/test/Conversion/AffineToStandard/lower-affine.mlir
index 07f7c64fe6ea5..063de9f8bea39 100644
--- a/mlir/test/Conversion/AffineToStandard/lower-affine.mlir
+++ b/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
 }