[Mlir-commits] [mlir] [MLIR][omp] Add omp.workshare op (PR #101443)
Ivan R. Ivanov
llvmlistbot at llvm.org
Wed Aug 21 07:07:26 PDT 2024
https://github.com/ivanradanov updated https://github.com/llvm/llvm-project/pull/101443
>From 232aa13f7b067c4135758dfe4341062f2ffd2b27 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 d63fdd88f79104..e360018cde28cf 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -287,6 +287,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 4c943ebbe3144f..b3489004a7ac2b 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1689,6 +1689,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 a51daa965c1d887705ed9e4309006e307ecf0c58 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 e360018cde28cf..ea302878923745 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -309,6 +309,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 5555acc4475c163b3b4f66f14f0cd2f57cb0db92 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 ea302878923745..dbb25f7cc7dd00 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -291,7 +291,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 616b2b449d1e5f9fc43b02ecc2af4079081fd7a6 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 8daabbeee2d2bcc3d06f109cb4926210939c4a4b 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 dbb25f7cc7dd00..94806dc4cc5e34 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -307,8 +307,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 b3489004a7ac2b..230af836c1531e 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1698,10 +1698,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 1b9784058bb23fa4687591db32e9cd47ca8fcc45 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 94806dc4cc5e34..b41c53ef282047 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -314,10 +314,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 9ac97e069addd2..8ef49ddc807b6a 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -2810,3 +2810,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 45a8aa4a12ab2456e87ca6c483701cdc256c0080 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 b41c53ef282047..f8d7e934edbc61 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -302,6 +302,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 = [
@@ -314,10 +318,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 230af836c1531e..f4acbd97ca6d1a 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1698,6 +1698,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 c76b07ec94a597..a67dcb34427062 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -2545,3 +2545,45 @@ func.func @omp_taskloop_invalid_composite(%lb: index, %ub: index, %step: index)
} {omp.composite}
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 8ef49ddc807b6a..1e0ccfb79ed54e 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -2847,11 +2847,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
}
@@ -2860,14 +2864,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 Mlir-commits
mailing list