-
Notifications
You must be signed in to change notification settings - Fork 14.3k
[MLIR][omp] Add omp.workshare op #101443
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[MLIR][omp] Add omp.workshare op #101443
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for your work on this. Please could you also add a parse/print test in mlir/test/Dialect/OpenMP/ops.mlir
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Please add some roundtrip tests in mlir/test/Dialect/OpenMP/ops.mlir.
The questions here would be:
- Should we have an alternative representation like
hlfir.elemental
that exposes the implicit loop of the constructs in it. - Should enforce additional constraints on the operations inside the body of the workshare construct? Like should they have loop interfaces?
cc551ca
to
7376fd4
Compare
@llvm/pr-subscribers-flang-driver @llvm/pr-subscribers-flang-openmp Author: Ivan R. Ivanov (ivanradanov) Changes1/4 in stack for workshare implementation 1/4 #101443 Full diff: https://github.com/llvm/llvm-project/pull/101443.diff 5 Files Affected:
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h
index 38e4d8f245e4fa..896ca9581c3fc8 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h
@@ -316,6 +316,8 @@ using TeamsOperands =
detail::Clauses<AllocateClauseOps, IfClauseOps, NumTeamsClauseOps,
PrivateClauseOps, ReductionClauseOps, ThreadLimitClauseOps>;
+using WorkshareOperands = detail::Clauses<NowaitClauseOps>;
+
using WsloopOperands =
detail::Clauses<AllocateClauseOps, LinearClauseOps, NowaitClauseOps,
OrderClauseOps, OrderedClauseOps, PrivateClauseOps,
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 68f92e6952694b..74497def8fd1ab 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -286,6 +286,49 @@ def SingleOp : OpenMP_Op<"single", traits = [
let hasVerifier = 1;
}
+//===----------------------------------------------------------------------===//
+// 2.8.3 Workshare Construct
+//===----------------------------------------------------------------------===//
+
+def WorkshareOp : OpenMP_Op<"workshare", traits = [
+ RecursiveMemoryEffects,
+ ], clauses = [
+ OpenMP_NowaitClause,
+ ], singleRegion = true> {
+ let summary = "workshare directive";
+ let description = [{
+ The workshare construct divides the execution of the enclosed structured
+ block into separate units of work, and causes the threads of the team to
+ share the work such that each unit is executed only once by one thread, in
+ the context of its implicit task
+
+ This operation is used for the intermediate representation of the workshare
+ block before the work gets divided between the threads. See the flang
+ LowerWorkshare pass for details.
+ }] # clausesDescription;
+
+ let builders = [
+ OpBuilder<(ins CArg<"const WorkshareOperands &">:$clauses)>
+ ];
+}
+
+def WorkshareLoopWrapperOp : OpenMP_Op<"workshare_loop_wrapper", traits = [
+ DeclareOpInterfaceMethods<LoopWrapperInterface>,
+ RecursiveMemoryEffects, SingleBlock
+ ], singleRegion = true> {
+ let summary = "contains loop nests to be parallelized by workshare";
+ let description = [{
+ This operation wraps a loop nest that is marked for dividing into units of
+ work by an encompassing omp.workshare operation.
+ }];
+
+ let builders = [
+ OpBuilder<(ins), [{ build($_builder, $_state, {}); }]>
+ ];
+ let assemblyFormat = "$region attr-dict";
+ let hasVerifier = 1;
+}
+
//===----------------------------------------------------------------------===//
// Loop Nest
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 11780f84697b15..90f9a19ebe32b5 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1683,6 +1683,29 @@ LogicalResult SingleOp::verify() {
getCopyprivateSyms());
}
+//===----------------------------------------------------------------------===//
+// WorkshareOp
+//===----------------------------------------------------------------------===//
+
+void WorkshareOp::build(OpBuilder &builder, OperationState &state,
+ const WorkshareOperands &clauses) {
+ WorkshareOp::build(builder, state, clauses.nowait);
+}
+
+//===----------------------------------------------------------------------===//
+// WorkshareLoopWrapperOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult WorkshareLoopWrapperOp::verify() {
+ if (!isWrapper())
+ return emitOpError() << "must be a loop wrapper";
+ if (getNestedWrapper())
+ return emitError() << "nested wrappers not supported";
+ if (!(*this)->getParentOfType<WorkshareOp>())
+ return emitError() << "must be nested in an omp.workshare";
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// WsloopOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 1d1d93f0977588..ee7c448c467cf5 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -2383,3 +2383,45 @@ func.func @masked_arg_count_mismatch(%arg0: i32, %arg1: i32) {
}) : (i32, i32) -> ()
return
}
+
+// -----
+func.func @nested_wrapper(%idx : index) {
+ omp.workshare {
+ // expected-error @below {{nested wrappers not supported}}
+ omp.workshare_loop_wrapper {
+ omp.simd {
+ omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) {
+ omp.yield
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
+
+// -----
+func.func @not_wrapper() {
+ omp.workshare {
+ // expected-error @below {{must be a loop wrapper}}
+ omp.workshare_loop_wrapper {
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
+
+// -----
+func.func @missing_workshare(%idx : index) {
+ // expected-error @below {{must be nested in an omp.workshare}}
+ omp.workshare_loop_wrapper {
+ omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) {
+ omp.yield
+ }
+ omp.terminator
+ }
+ return
+}
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index d2924998f41b87..24363334b4cb54 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -2789,3 +2789,72 @@ func.func @omp_target_private(%map1: memref<?xi32>, %map2: memref<?xi32>, %priv_
return
}
+
+// CHECK-LABEL: func @omp_workshare
+func.func @omp_workshare() {
+ // CHECK: omp.workshare {
+ omp.workshare {
+ "test.payload"() : () -> ()
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+ return
+}
+
+// CHECK-LABEL: func @omp_workshare_nowait
+func.func @omp_workshare_nowait() {
+ // CHECK: omp.workshare nowait {
+ omp.workshare nowait {
+ "test.payload"() : () -> ()
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+ return
+}
+
+// CHECK-LABEL: func @omp_workshare_multiple_blocks
+func.func @omp_workshare_multiple_blocks() {
+ // CHECK: omp.workshare {
+ omp.workshare {
+ cf.br ^bb2
+ ^bb2:
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+ return
+}
+
+// CHECK-LABEL: func @omp_workshare_loop_wrapper
+func.func @omp_workshare_loop_wrapper(%idx : index) {
+ // CHECK-NEXT: omp.workshare {
+ omp.workshare {
+ // CHECK-NEXT: omp.workshare_loop_wrapper
+ omp.workshare_loop_wrapper {
+ // CHECK-NEXT: omp.loop_nest
+ omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) {
+ omp.yield
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
+
+// CHECK-LABEL: func @omp_workshare_loop_wrapper_attrs
+func.func @omp_workshare_loop_wrapper_attrs(%idx : index) {
+ // CHECK-NEXT: omp.workshare {
+ omp.workshare {
+ // CHECK-NEXT: omp.workshare_loop_wrapper {
+ omp.workshare_loop_wrapper {
+ // CHECK-NEXT: omp.loop_nest
+ omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) {
+ omp.yield
+ }
+ omp.terminator
+ // CHECK: } {attr_in_dict}
+ } {attr_in_dict}
+ omp.terminator
+ }
+ return
+}
|
@llvm/pr-subscribers-mlir-openmp Author: Ivan R. Ivanov (ivanradanov) Changes1/4 in stack for workshare implementation 1/4 #101443 Full diff: https://github.com/llvm/llvm-project/pull/101443.diff 5 Files Affected:
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h
index 38e4d8f245e4fa..896ca9581c3fc8 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h
@@ -316,6 +316,8 @@ using TeamsOperands =
detail::Clauses<AllocateClauseOps, IfClauseOps, NumTeamsClauseOps,
PrivateClauseOps, ReductionClauseOps, ThreadLimitClauseOps>;
+using WorkshareOperands = detail::Clauses<NowaitClauseOps>;
+
using WsloopOperands =
detail::Clauses<AllocateClauseOps, LinearClauseOps, NowaitClauseOps,
OrderClauseOps, OrderedClauseOps, PrivateClauseOps,
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 68f92e6952694b..74497def8fd1ab 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -286,6 +286,49 @@ def SingleOp : OpenMP_Op<"single", traits = [
let hasVerifier = 1;
}
+//===----------------------------------------------------------------------===//
+// 2.8.3 Workshare Construct
+//===----------------------------------------------------------------------===//
+
+def WorkshareOp : OpenMP_Op<"workshare", traits = [
+ RecursiveMemoryEffects,
+ ], clauses = [
+ OpenMP_NowaitClause,
+ ], singleRegion = true> {
+ let summary = "workshare directive";
+ let description = [{
+ The workshare construct divides the execution of the enclosed structured
+ block into separate units of work, and causes the threads of the team to
+ share the work such that each unit is executed only once by one thread, in
+ the context of its implicit task
+
+ This operation is used for the intermediate representation of the workshare
+ block before the work gets divided between the threads. See the flang
+ LowerWorkshare pass for details.
+ }] # clausesDescription;
+
+ let builders = [
+ OpBuilder<(ins CArg<"const WorkshareOperands &">:$clauses)>
+ ];
+}
+
+def WorkshareLoopWrapperOp : OpenMP_Op<"workshare_loop_wrapper", traits = [
+ DeclareOpInterfaceMethods<LoopWrapperInterface>,
+ RecursiveMemoryEffects, SingleBlock
+ ], singleRegion = true> {
+ let summary = "contains loop nests to be parallelized by workshare";
+ let description = [{
+ This operation wraps a loop nest that is marked for dividing into units of
+ work by an encompassing omp.workshare operation.
+ }];
+
+ let builders = [
+ OpBuilder<(ins), [{ build($_builder, $_state, {}); }]>
+ ];
+ let assemblyFormat = "$region attr-dict";
+ let hasVerifier = 1;
+}
+
//===----------------------------------------------------------------------===//
// Loop Nest
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 11780f84697b15..90f9a19ebe32b5 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1683,6 +1683,29 @@ LogicalResult SingleOp::verify() {
getCopyprivateSyms());
}
+//===----------------------------------------------------------------------===//
+// WorkshareOp
+//===----------------------------------------------------------------------===//
+
+void WorkshareOp::build(OpBuilder &builder, OperationState &state,
+ const WorkshareOperands &clauses) {
+ WorkshareOp::build(builder, state, clauses.nowait);
+}
+
+//===----------------------------------------------------------------------===//
+// WorkshareLoopWrapperOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult WorkshareLoopWrapperOp::verify() {
+ if (!isWrapper())
+ return emitOpError() << "must be a loop wrapper";
+ if (getNestedWrapper())
+ return emitError() << "nested wrappers not supported";
+ if (!(*this)->getParentOfType<WorkshareOp>())
+ return emitError() << "must be nested in an omp.workshare";
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// WsloopOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 1d1d93f0977588..ee7c448c467cf5 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -2383,3 +2383,45 @@ func.func @masked_arg_count_mismatch(%arg0: i32, %arg1: i32) {
}) : (i32, i32) -> ()
return
}
+
+// -----
+func.func @nested_wrapper(%idx : index) {
+ omp.workshare {
+ // expected-error @below {{nested wrappers not supported}}
+ omp.workshare_loop_wrapper {
+ omp.simd {
+ omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) {
+ omp.yield
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
+
+// -----
+func.func @not_wrapper() {
+ omp.workshare {
+ // expected-error @below {{must be a loop wrapper}}
+ omp.workshare_loop_wrapper {
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
+
+// -----
+func.func @missing_workshare(%idx : index) {
+ // expected-error @below {{must be nested in an omp.workshare}}
+ omp.workshare_loop_wrapper {
+ omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) {
+ omp.yield
+ }
+ omp.terminator
+ }
+ return
+}
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index d2924998f41b87..24363334b4cb54 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -2789,3 +2789,72 @@ func.func @omp_target_private(%map1: memref<?xi32>, %map2: memref<?xi32>, %priv_
return
}
+
+// CHECK-LABEL: func @omp_workshare
+func.func @omp_workshare() {
+ // CHECK: omp.workshare {
+ omp.workshare {
+ "test.payload"() : () -> ()
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+ return
+}
+
+// CHECK-LABEL: func @omp_workshare_nowait
+func.func @omp_workshare_nowait() {
+ // CHECK: omp.workshare nowait {
+ omp.workshare nowait {
+ "test.payload"() : () -> ()
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+ return
+}
+
+// CHECK-LABEL: func @omp_workshare_multiple_blocks
+func.func @omp_workshare_multiple_blocks() {
+ // CHECK: omp.workshare {
+ omp.workshare {
+ cf.br ^bb2
+ ^bb2:
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+ return
+}
+
+// CHECK-LABEL: func @omp_workshare_loop_wrapper
+func.func @omp_workshare_loop_wrapper(%idx : index) {
+ // CHECK-NEXT: omp.workshare {
+ omp.workshare {
+ // CHECK-NEXT: omp.workshare_loop_wrapper
+ omp.workshare_loop_wrapper {
+ // CHECK-NEXT: omp.loop_nest
+ omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) {
+ omp.yield
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
+
+// CHECK-LABEL: func @omp_workshare_loop_wrapper_attrs
+func.func @omp_workshare_loop_wrapper_attrs(%idx : index) {
+ // CHECK-NEXT: omp.workshare {
+ omp.workshare {
+ // CHECK-NEXT: omp.workshare_loop_wrapper {
+ omp.workshare_loop_wrapper {
+ // CHECK-NEXT: omp.loop_nest
+ omp.loop_nest (%iv) : index = (%idx) to (%idx) step (%idx) {
+ omp.yield
+ }
+ omp.terminator
+ // CHECK: } {attr_in_dict}
+ } {attr_in_dict}
+ omp.terminator
+ }
+ return
+}
|
ce3ad06
to
7065222
Compare
d030f24
to
72e19fa
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM. Thanks for the updates
72e19fa
to
45a8aa4
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thank you, this LGTM as well. Just a small nit about the name for WorkshareLoopWrapperOp
.
8503797
to
d343f3a
Compare
45a5069
to
f8d7b47
Compare
Add custom omp loop wrapper Add recursive memory effects trait to workshare Remove stray include Remove omp.workshare verifier Add assembly format for wrapper and add test Add verification and descriptions wrong replace Fix wsloopwrapperop Fix op tests
d001eec
to
fa6e0dd
Compare
1/4 in stack for workshare implementation
1/4 #101443
2/4 #101444
3/4 #101445
4/4 #101446