[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()) {
+      if (failed(getBackwardSlice(definingOp, &slice, options)))
+        return failure();
+    }
+  return success();
+}
+
+static void replaceInitUsesInPromotedSlice(RewriterBase &rewriter,
+                                           const SetVector<Operation *> &slice,
+                                           ValueRange oldValues,
+                                           ValueRange newValues) {
+  assert(oldValues.size() == newValues.size() &&
+         "expected one replacement value per old value");
+
+  for (auto [oldValue, newValue] : llvm::zip(oldValues, newValues)) {
+    rewriter.replaceUsesWithIf(oldValue, newValue, [&](OpOperand &use) {
+      return slice.contains(use.getOwner());
+    });
+  }
+}
+
+static Operation *getAncestorInBlock(Operation *op, Block *block) {
+  while (op && op->getBlock() != block)
+    op = op->getParentOp();
+  return op;
+}
+
+// Check that all result users remain dominated after sibling fusion moves
+// `definingOp` before `insertPoint`.
+static bool resultUsersRemainDominated(Operation *opWithUsers,
+                                       Operation *definingOp,
+                                       Operation *insertPoint) {
+  Block *block = insertPoint->getBlock();
+  bool definingOpWillMove = definingOp != insertPoint;
+  for (Value result : opWithUsers->getResults()) {
+    if (result.use_empty())
+      continue;
+    for (OpOperand &use : result.getUses()) {
+      Operation *userAncestor = getAncestorInBlock(use.getOwner(), block);
+      if (!userAncestor)
+        return false;
+      if (definingOpWillMove && userAncestor->isBeforeInBlock(insertPoint))
+        return false;
+      if (!definingOpWillMove && !definingOp->isBeforeInBlock(userAncestor))
+        return false;
+    }
+  }
+  return true;
+}
+
+// Carry values from the cloned sibling loop through the destination loop nest.
+// Each newly added iter_arg is used as the corresponding init of the loop that
+// currently defines the forwarded values.
+static FailureOr<SmallVector<Value>>
+forwardValuesThroughAffineForNest(ValueRange values, ValueRange inits,
+                                  AffineForOp &outerForOp) {
+  IRRewriter rewriter(outerForOp->getContext());
+  SmallVector<Value> forwardedValues(values.begin(), values.end());
----------------
ftynse wrote:

Nit: `llvm::to_vector`.

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