[flang-commits] [flang] cdf59b3 - [Flang][OpenMP] Add lowering and translation for the threadset clause (#215513)

via flang-commits flang-commits at lists.llvm.org
Wed Sep 9 20:51:58 PDT 2026


Author: Ritanya-B-Bharadwaj
Date: 2026-09-10T09:21:51+05:30
New Revision: cdf59b384a48610a6450011279005b1e34f8572d

URL: https://github.com/llvm/llvm-project/commit/cdf59b384a48610a6450011279005b1e34f8572d
DIFF: https://github.com/llvm/llvm-project/commit/cdf59b384a48610a6450011279005b1e34f8572d.diff

LOG: [Flang][OpenMP] Add lowering and translation for the threadset clause (#215513)

Adds Flang lowering, and LLVM IR translation for the OpenMP 6.0 `threadset` clause (spec 14.8) on `task` and `taskloop`.

- Add a `ThreadsetPolicy` enum (`omp_pool`/`omp_team`) and clause on `omp.task` and `omp.taskloop.context`, lowered from Flang.
- Translate `threadset(omp_pool)` to the free-agent task flag (`0x80`); `omp_team` leaves it unset, matching clang.
- Add semantics, lowering, and translation tests.

Related: https://github.com/llvm/llvm-project/pull/135807,
https://github.com/llvm/llvm-project/pull/144409,
https://github.com/llvm/llvm-project/pull/169856

Added: 
    flang/test/Lower/OpenMP/threadset.f90
    flang/test/Semantics/OpenMP/threadset-clause-v60.f90
    mlir/test/Target/LLVMIR/openmp-task-threadset.mlir
    mlir/test/Target/LLVMIR/openmp-taskloop-threadset.mlir

Modified: 
    flang/docs/OpenMPSupport.md
    flang/lib/Lower/OpenMP/ClauseProcessor.cpp
    flang/lib/Lower/OpenMP/ClauseProcessor.h
    flang/lib/Lower/OpenMP/OpenMP.cpp
    flang/test/Semantics/OpenMP/threadset-clause.f90
    llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
    llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
    mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
    mlir/include/mlir/Dialect/OpenMP/OpenMPEnums.td
    mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
    mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
    mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
    mlir/test/Dialect/OpenMP/ops.mlir

Removed: 
    flang/test/Lower/OpenMP/Todo/threadset.f90


################################################################################
diff  --git a/flang/docs/OpenMPSupport.md b/flang/docs/OpenMPSupport.md
index 5a5285421e724..ed14e6a2ab99f 100644
--- a/flang/docs/OpenMPSupport.md
+++ b/flang/docs/OpenMPSupport.md
@@ -245,7 +245,7 @@ Parser/Semantics, MLIR, Lowering, or the OpenMPIRBuilder.
 
 | Feature | Status | Claimed By | Notes | Reviews |
 |:--------|:-------|:-----------|:------|:--------|
-| threadset clause | <span class="part">partial</span> | | Semantics coverage exists (`flang/test/Semantics/OpenMP/threadset-clause.f90`) with ongoing lowering/runtime validation. | [llvm/llvm-project#169856](https://github.com/llvm/llvm-project/pull/169856) |
+| threadset clause | <span class="part">partial</span> | | Frontend (parsing + semantics) and lowering coverage exists (`flang/test/Semantics/OpenMP/threadset-clause.f90`, `flang/test/Semantics/OpenMP/threadset-clause-v60.f90`, `flang/test/Lower/OpenMP/threadset.f90`); the clause is represented on `omp.task`/`omp.taskloop.context` and translated to the runtime free-agent task flag. | [llvm/llvm-project#169856](https://github.com/llvm/llvm-project/pull/169856) |
 | groupprivate directive | <span class="part">partial</span> |  | Semantics and lowering coverage exists (`flang/test/Semantics/OpenMP/groupprivate.f90`, `flang/test/Lower/OpenMP/groupprivate.f90`, `flang/test/Lower/OpenMP/groupprivate-modfile.f90`); lowering in non-teams contexts not fully handled. | [llvm/llvm-project#166199](https://github.com/llvm/llvm-project/pull/166199), [llvm/llvm-project#166214](https://github.com/llvm/llvm-project/pull/166214), [llvm/llvm-project#180934](https://github.com/llvm/llvm-project/pull/180934) |
 | recording of task graphs | <span class="progress">in progress</span> | | Semantics coverage exists (`flang/test/Semantics/OpenMP/taskgraph.f90`). | |
 | workdistribute construct | <span class="part">partial</span> | | Semantics/lowering/transform coverage exists (`flang/test/Semantics/OpenMP/workdistribute01.f90`, `flang/test/Lower/OpenMP/workdistribute.f90`, `flang/test/Transforms/OpenMP/lower-workdistribute-fission.mlir`) including `target teams` placement updates; some team-nesting combinations still intentionally diagnose as unsupported. | [llvm/llvm-project#154377](https://github.com/llvm/llvm-project/pull/154377), [llvm/llvm-project#154378](https://github.com/llvm/llvm-project/pull/154378), [llvm/llvm-project#140523](https://github.com/llvm/llvm-project/pull/140523), [llvm/llvm-project#199006](https://github.com/llvm/llvm-project/pull/199006) |

diff  --git a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
index 22506b14b0fa6..b871f630d3dae 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
@@ -985,6 +985,22 @@ bool ClauseProcessor::processThreadLimit(
   return false;
 }
 
+bool ClauseProcessor::processThreadset(
+    mlir::omp::ThreadsetClauseOps &result) const {
+  using Threadset = omp::clause::Threadset;
+  if (auto *clause = findUniqueClause<Threadset>()) {
+    fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
+    mlir::omp::ThreadsetPolicy policy =
+        clause->v == Threadset::ThreadsetPolicy::Omp_Pool
+            ? mlir::omp::ThreadsetPolicy::omp_pool
+            : mlir::omp::ThreadsetPolicy::omp_team;
+    result.threadset =
+        mlir::omp::ThreadsetPolicyAttr::get(firOpBuilder.getContext(), policy);
+    return true;
+  }
+  return false;
+}
+
 bool ClauseProcessor::processUntied(mlir::omp::UntiedClauseOps &result) const {
   return markClauseOccurrence<omp::clause::Untied>(result.untied);
 }

diff  --git a/flang/lib/Lower/OpenMP/ClauseProcessor.h b/flang/lib/Lower/OpenMP/ClauseProcessor.h
index 6bfa939f52c0b..6b0e0cf606f48 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.h
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.h
@@ -125,6 +125,7 @@ class ClauseProcessor {
   bool processSimd(mlir::omp::OrderedRegionOperands &result) const;
   bool processThreadLimit(lower::StatementContext &stmtCtx,
                           mlir::omp::ThreadLimitClauseOps &result) const;
+  bool processThreadset(mlir::omp::ThreadsetClauseOps &result) const;
   bool processUntied(mlir::omp::UntiedClauseOps &result) const;
   bool processDetach(mlir::omp::DetachClauseOps &result) const;
   // 'Repeatable' clauses: They can appear multiple times in the clause list.

diff  --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 8ea1041adcbbf..89f768c6dc811 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -2746,6 +2746,7 @@ static void genTaskClauses(lower::AbstractConverter &converter,
   cp.processInReduction(loc, clauseOps, inReductionObjects);
   cp.processMergeable(clauseOps);
   cp.processPriority(stmtCtx, clauseOps);
+  cp.processThreadset(clauseOps);
   cp.processUntied(clauseOps);
   cp.processDetach(clauseOps);
 }
@@ -2779,6 +2780,7 @@ static void genTaskloopClauses(
   cp.processNumTasks(stmtCtx, clauseOps);
   cp.processPriority(stmtCtx, clauseOps);
   cp.processReduction(loc, clauseOps, reductionObjects);
+  cp.processThreadset(clauseOps);
   cp.processUntied(clauseOps);
 }
 
@@ -8196,6 +8198,7 @@ static void genOMP(lower::AbstractConverter &converter, lower::SymMap &symTable,
         !std::holds_alternative<clause::Simd>(clause.u) &&
         !std::holds_alternative<clause::ThreadLimit>(clause.u) &&
         !std::holds_alternative<clause::Threads>(clause.u) &&
+        !std::holds_alternative<clause::Threadset>(clause.u) &&
         !std::holds_alternative<clause::UseDeviceAddr>(clause.u) &&
         !std::holds_alternative<clause::UseDevicePtr>(clause.u) &&
         !std::holds_alternative<clause::InReduction>(clause.u) &&

diff  --git a/flang/test/Lower/OpenMP/Todo/threadset.f90 b/flang/test/Lower/OpenMP/Todo/threadset.f90
deleted file mode 100644
index b022baf02654b..0000000000000
--- a/flang/test/Lower/OpenMP/Todo/threadset.f90
+++ /dev/null
@@ -1,10 +0,0 @@
-! RUN: %not_todo_cmd %flang_fc1 -emit-fir -fopenmp -fopenmp-version=60 -o - %s 2>&1 | FileCheck %s
-
-! CHECK: not yet implemented: THREADSET clause is not implemented yet
-
-subroutine f00(x)
-  integer :: x(10)
-  !$omp task threadset(omp_pool)
-  x = x + 1
-  !$omp end task
-end

diff  --git a/flang/test/Lower/OpenMP/threadset.f90 b/flang/test/Lower/OpenMP/threadset.f90
new file mode 100644
index 0000000000000..c758e160afc01
--- /dev/null
+++ b/flang/test/Lower/OpenMP/threadset.f90
@@ -0,0 +1,37 @@
+!RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=60 %s -o - | FileCheck %s
+
+!CHECK-LABEL: func @_QPtask_threadset_pool
+subroutine task_threadset_pool()
+  !CHECK: omp.task threadset(omp_pool) {
+  !$omp task threadset(omp_pool)
+  !CHECK: omp.terminator
+  !$omp end task
+end subroutine task_threadset_pool
+
+!CHECK-LABEL: func @_QPtask_threadset_team
+subroutine task_threadset_team()
+  !CHECK: omp.task threadset(omp_team) {
+  !$omp task threadset(omp_team)
+  !CHECK: omp.terminator
+  !$omp end task
+end subroutine task_threadset_team
+
+!CHECK-LABEL: func @_QPtaskloop_threadset_pool
+subroutine taskloop_threadset_pool()
+  integer :: i
+  !CHECK: omp.taskloop.context threadset(omp_pool)
+  !$omp taskloop threadset(omp_pool)
+  do i = 1, 10
+  end do
+  !$omp end taskloop
+end subroutine taskloop_threadset_pool
+
+!CHECK-LABEL: func @_QPtaskloop_threadset_team
+subroutine taskloop_threadset_team()
+  integer :: i
+  !CHECK: omp.taskloop.context threadset(omp_team)
+  !$omp taskloop threadset(omp_team)
+  do i = 1, 10
+  end do
+  !$omp end taskloop
+end subroutine taskloop_threadset_team

diff  --git a/flang/test/Semantics/OpenMP/threadset-clause-v60.f90 b/flang/test/Semantics/OpenMP/threadset-clause-v60.f90
new file mode 100644
index 0000000000000..f5542e4941cc2
--- /dev/null
+++ b/flang/test/Semantics/OpenMP/threadset-clause-v60.f90
@@ -0,0 +1,38 @@
+!RUN: %python %S/../test_errors.py %s %flang -fopenmp -fopenmp-version=60
+
+subroutine f00(x)
+  integer :: x(10)
+  integer :: i
+
+! Valid uses on TASK.
+!$omp task threadset(omp_pool)
+  x = x + 1
+!$omp end task
+
+!$omp task threadset(omp_team)
+  x = x + 1
+!$omp end task
+
+! Valid uses on TASKLOOP.
+!$omp taskloop threadset(omp_pool)
+  do i = 1, 10
+  end do
+!$omp end taskloop
+
+!$omp taskloop threadset(omp_team)
+  do i = 1, 10
+  end do
+!$omp end taskloop
+
+! At most one THREADSET clause is allowed on the directive.
+!ERROR: At most one THREADSET clause can appear on TASK directive
+!$omp task threadset(omp_pool) threadset(omp_team)
+  x = x + 1
+!$omp end task
+
+! THREADSET is only allowed on TASK and TASKLOOP.
+!ERROR: THREADSET clause is not allowed on PARALLEL directive
+!$omp parallel threadset(omp_pool)
+  x = x + 1
+!$omp end parallel
+end

diff  --git a/flang/test/Semantics/OpenMP/threadset-clause.f90 b/flang/test/Semantics/OpenMP/threadset-clause.f90
index 9c811d440fb6a..65a7ed74037cc 100644
--- a/flang/test/Semantics/OpenMP/threadset-clause.f90
+++ b/flang/test/Semantics/OpenMP/threadset-clause.f90
@@ -2,8 +2,15 @@
 
 subroutine f00(x)
   integer :: x(10)
+  integer :: i
 !ERROR: THREADSET clause is not allowed on TASK directive in OpenMP v4.5, try -fopenmp-version=60
 !$omp task threadset(omp_pool)
   x = x + 1
 !$omp end task
+
+!ERROR: THREADSET clause is not allowed on TASKLOOP directive in OpenMP v4.5, try -fopenmp-version=60
+!$omp taskloop threadset(omp_team)
+  do i = 1, 10
+  end do
+!$omp end taskloop
 end

diff  --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index 5f3da79cf53b1..e4fffae9a7bbe 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -1628,6 +1628,8 @@ class OpenMPIRBuilder {
   /// \param TaskContextStructPtrVal If non-null, a pointer to  to be placed
   ///                                immediately after the {lower bound, upper
   ///                                bound, step} values in the task data.
+  /// \param FreeAgent If `true`, the generated tasks are eligible to be
+  ///                  executed by a free-agent thread (threadset(omp_pool)).
   LLVM_ABI InsertPointOrErrorTy createTaskloop(
       const LocationDescription &Loc, InsertPointTy AllocaIP,
       ArrayRef<BasicBlock *> DeallocBlocks, BodyGenCallbackTy BodyGenCB,
@@ -1637,7 +1639,7 @@ class OpenMPIRBuilder {
       int Sched = 0, Value *Final = nullptr, bool Mergeable = false,
       Value *Priority = nullptr, uint64_t NumOfCollapseLoops = 1,
       TaskDupCallbackTy DupCB = nullptr,
-      Value *TaskContextStructPtrVal = nullptr);
+      Value *TaskContextStructPtrVal = nullptr, bool FreeAgent = false);
 
   /// Generator for `#omp task`
   ///
@@ -1664,13 +1666,16 @@ class OpenMPIRBuilder {
   /// \param Mergeable	 If the given task is `mergeable`
   /// \param priority `priority-value' specifies the execution order of the
   ///                 tasks that is generated by the construct
+  /// \param FreeAgent If `true`, the task is eligible to be executed by a
+  ///                  free-agent thread (threadset(omp_pool)).
   LLVM_ABI InsertPointOrErrorTy createTask(
       const LocationDescription &Loc, InsertPointTy AllocaIP,
       ArrayRef<BasicBlock *> DeallocBlocks, BodyGenCallbackTy BodyGenCB,
       bool Tied = true, Value *Final = nullptr, Value *IfCondition = nullptr,
       const DependenciesInfo &Dependencies = {},
       const AffinityData &Affinities = {}, bool Mergeable = false,
-      Value *EventHandle = nullptr, Value *Priority = nullptr);
+      Value *EventHandle = nullptr, Value *Priority = nullptr,
+      bool FreeAgent = false);
 
   /// Generator for the taskgroup construct
   ///

diff  --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 8d6a35a368b4f..1a9ed6d9f69ea 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -2549,7 +2549,7 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTaskloop(
     Value *LBVal, Value *UBVal, Value *StepVal, bool Untied, Value *IfCond,
     Value *GrainSize, bool NoGroup, int Sched, Value *Final, bool Mergeable,
     Value *Priority, uint64_t NumOfCollapseLoops, TaskDupCallbackTy DupCB,
-    Value *TaskContextStructPtrVal) {
+    Value *TaskContextStructPtrVal, bool FreeAgent) {
 
   if (!updateToLocation(Loc))
     return InsertPointTy();
@@ -2629,7 +2629,8 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTaskloop(
                        TaskloopAllocaBB, CLI, TaskDupFn, ToBeDeleted, IfCond,
                        GrainSize, NoGroup, Sched, FakeLB, FakeUB, FakeStep,
                        FakeSharedsTy, Final, Mergeable, Priority,
-                       NumOfCollapseLoops](Function &OutlinedFn) mutable {
+                       NumOfCollapseLoops,
+                       FreeAgent](Function &OutlinedFn) mutable {
     // Replace the Stale CI by appropriate RTL function call.
     assert(OutlinedFn.hasOneUse() &&
            "there must be a single user for the outlined function");
@@ -2671,6 +2672,8 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTaskloop(
     // Task is not mergeable if (Flags & 4) == 0.
     // Task is priority if (Flags & 32) == 32.
     // Task is not priority if (Flags & 32) == 0.
+    // Task is free-agent eligible if (Flags & 128) == 128.
+    // Task is not free-agent eligible if (Flags & 128) == 0.
     Value *Flags = Builder.getInt32(Untied ? 0 : 1);
     if (Final)
       Flags = Builder.CreateOr(Builder.getInt32(2), Flags);
@@ -2678,6 +2681,8 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTaskloop(
       Flags = Builder.CreateOr(Builder.getInt32(4), Flags);
     if (Priority)
       Flags = Builder.CreateOr(Builder.getInt32(32), Flags);
+    if (FreeAgent)
+      Flags = Builder.CreateOr(Builder.getInt32(128), Flags);
 
     Value *TaskSize = Builder.getInt64(
         divideCeil(M.getDataLayout().getTypeSizeInBits(Task), 8));
@@ -2897,7 +2902,7 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTask(
     ArrayRef<BasicBlock *> DeallocBlocks, BodyGenCallbackTy BodyGenCB,
     bool Tied, Value *Final, Value *IfCondition,
     const DependenciesInfo &Dependencies, const AffinityData &Affinities,
-    bool Mergeable, Value *EventHandle, Value *Priority) {
+    bool Mergeable, Value *EventHandle, Value *Priority, bool FreeAgent) {
 
   if (!updateToLocation(Loc))
     return InsertPointTy();
@@ -2945,7 +2950,7 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTask(
       Builder, AllocaIP, ToBeDeleted, TaskAllocaIP, "global.tid", false));
 
   OI->PostOutlineCB = [this, Ident, Tied, Final, IfCondition, Dependencies,
-                       Affinities, Mergeable, Priority, EventHandle,
+                       Affinities, Mergeable, Priority, EventHandle, FreeAgent,
                        TaskAllocaBB,
                        ToBeDeleted](Function &OutlinedFn) mutable {
     // Replace the Stale CI by appropriate RTL function call.
@@ -2978,6 +2983,8 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTask(
     // Task is not detachable iff (Flags & 64) == 0.
     // Task is priority iff (Flags & 32) == 32.
     // Task is not priority iff (Flags & 32) == 0.
+    // Task is free-agent eligible iff (Flags & 128) == 128.
+    // Task is not free-agent eligible iff (Flags & 128) == 0.
     // TODO: Handle the other flags.
     Value *Flags = Builder.getInt32(Tied);
     auto *ConstIfCondition = dyn_cast_or_null<ConstantInt>(IfCondition);
@@ -2994,6 +3001,8 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTask(
       Flags = Builder.CreateOr(Builder.getInt32(64), Flags);
     if (Priority)
       Flags = Builder.CreateOr(Builder.getInt32(32), Flags);
+    if (FreeAgent)
+      Flags = Builder.CreateOr(Builder.getInt32(128), Flags);
 
     // Argument - `sizeof_kmp_task_t` (TaskSize)
     // Tasksize refers to the size in bytes of kmp_task_t data structure

diff  --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
index bd9f1321f077b..55a8bba40404b 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td
@@ -1716,6 +1716,33 @@ class OpenMP_ThreadLimitClauseSkip<
 
 def OpenMP_ThreadLimitClause : OpenMP_ThreadLimitClauseSkip<>;
 
+//===----------------------------------------------------------------------===//
+// V6.0: [14.8] `threadset` clause
+//===----------------------------------------------------------------------===//
+
+class OpenMP_ThreadsetClauseSkip<
+    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
+    OptionalAttr<ThreadsetPolicyAttr>:$threadset
+  );
+
+  let optAssemblyFormat = [{
+    `threadset` `(` custom<ClauseAttr>($threadset) `)`
+  }];
+
+  let description = [{
+    The `threadset` clause specifies the set of threads that may execute the
+    generated task. When the policy is `omp_pool`, the task is eligible to be
+    executed by a free-agent thread. When the policy is `omp_team`, the task
+    may only be executed by a thread in the current team.
+  }];
+}
+
+def OpenMP_ThreadsetClause : OpenMP_ThreadsetClauseSkip<>;
+
 //===----------------------------------------------------------------------===//
 // V5.2: [9.1.1] `sizes` clause
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPEnums.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPEnums.td
index f1b9c033ea0b9..50b02c0d9e61c 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPEnums.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPEnums.td
@@ -376,6 +376,20 @@ def FallbackModifierAttr : OpenMP_EnumAttr<FallbackModifier,
                                             "fallback_modifier">;
 
 //===----------------------------------------------------------------------===//
+// threadset_policy enum.
+//===----------------------------------------------------------------------===//
+
+def ThreadsetOmpPool : I32EnumAttrCase<"omp_pool", 0>;
+def ThreadsetOmpTeam : I32EnumAttrCase<"omp_team", 1>;
+
+def ThreadsetPolicy : OpenMP_I32EnumAttr<
+    "ThreadsetPolicy",
+    "threadset policy", [
+      ThreadsetOmpPool,
+      ThreadsetOmpTeam
+    ]>;
+
+def ThreadsetPolicyAttr : OpenMP_EnumAttr<ThreadsetPolicy, "threadset_policy">;
 // interop_type for interop construct
 //===----------------------------------------------------------------------===//
 

diff  --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 0df33d96319f5..b2ad6abd40ed3 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -1037,7 +1037,7 @@ def TaskOp : OpenMP_Op<"task", traits = [
     OpenMP_FinalClause, OpenMP_IfClause,
     OpenMP_InReductionClauseSkip<assemblyFormat = true>,
     OpenMP_MergeableClause, OpenMP_PriorityClause, OpenMP_PrivateClause,
-    OpenMP_UntiedClause, OpenMP_DetachClause
+    OpenMP_ThreadsetClause, OpenMP_UntiedClause, OpenMP_DetachClause
   ], singleRegion = true> {
   let summary = "task construct";
   let description = [{
@@ -1084,7 +1084,7 @@ def TaskloopContextOp : OpenMP_Op<"taskloop.context", traits = [
     OpenMP_IfClause, OpenMP_InReductionClauseSkip<assemblyFormat = true>,
     OpenMP_MergeableClause, OpenMP_NogroupClause, OpenMP_NumTasksClause,
     OpenMP_PriorityClause, OpenMP_PrivateClause, OpenMP_ReductionClause,
-    OpenMP_UntiedClause
+    OpenMP_ThreadsetClause, OpenMP_UntiedClause
   ], singleRegion = true> {
   let summary = "OutlinableOpenMPOpInterface wrapper for taskloop construct";
   let description = [{

diff  --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 1f210ef60fe39..a74d0eb51564c 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -3843,19 +3843,20 @@ LogicalResult DeclareReductionOp::verifyRegions() {
 void TaskOp::build(OpBuilder &builder, OperationState &state,
                    const TaskOperands &clauses) {
   MLIRContext *ctx = builder.getContext();
-  TaskOp::build(
-      builder, state, clauses.iterated, clauses.affinityVars,
-      clauses.allocateVars, clauses.allocatorVars,
-      makeDenseI64ArrayAttr(ctx, clauses.allocateAlignments),
-      makeDenseI64ArrayAttr(ctx, clauses.allocatePrivateIndices),
-      makeArrayAttr(ctx, clauses.dependKinds), clauses.dependVars,
-      makeArrayAttr(ctx, clauses.dependIteratedKinds), clauses.dependIterated,
-      clauses.final, clauses.ifExpr, clauses.inReductionVars,
-      makeDenseBoolArrayAttr(ctx, clauses.inReductionByref),
-      makeArrayAttr(ctx, clauses.inReductionSyms), clauses.mergeable,
-      clauses.priority, /*private_vars=*/clauses.privateVars,
-      /*private_syms=*/makeArrayAttr(ctx, clauses.privateSyms),
-      clauses.privateNeedsBarrier, clauses.untied, clauses.eventHandle);
+  TaskOp::build(builder, state, clauses.iterated, clauses.affinityVars,
+                clauses.allocateVars, clauses.allocatorVars,
+                makeDenseI64ArrayAttr(ctx, clauses.allocateAlignments),
+                makeDenseI64ArrayAttr(ctx, clauses.allocatePrivateIndices),
+                makeArrayAttr(ctx, clauses.dependKinds), clauses.dependVars,
+                makeArrayAttr(ctx, clauses.dependIteratedKinds),
+                clauses.dependIterated, clauses.final, clauses.ifExpr,
+                clauses.inReductionVars,
+                makeDenseBoolArrayAttr(ctx, clauses.inReductionByref),
+                makeArrayAttr(ctx, clauses.inReductionSyms), clauses.mergeable,
+                clauses.priority, /*private_vars=*/clauses.privateVars,
+                /*private_syms=*/makeArrayAttr(ctx, clauses.privateSyms),
+                clauses.privateNeedsBarrier, clauses.threadset, clauses.untied,
+                clauses.eventHandle);
 }
 
 LogicalResult TaskOp::verify() {
@@ -3925,7 +3926,8 @@ void TaskloopContextOp::build(OpBuilder &builder, OperationState &state,
       /*private_syms=*/makeArrayAttr(ctx, clauses.privateSyms),
       clauses.privateNeedsBarrier, clauses.reductionMod, clauses.reductionVars,
       makeDenseBoolArrayAttr(ctx, clauses.reductionByref),
-      makeArrayAttr(ctx, clauses.reductionSyms), clauses.untied);
+      makeArrayAttr(ctx, clauses.reductionSyms), clauses.threadset,
+      clauses.untied);
   state.addAttribute("omp.combined", UnitAttr::get(ctx));
 }
 

diff  --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index e6e784f37d6b2..1d3827783b31f 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -3521,7 +3521,8 @@ convertOmpTaskOp(omp::TaskOp taskOp, llvm::IRBuilderBase &builder,
           moduleTranslation.lookupValue(taskOp.getIfExpr()), dependencies, ad,
           taskOp.getMergeable(),
           moduleTranslation.lookupValue(taskOp.getEventHandle()),
-          moduleTranslation.lookupValue(taskOp.getPriority()));
+          moduleTranslation.lookupValue(taskOp.getPriority()),
+          taskOp.getThreadset() == omp::ThreadsetPolicy::omp_pool);
 
   if (failed(handleError(afterIP, *taskOp)))
     return failure();
@@ -4158,7 +4159,8 @@ convertOmpTaskloopContextOp(omp::TaskloopContextOp contextOp,
           contextOp.getMergeable(),
           moduleTranslation.lookupValue(contextOp.getPriority()),
           loopOp.getCollapseNumLoops(), taskDupOrNull,
-          taskStructMgr.getStructPtr());
+          taskStructMgr.getStructPtr(),
+          contextOp.getThreadset() == omp::ThreadsetPolicy::omp_pool);
 
   if (failed(handleError(afterIP, opInst)))
     return failure();

diff  --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 92cc73f59ae58..d20ed5871c8d1 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -2486,6 +2486,23 @@ func.func @omp_task(%bool_var: i1, %i64_var: i64, %i32_var: i32, %data_var: memr
     omp.terminator
   }
 
+  // Checking `threadset` clause
+  // CHECK: omp.task threadset(omp_pool) {
+  omp.task threadset(omp_pool) {
+    // CHECK: "test.foo"() : () -> ()
+    "test.foo"() : () -> ()
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
+  // CHECK: omp.task threadset(omp_team) {
+  omp.task threadset(omp_team) {
+    // CHECK: "test.foo"() : () -> ()
+    "test.foo"() : () -> ()
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
   // Checking `in_reduction` clause
   %c1 = arith.constant 1 : i32
   // CHECK: %[[redn_var1:.*]] = llvm.alloca %{{.*}} x f32 : (i32) -> !llvm.ptr
@@ -3015,6 +3032,30 @@ func.func @omp_taskloop(%lb: i32, %ub: i32, %step: i32) -> () {
     omp.terminator
   } {omp.combined}
 
+  // CHECK: omp.taskloop.context threadset(omp_pool) {
+  omp.taskloop.context threadset(omp_pool) {
+    // CHECK: omp.taskloop.wrapper {
+    omp.taskloop.wrapper {
+      omp.loop_nest (%i, %j) : i32 = (%lb, %ub) to (%ub, %lb) step (%step, %step) {
+        // CHECK: omp.yield
+        omp.yield
+      }
+    }
+    omp.terminator
+  } {omp.combined}
+
+  // CHECK: omp.taskloop.context threadset(omp_team) {
+  omp.taskloop.context threadset(omp_team) {
+    // CHECK: omp.taskloop.wrapper {
+    omp.taskloop.wrapper {
+      omp.loop_nest (%i, %j) : i32 = (%lb, %ub) to (%ub, %lb) step (%step, %step) {
+        // CHECK: omp.yield
+        omp.yield
+      }
+    }
+    omp.terminator
+  } {omp.combined}
+
   %testf32 = "test.f32"() : () -> (!llvm.ptr)
   %testf32_2 = "test.f32"() : () -> (!llvm.ptr)
   // CHECK: omp.taskloop.context in_reduction(@add_f32 %{{.+}} -> %{{.+}}, @add_f32 %{{.+}} -> %{{.+}} : !llvm.ptr, !llvm.ptr) {

diff  --git a/mlir/test/Target/LLVMIR/openmp-task-threadset.mlir b/mlir/test/Target/LLVMIR/openmp-task-threadset.mlir
new file mode 100644
index 0000000000000..26aa1fb8ae0f2
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-task-threadset.mlir
@@ -0,0 +1,26 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// Check that threadset(omp_pool) sets the free-agent task flag (0x80 = 128) and
+// threadset(omp_team) does not.
+
+llvm.func @task_threadset_pool() {
+  omp.task threadset(omp_pool) {
+    omp.terminator
+  }
+  llvm.return
+}
+
+// CHECK-LABEL: define void @task_threadset_pool()
+// CHECK: call ptr @__kmpc_omp_task_alloc(ptr @{{.+}}, i32 %{{.+}}, i32 129, i64 {{.+}}, i64 {{.+}}, ptr @{{.+}})
+
+// -----
+
+llvm.func @task_threadset_team() {
+  omp.task threadset(omp_team) {
+    omp.terminator
+  }
+  llvm.return
+}
+
+// CHECK-LABEL: define void @task_threadset_team()
+// CHECK: call ptr @__kmpc_omp_task_alloc(ptr @{{.+}}, i32 %{{.+}}, i32 1, i64 {{.+}}, i64 {{.+}}, ptr @{{.+}})

diff  --git a/mlir/test/Target/LLVMIR/openmp-taskloop-threadset.mlir b/mlir/test/Target/LLVMIR/openmp-taskloop-threadset.mlir
new file mode 100644
index 0000000000000..332fee75ecdb6
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-taskloop-threadset.mlir
@@ -0,0 +1,50 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// Check that threadset(omp_pool) sets the free-agent task flag (0x80 = 128) and
+// threadset(omp_team) does not.
+
+omp.private {type = private} @_QFtestEi_private_i32 : i32
+
+llvm.func @_QPtest_pool() {
+  %0 = llvm.mlir.constant(1 : i64) : i64
+  %1 = llvm.alloca %0 x i32 {bindc_name = "i"} : (i64) -> !llvm.ptr
+  %7 = llvm.mlir.constant(1 : i32) : i32
+  %8 = llvm.mlir.constant(5 : i32) : i32
+  %9 = llvm.mlir.constant(1 : i32) : i32
+  omp.taskloop.context threadset(omp_pool) private(@_QFtestEi_private_i32 %1 -> %arg1 : !llvm.ptr) {
+    omp.taskloop.wrapper {
+      omp.loop_nest (%arg2) : i32 = (%7) to (%8) inclusive step (%9) {
+        llvm.store %arg2, %arg1 : i32, !llvm.ptr
+        omp.yield
+      }
+    }
+    omp.terminator
+  } {omp.combined}
+  llvm.return
+}
+
+// CHECK-LABEL: define void @_QPtest_pool()
+// CHECK: call ptr @__kmpc_omp_task_alloc(ptr @{{.+}}, i32 %{{.+}}, i32 129, i64 {{.+}}, i64 {{.+}}, ptr @_QPtest_pool..omp_par)
+
+// -----
+
+llvm.func @_QPtest_team() {
+  %0 = llvm.mlir.constant(1 : i64) : i64
+  %1 = llvm.alloca %0 x i32 {bindc_name = "i"} : (i64) -> !llvm.ptr
+  %7 = llvm.mlir.constant(1 : i32) : i32
+  %8 = llvm.mlir.constant(5 : i32) : i32
+  %9 = llvm.mlir.constant(1 : i32) : i32
+  omp.taskloop.context threadset(omp_team) private(@_QFtestEi_private_i32 %1 -> %arg1 : !llvm.ptr) {
+    omp.taskloop.wrapper {
+      omp.loop_nest (%arg2) : i32 = (%7) to (%8) inclusive step (%9) {
+        llvm.store %arg2, %arg1 : i32, !llvm.ptr
+        omp.yield
+      }
+    }
+    omp.terminator
+  } {omp.combined}
+  llvm.return
+}
+
+// CHECK-LABEL: define void @_QPtest_team()
+// CHECK: call ptr @__kmpc_omp_task_alloc(ptr @{{.+}}, i32 %{{.+}}, i32 1, i64 {{.+}}, i64 {{.+}}, ptr @_QPtest_team..omp_par)


        


More information about the flang-commits mailing list