[OpenMP] Add workdistribute construct in openMP dialect and in llvm frontend #154376
+253 −0
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
This was referenced Aug 19, 2025
llvmbot added mlir mlir:openmp flang:openmp clang:openmp 
Member
| @llvm/pr-subscribers-mlir-openmp @llvm/pr-subscribers-mlir Author: Chaitanya (skc7) ChangesThis PR adds workdistribute mlir op in omp dialect and also in llvm frontend. This is part 1 of workdistribute frontend support. Part1 : #154376 [OpenMP] Add workdistribute construct in openMP dialect and in llvm frontend Full diff: https://github.com/llvm/llvm-project/pull/154376.diff 5 Files Affected:
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMP.td b/llvm/include/llvm/Frontend/OpenMP/OMP.td index 202f684d808bc..6eaf19eed40c8 100644 --- a/llvm/include/llvm/Frontend/OpenMP/OMP.td +++ b/llvm/include/llvm/Frontend/OpenMP/OMP.td @@ -1322,6 +1322,17 @@ 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; + let languages = [L_Fortran]; +} +def OMP_EndWorkdistribute : Directive<[Spelling<"end workdistribute">]> { + let leafConstructs = OMP_Workdistribute.leafConstructs; + let association = OMP_Workdistribute.association; + let category = OMP_Workdistribute.category; + let languages = [L_Fortran]; +} //===----------------------------------------------------------------------===// // Definitions of OpenMP compound directives @@ -2482,6 +2493,35 @@ 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; + let languages = [L_Fortran]; +} def OMP_target_teams_loop : Directive<[Spelling<"target teams loop">]> { let allowedClauses = [ VersionedClause<OMPC_Allocate>, @@ -2723,6 +2763,25 @@ 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; + let languages = [L_Fortran]; +} 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 c956d69781b3d..2548a8ab4aac6 100644 --- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td +++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td @@ -2209,4 +2209,27 @@ def TargetFreeMemOp : OpenMP_Op<"target_freemem", let assemblyFormat = "$device `,` $heapref attr-dict `:` type($device) `,` qualified(type($heapref))"; } +//===----------------------------------------------------------------------===// +// 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 hasVerifier = 1; + let assemblyFormat = "$region attr-dict"; +} + #endif // OPENMP_OPS diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp index fa94219016c1e..6e43f28e8d93d 100644 --- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp +++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp @@ -3975,6 +3975,58 @@ llvm::LogicalResult omp::TargetAllocMemOp::verify() { return mlir::success(); } +//===----------------------------------------------------------------------===// +// WorkdistributeOp +//===----------------------------------------------------------------------===// + +LogicalResult WorkdistributeOp::verify() { + // Check that region exists and is not empty + Region ®ion = getRegion(); + if (region.empty()) + return emitOpError("region cannot be empty"); + // Verify single entry point. + Block &entryBlock = region.front(); + if (entryBlock.empty()) + return emitOpError("region must contain a structured block"); + // Verify single exit point. + bool hasTerminator = false; + for (Block &block : region) { + if (isa<TerminatorOp>(block.back())) { + if (hasTerminator) { + return emitOpError("region must have exactly one terminator"); + } + hasTerminator = true; + } + } + if (!hasTerminator) { + return emitOpError("region must be terminated with omp.terminator"); + } + auto walkResult = region.walk([&](Operation *op) -> WalkResult { + // No implicit barrier at end + if (isa<BarrierOp>(op)) { + return emitOpError( + "explicit barriers are not allowed in workdistribute region"); + } + // Check for invalid nested constructs + if (isa<ParallelOp>(op)) { + return emitOpError( + "nested parallel constructs not allowed in workdistribute"); + } + if (isa<TeamsOp>(op)) { + return emitOpError( + "nested teams constructs not allowed in workdistribute"); + } + return WalkResult::advance(); + }); + if (walkResult.wasInterrupted()) + return failure(); + + 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..986c3844d0bb9 100644 --- a/mlir/test/Dialect/OpenMP/invalid.mlir +++ b/mlir/test/Dialect/OpenMP/invalid.mlir @@ -3017,3 +3017,110 @@ func.func @invalid_allocate_allocator(%arg0 : memref<i32>) -> () { return } + +// ----- +func.func @invalid_workdistribute_empty_region() -> () { + omp.teams { + // expected-error @below {{region cannot be empty}} + omp.workdistribute { + } + omp.terminator + } + return +} + +// ----- +func.func @invalid_workdistribute_no_terminator() -> () { + omp.teams { + // expected-error @below {{region must be terminated with omp.terminator}} + omp.workdistribute { + %c0 = arith.constant 0 : i32 + } + omp.terminator + } + return +} + +// ----- +func.func @invalid_workdistribute_wrong_terminator() -> () { + omp.teams { + // expected-error @below {{region must be terminated with omp.terminator}} + omp.workdistribute { + %c0 = arith.constant 0 : i32 + func.return + } + omp.terminator + } + return +} + +// ----- +func.func @invalid_workdistribute_multiple_terminators() -> () { + omp.teams { + // expected-error @below {{region must have exactly one terminator}} + omp.workdistribute { + %cond = arith.constant true + cf.cond_br %cond, ^bb1, ^bb2 + ^bb1: + omp.terminator + ^bb2: + omp.terminator + } + omp.terminator + } + return +} + +// ----- +func.func @invalid_workdistribute_with_barrier() -> () { + omp.teams { + // expected-error @below {{explicit barriers are not allowed in workdistribute region}} + omp.workdistribute { + %c0 = arith.constant 0 : i32 + omp.barrier + omp.terminator + } + omp.terminator + } + return +} + +// ----- +func.func @invalid_workdistribute_nested_parallel() -> () { + omp.teams { + // expected-error @below {{nested parallel constructs not allowed in workdistribute}} + omp.workdistribute { + omp.parallel { + omp.terminator + } + omp.terminator + } + omp.terminator + } + return +} + +// ----- +// Test: nested teams not allowed in workdistribute +func.func @invalid_workdistribute_nested_teams() -> () { + omp.teams { + // expected-error @below {{nested teams constructs not allowed in workdistribute}} + omp.workdistribute { + omp.teams { + omp.terminator + } + omp.terminator + } + omp.terminator + } + return +} + +// ----- +func.func @invalid_workdistribute() -> () { +// expected-error @below {{workdistribute must be nested under teams}} + omp.workdistribute { + omp.terminator + } + return +} diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir index 8c846cde1a3ca..3c2e0a3b7cc15 100644 --- a/mlir/test/Dialect/OpenMP/ops.mlir +++ b/mlir/test/Dialect/OpenMP/ops.mlir @@ -3238,3 +3238,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 +} |
Member
| @llvm/pr-subscribers-flang-openmp Author: Chaitanya (skc7) ChangesThis PR adds workdistribute mlir op in omp dialect and also in llvm frontend. This is part 1 of workdistribute frontend support. Part1 : #154376 [OpenMP] Add workdistribute construct in openMP dialect and in llvm frontend Full diff: https://github.com/llvm/llvm-project/pull/154376.diff 5 Files Affected:
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMP.td b/llvm/include/llvm/Frontend/OpenMP/OMP.td index 202f684d808bc..6eaf19eed40c8 100644 --- a/llvm/include/llvm/Frontend/OpenMP/OMP.td +++ b/llvm/include/llvm/Frontend/OpenMP/OMP.td @@ -1322,6 +1322,17 @@ 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; + let languages = [L_Fortran]; +} +def OMP_EndWorkdistribute : Directive<[Spelling<"end workdistribute">]> { + let leafConstructs = OMP_Workdistribute.leafConstructs; + let association = OMP_Workdistribute.association; + let category = OMP_Workdistribute.category; + let languages = [L_Fortran]; +} //===----------------------------------------------------------------------===// // Definitions of OpenMP compound directives @@ -2482,6 +2493,35 @@ 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; + let languages = [L_Fortran]; +} def OMP_target_teams_loop : Directive<[Spelling<"target teams loop">]> { let allowedClauses = [ VersionedClause<OMPC_Allocate>, @@ -2723,6 +2763,25 @@ 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; + let languages = [L_Fortran]; +} 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 c956d69781b3d..2548a8ab4aac6 100644 --- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td +++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td @@ -2209,4 +2209,27 @@ def TargetFreeMemOp : OpenMP_Op<"target_freemem", let assemblyFormat = "$device `,` $heapref attr-dict `:` type($device) `,` qualified(type($heapref))"; } +//===----------------------------------------------------------------------===// +// 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 hasVerifier = 1; + let assemblyFormat = "$region attr-dict"; +} + #endif // OPENMP_OPS diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp index fa94219016c1e..6e43f28e8d93d 100644 --- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp +++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp @@ -3975,6 +3975,58 @@ llvm::LogicalResult omp::TargetAllocMemOp::verify() { return mlir::success(); } +//===----------------------------------------------------------------------===// +// WorkdistributeOp +//===----------------------------------------------------------------------===// + +LogicalResult WorkdistributeOp::verify() { + // Check that region exists and is not empty + Region ®ion = getRegion(); + if (region.empty()) + return emitOpError("region cannot be empty"); + // Verify single entry point. + Block &entryBlock = region.front(); + if (entryBlock.empty()) + return emitOpError("region must contain a structured block"); + // Verify single exit point. + bool hasTerminator = false; + for (Block &block : region) { + if (isa<TerminatorOp>(block.back())) { + if (hasTerminator) { + return emitOpError("region must have exactly one terminator"); + } + hasTerminator = true; + } + } + if (!hasTerminator) { + return emitOpError("region must be terminated with omp.terminator"); + } + auto walkResult = region.walk([&](Operation *op) -> WalkResult { + // No implicit barrier at end + if (isa<BarrierOp>(op)) { + return emitOpError( + "explicit barriers are not allowed in workdistribute region"); + } + // Check for invalid nested constructs + if (isa<ParallelOp>(op)) { + return emitOpError( + "nested parallel constructs not allowed in workdistribute"); + } + if (isa<TeamsOp>(op)) { + return emitOpError( + "nested teams constructs not allowed in workdistribute"); + } + return WalkResult::advance(); + }); + if (walkResult.wasInterrupted()) + return failure(); + + 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..986c3844d0bb9 100644 --- a/mlir/test/Dialect/OpenMP/invalid.mlir +++ b/mlir/test/Dialect/OpenMP/invalid.mlir @@ -3017,3 +3017,110 @@ func.func @invalid_allocate_allocator(%arg0 : memref<i32>) -> () { return } + +// ----- +func.func @invalid_workdistribute_empty_region() -> () { + omp.teams { + // expected-error @below {{region cannot be empty}} + omp.workdistribute { + } + omp.terminator + } + return +} + +// ----- +func.func @invalid_workdistribute_no_terminator() -> () { + omp.teams { + // expected-error @below {{region must be terminated with omp.terminator}} + omp.workdistribute { + %c0 = arith.constant 0 : i32 + } + omp.terminator + } + return +} + +// ----- +func.func @invalid_workdistribute_wrong_terminator() -> () { + omp.teams { + // expected-error @below {{region must be terminated with omp.terminator}} + omp.workdistribute { + %c0 = arith.constant 0 : i32 + func.return + } + omp.terminator + } + return +} + +// ----- +func.func @invalid_workdistribute_multiple_terminators() -> () { + omp.teams { + // expected-error @below {{region must have exactly one terminator}} + omp.workdistribute { + %cond = arith.constant true + cf.cond_br %cond, ^bb1, ^bb2 + ^bb1: + omp.terminator + ^bb2: + omp.terminator + } + omp.terminator + } + return +} + +// ----- +func.func @invalid_workdistribute_with_barrier() -> () { + omp.teams { + // expected-error @below {{explicit barriers are not allowed in workdistribute region}} + omp.workdistribute { + %c0 = arith.constant 0 : i32 + omp.barrier + omp.terminator + } + omp.terminator + } + return +} + +// ----- +func.func @invalid_workdistribute_nested_parallel() -> () { + omp.teams { + // expected-error @below {{nested parallel constructs not allowed in workdistribute}} + omp.workdistribute { + omp.parallel { + omp.terminator + } + omp.terminator + } + omp.terminator + } + return +} + +// ----- +// Test: nested teams not allowed in workdistribute +func.func @invalid_workdistribute_nested_teams() -> () { + omp.teams { + // expected-error @below {{nested teams constructs not allowed in workdistribute}} + omp.workdistribute { + omp.teams { + omp.terminator + } + omp.terminator + } + omp.terminator + } + return +} + +// ----- +func.func @invalid_workdistribute() -> () { +// expected-error @below {{workdistribute must be nested under teams}} + omp.workdistribute { + omp.terminator + } + return +} diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir index 8c846cde1a3ca..3c2e0a3b7cc15 100644 --- a/mlir/test/Dialect/OpenMP/ops.mlir +++ b/mlir/test/Dialect/OpenMP/ops.mlir @@ -3238,3 +3238,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 +} |
skc7 added a commit to ROCm/llvm-project that referenced this pull request Oct 31, 2025
…lvm#154376) This PR adds workdistribute mlir op in omp dialect and also in llvm frontend. The work in this PR is c-p and updated from @ivanradanov commits from coexecute implementation: flang_workdistribute_iwomp_2024
skc7 added a commit to ROCm/llvm-project that referenced this pull request Oct 31, 2025
…lvm#154376) This PR adds workdistribute mlir op in omp dialect and also in llvm frontend. The work in this PR is c-p and updated from @ivanradanov commits from coexecute implementation: flang_workdistribute_iwomp_2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Add this suggestion to a batch that can be applied as a single commit. This suggestion is invalid because no changes were made to the code. Suggestions cannot be applied while the pull request is closed. Suggestions cannot be applied while viewing a subset of changes. Only one suggestion per line can be applied in a batch. Add this suggestion to a batch that can be applied as a single commit. Applying suggestions on deleted lines is not supported. You must change the existing code in this line in order to create a valid suggestion. Outdated suggestions cannot be applied. This suggestion has been applied or marked resolved. Suggestions cannot be applied from pending reviews. Suggestions cannot be applied on multi-line comments. Suggestions cannot be applied while the pull request is queued to merge. Suggestion cannot be applied right now. Please check back later.
This PR adds workdistribute mlir op in omp dialect and also in llvm frontend.
This is part 1 of workdistribute frontend support.
Part1 : #154376 [OpenMP] Add workdistribute construct in openMP dialect and in llvm frontend
Part2 : #154377 [flang][openmp] Add parser/semantic support for workdistribute
Part3 : #154378 [flang][openmp] Add Lowering to omp mlir for workdistribute construct