[Mlir-commits] [mlir] [mlir][OpenMP] Don't use label prefixes on linear variable rewrite (PR #200900)
Leandro Lupori
llvmlistbot at llvm.org
Fri Jun 19 15:35:41 PDT 2026
https://github.com/luporl updated https://github.com/llvm/llvm-project/pull/200900
>From 2e1e984f3fcca45bfdc010513ea2a8800a27eb09 Mon Sep 17 00:00:00 2001
From: Leandro Lupori <leandro.lupori at linaro.org>
Date: Mon, 1 Jun 2026 17:54:55 +0000
Subject: [PATCH 1/2] [mlir][OpenMP] Don't use label prefixes on linear
variable rewrite
This is a follow-up to #194623. After that PR, matching specific label
prefixes became unneccessary. In fact, doing so could potentially lead
to missed linear variables in the rewrite, if they appear in basic
blocks with unexpected label prefixes.
---
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 52 +++++--------------
1 file changed, 13 insertions(+), 39 deletions(-)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index deefa29c157f3..a991d383bf8a2 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -284,42 +284,36 @@ class LinearClauseProcessor {
}
}
- // Rewrite all uses of the original variable, in the basic blocks whose names
- // start with `prefix`, with the linear variable in-place.
+ // Rewrite all uses of the original variable, in the basic blocks in the
+ // [startBB, endBB] interval, with the linear variable in-place.
void rewriteInPlace(llvm::IRBuilderBase &builder, llvm::BasicBlock *startBB,
- llvm::BasicBlock *endBB, llvm::StringRef prefix,
- size_t varIndex) {
+ llvm::BasicBlock *endBB, size_t varIndex) {
llvm::SmallVector<llvm::BasicBlock *, 32> worklist;
- llvm::SmallPtrSet<llvm::BasicBlock *, 32> visited;
- llvm::SmallPtrSet<llvm::BasicBlock *, 32> matchingBBs;
+ llvm::SmallPtrSet<llvm::BasicBlock *, 32> collectedBBs;
assert(startBB && endBB && "Invalid startBB/endBB");
- // Traverse basic blocks from startBB to endBB and save those
- // whose names start with the specified prefix.
+ // Collect basic blocks from startBB to endBB.
worklist.push_back(startBB);
- visited.insert(startBB);
+ collectedBBs.insert(startBB);
while (!worklist.empty()) {
llvm::BasicBlock *bb = worklist.pop_back_val();
- if (bb->hasName() && bb->getName().starts_with(prefix))
- matchingBBs.insert(bb);
-
if (bb == endBB)
continue;
for (llvm::BasicBlock *succ : llvm::successors(bb)) {
- if (visited.insert(succ).second)
+ if (collectedBBs.insert(succ).second)
worklist.push_back(succ);
}
}
- // Rewrite all uses in the matching BBs.
+ // Rewrite all uses in the collected BBs.
llvm::SmallVector<llvm::User *> users(linearOrigVal[varIndex]->users());
for (auto *user : users) {
if (auto *userInst = dyn_cast<llvm::Instruction>(user)) {
- if (matchingBBs.contains(userInst->getParent()))
+ if (collectedBBs.contains(userInst->getParent()))
user->replaceUsesOfWith(linearOrigVal[varIndex],
linearLoopBodyTemps[varIndex]);
}
@@ -3897,8 +3891,7 @@ convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
return failure();
for (size_t index = 0; index < wsloopOp.getLinearVars().size(); index++)
linearClauseProcessor.rewriteInPlace(
- builder, sourceBlock->getSingleSuccessor(), *regionBlock,
- "omp.loop_nest.region", index);
+ builder, sourceBlock->getSingleSuccessor(), *regionBlock, index);
builder.restoreIP(oldIP);
}
@@ -4263,28 +4256,9 @@ convertOmpSimd(Operation &opInst, llvm::IRBuilderBase &builder,
linearClauseProcessor.emitStoresForLinearVar(builder);
- // Check if this SIMD loop contains ordered regions
- bool hasOrderedRegions = false;
- simdOp.getRegion().walk([&](omp::OrderedRegionOp orderedOp) {
- hasOrderedRegions = true;
- return WalkResult::interrupt();
- });
-
- for (size_t index = 0; index < simdOp.getLinearVars().size(); index++) {
- llvm::BasicBlock *startBB = sourceBlock->getSingleSuccessor();
- llvm::BasicBlock *endBB = *regionBlock;
- linearClauseProcessor.rewriteInPlace(builder, startBB, endBB,
- "omp.loop_nest.region", index);
-
- if (hasOrderedRegions) {
- // Also rewrite uses in ordered regions so they read the current value
- linearClauseProcessor.rewriteInPlace(builder, startBB, endBB,
- "omp.ordered.region", index);
- // Also rewrite uses in finalize blocks (code after ordered regions)
- linearClauseProcessor.rewriteInPlace(builder, startBB, endBB,
- "omp_region.finalize", index);
- }
- }
+ for (size_t index = 0; index < simdOp.getLinearVars().size(); index++)
+ linearClauseProcessor.rewriteInPlace(
+ builder, sourceBlock->getSingleSuccessor(), *regionBlock, index);
// We now need to reduce the per-simd-lane reduction variable into the
// original variable. This works a bit differently to other reductions (e.g.
>From 79f681e9bff42e11281efbb93fa7e84d45491ab1 Mon Sep 17 00:00:00 2001
From: Leandro Lupori <leandro.lupori at linaro.org>
Date: Fri, 19 Jun 2026 19:13:29 -0300
Subject: [PATCH 2/2] Rewrite only the blocks between the loop body and latch
---
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 16 ++++-----
.../Target/LLVMIR/openmp-wsloop-linear.mlir | 33 +++++++++++++++++++
2 files changed, 41 insertions(+), 8 deletions(-)
create mode 100644 mlir/test/Target/LLVMIR/openmp-wsloop-linear.mlir
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index a991d383bf8a2..76f77c3c5788b 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -3827,7 +3827,6 @@ convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
linearClauseProcessor.initLinearStep(moduleTranslation, linearStep);
}
- llvm::BasicBlock *sourceBlock = builder.GetInsertBlock();
llvm::Expected<llvm::BasicBlock *> regionBlock = convertOmpOpRegions(
wsloopOp.getRegion(), "omp.wsloop.region", builder, moduleTranslation);
@@ -3868,6 +3867,10 @@ convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
}
}
+ for (size_t index = 0; index < wsloopOp.getLinearVars().size(); index++)
+ linearClauseProcessor.rewriteInPlace(builder, loopInfo->getBody(),
+ loopInfo->getLatch(), index);
+
llvm::OpenMPIRBuilder::InsertPointOrErrorTy wsloopIP =
ompBuilder->applyWorkshareLoop(
ompLoc.DL, loopInfo, allocaIP, loopNeedsBarrier,
@@ -3889,9 +3892,6 @@ convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
loopInfo->getLastIter());
if (failed(handleError(afterBarrierIP, *loopOp)))
return failure();
- for (size_t index = 0; index < wsloopOp.getLinearVars().size(); index++)
- linearClauseProcessor.rewriteInPlace(
- builder, sourceBlock->getSingleSuccessor(), *regionBlock, index);
builder.restoreIP(oldIP);
}
@@ -4248,6 +4248,10 @@ convertOmpSimd(Operation &opInst, llvm::IRBuilderBase &builder,
}
builder.SetInsertPoint(*regionBlock, (*regionBlock)->begin());
+ for (size_t index = 0; index < simdOp.getLinearVars().size(); index++)
+ linearClauseProcessor.rewriteInPlace(builder, loopInfo->getBody(),
+ loopInfo->getLatch(), index);
+
ompBuilder->applySimd(loopInfo, alignedVars,
simdOp.getIfExpr()
? moduleTranslation.lookupValue(simdOp.getIfExpr())
@@ -4256,10 +4260,6 @@ convertOmpSimd(Operation &opInst, llvm::IRBuilderBase &builder,
linearClauseProcessor.emitStoresForLinearVar(builder);
- for (size_t index = 0; index < simdOp.getLinearVars().size(); index++)
- linearClauseProcessor.rewriteInPlace(
- builder, sourceBlock->getSingleSuccessor(), *regionBlock, index);
-
// We now need to reduce the per-simd-lane reduction variable into the
// original variable. This works a bit differently to other reductions (e.g.
// wsloop) because we don't need to call into the OpenMP runtime to handle
diff --git a/mlir/test/Target/LLVMIR/openmp-wsloop-linear.mlir b/mlir/test/Target/LLVMIR/openmp-wsloop-linear.mlir
new file mode 100644
index 0000000000000..f47ae6e11170f
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-wsloop-linear.mlir
@@ -0,0 +1,33 @@
+// Ensure that omp.wsloop with the linear clause is translated correctly.
+// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck %s
+
+// -----
+
+// %.linear_result must appear only in the loop body and in the
+// linear_lastiter_exit block, where it's used to update the original
+// variable.
+// CHECK-LABEL: void @wsloop_linear_post_use({{.*}})
+// CHECK: %.linear_result = alloca i32
+// CHECK-NOT: %.linear_result
+// CHECK: omp_loop.body:
+// CHECK: %.linear_result
+// CHECK: omp_loop.inc:
+// CHECK-NOT: %.linear_result
+// CHECK: omp_loop.linear_lastiter_exit:
+// CHECK: %.linear_result
+// CHECK: omp_loop.linear_exit:
+// CHECK-NOT: %.linear_result
+
+llvm.func @wsloop_linear_post_use(%lb : i32, %ub : i32, %step : i32,
+ %x : !llvm.ptr, %out : !llvm.ptr) {
+ omp.wsloop linear(%x : !llvm.ptr = %step : i32) {
+ omp.loop_nest (%iv) : i32 = (%lb) to (%ub) step (%step) {
+ %cur = llvm.load %x : !llvm.ptr -> i32
+ llvm.store %cur, %out : i32, !llvm.ptr
+ omp.yield
+ }
+ } {linear_var_types = [i32]}
+ %after = llvm.load %x : !llvm.ptr -> i32
+ llvm.store %after, %out : i32, !llvm.ptr
+ llvm.return
+}
More information about the Mlir-commits
mailing list