[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

265

Nit: Alignment

265

You can remove private, firstprivate and lastprivate for now.

269

Nit: Alignment

287

Unrelated to this patch.

299–300

Unrelated to this patch.

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

799–916 ↗

(On Diff #357051)

These are all unrelated to this patch. Please remove.

mlir/test/Dialect/OpenMP/ops.mlir

299

Removal of this is unrelated to this patch.

[clementval]

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

239

Line should be wrapped at 80.

240

Same here.

250

80

mlir/test/Dialect/OpenMP/ops.mlir

286

Unrelated to this patch

[clementval]

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

269

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

269

Alignment is still off.

[abidmalikwaterloo]

Made changes based on the reviewers' comments.

[abidmalikwaterloo]

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

265

Any special reason for this?

269

We can make it AnyInteger. or IntORIndex

[clementval]

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

272

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

265

Any reason for removing them?

272

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

272

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

272

[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

272

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

272

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

272

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

232

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

232

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

253–255

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

279

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.

293

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 ↗

(On Diff #428955)

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

279

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

253–255

Ok

293

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

265

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

279

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.

293

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

279

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

293

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 ↗

(On Diff #428955)

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

279

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

293

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 ↗

(On Diff #428955)

Yup. We probably don't need these functions.

565 ↗

(On Diff #428955)

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

293

done

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

226 ↗

(On Diff #428955)

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 ↗

(On Diff #441105)

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

219 ↗

(On Diff #441105)

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

215–216

[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

215–216

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

[clementval]

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

215–216

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

255

Do you mean static_dist_schedule here?

260–261

We can remove collapse for now.

278

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 ↗

(On Diff #473706)

220 ↗

(On Diff #473706)

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

129

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

215–216

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

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

211 ↗

(On Diff #473706)

Nit: Please make this change.

220 ↗

(On Diff #473706)

Nit: Please make this change as well.

mlir/test/Dialect/OpenMP/ops.mlir

129

Nit: please add a pretty-print test.

[abidmalikwaterloo]

mlir/test/Dialect/OpenMP/ops.mlir

129

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

129

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

def WsLoopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments,
   let extraClassDeclaration = [{
     /// Returns the number of loops in the workshape loop nest.
     unsigned getNumLoops() { return lowerBound().size(); }
   }];
   let parser = [{ return parseWsLoopOp(parser, result); }];
   let printer = [{ return printWsLoopOp(p, *this); }];
 }
 
 def YieldOp : OpenMP_Op<"yield", [NoSideEffect, ReturnLike, Terminator,
                                   HasParent<"WsLoopOp">]> {
   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.
     If "omp.yield" has any operands, the operands must match the parent
     operation's results.
   }];
 
   let arguments = (ins Variadic<AnyType>:$results);
 
   let assemblyFormat = [{ ( `(` $results^ `:` type($results) `)` )? attr-dict}];
 }
 
+//===---------------------------------------------------------------------===//
+// 2.9.4 distribute loop 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.
+
+
+    `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,
+             Variadic<AnyType>:$private_var,
+             Variadic<AnyType>:$firstprivate_var,
+             Variadic<AnyType>:$lastprivate_var,
+             OptionalAttr<ScheduleKind>:$dist_schedule_var,
+             Optional<AnyType>:$schedule_chunk,
+             Confined<OptionalAttr<I64Attr>, [IntMinValue<1>]>:$collapse,
+             Variadic<AnyType>:$allocate,
+             Variadic<AnyType>:$allocators,
+             UnitAttr:$inclusive);
+
+  let regions = (region AnyRegion:$region);
+}
+
 //===----------------------------------------------------------------------===//
 // 2.10.4 taskyield Construct
 //===----------------------------------------------------------------------===//
 
 def TaskyieldOp : OpenMP_Op<"taskyield"> {
   let summary = "taskyield construct";
   let description = [{
     The taskyield construct specifies that the current task can be suspended
     in favor of execution of a different task.
   }];
 
   let assemblyFormat = "attr-dict";
 }
 
 //===----------------------------------------------------------------------===//
 // 2.13.7 flush Construct
 //===----------------------------------------------------------------------===//
 def FlushOp : OpenMP_Op<"flush"> {
   let summary = "flush construct";
   let description = [{
     The flush construct executes the OpenMP flush operation. This operation
     makes a thread’s temporary view of memory consistent with memory 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:  private, map, is_device_ptr, firstprivate, 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 builders = [
        OpBuilder<(ins "IntegerAttr":$if_expr,
                        "IntegerAttr":$device,
                        "IntegerAttr":$thread_limit,
                        "UnitAttr":$nowait)>
   ];
 }
 
 
 //===----------------------------------------------------------------------===//
 // 2.16 master Construct
 //===----------------------------------------------------------------------===//
 def MasterOp : OpenMP_Op<"master"> {
   let summary = "master construct";
   let description = [{

mlir/test/Dialect/OpenMP/ops.mlir

  omp.parallel num_threads(%num_threads : si32) if(%if_cond: i1)
               private(%data_var : memref<i32>) proc_bind(close) {
    omp.terminator
  }
 
  return
 }
 
 // CHECK-LABEL: omp_wsloop
 func @omp_wsloop(%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>, %{{.*}} : memref<i32>) collapse(2) ordered(1)
   "omp.wsloop" (%lb, %ub, %step, %data_var, %data_var) ({
     ^bb0(%iv: index):
       omp.yield
   }) {operand_segment_sizes = dense<[1,1,1,2,0,0,0,0,0]> : vector<9xi32>, collapse_val = 2, ordered_val = 1} :
     (index, index, index, memref<i32>, memref<i32>) -> ()
omp.wsloop (%iv1, %iv2) : i32 = (%lb1, %lb2) to (%ub1, %ub2) step (%step1, %step2) {
     omp.yield
   }
 
   return
 }
 
 
 // CHECK-LABEL: omp_target
 func @omp_target(%if_cond : i1, %device : si32,  %num_threads : si32) -> () {
 
     // Test with optional operands; if_expr, device, thread_limit, and nowait.
     // CHECK: omp.target
     "omp.target"(%if_cond, %device, %num_threads) ({
        // CHECK: omp.terminator
        omp.terminator
     }) {operand_segment_sizes = dense<[1,1,1]>: vector<3xi32>, nowait } : ( i1, si32, si32  ) -> ()
 
     // CHECK: omp.barrier
     omp.barrier
 
     return
 }
+
+// CHECK-LABEL: omp_distribute
+func @omp_distribute(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %chunk_var : si32) -> () {
+
+    // CHECK: omp.distribute
+    "omp.distribute" (%lb, %ub, %step, %data_var,%data_var, %data_var, %data_var, %data_var, %chunk_var) ({
+  }) {operand_segment_sizes = dense<[1,1,1,1,1,1,1,1,1]> : vector<9xi32>, dist_schedule_val = "Static", collapse_val = 3} : (index, index, index, memref<i32>, memref<i32>, memref<i32>, memref<i32>, memref<i32>, si32) -> ()
+
+    // CHECK: omp.barrier
+   omp.barrier
+
+  return
+
+}