[MLIR][OpenMP] Distribute Construct Operation

[abidmalikwaterloo]

Summary

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

[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.

527

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

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

Nit: please add a pretty-print test.

[abidmalikwaterloo]

mlir/test/Dialect/OpenMP/ops.mlir

124

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

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.