[Mlir-commits] [mlir] [mlir][vector] Fold `in_bounds` for transfers with loop-derived indices (PR #215340)

[Matthias Springer]

================
@@ -5373,23 +5373,39 @@ VectorType TransferReadOp::getVectorType() {
 
 template <typename TransferOp>
 static bool isInBounds(TransferOp op, int64_t resultIdx, int64_t indicesIdx) {
-  // TODO: support more aggressive createOrFold on:
-  // op.getIndices()[indicesIdx] + vectorType < dim(op.getSource(), indicesIdx)
   if (op.getShapedType().isDynamicDim(indicesIdx))
     return false;
   // Scalable dimensions are `vscale` times larger at runtime, so the static
   // size is only a lower bound and cannot prove that the transfer fits.
   if (op.getVectorType().getScalableDims()[resultIdx])
     return false;
   Value index = op.getIndices()[indicesIdx];
-  std::optional<int64_t> cstOp = getConstantIntValue(index);
-  if (!cstOp.has_value())
-    return false;
 
   int64_t sourceSize = op.getShapedType().getDimSize(indicesIdx);
   int64_t vectorSize = op.getVectorType().getDimSize(resultIdx);
+  // Largest index at which a full vector still fits. Computed as a subtraction
+  // so that adding to a large index cannot overflow.
+  int64_t maxStart = sourceSize - vectorSize;
+
+  // `in_bounds` guarantees that the transfer, including its starting point,
+  // stays within the source, so a negative index is out of bounds.
+  if (std::optional<int64_t> cstOp = getConstantIntValue(index))
+    return *cstOp >= 0 && *cstOp <= maxStart;
+
+  // The index is typically a loop induction variable or an affine expression
+  // thereof. The transfer is in bounds if even the largest index the enclosing
+  // loops can produce leaves room for a full vector.
+  FailureOr<int64_t> maxIndex = ValueBoundsConstraintSet::computeConstantBound(
+      presburger::BoundType::UB, index, /*stopCondition=*/nullptr,
+      ValueBoundsOptions{/*closedUB=*/true});
+  if (failed(maxIndex) || *maxIndex > maxStart)
+    return false;
 
-  return cstOp.value() + vectorSize <= sourceSize;
+  // The starting point must be in bounds as well. Queried only once the upper
+  // bound holds, so that indices that fail it pay for one query, not two.
+  FailureOr<int64_t> minIndex = ValueBoundsConstraintSet::computeConstantBound(
----------------
matthias-springer wrote:

This is building a constraint set, which can be quite expensive. `isInBounds` is called from the folder and therefore by the canonicalizer pass. This could severely blow up compilation time. The optimization itself makes sense, but I would create a separate pass for it.


https://github.com/llvm/llvm-project/pull/215340