[Mlir-commits] [mlir] [MLIR][Affine] Fix scalrep and underlying analysis utility (PR #130251)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Mar 6 22:04:07 PST 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-affine
Author: Uday Bondhugula (bondhugula)
<details>
<summary>Changes</summary>
Fixes: https://github.com/llvm/llvm-project/issues/53034
Properly check the scopes of affine operations across which we are
performing affine analysis. Fixes transforms and analyses in the
presence of non-affine regions.
---
Full diff: https://github.com/llvm/llvm-project/pull/130251.diff
6 Files Affected:
- (modified) mlir/include/mlir/Dialect/Affine/IR/AffineOps.h (+7)
- (modified) mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp (+2-1)
- (modified) mlir/lib/Dialect/Affine/Analysis/Utils.cpp (+2-2)
- (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+10)
- (modified) mlir/lib/Dialect/Affine/Utils/Utils.cpp (+4-3)
- (modified) mlir/test/Dialect/Affine/scalrep.mlir (+36-5)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
index 7c950623f77f4..bbf38f2293448 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
@@ -46,6 +46,13 @@ bool isTopLevelValue(Value value, Region *region);
/// trait `AffineScope`; `nullptr` if there is no such region.
Region *getAffineScope(Operation *op);
+/// Returns the closest region enclosing `op` that is held by a non-affine
+/// operation; `nullptr` if there is no such region. This method is meant to
+/// be used by affine analysis methods (e.g. dependence analysis) which are
+/// only meaningful when performed among/between operations from the same
+/// analysis scope.
+Region *getAffineAnalysisScope(Operation *op);
+
/// AffineDmaStartOp starts a non-blocking DMA operation that transfers data
/// from a source memref to a destination memref. The source and destination
/// memref need not be of the same dimensionality, but need to have the same
diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
index 9b776900c379a..84b76d33c3e67 100644
--- a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
@@ -626,7 +626,8 @@ DependenceResult mlir::affine::checkMemrefAccessDependence(
// We can't analyze further if the ops lie in different affine scopes or have
// no common block in an affine scope.
- if (getAffineScope(srcAccess.opInst) != getAffineScope(dstAccess.opInst))
+ if (getAffineAnalysisScope(srcAccess.opInst) !=
+ getAffineAnalysisScope(dstAccess.opInst))
return DependenceResult::Failure;
if (!getCommonBlockInAffineScope(srcAccess.opInst, dstAccess.opInst))
return DependenceResult::Failure;
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 56a7d8de1b7f8..cf9eaa9e7d66d 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -2307,8 +2307,8 @@ FailureOr<AffineValueMap> mlir::affine::simplifyConstrainedMinMaxOp(
Block *mlir::affine::findInnermostCommonBlockInScope(Operation *a,
Operation *b) {
- Region *aScope = mlir::affine::getAffineScope(a);
- Region *bScope = mlir::affine::getAffineScope(b);
+ Region *aScope = getAffineAnalysisScope(a);
+ Region *bScope = getAffineAnalysisScope(b);
if (aScope != bScope)
return nullptr;
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 06493de2b09d4..18e980cd8b88e 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -270,6 +270,16 @@ Region *mlir::affine::getAffineScope(Operation *op) {
return nullptr;
}
+Region *mlir::affine::getAffineAnalysisScope(Operation *op) {
+ Operation *curOp = op;
+ while (auto *parentOp = curOp->getParentOp()) {
+ if (!isa<AffineForOp, AffineIfOp, AffineParallelOp>(parentOp))
+ return curOp->getParentRegion();
+ curOp = parentOp;
+ }
+ return nullptr;
+}
+
// A Value can be used as a dimension id iff it meets one of the following
// conditions:
// *) It is valid as a symbol.
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 7ef016f88be37..009b1b2c22c94 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -636,7 +636,8 @@ static bool mustReachAtInnermost(const MemRefAccess &srcAccess,
const MemRefAccess &destAccess) {
// Affine dependence analysis is possible only if both ops in the same
// AffineScope.
- if (getAffineScope(srcAccess.opInst) != getAffineScope(destAccess.opInst))
+ if (getAffineAnalysisScope(srcAccess.opInst) !=
+ getAffineAnalysisScope(destAccess.opInst))
return false;
unsigned nsLoops =
@@ -659,9 +660,9 @@ static bool mayHaveEffect(Operation *srcMemOp, Operation *destMemOp,
// AffineScope. Also, we can only check if our affine scope is isolated from
// above; otherwise, values can from outside of the affine scope that the
// check below cannot analyze.
- Region *srcScope = getAffineScope(srcMemOp);
+ Region *srcScope = getAffineAnalysisScope(srcMemOp);
if (srcAccess.memref == destAccess.memref &&
- srcScope == getAffineScope(destMemOp)) {
+ srcScope == getAffineAnalysisScope(destMemOp)) {
unsigned nsLoops = getNumCommonSurroundingLoops(*srcMemOp, *destMemOp);
FlatAffineValueConstraints dependenceConstraints;
for (unsigned d = nsLoops + 1; d > minSurroundingLoops; d--) {
diff --git a/mlir/test/Dialect/Affine/scalrep.mlir b/mlir/test/Dialect/Affine/scalrep.mlir
index fdfe3bfb62f95..092597860c8d9 100644
--- a/mlir/test/Dialect/Affine/scalrep.mlir
+++ b/mlir/test/Dialect/Affine/scalrep.mlir
@@ -738,12 +738,11 @@ func.func @with_inner_ops(%arg0: memref<?xf64>, %arg1: memref<?xf64>, %arg2: i1)
return
}
-// CHECK: %[[pi:.+]] = arith.constant 3.140000e+00 : f64
-// CHECK: %{{.*}} = scf.if %arg2 -> (f64) {
-// CHECK: scf.yield %{{.*}} : f64
+// Semantics of non-affine region ops would be unknown.
+
// CHECK: } else {
-// CHECK: scf.yield %[[pi]] : f64
-// CHECK: }
+// CHECK-NEXT: %[[Y:.*]] = affine.load
+// CHECK-NEXT: scf.yield %[[Y]] : f64
// Check if scalar replacement works correctly when affine memory ops are in the
// body of an scf.for.
@@ -952,3 +951,35 @@ func.func @consecutive_store() {
}
return
}
+
+// CHECK-LABEL: func @scf_for_if
+func.func @scf_for_if(%arg0: memref<?xi32>, %arg1: i32) -> i32 attributes {llvm.linkage = #llvm.linkage<external>} {
+ %c1 = arith.constant 1 : index
+ %c0 = arith.constant 0 : index
+ %c0_i32 = arith.constant 0 : i32
+ %c5_i32 = arith.constant 5 : i32
+ %c10_i32 = arith.constant 10 : i32
+ %0 = memref.alloca() : memref<1xi32>
+ %1 = llvm.mlir.undef : i32
+ affine.store %1, %0[0] : memref<1xi32>
+ affine.store %c0_i32, %0[0] : memref<1xi32>
+ %2 = arith.index_cast %arg1 : i32 to index
+ scf.for %arg2 = %c0 to %2 step %c1 {
+ %4 = memref.load %arg0[%arg2] : memref<?xi32>
+ %5 = arith.muli %4, %c5_i32 : i32
+ %6 = arith.cmpi sgt, %5, %c10_i32 : i32
+ // CHECK: scf.if
+ scf.if %6 {
+ // No forwarding should happen here since we have an scf.for around and we
+ // can't analyze the flow of values.
+ // CHECK: affine.load
+ %7 = affine.load %0[0] : memref<1xi32>
+ %8 = arith.addi %5, %7 : i32
+ // CHECK: affine.store
+ affine.store %8, %0[0] : memref<1xi32>
+ }
+ }
+ // CHECK: affine.load
+ %3 = affine.load %0[0] : memref<1xi32>
+ return %3 : i32
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/130251
More information about the Mlir-commits
mailing list