[Flang][OpenMP] Lower schedule modifiers for worksharing loop

[kiranchandramohan]

Summary

Add support for lowering the schedule modifiers (simd, monotonic,
non-monotonic) in worksharing loops.

Note: This is part of upstreaming from the fir-dev branch of
https://github.com/flang-compiler/f18-llvm-project.

Co-authored-by: Mats Petersson <[e-mail]>
Co-authored-by: Jean Perier <[e-mail]>
Co-authored-by: Eric Schweitz <[e-mail]>
Co-authored-by: V Donaldson <[e-mail]>

[kiranchandramohan]

flang/lib/Lower/OpenMP.cpp

852

Check: Should None be represented in the IR?

[peixin]

LGTM

flang/lib/Lower/OpenMP.cpp

852

It does not matter. Both are OK. It only checks if there is simd/monotonic/nonmonotonic modifer there. https://github.com/llvm/llvm-project/blob/e6a76a49356efd11f5f36690181f0f60cecb2e01/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp#L822-L828

Also, you don't need to worry about the default behaviors if some clauses such as ordered clause exist or not. OMPIRBuilder will handle it.

[kiranchandramohan]

flang/lib/Lower/OpenMP.cpp

852

Thanks @peixin.

flang/lib/Lower/OpenMP.cpp

auto taskOp = firOpBuilder.create<mlir::omp::TaskOp>(
         /*in_reductions=*/nullptr, priorityClauseOperand, allocateOperands,
         allocatorOperands);
     createBodyOfOp(taskOp, converter, currentLocation, eval, &opClauseList);
   } else {
     TODO(converter.getCurrentLocation(), "Unhandled block directive");
   }
 }
 
