[llvm-branch-commits] [mlir] [OpenMP][MLIR] Add thread_limit mlir->llvm lowering (PR #179608)

via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Tue Oct 6 23:15:17 PDT 2026


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

>From 56123a6426ff5c5a1318a06be8a0a3041a5b715c Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Tue, 3 Feb 2026 15:29:10 +0530
Subject: [PATCH 1/6] [OpenMP][MLIR] Add num_threads mlir->llvm lowering

---
 .../LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp     | 2 +-
 mlir/test/Target/LLVMIR/openmp-todo.mlir                    | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 1a3f9ef22d2724..ba628798c7fa85 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -9162,7 +9162,7 @@ static void extractHostEvalClauses(
                   break;
                 }
               }
-            } else {
+            } else
               llvm_unreachable("unsupported host_eval use");
             }
           })
diff --git a/mlir/test/Target/LLVMIR/openmp-todo.mlir b/mlir/test/Target/LLVMIR/openmp-todo.mlir
index 320f678594bf18..a29e9f836e4bf4 100644
--- a/mlir/test/Target/LLVMIR/openmp-todo.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-todo.mlir
@@ -670,10 +670,10 @@ llvm.func @teams_num_teams_multi_dim(%lb : i32, %ub : i32) {
 
 // -----
 
-llvm.func @parallel_num_threads_multi_dim_standalone(%lb : i32, %ub : i32) {
-  // expected-error at below {{not yet implemented: Unhandled clause num_threads with multi-dimensional values outside target region in omp.parallel operation}}
+llvm.func @parallel_num_threads_too_many_dims(%lb : i32, %ub : i32) {
+  // expected-error at below {{not yet implemented: Unhandled clause num_threads with more than 3 dimensions in omp.parallel operation}}
   // expected-error at below {{LLVM Translation failed for operation: omp.parallel}}
-  omp.parallel num_threads(%lb, %ub : i32, i32) {
+  omp.parallel num_threads(%lb, %ub, %lb, %ub : i32, i32, i32, i32) {
     omp.terminator
   }
   llvm.return

>From 96551d533f17c5749f049fc93ea7ca3927f081f9 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Wed, 4 Feb 2026 10:21:06 +0530
Subject: [PATCH 2/6] [OpenMP][MLIR] Add thread_limit mlir->llvm lowering

---
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      | 129 +++++++++++-------
 .../LLVMIR/openmp-target-launch-host.mlir     |   6 +-
 mlir/test/Target/LLVMIR/openmp-teams.mlir     |  50 ++++---
 mlir/test/Target/LLVMIR/openmp-todo.mlir      |   6 +-
 4 files changed, 117 insertions(+), 74 deletions(-)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index ba628798c7fa85..274fa7b5578802 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -531,8 +531,8 @@ static LogicalResult checkImplementationStatus(Operation &op) {
   };
 
   auto checkThreadLimit = [&todo](auto op, LogicalResult &result) {
-    if (op.hasThreadLimitMultiDim())
-      result = todo("thread_limit with multi-dimensional values");
+    if (op.getThreadLimitDimsCount() > 3)
+      result = todo("thread_limit with more than 3 dimensions");
   };
   auto checkMap = [&todo](auto op, LogicalResult &result) {
     if (!op.getMapIterated().empty())
@@ -9126,12 +9126,14 @@ static llvm::IRBuilderBase::InsertPoint createDeviceArgumentAccessor(
 ///
 /// Loop bounds and steps are only optionally populated, if output vectors are
 /// provided.
-static void extractHostEvalClauses(
-    omp::TargetOp targetOp, llvm::SmallVectorImpl<Value> &numThreadsVars,
-    Value &numTeamsLower, Value &numTeamsUpper, Value &threadLimit,
-    llvm::SmallVectorImpl<Value> *lowerBounds = nullptr,
-    llvm::SmallVectorImpl<Value> *upperBounds = nullptr,
-    llvm::SmallVectorImpl<Value> *steps = nullptr) {
+static void
+extractHostEvalClauses(omp::TargetOp targetOp,
+                       llvm::SmallVectorImpl<Value> &numThreadsVars,
+                       Value &numTeamsLower, Value &numTeamsUpper,
+                       llvm::SmallVectorImpl<Value> &threadLimitVars,
+                       llvm::SmallVectorImpl<Value> *lowerBounds = nullptr,
+                       llvm::SmallVectorImpl<Value> *upperBounds = nullptr,
+                       llvm::SmallVectorImpl<Value> *steps = nullptr) {
   auto blockArgIface = llvm::cast<omp::BlockArgOpenMPOpInterface>(*targetOp);
   for (auto item : llvm::zip_equal(targetOp.getHostEvalVars(),
                                    blockArgIface.getHostEvalBlockArgs())) {
@@ -9145,10 +9147,18 @@ static void extractHostEvalClauses(
             else if (llvm::is_contained(teamsOp.getNumTeamsUpperVars(),
                                         blockArg))
               numTeamsUpper = hostEvalVar;
-            else if (!teamsOp.getThreadLimitVars().empty() &&
-                     teamsOp.getThreadLimit(0) == blockArg)
-              threadLimit = hostEvalVar;
-            else
+            else if (llvm::is_contained(teamsOp.getThreadLimitVars(),
+                                        blockArg)) {
+              for (auto [i, limitVar] :
+                   llvm::enumerate(teamsOp.getThreadLimitVars())) {
+                if (limitVar == blockArg) {
+                  if (threadLimitVars.size() <= i)
+                    threadLimitVars.resize(i + 1);
+                  threadLimitVars[i] = hostEvalVar;
+                  break;
+                }
+              }
+            } else
               llvm_unreachable("unsupported host_eval use");
           })
           .Case([&](omp::ParallelOp parallelOp) {
@@ -9271,11 +9281,11 @@ initTargetDefaultAttrs(omp::TargetOp targetOp, Operation *capturedOp,
                        bool isTargetDevice, bool isGPU) {
   // TODO: Handle constant 'if' clauses.
 
-  Value numTeamsLower, numTeamsUpper, threadLimit;
-  llvm::SmallVector<Value> numThreadsVars;
+  Value numTeamsLower, numTeamsUpper;
+  llvm::SmallVector<Value> numThreadsVars, threadLimitVars;
   if (!isTargetDevice) {
     extractHostEvalClauses(targetOp, numThreadsVars, numTeamsLower,
-                           numTeamsUpper, threadLimit);
+                           numTeamsUpper, threadLimitVars);
   } else {
     // In the target device, values for these clauses are not passed as
     // host_eval, but instead evaluated prior to entry to the region. This
@@ -9285,8 +9295,9 @@ initTargetDefaultAttrs(omp::TargetOp targetOp, Operation *capturedOp,
       // Handle num_teams upper bounds (only first value for now)
       if (!teamsOp.getNumTeamsUpperVars().empty())
         numTeamsUpper = teamsOp.getNumTeams(0);
-      if (!teamsOp.getThreadLimitVars().empty())
-        threadLimit = teamsOp.getThreadLimit(0);
+      threadLimitVars.reserve(teamsOp.getThreadLimitVars().size());
+      for (auto limitVar : teamsOp.getThreadLimitVars())
+        threadLimitVars.push_back(limitVar);
     }
 
     if (auto parallelOp = castOrGetParentOfType<omp::ParallelOp>(capturedOp)) {
@@ -9333,33 +9344,45 @@ initTargetDefaultAttrs(omp::TargetOp targetOp, Operation *capturedOp,
       result = 0;
   };
 
-  // Extract 'thread_limit' clause from 'target' and 'teams' directives.
-  int32_t targetThreadLimitVal = -1, teamsThreadLimitVal = -1;
-  if (!targetOp.getThreadLimitVars().empty())
-    setMaxValueFromClause(targetOp.getThreadLimit(0), targetThreadLimitVal);
-  setMaxValueFromClause(threadLimit, teamsThreadLimitVal);
+  // Extract multi-dimensional 'thread_limit' clause from 'target' and 'teams'.
+  llvm::SmallVector<int32_t, 3> targetThreadLimitVals(3, -1);
+  llvm::SmallVector<int32_t, 3> teamsThreadLimitVals(3, -1);
+  for (auto [i, limitVar] : llvm::enumerate(targetOp.getThreadLimitVars())) {
+    if (i < 3)
+      setMaxValueFromClause(limitVar, targetThreadLimitVals[i]);
+  }
+  for (auto [i, limitVar] : llvm::enumerate(threadLimitVars)) {
+    if (i < 3)
+      setMaxValueFromClause(limitVar, teamsThreadLimitVals[i]);
+  }
 
-  // Extract 'max_threads' clause from 'parallel' or set to 1 if it's SIMD.
-  int32_t maxThreadsVal = -1;
+  // Extract multi-dimensional 'num_threads' clause from 'parallel' or set to 1
+  // if it's SIMD.
+  llvm::SmallVector<int32_t, 3> maxThreadsVals(3, -1);
   if (castOrGetParentOfType<omp::ParallelOp>(capturedOp)) {
-    // For multi-dimensional num_threads, only use the first dimension for now
-    if (!numThreadsVars.empty())
-      setMaxValueFromClause(numThreadsVars[0], maxThreadsVal);
+    for (auto [i, threadsVar] : llvm::enumerate(numThreadsVars)) {
+      if (i < 3)
+        setMaxValueFromClause(threadsVar, maxThreadsVals[i]);
+    }
   } else if (castOrGetParentOfType<omp::SimdOp>(capturedOp,
-                                                /*immediateParent=*/true))
-    maxThreadsVal = 1;
+                                                /*immediateParent=*/true)) {
+    maxThreadsVals[0] = 1;
+  }
 
   // For max values, < 0 means unset, == 0 means set but unknown. Select the
-  // minimum value between 'max_threads' and 'thread_limit' clauses that were
-  // set.
-  int32_t combinedMaxThreadsVal = targetThreadLimitVal;
-  if (combinedMaxThreadsVal < 0 ||
-      (teamsThreadLimitVal >= 0 && teamsThreadLimitVal < combinedMaxThreadsVal))
-    combinedMaxThreadsVal = teamsThreadLimitVal;
-
-  if (combinedMaxThreadsVal < 0 ||
-      (maxThreadsVal >= 0 && maxThreadsVal < combinedMaxThreadsVal))
-    combinedMaxThreadsVal = maxThreadsVal;
+  // minimum value between 'num_threads' and 'thread_limit' clauses that were
+  // set, for each dimension.
+  llvm::SmallVector<int32_t, 3> combinedMaxThreadsVals(3, -1);
+  for (size_t i = 0; i < 3; ++i) {
+    int32_t combined = targetThreadLimitVals[i];
+    if (combined < 0 ||
+        (teamsThreadLimitVals[i] >= 0 && teamsThreadLimitVals[i] < combined))
+      combined = teamsThreadLimitVals[i];
+    if (combined < 0 ||
+        (maxThreadsVals[i] >= 0 && maxThreadsVals[i] < combined))
+      combined = maxThreadsVals[i];
+    combinedMaxThreadsVals[i] = combined;
+  }
 
   int32_t reductionDataSize = 0;
   if (isGPU && capturedOp) {
@@ -9405,18 +9428,22 @@ initTargetRuntimeAttrs(llvm::IRBuilderBase &builder,
   omp::LoopNestOp loopOp = castOrGetParentOfType<omp::LoopNestOp>(capturedOp);
   unsigned numLoops = loopOp ? loopOp.getNumLoops() : 0;
 
-  Value numTeamsLower, numTeamsUpper, teamsThreadLimit;
-  llvm::SmallVector<Value> numThreadsVars;
+  Value numTeamsLower, numTeamsUpper;
+  llvm::SmallVector<Value> numThreadsVars, threadLimitVars;
   llvm::SmallVector<Value> lowerBounds(numLoops), upperBounds(numLoops),
       steps(numLoops);
   extractHostEvalClauses(targetOp, numThreadsVars, numTeamsLower, numTeamsUpper,
-                         teamsThreadLimit, &lowerBounds, &upperBounds, &steps);
+                         threadLimitVars, &lowerBounds, &upperBounds, &steps);
 
   // TODO: Handle constant 'if' clauses.
+  // Resize to 3 dimensions to match TargetKernelDefaultAttrs
+  attrs.TargetThreadLimit.resize(3);
   if (!targetOp.getThreadLimitVars().empty()) {
-    Value targetThreadLimit = targetOp.getThreadLimit(0);
-    attrs.TargetThreadLimit.front() =
-        moduleTranslation.lookupValue(targetThreadLimit);
+    for (auto [i, limitVar] : llvm::enumerate(targetOp.getThreadLimitVars())) {
+      if (limitVar) {
+        attrs.TargetThreadLimit[i] = moduleTranslation.lookupValue(limitVar);
+      }
+    }
   }
 
   // The __kmpc_push_num_teams_51 function expects int32 as the arguments.  So,
@@ -9430,9 +9457,15 @@ initTargetRuntimeAttrs(llvm::IRBuilderBase &builder,
     attrs.MaxTeams.front() = builder.CreateSExtOrTrunc(
         moduleTranslation.lookupValue(numTeamsUpper), builder.getInt32Ty());
 
-  if (teamsThreadLimit)
-    attrs.TeamsThreadLimit.front() = builder.CreateSExtOrTrunc(
-        moduleTranslation.lookupValue(teamsThreadLimit), builder.getInt32Ty());
+  attrs.TeamsThreadLimit.resize(3);
+  if (!threadLimitVars.empty()) {
+    for (auto [i, limitVar] : llvm::enumerate(threadLimitVars)) {
+      if (limitVar) {
+        attrs.TeamsThreadLimit[i] = builder.CreateSExtOrTrunc(
+            moduleTranslation.lookupValue(limitVar), builder.getInt32Ty());
+      }
+    }
+  }
 
   // One runtime value per num_threads dimension.
   if (!numThreadsVars.empty()) {
diff --git a/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir b/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir
index 9521dbf9d8e5a0..6d8a7eacc35fb3 100644
--- a/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir
@@ -2,13 +2,13 @@
 
 // CHECK: define void @main(i32 %[[NUM_TEAMS_ARG:.*]])
 // CHECK: %[[KERNEL_ARGS:.*]] = alloca %struct.__tgt_kernel_arguments
-// CHECK: %[[NUM_TEAMS:.*]] = insertvalue [3 x i32] zeroinitializer, i32 %[[NUM_TEAMS_ARG]], 0
+// CHECK: insertvalue [3 x i32] zeroinitializer, i32 %[[NUM_TEAMS_ARG]], 0
 
 // CHECK: %[[NUM_TEAMS_KARG:.*]] = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %[[KERNEL_ARGS]], i32 0, i32 10
-// CHECK: store [3 x i32] %[[NUM_TEAMS]], ptr %[[NUM_TEAMS_KARG]], align 4
+// CHECK-NEXT: store [3 x i32] %{{.*}}, ptr %[[NUM_TEAMS_KARG]], align 4
 
 // CHECK: %[[NUM_THREADS_ARG:.*]] = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %[[KERNEL_ARGS]], i32 0, i32 11
-// CHECK: store [3 x i32] [i32 10, i32 0, i32 0], ptr %[[NUM_THREADS_ARG]], align 4
+// CHECK-NEXT: store [3 x i32] [i32 10, i32 0, i32 0], ptr %[[NUM_THREADS_ARG]], align 4
 
 // CHECK: %{{.*}} = call i32 @__tgt_target_kernel(ptr {{.*}}, i64 -1, i32 %[[NUM_TEAMS_ARG]], i32 [[NUM_THREADS:10]], ptr @.[[OUTLINED_FN:.*]].region_id, ptr %[[KERNEL_ARGS]])
 // CHECK: call void @[[OUTLINED_FN]](i32 %[[NUM_TEAMS_ARG]], ptr null)
diff --git a/mlir/test/Target/LLVMIR/openmp-teams.mlir b/mlir/test/Target/LLVMIR/openmp-teams.mlir
index a7bc92f8b99b9e..adca15d1c5fcc0 100644
--- a/mlir/test/Target/LLVMIR/openmp-teams.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-teams.mlir
@@ -314,26 +314,36 @@ llvm.func @teams_if_with_num_teams(%condition: i1, %numTeamsLower: i32, %numTeam
 
 // -----
 
-// Check that the thread and bound id arguments of the outlined function are
-// generic pointers even when allocas are created in a non-zero address space.
-// The cast that produces them is deleted along with the fake allocas once the
-// region has been outlined, so the signature below is the only place it can be
-// observed.
-module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui32>>} {
-    llvm.func @foo()
-
-    // CHECK-LABEL: @omp_teams_alloca_addrspace
-    // CHECK: call void {{.*}} @__kmpc_fork_teams(ptr @{{.+}}, i32 0, ptr @[[OUTLINED_FN:.+]])
-    llvm.func @omp_teams_alloca_addrspace() {
-        omp.teams {
-            llvm.call @foo() : () -> ()
-            omp.terminator
-        }
-        llvm.return
+llvm.func @duringTeams()
+
+// CHECK-LABEL: @omp_teams_thread_limit_2d
+// CHECK-SAME: (i32 [[LIMIT_X:.+]], i32 [[LIMIT_Y:.+]])
+llvm.func @omp_teams_thread_limit_2d(%limitX: i32, %limitY: i32) {
+    // Multi-dimensional thread_limit: all dimensions are passed
+    // CHECK: [[THREAD_NUM:%.+]] = call i32 @__kmpc_global_thread_num
+    // CHECK-NEXT: call void @__kmpc_push_num_teams_51({{.+}}, i32 [[THREAD_NUM]], i32 0, i32 0, i32 [[LIMIT_X]])
+    // CHECK: call void (ptr, i32, ptr, ...) @__kmpc_fork_teams(ptr @{{[0-9]+}}, i32 0, ptr [[OUTLINED_FN:.+]])
+    omp.teams thread_limit(%limitX, %limitY : i32, i32) {
+        llvm.call @duringTeams() : () -> ()
+        omp.terminator
     }
+    llvm.return
+}
+
+// -----
+
+llvm.func @duringTeams()
 
-    // CHECK:      define internal void @[[OUTLINED_FN]]
-    // CHECK-SAME: (ptr %global.tid.ptr, ptr %bound.tid.ptr)
-    // CHECK:   call void @foo()
-    // CHECK:   ret void
+// CHECK-LABEL: @omp_teams_thread_limit_3d
+// CHECK-SAME: (i32 [[LIMIT_X:.+]], i64 [[LIMIT_Y:.+]], i16 [[LIMIT_Z:.+]])
+llvm.func @omp_teams_thread_limit_3d(%limitX: i32, %limitY: i64, %limitZ: i16) {
+    // Multi-dimensional thread_limit with mixed types: all dimensions are passed
+    // CHECK: [[THREAD_NUM:%.+]] = call i32 @__kmpc_global_thread_num
+    // CHECK-NEXT: call void @__kmpc_push_num_teams_51({{.+}}, i32 [[THREAD_NUM]], i32 0, i32 0, i32 [[LIMIT_X]])
+    // CHECK: call void (ptr, i32, ptr, ...) @__kmpc_fork_teams(ptr @{{[0-9]+}}, i32 0, ptr [[OUTLINED_FN:.+]])
+    omp.teams thread_limit(%limitX, %limitY, %limitZ : i32, i64, i16) {
+        llvm.call @duringTeams() : () -> ()
+        omp.terminator
+    }
+    llvm.return
 }
diff --git a/mlir/test/Target/LLVMIR/openmp-todo.mlir b/mlir/test/Target/LLVMIR/openmp-todo.mlir
index a29e9f836e4bf4..7138d8515fa514 100644
--- a/mlir/test/Target/LLVMIR/openmp-todo.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-todo.mlir
@@ -692,10 +692,10 @@ llvm.func @parallel_num_threads_too_many_dims(%lb : i32, %ub : i32) {
 
 // -----
 
-llvm.func @teams_thread_limit_multi_dim(%lb : i32, %ub : i32) {
-  // expected-error at below {{not yet implemented: Unhandled clause thread_limit with multi-dimensional values in omp.teams operation}}
+llvm.func @teams_thread_limit_too_many_dims(%lb : i32, %ub : i32) {
+  // expected-error at below {{not yet implemented: Unhandled clause thread_limit with more than 3 dimensions in omp.teams operation}}
   // expected-error at below {{LLVM Translation failed for operation: omp.teams}}
-  omp.teams thread_limit(%lb, %ub : i32, i32) {
+  omp.teams thread_limit(%lb, %ub, %lb, %ub : i32, i32, i32, i32) {
     omp.terminator
   }
   llvm.return

>From 48228c15853972b8246f0a375354f9ccfbd13490 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Thu, 9 Apr 2026 16:23:40 +0530
Subject: [PATCH 3/6] update

---
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      | 91 +++++++++++--------
 .../LLVMIR/openmp-target-launch-device.mlir   | 23 ++++-
 .../LLVMIR/openmp-target-launch-host.mlir     | 27 +++++-
 3 files changed, 98 insertions(+), 43 deletions(-)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 274fa7b5578802..7579afcb5e2d57 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -9344,26 +9344,27 @@ initTargetDefaultAttrs(omp::TargetOp targetOp, Operation *capturedOp,
       result = 0;
   };
 
-  // Extract multi-dimensional 'thread_limit' clause from 'target' and 'teams'.
-  llvm::SmallVector<int32_t, 3> targetThreadLimitVals(3, -1);
-  llvm::SmallVector<int32_t, 3> teamsThreadLimitVals(3, -1);
-  for (auto [i, limitVar] : llvm::enumerate(targetOp.getThreadLimitVars())) {
-    if (i < 3)
-      setMaxValueFromClause(limitVar, targetThreadLimitVals[i]);
-  }
-  for (auto [i, limitVar] : llvm::enumerate(threadLimitVars)) {
-    if (i < 3)
-      setMaxValueFromClause(limitVar, teamsThreadLimitVals[i]);
-  }
-
-  // Extract multi-dimensional 'num_threads' clause from 'parallel' or set to 1
-  // if it's SIMD.
-  llvm::SmallVector<int32_t, 3> maxThreadsVals(3, -1);
+  // Extract 'thread_limit' clause from 'target' and 'teams'. The number of
+  // dimensions is determined by the clauses present (the >3 dims check in
+  // checkImplementationStatus guards against unsupported counts).
+  size_t numTargetDims = targetOp.getThreadLimitVars().size();
+  size_t numTeamsDims = threadLimitVars.size();
+  size_t numParallelDims = numThreadsVars.size();
+  size_t numDims =
+      std::max({numTargetDims, numTeamsDims, numParallelDims, size_t(1)});
+
+  llvm::SmallVector<int32_t, 3> targetThreadLimitVals(numDims, -1);
+  llvm::SmallVector<int32_t, 3> teamsThreadLimitVals(numDims, -1);
+  for (auto [i, limitVar] : llvm::enumerate(targetOp.getThreadLimitVars()))
+    setMaxValueFromClause(limitVar, targetThreadLimitVals[i]);
+  for (auto [i, limitVar] : llvm::enumerate(threadLimitVars))
+    setMaxValueFromClause(limitVar, teamsThreadLimitVals[i]);
+
+  // Extract 'num_threads' clause from 'parallel' or set to 1 if it's SIMD.
+  llvm::SmallVector<int32_t, 3> maxThreadsVals(numDims, -1);
   if (castOrGetParentOfType<omp::ParallelOp>(capturedOp)) {
-    for (auto [i, threadsVar] : llvm::enumerate(numThreadsVars)) {
-      if (i < 3)
-        setMaxValueFromClause(threadsVar, maxThreadsVals[i]);
-    }
+    for (auto [i, threadsVar] : llvm::enumerate(numThreadsVars))
+      setMaxValueFromClause(threadsVar, maxThreadsVals[i]);
   } else if (castOrGetParentOfType<omp::SimdOp>(capturedOp,
                                                 /*immediateParent=*/true)) {
     maxThreadsVals[0] = 1;
@@ -9372,8 +9373,8 @@ initTargetDefaultAttrs(omp::TargetOp targetOp, Operation *capturedOp,
   // For max values, < 0 means unset, == 0 means set but unknown. Select the
   // minimum value between 'num_threads' and 'thread_limit' clauses that were
   // set, for each dimension.
-  llvm::SmallVector<int32_t, 3> combinedMaxThreadsVals(3, -1);
-  for (size_t i = 0; i < 3; ++i) {
+  llvm::SmallVector<int32_t, 3> combinedMaxThreadsVals(numDims, -1);
+  for (size_t i = 0; i < numDims; ++i) {
     int32_t combined = targetThreadLimitVals[i];
     if (combined < 0 ||
         (teamsThreadLimitVals[i] >= 0 && teamsThreadLimitVals[i] < combined))
@@ -9429,21 +9430,22 @@ initTargetRuntimeAttrs(llvm::IRBuilderBase &builder,
   unsigned numLoops = loopOp ? loopOp.getNumLoops() : 0;
 
   Value numTeamsLower, numTeamsUpper;
-  llvm::SmallVector<Value> numThreadsVars, threadLimitVars;
+  llvm::SmallVector<Value> numThreadsVars, teamsThreadLimitVars;
   llvm::SmallVector<Value> lowerBounds(numLoops), upperBounds(numLoops),
       steps(numLoops);
   extractHostEvalClauses(targetOp, numThreadsVars, numTeamsLower, numTeamsUpper,
-                         threadLimitVars, &lowerBounds, &upperBounds, &steps);
+                         teamsThreadLimitVars, &lowerBounds, &upperBounds,
+                         &steps);
 
   // TODO: Handle constant 'if' clauses.
-  // Resize to 3 dimensions to match TargetKernelDefaultAttrs
-  attrs.TargetThreadLimit.resize(3);
   if (!targetOp.getThreadLimitVars().empty()) {
-    for (auto [i, limitVar] : llvm::enumerate(targetOp.getThreadLimitVars())) {
-      if (limitVar) {
-        attrs.TargetThreadLimit[i] = moduleTranslation.lookupValue(limitVar);
-      }
-    }
+    attrs.TargetThreadLimit.clear();
+    llvm::transform(targetOp.getThreadLimitVars(),
+                    std::back_inserter(attrs.TargetThreadLimit),
+                    [&](Value limitVar) -> llvm::Value * {
+                      return limitVar ? moduleTranslation.lookupValue(limitVar)
+                                      : nullptr;
+                    });
   }
 
   // The __kmpc_push_num_teams_51 function expects int32 as the arguments.  So,
@@ -9457,14 +9459,17 @@ initTargetRuntimeAttrs(llvm::IRBuilderBase &builder,
     attrs.MaxTeams.front() = builder.CreateSExtOrTrunc(
         moduleTranslation.lookupValue(numTeamsUpper), builder.getInt32Ty());
 
-  attrs.TeamsThreadLimit.resize(3);
-  if (!threadLimitVars.empty()) {
-    for (auto [i, limitVar] : llvm::enumerate(threadLimitVars)) {
-      if (limitVar) {
-        attrs.TeamsThreadLimit[i] = builder.CreateSExtOrTrunc(
-            moduleTranslation.lookupValue(limitVar), builder.getInt32Ty());
-      }
-    }
+  if (!teamsThreadLimitVars.empty()) {
+    attrs.TeamsThreadLimit.clear();
+    llvm::transform(teamsThreadLimitVars,
+                    std::back_inserter(attrs.TeamsThreadLimit),
+                    [&](Value limitVar) -> llvm::Value * {
+                      return limitVar
+                                 ? builder.CreateSExtOrTrunc(
+                                       moduleTranslation.lookupValue(limitVar),
+                                       builder.getInt32Ty())
+                                 : nullptr;
+                    });
   }
 
   // One runtime value per num_threads dimension.
@@ -9475,6 +9480,16 @@ initTargetRuntimeAttrs(llvm::IRBuilderBase &builder,
           numThreadsVar ? moduleTranslation.lookupValue(numThreadsVar)
                         : nullptr);
   }
+  // Ensure TargetThreadLimit and TeamsThreadLimit have matching sizes
+  // for zip_equal in OMPIRBuilder.
+  size_t maxDims =
+      std::max(attrs.TargetThreadLimit.size(), attrs.TeamsThreadLimit.size());
+  attrs.TargetThreadLimit.resize(maxDims);
+  attrs.TeamsThreadLimit.resize(maxDims);
+
+  // Handle multi-dimensional num_threads (only first value for now)
+  if (!numThreadsVars.empty())
+    attrs.MaxThreads = moduleTranslation.lookupValue(numThreadsVars[0]);
 
   if (targetOp.hasHostEvalTripCount()) {
     llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
diff --git a/mlir/test/Target/LLVMIR/openmp-target-launch-device.mlir b/mlir/test/Target/LLVMIR/openmp-target-launch-device.mlir
index e404d2bb02bcb7..5ff2b3b182ff52 100644
--- a/mlir/test/Target/LLVMIR/openmp-target-launch-device.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-target-launch-device.mlir
@@ -14,6 +14,13 @@
 // CHECK-SAME: %struct.ConfigurationEnvironmentTy { i8 1, i8 1, i8 [[EXEC_MODE2:1]], i32 [[MIN_THREADS2:1]], i32 [[MAX_THREADS2:94]], i32 [[MIN_TEAMS2:40]], i32 [[MAX_TEAMS2:40]], i32 0 },
 // CHECK-SAME: ptr @{{.*}}, ptr @{{.*}} }
 
+// Multi-dim thread_limit: min(target=20, teams_x=10) = 10.
+// CHECK:      @[[EXEC_MODE3:.*]] = weak protected constant i8 1
+// CHECK:      @llvm.compiler.used{{.*}} = appending global [1 x ptr] [ptr @[[EXEC_MODE3]]], section "llvm.metadata"
+// CHECK:      @[[KERNEL3_ENV:.*_kernel_environment]] = weak_odr protected constant %struct.KernelEnvironmentTy {
+// CHECK-SAME: %struct.ConfigurationEnvironmentTy { i8 1, i8 1, i8 [[EXEC_MODE3:1]], i32 [[MIN_THREADS3:1]], i32 [[MAX_THREADS3:10]], i32 0, i32 0, i32 0, i32 0 },
+// CHECK-SAME: ptr @{{.*}}, ptr @{{.*}} }
+
 module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui32>>, llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_target_device = true, omp.is_gpu = true} {
   llvm.func @main(%num_teams : !llvm.ptr) {
     // CHECK: define weak_odr protected amdgpu_kernel void @__omp_offloading_{{.*}}_main_l{{[0-9]+}}(ptr %[[NUM_TEAMS_ARG:.*]], ptr %[[KERNEL_ARGS:.*]]) #[[ATTRS1:[0-9]+]]
@@ -38,7 +45,21 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memo
         omp.terminator
       }
       omp.terminator
-    } {omp.combined}
+    }
+
+    // Multi-dim thread_limit: first dim constant, second dim constant.
+    // MaxThreads uses the first dim combined value: min(20, 10) = 10.
+    // CHECK: define weak_odr protected amdgpu_kernel void @__omp_offloading_{{.*}}_main_l{{[0-9]+}}(ptr %[[KERNEL_ARGS:.*]]) #[[ATTRS1]]
+    // CHECK: %{{.*}} = call i32 @__kmpc_target_init(ptr @[[KERNEL3_ENV]], ptr %[[KERNEL_ARGS]])
+    %target_threads3 = llvm.mlir.constant(20) : i32
+    omp.target thread_limit(%target_threads3 : i32) {
+      %teams_threads_x = llvm.mlir.constant(10) : i32
+      %teams_threads_y = llvm.mlir.constant(5) : i32
+      omp.teams thread_limit(%teams_threads_x, %teams_threads_y : i32, i32) {
+        omp.terminator
+      }
+      omp.terminator
+    }
     llvm.return
   }
 }
diff --git a/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir b/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir
index 6d8a7eacc35fb3..ac7964ee65bf1e 100644
--- a/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir
@@ -2,18 +2,26 @@
 
 // CHECK: define void @main(i32 %[[NUM_TEAMS_ARG:.*]])
 // CHECK: %[[KERNEL_ARGS:.*]] = alloca %struct.__tgt_kernel_arguments
-// CHECK: insertvalue [3 x i32] zeroinitializer, i32 %[[NUM_TEAMS_ARG]], 0
+// CHECK: %[[NUM_TEAMS:.*]] = insertvalue [3 x i32] zeroinitializer, i32 %[[NUM_TEAMS_ARG]], 0
 
 // CHECK: %[[NUM_TEAMS_KARG:.*]] = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %[[KERNEL_ARGS]], i32 0, i32 10
-// CHECK-NEXT: store [3 x i32] %{{.*}}, ptr %[[NUM_TEAMS_KARG]], align 4
+// CHECK: store [3 x i32] %[[NUM_TEAMS]], ptr %[[NUM_TEAMS_KARG]], align 4
 
 // CHECK: %[[NUM_THREADS_ARG:.*]] = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %[[KERNEL_ARGS]], i32 0, i32 11
-// CHECK-NEXT: store [3 x i32] [i32 10, i32 0, i32 0], ptr %[[NUM_THREADS_ARG]], align 4
+// CHECK: store [3 x i32] [i32 10, i32 0, i32 0], ptr %[[NUM_THREADS_ARG]], align 4
 
 // CHECK: %{{.*}} = call i32 @__tgt_target_kernel(ptr {{.*}}, i64 -1, i32 %[[NUM_TEAMS_ARG]], i32 [[NUM_THREADS:10]], ptr @.[[OUTLINED_FN:.*]].region_id, ptr %[[KERNEL_ARGS]])
 // CHECK: call void @[[OUTLINED_FN]](i32 %[[NUM_TEAMS_ARG]], ptr null)
 
-// CHECK: define internal void @[[OUTLINED_FN]](i32 %[[NUM_TEAMS_OUTLINED:.*]], ptr %{{.*}})
+// Multi-dim thread_limit: first dim is constant (10), second dim is runtime variable.
+// The NumThreads [3 x i32] array should have dim0=10, dim1=%thread_limit_y, dim2=0.
+// CHECK: define void @main_multidim_thread_limit(i32 %[[TL_Y:.*]])
+// CHECK: %[[KERNEL_ARGS2:.*]] = alloca %struct.__tgt_kernel_arguments
+// CHECK: %[[NT_ARR:.*]] = insertvalue [3 x i32] [i32 10, i32 0, i32 0], i32 %[[TL_Y]], 1
+// CHECK: %[[NT_GEP:.*]] = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %[[KERNEL_ARGS2]], i32 0, i32 11
+// CHECK-NEXT: store [3 x i32] %[[NT_ARR]], ptr %[[NT_GEP]], align 4
+// CHECK: call i32 @__tgt_target_kernel(ptr {{.*}}, i64 -1, i32 0, i32 10, ptr @.{{.*}}.region_id, ptr %[[KERNEL_ARGS2]])
+// CHECK: define internal void @[[OUTLINED_FN]](i32 %[[NUM_TEAMS_OUTLINED:.*]])
 // CHECK: call void @__kmpc_push_num_teams_51(ptr {{.*}}, i32 {{.*}}, i32 %[[NUM_TEAMS_OUTLINED]], i32 %[[NUM_TEAMS_OUTLINED]], i32 [[NUM_THREADS]])
 module attributes {omp.is_target_device = false, omp.target_triples = ["amdgcn-amd-amdhsa"]} {
   llvm.func @main(%num_teams : i32) {
@@ -28,4 +36,15 @@ module attributes {omp.is_target_device = false, omp.target_triples = ["amdgcn-a
     }
     llvm.return
   }
+
+  llvm.func @main_multidim_thread_limit(%thread_limit_y : i32) {
+    %teams_threads_x = llvm.mlir.constant(10) : i32
+    omp.target host_eval(%teams_threads_x -> %arg_tlx, %thread_limit_y -> %arg_tly : i32, i32) {
+      omp.teams thread_limit(%arg_tlx, %arg_tly : i32, i32) {
+        omp.terminator
+      }
+      omp.terminator
+    }
+    llvm.return
+  }
 }

>From 5584037163424ef6587e129e99d04356fe9a4289 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Thu, 16 Apr 2026 15:57:40 +0530
Subject: [PATCH 4/6] update 2

---
 .../Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp    | 13 +++++++------
 .../Target/LLVMIR/openmp-target-launch-device.mlir  |  5 ++---
 .../Target/LLVMIR/openmp-target-launch-host.mlir    |  2 +-
 mlir/test/Target/LLVMIR/openmp-teams.mlir           |  2 --
 4 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 7579afcb5e2d57..c57ab47cfc8982 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -9142,13 +9142,13 @@ extractHostEvalClauses(omp::TargetOp targetOp,
     for (Operation *user : blockArg.getUsers()) {
       llvm::TypeSwitch<Operation *>(user)
           .Case([&](omp::TeamsOp teamsOp) {
-            if (teamsOp.getNumTeamsLower() == blockArg)
+            if (teamsOp.getNumTeamsLower() == blockArg) {
               numTeamsLower = hostEvalVar;
-            else if (llvm::is_contained(teamsOp.getNumTeamsUpperVars(),
-                                        blockArg))
+            } else if (llvm::is_contained(teamsOp.getNumTeamsUpperVars(),
+                                          blockArg)) {
               numTeamsUpper = hostEvalVar;
-            else if (llvm::is_contained(teamsOp.getThreadLimitVars(),
-                                        blockArg)) {
+            } else if (llvm::is_contained(teamsOp.getThreadLimitVars(),
+                                          blockArg)) {
               for (auto [i, limitVar] :
                    llvm::enumerate(teamsOp.getThreadLimitVars())) {
                 if (limitVar == blockArg) {
@@ -9158,8 +9158,9 @@ extractHostEvalClauses(omp::TargetOp targetOp,
                   break;
                 }
               }
-            } else
+            } else {
               llvm_unreachable("unsupported host_eval use");
+            }
           })
           .Case([&](omp::ParallelOp parallelOp) {
             if (llvm::is_contained(parallelOp.getNumThreadsVars(), blockArg)) {
diff --git a/mlir/test/Target/LLVMIR/openmp-target-launch-device.mlir b/mlir/test/Target/LLVMIR/openmp-target-launch-device.mlir
index 5ff2b3b182ff52..b74eb63dfa905b 100644
--- a/mlir/test/Target/LLVMIR/openmp-target-launch-device.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-target-launch-device.mlir
@@ -14,7 +14,8 @@
 // CHECK-SAME: %struct.ConfigurationEnvironmentTy { i8 1, i8 1, i8 [[EXEC_MODE2:1]], i32 [[MIN_THREADS2:1]], i32 [[MAX_THREADS2:94]], i32 [[MIN_TEAMS2:40]], i32 [[MAX_TEAMS2:40]], i32 0 },
 // CHECK-SAME: ptr @{{.*}}, ptr @{{.*}} }
 
-// Multi-dim thread_limit: min(target=20, teams_x=10) = 10.
+// Multi-dim thread_limit: first dim constant (10), second dim constant (5).
+// MaxThreads uses the first dim combined value: min(target=20, teams_x=10) = 10.
 // CHECK:      @[[EXEC_MODE3:.*]] = weak protected constant i8 1
 // CHECK:      @llvm.compiler.used{{.*}} = appending global [1 x ptr] [ptr @[[EXEC_MODE3]]], section "llvm.metadata"
 // CHECK:      @[[KERNEL3_ENV:.*_kernel_environment]] = weak_odr protected constant %struct.KernelEnvironmentTy {
@@ -47,8 +48,6 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memo
       omp.terminator
     }
 
-    // Multi-dim thread_limit: first dim constant, second dim constant.
-    // MaxThreads uses the first dim combined value: min(20, 10) = 10.
     // CHECK: define weak_odr protected amdgpu_kernel void @__omp_offloading_{{.*}}_main_l{{[0-9]+}}(ptr %[[KERNEL_ARGS:.*]]) #[[ATTRS1]]
     // CHECK: %{{.*}} = call i32 @__kmpc_target_init(ptr @[[KERNEL3_ENV]], ptr %[[KERNEL_ARGS]])
     %target_threads3 = llvm.mlir.constant(20) : i32
diff --git a/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir b/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir
index ac7964ee65bf1e..c615453a2870d8 100644
--- a/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir
@@ -21,7 +21,7 @@
 // CHECK: %[[NT_GEP:.*]] = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %[[KERNEL_ARGS2]], i32 0, i32 11
 // CHECK-NEXT: store [3 x i32] %[[NT_ARR]], ptr %[[NT_GEP]], align 4
 // CHECK: call i32 @__tgt_target_kernel(ptr {{.*}}, i64 -1, i32 0, i32 10, ptr @.{{.*}}.region_id, ptr %[[KERNEL_ARGS2]])
-// CHECK: define internal void @[[OUTLINED_FN]](i32 %[[NUM_TEAMS_OUTLINED:.*]])
+// CHECK: define internal void @[[OUTLINED_FN]](i32 %[[NUM_TEAMS_OUTLINED:.*]], ptr %{{.*}})
 // CHECK: call void @__kmpc_push_num_teams_51(ptr {{.*}}, i32 {{.*}}, i32 %[[NUM_TEAMS_OUTLINED]], i32 %[[NUM_TEAMS_OUTLINED]], i32 [[NUM_THREADS]])
 module attributes {omp.is_target_device = false, omp.target_triples = ["amdgcn-amd-amdhsa"]} {
   llvm.func @main(%num_teams : i32) {
diff --git a/mlir/test/Target/LLVMIR/openmp-teams.mlir b/mlir/test/Target/LLVMIR/openmp-teams.mlir
index adca15d1c5fcc0..126d3e652a6e15 100644
--- a/mlir/test/Target/LLVMIR/openmp-teams.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-teams.mlir
@@ -319,7 +319,6 @@ llvm.func @duringTeams()
 // CHECK-LABEL: @omp_teams_thread_limit_2d
 // CHECK-SAME: (i32 [[LIMIT_X:.+]], i32 [[LIMIT_Y:.+]])
 llvm.func @omp_teams_thread_limit_2d(%limitX: i32, %limitY: i32) {
-    // Multi-dimensional thread_limit: all dimensions are passed
     // CHECK: [[THREAD_NUM:%.+]] = call i32 @__kmpc_global_thread_num
     // CHECK-NEXT: call void @__kmpc_push_num_teams_51({{.+}}, i32 [[THREAD_NUM]], i32 0, i32 0, i32 [[LIMIT_X]])
     // CHECK: call void (ptr, i32, ptr, ...) @__kmpc_fork_teams(ptr @{{[0-9]+}}, i32 0, ptr [[OUTLINED_FN:.+]])
@@ -337,7 +336,6 @@ llvm.func @duringTeams()
 // CHECK-LABEL: @omp_teams_thread_limit_3d
 // CHECK-SAME: (i32 [[LIMIT_X:.+]], i64 [[LIMIT_Y:.+]], i16 [[LIMIT_Z:.+]])
 llvm.func @omp_teams_thread_limit_3d(%limitX: i32, %limitY: i64, %limitZ: i16) {
-    // Multi-dimensional thread_limit with mixed types: all dimensions are passed
     // CHECK: [[THREAD_NUM:%.+]] = call i32 @__kmpc_global_thread_num
     // CHECK-NEXT: call void @__kmpc_push_num_teams_51({{.+}}, i32 [[THREAD_NUM]], i32 0, i32 0, i32 [[LIMIT_X]])
     // CHECK: call void (ptr, i32, ptr, ...) @__kmpc_fork_teams(ptr @{{[0-9]+}}, i32 0, ptr [[OUTLINED_FN:.+]])

>From 6b2f4a22ef3831ed787eb126a819fdfc7d81afe8 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Mon, 31 Aug 2026 19:23:54 +0530
Subject: [PATCH 5/6] Update logic with initSentinel and padTrailingDims

---
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      | 58 +++++++++++++------
 .../LLVMIR/openmp-target-launch-device.mlir   | 27 +++++++--
 .../LLVMIR/openmp-target-launch-host.mlir     | 27 ++++++++-
 mlir/test/Target/LLVMIR/openmp-todo.mlir      | 11 ----
 4 files changed, 87 insertions(+), 36 deletions(-)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index c57ab47cfc8982..6e91551b9feb09 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -9173,7 +9173,7 @@ extractHostEvalClauses(omp::TargetOp targetOp,
                   break;
                 }
               }
-            } else
+            } else {
               llvm_unreachable("unsupported host_eval use");
             }
           })
@@ -9336,13 +9336,16 @@ initTargetDefaultAttrs(omp::TargetOp targetOp, Operation *capturedOp,
     if (!clauseValue)
       return;
 
-    if (auto val = extractConstInteger(clauseValue))
+    if (auto val = extractConstInteger(clauseValue)) {
       result = *val;
+      if (result < 0)
+        result = 0;
+      return;
+    }
 
-    // Found an applicable clause, so it's not undefined. Mark as unknown
-    // because it's not constant.
-    if (result < 0)
-      result = 0;
+    // Applicable but non-constant, so "set but unknown". Must also override a
+    // sentinel of 1, or a runtime bound would clamp the launch to one thread.
+    result = 0;
   };
 
   // Extract 'thread_limit' clause from 'target' and 'teams'. The number of
@@ -9354,16 +9357,26 @@ initTargetDefaultAttrs(omp::TargetOp targetOp, Operation *capturedOp,
   size_t numDims =
       std::max({numTargetDims, numTeamsDims, numParallelDims, size_t(1)});
 
-  llvm::SmallVector<int32_t, 3> targetThreadLimitVals(numDims, -1);
-  llvm::SmallVector<int32_t, 3> teamsThreadLimitVals(numDims, -1);
+  // A clause specified with fewer values than the kernel rank leaves its
+  // trailing dims implicitly 1 (not "unset"), so they still clamp the per-dim
+  // min below. Sentinel: 1 if the clause is present, -1 if entirely absent.
+  auto initSentinel = [](bool isSpecified) -> int32_t {
+    return isSpecified ? 1 : -1;
+  };
+  llvm::SmallVector<int32_t, 3> targetThreadLimitVals(
+      numDims, initSentinel(!targetOp.getThreadLimitVars().empty()));
+  llvm::SmallVector<int32_t, 3> teamsThreadLimitVals(
+      numDims, initSentinel(!threadLimitVars.empty()));
   for (auto [i, limitVar] : llvm::enumerate(targetOp.getThreadLimitVars()))
     setMaxValueFromClause(limitVar, targetThreadLimitVals[i]);
   for (auto [i, limitVar] : llvm::enumerate(threadLimitVars))
     setMaxValueFromClause(limitVar, teamsThreadLimitVals[i]);
 
   // Extract 'num_threads' clause from 'parallel' or set to 1 if it's SIMD.
-  llvm::SmallVector<int32_t, 3> maxThreadsVals(numDims, -1);
-  if (castOrGetParentOfType<omp::ParallelOp>(capturedOp)) {
+  auto parallelOp = castOrGetParentOfType<omp::ParallelOp>(capturedOp);
+  llvm::SmallVector<int32_t, 3> maxThreadsVals(
+      numDims, initSentinel(parallelOp && !numThreadsVars.empty()));
+  if (parallelOp) {
     for (auto [i, threadsVar] : llvm::enumerate(numThreadsVars))
       setMaxValueFromClause(threadsVar, maxThreadsVals[i]);
   } else if (castOrGetParentOfType<omp::SimdOp>(capturedOp,
@@ -9412,7 +9425,7 @@ initTargetDefaultAttrs(omp::TargetOp targetOp, Operation *capturedOp,
   attrs.MinTeams.front() = minTeamsVal;
   attrs.MaxTeams.front() = maxTeamsVal;
   attrs.MinThreads.front() = 1;
-  attrs.MaxThreads.front() = combinedMaxThreadsVal;
+  attrs.MaxThreads = combinedMaxThreadsVals;
   attrs.ReductionDataSize = reductionDataSize;
 }
 
@@ -9481,16 +9494,23 @@ initTargetRuntimeAttrs(llvm::IRBuilderBase &builder,
           numThreadsVar ? moduleTranslation.lookupValue(numThreadsVar)
                         : nullptr);
   }
-  // Ensure TargetThreadLimit and TeamsThreadLimit have matching sizes
-  // for zip_equal in OMPIRBuilder.
+  // Sizes must match for zip_equal in OMPIRBuilder. Trailing dims of a
+  // specified-but-shorter clause are implicitly 1 and still clamp, whereas an
+  // absent clause stays null so it imposes no constraint.
   size_t maxDims =
       std::max(attrs.TargetThreadLimit.size(), attrs.TeamsThreadLimit.size());
-  attrs.TargetThreadLimit.resize(maxDims);
-  attrs.TeamsThreadLimit.resize(maxDims);
-
-  // Handle multi-dimensional num_threads (only first value for now)
-  if (!numThreadsVars.empty())
-    attrs.MaxThreads = moduleTranslation.lookupValue(numThreadsVars[0]);
+  auto padTrailingDims = [&](llvm::SmallVectorImpl<llvm::Value *> &vec,
+                             bool isSpecified) {
+    if (vec.size() >= maxDims)
+      return;
+    if (isSpecified)
+      vec.append(maxDims - vec.size(), builder.getInt32(1));
+    else
+      vec.resize(maxDims);
+  };
+  padTrailingDims(attrs.TargetThreadLimit,
+                  !targetOp.getThreadLimitVars().empty());
+  padTrailingDims(attrs.TeamsThreadLimit, !teamsThreadLimitVars.empty());
 
   if (targetOp.hasHostEvalTripCount()) {
     llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
diff --git a/mlir/test/Target/LLVMIR/openmp-target-launch-device.mlir b/mlir/test/Target/LLVMIR/openmp-target-launch-device.mlir
index b74eb63dfa905b..2a7ec81a01259e 100644
--- a/mlir/test/Target/LLVMIR/openmp-target-launch-device.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-target-launch-device.mlir
@@ -19,7 +19,15 @@
 // CHECK:      @[[EXEC_MODE3:.*]] = weak protected constant i8 1
 // CHECK:      @llvm.compiler.used{{.*}} = appending global [1 x ptr] [ptr @[[EXEC_MODE3]]], section "llvm.metadata"
 // CHECK:      @[[KERNEL3_ENV:.*_kernel_environment]] = weak_odr protected constant %struct.KernelEnvironmentTy {
-// CHECK-SAME: %struct.ConfigurationEnvironmentTy { i8 1, i8 1, i8 [[EXEC_MODE3:1]], i32 [[MIN_THREADS3:1]], i32 [[MAX_THREADS3:10]], i32 0, i32 0, i32 0, i32 0 },
+// CHECK-SAME: %struct.ConfigurationEnvironmentTy { i8 1, i8 1, i8 [[EXEC_MODE3:1]], i32 [[MIN_THREADS3:1]], i32 [[MAX_THREADS3:10]], i32 0, i32 0, i32 0 },
+// CHECK-SAME: ptr @{{.*}}, ptr @{{.*}} }
+
+// Non-constant thread_limit: no compile-time bound, so MaxThreads stays 0
+// ("set but unknown") rather than being left at the specified-clause sentinel 1.
+// CHECK:      @[[EXEC_MODE4:.*]] = weak protected constant i8 1
+// CHECK:      @llvm.compiler.used{{.*}} = appending global [1 x ptr] [ptr @[[EXEC_MODE4]]], section "llvm.metadata"
+// CHECK:      @[[KERNEL4_ENV:.*_kernel_environment]] = weak_odr protected constant %struct.KernelEnvironmentTy {
+// CHECK-SAME: %struct.ConfigurationEnvironmentTy { i8 1, i8 1, i8 1, i32 1, i32 0, i32 0, i32 0, i32 0 },
 // CHECK-SAME: ptr @{{.*}}, ptr @{{.*}} }
 
 module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui32>>, llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_target_device = true, omp.is_gpu = true} {
@@ -46,19 +54,30 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memo
         omp.terminator
       }
       omp.terminator
-    }
+    } {omp.combined}
 
     // CHECK: define weak_odr protected amdgpu_kernel void @__omp_offloading_{{.*}}_main_l{{[0-9]+}}(ptr %[[KERNEL_ARGS:.*]]) #[[ATTRS1]]
     // CHECK: %{{.*}} = call i32 @__kmpc_target_init(ptr @[[KERNEL3_ENV]], ptr %[[KERNEL_ARGS]])
     %target_threads3 = llvm.mlir.constant(20) : i32
-    omp.target thread_limit(%target_threads3 : i32) {
+    omp.target kernel_type(generic) thread_limit(%target_threads3 : i32) {
       %teams_threads_x = llvm.mlir.constant(10) : i32
       %teams_threads_y = llvm.mlir.constant(5) : i32
       omp.teams thread_limit(%teams_threads_x, %teams_threads_y : i32, i32) {
         omp.terminator
       }
       omp.terminator
-    }
+    } {omp.combined}
+
+    // CHECK: define weak_odr protected amdgpu_kernel void @__omp_offloading_{{.*}}_main_l{{[0-9]+}}(ptr %{{.*}}, ptr %[[KERNEL_ARGS4:.*]]){{.*}}
+    // CHECK: %{{.*}} = call i32 @__kmpc_target_init(ptr @[[KERNEL4_ENV]], ptr %[[KERNEL_ARGS4]])
+    %1 = omp.map.info var_ptr(%num_teams : !llvm.ptr, i32) map_clauses(to) capture(ByCopy) -> !llvm.ptr
+    omp.target kernel_type(generic) map_entries(%1 -> %arg_limit : !llvm.ptr) {
+      %runtime_threads = llvm.load %arg_limit : !llvm.ptr -> i32
+      omp.teams thread_limit(%runtime_threads : i32) {
+        omp.terminator
+      }
+      omp.terminator
+    } {omp.combined}
     llvm.return
   }
 }
diff --git a/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir b/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir
index c615453a2870d8..cd6747dd4007f8 100644
--- a/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir
@@ -20,7 +20,16 @@
 // CHECK: %[[NT_ARR:.*]] = insertvalue [3 x i32] [i32 10, i32 0, i32 0], i32 %[[TL_Y]], 1
 // CHECK: %[[NT_GEP:.*]] = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %[[KERNEL_ARGS2]], i32 0, i32 11
 // CHECK-NEXT: store [3 x i32] %[[NT_ARR]], ptr %[[NT_GEP]], align 4
-// CHECK: call i32 @__tgt_target_kernel(ptr {{.*}}, i64 -1, i32 0, i32 10, ptr @.{{.*}}.region_id, ptr %[[KERNEL_ARGS2]])
+// CHECK: call i32 @__tgt_target_kernel(ptr {{.*}}, i64 -1, i32 -1, i32 10, ptr @.{{.*}}.region_id, ptr %[[KERNEL_ARGS2]])
+
+// Mixed-rank thread_limit: target is 1-D (20), teams is 2-D (10, 5). The target's
+// unspecified trailing dim is implicitly 1, so dim1 clamps to min(1, 5) = 1.
+// CHECK: define void @main_mixed_rank_thread_limit()
+// CHECK: %[[KERNEL_ARGS3:.*]] = alloca %struct.__tgt_kernel_arguments
+// CHECK: %[[NT_GEP3:.*]] = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %[[KERNEL_ARGS3]], i32 0, i32 11
+// CHECK-NEXT: store [3 x i32] [i32 10, i32 1, i32 0], ptr %[[NT_GEP3]], align 4
+// CHECK: call i32 @__tgt_target_kernel(ptr {{.*}}, i64 -1, i32 -1, i32 10, ptr @.{{.*}}.region_id, ptr %[[KERNEL_ARGS3]])
+
 // CHECK: define internal void @[[OUTLINED_FN]](i32 %[[NUM_TEAMS_OUTLINED:.*]], ptr %{{.*}})
 // CHECK: call void @__kmpc_push_num_teams_51(ptr {{.*}}, i32 {{.*}}, i32 %[[NUM_TEAMS_OUTLINED]], i32 %[[NUM_TEAMS_OUTLINED]], i32 [[NUM_THREADS]])
 module attributes {omp.is_target_device = false, omp.target_triples = ["amdgcn-amd-amdhsa"]} {
@@ -39,7 +48,21 @@ module attributes {omp.is_target_device = false, omp.target_triples = ["amdgcn-a
 
   llvm.func @main_multidim_thread_limit(%thread_limit_y : i32) {
     %teams_threads_x = llvm.mlir.constant(10) : i32
-    omp.target host_eval(%teams_threads_x -> %arg_tlx, %thread_limit_y -> %arg_tly : i32, i32) {
+    omp.target kernel_type(generic) host_eval(%teams_threads_x -> %arg_tlx, %thread_limit_y -> %arg_tly : i32, i32) {
+      omp.teams thread_limit(%arg_tlx, %arg_tly : i32, i32) {
+        omp.terminator
+      }
+      omp.terminator
+    }
+    llvm.return
+  }
+
+  llvm.func @main_mixed_rank_thread_limit() {
+    %target_threads = llvm.mlir.constant(20) : i32
+    %teams_threads_x = llvm.mlir.constant(10) : i32
+    %teams_threads_y = llvm.mlir.constant(5) : i32
+    omp.target kernel_type(generic) thread_limit(%target_threads : i32)
+               host_eval(%teams_threads_x -> %arg_tlx, %teams_threads_y -> %arg_tly : i32, i32) {
       omp.teams thread_limit(%arg_tlx, %arg_tly : i32, i32) {
         omp.terminator
       }
diff --git a/mlir/test/Target/LLVMIR/openmp-todo.mlir b/mlir/test/Target/LLVMIR/openmp-todo.mlir
index 7138d8515fa514..acb19badaad4c6 100644
--- a/mlir/test/Target/LLVMIR/openmp-todo.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-todo.mlir
@@ -681,17 +681,6 @@ llvm.func @parallel_num_threads_too_many_dims(%lb : i32, %ub : i32) {
 
 // -----
 
-llvm.func @parallel_num_threads_too_many_dims(%lb : i32, %ub : i32) {
-  // expected-error at below {{not yet implemented: Unhandled clause num_threads with more than 3 dimensions in omp.parallel operation}}
-  // expected-error at below {{LLVM Translation failed for operation: omp.parallel}}
-  omp.parallel num_threads(%lb, %ub, %lb, %ub : i32, i32, i32, i32) {
-    omp.terminator
-  }
-  llvm.return
-}
-
-// -----
-
 llvm.func @teams_thread_limit_too_many_dims(%lb : i32, %ub : i32) {
   // expected-error at below {{not yet implemented: Unhandled clause thread_limit with more than 3 dimensions in omp.teams operation}}
   // expected-error at below {{LLVM Translation failed for operation: omp.teams}}

>From a757007b2d8534da462a9f24b005ce3c95acf74c Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Wed, 7 Oct 2026 11:36:28 +0530
Subject: [PATCH 6/6] update after rebase to latest main

---
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      |  9 +++-
 .../LLVMIR/openmp-target-launch-device.mlir   | 15 +++---
 .../LLVMIR/openmp-target-launch-host.mlir     |  8 ++--
 mlir/test/Target/LLVMIR/openmp-teams.mlir     | 48 ++++++++-----------
 mlir/test/Target/LLVMIR/openmp-todo.mlir      | 22 +++++++++
 5 files changed, 62 insertions(+), 40 deletions(-)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 6e91551b9feb09..d680d3690601a6 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -531,8 +531,15 @@ static LogicalResult checkImplementationStatus(Operation &op) {
   };
 
   auto checkThreadLimit = [&todo](auto op, LogicalResult &result) {
-    if (op.getThreadLimitDimsCount() > 3)
+    if (op.getThreadLimitDimsCount() > 3) {
       result = todo("thread_limit with more than 3 dimensions");
+      return;
+    }
+
+    if (op.hasThreadLimitMultiDim() && !isa<omp::TargetOp>(op.getOperation()) &&
+        !op->template getParentOfType<omp::TargetOp>())
+      result = todo(
+          "thread_limit with multi-dimensional values outside target region");
   };
   auto checkMap = [&todo](auto op, LogicalResult &result) {
     if (!op.getMapIterated().empty())
diff --git a/mlir/test/Target/LLVMIR/openmp-target-launch-device.mlir b/mlir/test/Target/LLVMIR/openmp-target-launch-device.mlir
index 2a7ec81a01259e..7870446084337a 100644
--- a/mlir/test/Target/LLVMIR/openmp-target-launch-device.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-target-launch-device.mlir
@@ -3,8 +3,8 @@
 // CHECK:      @[[EXEC_MODE1:.*]] = weak protected constant i8 1
 // CHECK:      @llvm.compiler.used{{.*}} = appending global [1 x ptr] [ptr @[[EXEC_MODE1]]], section "llvm.metadata"
 // CHECK:      @[[KERNEL1_ENV:.*_kernel_environment]] = weak_odr protected constant %struct.KernelEnvironmentTy {
-// Both kernels below are generic, so their blocks carry one warp for the main
-// thread on top of the thread_limit they ask for: 10 + 64 and 30 + 64.
+// All kernels below are generic, so a known thread_limit gets one extra warp
+// for the main thread: 10 + 64, 30 + 64 and 10 + 64.
 // CHECK-SAME: %struct.ConfigurationEnvironmentTy { i8 1, i8 1, i8 [[EXEC_MODE1:1]], i32 [[MIN_THREADS1:1]], i32 [[MAX_THREADS1:74]], i32 [[MIN_TEAMS1:1]], i32 [[MAX_TEAMS1:-1]], i32 0 },
 // CHECK-SAME: ptr @{{.*}}, ptr @{{.*}} }
 
@@ -15,11 +15,12 @@
 // CHECK-SAME: ptr @{{.*}}, ptr @{{.*}} }
 
 // Multi-dim thread_limit: first dim constant (10), second dim constant (5).
-// MaxThreads uses the first dim combined value: min(target=20, teams_x=10) = 10.
+// MaxThreads uses the first dim combined value: min(target=20, teams_x=10) = 10,
+// plus the generic-mode warp.
 // CHECK:      @[[EXEC_MODE3:.*]] = weak protected constant i8 1
 // CHECK:      @llvm.compiler.used{{.*}} = appending global [1 x ptr] [ptr @[[EXEC_MODE3]]], section "llvm.metadata"
 // CHECK:      @[[KERNEL3_ENV:.*_kernel_environment]] = weak_odr protected constant %struct.KernelEnvironmentTy {
-// CHECK-SAME: %struct.ConfigurationEnvironmentTy { i8 1, i8 1, i8 [[EXEC_MODE3:1]], i32 [[MIN_THREADS3:1]], i32 [[MAX_THREADS3:10]], i32 0, i32 0, i32 0 },
+// CHECK-SAME: %struct.ConfigurationEnvironmentTy { i8 1, i8 1, i8 [[EXEC_MODE3:1]], i32 [[MIN_THREADS3:1]], i32 [[MAX_THREADS3:74]], i32 0, i32 0, i32 0 },
 // CHECK-SAME: ptr @{{.*}}, ptr @{{.*}} }
 
 // Non-constant thread_limit: no compile-time bound, so MaxThreads stays 0
@@ -58,10 +59,10 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memo
 
     // CHECK: define weak_odr protected amdgpu_kernel void @__omp_offloading_{{.*}}_main_l{{[0-9]+}}(ptr %[[KERNEL_ARGS:.*]]) #[[ATTRS1]]
     // CHECK: %{{.*}} = call i32 @__kmpc_target_init(ptr @[[KERNEL3_ENV]], ptr %[[KERNEL_ARGS]])
-    %target_threads3 = llvm.mlir.constant(20) : i32
+    %target_threads3 = llvm.mlir.constant(20 : i32) : i32
     omp.target kernel_type(generic) thread_limit(%target_threads3 : i32) {
-      %teams_threads_x = llvm.mlir.constant(10) : i32
-      %teams_threads_y = llvm.mlir.constant(5) : i32
+      %teams_threads_x = llvm.mlir.constant(10 : i32) : i32
+      %teams_threads_y = llvm.mlir.constant(5 : i32) : i32
       omp.teams thread_limit(%teams_threads_x, %teams_threads_y : i32, i32) {
         omp.terminator
       }
diff --git a/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir b/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir
index cd6747dd4007f8..83241185455009 100644
--- a/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-target-launch-host.mlir
@@ -47,7 +47,7 @@ module attributes {omp.is_target_device = false, omp.target_triples = ["amdgcn-a
   }
 
   llvm.func @main_multidim_thread_limit(%thread_limit_y : i32) {
-    %teams_threads_x = llvm.mlir.constant(10) : i32
+    %teams_threads_x = llvm.mlir.constant(10 : i32) : i32
     omp.target kernel_type(generic) host_eval(%teams_threads_x -> %arg_tlx, %thread_limit_y -> %arg_tly : i32, i32) {
       omp.teams thread_limit(%arg_tlx, %arg_tly : i32, i32) {
         omp.terminator
@@ -58,9 +58,9 @@ module attributes {omp.is_target_device = false, omp.target_triples = ["amdgcn-a
   }
 
   llvm.func @main_mixed_rank_thread_limit() {
-    %target_threads = llvm.mlir.constant(20) : i32
-    %teams_threads_x = llvm.mlir.constant(10) : i32
-    %teams_threads_y = llvm.mlir.constant(5) : i32
+    %target_threads = llvm.mlir.constant(20 : i32) : i32
+    %teams_threads_x = llvm.mlir.constant(10 : i32) : i32
+    %teams_threads_y = llvm.mlir.constant(5 : i32) : i32
     omp.target kernel_type(generic) thread_limit(%target_threads : i32)
                host_eval(%teams_threads_x -> %arg_tlx, %teams_threads_y -> %arg_tly : i32, i32) {
       omp.teams thread_limit(%arg_tlx, %arg_tly : i32, i32) {
diff --git a/mlir/test/Target/LLVMIR/openmp-teams.mlir b/mlir/test/Target/LLVMIR/openmp-teams.mlir
index 126d3e652a6e15..a7bc92f8b99b9e 100644
--- a/mlir/test/Target/LLVMIR/openmp-teams.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-teams.mlir
@@ -314,34 +314,26 @@ llvm.func @teams_if_with_num_teams(%condition: i1, %numTeamsLower: i32, %numTeam
 
 // -----
 
-llvm.func @duringTeams()
-
-// CHECK-LABEL: @omp_teams_thread_limit_2d
-// CHECK-SAME: (i32 [[LIMIT_X:.+]], i32 [[LIMIT_Y:.+]])
-llvm.func @omp_teams_thread_limit_2d(%limitX: i32, %limitY: i32) {
-    // CHECK: [[THREAD_NUM:%.+]] = call i32 @__kmpc_global_thread_num
-    // CHECK-NEXT: call void @__kmpc_push_num_teams_51({{.+}}, i32 [[THREAD_NUM]], i32 0, i32 0, i32 [[LIMIT_X]])
-    // CHECK: call void (ptr, i32, ptr, ...) @__kmpc_fork_teams(ptr @{{[0-9]+}}, i32 0, ptr [[OUTLINED_FN:.+]])
-    omp.teams thread_limit(%limitX, %limitY : i32, i32) {
-        llvm.call @duringTeams() : () -> ()
-        omp.terminator
+// Check that the thread and bound id arguments of the outlined function are
+// generic pointers even when allocas are created in a non-zero address space.
+// The cast that produces them is deleted along with the fake allocas once the
+// region has been outlined, so the signature below is the only place it can be
+// observed.
+module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui32>>} {
+    llvm.func @foo()
+
+    // CHECK-LABEL: @omp_teams_alloca_addrspace
+    // CHECK: call void {{.*}} @__kmpc_fork_teams(ptr @{{.+}}, i32 0, ptr @[[OUTLINED_FN:.+]])
+    llvm.func @omp_teams_alloca_addrspace() {
+        omp.teams {
+            llvm.call @foo() : () -> ()
+            omp.terminator
+        }
+        llvm.return
     }
-    llvm.return
-}
-
-// -----
-
-llvm.func @duringTeams()
 
-// CHECK-LABEL: @omp_teams_thread_limit_3d
-// CHECK-SAME: (i32 [[LIMIT_X:.+]], i64 [[LIMIT_Y:.+]], i16 [[LIMIT_Z:.+]])
-llvm.func @omp_teams_thread_limit_3d(%limitX: i32, %limitY: i64, %limitZ: i16) {
-    // CHECK: [[THREAD_NUM:%.+]] = call i32 @__kmpc_global_thread_num
-    // CHECK-NEXT: call void @__kmpc_push_num_teams_51({{.+}}, i32 [[THREAD_NUM]], i32 0, i32 0, i32 [[LIMIT_X]])
-    // CHECK: call void (ptr, i32, ptr, ...) @__kmpc_fork_teams(ptr @{{[0-9]+}}, i32 0, ptr [[OUTLINED_FN:.+]])
-    omp.teams thread_limit(%limitX, %limitY, %limitZ : i32, i64, i16) {
-        llvm.call @duringTeams() : () -> ()
-        omp.terminator
-    }
-    llvm.return
+    // CHECK:      define internal void @[[OUTLINED_FN]]
+    // CHECK-SAME: (ptr %global.tid.ptr, ptr %bound.tid.ptr)
+    // CHECK:   call void @foo()
+    // CHECK:   ret void
 }
diff --git a/mlir/test/Target/LLVMIR/openmp-todo.mlir b/mlir/test/Target/LLVMIR/openmp-todo.mlir
index acb19badaad4c6..fca9cf5344674e 100644
--- a/mlir/test/Target/LLVMIR/openmp-todo.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-todo.mlir
@@ -670,6 +670,17 @@ llvm.func @teams_num_teams_multi_dim(%lb : i32, %ub : i32) {
 
 // -----
 
+llvm.func @parallel_num_threads_multi_dim_standalone(%lb : i32, %ub : i32) {
+  // expected-error at below {{not yet implemented: Unhandled clause num_threads with multi-dimensional values outside target region in omp.parallel operation}}
+  // expected-error at below {{LLVM Translation failed for operation: omp.parallel}}
+  omp.parallel num_threads(%lb, %ub : i32, i32) {
+    omp.terminator
+  }
+  llvm.return
+}
+
+// -----
+
 llvm.func @parallel_num_threads_too_many_dims(%lb : i32, %ub : i32) {
   // expected-error at below {{not yet implemented: Unhandled clause num_threads with more than 3 dimensions in omp.parallel operation}}
   // expected-error at below {{LLVM Translation failed for operation: omp.parallel}}
@@ -692,6 +703,17 @@ llvm.func @teams_thread_limit_too_many_dims(%lb : i32, %ub : i32) {
 
 // -----
 
+llvm.func @teams_thread_limit_multi_dim_standalone(%x : i32, %y : i32) {
+  // expected-error at below {{not yet implemented: Unhandled clause thread_limit with multi-dimensional values outside target region in omp.teams operation}}
+  // expected-error at below {{LLVM Translation failed for operation: omp.teams}}
+  omp.teams thread_limit(%x, %y : i32, i32) {
+    omp.terminator
+  }
+  llvm.return
+}
+
+// -----
+
 llvm.func @teams_dyn_groupprivate(%dyn_size : i32) {
   // expected-error at below {{not yet implemented: Unhandled clause dyn_groupprivate in omp.teams operation}}
   // expected-error at below {{LLVM Translation failed for operation: omp.teams}}



More information about the llvm-branch-commits mailing list