[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

506

Nit: Alignment

506

You can remove private, firstprivate and lastprivate for now.

510

Nit: Alignment

840

Unrelated to this patch.

845–853

Unrelated to this patch.

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

881–1012

These are all unrelated to this patch. Please remove.

mlir/test/Dialect/OpenMP/ops.mlir

468–526

Removal of this is unrelated to this patch.

[clementval]

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

480

Line should be wrapped at 80.

481

Same here.

491

80

mlir/test/Dialect/OpenMP/ops.mlir

463

Unrelated to this patch

[clementval]

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

510

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

510

Alignment is still off.

[abidmalikwaterloo]

Made changes based on the reviewers' comments.

[abidmalikwaterloo]

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

506

Any special reason for this?

510

We can make it AnyInteger. or IntORIndex

[clementval]

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

513

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

506

Any reason for removing them?

513

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

513

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

513

[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

513

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

513

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

513

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

473

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

473

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 ↗

(On Diff #428955)

unrelated change?

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

494–496

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

520

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.

534

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

300

nit: schedule spelling.

[abidmalikwaterloo]

clang/lib/Testing/CMakeLists.txt

29 ↗

(On Diff #428955)

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

520

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

494–496

Ok

534

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

[shraiysh]

clang/lib/Testing/CMakeLists.txt

29 ↗

(On Diff #428955)

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

506

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

520

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.

534

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 ↗

(On Diff #428955)

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

520

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

534

At this stage, we can add:

LogicalResult DistributeOp::verify(){

return success();

}

We can add the check later!

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

300

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

[shraiysh]

clang/lib/Testing/CMakeLists.txt

29 ↗

(On Diff #428955)

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

520

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

534

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

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

300

Yup. We probably don't need these functions.

613

Change this comment to Loop control.

[abidmalikwaterloo]

clang/lib/Testing/CMakeLists.txt

29 ↗

(On Diff #428955)

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

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

534

done

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

300

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

289

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

293

no brace

[abidmalikwaterloo]

Updated DistributeOp::verify as per reviewer's comments

[abidmalikwaterloo]

update the changes

[clementval]

clang/lib/Testing/CMakeLists.txt

30 ↗

(On Diff #442321)

You still have these changes that are unrelated.

[abidmalikwaterloo]

Removed unrelated changes in CMakefile.txt

[clementval]

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

455

[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

455

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

[clementval]

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

455

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

496

Do you mean static_dist_schedule here?

501–502

We can remove collapse for now.

519

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

285

294

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

136

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

455

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

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

285

Nit: Please make this change.

294

Nit: Please make this change as well.

mlir/test/Dialect/OpenMP/ops.mlir

136

Nit: please add a pretty-print test.

[abidmalikwaterloo]

mlir/test/Dialect/OpenMP/ops.mlir

136

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

136

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.

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

let summary = "simd loop construct";
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 
 def YieldOp : OpenMP_Op<"yield",
     [Pure, ReturnLike, Terminator,
      ParentOneOf<["WsLoopOp", "ReductionDeclareOp",
-     "AtomicUpdateOp", "SimdLoopOp"]>]> {
+     "AtomicUpdateOp", "SimdLoopOp", "DistributeOp"]>]> {
   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:  private_var, firstprivate_var and lastprivate_var variables 
+    // will be defined later
+    
+    The `static_dist_schedule` 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.
+
+    // TODO: collapse
+  }];
+
+  let arguments = (ins Variadic<IntLikeType>:$lowerBound,
+             Variadic<IntLikeType>:$upperBound,
+             Variadic<IntLikeType>:$step,
+             UnitAttr:$static_dist_schedule,
+             Optional<IntLikeType>:$schedule_chunk,
+             Variadic<AnyType>:$allocate_vars,
+             Variadic<AnyType>:$allocators_vars,
+             UnitAttr:$inclusive);
+           
+  let assemblyFormat = [{
+    oilist(`dist_schedule` `(` `static`
+        ($static_dist_schedule^)? 
+        ( $schedule_chunk^ `:` type($schedule_chunk))? `)`
+          | `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);
+  let hasVerifier = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // 2.10.1 task Construct
 //===----------------------------------------------------------------------===//
 
 def TaskOp : OpenMP_Op<"task", [AttrSizedOperandSegments,
                        OutlineableOpenMPOpInterface, AutomaticAllocationScope,
                        ReductionClauseInterface]> {
   let summary = "task construct";
   let description = [{
     The task construct defines an explicit task.
 
     For definitions of "undeferred task", "included task", "final task" and
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

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();
 }
 
+//===---------------------------------------------------------------------===//
+// Verifier for Dristribute Op
+//===---------------------------------------------------------------------===//
+
+LogicalResult DistributeOp::verify(){
+  if (this->getLowerBound().empty()) 
+    return emitOpError() << "empty lowerbound for distribute loop operation";
+  
+  if (this->getUpperBound().empty()) 
+    return emitOpError() << "empty upperbound for distribute loop operation";
+  
+  if (this->getLowerBound().size() != this->getUpperBound().size())
+    return emitOpError() << "upperbound and lowerbound sizes are not equal";
+    
+  if (this->getLowerBound().size() != this->getStep().size())
+    return emitOpError() << "upperbound, lowerbound and step sizes are not equal";
+      
+  return success(); 
+}
+
 /// 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 (getAllocateVars().size() != getAllocatorsVars().size())
     return emitError(
         "expected equal sizes for allocate and allocator variables");
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
-// WsLoopOp
+// LoopControl
 //===----------------------------------------------------------------------===//
 
 /// 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
 parseLoopControl(OpAsmParser &parser, Region &region,
                  SmallVectorImpl<OpAsmParser::UnresolvedOperand> &lowerBound,
if (!container.getOrderedValAttr() ||
       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 = getMemoryOrderVal()) {
     if (*mo == ClauseMemoryOrderKind::Acq_rel ||
         *mo == ClauseMemoryOrderKind::Release) {
       return emitError(
           "memory-order must not be acq_rel or release for atomic reads");
     }
   }
   if (getX() == getV())
     return emitError(
         "read and write must not be to the same location for atomic reads");
   return verifySynchronizationHint(*this, getHintVal());
 }
 
 //===----------------------------------------------------------------------===//
 // Verifier for AtomicWriteOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult AtomicWriteOp::verify() {
   if (auto mo = getMemoryOrderVal()) {
     if (*mo == ClauseMemoryOrderKind::Acq_rel ||
         *mo == ClauseMemoryOrderKind::Acquire) {
       return emitError(
           "memory-order must not be acq_rel or acquire for atomic writes");
     }
   }
   if (getAddress().getType().cast<PointerLikeType>().getElementType() !=
       getValue().getType())
     return emitError("address must dereference to value type");
   return verifySynchronizationHint(*this, getHintVal());
 }
 
 //===----------------------------------------------------------------------===//
 // Verifier for AtomicUpdateOp
 //===----------------------------------------------------------------------===//
 
 bool AtomicUpdateOp::isNoOp() {
   YieldOp yieldOp = dyn_cast<omp::YieldOp>(getFirstOp());
   return (yieldOp &&
           yieldOp.getResults().front() == getRegion().front().getArgument(0));
 }
 
 Value AtomicUpdateOp::getWriteOpVal() {
   YieldOp yieldOp = dyn_cast<omp::YieldOp>(getFirstOp());
   if (yieldOp &&
       yieldOp.getResults().front() != getRegion().front().getArgument(0))
     return yieldOp.getResults().front();
   return nullptr;
 }
 
 LogicalResult AtomicUpdateOp::canonicalize(AtomicUpdateOp op,
                                            PatternRewriter &rewriter) {
   if (op.isNoOp()) {
     rewriter.eraseOp(op);
     return success();
   }
   if (Value writeVal = op.getWriteOpVal()) {
     rewriter.replaceOpWithNewOp<AtomicWriteOp>(op, op.getX(), writeVal,
                                                op.getHintValAttr(),
                                                op.getMemoryOrderValAttr());
     return success();
   }
   return failure();
 }
 
 LogicalResult AtomicUpdateOp::verify() {
   if (auto mo = getMemoryOrderVal()) {
     if (*mo == ClauseMemoryOrderKind::Acq_rel ||
         *mo == ClauseMemoryOrderKind::Acquire) {
       return emitError(
           "memory-order must not be acq_rel or acquire for atomic updates");
     }
   }
 
   if (getRegion().getNumArguments() != 1)
     return emitError("the region must accept exactly one argument");
 
   if (getX().getType().cast<PointerLikeType>().getElementType() !=
       getRegion().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, getHintVal());
 }
 
 LogicalResult AtomicUpdateOp::verifyRegions() {
 
   YieldOp yieldOp = *getRegion().getOps<YieldOp>().begin();
 
   if (yieldOp.getResults().size() != 1)
     return emitError("only updated value must be returned");
   if (yieldOp.getResults().front().getType() !=
       getRegion().getArgument(0).getType())
     return emitError("input and yielded value must have the same type");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Verifier for AtomicCaptureOp
 //===----------------------------------------------------------------------===//
 
 Operation *AtomicCaptureOp::getFirstOp() {
   return &getRegion().front().getOperations().front();
 }
 
 Operation *AtomicCaptureOp::getSecondOp() {
   auto &ops = getRegion().front().getOperations();
   return ops.getNextNode(ops.front());
 }
 
 AtomicReadOp AtomicCaptureOp::getAtomicReadOp() {
   if (auto op = dyn_cast<AtomicReadOp>(getFirstOp()))
     return op;
   return dyn_cast<AtomicReadOp>(getSecondOp());
 }
 
 AtomicWriteOp AtomicCaptureOp::getAtomicWriteOp() {
   if (auto op = dyn_cast<AtomicWriteOp>(getFirstOp()))
     return op;
   return dyn_cast<AtomicWriteOp>(getSecondOp());
 }
 
 AtomicUpdateOp AtomicCaptureOp::getAtomicUpdateOp() {
   if (auto op = dyn_cast<AtomicUpdateOp>(getFirstOp()))
     return op;
   return dyn_cast<AtomicUpdateOp>(getSecondOp());
 }
 
 LogicalResult AtomicCaptureOp::verify() {
   return verifySynchronizationHint(*this, getHintVal());
 }
 
 LogicalResult AtomicCaptureOp::verifyRegions() {
   Block::OpListType &ops = getRegion().front().getOperations();
   if (ops.size() != 3)

mlir/test/Dialect/OpenMP/ops.mlir

  // CHECK omp.parallel if(%{{.*}}) num_threads(%{{.*}} : i32) private(%{{.*}} : memref<i32>) proc_bind(close)
  omp.parallel num_threads(%num_threads : i32) if(%if_cond: i1) proc_bind(close) {
    omp.terminator
  }
 
   return
 }
 
+// CHECK-LABEL: omp_DistributeOp
+func.func @omp_DistributeOp(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %chunk_var : i32) -> () {
+// CHECK: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+  "omp.distribute" (%lb, %ub, %step) ({
+    ^bb0(%iv: index):
+     omp.yield 
+  }) {operand_segment_sizes = array<i32: 1,1,1,0,0,0>} :
+    (index, index, index) -> ()
+ 
+ return
+ }
+ 
+// CHECK-LABEL: omp_distribute_pretty
+ func.func @omp_distribute_pretty(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %chunk_var : i32) -> () {
+// CHECK: omp.distribute
+// CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+  omp.distribute 
+  for (%iv) : index = (%lb) to (%ub) step (%step) {
+    omp.yield
+  }
+
+ return
+ }
+ 
 // CHECK-LABEL: omp_wsloop
 func.func @omp_wsloop(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %linear_var : i32, %chunk_var : i32) -> () {
 
   // CHECK: omp.wsloop ordered(1)
   // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
   "omp.wsloop" (%lb, %ub, %step) ({
     ^bb0(%iv: index):
       omp.yield
   }) {operand_segment_sizes = array<i32: 1,1,1,0,0,0,0>, ordered_val = 1} :
     (index, index, index) -> ()
 
   // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static)
-  // CHECK-SAMe: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+  // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
   "omp.wsloop" (%lb, %ub, %step, %data_var, %linear_var) ({
     ^bb0(%iv: index):
       omp.yield
   }) {operand_segment_sizes = array<i32: 1,1,1,1,1,0,0>, schedule_val = #omp<schedulekind static>} :
     (index, index, index, memref<i32>, i32) -> ()
 
   // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>, %{{.*}} = %{{.*}} : memref<i32>) schedule(static)
   // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
func.func @omp_simdloop_pretty_multiple(%lb1 : index, %ub1 : index, %step1 : index, %lb2 : index, %ub2 : index, %step2 : index) -> () {
   // CHECK: omp.simdloop for (%{{.*}}, %{{.*}}) : index = (%{{.*}}, %{{.*}}) to (%{{.*}}, %{{.*}}) step (%{{.*}}, %{{.*}})
   omp.simdloop for (%iv1, %iv2) : index = (%lb1, %lb2) to (%ub1, %ub2) step (%step1, %step2) {
     omp.yield
   }
   return
 }
 
 // CHECK-LABEL: omp_target
 func.func @omp_target(%if_cond : i1, %device : si32,  %num_threads : i32) -> () {
 
     // Test with optional operands; if_expr, device, thread_limit, private, firstprivate and nowait.
     // CHECK: omp.target if({{.*}}) device({{.*}}) thread_limit({{.*}}) nowait
     "omp.target"(%if_cond, %device, %num_threads) ({
        // CHECK: omp.terminator
        omp.terminator
     }) {nowait, operand_segment_sizes = array<i32: 1,1,1>} : ( i1, si32, i32 ) -> ()
 
func.func @omp_target_pretty(%if_cond : i1, %device : si32,  %num_threads : i32) -> () {
 
     // CHECK: omp.target if({{.*}}) device({{.*}}) nowait
     omp.target if(%if_cond) device(%device : si32) thread_limit(%num_threads : i32) nowait {
       omp.terminator
     }
 
     return
 }
 
 // CHECK: omp.reduction.declare
 // CHECK-LABEL: @add_f32
 // CHECK: : f32
 // CHECK: init
 // CHECK: ^{{.+}}(%{{.+}}: f32):
 // CHECK:   omp.yield
 // CHECK: combiner
 // CHECK: ^{{.+}}(%{{.+}}: f32, %{{.+}}: f32):
 // CHECK:   omp.yield
 // CHECK: atomic
 // CHECK: ^{{.+}}(%{{.+}}: !llvm.ptr<f32>, %{{.+}}: !llvm.ptr<f32>):
 // CHECK:  omp.yield
 omp.reduction.declare @add_f32 : f32
 init {
 ^bb0(%arg: f32):
   %0 = arith.constant 0.0 : f32
   omp.yield (%0 : f32)
 }
 combiner {
 ^bb1(%arg0: f32, %arg1: f32):
   %1 = arith.addf %arg0, %arg1 : f32
   omp.yield (%1 : f32)
 }
 atomic {
 ^bb2(%arg2: !llvm.ptr<f32>, %arg3: !llvm.ptr<f32>):
   %2 = llvm.load %arg3 : !llvm.ptr<f32>
   llvm.atomicrmw fadd %arg2, %2 monotonic : f32
   omp.yield
 }
 
 // CHECK-LABEL: func @wsloop_reduction
 func.func @wsloop_reduction(%lb : index, %ub : index, %step : index) {
   %c1 = arith.constant 1 : i32
   %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<f32>
   // CHECK: reduction(@add_f32 -> %{{.+}} : !llvm.ptr<f32>)
   omp.wsloop reduction(@add_f32 -> %0 : !llvm.ptr<f32>)
   for (%iv) : index = (%lb) to (%ub) step (%step) {
     %1 = arith.constant 2.0 : f32
     // CHECK: omp.reduction %{{.+}}, %{{.+}}
     omp.reduction %1, %0 : !llvm.ptr<f32>
     omp.yield
   }
   return
 }
 
 // CHECK-LABEL: func @parallel_reduction
 func.func @parallel_reduction() {
   %c1 = arith.constant 1 : i32
   %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<f32>
   // CHECK: omp.parallel reduction(@add_f32 -> {{.+}} : !llvm.ptr<f32>)
   omp.parallel reduction(@add_f32 -> %0 : !llvm.ptr<f32>) {
     %1 = arith.constant 2.0 : f32
     // CHECK: omp.reduction %{{.+}}, %{{.+}}
     omp.reduction %1, %0 : !llvm.ptr<f32>
     omp.terminator
   }
   return
 }
 
 // CHECK: func @parallel_wsloop_reduction
 func.func @parallel_wsloop_reduction(%lb : index, %ub : index, %step : index) {
   %c1 = arith.constant 1 : i32
   %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<f32>
   // CHECK: omp.parallel reduction(@add_f32 -> %{{.+}} : !llvm.ptr<f32>) {
   omp.parallel reduction(@add_f32 -> %0 : !llvm.ptr<f32>) {
     // CHECK: omp.wsloop for (%{{.+}}) : index = (%{{.+}}) to (%{{.+}}) step (%{{.+}})