diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h --- a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h +++ b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h @@ -39,6 +39,10 @@ /// from the outermost 'affine.for' operation to the innermost one. void getAffineForIVs(Operation &op, SmallVectorImpl *loops); +/// Populates 'loops' with IVs of the surrounding affine.for and affine.parallel +/// ops ordered from the outermost one to the innermost. +void getAffineIVs(Operation &op, SmallVectorImpl &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 --- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h +++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h @@ -441,10 +441,21 @@ /// 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 forInsts, diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp --- a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp +++ b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp @@ -287,14 +287,6 @@ 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(); - if (!ivArg) - return false; - return isa(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 --- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp +++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp @@ -454,19 +454,17 @@ unsigned rank = access.getRank(); LLVM_DEBUG(llvm::dbgs() << "MemRefRegion::compute: " << *op - << "depth: " << loopDepth << "\n";); + << "\ndepth: " << loopDepth << "\n";); // 0-d memrefs. if (rank == 0) { - SmallVector ivs; - getAffineForIVs(*op, &ivs); + SmallVector ivs; + getAffineIVs(*op, ivs); assert(loopDepth <= ivs.size() && "invalid 'loopDepth'"); // The first 'loopDepth' IVs are symbols for this region. ivs.resize(loopDepth); - SmallVector 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,26 @@ // 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(op)) { + Value symbol = operand; + if (Operation *op = symbol.getDefiningOp()) { + if (auto constOp = dyn_cast(op)) cst.addBound(FlatAffineValueConstraints::EQ, symbol, constOp.value()); - } } + } else { + LLVM_DEBUG(llvm::dbgs() << "unknown affine dimensional value"); + return failure(); } } @@ -552,16 +553,14 @@ // Eliminate any loop IVs other than the outermost 'loopDepth' IVs, on which // this memref region is symbolic. - SmallVector enclosingIVs; - getAffineForIVs(*op, &enclosingIVs); + SmallVector enclosingIVs; + getAffineIVs(*op, enclosingIVs); assert(loopDepth <= enclosingIVs.size() && "invalid loop depth"); enclosingIVs.resize(loopDepth); SmallVector 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 +1263,19 @@ [](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 &loops) { +void mlir::getAffineIVs(Operation &op, SmallVectorImpl &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(currOp)) - loops.push_back(currAffineForOp.getInductionVar()); + ivs.push_back(currAffineForOp.getInductionVar()); else if (auto parOp = dyn_cast(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 --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -2504,8 +2504,14 @@ 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(); if (!ivArg || !ivArg.getOwner()) @@ -2517,6 +2523,17 @@ return AffineForOp(); } +AffineParallelOp mlir::getAffineParallelInductionVarOwner(Value val) { + auto ivArg = val.dyn_cast(); + if (!ivArg || !ivArg.getOwner()) + return nullptr; + Operation *containingOp = ivArg.getOwner()->getParentOp(); + auto parallelOp = dyn_cast(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 forInsts, diff --git a/mlir/test/Transforms/memref-bound-check.mlir b/mlir/test/Transforms/memref-bound-check.mlir --- a/mlir/test/Transforms/memref-bound-check.mlir +++ b/mlir/test/Transforms/memref-bound-check.mlir @@ -293,3 +293,14 @@ } 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@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 --- a/mlir/test/lib/Dialect/Affine/TestAffineDataCopy.cpp +++ b/mlir/test/lib/Dialect/Affine/TestAffineDataCopy.cpp @@ -60,7 +60,8 @@ // Gather all AffineForOps by loop depth. std::vector> 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.