[MLIR][OpenMP] Distribute Construct Operation

[abidmalikwaterloo]

Summary

OpenMP distribute construct for MLIR operation and added a test case.

[kiranchandramohan]

Thanks @abidmalikwaterloo for this patch. This patch needs a cleanup and should not have any unrelated changes. I have added some comments.

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

500

Nit: Alignment

500

You can remove private, firstprivate and lastprivate for now.

504

Nit: Alignment

607

Unrelated to this patch.

612–620

Unrelated to this patch.

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

821–874

These are all unrelated to this patch. Please remove.

mlir/test/Dialect/OpenMP/ops.mlir

303–312 ↗

(On Diff #357051)

Removal of this is unrelated to this patch.

[clementval]

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

474

Line should be wrapped at 80.

475

Same here.

485

80

mlir/test/Dialect/OpenMP/ops.mlir

286 ↗

(On Diff #357051)

Unrelated to this patch

[clementval]

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

504

Maybe too permissive for a chunk_size?

[abidmalikwaterloo]

This patch has the same problem as the patch https://reviews.llvm.org/D105581/new/. The branching needs to be done from the main. The parsing part should not be there.

[abidmalikwaterloo]

rebase the patch branch with the main branch

[clementval]

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

504

Alignment is still off.

[abidmalikwaterloo]

Made changes based on the reviewers' comments.

[abidmalikwaterloo]

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

500

Any special reason for this?

504

We can make it AnyInteger. or IntORIndex

[clementval]

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

507

Does it make sense to have a min value of 0. I guess if the collapse attr is present it value is greater than 0 and probably even greater than 1.

[abidmalikwaterloo]

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

500

Any reason for removing them?

507

The loop operation is using the same value. However, I agreed that it should be >= 1.

[abidmalikwaterloo]

ping!

[clementval]

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

507

So do you want to keep 0 here?

[shraiysh]

Are we going to add pretty printers and parsers in a separate patch? Also, there are a few restrictions in the OpenMP reference, like

• The region corresponding to the distribute construct must be strictly nested inside a teams region.
• A list item may appear in a firstprivate or lastprivate clause but not both.

Maybe we can verify these (and others, if possible) too.

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

507

[suggestion] Maybe have it as a default valued attr with minimum of 1 like Confined<DefaultValuedAttr<I64Attr, "1">, [IntMinValue<1>]>.

[abidmalikwaterloo]

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

507

Yes, I am thinking of this. However, I was thinking should we do the same for the parallel operations?

[abidmalikwaterloo]

change the collapse defualt value from 0 to 1

[abidmalikwaterloo]

change the default collapse value from 0 to 1

[clementval]

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

507

Yeah you can probably do it for the parallel op as well but it should be a separate patch.

[abidmalikwaterloo]

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

507

Will do.

[clementval]

Are you still working on this?

[abidmalikwaterloo]

I was busy. I will try to finalize this week.

[peixin]

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

467

Should this be 2.9.4.1 distribute Construct since you are implementing distribute construct rather than distribute loop constructs?

[abidmalikwaterloo]

I have time now I will start looking them and would like finish them ASAP

[abidmalikwaterloo]

Based on the feedback from other patches. This also needs an assembly format!

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

467

will update.

[abidmalikwaterloo]

Added the assembly format, dist_schedule parser support

[shraiysh]

Thanks for working on this. A couple comments. There are no testcases. Please add testcases.

clang/lib/Testing/CMakeLists.txt

29

unrelated change?

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

488–490

If they are not added, please remove them from documentation too. It would be misleading otherwise.

514

dist_schedule is a dummy clause, where kind is only allowed to be static according to the standard. I don't think that requires a custom function, we can instead have something like the following -

arguments = UnitAttr:$static_dist_schedule, Optional<IntLikeType>:$schedule_chunk

assemblyFormat = `dist_schedule` `(` (`static` $static_dist_schedule^)? (`:` $schedule_chunk^)? `)`

Would that work? Let me know if there are any suggestions.

528

I think we need a verifier for this too. There are a couple semantic checks which we can do in verifier.

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

226

nit: schedule spelling.

[abidmalikwaterloo]

clang/lib/Testing/CMakeLists.txt

29

When I rebase, these changes were highlighted in the main branch which was missing in the patch as it was too old.

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

514

My only concern is; will this be a dummy clause with the static scheduler forever? I am pretty sure dist_schedule will have a dynamic or a user defined scheduling strategy as well to improve the performance of a given application

[abidmalikwaterloo]

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

488–490

Ok

528

Can you say more about the semantic checks you have in mind?

[shraiysh]

clang/lib/Testing/CMakeLists.txt

29

Hmm, these are in llvm-project/main right now but this means that the patch has not been rebased properly. These changes should not be a part of this patch :/ Can you try rebasing again? Otherwise this will cause issues while/after landing this patch.

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

500

They are being removed because we do not have a clean way of handling these clauses in MLIR for any generic frontend.

514

If and when it changes in the standard, at that time we can change the parsing/printing accordingly. Till then such a function seems unnecessary and a possible source of errors because it accepts invalid OpenMP code.

528

The following restriction from the standard can be added to the verifier/Operation definition -

The region corresponding to the distribute construct must be strictly nested inside a teams region.

The other restrictions - I am okay with not adding them because I don't know how they would be added. Needless to say if we figure out how to add them, then we should do it.

[abidmalikwaterloo]

clang/lib/Testing/CMakeLists.txt

29

git rebase origin arcpatch-D105584 gives " The current branch is up to date". It means the patch is up to date. Should I remove them manually?

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

514

Should it be then? dist_schedule ( (static ) (: $schedule_chunk^)? )
$static_dist_schedule seems redundant???

528

At this stage, we can add:

LogicalResult DistributeOp::verify(){

return success();

}

We can add the check later!

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

226

I do not think we need this if we treat the schedule as a variable?

[shraiysh]

clang/lib/Testing/CMakeLists.txt

29

I'm not sure what the issue is, but this change should not be reflected here if the patch is properly rebased with main.

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

514

AFAIK, the keyword static is optional. Without the presence of $static_dist_schedule, how are you going to store this flag?

528

Why not add the check now itself. Can it not be implemented? It's a small check :/

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

226

Yup. We probably don't need these functions.

565

Change this comment to Loop control.

[abidmalikwaterloo]

clang/lib/Testing/CMakeLists.txt

29

I will remove them manually. I think rebase is correct.

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

528

done

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

226

done

[abidmalikwaterloo]

// CHECK-LABEL: omp_DistributeOp
func.func @omp_DistributeOp(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %chunk_var : i32) -> () {
   // CHECK: omp.DistributeOp collapse(2)
  "omp.DistributeOp" (%lb, %ub, %step) ({
    ^bb0(%iv: index):
     omp.yield
  }) {operand_segment_sizes = dense<[1,1,1,0,0]> : vector<5xi32>, collapse_val = 2} :
    (index, index, index) -> ()
 
 return
 }

Is this a valid test case for the operation?

[clementval]

In D105584#3617456, @abidmalikwaterloo wrote:

// CHECK-LABEL: omp_DistributeOp
func.func @omp_DistributeOp(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %chunk_var : i32) -> () {
   // CHECK: omp.DistributeOp collapse(2)
  "omp.DistributeOp" (%lb, %ub, %step) ({
    ^bb0(%iv: index):
     omp.yield
  }) {operand_segment_sizes = dense<[1,1,1,0,0]> : vector<5xi32>, collapse_val = 2} :
    (index, index, index) -> ()
 
 return
 }

Is this a valid test case for the operation?

Should it be "omp.distribute" (%lb, %ub, %step) ({?

[abidmalikwaterloo]

In D105584#3617666, @clementval wrote:

In D105584#3617456, @abidmalikwaterloo wrote:

// CHECK-LABEL: omp_DistributeOp
func.func @omp_DistributeOp(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %chunk_var : i32) -> () {
   // CHECK: omp.DistributeOp collapse(2)
  "omp.DistributeOp" (%lb, %ub, %step) ({
    ^bb0(%iv: index):
     omp.yield
  }) {operand_segment_sizes = dense<[1,1,1,0,0]> : vector<5xi32>, collapse_val = 2} :
    (index, index, index) -> ()
 
 return
 }

Is this a valid test case for the operation?

Should it be "omp.distribute" (%lb, %ub, %step) ({?

Yes. It should be. The rest should be okay? But, the test is failing.

[abidmalikwaterloo]

In D105584#3618040, @abidmalikwaterloo wrote:

In D105584#3617666, @clementval wrote:

In D105584#3617456, @abidmalikwaterloo wrote:

// CHECK-LABEL: omp_DistributeOp
func.func @omp_DistributeOp(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %chunk_var : i32) -> () {
   // CHECK: omp.DistributeOp collapse(2)
  "omp.DistributeOp" (%lb, %ub, %step) ({
    ^bb0(%iv: index):
     omp.yield
  }) {operand_segment_sizes = dense<[1,1,1,0,0]> : vector<5xi32>, collapse_val = 2} :
    (index, index, index) -> ()
 
 return
 }

Is this a valid test case for the operation?

Should it be "omp.distribute" (%lb, %ub, %step) ({?

Yes. It should be. The rest should be okay? But, the test is failing.

It is working. I will up load the patch. Thanks

[abidmalikwaterloo]

Update the patch for assembly format. Change the WSLoopControl
to LoopControl. Added DistributeOp to parent of Yield OP. Added
one test case

[clementval]

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

215

No brace - https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

219

no brace

[abidmalikwaterloo]

Updated DistributeOp::verify as per reviewer's comments

[abidmalikwaterloo]

update the changes

[clementval]

clang/lib/Testing/CMakeLists.txt

30

You still have these changes that are unrelated.

[abidmalikwaterloo]

Removed unrelated changes in CMakefile.txt

[clementval]

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

449

[clementval]

Are you still working on this? Otherwise please close it.

[abidmalikwaterloo]

Seems that everything has been taken care of. It was accidentally out of my radar (do not know why). I will go through it once again and update it.

[abidmalikwaterloo]

Ping!
Does this need any further change? Seems okay to me.

[clementval]

I still see some review comments not done.

[abidmalikwaterloo]

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

449

What's the problem here? It seems fine to me.

[clementval]

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

449

syntax doesn't match the existing code. Change is highlighted.

[abidmalikwaterloo]

corrected the syntex of the code

[kiranchandramohan]

Patch probably needs a rebase. A few more minor things to fix. Looks mostly ready.

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

490

Do you mean static_dist_schedule here?

495–496

We can remove collapse for now.

513

In https://reviews.llvm.org/D128338 we removed collapse from all constructs. Currently collapse is modelled by having multiple entries in the lowerBound, upperBound and step.

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

211

220

Nit: Check that their sizes are equal as well. And also if step is present then its size matches lowerbound/upperbound.

mlir/test/Dialect/OpenMP/ops.mlir

124 ↗

(On Diff #473706)

Add a pretty-print test as well.

[abidmalikwaterloo]

In D105584#3917238, @kiranchandramohan wrote:

Patch probably needs a rebase. A few more minor things to fix. Looks mostly ready.

With which branch should I rebase? I have taken care rest of the comments. Thanks!

[kiranchandramohan]

In D105584#3954144, @abidmalikwaterloo wrote:

In D105584#3917238, @kiranchandramohan wrote:

Patch probably needs a rebase. A few more minor things to fix. Looks mostly ready.

With which branch should I rebase? I have taken care rest of the comments. Thanks!

Rebase to the main branch.

[abidmalikwaterloo]

Update the patch based on the feedback and comments.

[kiranchandramohan]

@abidmalikwaterloo It seems you missed a few of the previous comments. Please fix these so that we can approve.

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

449

Nit: Please fix this (adding a space before "DistributeOp").

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

211

Nit: Please make this change.

220

Nit: Please make this change as well.

mlir/test/Dialect/OpenMP/ops.mlir

124 ↗

(On Diff #473706)

Nit: please add a pretty-print test.

[abidmalikwaterloo]

mlir/test/Dialect/OpenMP/ops.mlir

124 ↗

(On Diff #473706)

Just need clarification. Do you mean something similar to the following:

// CHECK-LABEL: omp_wsloop_pretty
func.func @omp_wsloop_pretty(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %linear_var : i32, %chunk_var : i32, %chunk_var2 : i16) -> () {

  // CHECK: omp.wsloop ordered(2)
  // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
  omp.wsloop ordered(2)
  for (%iv) : index = (%lb) to (%ub) step (%step) {
    omp.yield
  }

  // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static)
  // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
  omp.wsloop schedule(static) linear(%data_var = %linear_var : memref<i32>)
  for (%iv) : index = (%lb) to (%ub) step (%step) {
    omp.yield
  }

  // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static = %{{.*}} : i32) ordered(2)
  // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
  omp.wsloop ordered(2) linear(%data_var = %linear_var : memref<i32>) schedule(static = %chunk_var : i32)
  for (%iv) : index = (%lb) to (%ub) step (%step) {
    omp.yield
  }

  // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}} : i32, nonmonotonic) ordered(2)
  // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
  omp.wsloop ordered(2) linear(%data_var = %linear_var : memref<i32>) schedule(dynamic = %chunk_var : i32, nonmonotonic)
  for (%iv) : index = (%lb) to (%ub) step (%step)  {
    omp.yield
  }

  // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}} : i16, monotonic) ordered(2)
  // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
  omp.wsloop ordered(2) linear(%data_var = %linear_var : memref<i32>) schedule(dynamic = %chunk_var2 : i16, monotonic)
  for (%iv) : index = (%lb) to (%ub) step (%step) {
    omp.yield
  }

  // CHECK: omp.wsloop for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
  omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
    omp.yield
  }

  // CHECK: omp.wsloop for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) inclusive step (%{{.*}})
  omp.wsloop for (%iv) : index = (%lb) to (%ub) inclusive step (%step) {
    omp.yield
  }

  // CHECK: omp.wsloop nowait
  // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
  omp.wsloop nowait
  for (%iv) : index = (%lb) to (%ub) step (%step) {
    omp.yield
  }

  // CHECK: omp.wsloop nowait order(concurrent)
  // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
  omp.wsloop order(concurrent) nowait
  for (%iv) : index = (%lb) to (%ub) step (%step) {
    omp.yield
  }

  return
}

[abidmalikwaterloo]

Ping!

[abidmalikwaterloo]

ping!

[kiranchandramohan]

mlir/test/Dialect/OpenMP/ops.mlir

124 ↗

(On Diff #473706)

Yes.

[abidmalikwaterloo]

Added a pretty test for omp.distrubute in ops.mlir

[kiranchandramohan]

Since the plan is to switch to the Canonical Style operation, may be this can wait till the omp.canonical_loop changes are in?

[abidmalikwaterloo]

In D105584#4081228, @kiranchandramohan wrote:

Since the plan is to switch to the Canonical Style operation, maybe this can wait till the omp.canonical_loop changes are in?

Ok..make sense

[kiranchandramohan]

Use omp.canonical_loop once it is in.

clang/lib/Testing/CMakeLists.txt

clang_target_link_libraries(clangTesting
   clangFrontend
   clangLex
   clangSerialization
   )
 
 target_link_libraries(clangTesting
   PRIVATE
   llvm_gtest
+  clangBasic
+  clangFrontend
   )

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

oilist(`linear` `(`
           |`collapse` `(` $collapse_val `)`
           |`nowait` $nowait
           |`ordered` `(` $ordered_val `)`
           |`order` `(` custom<ClauseAttr>($order_val) `)`
           |`reduction` `(`
               custom<ReductionVarList>(
                 $reduction_vars, type($reduction_vars), $reductions
               ) `)`
-    ) `for` custom<WsLoopControl>($region, $lowerBound, $upperBound, $step,
+    ) `for` custom<LoopControl>($region, $lowerBound, $upperBound, $step,
                                   type($step), $inclusive) attr-dict
   }];
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Simd construct [2.9.3.1] 
 //===----------------------------------------------------------------------===//
let summary = "simd loop construct";
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1; 
 }
 
 
 def YieldOp : OpenMP_Op<"yield",
     [NoSideEffect, ReturnLike, Terminator,
      ParentOneOf<["WsLoopOp", "ReductionDeclareOp", 
      "AtomicUpdateOp", "SimdLoopOp"]>]> {
   let summary = "loop yield and termination operation";
   let description = [{
     "omp.yield" yields SSA values from the OpenMP dialect op region and
     terminates the region. The semantics of how the values are yielded is
     defined by the parent operation.
   }];
 
   let arguments = (ins Variadic<AnyType>:$results);
 
   let builders = [
     OpBuilder<(ins), [{ build($_builder, $_state, {}); }]>
   ];
 
   let assemblyFormat = [{ ( `(` $results^ `:` type($results) `)` )? attr-dict}];
 }
 
 //===----------------------------------------------------------------------===//
+// 2.9.4.1 distribute Construct
+//===----------------------------------------------------------------------===//
+
+def DistributeOp : OpenMP_Op<"distribute", [AttrSizedOperandSegments,
+                        AllTypesMatch<["lowerBound", "upperBound", "step"]>]> {
+  let summary = "distribute loop construct";
+  let description = [{ 
+    The distribute construct specifies that the iterations of one or more loop
+    will be executed by the initial teams in the context of their implicit 
+    tasks. The iterations are distributed across the initial threads of all 
+    initial teams that execute the teams region to which the distribute region 
+    binds.
+
+    The distribute loop construct specifies that the iterations of the loop(s)
+    will be executed in parallel by threads in the current context. These
+    iterations are spread across threads that already exist in the enclosing
+    region. The lower and upper bounds specify a half-open range: the
+    range includes the lower bound but does not include the upper bound. If the
+    `inclusive` attribute is specified then the upper bound is also included.
+
+    // TODO:  the following variables will be defined later
+    `private_var`, `firstprivate_var`, and `lastprivate_var` arguments are 
+    variadic list of operands that specify the data sharing attributes of the 
+    list of values 
+
+    The optional `dist_schedule_var` attribute specifies the  schedule for this
+    loop, determining how the loop is distributed across the parallel threads.
+    The optional `schedule_chunk` associated with this determines further
+    controls this distribution.
+
+    The optional `collapse` attribute specifies the number of loops which
+    are collapsed to form the distribute loop.
+ }];
+
+  let arguments = (ins Variadic<IntLikeType>:$lowerBound,
+             Variadic<IntLikeType>:$upperBound,
+             Variadic<IntLikeType>:$step,
+             OptionalAttr<ScheduleKindAttr>:$dist_schedule_var,
+             Optional<AnyType>:$schedule_chunk,
+             Confined<OptionalAttr<I64Attr>, [IntMinValue<1>]>:$collapse_val,
+             Variadic<AnyType>:$allocate_vars,
+             Variadic<AnyType>:$allocators_vars,
+             UnitAttr:$inclusive);
+  
+         
+  let assemblyFormat = [{
+    oilist(`schedule` `(`
+              custom<DistributeScheduleClause>(
+               $dist_schedule_var,
+               $schedule_chunk, type($schedule_chunk)) `)` 
+          |`collapse` `(` $collapse_val `)`
+          | `allocate` `(`
+              custom<AllocateAndAllocator>(
+                $allocate_vars, type($allocate_vars),
+                $allocators_vars, type($allocators_vars)
+              ) `)`
+             
+    ) `for` custom<LoopControl>($region, $lowerBound, $upperBound, $step,
+                                  type($step), $inclusive) attr-dict
+  }];           
+
+  let regions = (region AnyRegion:$region);
+}
+
+//===----------------------------------------------------------------------===//
 // 2.10.1 task Construct
 //===----------------------------------------------------------------------===//
 
 def TaskOp : OpenMP_Op<"task", [AttrSizedOperandSegments,
                        OutlineableOpenMPOpInterface, AutomaticAllocationScope,
                        ReductionClauseInterface]> {
   let summary = "task construct";
   let description = [{
let description = [{
 
     The optional $nowait elliminates the implicit barrier so the parent task can make progress
     even if the target task is not yet completed.
 
     TODO:  map, is_device_ptr, depend, defaultmap, in_reduction
 
   }];
 
   let arguments = (ins Optional<I1>:$if_expr,
                        Optional<AnyInteger>:$device,
                        Optional<AnyInteger>:$thread_limit,
                        UnitAttr:$nowait);
 
   let regions = (region AnyRegion:$region);
 
   let assemblyFormat = [{
     oilist( `if` `(` $if_expr `)`
           | `device` `(` $device `:` type($device) `)`
           | `thread_limit` `(` $thread_limit `:` type($thread_limit) `)`
           | `nowait` $nowait
     ) $region attr-dict
   }];
 }
 
 
 //===----------------------------------------------------------------------===//
 // 2.16 master Construct
 //===----------------------------------------------------------------------===//
 def MasterOp : OpenMP_Op<"master"> {
   let summary = "master construct";

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

 //===- OpenMPDialect.cpp - MLIR Dialect for OpenMP implementation ---------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
if (modifiers.size() == 1) {
     if (symbolizeScheduleModifier(modifiers[0]) == ScheduleModifier::simd ||
         symbolizeScheduleModifier(modifiers[1]) != ScheduleModifier::simd)
       return parser.emitError(parser.getNameLoc())
              << " incorrect modifier order";
   }
   return success();
 }
 
+//===---------------------------------------------------------------------===//
+// Parser for Distribute Schedule Clause 
+//===---------------------------------------------------------------------===//
+
+static ParseResult
+parseDistributeScheduleClause(
+    OpAsmParser &parser, ClauseScheduleKindAttr &scheduleAttr,
+    Optional<OpAsmParser::UnresolvedOperand> &chunkSize, Type &chunkType) {
+    
+    StringRef keyword;
+    if (parser.parseKeyword(&keyword))
+      return failure();
+    llvm::Optional<mlir::omp::ClauseScheduleKind> schedule =
+      symbolizeClauseScheduleKind(keyword);
+    if (!schedule)
+      return parser.emitError(parser.getNameLoc()) <<
+            " expected scheudle kind";
+    
+    scheduleAttr = ClauseScheduleKindAttr::get(parser.getContext(), *schedule);
+    
+    switch (*schedule) {
+    case ClauseScheduleKind::Static:
+    case ClauseScheduleKind::Dynamic:
+    case ClauseScheduleKind::Guided:
+      if (succeeded(parser.parseOptionalEqual())) {
+        chunkSize = OpAsmParser::UnresolvedOperand{};
+        if (parser.parseOperand(*chunkSize) || 
+            parser.parseColonType(chunkType))
+          return failure();
+      } else {
+          chunkSize = llvm::NoneType::None;
+      }   
+      break;
+    case ClauseScheduleKind::Auto:
+    case ClauseScheduleKind::Runtime:
+      chunkSize = llvm::NoneType::None;
+    } 
+        
+    return success();
+}
+
+// Print distribute schedule clause
+static void printDistributeScheduleClause(OpAsmPrinter &p, Operation *op,
+                                ClauseScheduleKindAttr schedAttr,
+                                Value scheduleChunkVar,
+                                Type scheduleChunkType) {
+    p << stringifyClauseScheduleKind(schedAttr.getValue());
+    if (scheduleChunkVar)
+      p << " = " << scheduleChunkVar << " : " << scheduleChunkVar.getType();
+  
+}
+
 /// schedule ::= `schedule` `(` 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
LogicalResult SingleOp::verify() {
   if (allocate_vars().size() != allocators_vars().size())
     return emitError(
         "expected equal sizes for allocate and allocator variables");
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // WsLoopOp
 //===----------------------------------------------------------------------===//
 
 /// loop-control ::= `(` ssa-id-list `)` `:` type `=`  loop-bounds
 /// loop-bounds := `(` ssa-id-list `)` to `(` ssa-id-list `)` inclusive? steps
 /// steps := `step` `(`ssa-id-list`)`
 ParseResult
-parseWsLoopControl(OpAsmParser &parser, Region &region,
+parseLoopControl(OpAsmParser &parser, Region &region,
                    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &lowerBound,
                    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &upperBound,
                    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &steps,
                    SmallVectorImpl<Type> &loopVarTypes, UnitAttr &inclusive) {
   // Parse an opening `(` followed by induction variables followed by `)`
   SmallVector<OpAsmParser::UnresolvedOperand> ivs;
   if (parser.parseRegionArgumentList(ivs, /*requiredOperandCount=*/-1,
                                      OpAsmParser::Delimiter::Paren))
parseLoopControl(OpAsmParser &parser, Region &region,
   // Now parse the body.
   loopVarTypes = SmallVector<Type>(numIVs, loopVarType);
   SmallVector<OpAsmParser::UnresolvedOperand> blockArgs(ivs);
   if (parser.parseRegion(region, blockArgs, loopVarTypes))
     return failure();
   return success();
 }
 
-void printWsLoopControl(OpAsmPrinter &p, Operation *op, Region &region,
+void printLoopControl(OpAsmPrinter &p, Operation *op, Region &region,
                         ValueRange lowerBound, ValueRange upperBound,
                         ValueRange steps, TypeRange loopVarTypes,
                         UnitAttr inclusive) {
   auto args = region.front().getArguments();
   p << " (" << args << ") : " << args[0].getType() << " = (" << lowerBound
     << ") to (" << upperBound << ") ";
   if (inclusive)
     p << "inclusive ";
if (!container.ordered_valAttr() ||
       return emitOpError() << "ordered region must be closely nested inside "
                            << "a worksharing-loop region with an ordered "
                            << "clause without parameter present";
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Verifier for AtomicReadOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult AtomicReadOp::verify() {
   if (auto mo = memory_order_val()) {
     if (*mo == ClauseMemoryOrderKind::Acq_rel ||
         *mo == ClauseMemoryOrderKind::Release) {
       return emitError(
           "memory-order must not be acq_rel or release for atomic reads");
     }
   }
   if (x() == v())
     return emitError(
         "read and write must not be to the same location for atomic reads");
   return verifySynchronizationHint(*this, hint_val());
 }
 
 //===----------------------------------------------------------------------===//
 // Verifier for AtomicWriteOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult AtomicWriteOp::verify() {
   if (auto mo = memory_order_val()) {
     if (*mo == ClauseMemoryOrderKind::Acq_rel ||
         *mo == ClauseMemoryOrderKind::Acquire) {
       return emitError(
           "memory-order must not be acq_rel or acquire for atomic writes");
     }
   }
   return verifySynchronizationHint(*this, hint_val());
 }
 
 //===----------------------------------------------------------------------===//
 // Verifier for AtomicUpdateOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult AtomicUpdateOp::verify() {
   if (auto mo = memory_order_val()) {
     if (*mo == ClauseMemoryOrderKind::Acq_rel ||
         *mo == ClauseMemoryOrderKind::Acquire) {
       return emitError(
           "memory-order must not be acq_rel or acquire for atomic updates");
     }
   }
 
   if (x().getType().cast<PointerLikeType>().getElementType() !=
       region().getArgument(0).getType()) {
     return emitError("the type of the operand must be a pointer type whose "
                      "element type is the same as that of the region argument");
   }
 
   return verifySynchronizationHint(*this, hint_val());
 }
 
 LogicalResult AtomicUpdateOp::verifyRegions() {
   if (region().getNumArguments() != 1)
     return emitError("the region must accept exactly one argument");
 
   if (region().front().getOperations().size() < 2)
     return emitError() << "the update region must have at least two operations "
                           "(binop and terminator)";