Skip to content

Commit f73bf74

Browse files
ergawyergawy
committed
[flang] Basic lowering of fir.do_concurrent.loop locality specifiers
1 parent 261d867 commit f73bf74

File tree

2 files changed

+125
-1
lines changed

2 files changed

+125
-1
lines changed

‎flang/lib/Optimizer/Transforms/SimplifyFIROperations.cpp‎

Lines changed: 63 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,20 @@ mlir::LogicalResult BoxTotalElementsConversion::matchAndRewrite(
149149

150150
class DoConcurrentConversion
151151
: public mlir::OpRewritePattern<fir::DoConcurrentOp> {
152+
/// Looks up from the operation from and returns the PrivateClauseOp with
153+
/// name symbolName
154+
///
155+
/// TODO Copied from OpenMPToLLVMIRTranslation.cpp, move to a shared location.
156+
/// Maybe a static function on the `PrivateClauseOp`.
157+
static mlir::omp::PrivateClauseOp
158+
findPrivatizer(mlir::Operation *from, mlir::SymbolRefAttr symbolName) {
159+
mlir::omp::PrivateClauseOp privatizer =
160+
mlir::SymbolTable::lookupNearestSymbolFrom<mlir::omp::PrivateClauseOp>(
161+
from, symbolName);
162+
assert(privatizer && "privatizer not found in the symbol table");
163+
return privatizer;
164+
}
165+
152166
public:
153167
using mlir::OpRewritePattern<fir::DoConcurrentOp>::OpRewritePattern;
154168

@@ -162,7 +176,55 @@ class DoConcurrentConversion
162176
assert(loop.getRegion().hasOneBlock());
163177
mlir::Block &loopBlock = loop.getRegion().getBlocks().front();
164178

165-
// Collect iteration variable(s) allocations do that we can move them
179+
// Handle privatization
180+
if (!loop.getPrivateVars().empty()) {
181+
mlir::OpBuilder::InsertionGuard guard(rewriter);
182+
rewriter.setInsertionPointToStart(&loop.getRegion().front());
183+
184+
std::optional<mlir::ArrayAttr> privateSyms = loop.getPrivateSyms();
185+
186+
for (auto [privateVar, privateArg, privatizerSym] :
187+
llvm::zip_equal(loop.getPrivateVars(), loop.getRegionPrivateArgs(),
188+
*privateSyms)) {
189+
mlir::SymbolRefAttr privatizerName =
190+
llvm::cast<mlir::SymbolRefAttr>(privatizerSym);
191+
mlir::omp::PrivateClauseOp privatizer =
192+
findPrivatizer(loop, privatizerName);
193+
194+
mlir::Value localAlloc =
195+
rewriter.create<fir::AllocaOp>(loop.getLoc(), privatizer.getType());
196+
197+
if (privatizer.getDataSharingType() ==
198+
mlir::omp::DataSharingClauseType::FirstPrivate) {
199+
// It is reasonable to make this assumption since, at this stage,
200+
// control-flow ops are not converted yet. Therefore, things like `if`
201+
// conditions will still be represented by their encapsulating `fir`
202+
// dialect ops.
203+
assert(privatizer.getCopyRegion().hasOneBlock() &&
204+
"Expected privatizer to have a single block.");
205+
mlir::Block *beforeLocalInit = rewriter.getInsertionBlock();
206+
mlir::Block *afterLocalInit = rewriter.splitBlock(
207+
rewriter.getInsertionBlock(), rewriter.getInsertionPoint());
208+
rewriter.cloneRegionBefore(privatizer.getCopyRegion(),
209+
afterLocalInit);
210+
mlir::Block *copyRegionBody = beforeLocalInit->getNextNode();
211+
212+
rewriter.eraseOp(copyRegionBody->getTerminator());
213+
rewriter.mergeBlocks(afterLocalInit, copyRegionBody);
214+
rewriter.mergeBlocks(copyRegionBody, beforeLocalInit,
215+
{privateVar, privateArg});
216+
}
217+
218+
rewriter.replaceAllUsesWith(privateArg, localAlloc);
219+
}
220+
221+
loop.getRegion().front().eraseArguments(loop.getNumInductionVars(),
222+
loop.getNumPrivateOperands());
223+
loop.getPrivateVarsMutable().clear();
224+
loop.setPrivateSymsAttr(nullptr);
225+
}
226+
227+
// Collect iteration variable(s) allocations so that we can move them
166228
// outside the `fir.do_concurrent` wrapper.
167229
llvm::SmallVector<mlir::Operation *> opsToMove;
168230
for (mlir::Operation &op : llvm::drop_end(wrapperBlock))

‎flang/test/Transforms/do_concurrent-to-do_loop-unodered.fir‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,3 +121,65 @@ func.func @dc_2d_reduction(%i_lb: index, %i_ub: index, %i_st: index,
121121
// CHECK: }
122122
// CHECK: return
123123
// CHECK: }
124+
125+
// -----
126+
127+
omp.private {type = private} @local_privatizer : i32
128+
129+
omp.private {type = firstprivate} @local_init_privatizer : i32 copy {
130+
^bb0(%arg0: !fir.ref<i32>, %arg1: !fir.ref<i32>):
131+
%0 = fir.load %arg0 : !fir.ref<i32>
132+
fir.store %0 to %arg1 : !fir.ref<i32>
133+
omp.yield(%arg1 : !fir.ref<i32>)
134+
}
135+
136+
func.func @do_concurrent_locality_specs() {
137+
%3 = fir.alloca i32 {bindc_name = "local_init_var", uniq_name = "_QFdo_concurrentElocal_init_var"}
138+
%4:2 = hlfir.declare %3 {uniq_name = "_QFdo_concurrentElocal_init_var"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
139+
%5 = fir.alloca i32 {bindc_name = "local_var", uniq_name = "_QFdo_concurrentElocal_var"}
140+
%6:2 = hlfir.declare %5 {uniq_name = "_QFdo_concurrentElocal_var"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
141+
%c1 = arith.constant 1 : index
142+
%c10 = arith.constant 1 : index
143+
fir.do_concurrent {
144+
%9 = fir.alloca i32 {bindc_name = "i"}
145+
%10:2 = hlfir.declare %9 {uniq_name = "_QFdo_concurrentEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
146+
fir.do_concurrent.loop (%arg0) = (%c1) to (%c10) step (%c1) private(@local_privatizer %6#0 -> %arg1, @local_init_privatizer %4#0 -> %arg2 : !fir.ref<i32>, !fir.ref<i32>) {
147+
%11 = fir.convert %arg0 : (index) -> i32
148+
fir.store %11 to %10#0 : !fir.ref<i32>
149+
%13:2 = hlfir.declare %arg1 {uniq_name = "_QFdo_concurrentElocal_var"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
150+
%15:2 = hlfir.declare %arg2 {uniq_name = "_QFdo_concurrentElocal_init_var"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
151+
%17 = fir.load %10#0 : !fir.ref<i32>
152+
%c5_i32 = arith.constant 5 : i32
153+
%18 = arith.cmpi slt, %17, %c5_i32 : i32
154+
fir.if %18 {
155+
%c42_i32 = arith.constant 42 : i32
156+
hlfir.assign %c42_i32 to %13#0 : i32, !fir.ref<i32>
157+
} else {
158+
%c84_i32 = arith.constant 84 : i32
159+
hlfir.assign %c84_i32 to %15#0 : i32, !fir.ref<i32>
160+
}
161+
}
162+
}
163+
return
164+
}
165+
166+
// CHECK-LABEL: func.func @do_concurrent_locality_specs() {
167+
// CHECK: %[[LOC_INIT_DECL:.*]]:2 = hlfir.declare %{{.*}} {uniq_name = "{{.*}}Elocal_init_var"}
168+
// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} unordered {
169+
// Verify privatization of the `local` var.
170+
// CHECK: %[[PRIV_LOC_ALLOC:.*]] = fir.alloca i32
171+
172+
// Verify privatization of the `local_init` var.
173+
// CHECK: %[[PRIV_LOC_INIT_ALLOC:.*]] = fir.alloca i32
174+
// CHECK: %[[LOC_INIT_VAL:.*]] = fir.load %[[LOC_INIT_DECL]]#0 : !fir.ref<i32>
175+
// CHECK: fir.store %[[LOC_INIT_VAL]] to %[[PRIV_LOC_INIT_ALLOC]] : !fir.ref<i32>
176+
177+
// CHECK: %[[VAL_15:.*]]:2 = hlfir.declare %[[PRIV_LOC_ALLOC]]
178+
// CHECK: %[[VAL_16:.*]]:2 = hlfir.declare %[[PRIV_LOC_INIT_ALLOC]]
179+
180+
// CHECK: hlfir.assign %{{.*}} to %[[VAL_15]]#0 : i32, !fir.ref<i32>
181+
// CHECK: hlfir.assign %{{.*}} to %[[VAL_16]]#0 : i32, !fir.ref<i32>
182+
// CHECK: }
183+
// CHECK: return
184+
// CHECK: }
185+

0 commit comments

Comments
 (0)