[mlir][OpenMP]Support for modifiers in workshare loops

[Leporacanthicus]

Summary

Pass the modifiers from the Flang parser to FIR/MLIR workshare
loop operation.

Not yet supporting the SIMD modifier, which is a bit more work
than just adding it to the list of modifiers, so will go in a
separate patch.

This adds a new field to the WsLoopOp.

Also add test for dynamic WSLoop, checking that dynamic schedule calls
the init and next functions as expected.

[Meinersbur]

What are your plans to support the simd directive? Does schedule_modifiers() also become a list? A flags enum? Or do you split schedule_modifiers() into separate attributes such as monotonicity_schedule_modifier() and simd_schedule_modifier()?

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

416

Parameters should be SmallVectorImpl to not force the caller to a specific small-size. Same theoretically also applies to SmallString (SmallVectorTempl<char>) but that doesn't work with type arguments. (but it would for SmallString<8> &schedule).

It's a static function anyway, so treat this as a nitpicking.

[MatsPetersson]

In D111053#3041165, @Meinersbur wrote:

What are your plans to support the simd directive? Does schedule_modifiers() also become a list? A flags enum? Or do you split schedule_modifiers() into separate attributes such as mononoicity_schedule_modifier() and simd_schedule_modifier()?

Yes, it's in D111051 (linked to this one as "Child Revision"). The reviews are "backwards" because of my inability to use Arcanist - it decided that I only wanted to upload ONE of the two patches. Both of these patches already in fir-dev, I'm submitting them to llvm/main to make sure they don't accidentally get lost when fir-dev is catching up with llvm/main.

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

416

I'll take a look at this - I think there's quite a bit of "potential for improvement" here, and I'll see if it's better to do a whole other "clean things up" patch.

[ftynse]

Thanks! A couple of minor comments from me.

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

228

Nit: there's only one attribute that can be present, so I would use singular for the name.

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

440

The case with more modifiers is not handled.

745

Could we omit the "none" modifier in the format as, I believe, it is the default? This will avoid the change in most tests below.

[clementval]

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

137

I think @kiranchandramohan added a way to have this coming directly from the OMP.td file.

[Leporacanthicus]

Updating according to review comments

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

416

Fixed.

440

It is "handled" in the sense that it gives an error below for "expected parenthesis". The only valid case is for further modifiers is SIMD, which is covered in the next patch.

[kiranchandramohan]

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

137

This was introduced as part of https://reviews.llvm.org/D84347.

If you can define ClauseVal entries in llvm/include/llvm/Frontend/OpenMP/OMP.td then these entries will be autogenerated. Example for proc_bind given below.

def OMP_PROC_BIND_master : ClauseVal<"master",2,1> {}
def OMP_PROC_BIND_close : ClauseVal<"close",3,1> {}
def OMP_PROC_BIND_spread : ClauseVal<"spread",4,1> {}
def OMP_PROC_BIND_default : ClauseVal<"default",5,0> {}
def OMP_PROC_BIND_unknown : ClauseVal<"unknown",6,0> { let isDefault = 1; }
def OMPC_ProcBind : Clause<"proc_bind"> {
  let clangClass = "OMPProcBindClause";
  let flangClass = "OmpProcBindClause";
  let enumClauseValue = "ProcBindKind";
  let allowedClauseValues = [
    OMP_PROC_BIND_master,
    OMP_PROC_BIND_close,
    OMP_PROC_BIND_spread,
    OMP_PROC_BIND_default,
    OMP_PROC_BIND_unknown
  ];
}

[Leporacanthicus]

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

137

But a schedule modifier is not a clause on its own - it's part of the ScheduleClause - I first implemented it as a SheduleModifierClause, but I was told that it isn't a clause.

[Leporacanthicus]

Rebase to latest llvm/main

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

let description = [{
     operation.  These regions are not expected to return any value so the
     terminator takes no operands. The terminator op returns control to the
     enclosing op.
   }];
 
   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 ScheduleModifier : StrEnumAttr<"ScheduleModifier", "OpenMP Schedule Modifier",
