[flang-commits] [flang] [flang][Transforms] Add LiftSCFWhileToSCFFor pass (PR #213273)
Kareem Ergawy via flang-commits
flang-commits at lists.llvm.org
Fri Jul 31 06:16:09 PDT 2026
https://github.com/ergawy created https://github.com/llvm/llvm-project/pull/213273
Introduces `--lift-scf-while-to-scf-for`, which walks every scf.while produced by upstream `lift-cf-to-scf` and, when the canonical "trip-counter + induction-variable" shape is recognized, rewrites the loop as scf.for. Other loop-carried before-args are passed through as iter_args. Loops that do not match (data-dependent updates, used results, non-positive IV step, etc.) are left untouched and a diagnostic is emitted to stderr.
Recognition collects a LoopInfo describing the loop:
- trip counter: candidate the gating cmp uses, must have step -1
- induction variable: the single non-trip affine recurrence, must have positive step (negative step bails for now)
- iter_args: any other before-args that pass through unchanged
Rewrite materializes `ub_excl = ivInit + tripInit * ivStep` before the scf.while, casts lb/ub/step to index (so the resulting scf.for is compatible with downstream affine conversion), and clones the body from the predicate scf.if's continues branch. The IV recurrence result is remapped to `forIV + step` materialized at the top of the body so cloned ops referencing the next-iteration value keep working.
Companion to lift-cf-to-scf for rebuilding counted loops from recovered do-while loops.
>From 8f3ad1f08e1d2ad8248f19a238660ff8d5fbe3d2 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at gmail.com>
Date: Fri, 31 Jul 2026 02:37:29 -0700
Subject: [PATCH] [flang][Transforms] Add LiftSCFWhileToSCFFor pass
Introduces `--lift-scf-while-to-scf-for`, which walks every scf.while
produced by upstream `lift-cf-to-scf` and, when the canonical
"trip-counter + induction-variable" shape is recognized, rewrites the
loop as scf.for. Other loop-carried before-args are passed through as
iter_args. Loops that do not match (data-dependent updates, used
results, non-positive IV step, etc.) are left untouched and a
diagnostic is emitted to stderr.
Recognition collects a LoopInfo describing the loop:
- trip counter: candidate the gating cmp uses, must have step -1
- induction variable: the single non-trip affine recurrence, must
have positive step (negative step bails for now)
- iter_args: any other before-args that pass through unchanged
Rewrite materializes `ub_excl = ivInit + tripInit * ivStep` before
the scf.while, casts lb/ub/step to index (so the resulting scf.for
is compatible with downstream affine conversion), and clones the
body from the predicate scf.if's continues branch. The IV recurrence
result is remapped to `forIV + step` materialized at the top of the
body so cloned ops referencing the next-iteration value keep working.
Companion to lift-cf-to-scf for rebuilding counted loops from
recovered do-while loops.
Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply at anthropic.com>
---
.../flang/Optimizer/Transforms/Passes.td | 20 +
flang/lib/Optimizer/Transforms/CMakeLists.txt | 1 +
.../Transforms/LiftSCFWhileToSCFFor.cpp | 560 ++++++++++++++++++
.../iv-next-body-use.mlir | 54 ++
.../Fir/LiftSCFWhileToSCFFor/step-gt-one.mlir | 51 ++
.../Fir/LiftSCFWhileToSCFFor/step-one.mlir | 48 ++
.../LiftSCFWhileToSCFFor/two-iter-args.mlir | 72 +++
7 files changed, 806 insertions(+)
create mode 100644 flang/lib/Optimizer/Transforms/LiftSCFWhileToSCFFor.cpp
create mode 100644 flang/test/Fir/LiftSCFWhileToSCFFor/iv-next-body-use.mlir
create mode 100644 flang/test/Fir/LiftSCFWhileToSCFFor/step-gt-one.mlir
create mode 100644 flang/test/Fir/LiftSCFWhileToSCFFor/step-one.mlir
create mode 100644 flang/test/Fir/LiftSCFWhileToSCFFor/two-iter-args.mlir
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index 9a1357881eae0..93e8e2cad7bec 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -92,6 +92,26 @@ def SelectOpsConversion : Pass<"fir-select-ops-conversion", "::mlir::func::FuncO
"mlir::cf::ControlFlowDialect", "fir::FIROpsDialect"];
}
+def LiftSCFWhileToSCFFor : Pass<"lift-scf-while-to-scf-for",
+ "::mlir::func::FuncOp"> {
+ let summary = "Rewrite recognizable counted scf.while loops as scf.for";
+ let description = [{
+ Walks every scf.while operation in the function and, when the canonical
+ "trip-counter + induction variable" shape is recognized, rewrites the
+ loop as scf.for. Other loop-carried before-args are passed through as
+ iter_args. Loops that do not match the builtin heuristic (data-dependent
+ updates, used results, non-positive IV step, etc.) are left untouched
+ and a diagnostic is emitted to stderr.
+
+ The main use pattern today is to rebuild counted loops that upstream
+ `lift-cf-to-scf` had to leave as scf.while because it can only recover
+ do-while shape from an unstructured CFG.
+ }];
+
+ let dependentDialects = ["mlir::arith::ArithDialect",
+ "mlir::scf::SCFDialect"];
+}
+
def FIRToSCFPass : Pass<"fir-to-scf"> {
let summary = "Convert FIR structured control flow ops to SCF dialect.";
let description = [{
diff --git a/flang/lib/Optimizer/Transforms/CMakeLists.txt b/flang/lib/Optimizer/Transforms/CMakeLists.txt
index f26f8c5c64bb0..f0d80454922ec 100644
--- a/flang/lib/Optimizer/Transforms/CMakeLists.txt
+++ b/flang/lib/Optimizer/Transforms/CMakeLists.txt
@@ -42,6 +42,7 @@ add_flang_library(FIRTransforms
PolymorphicOpConversion.cpp
FunctionAttr.cpp
GenRuntimeCallsForTest.cpp
+ LiftSCFWhileToSCFFor.cpp
LoopInvariantCodeMotion.cpp
LoopVersioning.cpp
SelectOpsConversion.cpp
diff --git a/flang/lib/Optimizer/Transforms/LiftSCFWhileToSCFFor.cpp b/flang/lib/Optimizer/Transforms/LiftSCFWhileToSCFFor.cpp
new file mode 100644
index 0000000000000..0978e8ae80695
--- /dev/null
+++ b/flang/lib/Optimizer/Transforms/LiftSCFWhileToSCFFor.cpp
@@ -0,0 +1,560 @@
+//===- LiftSCFWhileToSCFFor.cpp - Rewrite counted scf.while as scf.for ----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Walks each scf.while produced by lift-cf-to-scf and, when the canonical
+// shape is recognized, rewrites it as an scf.for.
+//
+// Recognized shape:
+//
+// scf.while (%trip = %tripInit, %iv = %ivInit, %x = %xInit, ...)
+// : (i32, i32, T, ...) -> (i32, i32, T, ...) {
+// %cmp = arith.cmpi sgt, %trip, %c0 : i32
+// %r:M = scf.if %cmp -> (i32, i32, T, ..., i32) {
+// // ...body...
+// %tripNext = arith.subi %trip, %c1 : i32
+// %ivNext = arith.addi %iv, %ivStep : i32 // or arith.subi
+// %xNext = <any expr>
+// scf.yield %tripNext, %ivNext, %xNext, ..., %c1_marker : ...
+// } else {
+// scf.yield %poison, %poison, %poison, ..., %c0_marker : ...
+// }
+// %enter = arith.trunci %r#M-1 : i32 to i1
+// scf.condition(%enter) %r#0, %r#1, %r#2, ... : ...
+// } do {
+// ^bb0(%a0, %a1, %a2, ...):
+// scf.yield %a0, %a1, %a2, ... : ... // identity
+// }
+//
+// Rewrite (canonical case: tripStep = -1, ivStep > 0):
+//
+// %N = %tripInit
+// %scaledN = arith.muli %N, %ivStep : i32 // skipped when ivStep == 1
+// %ub = arith.addi %ivInit, %scaledN : i32 // or %ivInit + %N
+// scf.for %i = %ivInit to %ub step %ivStep iter_args(%x = %xInit) -> (T) {
+// // body cloned from the scf.if continues branch, with
+// // %iv -> %i, %x -> iter-arg, %trip unused
+// scf.yield %xNext : T
+// }
+//
+// Bailouts (no rewrite, just a stderr diagnostic):
+// - any result of the scf.while is used
+// - the scf.condition predicate doesn't trace to an scf.if's condition (we
+// need the if to extract the body from its continues branch)
+// - more than one induction-variable candidate (in addition to trip)
+// - zero induction-variable candidates
+// - ivStep is not positive (scf.for requires step > 0)
+// - trip step != -1 (N derivation would be wrong)
+// - the trip arg is referenced anywhere in the body except in the
+// recurrence and the gating cmp
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/Matchers.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace fir {
+#define GEN_PASS_DEF_LIFTSCFWHILETOSCFFOR
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+#define DEBUG_TYPE "lift-scf-while-to-scf-for"
+
+namespace {
+using namespace fir;
+using namespace mlir;
+
+// True if `x` is loop-invariant w.r.t. `whileOp` — its definition lives
+// strictly outside the op.
+static bool isLoopInvariant(Value x, scf::WhileOp whileOp) {
+ if (Operation *def = x.getDefiningOp()) {
+ return !whileOp->isProperAncestor(def);
+ }
+ Block *owner = cast<BlockArgument>(x).getOwner();
+ Operation *ownerOp = owner->getParentOp();
+ if (ownerOp == whileOp.getOperation()) {
+ return false;
+ }
+ return !whileOp->isProperAncestor(ownerOp);
+}
+
+// Returns the value that, in the next iteration of `whileOp`, becomes the
+// before-block argument at `argIdx`.
+static Value getNextIterValueInBefore(scf::WhileOp whileOp, unsigned argIdx) {
+ auto yieldOp = cast<scf::YieldOp>(whileOp.getAfterBody()->getTerminator());
+ Value yielded = yieldOp.getOperand(argIdx);
+
+ // This is the canonical case since the after region just passes through
+ // values computed in the before region.
+ if (auto afterArg = dyn_cast<BlockArgument>(yielded)) {
+ if (afterArg.getOwner() == whileOp.getAfterBody()) {
+ auto condOp = whileOp.getConditionOp();
+ return condOp.getArgs()[afterArg.getArgNumber()];
+ }
+ }
+
+ return yielded;
+}
+
+struct AffineRecurrence {
+ // `op` is null if we failed to find an affine recurrence for a value `v`.
+ Operation *op = nullptr;
+ Value step;
+ bool isSub = false;
+ explicit operator bool() const { return op != nullptr; }
+};
+
+// If `v` is `arg + c` or `arg - c` (with `arg` matching `beforeArg` and `c`
+// loop-invariant w.r.t. `whileOp`), return the recurrence op, the step value,
+// and whether the recurrence is a subtraction. The returned step is the raw
+// RHS; for arith.subi the caller should interpret it as a negative step.
+static AffineRecurrence matchAffineRecurrence(
+ Value v, Value beforeArg, scf::WhileOp whileOp) {
+ if (auto add = v.getDefiningOp<arith::AddIOp>()) {
+ if (add.getLhs() == beforeArg && isLoopInvariant(add.getRhs(), whileOp)) {
+ return {add.getOperation(), add.getRhs(), /*isSub=*/false};
+ }
+ if (add.getRhs() == beforeArg && isLoopInvariant(add.getLhs(), whileOp)) {
+ return {add.getOperation(), add.getLhs(), /*isSub=*/false};
+ }
+ }
+
+ if (auto sub = v.getDefiningOp<arith::SubIOp>()) {
+ if (sub.getLhs() == beforeArg && isLoopInvariant(sub.getRhs(), whileOp)) {
+ return {sub.getOperation(), sub.getRhs(), /*isSub=*/true};
+ }
+ }
+
+ return {};
+}
+
+// Locate the "continues" branch of an scf.if whose result feeds scf.condition.
+// We use the heuristic: the structurizer encodes "loop continues" as an extra
+// i32 result that gets trunc'd to the scf.condition predicate. The branch
+// whose yield at that result index is a non-zero constant is the continues
+// branch.
+static Block *getContinuesBranch(scf::IfOp ifOp, Value condPredicate) {
+ auto trunc = condPredicate.getDefiningOp<arith::TruncIOp>();
+ if (!trunc) {
+ return ifOp.thenBlock();
+ }
+
+ Value src = trunc.getIn();
+ auto opRes = dyn_cast<OpResult>(src);
+ if (!opRes || opRes.getOwner() != ifOp.getOperation()) {
+ return ifOp.thenBlock();
+ }
+
+ unsigned predIdx = opRes.getResultNumber();
+
+ auto thenYield = cast<scf::YieldOp>(ifOp.thenBlock()->getTerminator());
+ auto elseYield = cast<scf::YieldOp>(ifOp.elseBlock()->getTerminator());
+
+ auto isNonZeroConst = [](Value v) {
+ IntegerAttr attr;
+ return matchPattern(v, m_Constant(&attr)) && attr.getInt() != 0;
+ };
+
+ if (isNonZeroConst(thenYield.getOperand(predIdx))) {
+ return ifOp.thenBlock();
+ }
+
+ if (isNonZeroConst(elseYield.getOperand(predIdx))) {
+ return ifOp.elseBlock();
+ }
+
+ return ifOp.thenBlock();
+}
+
+//===----------------------------------------------------------------------===//
+// LoopInfo
+//===----------------------------------------------------------------------===//
+
+// Everything we need to drive the scf.while → scf.for rewrite when the
+// canonical shape is recognized.
+struct LoopInfo {
+ scf::WhileOp whileOp;
+
+ // The continues branch of the scf.if that gates the loop. Holds the body
+ // we'll clone into the new scf.for.
+ Block *continuesBlock = nullptr;
+
+ // Trip counter.
+ Value tripInit;
+ Operation *tripRecurrence = nullptr;
+
+ // Sole induction variable (becomes the scf.for IV).
+ unsigned ivArgIdx = 0;
+ Value ivInit; // lb
+ Value ivStep; // signed, must be positive (we reject subi for now)
+ Operation *ivRecurrence = nullptr;
+
+ // Other before-args that pass through as iter_args of the new scf.for.
+ // For each:
+ // - `argIdx` : index in whileOp.getBeforeArguments(); also the
+ // index into whileOp.getInits() that gives the
+ // initial value of the iter_arg.
+ // - `contYieldSlot` : operand index in the continues-branch yield that
+ // carries this iter_arg's next-iteration value.
+ // Equivalently, the index of the corresponding
+ // scf.if result.
+ // These are stored separately because the before-arg lane and the
+ // condition/scf.if-result lane are independent in scf.while: their counts
+ // and permutations need not match.
+ struct IterArg {
+ unsigned argIdx;
+ unsigned contYieldSlot;
+ };
+ SmallVector<IterArg> iterArgs;
+
+ LoopInfo(scf::WhileOp whileOp, Block *continuesBlock, Value tripInit,
+ Operation *tripRecurrence, unsigned ivArgIdx, Value ivInit,
+ Value ivStep, Operation *ivRecurrence,
+ SmallVector<IterArg> iterArgs)
+ : whileOp(whileOp), continuesBlock(continuesBlock), tripInit(tripInit),
+ tripRecurrence(tripRecurrence), ivArgIdx(ivArgIdx), ivInit(ivInit),
+ ivStep(ivStep), ivRecurrence(ivRecurrence),
+ iterArgs(std::move(iterArgs)) {}
+};
+
+// Try to build a LoopInfo for `whileOp`. On failure, emit a one-line stderr
+// diagnostic explaining why and return std::nullopt.
+static std::optional<LoopInfo> tryBuildLoopInfo(
+ scf::WhileOp whileOp, llvm::raw_ostream &os) {
+ auto bail = [&](StringRef reason) -> std::optional<LoopInfo> {
+ os << " [skip] " << reason << "\n";
+ return std::nullopt;
+ };
+
+ // Results of the scf.while must be unused.
+ for (Value r : whileOp.getResults()) {
+ if (!r.use_empty()) {
+ return bail("scf.while result has uses");
+ }
+ }
+
+ // The condition predicate must come from an scf.if whose continues branch
+ // we can use as the loop body.
+ Value condPredicate = whileOp.getConditionOp().getCondition();
+ auto trunc = condPredicate.getDefiningOp<arith::TruncIOp>();
+
+ // TODO: Maybe we can relax that in the future.
+ if (!trunc) {
+ return bail("scf.condition predicate is not arith.trunci");
+ }
+
+ auto ifOp = trunc.getIn().getDefiningOp<scf::IfOp>();
+
+ if (!ifOp) {
+ return bail("scf.condition predicate does not trace to an scf.if");
+ }
+
+ // The scf.if's predicate is the comparison that gates loop continuation.
+ arith::CmpIOp cmp = ifOp.getCondition().getDefiningOp<arith::CmpIOp>();
+
+ // TODO: Maybe we can make that more generic in the future.
+ if (!cmp) {
+ return bail("scf.if condition is not arith.cmpi");
+ }
+
+ // Find the IV-side of the gating cmp.
+ bool lhsInv = isLoopInvariant(cmp.getLhs(), whileOp);
+ bool rhsInv = isLoopInvariant(cmp.getRhs(), whileOp);
+ Value cmpIV;
+
+ if (rhsInv && !lhsInv) {
+ cmpIV = cmp.getLhs();
+ } else if (lhsInv && !rhsInv) {
+ cmpIV = cmp.getRhs();
+ } else {
+ return bail("gating compare has no invariant side");
+ }
+
+ // Classify each before-arg as trip / IV / iter_arg.
+ std::optional<unsigned> tripCounterArgIdx;
+ AffineRecurrence tripCounterRecurrence;
+ SmallVector<std::pair<unsigned, AffineRecurrence>> nonTripCounterRecurrences;
+ SmallVector<LoopInfo::IterArg> iterArgs;
+
+ Block *cont = getContinuesBranch(ifOp, condPredicate);
+ auto contYield = cast<scf::YieldOp>(cont->getTerminator());
+
+ for (auto [argIdx, beforeArg] :
+ llvm::enumerate(whileOp.getBeforeArguments())) {
+ Value nextInBefore = getNextIterValueInBefore(whileOp, (unsigned)argIdx);
+
+ // In the canonical lift-cf-to-scf shape every iteration value flows
+ // through one of `ifOp`'s results. We need the result index — both to
+ // match the recurrence in the continues branch (for the trip + IV) and
+ // to find the next-iteration value for pass-through iter_args.
+ auto opRes = dyn_cast<OpResult>(nextInBefore);
+ if (!opRes || opRes.getOwner() != ifOp.getOperation()) {
+ return bail("before-arg's next-iter value does not come from the "
+ "predicate scf.if");
+ }
+ unsigned contYieldSlot = opRes.getResultNumber();
+ Value branchSource = contYield.getOperand(contYieldSlot);
+ AffineRecurrence rec =
+ matchAffineRecurrence(branchSource, beforeArg, whileOp);
+
+ if (beforeArg == cmpIV) {
+ if (!rec) {
+ return bail("gating cmp's IV operand has no affine recurrence");
+ }
+
+ if (tripCounterArgIdx) {
+ return bail("multiple before-args match the gating cmp's IV operand");
+ }
+
+ tripCounterArgIdx = (unsigned)argIdx;
+ tripCounterRecurrence = rec;
+ continue;
+ }
+
+ if (rec) {
+ nonTripCounterRecurrences.push_back({(unsigned)argIdx, rec});
+ } else {
+ iterArgs.push_back({(unsigned)argIdx, contYieldSlot});
+ }
+ }
+
+ if (!tripCounterArgIdx) {
+ return bail("no before-arg matches the gating cmp's IV operand");
+ }
+
+ // Trip step must be -1 (canonical lift output).
+ IntegerAttr tripStepAttr;
+ bool tripStepIsOne =
+ matchPattern(tripCounterRecurrence.step, m_Constant(&tripStepAttr)) &&
+ tripStepAttr.getInt() == 1;
+
+ if (!tripCounterRecurrence.isSub || !tripStepIsOne) {
+ return bail("trip-counter step is not -1");
+ }
+
+ // Exactly one non-trip recurrence is the induction variable. Other
+ // recurrence-shaped before-args could in principle be supported as
+ // affine-update iter_args, but that's not yet implemented.
+ if (nonTripCounterRecurrences.empty()) {
+ return bail("no induction variable candidate found");
+ }
+
+ if (nonTripCounterRecurrences.size() > 1) {
+ return bail("multiple induction variable candidates (not yet supported)");
+ }
+
+ auto [ivIdx, ivRec] = nonTripCounterRecurrences.front();
+ if (ivRec.isSub) {
+ return bail("induction variable step is negative (not yet supported)");
+ }
+
+ LoopInfo info(whileOp, cont, whileOp.getInits()[*tripCounterArgIdx],
+ tripCounterRecurrence.op, ivIdx, whileOp.getInits()[ivIdx],
+ ivRec.step, ivRec.op, std::move(iterArgs));
+
+ // The trip before-arg may only be used by the gating cmp and its own
+ // recurrence. Any other use would have semantics we can't preserve after
+ // dropping the trip counter.
+ Value tripCounterArg = whileOp.getBeforeArguments()[*tripCounterArgIdx];
+
+ for (OpOperand &use : tripCounterArg.getUses()) {
+ Operation *user = use.getOwner();
+
+ if (user == cmp.getOperation() || user == tripCounterRecurrence.op) {
+ continue;
+ }
+
+ return bail("trip-counter before-arg has uses outside the cmp/recurrence");
+ }
+
+ // The trip recurrence's result must be used only by the scf.yield that
+ // propagates it back to the next iteration. Anything else would require us
+ // to materialize a trip-equivalent expression in the new scf.for, which we
+ // don't support.
+ for (OpOperand &use : tripCounterRecurrence.op->getResult(0).getUses()) {
+ if (!isa<scf::YieldOp>(use.getOwner())) {
+ return bail("trip-counter recurrence result has non-yield uses");
+ }
+ }
+
+ return info;
+}
+
+//===----------------------------------------------------------------------===//
+// Rewrite
+//===----------------------------------------------------------------------===//
+
+static void rewriteToSCFFor(const LoopInfo &info, llvm::raw_ostream &os) {
+ scf::WhileOp whileOp = info.whileOp;
+ OpBuilder builder(whileOp);
+ Location loc = whileOp.getLoc();
+
+ // The original IV is typed at whatever the source-level induction variable
+ // used (commonly i32 for Fortran). Downstream affine / index-based
+ // conversions expect index-typed scf.for bounds, so we always emit an
+ // index-typed scf.for. The original IV type is restored via
+ // arith.index_cast inside the body.
+ Type origType = info.ivInit.getType();
+ Type indexType = builder.getIndexType();
+ auto castTo = [&](Value v, Type target) -> Value {
+ if (v.getType() == target) {
+ return v;
+ }
+ return arith::IndexCastOp::create(builder, loc, target, v);
+ };
+
+ // 1. Materialize ub_excl in the original type just before the scf.while,
+ // then cast lb/ub/step to index.
+ //
+ // N = tripInit
+ // scaled = (ivStep == 1) ? N : arith.muli N, ivStep
+ // ub = arith.addi ivInit, scaled
+ IntegerAttr ivStepAttr;
+ bool ivStepIsOne = matchPattern(info.ivStep, m_Constant(&ivStepAttr)) &&
+ ivStepAttr.getInt() == 1;
+ Value scaled = ivStepIsOne
+ ? info.tripInit
+ : arith::MulIOp::create(builder, loc, info.tripInit, info.ivStep);
+ Value ub = arith::AddIOp::create(builder, loc, info.ivInit, scaled);
+
+ Value lbIdx = castTo(info.ivInit, indexType);
+ Value ubIdx = castTo(ub, indexType);
+ Value stepIdx = castTo(info.ivStep, indexType);
+
+ // 2. Collect iter_arg inits (from the iter_args we're passing through) and
+ // the corresponding "next" values (yielded inside the continues branch).
+ // The before-arg lane and the continues-yield lane are independent in
+ // scf.while, so we use the per-iter-arg `contYieldSlot` recorded during
+ // analysis to index `contYield` rather than the before-arg index.
+ auto contYield = cast<scf::YieldOp>(info.continuesBlock->getTerminator());
+ SmallVector<Value> iterInits;
+ SmallVector<Value> iterNexts;
+ iterInits.reserve(info.iterArgs.size());
+ iterNexts.reserve(info.iterArgs.size());
+ for (const LoopInfo::IterArg &ia : info.iterArgs) {
+ iterInits.push_back(whileOp.getInits()[ia.argIdx]);
+ iterNexts.push_back(contYield.getOperand(ia.contYieldSlot));
+ }
+
+ // 3. Create the scf.for with index-typed bounds. The default ForOp builder
+ // only auto-inserts a `scf.yield` terminator when `iterInits` is empty;
+ // for the non-empty case it expects the caller to either provide a body
+ // builder or insert the terminator itself. Insert a placeholder yield
+ // (operands populated in step 6) so the rest of the rewrite can position
+ // its inserts via `forBody->getTerminator()` uniformly.
+ auto forOp =
+ scf::ForOp::create(builder, loc, lbIdx, ubIdx, stepIdx, iterInits);
+
+ Block *forBody = forOp.getBody();
+ if (forBody->empty() || !forBody->back().hasTrait<OpTrait::IsTerminator>()) {
+ OpBuilder termBuilder(builder.getContext());
+ termBuilder.setInsertionPointToEnd(forBody);
+ scf::YieldOp::create(termBuilder, loc, ValueRange{});
+ }
+ builder.setInsertionPoint(forBody->getTerminator());
+
+ // 4. Build the value mapping:
+ // - Original IV before-arg → arith.index_cast of the scf.for IV (back to
+ // the original IV's type) so body ops keep working unchanged.
+ // - Original iter_args before-args → scf.for iter_arg block args.
+ // - Original trip before-arg → unused (validated to have no body uses).
+ // - Original trip recurrence result → unused (validated above).
+ //
+ // The IV recurrence is left in the body-cloning loop below: SSA puts
+ // it before any user, the `ivArg → ivAsOrig` mapping handles its IV
+ // operand, and the loop-invariant step passes through as-is. Cloning
+ // rather than re-creating preserves attributes (nsw, nuw, ...) and
+ // works whether the original was arith.addi or arith.subi.
+ IRMapping mapping;
+ Value ivAsOrig = castTo(forOp.getInductionVar(), origType);
+
+ mapping.map(whileOp.getBeforeArguments()[info.ivArgIdx], ivAsOrig);
+ for (auto [i, ia] : llvm::enumerate(info.iterArgs)) {
+ mapping.map(whileOp.getBeforeArguments()[ia.argIdx],
+ forBody->getArgument(1 + i));
+ }
+
+ // 5. Clone every op from the continues block into the scf.for body, except:
+ // - the trip recurrence (its operand %trip has no mapping; cloning it
+ // would leave a dangling reference to the erased before-block arg)
+ // - the scf.yield terminator (we emit our own below)
+ for (Operation &op : *info.continuesBlock) {
+ if (&op == info.tripRecurrence) {
+ continue;
+ }
+ if (isa<scf::YieldOp>(op)) {
+ continue;
+ }
+ builder.clone(op, mapping);
+ }
+
+ // 6. Update the scf.for's existing terminator with the remapped next-values
+ // for the iter_args.
+ SmallVector<Value> yieldOperands;
+ yieldOperands.reserve(iterNexts.size());
+ for (Value v : iterNexts) {
+ yieldOperands.push_back(mapping.lookupOrDefault(v));
+ }
+ forBody->getTerminator()->setOperands(yieldOperands);
+
+ // 7. The scf.while's results (trip exit value, iv exit value, iter_arg exit
+ // values) are all unused by precondition. We can simply erase the while.
+ whileOp.erase();
+
+ os << " [rewritten] scf.while → scf.for at ";
+ loc.print(os);
+ os << "\n";
+}
+
+//===----------------------------------------------------------------------===//
+// Pass driver
+//===----------------------------------------------------------------------===//
+
+class LiftSCFWhileToSCFFor
+ : public fir::impl::LiftSCFWhileToSCFForBase<LiftSCFWhileToSCFFor> {
+public:
+ using LiftSCFWhileToSCFForBase<
+ LiftSCFWhileToSCFFor>::LiftSCFWhileToSCFForBase;
+
+ void runOnOperation() override {
+ auto &os = llvm::errs();
+ // Collect first; rewrite separately so we don't mutate during walk.
+ SmallVector<scf::WhileOp> worklist;
+ getOperation()->walk([&](scf::WhileOp w) { worklist.push_back(w); });
+
+ os << "\n[lift-scf-while-to-scf-for] discovered " << worklist.size()
+ << " scf.while op(s) in " << getOperation()->getName() << "\n";
+
+ unsigned rewritten = 0;
+ for (scf::WhileOp whileOp : worklist) {
+ os << "\n[lift-scf-while-to-scf-for] scf.while at ";
+ whileOp.getLoc().print(os);
+ os << "\n";
+
+ std::optional<LoopInfo> info = tryBuildLoopInfo(whileOp, os);
+ if (!info) {
+ continue;
+ }
+ rewriteToSCFFor(*info, os);
+ ++rewritten;
+ }
+
+ if (!worklist.empty()) {
+ unsigned skipped = worklist.size() - rewritten;
+ os << "\n[lift-scf-while-to-scf-for] summary: " << rewritten
+ << " rewritten, " << skipped << " skipped (of " << worklist.size()
+ << " discovered)\n";
+ }
+ }
+};
+
+} // namespace
diff --git a/flang/test/Fir/LiftSCFWhileToSCFFor/iv-next-body-use.mlir b/flang/test/Fir/LiftSCFWhileToSCFFor/iv-next-body-use.mlir
new file mode 100644
index 0000000000000..7dafa6e2a4aa6
--- /dev/null
+++ b/flang/test/Fir/LiftSCFWhileToSCFFor/iv-next-body-use.mlir
@@ -0,0 +1,54 @@
+// The IV recurrence's result (%ivNext) has body uses beyond the
+// scf.yield — here a `memref.store`, modelling the `fir.declare_value` ops
+// mem2reg emits after promoting a Fortran scalar variable. The body-clone
+// loop walks the continues block in source order; the IV recurrence is
+// cloned as a regular body op (its operand %iv resolves to ivAsOrig via
+// the IRMapping, and the loop-invariant step passes through unmapped).
+// SSA ordering guarantees the clone is in the new body before any cloned
+// user of its result. Without the mapping, the cloned `memref.store` would
+// hold a dangling reference to the original %ivNext inside the erased
+// scf.while.
+
+// RUN: fir-opt %s --lift-scf-while-to-scf-for 2>/dev/null | FileCheck %s
+
+func.func @iv_next_with_body_use(%lb: i32, %N: i32, %sink: memref<i32>) {
+ %c0 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : i32
+ %poison = arith.constant 0 : i32
+ %r:2 = scf.while (%trip = %N, %iv = %lb) : (i32, i32) -> (i32, i32) {
+ %cmp = arith.cmpi sgt, %trip, %c0 : i32
+ %ifr:3 = scf.if %cmp -> (i32, i32, i32) {
+ %tripNext = arith.subi %trip, %c1 : i32
+ %ivNext = arith.addi %iv, %c1 : i32
+ // %ivNext is consumed both by the scf.yield AND by this store;
+ // the latter survives cloning and must point at the new ivNextInFor.
+ memref.store %ivNext, %sink[] : memref<i32>
+ scf.yield %tripNext, %ivNext, %c1 : i32, i32, i32
+ } else {
+ scf.yield %poison, %poison, %c0 : i32, i32, i32
+ }
+ %enter = arith.trunci %ifr#2 : i32 to i1
+ scf.condition(%enter) %ifr#0, %ifr#1 : i32, i32
+ } do {
+ ^bb0(%a0: i32, %a1: i32):
+ scf.yield %a0, %a1 : i32, i32
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @iv_next_with_body_use(
+// CHECK-SAME: %{{.*}}: i32, %[[N:.*]]: i32, %[[SINK:.*]]: memref<i32>)
+
+// CHECK-NOT: scf.while
+
+// Inside the new scf.for body, the IV is cast back to i32 (this is the
+// `ivAsOrig` value), then the IV-recurrence equivalent is materialized as
+// `arith.addi %ivAsOrig, %c1`. The cloned `memref.store` must reference
+// THAT materialized value, NOT a dangling pointer into the destroyed
+// scf.if.
+//
+// CHECK: scf.for %[[I:.*]] = %{{.*}} to %{{.*}} step %{{.*}}
+// CHECK: %[[IVi32:.*]] = arith.index_cast %[[I]] : index to i32
+// CHECK: %[[IVNEXT:.*]] = arith.addi %[[IVi32]], %{{.*}} : i32
+// CHECK: memref.store %[[IVNEXT]], %[[SINK]][] : memref<i32>
+// CHECK: }
diff --git a/flang/test/Fir/LiftSCFWhileToSCFFor/step-gt-one.mlir b/flang/test/Fir/LiftSCFWhileToSCFFor/step-gt-one.mlir
new file mode 100644
index 0000000000000..186ae78530501
--- /dev/null
+++ b/flang/test/Fir/LiftSCFWhileToSCFFor/step-gt-one.mlir
@@ -0,0 +1,51 @@
+// Counted scf.while with ivStep > 1.
+//
+// The pass should scale the trip count by the step before adding to the
+// initial value: `ub = ivInit + (tripInit * ivStep)`.
+
+// RUN: fir-opt %s --lift-scf-while-to-scf-for 2>/dev/null | FileCheck %s
+
+func.func @counted_step2(%lb: i32, %N: i32) {
+ %c0 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : i32
+ %c2 = arith.constant 2 : i32
+ %poison = arith.constant 0 : i32
+ %r:2 = scf.while (%trip = %N, %iv = %lb) : (i32, i32) -> (i32, i32) {
+ %cmp = arith.cmpi sgt, %trip, %c0 : i32
+ %ifr:3 = scf.if %cmp -> (i32, i32, i32) {
+ %tripNext = arith.subi %trip, %c1 : i32
+ %ivNext = arith.addi %iv, %c2 : i32
+ scf.yield %tripNext, %ivNext, %c1 : i32, i32, i32
+ } else {
+ scf.yield %poison, %poison, %c0 : i32, i32, i32
+ }
+ %enter = arith.trunci %ifr#2 : i32 to i1
+ scf.condition(%enter) %ifr#0, %ifr#1 : i32, i32
+ } do {
+ ^bb0(%a0: i32, %a1: i32):
+ scf.yield %a0, %a1 : i32, i32
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @counted_step2(
+// CHECK-SAME: %[[LB:.*]]: i32, %[[N:.*]]: i32)
+
+// scf.while is gone, scf.for takes over.
+// CHECK-NOT: scf.while
+
+// Step is the 2 constant.
+// CHECK: %[[C2:.*]] = arith.constant 2 : i32
+
+// ub = lb + (N * step), with step == 2.
+// CHECK: %[[SCALED:.*]] = arith.muli %[[N]], %[[C2]] : i32
+// CHECK: %[[UB:.*]] = arith.addi %[[LB]], %[[SCALED]] : i32
+
+// lb, ub and step cast to index for the scf.for.
+// CHECK: %[[LB_IDX:.*]] = arith.index_cast %[[LB]] : i32 to index
+// CHECK: %[[UB_IDX:.*]] = arith.index_cast %[[UB]] : i32 to index
+// CHECK: %[[STEP_IDX:.*]] = arith.index_cast %[[C2]] : i32 to index
+
+// scf.for over that range.
+// CHECK: scf.for %{{.*}} = %[[LB_IDX]] to %[[UB_IDX]] step %[[STEP_IDX]] {
+// CHECK: }
diff --git a/flang/test/Fir/LiftSCFWhileToSCFFor/step-one.mlir b/flang/test/Fir/LiftSCFWhileToSCFFor/step-one.mlir
new file mode 100644
index 0000000000000..1fb97686b1fab
--- /dev/null
+++ b/flang/test/Fir/LiftSCFWhileToSCFFor/step-one.mlir
@@ -0,0 +1,48 @@
+// Counted scf.while with ivStep == 1.
+//
+// The pass should emit the upper bound as `arith.addi ivInit, tripInit` —
+// no `arith.muli`, since multiplying tripInit by 1 is a no-op.
+
+// RUN: fir-opt %s --lift-scf-while-to-scf-for 2>/dev/null | FileCheck %s
+
+func.func @counted_step1(%lb: i32, %N: i32) {
+ %c0 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : i32
+ %poison = arith.constant 0 : i32
+ %r:2 = scf.while (%trip = %N, %iv = %lb) : (i32, i32) -> (i32, i32) {
+ %cmp = arith.cmpi sgt, %trip, %c0 : i32
+ %ifr:3 = scf.if %cmp -> (i32, i32, i32) {
+ %tripNext = arith.subi %trip, %c1 : i32
+ %ivNext = arith.addi %iv, %c1 : i32
+ scf.yield %tripNext, %ivNext, %c1 : i32, i32, i32
+ } else {
+ scf.yield %poison, %poison, %c0 : i32, i32, i32
+ }
+ %enter = arith.trunci %ifr#2 : i32 to i1
+ scf.condition(%enter) %ifr#0, %ifr#1 : i32, i32
+ } do {
+ ^bb0(%a0: i32, %a1: i32):
+ scf.yield %a0, %a1 : i32, i32
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @counted_step1(
+// CHECK-SAME: %[[LB:.*]]: i32, %[[N:.*]]: i32)
+
+// scf.while is gone, scf.for takes over.
+// CHECK-NOT: scf.while
+
+// ub = lb + N (no muli, since ivStep == 1)
+// CHECK-NOT: arith.muli
+// CHECK: %[[UB:.*]] = arith.addi %[[LB]], %[[N]] : i32
+
+// lb, ub and step cast to index for the scf.for.
+// CHECK: %[[LB_IDX:.*]] = arith.index_cast %[[LB]] : i32 to index
+// CHECK: %[[UB_IDX:.*]] = arith.index_cast %[[UB]] : i32 to index
+// CHECK: %[[STEP_IDX:.*]] = arith.index_cast %{{.*}} : i32 to index
+
+// scf.for over that range; no muli inside the body either.
+// CHECK: scf.for %{{.*}} = %[[LB_IDX]] to %[[UB_IDX]] step %[[STEP_IDX]] {
+// CHECK-NOT: arith.muli
+// CHECK: }
diff --git a/flang/test/Fir/LiftSCFWhileToSCFFor/two-iter-args.mlir b/flang/test/Fir/LiftSCFWhileToSCFFor/two-iter-args.mlir
new file mode 100644
index 0000000000000..1d2fca836b061
--- /dev/null
+++ b/flang/test/Fir/LiftSCFWhileToSCFFor/two-iter-args.mlir
@@ -0,0 +1,72 @@
+// Counted scf.while with one induction variable and two pass-through
+// iter_args. Verifies that:
+// * the rewrite preserves both iter_args in order, threading their inits
+// from the scf.while inits and their next-iteration values from the
+// correct contYield slots,
+// * uses of the original iter_arg before-args inside the body are
+// remapped to the corresponding scf.for iter_arg block arguments,
+// * the final scf.yield carries the cloned next-values for both
+// iter_args.
+
+// RUN: fir-opt %s --lift-scf-while-to-scf-for 2>/dev/null | FileCheck %s
+
+func.func @two_iter(%lb: i32, %N: i32, %x0: i32, %y0: i32) {
+ %c0 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : i32
+ %poison = arith.constant 0 : i32
+ %r:4 = scf.while (%trip = %N, %iv = %lb, %x = %x0, %y = %y0)
+ : (i32, i32, i32, i32) -> (i32, i32, i32, i32) {
+ %cmp = arith.cmpi sgt, %trip, %c0 : i32
+ %ifr:5 = scf.if %cmp -> (i32, i32, i32, i32, i32) {
+ %tripNext = arith.subi %trip, %c1 : i32
+ %ivNext = arith.addi %iv, %c1 : i32
+ // non-affine recurrence in %x (rhs is the IV, not a loop invariant)
+ %xNext = arith.muli %x, %iv : i32
+ // non-affine recurrence in %y (subi where rhs is the IV, not invariant)
+ %yNext = arith.subi %y, %iv : i32
+ scf.yield %tripNext, %ivNext, %xNext, %yNext, %c1
+ : i32, i32, i32, i32, i32
+ } else {
+ scf.yield %poison, %poison, %poison, %poison, %c0
+ : i32, i32, i32, i32, i32
+ }
+ %enter = arith.trunci %ifr#4 : i32 to i1
+ scf.condition(%enter) %ifr#0, %ifr#1, %ifr#2, %ifr#3 : i32, i32, i32, i32
+ } do {
+ ^bb0(%a0: i32, %a1: i32, %a2: i32, %a3: i32):
+ scf.yield %a0, %a1, %a2, %a3 : i32, i32, i32, i32
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @two_iter(
+// CHECK-SAME: %[[LB:.*]]: i32, %[[N:.*]]: i32, %[[X0:.*]]: i32, %[[Y0:.*]]: i32)
+
+// scf.while is replaced by scf.for.
+// CHECK-NOT: scf.while
+
+// Bounds materialization (ivStep == 1, so no muli; just `lb + N`).
+// CHECK-NOT: arith.muli {{%.*}}, {{%c1}}
+// CHECK: %[[UB:.*]] = arith.addi %[[LB]], %[[N]] : i32
+
+// Index-typed bounds.
+// CHECK: %[[LB_IDX:.*]] = arith.index_cast %[[LB]]
+// CHECK: %[[UB_IDX:.*]] = arith.index_cast %[[UB]]
+// CHECK: %[[STEP_IDX:.*]] = arith.index_cast %{{.*}}
+
+// scf.for with two iter_args, in the same order as the source's non-IV
+// before-args. The first iter_arg is initialized from %x0, the second
+// from %y0.
+// CHECK: scf.for %[[I:.*]] = %[[LB_IDX]] to %[[UB_IDX]] step %[[STEP_IDX]]
+// CHECK-SAME: iter_args(%[[XARG:.*]] = %[[X0]], %[[YARG:.*]] = %[[Y0]])
+// CHECK-SAME: -> (i32, i32)
+
+// Body: cast IV back to i32 and recompute the two non-affine next-values
+// using the iter_arg block arguments (not the original SSA names).
+// CHECK: %[[I_I32:.*]] = arith.index_cast %[[I]] : index to i32
+// CHECK: %[[XNEXT:.*]] = arith.muli %[[XARG]], %[[I_I32]] : i32
+// CHECK: %[[YNEXT:.*]] = arith.subi %[[YARG]], %[[I_I32]] : i32
+
+// Yield carries the next-iter values in the same iter_args order.
+// CHECK: scf.yield %[[XNEXT]], %[[YNEXT]] : i32, i32
+// CHECK: }
More information about the flang-commits
mailing list