[mlir][OpenMP] Add support for SIMD modifier

[Leporacanthicus]

Summary

Add support for SIMD modifier in OpenMP worksharing loops.

[Leporacanthicus]

Trying to get two commits in one arcanist command - may not work

[ftynse]

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

231–232

I'm not sure I understand the modeling here. The same enum attribute is listed twice with different names, and there is no verification of what can go where. It looks like one can have OMP_SCHEUDLE_MOD_SIMD as schedule_modifiers and OMP__SCHEDULE_MOD_None as simd_modifier. Or even monotonic in one and non-monotonic in the other.

Should SIMD be just a unit attribute?

Also consider renaming schedule_modifiers to schedule_modifier (singular) because it can only contain a single instance anyway.

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

272–275

Nit: you may need to update this comment

860

This needs to error out with a message if there is more than two modifiers.

We also need a verifier that checks that modifiers are what we expect them to be, e.g., there's no conflict between monotonic and non-monotonic, no repetition, etc. See the remark above.

990

It looks preferable to have the default case here (modifier is none and missing modifier are the same, aren't they?) and just avoid extra printing in that case altogether.

[Meinersbur]

There is an RFC for better support for unordered lists in MLIR assembly, maybe that could be used. However, as @Leporacanthicus mentioned in the OpenMP/flang conference call, D111051 and this are part of upstreaming the fir-dev branch. In the interest of not making the upstreaming process for fir-dev more difficult, we should consider to keep the delta low and defer larger changes after complete upstreaming of fir-dev.

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

231–232

Since simd is also a schedule modifier, maybe rename schedule_modifiers to monotonicity_modifier? I.e. we'd have monotonicity_[schedule_]modifier and simd_[schedule_]modifier.

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

990

[serious] I think there might be a real bug here. If no monotonicity modifier is used, but the simd modifier, then according this code will print schedule(static,simd). However, the parsing code will interpret the first modifier as monotonicity modifier, and try to interpret simd as one.

[Leporacanthicus]

Update to match changes in dependent commit (+ rebase)

[Leporacanthicus]

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

231–232

I haven't completely fixed this, but i have added error checking code in a separate commit [to keep changes here to a minimum -> helping the rebase on fir-dev to avoid too much problems]. https://reviews.llvm.org/D112526

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

272–275

Done as a separate commit. https://reviews.llvm.org/D112526

860

Done in a separate commit:
https://reviews.llvm.org/D112526

990

Done in a separate commit, see https://reviews.llvm.org/D112526

[ftynse]

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

231–232

Why can't that commit be folded into this change? It has not landed yet, and it would be significantly easier to review to bisect if necessary. This may require rebasing some downstream patchset, but that is the usual price to pay for developing downstream.

[Leporacanthicus]

BUMP.

[Meinersbur]

In D111051#3100634, @Leporacanthicus wrote:

BUMP.

Could you address @ftynse's question on why you create new reviews instead of updating this one?

[Leporacanthicus]

Squashed the "fixes" commit and the original commit

[shraiysh]

Thanks for this patch. I have a few comments @Leporacanthicus.

Could you please clarify the purpose of mlir/test/Target/LLVMIR/openmp-llvm-bad-schedule-modifier.mlir? As far as I can tell, these tests can be a part of llvm-project/mlir/test/Dialect/OpenMP/invalid.mlir.

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

231–232

+1 to SIMD being a unit attribute.

Would it be better to have an ArrayAttr of ScheculeModifier (similar to the spec)?

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

245

suggestion: renaming it to verifyScheduleClause and calling it from verifyWsLoopOp?

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

525–526

nit: indentation

537–538

nit: indentation

[ftynse]

LGTM when all comments are addresesd.

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

258

Nit: please use braces symmterically

269

Having two terms on the LHS isn't valid BNF AFAIK.

274

Ditto, I don't understand what this wants to say.

307

I believe there is a function symbolizeScheduleModifier or something similar that is generated from ODS, which returns Optional<ScheduleModifier> equal to None if the wrong keyword is used. This avoids magic strings and is also future-proof.

307–310

Nit: I would have put this into the validation function.

[Leporacanthicus]

Fix review comments

[Leporacanthicus]

Make simd_modifier a unit-attribute and fix formatting

[Leporacanthicus]

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

231–232

I've made the simd_modifier into a unit attribute.

This will need to be updated in fir-dev to match, I will make a patch for that when I get this one into llvm/main.

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

245

Rename: I've changed the name to verifyScheduelModifiers to represent what it does.
Calling it from verifyWsLoopOp: some of the error checking done here can't be done at that point - such as checking if there's more than two modifiers, because we no longer have the array of modifiers, only the two that has been stored in the WsLoopOp.

258

Rewritten it in a simpler way to reduce the need for parenthesis.

269

Darned clang-format has messed with my comments!

I'll fix, but may need to use shorter names.

274

Ditto clang-format reflowing comments.

307

Indeed there is. I'm changing this here, and then doing a similar change for the ScheduleKind which is being tested in form of strings above.

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

525–526

Fixing all of the incorrectly formatted tests (all my fault!)

[peixin]

LGTM.

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

695

The simd modifier info is not passed to applyStaticWorkshareLoop. That is, the schedule type of StaticBalancedChunked is not handled and checked. Considering this function is refactored in https://reviews.llvm.org/D114413. Maybe you can merge this patch, and let it to be handled in D114413.

[Leporacanthicus]

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

695

Thanks for the LGTM. I've pushed this patch, and if you are happy to do that in D114413, then do that - or we can fix it in a different patch.

[peixin]

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

695

OK. It is better to let one patch focus on one thing. Please remember to make this fix after D114413 is accepted and merged.

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

enum class OMPScheduleType {
   Static = 34, // static unspecialized
   DistributeChunked = 91,
   Distribute = 92,
   DynamicChunked = 35,
   GuidedChunked = 36, // guided unspecialized
   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 =
       (1 << 30), // Set if the nonmonotonic schedule modifier was present
   ModifierMask = ModifierMonotonic | ModifierNonmonotonic,
   LLVM_MARK_AS_BITMASK_ENUM(/* LargestValue */ ModifierMask)
 };
 

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

def TerminatorOp : OpenMP_Op<"terminator", [Terminator]> {
   }];
 
   let assemblyFormat = "attr-dict";
 }
 
 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";
 }
 
 //===----------------------------------------------------------------------===//
 // 2.9.2 Workshare Loop Construct
 //===----------------------------------------------------------------------===//
 
