[flang] [llvm] [mlir] [flang] Add support for workdistribute construct in flang frontend (PR #146029)
via llvm-commits
llvm-commits at lists.llvm.org
Wed Jul 16 07:01:49 PDT 2025
https://github.com/skc7 updated https://github.com/llvm/llvm-project/pull/146029
>From cb89a7f7d8edf63ef6d054737cd8615bec86ce27 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Fri, 27 Jun 2025 11:05:35 +0530
Subject: [PATCH 1/2] [flang] Add support for workdistribute construct in flang
frontend
---
.../flang/Semantics/openmp-directive-sets.h | 14 +++++
flang/lib/Lower/OpenMP/OpenMP.cpp | 26 +++++++-
flang/lib/Parser/openmp-parsers.cpp | 7 ++-
flang/lib/Semantics/resolve-directives.cpp | 8 ++-
flang/test/Lower/OpenMP/workdistribute.f90 | 59 +++++++++++++++++++
llvm/include/llvm/Frontend/OpenMP/OMP.td | 55 +++++++++++++++++
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 22 +++++++
7 files changed, 188 insertions(+), 3 deletions(-)
create mode 100644 flang/test/Lower/OpenMP/workdistribute.f90
diff --git a/flang/include/flang/Semantics/openmp-directive-sets.h b/flang/include/flang/Semantics/openmp-directive-sets.h
index dd610c9702c28..7ced6ed9b44d6 100644
--- a/flang/include/flang/Semantics/openmp-directive-sets.h
+++ b/flang/include/flang/Semantics/openmp-directive-sets.h
@@ -143,6 +143,7 @@ static const OmpDirectiveSet topTargetSet{
Directive::OMPD_target_teams_distribute_parallel_do_simd,
Directive::OMPD_target_teams_distribute_simd,
Directive::OMPD_target_teams_loop,
+ Directive::OMPD_target_teams_workdistribute,
};
static const OmpDirectiveSet allTargetSet{topTargetSet};
@@ -172,6 +173,7 @@ static const OmpDirectiveSet topTeamsSet{
Directive::OMPD_teams_distribute_parallel_do_simd,
Directive::OMPD_teams_distribute_simd,
Directive::OMPD_teams_loop,
+ Directive::OMPD_teams_workdistribute,
};
static const OmpDirectiveSet bottomTeamsSet{
@@ -187,9 +189,16 @@ static const OmpDirectiveSet allTeamsSet{
Directive::OMPD_target_teams_distribute_parallel_do_simd,
Directive::OMPD_target_teams_distribute_simd,
Directive::OMPD_target_teams_loop,
+ Directive::OMPD_target_teams_workdistribute,
} | topTeamsSet,
};
+static const OmpDirectiveSet allWorkdistributeSet{
+ Directive::OMPD_workdistribute,
+ Directive::OMPD_teams_workdistribute,
+ Directive::OMPD_target_teams_workdistribute,
+};
+
//===----------------------------------------------------------------------===//
// Directive sets for groups of multiple directives
//===----------------------------------------------------------------------===//
@@ -230,6 +239,9 @@ static const OmpDirectiveSet blockConstructSet{
Directive::OMPD_taskgroup,
Directive::OMPD_teams,
Directive::OMPD_workshare,
+ Directive::OMPD_target_teams_workdistribute,
+ Directive::OMPD_teams_workdistribute,
+ Directive::OMPD_workdistribute,
};
static const OmpDirectiveSet loopConstructSet{
@@ -294,6 +306,7 @@ static const OmpDirectiveSet workShareSet{
Directive::OMPD_scope,
Directive::OMPD_sections,
Directive::OMPD_single,
+ Directive::OMPD_workdistribute,
} | allDoSet,
};
@@ -376,6 +389,7 @@ static const OmpDirectiveSet nestedReduceWorkshareAllowedSet{
};
static const OmpDirectiveSet nestedTeamsAllowedSet{
+ Directive::OMPD_workdistribute,
Directive::OMPD_distribute,
Directive::OMPD_distribute_parallel_do,
Directive::OMPD_distribute_parallel_do_simd,
diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 4458f62eea95a..fcee143a19ef9 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -618,6 +618,16 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
cp.processCollapse(loc, eval, hostInfo->ops, hostInfo->iv);
break;
+ case OMPD_teams_workdistribute:
+ cp.processThreadLimit(stmtCtx, hostInfo.ops);
+ [[fallthrough]];
+ case OMPD_target_teams_workdistribute:
+ cp.processNumTeams(stmtCtx, hostInfo.ops);
+ processSingleNestedIf([](Directive nestedDir) {
+ return topDistributeSet.test(nestedDir) || topLoopSet.test(nestedDir);
+ });
+ break;
+
// Standalone 'target' case.
case OMPD_target: {
processSingleNestedIf(
@@ -2719,6 +2729,17 @@ genTeamsOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
queue, item, clauseOps);
}
+static mlir::omp::WorkdistributeOp genWorkdistributeOp(
+ lower::AbstractConverter &converter, lower::SymMap &symTable,
+ semantics::SemanticsContext &semaCtx, lower::pft::Evaluation &eval,
+ mlir::Location loc, const ConstructQueue &queue,
+ ConstructQueue::const_iterator item) {
+ return genOpWithBody<mlir::omp::WorkdistributeOp>(
+ OpWithBodyGenInfo(converter, symTable, semaCtx, loc, eval,
+ llvm::omp::Directive::OMPD_workdistribute),
+ queue, item);
+}
+
//===----------------------------------------------------------------------===//
// Code generation functions for the standalone version of constructs that can
// also be a leaf of a composite construct
@@ -3339,7 +3360,10 @@ static void genOMPDispatch(lower::AbstractConverter &converter,
TODO(loc, "Unhandled loop directive (" +
llvm::omp::getOpenMPDirectiveName(dir, version) + ")");
}
- // case llvm::omp::Directive::OMPD_workdistribute:
+ case llvm::omp::Directive::OMPD_workdistribute:
+ newOp = genWorkdistributeOp(converter, symTable, semaCtx, eval, loc, queue,
+ item);
+ break;
case llvm::omp::Directive::OMPD_workshare:
newOp = genWorkshareOp(converter, symTable, stmtCtx, semaCtx, eval, loc,
queue, item);
diff --git a/flang/lib/Parser/openmp-parsers.cpp b/flang/lib/Parser/openmp-parsers.cpp
index 76c9499410017..9a08df6e4a889 100644
--- a/flang/lib/Parser/openmp-parsers.cpp
+++ b/flang/lib/Parser/openmp-parsers.cpp
@@ -1567,12 +1567,17 @@ TYPE_PARSER(
"TARGET DATA" >> pure(llvm::omp::Directive::OMPD_target_data),
"TARGET_DATA" >> pure(llvm::omp::Directive::OMPD_target_data),
"TARGET PARALLEL" >> pure(llvm::omp::Directive::OMPD_target_parallel),
+ "TARGET TEAMS WORKDISTRIBUTE" >>
+ pure(llvm::omp::Directive::OMPD_target_teams_workdistribute),
"TARGET TEAMS" >> pure(llvm::omp::Directive::OMPD_target_teams),
"TARGET" >> pure(llvm::omp::Directive::OMPD_target),
"TASK"_id >> pure(llvm::omp::Directive::OMPD_task),
"TASKGROUP" >> pure(llvm::omp::Directive::OMPD_taskgroup),
+ "TEAMS WORKDISTRIBUTE" >>
+ pure(llvm::omp::Directive::OMPD_teams_workdistribute),
"TEAMS" >> pure(llvm::omp::Directive::OMPD_teams),
- "WORKSHARE" >> pure(llvm::omp::Directive::OMPD_workshare))))
+ "WORKSHARE" >> pure(llvm::omp::Directive::OMPD_workshare),
+ "WORKDISTRIBUTE" >> pure(llvm::omp::Directive::OMPD_workdistribute))))
TYPE_PARSER(sourced(construct<OmpBeginBlockDirective>(
sourced(Parser<OmpBlockDirective>{}), Parser<OmpClauseList>{})))
diff --git a/flang/lib/Semantics/resolve-directives.cpp b/flang/lib/Semantics/resolve-directives.cpp
index 151f4ccae634e..42d24e703889e 100644
--- a/flang/lib/Semantics/resolve-directives.cpp
+++ b/flang/lib/Semantics/resolve-directives.cpp
@@ -1680,10 +1680,13 @@ bool OmpAttributeVisitor::Pre(const parser::OpenMPBlockConstruct &x) {
case llvm::omp::Directive::OMPD_task:
case llvm::omp::Directive::OMPD_taskgroup:
case llvm::omp::Directive::OMPD_teams:
+ case llvm::omp::Directive::OMPD_workdistribute:
case llvm::omp::Directive::OMPD_workshare:
case llvm::omp::Directive::OMPD_parallel_workshare:
case llvm::omp::Directive::OMPD_target_teams:
+ case llvm::omp::Directive::OMPD_target_teams_workdistribute:
case llvm::omp::Directive::OMPD_target_parallel:
+ case llvm::omp::Directive::OMPD_teams_workdistribute:
PushContext(beginDir.source, beginDir.v);
break;
default:
@@ -1713,9 +1716,12 @@ void OmpAttributeVisitor::Post(const parser::OpenMPBlockConstruct &x) {
case llvm::omp::Directive::OMPD_target:
case llvm::omp::Directive::OMPD_task:
case llvm::omp::Directive::OMPD_teams:
+ case llvm::omp::Directive::OMPD_workdistribute:
case llvm::omp::Directive::OMPD_parallel_workshare:
case llvm::omp::Directive::OMPD_target_teams:
- case llvm::omp::Directive::OMPD_target_parallel: {
+ case llvm::omp::Directive::OMPD_target_parallel:
+ case llvm::omp::Directive::OMPD_target_teams_workdistribute:
+ case llvm::omp::Directive::OMPD_teams_workdistribute: {
bool hasPrivate;
for (const auto *allocName : allocateNames_) {
hasPrivate = false;
diff --git a/flang/test/Lower/OpenMP/workdistribute.f90 b/flang/test/Lower/OpenMP/workdistribute.f90
new file mode 100644
index 0000000000000..924205bb72e5e
--- /dev/null
+++ b/flang/test/Lower/OpenMP/workdistribute.f90
@@ -0,0 +1,59 @@
+! RUN: %flang_fc1 -emit-hlfir -fopenmp %s -o - | FileCheck %s
+
+! CHECK-LABEL: func @_QPtarget_teams_workdistribute
+subroutine target_teams_workdistribute()
+ ! CHECK: omp.target
+ ! CHECK: omp.teams
+ ! CHECK: omp.workdistribute
+ !$omp target teams workdistribute
+ ! CHECK: fir.call
+ call f1()
+ ! CHECK: omp.terminator
+ ! CHECK: omp.terminator
+ ! CHECK: omp.terminator
+ !$omp end target teams workdistribute
+end subroutine target_teams_workdistribute
+
+! CHECK-LABEL: func @_QPteams_workdistribute
+subroutine teams_workdistribute()
+ ! CHECK: omp.teams
+ ! CHECK: omp.workdistribute
+ !$omp teams workdistribute
+ ! CHECK: fir.call
+ call f1()
+ ! CHECK: omp.terminator
+ ! CHECK: omp.terminator
+ !$omp end teams workdistribute
+end subroutine teams_workdistribute
+
+! CHECK-LABEL: func @_QPtarget_teams_workdistribute_m
+subroutine target_teams_workdistribute_m()
+ ! CHECK: omp.target
+ ! CHECK: omp.teams
+ ! CHECK: omp.workdistribute
+ !$omp target
+ !$omp teams
+ !$omp workdistribute
+ ! CHECK: fir.call
+ call f1()
+ ! CHECK: omp.terminator
+ ! CHECK: omp.terminator
+ ! CHECK: omp.terminator
+ !$omp end workdistribute
+ !$omp end teams
+ !$omp end target
+end subroutine target_teams_workdistribute_m
+
+! CHECK-LABEL: func @_QPteams_workdistribute_m
+subroutine teams_workdistribute_m()
+ ! CHECK: omp.teams
+ ! CHECK: omp.workdistribute
+ !$omp teams
+ !$omp workdistribute
+ ! CHECK: fir.call
+ call f1()
+ ! CHECK: omp.terminator
+ ! CHECK: omp.terminator
+ !$omp end workdistribute
+ !$omp end teams
+end subroutine teams_workdistribute_m
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMP.td b/llvm/include/llvm/Frontend/OpenMP/OMP.td
index 1b94657dfae1e..64a9b29a3f1d0 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMP.td
+++ b/llvm/include/llvm/Frontend/OpenMP/OMP.td
@@ -1309,6 +1309,15 @@ def OMP_EndWorkshare : Directive<[Spelling<"end workshare">]> {
let category = OMP_Workshare.category;
let languages = [L_Fortran];
}
+def OMP_Workdistribute : Directive<[Spelling<"workdistribute">]> {
+ let association = AS_Block;
+ let category = CA_Executable;
+}
+def OMP_EndWorkdistribute : Directive<[Spelling<"end workdistribute">]> {
+ let leafConstructs = OMP_Workdistribute.leafConstructs;
+ let association = OMP_Workdistribute.association;
+ let category = OMP_Workdistribute.category;
+}
//===----------------------------------------------------------------------===//
// Definitions of OpenMP compound directives
@@ -2452,6 +2461,34 @@ def OMP_TargetTeamsDistributeSimd
let leafConstructs = [OMP_Target, OMP_Teams, OMP_Distribute, OMP_Simd];
let category = CA_Executable;
}
+def OMP_TargetTeamsWorkdistribute : Directive<[Spelling<"target teams workdistribute">]> {
+ let allowedClauses = [
+ VersionedClause<OMPC_Allocate>,
+ VersionedClause<OMPC_Depend>,
+ VersionedClause<OMPC_FirstPrivate>,
+ VersionedClause<OMPC_HasDeviceAddr, 51>,
+ VersionedClause<OMPC_If>,
+ VersionedClause<OMPC_IsDevicePtr>,
+ VersionedClause<OMPC_Map>,
+ VersionedClause<OMPC_OMPX_Attribute>,
+ VersionedClause<OMPC_Private>,
+ VersionedClause<OMPC_Reduction>,
+ VersionedClause<OMPC_Shared>,
+ VersionedClause<OMPC_UsesAllocators, 50>,
+ ];
+ let allowedOnceClauses = [
+ VersionedClause<OMPC_Default>,
+ VersionedClause<OMPC_DefaultMap>,
+ VersionedClause<OMPC_Device>,
+ VersionedClause<OMPC_NoWait>,
+ VersionedClause<OMPC_NumTeams>,
+ VersionedClause<OMPC_OMPX_DynCGroupMem>,
+ VersionedClause<OMPC_OMPX_Bare>,
+ VersionedClause<OMPC_ThreadLimit>,
+ ];
+ let leafConstructs = [OMP_Target, OMP_Teams, OMP_Workdistribute];
+ let category = CA_Executable;
+}
def OMP_target_teams_loop : Directive<[Spelling<"target teams loop">]> {
let allowedClauses = [
VersionedClause<OMPC_Allocate>,
@@ -2682,6 +2719,24 @@ def OMP_TeamsDistributeSimd : Directive<[Spelling<"teams distribute simd">]> {
let leafConstructs = [OMP_Teams, OMP_Distribute, OMP_Simd];
let category = CA_Executable;
}
+def OMP_TeamsWorkdistribute : Directive<[Spelling<"teams workdistribute">]> {
+ let allowedClauses = [
+ VersionedClause<OMPC_Allocate>,
+ VersionedClause<OMPC_FirstPrivate>,
+ VersionedClause<OMPC_OMPX_Attribute>,
+ VersionedClause<OMPC_Private>,
+ VersionedClause<OMPC_Reduction>,
+ VersionedClause<OMPC_Shared>,
+ ];
+ let allowedOnceClauses = [
+ VersionedClause<OMPC_Default>,
+ VersionedClause<OMPC_If, 52>,
+ VersionedClause<OMPC_NumTeams>,
+ VersionedClause<OMPC_ThreadLimit>,
+ ];
+ let leafConstructs = [OMP_Teams, OMP_Workdistribute];
+ let category = CA_Executable;
+}
def OMP_teams_loop : Directive<[Spelling<"teams loop">]> {
let allowedClauses = [
VersionedClause<OMPC_Allocate>,
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 8cf18b43450ab..18a68ee637dd3 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -2113,4 +2113,26 @@ def AllocateDirOp : OpenMP_Op<"allocate_dir", clauses = [
let hasVerifier = 1;
}
+//===----------------------------------------------------------------------===//
+// workdistribute Construct
+//===----------------------------------------------------------------------===//
+
+def WorkdistributeOp : OpenMP_Op<"workdistribute"> {
+ let summary = "workdistribute directive";
+ let description = [{
+ workdistribute divides execution of the enclosed structured block into
+ separate units of work, each executed only once by each
+ initial thread in the league.
+ ```
+ !$omp target teams
+ !$omp workdistribute
+ y = a * x + y
+ !$omp end workdistribute
+ !$omp end target teams
+ ```
+ }];
+ let regions = (region AnyRegion:$region);
+ let assemblyFormat = "$region attr-dict";
+}
+
#endif // OPENMP_OPS
>From fe2a447e45038d1b8251b707e8f14c7ff7bf79ab Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Fri, 27 Jun 2025 16:33:58 +0530
Subject: [PATCH 2/2] [OpenMP] Add verifier and tests for workdistribute mlir
op.
---
flang/lib/Lower/OpenMP/OpenMP.cpp | 4 ++--
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 1 +
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 15 ++++++++++++
mlir/test/Dialect/OpenMP/invalid.mlir | 23 +++++++++++++++++++
mlir/test/Dialect/OpenMP/ops.mlir | 12 ++++++++++
5 files changed, 53 insertions(+), 2 deletions(-)
diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index fcee143a19ef9..37af546657644 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -619,10 +619,10 @@ static void processHostEvalClauses(lower::AbstractConverter &converter,
break;
case OMPD_teams_workdistribute:
- cp.processThreadLimit(stmtCtx, hostInfo.ops);
+ cp.processThreadLimit(stmtCtx, hostInfo->ops);
[[fallthrough]];
case OMPD_target_teams_workdistribute:
- cp.processNumTeams(stmtCtx, hostInfo.ops);
+ cp.processNumTeams(stmtCtx, hostInfo->ops);
processSingleNestedIf([](Directive nestedDir) {
return topDistributeSet.test(nestedDir) || topLoopSet.test(nestedDir);
});
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 18a68ee637dd3..189b96510ce99 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -2132,6 +2132,7 @@ def WorkdistributeOp : OpenMP_Op<"workdistribute"> {
```
}];
let regions = (region AnyRegion:$region);
+ let hasVerifier = 1;
let assemblyFormat = "$region attr-dict";
}
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 769aee64e1695..fd4f1891703f1 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -3878,6 +3878,21 @@ LogicalResult AllocateDirOp::verify() {
return success();
}
+//===----------------------------------------------------------------------===//
+// WorkdistributeOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult WorkdistributeOp::verify() {
+ Region ®ion = getRegion();
+ if (!region.hasOneBlock())
+ return emitOpError("region must contain exactly one block");
+
+ Operation *parentOp = (*this)->getParentOp();
+ if (!llvm::dyn_cast<TeamsOp>(parentOp))
+ return emitOpError("workdistribute must be nested under teams");
+ return success();
+}
+
#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 5088f2dfa7d7a..9723a263fca61 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -3017,3 +3017,26 @@ func.func @invalid_allocate_allocator(%arg0 : memref<i32>) -> () {
return
}
+
+// -----
+func.func @invalid_workdistribute() -> () {
+// expected-error @below {{workdistribute must be nested under teams}}
+ omp.workdistribute {
+ omp.terminator
+ }
+ return
+}
+
+// -----
+func.func @invalid_workdistribute_with_multiple_blocks() -> () {
+ omp.teams {
+ // expected-error @below {{region must contain exactly one block}}
+ omp.workdistribute {
+ cf.br ^bb1
+ ^bb1:
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 4c50ed3230976..f77035753ddb7 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -3230,3 +3230,15 @@ func.func @omp_allocate_dir(%arg0 : memref<i32>, %arg1 : memref<i32>) -> () {
return
}
+// CHECK-LABEL: func.func @omp_workdistribute
+func.func @omp_workdistribute() {
+ // CHECK: omp.teams
+ omp.teams {
+ // CHECK: omp.workdistribute
+ omp.workdistribute {
+ omp.terminator
+ }
+ omp.terminator
+ }
+ return
+}
More information about the llvm-commits
mailing list