[Mlir-commits] [llvm] [mlir] [mlir][OpenMP] Allow multi-block omp.iterator regions (PR #229198)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 5 12:49:14 PDT 2026
https://github.com/chichunchen created https://github.com/llvm/llvm-project/pull/229198
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:
```mlir
%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.
>From 0883581d60b021b74d251323bcb06e79bca4611b 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] [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 | 52 ++++--
mlir/test/Dialect/OpenMP/invalid.mlir | 71 +++++++-
mlir/test/Dialect/OpenMP/ops.mlir | 21 +++
.../LLVMIR/openmp-iterator-multi-block.mlir | 152 ++++++++++++++++++
9 files changed, 458 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 c92756092af33..fd6cd2fb22885 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 3a8125b835e53..243c06c297f14 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 901d7e46d5d3e..a4a5d541b1cab 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 6bc99cf765699..a839287121396 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 c6a424bd93214..d3a7fb47b9917 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 fce2096a0ab3f..18eba462c4cf2 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,39 @@ 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);
+ }
+
+ // No PHI forwards the yield (unlike convertOmpOpRegions): its value may lack
+ // an LLVM type, and the unique omp.yield dominates the continuation anyway.
+ 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, continuation->getFirstInsertionPt());
return mlir::success();
}
@@ -3072,7 +3094,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 +3102,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 c3fe4f60e4166..1dfd23ef031d8 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 9371d3599b72a..7416c3da3b395 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 0000000000000..80a817ef8fb72
--- /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
More information about the Mlir-commits
mailing list