[Mlir-commits] [mlir] [mlir][tensor]-Handle Dynamic Offset in BubbleUpSliceOpThroughCollapse (PR #178921)

[ofri frishman]

================
@@ -579,6 +580,170 @@ LogicalResult mlir::tensor::getCollapsedExtractSliceInfo(
   return success();
 }
 
+// Checks if the `ofr` is a multiple of the `factor`.
+// Handles both static integer and dynamic values
+// where the value is the result of an affine.apply.
+static bool isMultipleOf(OpFoldResult ofr, int64_t factor) {
+  std::optional<int64_t> staticValue = getConstantIntValue(ofr);
+  if (staticValue.has_value())
+    return staticValue.value() % factor == 0;
+
+  Value value = dyn_cast<Value>(ofr);
+  if (!value)
+    return false;
+  auto applyOp = value.getDefiningOp<affine::AffineApplyOp>();
+  if (!applyOp)
+    return false;
+  AffineMap map = applyOp.getAffineMap();
+  SmallVector<Value> operands(applyOp.getOperands());
+  affine::fullyComposeAffineMapAndOperands(&map, &operands);
+  map = simplifyAffineMap(map);
+  if (map.getNumResults() != 1)
+    return false;
+  return map.getResult(0).isMultipleOf(factor);
+}
+
+static LogicalResult computeExpandedSliceInfoForReassocGroup(
----------------
ofri-frishman wrote:

perhaps worth adding a few words to explain parameters and what the function does

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