[Mlir-commits] [llvm] [mlir] [mlir][OpenMP] Allow multi-block omp.iterator regions (PR #229198)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Oct 6 09:56:57 PDT 2026
https://github.com/chichunchen updated https://github.com/llvm/llvm-project/pull/229198
>From f6a258cb17ca385e6d007f3e605ddca27bd6b8ca Mon Sep 17 00:00:00 2001
From: "Chi Chun, Chen" <chichun.chen at hpe.com>
Date: Mon, 5 Oct 2026 14:23:55 -0500
Subject: [PATCH 1/2] [mlir][OpenMP] Allow multi-block omp.iterator regions
Frontends lower a locator inside an omp.iterator region like any other
expression, so the region can contain control flow. For example, Flang
lowers LEN_TRIM in `depend(iterator(i=1:n), in: a(i+len_trim(s)))` to a
loop, and later passes inline SUM or COUNT as loops. After control-flow
conversion the region has several blocks, which the single-block
omp.iterator rejects:
%it = omp.iterator(%i: i64) = (%c1 to %n step %c1) {
llvm.br ^scan(%slen : i64)
^scan(%k: i64): // LEN_TRIM skips trailing blanks
...
llvm.cond_br %blank, ^scan(%km1 : i64), ^done
^done:
... // address of a(i+len_trim(s))
omp.yield(%addr : !llvm.ptr)
} -> !omp.iterated<!llvm.ptr>
Allow the region to have several blocks. Exactly one of them ends in
omp.yield, and every other block branches within the region, as for
omp.private. Translation keeps the single-block path; for multi-block
regions it converts the blocks in dominance order and reads the yielded
value after the yield. OpenMPIRBuilder::createIteratorLoop now branches the
body's single unterminated block to the latch and rejects bodies that
branch out of the loop.
Assisted with Copilot and Claude Opus 5.
---
.../llvm/Frontend/OpenMP/OMPIRBuilder.h | 9 +-
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 54 +++++--
.../Frontend/OpenMPIRBuilderTest.cpp | 84 +++++++++-
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 14 +-
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 44 ++++-
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 54 +++++--
mlir/test/Dialect/OpenMP/invalid.mlir | 71 +++++++-
mlir/test/Dialect/OpenMP/ops.mlir | 21 +++
.../LLVMIR/openmp-iterator-multi-block.mlir | 152 ++++++++++++++++++
9 files changed, 460 insertions(+), 43 deletions(-)
create mode 100644 mlir/test/Target/LLVMIR/openmp-iterator-multi-block.mlir
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index c92756092af337..fd6cd2fb22885f 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -4291,10 +4291,11 @@ class OpenMPIRBuilder {
///
/// - The skeleton’s unconditional branch from the loop body is removed
/// before invoking \p BodyGen.
- /// - \p BodyGen may freely emit instructions and temporarily introduce
- /// control flow.
- /// - If the loop body does not end with a terminator after \p BodyGen
- /// returns, a branch to the latch is inserted to restore canonical form.
+ /// - \p BodyGen may freely emit instructions and introduce control flow
+ /// within the loop body.
+ /// - If the body leaves exactly one block without a terminator after
+ /// \p BodyGen returns, a branch to the latch is inserted there to restore
+ /// canonical form.
///
/// \param Loc The location where the iterator modifier was encountered.
/// \param TripCount Number of loop iterations.
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 3a8125b835e53d..243c06c297f14a 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -12896,22 +12896,52 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createIteratorLoop(
T->eraseFromParent();
InsertPointTy BodyIP = CLI->getBodyIP();
+ // Blocks numbered before BodyGen, other than the body and latch, are outside
+ // the loop body.
+ unsigned FirstBodyGenBlock = F->getMaxBlockNumber();
+ [[maybe_unused]] unsigned BlockNumberEpoch = F->getBlockNumberEpoch();
if (llvm::Error Err = BodyGen(BodyIP, CLI->getIndVar()))
return Err;
+ assert(F->getBlockNumberEpoch() == BlockNumberEpoch &&
+ "iterator bodygen must not renumber blocks");
- // Body must either fallthrough to the latch or branch directly to it.
- if (Instruction *BodyTerminator = CLI->getBody()->getTerminatorOrNull()) {
- auto *BodyBr = dyn_cast<UncondBrInst>(BodyTerminator);
- if (!BodyBr || BodyBr->getSuccessor() != CLI->getLatch()) {
- return make_error<StringError>(
- "iterator bodygen must terminate the canonical body with an "
- "unconditional branch to the loop latch",
- inconvertibleErrorCode());
+ // The body may span several blocks. Branch its single unterminated block to
+ // the latch; otherwise some block must already branch there.
+ BasicBlock *Latch = CLI->getLatch();
+ BasicBlock *OpenBB = nullptr;
+ bool ReachesLatch = false;
+ SmallVector<BasicBlock *> Worklist{CLI->getBody()};
+ SmallPtrSet<BasicBlock *, 8> Visited{CLI->getBody()};
+ while (!Worklist.empty()) {
+ BasicBlock *BB = Worklist.pop_back_val();
+ if (!BB->hasTerminator()) {
+ if (OpenBB)
+ return make_error<StringError>(
+ "iterator bodygen must leave at most one unterminated block",
+ inconvertibleErrorCode());
+ OpenBB = BB;
+ continue;
}
- } else {
- // Ensure we end the loop body by jumping to the latch.
- Builder.SetInsertPoint(CLI->getBody());
- Builder.CreateBr(CLI->getLatch());
+ for (BasicBlock *Succ : successors(BB)) {
+ if (Succ == Latch) {
+ ReachesLatch = true;
+ continue;
+ }
+ if (Succ->getNumber() < FirstBodyGenBlock && Succ != CLI->getBody())
+ return make_error<StringError>(
+ "iterator bodygen must not branch out of the loop body",
+ inconvertibleErrorCode());
+ if (Visited.insert(Succ).second)
+ Worklist.push_back(Succ);
+ }
+ }
+
+ if (OpenBB) {
+ Builder.SetInsertPoint(OpenBB);
+ Builder.CreateBr(Latch);
+ } else if (!ReachesLatch) {
+ return make_error<StringError>("iterator bodygen must reach the loop latch",
+ inconvertibleErrorCode());
}
// Link After -> ContBB
diff --git a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
index 901d7e46d5d3e6..a4a5d541b1cabe 100644
--- a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
+++ b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
@@ -8020,7 +8020,89 @@ TEST_F(OpenMPIRBuilderTest, CreateIteratorLoopInvalidLoopBody) {
OpenMPIRBuilder::LocationDescription Loc(Builder.saveIP(), DL);
OpenMPIRBuilder::InsertPointOrErrorTy AfterIP = OMPBuilder.createIteratorLoop(
Loc, Builder.getInt64(4), BodyGenCB, "iterator");
- ASSERT_TRUE(errorToBool(AfterIP.takeError()));
+ EXPECT_EQ(toString(AfterIP.takeError()),
+ "iterator bodygen must reach the loop latch");
+}
+
+TEST_F(OpenMPIRBuilderTest, CreateIteratorLoopMultiBlockBody) {
+ using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
+ OpenMPIRBuilder OMPBuilder(*M);
+ OMPBuilder.initialize();
+ F->setName("func");
+ IRBuilder<> Builder(BB);
+
+ BasicBlock *Merge = nullptr;
+ auto BodyGenCB = [&](InsertPointTy BodyIP, Value *LinearIV) -> Error {
+ Builder.restoreIP(BodyIP);
+ Value *IsZero = Builder.CreateICmpEQ(LinearIV, Builder.getInt64(0));
+ BasicBlock *Then = BasicBlock::Create(Ctx, "iterator.then", F);
+ BasicBlock *Else = BasicBlock::Create(Ctx, "iterator.else", F);
+ Merge = BasicBlock::Create(Ctx, "iterator.merge", F);
+ Builder.CreateCondBr(IsZero, Then, Else);
+ Builder.SetInsertPoint(Then);
+ Builder.CreateBr(Merge);
+ Builder.SetInsertPoint(Else);
+ Builder.CreateBr(Merge);
+ Builder.SetInsertPoint(Merge);
+ Builder.CreateAdd(LinearIV, Builder.getInt64(1));
+ return Error::success();
+ };
+
+ OpenMPIRBuilder::LocationDescription Loc(Builder.saveIP(), DL);
+ ASSERT_EXPECTED_INIT(InsertPointTy, AfterIP,
+ OMPBuilder.createIteratorLoop(Loc, Builder.getInt64(4),
+ BodyGenCB, "iterator"));
+ Builder.restoreIP(AfterIP);
+ Builder.CreateRetVoid();
+
+ auto *MergeBr = dyn_cast<UncondBrInst>(Merge->getTerminator());
+ ASSERT_NE(MergeBr, nullptr);
+ EXPECT_EQ(MergeBr->getSuccessor()->getName(), "omp_iterator.inc");
+ EXPECT_FALSE(verifyFunction(*F, &errs()));
+}
+
+TEST_F(OpenMPIRBuilderTest, CreateIteratorLoopInvalidMultiBlockBody) {
+ using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
+ OpenMPIRBuilder OMPBuilder(*M);
+ OMPBuilder.initialize();
+ F->setName("func");
+ IRBuilder<> Builder(BB);
+
+ // Two blocks are left open, so the latch branch has no unique home.
+ auto TwoOpenCB = [&](InsertPointTy BodyIP, Value *LinearIV) -> Error {
+ Builder.restoreIP(BodyIP);
+ Value *IsZero = Builder.CreateICmpEQ(LinearIV, Builder.getInt64(0));
+ BasicBlock *Then = BasicBlock::Create(Ctx, "iterator.then", F);
+ BasicBlock *Else = BasicBlock::Create(Ctx, "iterator.else", F);
+ Builder.CreateCondBr(IsZero, Then, Else);
+ return Error::success();
+ };
+ OpenMPIRBuilder::LocationDescription Loc(Builder.saveIP(), DL);
+ OpenMPIRBuilder::InsertPointOrErrorTy AfterIP = OMPBuilder.createIteratorLoop(
+ Loc, Builder.getInt64(4), TwoOpenCB, "iterator");
+ EXPECT_EQ(toString(AfterIP.takeError()),
+ "iterator bodygen must leave at most one unterminated block");
+}
+
+TEST_F(OpenMPIRBuilderTest, CreateIteratorLoopBodyEscapes) {
+ using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
+ OpenMPIRBuilder OMPBuilder(*M);
+ OMPBuilder.initialize();
+ F->setName("func");
+ IRBuilder<> Builder(BB);
+
+ // An unterminated block outside the loop must not receive the latch branch.
+ BasicBlock *Outside = BasicBlock::Create(Ctx, "outside", F);
+ auto EscapeCB = [&](InsertPointTy BodyIP, Value *LinearIV) -> Error {
+ Builder.restoreIP(BodyIP);
+ Builder.CreateBr(Outside);
+ return Error::success();
+ };
+ OpenMPIRBuilder::LocationDescription Loc(Builder.saveIP(), DL);
+ OpenMPIRBuilder::InsertPointOrErrorTy AfterIP = OMPBuilder.createIteratorLoop(
+ Loc, Builder.getInt64(4), EscapeCB, "iterator");
+ EXPECT_EQ(toString(AfterIP.takeError()),
+ "iterator bodygen must not branch out of the loop body");
}
TEST_F(OpenMPIRBuilderTest, CreateTaskgroup) {
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 6bc99cf7656992..a839287121396d 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -2727,7 +2727,7 @@ def AffinityEntryOp : OpenMP_Op<"affinity_entry", [Pure]> {
// Iterator Op
//===----------------------------------------------------------------------===//
-def IteratorOp : OpenMP_Op<"iterator", [AttrSizedOperandSegments, SingleBlock],
+def IteratorOp : OpenMP_Op<"iterator", [AttrSizedOperandSegments],
clauses = [OpenMP_LoopRelatedClause]> {
let summary = "OpenMP iterator modifier";
let description = [{
@@ -2736,6 +2736,9 @@ def IteratorOp : OpenMP_Op<"iterator", [AttrSizedOperandSegments, SingleBlock],
can be directly consumed by OpenMP clauses that accept iterator modifiers,
such as `affinity`, `map`, `to`, `from`, or `depend`.
+ The region may contain multiple blocks. Exactly one of them is terminated
+ by `omp.yield`; every other block branches to a block of the region.
+
Example:
%it = omp.iterator(%i, %j) =
(%lb_i to %ub_i step %st_i, %lb_j to %ub_j step %st_j) {
@@ -2744,7 +2747,7 @@ def IteratorOp : OpenMP_Op<"iterator", [AttrSizedOperandSegments, SingleBlock],
} -> !omp.iterated<ptr>
}];
- let regions = (region SizedRegion<1>:$region);
+ let regions = (region MinSizedRegion<1>:$region);
let results = (outs OpenMP_IteratedType:$iterated);
let assemblyFormat = [{
@@ -2754,7 +2757,14 @@ def IteratorOp : OpenMP_Op<"iterator", [AttrSizedOperandSegments, SingleBlock],
(`inclusive` $loop_inclusive^)? `->` qualified(type($iterated)) attr-dict
}];
+ let extraClassDeclaration = [{
+ /// Returns the region's unique `omp.yield`, or null if the region has none
+ /// or more than one.
+ YieldOp getYieldOp();
+ }] # clausesExtraClassDeclaration;
+
let hasVerifier = 1;
+ let hasRegionVerifier = 1;
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index c6a424bd932146..d3a7fb47b99179 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2613,10 +2613,11 @@ static LogicalResult verifyMapClause(Operation *op, OperandRange mapVars,
"'omp.iterator' ops";
// Check that the iterator body yields a value defined by omp.map.info.
- auto yieldOp =
- cast<mlir::omp::YieldOp>(iterOp.getRegion().front().getTerminator());
+ mlir::omp::YieldOp yieldOp = iterOp.getYieldOp();
auto yieldedMapInfo =
- yieldOp.getResults()[0].getDefiningOp<mlir::omp::MapInfoOp>();
+ yieldOp && yieldOp.getNumOperands() == 1
+ ? yieldOp.getResults()[0].getDefiningOp<mlir::omp::MapInfoOp>()
+ : nullptr;
if (!yieldedMapInfo)
return op->emitOpError() << "'map_iterated' iterator body must yield "
"a value defined by 'omp.map.info'";
@@ -5742,6 +5743,20 @@ static void printIteratorHeader(OpAsmPrinter &p, Operation *op, Region ®ion,
/*printBlockTerminators=*/true);
}
+YieldOp IteratorOp::getYieldOp() {
+ YieldOp yield;
+ for (Block &block : getRegion()) {
+ auto blockYield =
+ dyn_cast_if_present<YieldOp>(block.empty() ? nullptr : &block.back());
+ if (!blockYield)
+ continue;
+ if (yield)
+ return {};
+ yield = blockYield;
+ }
+ return yield;
+}
+
LogicalResult IteratorOp::verify() {
auto iteratedTy = llvm::dyn_cast<omp::IteratedType>(getIterated().getType());
if (!iteratedTy)
@@ -5771,11 +5786,15 @@ LogicalResult IteratorOp::verify() {
"greater than or equal to upper bound";
}
- Block &b = getRegion().front();
- auto yield = llvm::dyn_cast<omp::YieldOp>(b.getTerminator());
-
- if (!yield)
+ YieldOp yield = getYieldOp();
+ if (!yield) {
+ bool hasYield = llvm::any_of(getRegion(), [](Block &block) {
+ return !block.empty() && isa<YieldOp>(block.back());
+ });
+ if (hasYield)
+ return emitOpError() << "region must contain exactly one omp.yield";
return emitOpError() << "region must be terminated by omp.yield";
+ }
if (yield.getNumOperands() != 1)
return emitOpError()
@@ -5792,6 +5811,17 @@ LogicalResult IteratorOp::verify() {
return success();
}
+LogicalResult IteratorOp::verifyRegions() {
+ // Translation routes only the omp.yield block back to the iterator loop.
+ for (Block &block : getRegion()) {
+ Operation *terminator = block.getTerminator();
+ if (block.getNumSuccessors() == 0 && !isa<YieldOp>(terminator))
+ return mlir::emitError(terminator->getLoc())
+ << "expected exit block terminator to be an `omp.yield` op.";
+ }
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// GroupprivateOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index fce2096a0ab3fe..109460502d6f11 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -3032,11 +3032,10 @@ static void fillAffinityLocators(Operation::operand_range affinityVars,
}
}
-static mlir::LogicalResult
-convertIteratorRegion(llvm::Value *linearIV, IteratorInfo &iterInfo,
- mlir::Block &iteratorRegionBlock,
- llvm::IRBuilderBase &builder,
- LLVM::ModuleTranslation &moduleTranslation) {
+static mlir::LogicalResult convertIteratorRegion(
+ llvm::Value *linearIV, IteratorInfo &iterInfo, mlir::Region &iteratorRegion,
+ llvm::IRBuilderBase &builder, LLVM::ModuleTranslation &moduleTranslation) {
+ mlir::Block &entryBlock = iteratorRegion.front();
llvm::Value *tmp = linearIV;
for (int d = (int)iterInfo.getDims() - 1; d >= 0; --d) {
llvm::Value *trip = iterInfo.getTrips()[d];
@@ -3050,16 +3049,41 @@ convertIteratorRegion(llvm::Value *linearIV, IteratorInfo &iterInfo,
iterInfo.getLowerBounds()[d],
builder.CreateMul(idx, iterInfo.getSteps()[d]), "omp.it.phys_iv");
- moduleTranslation.mapValue(iteratorRegionBlock.getArgument(d), physIV);
+ moduleTranslation.mapValue(entryBlock.getArgument(d), physIV);
}
// Translate the iterator region into the loop body.
- moduleTranslation.mapBlock(&iteratorRegionBlock, builder.GetInsertBlock());
- if (mlir::failed(moduleTranslation.convertBlock(iteratorRegionBlock,
- /*ignoreArguments=*/true,
- builder))) {
- return mlir::failure();
+ if (iteratorRegion.hasOneBlock()) {
+ moduleTranslation.mapBlock(&entryBlock, builder.GetInsertBlock());
+ return moduleTranslation.convertBlock(entryBlock, /*ignoreArguments=*/true,
+ builder);
+ }
+
+ // Unlike convertOmpOpRegions, do not forward the yielded value through a PHI:
+ // it may have no LLVM type (e.g. `!omp.affinity_entry_ty`). The unique
+ // `omp.yield` block is the continuation's only predecessor, so callers can
+ // look up the yielded value directly.
+ llvm::BasicBlock *continuation =
+ splitBB(builder, true, "omp.iterator.region.cont");
+ llvm::Instruction *sourceBranch = builder.GetInsertBlock()->getTerminator();
+ for (mlir::Block &block : iteratorRegion)
+ moduleTranslation.mapBlock(
+ &block,
+ llvm::BasicBlock::Create(builder.getContext(), "omp.iterator.region",
+ continuation->getParent(), continuation));
+ sourceBranch->setSuccessor(0, moduleTranslation.lookupBlock(&entryBlock));
+
+ for (mlir::Block *block : getBlocksSortedByDominance(iteratorRegion)) {
+ llvm::IRBuilderBase::InsertPointGuard guard(builder);
+ if (failed(moduleTranslation.convertBlock(*block, block->isEntryBlock(),
+ builder)))
+ return mlir::failure();
+ if (isa<omp::YieldOp>(block->getTerminator()))
+ builder.CreateBr(continuation);
}
+ LLVM::detail::connectPHINodes(iteratorRegion, moduleTranslation);
+
+ builder.SetInsertPoint(continuation);
return mlir::success();
}
@@ -3072,7 +3096,6 @@ fillIteratorLoop(mlir::omp::IteratorOp itersOp, llvm::IRBuilderBase &builder,
IteratorInfo &iterInfo, llvm::StringRef loopName,
IteratorStoreEntryTy genStoreEntry) {
mlir::Region &itersRegion = itersOp.getRegion();
- mlir::Block &iteratorRegionBlock = itersRegion.front();
llvm::OpenMPIRBuilder::LocationDescription loc(builder);
@@ -3081,14 +3104,13 @@ fillIteratorLoop(mlir::omp::IteratorOp itersOp, llvm::IRBuilderBase &builder,
llvm::IRBuilderBase::InsertPointGuard guard(builder);
builder.restoreIP(bodyIP);
- if (failed(convertIteratorRegion(linearIV, iterInfo, iteratorRegionBlock,
- builder, moduleTranslation))) {
+ if (failed(convertIteratorRegion(linearIV, iterInfo, itersRegion, builder,
+ moduleTranslation))) {
return llvm::make_error<llvm::StringError>(
"failed to convert iterator region", llvm::inconvertibleErrorCode());
}
- auto yield =
- mlir::dyn_cast<mlir::omp::YieldOp>(iteratorRegionBlock.getTerminator());
+ mlir::omp::YieldOp yield = itersOp.getYieldOp();
assert(yield && yield.getResults().size() == 1 &&
"expect omp.yield in iterator region to have one result");
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index c3fe4f60e41667..1dfd23ef031d8d 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -4939,7 +4939,7 @@ func.func @iterator_missing_yield(%lb : index, %ub : index, %st : index) {
// -----
func.func @iterator_empty_body(%lb : index, %ub : index, %st : index) {
- // expected-error at +1 {{expects a non-empty block}}
+ // expected-error at +1 {{region must be terminated by omp.yield}}
%0 = omp.iterator(%i: index) = (%lb to %ub step %st) {
} -> !omp.iterated<index>
return
@@ -4967,6 +4967,60 @@ func.func @iterator_yield_type_mismatch(%lb : index, %ub : index, %st : index) {
// -----
+func.func @iterator_multi_block_yield_type_mismatch(%lb : index, %ub : index, %st : index) {
+ // expected-error at +1 {{omp.iterated element type ('i64') does not match omp.yield operand type ('index')}}
+ %0 = omp.iterator(%i: index) = (%lb to %ub step %st) {
+ cf.br ^bb1
+ ^bb1:
+ omp.yield(%i : index)
+ } -> !omp.iterated<i64>
+ return
+}
+
+// -----
+
+func.func @iterator_two_yields(%lb : index, %ub : index, %st : index, %c : i1) {
+ // expected-error at +1 {{region must contain exactly one omp.yield}}
+ %0 = omp.iterator(%i: index) = (%lb to %ub step %st) {
+ cf.cond_br %c, ^bb1, ^bb2
+ ^bb1:
+ omp.yield(%i : index)
+ ^bb2:
+ omp.yield(%i : index)
+ } -> !omp.iterated<index>
+ return
+}
+
+// -----
+
+func.func @iterator_return_exit(%lb : index, %ub : index, %st : index, %c : i1) {
+ %0 = omp.iterator(%i: index) = (%lb to %ub step %st) {
+ cf.cond_br %c, ^bb1, ^bb2
+ ^bb1:
+ omp.yield(%i : index)
+ ^bb2:
+ // expected-error @below {{expected exit block terminator to be an `omp.yield` op.}}
+ llvm.return
+ } -> !omp.iterated<index>
+ return
+}
+
+// -----
+
+func.func @iterator_terminator_exit(%lb : index, %ub : index, %st : index, %c : i1) {
+ %0 = omp.iterator(%i: index) = (%lb to %ub step %st) {
+ cf.cond_br %c, ^bb1, ^bb2
+ ^bb1:
+ omp.yield(%i : index)
+ ^bb2:
+ // expected-error @below {{expected exit block terminator to be an `omp.yield` op.}}
+ omp.terminator
+ } -> !omp.iterated<index>
+ return
+}
+
+// -----
+
func.func @map_iterated_not_iterator(%it : !omp.iterated<!llvm.ptr>) {
// expected-error @below {{'omp.target_update' op 'map_iterated' arguments must be defined by 'omp.iterator' ops}}
omp.target_update map_iterated(%it : !omp.iterated<!llvm.ptr>)
@@ -4987,6 +5041,21 @@ func.func @map_iterated_yield_not_map_info(%lb : index, %ub : index, %st : index
// -----
+func.func @map_iterated_multi_block_yield_not_map_info(%lb : index, %ub : index,
+ %st : index,
+ %addr : !llvm.ptr) {
+ %it = omp.iterator(%iv: index) = (%lb to %ub step %st) {
+ cf.br ^bb1
+ ^bb1:
+ omp.yield(%addr : !llvm.ptr)
+ } -> !omp.iterated<!llvm.ptr>
+ // expected-error @below {{'omp.target_enter_data' op 'map_iterated' iterator body must yield a value defined by 'omp.map.info'}}
+ omp.target_enter_data map_iterated(%it : !omp.iterated<!llvm.ptr>) {}
+ return
+}
+
+// -----
+
func.func @map_iterated_yield_not_map_info_exit(%lb : index, %ub : index, %st : index,
%addr : !llvm.ptr) {
%it = omp.iterator(%iv: index) = (%lb to %ub step %st) {
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 9371d3599b72ae..7416c3da3b3951 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -4270,6 +4270,27 @@ func.func @omp_iterator_negative_step(%s2 : !llvm.struct<(ptr, i64)>) -> () {
return
}
+// CHECK-LABEL: func.func @omp_iterator_multi_block
+func.func @omp_iterator_multi_block(%a : !llvm.ptr, %b : !llvm.ptr) -> () {
+ // CHECK: omp.iterator(%[[IV:.*]]: index) = (%{{.*}} to %{{.*}} step %{{.*}}) {
+ // CHECK: %[[C:.*]] = arith.cmpi eq, %[[IV]], %{{.*}} : index
+ // CHECK: cf.cond_br %[[C]], ^[[BB1:.*]](%{{.*}} : !llvm.ptr), ^[[BB1]](%{{.*}} : !llvm.ptr)
+ // CHECK: ^[[BB1]](%[[P:.*]]: !llvm.ptr):
+ // CHECK: omp.yield(%[[P]] : !llvm.ptr)
+ // CHECK: } -> !omp.iterated<!llvm.ptr>
+ %lb = arith.constant 1 : index
+ %ub = arith.constant 4 : index
+ %st = arith.constant 1 : index
+
+ %0 = omp.iterator(%iv: index) = (%lb to %ub step %st) {
+ %c = arith.cmpi eq, %iv, %lb : index
+ cf.cond_br %c, ^bb1(%a : !llvm.ptr), ^bb1(%b : !llvm.ptr)
+ ^bb1(%p : !llvm.ptr):
+ omp.yield(%p : !llvm.ptr)
+ } -> !omp.iterated<!llvm.ptr>
+ return
+}
+
// CHECK-LABEL: func.func @omp_task_affinity_iterator_1d
func.func @omp_task_affinity_iterator_1d(%lb : index, %ub : index, %step : index,
%addr : !llvm.ptr, %len : i64) -> () {
diff --git a/mlir/test/Target/LLVMIR/openmp-iterator-multi-block.mlir b/mlir/test/Target/LLVMIR/openmp-iterator-multi-block.mlir
new file mode 100644
index 00000000000000..80a817ef8fb728
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-iterator-multi-block.mlir
@@ -0,0 +1,152 @@
+// RUN: mlir-translate --mlir-to-llvmir %s | FileCheck %s
+
+// An omp.iterator region may span several blocks. Each entry stores the value
+// yielded by the region's single omp.yield.
+
+// The yielded address is selected by a branch.
+llvm.func @depend_diamond(%a : !llvm.ptr, %b : !llvm.ptr) {
+ %c1 = llvm.mlir.constant(1 : i64) : i64
+ %c4 = llvm.mlir.constant(4 : i64) : i64
+ %it = omp.iterator(%iv: i64) = (%c1 to %c4 step %c1) {
+ %first = llvm.icmp "eq" %iv, %c1 : i64
+ llvm.cond_br %first, ^bb1(%a : !llvm.ptr), ^bb1(%b : !llvm.ptr)
+ ^bb1(%p : !llvm.ptr):
+ omp.yield(%p : !llvm.ptr)
+ } -> !omp.iterated<!llvm.ptr>
+ omp.task depend(taskdependin -> %it : !omp.iterated<!llvm.ptr>) {
+ omp.terminator
+ }
+ llvm.return
+}
+
+// CHECK-LABEL: define void @depend_diamond
+// CHECK-SAME: (ptr %[[A:[0-9]+]], ptr %[[B:[0-9]+]])
+// CHECK: %[[DEPS:.*]] = tail call ptr @malloc(i64 80)
+// CHECK: omp_dep_iterator.body:
+// CHECK: %[[IV:.*]] = add i64 1, %{{.*}}
+// CHECK: br label %[[ENTRY:omp.iterator.region]]
+// CHECK: [[ENTRY]]:
+// CHECK: %[[FIRST:.*]] = icmp eq i64 %[[IV]], 1
+// CHECK: br i1 %[[FIRST]], label %[[MERGE:omp.iterator.region[0-9]+]], label %[[ELSE:omp.iterator.region[0-9]+]]
+// CHECK: [[MERGE]]:
+// CHECK: %[[P:.*]] = phi ptr [ %[[BV:.*]], %[[ELSE]] ], [ %[[A]], %[[ENTRY]] ]
+// CHECK: br label %omp.iterator.region.cont
+// CHECK: [[ELSE]]:
+// CHECK: %[[BV]] = phi ptr [ %[[B]], %[[ENTRY]] ]
+// CHECK: omp.iterator.region.cont:
+// CHECK: %[[ADDR:.*]] = ptrtoint ptr %[[P]] to i64
+// CHECK: store i64 %[[ADDR]], ptr
+// CHECK: br label %omp_dep_iterator.inc
+// CHECK: call i32 @__kmpc_omp_task_with_deps({{.*}}, i32 4, ptr %[[DEPS]], i32 0, ptr null)
+
+// The affinity length is defined before the branch, the address after it.
+llvm.func @affinity_split(%a : !llvm.ptr, %b : !llvm.ptr) {
+ %c1 = llvm.mlir.constant(1 : i64) : i64
+ %c3 = llvm.mlir.constant(3 : i64) : i64
+ %it = omp.iterator(%iv: i64) = (%c1 to %c3 step %c1) {
+ %len = llvm.mul %iv, %c3 : i64
+ %first = llvm.icmp "eq" %iv, %c1 : i64
+ llvm.cond_br %first, ^bb1, ^bb2
+ ^bb1:
+ llvm.br ^bb3(%a : !llvm.ptr)
+ ^bb2:
+ llvm.br ^bb3(%b : !llvm.ptr)
+ ^bb3(%p : !llvm.ptr):
+ %e = omp.affinity_entry %p, %len
+ : (!llvm.ptr, i64) -> !omp.affinity_entry_ty<!llvm.ptr, i64>
+ omp.yield(%e : !omp.affinity_entry_ty<!llvm.ptr, i64>)
+ } -> !omp.iterated<!omp.affinity_entry_ty<!llvm.ptr, i64>>
+ omp.task affinity(%it : !omp.iterated<!omp.affinity_entry_ty<!llvm.ptr, i64>>) {
+ omp.terminator
+ }
+ llvm.return
+}
+
+// CHECK-LABEL: define void @affinity_split
+// CHECK-SAME: (ptr %[[A:[0-9]+]], ptr %[[B:[0-9]+]])
+// CHECK: omp_iterator.body:
+// CHECK: %[[IV:.*]] = add i64 1, %{{.*}}
+// CHECK: omp.iterator.region:
+// CHECK: %[[LEN:.*]] = mul i64 %[[IV]], 3
+// CHECK: %[[P:.*]] = phi ptr [ %[[B]], %{{.*}} ], [ %[[A]], %{{.*}} ]
+// CHECK: omp.iterator.region.cont:
+// CHECK: %[[ADDR:.*]] = ptrtoint ptr %[[P]] to i64
+// CHECK: store i64 %[[ADDR]], ptr
+// CHECK: store i64 %[[LEN]], ptr
+// CHECK: br label %omp_iterator.inc
+// CHECK: call i32 @__kmpc_omp_reg_task_with_affinity({{.*}}, i32 3, ptr
+
+// A loop inside the region computes the yielded address.
+llvm.func @depend_inner_loop(%addr : !llvm.ptr) {
+ %c0 = llvm.mlir.constant(0 : i64) : i64
+ %c1 = llvm.mlir.constant(1 : i64) : i64
+ %c3 = llvm.mlir.constant(3 : i64) : i64
+ %it = omp.iterator(%iv: i64) = (%c0 to %c3 step %c1) {
+ llvm.br ^loop(%c0, %c0 : i64, i64)
+ ^loop(%k : i64, %acc : i64):
+ %done = llvm.icmp "sge" %k, %iv : i64
+ llvm.cond_br %done, ^exit, ^body
+ ^body:
+ %acc2 = llvm.add %acc, %k : i64
+ %k2 = llvm.add %k, %c1 : i64
+ llvm.br ^loop(%k2, %acc2 : i64, i64)
+ ^exit:
+ %p = llvm.getelementptr %addr[%acc] : (!llvm.ptr, i64) -> !llvm.ptr, i8
+ omp.yield(%p : !llvm.ptr)
+ } -> !omp.iterated<!llvm.ptr>
+ omp.task depend(taskdependin -> %it : !omp.iterated<!llvm.ptr>) {
+ omp.terminator
+ }
+ llvm.return
+}
+
+// CHECK-LABEL: define void @depend_inner_loop
+// CHECK-SAME: (ptr %[[ADDR:[0-9]+]])
+// CHECK: omp_dep_iterator.body:
+// CHECK: %[[IV:.*]] = add i64 0, %{{.*}}
+// CHECK: [[LOOP:omp.iterator.region[0-9]+]]:
+// CHECK: %[[K:.*]] = phi i64 [ %[[K2:.*]], %[[BODY:omp.iterator.region[0-9]+]] ], [ 0, %omp.iterator.region ]
+// CHECK: %[[ACC:.*]] = phi i64 [ %[[ACC2:.*]], %[[BODY]] ], [ 0, %omp.iterator.region ]
+// CHECK: %[[DONE:.*]] = icmp sge i64 %[[K]], %[[IV]]
+// CHECK: br i1 %[[DONE]], label %[[EXIT:omp.iterator.region[0-9]+]], label %[[BODY]]
+// CHECK: [[BODY]]:
+// CHECK: %[[ACC2]] = add i64 %[[ACC]], %[[K]]
+// CHECK: %[[K2]] = add i64 %[[K]], 1
+// CHECK: br label %[[LOOP]]
+// CHECK: [[EXIT]]:
+// CHECK: %[[P:.*]] = getelementptr i8, ptr %[[ADDR]], i64 %[[ACC]]
+// CHECK: omp.iterator.region.cont:
+// CHECK: %[[PI:.*]] = ptrtoint ptr %[[P]] to i64
+// CHECK: store i64 %[[PI]], ptr
+// CHECK: call i32 @__kmpc_omp_task_with_deps({{.*}}, i32 4, ptr
+
+// Both induction values are mapped in their declared order.
+llvm.func @depend_2d(%a : !llvm.ptr, %b : !llvm.ptr) {
+ %c1 = llvm.mlir.constant(1 : i64) : i64
+ %c2 = llvm.mlir.constant(2 : i64) : i64
+ %c3 = llvm.mlir.constant(3 : i64) : i64
+ %it = omp.iterator(%i: i64, %j: i64) = (%c1 to %c2 step %c1, %c1 to %c3 step %c1) {
+ %lt = llvm.icmp "slt" %i, %j : i64
+ llvm.cond_br %lt, ^bb1(%a : !llvm.ptr), ^bb1(%b : !llvm.ptr)
+ ^bb1(%p : !llvm.ptr):
+ omp.yield(%p : !llvm.ptr)
+ } -> !omp.iterated<!llvm.ptr>
+ omp.task depend(taskdependin -> %it : !omp.iterated<!llvm.ptr>) {
+ omp.terminator
+ }
+ llvm.return
+}
+
+// CHECK-LABEL: define void @depend_2d
+// CHECK-SAME: (ptr %[[A:[0-9]+]], ptr %[[B:[0-9]+]])
+// CHECK: omp_dep_iterator.body:
+// CHECK: %[[J:.*]] = add i64 1, %{{.*}}
+// CHECK: %[[I:.*]] = add i64 1, %{{.*}}
+// CHECK: omp.iterator.region:
+// CHECK: icmp slt i64 %[[I]], %[[J]]
+// CHECK: %[[P:.*]] = phi ptr [ %{{.*}}, %{{.*}} ], [ %[[A]], %omp.iterator.region ]
+// CHECK: phi ptr [ %[[B]], %omp.iterator.region ]
+// CHECK: omp.iterator.region.cont:
+// CHECK: %[[PI:.*]] = ptrtoint ptr %[[P]] to i64
+// CHECK: store i64 %[[PI]], ptr
+// CHECK: call i32 @__kmpc_omp_task_with_deps({{.*}}, i32 6, ptr
>From 81ad16e5e25fdbd0fbefa0dc358a962a0a1312cf Mon Sep 17 00:00:00 2001
From: "Chi Chun, Chen" <chichun.chen at hpe.com>
Date: Tue, 6 Oct 2026 11:12:00 -0500
Subject: [PATCH 2/2] Count empty iterator ranges and keep induction types
Translation converted iterator bounds to i64 and mapped the rebuilt i64
induction value onto the region argument. An i32 induction value forwarded
to a block argument therefore produced an ill-typed PHI, and a single-block
region stored an i64 into an i32 slot. The count (ub - lb) / step + 1 also
gave dynamic empty ranges such as 3:2:2 one entry, so the locator was
evaluated.
Count each range with the canonical trip-count helper on operands widened
by one bit, so empty ranges in either direction count as zero. Keep bounds
in their declared type and rebuild the induction value in that type. The
verifier now requires one region argument per range, typed like its
bounds, which translation relies on.
---
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 10 ++
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 46 ++++----
mlir/test/Dialect/OpenMP/invalid.mlir | 21 ++++
.../LLVMIR/openmp-iterator-multi-block.mlir | 54 ++++++++++
.../Target/LLVMIR/openmp-iterator-wide.mlir | 102 ++++++++++++++++++
mlir/test/Target/LLVMIR/openmp-iterator.mlir | 67 ++++++++++--
6 files changed, 270 insertions(+), 30 deletions(-)
create mode 100644 mlir/test/Target/LLVMIR/openmp-iterator-wide.mlir
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index d3a7fb47b99179..a46ae84df1784a 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -5762,6 +5762,16 @@ LogicalResult IteratorOp::verify() {
if (!iteratedTy)
return emitOpError() << "result must be omp.iterated<entry_ty>";
+ Block::BlockArgListType ivs = getRegion().front().getArguments();
+ if (getLoopLowerBounds().size() != ivs.size())
+ return emitOpError() << "number of range arguments and IVs do not match";
+
+ for (auto [lb, iv] : llvm::zip_equal(getLoopLowerBounds(), ivs)) {
+ if (lb.getType() != iv.getType())
+ return emitOpError()
+ << "range argument type does not match corresponding IV type";
+ }
+
for (auto [lb, ub, step] : llvm::zip_equal(
getLoopLowerBounds(), getLoopUpperBounds(), getLoopSteps())) {
if (matchPattern(step, m_Zero()))
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 109460502d6f11..527f4cfb3b3023 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -2849,8 +2849,9 @@ class TaskContextStructManager {
/// mlir::omp::IteratorOp for lowering to LLVM IR.
///
/// It computes the per-dimension trip counts and the total linearized trip
-/// count, casted to i64. These are used to build a canonical loop and to
-/// reconstruct the physical induction variables inside the loop body.
+/// count as i64. Bounds and steps keep their declared type so that the
+/// physical induction variables can be reconstructed in that type inside the
+/// loop body.
class IteratorInfo {
private:
llvm::SmallVector<llvm::Value *> lowerBounds;
@@ -2860,18 +2861,6 @@ class IteratorInfo {
unsigned dims;
llvm::Value *totalTrips;
- llvm::Value *lookUpAsI64(mlir::Value val, const LLVM::ModuleTranslation &mt,
- llvm::IRBuilderBase &builder) {
- llvm::Value *v = mt.lookupValue(val);
- if (!v)
- return nullptr;
- if (v->getType()->isIntegerTy(64))
- return v;
- if (v->getType()->isIntegerTy())
- return builder.CreateSExtOrTrunc(v, builder.getInt64Ty());
- return nullptr;
- }
-
public:
IteratorInfo(mlir::omp::IteratorOp itersOp,
mlir::LLVM::ModuleTranslation &moduleTranslation,
@@ -2883,12 +2872,12 @@ class IteratorInfo {
trips.resize(dims);
for (unsigned d = 0; d < dims; ++d) {
- llvm::Value *lb = lookUpAsI64(itersOp.getLoopLowerBounds()[d],
- moduleTranslation, builder);
- llvm::Value *ub = lookUpAsI64(itersOp.getLoopUpperBounds()[d],
- moduleTranslation, builder);
+ llvm::Value *lb =
+ moduleTranslation.lookupValue(itersOp.getLoopLowerBounds()[d]);
+ llvm::Value *ub =
+ moduleTranslation.lookupValue(itersOp.getLoopUpperBounds()[d]);
llvm::Value *st =
- lookUpAsI64(itersOp.getLoopSteps()[d], moduleTranslation, builder);
+ moduleTranslation.lookupValue(itersOp.getLoopSteps()[d]);
assert(lb && ub && st &&
"Expect lowerBounds, upperBounds, and steps in IteratorOp");
assert((!llvm::isa<llvm::ConstantInt>(st) ||
@@ -2899,11 +2888,18 @@ class IteratorInfo {
upperBounds[d] = ub;
steps[d] = st;
- // trips = ((ub - lb) / step) + 1 (inclusive ub, assume positive step)
- llvm::Value *diff = builder.CreateSub(ub, lb);
- llvm::Value *div = builder.CreateSDiv(diff, st);
- trips[d] = builder.CreateAdd(
- div, llvm::ConstantInt::get(builder.getInt64Ty(), 1));
+ // Use a direction-aware count so an empty range contributes no entries.
+ // Widen by one bit so the span between the bounds cannot overflow.
+ llvm::Type *countTy =
+ builder.getIntNTy(lb->getType()->getIntegerBitWidth() + 1);
+ llvm::Value *start = builder.CreateSExt(lb, countTy);
+ llvm::Value *stop = builder.CreateSExt(ub, countTy);
+ llvm::Value *step = builder.CreateSExt(st, countTy);
+ llvm::Value *count =
+ moduleTranslation.getOpenMPBuilder()->calculateCanonicalLoopTripCount(
+ builder, start, stop, step, /*IsSigned=*/true,
+ /*InclusiveStop=*/true);
+ trips[d] = builder.CreateZExtOrTrunc(count, builder.getInt64Ty());
}
totalTrips = llvm::ConstantInt::get(builder.getInt64Ty(), 1);
@@ -3045,6 +3041,8 @@ static mlir::LogicalResult convertIteratorRegion(
tmp = builder.CreateUDiv(tmp, trip);
// physIV_d = lb_d + idx_d * step_d
+ idx =
+ builder.CreateZExtOrTrunc(idx, iterInfo.getLowerBounds()[d]->getType());
llvm::Value *physIV = builder.CreateAdd(
iterInfo.getLowerBounds()[d],
builder.CreateMul(idx, iterInfo.getSteps()[d]), "omp.it.phys_iv");
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 1dfd23ef031d8d..e3c32b0104a4ca 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -4886,6 +4886,27 @@ func.func @iterator_bad_result_type(%lb : index, %ub : index, %st : index) {
// -----
+func.func @iterator_iv_count_mismatch(%lb : i64, %ub : i64, %st : i64, %x : !llvm.ptr) {
+ // expected-error at +1 {{number of range arguments and IVs do not match}}
+ %0 = "omp.iterator"(%lb, %ub, %st) <{operandSegmentSizes = array<i32: 1, 1, 1>}> ({
+ omp.yield(%x : !llvm.ptr)
+ }) : (i64, i64, i64) -> !omp.iterated<!llvm.ptr>
+ return
+}
+
+// -----
+
+func.func @iterator_iv_type_mismatch(%lb : i32, %ub : i32, %st : i32, %x : !llvm.ptr) {
+ // expected-error at +1 {{range argument type does not match corresponding IV type}}
+ %0 = "omp.iterator"(%lb, %ub, %st) <{operandSegmentSizes = array<i32: 1, 1, 1>}> ({
+ ^bb0(%i: i64):
+ omp.yield(%x : !llvm.ptr)
+ }) : (i32, i32, i32) -> !omp.iterated<!llvm.ptr>
+ return
+}
+
+// -----
+
func.func @iterator_zero_step(%s2 : !llvm.struct<(ptr, i64)>) {
%lb = arith.constant 1 : index
%ub = arith.constant 4 : index
diff --git a/mlir/test/Target/LLVMIR/openmp-iterator-multi-block.mlir b/mlir/test/Target/LLVMIR/openmp-iterator-multi-block.mlir
index 80a817ef8fb728..3487897912eb51 100644
--- a/mlir/test/Target/LLVMIR/openmp-iterator-multi-block.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-iterator-multi-block.mlir
@@ -150,3 +150,57 @@ llvm.func @depend_2d(%a : !llvm.ptr, %b : !llvm.ptr) {
// CHECK: %[[PI:.*]] = ptrtoint ptr %[[P]] to i64
// CHECK: store i64 %[[PI]], ptr
// CHECK: call i32 @__kmpc_omp_task_with_deps({{.*}}, i32 6, ptr
+
+// A block argument fed by the induction value keeps its declared type.
+llvm.func @depend_narrow(%x : !llvm.ptr) {
+ %lo = llvm.mlir.constant(1 : i32) : i32
+ %hi = llvm.mlir.constant(2 : i32) : i32
+ %it = omp.iterator(%i: i32) = (%lo to %hi step %lo) {
+ llvm.br ^next(%i : i32)
+ ^next(%v: i32):
+ llvm.store %v, %x : i32, !llvm.ptr
+ omp.yield(%x : !llvm.ptr)
+ } -> !omp.iterated<!llvm.ptr>
+ omp.taskwait depend(taskdependin -> %it : !omp.iterated<!llvm.ptr>)
+ llvm.return
+}
+
+// CHECK-LABEL: define void @depend_narrow
+// CHECK: omp_dep_iterator.body:
+// CHECK: %[[IDX:.*]] = trunc i64 %{{.*}} to i32
+// CHECK: %[[OFF:.*]] = mul i32 %[[IDX]], 1
+// CHECK: %[[IV:.*]] = add i32 1, %[[OFF]]
+// CHECK: %[[V:.*]] = phi i32 [ %[[IV]], %omp.iterator.region ]
+// CHECK: store i32 %[[V]], ptr
+
+// An empty dynamic range, e.g. m = 2, must not evaluate the region.
+llvm.func @depend_empty_dynamic(%a : !llvm.ptr, %m : i32, %d : i32) {
+ %c2 = llvm.mlir.constant(2 : i32) : i32
+ %c3 = llvm.mlir.constant(3 : i32) : i32
+ %it = omp.iterator(%i: i32) = (%c3 to %m step %c2) {
+ %q = llvm.sdiv %i, %d : i32
+ llvm.br ^bb1(%q : i32)
+ ^bb1(%k : i32):
+ %p = llvm.getelementptr %a[%k] : (!llvm.ptr, i32) -> !llvm.ptr, i32
+ omp.yield(%p : !llvm.ptr)
+ } -> !omp.iterated<!llvm.ptr>
+ omp.task depend(taskdependin -> %it : !omp.iterated<!llvm.ptr>) {
+ omp.terminator
+ }
+ llvm.return
+}
+
+// CHECK-LABEL: define void @depend_empty_dynamic
+// CHECK-SAME: (ptr %{{.*}}, i32 %[[M:[0-9]+]], i32 %{{.*}})
+// CHECK: %[[STOP:.*]] = sext i32 %[[M]] to i33
+// CHECK: %[[LB:.*]] = select i1 false, i33 %[[STOP]], i33 3
+// CHECK: %[[UB:.*]] = select i1 false, i33 3, i33 %[[STOP]]
+// CHECK: %[[EMPTY:.*]] = icmp slt i33 %[[UB]], %[[LB]]
+// CHECK: %[[COUNT:.*]] = select i1 %[[EMPTY]], i33 0, i33 %{{.*}}
+// CHECK: %[[TRIPS:.*]] = zext i33 %[[COUNT]] to i64
+// CHECK: %[[TOTAL:.*]] = mul i64 1, %[[TRIPS]]
+// CHECK: omp_dep_iterator.cond:
+// CHECK: icmp ult i64 %omp_dep_iterator.iv, %[[TOTAL]]
+// CHECK: omp_dep_iterator.body:
+// CHECK: omp.iterator.region:
+// CHECK: sdiv i32 %omp.it.phys_iv,
diff --git a/mlir/test/Target/LLVMIR/openmp-iterator-wide.mlir b/mlir/test/Target/LLVMIR/openmp-iterator-wide.mlir
new file mode 100644
index 00000000000000..93791a3bbf0f53
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-iterator-wide.mlir
@@ -0,0 +1,102 @@
+// RUN: mlir-translate --mlir-to-llvmir %s | FileCheck %s
+// RUN: mlir-translate --mlir-to-llvmir %s | opt -passes=verify -disable-output
+
+// Bounds and induction values keep their declared width; only the trip count
+// is converted to i64.
+llvm.func @dynamic(%x: !llvm.ptr, %lb: i128, %ub: i128, %st: i128) {
+ %it = omp.iterator(%i: i128) = (%lb to %ub step %st) {
+ llvm.store %i, %x : i128, !llvm.ptr
+ omp.yield(%x : !llvm.ptr)
+ } -> !omp.iterated<!llvm.ptr>
+ omp.taskwait depend(taskdependin -> %it : !omp.iterated<!llvm.ptr>)
+ llvm.return
+}
+// CHECK-LABEL: define void @dynamic(
+// CHECK: sext i128 %{{.*}} to i129
+// CHECK: sext i128 %{{.*}} to i129
+// CHECK: sext i128 %{{.*}} to i129
+// CHECK: icmp slt i129
+// CHECK: sub nsw i129
+// CHECK: udiv i129
+// CHECK: select i1 %{{.*}}, i129 0, i129 %{{.*}}
+// CHECK: trunc i129 %{{.*}} to i64
+// CHECK: %[[IDX:.*]] = zext i64 %{{.*}} to i128
+// CHECK: %[[OFFSET:.*]] = mul i128 %[[IDX]], %{{.*}}
+// CHECK: %[[IV:.*]] = add i128 %{{.*}}, %[[OFFSET]]
+// CHECK: store i128 %[[IV]], ptr
+
+// The unsigned linear index must also convert to narrower region arguments.
+llvm.func @narrow(%x: !llvm.ptr, %lb: i32, %ub: i32, %st: i32) {
+ %it = omp.iterator(%i: i32) = (%lb to %ub step %st) {
+ llvm.store %i, %x : i32, !llvm.ptr
+ omp.yield(%x : !llvm.ptr)
+ } -> !omp.iterated<!llvm.ptr>
+ omp.taskwait depend(taskdependin -> %it : !omp.iterated<!llvm.ptr>)
+ llvm.return
+}
+// CHECK-LABEL: define void @narrow(
+// CHECK: sext i32 %{{.*}} to i33
+// CHECK: zext i33 %{{.*}} to i64
+// CHECK: %[[IDX:.*]] = trunc i64 %{{.*}} to i32
+// CHECK: %[[OFFSET:.*]] = mul i32 %[[IDX]], %{{.*}}
+// CHECK: %[[IV:.*]] = add i32 %{{.*}}, %[[OFFSET]]
+// CHECK: store i32 %[[IV]], ptr
+
+llvm.func @nonempty_positive(%x: !llvm.ptr) {
+ %lb = llvm.mlir.constant(18446744073709551617 : i128) : i128
+ %ub = llvm.mlir.constant(18446744073709551619 : i128) : i128
+ %st = llvm.mlir.constant(1 : i128) : i128
+ %it = omp.iterator(%i: i128) = (%lb to %ub step %st) {
+ llvm.store %i, %x : i128, !llvm.ptr
+ omp.yield(%x : !llvm.ptr)
+ } -> !omp.iterated<!llvm.ptr>
+ omp.taskwait depend(taskdependin -> %it : !omp.iterated<!llvm.ptr>)
+ llvm.return
+}
+// CHECK-LABEL: define void @nonempty_positive(
+// CHECK: icmp ult i64 %{{.*}}, 3
+// CHECK: zext i64 %{{.*}} to i128
+// CHECK: mul i128 %{{.*}}, 1
+// CHECK: add i128 18446744073709551617, %{{.*}}
+// CHECK: store i128
+// CHECK: call void @__kmpc_omp_taskwait_deps_51(
+// CHECK-SAME: i32 3, ptr %{{.*}}, i32 0, ptr null, i32 0)
+
+llvm.func @nonempty_negative(%x: !llvm.ptr) {
+ %lb = llvm.mlir.constant(18446744073709551619 : i128) : i128
+ %ub = llvm.mlir.constant(18446744073709551617 : i128) : i128
+ %st = llvm.mlir.constant(-1 : i128) : i128
+ %it = omp.iterator(%i: i128) = (%lb to %ub step %st) {
+ llvm.store %i, %x : i128, !llvm.ptr
+ omp.yield(%x : !llvm.ptr)
+ } -> !omp.iterated<!llvm.ptr>
+ omp.taskwait depend(taskdependin -> %it : !omp.iterated<!llvm.ptr>)
+ llvm.return
+}
+// CHECK-LABEL: define void @nonempty_negative(
+// CHECK: icmp ult i64 %{{.*}}, 3
+// CHECK: zext i64 %{{.*}} to i128
+// CHECK: mul i128 %{{.*}}, -1
+// CHECK: add i128 18446744073709551619, %{{.*}}
+// CHECK: store i128
+// CHECK: call void @__kmpc_omp_taskwait_deps_51(
+// CHECK-SAME: i32 3, ptr %{{.*}}, i32 0, ptr null, i32 0)
+
+llvm.func @wide_step(%x: !llvm.ptr) {
+ %lb = llvm.mlir.constant(0 : i128) : i128
+ %ub = llvm.mlir.constant(36893488147419103232 : i128) : i128
+ %st = llvm.mlir.constant(18446744073709551616 : i128) : i128
+ %it = omp.iterator(%i: i128) = (%lb to %ub step %st) {
+ llvm.store %i, %x : i128, !llvm.ptr
+ omp.yield(%x : !llvm.ptr)
+ } -> !omp.iterated<!llvm.ptr>
+ omp.taskwait depend(taskdependin -> %it : !omp.iterated<!llvm.ptr>)
+ llvm.return
+}
+// CHECK-LABEL: define void @wide_step(
+// CHECK: icmp ult i64 %{{.*}}, 3
+// CHECK: zext i64 %{{.*}} to i128
+// CHECK: mul i128 %{{.*}}, 18446744073709551616
+// CHECK: store i128
+// CHECK: call void @__kmpc_omp_taskwait_deps_51(
+// CHECK-SAME: i32 3, ptr %{{.*}}, i32 0, ptr null, i32 0)
diff --git a/mlir/test/Target/LLVMIR/openmp-iterator.mlir b/mlir/test/Target/LLVMIR/openmp-iterator.mlir
index 7fdc6f15b03115..ad8d603d7d266d 100644
--- a/mlir/test/Target/LLVMIR/openmp-iterator.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-iterator.mlir
@@ -260,9 +260,12 @@ llvm.func @task_affinity_iterator_dynamic_tripcount(
}
// CHECK-LABEL: define internal void @task_affinity_iterator_dynamic_tripcount
-// CHECK: [[DIFF:%.*]] = sub i64 {{.*}}, {{.*}}
-// CHECK: [[DIV:%.*]] = sdiv i64 [[DIFF]], {{.*}}
-// CHECK: [[TRIPS:%.*]] = add i64 [[DIV]], 1
+// CHECK: [[DIFF:%.*]] = sub nsw i65 {{.*}}, {{.*}}
+// CHECK: [[EMPTY:%.*]] = icmp slt i65 {{.*}}, {{.*}}
+// CHECK: [[DIV:%.*]] = udiv i65 [[DIFF]], {{.*}}
+// CHECK: [[COUNT:%.*]] = add i65 [[DIV]], 1
+// CHECK: [[CLAMPED:%.*]] = select i1 [[EMPTY]], i65 0, i65 [[COUNT]]
+// CHECK: [[TRIPS:%.*]] = trunc i65 [[CLAMPED]] to i64
// CHECK: [[SCALED:%.*]] = mul i64 1, [[TRIPS]]
// CHECK: [[AFFLIST:%.*]] = alloca { i64, i64, i32 }, i64 [[SCALED]]
@@ -409,9 +412,20 @@ llvm.func @omp_task_depend_iterator_dynamic(%addr : !llvm.ptr,
// CHECK-LABEL: define void @omp_task_depend_iterator_dynamic
//
// Tripcount computation from dynamic bounds
-// CHECK: %[[DIFF:.*]] = sub i64 %{{.*}}, %{{.*}}
-// CHECK: %[[DIV:.*]] = sdiv i64 %[[DIFF]], %{{.*}}
-// CHECK: %[[TRIPS:.*]] = add i64 %[[DIV]], 1
+// CHECK: %[[START:.*]] = sext i64 %{{.*}} to i65
+// CHECK: %[[STOP:.*]] = sext i64 %{{.*}} to i65
+// CHECK: %[[STEP:.*]] = sext i64 %{{.*}} to i65
+// CHECK: %[[NEG:.*]] = icmp slt i65 %[[STEP]], 0
+// CHECK: %[[MAG:.*]] = sub i65 0, %[[STEP]]
+// CHECK: %[[INCR:.*]] = select i1 %[[NEG]], i65 %[[MAG]], i65 %[[STEP]]
+// CHECK: %[[LOW:.*]] = select i1 %[[NEG]], i65 %[[STOP]], i65 %[[START]]
+// CHECK: %[[HIGH:.*]] = select i1 %[[NEG]], i65 %[[START]], i65 %[[STOP]]
+// CHECK: %[[DIFF:.*]] = sub nsw i65 %[[HIGH]], %[[LOW]]
+// CHECK: %[[EMPTY:.*]] = icmp slt i65 %[[HIGH]], %[[LOW]]
+// CHECK: %[[DIV:.*]] = udiv i65 %[[DIFF]], %[[INCR]]
+// CHECK: %[[COUNT:.*]] = add i65 %[[DIV]], 1
+// CHECK: %[[CLAMPED:.*]] = select i1 %[[EMPTY]], i65 0, i65 %[[COUNT]]
+// CHECK: %[[TRIPS:.*]] = trunc i65 %[[CLAMPED]] to i64
// CHECK: %[[SCALED:.*]] = mul i64 1, %[[TRIPS]]
// Dynamic total = 0 + scaled trip count
// CHECK: %[[TOTAL:.*]] = add i64 0, %[[SCALED]]
@@ -453,6 +467,47 @@ llvm.func @omp_task_depend_iterator_dynamic_mixed(%addr : !llvm.ptr,
// CHECK: call i32 @__kmpc_omp_task_with_deps(ptr @{{.*}}, i32 %{{.*}}, ptr %{{.*}}, i32 %[[NDEPS2]], ptr %[[DEP_ARR2]], i32 0, ptr null)
// CHECK: tail call void @free(ptr %[[DEP_ARR2]])
+// A signed i64 subtraction cannot represent this span. The two iterator
+// values are INT64_MIN and -1. The final increment to INT64_MAX - 1 is also
+// representable, as required by OpenMP 5.2 section 3.2.6.
+llvm.func @omp_task_depend_iterator_wide_span(%addr : !llvm.ptr) {
+ %min = llvm.mlir.constant(-9223372036854775808 : i64) : i64
+ %zero = llvm.mlir.constant(0 : i64) : i64
+ %max = llvm.mlir.constant(9223372036854775807 : i64) : i64
+ %it = omp.iterator(%iv: i64) = (%min to %zero step %max) {
+ omp.yield(%addr : !llvm.ptr)
+ } -> !omp.iterated<!llvm.ptr>
+ omp.task depend(taskdependin -> %it : !omp.iterated<!llvm.ptr>) {
+ omp.terminator
+ }
+ llvm.return
+}
+// CHECK-LABEL: define void @omp_task_depend_iterator_wide_span
+// CHECK: call ptr @malloc(i64 40)
+// CHECK: icmp ult i64 {{.*}}, 2
+// CHECK: call i32 @__kmpc_omp_task_with_deps(
+// CHECK-SAME: ptr {{.*}}, i32 {{.*}}, ptr {{.*}}, i32 2,
+
+// Exercise the most negative step with one iterator value, zero. The final
+// increment to INT64_MIN remains representable in the iterator type.
+llvm.func @omp_task_depend_iterator_min_step(%addr : !llvm.ptr) {
+ %zero = llvm.mlir.constant(0 : i64) : i64
+ %end = llvm.mlir.constant(-1 : i64) : i64
+ %min = llvm.mlir.constant(-9223372036854775808 : i64) : i64
+ %it = omp.iterator(%iv: i64) = (%zero to %end step %min) {
+ omp.yield(%addr : !llvm.ptr)
+ } -> !omp.iterated<!llvm.ptr>
+ omp.task depend(taskdependin -> %it : !omp.iterated<!llvm.ptr>) {
+ omp.terminator
+ }
+ llvm.return
+}
+// CHECK-LABEL: define void @omp_task_depend_iterator_min_step
+// CHECK: call ptr @malloc(i64 20)
+// CHECK: icmp ult i64 {{.*}}, 1
+// CHECK: call i32 @__kmpc_omp_task_with_deps(
+// CHECK-SAME: ptr {{.*}}, i32 {{.*}}, ptr {{.*}}, i32 1,
+
//--- target.mlir
// --------------------------------------------------------------------
More information about the Mlir-commits
mailing list