+                                         [OMP_SCHEDULE_MOD_None,
+                                          OMP_SCHEDULE_MOD_Monotonic,
+                                          OMP_SCHEDULE_MOD_Nonmonotonic]>
+{
+    let cppNamespace = "::mlir::omp";
+}
+
 //===----------------------------------------------------------------------===//
 // 2.9.2 Workshare Loop Construct
 //===----------------------------------------------------------------------===//
 
 def WsLoopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments,
                          AllTypesMatch<["lowerBound", "upperBound", "step"]>]> {
   let summary = "workshare loop construct";
   let description = [{
let arguments = (ins Variadic<IntLikeType>:$lowerBound,
              Variadic<AnyType>:$lastprivate_vars,
              Variadic<AnyType>:$linear_vars,
              Variadic<AnyType>:$linear_step_vars,
              Variadic<OpenMP_PointerLikeType>:$reduction_vars,
              OptionalAttr<TypedArrayAttrBase<SymbolRefAttr,
                             "array of symbol references">>:$reductions,
              OptionalAttr<ScheduleKind>:$schedule_val,
              Optional<AnyType>:$schedule_chunk_var,
+             OptionalAttr<ScheduleModifier>:$schedule_modifiers,
              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

 /// schedule ::= `schedule` `(` sched-list `)`
 /// sched-list ::= sched-val | sched-val sched-list
 /// 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`
 static ParseResult
 parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule,
+                    SmallVector<SmallString<12>> &modifiers,
                     Optional<OpAsmParser::OperandType> &chunkSize) {
   if (parser.parseLParen())
     return failure();
 
   StringRef keyword;
   if (parser.parseKeyword(&keyword))
     return failure();
 
   schedule = keyword;
   if (keyword == "static" || keyword == "dynamic" || keyword == "guided") {
     if (succeeded(parser.parseOptionalEqual())) {
       chunkSize = OpAsmParser::OperandType{};
       if (parser.parseOperand(*chunkSize))
         return failure();
     } else {
       chunkSize = llvm::NoneType::None;
     }
   } 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())) {
+    StringRef mod;
+    if (parser.parseKeyword(&mod))
+      return failure();
+    modifiers.push_back(mod);
+  }
+
   if (parser.parseRParen())
     return failure();
 
   return success();
 }
 
 /// reduction-init ::= `reduction` `(` reduction-entry-list `)`
 /// reduction-entry-list ::= reduction-entry
static ParseResult parseWsLoopOp(OpAsmParser &parser, OperationState &result) {
   SmallVector<Type> lastprivateTypes;
   SmallVector<OpAsmParser::OperandType> linears;
   SmallVector<Type> linearTypes;
   SmallVector<OpAsmParser::OperandType> linearSteps;
   SmallVector<SymbolRefAttr> reductionSymbols;
   SmallVector<OpAsmParser::OperandType> reductionVars;
   SmallVector<Type> reductionVarTypes;
   SmallString<8> schedule;
+  SmallVector<SmallString<12>> modifiers;
   Optional<OpAsmParser::OperandType> scheduleChunkSize;
 
   const StringRef opName = result.name.getStringRef();
   StringRef keyword;
 
   enum SegmentPos {
     lbPos = 0,
     ubPos,
if (keyword == "private") {
         return allowedOnce(parser, "linear", opName);
       if (parseLinearClause(parser, linears, linearTypes, linearSteps))
         return failure();
       segments[linearClausePos] = linears.size();
       segments[linearStepPos] = linearSteps.size();
     } else if (keyword == "schedule") {
       if (!schedule.empty())
         return allowedOnce(parser, "schedule", opName);
-      if (parseScheduleClause(parser, schedule, scheduleChunkSize))
+      if (parseScheduleClause(parser, schedule, modifiers, scheduleChunkSize))
         return failure();
       if (scheduleChunkSize) {
         segments[scheduleClausePos] = 1;
       }
     } else if (keyword == "collapse") {
       auto type = parser.getBuilder().getI64Type();
       mlir::IntegerAttr attr;
       if (parser.parseLParen() || parser.parseAttribute(attr, type) ||
if (segments[reductionVarPos]) {
     result.addAttribute("reductions",
                         parser.getBuilder().getArrayAttr(reductions));
   }
 
   if (!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_modifiers", mod);
+    }
     if (scheduleChunkSize) {
       auto chunkSizeType = parser.getBuilder().getI32Type();
       parser.resolveOperand(*scheduleChunkSize, chunkSizeType, result.operands);
     }
   }
 
   result.addAttribute("operand_segment_sizes",
                       parser.getBuilder().getI32VectorAttr(segments));
static void printWsLoopOp(OpAsmPrinter &p, WsLoopOp op) {
   }
 
   if (auto sched = op.schedule_val()) {
     auto schedLower = sched->lower();
     p << " schedule(" << schedLower;
     if (auto chunk = op.schedule_chunk_var()) {
       p << " = " << chunk;
     }
+    if (auto modifier = op.schedule_modifiers()) {
+      p << ", " << modifier;
+    }
     p << ")";
   }
 
   if (auto collapse = op.collapse_val())
     p << " collapse(" << collapse << ")";
 
   if (op.nowait())
     p << " nowait";

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

if (schedule == omp::ClauseScheduleKind::Static) {
     case omp::ClauseScheduleKind::Runtime:
       schedType = llvm::omp::OMPScheduleType::Runtime;
       break;
     default:
       llvm_unreachable("Unknown schedule value");
       break;
     }
 
-    ompBuilder->applyDynamicWorkshareLoop(ompLoc.DL, loopInfo, allocaIP,
-                                          schedType, !loop.nowait(), chunk);
+    if (loop.schedule_modifiers().hasValue()) {
+      omp::ScheduleModifier modifier =
+          *omp::symbolizeScheduleModifier(loop.schedule_modifiers().getValue());
+      switch (modifier) {
+      case omp::ScheduleModifier::monotonic:
+        schedType |= llvm::omp::OMPScheduleType::ModifierMonotonic;
+        break;
+      case omp::ScheduleModifier::nonmonotonic:
+        schedType |= llvm::omp::OMPScheduleType::ModifierNonmonotonic;
+        break;
+      default:
+        // Nothing to do here.
+        break;
+      }
+    }
+    afterIP = ompBuilder->applyDynamicWorkshareLoop(
+        ompLoc, loopInfo, allocaIP, schedType, !loop.nowait(), chunk);
   }
 
   // Continue building IR after the loop. Note that the LoopInfo returned by
   // `collapseLoops` points inside the outermost loop and is intended for
   // potential further loop transformations. Use the insertion point stored
   // before collapsing loops instead.
   builder.restoreIP(afterIP);
 

mlir/test/Dialect/OpenMP/ops.mlir

 func @omp_wsloop_pretty(%lb : index, %ub : index, %step : index,
                  %data_var : memref<i32>, %linear_var : i32, %chunk_var : i32) -> () {
 
   // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>)
   omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) private(%data_var : memref<i32>) collapse(2) ordered(2) {
     omp.yield
   }
 
-  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static)
-  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) schedule(static) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>) {
+  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static, none)
+  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.yield
+  }
+
+  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) lastprivate(%{{.*}} : memref<i32>) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static = %{{.*}}, none) collapse(3) ordered(2)
+  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) ordered(2) private(%data_var : memref<i32>)
+     firstprivate(%data_var : memref<i32>) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>)
+     schedule(static = %chunk_var, none) collapse(3) {
+    omp.yield
+  }
+
+  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) lastprivate(%{{.*}} : memref<i32>) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}}, nonmonotonic) collapse(3) ordered(2)
+  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) ordered(2) private(%data_var : memref<i32>)
+     firstprivate(%data_var : memref<i32>) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>)
+     schedule(dynamic = %chunk_var, nonmonotonic) collapse(3) {
     omp.yield
   }
 
