[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

3 files changed, 151 insertions, 7 deletions

diff --git a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
index 0924a5c..4c0056f 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 a0c47ad..bcfca70 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 26d581c..fb0644d 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) {