[Mlir-commits] [mlir] [mlir][vector] Check the indices when inferring in_bounds from a permutation map (PR #221268)

[Federico Bruzzone]

================
@@ -420,19 +420,69 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
                                 useInBoundsInsteadOfMasking);
 }
 
-/// Compute the in_bounds attribute for a transfer op given its permutation map
-/// and the source being accessed. Dimension i is in-bounds when the map result
-/// is an AffineDimExpr pointing to a static source dimension divisible by the
-/// vector size, or an AffineConstantExpr (broadcast).
-static SmallVector<bool> computeInBoundsFromPermutationMap(
-    AffineMap permutationMap, VectorType vectorType, ShapedType sourceType) {
+static bool isKnownMultipleOf(Value value, int64_t factor);
+
+/// Returns true if every result of `map` applied to `operands` is known to be
+/// a multiple of `factor`. Operands that are themselves multiples of `factor`
+/// are substituted with `factor * d` so that AffineExpr::isMultipleOf can see
+/// through them.
+static bool isKnownMultipleOf(AffineMap map, ValueRange operands,
+                              int64_t factor) {
+  MLIRContext *ctx = map.getContext();
+  SmallVector<AffineExpr> dimRepls, symRepls;
+  for (unsigned i = 0, e = map.getNumDims(); i < e; ++i) {
+    AffineExpr dim = getAffineDimExpr(i, ctx);
+    dimRepls.push_back(isKnownMultipleOf(operands[i], factor) ? dim * factor
+                                                              : dim);
+  }
+  for (unsigned i = 0, e = map.getNumSymbols(); i < e; ++i) {
+    AffineExpr sym = getAffineSymbolExpr(i, ctx);
+    symRepls.push_back(isKnownMultipleOf(operands[map.getNumDims() + i], factor)
+                           ? sym * factor
+                           : sym);
+  }
+  return llvm::all_of(map.getResults(), [&](AffineExpr expr) {
+    return expr.replaceDimsAndSymbols(dimRepls, symRepls).isMultipleOf(factor);
+  });
+}
+
+/// Returns true if `value` is known to be a multiple of `factor`: a constant,
+/// an affine.for induction variable whose step and lower bound are multiples
+/// of `factor`, or an affine.apply whose map preserves that property. Loop
+/// bounds and apply operands are followed recursively, so an induction
+/// variable whose lower bound is the induction variable of an outer loop
+/// stepping by a multiple of `factor` is recognized.
+static bool isKnownMultipleOf(Value value, int64_t factor) {
+  if (std::optional<int64_t> cst = getConstantIntValue(value))
+    return *cst % factor == 0;
+  if (affine::AffineForOp forOp = affine::getForInductionVarOwner(value))
+    return forOp.getStepAsInt() % factor == 0 &&
+           isKnownMultipleOf(forOp.getLowerBoundMap(),
+                             forOp.getLowerBoundOperands(), factor);
+  if (auto applyOp = value.getDefiningOp<affine::AffineApplyOp>())
+    return isKnownMultipleOf(applyOp.getAffineMap(), applyOp.getMapOperands(),
+                             factor);
+  return false;
+}
+
+/// Compute the in_bounds attribute for a transfer op given its permutation map,
+/// the source being accessed and the indices into it. Dimension i is in-bounds
+/// when the map result is an AffineDimExpr pointing to a static source
+/// dimension divisible by the vector size whose index is a known multiple of
+/// the vector size, or an AffineConstantExpr (broadcast).
+static SmallVector<bool>
+computeInBoundsFromPermutationMap(AffineMap permutationMap,
+                                  VectorType vectorType, ShapedType sourceType,
+                                  ValueRange indices) {
   SmallVector<bool> inBounds(vectorType.getRank(), false);
   for (unsigned i = 0; i < (unsigned)vectorType.getRank(); ++i) {
     AffineExpr expr = permutationMap.getResult(i);
     if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {
       unsigned memDim = dimExpr.getPosition();
+      int64_t vectorSize = vectorType.getDimSize(i);
----------------
FedericoBruzzone wrote:

`getDimSize(i)` is not the real size for a scalable dim...

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