[Mlir-commits] [mlir] [mlir][affine] Preserve used sibling loop results during fusion (PR #216564)

[Oleksandr Alex Zinenko]

================
@@ -124,6 +126,168 @@ static bool canRemoveSrcNodeAfterFusion(
   return true;
 }
 
+static void replaceInits(RewriterBase &rewriter, AffineForOp forOp,
+                         ArrayRef<unsigned> initIndices, ValueRange newInits) {
+  assert(initIndices.size() == newInits.size() &&
+         "expected one replacement value per init operand index");
+
+  SmallVector<Value> inits(forOp.getInits());
+  assert(llvm::all_of(initIndices,
+                      [&](unsigned index) { return index < inits.size(); }) &&
+         "expected replacement indices to fit in affine.for init operands");
+
+  for (auto [initIndex, newInit] : llvm::zip(initIndices, newInits))
+    inits[initIndex] = newInit;
+  rewriter.modifyOpInPlace(forOp,
+                           [&]() { forOp.getInitsMutable().assign(inits); });
+}
+
+static AffineForOp
+getAffineForOpWithResults(ValueRange values,
+                          SmallVectorImpl<unsigned> &indices) {
+  assert(!values.empty() && "expected values to forward");
+  auto firstResult = dyn_cast<OpResult>(values.front());
+  if (!firstResult)
+    return nullptr;
+
+  auto forOp = dyn_cast<AffineForOp>(firstResult.getOwner());
+  if (!forOp)
+    return nullptr;
+  assert(forOp.getNumResults() == forOp.getInits().size() &&
+         "expected one init operand per affine.for result");
+
+  indices.clear();
+  for (Value value : values) {
+    auto result = dyn_cast<OpResult>(value);
+    if (!result || result.getOwner() != forOp) {
+      indices.clear();
+      return nullptr;
+    }
+    indices.push_back(result.getResultNumber());
+  }
+  return forOp;
+}
+
+static LogicalResult getBackwardSliceForValues(ValueRange values,
+                                               AffineForOp parentForOp,
+                                               SetVector<Operation *> &slice) {
+  BackwardSliceOptions options;
+  options.inclusive = true;
+  options.omitBlockArguments = true;
+  options.filter = [&](Operation *op) {
+    return parentForOp->isProperAncestor(op);
+  };
+
+  for (Value value : values)
+    if (Operation *definingOp = value.getDefiningOp()) {
----------------
ftynse wrote:

There is an overload of `getBackwardSlice` that takes `Value`.

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