[Mlir-commits] [mlir] [SCF] Added canonicalizer for recursively dead uses of iter_args. (PR #191085)

[Mehdi Amini]

================
@@ -1002,11 +1002,153 @@ struct ForOpTensorCastFolder : public OpRewritePattern<ForOp> {
     return failure();
   }
 };
+
+/// Remove iter_arg/result pairs from scf.for when the result is unused and
+/// the corresponding iter_arg block argument is "effectively unused" --
+/// meaning it has no uses, or its only uses are as init operands for nested
+/// scf.for iter_args whose block arguments are also effectively unused.
+///
+/// This handles cases that the generic RemoveDeadRegionBranchOpSuccessorInputs
+/// pattern cannot, specifically when inner loop results are used by outer
+/// loop yields creating cross-loop use chains that appear live but are
+/// semantically dead.
+///
+/// Example:
+///   %r = scf.for %i = %lb to %ub step %s iter_args(%a = %init) -> (f32) {
+///     %inner = scf.for %j = %lb to %ub step %s
+///         iter_args(%b = %a) -> (f32) {
+///       // %b is unused in the body
+///       scf.yield %val : f32
+///     }
+///     scf.yield %inner : f32
+///   }
+///   // %r is unused
+///
+/// After canonicalization:
+///   scf.for %i = %lb to %ub step %s {
+///     scf.for %j = %lb to %ub step %s {
+///       // body without iter_args
+///     }
+///   }
+struct ForOpUnusedIterArgElimination : public OpRewritePattern<ForOp> {
+  using OpRewritePattern<ForOp>::OpRewritePattern;
+
+  /// Check if a block argument is effectively unused. A block argument is
+  /// effectively unused if it has no uses, or all its uses are init operands
+  /// for nested scf.for iter_args where: (a) the inner block arg is also
+  /// effectively unused, and (b) the inner for's result at that position is
+  /// only used as yield operands at positions we are removing from parentFor.
+  static bool isBlockArgEffectivelyUnused(Value blockArg, ForOp parentFor,
+                                          const BitVector &parentCandidates) {
+    if (blockArg.use_empty())
+      return true;
+
+    for (OpOperand &use : blockArg.getUses()) {
+      auto innerFor = dyn_cast<ForOp>(use.getOwner());
----------------
joker-eph wrote:

That seems completely ad-hoc to me: we should rely on inside-out processing to simplify first the inner ops and then the outer ones.  Hardcoding ForOp makes this all not composable.

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