[llvm-branch-commits] [flang] [llvm] [Flang] Add lowering from flang to mlir for num_threads (PR #175792)

[via llvm-branch-commits]

https://github.com/skc7 updated https://github.com/llvm/llvm-project/pull/175792

>From d13ee30093a64a6da59e35172b9b58b30516cfcd 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] [Flang] Add lowering from flang to mlir for num_threads

---
 flang/lib/Lower/OpenMP/ClauseProcessor.cpp   |  8 ++-
 flang/lib/Lower/OpenMP/Clauses.cpp           | 13 ++++-
 flang/test/Lower/OpenMP/num-threads-dims.f90 | 61 ++++++++++++++++++++
 llvm/include/llvm/Frontend/OpenMP/ClauseT.h  |  4 +-
 4 files changed, 79 insertions(+), 7 deletions(-)
 create mode 100644 flang/test/Lower/OpenMP/num-threads-dims.f90

diff --git a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
index 159fe371c7ef3..5d2ebe6e1572a 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
@@ -515,9 +515,11 @@ bool ClauseProcessor::processNumThreads(
     lower::StatementContext &stmtCtx,
     mlir::omp::NumThreadsClauseOps &result) const {
   if (auto *clause = findUniqueClause<omp::clause::NumThreads>()) {
-    // OMPIRBuilder expects `NUM_THREADS` clause as a `Value`.
-    result.numThreadsVars.push_back(
-        fir::getBase(converter.genExprValue(clause->v, stmtCtx)));
+    // OMPIRBuilder expects `NUM_THREADS` clause as a list of Values.
+    for (const auto &expr : clause->v) {
+      result.numThreadsVars.push_back(
+          fir::getBase(converter.genExprValue(expr, stmtCtx)));
+    }
     return true;
   }
   return false;
diff --git a/flang/lib/Lower/OpenMP/Clauses.cpp b/flang/lib/Lower/OpenMP/Clauses.cpp
index c739249bff211..cf03d4af788a5 100644
--- a/flang/lib/Lower/OpenMP/Clauses.cpp
+++ b/flang/lib/Lower/OpenMP/Clauses.cpp
@@ -1290,9 +1290,16 @@ NumTeams make(const parser::OmpClause::NumTeams &inp,
 NumThreads make(const parser::OmpClause::NumThreads &inp,
                 semantics::SemanticsContext &semaCtx) {
   // inp.v -> parser::OmpNumThreadsClause
-  auto &t1 = std::get<std::list<parser::ScalarIntExpr>>(inp.v.t);
-  assert(!t1.empty());
-  return NumThreads{/*Nthreads=*/makeExpr(t1.front(), semaCtx)};
+  // With dims modifier (OpenMP 6.1): multiple values for multi-dimensional grid
+  // Without dims modifier: single value
+  auto &values = std::get<std::list<parser::ScalarIntExpr>>(inp.v.t);
+  assert(!values.empty());
+
+  List<NumThreads::Nthreads> v;
+  for (const auto &val : values) {
+    v.push_back(makeExpr(val, semaCtx));
+  }
+  return NumThreads{/*Nthreads=*/v};
 }
 
 // OmpxAttribute: empty
diff --git a/flang/test/Lower/OpenMP/num-threads-dims.f90 b/flang/test/Lower/OpenMP/num-threads-dims.f90
new file mode 100644
index 0000000000000..f3a8d706b7283
--- /dev/null
+++ b/flang/test/Lower/OpenMP/num-threads-dims.f90
@@ -0,0 +1,61 @@
+! RUN: %flang_fc1 -emit-hlfir %openmp_flags -fopenmp-version=61 %s -o - | FileCheck %s
+
+!===============================================================================
+! `num_threads` clause with dims modifier (OpenMP 6.1)
+!===============================================================================
+
+! CHECK-LABEL: func @_QPparallel_numthreads_dims4
+subroutine parallel_numthreads_dims4()
+  ! CHECK: omp.parallel
+  ! CHECK-SAME: num_threads(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : i32, i32, i32, i32)
+  !$omp parallel num_threads(dims(4): 4, 5, 6, 7)
+  call f1()
+  ! CHECK: omp.terminator
+  !$omp end parallel
+end subroutine parallel_numthreads_dims4
+
+! CHECK-LABEL: func @_QPparallel_numthreads_dims2
+subroutine parallel_numthreads_dims2()
+  ! CHECK: omp.parallel
+  ! CHECK-SAME: num_threads(%{{.*}}, %{{.*}} : i32, i32)
+  !$omp parallel num_threads(dims(2): 8, 4)
+  call f1()
+  ! CHECK: omp.terminator
+  !$omp end parallel
+end subroutine parallel_numthreads_dims2
+
+! CHECK-LABEL: func @_QPparallel_numthreads_dims_var
+subroutine parallel_numthreads_dims_var(a, b, c)
+  integer, intent(in) :: a, b, c
+  ! CHECK: omp.parallel
+  ! CHECK-SAME: num_threads(%{{.*}}, %{{.*}}, %{{.*}} : i32, i32, i32)
+  !$omp parallel num_threads(dims(3): a, b, c)
+  call f1()
+  ! CHECK: omp.terminator
+  !$omp end parallel
+end subroutine parallel_numthreads_dims_var
+
+!===============================================================================
+! `num_threads` clause without dims modifier (legacy)
+!===============================================================================
+
+! CHECK-LABEL: func @_QPparallel_numthreads_legacy
+subroutine parallel_numthreads_legacy(n)
+  integer, intent(in) :: n
+  ! CHECK: omp.parallel
+  ! CHECK-SAME: num_threads(%{{.*}} : i32)
+  !$omp parallel num_threads(n)
+  call f1()
+  ! CHECK: omp.terminator
+  !$omp end parallel
+end subroutine parallel_numthreads_legacy
+
+! CHECK-LABEL: func @_QPparallel_numthreads_const
+subroutine parallel_numthreads_const()
+  ! CHECK: omp.parallel
+  ! CHECK-SAME: num_threads(%{{.*}} : i32)
+  !$omp parallel num_threads(16)
+  call f1()
+  ! CHECK: omp.terminator
+  !$omp end parallel
+end subroutine parallel_numthreads_const
diff --git a/llvm/include/llvm/Frontend/OpenMP/ClauseT.h b/llvm/include/llvm/Frontend/OpenMP/ClauseT.h
index 05ee1ae36a23d..e9f972b5b313e 100644
--- a/llvm/include/llvm/Frontend/OpenMP/ClauseT.h
+++ b/llvm/include/llvm/Frontend/OpenMP/ClauseT.h
@@ -1011,11 +1011,13 @@ struct NumTeamsT {
 };
 
 // V5.2: [10.1.2] `num_threads` clause
+// V6.1: Extended with dims modifier support
 template <typename T, typename I, typename E> //
 struct NumThreadsT {
   using Nthreads = E;
+  // Changed to list to support dims modifier with multiple values (OpenMP 6.1)
   using WrapperTrait = std::true_type;
-  Nthreads v;
+  ListT<Nthreads> v;
 };
 
 template <typename T, typename I, typename E> //