Skip to content

Commit e12b46f

Browse files
authored
[flang] support fir.alloca operations inside of omp reduction ops (#84952)
Advise to place the alloca at the start of the first block of whichever region (init or combiner) we are currently inside. It probably isn't safe to put an alloca inside of a combiner region because this will be executed multiple times. But that would be a bug to fix in Lower/OpenMP.cpp, not here. OpenMP array reductions 1/6 Next PR: #84953
1 parent f623adb commit e12b46f

File tree

3 files changed

+47
-2
lines changed

3 files changed

+47
-2
lines changed

flang/lib/Optimizer/Builder/FIRBuilder.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,8 @@ mlir::Block *fir::FirOpBuilder::getAllocaBlock() {
208208
.getParentOfType<mlir::omp::OutlineableOpenMPOpInterface>()) {
209209
return ompOutlineableIface.getAllocaBlock();
210210
}
211+
if (mlir::isa<mlir::omp::ReductionDeclareOp>(getRegion().getParentOp()))
212+
return &getRegion().front();
211213
if (auto accRecipeIface =
212214
getRegion().getParentOfType<mlir::acc::RecipeInterface>()) {
213215
return accRecipeIface.getAllocaBlock(getRegion());

flang/lib/Optimizer/CodeGen/CodeGen.cpp

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -410,8 +410,15 @@ class FIROpConversion : public mlir::ConvertOpToLLVMPattern<FromOp> {
410410
mlir::ConversionPatternRewriter &rewriter) const {
411411
auto thisPt = rewriter.saveInsertionPoint();
412412
mlir::Operation *parentOp = rewriter.getInsertionBlock()->getParentOp();
413-
mlir::Block *insertBlock = getBlockForAllocaInsert(parentOp);
414-
rewriter.setInsertionPointToStart(insertBlock);
413+
if (mlir::isa<mlir::omp::ReductionDeclareOp>(parentOp)) {
414+
// ReductionDeclareOp has multiple child regions. We want to get the first
415+
// block of whichever of those regions we are currently in
416+
mlir::Region *parentRegion = rewriter.getInsertionBlock()->getParent();
417+
rewriter.setInsertionPointToStart(&parentRegion->front());
418+
} else {
419+
mlir::Block *insertBlock = getBlockForAllocaInsert(parentOp);
420+
rewriter.setInsertionPointToStart(insertBlock);
421+
}
415422
auto size = genI32Constant(loc, rewriter, 1);
416423
unsigned allocaAs = getAllocaAddressSpace(rewriter);
417424
unsigned programAs = getProgramAddressSpace(rewriter);
Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
// RUN: tco %s | FileCheck %s
2+
3+
// the fir.embox in the init region is turned into an alloca for the box. Test
4+
// that CodeGen.cpp knows where to place an alloca when it is inside of an
5+
// omp.reduction.declare
6+
7+
// regretably this has to be nonsense IR because we need the subsequent patches
8+
// to process anything useful
9+
10+
omp.reduction.declare @test_reduction : !fir.ref<!fir.box<i32>> init {
11+
^bb0(%arg0: !fir.ref<!fir.box<i32>>):
12+
%0 = fir.alloca !fir.box<i32>
13+
%1 = fir.alloca i32
14+
%2 = fir.embox %1 : (!fir.ref<i32>) -> !fir.box<i32>
15+
16+
// use the embox for something so it isn't removed
17+
fir.store %2 to %0 : !fir.ref<!fir.box<i32>>
18+
19+
omp.yield(%0 : !fir.ref<!fir.box<i32>>)
20+
} combiner {
21+
^bb0(%arg0: !fir.ref<!fir.box<i32>>, %arg1: !fir.ref<!fir.box<i32>>):
22+
%0 = fir.undefined !fir.ref<!fir.box<i32>>
23+
omp.yield(%0 : !fir.ref<!fir.box<i32>>)
24+
}
25+
26+
func.func @_QQmain() attributes {fir.bindc_name = "reduce"} {
27+
%4 = fir.alloca !fir.box<i32>
28+
omp.parallel byref reduction(@test_reduction %4 -> %arg0 : !fir.ref<!fir.box<i32>>) {
29+
omp.terminator
30+
}
31+
return
32+
}
33+
34+
// basically we are testing that there isn't a crash
35+
// CHECK-LABEL: define void @_QQmain
36+
// CHECK-NEXT: alloca { ptr, i64, i32, i8, i8, i8, i8 }, i64 1, align 8

0 commit comments

Comments
 (0)