[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.
----------------
FedericoBruzzone wrote:

Could a comment say why the upper bound does not matter?

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