[Mlir-commits] [mlir] f36b016 - [MLIR][OpenMP] Add MLIR operation for OpenMP teams
Sergio Afonso
llvmlistbot at llvm.org
Mon Jul 10 04:20:18 PDT 2023
Author: Sergio Afonso
Date: 2023-07-10T12:19:38+01:00
New Revision: f36b0169b8821e431644cb240f081538eb2b55ef
URL: https://github.com/llvm/llvm-project/commit/f36b0169b8821e431644cb240f081538eb2b55ef
DIFF: https://github.com/llvm/llvm-project/commit/f36b0169b8821e431644cb240f081538eb2b55ef.diff
LOG: [MLIR][OpenMP] Add MLIR operation for OpenMP teams
This patch adds an operation definition for the OpenMP 'teams' construct.
Differential Revision: https://reviews.llvm.org/D154441
Added:
Modified:
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
mlir/test/Dialect/OpenMP/invalid.mlir
mlir/test/Dialect/OpenMP/ops.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 33ca435119256a..aa142ed6aab3eb 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -218,6 +218,66 @@ def TerminatorOp : OpenMP_Op<"terminator", [Terminator, Pure]> {
let assemblyFormat = "attr-dict";
}
+//===----------------------------------------------------------------------===//
+// 2.7 teams Construct
+//===----------------------------------------------------------------------===//
+def TeamsOp : OpenMP_Op<"teams", [
+ AttrSizedOperandSegments, RecursiveMemoryEffects,
+ ReductionClauseInterface]> {
+ let summary = "teams construct";
+ let description = [{
+ The teams construct defines a region of code that triggers the creation of a
+ league of teams. Once created, the number of teams remains constant for the
+ duration of its code region.
+
+ The optional $num_teams_upper and $num_teams_lower specify the limit on the
+ number of teams to be created. If only the upper bound is specified, it acts
+ as if the lower bound was set to the same value. It is not supported to set
+ $num_teams_lower if $num_teams_upper is not specified. They define a closed
+ range, where both the lower and upper bounds are included.
+
+ If the $if_expr is present and it evaluates to `false`, the number of teams
+ created is one.
+
+ The optional $thread_limit specifies the limit on the number of threads.
+
+ The $allocators_vars and $allocate_vars parameters are a variadic list of
+ values that specify the memory allocator to be used to obtain storage for
+ private values.
+ }];
+
+ let arguments = (ins Optional<AnyInteger>:$num_teams_lower,
+ Optional<AnyInteger>:$num_teams_upper,
+ Optional<I1>:$if_expr,
+ Optional<AnyInteger>:$thread_limit,
+ Variadic<AnyType>:$allocate_vars,
+ Variadic<AnyType>:$allocators_vars,
+ Variadic<OpenMP_PointerLikeType>:$reduction_vars,
+ OptionalAttr<SymbolRefArrayAttr>:$reductions);
+
+ let regions = (region AnyRegion:$region);
+
+ let assemblyFormat = [{
+ oilist(
+ `num_teams` `(` ( $num_teams_lower^ `:` type($num_teams_lower) )? `to`
+ $num_teams_upper `:` type($num_teams_upper) `)`
+ | `if` `(` $if_expr `)`
+ | `thread_limit` `(` $thread_limit `:` type($thread_limit) `)`
+ | `reduction` `(`
+ custom<ReductionVarList>(
+ $reduction_vars, type($reduction_vars), $reductions
+ ) `)`
+ | `allocate` `(`
+ custom<AllocateAndAllocator>(
+ $allocate_vars, type($allocate_vars),
+ $allocators_vars, type($allocators_vars)
+ ) `)`
+ ) $region attr-dict
+ }];
+
+ let hasVerifier = 1;
+}
+
def OMP_ScheduleModNone : I32EnumAttrCase<"none", 0>;
def OMP_ScheduleModMonotonic : I32EnumAttrCase<"monotonic", 1>;
def OMP_ScheduleModNonmonotonic : I32EnumAttrCase<"nonmonotonic", 2>;
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 4e3d899c062c63..35f6d1c1239645 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -864,6 +864,48 @@ LogicalResult ParallelOp::verify() {
return verifyReductionVarList(*this, getReductions(), getReductionVars());
}
+//===----------------------------------------------------------------------===//
+// TeamsOp
+//===----------------------------------------------------------------------===//
+
+static bool opInGlobalImplicitParallelRegion(Operation *op) {
+ while ((op = op->getParentOp()))
+ if (isa<OpenMPDialect>(op->getDialect()))
+ return false;
+ return true;
+}
+
+LogicalResult TeamsOp::verify() {
+ // Check parent region
+ // TODO If nested inside of a target region, also check that it does not
+ // contain any statements, declarations or directives other than this
+ // omp.teams construct. The issue is how to support the initialization of
+ // this operation's own arguments (allow SSA values across omp.target?).
+ Operation *op = getOperation();
+ if (!isa<TargetOp>(op->getParentOp()) &&
+ !opInGlobalImplicitParallelRegion(op))
+ return emitError("expected to be nested inside of omp.target or not nested "
+ "in any OpenMP dialect operations");
+
+ // Check for num_teams clause restrictions
+ if (auto numTeamsLowerBound = getNumTeamsLower()) {
+ auto numTeamsUpperBound = getNumTeamsUpper();
+ if (!numTeamsUpperBound)
+ return emitError("expected num_teams upper bound to be defined if the "
+ "lower bound is defined");
+ if (numTeamsLowerBound.getType() != numTeamsUpperBound.getType())
+ return emitError(
+ "expected num_teams upper bound and lower bound to be the same type");
+ }
+
+ // Check for allocate clause restrictions
+ if (getAllocateVars().size() != getAllocatorsVars().size())
+ return emitError(
+ "expected equal sizes for allocate and allocator variables");
+
+ return verifyReductionVarList(*this, getReductions(), getReductionVars());
+}
+
//===----------------------------------------------------------------------===//
// Verifier for SectionsOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index a3f0048a11c0b7..fc65fb77ffc88f 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1103,6 +1103,58 @@ func.func @omp_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
// -----
+func.func @omp_teams_parent() {
+ omp.parallel {
+ // expected-error @below {{expected to be nested inside of omp.target or not nested in any OpenMP dialect operations}}
+ omp.teams {
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
+
+// -----
+
+func.func @omp_teams_allocate(%data_var : memref<i32>) {
+ omp.target {
+ // expected-error @below {{expected equal sizes for allocate and allocator variables}}
+ "omp.teams" (%data_var) ({
+ omp.terminator
+ }) {operand_segment_sizes = array<i32: 0,0,0,0,1,0,0>} : (memref<i32>) -> ()
+ omp.terminator
+ }
+ return
+}
+
+// -----
+
+func.func @omp_teams_num_teams1(%lb : i32) {
+ omp.target {
+ // expected-error @below {{expected num_teams upper bound to be defined if the lower bound is defined}}
+ "omp.teams" (%lb) ({
+ omp.terminator
+ }) {operand_segment_sizes = array<i32: 1,0,0,0,0,0,0>} : (i32) -> ()
+ omp.terminator
+ }
+ return
+}
+
+// -----
+
+func.func @omp_teams_num_teams2(%lb : i32, %ub : i16) {
+ omp.target {
+ // expected-error @below {{expected num_teams upper bound and lower bound to be the same type}}
+ omp.teams num_teams(%lb : i32 to %ub : i16) {
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
+
+// -----
+
func.func @omp_sections(%data_var : memref<i32>) -> () {
// expected-error @below {{expected equal sizes for allocate and allocator variables}}
"omp.sections" (%data_var) ({
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 93df164b9764b9..2f0d224a3fef7e 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -616,6 +616,76 @@ func.func @parallel_wsloop_reduction(%lb : index, %ub : index, %step : index) {
return
}
+// CHECK-LABEL: omp_teams
+func.func @omp_teams(%lb : i32, %ub : i32, %if_cond : i1, %num_threads : i32,
+ %data_var : memref<i32>) -> () {
+ // Test nesting inside of omp.target
+ omp.target {
+ // CHECK: omp.teams
+ omp.teams {
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+
+ // CHECK: omp.teams
+ omp.teams {
+ %0 = arith.constant 1 : i32
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+
+ // Test num teams.
+ // CHECK: omp.teams num_teams(%{{.+}} : i32 to %{{.+}} : i32)
+ omp.teams num_teams(%lb : i32 to %ub : i32) {
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+
+ // CHECK: omp.teams num_teams( to %{{.+}} : i32)
+ omp.teams num_teams(to %ub : i32) {
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+
+ // Test if.
+ // CHECK: omp.teams if(%{{.+}})
+ omp.teams if(%if_cond) {
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+
+ // Test thread limit.
+ // CHECK: omp.teams thread_limit(%{{.+}} : i32)
+ omp.teams thread_limit(%num_threads : i32) {
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+
+ // Test reduction.
+ %c1 = arith.constant 1 : i32
+ %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<f32>
+ // CHECK: omp.teams reduction(@add_f32 -> %{{.+}} : !llvm.ptr<f32>) {
+ omp.teams reduction(@add_f32 -> %0 : !llvm.ptr<f32>) {
+ %1 = arith.constant 2.0 : f32
+ // CHECK: omp.reduction %{{.+}}, %{{.+}}
+ omp.reduction %1, %0 : f32, !llvm.ptr<f32>
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+
+ // Test allocate.
+ // CHECK: omp.teams allocate(%{{.+}} : memref<i32> -> %{{.+}} : memref<i32>)
+ omp.teams allocate(%data_var : memref<i32> -> %data_var : memref<i32>) {
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+
+ return
+}
+
// CHECK-LABEL: func @sections_reduction
func.func @sections_reduction() {
%c1 = arith.constant 1 : i32
More information about the Mlir-commits
mailing list