diff --git a/mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h --- a/mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h +++ b/mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h @@ -51,15 +51,21 @@ /// Reify a bound for the given index-typed value or shape dimension size in /// terms of the owning op's operands. `dim` must be `nullopt` if and only if -/// `value` is index-typed. LB and EQ bounds are closed, UB bounds are open. +/// `value` is index-typed. +/// +/// By default, lower/equal bounds are closed and upper bounds are open. If +/// `closedUB` is set to "true", upper bounds are also closed. FailureOr reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type, Value value, - std::optional dim); + std::optional dim, + bool closedUB = false); /// Reify a bound for the given index-typed value or shape dimension size in /// terms of SSA values for which `stopCondition` is met. `dim` must be -/// `nullopt` if and only if `value` is index-typed. LB and EQ bounds are -/// closed, UB bounds are open. +/// `nullopt` if and only if `value` is index-typed. +/// +/// By default, lower/equal bounds are closed and upper bounds are open. If +/// `closedUB` is set to "true", upper bounds are also closed. /// /// Example: /// %0 = arith.addi %a, %b : index @@ -74,7 +80,8 @@ FailureOr reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type, Value value, std::optional dim, - ValueBoundsConstraintSet::StopConditionFn stopCondition); + ValueBoundsConstraintSet::StopConditionFn stopCondition, + bool closedUB = false); } // namespace mlir diff --git a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h --- a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h +++ b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h @@ -94,44 +94,6 @@ /// constructing the necessary DimOp operators. SmallVector getDynOperands(Location loc, Value val, OpBuilder &b); -/// Computes an upper bound for the result `value` of an index computation. -/// Translates AffineMinOps and AffineApplyOps along the use-def chains of the -/// index computation to affine constraints and projects out intermediate -/// values. The method sets `boundMap` to an affine map that given -/// `boundOperands` evaluates to an upper bound for the index computation. -/// -/// If constantRequired is true, only returns the constant bounds (potentially -/// over-approximating) and fails when not possible. -/// -/// Example: -/// ``` -/// %dim0 = dim %tensor, %c0 -/// %dim1 = dim %tensor, %c1 -/// %0 = affine.min affine.map<(d0) -> (40, d0)> (%dim0) -/// %1 = affine.apply affine.map<(d0, d1) -> (d0 + d1)> (%0, %dim1) -/// ``` -/// getUpperBoundForIndex(%1, boundMap, boundOperands) -/// set the output parameters to: -/// - boundMap = affine.map<(d0) -> (d0 + 40)> -/// - boundOperands = [%dim1] -void getUpperBoundForIndex(Value value, AffineMap &boundMap, - SmallVectorImpl &boundOperands, - bool constantRequired = false); - -/// Returns a constant upper bound for the result `value` of an index -/// computation. Calls `getUpperBoundForIndex` and returns a constant upper -/// bound if the result of `boundMap` is a constant expression and failure -/// otherwise. -/// -/// Example: -/// ``` -/// %0 = affine.min affine.map<(d0) -> (40, d0)> (%d0) -/// %1 = affine.apply affine.map<(d0) -> (d0 + 2)> (%0) -/// ``` -/// getConstantUpperBoundForIndex(%1) returns 42 -/// (boundsMap = affine.map<() -> (42)>) -FailureOr getConstantUpperBoundForIndex(Value value); - /// Create a tensor::PadOp that pads `source` to the size of the statically /// sized `type` whose static sizes are assumed to be greater than the dynamic /// `source` size. The padding introduces trailing `pad` values until the diff --git a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h --- a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h +++ b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h @@ -101,11 +101,15 @@ /// chain) of the given value is visited in a worklist-driven manner and the /// constraint set is populated according to `ValueBoundsOpInterface` for each /// visited value. + /// + /// By default, lower/equal bounds are closed and upper bounds are open. If + /// `closedUB` is set to "true", upper bounds are also closed. static LogicalResult computeBound(AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type, Value value, std::optional dim, - StopConditionFn stopCondition); + StopConditionFn stopCondition, + bool closedUB = false); /// Compute a bound in terms of the values/dimensions in `dependencies`. The /// computed bound consists of only constant terms and dependent values (or @@ -114,7 +118,8 @@ ValueDimList &mapOperands, presburger::BoundType type, Value value, std::optional dim, - ValueDimList dependencies); + ValueDimList dependencies, + bool closedUB = false); /// Compute a constant bound for the given index-typed value or shape /// dimension size. @@ -129,10 +134,14 @@ /// The stop condition is optional: If none is specified, the backward slice /// is traversed in a breadth-first manner until a constant bound could be /// computed. + /// + /// By default, lower/equal bounds are closed and upper bounds are open. If + /// `closedUB` is set to "true", upper bounds are also closed. static FailureOr computeConstantBound(presburger::BoundType type, Value value, std::optional dim = std::nullopt, - StopConditionFn stopCondition = nullptr); + StopConditionFn stopCondition = nullptr, + bool closedUB = false); /// Add a bound for the given index-typed value or shaped value. This function /// returns a builder that adds the bound. diff --git a/mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp b/mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp --- a/mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp +++ b/mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp @@ -15,10 +15,9 @@ using namespace mlir; -FailureOr mlir::reifyValueBound(OpBuilder &b, Location loc, - presburger::BoundType type, - Value value, - std::optional dim) { +FailureOr +mlir::reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type, + Value value, std::optional dim, bool closedUB) { // We are trying to reify a bound for `value`. Construct a stop condition that // evaluates to "true" for any SSA value expect for `value`. I.e., the bound // will be computed in terms of any SSA values except for `value`. The first @@ -27,18 +26,19 @@ // Reify in terms of SSA values that are different from `value`. return v != value; }; - return reifyValueBound(b, loc, type, value, dim, stopCondition); + return reifyValueBound(b, loc, type, value, dim, stopCondition, closedUB); } FailureOr mlir::reifyValueBound( OpBuilder &b, Location loc, presburger::BoundType type, Value value, std::optional dim, - function_ref)> stopCondition) { + function_ref)> stopCondition, + bool closedUB) { // Compute bound. AffineMap boundMap; ValueDimList mapOperands; - if (failed(ValueBoundsConstraintSet::computeBound(boundMap, mapOperands, type, - value, dim, stopCondition))) + if (failed(ValueBoundsConstraintSet::computeBound( + boundMap, mapOperands, type, value, dim, stopCondition, closedUB))) return failure(); // Materialize tensor.dim/memref.dim ops. diff --git a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt --- a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt +++ b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt @@ -40,6 +40,7 @@ LINK_LIBS PUBLIC MLIRAffineDialect + MLIRAffineTransforms MLIRAffineUtils MLIRAnalysis MLIRArithDialect @@ -69,6 +70,7 @@ MLIRTensorTransforms MLIRTransforms MLIRTransformUtils + MLIRValueBoundsOpInterface MLIRVectorDialect MLIRVectorTransforms MLIRVectorUtils diff --git a/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp b/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp @@ -10,8 +10,10 @@ // //===----------------------------------------------------------------------===// +#include "mlir/Analysis/Presburger/IntegerRelation.h" #include "mlir/Analysis/SliceAnalysis.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Affine/Transforms/Transforms.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/Linalg/Transforms/Hoisting.h" @@ -459,11 +461,22 @@ // of the enclosing loops. for (auto forOp : packingLoops) { // Compute an upper bound `ubVal` for the upper bound of `forOp`. - AffineMap boundMap; - SmallVector boundOperands; - getUpperBoundForIndex(forOp.getUpperBound(), boundMap, boundOperands); - Value ubVal = - rewriter.createOrFold(loc, boundMap, boundOperands); + FailureOr loopUb = reifyValueBound( + rewriter, loc, presburger::BoundType::UB, forOp.getUpperBound(), + /*dim=*/std::nullopt, /*stopCondition=*/ + [&](Value v, std::optional d) { + if (v == forOp.getUpperBound()) + return false; + // Compute a bound that is independent of any affine op results. + Operation *op = v.getDefiningOp(); + if (!op) + return true; + return !isa(op); + }, + /*closedUB=*/true); + assert(succeeded(loopUb) && "could not get upper bound"); + Value ubVal = getValueOrCreateConstantIndexOp(rewriter, loc, *loopUb); + // Compute the maximal packing loop length as (ub - lb).ceilDiv(step) and // store the result to `dynamicTensorSizes`. // TODO: instead of using the lower bound of `forOp` directly, implement a diff --git a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp @@ -23,6 +23,7 @@ #include "mlir/IR/AffineExprVisitor.h" #include "mlir/IR/AffineMap.h" #include "mlir/IR/ImplicitLocOpBuilder.h" +#include "mlir/Interfaces/ValueBoundsOpInterface.h" #include "mlir/Support/LLVM.h" #include "mlir/Transforms/FoldUtils.h" #include "llvm/ADT/MapVector.h" @@ -234,7 +235,9 @@ Value materializedSize = getValueOrCreateConstantIndexOp(b, loc, rangeValue.size); FailureOr upperBound = - getConstantUpperBoundForIndex(materializedSize); + ValueBoundsConstraintSet::computeConstantBound( + presburger::BoundType::UB, materializedSize, /*dim=*/std::nullopt, + /*stopCondition=*/nullptr, /*closedUB=*/true); size = failed(upperBound) ? materializedSize : b.create(loc, *upperBound); diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp @@ -27,6 +27,7 @@ #include "mlir/Dialect/Vector/IR/VectorOps.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/Matchers.h" +#include "mlir/Interfaces/ValueBoundsOpInterface.h" #include "mlir/Pass/Pass.h" #include "mlir/Support/LLVM.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" @@ -139,7 +140,9 @@ } // Otherwise, try to compute a constant upper bound for the size value. FailureOr upperBound = - getConstantUpperBoundForIndex(en.value().get()); + ValueBoundsConstraintSet::computeConstantBound( + presburger::BoundType::UB, en.value().get(), + /*dim=*/std::nullopt, /*stopCondition=*/nullptr, /*closedUB=*/true); if (failed(upperBound)) { LLVM_DEBUG(DBGS() << "count not compute a bonding box for padding"); return rewriter.notifyMatchFailure( diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp --- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp @@ -297,124 +297,6 @@ return dynOperands; } -void getUpperBoundForIndex(Value value, AffineMap &boundMap, - SmallVectorImpl &boundOperands, - bool constantRequired) { - // Initialize `boundMap` and `boundOperands` to the identity returning - // `value`. This combination is the default result of the method if no - // simplification is possible. - assert(value.getType().isIndex() && "expect value to have index type"); - boundMap = AffineMap::getMultiDimIdentityMap(1, value.getContext()); - boundOperands.assign({value}); - canonicalizeMapAndOperands(&boundMap, &boundOperands); - - // Continue only if there is an affine index computation to simplify. - Operation *definingOp = value.getDefiningOp(); - if (!definingOp || !isa(definingOp)) - return; - - // Get the backward slice containing the affine index computation. - SetVector backwardSlice; - getBackwardSlice(definingOp, &backwardSlice, [](Operation *op) { - return isa(op); - }); - backwardSlice.insert(definingOp); - - // Setup a system of affine constraints that describe the index computation. - FlatAffineValueConstraints constraints; - - // Helper to find or create an identifier for the given value. - auto findOrCreateId = [&](Value value) { - if (!constraints.containsVar(value)) { - constraints.appendDimVar(value); - return true; - } - unsigned pos; - constraints.findVar(value, &pos); - return pos < constraints.getNumDimVars(); - }; - // Helper to get the position for the given value. - auto getPosition = [&](Value value) { - unsigned pos; - bool exists = constraints.findVar(value, &pos); - (void)exists; - assert(exists && "expect to find the identifier"); - return pos; - }; - - // Add the affine operations in `backwardSlice` to the constraints. - for (Operation *op : llvm::reverse(backwardSlice)) { - // Add an identifier for all op results and operands. - if (!(llvm::all_of(op->getResults(), findOrCreateId) && - llvm::all_of(op->getOperands(), findOrCreateId))) - return; - - // Add AffineApplyOps to the constraints. - if (auto applyOp = dyn_cast(op)) { - AffineMap map = constraints.computeAlignedMap(applyOp.getAffineMap(), - applyOp.getOperands()); - if (failed(constraints.addBound(BoundType::EQ, - getPosition(applyOp.getResult()), map))) - return; - continue; - } - // Add AffineMinOps to the constraints. - auto minOp = cast(op); - AffineMap map = constraints.computeAlignedMap(minOp.getAffineMap(), - minOp.getOperands()); - if (failed(constraints.addBound(BoundType::UB, - getPosition(minOp.getResult()), map, - /*isClosedBound=*/true))) - return; - } - - // Obtain an upper bound for the affine index computation by projecting out - // all temporary results and expressing the upper bound for `value` in terms - // of the terminals of the index computation. - unsigned pos = getPosition(value); - if (constantRequired) { - auto ubConst = constraints.getConstantBound64(BoundType::UB, pos); - if (!ubConst) - return; - - boundMap = AffineMap::getConstantMap(*ubConst, value.getContext()); - return; - } - - SmallVector lowerBounds(1), upperBounds(1); - constraints.getSliceBounds(pos, 1, value.getContext(), &lowerBounds, - &upperBounds, - /*getClosedUB=*/true); - // Verify `upperBounds[0]` is valid and has at least one result. - if (!upperBounds[0] || upperBounds[0].getNumResults() == 0) - return; - - // Set `boundMap` and `boundOperands` to the computed upper bound. - boundMap = upperBounds[0]; - constraints.getAllValues(&boundOperands); - erase_value(boundOperands, value); - canonicalizeMapAndOperands(&boundMap, &boundOperands); -} - -FailureOr getConstantUpperBoundForIndex(Value value) { - // Compute an upper bound for `value`. - AffineMap boundMap; - SmallVector boundOperands; - getUpperBoundForIndex(value, boundMap, boundOperands, - /*constantRequired=*/true); - - // Search the results of `boundMap` for constant upper bounds. - SmallVector constantBounds; - for (AffineExpr result : boundMap.getResults()) - if (auto constExpr = result.dyn_cast()) - constantBounds.push_back(constExpr.getValue()); - - // Return the minimal upper bound or failure if none is found. - if (constantBounds.empty()) - return failure(); - return *std::min_element(constantBounds.begin(), constantBounds.end()); -} - Value makeComposedPadHighOp(OpBuilder &b, Location loc, RankedTensorType type, Value source, Value pad, bool nofold) { // Exit if `source` is not defined by an ExtractSliceOp. diff --git a/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt b/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt --- a/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt +++ b/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt @@ -71,4 +71,5 @@ MLIRTensorDialect MLIRTensorUtils MLIRTilingInterface + MLIRValueBoundsOpInterface ) diff --git a/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp --- a/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp +++ b/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp @@ -17,6 +17,7 @@ #include "mlir/Dialect/Tensor/Utils/Utils.h" #include "mlir/Dialect/Utils/IndexingUtils.h" #include "mlir/Interfaces/TilingInterface.h" +#include "mlir/Interfaces/ValueBoundsOpInterface.h" using namespace mlir; using namespace mlir::tensor; @@ -263,8 +264,10 @@ OpFoldResult innerTileSize = dimAndTileMapping[tileDim]; info.isAlignedToInnerTileSize = false; - FailureOr cstSize = linalg::getConstantUpperBoundForIndex( - getValueOrCreateConstantIndexOp(b, loc, tileSize)); + FailureOr cstSize = ValueBoundsConstraintSet::computeConstantBound( + presburger::BoundType::UB, + getValueOrCreateConstantIndexOp(b, loc, tileSize), /*dim=*/std::nullopt, + /*stopCondition=*/nullptr, /*closedUB=*/true); std::optional cstInnerSize = getConstantIntValue(innerTileSize); if (!failed(cstSize) && cstInnerSize) { if (*cstSize % *cstInnerSize == 0) diff --git a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp --- a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp +++ b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp @@ -219,13 +219,15 @@ LogicalResult ValueBoundsConstraintSet::computeBound( AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type, - Value value, std::optional dim, StopConditionFn stopCondition) { + Value value, std::optional dim, StopConditionFn stopCondition, + bool closedUB) { #ifndef NDEBUG assertValidValueDim(value, dim); assert(!stopCondition(value, dim) && "stop condition should not be satisfied for starting point"); #endif // NDEBUG + int64_t ubAdjustment = closedUB ? 0 : 1; Builder b(value.getContext()); mapOperands.clear(); @@ -233,6 +235,9 @@ // Special case: If the stop condition is satisfied for the input // value/dimension, directly return it. mapOperands.push_back(std::make_pair(value, dim)); + AffineExpr bound = b.getAffineDimExpr(0); + if (type == BoundType::UB) + bound = bound + ubAdjustment; resultMap = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, b.getAffineDimExpr(0)); return success(); @@ -288,9 +293,9 @@ if (type == BoundType::EQ || type == BoundType::LB) { bound = lb[0]; } else { - // Computed UB is a closed bound. Turn into an open bound. + // Computed UB is a closed bound. bound = AffineMap::get(ub[0].getNumDims(), ub[0].getNumSymbols(), - ub[0].getResult(0) + 1); + ub[0].getResult(0) + ubAdjustment); } // Gather all SSA values that are used in the computed bound. @@ -353,17 +358,19 @@ LogicalResult ValueBoundsConstraintSet::computeBound( AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type, - Value value, std::optional dim, ValueDimList dependencies) { - return computeBound(resultMap, mapOperands, type, value, dim, - [&](Value v, std::optional d) { - return llvm::is_contained(dependencies, - std::make_pair(v, d)); - }); + Value value, std::optional dim, ValueDimList dependencies, + bool closedUB) { + return computeBound( + resultMap, mapOperands, type, value, dim, + [&](Value v, std::optional d) { + return llvm::is_contained(dependencies, std::make_pair(v, d)); + }, + closedUB); } FailureOr ValueBoundsConstraintSet::computeConstantBound( presburger::BoundType type, Value value, std::optional dim, - StopConditionFn stopCondition) { + StopConditionFn stopCondition, bool closedUB) { #ifndef NDEBUG assertValidValueDim(value, dim); #endif // NDEBUG @@ -384,8 +391,9 @@ } // Compute constant bound for `valueDim`. + int64_t ubAdjustment = closedUB ? 0 : 1; if (auto bound = cstr.cstr.getConstantBound64(type, pos)) - return type == BoundType::UB ? *bound + 1 : *bound; + return type == BoundType::UB ? *bound + ubAdjustment : *bound; return failure(); } diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel --- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel @@ -5721,6 +5721,7 @@ ":TensorDialect", ":TensorUtils", ":TilingInterface", + ":ValueBoundsOpInterface", "//llvm:Support", ], ) @@ -8795,6 +8796,7 @@ deps = [ ":AffineAnalysis", ":AffineDialect", + ":AffineTransforms", ":AffineUtils", ":Analysis", ":ArithDialect", @@ -8832,6 +8834,7 @@ ":TilingInterface", ":TransformUtils", ":Transforms", + ":ValueBoundsOpInterface", ":VectorDialect", ":VectorToSCF", ":VectorTransforms",