let arguments = (ins Variadic<IntLikeType>:$lowerBound,
              Variadic<AnyType>:$firstprivate_vars,
              Variadic<AnyType>:$lastprivate_vars,
              Variadic<AnyType>:$linear_vars,
              Variadic<AnyType>:$linear_step_vars,
              Variadic<OpenMP_PointerLikeType>:$reduction_vars,
              OptionalAttr<SymbolRefArrayAttr>:$reductions,
              OptionalAttr<ScheduleKind>:$schedule_val,
              Optional<AnyType>:$schedule_chunk_var,
              OptionalAttr<ScheduleModifier>:$schedule_modifier,
+             OptionalAttr<ScheduleModifier>:$simd_modifier,
              Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$collapse_val,
              UnitAttr:$nowait,
              Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$ordered_val,
              OptionalAttr<OrderKind>:$order_val,
              UnitAttr:$inclusive);
 
   let skipDefaultBuilders = 1;
 

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

for (unsigned i = 0; i < linearVarsSize; ++i) {
     std::string separator = i == linearVarsSize - 1 ? ") " : ", ";
     p << linearVars[i];
     if (linearStepVars.size() > i)
       p << " = " << linearStepVars[i];
     p << " : " << linearVars[i].getType() << separator;
   }
 }
 
+static ParseResult
+validateModifiers(OpAsmParser &parser,
+                  SmallVectorImpl<SmallString<12>> &modifiers) {
+  if (modifiers.size() > 2)
+    return parser.emitError(parser.getNameLoc()) << " unexpected modifier(s)";
+  // If we have one modifier that is "simd", then stick a "none" modiifer in
+  // index 0.
+  if (modifiers.size() == 1) {
+    if (modifiers[0] == "simd") {
+      modifiers.push_back(modifiers[0]);
+      modifiers[0] = "none";
+    }
+  } else if (modifiers.size() == 2 &&
+             (modifiers[0] == "simd" ||
+              (modifiers[1] != "simd" && modifiers[1] != "none")))
+    return parser.emitError(parser.getNameLoc()) << " incorrect modifier order";
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Parser and printer for Schedule Clause
 //===----------------------------------------------------------------------===//
 
 /// schedule ::= `schedule` `(` sched-list `)`
-/// sched-list ::= sched-val | sched-val sched-list
-/// sched-val ::= sched-with-chunk | sched-wo-chunk
+/// 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-modifier-val | sched-modifier-val `,`
+/// sched-modifier-val sched-modifier ::=  `monotonic` | `nonmonotonic` | `simd`
+/// | `none`
 static ParseResult
 parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule,
                     SmallVectorImpl<SmallString<12>> &modifiers,
                     Optional<OpAsmParser::OperandType> &chunkSize) {
   if (parser.parseLParen())
     return failure();
 
   StringRef keyword;
if (keyword == "static" || keyword == "dynamic" || keyword == "guided") {
     }
   } else if (keyword == "auto" || keyword == "runtime") {
     chunkSize = llvm::NoneType::None;
   } else {
     return parser.emitError(parser.getNameLoc()) << " expected schedule kind";
   }
 
   // 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();