+static mlir::omp::ScheduleModifier
+translateModifier(const Fortran::parser::OmpScheduleModifierType &m) {
+  switch (m.v) {
+  case Fortran::parser::OmpScheduleModifierType::ModType::Monotonic:
+    return mlir::omp::ScheduleModifier::monotonic;
+  case Fortran::parser::OmpScheduleModifierType::ModType::Nonmonotonic:
+    return mlir::omp::ScheduleModifier::nonmonotonic;
+  case Fortran::parser::OmpScheduleModifierType::ModType::Simd:
+    return mlir::omp::ScheduleModifier::simd;
+  }
+  return mlir::omp::ScheduleModifier::none;
+}
+
+static mlir::omp::ScheduleModifier
+getScheduleModifier(const Fortran::parser::OmpScheduleClause &x) {
+  const auto &modifier =
+      std::get<std::optional<Fortran::parser::OmpScheduleModifier>>(x.t);
+  // The input may have the modifier any order, so we look for one that isn't
+  // SIMD. If modifier is not set at all, fall down to the bottom and return
+  // "none".
+  if (modifier) {
+    const auto &modType1 =
+        std::get<Fortran::parser::OmpScheduleModifier::Modifier1>(modifier->t);
+    if (modType1.v.v ==
+        Fortran::parser::OmpScheduleModifierType::ModType::Simd) {
+      const auto &modType2 = std::get<
+          std::optional<Fortran::parser::OmpScheduleModifier::Modifier2>>(
+          modifier->t);
+      if (modType2 &&
+          modType2->v.v !=
+              Fortran::parser::OmpScheduleModifierType::ModType::Simd)
+        return translateModifier(modType2->v);
+
+      return mlir::omp::ScheduleModifier::none;
+    }
+
+    return translateModifier(modType1.v);
+  }
+  return mlir::omp::ScheduleModifier::none;
+}
+
+static mlir::omp::ScheduleModifier
+getSIMDModifier(const Fortran::parser::OmpScheduleClause &x) {
+  const auto &modifier =
+      std::get<std::optional<Fortran::parser::OmpScheduleModifier>>(x.t);
+  // Either of the two possible modifiers in the input can be the SIMD modifier,
+  // so look in either one, and return simd if we find one. Not found = return
+  // "none".
+  if (modifier) {
+    const auto &modType1 =
+        std::get<Fortran::parser::OmpScheduleModifier::Modifier1>(modifier->t);
+    if (modType1.v.v == Fortran::parser::OmpScheduleModifierType::ModType::Simd)
+      return mlir::omp::ScheduleModifier::simd;
+
+    const auto &modType2 = std::get<
+        std::optional<Fortran::parser::OmpScheduleModifier::Modifier2>>(
+        modifier->t);
+    if (modType2 && modType2->v.v ==
+                        Fortran::parser::OmpScheduleModifierType::ModType::Simd)
+      return mlir::omp::ScheduleModifier::simd;
+  }
+  return mlir::omp::ScheduleModifier::none;
+}
+
 static void genOMP(Fortran::lower::AbstractConverter &converter,
                    Fortran::lower::pft::Evaluation &eval,
                    const Fortran::parser::OpenMPLoopConstruct &loopConstruct) {
 
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   mlir::Location currentLocation = converter.getCurrentLocation();
   llvm::SmallVector<mlir::Value> lowerBound, upperBound, step, linearVars,
       linearStepVars, reductionVars;
if (const auto &orderedClause =
         wsLoopOp.schedule_valAttr(omp::ClauseScheduleKindAttr::get(
             context, omp::ClauseScheduleKind::Auto));
         break;
       case Fortran::parser::OmpScheduleClause::ScheduleType::Runtime:
         wsLoopOp.schedule_valAttr(omp::ClauseScheduleKindAttr::get(
             context, omp::ClauseScheduleKind::Runtime));
         break;
       }
+      mlir::omp::ScheduleModifier scheduleModifier =
+          getScheduleModifier(scheduleClause->v);
+      if (scheduleModifier != mlir::omp::ScheduleModifier::none)
+        wsLoopOp.schedule_modifierAttr(
+            omp::ScheduleModifierAttr::get(context, scheduleModifier));
+      if (getSIMDModifier(scheduleClause->v) !=
+          mlir::omp::ScheduleModifier::none)
+        wsLoopOp.simd_modifierAttr(firOpBuilder.getUnitAttr());
     }
   }
   // In FORTRAN `nowait` clause occur at the end of `omp do` directive.
   // i.e
   // !$omp do
   // <...>
   // !$omp end do nowait
   if (const auto &endClauseList =

flang/test/Lower/OpenMP/omp-wsloop-monotonic.f90

+! This test checks lowering of OpenMP DO Directive (Worksharing) with
+! monotonic schedule modifier.
+
+! RUN: bbc -fopenmp -emit-fir %s -o - | FileCheck %s
+
+program wsloop_dynamic
+  integer :: i
+!CHECK-LABEL: func @_QQmain()
+
+!$OMP PARALLEL
+!CHECK:  omp.parallel {
+
+!$OMP DO SCHEDULE(monotonic:dynamic)
+!CHECK:     %[[ALLOCA_IV:.*]] = fir.alloca i32 {{{.*}}, pinned}
+!CHECK:     %[[WS_LB:.*]] = arith.constant 1 : i32
+!CHECK:     %[[WS_UB:.*]] = arith.constant 9 : i32
+!CHECK:     %[[WS_STEP:.*]] = arith.constant 1 : i32
+!CHECK:     omp.wsloop schedule(dynamic, monotonic) nowait for (%[[I:.*]]) : i32 = (%[[WS_LB]]) to (%[[WS_UB]]) inclusive step (%[[WS_STEP]])
+!CHECK:       fir.store %[[I]] to %[[ALLOCA_IV:.*]] : !fir.ref<i32>
+
+  do i=1, 9
+    print*, i
+!CHECK:    %[[RTBEGIN:.*]] = fir.call @_FortranAioBeginExternalListOutput
+!CHECK:    %[[LOAD:.*]] = fir.load %[[ALLOCA_IV]] : !fir.ref<i32>
+!CHECK:    fir.call @_FortranAioOutputInteger32(%[[RTBEGIN]], %[[LOAD]]) : (!fir.ref<i8>, i32) -> i1
+!CHECK:    fir.call @_FortranAioEndIoStatement(%[[RTBEGIN]]) : (!fir.ref<i8>) -> i32
+  end do
+!CHECK:       omp.yield
+!CHECK:       omp.terminator
+!CHECK:     }
+
+!$OMP END DO NOWAIT
+!$OMP END PARALLEL
+end

flang/test/Lower/OpenMP/omp-wsloop-nonmonotonic.f90

+! This test checks lowering of OpenMP DO Directive(Worksharing) with
+! non-monotonic schedule modifier.
+
+! RUN: bbc -fopenmp -emit-fir %s -o - | FileCheck %s
+
+program wsloop_dynamic
+  integer :: i
+!CHECK-LABEL: func @_QQmain()
+
+
+!$OMP PARALLEL
+!CHECK:  omp.parallel {
+
+!$OMP DO SCHEDULE(nonmonotonic:dynamic)
+!CHECK:     %[[ALLOCA_IV:.*]] = fir.alloca i32 {{{.*}}, pinned}
+!CHECK:     %[[WS_LB:.*]] = arith.constant 1 : i32
+!CHECK:     %[[WS_UB:.*]] = arith.constant 9 : i32
+!CHECK:     %[[WS_STEP:.*]] = arith.constant 1 : i32
+!CHECK:     omp.wsloop schedule(dynamic, nonmonotonic) nowait for (%[[I:.*]]) : i32 = (%[[WS_LB]]) to (%[[WS_UB]]) inclusive step (%[[WS_STEP]])
+!CHECK:       fir.store %[[I]] to %[[ALLOCA_IV]] : !fir.ref<i32>
+
+  do i=1, 9
+    print*, i
+!CHECK:    %[[RTBEGIN:.*]] = fir.call @_FortranAioBeginExternalListOutput
+!CHECK:    %[[LOAD:.*]] = fir.load %[[ALLOCA_IV]] : !fir.ref<i32>
+!CHECK:    fir.call @_FortranAioOutputInteger32(%[[RTBEGIN]], %[[LOAD]]) : (!fir.ref<i8>, i32) -> i1
+!CHECK:    fir.call @_FortranAioEndIoStatement(%[[RTBEGIN]]) : (!fir.ref<i8>) -> i32
+  end do
+!CHECK:       omp.yield
+!CHECK:         }
+!CHECK:       omp.terminator
+!CHECK:     }
+
+!$OMP END DO NOWAIT
+!$OMP END PARALLEL
+end

