[Flang] Upstream the lowering of the while loop

[rovka]

Summary

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.

Co-authored-by: Eric Schweitz <[e-mail]>
Co-authored-by: V Donaldson <[e-mail]>

[kiranchandramohan]

Many thanks @rovka for the patch. Looks Good. Just a couple of things.
-> The clang-format check seems to be failing.
-> Could you add some information to the commit summary detailing the code generated for the do while?

flang/lib/Lower/Bridge.cpp

auto createNextBeginBlock = [&]() {
       // blocks, created in outermost to innermost order.
       return beginBlock = beginBlock->splitBlock(beginBlock->end());
     };
     mlir::Block *headerBlock =
         unstructuredContext ? createNextBeginBlock() : nullptr;
     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;
+    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,
           bounds->step);
       if (unstructuredContext) {
         maybeStartBlock(preheaderBlock);
         info.headerBlock = headerBlock;
         info.bodyBlock = bodyBlock;
         info.exitBlock = exitBlock;
       }
     } else {
       TODO(toLocation(), "infinite/unstructured loop/concurrent loop");
     }
 
     // Increment loop begin code.  (TODO: Infinite/while code was already
     // generated.)
+    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)
+    if (whileCondition)
+      genFIRBranch(headerBlock);
+    else
       genFIRIncrementLoopEnd(incrementLoopNestInfo);
   }
 
   /// Generate FIR to begin a structured or unstructured increment loop nest.
   void genFIRIncrementLoopBegin(IncrementLoopNestInfo &incrementLoopNestInfo) {
     assert(!incrementLoopNestInfo.empty() && "empty loop nest");
     mlir::Location loc = toLocation();
     auto genControlValue = [&](const Fortran::lower::SomeExpr *expr,
                                const IncrementLoopInfo &info) {

flang/test/Lower/mixed_loops.f90

+! 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

+! 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