[MLIR][OpenMP] Prevent loop wrapper translation crashes #115475
Conversation
This patch updates the `convertOmpOpRegions` translation function to prevent calling it for a loop wrapper region from causing a compiler crash due to a lack of terminator operations. This problem is currently not triggered because there are no cases for which the region of a loop wrapper is passed to that function. This might have to change in order to support composite construct translation to LLVM IR.
skatrak
requested review from
ergawy,
Meinersbur,
jsjodin,
tblah,
kiranchandramohan,
mjklemm,
raghavendhra,
agozillon,
DominikAdamski,
TIFitis,
kparzysz and
bhandarkar-pranav
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-flang-openmp Author: Sergio Afonso (skatrak) ChangesThis patch updates the This problem is currently not triggered because there are no cases for which the region of a loop wrapper is passed to that function. This might have to change in order to support composite construct translation to LLVM IR. Full diff: https://github.com/llvm/llvm-project/pull/115475.diff 1 Files Affected:
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index da11ee9960e1f9..b507fa656d601f 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -391,6 +391,8 @@ static llvm::Expected<llvm::BasicBlock *> convertOmpOpRegions(
Region ®ion, StringRef blockName, llvm::IRBuilderBase &builder,
LLVM::ModuleTranslation &moduleTranslation,
SmallVectorImpl<llvm::PHINode *> *continuationBlockPHIs = nullptr) {
+ bool isLoopWrapper = isa<omp::LoopWrapperInterface>(region.getParentOp());
+
llvm::BasicBlock *continuationBlock =
splitBB(builder, true, "omp.region.cont");
llvm::BasicBlock *sourceBlock = builder.GetInsertBlock();
@@ -407,30 +409,34 @@ static llvm::Expected<llvm::BasicBlock *> convertOmpOpRegions(
// Terminators (namely YieldOp) may be forwarding values to the region that
// need to be available in the continuation block. Collect the types of these
- // operands in preparation of creating PHI nodes.
+ // operands in preparation of creating PHI nodes. This is skipped for loop
+ // wrapper operations, for which we know in advance they have no terminators.
SmallVector<llvm::Type *> continuationBlockPHITypes;
- bool operandsProcessed = false;
unsigned numYields = 0;
- for (Block &bb : region.getBlocks()) {
- if (omp::YieldOp yield = dyn_cast<omp::YieldOp>(bb.getTerminator())) {
- if (!operandsProcessed) {
- for (unsigned i = 0, e = yield->getNumOperands(); i < e; ++i) {
- continuationBlockPHITypes.push_back(
- moduleTranslation.convertType(yield->getOperand(i).getType()));
- }
- operandsProcessed = true;
- } else {
- assert(continuationBlockPHITypes.size() == yield->getNumOperands() &&
- "mismatching number of values yielded from the region");
- for (unsigned i = 0, e = yield->getNumOperands(); i < e; ++i) {
- llvm::Type *operandType =
- moduleTranslation.convertType(yield->getOperand(i).getType());
- (void)operandType;
- assert(continuationBlockPHITypes[i] == operandType &&
- "values of mismatching types yielded from the region");
+
+ if (!isLoopWrapper) {
+ bool operandsProcessed = false;
+ for (Block &bb : region.getBlocks()) {
+ if (omp::YieldOp yield = dyn_cast<omp::YieldOp>(bb.getTerminator())) {
+ if (!operandsProcessed) {
+ for (unsigned i = 0, e = yield->getNumOperands(); i < e; ++i) {
+ continuationBlockPHITypes.push_back(
+ moduleTranslation.convertType(yield->getOperand(i).getType()));
+ }
+ operandsProcessed = true;
+ } else {
+ assert(continuationBlockPHITypes.size() == yield->getNumOperands() &&
+ "mismatching number of values yielded from the region");
+ for (unsigned i = 0, e = yield->getNumOperands(); i < e; ++i) {
+ llvm::Type *operandType =
+ moduleTranslation.convertType(yield->getOperand(i).getType());
+ (void)operandType;
+ assert(continuationBlockPHITypes[i] == operandType &&
+ "values of mismatching types yielded from the region");
+ }
}
+ numYields++;
}
- numYields++;
}
}
@@ -468,6 +474,13 @@ static llvm::Expected<llvm::BasicBlock *> convertOmpOpRegions(
moduleTranslation.convertBlock(*bb, bb->isEntryBlock(), builder)))
return llvm::make_error<PreviouslyReportedError>();
+ // Create a direct branch here for loop wrappers to prevent their lack of a
+ // terminator from causing a crash below.
+ if (isLoopWrapper) {
+ builder.CreateBr(continuationBlock);
+ continue;
+ }
+
// Special handling for `omp.yield` and `omp.terminator` (we may have more
// than one): they return the control to the parent OpenMP dialect operation
// so replace them with the branch to the continuation block. We handle this
|
|
@llvm/pr-subscribers-mlir-openmp Author: Sergio Afonso (skatrak) ChangesThis patch updates the This problem is currently not triggered because there are no cases for which the region of a loop wrapper is passed to that function. This might have to change in order to support composite construct translation to LLVM IR. Full diff: https://github.com/llvm/llvm-project/pull/115475.diff 1 Files Affected:
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index da11ee9960e1f9..b507fa656d601f 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -391,6 +391,8 @@ static llvm::Expected<llvm::BasicBlock *> convertOmpOpRegions(
Region ®ion, StringRef blockName, llvm::IRBuilderBase &builder,
LLVM::ModuleTranslation &moduleTranslation,
SmallVectorImpl<llvm::PHINode *> *continuationBlockPHIs = nullptr) {
+ bool isLoopWrapper = isa<omp::LoopWrapperInterface>(region.getParentOp());
+
llvm::BasicBlock *continuationBlock =
splitBB(builder, true, "omp.region.cont");
llvm::BasicBlock *sourceBlock = builder.GetInsertBlock();
@@ -407,30 +409,34 @@ static llvm::Expected<llvm::BasicBlock *> convertOmpOpRegions(
// Terminators (namely YieldOp) may be forwarding values to the region that
// need to be available in the continuation block. Collect the types of these
- // operands in preparation of creating PHI nodes.
+ // operands in preparation of creating PHI nodes. This is skipped for loop
+ // wrapper operations, for which we know in advance they have no terminators.
SmallVector<llvm::Type *> continuationBlockPHITypes;
- bool operandsProcessed = false;
unsigned numYields = 0;
- for (Block &bb : region.getBlocks()) {
- if (omp::YieldOp yield = dyn_cast<omp::YieldOp>(bb.getTerminator())) {
- if (!operandsProcessed) {
- for (unsigned i = 0, e = yield->getNumOperands(); i < e; ++i) {
- continuationBlockPHITypes.push_back(
- moduleTranslation.convertType(yield->getOperand(i).getType()));
- }
- operandsProcessed = true;
- } else {
- assert(continuationBlockPHITypes.size() == yield->getNumOperands() &&
- "mismatching number of values yielded from the region");
- for (unsigned i = 0, e = yield->getNumOperands(); i < e; ++i) {
- llvm::Type *operandType =
- moduleTranslation.convertType(yield->getOperand(i).getType());
- (void)operandType;
- assert(continuationBlockPHITypes[i] == operandType &&
- "values of mismatching types yielded from the region");
+
+ if (!isLoopWrapper) {
+ bool operandsProcessed = false;
+ for (Block &bb : region.getBlocks()) {
+ if (omp::YieldOp yield = dyn_cast<omp::YieldOp>(bb.getTerminator())) {
+ if (!operandsProcessed) {
+ for (unsigned i = 0, e = yield->getNumOperands(); i < e; ++i) {
+ continuationBlockPHITypes.push_back(
+ moduleTranslation.convertType(yield->getOperand(i).getType()));
+ }
+ operandsProcessed = true;
+ } else {
+ assert(continuationBlockPHITypes.size() == yield->getNumOperands() &&
+ "mismatching number of values yielded from the region");
+ for (unsigned i = 0, e = yield->getNumOperands(); i < e; ++i) {
+ llvm::Type *operandType =
+ moduleTranslation.convertType(yield->getOperand(i).getType());
+ (void)operandType;
+ assert(continuationBlockPHITypes[i] == operandType &&
+ "values of mismatching types yielded from the region");
+ }
}
+ numYields++;
}
- numYields++;
}
}
@@ -468,6 +474,13 @@ static llvm::Expected<llvm::BasicBlock *> convertOmpOpRegions(
moduleTranslation.convertBlock(*bb, bb->isEntryBlock(), builder)))
return llvm::make_error<PreviouslyReportedError>();
+ // Create a direct branch here for loop wrappers to prevent their lack of a
+ // terminator from causing a crash below.
+ if (isLoopWrapper) {
+ builder.CreateBr(continuationBlock);
+ continue;
+ }
+
// Special handling for `omp.yield` and `omp.terminator` (we may have more
// than one): they return the control to the parent OpenMP dialect operation
// so replace them with the branch to the continuation block. We handle this
|
There was a problem hiding this comment.
Reviewable status: 0 of 1 files reviewed, 1 unresolved discussion (waiting on @agozillon, @bhandarkar-pranav, @DominikAdamski, @ergawy, @jsjodin, @kparzysz, @Meinersbur, @mjklemm, @raghavendhra, @skatrak, @tblah, and @TIFitis)
mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp line 482 at r1 (raw file):
builder.CreateBr(continuationBlock); continue; }
Can we write a test for this?
|
Thanks @kiranchandramohan for the quick review!
At the moment we can't because this code path is not triggered by anything. Translation functions for loop wrappers currently call this function passing the region of the nested I think it's preferable to have this implemented before needing it rather than having it crash later on, since I think it was an oversight from when I implemented #112229. But let me know if you don't agree and I can keep it around until some composite construct lowering ends up triggering the issue. |
The danger with this is that someone can remove this piece of code in refactoring and think it was accidentally introduced since no test fails. Ideally, such code can be introduced in the first patch that needs it. You can organize such patches as two commits, where the first commit is this patch and the second commit is the one that needs it Note: I did not intend to request changes. I was trying reviewable.io and did not realise that it requested changes. |
Thank you both for your reviews. I found this bug through our current downstream support for 'distribute parallel do'. However, it will take a bit of time to get to upstreaming it, since other pieces must make it first, and there's a chance that we might want to rethink how we do it before then (potentially circumventing this code path). So I'll keep this PR around until we're ready to upstream the code that would trigger this failure, and I'll close it if we end up not needing it at that point. |
This patch updates the
convertOmpOpRegionstranslation function to prevent calling it for a loop wrapper region from causing a compiler crash due to a lack of terminator operations.This problem is currently not triggered because there are no cases for which the region of a loop wrapper is passed to that function. This might have to change in order to support composite construct translation to LLVM IR.