[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 created https://github.com/llvm/llvm-project/pull/175791

None

>From 1797a81185d0807391fb58e6538a37b7170d04e6 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Tue, 13 Jan 2026 21:44:02 +0530
Subject: [PATCH] [Flang] Add lowering for flang to mlir for thread_limit

---
 flang/lib/Lower/OpenMP/ClauseProcessor.cpp    | 20 +++++-
 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, 92 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 18bab01d94365..a8f4ffd1e4f40 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
@@ -668,9 +668,23 @@ 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.threadLimitDimsValues.push_back(threadLimitVal);
+    // The thread_limit clause accepts a list of values.
+    // With dims modifier (OpenMP 6.1): multiple values for multi-dimensional
+    // Without dims modifier: single value
+    assert(!clause->v.empty());
+
+    // If multiple values, this indicates dims modifier is present
+    if (clause->v.size() > 1) {
+      fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
+      result.threadLimitNumDims =
+          firOpBuilder.getI64IntegerAttr(clause->v.size());
+    }
+
+    // Populate all values
+    for (const auto &val : clause->v) {
+      result.threadLimitDimsValues.push_back(
+          fir::getBase(converter.genExprValue(val, stmtCtx)));
+    }
     return true;
   }
   return false;
diff --git a/flang/lib/Lower/OpenMP/Clauses.cpp b/flang/lib/Lower/OpenMP/Clauses.cpp
index a2716fb22a75c..18c2e60c4936d 100644
--- a/flang/lib/Lower/OpenMP/Clauses.cpp
+++ b/flang/lib/Lower/OpenMP/Clauses.cpp
@@ -1559,9 +1559,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..e3f7f2250ff33
--- /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(dims(3): %{{.*}}, %{{.*}}, %{{.*}} : 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(dims(2): %{{.*}}, %{{.*}} : 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(dims(3): %{{.*}}, %{{.*}}, %{{.*}} : 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 7543f27136e7d..5dfa11d2fd3d1 100644
--- a/llvm/include/llvm/Frontend/OpenMP/ClauseT.h
+++ b/llvm/include/llvm/Frontend/OpenMP/ClauseT.h
@@ -1240,10 +1240,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