Skip to content

Commit

Permalink
[MLIR][OpenMP] Named recipe op's block args accessors (NFC) (#112192)
Browse files Browse the repository at this point in the history 
This patch adds extra class declarations to the `omp.declare_reduction`
and `omp.private` operations to access the entry block arguments defined
by their regions. Some existing accesses to these arguments are updated
to use the new named methods to improve code readability.
  • Loading branch information
@skatrak
skatrak authored Oct 15, 2024
1 parent 790d986 commit 7ec3209
Show file tree
Hide file tree
Showing 2 changed files with 63 additions and 28 deletions.
61 changes: 49 additions & 12 deletions mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,24 @@ def PrivateClauseOp : OpenMP_Op<"private", [IsolatedFromAbove, RecipeInterface]>
CArg<"TypeAttr">:$type)>
];

let extraClassDeclaration = [{
BlockArgument getAllocMoldArg() {
return getAllocRegion().getArgument(0);
}
BlockArgument getCopyMoldArg() {
auto &region = getCopyRegion();
return region.empty() ? nullptr : region.getArgument(0);
}
BlockArgument getCopyPrivateArg() {
auto &region = getCopyRegion();
return region.empty() ? nullptr : region.getArgument(1);
}
BlockArgument getDeallocMoldArg() {
auto &region = getDeallocRegion();
return region.empty() ? nullptr : region.getArgument(0);
}
}];

let hasVerifier = 1;
}

