[Mlir-commits] [mlir] [OpenMP][MLIR] Add num_threads mlir->llvm lowering (PR #179420)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Oct 6 23:09:21 PDT 2026
https://github.com/skc7 updated https://github.com/llvm/llvm-project/pull/179420
>From 8dca1e9b1f28ba8453ccd12df62d0d9903dc8dd1 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/3] [OpenMP][MLIR] Add num_threads mlir->llvm lowering
---
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 67 ++++++++++++-------
mlir/test/Target/LLVMIR/openmp-todo.mlir | 6 +-
2 files changed, 46 insertions(+), 27 deletions(-)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index ba9b30e69aacef..e456a997882e28 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -519,8 +519,9 @@ static LogicalResult checkImplementationStatus(Operation &op) {
result = todo("num_teams with multi-dimensional values");
};
auto checkNumThreads = [&todo](auto op, LogicalResult &result) {
- if (op.hasNumThreadsMultiDim())
- result = todo("num_threads with multi-dimensional values");
+ // Check that we don't exceed the maximum supported dimensions (3)
+ if (op.getNumThreadsDimsCount() > 3)
+ result = todo("num_threads with more than 3 dimensions");
};
auto checkThreadLimit = [&todo](auto op, LogicalResult &result) {
@@ -9119,13 +9120,12 @@ static llvm::IRBuilderBase::InsertPoint createDeviceArgumentAccessor(
///
/// Loop bounds and steps are only optionally populated, if output vectors are
/// provided.
-static void
-extractHostEvalClauses(omp::TargetOp targetOp, Value &numThreads,
- 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, Value &threadLimit,
+ 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())) {
@@ -9146,10 +9146,17 @@ extractHostEvalClauses(omp::TargetOp targetOp, Value &numThreads,
llvm_unreachable("unsupported host_eval use");
})
.Case([&](omp::ParallelOp parallelOp) {
- if (!parallelOp.getNumThreadsVars().empty() &&
- parallelOp.getNumThreads(0) == blockArg)
- numThreads = hostEvalVar;
- else
+ if (llvm::is_contained(parallelOp.getNumThreadsVars(), blockArg)) {
+ for (auto [i, threadsVar] :
+ llvm::enumerate(parallelOp.getNumThreadsVars())) {
+ if (threadsVar == blockArg) {
+ if (numThreadsVars.size() <= i)
+ numThreadsVars.resize(i + 1);
+ numThreadsVars[i] = hostEvalVar;
+ break;
+ }
+ }
+ } else
llvm_unreachable("unsupported host_eval use");
})
.Case([&](omp::LoopNestOp loopOp) {
@@ -9257,10 +9264,11 @@ initTargetDefaultAttrs(omp::TargetOp targetOp, Operation *capturedOp,
bool isTargetDevice, bool isGPU) {
// TODO: Handle constant 'if' clauses.
- Value numThreads, numTeamsLower, numTeamsUpper, threadLimit;
+ Value numTeamsLower, numTeamsUpper, threadLimit;
+ llvm::SmallVector<Value> numThreadsVars;
if (!isTargetDevice) {
- extractHostEvalClauses(targetOp, numThreads, numTeamsLower, numTeamsUpper,
- threadLimit);
+ extractHostEvalClauses(targetOp, numThreadsVars, numTeamsLower,
+ numTeamsUpper, threadLimit);
} 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
@@ -9275,8 +9283,10 @@ initTargetDefaultAttrs(omp::TargetOp targetOp, Operation *capturedOp,
}
if (auto parallelOp = castOrGetParentOfType<omp::ParallelOp>(capturedOp)) {
- if (!parallelOp.getNumThreadsVars().empty())
- numThreads = parallelOp.getNumThreads(0);
+ // Handle multi-dimensional num_threads
+ numThreadsVars.reserve(parallelOp.getNumThreadsVars().size());
+ for (auto threadsVar : parallelOp.getNumThreadsVars())
+ numThreadsVars.push_back(threadsVar);
}
}
@@ -9324,10 +9334,12 @@ initTargetDefaultAttrs(omp::TargetOp targetOp, Operation *capturedOp,
// Extract 'max_threads' clause from 'parallel' or set to 1 if it's SIMD.
int32_t maxThreadsVal = -1;
- if (castOrGetParentOfType<omp::ParallelOp>(capturedOp))
- setMaxValueFromClause(numThreads, maxThreadsVal);
- else if (castOrGetParentOfType<omp::SimdOp>(capturedOp,
- /*immediateParent=*/true))
+ if (castOrGetParentOfType<omp::ParallelOp>(capturedOp)) {
+ // For multi-dimensional num_threads, only use the first dimension for now
+ if (!numThreadsVars.empty())
+ setMaxValueFromClause(numThreadsVars[0], maxThreadsVal);
+ } else if (castOrGetParentOfType<omp::SimdOp>(capturedOp,
+ /*immediateParent=*/true))
maxThreadsVal = 1;
// For max values, < 0 means unset, == 0 means set but unknown. Select the
@@ -9386,10 +9398,11 @@ initTargetRuntimeAttrs(llvm::IRBuilderBase &builder,
omp::LoopNestOp loopOp = castOrGetParentOfType<omp::LoopNestOp>(capturedOp);
unsigned numLoops = loopOp ? loopOp.getNumLoops() : 0;
- Value numThreads, numTeamsLower, numTeamsUpper, teamsThreadLimit;
+ Value numTeamsLower, numTeamsUpper, teamsThreadLimit;
+ llvm::SmallVector<Value> numThreadsVars;
llvm::SmallVector<Value> lowerBounds(numLoops), upperBounds(numLoops),
steps(numLoops);
- extractHostEvalClauses(targetOp, numThreads, numTeamsLower, numTeamsUpper,
+ extractHostEvalClauses(targetOp, numThreadsVars, numTeamsLower, numTeamsUpper,
teamsThreadLimit, &lowerBounds, &upperBounds, &steps);
// TODO: Handle constant 'if' clauses.
@@ -9414,8 +9427,14 @@ initTargetRuntimeAttrs(llvm::IRBuilderBase &builder,
attrs.TeamsThreadLimit.front() = builder.CreateSExtOrTrunc(
moduleTranslation.lookupValue(teamsThreadLimit), builder.getInt32Ty());
+<<<<<<< HEAD
if (numThreads)
attrs.MaxThreads.front() = moduleTranslation.lookupValue(numThreads);
+=======
+ // Handle multi-dimensional num_threads (only first value for now)
+ if (!numThreadsVars.empty())
+ attrs.MaxThreads = moduleTranslation.lookupValue(numThreadsVars[0]);
+>>>>>>> [OpenMP][MLIR] Add num_threads mlir->llvm lowering
if (targetOp.hasHostEvalTripCount()) {
llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
diff --git a/mlir/test/Target/LLVMIR/openmp-todo.mlir b/mlir/test/Target/LLVMIR/openmp-todo.mlir
index 327d0915388636..949b24d6745257 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(%lb : i32, %ub : i32) {
- // expected-error at below {{not yet implemented: Unhandled clause num_threads with multi-dimensional values 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 4cbf8af182bc0fcb138c402734eac3c883e71d23 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Fri, 10 Apr 2026 11:31:59 +0530
Subject: [PATCH 2/3] add todo test
---
.../Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp | 13 ++++++++++---
mlir/test/Target/LLVMIR/openmp-todo.mlir | 11 +++++++++++
2 files changed, 21 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index e456a997882e28..1026a41c4f7ccf 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -519,9 +519,15 @@ static LogicalResult checkImplementationStatus(Operation &op) {
result = todo("num_teams with multi-dimensional values");
};
auto checkNumThreads = [&todo](auto op, LogicalResult &result) {
- // Check that we don't exceed the maximum supported dimensions (3)
- if (op.getNumThreadsDimsCount() > 3)
+ if (op.getNumThreadsDimsCount() > 3) {
result = todo("num_threads with more than 3 dimensions");
+ return;
+ }
+
+ if (op.hasNumThreadsMultiDim() &&
+ !op->template getParentOfType<omp::TargetOp>())
+ result = todo(
+ "num_threads with multi-dimensional values outside target region");
};
auto checkThreadLimit = [&todo](auto op, LogicalResult &result) {
@@ -9156,8 +9162,9 @@ static void extractHostEvalClauses(
break;
}
}
- } else
+ } else {
llvm_unreachable("unsupported host_eval use");
+ }
})
.Case([&](omp::LoopNestOp loopOp) {
auto processBounds =
diff --git a/mlir/test/Target/LLVMIR/openmp-todo.mlir b/mlir/test/Target/LLVMIR/openmp-todo.mlir
index 949b24d6745257..320f678594bf18 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}}
>From f4ef7d5b7abb7579b0b8da8e6fc750f6e5e7dabb Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Mon, 31 Aug 2026 12:59:01 +0530
Subject: [PATCH 3/3] Update runtime MaxThreads attr
---
.../Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp | 16 ++++++++--------
1 file changed, 8 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 1026a41c4f7ccf..1a3f9ef22d2724 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -9434,14 +9434,14 @@ initTargetRuntimeAttrs(llvm::IRBuilderBase &builder,
attrs.TeamsThreadLimit.front() = builder.CreateSExtOrTrunc(
moduleTranslation.lookupValue(teamsThreadLimit), builder.getInt32Ty());
-<<<<<<< HEAD
- if (numThreads)
- attrs.MaxThreads.front() = moduleTranslation.lookupValue(numThreads);
-=======
- // Handle multi-dimensional num_threads (only first value for now)
- if (!numThreadsVars.empty())
- attrs.MaxThreads = moduleTranslation.lookupValue(numThreadsVars[0]);
->>>>>>> [OpenMP][MLIR] Add num_threads mlir->llvm lowering
+ // One runtime value per num_threads dimension.
+ if (!numThreadsVars.empty()) {
+ attrs.MaxThreads.clear();
+ for (Value numThreadsVar : numThreadsVars)
+ attrs.MaxThreads.push_back(
+ numThreadsVar ? moduleTranslation.lookupValue(numThreadsVar)
+ : nullptr);
+ }
if (targetOp.hasHostEvalTripCount()) {
llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
More information about the Mlir-commits
mailing list