[flang-commits] [flang] [mlir] [Flang][OpenMP] Lower scan directive and inscan reduction modifier (PR #206747)
via flang-commits
flang-commits at lists.llvm.org
Tue Jun 30 08:07:18 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-flang-fir-hlfir
Author: CHANDRA GHALE (chandraghale)
<details>
<summary>Changes</summary>
This builds on top of [https://github.com/llvm/llvm-project/pull/167031](https://github.com/llvm/llvm-project/pull/167031) to complete scan reduction lowering for OpenMP workshare loops. On top of that PR, it resets each reduction variable to its identity at the start of every input-loop iteration; otherwise the scan buffer holds a running sum and emitScanReduction prefix-sums it twice, producing incorrect results.
SIMD, composite/nested loops, collapse, tiling, the linear clause and the exclusive clause are not enabled for scan reductions in this PR.
---
Patch is 43.92 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/206747.diff
6 Files Affected:
- (modified) flang/lib/Lower/OpenMP/OpenMP.cpp (+52-5)
- (added) flang/test/Lower/OpenMP/Todo/nested-wsloop-scan.f90 (+34)
- (added) flang/test/Lower/OpenMP/Todo/wsloop-scan-collapse.f90 (+29)
- (modified) mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp (+403-50)
- (added) mlir/test/Target/LLVMIR/openmp-reduction-scan.mlir (+123)
- (modified) mlir/test/Target/LLVMIR/openmp-todo.mlir (+67-2)
``````````diff
diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 99ce48206c33b..186ad074cb885 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -2750,12 +2750,59 @@ genParallelOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
static mlir::omp::ScanOp
genScanOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
- semantics::SemanticsContext &semaCtx, mlir::Location loc,
- const ConstructQueue &queue, ConstructQueue::const_iterator item) {
+ semantics::SemanticsContext &semaCtx, lower::pft::Evaluation &eval,
+ mlir::Location loc, const ConstructQueue &queue,
+ ConstructQueue::const_iterator item) {
mlir::omp::ScanOperands clauseOps;
genScanClauses(converter, semaCtx, item->clauses, loc, clauseOps);
- return mlir::omp::ScanOp::create(converter.getFirOpBuilder(),
- converter.getCurrentLocation(), clauseOps);
+ mlir::omp::ScanOp scanOp = mlir::omp::ScanOp::create(
+ converter.getFirOpBuilder(), converter.getCurrentLocation(), clauseOps);
+
+ /// Scan reduction is not implemented with nested workshare loops, linear
+ /// clause, tiling
+ mlir::omp::LoopNestOp loopNestOp =
+ scanOp->getParentOfType<mlir::omp::LoopNestOp>();
+ llvm::SmallVector<mlir::omp::LoopWrapperInterface> loopWrappers;
+ loopNestOp.gatherWrappers(loopWrappers);
+ mlir::Operation *loopWrapperOp = loopWrappers.front().getOperation();
+ if (llvm::isa<mlir::omp::SimdOp>(loopWrapperOp))
+ TODO(loc, "unsupported simd");
+ if (loopWrappers.size() > 1)
+ TODO(loc, "unsupported composite");
+ mlir::omp::WsloopOp wsLoopOp = llvm::cast<mlir::omp::WsloopOp>(loopWrapperOp);
+ bool isNested =
+ (loopNestOp.getNumLoops() > 1) ||
+ (wsLoopOp && (wsLoopOp->getParentOfType<mlir::omp::WsloopOp>()));
+ if (isNested)
+ TODO(loc, "Scan directive inside nested workshare loops");
+ if (wsLoopOp && !wsLoopOp.getLinearVars().empty())
+ TODO(loc, "Scan directive with linear clause");
+ if (loopNestOp.getTileSizes())
+ TODO(loc, "Scan directive with loop tiling");
+
+ // All loop indices should be loaded after the scan construct as otherwise,
+ // it would result in using the index variable across scan directive.
+ // (`Intra-iteration dependences from a statement in the structured
+ // block sequence that precede a scan directive to a statement in the
+ // structured block sequence that follows a scan directive must not exist,
+ // except for dependences for the list items specified in an inclusive or
+ // exclusive clause.`).
+ // TODO: Nested loops are not handled.
+ mlir::Region ®ion = loopNestOp->getRegion(0);
+ mlir::Value indexVal = fir::getBase(region.getArgument(0));
+ lower::pft::Evaluation *doConstructEval = eval.parentConstruct;
+ fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
+ lower::pft::Evaluation *doLoop = &doConstructEval->getFirstNestedEvaluation();
+ auto *doStmt = doLoop->getIf<parser::NonLabelDoStmt>();
+ assert(doStmt && "Expected do loop to be in the nested evaluation");
+ const auto &loopControl =
+ std::get<std::optional<parser::LoopControl>>(doStmt->t);
+ const parser::LoopControl::Bounds *bounds =
+ std::get_if<parser::LoopControl::Bounds>(&loopControl->u);
+ mlir::Operation *storeOp =
+ setLoopVar(converter, loc, indexVal, bounds->Name().thing.symbol);
+ firOpBuilder.setInsertionPointAfter(storeOp);
+ return scanOp;
}
static mlir::omp::SectionsOp
@@ -4079,7 +4126,7 @@ static void genOMPDispatch(lower::AbstractConverter &converter,
loc, queue, item);
break;
case llvm::omp::Directive::OMPD_scan:
- newOp = genScanOp(converter, symTable, semaCtx, loc, queue, item);
+ newOp = genScanOp(converter, symTable, semaCtx, eval, loc, queue, item);
break;
case llvm::omp::Directive::OMPD_section:
llvm_unreachable("genOMPDispatch: OMPD_section");
diff --git a/flang/test/Lower/OpenMP/Todo/nested-wsloop-scan.f90 b/flang/test/Lower/OpenMP/Todo/nested-wsloop-scan.f90
new file mode 100644
index 0000000000000..414e2ef2d5aed
--- /dev/null
+++ b/flang/test/Lower/OpenMP/Todo/nested-wsloop-scan.f90
@@ -0,0 +1,34 @@
+! Tests scan reduction behavior when used in nested workshare loops
+
+! RUN: %not_todo_cmd %flang_fc1 -emit-hlfir -fopenmp -o - %s 2>&1 | FileCheck %s
+
+program nested_scan_example
+ implicit none
+ integer, parameter :: n = 4, m = 5
+ integer :: a(n, m), b(n, m)
+ integer :: i, j
+ integer :: row_sum, col_sum
+
+ do i = 1, n
+ do j = 1, m
+ a(i, j) = i + j
+ end do
+ end do
+
+ !$omp parallel do reduction(inscan, +: row_sum) private(col_sum, j)
+ do i = 1, n
+ row_sum = row_sum + i
+ !$omp scan inclusive(row_sum)
+
+ col_sum = 0
+ !$omp parallel do reduction(inscan, +: col_sum)
+ do j = 1, m
+ col_sum = col_sum + a(i, j)
+ !CHECK: not yet implemented: Scan directive inside nested workshare loops
+ !$omp scan inclusive(col_sum)
+ b(i, j) = col_sum + row_sum
+ end do
+ !$omp end parallel do
+ end do
+ !$omp end parallel do
+end program nested_scan_example
diff --git a/flang/test/Lower/OpenMP/Todo/wsloop-scan-collapse.f90 b/flang/test/Lower/OpenMP/Todo/wsloop-scan-collapse.f90
new file mode 100644
index 0000000000000..b8e6e831884ab
--- /dev/null
+++ b/flang/test/Lower/OpenMP/Todo/wsloop-scan-collapse.f90
@@ -0,0 +1,29 @@
+! Tests scan reduction behavior when used in nested workshare loops
+
+! RUN: %not_todo_cmd %flang_fc1 -emit-hlfir -fopenmp -o - %s 2>&1 | FileCheck %s
+
+program nested_loop_example
+ implicit none
+ integer :: i, j, x
+ integer, parameter :: N = 100, M = 200
+ real :: A(N, M), B(N, M)
+ x = 0
+
+ do i = 1, N
+ do j = 1, M
+ A(i, j) = i * j
+ end do
+ end do
+
+ !$omp parallel do collapse(2) reduction(inscan, +:x)
+ do i = 1, N
+ do j = 1, M
+ x = x + A(i,j)
+ !CHECK: not yet implemented: Scan directive inside nested workshare loops
+ !$omp scan inclusive(x)
+ B(i, j) = x
+ end do
+ end do
+ !$omp end parallel do
+
+end program nested_loop_example
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 77cc7a388a984..a204937798e1c 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -41,6 +41,8 @@
#include "llvm/TargetParser/Triple.h"
#include "llvm/Transforms/Utils/ModuleUtils.h"
+#include <cassert>
+#include <cstddef>
#include <cstdint>
#include <iterator>
#include <numeric>
@@ -84,6 +86,10 @@ class OpenMPAllocStackFrame
: allocInsertPoint(allocaIP), deallocBlocks(deallocBlocks) {}
llvm::OpenMPIRBuilder::InsertPointTy allocInsertPoint;
llvm::SmallVector<llvm::BasicBlock *> deallocBlocks;
+ /// Set to true when this alloca frame encloses an omp.parallel operation.
+ /// The alloca insertion point of a function in which a parallel op is
+ /// defined may be used to allocate the temporary buffer for scan reductions.
+ bool containsParallelOp = false;
};
/// Stack frame to hold a \see llvm::CanonicalLoopInfo representing the
@@ -93,7 +99,17 @@ class OpenMPLoopInfoStackFrame
: public StateStackFrameBase<OpenMPLoopInfoStackFrame> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpenMPLoopInfoStackFrame)
+ /// For constructs like `scan`, a single `omp.loop_nest` is split into an
+ /// input loop and a scan loop. In that case `loopInfo` holds the input loop
+ /// info and `scanloopInfo` holds the scan loop info.
llvm::CanonicalLoopInfo *loopInfo = nullptr;
+ llvm::CanonicalLoopInfo *scanloopInfo = nullptr;
+ llvm::ScanInfo *scanInfo = nullptr;
+ /// Map reduction variables to their LLVM types. Populated when the reduction
+ /// clause is processed and used when a `scan` directive is encountered in the
+ /// loop body.
+ std::unique_ptr<llvm::DenseMap<llvm::Value *, llvm::Type *>>
+ reductionVarToType = nullptr;
};
/// Custom error class to signal translation errors that don't need reporting,
@@ -365,6 +381,10 @@ static LogicalResult checkImplementationStatus(Operation &op) {
if (!op.getDependVars().empty() || op.getDependKinds())
result = todo("depend");
};
+ auto checkExclusive = [&todo](auto op, LogicalResult &result) {
+ if (!op.getExclusiveVars().empty())
+ result = todo("exclusive");
+ };
auto checkHint = [](auto op, LogicalResult &) {
if (op.getHint())
op.emitWarning("hint clause discarded");
@@ -392,8 +412,17 @@ static LogicalResult checkImplementationStatus(Operation &op) {
op.getReductionSyms())
result = todo("reduction");
if (op.getReductionMod() &&
- op.getReductionMod().value() != omp::ReductionModifier::defaultmod)
- result = todo("reduction with modifier");
+ op.getReductionMod().value() != omp::ReductionModifier::defaultmod) {
+ // The `inscan` modifier on a worksharing-loop is supported (it is
+ // translated as a scan reduction). All other reduction modifiers, and
+ // `inscan` on any other operation, are not yet implemented.
+ if (isa<omp::WsloopOp>(op)) {
+ if (op.getReductionMod().value() == omp::ReductionModifier::task)
+ result = todo("reduction with task modifier");
+ } else {
+ result = todo("reduction with modifier");
+ }
+ }
};
auto checkTaskReductionByref = [&todo](auto op, LogicalResult &result) {
if (auto byrefAttr = op.getTaskReductionByref())
@@ -448,6 +477,7 @@ static LogicalResult checkImplementationStatus(Operation &op) {
checkAllocate(op, result);
checkOrder(op, result);
})
+ .Case([&](omp::ScanOp op) { checkExclusive(op, result); })
.Case([&](omp::SectionsOp op) {
checkAllocate(op, result);
checkPrivate(op, result);
@@ -638,6 +668,97 @@ findCurrentLoopInfo(LLVM::ModuleTranslation &moduleTranslation) {
return loopInfo;
}
+/// Find the scan loop information structure for the scan loop nest being
+/// translated. It will return a `null` value unless called from the
+/// translation function for a loop wrapper operation after successfully
+/// translating its body.
+static llvm::CanonicalLoopInfo *
+findCurrentScanLoopInfo(LLVM::ModuleTranslation &moduleTranslation) {
+ llvm::CanonicalLoopInfo *scanLoopInfo = nullptr;
+ moduleTranslation.stackWalk<OpenMPLoopInfoStackFrame>(
+ [&](OpenMPLoopInfoStackFrame &frame) {
+ scanLoopInfo = frame.scanloopInfo;
+ return WalkResult::interrupt();
+ });
+ return scanLoopInfo;
+}
+
+/// Find the `ScanInfo` stored on the loop stack frame. Upon encountering an
+/// `inscan` reduction modifier, `scanInfoInitialize` initializes the
+/// `ScanInfo`, which is then used when a `scan` directive is encountered in the
+/// body of the loop nest.
+static llvm::ScanInfo *
+findScanInfo(LLVM::ModuleTranslation &moduleTranslation) {
+ llvm::ScanInfo *scanInfo = nullptr;
+ moduleTranslation.stackWalk<OpenMPLoopInfoStackFrame>(
+ [&](OpenMPLoopInfoStackFrame &frame) {
+ scanInfo = frame.scanInfo;
+ return WalkResult::interrupt();
+ });
+ return scanInfo;
+}
+
+/// The types of reduction variables are used for lowering a `scan` directive
+/// that appears in the body of the loop. The types are stored in the loop frame
+/// when the reduction clause is encountered and used when the `scan` directive
+/// is encountered.
+static llvm::DenseMap<llvm::Value *, llvm::Type *> *
+findReductionVarTypes(LLVM::ModuleTranslation &moduleTranslation) {
+ llvm::DenseMap<llvm::Value *, llvm::Type *> *reductionVarToType = nullptr;
+ moduleTranslation.stackWalk<OpenMPLoopInfoStackFrame>(
+ [&](OpenMPLoopInfoStackFrame &frame) {
+ if (!frame.reductionVarToType)
+ frame.reductionVarToType =
+ std::make_unique<llvm::DenseMap<llvm::Value *, llvm::Type *>>();
+ reductionVarToType = frame.reductionVarToType.get();
+ return WalkResult::interrupt();
+ });
+ return reductionVarToType;
+}
+
+/// Scan reduction requires a shared buffer to be allocated to perform the
+/// reduction. The allocation needs to be done outside the parallel region in
+/// which the scan operation is used.
+static llvm::OpenMPIRBuilder::InsertPointTy
+findParallelAllocaIP(llvm::IRBuilderBase &builder,
+ LLVM::ModuleTranslation &moduleTranslation) {
+ // If there is an alloca insertion point on the stack belonging to a frame
+ // that encloses a parallel op, use it.
+ llvm::OpenMPIRBuilder::InsertPointTy allocaInsertPoint;
+ WalkResult walkResult = moduleTranslation.stackWalk<OpenMPAllocStackFrame>(
+ [&](OpenMPAllocStackFrame &frame) {
+ if (frame.containsParallelOp) {
+ allocaInsertPoint = frame.allocInsertPoint;
+ return WalkResult::interrupt();
+ }
+ return WalkResult::skip();
+ });
+ if (walkResult.wasInterrupted())
+ return allocaInsertPoint;
+ // Otherwise, insert into the entry block of the surrounding function.
+ // If the current IRBuilder InsertPoint is the function's entry, it cannot
+ // also be used for alloca insertion which would result in insertion order
+ // confusion. Create a new BasicBlock for the Builder and use the entry block
+ // for the allocs.
+ // TODO: Create a dedicated alloca BasicBlock at function creation such that
+ // we do not need to move the current InsertPoint here.
+ if (builder.GetInsertBlock() ==
+ &builder.GetInsertBlock()->getParent()->getEntryBlock()) {
+ assert(builder.GetInsertPoint() == builder.GetInsertBlock()->end() &&
+ "Assuming end of basic block");
+ llvm::BasicBlock *entryBB = llvm::BasicBlock::Create(
+ builder.getContext(), "entry", builder.GetInsertBlock()->getParent(),
+ builder.GetInsertBlock()->getNextNode());
+ builder.CreateBr(entryBB);
+ builder.SetInsertPoint(entryBB);
+ }
+
+ llvm::BasicBlock &funcEntryBlock =
+ builder.GetInsertBlock()->getParent()->getEntryBlock();
+ return llvm::OpenMPIRBuilder::InsertPointTy(
+ &funcEntryBlock, funcEntryBlock.getFirstInsertionPt());
+}
+
/// Converts the given region that appears within an OpenMP dialect operation to
/// LLVM IR, creating a branch from the `sourceBlock` to the entry block of the
/// region, and a branch from any block with an successor-less OpenMP terminator
@@ -1402,7 +1523,8 @@ initReductionVars(OP op, ArrayRef<BlockArgument> reductionArgs,
SmallVectorImpl<llvm::Value *> &privateReductionVariables,
DenseMap<Value, llvm::Value *> &reductionVariableMap,
llvm::ArrayRef<bool> isByRef,
- SmallVectorImpl<DeferredStore> &deferredStores) {
+ SmallVectorImpl<DeferredStore> &deferredStores,
+ bool isInScanRegion = false) {
if (op.getNumReductionVars() == 0)
return success();
@@ -1439,11 +1561,17 @@ initReductionVars(OP op, ArrayRef<BlockArgument> reductionArgs,
for (auto [data, addr] : deferredStores)
builder.CreateStore(data, addr);
+ llvm::DenseMap<llvm::Value *, llvm::Type *> *reductionVarToType =
+ findReductionVarTypes(moduleTranslation);
// 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 *, 1> phis;
+ llvm::Type *reductionType =
+ moduleTranslation.convertType(reductionDecls[i].getType());
+ if (isInScanRegion && reductionVarToType != nullptr)
+ (*reductionVarToType)[privateReductionVariables[i]] = reductionType;
// map block argument to initializer region
mapInitializationArgs(op, moduleTranslation, builder, reductionDecls,
@@ -4327,11 +4455,15 @@ convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
return failure();
assert(afterAllocas.get()->getSinglePredecessor());
+ bool isInScanRegion =
+ wsloopOp.getReductionMod() && (wsloopOp.getReductionMod().value() ==
+ mlir::omp::ReductionModifier::inscan);
if (failed(initReductionVars(wsloopOp, reductionArgs, builder,
moduleTranslation,
afterAllocas.get()->getSinglePredecessor(),
reductionDecls, privateReductionVariables,
- reductionVariableMap, isByRef, deferredStores)))
+ reductionVariableMap, isByRef, deferredStores,
+ isInScanRegion)))
return failure();
// TODO: Handle doacross loops when the ordered clause has a parameter.
@@ -4339,6 +4471,10 @@ convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
std::optional<omp::ScheduleModifier> scheduleMod = wsloopOp.getScheduleMod();
bool isSimd = wsloopOp.getScheduleSimd();
bool loopNeedsBarrier = !wsloopOp.getNowait();
+ // TODO: Linear clause support needs to be enabled for scan reduction.
+ if (isInScanRegion)
+ assert(wsloopOp.getLinearVars().empty() &&
+ "Linear clause support is not enabled with scan reduction");
// The only legal way for the direct parent to be omp.distribute is that this
// represents 'distribute parallel do'. Otherwise, this is a regular
@@ -4377,21 +4513,93 @@ convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
if (failed(handleError(regionBlock, opInst)))
return failure();
+ // Generates the loop body for a worksharing loop, including linear-variable
+ // handling and the call into the OpenMPIRBuilder's worksharing-loop helper.
+ // For scan reductions this lambda is invoked twice: once for the input loop
+ // and once for the scan loop.
+ const auto &&wsloopCodeGen = [&](llvm::CanonicalLoopInfo *loopInfo,
+ bool noLoopMode,
+ bool inputScanLoop) -> LogicalResult {
+ // Emit Initialization and Update IR for linear variables
+ if (!wsloopOp.getLinearVars().empty()) {
+ linearClauseProcessor.initLinearVar(builder, moduleTranslation,
+ loopInfo->getPreheader());
+ llvm::OpenMPIRBuilder::InsertPointOrErrorTy afterBarrierIP =
+ moduleTranslation.getOpenMPBuilder()->createBarrier(
+ builder.saveIP(), llvm::omp::OMPD_barrier);
+ if (failed(handleError(afterBarrierIP, *loopOp)))
+ return failure();
+ builder.restoreIP(*afterBarrierIP);
+ linearClauseProcessor.updateLinearVar(builder, loopInfo->getBody(),
+ loopInfo->getIndVar());
+ linearClauseProcessor.splitLinearFiniBB(builder, loopInfo->getExit());
+ }
+
+ builder.SetInsertPoint(*regionBlock, (*regionBlock)->begin());
+
+ llvm::OpenMPIRBuilder::InsertPointOrErrorTy wsloopIP =
+ ompBuilder->applyWorkshareLoop(
+ ompLoc.DL, loopInfo, allocaIP, loopNeedsBarrier,
+ convertToScheduleKind(schedule), chunk, isSimd,
+ scheduleMod == omp::ScheduleModifier::monotonic,
+ scheduleMod == omp::ScheduleModifier::nonmonotonic, isOrdered,
+ workshareLoopType, noLoopMode, hasDistSchedule, distScheduleChunk);
+
+ if (failed(handleError(wsloopIP, opInst)))
+ return failure();
+
+ // Emit finalization and in-place rewrites for linear vars.
+ if (!wsloopOp.getLinearVars().empty()) {
+ llvm::OpenMPIRBuilder::InsertPointTy oldIP = builder.saveIP();
+ assert(loopInfo->getLastIter() &&
+ "`lastiter` in CanonicalLoopInfo is nullptr");
+ llvm::OpenMPIRBuilder::InsertPointOrErrorTy afterBarrierIP =
+ linearClauseProcessor.finalizeLinearVar(builder, moduleTranslation,
+ loopInfo->getLastIter());
+ if (failed(handleError(afterBarrierIP, *loopOp)))
+ return failure();
+ for (size_t index = 0; index < wsloopOp.getLinearVars().size(); index++)
+ linearClauseProcessor.rewriteInPlace(
+ builder, sourceBlock->getSingleSuccessor(), *regionBlock,
...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/206747
More information about the flang-commits
mailing list