[mlir][OpenMP] Add support for SIMD modifier

Add support for SIMD modifier in OpenMP worksharing loops.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D111051

llvm/include/llvm/Frontend/OpenMP/OMPConstants.h

@@ -120,6 +120,10 @@ enum class OMPScheduleType {
   Runtime = 37,
   Auto = 38, // auto
 
+  StaticBalancedChunked = 45, // static with chunk adjustment (e.g., simd)
+  GuidedSimd = 46,            // guided with chunk adjustment
+  RuntimeSimd = 47,           // runtime with chunk adjustment
+
   ModifierMonotonic =
       (1 << 29), // Set if the monotonic schedule modifier was present
   ModifierNonmonotonic =

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

@@ -139,11 +139,13 @@ def TerminatorOp : OpenMP_Op<"terminator", [Terminator]> {
 def OMP_SCHEDULE_MOD_None : StrEnumAttrCase<"none", 0>;
 def OMP_SCHEDULE_MOD_Monotonic : StrEnumAttrCase<"monotonic", 1>;
 def OMP_SCHEDULE_MOD_Nonmonotonic : StrEnumAttrCase<"nonmonotonic", 2>;
+def OMP_SCHEDULE_MOD_SIMD : StrEnumAttrCase<"simd", 3>;
 
 def ScheduleModifier : StrEnumAttr<"ScheduleModifier", "OpenMP Schedule Modifier",
                                          [OMP_SCHEDULE_MOD_None,
                                           OMP_SCHEDULE_MOD_Monotonic,
-                                          OMP_SCHEDULE_MOD_Nonmonotonic]>
+                                          OMP_SCHEDULE_MOD_Nonmonotonic,
+                                          OMP_SCHEDULE_MOD_SIMD]>
 {
     let cppNamespace = "::mlir::omp";
 }
@@ -289,6 +291,7 @@ def WsLoopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments,
              OptionalAttr<ScheduleKind>:$schedule_val,
              Optional<AnyType>:$schedule_chunk_var,
              OptionalAttr<ScheduleModifier>:$schedule_modifier,
+             UnitAttr:$simd_modifier,
              Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$collapse_val,
              UnitAttr:$nowait,
              Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$ordered_val,

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

@@ -244,12 +244,51 @@ static void printLinearClause(OpAsmPrinter &p, OperandRange linearVars,
 // Parser and printer for Schedule Clause
 //===----------------------------------------------------------------------===//
 
+static ParseResult
+verifyScheduleModifiers(OpAsmParser &parser,
+                        SmallVectorImpl<SmallString<12>> &modifiers) {
+  if (modifiers.size() > 2)
+    return parser.emitError(parser.getNameLoc()) << " unexpected modifier(s)";
+  for (auto mod : modifiers) {
+    // Translate the string. If it has no value, then it was not a valid
+    // modifier!
+    auto symbol = symbolizeScheduleModifier(mod);
+    if (!symbol.hasValue())
+      return parser.emitError(parser.getNameLoc())
+             << " unknown modifier type: " << mod;
+  }
+
+  // If we have one modifier that is "simd", then stick a "none" modiifer in
+  // index 0.
+  if (modifiers.size() == 1) {
+    if (symbolizeScheduleModifier(modifiers[0]) ==
+        mlir::omp::ScheduleModifier::simd) {
+      modifiers.push_back(modifiers[0]);
+      modifiers[0] =
+          stringifyScheduleModifier(mlir::omp::ScheduleModifier::none);
+    }
+  } else if (modifiers.size() == 2) {
+    // If there are two modifier:
+    // First modifier should not be simd, second one should be simd
+    if (symbolizeScheduleModifier(modifiers[0]) ==
+            mlir::omp::ScheduleModifier::simd ||
+        symbolizeScheduleModifier(modifiers[1]) !=
+            mlir::omp::ScheduleModifier::simd)
+      return parser.emitError(parser.getNameLoc())
+             << " incorrect modifier order";
+  }
+  return success();
+}
+
 /// schedule ::= `schedule` `(` sched-list `)`
-/// sched-list ::= sched-val | sched-val sched-list
+/// sched-list ::= sched-val | sched-val sched-list |
+///                sched-val `,` sched-modifier
 /// sched-val ::= sched-with-chunk | sched-wo-chunk
 /// sched-with-chunk ::= sched-with-chunk-types (`=` ssa-id-and-type)?
 /// sched-with-chunk-types ::= `static` | `dynamic` | `guided`
 /// sched-wo-chunk ::=  `auto` | `runtime`
+/// sched-modifier ::=  sched-mod-val | sched-mod-val `,` sched-mod-val
+/// sched-mod-val ::=  `monotonic` | `nonmonotonic` | `simd` | `none`
 static ParseResult
 parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule,
                     SmallVectorImpl<SmallString<12>> &modifiers,
@@ -277,7 +316,7 @@ parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule,
   }
 
   // If there is a comma, we have one or more modifiers..
-  if (succeeded(parser.parseOptionalComma())) {
+  while (succeeded(parser.parseOptionalComma())) {
     StringRef mod;
     if (parser.parseKeyword(&mod))
       return failure();
@@ -287,19 +326,24 @@ parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule,
   if (parser.parseRParen())
     return failure();
 
+  if (verifyScheduleModifiers(parser, modifiers))
+    return failure();
+
   return success();
 }
 
 /// Print schedule clause
 static void printScheduleClause(OpAsmPrinter &p, StringRef &sched,
-                                llvm::Optional<StringRef> modifier,
+                                llvm::Optional<StringRef> modifier, bool simd,
                                 Value scheduleChunkVar) {
   std::string schedLower = sched.lower();
   p << "schedule(" << schedLower;
   if (scheduleChunkVar)
     p << " = " << scheduleChunkVar;
-  if (modifier && modifier.getValue() != "none")
+  if (modifier && modifier.hasValue())
     p << ", " << modifier;
+  if (simd)
+    p << ", simd";
   p << ") ";
 }
 
@@ -822,6 +866,13 @@ static ParseResult parseClauses(OpAsmParser &parser, OperationState &result,
     if (modifiers.size() > 0) {
       auto mod = parser.getBuilder().getStringAttr(modifiers[0]);
       result.addAttribute("schedule_modifier", mod);
+      // Only SIMD attribute is allowed here!
+      if (modifiers.size() > 1) {
+        assert(symbolizeScheduleModifier(modifiers[1]) ==
+               mlir::omp::ScheduleModifier::simd);
+        auto attr = UnitAttr::get(parser.getBuilder().getContext());
+        result.addAttribute("simd_modifier", attr);
+      }
     }
     if (scheduleChunkSize) {
       auto chunkSizeType = parser.getBuilder().getI32Type();
@@ -1026,7 +1077,7 @@ static void printWsLoopOp(OpAsmPrinter &p, WsLoopOp op) {
 
   if (auto sched = op.schedule_val())
     printScheduleClause(p, sched.getValue(), op.schedule_modifier(),
-                        op.schedule_chunk_var());
+                        op.simd_modifier(), op.schedule_chunk_var());
 
   if (auto collapse = op.collapse_val())
     p << "collapse(" << collapse << ") ";

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

@@ -684,6 +684,9 @@ convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
       ompBuilder->collapseLoops(diLoc, loopInfos, {});
 
   allocaIP = findAllocaInsertPoint(builder, moduleTranslation);
+
+  bool isSimd = loop.simd_modifier();
+
   if (schedule == omp::ClauseScheduleKind::Static) {
     ompBuilder->applyStaticWorkshareLoop(ompLoc.DL, loopInfo, allocaIP,
                                          !loop.nowait(), chunk);
@@ -694,12 +697,18 @@ convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
       schedType = llvm::omp::OMPScheduleType::DynamicChunked;
       break;
     case omp::ClauseScheduleKind::Guided:
+      if (isSimd)
+        schedType = llvm::omp::OMPScheduleType::GuidedSimd;
+      else
         schedType = llvm::omp::OMPScheduleType::GuidedChunked;
       break;
     case omp::ClauseScheduleKind::Auto:
       schedType = llvm::omp::OMPScheduleType::Auto;
       break;
     case omp::ClauseScheduleKind::Runtime:
+      if (isSimd)
+        schedType = llvm::omp::OMPScheduleType::RuntimeSimd;
+      else
         schedType = llvm::omp::OMPScheduleType::Runtime;
       break;
     default:

mlir/test/Dialect/OpenMP/ops.mlir

@@ -197,7 +197,7 @@ func @omp_wsloop_pretty(%lb : index, %ub : index, %step : index,
   }
 
   // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static)
-  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) schedule(static, none) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>) {
+  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) schedule(static) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>) {
     omp.yield
   }
 

mlir/test/Target/LLVMIR/openmp-llvm-bad-schedule-modifier.mlir

@@ -0,0 +1,49 @@
+// RUN: not mlir-translate -mlir-to-llvmir -split-input-file %s 2>&1 | FileCheck %s
+
+llvm.func @test_omp_wsloop_dynamic_bad_modifier(%lb : i64, %ub : i64, %step : i64) -> () {
+  omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, ginandtonic) {
+    // CHECK: unknown modifier type: ginandtonic
+    omp.yield
+  }
+  llvm.return
+}
+
+// -----
+
+llvm.func @test_omp_wsloop_dynamic_many_modifier(%lb : i64, %ub : i64, %step : i64) -> () {
+  omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, monotonic, monotonic, monotonic) {
+    // CHECK: unexpected modifier(s)
+    omp.yield
+  }
+  llvm.return
+}
+
+// -----
+
+llvm.func @test_omp_wsloop_dynamic_wrong_modifier(%lb : i64, %ub : i64, %step : i64) -> () {
+  omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, simd, monotonic) {
+    // CHECK: incorrect modifier order
+    omp.yield
+  }
+  llvm.return
+}
+
+// -----
+
+llvm.func @test_omp_wsloop_dynamic_wrong_modifier2(%lb : i64, %ub : i64, %step : i64) -> () {
+  omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, monotonic, monotonic) {
+    // CHECK: incorrect modifier order
+    omp.yield
+  }
+  llvm.return
+}
+
+// -----
+
+llvm.func @test_omp_wsloop_dynamic_wrong_modifier3(%lb : i64, %ub : i64, %step : i64) -> () {
+  omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, simd, simd) {
+    // CHECK: incorrect modifier order
+    omp.yield
+  }
+  llvm.return
+}

mlir/test/Target/LLVMIR/openmp-llvm.mlir

@@ -517,6 +517,30 @@ llvm.func @test_omp_wsloop_dynamic_monotonic(%lb : i64, %ub : i64, %step : i64)
   llvm.return
 }
 
+llvm.func @test_omp_wsloop_runtime_simd(%lb : i64, %ub : i64, %step : i64) -> () {
+  omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(runtime, simd) {
+    // CHECK: call void @__kmpc_dispatch_init_8u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 47
+    // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u
+    // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0
+    // CHECK  br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}}
+    llvm.call @body(%iv) : (i64) -> ()
+    omp.yield
+  }
+  llvm.return
+}
+
+llvm.func @test_omp_wsloop_guided_simd(%lb : i64, %ub : i64, %step : i64) -> () {
+  omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(guided, simd) {
+    // CHECK: call void @__kmpc_dispatch_init_8u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 46
+    // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u
+    // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0
+    // CHECK  br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}}
+    llvm.call @body(%iv) : (i64) -> ()
+    omp.yield
+  }
+  llvm.return
+}
+
 // -----
 
 omp.critical.declare @mutex hint(contended)