-  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) lastprivate(%{{.*}} : memref<i32>) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static = %{{.*}}) collapse(3) ordered(2)
+  // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) lastprivate(%{{.*}} : memref<i32>) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}}, monotonic) collapse(3) ordered(2)
   omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) ordered(2) private(%data_var : memref<i32>)
      firstprivate(%data_var : memref<i32>) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>)
-     schedule(static = %chunk_var) collapse(3) {
+     schedule(dynamic = %chunk_var, monotonic) collapse(3) {
     omp.yield
   }
 
   // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private({{.*}} : memref<i32>)
   omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) private(%data_var : memref<i32>) {
     omp.yield
   }
 

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

^bb0(%arg1: i64):
     omp.yield
   }) {inclusive, operand_segment_sizes = dense<[1, 1, 1, 0, 0, 0, 0, 0, 0, 0]> : vector<10xi32>} : (i64, i64, i64) -> ()
   llvm.return
 }
 
 llvm.func @body(i64)
 
 llvm.func @test_omp_wsloop_dynamic(%lb : i64, %ub : i64, %step : i64) -> () {
- omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic) {
+ omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, none) {
   // CHECK: call void @__kmpc_dispatch_init_8u
   // 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_auto(%lb : i64, %ub : i64, %step : i64) -> () {
- omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(auto) {
+ omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(auto, none) {
   // CHECK: call void @__kmpc_dispatch_init_8u
   // 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_runtime(%lb : i64, %ub : i64, %step : i64) -> () {
- omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(runtime) {
+ omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(runtime, none) {
   // CHECK: call void @__kmpc_dispatch_init_8u
   // 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(%lb : i64, %ub : i64, %step : i64) -> () {
- omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(guided) {
+ omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(guided, none) {
   // CHECK: call void @__kmpc_dispatch_init_8u
   // 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
// CHECK: store
     omp.terminator
   }
   // CHECK: call void @__kmpc_end_critical({{.*}}critical_user_mutex.var{{.*}})
   llvm.return
 }
 
 // -----
 
+llvm.func @test_omp_wsloop_dynamic_nonmonotonic(%lb : i64, %ub : i64, %step : i64) -> () {
+ omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, nonmonotonic) {
+  // CHECK: call void @__kmpc_dispatch_init_8u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 1073741859
+  // 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_dynamic_monotonic(%lb : i64, %ub : i64, %step : i64) -> () {
+ omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic, monotonic) {
+  // CHECK: call void @__kmpc_dispatch_init_8u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 536870947
+  // 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
+}
+
+// -----
+
 // Check that the loop bounds are emitted in the correct location in case of
 // collapse. This only checks the overall shape of the IR, detailed checking
 // is done by the OpenMPIRBuilder.
 
 // CHECK-LABEL: @collapse_wsloop
 // CHECK: i32* noalias %[[TIDADDR:[0-9A-Za-z.]*]]
 // CHECK: load i32, i32* %[[TIDADDR]]
 // CHECK: store