flang/test/Lower/OpenMP/omp-wsloop-simd.f90

+! This test checks lowering of OpenMP DO Directive(Worksharing) with
+! simd schedule modifier.
+
+! RUN: bbc -fopenmp -emit-fir %s -o - | FileCheck %s
+
+program wsloop_dynamic
+  integer :: i
+!CHECK-LABEL: func @_QQmain()
+
+!$OMP PARALLEL
+!CHECK:  omp.parallel {
+
+!$OMP DO SCHEDULE(simd: runtime)
+!CHECK:     %[[WS_LB:.*]] = arith.constant 1 : i32
+!CHECK:     %[[WS_UB:.*]] = arith.constant 9 : i32
+!CHECK:     %[[WS_STEP:.*]] = arith.constant 1 : i32
+!CHECK:     omp.wsloop schedule(runtime, simd) nowait for (%[[I:.*]]) : i32 = (%[[WS_LB]]) to (%[[WS_UB]]) inclusive step (%[[WS_STEP]])
+!CHECK:       fir.store %[[I]] to %[[STORE:.*]] : !fir.ref<i32>
+
+  do i=1, 9
+    print*, i
+!CHECK:    %[[RTBEGIN:.*]] = fir.call @_FortranAioBeginExternalListOutput
+!CHECK:    %[[LOAD:.*]] = fir.load %[[STORE]] : !fir.ref<i32>
+!CHECK:    fir.call @_FortranAioOutputInteger32(%[[RTBEGIN]], %[[LOAD]]) : (!fir.ref<i8>, i32) -> i1
+!CHECK:    fir.call @_FortranAioEndIoStatement(%[[RTBEGIN]]) : (!fir.ref<i8>) -> i32
+  end do
+!CHECK:       omp.yield
+!CHECK:         }
+!CHECK:       omp.terminator
+!CHECK:     }
+
+!$OMP END DO NOWAIT
+!$OMP END PARALLEL
+end