Scf for loop flatten

class FlattenNestedLoopsPattern(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: scf.ForOp, rewriter: PatternRewriter) -> None:
        outer_body = op.body.block
        if not isinstance(inner_loop := outer_body.first_op, scf.ForOp):
            # Outer loop must contain inner loop
            return
        if (
            inner_loop
            is not (outer_yield_op := cast(scf.YieldOp, outer_body.last_op)).prev_op
        ):
            # Outer loop must contain only inner loop and yield
            return

        if op.iter_args:
            if not inner_loop.iter_args:
                return
            if len(op.iter_args) != len(inner_loop.iter_args):
                return
            if not all(
                lhs is rhs
                for (lhs, rhs) in zip(outer_body.args[1:], inner_loop.iter_args)
            ):
                return
            if not all(
                lhs is rhs
                for (lhs, rhs) in zip(inner_loop.results, outer_yield_op.operands)
            ):
                return
        elif inner_loop.iter_args:
            return

        if (inner_lb := const_evaluate_operand(inner_loop.lb)) is None:
            return
        if (inner_ub := const_evaluate_operand(inner_loop.ub)) is None:
            return
        if (outer_step := const_evaluate_operand(op.step)) is None:
            return
        if (inner_step := const_evaluate_operand(inner_loop.step)) is None:
            return
        if outer_step <= 0 or inner_step <= 0:
            return

        inner_span = inner_ub - inner_lb
        if inner_span < 0 or inner_span % inner_step != 0:
            return

        if (outer_lb := const_evaluate_operand(op.lb)) is None:
            return
        if (outer_ub := const_evaluate_operand(op.ub)) is None:
            return
        outer_span = outer_ub - outer_lb
        if outer_span < 0 or outer_span % outer_step != 0:
            return

        outer_index = outer_body.args[0]
        inner_index = inner_loop.body.block.args[0]

        if outer_index.uses or inner_index.uses:
            if inner_lb != 0:
                return

            if inner_ub != outer_step:
                return

            if outer_step % inner_step:
                return

            # If either induction variable is used, we can only fold if used exactly once
            if not outer_index.has_one_use() or not inner_index.has_one_use():
                return

            outer_user = next(iter(outer_index.uses)).operation
            inner_user = next(iter(inner_index.uses)).operation
            if outer_user is not inner_user:
                return

            user = outer_user

            if not isinstance(user, arith.AddiOp):
                return

            # We can fuse
            user.result.replace_all_uses_with(inner_index)
            rewriter.erase(user)
            new_ub = op.ub
            new_step = inner_loop.step
        else:
            if outer_lb != 0:
                # Do not currently handle lb != 0
                return

            factor = inner_span // inner_step
            factor_op = arith.ConstantOp(
                builtin.IntegerAttr(factor, builtin.IndexType())
            )
            new_ub_op = arith.MuliOp(op.ub, factor_op.result)
            rewriter.insert((factor_op, new_ub_op))
            new_ub = new_ub_op.result
            new_step = op.step

        moved_region = rewriter.move_region_contents_to_new_regions(inner_loop.body)
        rewriter.erase(outer_yield_op)
        rewriter.erase(inner_loop)

        rewriter.replace(
            op,
            scf.ForOp(
                op.lb,
                new_ub,
                new_step,
                op.iter_args,
                moved_region,
            ),
        )

class ScfForLoopFlattenPass(ModulePass):
    """
    Folds perfect loop nests when their statically known, positive-step ranges
    tile exactly. If induction variables are used, they must have the existing
    matching add operation and the inner range must match the outer step.
    Carried arguments are supported when both loops have matching arguments and
    the inner and outer yields forward them in order.
    """

    name = "scf-for-loop-flatten"

    def apply(self, ctx: Context, op: builtin.ModuleOp) -> None:
        PatternRewriteWalker(FlattenNestedLoopsPattern()).rewrite_module(op)