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

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. Reviewed By: peixin Differential Revision: https://reviews.llvm.org/D127311 Co-authored-by: Mats Petersson <mats.petersson@arm.com> Co-authored-by: Jean Perier <jperier@nvidia.com> Co-authored-by: Eric Schweitz <eschweitz@nvidia.com> Co-authored-by: V Donaldson <vdonaldson@nvidia.com>
llvm · Jun 9, 2022 · 8b951e6 · 8b951e6
1 parent 5f0d4f2
commit 8b951e6
Show file tree

Hide file tree

Showing 4 changed files with 176 additions and 0 deletions.
diff --git a/flang/lib/Lower/OpenMP.cpp b/flang/lib/Lower/OpenMP.cpp
@@ -628,6 +628,70 @@ genOMP(Fortran::lower::AbstractConverter &converter,
   }
 }
 
+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) {
@@ -783,6 +847,14 @@ static void genOMP(Fortran::lower::AbstractConverter &converter,
             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.

diff --git a/flang/test/Lower/OpenMP/omp-wsloop-monotonic.f90 b/flang/test/Lower/OpenMP/omp-wsloop-monotonic.f90
@@ -0,0 +1,34 @@
+! 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
diff --git a/flang/test/Lower/OpenMP/omp-wsloop-nonmonotonic.f90 b/flang/test/Lower/OpenMP/omp-wsloop-nonmonotonic.f90
@@ -0,0 +1,36 @@
+! 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
diff --git a/flang/test/Lower/OpenMP/omp-wsloop-simd.f90 b/flang/test/Lower/OpenMP/omp-wsloop-simd.f90
@@ -0,0 +1,34 @@
+! 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