From ec563f0fac7c9107f340bc258f207a5253432db7 Mon Sep 17 00:00:00 2001 From: Uday Bondhugula Date: Fri, 7 Mar 2025 10:18:33 +0530 Subject: [PATCH] [MLIR][Affine] Fix scalrep and underlying analysis utility 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. --- .../mlir/Dialect/Affine/IR/AffineOps.h | 7 ++++ .../Affine/Analysis/AffineAnalysis.cpp | 3 +- mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 4 +- mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 10 +++++ mlir/lib/Dialect/Affine/Utils/Utils.cpp | 7 ++-- mlir/test/Dialect/Affine/scalrep.mlir | 41 ++++++++++++++++--- 6 files changed, 61 insertions(+), 11 deletions(-) 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 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(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, %arg1: memref, %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, %arg1: i32) -> i32 attributes {llvm.linkage = #llvm.linkage} { + %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 + %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 +}