[Mlir-commits] [mlir] ea53a69 - [linalg][mlir] Replace getSmallestBoundingIndex in padding (NFC).
Tobias Gysi
llvmlistbot at llvm.org
Wed Nov 10 07:16:36 PST 2021
Author: Tobias Gysi
Date: 2021-11-10T15:12:51Z
New Revision: ea53a6938b12894dc883b85cf86da03ecec01b0f
URL: https://github.com/llvm/llvm-project/commit/ea53a6938b12894dc883b85cf86da03ecec01b0f
DIFF: https://github.com/llvm/llvm-project/commit/ea53a6938b12894dc883b85cf86da03ecec01b0f.diff
LOG: [linalg][mlir] Replace getSmallestBoundingIndex in padding (NFC).
Replace the getSmallestBoundingIndex method used in padding by getConstantUpperBoundForIndex that uses flat affine constraints to compute a constant upper bound.
Depends On D113398
Reviewed By: nicolasvasilache
Differential Revision: https://reviews.llvm.org/D113546
Added:
Modified:
mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
mlir/lib/Dialect/Linalg/Utils/Utils.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
index 0924a5c59dcc7..4c0056f2a3e79 100644
--- a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
@@ -60,6 +60,40 @@ SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b);
/// Otherwise return nullptr.
IntegerAttr getSmallestBoundingIndex(Value size);
+/// 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.
+///
+/// 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<Value> &boundOperands);
+
+/// 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<int64_t> getConstantUpperBoundForIndex(Value value);
+
/// Create an ExtractSliceOp and, if `source` is defined by an ExtractSliceOp,
/// fold it by adding the offsets.
///
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index a0c47ad9aeaaf..bcfca706900ed 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -171,16 +171,20 @@ static LogicalResult padOperandToSmallestStaticBoundingBox(
staticSizes.reserve(opToPad.getRank(opOperand));
auto shapedOp = cast<OffsetSizeAndStrideOpInterface>(sliceOp.getOperation());
for (auto size : shapedOp.getMixedSizes()) {
- auto indexAttr = size.is<Attribute>()
- ? size.get<Attribute>().dyn_cast<IntegerAttr>()
- : linalg::getSmallestBoundingIndex(size.get<Value>());
- // SmallestBoundingIndex must exist for all sizes.
- // For now return an error if we can't find it.
- if (!indexAttr) {
+ // If the size is an attribute add it directly to `staticSizes`.
+ if (size.is<Attribute>()) {
+ staticSizes.push_back(
+ size.get<Attribute>().dyn_cast<IntegerAttr>().getInt());
+ continue;
+ }
+ // Otherwise, try to compute a constant upper bound for the size value.
+ FailureOr<int64_t> upperBound =
+ getConstantUpperBoundForIndex(size.get<Value>());
+ if (failed(upperBound)) {
LLVM_DEBUG(DBGS() << "No constant bounding box can be found for padding");
return failure();
}
- staticSizes.push_back(indexAttr.getInt());
+ staticSizes.push_back(upperBound.getValue());
}
auto staticTensorType = RankedTensorType::get(
staticSizes, getElementTypeOrSelf(opOperand->get()));
diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
index 26d581cb1b5c6..fb0644dae81e1 100644
--- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
@@ -12,7 +12,10 @@
#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Analysis/AffineStructures.h"
+#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
#include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
@@ -209,6 +212,109 @@ IntegerAttr getSmallestBoundingIndex(Value size) {
return nullptr;
}
+void getUpperBoundForIndex(Value value, AffineMap &boundMap,
+ SmallVectorImpl<Value> &boundOperands) {
+ // 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<AffineApplyOp, AffineMinOp>(definingOp))
+ return;
+
+ // Get the backward slice containing the affine index computation.
+ SetVector<Operation *> backwardSlice;
+ getBackwardSlice(definingOp, &backwardSlice, [](Operation *op) {
+ return isa<AffineApplyOp, AffineMinOp>(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.containsId(value)) {
+ constraints.appendDimId(value);
+ return true;
+ }
+ unsigned pos;
+ constraints.findId(value, &pos);
+ return pos < constraints.getNumDimIds();
+ };
+ // Helper to get the position for the given value.
+ auto getPosition = [&](Value value) {
+ unsigned pos;
+ bool exists = constraints.findId(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<AffineApplyOp>(op)) {
+ AffineValueMap valueMap(applyOp.getAffineMap(), applyOp.getOperands(),
+ applyOp.getResult());
+ if (failed(constraints.composeMap(&valueMap)))
+ return;
+ continue;
+ }
+ // Add AffineMinOps to the constraints.
+ auto minOp = cast<AffineMinOp>(op);
+ AffineMap map = constraints.computeAlignedMap(minOp.getAffineMap(),
+ minOp.getOperands());
+ if (failed(constraints.addBound(FlatAffineConstraints::UB,
+ getPosition(minOp.getResult()), map)))
+ 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.
+ SmallVector<AffineMap> lowerBounds(1), upperBounds(1);
+ constraints.getSliceBounds(getPosition(value), 1, value.getContext(),
+ &lowerBounds, &upperBounds);
+
+ // 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<int64_t> getConstantUpperBoundForIndex(Value value) {
+ // Compute an upper bound for `value`.
+ AffineMap boundMap;
+ SmallVector<Value> boundOperands;
+ getUpperBoundForIndex(value, boundMap, boundOperands);
+
+ // Search the results of `boundMap` for constant upper bounds.
+ SmallVector<int64_t> constantBounds;
+ for (AffineExpr result : boundMap.getResults())
+ if (auto constExpr = result.dyn_cast<AffineConstantExpr>())
+ 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());
+}
+
tensor::ExtractSliceOp makeComposedExtractSliceOp(
OpBuilder &b, Location loc, Value source, ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides) {
More information about the Mlir-commits
mailing list