[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

256

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

256

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

229

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

280

The case with more modifiers is not handled.

936

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

139

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

256

Fixed.

280

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

139

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

139

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>:$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,
              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,
+                    SmallVectorImpl<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();
 }
 
 /// Print schedule clause
 static void printScheduleClause(OpAsmPrinter &p, StringRef &sched,
+                                llvm::Optional<StringRef> modifier,
                                 Value scheduleChunkVar) {
   std::string schedLower = sched.lower();
   p << "(" << schedLower;
   if (scheduleChunkVar)
     p << " = " << scheduleChunkVar;
+  if (modifier && modifier.getValue() != "none")
+    p << ", " << modifier;
   p << ") ";
 }
 
 //===----------------------------------------------------------------------===//
 // Parser, printer and verifier for ReductionVarList
 //===----------------------------------------------------------------------===//
 
 /// reduction ::= `reduction` `(` reduction-entry-list `)`
static ParseResult parseClauses(OpAsmParser &parser, OperationState &result,
   SmallVector<OpAsmParser::OperandType> reductionVars;
   SmallVector<Type> reductionVarTypes;
 
   SmallVector<OpAsmParser::OperandType> linears;
   SmallVector<Type> linearTypes;
   SmallVector<OpAsmParser::OperandType> linearSteps;
 
   SmallString<8> schedule;
+  SmallVector<SmallString<12>> modifiers;
   Optional<OpAsmParser::OperandType> scheduleChunkSize;
 
   // Compute the position of clauses in operand segments
   int currPos = 0;
   for (ClauseType clause : clauses) {
 
     // Skip the following clauses - they do not take any position in operand
     // segments
while (succeeded(parser.parseOptionalKeyword(&clauseKeyword))) {
     } else if (clauseKeyword == "linear") {
       if (checkAllowed(linearClause) ||
           parseLinearClause(parser, linears, linearTypes, linearSteps))
         return failure();
       clauseSegments[pos[linearClause]] = linears.size();
       clauseSegments[pos[linearClause] + 1] = linearSteps.size();
     } else if (clauseKeyword == "schedule") {
       if (checkAllowed(scheduleClause) ||
-          parseScheduleClause(parser, schedule, scheduleChunkSize))
+          parseScheduleClause(parser, schedule, modifiers, scheduleChunkSize))
         return failure();
       if (scheduleChunkSize) {
         clauseSegments[pos[scheduleClause]] = 1;
       }
     } else if (clauseKeyword == "collapse") {
       auto type = parser.getBuilder().getI64Type();
       mlir::IntegerAttr attr;
       if (checkAllowed(collapseClause) || parser.parseLParen() ||
if (failed(parser.resolveOperands(linears, linearTypes, linears[0].location,
       return failure();
   }
 
   // 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 (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_chunk_var());
+    printScheduleClause(p, sched.getValue(), op.schedule_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

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_modifier().hasValue()) {
+      omp::ScheduleModifier modifier =
+          *omp::symbolizeScheduleModifier(loop.schedule_modifier().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.DL, 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>) {
+  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 = %{{.*}}) 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) {
     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(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(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
   }
 
   return
 }
 

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

omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(guided) {
   // 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_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
+}
+
 // -----
 
 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