[llvm-branch-commits] [flang] [llvm] [Flang] Add lowering for flang to mlir for thread_limit (PR #175791)

[via llvm-branch-commits]

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

>From bf4e81aac26519e2040e4bc65c869956a3261e8e 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 for flang to mlir for thread_limit

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

diff --git a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
index 2ce00de8a2cc5..547656143291c 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
@@ -668,9 +668,10 @@ bool ClauseProcessor::processThreadLimit(
     lower::StatementContext &stmtCtx,
     mlir::omp::ThreadLimitClauseOps &result) const {
   if (auto *clause = findUniqueClause<omp::clause::ThreadLimit>()) {
-    mlir::Value threadLimitVal =
-        fir::getBase(converter.genExprValue(clause->v, stmtCtx));
-    result.threadLimitVars.push_back(threadLimitVal);
+    result.threadLimitVars.reserve(clause->v.size());
+    for (const ExprTy &vv : clause->v)
+      result.threadLimitVars.push_back(
+          fir::getBase(converter.genExprValue(vv, stmtCtx)));
     return true;
   }
   return false;
diff --git a/flang/lib/Lower/OpenMP/Clauses.cpp b/flang/lib/Lower/OpenMP/Clauses.cpp
index c739249bff211..dcf122f4dee26 100644
--- a/flang/lib/Lower/OpenMP/Clauses.cpp
+++ b/flang/lib/Lower/OpenMP/Clauses.cpp
@@ -1542,9 +1542,16 @@ TaskReduction make(const parser::OmpClause::TaskReduction &inp,
 ThreadLimit make(const parser::OmpClause::ThreadLimit &inp,
                  semantics::SemanticsContext &semaCtx) {
   // inp.v -> parser::OmpThreadLimitClause
-  auto &t1 = std::get<std::list<parser::ScalarIntExpr>>(inp.v.t);
-  assert(!t1.empty());
-  return ThreadLimit{/*Threadlim=*/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<ThreadLimit::Threadlim> v;
+  for (const auto &val : values) {
+    v.push_back(makeExpr(val, semaCtx));
+  }
+  return ThreadLimit{/*Threadlim=*/v};
 }
 
 Threadset make(const parser::OmpClause::Threadset &inp,
diff --git a/flang/test/Lower/OpenMP/thread-limit-dims.f90 b/flang/test/Lower/OpenMP/thread-limit-dims.f90
new file mode 100644
index 0000000000000..d9ca17fd5f75c
--- /dev/null
+++ b/flang/test/Lower/OpenMP/thread-limit-dims.f90
@@ -0,0 +1,62 @@
+! RUN: %flang_fc1 -emit-hlfir %openmp_flags -fopenmp-version=61 %s -o - | FileCheck %s
+
+!===============================================================================
+! `thread_limit` clause with dims modifier (OpenMP 6.1)
+!===============================================================================
+
+! CHECK-LABEL: func @_QPteams_threadlimit_dims3
+subroutine teams_threadlimit_dims3()
+  ! CHECK: omp.teams
+  ! CHECK-SAME: thread_limit(%{{.*}}, %{{.*}}, %{{.*}} : i32, i32, i32)
+  !$omp teams thread_limit(dims(3): 16, 8, 4)
+  call f1()
+  ! CHECK: omp.terminator
+  !$omp end teams
+end subroutine teams_threadlimit_dims3
+
+! CHECK-LABEL: func @_QPteams_threadlimit_dims2
+subroutine teams_threadlimit_dims2()
+  ! CHECK: omp.teams
+  ! CHECK-SAME: thread_limit(%{{.*}}, %{{.*}} : i32, i32)
+  !$omp teams thread_limit(dims(2): 32, 16)
+  call f1()
+  ! CHECK: omp.terminator
+  !$omp end teams
+end subroutine teams_threadlimit_dims2
+
+! CHECK-LABEL: func @_QPteams_threadlimit_dims_var
+subroutine teams_threadlimit_dims_var(a, b, c)
+  integer, intent(in) :: a, b, c
+  ! CHECK: omp.teams
+  ! CHECK-SAME: thread_limit(%{{.*}}, %{{.*}}, %{{.*}} : i32, i32, i32)
+  !$omp teams thread_limit(dims(3): a, b, c)
+  call f1()
+  ! CHECK: omp.terminator
+  !$omp end teams
+end subroutine teams_threadlimit_dims_var
+
+!===============================================================================
+! `thread_limit` clause without dims modifier (legacy)
+!===============================================================================
+
+! CHECK-LABEL: func @_QPteams_threadlimit_legacy
+subroutine teams_threadlimit_legacy(n)
+  integer, intent(in) :: n
+  ! CHECK: omp.teams
+  ! CHECK-SAME: thread_limit(%{{.*}} : i32)
+  !$omp teams thread_limit(n)
+  call f1()
+  ! CHECK: omp.terminator
+  !$omp end teams
+end subroutine teams_threadlimit_legacy
+
+! CHECK-LABEL: func @_QPteams_threadlimit_const
+subroutine teams_threadlimit_const()
+  ! CHECK: omp.teams
+  ! CHECK-SAME: thread_limit(%{{.*}} : i32)
+  !$omp teams thread_limit(64)
+  call f1()
+  ! CHECK: omp.terminator
+  !$omp end teams
+end subroutine teams_threadlimit_const
+
diff --git a/llvm/include/llvm/Frontend/OpenMP/ClauseT.h b/llvm/include/llvm/Frontend/OpenMP/ClauseT.h
index 05ee1ae36a23d..9f38961bc2a06 100644
--- a/llvm/include/llvm/Frontend/OpenMP/ClauseT.h
+++ b/llvm/include/llvm/Frontend/OpenMP/ClauseT.h
@@ -1231,10 +1231,12 @@ struct TaskReductionT {
 
 // V5.2: [13.3] `thread_limit` clause
 template <typename T, typename I, typename E> //
+// V6.1: Extended with dims modifier support
 struct ThreadLimitT {
   using Threadlim = E;
+  // Changed to list to support dims modifier with multiple values (OpenMP 6.1)
   using WrapperTrait = std::true_type;
-  Threadlim v;
+  ListT<Threadlim> v;
 };
 
 // V5.2: [15.10.3] `parallelization-level` clauses