[Mlir-commits] [mlir] [mlir][affine] Add arith.subi support to memref to affine raise (PR #222880)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mattéo Rizza Murgier (matteo-rm)

<details>
<summary>Changes</summary>

This PR adds support for `arith.subi` to `toAffineExpr` to e.g. allow matching expressions like `i-1` which appear with one-indexed arrays.

---
Full diff: https://github.com/llvm/llvm-project/pull/222880.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Transforms/RaiseMemrefDialect.cpp (+9-6) 
- (modified) mlir/test/Dialect/Affine/raise-memref.mlir (+59-1) 


``````````diff
diff --git a/mlir/lib/Dialect/Affine/Transforms/RaiseMemrefDialect.cpp b/mlir/lib/Dialect/Affine/Transforms/RaiseMemrefDialect.cpp
index 29cc256cb90f9..9c651e9a86910 100644
--- a/mlir/lib/Dialect/Affine/Transforms/RaiseMemrefDialect.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/RaiseMemrefDialect.cpp
@@ -75,14 +75,14 @@ static AffineExpr toAffineExpr(Value value,
   }
 
   Operation *definingOp = value.getDefiningOp();
-  if (llvm::isa_and_nonnull<arith::AddIOp>(definingOp) ||
-      llvm::isa_and_nonnull<arith::MulIOp>(definingOp)) {
+  if (llvm::isa_and_nonnull<arith::AddIOp, arith::SubIOp, arith::MulIOp>(
+          definingOp)) {
     // TODO: replace recursion with explicit stack.
     // For the moment this can be tolerated as we only recurse on
-    // arith.addi and arith.muli, so there cannot be any infinite
-    // recursion. The depth of these expressions should be in most
-    // cases very manageable, as affine expressions should be as
-    // simple as `a + b * c`.
+    // arith.addi, arith.subi and arith.muli, so there cannot be
+    // any infinite recursion. The depth of these expressions
+    // should be in most cases very manageable, as affine
+    // expressions should be as simple as `a + b * c`.
     AffineExpr lhsE =
         toAffineExpr(definingOp->getOperand(0), affineDims, affineSymbols);
     AffineExpr rhsE =
@@ -92,6 +92,9 @@ static AffineExpr toAffineExpr(Value value,
       AffineExprKind kind;
       if (isa<arith::AddIOp>(definingOp)) {
         kind = mlir::AffineExprKind::Add;
+      } else if (isa<arith::SubIOp>(definingOp)) {
+        kind = mlir::AffineExprKind::Add;
+        rhsE = -rhsE;
       } else {
         kind = mlir::AffineExprKind::Mul;
 
diff --git a/mlir/test/Dialect/Affine/raise-memref.mlir b/mlir/test/Dialect/Affine/raise-memref.mlir
index 8dc24d99db3e2..4f3497999a52f 100644
--- a/mlir/test/Dialect/Affine/raise-memref.mlir
+++ b/mlir/test/Dialect/Affine/raise-memref.mlir
@@ -117,7 +117,6 @@ func.func @symbols(%N : index) {
 // CHECK:                  memref.store %[[lhs5]], %{{.*}}[%[[a1]], %[[lhs7]]] :
 // CHECK:                  affine.yield %[[lhs6]]
 
-
 // CHECK-LABEL:    func @non_affine(
 func.func @non_affine(%N : index) {
   %2 = memref.alloc() : memref<1024x1024xf32>
@@ -136,3 +135,62 @@ func.func @non_affine(%N : index) {
 // CHECK:                  %[[ij:.*]] = arith.muli %[[i]], %[[j]]
 // CHECK:                  %[[v:.*]] = memref.load %{{.*}}[%[[i]], %[[ij]]]
 // CHECK:                  memref.store %[[v]], %{{.*}}[%[[ij]], %[[ij]]]
+
+// CHECK-LABEL:    func @subtract_dims(
+func.func @subtract_dims(%N : index) {
+  %0 = memref.alloc() : memref<1024xf32>
+  affine.for %i = 0 to %N {
+    affine.for %j = 0 to %N {
+      %idx = arith.subi %i, %j : index
+      %v = memref.load %0[%idx] : memref<1024xf32>
+      memref.store %v, %0[%idx] : memref<1024xf32>
+    }
+  }
+  return
+}
+
+// CHECK:          affine.for %[[i:.*]] =
+// CHECK:            affine.for %[[j:.*]] =
+// CHECK:              %[[v:.*]] = affine.load %{{.*}}[%[[i]] - %[[j]]]
+// CHECK:              affine.store %[[v]], %{{.*}}[%[[i]] - %[[j]]]
+// CHECK-NOT:          memref.load
+// CHECK-NOT:          memref.store
+
+// CHECK-LABEL:    func @subtract_constant(
+func.func @subtract_constant(%N : index) {
+  %c1 = arith.constant 1 : index
+  %0 = memref.alloc() : memref<1024xf32>
+  affine.for %i = 1 to %N {
+    %idx = arith.subi %i, %c1 : index
+    %v = memref.load %0[%idx] : memref<1024xf32>
+    memref.store %v, %0[%idx] : memref<1024xf32>
+  }
+  return
+}
+
+// CHECK:          affine.for %[[i:.*]] =
+// CHECK:            %[[v:.*]] = affine.load %{{.*}}[%[[i]] - 1]
+// CHECK:            affine.store %[[v]], %{{.*}}[%[[i]] - 1]
+// CHECK-NOT:        memref.load
+
+// CHECK-LABEL:    func @subtract_from_non_affine(
+func.func @subtract_from_non_affine(%N : index) {
+  %c1 = arith.constant 1 : index
+  %0 = memref.alloc() : memref<1024xf32>
+  affine.for %i = 0 to %N {
+    affine.for %j = 0 to %N {
+      %ij = arith.muli %i, %j : index
+      %idx = arith.subi %ij, %c1 : index
+      %v = memref.load %0[%idx] : memref<1024xf32>
+      memref.store %v, %0[%idx] : memref<1024xf32>
+    }
+  }
+  return
+}
+
+// CHECK:          affine.for %[[i:.*]] =
+// CHECK:            affine.for %[[j:.*]] =
+// CHECK:              %[[ij:.*]] = arith.muli %[[i]], %[[j]]
+// CHECK:              %[[idx:.*]] = arith.subi %[[ij]]
+// CHECK:              %[[v:.*]] = memref.load %{{.*}}[%[[idx]]]
+// CHECK:              memref.store %[[v]], %{{.*}}[%[[idx]]]

``````````

</details>


https://github.com/llvm/llvm-project/pull/222880