[flang] [llvm] [mlir] [OpenMP][MLIR] Add num_teams clause with dims modifier support (PR #169883)
via llvm-commits
llvm-commits at lists.llvm.org
Thu Jan 22 08:12:15 PST 2026
https://github.com/skc7 updated https://github.com/llvm/llvm-project/pull/169883
>From 675c0df8219c644d56fe8b8e0855ab03803783c1 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Fri, 28 Nov 2025 13:37:14 +0530
Subject: [PATCH 01/14] [OpenMP][MLIR] Add num_teams clause with dims modifier
support
---
.../mlir/Dialect/OpenMP/OpenMPClauses.td | 72 +++++++++++++++++++
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 3 +-
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 64 +++++++++++++++++
mlir/test/Dialect/OpenMP/invalid.mlir | 4 +-
mlir/test/Dialect/OpenMP/ops.mlir | 6 ++
5 files changed, 146 insertions(+), 3 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index 05e2ee4e5632b..1341bba3ad85e 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
@@ -1532,4 +1532,76 @@ class OpenMP_UseDevicePtrClauseSkip<
def OpenMP_UseDevicePtrClause : OpenMP_UseDevicePtrClauseSkip<>;
+//===----------------------------------------------------------------------===//
+// V6.2: Multidimensional `num_teams` clause with dims modifier
+//===----------------------------------------------------------------------===//
+
+class OpenMP_NumTeamsMultiDimClauseSkip<
+ bit traits = false, bit arguments = false, bit assemblyFormat = false,
+ bit description = false, bit extraClassDeclaration = false
+ > : OpenMP_Clause<traits, arguments, assemblyFormat, description,
+ extraClassDeclaration> {
+ let arguments = (ins
+ ConfinedAttr<OptionalAttr<I64Attr>, [IntPositive]>:$num_teams_dims,
+ Variadic<AnyInteger>:$num_teams_values
+ );
+
+ let optAssemblyFormat = [{
+ `num_teams_multi_dim` `(` custom<NumTeamsMultiDimClause>($num_teams_dims,
+ $num_teams_values,
+ type($num_teams_values)) `)`
+ }];
+
+ let description = [{
+ The `num_teams_multi_dim` clause with dims modifier support specifies the limit on
+ the number of teams to be created in a multidimensional team space.
+
+ The dims modifier for the num_teams_multi_dim clause specifies the number of
+ dimensions for the league space (team space) that the clause arranges.
+ The dimensions argument in the dims modifier specifies the number of
+ dimensions and determines the length of the list argument. The list items
+ are specified in ascending order according to the ordinal number of the
+ dimensions (dimension 0, 1, 2, ..., N-1).
+
+ - If `dims` is not specified: The space is unidimensional (1D) with a single value
+ - If `dims(1)` is specified: The space is explicitly unidimensional (1D)
+ - If `dims(N)` where N > 1: The space is strictly multidimensional (N-D)
+
+ **Examples:**
+ - `num_teams_multi_dim(dims(3): %nt0, %nt1, %nt2 : i32, i32, i32)` creates a
+ 3-dimensional team space with limits nt0, nt1, nt2 for dimensions 0, 1, 2.
+ - `num_teams_multi_dim(%nt : i32)` creates a unidimensional team space with limit nt.
+ }];
+
+ let extraClassDeclaration = [{
+ /// Returns true if the dims modifier is explicitly present
+ bool hasDimsModifier() {
+ return getNumTeamsDims().has_value();
+ }
+
+ /// Returns the number of dimensions specified by dims modifier
+ /// Returns 1 if dims modifier is not present (unidimensional by default)
+ unsigned getNumDimensions() {
+ if (!hasDimsModifier())
+ return 1;
+ return static_cast<unsigned>(*getNumTeamsDims());
+ }
+
+ /// Returns all dimension values as an operand range
+ ::mlir::OperandRange getDimensionValues() {
+ return getNumTeamsValues();
+ }
+
+ /// Returns the value for a specific dimension index
+ /// Index must be less than getNumDimensions()
+ ::mlir::Value getDimensionValue(unsigned index) {
+ assert(index < getDimensionValues().size() &&
+ "Dimension index out of bounds");
+ return getDimensionValues()[index];
+ }
+ }];
+}
+
+def OpenMP_NumTeamsMultiDimClause : OpenMP_NumTeamsMultiDimClauseSkip<>;
+
#endif // OPENMP_CLAUSES
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index d4e8cecda2601..76eeb0bd70ec3 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -241,7 +241,8 @@ def TeamsOp : OpenMP_Op<"teams", traits = [
AttrSizedOperandSegments, RecursiveMemoryEffects, OutlineableOpenMPOpInterface
], clauses = [
OpenMP_AllocateClause, OpenMP_IfClause, OpenMP_NumTeamsClause,
- OpenMP_PrivateClause, OpenMP_ReductionClause, OpenMP_ThreadLimitClause
+ OpenMP_NumTeamsMultiDimClause, OpenMP_PrivateClause, OpenMP_ReductionClause,
+ OpenMP_ThreadLimitClause
], singleRegion = true> {
let summary = "teams construct";
let description = [{
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index cf25fc53e0238..783fb32848136 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2626,6 +2626,7 @@ void TeamsOp::build(OpBuilder &builder, OperationState &state,
// TODO Store clauses in op: privateVars, privateSyms, privateNeedsBarrier
TeamsOp::build(builder, state, clauses.allocateVars, clauses.allocatorVars,
clauses.ifExpr, clauses.numTeamsLower, clauses.numTeamsUpper,
+ clauses.numTeamsDims, clauses.numTeamsValues,
/*private_vars=*/{}, /*private_syms=*/nullptr,
/*private_needs_barrier=*/nullptr, clauses.reductionMod,
clauses.reductionVars,
@@ -4488,6 +4489,69 @@ void DeclareSimdOp::build(OpBuilder &odsBuilder, OperationState &odsState,
clauses.linearVarTypes, clauses.simdlen);
}
+//===----------------------------------------------------------------------===//
+// Parser and printer for NumTeamsMultiDim Clause (with dims modifier)
+//===----------------------------------------------------------------------===//
+// num_teams_multidim ::= `num_teams` `(` [`dims` `(` dim-count `)` `:`] values
+// `)` Example: num_teams(dims(3): %v0, %v1, %v2 : i32, i32, i32) Or:
+// num_teams(%v : i32)
+static ParseResult parseNumTeamsMultiDimClause(
+ OpAsmParser &parser, IntegerAttr &dimsAttr,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
+ SmallVectorImpl<Type> &types) {
+ std::optional<int64_t> dims;
+ // Try to parse optional dims modifier: dims(N):
+ if (succeeded(parser.parseOptionalKeyword("dims"))) {
+ int64_t dimsValue;
+ if (parser.parseLParen() || parser.parseInteger(dimsValue) ||
+ parser.parseRParen() || parser.parseColon()) {
+ return failure();
+ }
+ dims = dimsValue;
+ }
+ // Parse the operand list
+ if (parser.parseOperandList(values))
+ return failure();
+ // Parse colon and types
+ if (parser.parseColon() || parser.parseTypeList(types))
+ return failure();
+
+ // Verify dims matches number of values if specified
+ if (dims.has_value() && values.size() != static_cast<size_t>(*dims)) {
+ return parser.emitError(parser.getCurrentLocation())
+ << "dims(" << *dims << ") specified but " << values.size()
+ << " values provided";
+ }
+
+ // If dims not specified but we have values, it's implicitly unidimensional
+ if (!dims.has_value() && values.size() != 1) {
+ return parser.emitError(parser.getCurrentLocation())
+ << "expected 1 value without dims modifier, got " << values.size();
+ }
+
+ // Convert to IntegerAttr
+ if (dims.has_value()) {
+ dimsAttr = parser.getBuilder().getI64IntegerAttr(*dims);
+ }
+ return success();
+}
+
+static void printNumTeamsMultiDimClause(OpAsmPrinter &p, Operation *op,
+ IntegerAttr dimsAttr,
+ OperandRange values, TypeRange types) {
+ // Print dims modifier if present
+ if (dimsAttr) {
+ p << "dims(" << dimsAttr.getInt() << "): ";
+ }
+
+ // Print operands
+ p.printOperands(values);
+
+ // Print types
+ p << " : ";
+ llvm::interleaveComma(types, p);
+}
+
#define GET_ATTRDEF_CLASSES
#include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 1350c5e2ee8d2..cd06011c2cbc4 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1438,7 +1438,7 @@ func.func @omp_teams_allocate(%data_var : memref<i32>) {
// expected-error @below {{expected equal sizes for allocate and allocator variables}}
"omp.teams" (%data_var) ({
omp.terminator
- }) {operandSegmentSizes = array<i32: 1,0,0,0,0,0,0,0>} : (memref<i32>) -> ()
+ }) {operandSegmentSizes = array<i32: 1,0,0,0,0,0,0,0,0>} : (memref<i32>) -> ()
omp.terminator
}
return
@@ -1451,7 +1451,7 @@ func.func @omp_teams_num_teams1(%lb : i32) {
// expected-error @below {{expected num_teams upper bound to be defined if the lower bound is defined}}
"omp.teams" (%lb) ({
omp.terminator
- }) {operandSegmentSizes = array<i32: 0,0,0,1,0,0,0,0>} : (i32) -> ()
+ }) {operandSegmentSizes = array<i32: 0,0,0,1,0,0,0,0,0>} : (i32) -> ()
omp.terminator
}
return
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 902d962138fc9..82f67d961fa10 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -1109,6 +1109,12 @@ func.func @omp_teams(%lb : i32, %ub : i32, %if_cond : i1, %num_threads : i32,
omp.terminator
}
+ // CHECK: omp.teams num_teams_multi_dim(dims(3): %{{.*}}, %{{.*}}, %{{.*}} : i32, i32, i32)
+ omp.teams num_teams_multi_dim(dims(3): %lb, %ub, %ub : i32, i32, i32) {
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+
// Test if.
// CHECK: omp.teams if(%{{.+}})
omp.teams if(%if_cond) {
>From 465cc980bff5770a9cfb7801b9079ae3beb289c5 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Fri, 28 Nov 2025 14:20:39 +0530
Subject: [PATCH 02/14] fix comment
---
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 5 ++---
1 file changed, 2 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 783fb32848136..1962287a6ed83 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -4492,9 +4492,8 @@ void DeclareSimdOp::build(OpBuilder &odsBuilder, OperationState &odsState,
//===----------------------------------------------------------------------===//
// Parser and printer for NumTeamsMultiDim Clause (with dims modifier)
//===----------------------------------------------------------------------===//
-// num_teams_multidim ::= `num_teams` `(` [`dims` `(` dim-count `)` `:`] values
-// `)` Example: num_teams(dims(3): %v0, %v1, %v2 : i32, i32, i32) Or:
-// num_teams(%v : i32)
+// num_teams_multi_dim(dims(3): %v0, %v1, %v2 : i32, i32, i32) Or:
+// num_teams_multi_dim(%v : i32)
static ParseResult parseNumTeamsMultiDimClause(
OpAsmParser &parser, IntegerAttr &dimsAttr,
SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
>From c14791d711a4d1d0943beb5c88fe36adb8732458 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Mon, 1 Dec 2025 12:56:28 +0530
Subject: [PATCH 03/14] Comments fix
---
mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td | 2 +-
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 3 ++-
2 files changed, 3 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index 1341bba3ad85e..8ea3dedf1d3ac 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
@@ -1533,7 +1533,7 @@ class OpenMP_UseDevicePtrClauseSkip<
def OpenMP_UseDevicePtrClause : OpenMP_UseDevicePtrClauseSkip<>;
//===----------------------------------------------------------------------===//
-// V6.2: Multidimensional `num_teams` clause with dims modifier
+// V6.1: Multidimensional `num_teams` clause with dims modifier
//===----------------------------------------------------------------------===//
class OpenMP_NumTeamsMultiDimClauseSkip<
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 1962287a6ed83..ac5b2bbf942d5 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -4525,7 +4525,8 @@ static ParseResult parseNumTeamsMultiDimClause(
// If dims not specified but we have values, it's implicitly unidimensional
if (!dims.has_value() && values.size() != 1) {
return parser.emitError(parser.getCurrentLocation())
- << "expected 1 value without dims modifier, got " << values.size();
+ << "expected 1 value without dims modifier, but got "
+ << values.size() << " values";
}
// Convert to IntegerAttr
>From 63be1ad33e812e4406a1e7eb2939af8f3d3b9646 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Mon, 1 Dec 2025 15:02:21 +0530
Subject: [PATCH 04/14] Use DimsModifier for custom assembly parser and printer
---
.../mlir/Dialect/OpenMP/OpenMPClauses.td | 4 ++--
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 20 +++++++++----------
2 files changed, 12 insertions(+), 12 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index 8ea3dedf1d3ac..468c5ce132aaf 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
@@ -1533,7 +1533,7 @@ class OpenMP_UseDevicePtrClauseSkip<
def OpenMP_UseDevicePtrClause : OpenMP_UseDevicePtrClauseSkip<>;
//===----------------------------------------------------------------------===//
-// V6.1: Multidimensional `num_teams` clause with dims modifier
+// V6.1: `num_teams` clause with dims modifier
//===----------------------------------------------------------------------===//
class OpenMP_NumTeamsMultiDimClauseSkip<
@@ -1547,7 +1547,7 @@ class OpenMP_NumTeamsMultiDimClauseSkip<
);
let optAssemblyFormat = [{
- `num_teams_multi_dim` `(` custom<NumTeamsMultiDimClause>($num_teams_dims,
+ `num_teams_multi_dim` `(` custom<DimsModifier>($num_teams_dims,
$num_teams_values,
type($num_teams_values)) `)`
}];
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index ac5b2bbf942d5..8fb0e4992f5f3 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -4490,14 +4490,14 @@ void DeclareSimdOp::build(OpBuilder &odsBuilder, OperationState &odsState,
}
//===----------------------------------------------------------------------===//
-// Parser and printer for NumTeamsMultiDim Clause (with dims modifier)
+// Parser and printer for Clauses with dims modifier
//===----------------------------------------------------------------------===//
-// num_teams_multi_dim(dims(3): %v0, %v1, %v2 : i32, i32, i32) Or:
-// num_teams_multi_dim(%v : i32)
-static ParseResult parseNumTeamsMultiDimClause(
- OpAsmParser &parser, IntegerAttr &dimsAttr,
- SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
- SmallVectorImpl<Type> &types) {
+// clause_name(dims(3): %v0, %v1, %v2 : i32, i32, i32)
+// clause_name(%v : i32)
+static ParseResult
+parseDimsModifier(OpAsmParser &parser, IntegerAttr &dimsAttr,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
+ SmallVectorImpl<Type> &types) {
std::optional<int64_t> dims;
// Try to parse optional dims modifier: dims(N):
if (succeeded(parser.parseOptionalKeyword("dims"))) {
@@ -4536,9 +4536,9 @@ static ParseResult parseNumTeamsMultiDimClause(
return success();
}
-static void printNumTeamsMultiDimClause(OpAsmPrinter &p, Operation *op,
- IntegerAttr dimsAttr,
- OperandRange values, TypeRange types) {
+static void printDimsModifier(OpAsmPrinter &p, Operation *op,
+ IntegerAttr dimsAttr, OperandRange values,
+ TypeRange types) {
// Print dims modifier if present
if (dimsAttr) {
p << "dims(" << dimsAttr.getInt() << "): ";
>From 079b32197c0b60902117d5b53a4686c704df7c14 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Wed, 10 Dec 2025 13:13:17 +0530
Subject: [PATCH 05/14] use dims modifer in main num_teams clause itself
instead of creating new clause
---
.../mlir/Dialect/OpenMP/OpenMPClauses.td | 128 +++++------
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 3 +-
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 204 ++++++++++++++----
mlir/test/Dialect/OpenMP/invalid.mlir | 77 ++++++-
mlir/test/Dialect/OpenMP/ops.mlir | 4 +-
5 files changed, 284 insertions(+), 132 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index 468c5ce132aaf..0b17a62c88d92 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
@@ -974,22 +974,62 @@ class OpenMP_NumTeamsClauseSkip<
> : OpenMP_Clause<traits, arguments, assemblyFormat, description,
extraClassDeclaration> {
let arguments = (ins
+ ConfinedAttr<OptionalAttr<I64Attr>, [IntPositive]>:$num_teams_dims,
+ Variadic<AnyInteger>:$num_teams_values,
Optional<AnyInteger>:$num_teams_lower,
Optional<AnyInteger>:$num_teams_upper
);
let optAssemblyFormat = [{
- `num_teams` `(` ( $num_teams_lower^ `:` type($num_teams_lower) )? `to`
- $num_teams_upper `:` type($num_teams_upper) `)`
+ `num_teams` `(` custom<NumTeamsClause>(
+ $num_teams_dims, $num_teams_values, type($num_teams_values),
+ $num_teams_lower, type($num_teams_lower),
+ $num_teams_upper, type($num_teams_upper)
+ ) `)`
}];
let description = [{
- The optional `num_teams_upper` and `num_teams_lower` arguments 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 allowed 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.
+ The `num_teams` clause specifies the bounds on the league space formed by the
+ construct on which it appears.
+
+ With dims modifier: (OpenMP 6.1 requirement)
+ - Uses `num_teams_dims` (dimension count) and `num_teams_values` (upper bounds list)
+ - Specifies upper bounds for each dimension (all must have same type)
+ - Format: `num_teams(dims(N): upper_bound_0, ..., upper_bound_N-1 : type)`
+ - Example: `num_teams(dims(3): %ub0, %ub1, %ub2 : i32)`
+
+ Without dims modifier:
+ - Uses `num_teams_upper` and optional `num_teams_lower`
+ - If lower bound not specified, it defaults to upper bound value
+ - Format: `num_teams(lower : type to upper : type)` or `num_teams(to upper : type)`
+ - Example: `num_teams(%lb : i32 to %ub : i32)` or `num_teams(to %ub : i32)`
+ }];
+
+ let extraClassDeclaration = [{
+ /// Returns true if the dims modifier is explicitly present
+ bool hasDimsModifier() {
+ return getNumTeamsDims().has_value();
+ }
+
+ /// Returns the number of dimensions specified by dims modifier
+ unsigned getNumDimensions() {
+ if (!hasDimsModifier())
+ return 1;
+ return static_cast<unsigned>(*getNumTeamsDims());
+ }
+
+ /// Returns all dimension values as an operand range
+ ::mlir::OperandRange getDimensionValues() {
+ return getNumTeamsValues();
+ }
+
+ /// Returns the value for a specific dimension index
+ /// Index must be less than getNumDimensions()
+ ::mlir::Value getDimensionValue(unsigned index) {
+ assert(index < getDimensionValues().size() &&
+ "Dimension index out of bounds");
+ return getDimensionValues()[index];
+ }
}];
}
@@ -1532,76 +1572,4 @@ class OpenMP_UseDevicePtrClauseSkip<
def OpenMP_UseDevicePtrClause : OpenMP_UseDevicePtrClauseSkip<>;
-//===----------------------------------------------------------------------===//
-// V6.1: `num_teams` clause with dims modifier
-//===----------------------------------------------------------------------===//
-
-class OpenMP_NumTeamsMultiDimClauseSkip<
- bit traits = false, bit arguments = false, bit assemblyFormat = false,
- bit description = false, bit extraClassDeclaration = false
- > : OpenMP_Clause<traits, arguments, assemblyFormat, description,
- extraClassDeclaration> {
- let arguments = (ins
- ConfinedAttr<OptionalAttr<I64Attr>, [IntPositive]>:$num_teams_dims,
- Variadic<AnyInteger>:$num_teams_values
- );
-
- let optAssemblyFormat = [{
- `num_teams_multi_dim` `(` custom<DimsModifier>($num_teams_dims,
- $num_teams_values,
- type($num_teams_values)) `)`
- }];
-
- let description = [{
- The `num_teams_multi_dim` clause with dims modifier support specifies the limit on
- the number of teams to be created in a multidimensional team space.
-
- The dims modifier for the num_teams_multi_dim clause specifies the number of
- dimensions for the league space (team space) that the clause arranges.
- The dimensions argument in the dims modifier specifies the number of
- dimensions and determines the length of the list argument. The list items
- are specified in ascending order according to the ordinal number of the
- dimensions (dimension 0, 1, 2, ..., N-1).
-
- - If `dims` is not specified: The space is unidimensional (1D) with a single value
- - If `dims(1)` is specified: The space is explicitly unidimensional (1D)
- - If `dims(N)` where N > 1: The space is strictly multidimensional (N-D)
-
- **Examples:**
- - `num_teams_multi_dim(dims(3): %nt0, %nt1, %nt2 : i32, i32, i32)` creates a
- 3-dimensional team space with limits nt0, nt1, nt2 for dimensions 0, 1, 2.
- - `num_teams_multi_dim(%nt : i32)` creates a unidimensional team space with limit nt.
- }];
-
- let extraClassDeclaration = [{
- /// Returns true if the dims modifier is explicitly present
- bool hasDimsModifier() {
- return getNumTeamsDims().has_value();
- }
-
- /// Returns the number of dimensions specified by dims modifier
- /// Returns 1 if dims modifier is not present (unidimensional by default)
- unsigned getNumDimensions() {
- if (!hasDimsModifier())
- return 1;
- return static_cast<unsigned>(*getNumTeamsDims());
- }
-
- /// Returns all dimension values as an operand range
- ::mlir::OperandRange getDimensionValues() {
- return getNumTeamsValues();
- }
-
- /// Returns the value for a specific dimension index
- /// Index must be less than getNumDimensions()
- ::mlir::Value getDimensionValue(unsigned index) {
- assert(index < getDimensionValues().size() &&
- "Dimension index out of bounds");
- return getDimensionValues()[index];
- }
- }];
-}
-
-def OpenMP_NumTeamsMultiDimClause : OpenMP_NumTeamsMultiDimClauseSkip<>;
-
#endif // OPENMP_CLAUSES
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 76eeb0bd70ec3..d4e8cecda2601 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -241,8 +241,7 @@ def TeamsOp : OpenMP_Op<"teams", traits = [
AttrSizedOperandSegments, RecursiveMemoryEffects, OutlineableOpenMPOpInterface
], clauses = [
OpenMP_AllocateClause, OpenMP_IfClause, OpenMP_NumTeamsClause,
- OpenMP_NumTeamsMultiDimClause, OpenMP_PrivateClause, OpenMP_ReductionClause,
- OpenMP_ThreadLimitClause
+ OpenMP_PrivateClause, OpenMP_ReductionClause, OpenMP_ThreadLimitClause
], singleRegion = true> {
let summary = "teams construct";
let description = [{
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 8fb0e4992f5f3..17e286beaf51c 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2625,8 +2625,8 @@ void TeamsOp::build(OpBuilder &builder, OperationState &state,
MLIRContext *ctx = builder.getContext();
// TODO Store clauses in op: privateVars, privateSyms, privateNeedsBarrier
TeamsOp::build(builder, state, clauses.allocateVars, clauses.allocatorVars,
- clauses.ifExpr, clauses.numTeamsLower, clauses.numTeamsUpper,
- clauses.numTeamsDims, clauses.numTeamsValues,
+ clauses.ifExpr, clauses.numTeamsDims, clauses.numTeamsValues,
+ clauses.numTeamsLower, clauses.numTeamsUpper,
/*private_vars=*/{}, /*private_syms=*/nullptr,
/*private_needs_barrier=*/nullptr, clauses.reductionMod,
clauses.reductionVars,
@@ -2648,14 +2648,57 @@ LogicalResult TeamsOp::verify() {
"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())
+ auto numTeamsDims = getNumTeamsDims();
+ auto numTeamsValues = getNumTeamsValues();
+ auto numTeamsLower = getNumTeamsLower();
+ auto numTeamsUpper = getNumTeamsUpper();
+
+ // Cannot use both dims modifier and unidimensional style
+ if (numTeamsDims.has_value() && (numTeamsLower || numTeamsUpper)) {
+ return emitError(
+ "num_teams with dims modifier cannot be used together with "
+ "lower/upper bounds (unidimensional style)");
+ }
+
+ // With dims modifier (multidimensional)
+ if (numTeamsDims.has_value()) {
+ if (numTeamsValues.empty()) {
+ return emitError(
+ "num_teams dims modifier requires values to be specified");
+ }
+
+ if (numTeamsValues.size() != static_cast<size_t>(*numTeamsDims)) {
+ return emitError("num_teams dims(")
+ << *numTeamsDims << ") specified but " << numTeamsValues.size()
+ << " values provided";
+ }
+
+ // All values must have the same type
+ if (!numTeamsValues.empty()) {
+ Type firstType = numTeamsValues.front().getType();
+ for (auto value : numTeamsValues) {
+ if (value.getType() != firstType) {
+ return emitError(
+ "num_teams dims modifier requires all values to have "
+ "the same type");
+ }
+ }
+ }
+ } else {
+ // Without dims modifier
+ if (!numTeamsValues.empty()) {
return emitError(
- "expected num_teams upper bound and lower bound to be the same type");
+ "num_teams values can only be specified with dims modifier");
+ }
+
+ if (numTeamsLower) {
+ if (!numTeamsUpper)
+ return emitError("expected num_teams upper bound to be defined if the "
+ "lower bound is defined");
+ if (numTeamsLower.getType() != numTeamsUpper.getType())
+ return emitError("expected num_teams upper bound and lower bound to be "
+ "the same type");
+ }
}
// Check for allocate clause restrictions
@@ -4490,66 +4533,133 @@ void DeclareSimdOp::build(OpBuilder &odsBuilder, OperationState &odsState,
}
//===----------------------------------------------------------------------===//
-// Parser and printer for Clauses with dims modifier
+// Helper: Parse dims modifier with values
+//===----------------------------------------------------------------------===//
+// Parses: dims(N): values : type (single type for all values)
+static ParseResult parseDimsModifierWithValues(
+ OpAsmParser &parser, IntegerAttr &dimsAttr,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
+ SmallVectorImpl<Type> &types) {
+ if (failed(parser.parseOptionalKeyword("dims"))) {
+ return failure();
+ }
+
+ // Parse (N): values : type
+ int64_t dimsValue;
+ if (parser.parseLParen() || parser.parseInteger(dimsValue) ||
+ parser.parseRParen() || parser.parseColon()) {
+ return failure();
+ }
+
+ if (parser.parseOperandList(values) || parser.parseColon()) {
+ return failure();
+ }
+
+ // Parse single type (all values have same type)
+ Type valueType;
+ if (parser.parseType(valueType)) {
+ return failure();
+ }
+
+ // Fill types vector with same type for all values
+ types.assign(values.size(), valueType);
+
+ dimsAttr = parser.getBuilder().getI64IntegerAttr(dimsValue);
+ return success();
+}
+
+//===----------------------------------------------------------------------===//
+// Parser and printer for num_teams clause with dims modifier
//===----------------------------------------------------------------------===//
-// clause_name(dims(3): %v0, %v1, %v2 : i32, i32, i32)
-// clause_name(%v : i32)
static ParseResult
-parseDimsModifier(OpAsmParser &parser, IntegerAttr &dimsAttr,
- SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
- SmallVectorImpl<Type> &types) {
- std::optional<int64_t> dims;
- // Try to parse optional dims modifier: dims(N):
- if (succeeded(parser.parseOptionalKeyword("dims"))) {
- int64_t dimsValue;
- if (parser.parseLParen() || parser.parseInteger(dimsValue) ||
- parser.parseRParen() || parser.parseColon()) {
+parseNumTeamsClause(OpAsmParser &parser, IntegerAttr &dimsAttr,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
+ SmallVectorImpl<Type> &types,
+ std::optional<OpAsmParser::UnresolvedOperand> &lowerBound,
+ Type &lowerBoundType,
+ std::optional<OpAsmParser::UnresolvedOperand> &upperBound,
+ Type &upperBoundType) {
+
+ // Format: num_teams(dims(N): values : type)
+ if (succeeded(parseDimsModifierWithValues(parser, dimsAttr, values, types))) {
+ return success();
+ }
+
+ // Format: num_teams(to upper : type)
+ if (succeeded(parser.parseOptionalKeyword("to"))) {
+ OpAsmParser::UnresolvedOperand upperOperand;
+ if (parser.parseOperand(upperOperand) || parser.parseColon() ||
+ parser.parseType(upperBoundType)) {
return failure();
}
- dims = dimsValue;
+ upperBound = upperOperand;
+ return success();
}
- // Parse the operand list
- if (parser.parseOperandList(values))
- return failure();
- // Parse colon and types
- if (parser.parseColon() || parser.parseTypeList(types))
- return failure();
- // Verify dims matches number of values if specified
- if (dims.has_value() && values.size() != static_cast<size_t>(*dims)) {
- return parser.emitError(parser.getCurrentLocation())
- << "dims(" << *dims << ") specified but " << values.size()
- << " values provided";
+ // Format: num_teams(lower : type to upper : type)
+ OpAsmParser::UnresolvedOperand lowerOperand;
+ if (parser.parseOperand(lowerOperand) || parser.parseColon() ||
+ parser.parseType(lowerBoundType)) {
+ return failure();
}
- // If dims not specified but we have values, it's implicitly unidimensional
- if (!dims.has_value() && values.size() != 1) {
+ if (failed(parser.parseKeyword("to"))) {
return parser.emitError(parser.getCurrentLocation())
- << "expected 1 value without dims modifier, but got "
- << values.size() << " values";
+ << "expected 'to' keyword in num_teams clause";
}
- // Convert to IntegerAttr
- if (dims.has_value()) {
- dimsAttr = parser.getBuilder().getI64IntegerAttr(*dims);
+ OpAsmParser::UnresolvedOperand upperOperand;
+ if (parser.parseOperand(upperOperand) || parser.parseColon() ||
+ parser.parseType(upperBoundType)) {
+ return failure();
}
+
+ lowerBound = lowerOperand;
+ upperBound = upperOperand;
return success();
}
-static void printDimsModifier(OpAsmPrinter &p, Operation *op,
- IntegerAttr dimsAttr, OperandRange values,
- TypeRange types) {
- // Print dims modifier if present
+//===----------------------------------------------------------------------===//
+// Helper: Print dims modifier with values
+//===----------------------------------------------------------------------===//
+// Prints: dims(N): values : type (single type for all values)
+static void printDimsModifierWithValues(OpAsmPrinter &p, IntegerAttr dimsAttr,
+ OperandRange values, TypeRange types) {
if (dimsAttr) {
p << "dims(" << dimsAttr.getInt() << "): ";
}
- // Print operands
p.printOperands(values);
- // Print types
+ // Print single type
p << " : ";
- llvm::interleaveComma(types, p);
+ if (!types.empty()) {
+ p << types.front();
+ }
+}
+
+static void printNumTeamsClause(OpAsmPrinter &p, Operation *op,
+ IntegerAttr dimsAttr, OperandRange values,
+ TypeRange types, Value lowerBound,
+ Type lowerBoundType, Value upperBound,
+ Type upperBoundType) {
+ if (!values.empty()) {
+ // Multidimensional: dims(N): values : type
+ printDimsModifierWithValues(p, dimsAttr, values, types);
+ } else if (upperBound) {
+ if (lowerBound) {
+ // Both bounds: lower : type to upper : type
+ p.printOperand(lowerBound);
+ p << " : " << lowerBoundType << " to ";
+ p.printOperand(upperBound);
+ p << " : " << upperBoundType;
+ } else {
+ // Upper only: to upper : type
+ p << " to ";
+ p.printOperand(upperBound);
+ p << " : " << upperBoundType;
+ }
+ }
}
#define GET_ATTRDEF_CLASSES
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index cd06011c2cbc4..a331b45a7eb78 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1451,7 +1451,82 @@ func.func @omp_teams_num_teams1(%lb : i32) {
// expected-error @below {{expected num_teams upper bound to be defined if the lower bound is defined}}
"omp.teams" (%lb) ({
omp.terminator
- }) {operandSegmentSizes = array<i32: 0,0,0,1,0,0,0,0,0>} : (i32) -> ()
+ }) {operandSegmentSizes = array<i32: 0,0,0,0,1,0,0,0,0>} : (i32) -> ()
+ omp.terminator
+ }
+ return
+}
+
+// -----
+
+func.func @omp_teams_num_teams_dims_mismatch() {
+ omp.target {
+ %v0 = arith.constant 1 : i32
+ %v1 = arith.constant 2 : i32
+ // expected-error @below {{num_teams dims(3) specified but 2 values provided}}
+ "omp.teams" (%v0, %v1) ({
+ omp.terminator
+ }) {num_teams_dims = 3 : i64, operandSegmentSizes = array<i32: 0,0,0,2,0,0,0,0,0>} : (i32, i32) -> ()
+ omp.terminator
+ }
+ return
+}
+
+// -----
+
+func.func @omp_teams_num_teams_dims_with_bounds() {
+ omp.target {
+ %v0 = arith.constant 1 : i32
+ %v1 = arith.constant 2 : i32
+ %lb = arith.constant 3 : i32
+ %ub = arith.constant 4 : i32
+ // expected-error @below {{num_teams with dims modifier cannot be used together with lower/upper bounds (unidimensional style)}}
+ "omp.teams" (%v0, %v1, %lb, %ub) ({
+ omp.terminator
+ }) {num_teams_dims = 2 : i64, operandSegmentSizes = array<i32: 0,0,0,2,1,1,0,0,0>} : (i32, i32, i32, i32) -> ()
+ omp.terminator
+ }
+ return
+}
+
+// -----
+
+func.func @omp_teams_num_teams_values_without_dims() {
+ omp.target {
+ %v0 = arith.constant 1 : i32
+ %v1 = arith.constant 2 : i32
+ // expected-error @below {{num_teams values can only be specified with dims modifier}}
+ "omp.teams" (%v0, %v1) ({
+ omp.terminator
+ }) {operandSegmentSizes = array<i32: 0,0,0,2,0,0,0,0,0>} : (i32, i32) -> ()
+ omp.terminator
+ }
+ return
+}
+
+// -----
+
+func.func @omp_teams_num_teams_dims_no_values() {
+ omp.target {
+ // expected-error @below {{num_teams dims modifier requires values to be specified}}
+ "omp.teams" () ({
+ omp.terminator
+ }) {num_teams_dims = 2 : i64, operandSegmentSizes = array<i32: 0,0,0,0,0,0,0,0,0>} : () -> ()
+ omp.terminator
+ }
+ return
+}
+
+// -----
+
+func.func @omp_teams_num_teams_dims_type_mismatch() {
+ omp.target {
+ %v0 = arith.constant 1 : i32
+ %v1 = arith.constant 2 : i64
+ // expected-error @below {{num_teams dims modifier requires all values to have the same type}}
+ "omp.teams" (%v0, %v1) ({
+ omp.terminator
+ }) {num_teams_dims = 2 : i64, operandSegmentSizes = array<i32: 0,0,0,2,0,0,0,0,0>} : (i32, i64) -> ()
omp.terminator
}
return
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 82f67d961fa10..998ddcd07fd45 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -1109,8 +1109,8 @@ func.func @omp_teams(%lb : i32, %ub : i32, %if_cond : i1, %num_threads : i32,
omp.terminator
}
- // CHECK: omp.teams num_teams_multi_dim(dims(3): %{{.*}}, %{{.*}}, %{{.*}} : i32, i32, i32)
- omp.teams num_teams_multi_dim(dims(3): %lb, %ub, %ub : i32, i32, i32) {
+ // CHECK: omp.teams num_teams(dims(3): %{{.*}}, %{{.*}}, %{{.*}} : i32)
+ omp.teams num_teams(dims(3): %lb, %ub, %ub : i32) {
// CHECK: omp.terminator
omp.terminator
}
>From 559e299e82b1e083d3638a38b7b2f6849d38811c Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Wed, 10 Dec 2025 15:40:37 +0530
Subject: [PATCH 06/14] Mark mlir->llvmir for num_teams with dims as NYI
---
.../Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp | 12 ++++++++++++
1 file changed, 12 insertions(+)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 4e7942e382c8b..d6aa87b153008 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -2024,6 +2024,10 @@ convertOmpTeams(omp::TeamsOp op, llvm::IRBuilderBase &builder,
if (failed(checkImplementationStatus(*op)))
return failure();
+ if (op.getNumTeamsDims().has_value() || !op.getNumTeamsValues().empty()) {
+ return op.emitError("Lowering of num_teams with dims modifier is NYI.");
+ }
+
DenseMap<Value, llvm::Value *> reductionVariableMap;
unsigned numReductionVars = op.getNumReductionVars();
SmallVector<omp::DeclareReductionOp> reductionDecls;
@@ -6035,6 +6039,10 @@ extractHostEvalClauses(omp::TargetOp targetOp, Value &numThreads,
for (Operation *user : blockArg.getUsers()) {
llvm::TypeSwitch<Operation *>(user)
.Case([&](omp::TeamsOp teamsOp) {
+ // num_teams dims and values are not yet supported
+ assert(!teamsOp.getNumTeamsDims().has_value() &&
+ teamsOp.getNumTeamsValues().empty() &&
+ "Lowering of num_teams with dims modifier is NYI.");
if (teamsOp.getNumTeamsLower() == blockArg)
numTeamsLower = hostEvalVar;
else if (teamsOp.getNumTeamsUpper() == blockArg)
@@ -6157,6 +6165,10 @@ initTargetDefaultAttrs(omp::TargetOp targetOp, Operation *capturedOp,
// host_eval, but instead evaluated prior to entry to the region. This
// ensures values are mapped and available inside of the target region.
if (auto teamsOp = castOrGetParentOfType<omp::TeamsOp>(capturedOp)) {
+ // num_teams dims and values are not yet supported
+ assert(!teamsOp.getNumTeamsDims().has_value() &&
+ teamsOp.getNumTeamsValues().empty() &&
+ "Lowering of num_teams with dims modifier is NYI.");
numTeamsLower = teamsOp.getNumTeamsLower();
numTeamsUpper = teamsOp.getNumTeamsUpper();
threadLimit = teamsOp.getThreadLimit();
>From cf1667784b7a984b3a73dc2807d3e2e8e9bb0f36 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Thu, 11 Dec 2025 17:08:24 +0530
Subject: [PATCH 07/14] few more fixes
---
.../mlir/Dialect/OpenMP/OpenMPClauses.td | 35 ++---
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 142 ++++++++++--------
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 8 +-
mlir/test/Dialect/OpenMP/invalid.mlir | 18 +--
4 files changed, 108 insertions(+), 95 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index 0b17a62c88d92..2beb6690f1736 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
@@ -974,15 +974,15 @@ class OpenMP_NumTeamsClauseSkip<
> : OpenMP_Clause<traits, arguments, assemblyFormat, description,
extraClassDeclaration> {
let arguments = (ins
- ConfinedAttr<OptionalAttr<I64Attr>, [IntPositive]>:$num_teams_dims,
- Variadic<AnyInteger>:$num_teams_values,
+ ConfinedAttr<OptionalAttr<I64Attr>, [IntPositive]>:$num_teams_num_dims,
+ Variadic<AnyInteger>:$num_teams_dims_values,
Optional<AnyInteger>:$num_teams_lower,
Optional<AnyInteger>:$num_teams_upper
);
let optAssemblyFormat = [{
`num_teams` `(` custom<NumTeamsClause>(
- $num_teams_dims, $num_teams_values, type($num_teams_values),
+ $num_teams_num_dims, $num_teams_dims_values, type($num_teams_dims_values),
$num_teams_lower, type($num_teams_lower),
$num_teams_upper, type($num_teams_upper)
) `)`
@@ -993,7 +993,7 @@ class OpenMP_NumTeamsClauseSkip<
construct on which it appears.
With dims modifier: (OpenMP 6.1 requirement)
- - Uses `num_teams_dims` (dimension count) and `num_teams_values` (upper bounds list)
+ - Uses `num_teams_num_dims` (dimension count) and `num_teams_dims_values` (upper bounds list)
- Specifies upper bounds for each dimension (all must have same type)
- Format: `num_teams(dims(N): upper_bound_0, ..., upper_bound_N-1 : type)`
- Example: `num_teams(dims(3): %ub0, %ub1, %ub2 : i32)`
@@ -1007,28 +1007,23 @@ class OpenMP_NumTeamsClauseSkip<
let extraClassDeclaration = [{
/// Returns true if the dims modifier is explicitly present
- bool hasDimsModifier() {
- return getNumTeamsDims().has_value();
+ bool hasNumTeamsDimsModifier() {
+ return getNumTeamsNumDims().has_value() && getNumTeamsNumDims().value();
}
- /// Returns the number of dimensions specified by dims modifier
- unsigned getNumDimensions() {
- if (!hasDimsModifier())
+ /// Returns the number of dimensions specified for num_teams
+ unsigned getNumTeamsDimsCount() {
+ if (!hasNumTeamsDimsModifier())
return 1;
- return static_cast<unsigned>(*getNumTeamsDims());
- }
-
- /// Returns all dimension values as an operand range
- ::mlir::OperandRange getDimensionValues() {
- return getNumTeamsValues();
+ return static_cast<unsigned>(*getNumTeamsNumDims());
}
/// Returns the value for a specific dimension index
- /// Index must be less than getNumDimensions()
- ::mlir::Value getDimensionValue(unsigned index) {
- assert(index < getDimensionValues().size() &&
- "Dimension index out of bounds");
- return getDimensionValues()[index];
+ /// Index must be less than getNumTeamsDimsCount()
+ ::mlir::Value getNumTeamsDimsValue(unsigned index) {
+ assert(index < getNumTeamsDimsCount() &&
+ "Num teams dims index out of bounds");
+ return getNumTeamsDimsValues()[index];
}
}];
}
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 17e286beaf51c..23ef98d307f83 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2197,6 +2197,40 @@ LogicalResult TargetUpdateOp::verify() {
// TargetOp
//===----------------------------------------------------------------------===//
+// Helper: Verify dims modifier
+static LogicalResult verifyDimsModifier(Operation *op,
+ std::optional<IntegerAttr> numDimsAttr,
+ OperandRange dimsValues) {
+ if (numDimsAttr.has_value() && numDimsAttr.value()) {
+ if (dimsValues.empty()) {
+ return op->emitError("dims modifier requires values to be specified");
+ }
+
+ if (dimsValues.size() != static_cast<size_t>(numDimsAttr->getInt())) {
+ return op->emitError("dims(")
+ << numDimsAttr->getInt() << ") specified but " << dimsValues.size()
+ << " values provided";
+ }
+
+ if (!dimsValues.empty()) {
+ Type firstType = dimsValues.front().getType();
+ for (auto value : dimsValues) {
+ if (value.getType() != firstType) {
+ return op->emitError(
+ "dims modifier requires all values to have the same type");
+ }
+ }
+ }
+ return success();
+ } else {
+ if (!dimsValues.empty()) {
+ return op->emitError(
+ "dims values can only be specified with dims modifier");
+ }
+ }
+ return success();
+}
+
void TargetOp::build(OpBuilder &builder, OperationState &state,
const TargetOperands &clauses) {
MLIRContext *ctx = builder.getContext();
@@ -2624,15 +2658,49 @@ void TeamsOp::build(OpBuilder &builder, OperationState &state,
const TeamsOperands &clauses) {
MLIRContext *ctx = builder.getContext();
// TODO Store clauses in op: privateVars, privateSyms, privateNeedsBarrier
- TeamsOp::build(builder, state, clauses.allocateVars, clauses.allocatorVars,
- clauses.ifExpr, clauses.numTeamsDims, clauses.numTeamsValues,
- clauses.numTeamsLower, clauses.numTeamsUpper,
- /*private_vars=*/{}, /*private_syms=*/nullptr,
- /*private_needs_barrier=*/nullptr, clauses.reductionMod,
- clauses.reductionVars,
- makeDenseBoolArrayAttr(ctx, clauses.reductionByref),
- makeArrayAttr(ctx, clauses.reductionSyms),
- clauses.threadLimit);
+ TeamsOp::build(
+ builder, state, clauses.allocateVars, clauses.allocatorVars,
+ clauses.ifExpr, clauses.numTeamsNumDims, clauses.numTeamsDimsValues,
+ clauses.numTeamsLower, clauses.numTeamsUpper,
+ /*private_vars=*/{}, /*private_syms=*/nullptr,
+ /*private_needs_barrier=*/nullptr, clauses.reductionMod,
+ clauses.reductionVars,
+ makeDenseBoolArrayAttr(ctx, clauses.reductionByref),
+ makeArrayAttr(ctx, clauses.reductionSyms), clauses.threadLimit);
+}
+
+// Helper: Verify num_teams clause
+static LogicalResult
+verifyNumTeamsClause(Operation *op, std::optional<IntegerAttr> numTeamsNumDims,
+ OperandRange numTeamsDimsValues, Value numTeamsLower,
+ Value numTeamsUpper) {
+ bool hasDimsModifier = numTeamsNumDims.has_value() && numTeamsNumDims.value();
+
+ // Cannot use both dims modifier and unidimensional style
+ if (hasDimsModifier && (numTeamsLower || numTeamsUpper)) {
+ return op->emitError(
+ "num_teams with dims modifier cannot be used together with "
+ "lower/upper bounds");
+ }
+
+ // With dims modifier
+ if (failed(verifyDimsModifier(op, numTeamsNumDims, numTeamsDimsValues)))
+ return failure();
+
+ // Without dims modifier
+ if (!hasDimsModifier) {
+ if (numTeamsLower) {
+ if (!numTeamsUpper)
+ return op->emitError(
+ "expected num_teams upper bound to be defined if the "
+ "lower bound is defined");
+ if (numTeamsLower.getType() != numTeamsUpper.getType())
+ return op->emitError(
+ "expected num_teams upper bound and lower bound to be "
+ "the same type");
+ }
+ }
+ return success();
}
LogicalResult TeamsOp::verify() {
@@ -2648,58 +2716,10 @@ LogicalResult TeamsOp::verify() {
"in any OpenMP dialect operations");
// Check for num_teams clause restrictions
- auto numTeamsDims = getNumTeamsDims();
- auto numTeamsValues = getNumTeamsValues();
- auto numTeamsLower = getNumTeamsLower();
- auto numTeamsUpper = getNumTeamsUpper();
-
- // Cannot use both dims modifier and unidimensional style
- if (numTeamsDims.has_value() && (numTeamsLower || numTeamsUpper)) {
- return emitError(
- "num_teams with dims modifier cannot be used together with "
- "lower/upper bounds (unidimensional style)");
- }
-
- // With dims modifier (multidimensional)
- if (numTeamsDims.has_value()) {
- if (numTeamsValues.empty()) {
- return emitError(
- "num_teams dims modifier requires values to be specified");
- }
-
- if (numTeamsValues.size() != static_cast<size_t>(*numTeamsDims)) {
- return emitError("num_teams dims(")
- << *numTeamsDims << ") specified but " << numTeamsValues.size()
- << " values provided";
- }
-
- // All values must have the same type
- if (!numTeamsValues.empty()) {
- Type firstType = numTeamsValues.front().getType();
- for (auto value : numTeamsValues) {
- if (value.getType() != firstType) {
- return emitError(
- "num_teams dims modifier requires all values to have "
- "the same type");
- }
- }
- }
- } else {
- // Without dims modifier
- if (!numTeamsValues.empty()) {
- return emitError(
- "num_teams values can only be specified with dims modifier");
- }
-
- if (numTeamsLower) {
- if (!numTeamsUpper)
- return emitError("expected num_teams upper bound to be defined if the "
- "lower bound is defined");
- if (numTeamsLower.getType() != numTeamsUpper.getType())
- return emitError("expected num_teams upper bound and lower bound to be "
- "the same type");
- }
- }
+ if (failed(verifyNumTeamsClause(
+ op, this->getNumTeamsNumDimsAttr(), this->getNumTeamsDimsValues(),
+ this->getNumTeamsLower(), this->getNumTeamsUpper())))
+ return failure();
// Check for allocate clause restrictions
if (getAllocateVars().size() != getAllocatorVars().size())
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index d6aa87b153008..4612f6d51f579 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -2024,7 +2024,7 @@ convertOmpTeams(omp::TeamsOp op, llvm::IRBuilderBase &builder,
if (failed(checkImplementationStatus(*op)))
return failure();
- if (op.getNumTeamsDims().has_value() || !op.getNumTeamsValues().empty()) {
+ if (op.hasNumTeamsDimsModifier()) {
return op.emitError("Lowering of num_teams with dims modifier is NYI.");
}
@@ -6040,8 +6040,7 @@ extractHostEvalClauses(omp::TargetOp targetOp, Value &numThreads,
llvm::TypeSwitch<Operation *>(user)
.Case([&](omp::TeamsOp teamsOp) {
// num_teams dims and values are not yet supported
- assert(!teamsOp.getNumTeamsDims().has_value() &&
- teamsOp.getNumTeamsValues().empty() &&
+ assert(!teamsOp.hasNumTeamsDimsModifier() &&
"Lowering of num_teams with dims modifier is NYI.");
if (teamsOp.getNumTeamsLower() == blockArg)
numTeamsLower = hostEvalVar;
@@ -6166,8 +6165,7 @@ initTargetDefaultAttrs(omp::TargetOp targetOp, Operation *capturedOp,
// ensures values are mapped and available inside of the target region.
if (auto teamsOp = castOrGetParentOfType<omp::TeamsOp>(capturedOp)) {
// num_teams dims and values are not yet supported
- assert(!teamsOp.getNumTeamsDims().has_value() &&
- teamsOp.getNumTeamsValues().empty() &&
+ assert(!teamsOp.hasNumTeamsDimsModifier() &&
"Lowering of num_teams with dims modifier is NYI.");
numTeamsLower = teamsOp.getNumTeamsLower();
numTeamsUpper = teamsOp.getNumTeamsUpper();
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index a331b45a7eb78..9a1eb9baca90c 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1463,10 +1463,10 @@ func.func @omp_teams_num_teams_dims_mismatch() {
omp.target {
%v0 = arith.constant 1 : i32
%v1 = arith.constant 2 : i32
- // expected-error @below {{num_teams dims(3) specified but 2 values provided}}
+ // expected-error @below {{dims(3) specified but 2 values provided}}
"omp.teams" (%v0, %v1) ({
omp.terminator
- }) {num_teams_dims = 3 : i64, operandSegmentSizes = array<i32: 0,0,0,2,0,0,0,0,0>} : (i32, i32) -> ()
+ }) {num_teams_num_dims = 3 : i64, operandSegmentSizes = array<i32: 0,0,0,2,0,0,0,0,0>} : (i32, i32) -> ()
omp.terminator
}
return
@@ -1480,10 +1480,10 @@ func.func @omp_teams_num_teams_dims_with_bounds() {
%v1 = arith.constant 2 : i32
%lb = arith.constant 3 : i32
%ub = arith.constant 4 : i32
- // expected-error @below {{num_teams with dims modifier cannot be used together with lower/upper bounds (unidimensional style)}}
+ // expected-error @below {{num_teams with dims modifier cannot be used together with lower/upper bounds}}
"omp.teams" (%v0, %v1, %lb, %ub) ({
omp.terminator
- }) {num_teams_dims = 2 : i64, operandSegmentSizes = array<i32: 0,0,0,2,1,1,0,0,0>} : (i32, i32, i32, i32) -> ()
+ }) {num_teams_num_dims = 2 : i64, operandSegmentSizes = array<i32: 0,0,0,2,1,1,0,0,0>} : (i32, i32, i32, i32) -> ()
omp.terminator
}
return
@@ -1495,7 +1495,7 @@ func.func @omp_teams_num_teams_values_without_dims() {
omp.target {
%v0 = arith.constant 1 : i32
%v1 = arith.constant 2 : i32
- // expected-error @below {{num_teams values can only be specified with dims modifier}}
+ // expected-error @below {{dims values can only be specified with dims modifier}}
"omp.teams" (%v0, %v1) ({
omp.terminator
}) {operandSegmentSizes = array<i32: 0,0,0,2,0,0,0,0,0>} : (i32, i32) -> ()
@@ -1508,10 +1508,10 @@ func.func @omp_teams_num_teams_values_without_dims() {
func.func @omp_teams_num_teams_dims_no_values() {
omp.target {
- // expected-error @below {{num_teams dims modifier requires values to be specified}}
+ // expected-error @below {{dims modifier requires values to be specified}}
"omp.teams" () ({
omp.terminator
- }) {num_teams_dims = 2 : i64, operandSegmentSizes = array<i32: 0,0,0,0,0,0,0,0,0>} : () -> ()
+ }) {num_teams_num_dims = 2 : i64, operandSegmentSizes = array<i32: 0,0,0,0,0,0,0,0,0>} : () -> ()
omp.terminator
}
return
@@ -1523,10 +1523,10 @@ func.func @omp_teams_num_teams_dims_type_mismatch() {
omp.target {
%v0 = arith.constant 1 : i32
%v1 = arith.constant 2 : i64
- // expected-error @below {{num_teams dims modifier requires all values to have the same type}}
+ // expected-error @below {{dims modifier requires all values to have the same type}}
"omp.teams" (%v0, %v1) ({
omp.terminator
- }) {num_teams_dims = 2 : i64, operandSegmentSizes = array<i32: 0,0,0,2,0,0,0,0,0>} : (i32, i64) -> ()
+ }) {num_teams_num_dims = 2 : i64, operandSegmentSizes = array<i32: 0,0,0,2,0,0,0,0,0>} : (i32, i64) -> ()
omp.terminator
}
return
>From df3577971fb0ba415f617fed4bbd9eaa9a06a6fa Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Tue, 16 Dec 2025 19:58:14 +0530
Subject: [PATCH 08/14] allow single dims_values
---
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 23ef98d307f83..d097b9372aeab 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2223,7 +2223,7 @@ static LogicalResult verifyDimsModifier(Operation *op,
}
return success();
} else {
- if (!dimsValues.empty()) {
+ if (dimsValues.size() > 1) {
return op->emitError(
"dims values can only be specified with dims modifier");
}
>From 5845aa4458ece0308d9900493f7c359bca50cf74 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Fri, 19 Dec 2025 09:51:21 +0530
Subject: [PATCH 09/14] Make dims_values as IntLikeType
---
mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td | 2 +-
.../Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp | 11 +++++++----
2 files changed, 8 insertions(+), 5 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index 2beb6690f1736..b949e2629a095 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
@@ -975,7 +975,7 @@ class OpenMP_NumTeamsClauseSkip<
extraClassDeclaration> {
let arguments = (ins
ConfinedAttr<OptionalAttr<I64Attr>, [IntPositive]>:$num_teams_num_dims,
- Variadic<AnyInteger>:$num_teams_dims_values,
+ Variadic<IntLikeType>:$num_teams_dims_values,
Optional<AnyInteger>:$num_teams_lower,
Optional<AnyInteger>:$num_teams_upper
);
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 4612f6d51f579..fb5a27d5e15c7 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -2025,7 +2025,8 @@ convertOmpTeams(omp::TeamsOp op, llvm::IRBuilderBase &builder,
return failure();
if (op.hasNumTeamsDimsModifier()) {
- return op.emitError("Lowering of num_teams with dims modifier is NYI.");
+ return op.emitError(
+ "Lowering of num_teams with dims modifier is not yet implemented.");
}
DenseMap<Value, llvm::Value *> reductionVariableMap;
@@ -6041,7 +6042,8 @@ extractHostEvalClauses(omp::TargetOp targetOp, Value &numThreads,
.Case([&](omp::TeamsOp teamsOp) {
// num_teams dims and values are not yet supported
assert(!teamsOp.hasNumTeamsDimsModifier() &&
- "Lowering of num_teams with dims modifier is NYI.");
+ "Lowering of num_teams with dims modifier is not yet "
+ "implemented.");
if (teamsOp.getNumTeamsLower() == blockArg)
numTeamsLower = hostEvalVar;
else if (teamsOp.getNumTeamsUpper() == blockArg)
@@ -6165,8 +6167,9 @@ initTargetDefaultAttrs(omp::TargetOp targetOp, Operation *capturedOp,
// ensures values are mapped and available inside of the target region.
if (auto teamsOp = castOrGetParentOfType<omp::TeamsOp>(capturedOp)) {
// num_teams dims and values are not yet supported
- assert(!teamsOp.hasNumTeamsDimsModifier() &&
- "Lowering of num_teams with dims modifier is NYI.");
+ assert(
+ !teamsOp.hasNumTeamsDimsModifier() &&
+ "Lowering of num_teams with dims modifier is not yet implemented.");
numTeamsLower = teamsOp.getNumTeamsLower();
numTeamsUpper = teamsOp.getNumTeamsUpper();
threadLimit = teamsOp.getThreadLimit();
>From bc479b54119a831afa74e5137eb5bbeb3b30ddb0 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Fri, 16 Jan 2026 08:48:30 +0530
Subject: [PATCH 10/14] Update num_teams to have just the list and no dims(N)
syntax
---
.../mlir/Dialect/OpenMP/OpenMPClauses.td | 40 +--
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 246 +++++++-----------
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 18 +-
mlir/test/Dialect/OpenMP/invalid.mlir | 64 +----
mlir/test/Dialect/OpenMP/ops.mlir | 15 +-
mlir/test/Target/LLVMIR/openmp-todo.mlir | 11 +
6 files changed, 141 insertions(+), 253 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index b949e2629a095..d4640f254ed1f 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
@@ -974,15 +974,14 @@ class OpenMP_NumTeamsClauseSkip<
> : OpenMP_Clause<traits, arguments, assemblyFormat, description,
extraClassDeclaration> {
let arguments = (ins
- ConfinedAttr<OptionalAttr<I64Attr>, [IntPositive]>:$num_teams_num_dims,
- Variadic<IntLikeType>:$num_teams_dims_values,
+ Variadic<IntLikeType>:$num_teams_vals,
Optional<AnyInteger>:$num_teams_lower,
Optional<AnyInteger>:$num_teams_upper
);
let optAssemblyFormat = [{
`num_teams` `(` custom<NumTeamsClause>(
- $num_teams_num_dims, $num_teams_dims_values, type($num_teams_dims_values),
+ $num_teams_vals, type($num_teams_vals),
$num_teams_lower, type($num_teams_lower),
$num_teams_upper, type($num_teams_upper)
) `)`
@@ -992,13 +991,14 @@ class OpenMP_NumTeamsClauseSkip<
The `num_teams` clause specifies the bounds on the league space formed by the
construct on which it appears.
- With dims modifier: (OpenMP 6.1 requirement)
- - Uses `num_teams_num_dims` (dimension count) and `num_teams_dims_values` (upper bounds list)
- - Specifies upper bounds for each dimension (all must have same type)
- - Format: `num_teams(dims(N): upper_bound_0, ..., upper_bound_N-1 : type)`
- - Example: `num_teams(dims(3): %ub0, %ub1, %ub2 : i32)`
+ Multi-dimensional (OpenMP 6.1 dims modifier):
+ - Uses `num_teams_vals` with multiple values (one per dimension)
+ - Values may have different types; they will be cast at LLVM IR translation
+ - Format: `num_teams(%v0, %v1, ... : type0, type1, ...)`
+ - Example: `num_teams(%ub0, %ub1, %ub2 : i32, i64, i32)`
+ - The number of dimensions is implicitly `num_teams_vals.size()`
- Without dims modifier:
+ Uni-dimensional (legacy format):
- Uses `num_teams_upper` and optional `num_teams_lower`
- If lower bound not specified, it defaults to upper bound value
- Format: `num_teams(lower : type to upper : type)` or `num_teams(to upper : type)`
@@ -1006,24 +1006,24 @@ class OpenMP_NumTeamsClauseSkip<
}];
let extraClassDeclaration = [{
- /// Returns true if the dims modifier is explicitly present
- bool hasNumTeamsDimsModifier() {
- return getNumTeamsNumDims().has_value() && getNumTeamsNumDims().value();
+ /// Returns true if using the multi-dimensional values format
+ bool hasNumTeamsMultiDim() {
+ return !getNumTeamsVals().empty();
}
/// Returns the number of dimensions specified for num_teams
unsigned getNumTeamsDimsCount() {
- if (!hasNumTeamsDimsModifier())
- return 1;
- return static_cast<unsigned>(*getNumTeamsNumDims());
+ if (hasNumTeamsMultiDim())
+ return getNumTeamsVals().size();
+ return getNumTeamsUpper() ? 1 : 0;
}
- /// Returns the value for a specific dimension index
- /// Index must be less than getNumTeamsDimsCount()
- ::mlir::Value getNumTeamsDimsValue(unsigned index) {
- assert(index < getNumTeamsDimsCount() &&
+ /// Returns the value for a specific dimension index (multi-dim format)
+ /// Index must be less than getNumTeamsVals().size()
+ ::mlir::Value getNumTeamsVal(unsigned index) {
+ assert(index < getNumTeamsVals().size() &&
"Num teams dims index out of bounds");
- return getNumTeamsDimsValues()[index];
+ return getNumTeamsVals()[index];
}
}];
}
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index d097b9372aeab..25bf4e70d9a83 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2197,40 +2197,6 @@ LogicalResult TargetUpdateOp::verify() {
// TargetOp
//===----------------------------------------------------------------------===//
-// Helper: Verify dims modifier
-static LogicalResult verifyDimsModifier(Operation *op,
- std::optional<IntegerAttr> numDimsAttr,
- OperandRange dimsValues) {
- if (numDimsAttr.has_value() && numDimsAttr.value()) {
- if (dimsValues.empty()) {
- return op->emitError("dims modifier requires values to be specified");
- }
-
- if (dimsValues.size() != static_cast<size_t>(numDimsAttr->getInt())) {
- return op->emitError("dims(")
- << numDimsAttr->getInt() << ") specified but " << dimsValues.size()
- << " values provided";
- }
-
- if (!dimsValues.empty()) {
- Type firstType = dimsValues.front().getType();
- for (auto value : dimsValues) {
- if (value.getType() != firstType) {
- return op->emitError(
- "dims modifier requires all values to have the same type");
- }
- }
- }
- return success();
- } else {
- if (dimsValues.size() > 1) {
- return op->emitError(
- "dims values can only be specified with dims modifier");
- }
- }
- return success();
-}
-
void TargetOp::build(OpBuilder &builder, OperationState &state,
const TargetOperands &clauses) {
MLIRContext *ctx = builder.getContext();
@@ -2658,48 +2624,43 @@ void TeamsOp::build(OpBuilder &builder, OperationState &state,
const TeamsOperands &clauses) {
MLIRContext *ctx = builder.getContext();
// TODO Store clauses in op: privateVars, privateSyms, privateNeedsBarrier
- TeamsOp::build(
- builder, state, clauses.allocateVars, clauses.allocatorVars,
- clauses.ifExpr, clauses.numTeamsNumDims, clauses.numTeamsDimsValues,
- clauses.numTeamsLower, clauses.numTeamsUpper,
- /*private_vars=*/{}, /*private_syms=*/nullptr,
- /*private_needs_barrier=*/nullptr, clauses.reductionMod,
- clauses.reductionVars,
- makeDenseBoolArrayAttr(ctx, clauses.reductionByref),
- makeArrayAttr(ctx, clauses.reductionSyms), clauses.threadLimit);
-}
-
-// Helper: Verify num_teams clause
-static LogicalResult
-verifyNumTeamsClause(Operation *op, std::optional<IntegerAttr> numTeamsNumDims,
- OperandRange numTeamsDimsValues, Value numTeamsLower,
- Value numTeamsUpper) {
- bool hasDimsModifier = numTeamsNumDims.has_value() && numTeamsNumDims.value();
-
- // Cannot use both dims modifier and unidimensional style
- if (hasDimsModifier && (numTeamsLower || numTeamsUpper)) {
- return op->emitError(
- "num_teams with dims modifier cannot be used together with "
- "lower/upper bounds");
+ TeamsOp::build(builder, state, clauses.allocateVars, clauses.allocatorVars,
+ clauses.ifExpr, clauses.numTeamsVals, clauses.numTeamsLower,
+ clauses.numTeamsUpper,
+ /*private_vars=*/{}, /*private_syms=*/nullptr,
+ /*private_needs_barrier=*/nullptr, clauses.reductionMod,
+ clauses.reductionVars,
+ makeDenseBoolArrayAttr(ctx, clauses.reductionByref),
+ makeArrayAttr(ctx, clauses.reductionSyms),
+ clauses.threadLimit);
+}
+
+// Verify num_teams clause
+static LogicalResult verifyNumTeamsClause(Operation *op,
+ OperandRange numTeamsVals,
+ Value numTeamsLower,
+ Value numTeamsUpper) {
+ bool hasLegacyOperands = numTeamsLower || numTeamsUpper;
+
+ // Cannot use both multi-dimensional and legacy format simultaneously
+ if (!numTeamsVals.empty() && hasLegacyOperands) {
+ return op->emitError()
+ << "num_teams multi-dimensional values cannot be used together with "
+ "legacy lower/upper bounds";
}
- // With dims modifier
- if (failed(verifyDimsModifier(op, numTeamsNumDims, numTeamsDimsValues)))
- return failure();
-
- // Without dims modifier
- if (!hasDimsModifier) {
- if (numTeamsLower) {
- if (!numTeamsUpper)
- return op->emitError(
- "expected num_teams upper bound to be defined if the "
- "lower bound is defined");
- if (numTeamsLower.getType() != numTeamsUpper.getType())
- return op->emitError(
- "expected num_teams upper bound and lower bound to be "
- "the same type");
- }
+ // If lower is specified, upper must also be specified
+ if (numTeamsLower) {
+ if (!numTeamsUpper)
+ return op->emitError(
+ "expected num_teams upper bound to be defined if the "
+ "lower bound is defined");
+ if (numTeamsLower.getType() != numTeamsUpper.getType())
+ return op->emitError(
+ "expected num_teams upper bound and lower bound to be "
+ "the same type");
}
+
return success();
}
@@ -2716,9 +2677,9 @@ LogicalResult TeamsOp::verify() {
"in any OpenMP dialect operations");
// Check for num_teams clause restrictions
- if (failed(verifyNumTeamsClause(
- op, this->getNumTeamsNumDimsAttr(), this->getNumTeamsDimsValues(),
- this->getNumTeamsLower(), this->getNumTeamsUpper())))
+ if (failed(verifyNumTeamsClause(op, this->getNumTeamsVals(),
+ this->getNumTeamsLower(),
+ this->getNumTeamsUpper())))
return failure();
// Check for allocate clause restrictions
@@ -4553,46 +4514,14 @@ void DeclareSimdOp::build(OpBuilder &odsBuilder, OperationState &odsState,
}
//===----------------------------------------------------------------------===//
-// Helper: Parse dims modifier with values
-//===----------------------------------------------------------------------===//
-// Parses: dims(N): values : type (single type for all values)
-static ParseResult parseDimsModifierWithValues(
- OpAsmParser &parser, IntegerAttr &dimsAttr,
- SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
- SmallVectorImpl<Type> &types) {
- if (failed(parser.parseOptionalKeyword("dims"))) {
- return failure();
- }
-
- // Parse (N): values : type
- int64_t dimsValue;
- if (parser.parseLParen() || parser.parseInteger(dimsValue) ||
- parser.parseRParen() || parser.parseColon()) {
- return failure();
- }
-
- if (parser.parseOperandList(values) || parser.parseColon()) {
- return failure();
- }
-
- // Parse single type (all values have same type)
- Type valueType;
- if (parser.parseType(valueType)) {
- return failure();
- }
-
- // Fill types vector with same type for all values
- types.assign(values.size(), valueType);
-
- dimsAttr = parser.getBuilder().getI64IntegerAttr(dimsValue);
- return success();
-}
-
-//===----------------------------------------------------------------------===//
-// Parser and printer for num_teams clause with dims modifier
+// Parser and printer for num_teams clause
//===----------------------------------------------------------------------===//
+// Parses num_teams clause in two formats:
+// 1. Multi-dimensional: num_teams(%v0, %v1, ... : type0, type1, ...)
+// 2. Uni-dimensional (legacy): num_teams(%lb : type to %ub : type)
+// or num_teams(to %ub : type)
static ParseResult
-parseNumTeamsClause(OpAsmParser &parser, IntegerAttr &dimsAttr,
+parseNumTeamsClause(OpAsmParser &parser,
SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
SmallVectorImpl<Type> &types,
std::optional<OpAsmParser::UnresolvedOperand> &lowerBound,
@@ -4600,11 +4529,6 @@ parseNumTeamsClause(OpAsmParser &parser, IntegerAttr &dimsAttr,
std::optional<OpAsmParser::UnresolvedOperand> &upperBound,
Type &upperBoundType) {
- // Format: num_teams(dims(N): values : type)
- if (succeeded(parseDimsModifierWithValues(parser, dimsAttr, values, types))) {
- return success();
- }
-
// Format: num_teams(to upper : type)
if (succeeded(parser.parseOptionalKeyword("to"))) {
OpAsmParser::UnresolvedOperand upperOperand;
@@ -4616,56 +4540,68 @@ parseNumTeamsClause(OpAsmParser &parser, IntegerAttr &dimsAttr,
return success();
}
- // Format: num_teams(lower : type to upper : type)
- OpAsmParser::UnresolvedOperand lowerOperand;
- if (parser.parseOperand(lowerOperand) || parser.parseColon() ||
- parser.parseType(lowerBoundType)) {
+ // Try to parse multi-dimensional format: values : types
+ // or legacy format: lower : type to upper : type
+ OpAsmParser::UnresolvedOperand firstOperand;
+ if (parser.parseOperand(firstOperand))
return failure();
- }
- if (failed(parser.parseKeyword("to"))) {
- return parser.emitError(parser.getCurrentLocation())
- << "expected 'to' keyword in num_teams clause";
+ // Check if there is a comma (multi-dimensional) or colon
+ if (succeeded(parser.parseOptionalComma())) {
+ values.push_back(firstOperand);
+ if (parser.parseOperandList(values) || parser.parseColon() ||
+ parser.parseTypeList(types)) {
+ return failure();
+ }
+ return success();
}
- OpAsmParser::UnresolvedOperand upperOperand;
- if (parser.parseOperand(upperOperand) || parser.parseColon() ||
- parser.parseType(upperBoundType)) {
+ // Must have a colon next
+ if (parser.parseColon())
return failure();
- }
- lowerBound = lowerOperand;
- upperBound = upperOperand;
- return success();
-}
+ Type firstType;
+ if (parser.parseType(firstType))
+ return failure();
-//===----------------------------------------------------------------------===//
-// Helper: Print dims modifier with values
-//===----------------------------------------------------------------------===//
-// Prints: dims(N): values : type (single type for all values)
-static void printDimsModifierWithValues(OpAsmPrinter &p, IntegerAttr dimsAttr,
- OperandRange values, TypeRange types) {
- if (dimsAttr) {
- p << "dims(" << dimsAttr.getInt() << "): ";
- }
+ // Check for 'to' keyword (legacy format)
+ if (succeeded(parser.parseOptionalKeyword("to"))) {
+ lowerBound = firstOperand;
+ lowerBoundType = firstType;
- p.printOperands(values);
+ OpAsmParser::UnresolvedOperand upperOperand;
+ if (parser.parseOperand(upperOperand) || parser.parseColon() ||
+ parser.parseType(upperBoundType)) {
+ return failure();
+ }
+ upperBound = upperOperand;
+ return success();
+ }
- // Print single type
- p << " : ";
- if (!types.empty()) {
- p << types.front();
+ // Check for comma after type (multi-dimensional with types)
+ if (succeeded(parser.parseOptionalComma())) {
+ values.push_back(firstOperand);
+ types.push_back(firstType);
+ if (parser.parseTypeList(types))
+ return failure();
+ return success();
}
+
+ // Single value multi-dimensional format: %v0 : type0
+ values.push_back(firstOperand);
+ types.push_back(firstType);
+ return success();
}
static void printNumTeamsClause(OpAsmPrinter &p, Operation *op,
- IntegerAttr dimsAttr, OperandRange values,
- TypeRange types, Value lowerBound,
- Type lowerBoundType, Value upperBound,
- Type upperBoundType) {
+ OperandRange values, TypeRange types,
+ Value lowerBound, Type lowerBoundType,
+ Value upperBound, Type upperBoundType) {
if (!values.empty()) {
- // Multidimensional: dims(N): values : type
- printDimsModifierWithValues(p, dimsAttr, values, types);
+ // Multi-dimensional: %v0, %v1, ... : type0, type1, ...
+ llvm::interleaveComma(values, p, [&](Value v) { p.printOperand(v); });
+ p << " : ";
+ llvm::interleaveComma(types, p);
} else if (upperBound) {
if (lowerBound) {
// Both bounds: lower : type to upper : type
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index fb5a27d5e15c7..26ed835e856ed 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -376,6 +376,10 @@ static LogicalResult checkImplementationStatus(Operation &op) {
op.getTaskReductionSyms())
result = todo("task_reduction");
};
+ auto checkNumTeamsMultiDim = [&todo](auto op, LogicalResult &result) {
+ if (op.hasNumTeamsMultiDim())
+ result = todo("num_teams with multi-dimensional values");
+ };
LogicalResult result = success();
llvm::TypeSwitch<Operation &>(op)
@@ -400,6 +404,7 @@ static LogicalResult checkImplementationStatus(Operation &op) {
.Case([&](omp::TeamsOp op) {
checkAllocate(op, result);
checkPrivate(op, result);
+ checkNumTeamsMultiDim(op, result);
})
.Case([&](omp::TaskOp op) {
checkAllocate(op, result);
@@ -2024,11 +2029,6 @@ convertOmpTeams(omp::TeamsOp op, llvm::IRBuilderBase &builder,
if (failed(checkImplementationStatus(*op)))
return failure();
- if (op.hasNumTeamsDimsModifier()) {
- return op.emitError(
- "Lowering of num_teams with dims modifier is not yet implemented.");
- }
-
DenseMap<Value, llvm::Value *> reductionVariableMap;
unsigned numReductionVars = op.getNumReductionVars();
SmallVector<omp::DeclareReductionOp> reductionDecls;
@@ -6040,10 +6040,6 @@ extractHostEvalClauses(omp::TargetOp targetOp, Value &numThreads,
for (Operation *user : blockArg.getUsers()) {
llvm::TypeSwitch<Operation *>(user)
.Case([&](omp::TeamsOp teamsOp) {
- // num_teams dims and values are not yet supported
- assert(!teamsOp.hasNumTeamsDimsModifier() &&
- "Lowering of num_teams with dims modifier is not yet "
- "implemented.");
if (teamsOp.getNumTeamsLower() == blockArg)
numTeamsLower = hostEvalVar;
else if (teamsOp.getNumTeamsUpper() == blockArg)
@@ -6166,10 +6162,6 @@ initTargetDefaultAttrs(omp::TargetOp targetOp, Operation *capturedOp,
// host_eval, but instead evaluated prior to entry to the region. This
// ensures values are mapped and available inside of the target region.
if (auto teamsOp = castOrGetParentOfType<omp::TeamsOp>(capturedOp)) {
- // num_teams dims and values are not yet supported
- assert(
- !teamsOp.hasNumTeamsDimsModifier() &&
- "Lowering of num_teams with dims modifier is not yet implemented.");
numTeamsLower = teamsOp.getNumTeamsLower();
numTeamsUpper = teamsOp.getNumTeamsUpper();
threadLimit = teamsOp.getThreadLimit();
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 9a1eb9baca90c..d451b14e8bfc9 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1459,74 +1459,16 @@ func.func @omp_teams_num_teams1(%lb : i32) {
// -----
-func.func @omp_teams_num_teams_dims_mismatch() {
- omp.target {
- %v0 = arith.constant 1 : i32
- %v1 = arith.constant 2 : i32
- // expected-error @below {{dims(3) specified but 2 values provided}}
- "omp.teams" (%v0, %v1) ({
- omp.terminator
- }) {num_teams_num_dims = 3 : i64, operandSegmentSizes = array<i32: 0,0,0,2,0,0,0,0,0>} : (i32, i32) -> ()
- omp.terminator
- }
- return
-}
-
-// -----
-
-func.func @omp_teams_num_teams_dims_with_bounds() {
+func.func @omp_teams_num_teams_multidim_with_bounds() {
omp.target {
%v0 = arith.constant 1 : i32
%v1 = arith.constant 2 : i32
%lb = arith.constant 3 : i32
%ub = arith.constant 4 : i32
- // expected-error @below {{num_teams with dims modifier cannot be used together with lower/upper bounds}}
+ // expected-error @below {{num_teams multi-dimensional values cannot be used together with legacy lower/upper bounds}}
"omp.teams" (%v0, %v1, %lb, %ub) ({
omp.terminator
- }) {num_teams_num_dims = 2 : i64, operandSegmentSizes = array<i32: 0,0,0,2,1,1,0,0,0>} : (i32, i32, i32, i32) -> ()
- omp.terminator
- }
- return
-}
-
-// -----
-
-func.func @omp_teams_num_teams_values_without_dims() {
- omp.target {
- %v0 = arith.constant 1 : i32
- %v1 = arith.constant 2 : i32
- // expected-error @below {{dims values can only be specified with dims modifier}}
- "omp.teams" (%v0, %v1) ({
- omp.terminator
- }) {operandSegmentSizes = array<i32: 0,0,0,2,0,0,0,0,0>} : (i32, i32) -> ()
- omp.terminator
- }
- return
-}
-
-// -----
-
-func.func @omp_teams_num_teams_dims_no_values() {
- omp.target {
- // expected-error @below {{dims modifier requires values to be specified}}
- "omp.teams" () ({
- omp.terminator
- }) {num_teams_num_dims = 2 : i64, operandSegmentSizes = array<i32: 0,0,0,0,0,0,0,0,0>} : () -> ()
- omp.terminator
- }
- return
-}
-
-// -----
-
-func.func @omp_teams_num_teams_dims_type_mismatch() {
- omp.target {
- %v0 = arith.constant 1 : i32
- %v1 = arith.constant 2 : i64
- // expected-error @below {{dims modifier requires all values to have the same type}}
- "omp.teams" (%v0, %v1) ({
- omp.terminator
- }) {num_teams_num_dims = 2 : i64, operandSegmentSizes = array<i32: 0,0,0,2,0,0,0,0,0>} : (i32, i64) -> ()
+ }) {operandSegmentSizes = array<i32: 0,0,0,2,1,1,0,0,0>} : (i32, i32, i32, i32) -> ()
omp.terminator
}
return
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 998ddcd07fd45..49a88e0443e60 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -1077,7 +1077,7 @@ func.func @parallel_wsloop_reduction(%lb : index, %ub : index, %step : index) {
// CHECK-LABEL: omp_teams
func.func @omp_teams(%lb : i32, %ub : i32, %if_cond : i1, %num_threads : i32,
- %data_var : memref<i32>) -> () {
+ %data_var : memref<i32>, %ub64 : i64, %ub16 : i16) -> () {
// Test nesting inside of omp.target
omp.target {
// CHECK: omp.teams
@@ -1109,8 +1109,15 @@ func.func @omp_teams(%lb : i32, %ub : i32, %if_cond : i1, %num_threads : i32,
omp.terminator
}
- // CHECK: omp.teams num_teams(dims(3): %{{.*}}, %{{.*}}, %{{.*}} : i32)
- omp.teams num_teams(dims(3): %lb, %ub, %ub : i32) {
+ // CHECK: omp.teams num_teams(%{{.*}}, %{{.*}}, %{{.*}} : i32, i32, i32)
+ omp.teams num_teams(%lb, %ub, %ub : i32, i32, i32) {
+ // CHECK: omp.terminator
+ omp.terminator
+ }
+
+ // Test num_teams with mixed types.
+ // CHECK: omp.teams num_teams(%{{.*}}, %{{.*}}, %{{.*}} : i32, i64, i16)
+ omp.teams num_teams(%lb, %ub64, %ub16 : i32, i64, i16) {
// CHECK: omp.terminator
omp.terminator
}
@@ -3080,7 +3087,7 @@ func.func @omp_target_host_eval(%x : i32) {
// CHECK: omp.teams num_teams( to %[[HOST_ARG]] : i32)
// CHECK-SAME: thread_limit(%[[HOST_ARG]] : i32)
omp.target host_eval(%x -> %arg0 : i32) {
- omp.teams num_teams(to %arg0 : i32) thread_limit(%arg0 : i32) {
+ omp.teams num_teams( to %arg0 : i32) thread_limit(%arg0 : i32) {
omp.terminator
}
omp.terminator
diff --git a/mlir/test/Target/LLVMIR/openmp-todo.mlir b/mlir/test/Target/LLVMIR/openmp-todo.mlir
index cb7ea47ea035d..1ea56fdd0bf16 100644
--- a/mlir/test/Target/LLVMIR/openmp-todo.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-todo.mlir
@@ -432,6 +432,17 @@ llvm.func @teams_private(%x : !llvm.ptr) {
// -----
+llvm.func @teams_num_teams_multi_dim(%lb : i32, %ub : i32) {
+ // expected-error at below {{not yet implemented: Unhandled clause num_teams with multi-dimensional values in omp.teams operation}}
+ // expected-error at below {{LLVM Translation failed for operation: omp.teams}}
+ omp.teams num_teams(%lb, %ub, %ub : i32, i32, i32) {
+ omp.terminator
+ }
+ llvm.return
+}
+
+// -----
+
llvm.func @wsloop_allocate(%lb : i32, %ub : i32, %step : i32, %x : !llvm.ptr) {
// expected-error at below {{not yet implemented: Unhandled clause allocate in omp.wsloop operation}}
// expected-error at below {{LLVM Translation failed for operation: omp.wsloop}}
>From b8e7cc44d614abfd57f834928858101acd318b23 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Sat, 17 Jan 2026 10:50:56 +0530
Subject: [PATCH 11/14] name change to getNumTeams()
---
mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td | 6 +++---
.../LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp | 4 ++--
2 files changed, 5 insertions(+), 5 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index d4640f254ed1f..b612d4e136baf 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
@@ -1020,10 +1020,10 @@ class OpenMP_NumTeamsClauseSkip<
/// Returns the value for a specific dimension index (multi-dim format)
/// Index must be less than getNumTeamsVals().size()
- ::mlir::Value getNumTeamsVal(unsigned index) {
- assert(index < getNumTeamsVals().size() &&
+ ::mlir::Value getNumTeams(unsigned dim = 0) {
+ assert(dim < getNumTeamsDimsCount() &&
"Num teams dims index out of bounds");
- return getNumTeamsVals()[index];
+ return getNumTeamsVals()[dim];
}
}];
}
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 26ed835e856ed..0b7bf64cefe4c 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -376,7 +376,7 @@ static LogicalResult checkImplementationStatus(Operation &op) {
op.getTaskReductionSyms())
result = todo("task_reduction");
};
- auto checkNumTeamsMultiDim = [&todo](auto op, LogicalResult &result) {
+ auto checkNumTeams = [&todo](auto op, LogicalResult &result) {
if (op.hasNumTeamsMultiDim())
result = todo("num_teams with multi-dimensional values");
};
@@ -404,7 +404,7 @@ static LogicalResult checkImplementationStatus(Operation &op) {
.Case([&](omp::TeamsOp op) {
checkAllocate(op, result);
checkPrivate(op, result);
- checkNumTeamsMultiDim(op, result);
+ checkNumTeams(op, result);
})
.Case([&](omp::TaskOp op) {
checkAllocate(op, result);
>From fab5e028ae97d848029f02a32cd307b80e07f999 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Tue, 20 Jan 2026 10:46:26 +0530
Subject: [PATCH 12/14] rename num_teams_vals to num_teams_vars
---
.../include/mlir/Dialect/OpenMP/OpenMPClauses.td | 16 ++++++++--------
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 8 ++++----
2 files changed, 12 insertions(+), 12 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index b612d4e136baf..8fc2a4bb3ec96 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
@@ -974,14 +974,14 @@ class OpenMP_NumTeamsClauseSkip<
> : OpenMP_Clause<traits, arguments, assemblyFormat, description,
extraClassDeclaration> {
let arguments = (ins
- Variadic<IntLikeType>:$num_teams_vals,
+ Variadic<IntLikeType>:$num_teams_vars,
Optional<AnyInteger>:$num_teams_lower,
Optional<AnyInteger>:$num_teams_upper
);
let optAssemblyFormat = [{
`num_teams` `(` custom<NumTeamsClause>(
- $num_teams_vals, type($num_teams_vals),
+ $num_teams_vars, type($num_teams_vars),
$num_teams_lower, type($num_teams_lower),
$num_teams_upper, type($num_teams_upper)
) `)`
@@ -992,11 +992,11 @@ class OpenMP_NumTeamsClauseSkip<
construct on which it appears.
Multi-dimensional (OpenMP 6.1 dims modifier):
- - Uses `num_teams_vals` with multiple values (one per dimension)
+ - Uses `num_teams_vars` with multiple values (one per dimension)
- Values may have different types; they will be cast at LLVM IR translation
- Format: `num_teams(%v0, %v1, ... : type0, type1, ...)`
- Example: `num_teams(%ub0, %ub1, %ub2 : i32, i64, i32)`
- - The number of dimensions is implicitly `num_teams_vals.size()`
+ - The number of dimensions is implicitly `num_teams_vars.size()`
Uni-dimensional (legacy format):
- Uses `num_teams_upper` and optional `num_teams_lower`
@@ -1008,22 +1008,22 @@ class OpenMP_NumTeamsClauseSkip<
let extraClassDeclaration = [{
/// Returns true if using the multi-dimensional values format
bool hasNumTeamsMultiDim() {
- return !getNumTeamsVals().empty();
+ return !getNumTeamsVars().empty();
}
/// Returns the number of dimensions specified for num_teams
unsigned getNumTeamsDimsCount() {
if (hasNumTeamsMultiDim())
- return getNumTeamsVals().size();
+ return getNumTeamsVars().size();
return getNumTeamsUpper() ? 1 : 0;
}
/// Returns the value for a specific dimension index (multi-dim format)
- /// Index must be less than getNumTeamsVals().size()
+ /// Index must be less than getNumTeamsVars().size()
::mlir::Value getNumTeams(unsigned dim = 0) {
assert(dim < getNumTeamsDimsCount() &&
"Num teams dims index out of bounds");
- return getNumTeamsVals()[dim];
+ return getNumTeamsVars()[dim];
}
}];
}
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 25bf4e70d9a83..0375e52b1e034 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2625,7 +2625,7 @@ void TeamsOp::build(OpBuilder &builder, OperationState &state,
MLIRContext *ctx = builder.getContext();
// TODO Store clauses in op: privateVars, privateSyms, privateNeedsBarrier
TeamsOp::build(builder, state, clauses.allocateVars, clauses.allocatorVars,
- clauses.ifExpr, clauses.numTeamsVals, clauses.numTeamsLower,
+ clauses.ifExpr, clauses.numTeamsVars, clauses.numTeamsLower,
clauses.numTeamsUpper,
/*private_vars=*/{}, /*private_syms=*/nullptr,
/*private_needs_barrier=*/nullptr, clauses.reductionMod,
@@ -2637,13 +2637,13 @@ void TeamsOp::build(OpBuilder &builder, OperationState &state,
// Verify num_teams clause
static LogicalResult verifyNumTeamsClause(Operation *op,
- OperandRange numTeamsVals,
+ OperandRange numTeamsVars,
Value numTeamsLower,
Value numTeamsUpper) {
bool hasLegacyOperands = numTeamsLower || numTeamsUpper;
// Cannot use both multi-dimensional and legacy format simultaneously
- if (!numTeamsVals.empty() && hasLegacyOperands) {
+ if (!numTeamsVars.empty() && hasLegacyOperands) {
return op->emitError()
<< "num_teams multi-dimensional values cannot be used together with "
"legacy lower/upper bounds";
@@ -2677,7 +2677,7 @@ LogicalResult TeamsOp::verify() {
"in any OpenMP dialect operations");
// Check for num_teams clause restrictions
- if (failed(verifyNumTeamsClause(op, this->getNumTeamsVals(),
+ if (failed(verifyNumTeamsClause(op, this->getNumTeamsVars(),
this->getNumTeamsLower(),
this->getNumTeamsUpper())))
return failure();
>From a173e790953f176bd2b0561c697b66bc6b2dcef6 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Thu, 22 Jan 2026 10:34:48 +0530
Subject: [PATCH 13/14] Use only num_teams_lower and num_teams_upper_vars for
clause
---
flang/lib/Lower/OpenMP/ClauseProcessor.cpp | 30 +++-
flang/lib/Lower/OpenMP/Clauses.cpp | 25 ++-
flang/lib/Lower/OpenMP/OpenMP.cpp | 14 +-
flang/test/Lower/OpenMP/num-teams-dims.f90 | 51 ++++++
llvm/include/llvm/Frontend/OpenMP/ClauseT.h | 3 +-
.../mlir/Dialect/OpenMP/OpenMPClauses.td | 36 ++--
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 162 +++---------------
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 11 +-
mlir/test/Dialect/OpenMP/invalid.mlir | 12 +-
mlir/test/Dialect/OpenMP/ops.mlir | 8 +-
mlir/test/Target/LLVMIR/openmp-todo.mlir | 2 +-
11 files changed, 157 insertions(+), 197 deletions(-)
create mode 100644 flang/test/Lower/OpenMP/num-teams-dims.f90
diff --git a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
index 2f531efaf09aa..f296b35d6ce30 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
@@ -495,17 +495,29 @@ bool ClauseProcessor::processSizes(StatementContext &stmtCtx,
bool ClauseProcessor::processNumTeams(
lower::StatementContext &stmtCtx,
mlir::omp::NumTeamsClauseOps &result) const {
- // TODO Get lower and upper bounds for num_teams when parser is updated to
- // accept both.
if (auto *clause = findUniqueClause<omp::clause::NumTeams>()) {
// The num_teams directive accepts a list of team lower/upper bounds.
- // This is an extension to support grid specification for ompx_bare.
- // Here, only expect a single element in the list.
- assert(clause->v.size() == 1);
- // auto lowerBound = std::get<std::optional<ExprTy>>(clause->v[0]->t);
- auto &upperBound = std::get<ExprTy>(clause->v[0].t);
- result.numTeamsUpper =
- fir::getBase(converter.genExprValue(upperBound, stmtCtx));
+ // With dims modifier (OpenMP 6.1): multiple values for multi-dimensional
+ // Without dims modifier: single value (upper bound only)
+ assert(!clause->v.empty());
+
+ // For single value case, also handle the optional lower bound
+ if (clause->v.size() == 1) {
+ auto &lowerBound = std::get<std::optional<ExprTy>>(clause->v[0].t);
+ if (lowerBound) {
+ result.numTeamsLower =
+ fir::getBase(converter.genExprValue(*lowerBound, stmtCtx));
+ }
+ }
+
+ // Extract upper bounds from all Range elements
+ result.numTeamsUpperVars.reserve(clause->v.size());
+ for (const auto &range : clause->v) {
+ auto &upperBound = std::get<ExprTy>(range.t);
+ result.numTeamsUpperVars.push_back(
+ fir::getBase(converter.genExprValue(upperBound, stmtCtx)));
+ }
+
return true;
}
return false;
diff --git a/flang/lib/Lower/OpenMP/Clauses.cpp b/flang/lib/Lower/OpenMP/Clauses.cpp
index c739249bff211..8d3eb093147d1 100644
--- a/flang/lib/Lower/OpenMP/Clauses.cpp
+++ b/flang/lib/Lower/OpenMP/Clauses.cpp
@@ -1280,10 +1280,27 @@ NumTasks make(const parser::OmpClause::NumTasks &inp,
NumTeams make(const parser::OmpClause::NumTeams &inp,
semantics::SemanticsContext &semaCtx) {
// inp.v -> parser::OmpNumTeamsClause
- auto &t1 = std::get<std::list<parser::ScalarIntExpr>>(inp.v.t);
- assert(!t1.empty());
- List<NumTeams::Range> v{{{/*LowerBound=*/std::nullopt,
- /*UpperBound=*/makeExpr(t1.front(), semaCtx)}}};
+ auto &mods = semantics::OmpGetModifiers(inp.v);
+ auto *dims = semantics::OmpGetUniqueModifier<parser::OmpDimsModifier>(mods);
+ auto *lowerBound =
+ semantics::OmpGetUniqueModifier<parser::OmpLowerBound>(mods);
+ auto &values = std::get<std::list<parser::ScalarIntExpr>>(inp.v.t);
+ assert(!values.empty());
+
+ // With dims modifier: create Range for each value (all upper bounds)
+ if (dims) {
+ List<NumTeams::Range> v;
+ for (const auto &val : values) {
+ v.push_back(NumTeams::Range{{/*LowerBound=*/std::nullopt,
+ /*UpperBound=*/makeExpr(val, semaCtx)}});
+ }
+ return NumTeams{/*List=*/v};
+ }
+
+ // Without dims modifier: single element with optional lower bound
+ auto lb = maybeApplyToV(makeExprFn(semaCtx), lowerBound);
+ List<NumTeams::Range> v{{{/*LowerBound=*/lb,
+ /*UpperBound=*/makeExpr(values.front(), semaCtx)}}};
return NumTeams{/*List=*/v};
}
diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 0764693f748a5..bf295976b76c3 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -96,8 +96,8 @@ class HostEvalInfo {
if (ops.numTeamsLower)
vars.push_back(ops.numTeamsLower);
- if (ops.numTeamsUpper)
- vars.push_back(ops.numTeamsUpper);
+ for (auto numTeamsUpper : ops.numTeamsUpperVars)
+ vars.push_back(numTeamsUpper);
if (ops.numThreads)
vars.push_back(ops.numThreads);
@@ -115,7 +115,7 @@ class HostEvalInfo {
assert(args.size() ==
ops.loopLowerBounds.size() + ops.loopUpperBounds.size() +
ops.loopSteps.size() + (ops.numTeamsLower ? 1 : 0) +
- (ops.numTeamsUpper ? 1 : 0) + (ops.numThreads ? 1 : 0) +
+ ops.numTeamsUpperVars.size() + (ops.numThreads ? 1 : 0) +
(ops.threadLimit ? 1 : 0) &&
"invalid block argument list");
int argIndex = 0;
@@ -131,8 +131,8 @@ class HostEvalInfo {
if (ops.numTeamsLower)
ops.numTeamsLower = args[argIndex++];
- if (ops.numTeamsUpper)
- ops.numTeamsUpper = args[argIndex++];
+ for (size_t i = 0; i < ops.numTeamsUpperVars.size(); ++i)
+ ops.numTeamsUpperVars[i] = args[argIndex++];
if (ops.numThreads)
ops.numThreads = args[argIndex++];
@@ -185,11 +185,11 @@ class HostEvalInfo {
/// \returns whether an update was performed. If not, these clauses were not
/// evaluated in the host device.
bool apply(mlir::omp::TeamsOperands &clauseOps) {
- if (!ops.numTeamsLower && !ops.numTeamsUpper && !ops.threadLimit)
+ if (!ops.numTeamsLower && ops.numTeamsUpperVars.empty() && !ops.threadLimit)
return false;
clauseOps.numTeamsLower = ops.numTeamsLower;
- clauseOps.numTeamsUpper = ops.numTeamsUpper;
+ clauseOps.numTeamsUpperVars = ops.numTeamsUpperVars;
clauseOps.threadLimit = ops.threadLimit;
return true;
}
diff --git a/flang/test/Lower/OpenMP/num-teams-dims.f90 b/flang/test/Lower/OpenMP/num-teams-dims.f90
new file mode 100644
index 0000000000000..d41c3b41fbbea
--- /dev/null
+++ b/flang/test/Lower/OpenMP/num-teams-dims.f90
@@ -0,0 +1,51 @@
+! RUN: %flang_fc1 -emit-hlfir %openmp_flags -fopenmp-version=61 %s -o - | FileCheck %s
+
+!===============================================================================
+! `num_teams` clause with dims modifier (OpenMP 6.1)
+!===============================================================================
+
+! CHECK-LABEL: func @_QPteams_numteams_dims2
+subroutine teams_numteams_dims2()
+ ! CHECK: omp.teams
+ ! CHECK-SAME: num_teams( to %{{.*}}, %{{.*}} : i32, i32)
+ !$omp teams num_teams(dims(2): 10, 4)
+ call f1()
+ ! CHECK: omp.terminator
+ !$omp end teams
+end subroutine teams_numteams_dims2
+
+! CHECK-LABEL: func @_QPteams_numteams_dims3
+subroutine teams_numteams_dims3()
+ ! CHECK: omp.teams
+ ! CHECK-SAME: num_teams( to %{{.*}}, %{{.*}}, %{{.*}} : i32, i32, i32)
+ !$omp teams num_teams(dims(3): 8, 4, 2)
+ call f1()
+ ! CHECK: omp.terminator
+ !$omp end teams
+end subroutine teams_numteams_dims3
+
+! CHECK-LABEL: func @_QPteams_numteams_dims_var
+subroutine teams_numteams_dims_var(a, b, c)
+ integer, intent(in) :: a, b, c
+ ! CHECK: omp.teams
+ ! CHECK-SAME: num_teams( to %{{.*}}, %{{.*}}, %{{.*}} : i32, i32, i32)
+ !$omp teams num_teams(dims(3): a, b, c)
+ call f1()
+ ! CHECK: omp.terminator
+ !$omp end teams
+end subroutine teams_numteams_dims_var
+
+!===============================================================================
+! `num_teams` clause with lower bound (legacy, without dims)
+!===============================================================================
+
+! CHECK-LABEL: func @_QPteams_numteams_lower_upper
+subroutine teams_numteams_lower_upper(lower, upper)
+ integer, intent(in) :: lower, upper
+ ! CHECK: omp.teams
+ ! CHECK-SAME: num_teams(%{{.*}} : i32 to %{{.*}} : i32)
+ !$omp teams num_teams(lower: upper)
+ call f1()
+ ! CHECK: omp.terminator
+ !$omp end teams
+end subroutine teams_numteams_lower_upper
diff --git a/llvm/include/llvm/Frontend/OpenMP/ClauseT.h b/llvm/include/llvm/Frontend/OpenMP/ClauseT.h
index 05ee1ae36a23d..c38eb27e5cef4 100644
--- a/llvm/include/llvm/Frontend/OpenMP/ClauseT.h
+++ b/llvm/include/llvm/Frontend/OpenMP/ClauseT.h
@@ -992,6 +992,7 @@ struct NumTasksT {
};
// V5.2: [10.2.1] `num_teams` clause
+// V6.1: Extended with dims modifier support
template <typename T, typename I, typename E> //
struct NumTeamsT {
using LowerBound = E;
@@ -1004,7 +1005,7 @@ struct NumTeamsT {
};
// The name List is not a spec name. The list is an extension to allow
- // specifying a grid with connection with the ompx_bare clause.
+ // specifying a grid with the dims modifier (OpenMP 6.1) or ompx_bare clause.
using List = ListT<Range>;
using WrapperTrait = std::true_type;
List v;
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index 8fc2a4bb3ec96..9c9702a1ba8cf 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
@@ -974,17 +974,13 @@ class OpenMP_NumTeamsClauseSkip<
> : OpenMP_Clause<traits, arguments, assemblyFormat, description,
extraClassDeclaration> {
let arguments = (ins
- Variadic<IntLikeType>:$num_teams_vars,
- Optional<AnyInteger>:$num_teams_lower,
- Optional<AnyInteger>:$num_teams_upper
+ Optional<IntLikeType>:$num_teams_lower,
+ Variadic<IntLikeType>:$num_teams_upper_vars
);
let optAssemblyFormat = [{
- `num_teams` `(` custom<NumTeamsClause>(
- $num_teams_vars, type($num_teams_vars),
- $num_teams_lower, type($num_teams_lower),
- $num_teams_upper, type($num_teams_upper)
- ) `)`
+ `num_teams` `(` ( $num_teams_lower^ `:` type($num_teams_lower) )? `to`
+ $num_teams_upper_vars `:` type($num_teams_upper_vars) `)`
}];
let description = [{
@@ -992,38 +988,36 @@ class OpenMP_NumTeamsClauseSkip<
construct on which it appears.
Multi-dimensional (OpenMP 6.1 dims modifier):
- - Uses `num_teams_vars` with multiple values (one per dimension)
+ - Uses `num_teams_upper_vars` with multiple values (one per dimension)
- Values may have different types; they will be cast at LLVM IR translation
- - Format: `num_teams(%v0, %v1, ... : type0, type1, ...)`
- - Example: `num_teams(%ub0, %ub1, %ub2 : i32, i64, i32)`
- - The number of dimensions is implicitly `num_teams_vars.size()`
+ - Format: `num_teams(to %v0, %v1, ... : type0, type1, ...)`
+ - Example: `num_teams(to %ub0, %ub1, %ub2 : i32, i64, i32)`
+ - The number of dimensions is implicitly `num_teams_upper_vars.size()`
Uni-dimensional (legacy format):
- - Uses `num_teams_upper` and optional `num_teams_lower`
+ - Uses `num_teams_upper_vars` with one value and optional `num_teams_lower`
- If lower bound not specified, it defaults to upper bound value
- - Format: `num_teams(lower : type to upper : type)` or `num_teams(to upper : type)`
+ - Format: `num_teams(%lb : type to %ub : type)` or `num_teams(to %ub : type)`
- Example: `num_teams(%lb : i32 to %ub : i32)` or `num_teams(to %ub : i32)`
}];
let extraClassDeclaration = [{
/// Returns true if using the multi-dimensional values format
bool hasNumTeamsMultiDim() {
- return !getNumTeamsVars().empty();
+ return getNumTeamsUpperVars().size() > 1;
}
/// Returns the number of dimensions specified for num_teams
unsigned getNumTeamsDimsCount() {
- if (hasNumTeamsMultiDim())
- return getNumTeamsVars().size();
- return getNumTeamsUpper() ? 1 : 0;
+ return getNumTeamsUpperVars().size();
}
- /// Returns the value for a specific dimension index (multi-dim format)
- /// Index must be less than getNumTeamsVars().size()
+ /// Returns the value for a specific dimension index
+ /// Index must be less than getNumTeamsUpperVars().size()
::mlir::Value getNumTeams(unsigned dim = 0) {
assert(dim < getNumTeamsDimsCount() &&
"Num teams dims index out of bounds");
- return getNumTeamsVars()[dim];
+ return getNumTeamsUpperVars()[dim];
}
}];
}
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 0375e52b1e034..fb2e1ed14a9f4 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2240,10 +2240,10 @@ LogicalResult TargetOp::verifyRegions() {
cast<BlockArgOpenMPOpInterface>(getOperation()).getHostEvalBlockArgs()) {
for (Operation *user : hostEvalArg.getUsers()) {
if (auto teamsOp = dyn_cast<TeamsOp>(user)) {
- if (llvm::is_contained({teamsOp.getNumTeamsLower(),
- teamsOp.getNumTeamsUpper(),
- teamsOp.getThreadLimit()},
- hostEvalArg))
+ // Check if used in num_teams_lower or any of num_teams_upper_vars
+ if (hostEvalArg == teamsOp.getNumTeamsLower() ||
+ llvm::is_contained(teamsOp.getNumTeamsUpperVars(), hostEvalArg) ||
+ hostEvalArg == teamsOp.getThreadLimit())
continue;
return emitOpError() << "host_eval argument only legal as 'num_teams' "
@@ -2370,7 +2370,7 @@ Operation *TargetOp::getInnermostCapturedOmpOp() {
static bool canPromoteToNoLoop(Operation *capturedOp, TeamsOp teamsOp,
WsloopOp *wsLoopOp) {
// num_teams clause can break no-loop teams/threads assumption.
- if (teamsOp.getNumTeamsUpper())
+ if (!teamsOp.getNumTeamsUpperVars().empty())
return false;
// Reduction kernels are slower in no-loop mode.
@@ -2624,38 +2624,26 @@ void TeamsOp::build(OpBuilder &builder, OperationState &state,
const TeamsOperands &clauses) {
MLIRContext *ctx = builder.getContext();
// TODO Store clauses in op: privateVars, privateSyms, privateNeedsBarrier
- TeamsOp::build(builder, state, clauses.allocateVars, clauses.allocatorVars,
- clauses.ifExpr, clauses.numTeamsVars, clauses.numTeamsLower,
- clauses.numTeamsUpper,
- /*private_vars=*/{}, /*private_syms=*/nullptr,
- /*private_needs_barrier=*/nullptr, clauses.reductionMod,
- clauses.reductionVars,
- makeDenseBoolArrayAttr(ctx, clauses.reductionByref),
- makeArrayAttr(ctx, clauses.reductionSyms),
- clauses.threadLimit);
+ TeamsOp::build(
+ builder, state, clauses.allocateVars, clauses.allocatorVars,
+ clauses.ifExpr, clauses.numTeamsLower, clauses.numTeamsUpperVars,
+ /*private_vars=*/{}, /*private_syms=*/nullptr,
+ /*private_needs_barrier=*/nullptr, clauses.reductionMod,
+ clauses.reductionVars,
+ makeDenseBoolArrayAttr(ctx, clauses.reductionByref),
+ makeArrayAttr(ctx, clauses.reductionSyms), clauses.threadLimit);
}
// Verify num_teams clause
-static LogicalResult verifyNumTeamsClause(Operation *op,
- OperandRange numTeamsVars,
- Value numTeamsLower,
- Value numTeamsUpper) {
- bool hasLegacyOperands = numTeamsLower || numTeamsUpper;
-
- // Cannot use both multi-dimensional and legacy format simultaneously
- if (!numTeamsVars.empty() && hasLegacyOperands) {
- return op->emitError()
- << "num_teams multi-dimensional values cannot be used together with "
- "legacy lower/upper bounds";
- }
-
- // If lower is specified, upper must also be specified
+static LogicalResult verifyNumTeamsClause(Operation *op, Value numTeamsLower,
+ OperandRange numTeamsUpperVars) {
+ // If lower is specified, upper must have exactly one value
if (numTeamsLower) {
- if (!numTeamsUpper)
+ if (numTeamsUpperVars.size() != 1)
return op->emitError(
- "expected num_teams upper bound to be defined if the "
- "lower bound is defined");
- if (numTeamsLower.getType() != numTeamsUpper.getType())
+ "expected exactly one num_teams upper bound when lower bound is "
+ "specified");
+ if (numTeamsLower.getType() != numTeamsUpperVars[0].getType())
return op->emitError(
"expected num_teams upper bound and lower bound to be "
"the same type");
@@ -2677,9 +2665,8 @@ LogicalResult TeamsOp::verify() {
"in any OpenMP dialect operations");
// Check for num_teams clause restrictions
- if (failed(verifyNumTeamsClause(op, this->getNumTeamsVars(),
- this->getNumTeamsLower(),
- this->getNumTeamsUpper())))
+ if (failed(verifyNumTeamsClause(op, this->getNumTeamsLower(),
+ this->getNumTeamsUpperVars())))
return failure();
// Check for allocate clause restrictions
@@ -4513,111 +4500,6 @@ void DeclareSimdOp::build(OpBuilder &odsBuilder, OperationState &odsState,
clauses.linearVarTypes, clauses.simdlen);
}
-//===----------------------------------------------------------------------===//
-// Parser and printer for num_teams clause
-//===----------------------------------------------------------------------===//
-// Parses num_teams clause in two formats:
-// 1. Multi-dimensional: num_teams(%v0, %v1, ... : type0, type1, ...)
-// 2. Uni-dimensional (legacy): num_teams(%lb : type to %ub : type)
-// or num_teams(to %ub : type)
-static ParseResult
-parseNumTeamsClause(OpAsmParser &parser,
- SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
- SmallVectorImpl<Type> &types,
- std::optional<OpAsmParser::UnresolvedOperand> &lowerBound,
- Type &lowerBoundType,
- std::optional<OpAsmParser::UnresolvedOperand> &upperBound,
- Type &upperBoundType) {
-
- // Format: num_teams(to upper : type)
- if (succeeded(parser.parseOptionalKeyword("to"))) {
- OpAsmParser::UnresolvedOperand upperOperand;
- if (parser.parseOperand(upperOperand) || parser.parseColon() ||
- parser.parseType(upperBoundType)) {
- return failure();
- }
- upperBound = upperOperand;
- return success();
- }
-
- // Try to parse multi-dimensional format: values : types
- // or legacy format: lower : type to upper : type
- OpAsmParser::UnresolvedOperand firstOperand;
- if (parser.parseOperand(firstOperand))
- return failure();
-
- // Check if there is a comma (multi-dimensional) or colon
- if (succeeded(parser.parseOptionalComma())) {
- values.push_back(firstOperand);
- if (parser.parseOperandList(values) || parser.parseColon() ||
- parser.parseTypeList(types)) {
- return failure();
- }
- return success();
- }
-
- // Must have a colon next
- if (parser.parseColon())
- return failure();
-
- Type firstType;
- if (parser.parseType(firstType))
- return failure();
-
- // Check for 'to' keyword (legacy format)
- if (succeeded(parser.parseOptionalKeyword("to"))) {
- lowerBound = firstOperand;
- lowerBoundType = firstType;
-
- OpAsmParser::UnresolvedOperand upperOperand;
- if (parser.parseOperand(upperOperand) || parser.parseColon() ||
- parser.parseType(upperBoundType)) {
- return failure();
- }
- upperBound = upperOperand;
- return success();
- }
-
- // Check for comma after type (multi-dimensional with types)
- if (succeeded(parser.parseOptionalComma())) {
- values.push_back(firstOperand);
- types.push_back(firstType);
- if (parser.parseTypeList(types))
- return failure();
- return success();
- }
-
- // Single value multi-dimensional format: %v0 : type0
- values.push_back(firstOperand);
- types.push_back(firstType);
- return success();
-}
-
-static void printNumTeamsClause(OpAsmPrinter &p, Operation *op,
- OperandRange values, TypeRange types,
- Value lowerBound, Type lowerBoundType,
- Value upperBound, Type upperBoundType) {
- if (!values.empty()) {
- // Multi-dimensional: %v0, %v1, ... : type0, type1, ...
- llvm::interleaveComma(values, p, [&](Value v) { p.printOperand(v); });
- p << " : ";
- llvm::interleaveComma(types, p);
- } else if (upperBound) {
- if (lowerBound) {
- // Both bounds: lower : type to upper : type
- p.printOperand(lowerBound);
- p << " : " << lowerBoundType << " to ";
- p.printOperand(upperBound);
- p << " : " << upperBoundType;
- } else {
- // Upper only: to upper : type
- p << " to ";
- p.printOperand(upperBound);
- p << " : " << upperBoundType;
- }
- }
-}
-
#define GET_ATTRDEF_CLASSES
#include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 0b7bf64cefe4c..71ab9e791b8b8 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -2071,8 +2071,8 @@ convertOmpTeams(omp::TeamsOp op, llvm::IRBuilderBase &builder,
numTeamsLower = moduleTranslation.lookupValue(numTeamsLowerVar);
llvm::Value *numTeamsUpper = nullptr;
- if (Value numTeamsUpperVar = op.getNumTeamsUpper())
- numTeamsUpper = moduleTranslation.lookupValue(numTeamsUpperVar);
+ if (!op.getNumTeamsUpperVars().empty())
+ numTeamsUpper = moduleTranslation.lookupValue(op.getNumTeams(0));
llvm::Value *threadLimit = nullptr;
if (Value threadLimitVar = op.getThreadLimit())
@@ -6042,7 +6042,8 @@ extractHostEvalClauses(omp::TargetOp targetOp, Value &numThreads,
.Case([&](omp::TeamsOp teamsOp) {
if (teamsOp.getNumTeamsLower() == blockArg)
numTeamsLower = hostEvalVar;
- else if (teamsOp.getNumTeamsUpper() == blockArg)
+ else if (llvm::is_contained(teamsOp.getNumTeamsUpperVars(),
+ blockArg))
numTeamsUpper = hostEvalVar;
else if (teamsOp.getThreadLimit() == blockArg)
threadLimit = hostEvalVar;
@@ -6163,7 +6164,9 @@ initTargetDefaultAttrs(omp::TargetOp targetOp, Operation *capturedOp,
// ensures values are mapped and available inside of the target region.
if (auto teamsOp = castOrGetParentOfType<omp::TeamsOp>(capturedOp)) {
numTeamsLower = teamsOp.getNumTeamsLower();
- numTeamsUpper = teamsOp.getNumTeamsUpper();
+ // Handle num_teams upper bounds (only first value for now)
+ if (!teamsOp.getNumTeamsUpperVars().empty())
+ numTeamsUpper = teamsOp.getNumTeams(0);
threadLimit = teamsOp.getThreadLimit();
}
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index d451b14e8bfc9..ba431c3522a31 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1438,7 +1438,7 @@ func.func @omp_teams_allocate(%data_var : memref<i32>) {
// expected-error @below {{expected equal sizes for allocate and allocator variables}}
"omp.teams" (%data_var) ({
omp.terminator
- }) {operandSegmentSizes = array<i32: 1,0,0,0,0,0,0,0,0>} : (memref<i32>) -> ()
+ }) {operandSegmentSizes = array<i32: 1,0,0,0,0,0,0,0>} : (memref<i32>) -> ()
omp.terminator
}
return
@@ -1448,10 +1448,10 @@ func.func @omp_teams_allocate(%data_var : memref<i32>) {
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}}
+ // expected-error @below {{expected exactly one num_teams upper bound when lower bound is specified}}
"omp.teams" (%lb) ({
omp.terminator
- }) {operandSegmentSizes = array<i32: 0,0,0,0,1,0,0,0,0>} : (i32) -> ()
+ }) {operandSegmentSizes = array<i32: 0,0,0,1,0,0,0,0>} : (i32) -> ()
omp.terminator
}
return
@@ -1465,10 +1465,10 @@ func.func @omp_teams_num_teams_multidim_with_bounds() {
%v1 = arith.constant 2 : i32
%lb = arith.constant 3 : i32
%ub = arith.constant 4 : i32
- // expected-error @below {{num_teams multi-dimensional values cannot be used together with legacy lower/upper bounds}}
- "omp.teams" (%v0, %v1, %lb, %ub) ({
+ // expected-error @below {{expected exactly one num_teams upper bound when lower bound is specified}}
+ "omp.teams" (%lb, %v0, %v1) ({
omp.terminator
- }) {operandSegmentSizes = array<i32: 0,0,0,2,1,1,0,0,0>} : (i32, i32, i32, i32) -> ()
+ }) {operandSegmentSizes = array<i32: 0,0,0,1,2,0,0,0>} : (i32, i32, i32) -> ()
omp.terminator
}
return
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 49a88e0443e60..c27371596c2cd 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -1109,15 +1109,15 @@ func.func @omp_teams(%lb : i32, %ub : i32, %if_cond : i1, %num_threads : i32,
omp.terminator
}
- // CHECK: omp.teams num_teams(%{{.*}}, %{{.*}}, %{{.*}} : i32, i32, i32)
- omp.teams num_teams(%lb, %ub, %ub : i32, i32, i32) {
+ // CHECK: omp.teams num_teams( to %{{.*}}, %{{.*}}, %{{.*}} : i32, i32, i32)
+ omp.teams num_teams(to %lb, %ub, %ub : i32, i32, i32) {
// CHECK: omp.terminator
omp.terminator
}
// Test num_teams with mixed types.
- // CHECK: omp.teams num_teams(%{{.*}}, %{{.*}}, %{{.*}} : i32, i64, i16)
- omp.teams num_teams(%lb, %ub64, %ub16 : i32, i64, i16) {
+ // CHECK: omp.teams num_teams( to %{{.*}}, %{{.*}}, %{{.*}} : i32, i64, i16)
+ omp.teams num_teams(to %lb, %ub64, %ub16 : i32, i64, i16) {
// CHECK: omp.terminator
omp.terminator
}
diff --git a/mlir/test/Target/LLVMIR/openmp-todo.mlir b/mlir/test/Target/LLVMIR/openmp-todo.mlir
index 1ea56fdd0bf16..09919ae4f5267 100644
--- a/mlir/test/Target/LLVMIR/openmp-todo.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-todo.mlir
@@ -435,7 +435,7 @@ llvm.func @teams_private(%x : !llvm.ptr) {
llvm.func @teams_num_teams_multi_dim(%lb : i32, %ub : i32) {
// expected-error at below {{not yet implemented: Unhandled clause num_teams with multi-dimensional values in omp.teams operation}}
// expected-error at below {{LLVM Translation failed for operation: omp.teams}}
- omp.teams num_teams(%lb, %ub, %ub : i32, i32, i32) {
+ omp.teams num_teams(to %ub, %ub, %ub : i32, i32, i32) {
omp.terminator
}
llvm.return
>From e9e8650ea358c8322333880f767df0177b821794 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Thu, 22 Jan 2026 21:41:28 +0530
Subject: [PATCH 14/14] Remove space in test
---
mlir/test/Dialect/OpenMP/ops.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index c27371596c2cd..1c9fba3a61670 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -3087,7 +3087,7 @@ func.func @omp_target_host_eval(%x : i32) {
// CHECK: omp.teams num_teams( to %[[HOST_ARG]] : i32)
// CHECK-SAME: thread_limit(%[[HOST_ARG]] : i32)
omp.target host_eval(%x -> %arg0 : i32) {
- omp.teams num_teams( to %arg0 : i32) thread_limit(%arg0 : i32) {
+ omp.teams num_teams(to %arg0 : i32) thread_limit(%arg0 : i32) {
omp.terminator
}
omp.terminator
More information about the llvm-commits
mailing list