[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
Mon Aug 19 01:41:05 PDT 2024
https://github.com/ivanradanov updated https://github.com/llvm/llvm-project/pull/101443
>From 604b0293e0574e9d697d4071c2b853a5a27af1e1 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 f2fd4f278c23ec99dae3ac44e1c05fcb629f707d 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 22c66e6db3997e38254d9848661a38627cd7bb19 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 5d6094bc98a5a63cf1db6ef9eb2a34337acb86ed 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 e41c776df6dd0af31b6739448323dc704c9716f7 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 94a2e49f1ce0e8d95287041f64909ec23bd189ed 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 72e19faffb72d3aafed907efd0613416d06cbce0 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