[Mlir-commits] [mlir] [mlir][scf]: Avoid inserting affine.min when tiling dynamic operation sizes if possible (PR #113819)

[llvmlistbot at llvm.org]

================
@@ -186,18 +188,49 @@ static void checkSafeToTileToForall(TilingInterface op,
   }
 }
 
+/// Collect divider of the `ofr`.
+static void collectDividers(OpFoldResult ofr,
+                            SmallVector<OpFoldResult> &dividers) {
+  dividers.push_back(ofr);
+  if (ofr.is<Attribute>())
+    return;
+  auto mulOp = cast<Value>(ofr).getDefiningOp<arith::MulIOp>();
+  if (!mulOp)
+    return;
+
+  // Given `ofr` = `x` * `y`, all dividers of `x` and `y` are dividers of `ofr`.
+  collectDividers(mulOp.getLhs(), dividers);
+  collectDividers(mulOp.getRhs(), dividers);
+}
+
 /// Check if `stride` evenly divides the trip count `size - offset`.
 static bool tileDividesIterationDomain(Range loopRange) {
+  std::optional<int64_t> strideAsInt = getConstantIntValue(loopRange.stride);
   std::optional<int64_t> offsetAsInt = getConstantIntValue(loopRange.offset);
-  if (!offsetAsInt)
-    return false;
   std::optional<int64_t> sizeAsInt = getConstantIntValue(loopRange.size);
-  if (!sizeAsInt)
-    return false;
-  std::optional<int64_t> strideAsInt = getConstantIntValue(loopRange.stride);
-  if (!strideAsInt)
-    return false;
-  return ((sizeAsInt.value() - offsetAsInt.value()) % strideAsInt.value() == 0);
+  if (strideAsInt && offsetAsInt && sizeAsInt)
+    // `stride`/`size`/`offset` are static, checking (size - offset) % stride =
+    // 0.
+    return ((sizeAsInt.value() - offsetAsInt.value()) % strideAsInt.value() ==
+            0);
+
+  // At least `stride`/`size`/`offset` is dynamic.
+  SmallVector<OpFoldResult> dividersOfSize, dividersOfOffset;
+  collectDividers(loopRange.size, dividersOfSize);
+  collectDividers(loopRange.offset, dividersOfOffset);
+
+  // Return true if `stride` divides one of the dividers of both `size` and
+  // `offset`.
+  auto isStrideDividesDivider = [&](OpFoldResult divider) {
+    if (!strideAsInt)
+      // `stride` is dynamic.
+      return divider == loopRange.stride;
+
+    std::optional<int64_t> dividerAsInt = getConstantIntValue(divider);
+    return dividerAsInt && *dividerAsInt % *strideAsInt == 0;
+  };
+  return llvm::any_of(dividersOfSize, isStrideDividesDivider) &&
----------------
MaheshRavishankar wrote:

Dont you think this logic is more complicated that you want it to be. I'd like to scope this a bit more narrowly. Just look for the immediate operations being a `arith.mul` and avoid the recursion.

The more general solution should be using something the `IntegerDivisibilityAnalysis` ([here](https://github.com/iree-org/iree/blob/main/compiler/src/iree/compiler/Dialect/Util/Analysis/IntegerDivisibilityAnalysis.h)) that is very similar to the range inference analysis that can be used to fold this away. Doing this in transformations like this becomes unmaintainable in the long run.

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