[flang] Upstream lowering of real control loops

[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 a test.

Real-valued control variables are handled the same as integer ones, the
only difference is that they need to use floating point instructions
instead of the integer equivalents.

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

[kiranchandramohan]

LG. Maybe also add the info in the summary that loops with real control are always unstructured.

[rovka]

Committed as https://github.com/llvm/llvm-project/commit/a159128811ae1d0a781dbc8590fa1e0f26642f82

flang/lib/Lower/Bridge.cpp

struct IncrementLoopInfo {
   bool isUnordered; // do concurrent, forall
   mlir::Value loopVariable = nullptr;
   mlir::Value stepValue = nullptr; // possible uses in multiple blocks
 
   // Data members for structured loops.
   fir::DoLoopOp doLoop = nullptr;
 
   // Data members for unstructured loops.
+  bool hasRealControl = false;
   mlir::Value tripVariable = nullptr;
   mlir::Block *headerBlock = nullptr; // loop entry and test block
   mlir::Block *bodyBlock = nullptr;   // first loop body block
   mlir::Block *exitBlock = nullptr;   // loop exit target block
 };
 
 /// Helper class to generate the runtime type info global data. This data
 /// is required to describe the derived type to the runtime so that it can
if ((whileCondition = std::get_if<Fortran::parser::ScalarLogicalExpr>(
                    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.hasRealControl = info.loopVariableSym.GetType()->IsNumeric(
+            Fortran::common::TypeCategory::Real);
         info.headerBlock = headerBlock;
         info.bodyBlock = bodyBlock;
         info.exitBlock = exitBlock;
       }
     } else {
       TODO(toLocation(), "infinite/unstructured loop/concurrent loop");
     }
 
void genFIRIncrementLoopBegin(IncrementLoopNestInfo &incrementLoopNestInfo) {
     auto genControlValue = [&](const Fortran::lower::SomeExpr *expr,
                                const IncrementLoopInfo &info) {
       mlir::Type controlType = info.isStructured() ? builder->getIndexType()
                                                    : info.getLoopVariableType();
       Fortran::lower::StatementContext stmtCtx;
       if (expr)
         return builder->createConvert(loc, controlType,
                                       createFIRExpr(loc, expr, stmtCtx));
+
+      if (info.hasRealControl)
+        return builder->createRealConstant(loc, controlType, 1u);
       return builder->createIntegerConstant(loc, controlType, 1); // step
     };
     for (IncrementLoopInfo &info : incrementLoopNestInfo) {
       info.loopVariable = genLoopVariableAddress(loc, info.loopVariableSym);
       mlir::Value lowerValue = genControlValue(info.lowerExpr, info);
       mlir::Value upperValue = genControlValue(info.upperExpr, info);
       info.stepValue = genControlValue(info.stepExpr, info);
 
for (IncrementLoopInfo &info : incrementLoopNestInfo) {
         builder->create<fir::StoreOp>(loc, value, info.loopVariable);
         // TODO: Mask expr
         // TODO: handle Locality Spec
         continue;
       }
 
       // Unstructured loop preheader - initialize tripVariable and loopVariable.
       mlir::Value tripCount;
+      if (info.hasRealControl) {
+        auto diff1 =
+            builder->create<mlir::arith::SubFOp>(loc, upperValue, lowerValue);
+        auto diff2 =
+            builder->create<mlir::arith::AddFOp>(loc, diff1, info.stepValue);
+        tripCount =
+            builder->create<mlir::arith::DivFOp>(loc, diff2, info.stepValue);
+        tripCount =
+            builder->createConvert(loc, builder->getIndexType(), tripCount);
+
+      } else {
         auto diff1 =
             builder->create<mlir::arith::SubIOp>(loc, upperValue, lowerValue);
         auto diff2 =
             builder->create<mlir::arith::AddIOp>(loc, diff1, info.stepValue);
         tripCount =
             builder->create<mlir::arith::DivSIOp>(loc, diff2, info.stepValue);
+      }
       info.tripVariable = builder->createTemporary(loc, tripCount.getType());
       builder->create<fir::StoreOp>(loc, tripCount, info.tripVariable);
       builder->create<fir::StoreOp>(loc, lowerValue, info.loopVariable);
 
       // Unstructured loop header - generate loop condition and mask.
       startBlock(info.headerBlock);
       tripCount = builder->create<fir::LoadOp>(loc, info.tripVariable);
       mlir::Value zero =
for (auto it = incrementLoopNestInfo.rbegin(),
       // Unstructured loop - decrement tripVariable and step loopVariable.
       mlir::Value tripCount =
           builder->create<fir::LoadOp>(loc, info.tripVariable);
       mlir::Value one =
           builder->createIntegerConstant(loc, tripCount.getType(), 1);
       tripCount = builder->create<mlir::arith::SubIOp>(loc, tripCount, one);
       builder->create<fir::StoreOp>(loc, tripCount, info.tripVariable);
       mlir::Value value = builder->create<fir::LoadOp>(loc, info.loopVariable);
-      value = builder->create<mlir::arith::AddIOp>(loc, value, info.stepValue);
+      if (info.hasRealControl)
+        value =
+            builder->create<mlir::arith::AddFOp>(loc, value, info.stepValue);
+      else
+        value =
+            builder->create<mlir::arith::AddIOp>(loc, value, info.stepValue);
       builder->create<fir::StoreOp>(loc, value, info.loopVariable);
 
       genFIRBranch(info.headerBlock);
       if (&info != &incrementLoopNestInfo.front()) // not outermost
         startBlock(info.exitBlock); // latch block of enclosing dimension
     }
   }
 

flang/test/Lower/do_loop.f90

do i=s,e,st
     ! CHECK: fir.store %[[LI_CVT]] to %[[I_REF]] : !fir.ref<i64>
     ! CHECK: %[[LI_NEXT:.*]] = arith.addi %[[LI]], %[[ST_CVT]] : index
     ! CHECK: fir.result %[[LI_NEXT]] : index
   end do
   ! CHECK: }
   ! CHECK: %[[I_RES_CVT:.*]] = fir.convert %[[I_RES]] : (index) -> i64
   ! CHECK: fir.store %[[I_RES_CVT]] to %[[I_REF]] : !fir.ref<i64>
 end subroutine
+
+! Test real loop control.
+! CHECK-LABEL: loop_with_real_control
+! CHECK-SAME: (%[[S_REF:.*]]: !fir.ref<f32> {fir.bindc_name = "s"}, %[[E_REF:.*]]: !fir.ref<f32> {fir.bindc_name = "e"}, %[[ST_REF:.*]]: !fir.ref<f32> {fir.bindc_name = "st"}) {
+subroutine loop_with_real_control(s,e,st)
+  ! CHECK-DAG: %[[INDEX_REF:.*]] = fir.alloca index
+  ! CHECK-DAG: %[[X_REF:.*]] = fir.alloca f32 {bindc_name = "x", uniq_name = "_QFloop_with_real_controlEx"}
+  ! CHECK-DAG: %[[S:.*]] = fir.load %[[S_REF]] : !fir.ref<f32>
+  ! CHECK-DAG: %[[E:.*]] = fir.load %[[E_REF]] : !fir.ref<f32>
+  ! CHECK-DAG: %[[ST:.*]] = fir.load %[[ST_REF]] : !fir.ref<f32>
+  real :: x, s, e, st
+
+  ! CHECK: %[[DIFF:.*]] = arith.subf %[[E]], %[[S]] : f32
+  ! CHECK: %[[RANGE:.*]] = arith.addf %[[DIFF]], %[[ST]] : f32
+  ! CHECK: %[[HIGH:.*]] = arith.divf %[[RANGE]], %[[ST]] : f32
+  ! CHECK: %[[HIGH_INDEX:.*]] = fir.convert %[[HIGH]] : (f32) -> index
+  ! CHECK: fir.store %[[HIGH_INDEX]] to %[[INDEX_REF]] : !fir.ref<index>
+  ! CHECK: fir.store %[[S]] to %[[X_REF]] : !fir.ref<f32>
+
+  ! CHECK: br ^[[HDR:.*]]
+  ! CHECK: ^[[HDR]]:  // 2 preds: ^{{.*}}, ^[[EXIT:.*]]
+  ! CHECK-DAG: %[[INDEX:.*]] = fir.load %[[INDEX_REF]] : !fir.ref<index>
+  ! CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+  ! CHECK: %[[COND:.*]] = arith.cmpi sgt, %[[INDEX]], %[[C0]] : index
+  ! CHECK: cond_br %[[COND]], ^[[BODY:.*]], ^[[EXIT:.*]]
+  do x=s,e,st
+    ! CHECK: ^[[BODY]]:  // pred: ^[[HDR]]
+    ! CHECK-DAG: %[[INDEX2:.*]] = fir.load %[[INDEX_REF]] : !fir.ref<index>
+    ! CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+    ! CHECK: %[[INC:.*]] = arith.subi %[[INDEX2]], %[[C1]] : index
+    ! CHECK: fir.store %[[INC]] to %[[INDEX_REF]] : !fir.ref<index>
+    ! CHECK: %[[X2:.*]] = fir.load %[[X_REF]] : !fir.ref<f32>
+    ! CHECK: %[[XINC:.*]] = arith.addf %[[X2]], %[[ST]] : f32
+    ! CHECK: fir.store %[[XINC]] to %[[X_REF]] : !fir.ref<f32>
+    ! CHECK: br ^[[HDR]]
+  end do
+end subroutine