+    if (mod != "none" && mod != "monotonic" && mod != "nonmonotonic" &&
+        mod != "simd")
+      return parser.emitError(parser.getNameLoc())
+             << " unknown modifier type: " << mod;
     modifiers.push_back(mod);
   }
 
+  if (validateModifiers(parser, modifiers))
+    return failure();
+
   if (parser.parseRParen())
     return failure();
 
   return success();
 }
 
 /// Print schedule clause
 static void printScheduleClause(OpAsmPrinter &p, StringRef &sched,
                                 llvm::Optional<StringRef> modifier,
+                                llvm::Optional<StringRef> simd,
                                 Value scheduleChunkVar) {
   std::string schedLower = sched.lower();
   p << "(" << schedLower;
   if (scheduleChunkVar)
     p << " = " << scheduleChunkVar;
   if (modifier && modifier.getValue() != "none")
     p << ", " << modifier;
+  if (simd && simd.getValue() != "none")
+    p << ", " << simd;
   p << ") ";
 }
 
 //===----------------------------------------------------------------------===//
 // Parser, printer and verifier for ReductionVarList
 //===----------------------------------------------------------------------===//
 
 /// reduction ::= `reduction` `(` reduction-entry-list `)`
static ParseResult parseClauses(OpAsmParser &parser, OperationState &result,
   // Add schedule parameters
   if (done[scheduleClause] && !schedule.empty()) {
     schedule[0] = llvm::toUpper(schedule[0]);
     auto attr = parser.getBuilder().getStringAttr(schedule);
     result.addAttribute("schedule_val", attr);
     if (modifiers.size() > 0) {
       auto mod = parser.getBuilder().getStringAttr(modifiers[0]);
       result.addAttribute("schedule_modifier", mod);
+      if (modifiers.size() > 1) {
+        mod = parser.getBuilder().getStringAttr(modifiers[1]);
+        result.addAttribute("simd_modifier", mod);
+      }
     }
     if (scheduleChunkSize) {
       auto chunkSizeType = parser.getBuilder().getI32Type();
       parser.resolveOperand(*scheduleChunkSize, chunkSizeType, result.operands);
     }
   }
 
   segments.insert(segments.end(), clauseSegments.begin(), clauseSegments.end());
static void printWsLoopOp(OpAsmPrinter &p, WsLoopOp op) {
 
   if (op.linear_vars().size()) {
     p << "linear";
     printLinearClause(p, op.linear_vars(), op.linear_step_vars());
   }
 
   if (auto sched = op.schedule_val()) {
     p << "schedule";
     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 << ") ";
 
   if (op.nowait())
     p << "nowait ";
 

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

convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
 
   // Collapse loops. Store the insertion point because LoopInfos may get
   // invalidated.
   llvm::IRBuilderBase::InsertPoint afterIP = loopInfos.front()->getAfterIP();
   llvm::CanonicalLoopInfo *loopInfo =
       ompBuilder->collapseLoops(diLoc, loopInfos, {});
 
   allocaIP = findAllocaInsertPoint(builder, moduleTranslation);
+
+  bool isSimd = false;
+  if (auto simd = loop.simd_modifier()) {
+    omp::ScheduleModifier modifier = *omp::symbolizeScheduleModifier(*simd);
+    isSimd = (modifier == omp::ScheduleModifier::simd);
+  }
+
   if (schedule == omp::ClauseScheduleKind::Static) {
     ompBuilder->applyStaticWorkshareLoop(ompLoc.DL, loopInfo, allocaIP,
                                          !loop.nowait(), chunk);
   } else {
     llvm::omp::OMPScheduleType schedType;
     switch (schedule) {
     case omp::ClauseScheduleKind::Dynamic:
       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:
       llvm_unreachable("Unknown schedule value");
       break;
     }
 
     if (loop.schedule_modifier().hasValue()) {
       omp::ScheduleModifier modifier =

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

+// 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

omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, monotonic) {
   // 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_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)
 
 // CHECK-LABEL: @omp_critical
 llvm.func @omp_critical(%x : !llvm.ptr<i32>, %xval : i32) -> () {
   // CHECK: call void @__kmpc_critical({{.*}}critical_user_.var{{.*}})
   // CHECK: br label %omp.critical.region