[Flang][OpenMP] Implementation of lowering of SIMD construct.

[arnamoy10]

Summary

This patch adds code so that using bbc we are able to see an end-to-end lowering of simd construct in action.

[arnamoy10]

Fixing the test case

[kiranchandramohan]

Looks OK. Have a few comments.

flang/lib/Lower/OpenMP.cpp

326

Nit: Can this be added to the previous if using ||?

787

Can you add Todos for the clauses or add tests for supported clauses?

flang/test/Lower/OpenMP/simd.f90

4

We are aligning with the rest of the FIR lowering tests and only testing the lowering to FIR+OpenMP dialects here. Please change this test to do just that and remove the LLVM IR check.

7

Nit: align this with rest of the code.

21

Nit: no new line.

[kiranchandramohan]

flang/test/Lower/OpenMP/simd.f90

4

Also add a FIR+OpenMP -> LLVM+OpenMP dialect test in https://github.com/llvm/llvm-project/blob/main/flang/test/Fir/convert-to-llvm-openmp-and-fir.fir

[shraiysh]

flang/lib/Lower/OpenMP.cpp

787

If D125302 gets merged before this, please add tests for collapse too.

[peixin]

flang/lib/Lower/OpenMP.cpp

714

It's time to refactor this function. It's better not to share all the code of lowering clauses for different LoopDirective since they does have different clauses.

Considering the following:

if (LoopDirective == llvm::omp::OMPD_do) {
  genWsLoop(...);
} else if (LoopDirective == llvm::omp::OMPD_simd) {
  genSIMDLoop(...);
} else {
  TODO
}

Or a switch case statement?

The clauses operands collection should be extracted into some functions so that they can be shared.

787

@shraiysh @kiranchandramohan Please notice that no clause is supported for SIMD in MLIR for now.

It's better not to refactor the code for wsloop before the upstream of it is done. So it seems to be better give TODOs when getting any clause in SIMD construct.

When the upstream for wsloop is done, it's better to refactor this code as I suggested above.

[arnamoy10]

Addressing reviewers comments. No test case for clauses could be added as we currently support simd without any clauses

[arnamoy10]

Addressing reviewers comments. No test case for clauses could be added as we currently support simd without any clauses

flang/lib/Lower/OpenMP.cpp

714

Not refactoring in this patch, waiting until wsloop comes, as suggested below

787

Did, thank you

[shraiysh]

LGTM. Minor changes requested. Thank you!

flang/test/Fir/convert-to-llvm-openmp-and-fir.fir

74

Split the testcase by adding

// -----

117

nit: please add a new line

flang/test/Lower/OpenMP/simd.f90

8

nit: 2-space indentation for the subroutine body.

[arnamoy10]

Rebasing and fixing minor stuff based on comments.

[kiranchandramohan]

It seems you missed the portion of the original test (Fortran -> OpenMP + FIR).

[arnamoy10]

Adding testcase, thanks @kiranchandramohan

[peixin]

LGTM

[kiranchandramohan]

LG.

Not for this patch: Maybe we should also adopt the base flang style of hard TODOs for unimplemented features/clauses.

[arnamoy10]

Rebase before land

[arnamoy10]

Fixing fortran->FIR testcase after rebasing.

flang/lib/Lower/OpenMP.cpp

