[Mlir-commits] [mlir] 3e497a1 - [MLIR] Update/fix memref region computation for affine.parallel ops
Uday Bondhugula
llvmlistbot at llvm.org
Tue Mar 14 18:19:11 PDT 2023
Author: Uday Bondhugula
Date: 2023-03-15T06:40:24+05:30
New Revision: 3e497a114730cebdf96f3556af081d876b6ac8b3
URL: https://github.com/llvm/llvm-project/commit/3e497a114730cebdf96f3556af081d876b6ac8b3
DIFF: https://github.com/llvm/llvm-project/commit/3e497a114730cebdf96f3556af081d876b6ac8b3.diff
LOG: [MLIR] Update/fix memref region computation for affine.parallel ops
When the affine.parallel op was introduced, affine utilities weren't
extended to handle it. Extending these is straightforward and natural
given that addAffineParallelOpDomain has also been added.
Update/complete memref region compute to account for affine.parallel
ops. Handle failure cleanly.
Add and expose utilities missing for affine.parallel to be consistent
with affine.for.
All of these allow various affine passes to work with a combination of
affine.parallel and affine.for ops.
Differential Revision: https://reviews.llvm.org/D145669
Added:
Modified:
mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
mlir/lib/Dialect/Affine/Analysis/Utils.cpp
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
mlir/test/Transforms/memref-bound-check.mlir
mlir/test/lib/Dialect/Affine/TestAffineDataCopy.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
index 72850bf9ab51..7f6eced071b1 100644
--- a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
@@ -39,6 +39,10 @@ class Value;
/// from the outermost 'affine.for' operation to the innermost one.
void getAffineForIVs(Operation &op, SmallVectorImpl<AffineForOp> *loops);
+/// Populates 'ivs' with IVs of the surrounding affine.for and affine.parallel
+/// ops ordered from the outermost one to the innermost.
+void getAffineIVs(Operation &op, SmallVectorImpl<Value> &ivs);
+
/// Populates 'ops' with affine operations enclosing `op` ordered from outermost
/// to innermost. affine.for, affine.if, or affine.parallel ops comprise such
/// surrounding affine ops.
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
index 7bf78cd369f2..cda38a558064 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
@@ -441,10 +441,21 @@ namespace mlir {
/// AffineForOp.
bool isAffineForInductionVar(Value val);
+/// Returns true if `val` is the induction variable of an AffineParallelOp.
+bool isAffineParallelInductionVar(Value val);
+
+/// Returns true if the provided value is the induction variable of an
+/// AffineForOp or AffineParallelOp.
+bool isAffineInductionVar(Value val);
+
/// Returns the loop parent of an induction variable. If the provided value is
/// not an induction variable, then return nullptr.
AffineForOp getForInductionVarOwner(Value val);
+/// Returns true if the provided value is among the induction variables of an
+/// AffineParallelOp.
+AffineParallelOp getAffineParallelInductionVarOwner(Value val);
+
/// Extracts the induction variables from a list of AffineForOps and places them
/// in the output argument `ivs`.
void extractForInductionVars(ArrayRef<AffineForOp> forInsts,
diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
index e702d2f37f3b..9e0dbee83b76 100644
--- a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
@@ -287,14 +287,6 @@ static LogicalResult getOpIndexSet(Operation *op,
return getIndexSet(ops, indexSet);
}
-/// Returns true if `val` is an induction of an affine.parallel op.
-static bool isAffineParallelInductionVar(Value val) {
- auto ivArg = val.dyn_cast<BlockArgument>();
- if (!ivArg)
- return false;
- return isa<AffineParallelOp>(ivArg.getOwner()->getParentOp());
-}
-
// Returns the number of outer loop common to 'src/dstDomain'.
// Loops common to 'src/dst' domains are added to 'commonLoops' if non-null.
static unsigned
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 581ffb0ad2f6..0ae40826de8b 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -454,19 +454,17 @@ LogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
unsigned rank = access.getRank();
LLVM_DEBUG(llvm::dbgs() << "MemRefRegion::compute: " << *op
- << "depth: " << loopDepth << "\n";);
+ << "\ndepth: " << loopDepth << "\n";);
// 0-d memrefs.
if (rank == 0) {
- SmallVector<AffineForOp, 4> ivs;
- getAffineForIVs(*op, &ivs);
+ SmallVector<Value, 4> ivs;
+ getAffineIVs(*op, ivs);
assert(loopDepth <= ivs.size() && "invalid 'loopDepth'");
// The first 'loopDepth' IVs are symbols for this region.
ivs.resize(loopDepth);
- SmallVector<Value, 4> regionSymbols;
- extractForInductionVars(ivs, ®ionSymbols);
// A 0-d memref has a 0-d region.
- cst.reset(rank, loopDepth, /*numLocals=*/0, regionSymbols);
+ cst.reset(rank, loopDepth, /*numLocals=*/0, ivs);
return success();
}
@@ -503,23 +501,24 @@ LogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
// Add inequalities for loop lower/upper bounds.
for (unsigned i = 0; i < numDims + numSymbols; ++i) {
auto operand = operands[i];
- if (auto loop = getForInductionVarOwner(operand)) {
+ if (auto affineFor = getForInductionVarOwner(operand)) {
// Note that cst can now have more dimensions than accessMap if the
// bounds expressions involve outer loops or other symbols.
// TODO: rewrite this to use getInstIndexSet; this way
// conditionals will be handled when the latter supports it.
- if (failed(cst.addAffineForOpDomain(loop)))
+ if (failed(cst.addAffineForOpDomain(affineFor)))
return failure();
- } else {
- // Has to be a valid symbol.
- auto symbol = operand;
- assert(isValidSymbol(symbol));
+ } else if (auto parallelOp = getAffineParallelInductionVarOwner(operand)) {
+ if (failed(cst.addAffineParallelOpDomain(parallelOp)))
+ return failure();
+ } else if (isValidSymbol(operand)) {
// Check if the symbol is a constant.
- if (auto *op = symbol.getDefiningOp()) {
- if (auto constOp = dyn_cast<arith::ConstantIndexOp>(op)) {
- cst.addBound(FlatAffineValueConstraints::EQ, symbol, constOp.value());
- }
- }
+ Value symbol = operand;
+ if (auto constVal = getConstantIntValue(symbol))
+ cst.addBound(FlatAffineValueConstraints::EQ, symbol, constVal.value());
+ } else {
+ LLVM_DEBUG(llvm::dbgs() << "unknown affine dimensional value");
+ return failure();
}
}
@@ -552,16 +551,14 @@ LogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
// Eliminate any loop IVs other than the outermost 'loopDepth' IVs, on which
// this memref region is symbolic.
- SmallVector<AffineForOp, 4> enclosingIVs;
- getAffineForIVs(*op, &enclosingIVs);
+ SmallVector<Value, 4> enclosingIVs;
+ getAffineIVs(*op, enclosingIVs);
assert(loopDepth <= enclosingIVs.size() && "invalid loop depth");
enclosingIVs.resize(loopDepth);
SmallVector<Value, 4> vars;
cst.getValues(cst.getNumDimVars(), cst.getNumDimAndSymbolVars(), &vars);
- for (auto var : vars) {
- AffineForOp iv;
- if ((iv = getForInductionVarOwner(var)) &&
- !llvm::is_contained(enclosingIVs, iv)) {
+ for (Value var : vars) {
+ if ((isAffineInductionVar(var)) && !llvm::is_contained(enclosingIVs, var)) {
cst.projectOut(var);
}
}
@@ -1264,21 +1261,19 @@ bool MemRefAccess::operator==(const MemRefAccess &rhs) const {
[](AffineExpr e) { return e == 0; });
}
-/// Populates 'loops' with IVs of the surrounding affine.for and affine.parallel
-/// ops ordered from the outermost one to the innermost.
-static void getAffineIVs(Operation &op, SmallVectorImpl<Value> &loops) {
+void mlir::getAffineIVs(Operation &op, SmallVectorImpl<Value> &ivs) {
auto *currOp = op.getParentOp();
AffineForOp currAffineForOp;
- // Traverse up the hierarchy collecting all 'affine.for' operation while
- // skipping over 'affine.if' operations.
+ // Traverse up the hierarchy collecting all 'affine.for' and affine.parallel
+ // operation while skipping over 'affine.if' operations.
while (currOp) {
if (AffineForOp currAffineForOp = dyn_cast<AffineForOp>(currOp))
- loops.push_back(currAffineForOp.getInductionVar());
+ ivs.push_back(currAffineForOp.getInductionVar());
else if (auto parOp = dyn_cast<AffineParallelOp>(currOp))
- llvm::append_range(loops, parOp.getIVs());
+ llvm::append_range(ivs, parOp.getIVs());
currOp = currOp->getParentOp();
}
- std::reverse(loops.begin(), loops.end());
+ std::reverse(ivs.begin(), ivs.end());
}
/// Returns the number of surrounding loops common to 'loopsA' and 'loopsB',
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 8b46f530ae39..5ed434fa1e80 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2504,8 +2504,14 @@ bool mlir::isAffineForInductionVar(Value val) {
return getForInductionVarOwner(val) != AffineForOp();
}
-/// Returns the loop parent of an induction variable. If the provided value is
-/// not an induction variable, then return nullptr.
+bool mlir::isAffineParallelInductionVar(Value val) {
+ return getAffineParallelInductionVarOwner(val) != nullptr;
+}
+
+bool mlir::isAffineInductionVar(Value val) {
+ return isAffineForInductionVar(val) || isAffineParallelInductionVar(val);
+}
+
AffineForOp mlir::getForInductionVarOwner(Value val) {
auto ivArg = val.dyn_cast<BlockArgument>();
if (!ivArg || !ivArg.getOwner())
@@ -2517,6 +2523,17 @@ AffineForOp mlir::getForInductionVarOwner(Value val) {
return AffineForOp();
}
+AffineParallelOp mlir::getAffineParallelInductionVarOwner(Value val) {
+ auto ivArg = val.dyn_cast<BlockArgument>();
+ if (!ivArg || !ivArg.getOwner())
+ return nullptr;
+ Operation *containingOp = ivArg.getOwner()->getParentOp();
+ auto parallelOp = dyn_cast<AffineParallelOp>(containingOp);
+ if (parallelOp && llvm::is_contained(parallelOp.getIVs(), val))
+ return parallelOp;
+ return nullptr;
+}
+
/// Extracts the induction variables from a list of AffineForOps and returns
/// them.
void mlir::extractForInductionVars(ArrayRef<AffineForOp> forInsts,
diff --git a/mlir/test/Transforms/memref-bound-check.mlir b/mlir/test/Transforms/memref-bound-check.mlir
index 5c707e6900a4..fce6bdbca4aa 100644
--- a/mlir/test/Transforms/memref-bound-check.mlir
+++ b/mlir/test/Transforms/memref-bound-check.mlir
@@ -293,3 +293,14 @@ func.func @zero_d_memref() {
}
return
}
+
+// CHECK-LABEL: func @affine_parallel
+func.func @affine_parallel(%M: memref<2048x2048xf64>) {
+ affine.parallel (%i) = (0) to (3000) {
+ affine.for %j = 0 to 2048 {
+ affine.load %M[%i, %j] : memref<2048x2048xf64>
+ // expected-error at above {{'affine.load' op memref out of upper bound access along dimension #1}}
+ }
+ }
+ return
+}
diff --git a/mlir/test/lib/Dialect/Affine/TestAffineDataCopy.cpp b/mlir/test/lib/Dialect/Affine/TestAffineDataCopy.cpp
index 62671dde34b0..c602095bbfee 100644
--- a/mlir/test/lib/Dialect/Affine/TestAffineDataCopy.cpp
+++ b/mlir/test/lib/Dialect/Affine/TestAffineDataCopy.cpp
@@ -60,7 +60,8 @@ void TestAffineDataCopy::runOnOperation() {
// Gather all AffineForOps by loop depth.
std::vector<SmallVector<AffineForOp, 2>> depthToLoops;
gatherLoops(getOperation(), depthToLoops);
- assert(!depthToLoops.empty() && "Loop nest not found");
+ if (depthToLoops.empty())
+ return;
// Only support tests with a single loop nest and a single innermost loop
// for now.
More information about the Mlir-commits
mailing list