[Mlir-commits] [mlir] [mlir][OpenMP] Lower REDUCTION clause for SECTIONS construct (PR #97859)
Tom Eccles
llvmlistbot at llvm.org
Fri Jul 5 14:12:43 PDT 2024
https://github.com/tblah created https://github.com/llvm/llvm-project/pull/97859
This shares code with WsloopOp (the changes to Wsloop should be NFC). OpenMPIRBuilder basically implements SECTIONS as a wsloop over a case statement with each SECTION as a case for a particular loopiv value.
Unfortunately it proved very difficult to share code between these and ParallelOp. ParallelOp does quite a few things differently (doing more work inside of the bodygen callback and laying out blocks differently). Aligning reduction implementations for wsloop and parallel will probably involve functional changes to both, so I won't attempt that in this commit.
>From 5c929cfe0549384a6b6fc469e0553776f3ef9b18 Mon Sep 17 00:00:00 2001
From: Tom Eccles <tom.eccles at arm.com>
Date: Thu, 4 Jul 2024 15:13:21 +0000
Subject: [PATCH] [mlir][OpenMP] Lower REDUCTION clause for SECTIONS construct
This shares code with WsloopOp (the changes to Wsloop should be NFC).
OpenMPIRBuilder basically implements SECTIONS as a wsloop over a case
statement with each SECTION as a case for a particular loopiv value.
Unfortunately it proved very difficult to share code between these and
ParallelOp. ParallelOp does quite a few things differently (doing more
work inside of the bodygen callback and laying out blocks differently).
Aligning reduction implementations for wsloop and parallel will probably
involve functional changes to both, so I won't attempt that in this
commit.
---
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 510 ++++++++++--------
mlir/test/Target/LLVMIR/openmp-llvm.mlir | 2 +
.../LLVMIR/openmp-reduction-sections.mlir | 168 ++++++
3 files changed, 467 insertions(+), 213 deletions(-)
create mode 100644 mlir/test/Target/LLVMIR/openmp-reduction-sections.mlir
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 35971fbacbf91..3dd87632a82c6 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -555,6 +555,239 @@ convertOmpOrderedRegion(Operation &opInst, llvm::IRBuilderBase &builder,
return bodyGenStatus;
}
+/// Allocate space for privatized reduction variables.
+template <typename T>
+static void allocByValReductionVars(
+ T loop, ArrayRef<BlockArgument> reductionArgs, llvm::IRBuilderBase &builder,
+ LLVM::ModuleTranslation &moduleTranslation,
+ llvm::OpenMPIRBuilder::InsertPointTy &allocaIP,
+ SmallVectorImpl<omp::DeclareReductionOp> &reductionDecls,
+ SmallVectorImpl<llvm::Value *> &privateReductionVariables,
+ DenseMap<Value, llvm::Value *> &reductionVariableMap,
+ llvm::ArrayRef<bool> isByRefs) {
+ llvm::IRBuilderBase::InsertPointGuard guard(builder);
+ builder.SetInsertPoint(allocaIP.getBlock()->getTerminator());
+
+ for (std::size_t i = 0; i < loop.getNumReductionVars(); ++i) {
+ if (isByRefs[i])
+ continue;
+ llvm::Value *var = builder.CreateAlloca(
+ moduleTranslation.convertType(reductionDecls[i].getType()));
+ moduleTranslation.mapValue(reductionArgs[i], var);
+ privateReductionVariables[i] = var;
+ reductionVariableMap.try_emplace(loop.getReductionVars()[i], var);
+ }
+}
+
+/// Map input argument to all reduction initialization regions
+template <typename T>
+static void
+mapInitializationArg(T loop, LLVM::ModuleTranslation &moduleTranslation,
+ SmallVectorImpl<omp::DeclareReductionOp> &reductionDecls,
+ unsigned i) {
+ // map input argument to the initialization region
+ mlir::omp::DeclareReductionOp &reduction = reductionDecls[i];
+ Region &initializerRegion = reduction.getInitializerRegion();
+ Block &entry = initializerRegion.front();
+ assert(entry.getNumArguments() == 1 &&
+ "the initialization region has one argument");
+
+ mlir::Value mlirSource = loop.getReductionVars()[i];
+ llvm::Value *llvmSource = moduleTranslation.lookupValue(mlirSource);
+ assert(llvmSource && "lookup reduction var");
+ moduleTranslation.mapValue(entry.getArgument(0), llvmSource);
+}
+
+/// Collect reduction info
+template <typename T>
+static void collectReductionInfo(
+ T loop, llvm::IRBuilderBase &builder,
+ LLVM::ModuleTranslation &moduleTranslation,
+ SmallVectorImpl<omp::DeclareReductionOp> &reductionDecls,
+ SmallVectorImpl<OwningReductionGen> &owningReductionGens,
+ SmallVectorImpl<OwningAtomicReductionGen> &owningAtomicReductionGens,
+ const ArrayRef<llvm::Value *> privateReductionVariables,
+ SmallVectorImpl<llvm::OpenMPIRBuilder::ReductionInfo> &reductionInfos) {
+ unsigned numReductions = loop.getNumReductionVars();
+
+ for (unsigned i = 0; i < numReductions; ++i) {
+ owningReductionGens.push_back(
+ makeReductionGen(reductionDecls[i], builder, moduleTranslation));
+ owningAtomicReductionGens.push_back(
+ makeAtomicReductionGen(reductionDecls[i], builder, moduleTranslation));
+ }
+
+ // Collect the reduction information.
+ reductionInfos.reserve(numReductions);
+ for (unsigned i = 0; i < numReductions; ++i) {
+ llvm::OpenMPIRBuilder::ReductionGenAtomicCBTy atomicGen = nullptr;
+ if (owningAtomicReductionGens[i])
+ atomicGen = owningAtomicReductionGens[i];
+ llvm::Value *variable =
+ moduleTranslation.lookupValue(loop.getReductionVars()[i]);
+ reductionInfos.push_back(
+ {moduleTranslation.convertType(reductionDecls[i].getType()), variable,
+ privateReductionVariables[i],
+ /*EvaluationKind=*/llvm::OpenMPIRBuilder::EvalKind::Scalar,
+ owningReductionGens[i],
+ /*ReductionGenClang=*/nullptr, atomicGen});
+ }
+}
+
+/// handling of DeclareReductionOp's cleanup region
+static LogicalResult
+inlineOmpRegionCleanup(llvm::SmallVectorImpl<Region *> &cleanupRegions,
+ llvm::ArrayRef<llvm::Value *> privateVariables,
+ LLVM::ModuleTranslation &moduleTranslation,
+ llvm::IRBuilderBase &builder, StringRef regionName,
+ bool shouldLoadCleanupRegionArg = true) {
+ for (auto [i, cleanupRegion] : llvm::enumerate(cleanupRegions)) {
+ if (cleanupRegion->empty())
+ continue;
+
+ // map the argument to the cleanup region
+ Block &entry = cleanupRegion->front();
+
+ llvm::Instruction *potentialTerminator =
+ builder.GetInsertBlock()->empty() ? nullptr
+ : &builder.GetInsertBlock()->back();
+ if (potentialTerminator && potentialTerminator->isTerminator())
+ builder.SetInsertPoint(potentialTerminator);
+ llvm::Value *prviateVarValue =
+ shouldLoadCleanupRegionArg
+ ? builder.CreateLoad(
+ moduleTranslation.convertType(entry.getArgument(0).getType()),
+ privateVariables[i])
+ : privateVariables[i];
+
+ moduleTranslation.mapValue(entry.getArgument(0), prviateVarValue);
+
+ if (failed(inlineConvertOmpRegions(*cleanupRegion, regionName, builder,
+ moduleTranslation)))
+ return failure();
+
+ // clear block argument mapping in case it needs to be re-created with a
+ // different source for another use of the same reduction decl
+ moduleTranslation.forgetMapping(*cleanupRegion);
+ }
+ return success();
+}
+
+// TODO: not used by ParallelOp
+template <class OP>
+static LogicalResult createReductionsAndCleanup(
+ OP op, llvm::IRBuilderBase &builder,
+ LLVM::ModuleTranslation &moduleTranslation,
+ llvm::OpenMPIRBuilder::InsertPointTy &allocaIP,
+ SmallVectorImpl<omp::DeclareReductionOp> &reductionDecls,
+ ArrayRef<llvm::Value *> privateReductionVariables, ArrayRef<bool> isByRef) {
+ // Process the reductions if required.
+ if (op.getNumReductionVars() == 0)
+ return success();
+
+ llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
+
+ // Create the reduction generators. We need to own them here because
+ // ReductionInfo only accepts references to the generators.
+ SmallVector<OwningReductionGen> owningReductionGens;
+ SmallVector<OwningAtomicReductionGen> owningAtomicReductionGens;
+ SmallVector<llvm::OpenMPIRBuilder::ReductionInfo> reductionInfos;
+ collectReductionInfo(op, builder, moduleTranslation, reductionDecls,
+ owningReductionGens, owningAtomicReductionGens,
+ privateReductionVariables, reductionInfos);
+
+ // The call to createReductions below expects the block to have a
+ // terminator. Create an unreachable instruction to serve as terminator
+ // and remove it later.
+ llvm::UnreachableInst *tempTerminator = builder.CreateUnreachable();
+ builder.SetInsertPoint(tempTerminator);
+ llvm::OpenMPIRBuilder::InsertPointTy contInsertPoint =
+ ompBuilder->createReductions(builder.saveIP(), allocaIP, reductionInfos,
+ isByRef, op.getNowait());
+ if (!contInsertPoint.getBlock())
+ return op->emitOpError() << "failed to convert reductions";
+ auto nextInsertionPoint =
+ ompBuilder->createBarrier(contInsertPoint, llvm::omp::OMPD_for);
+ tempTerminator->eraseFromParent();
+ builder.restoreIP(nextInsertionPoint);
+
+ // after the construct, deallocate private reduction variables
+ SmallVector<Region *> reductionRegions;
+ llvm::transform(reductionDecls, std::back_inserter(reductionRegions),
+ [](omp::DeclareReductionOp reductionDecl) {
+ return &reductionDecl.getCleanupRegion();
+ });
+ return inlineOmpRegionCleanup(reductionRegions, privateReductionVariables,
+ moduleTranslation, builder,
+ "omp.reduction.cleanup");
+ return success();
+}
+
+static ArrayRef<bool> getIsByRef(std::optional<ArrayRef<bool>> attr) {
+ if (!attr)
+ return {};
+ return *attr;
+}
+
+// TODO: not used by omp.parallel
+template <typename OP>
+static LogicalResult allocAndInitializeReductionVars(
+ OP op, ArrayRef<BlockArgument> reductionArgs, llvm::IRBuilderBase &builder,
+ LLVM::ModuleTranslation &moduleTranslation,
+ llvm::OpenMPIRBuilder::InsertPointTy &allocaIP,
+ SmallVectorImpl<omp::DeclareReductionOp> &reductionDecls,
+ SmallVectorImpl<llvm::Value *> &privateReductionVariables,
+ DenseMap<Value, llvm::Value *> &reductionVariableMap,
+ llvm::ArrayRef<bool> isByRef) {
+ if (op.getNumReductionVars() == 0)
+ return success();
+
+ allocByValReductionVars(op, reductionArgs, builder, moduleTranslation,
+ allocaIP, reductionDecls, privateReductionVariables,
+ reductionVariableMap, isByRef);
+
+ // Before the loop, store the initial values of reductions into reduction
+ // variables. Although this could be done after allocas, we don't want to mess
+ // up with the alloca insertion point.
+ for (unsigned i = 0; i < op.getNumReductionVars(); ++i) {
+ SmallVector<llvm::Value *> phis;
+
+ // map block argument to initializer region
+ mapInitializationArg(op, moduleTranslation, reductionDecls, i);
+
+ if (failed(inlineConvertOmpRegions(reductionDecls[i].getInitializerRegion(),
+ "omp.reduction.neutral", builder,
+ moduleTranslation, &phis)))
+ return failure();
+ assert(phis.size() == 1 && "expected one value to be yielded from the "
+ "reduction neutral element declaration region");
+ if (isByRef[i]) {
+ // Allocate reduction variable (which is a pointer to the real reduction
+ // variable allocated in the inlined region)
+ llvm::Value *var = builder.CreateAlloca(
+ moduleTranslation.convertType(reductionDecls[i].getType()));
+ // Store the result of the inlined region to the allocated reduction var
+ // ptr
+ builder.CreateStore(phis[0], var);
+
+ privateReductionVariables[i] = var;
+ moduleTranslation.mapValue(reductionArgs[i], phis[0]);
+ reductionVariableMap.try_emplace(op.getReductionVars()[i], phis[0]);
+ } else {
+ // for by-ref case the store is inside of the reduction region
+ builder.CreateStore(phis[0], privateReductionVariables[i]);
+ // the rest was handled in allocByValReductionVars
+ }
+
+ // forget the mapping for the initializer region because we might need a
+ // different mapping if this reduction declaration is re-used for a
+ // different variable
+ moduleTranslation.forgetMapping(reductionDecls[i].getInitializerRegion());
+ }
+
+ return success();
+}
+
static LogicalResult
convertOmpSections(Operation &opInst, llvm::IRBuilderBase &builder,
LLVM::ModuleTranslation &moduleTranslation) {
@@ -565,13 +798,38 @@ convertOmpSections(Operation &opInst, llvm::IRBuilderBase &builder,
auto sectionsOp = cast<omp::SectionsOp>(opInst);
// TODO: Support the following clauses: private, firstprivate, lastprivate,
- // reduction, allocate
- if (!sectionsOp.getReductionVars().empty() || sectionsOp.getReductions() ||
- !sectionsOp.getAllocateVars().empty() ||
+ // allocate
+ if (!sectionsOp.getAllocateVars().empty() ||
!sectionsOp.getAllocatorsVars().empty())
return emitError(sectionsOp.getLoc())
- << "reduction and allocate clauses are not supported for sections "
- "construct";
+ << "allocate clause is not supported for sections construct";
+
+ llvm::ArrayRef<bool> isByRef = getIsByRef(sectionsOp.getReductionVarsByref());
+ assert(isByRef.size() == sectionsOp.getNumReductionVars());
+
+ SmallVector<omp::DeclareReductionOp> reductionDecls;
+ collectReductionDecls(sectionsOp, reductionDecls);
+ llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
+ findAllocaInsertPoint(builder, moduleTranslation);
+
+ SmallVector<llvm::Value *> privateReductionVariables(
+ sectionsOp.getNumReductionVars());
+ DenseMap<Value, llvm::Value *> reductionVariableMap;
+
+ MutableArrayRef<BlockArgument> reductionArgs =
+ sectionsOp.getRegion().getArguments();
+
+ if (failed(allocAndInitializeReductionVars(
+ sectionsOp, reductionArgs, builder, moduleTranslation, allocaIP,
+ reductionDecls, privateReductionVariables, reductionVariableMap,
+ isByRef)))
+ return failure();
+
+ // Store the mapping between reduction variables and their private copies on
+ // ModuleTranslation stack. It can be then recovered when translating
+ // omp.reduce operations in a separate call.
+ LLVM::ModuleTranslation::SaveStack<OpenMPVarMappingStackFrame> mappingGuard(
+ moduleTranslation, reductionVariableMap);
LogicalResult bodyGenStatus = success();
SmallVector<StorableBodyGenCallbackTy> sectionCBs;
@@ -582,9 +840,24 @@ convertOmpSections(Operation &opInst, llvm::IRBuilderBase &builder,
continue;
Region ®ion = sectionOp.getRegion();
- auto sectionCB = [®ion, &builder, &moduleTranslation, &bodyGenStatus](
- InsertPointTy allocaIP, InsertPointTy codeGenIP) {
+ auto sectionCB = [§ionsOp, ®ion, &builder, &moduleTranslation,
+ &bodyGenStatus](InsertPointTy allocaIP,
+ InsertPointTy codeGenIP) {
builder.restoreIP(codeGenIP);
+
+ // map the omp.section reduction block argument to the omp.sections block
+ // arguments
+ // TODO: this assumes that the only block arguments are reduction
+ // variables
+ assert(region.getNumArguments() ==
+ sectionsOp.getRegion().getNumArguments());
+ for (auto [sectionsArg, sectionArg] : llvm::zip_equal(
+ sectionsOp.getRegion().getArguments(), region.getArguments())) {
+ llvm::Value *llvmVal = moduleTranslation.lookupValue(sectionsArg);
+ assert(llvmVal);
+ moduleTranslation.mapValue(sectionArg, llvmVal);
+ }
+
convertOmpOpRegions(region, "omp.section.region", builder,
moduleTranslation, bodyGenStatus);
};
@@ -613,13 +886,19 @@ convertOmpSections(Operation &opInst, llvm::IRBuilderBase &builder,
// called for variables which have destructors/finalizers.
auto finiCB = [&](InsertPointTy codeGenIP) {};
- llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
- findAllocaInsertPoint(builder, moduleTranslation);
+ allocaIP = findAllocaInsertPoint(builder, moduleTranslation);
llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createSections(
ompLoc, allocaIP, sectionCBs, privCB, finiCB, false,
sectionsOp.getNowait()));
- return bodyGenStatus;
+
+ if (failed(bodyGenStatus))
+ return bodyGenStatus;
+
+ // Process the reductions if required.
+ return createReductionsAndCleanup(sectionsOp, builder, moduleTranslation,
+ allocaIP, reductionDecls,
+ privateReductionVariables, isByRef);
}
/// Converts an OpenMP single construct into LLVM IR using OpenMPIRBuilder.
@@ -769,131 +1048,6 @@ convertOmpTaskgroupOp(omp::TaskgroupOp tgOp, llvm::IRBuilderBase &builder,
ompLoc, allocaIP, bodyCB));
return bodyGenStatus;
}
-
-/// Allocate space for privatized reduction variables.
-template <typename T>
-static void allocByValReductionVars(
- T loop, ArrayRef<BlockArgument> reductionArgs, llvm::IRBuilderBase &builder,
- LLVM::ModuleTranslation &moduleTranslation,
- llvm::OpenMPIRBuilder::InsertPointTy &allocaIP,
- SmallVectorImpl<omp::DeclareReductionOp> &reductionDecls,
- SmallVectorImpl<llvm::Value *> &privateReductionVariables,
- DenseMap<Value, llvm::Value *> &reductionVariableMap,
- llvm::ArrayRef<bool> isByRefs) {
- llvm::IRBuilderBase::InsertPointGuard guard(builder);
- builder.SetInsertPoint(allocaIP.getBlock()->getTerminator());
-
- for (std::size_t i = 0; i < loop.getNumReductionVars(); ++i) {
- if (isByRefs[i])
- continue;
- llvm::Value *var = builder.CreateAlloca(
- moduleTranslation.convertType(reductionDecls[i].getType()));
- moduleTranslation.mapValue(reductionArgs[i], var);
- privateReductionVariables[i] = var;
- reductionVariableMap.try_emplace(loop.getReductionVars()[i], var);
- }
-}
-
-/// Map input argument to all reduction initialization regions
-template <typename T>
-static void
-mapInitializationArg(T loop, LLVM::ModuleTranslation &moduleTranslation,
- SmallVectorImpl<omp::DeclareReductionOp> &reductionDecls,
- unsigned i) {
- // map input argument to the initialization region
- mlir::omp::DeclareReductionOp &reduction = reductionDecls[i];
- Region &initializerRegion = reduction.getInitializerRegion();
- Block &entry = initializerRegion.front();
- assert(entry.getNumArguments() == 1 &&
- "the initialization region has one argument");
-
- mlir::Value mlirSource = loop.getReductionVars()[i];
- llvm::Value *llvmSource = moduleTranslation.lookupValue(mlirSource);
- assert(llvmSource && "lookup reduction var");
- moduleTranslation.mapValue(entry.getArgument(0), llvmSource);
-}
-
-/// Collect reduction info
-template <typename T>
-static void collectReductionInfo(
- T loop, llvm::IRBuilderBase &builder,
- LLVM::ModuleTranslation &moduleTranslation,
- SmallVector<omp::DeclareReductionOp> &reductionDecls,
- SmallVector<OwningReductionGen> &owningReductionGens,
- SmallVector<OwningAtomicReductionGen> &owningAtomicReductionGens,
- const SmallVector<llvm::Value *> &privateReductionVariables,
- SmallVector<llvm::OpenMPIRBuilder::ReductionInfo> &reductionInfos) {
- unsigned numReductions = loop.getNumReductionVars();
-
- for (unsigned i = 0; i < numReductions; ++i) {
- owningReductionGens.push_back(
- makeReductionGen(reductionDecls[i], builder, moduleTranslation));
- owningAtomicReductionGens.push_back(
- makeAtomicReductionGen(reductionDecls[i], builder, moduleTranslation));
- }
-
- // Collect the reduction information.
- reductionInfos.reserve(numReductions);
- for (unsigned i = 0; i < numReductions; ++i) {
- llvm::OpenMPIRBuilder::ReductionGenAtomicCBTy atomicGen = nullptr;
- if (owningAtomicReductionGens[i])
- atomicGen = owningAtomicReductionGens[i];
- llvm::Value *variable =
- moduleTranslation.lookupValue(loop.getReductionVars()[i]);
- reductionInfos.push_back(
- {moduleTranslation.convertType(reductionDecls[i].getType()), variable,
- privateReductionVariables[i],
- /*EvaluationKind=*/llvm::OpenMPIRBuilder::EvalKind::Scalar,
- owningReductionGens[i],
- /*ReductionGenClang=*/nullptr, atomicGen});
- }
-}
-
-/// handling of DeclareReductionOp's cleanup region
-static LogicalResult
-inlineOmpRegionCleanup(llvm::SmallVectorImpl<Region *> &cleanupRegions,
- llvm::ArrayRef<llvm::Value *> privateVariables,
- LLVM::ModuleTranslation &moduleTranslation,
- llvm::IRBuilderBase &builder, StringRef regionName,
- bool shouldLoadCleanupRegionArg = true) {
- for (auto [i, cleanupRegion] : llvm::enumerate(cleanupRegions)) {
- if (cleanupRegion->empty())
- continue;
-
- // map the argument to the cleanup region
- Block &entry = cleanupRegion->front();
-
- llvm::Instruction *potentialTerminator =
- builder.GetInsertBlock()->empty() ? nullptr
- : &builder.GetInsertBlock()->back();
- if (potentialTerminator && potentialTerminator->isTerminator())
- builder.SetInsertPoint(potentialTerminator);
- llvm::Value *prviateVarValue =
- shouldLoadCleanupRegionArg
- ? builder.CreateLoad(
- moduleTranslation.convertType(entry.getArgument(0).getType()),
- privateVariables[i])
- : privateVariables[i];
-
- moduleTranslation.mapValue(entry.getArgument(0), prviateVarValue);
-
- if (failed(inlineConvertOmpRegions(*cleanupRegion, regionName, builder,
- moduleTranslation)))
- return failure();
-
- // clear block argument mapping in case it needs to be re-created with a
- // different source for another use of the same reduction decl
- moduleTranslation.forgetMapping(*cleanupRegion);
- }
- return success();
-}
-
-static ArrayRef<bool> getIsByRef(std::optional<ArrayRef<bool>> attr) {
- if (!attr)
- return {};
- return *attr;
-}
-
/// Converts an OpenMP workshare loop into LLVM IR using OpenMPIRBuilder.
static LogicalResult
convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
@@ -930,48 +1084,11 @@ convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
MutableArrayRef<BlockArgument> reductionArgs =
wsloopOp.getRegion().getArguments();
- allocByValReductionVars(wsloopOp, reductionArgs, builder, moduleTranslation,
- allocaIP, reductionDecls, privateReductionVariables,
- reductionVariableMap, isByRef);
-
- // Before the loop, store the initial values of reductions into reduction
- // variables. Although this could be done after allocas, we don't want to mess
- // up with the alloca insertion point.
- for (unsigned i = 0; i < wsloopOp.getNumReductionVars(); ++i) {
- SmallVector<llvm::Value *> phis;
-
- // map block argument to initializer region
- mapInitializationArg(wsloopOp, moduleTranslation, reductionDecls, i);
-
- if (failed(inlineConvertOmpRegions(reductionDecls[i].getInitializerRegion(),
- "omp.reduction.neutral", builder,
- moduleTranslation, &phis)))
- return failure();
- assert(phis.size() == 1 && "expected one value to be yielded from the "
- "reduction neutral element declaration region");
- if (isByRef[i]) {
- // Allocate reduction variable (which is a pointer to the real reduction
- // variable allocated in the inlined region)
- llvm::Value *var = builder.CreateAlloca(
- moduleTranslation.convertType(reductionDecls[i].getType()));
- // Store the result of the inlined region to the allocated reduction var
- // ptr
- builder.CreateStore(phis[0], var);
-
- privateReductionVariables[i] = var;
- moduleTranslation.mapValue(reductionArgs[i], phis[0]);
- reductionVariableMap.try_emplace(wsloopOp.getReductionVars()[i], phis[0]);
- } else {
- // for by-ref case the store is inside of the reduction region
- builder.CreateStore(phis[0], privateReductionVariables[i]);
- // the rest was handled in allocByValReductionVars
- }
-
- // forget the mapping for the initializer region because we might need a
- // different mapping if this reduction declaration is re-used for a
- // different variable
- moduleTranslation.forgetMapping(reductionDecls[i].getInitializerRegion());
- }
+ if (failed(allocAndInitializeReductionVars(
+ wsloopOp, reductionArgs, builder, moduleTranslation, allocaIP,
+ reductionDecls, privateReductionVariables, reductionVariableMap,
+ isByRef)))
+ return failure();
// Store the mapping between reduction variables and their private copies on
// ModuleTranslation stack. It can be then recovered when translating
@@ -1064,42 +1181,9 @@ convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
builder.restoreIP(afterIP);
// Process the reductions if required.
- if (wsloopOp.getNumReductionVars() == 0)
- return success();
-
- // Create the reduction generators. We need to own them here because
- // ReductionInfo only accepts references to the generators.
- SmallVector<OwningReductionGen> owningReductionGens;
- SmallVector<OwningAtomicReductionGen> owningAtomicReductionGens;
- SmallVector<llvm::OpenMPIRBuilder::ReductionInfo> reductionInfos;
- collectReductionInfo(wsloopOp, builder, moduleTranslation, reductionDecls,
- owningReductionGens, owningAtomicReductionGens,
- privateReductionVariables, reductionInfos);
-
- // The call to createReductions below expects the block to have a
- // terminator. Create an unreachable instruction to serve as terminator
- // and remove it later.
- llvm::UnreachableInst *tempTerminator = builder.CreateUnreachable();
- builder.SetInsertPoint(tempTerminator);
- llvm::OpenMPIRBuilder::InsertPointTy contInsertPoint =
- ompBuilder->createReductions(builder.saveIP(), allocaIP, reductionInfos,
- isByRef, wsloopOp.getNowait());
- if (!contInsertPoint.getBlock())
- return wsloopOp->emitOpError() << "failed to convert reductions";
- auto nextInsertionPoint =
- ompBuilder->createBarrier(contInsertPoint, llvm::omp::OMPD_for);
- tempTerminator->eraseFromParent();
- builder.restoreIP(nextInsertionPoint);
-
- // after the workshare loop, deallocate private reduction variables
- SmallVector<Region *> reductionRegions;
- llvm::transform(reductionDecls, std::back_inserter(reductionRegions),
- [](omp::DeclareReductionOp reductionDecl) {
- return &reductionDecl.getCleanupRegion();
- });
- return inlineOmpRegionCleanup(reductionRegions, privateReductionVariables,
- moduleTranslation, builder,
- "omp.reduction.cleanup");
+ return createReductionsAndCleanup(wsloopOp, builder, moduleTranslation,
+ allocaIP, reductionDecls,
+ privateReductionVariables, isByRef);
}
/// A RAII class that on construction replaces the region arguments of the
diff --git a/mlir/test/Target/LLVMIR/openmp-llvm.mlir b/mlir/test/Target/LLVMIR/openmp-llvm.mlir
index 321de67aa48a1..3a1fc5efbd823 100644
--- a/mlir/test/Target/LLVMIR/openmp-llvm.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-llvm.mlir
@@ -2093,6 +2093,8 @@ llvm.func @omp_sections_empty() -> () {
omp.sections {
omp.terminator
}
+ // CHECK-NEXT: br label %entry
+ // CHECK: entry:
// CHECK-NEXT: ret void
llvm.return
}
diff --git a/mlir/test/Target/LLVMIR/openmp-reduction-sections.mlir b/mlir/test/Target/LLVMIR/openmp-reduction-sections.mlir
new file mode 100644
index 0000000000000..d9a1597dd3e84
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-reduction-sections.mlir
@@ -0,0 +1,168 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+omp.declare_reduction @add_reduction_f32 : f32 init {
+^bb0(%arg0: f32):
+ %0 = llvm.mlir.constant(0.000000e+00 : f32) : f32
+ omp.yield(%0 : f32)
+} combiner {
+^bb0(%arg0: f32, %arg1: f32):
+ %0 = llvm.fadd %arg0, %arg1 {fastmathFlags = #llvm.fastmath<contract>} : f32
+ omp.yield(%0 : f32)
+}
+llvm.func @sections_(%arg0: !llvm.ptr {fir.bindc_name = "x"}) attributes {fir.internal_name = "_QPsections"} {
+ %0 = llvm.mlir.constant(2.000000e+00 : f32) : f32
+ %1 = llvm.mlir.constant(1.000000e+00 : f32) : f32
+ omp.parallel {
+ omp.sections reduction(@add_reduction_f32 -> %arg0 : !llvm.ptr) {
+ ^bb0(%arg1: !llvm.ptr):
+ omp.section {
+ ^bb0(%arg2: !llvm.ptr):
+ %2 = llvm.load %arg2 : !llvm.ptr -> f32
+ %3 = llvm.fadd %2, %1 {fastmathFlags = #llvm.fastmath<contract>} : f32
+ llvm.store %3, %arg2 : f32, !llvm.ptr
+ omp.terminator
+ }
+ omp.section {
+ ^bb0(%arg2: !llvm.ptr):
+ %2 = llvm.load %arg2 : !llvm.ptr -> f32
+ %3 = llvm.fadd %2, %0 {fastmathFlags = #llvm.fastmath<contract>} : f32
+ llvm.store %3, %arg2 : f32, !llvm.ptr
+ omp.terminator
+ }
+ omp.terminator
+ }
+ omp.terminator
+ }
+ llvm.return
+}
+
+// CHECK-LABEL: define void @sections_
+// CHECK: %[[VAL_0:.*]] = alloca { ptr }, align 8
+// CHECK: br label %[[VAL_1:.*]]
+// CHECK: entry: ; preds = %[[VAL_2:.*]]
+// CHECK: %[[VAL_3:.*]] = call i32 @__kmpc_global_thread_num(ptr @1)
+// CHECK: br label %[[VAL_4:.*]]
+// CHECK: omp_parallel: ; preds = %[[VAL_1]]
+// CHECK: %[[VAL_5:.*]] = getelementptr { ptr }, ptr %[[VAL_0]], i32 0, i32 0
+// CHECK: store ptr %[[VAL_6:.*]], ptr %[[VAL_5]], align 8
+// CHECK: call void (ptr, i32, ptr, ...) @__kmpc_fork_call(ptr @1, i32 1, ptr @sections_..omp_par, ptr %[[VAL_0]])
+// CHECK: br label %[[VAL_7:.*]]
+// CHECK: omp.par.outlined.exit: ; preds = %[[VAL_4]]
+// CHECK: br label %[[VAL_8:.*]]
+// CHECK: omp.par.exit.split: ; preds = %[[VAL_7]]
+// CHECK: ret void
+
+// CHECK-LABEL: define internal void @sections_..omp_par
+// CHECK: omp.par.entry:
+// CHECK: %[[VAL_9:.*]] = getelementptr { ptr }, ptr %[[VAL_10:.*]], i32 0, i32 0
+// CHECK: %[[VAL_11:.*]] = load ptr, ptr %[[VAL_9]], align 8
+// CHECK: %[[VAL_12:.*]] = alloca i32, align 4
+// CHECK: %[[VAL_13:.*]] = alloca i32, align 4
+// CHECK: %[[VAL_14:.*]] = alloca i32, align 4
+// CHECK: %[[VAL_15:.*]] = alloca i32, align 4
+// CHECK: %[[VAL_16:.*]] = alloca i32, align 4
+// CHECK: %[[VAL_17:.*]] = load i32, ptr %[[VAL_18:.*]], align 4
+// CHECK: store i32 %[[VAL_17]], ptr %[[VAL_16]], align 4
+// CHECK: %[[VAL_19:.*]] = load i32, ptr %[[VAL_16]], align 4
+// CHECK: %[[VAL_20:.*]] = alloca float, align 4
+// CHECK: %[[VAL_21:.*]] = alloca [1 x ptr], align 8
+// CHECK: br label %[[VAL_22:.*]]
+// CHECK: omp.reduction.init: ; preds = %[[VAL_23:.*]]
+// CHECK: br label %[[VAL_24:.*]]
+// CHECK: omp.par.region: ; preds = %[[VAL_22]]
+// CHECK: br label %[[VAL_25:.*]]
+// CHECK: omp.par.region1: ; preds = %[[VAL_24]]
+// CHECK: store float 0.000000e+00, ptr %[[VAL_20]], align 4
+// CHECK: br label %[[VAL_26:.*]]
+// CHECK: omp_section_loop.preheader: ; preds = %[[VAL_25]]
+// CHECK: store i32 0, ptr %[[VAL_13]], align 4
+// CHECK: store i32 1, ptr %[[VAL_14]], align 4
+// CHECK: store i32 1, ptr %[[VAL_15]], align 4
+// CHECK: %[[VAL_27:.*]] = call i32 @__kmpc_global_thread_num(ptr @1)
+// CHECK: call void @__kmpc_for_static_init_4u(ptr @1, i32 %[[VAL_27]], i32 34, ptr %[[VAL_12]], ptr %[[VAL_13]], ptr %[[VAL_14]], ptr %[[VAL_15]], i32 1, i32 0)
+// CHECK: %[[VAL_28:.*]] = load i32, ptr %[[VAL_13]], align 4
+// CHECK: %[[VAL_29:.*]] = load i32, ptr %[[VAL_14]], align 4
+// CHECK: %[[VAL_30:.*]] = sub i32 %[[VAL_29]], %[[VAL_28]]
+// CHECK: %[[VAL_31:.*]] = add i32 %[[VAL_30]], 1
+// CHECK: br label %[[VAL_32:.*]]
+// CHECK: omp_section_loop.header: ; preds = %[[VAL_33:.*]], %[[VAL_26]]
+// CHECK: %[[VAL_34:.*]] = phi i32 [ 0, %[[VAL_26]] ], [ %[[VAL_35:.*]], %[[VAL_33]] ]
+// CHECK: br label %[[VAL_36:.*]]
+// CHECK: omp_section_loop.cond: ; preds = %[[VAL_32]]
+// CHECK: %[[VAL_37:.*]] = icmp ult i32 %[[VAL_34]], %[[VAL_31]]
+// CHECK: br i1 %[[VAL_37]], label %[[VAL_38:.*]], label %[[VAL_39:.*]]
+// CHECK: omp_section_loop.exit: ; preds = %[[VAL_36]]
+// CHECK: call void @__kmpc_for_static_fini(ptr @1, i32 %[[VAL_27]])
+// CHECK: %[[VAL_40:.*]] = call i32 @__kmpc_global_thread_num(ptr @1)
+// CHECK: call void @__kmpc_barrier(ptr @2, i32 %[[VAL_40]])
+// CHECK: br label %[[VAL_41:.*]]
+// CHECK: omp_section_loop.after: ; preds = %[[VAL_39]]
+// CHECK: br label %[[VAL_42:.*]]
+// CHECK: omp_section_loop.aftersections.fini: ; preds = %[[VAL_41]]
+// CHECK: %[[VAL_43:.*]] = getelementptr inbounds [1 x ptr], ptr %[[VAL_21]], i64 0, i64 0
+// CHECK: store ptr %[[VAL_20]], ptr %[[VAL_43]], align 8
+// CHECK: %[[VAL_44:.*]] = call i32 @__kmpc_global_thread_num(ptr @1)
+// CHECK: %[[VAL_45:.*]] = call i32 @__kmpc_reduce(ptr @1, i32 %[[VAL_44]], i32 1, i64 8, ptr %[[VAL_21]], ptr @.omp.reduction.func, ptr @.gomp_critical_user_.reduction.var)
+// CHECK: switch i32 %[[VAL_45]], label %[[VAL_46:.*]] [
+// CHECK: i32 1, label %[[VAL_47:.*]]
+// CHECK: i32 2, label %[[VAL_48:.*]]
+// CHECK: ]
+// CHECK: reduce.switch.atomic: ; preds = %[[VAL_42]]
+// CHECK: unreachable
+// CHECK: reduce.switch.nonatomic: ; preds = %[[VAL_42]]
+// CHECK: %[[VAL_49:.*]] = load float, ptr %[[VAL_11]], align 4
+// CHECK: %[[VAL_50:.*]] = load float, ptr %[[VAL_20]], align 4
+// CHECK: %[[VAL_51:.*]] = fadd contract float %[[VAL_49]], %[[VAL_50]]
+// CHECK: store float %[[VAL_51]], ptr %[[VAL_11]], align 4
+// CHECK: call void @__kmpc_end_reduce(ptr @1, i32 %[[VAL_44]], ptr @.gomp_critical_user_.reduction.var)
+// CHECK: br label %[[VAL_46]]
+// CHECK: reduce.finalize: ; preds = %[[VAL_47]], %[[VAL_42]]
+// CHECK: %[[VAL_52:.*]] = call i32 @__kmpc_global_thread_num(ptr @1)
+// CHECK: call void @__kmpc_barrier(ptr @2, i32 %[[VAL_52]])
+// CHECK: br label %[[VAL_53:.*]]
+// CHECK: omp.region.cont: ; preds = %[[VAL_46]]
+// CHECK: br label %[[VAL_54:.*]]
+// CHECK: omp.par.pre_finalize: ; preds = %[[VAL_53]]
+// CHECK: br label %[[VAL_55:.*]]
+// CHECK: omp_section_loop.body: ; preds = %[[VAL_36]]
+// CHECK: %[[VAL_56:.*]] = add i32 %[[VAL_34]], %[[VAL_28]]
+// CHECK: %[[VAL_57:.*]] = mul i32 %[[VAL_56]], 1
+// CHECK: %[[VAL_58:.*]] = add i32 %[[VAL_57]], 0
+// CHECK: switch i32 %[[VAL_58]], label %[[VAL_59:.*]] [
+// CHECK: i32 0, label %[[VAL_60:.*]]
+// CHECK: i32 1, label %[[VAL_61:.*]]
+// CHECK: ]
+// CHECK: omp_section_loop.body.case3: ; preds = %[[VAL_38]]
+// CHECK: br label %[[VAL_62:.*]]
+// CHECK: omp.section.region5: ; preds = %[[VAL_61]]
+// CHECK: %[[VAL_63:.*]] = load float, ptr %[[VAL_20]], align 4
+// CHECK: %[[VAL_64:.*]] = fadd contract float %[[VAL_63]], 2.000000e+00
+// CHECK: store float %[[VAL_64]], ptr %[[VAL_20]], align 4
+// CHECK: br label %[[VAL_65:.*]]
+// CHECK: omp.region.cont4: ; preds = %[[VAL_62]]
+// CHECK: br label %[[VAL_59]]
+// CHECK: omp_section_loop.body.case: ; preds = %[[VAL_38]]
+// CHECK: br label %[[VAL_66:.*]]
+// CHECK: omp.section.region: ; preds = %[[VAL_60]]
+// CHECK: %[[VAL_67:.*]] = load float, ptr %[[VAL_20]], align 4
+// CHECK: %[[VAL_68:.*]] = fadd contract float %[[VAL_67]], 1.000000e+00
+// CHECK: store float %[[VAL_68]], ptr %[[VAL_20]], align 4
+// CHECK: br label %[[VAL_69:.*]]
+// CHECK: omp.region.cont2: ; preds = %[[VAL_66]]
+// CHECK: br label %[[VAL_59]]
+// CHECK: omp_section_loop.body.sections.after: ; preds = %[[VAL_65]], %[[VAL_69]], %[[VAL_38]]
+// CHECK: br label %[[VAL_33]]
+// CHECK: omp_section_loop.inc: ; preds = %[[VAL_59]]
+// CHECK: %[[VAL_35]] = add nuw i32 %[[VAL_34]], 1
+// CHECK: br label %[[VAL_32]]
+// CHECK: omp.par.outlined.exit.exitStub: ; preds = %[[VAL_54]]
+// CHECK: ret void
+// CHECK: %[[VAL_70:.*]] = getelementptr inbounds [1 x ptr], ptr %[[VAL_71:.*]], i64 0, i64 0
+// CHECK: %[[VAL_72:.*]] = load ptr, ptr %[[VAL_70]], align 8
+// CHECK: %[[VAL_73:.*]] = load float, ptr %[[VAL_72]], align 4
+// CHECK: %[[VAL_74:.*]] = getelementptr inbounds [1 x ptr], ptr %[[VAL_75:.*]], i64 0, i64 0
+// CHECK: %[[VAL_76:.*]] = load ptr, ptr %[[VAL_74]], align 8
+// CHECK: %[[VAL_77:.*]] = load float, ptr %[[VAL_76]], align 4
+// CHECK: %[[VAL_78:.*]] = fadd contract float %[[VAL_73]], %[[VAL_77]]
+// CHECK: store float %[[VAL_78]], ptr %[[VAL_72]], align 4
+// CHECK: ret void
More information about the Mlir-commits
mailing list