[llvm-branch-commits] [mlir] [MLIR][omp] Add omp.workshare op (PR #101443)
Ivan R. Ivanov via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Sun Aug 18 23:22:11 PDT 2024
https://github.com/ivanradanov updated https://github.com/llvm/llvm-project/pull/101443
>From e453738aaa1bd74e2462025075ea4ac12868caac Mon Sep 17 00:00:00 2001
From: Ivan Radanov Ivanov <ivanov.i.aa at m.titech.ac.jp>
Date: Wed, 31 Jul 2024 14:09:09 +0900
Subject: [PATCH 1/7] [MLIR][omp] Add omp.workshare op
---
.../Dialect/OpenMP/OpenMPClauseOperands.h | 3 +++
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 22 +++++++++++++++++++
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 13 +++++++++++
3 files changed, 38 insertions(+)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h
index 38e4d8f245e4fa..d14e5e17afbb08 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h
@@ -17,6 +17,7 @@
#include "mlir/IR/BuiltinAttributes.h"
#include "llvm/ADT/SmallVector.h"
+#include <mlir/Dialect/OpenMP/OpenMPDialect.h>
#include "mlir/Dialect/OpenMP/OpenMPOpsEnums.h.inc"
@@ -316,6 +317,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..5199ff50abb959 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -286,6 +286,28 @@ def SingleOp : OpenMP_Op<"single", traits = [
let hasVerifier = 1;
}
+//===----------------------------------------------------------------------===//
+// 2.8.3 Workshare Construct
+//===----------------------------------------------------------------------===//
+
+def WorkshareOp : OpenMP_Op<"workshare", 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
+ }] # clausesDescription;
+
+ let builders = [
+ OpBuilder<(ins CArg<"const WorkshareOperands &">:$clauses)>
+ ];
+
+ 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..9a189eb2059e01 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1683,6 +1683,19 @@ LogicalResult SingleOp::verify() {
getCopyprivateSyms());
}
+//===----------------------------------------------------------------------===//
+// WorkshareOp
+//===----------------------------------------------------------------------===//
+
+void WorkshareOp::build(OpBuilder &builder, OperationState &state,
+ const WorkshareOperands &clauses) {
+ WorkshareOp::build(builder, state, clauses.nowait);
+}
+
+LogicalResult WorkshareOp::verify() {
+ return (*this)->getRegion(0).getBlocks().size() == 1 ? success() : failure();
+}
+
//===----------------------------------------------------------------------===//
// WsloopOp
//===----------------------------------------------------------------------===//
>From 13f5fa2b7628e78b308a3312fd045d87dd89d458 Mon Sep 17 00:00:00 2001
From: Ivan Radanov Ivanov <ivanov.i.aa at m.titech.ac.jp>
Date: Fri, 2 Aug 2024 16:10:25 +0900
Subject: [PATCH 2/7] Add custom omp loop wrapper
---
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 11 +++++++++++
1 file changed, 11 insertions(+)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 5199ff50abb959..76f0c472cfdb14 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -308,6 +308,17 @@ def WorkshareOp : OpenMP_Op<"workshare", clauses = [
let hasVerifier = 1;
}
+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 builders = [
+ OpBuilder<(ins), [{ build($_builder, $_state, {}); }]>
+ ];
+}
+
//===----------------------------------------------------------------------===//
// Loop Nest
//===----------------------------------------------------------------------===//
>From f0915f426f03e3e867953de36e68fd0ccaf5b0ed Mon Sep 17 00:00:00 2001
From: Ivan Radanov Ivanov <ivanov.i.aa at m.titech.ac.jp>
Date: Fri, 2 Aug 2024 16:08:58 +0900
Subject: [PATCH 3/7] Add recursive memory effects trait to workshare
---
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 4 +++-
1 file changed, 3 insertions(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 76f0c472cfdb14..7d1c80333855e7 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -290,7 +290,9 @@ def SingleOp : OpenMP_Op<"single", traits = [
// 2.8.3 Workshare Construct
//===----------------------------------------------------------------------===//
-def WorkshareOp : OpenMP_Op<"workshare", clauses = [
+def WorkshareOp : OpenMP_Op<"workshare", traits = [
+ RecursiveMemoryEffects,
+ ], clauses = [
OpenMP_NowaitClause,
], singleRegion = true> {
let summary = "workshare directive";
>From 0552b8e0592488b66a1c0195065a8b7c2cc92fb6 Mon Sep 17 00:00:00 2001
From: Ivan Radanov Ivanov <ivanov.i.aa at m.titech.ac.jp>
Date: Sun, 4 Aug 2024 17:04:07 +0900
Subject: [PATCH 4/7] Remove stray include
---
mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h | 1 -
1 file changed, 1 deletion(-)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h
index d14e5e17afbb08..896ca9581c3fc8 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h
@@ -17,7 +17,6 @@
#include "mlir/IR/BuiltinAttributes.h"
#include "llvm/ADT/SmallVector.h"
-#include <mlir/Dialect/OpenMP/OpenMPDialect.h>
#include "mlir/Dialect/OpenMP/OpenMPOpsEnums.h.inc"
>From e6e69596be16212ba997a698a290de99e76dcc00 Mon Sep 17 00:00:00 2001
From: Ivan Radanov Ivanov <ivanov.i.aa at m.titech.ac.jp>
Date: Sun, 4 Aug 2024 21:56:13 +0900
Subject: [PATCH 5/7] Remove omp.workshare verifier
---
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 2 --
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 4 ----
2 files changed, 6 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 7d1c80333855e7..863cd81923c87d 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -306,8 +306,6 @@ def WorkshareOp : OpenMP_Op<"workshare", traits = [
let builders = [
OpBuilder<(ins CArg<"const WorkshareOperands &">:$clauses)>
];
-
- let hasVerifier = 1;
}
def WorkshareLoopWrapperOp : OpenMP_Op<"workshare_loop_wrapper", traits = [
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 9a189eb2059e01..6c1b77077bdbaf 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1692,10 +1692,6 @@ void WorkshareOp::build(OpBuilder &builder, OperationState &state,
WorkshareOp::build(builder, state, clauses.nowait);
}
-LogicalResult WorkshareOp::verify() {
- return (*this)->getRegion(0).getBlocks().size() == 1 ? success() : failure();
-}
-
//===----------------------------------------------------------------------===//
// WsloopOp
//===----------------------------------------------------------------------===//
>From 7376fd432180a6c268c5c7fad58538de81ea6093 Mon Sep 17 00:00:00 2001
From: Ivan Radanov Ivanov <ivanov.i.aa at m.titech.ac.jp>
Date: Tue, 6 Aug 2024 13:41:22 +0900
Subject: [PATCH 6/7] Add assembly format for wrapper and add test
---
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 2 +-
mlir/test/Dialect/OpenMP/ops.mlir | 61 +++++++++++++++++++
2 files changed, 62 insertions(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 863cd81923c87d..0f29e911cb2f20 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -313,10 +313,10 @@ def WorkshareLoopWrapperOp : OpenMP_Op<"workshare_loop_wrapper", traits = [
RecursiveMemoryEffects, SingleBlock
], singleRegion = true> {
let summary = "contains loop nests to be parallelized by workshare";
-
let builders = [
OpBuilder<(ins), [{ build($_builder, $_state, {}); }]>
];
+ let assemblyFormat = "$region attr-dict";
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index d2924998f41b87..981e3fbb0306bf 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -2789,3 +2789,64 @@ 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_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
+ }
+ return
+}
+
+// CHECK-LABEL: func @omp_workshare_loop_wrapper_attrs
+func.func @omp_workshare_loop_wrapper_attrs(%idx : index) {
+ // 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}
+ return
+}
>From d030f24253b929fe85f935c32b2ae3459897c4a8 Mon Sep 17 00:00:00 2001
From: Ivan Radanov Ivanov <ivanov.i.aa at m.titech.ac.jp>
Date: Mon, 19 Aug 2024 14:42:35 +0900
Subject: [PATCH 7/7] Add verification and descriptions
---
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 10 +++++
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 14 +++++++
mlir/test/Dialect/OpenMP/invalid.mlir | 42 +++++++++++++++++++
mlir/test/Dialect/OpenMP/ops.mlir | 34 +++++++++------
4 files changed, 87 insertions(+), 13 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 0f29e911cb2f20..74497def8fd1ab 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -301,6 +301,10 @@ def WorkshareOp : OpenMP_Op<"workshare", traits = [
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 = [
@@ -313,10 +317,16 @@ def WorkshareLoopWrapperOp : OpenMP_Op<"workshare_loop_wrapper", traits = [
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;
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 6c1b77077bdbaf..90f9a19ebe32b5 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1692,6 +1692,20 @@ void WorkshareOp::build(OpBuilder &builder, OperationState &state,
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 981e3fbb0306bf..24363334b4cb54 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -2826,11 +2826,15 @@ func.func @omp_workshare_multiple_blocks() {
// CHECK-LABEL: func @omp_workshare_loop_wrapper
func.func @omp_workshare_loop_wrapper(%idx : index) {
- // 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
+ // 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
}
@@ -2839,14 +2843,18 @@ func.func @omp_workshare_loop_wrapper(%idx : index) {
// CHECK-LABEL: func @omp_workshare_loop_wrapper_attrs
func.func @omp_workshare_loop_wrapper_attrs(%idx : index) {
- // 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
- }
+ // 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
- // CHECK: } {attr_in_dict}
- } {attr_in_dict}
+ }
return
}
More information about the llvm-branch-commits
mailing list