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

[Dhairyashil R G]

================
@@ -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(
----------------
dhairyashilRG wrote:

Agreed, a separate pass is the right call. I'd offered gating in the description, but that was wrong: unlike `isDisjointTransferIndices`, which is reached from transforms that pass `testDynamicValueUsingBounds`, `fold()` has no options to gate on. And a failed fold leaves the op unchanged, so every canonicalizer run in a pipeline rebuilds the constraint set again, and the failing cases (dynamic bounds, or a bound that doesn't leave room) are the ones that recur.

Plan: `-vector-infer-in-bounds` in `Vector`/`Transforms`, modelled on `eliminateVectorMasks`, a walk over transfer ops rather than patterns re-run to fixpoint. The folder keeps only its constant-index path (I'll split the negative-index `>= 0` check into its own small patch, since that part is cheap and belongs in the folder). Tests move to their own file, and the `hoisting.mlir` CHECK churn goes away.

One alternative I'd rather have your read on than guess: doing this at creation, in `createReadOrMaskedRead`/`createWriteOrMaskedWrite`, the FIXMEs there say the computation "ignores the read/write indices", which is exactly this gap. That's once per created op instead of repeatedly under canonicalize, and it has in-tree consumers in both the Linalg vectorizer and `-affine-super-vectorize`. Better home than an opt-in pass, or worse? cc @banach-space for the vector-side naming/placement.

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