diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h --- a/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h +++ b/mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h @@ -24,6 +24,7 @@ class AffineCondition; class AffineForOp; class AffineIfOp; +class AffineParallelOp; class AffineMap; class AffineValueMap; class IntegerSet; @@ -141,6 +142,13 @@ // TODO: add support for non-unit strides. LogicalResult addAffineForOpDomain(AffineForOp forOp); + /// Add constraints (lower and upper bounds) for the specified + /// 'affine.parallel' operation's Value using IR information stored in its + /// bound maps. Returns failure for the yet unimplemented/unsupported cases. + /// Asserts if the Value corresponding to the 'affine.parallel' operation + /// isn't found in the constraint system. + LogicalResult addAffineParallelOpDomain(AffineParallelOp parallelOp); + /// Adds constraints (lower and upper bounds) for each loop in the loop nest /// described by the bound maps `lbMaps` and `ubMaps` of a computation slice. /// Every pair (`lbMaps[i]`, `ubMaps[i]`) describes the bounds of a loop in 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 @@ -45,6 +45,10 @@ /// trait `AffineScope`; `nullptr` if there is no such region. Region *getAffineScope(Operation *op); +/// Returns whether the op is within an affine parallel op which can potentially +/// preventing the dependence analysis. +bool isInParallelOp(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 @@ -450,6 +454,11 @@ void extractForInductionVars(ArrayRef forInsts, SmallVectorImpl *ivs); +/// Extracts the induction variables from a list of either AffineForOp or +/// AffineParallelOp and places them in the output argument `ivs`. +void extractInductionVars(ArrayRef operations, + SmallVectorImpl *ivs); + /// Builds a perfect nest of affine.for loops, i.e., each loop except the /// innermost one contains only another loop and a terminator. The loops iterate /// from "lbs" to "ubs" with "steps". The body of the innermost loop is 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 @@ -240,19 +240,26 @@ LogicalResult mlir::getIndexSet(MutableArrayRef ops, FlatAffineValueConstraints *domain) { SmallVector indices; - SmallVector forOps; + SmallVector loopOps; + size_t numDims = 0; for (Operation *op : ops) { - if (!isa(op)) { - // TODO: Support affine.parallel ops. - LLVM_DEBUG(llvm::dbgs() << "getIndexSet only handles affine.for/if ops"); + if (!isa(op)) { + LLVM_DEBUG(llvm::dbgs() << "getIndexSet only handles affine.for/if/" + "parallel ops"); return failure(); } - if (AffineForOp forOp = dyn_cast(op)) - forOps.push_back(forOp); + if (AffineForOp forOp = dyn_cast(op)) { + loopOps.push_back(forOp); + // An AffineForOp retains only 1 induction variable. + numDims += 1; + } else if (AffineParallelOp parallelOp = dyn_cast(op)) { + loopOps.push_back(parallelOp); + numDims += parallelOp.getNumDims(); + } } - extractForInductionVars(forOps, &indices); - // Reset while associated Values in 'indices' to the domain. - domain->reset(forOps.size(), /*numSymbols=*/0, /*numLocals=*/0, indices); + extractInductionVars(loopOps, &indices); + // Reset while associating Values in 'indices' to the domain. + domain->reset(numDims, /*numSymbols=*/0, /*numLocals=*/0, indices); for (Operation *op : ops) { // Add constraints from forOp's bounds. if (AffineForOp forOp = dyn_cast(op)) { @@ -260,7 +267,9 @@ return failure(); } else if (AffineIfOp ifOp = dyn_cast(op)) { domain->addAffineIfOpDomain(ifOp); - } + } else if (AffineParallelOp parallelOp = dyn_cast(op)) + if (failed(domain->addAffineParallelOpDomain(parallelOp))) + return failure(); } return success(); } @@ -604,6 +613,10 @@ if (getAffineScope(srcAccess.opInst) != getAffineScope(dstAccess.opInst)) return DependenceResult::Failure; + // TODO: Support affine.parallel which does not specify the ordering. + if (isInParallelOp(srcAccess.opInst) || isInParallelOp(dstAccess.opInst)) + return DependenceResult::Failure; + // Create access relation from each MemRefAccess. FlatAffineRelation srcRel, dstRel; if (failed(srcAccess.getAccessRelation(srcRel))) diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp --- a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp +++ b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp @@ -639,6 +639,34 @@ forOp.getUpperBoundOperands()); } +LogicalResult FlatAffineValueConstraints::addAffineParallelOpDomain( + AffineParallelOp parallelOp) { + size_t ivPos = 0; + for (auto iv : parallelOp.getIVs()) { + LLVM_DEBUG(iv.dump();); + unsigned pos; + if (!findVar(iv, &pos)) { + assert(false && "variable expected for the IV value"); + return failure(); + } + + AffineMap lowerBound = parallelOp.getLowerBoundMap(ivPos); + if (lowerBound.isConstant()) + addBound(BoundType::LB, pos, lowerBound.getSingleConstantResult()); + else if (failed(addBound(BoundType::LB, pos, lowerBound, + parallelOp.getLowerBoundsOperands()))) + return failure(); + + auto upperBound = parallelOp.getUpperBoundMap(ivPos); + if (upperBound.isConstant()) + addBound(BoundType::UB, pos, upperBound.getSingleConstantResult()); + else if (failed(addBound(BoundType::UB, pos, upperBound, + parallelOp.getUpperBoundsOperands()))) + return failure(); + } + return success(); +} + LogicalResult FlatAffineValueConstraints::addDomainFromSliceMaps(ArrayRef lbMaps, ArrayRef ubMaps, 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 @@ -254,6 +254,16 @@ return nullptr; } +bool mlir::isInParallelOp(mlir::Operation *op) { + auto *curOp = op; + while (auto *parentOp = curOp->getParentOp()) { + if (isa(parentOp)) + return true; + curOp = parentOp; + } + return false; +} + // A Value can be used as a dimension id iff it meets one of the following // conditions: // *) It is valid as a symbol. @@ -2311,6 +2321,19 @@ ivs->push_back(forInst.getInductionVar()); } +void mlir::extractInductionVars(ArrayRef operations, + SmallVectorImpl *ivs) { + ivs->reserve(operations.size()); + for (Operation *op : operations) { + // Add constraints from forOp's bounds. + if (AffineForOp forOp = dyn_cast(op)) + ivs->push_back(forOp.getInductionVar()); + else if (AffineParallelOp parallelOp = dyn_cast(op)) + for (size_t i = 0; i < parallelOp.getBody()->getNumArguments(); i++) + ivs->push_back(parallelOp.getBody()->getArgument(i)); + } +} + /// Builds an affine loop nest, using "loopCreatorFn" to create individual loop /// operations. template