Expand Down Expand Up @@ -1601,22 +1619,41 @@ def DeclareReductionOp : OpenMP_Op<"declare_reduction", [IsolatedFromAbove,
"( `cleanup` $cleanupRegion^ )? ";

let extraClassDeclaration = [{
BlockArgument getAllocMoldArg() {
auto &region = getAllocRegion();
return region.empty() ? nullptr : region.getArgument(0);
}
BlockArgument getInitializerMoldArg() {
return getInitializerRegion().getArgument(0);
}
BlockArgument getInitializerAllocArg() {
return getAllocRegion().empty() ?
nullptr : getInitializerRegion().getArgument(1);
}
BlockArgument getReductionLhsArg() {
return getReductionRegion().getArgument(0);
}
BlockArgument getReductionRhsArg() {
return getReductionRegion().getArgument(1);
}
BlockArgument getAtomicReductionLhsArg() {
auto &region = getAtomicReductionRegion();
return region.empty() ? nullptr : region.getArgument(0);
}
BlockArgument getAtomicReductionRhsArg() {
auto &region = getAtomicReductionRegion();
return region.empty() ? nullptr : region.getArgument(1);
}
BlockArgument getCleanupAllocArg() {
auto &region = getCleanupRegion();
return region.empty() ? nullptr : region.getArgument(0);
}

PointerLikeType getAccumulatorType() {
if (getAtomicReductionRegion().empty())
return {};

return cast<PointerLikeType>(getAtomicReductionRegion().front().getArgument(0).getType());
}

Value getInitializerMoldArg() {
return getInitializerRegion().front().getArgument(0);
}

Value getInitializerAllocArg() {
if (getAllocRegion().empty() ||
getInitializerRegion().front().getNumArguments() != 2)
return {nullptr};
return getInitializerRegion().front().getArgument(1);
return cast<PointerLikeType>(getAtomicReductionLhsArg().getType());
}
}];
let hasRegionVerifier = 1;
Expand Down
30 changes: 14 additions & 16 deletions mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -480,12 +480,11 @@ makeReductionGen(omp::DeclareReductionOp decl, llvm::IRBuilderBase &builder,
[&, decl](llvm::OpenMPIRBuilder::InsertPointTy insertPoint,
llvm::Value *lhs, llvm::Value *rhs,
llvm::Value *&result) mutable {
Region &reductionRegion = decl.getReductionRegion();
moduleTranslation.mapValue(reductionRegion.front().getArgument(0), lhs);
moduleTranslation.mapValue(reductionRegion.front().getArgument(1), rhs);
moduleTranslation.mapValue(decl.getReductionLhsArg(), lhs);
moduleTranslation.mapValue(decl.getReductionRhsArg(), rhs);
builder.restoreIP(insertPoint);
SmallVector<llvm::Value *> phis;
if (failed(inlineConvertOmpRegions(reductionRegion,
if (failed(inlineConvertOmpRegions(decl.getReductionRegion(),
"omp.reduction.nonatomic.body",
builder, moduleTranslation, &phis)))
return llvm::OpenMPIRBuilder::InsertPointTy();
Expand Down Expand Up @@ -513,12 +512,11 @@ makeAtomicReductionGen(omp::DeclareReductionOp decl,
OwningAtomicReductionGen atomicGen =
[&, decl](llvm::OpenMPIRBuilder::InsertPointTy insertPoint, llvm::Type *,
llvm::Value *lhs, llvm::Value *rhs) mutable {
Region &atomicRegion = decl.getAtomicReductionRegion();
moduleTranslation.mapValue(atomicRegion.front().getArgument(0), lhs);
moduleTranslation.mapValue(atomicRegion.front().getArgument(1), rhs);
moduleTranslation.mapValue(decl.getAtomicReductionLhsArg(), lhs);
moduleTranslation.mapValue(decl.getAtomicReductionRhsArg(), rhs);
builder.restoreIP(insertPoint);
SmallVector<llvm::Value *> phis;
if (failed(inlineConvertOmpRegions(atomicRegion,
if (failed(inlineConvertOmpRegions(decl.getAtomicReductionRegion(),
"omp.reduction.atomic.body", builder,
moduleTranslation, &phis)))
return llvm::OpenMPIRBuilder::InsertPointTy();
Expand Down Expand Up @@ -1674,9 +1672,10 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
// argument of the `alloc` region and the second argument of the `copy`
// region to be the yielded value of the `alloc` region (this is the
// private clone of the privatized value).
copyCloneBuilder.mergeBlocks(
&*newCopyRegionFrontBlock, &*oldAllocBackBlock,
{allocRegion.getArgument(0), oldAllocYieldOp.getOperand(0)});
copyCloneBuilder.mergeBlocks(&*newCopyRegionFrontBlock,
&*oldAllocBackBlock,
{mlirPrivatizerClone.getAllocMoldArg(),
oldAllocYieldOp.getOperand(0)});

// 4. The old terminator of the `alloc` region is not needed anymore, so
// delete it.
Expand All @@ -1686,8 +1685,8 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
// Replace the privatizer block argument with mlir value being privatized.
// This way, the body of the privatizer will be changed from using the
// region/block argument to the value being privatized.
auto allocRegionArg = allocRegion.getArgument(0);
replaceAllUsesInRegionWith(allocRegionArg, mlirPrivVar, allocRegion);
replaceAllUsesInRegionWith(mlirPrivatizerClone.getAllocMoldArg(),
mlirPrivVar, allocRegion);

auto oldIP = builder.saveIP();
builder.restoreIP(allocaIP);
Expand Down Expand Up @@ -3480,10 +3479,9 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
" private allocatables is not supported yet");
bodyGenStatus = failure();
} else {
Region &allocRegion = privatizer.getAllocRegion();
BlockArgument allocRegionArg = allocRegion.getArgument(0);
moduleTranslation.mapValue(allocRegionArg,
moduleTranslation.mapValue(privatizer.getAllocMoldArg(),
moduleTranslation.lookupValue(privVar));
Region &allocRegion = privatizer.getAllocRegion();
SmallVector<llvm::Value *, 1> yieldedValues;
if (failed(inlineConvertOmpRegions(
allocRegion, "omp.targetop.privatizer", builder,
Expand Down

0 comments on commit 7ec3209

Please sign in to comment.