[flang] Upstream the lowering of the while loop

Upstream the code for handling while loops from the fir-dev branch at
https://github.com/flang-compiler/f18-llvm-project/tree/fir-dev/

Also add tests.

The while loop is lowered to a header block that checks the loop
condition and branches either to the exit block or to the body of the
loop. The body of the loop will unconditionally branch back to the
header.

Differential Revision: https://reviews.llvm.org/D126636

Co-authored-by: Eric Schweitz <eschweitz@nvidia.com>
Co-authored-by: V Donaldson <vdonaldson@nvidia.com>

flang/lib/Lower/Bridge.cpp

@@ -981,8 +981,16 @@ private:
     mlir::Block *bodyBlock = doStmtEval.lexicalSuccessor->block;
     mlir::Block *exitBlock = doStmtEval.parentConstruct->constructExit->block;
     IncrementLoopNestInfo incrementLoopNestInfo;
-    if (const auto *bounds = std::get_if<Fortran::parser::LoopControl::Bounds>(
+    const Fortran::parser::ScalarLogicalExpr *whileCondition = nullptr;
-            &loopControl->u)) {
+    if ((whileCondition = std::get_if<Fortran::parser::ScalarLogicalExpr>(
+             &loopControl->u))) {
+      assert(unstructuredContext && "while loop must be unstructured");
+      maybeStartBlock(preheaderBlock); // no block or empty block
+      startBlock(headerBlock);
+      genFIRConditionalBranch(*whileCondition, bodyBlock, exitBlock);
+    } else if (const auto *bounds =
+                   std::get_if<Fortran::parser::LoopControl::Bounds>(
+                       &loopControl->u)) {
       // Non-concurrent increment loop.
       IncrementLoopInfo &info = incrementLoopNestInfo.emplace_back(
           *bounds->name.thing.symbol, bounds->lower, bounds->upper,
@@ -999,15 +1007,19 @@ private:
 
     // Increment loop begin code.  (TODO: Infinite/while code was already
     // generated.)
-    genFIRIncrementLoopBegin(incrementLoopNestInfo);
+    if (!whileCondition)
+      genFIRIncrementLoopBegin(incrementLoopNestInfo);
 
     // Loop body code - NonLabelDoStmt and EndDoStmt code is generated here.
     // Their genFIR calls are nops except for block management in some cases.
     for (Fortran::lower::pft::Evaluation &e : eval.getNestedEvaluations())
       genFIR(e, unstructuredContext);
 
-    // Loop end code. (TODO: infinite/while loop)
+    // Loop end code. (TODO: infinite loop)
-    genFIRIncrementLoopEnd(incrementLoopNestInfo);
+    if (whileCondition)
+      genFIRBranch(headerBlock);
+    else
+      genFIRIncrementLoopEnd(incrementLoopNestInfo);
   }
 
   /// Generate FIR to begin a structured or unstructured increment loop nest.

flang/test/Lower/mixed_loops.f90

@@ -0,0 +1,118 @@
+! RUN: bbc -emit-fir -o - %s | FileCheck %s
+! RUN: %flang_fc1 -emit-fir -o - %s | FileCheck %s
+
+! Test while loop inside do loop.
+! CHECK-LABEL: while_inside_do_loop
+subroutine while_inside_do_loop
+  ! CHECK-DAG: %[[T_REF:.*]] = fir.alloca i32
+  ! CHECK-DAG: %[[I_REF:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFwhile_inside_do_loopEi"}
+  ! CHECK-DAG: %[[J_REF:.*]] = fir.alloca i32 {bindc_name = "j", uniq_name = "_QFwhile_inside_do_loopEj"}
+  integer :: i, j
+
+  ! CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32
+  ! CHECK-DAG: %[[C8:.*]] = arith.constant 8 : i32
+  ! CHECK-DAG: %[[C13:.*]] = arith.constant 13 : i32
+  ! CHECK: %[[DIFF:.*]] = arith.subi %[[C13]], %[[C8]] : i32
+  ! CHECK: %[[RANGE:.*]] = arith.addi %[[DIFF]], %[[C1]] : i32
+  ! CHECK: %[[HIGH:.*]] = arith.divsi %[[RANGE]], %[[C1]] : i32
+  ! CHECK: fir.store %[[HIGH]] to %[[T_REF]] : !fir.ref<i32>
+  ! CHECK: fir.store %[[C8]] to %[[I_REF]] : !fir.ref<i32>
+
+  ! CHECK: br ^[[HDR1:.*]]
+  ! CHECK: ^[[HDR1]]:  // 2 preds: ^{{.*}}, ^[[EXIT2:.*]]
+  ! CHECK-DAG: %[[T:.*]] = fir.load %[[T_REF]] : !fir.ref<i32>
+  ! CHECK-DAG: %[[C0:.*]] = arith.constant 0 : i32
+  ! CHECK: %[[COND:.*]] = arith.cmpi sgt, %[[T]], %[[C0]] : i32
+  ! CHECK: cond_br %[[COND]], ^[[BODY1:.*]], ^[[EXIT1:.*]]
+  do i=8,13
+    ! CHECK: ^[[BODY1]]:  // pred: ^[[HDR1]]
+    ! CHECK: %[[C3:.*]] = arith.constant 3 : i32
+    ! CHECK: fir.store %[[C3]] to %[[J_REF]] : !fir.ref<i32>
+    j=3
+
+    ! CHECK: br ^[[HDR2:.*]]
+    ! CHECK: ^[[HDR2]]:  // 2 preds: ^[[BODY1]], ^[[BODY2:.*]]
+    ! CHECK-DAG: %[[J:.*]] = fir.load %[[J_REF]] : !fir.ref<i32>
+    ! CHECK-DAG: %[[I:.*]] = fir.load %[[I_REF]] : !fir.ref<i32>
+    ! CHECK: %[[COND2:.*]] = arith.cmpi slt, %[[J]], %[[I]] : i32
+    ! CHECK: cond_br %[[COND2]], ^[[BODY2]], ^[[EXIT2]]
+    do while (j .lt. i)
+      ! CHECK: ^[[BODY2]]:  // pred: ^[[HDR2]]
+      ! CHECK-DAG: %[[J2:.*]] = fir.load %[[J_REF]] : !fir.ref<i32>
+      ! CHECK-DAG: %[[C2:.*]] = arith.constant 2 : i32
+      ! CHECK: %[[INC2:.*]] = arith.muli %[[C2]], %[[J2]] : i32
+      ! CHECK: fir.store %[[INC2]] to %[[J_REF]] : !fir.ref<i32>
+      j=j*2
+    ! CHECK: br ^[[HDR2]]
+    end do
+
+  ! CHECK: ^[[EXIT2]]: // pred: ^[[HDR2]]
+  ! CHECK-DAG: %[[T2:.*]] = fir.load %[[T_REF]] : !fir.ref<i32>
+  ! CHECK-DAG: %[[C1_AGAIN:.*]] = arith.constant 1 : i32
+  ! CHECK: %[[TDEC:.*]] = arith.subi %[[T2]], %[[C1_AGAIN]] : i32
+  ! CHECK: fir.store %[[TDEC]] to %[[T_REF]]
+  ! CHECK: %[[I3:.*]] = fir.load %[[I_REF]] : !fir.ref<i32>
+  ! CHECK: %[[IINC:.*]] = arith.addi %[[I3]], %[[C1]] : i32
+  ! CHECK: fir.store %[[IINC]] to %[[I_REF]] : !fir.ref<i32>
+  ! CHECK: br ^[[HDR1]]
+  end do
+
+  ! CHECK: ^[[EXIT1]]:  // pred: ^[[HDR1]]
+  ! CHECK: %[[IPRINT:.*]] = fir.load %[[I_REF]] : !fir.ref<i32>
+  ! CHECK: fir.call @_FortranAioOutputInteger32(%{{.*}}, %[[IPRINT]]) : (!fir.ref<i8>, i32) -> i1
+  ! CHECK: %[[JPRINT:.*]] = fir.load %[[J_REF]] : !fir.ref<i32>
+  ! CHECK: fir.call @_FortranAioOutputInteger32(%{{.*}}, %[[JPRINT]]) : (!fir.ref<i8>, i32) -> i1
+  print *, i, j
+end subroutine
+
+! Test do loop inside while loop.
+! CHECK-LABEL: do_inside_while_loop
+subroutine do_inside_while_loop
+  ! CHECK-DAG: %[[I_REF:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFdo_inside_while_loopEi"}
+  ! CHECK-DAG: %[[J_REF:.*]] = fir.alloca i32 {bindc_name = "j", uniq_name = "_QFdo_inside_while_loopEj"}
+  integer :: i, j
+
+    ! CHECK: %[[C3:.*]] = arith.constant 3 : i32
+    ! CHECK: fir.store %[[C3]] to %[[J_REF]] : !fir.ref<i32>
+    j=3
+
+    ! CHECK: br ^[[HDR1:.*]]
+    ! CHECK: ^[[HDR1]]:  // 2 preds: ^{{.*}}, ^[[BODY1:.*]]
+    ! CHECK-DAG: %[[J:.*]] = fir.load %[[J_REF]] : !fir.ref<i32>
+    ! CHECK-DAG: %[[UL:.*]] = arith.constant 21 : i32
+    ! CHECK: %[[COND:.*]] = arith.cmpi slt, %[[J]], %[[UL]] : i32
+    ! CHECK: cond_br %[[COND]], ^[[BODY1]], ^[[EXIT1:.*]]
+    do while (j .lt. 21)
+      ! CHECK: ^[[BODY1]]:  // pred: ^[[HDR1]]
+
+      ! CHECK-DAG: %[[C8_I32:.*]] = arith.constant 8 : i32
+      ! CHECK-DAG: %[[C8:.*]] = fir.convert %[[C8_I32]] : (i32) -> index
+      ! CHECK-DAG: %[[C13_I32:.*]] = arith.constant 13 : i32
+      ! CHECK-DAG: %[[C13:.*]] = fir.convert %[[C13_I32]] : (i32) -> index
+      ! CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+      ! CHECK: %[[RESULT:.*]] = fir.do_loop %[[IDX:.*]] = %[[C8]] to %[[C13]] step %[[C1]] -> index {
+        ! CHECK: %[[I32:.*]] = fir.convert %[[IDX]] : (index) -> i32
+        ! CHECK: fir.store %[[I32]] to %[[I_REF]] : !fir.ref<i32>
+        ! CHECK-DAG: %[[J2:.*]] = fir.load %[[J_REF]] : !fir.ref<i32>
+        ! CHECK-DAG: %[[C2:.*]] = arith.constant 2 : i32
+        ! CHECK: %[[JINC:.*]] = arith.muli %[[C2]], %[[J2]] : i32
+        ! CHECK: fir.store %[[JINC]] to %[[J_REF]] : !fir.ref<i32>
+        ! CHECK: %[[IINC:.*]] = arith.addi %[[IDX]], %[[C1]] : index
+        ! CHECK: fir.result %[[IINC]] : index
+      do i=8,13
+        j=j*2
+
+      ! CHECK: %[[IFINAL:.*]] = fir.convert %[[RESULT]] : (index) -> i32
+      ! CHECK: fir.store %[[IFINAL]] to %[[I_REF]] : !fir.ref<i32>
+      end do
+
+    ! CHECK: br ^[[HDR1]]
+    end do
+
+  ! CHECK: ^[[EXIT1]]:  // pred: ^[[HDR1]]
+  ! CHECK: %[[IPRINT:.*]] = fir.load %[[I_REF]] : !fir.ref<i32>
+  ! CHECK: fir.call @_FortranAioOutputInteger32(%{{.*}}, %[[IPRINT]]) : (!fir.ref<i8>, i32) -> i1
+  ! CHECK: %[[JPRINT:.*]] = fir.load %[[J_REF]] : !fir.ref<i32>
+  ! CHECK: fir.call @_FortranAioOutputInteger32(%{{.*}}, %[[JPRINT]]) : (!fir.ref<i8>, i32) -> i1
+  print *, i, j
+end subroutine

flang/test/Lower/while_loop.f90

@@ -0,0 +1,91 @@
+! RUN: bbc -emit-fir -o - %s | FileCheck %s
+! RUN: %flang_fc1 -emit-fir -o - %s | FileCheck %s
+
+! Test a simple while loop.
+! CHECK-LABEL: simple_loop
+subroutine simple_loop
+  ! CHECK: %[[I_REF:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFsimple_loopEi"}
+  integer :: i
+
+  ! CHECK: %[[C5:.*]] = arith.constant 5 : i32
+  ! CHECK: fir.store %[[C5]] to %[[I_REF]]
+  i = 5
+
+  ! CHECK: br ^[[BB1:.*]]
+  ! CHECK: ^[[BB1]]:  // 2 preds: ^{{.*}}, ^[[BB2:.*]]
+  ! CHECK-DAG: %[[I:.*]] = fir.load %[[I_REF]] : !fir.ref<i32>
+  ! CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32
+  ! CHECK: %[[COND:.*]] = arith.cmpi sgt, %[[I]], %[[C1]] : i32
+  ! CHECK: cond_br %[[COND]], ^[[BB2]], ^[[BB3:.*]]
+  ! CHECK: ^[[BB2]]:  // pred: ^[[BB1]]
+  ! CHECK-DAG: %[[I2:.*]] = fir.load %[[I_REF]] : !fir.ref<i32>
+  ! CHECK-DAG: %[[C2:.*]] = arith.constant 2 : i32
+  ! CHECK: %[[INC:.*]] = arith.subi %[[I2]], %[[C2]] : i32
+  ! CHECK: fir.store %[[INC]] to %[[I_REF]] : !fir.ref<i32>
+  ! CHECK: br ^[[BB1]]
+  do while (i .gt. 1)
+    i = i - 2
+  end do
+
+  ! CHECK: ^[[BB3]]:  // pred: ^[[BB1]]
+  ! CHECK: %[[I3:.*]] = fir.load %[[I_REF]] : !fir.ref<i32>
+  ! CHECK: fir.call @_FortranAioOutputInteger32(%{{.*}}, %[[I3]]) : (!fir.ref<i8>, i32) -> i1
+  print *, i
+end subroutine
+
+! Test 2 nested while loops.
+! CHECK-LABEL: while_inside_while_loop
+subroutine while_inside_while_loop
+  ! CHECK-DAG: %[[I_REF:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFwhile_inside_while_loopEi"}
+  ! CHECK-DAG: %[[J_REF:.*]] = fir.alloca i32 {bindc_name = "j", uniq_name = "_QFwhile_inside_while_loopEj"}
+  integer :: i, j
+
+  ! CHECK: %[[C13:.*]] = arith.constant 13 : i32
+  ! CHECK: fir.store %[[C13]] to %[[I_REF]]
+  i = 13
+
+  ! CHECK: br ^[[HDR1:.*]]
+  ! CHECK: ^[[HDR1]]:  // 2 preds: ^{{.*}}, ^[[EXIT2:.*]]
+  ! CHECK-DAG: %[[I:.*]] = fir.load %[[I_REF]] : !fir.ref<i32>
+  ! CHECK-DAG: %[[C8:.*]] = arith.constant 8 : i32
+  ! CHECK: %[[COND:.*]] = arith.cmpi sgt, %[[I]], %[[C8]] : i32
+  ! CHECK: cond_br %[[COND]], ^[[BODY1:.*]], ^[[EXIT1:.*]]
+  do while (i .gt. 8)
+    ! CHECK: ^[[BODY1]]:  // pred: ^[[HDR1]]
+    ! CHECK-DAG: %[[I2:.*]] = fir.load %[[I_REF]] : !fir.ref<i32>
+    ! CHECK-DAG: %[[C5:.*]] = arith.constant 5 : i32
+    ! CHECK: %[[INC:.*]] = arith.subi %[[I2]], %[[C5]] : i32
+    ! CHECK: fir.store %[[INC]] to %[[I_REF]] : !fir.ref<i32>
+    i = i - 5
+
+    ! CHECK: %[[C3:.*]] = arith.constant 3 : i32
+    ! CHECK: fir.store %[[C3]] to %[[J_REF]]
+    j = 3
+
+    ! CHECK: br ^[[HDR2:.*]]
+    ! CHECK: ^[[HDR2]]:  // 2 preds: ^[[BODY1]], ^[[BODY2:.*]]
+    ! CHECK-DAG: %[[J:.*]] = fir.load %[[J_REF]] : !fir.ref<i32>
+    ! CHECK-DAG: %[[I3:.*]] = fir.load %[[I_REF]] : !fir.ref<i32>
+    ! CHECK: %[[COND2:.*]] = arith.cmpi slt, %[[J]], %[[I3]] : i32
+    ! CHECK: cond_br %[[COND2]], ^[[BODY2]], ^[[EXIT2]]
+    do while (j .lt. i)
+      ! CHECK: ^[[BODY2]]:  // pred: ^[[HDR2]]
+      ! CHECK-DAG: %[[J2:.*]] = fir.load %[[J_REF]] : !fir.ref<i32>
+      ! CHECK-DAG: %[[C2:.*]] = arith.constant 2 : i32
+      ! CHECK: %[[INC2:.*]] = arith.muli %[[C2]], %[[J2]] : i32
+      ! CHECK: fir.store %[[INC2]] to %[[J_REF]] : !fir.ref<i32>
+      j = j * 2
+    ! CHECK: br ^[[HDR2]]
+    end do
+
+    ! CHECK: ^[[EXIT2]]: // pred: ^[[HDR2]]
+    ! CHECK: br ^[[HDR1]]
+  end do
+
+  ! CHECK: ^[[EXIT1]]:  // pred: ^[[HDR1]]
+  ! CHECK: %[[IPRINT:.*]] = fir.load %[[I_REF]] : !fir.ref<i32>
+  ! CHECK: fir.call @_FortranAioOutputInteger32(%{{.*}}, %[[IPRINT]]) : (!fir.ref<i8>, i32) -> i1
+  ! CHECK: %[[JPRINT:.*]] = fir.load %[[J_REF]] : !fir.ref<i32>
+  ! CHECK: fir.call @_FortranAioOutputInteger32(%{{.*}}, %[[JPRINT]]) : (!fir.ref<i8>, i32) -> i1
+  print *, i, j
+end subroutine