createBodyOfOp(Op &op, Fortran::lower::AbstractConverter &converter,
   firOpBuilder.setInsertionPointToEnd(&block);
 
   // If it is an unstructured region and is not the outer region of a combined
   // construct, create empty blocks for all evaluations.
   if (eval.lowerAsUnstructured() && !outerCombined)
     createEmptyRegionBlocks(firOpBuilder, eval.getNestedEvaluations());
 
   // Insert the terminator.
-  if constexpr (std::is_same_v<Op, omp::WsLoopOp>) {
+  if constexpr (std::is_same_v<Op, omp::WsLoopOp> ||
+                std::is_same_v<Op, omp::SimdLoopOp>) {
     mlir::ValueRange results;
     firOpBuilder.create<mlir::omp::YieldOp>(loc, results);
   } else {
     firOpBuilder.create<mlir::omp::TerminatorOp>(loc);
   }
 
   // Reset the insert point to before the terminator.
   if (storeOp)
     firOpBuilder.setInsertionPointAfter(storeOp);
   else
     firOpBuilder.setInsertionPointToStart(&block);
static void genOMP(Fortran::lower::AbstractConverter &converter,
 
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   mlir::Location currentLocation = converter.getCurrentLocation();
   llvm::SmallVector<mlir::Value> lowerBound, upperBound, step, linearVars,
       linearStepVars, reductionVars;
   mlir::Value scheduleChunkClauseOperand;
   mlir::Attribute scheduleClauseOperand, collapseClauseOperand,
       noWaitClauseOperand, orderedClauseOperand, orderClauseOperand;
-  const auto &wsLoopOpClauseList = std::get<Fortran::parser::OmpClauseList>(
+  const auto &loopOpClauseList = std::get<Fortran::parser::OmpClauseList>(
       std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t).t);
 
   const auto ompDirective =
       std::get<Fortran::parser::OmpLoopDirective>(
           std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t).t)
           .v;
   if (llvm::omp::OMPD_parallel_do == ompDirective) {
     createCombinedParallelOp<Fortran::parser::OmpBeginLoopDirective>(
         converter, eval,
         std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t));
-  } else if (llvm::omp::OMPD_do != ompDirective) {
+  } else if (llvm::omp::OMPD_do != ompDirective &&
+             llvm::omp::OMPD_simd != ompDirective) {
     TODO(converter.getCurrentLocation(), "Construct enclosing do loop");
   }
 
   // Collect the loops to collapse.
   auto *doConstructEval = &eval.getFirstNestedEvaluation();
 
   std::int64_t collapseValue =
-      Fortran::lower::getCollapseValue(wsLoopOpClauseList);
+      Fortran::lower::getCollapseValue(loopOpClauseList);
   std::size_t loopVarTypeSize = 0;
   SmallVector<const Fortran::semantics::Symbol *> iv;
   do {
     auto *doLoop = &doConstructEval->getFirstNestedEvaluation();
     auto *doStmt = doLoop->getIf<Fortran::parser::NonLabelDoStmt>();
     assert(doStmt && "Expected do loop to be in the nested evaluation");
     const auto &loopControl =
         std::get<std::optional<Fortran::parser::LoopControl>>(doStmt->t);
do {
     loopVarTypeSize = std::max(loopVarTypeSize,
                                bounds->name.thing.symbol->GetUltimate().size());
 
     collapseValue--;
     doConstructEval =
         &*std::next(doConstructEval->getNestedEvaluations().begin());
   } while (collapseValue > 0);
 
-  for (const auto &clause : wsLoopOpClauseList.v) {
+  for (const auto &clause : loopOpClauseList.v) {
     if (const auto &scheduleClause =
             std::get_if<Fortran::parser::OmpClause::Schedule>(&clause.u)) {
       if (const auto &chunkExpr =
               std::get<std::optional<Fortran::parser::ScalarIntExpr>>(
                   scheduleClause->v.t)) {
         if (const auto *expr = Fortran::semantics::GetExpr(*chunkExpr)) {
           Fortran::lower::StatementContext stmtCtx;
           scheduleChunkClauseOperand =
for (unsigned it = 0; it < (unsigned)lowerBound.size(); it++) {
     lowerBound[it] = firOpBuilder.createConvert(currentLocation, loopVarType,
                                                 lowerBound[it]);
     upperBound[it] = firOpBuilder.createConvert(currentLocation, loopVarType,
                                                 upperBound[it]);
     step[it] =
         firOpBuilder.createConvert(currentLocation, loopVarType, step[it]);
   }
 
+  // 2.9.3.1 SIMD construct
+  // TODO: Support all the clauses
+  if (llvm::omp::OMPD_simd == ompDirective) {
+    TypeRange resultType;
+    auto SimdLoopOp = firOpBuilder.create<mlir::omp::SimdLoopOp>(
+        currentLocation, resultType, lowerBound, upperBound, step);
+    createBodyOfOp<omp::SimdLoopOp>(SimdLoopOp, converter, currentLocation,
+                                    eval, &loopOpClauseList, iv);
+    return;
+  }
+
   // FIXME: Add support for following clauses:
   // 1. linear
   // 2. order
   auto wsLoopOp = firOpBuilder.create<mlir::omp::WsLoopOp>(
       currentLocation, lowerBound, upperBound, step, linearVars, linearStepVars,
       reductionVars, /*reductions=*/nullptr,
       scheduleClauseOperand.dyn_cast_or_null<omp::ClauseScheduleKindAttr>(),
       scheduleChunkClauseOperand, /*schedule_modifiers=*/nullptr,
       /*simd_modifier=*/nullptr,
       collapseClauseOperand.dyn_cast_or_null<IntegerAttr>(),
       noWaitClauseOperand.dyn_cast_or_null<UnitAttr>(),
       orderedClauseOperand.dyn_cast_or_null<IntegerAttr>(),
       orderClauseOperand.dyn_cast_or_null<omp::ClauseOrderKindAttr>(),
       /*inclusive=*/firOpBuilder.getUnitAttr());
 
   // Handle attribute based clauses.
-  for (const Fortran::parser::OmpClause &clause : wsLoopOpClauseList.v) {
+  for (const Fortran::parser::OmpClause &clause : loopOpClauseList.v) {
     if (const auto &orderedClause =
             std::get_if<Fortran::parser::OmpClause::Ordered>(&clause.u)) {
       if (orderedClause->v.has_value()) {
         const auto *expr = Fortran::semantics::GetExpr(orderedClause->v);
         const std::optional<std::int64_t> orderedClauseValue =
             Fortran::evaluate::ToInt64(*expr);
         wsLoopOp.ordered_valAttr(
             firOpBuilder.getI64IntegerAttr(*orderedClauseValue));
if (const auto &endClauseList =
     const auto &clauseList =
         std::get<Fortran::parser::OmpClauseList>((*endClauseList).t);
     for (const Fortran::parser::OmpClause &clause : clauseList.v)
       if (std::get_if<Fortran::parser::OmpClause::Nowait>(&clause.u))
         wsLoopOp.nowaitAttr(firOpBuilder.getUnitAttr());
   }
 
   createBodyOfOp<omp::WsLoopOp>(wsLoopOp, converter, currentLocation, eval,
-                                &wsLoopOpClauseList, iv);
+                                &loopOpClauseList, iv);
 }
 
 static void
 genOMP(Fortran::lower::AbstractConverter &converter,
        Fortran::lower::pft::Evaluation &eval,
        const Fortran::parser::OpenMPCriticalConstruct &criticalConstruct) {
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   mlir::Location currentLocation = converter.getCurrentLocation();

flang/test/Fir/convert-to-llvm-openmp-and-fir.fir

 // CHECK:     %[[N:.*]] = llvm.load %[[N_REF]] : !llvm.ptr<i32>
 // CHECK:     llvm.store %[[N]], %[[X_REF]] : !llvm.ptr<i32>
 // CHECK:     omp.terminator
 // CHECK:   }
 // CHECK:   omp.terminator
 // CHECK: }
 // CHECK: llvm.return
 // CHECK: }
 
 
 // -----
 
 func.func @_QPsb(%arr: !fir.box<!fir.array<?xi32>> {fir.bindc_name = "arr"}) {
   %0 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFsbEi"}
   omp.parallel   {
     %c1 = arith.constant 1 : i32
     %c50 = arith.constant 50 : i32

 
 func.func private @foo()
 func.func private @bar()
 
 func.func @sections_no_data() {
   omp.sections {
     omp.section {
       fir.call @foo() : () -> ()
       omp.terminator
     }
     omp.section {
       fir.call @bar() : () -> ()
       omp.terminator
     }
     omp.terminator
   }
   return

 // CHECK:     omp.terminator
 // CHECK:   }
 // CHECK:   omp.section {
 // CHECK:     llvm.call @bar(%[[ARG0]], %[[ARG1]]) : (!llvm.ptr<i32>, !llvm.ptr<i32>) -> ()
 // CHECK:     omp.terminator
 // CHECK:   }
 // CHECK:   omp.terminator
 // CHECK: }
+
+// -----
+
+func.func @_QPsimd1(%arg0: !fir.ref<i32> {fir.bindc_name = "n"}, %arg1: !fir.ref<!fir.array<?xi32>> {fir.bindc_name = "arr"}) {
+  %c1_i64 = arith.constant 1 : i64
+  %c1_i32 = arith.constant 1 : i32
+  %0 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFsbEi"}
+  omp.parallel  {
+    %1 = fir.alloca i32 {adapt.valuebyref, pinned}
+    %2 = fir.load %arg0 : !fir.ref<i32>
+    omp.simdloop (%arg2) : i32 = (%c1_i32) to (%2) step (%c1_i32)  {
+      fir.store %arg2 to %1 : !fir.ref<i32>
+      %3 = fir.load %1 : !fir.ref<i32>
+      %4 = fir.convert %3 : (i32) -> i64
+      %5 = arith.subi %4, %c1_i64 : i64
+      %6 = fir.coordinate_of %arg1, %5 : (!fir.ref<!fir.array<?xi32>>, i64) -> !fir.ref<i32>
+      fir.store %3 to %6 : !fir.ref<i32>
+      omp.yield
+    }
+    omp.terminator
+  }
+  return
+}
+
+// CHECK-LABEL:  _QPsimd1
+// CHECK-SAME: %[[N_REF:.*]]: !llvm.ptr<i32> {fir.bindc_name = "n"}, %[[ARR_REF:.*]]: !llvm.ptr<i32> {fir.bindc_name = "arr"}) {
+// CHECK:    %[[ONE_1:.*]] = llvm.mlir.constant(1 : i64) : i64
+// CHECK:    %[[ONE_2:.*]] = llvm.mlir.constant(1 : i32) : i32
+// CHECK: omp.parallel   {
+// CHECK:      %[[ONE_3:.*]] = llvm.mlir.constant(1 : i64) : i64
+// CHECK:      %[[I_VAR:.*]] = llvm.alloca %[[ONE_3]] x i32 {adapt.valuebyref, in_type = i32, operand_segment_sizes = dense<0> : vector<2xi32>, pinned} : (i64) -> !llvm.ptr<i32>
+// CHECK:      %[[N:.*]] = llvm.load %[[N_REF]] : !llvm.ptr<i32>
+// CHECK: omp.simdloop
+// CHECK-SAME: (%[[I:.*]]) : i32 = (%[[ONE_2]]) to (%[[N]]) step (%[[ONE_2]]) {
+// CHECK:   llvm.store %[[I]], %[[I_VAR]] : !llvm.ptr<i32>
+// CHECK:   %[[I1:.*]] = llvm.load %[[I_VAR]] : !llvm.ptr<i32>
+// CHECK:   %[[I1_EXT:.*]] = llvm.sext %[[I1]] : i32 to i64
+// CHECK:   %[[I_CSTYLE:.*]] = llvm.sub %[[I1_EXT]], %[[ONE_1]]  : i64
+// CHECK:   %[[ARR_I_REF:.*]] = llvm.getelementptr %[[ARR_REF]][%[[I_CSTYLE]]] : (!llvm.ptr<i32>, i64) -> !llvm.ptr<i32>
+// CHECK:   llvm.store %[[I1]], %[[ARR_I_REF]] : !llvm.ptr<i32>
+// CHECK: omp.yield
+// CHECK: }
+// CHECK: omp.terminator
+// CHECK: }
+// CHECK: llvm.return
+// CHECK: }

flang/test/Lower/OpenMP/simd.f90

+! Tests for 2.9.3.1 Simd
+
+! RUN: bbc -fopenmp -emit-fir %s -o - | FileCheck %s
+
+!CHECK-LABEL: func @_QPsimdloop()
+subroutine simdloop
+integer :: i
+  !$OMP SIMD
+  ! CHECK: %[[LB:.*]] = arith.constant 1 : i32
+  ! CHECK-NEXT: %[[UB:.*]] = arith.constant 9 : i32
+  ! CHECK-NEXT: %[[STEP:.*]] = arith.constant 1 : i32
+  ! CHECK-NEXT: omp.simdloop (%[[I:.*]]) : i32 = (%[[LB]]) to (%[[UB]]) step (%[[STEP]]) { 
+  do i=1, 9
+    ! CHECK: fir.store %[[I]] to %[[LOCAL:.*]] : !fir.ref<i32>
+    ! CHECK: %[[LD:.*]] = fir.load %[[LOCAL]] : !fir.ref<i32>
+    ! CHECK: fir.call @_FortranAioOutputInteger32({{.*}}, %[[LD]]) : (!fir.ref<i8>, i32) -> i1
+    print*, i
+  end do
+  !$OMP END SIMD 
+end subroutine