[flang-commits] [flang] [flang] change yielded iv value to be `iv + step` (PR #174124)

[Susan Tan ス-ザン タン via flang-commits]

https://github.com/SusanTan created https://github.com/llvm/llvm-project/pull/174124

In cases where induction variables are used after the loop,  like 

```
write(*,*) (a(j),j=1,10) 
print *, j
```

the incremented value should be used. Updating the FIRToSCF pass to support this.

>From 0ed1313746db1b494d181577da30b0ef5d19311d Mon Sep 17 00:00:00 2001
From: Susan Tan <zujunt at nvidia.com>
Date: Wed, 31 Dec 2025 14:15:22 -0800
Subject: [PATCH] change yielded iv

---
 flang/lib/Optimizer/Transforms/FIRToSCF.cpp | 30 +++++++++++--
 flang/test/Fir/FirToSCF/do-loop.fir         | 50 +++++++++++++++++++++
 2 files changed, 77 insertions(+), 3 deletions(-)

diff --git a/flang/lib/Optimizer/Transforms/FIRToSCF.cpp b/flang/lib/Optimizer/Transforms/FIRToSCF.cpp
index d38bedcf10e87..04ba0537a6434 100644
--- a/flang/lib/Optimizer/Transforms/FIRToSCF.cpp
+++ b/flang/lib/Optimizer/Transforms/FIRToSCF.cpp
@@ -89,12 +89,36 @@ struct DoLoopConversion : public mlir::OpRewritePattern<fir::DoLoopOp> {
     mlir::Value iv = mlir::arith::MulIOp::create(
         rewriter, loc, scfLoopLikeOp.getSingleInductionVar().value(), step);
     iv = mlir::arith::AddIOp::create(rewriter, loc, low, iv);
+    mlir::Value firIV = doLoopOp.getInductionVar();
+    firIV.replaceAllUsesWith(iv);
+
+    mlir::Value finalValue;
+    if (hasFinalValue) {
+      // Prefer re-using an existing `arith.addi` in the moved loop body if it
+      // already computes the next `iv + step`.
+      if (!results.empty()) {
+        if (auto addOp = results.front().getDefiningOp<mlir::arith::AddIOp>()) {
+          mlir::Value lhs = addOp.getLhs();
+          mlir::Value rhs = addOp.getRhs();
+          if ((lhs == iv && rhs == step) || (lhs == step && rhs == iv))
+            finalValue = results.front();
+        }
+      }
+      if (!finalValue)
+        finalValue = mlir::arith::AddIOp::create(rewriter, loc, iv, step);
+    }
 
-    if (!results.empty()) {
+    if (hasFinalValue || !results.empty()) {
       rewriter.setInsertionPointToEnd(&scfLoopBody);
-      mlir::scf::YieldOp::create(rewriter, resultOp->getLoc(), results);
+      llvm::SmallVector<mlir::Value> yieldOperands;
+      if (hasFinalValue) {
+        yieldOperands.push_back(finalValue);
+        llvm::append_range(yieldOperands, results.drop_front());
+      } else {
+        llvm::append_range(yieldOperands, results);
+      }
+      mlir::scf::YieldOp::create(rewriter, resultOp->getLoc(), yieldOperands);
     }
-    doLoopOp.getInductionVar().replaceAllUsesWith(iv);
     rewriter.replaceAllUsesWith(
         doLoopOp.getRegionIterArgs(),
         hasFinalValue ? scfLoopLikeOp.getRegionIterArgs().drop_front()
diff --git a/flang/test/Fir/FirToSCF/do-loop.fir b/flang/test/Fir/FirToSCF/do-loop.fir
index 8862a4c2969e8..b61d4230b57fd 100644
--- a/flang/test/Fir/FirToSCF/do-loop.fir
+++ b/flang/test/Fir/FirToSCF/do-loop.fir
@@ -155,6 +155,56 @@ func.func @loop_with_final_value(%arg0: !fir.ref<!fir.array<100xi32>>, %arg1: !f
 
 // -----
 
+// CHECK-LABEL:   func.func @loop_with_final_value_yielding_iv() {
+// CHECK:           %[[C1:.*]] = arith.constant 1 : index
+// CHECK:           %[[C10:.*]] = arith.constant 10 : index
+// CHECK:           %[[TRIP:.*]] = arith.divsi
+// CHECK:           %[[C0:.*]] = arith.constant 0 : index
+// CHECK:           %[[ONE:.*]] = arith.constant 1 : index
+// CHECK:           %[[FOR:.*]] = scf.for %{{.*}} = %[[C0]] to %[[TRIP]] step %[[ONE]] iter_args(%{{.*}} = %[[C1]]) -> (index) {
+// CHECK:             %[[NEXT:.*]] = arith.addi %{{.*}}, %[[C1]]
+// CHECK:             scf.yield %[[NEXT]] : index
+// CHECK:           }
+// CHECK:           return
+// CHECK:         }
+func.func @loop_with_final_value_yielding_iv() {
+  %c1 = arith.constant 1 : index
+  %c10 = arith.constant 10 : index
+  %0:1 = fir.do_loop %i = %c1 to %c10 step %c1 -> index {
+    fir.result %i : index
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL:   func.func @loop_with_final_value_and_result() {
+// CHECK:           %[[C1:.*]] = arith.constant 1 : index
+// CHECK:           %[[C0I32:.*]] = arith.constant 0 : i32
+// CHECK:           %[[C10:.*]] = arith.constant 10 : index
+// CHECK:           %[[TRIP:.*]] = arith.divsi
+// CHECK:           %[[C0:.*]] = arith.constant 0 : index
+// CHECK:           %[[ONE:.*]] = arith.constant 1 : index
+// CHECK:           %[[LOOP:.*]]:2 = scf.for %{{.*}} = %[[C0]] to %[[TRIP]] step %[[ONE]] iter_args(%[[IVIN:.*]] = %[[C1]], %[[ACCIN:.*]] = %[[C0I32]]) -> (index, i32) {
+// CHECK:             %[[IVNEXT:.*]] = arith.addi %{{.*}}, %[[C1]] : index
+// CHECK:             %[[ACCOUT:.*]] = arith.addi %[[ACCIN]], %[[ACCIN]] : i32
+// CHECK:             scf.yield %[[IVNEXT]], %[[ACCOUT]] : index, i32
+// CHECK:           }
+// CHECK:           return
+// CHECK:         }
+func.func @loop_with_final_value_and_result() {
+  %c1 = arith.constant 1 : index
+  %c0_i32 = arith.constant 0 : i32
+  %c10 = arith.constant 10 : index
+  %0:2 = fir.do_loop %i = %c1 to %c10 step %c1 iter_args(%acc = %c0_i32) -> (index, i32) {
+    %acc2 = arith.addi %acc, %acc : i32
+    fir.result %i, %acc2 : index, i32
+  }
+  return
+}
+
+// -----
+
 // CHECK-LABEL:   func.func @loop_with_unordered_attr(
 // CHECK-SAME:      %[[ARG0:.*]]: !fir.ref<!fir.array<100xi32>>) {
 // CHECK:           %[[CONSTANT_0:.*]] = arith.constant 1 : index