diff --git a/flang/lib/Optimizer/Transforms/SimplifyFIROperations.cpp b/flang/lib/Optimizer/Transforms/SimplifyFIROperations.cpp index 6d106046b70f2..cb9e48cced2a1 100644 --- a/flang/lib/Optimizer/Transforms/SimplifyFIROperations.cpp +++ b/flang/lib/Optimizer/Transforms/SimplifyFIROperations.cpp @@ -149,6 +149,17 @@ mlir::LogicalResult BoxTotalElementsConversion::matchAndRewrite( class DoConcurrentConversion : public mlir::OpRewritePattern { + /// Looks up from the operation from and returns the LocalitySpecifierOp with + /// name symbolName + static fir::LocalitySpecifierOp + findLocalizer(mlir::Operation *from, mlir::SymbolRefAttr symbolName) { + fir::LocalitySpecifierOp localizer = + mlir::SymbolTable::lookupNearestSymbolFrom( + from, symbolName); + assert(localizer && "localizer not found in the symbol table"); + return localizer; + } + public: using mlir::OpRewritePattern::OpRewritePattern; @@ -162,7 +173,59 @@ class DoConcurrentConversion assert(loop.getRegion().hasOneBlock()); mlir::Block &loopBlock = loop.getRegion().getBlocks().front(); - // Collect iteration variable(s) allocations do that we can move them + // Handle localization + if (!loop.getLocalVars().empty()) { + mlir::OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointToStart(&loop.getRegion().front()); + + std::optional localSyms = loop.getLocalSyms(); + + for (auto [localVar, localArg, localizerSym] : llvm::zip_equal( + loop.getLocalVars(), loop.getRegionLocalArgs(), *localSyms)) { + mlir::SymbolRefAttr localizerName = + llvm::cast(localizerSym); + fir::LocalitySpecifierOp localizer = findLocalizer(loop, localizerName); + + if (!localizer.getInitRegion().empty() || + !localizer.getDeallocRegion().empty()) + TODO(localizer.getLoc(), "localizers with `init` and `dealloc` " + "regions are not handled yet."); + + // TODO Should this be a heap allocation instead? For now, we allocate + // on the stack for each loop iteration. + mlir::Value localAlloc = + rewriter.create(loop.getLoc(), localizer.getType()); + + if (localizer.getLocalitySpecifierType() == + fir::LocalitySpecifierType::LocalInit) { + // It is reasonable to make this assumption since, at this stage, + // control-flow ops are not converted yet. Therefore, things like `if` + // conditions will still be represented by their encapsulating `fir` + // dialect ops. + assert(localizer.getCopyRegion().hasOneBlock() && + "Expected localizer to have a single block."); + mlir::Block *beforeLocalInit = rewriter.getInsertionBlock(); + mlir::Block *afterLocalInit = rewriter.splitBlock( + rewriter.getInsertionBlock(), rewriter.getInsertionPoint()); + rewriter.cloneRegionBefore(localizer.getCopyRegion(), afterLocalInit); + mlir::Block *copyRegionBody = beforeLocalInit->getNextNode(); + + rewriter.eraseOp(copyRegionBody->getTerminator()); + rewriter.mergeBlocks(afterLocalInit, copyRegionBody); + rewriter.mergeBlocks(copyRegionBody, beforeLocalInit, + {localVar, localArg}); + } + + rewriter.replaceAllUsesWith(localArg, localAlloc); + } + + loop.getRegion().front().eraseArguments(loop.getNumInductionVars(), + loop.getNumLocalOperands()); + loop.getLocalVarsMutable().clear(); + loop.setLocalSymsAttr(nullptr); + } + + // Collect iteration variable(s) allocations so that we can move them // outside the `fir.do_concurrent` wrapper. llvm::SmallVector opsToMove; for (mlir::Operation &op : llvm::drop_end(wrapperBlock)) diff --git a/flang/test/Transforms/do_concurrent-to-do_loop-unodered.fir b/flang/test/Transforms/do_concurrent-to-do_loop-unodered.fir index d2ceafdda5b22..d9ef36b175598 100644 --- a/flang/test/Transforms/do_concurrent-to-do_loop-unodered.fir +++ b/flang/test/Transforms/do_concurrent-to-do_loop-unodered.fir @@ -121,3 +121,64 @@ func.func @dc_2d_reduction(%i_lb: index, %i_ub: index, %i_st: index, // CHECK: } // CHECK: return // CHECK: } + +// ----- + +fir.local {type = local} @local_localizer : i32 + +fir.local {type = local_init} @local_init_localizer : i32 copy { +^bb0(%arg0: !fir.ref, %arg1: !fir.ref): + %0 = fir.load %arg0 : !fir.ref + fir.store %0 to %arg1 : !fir.ref + fir.yield(%arg1 : !fir.ref) +} + +func.func @do_concurrent_locality_specs() { + %3 = fir.alloca i32 {bindc_name = "local_init_var", uniq_name = "_QFdo_concurrentElocal_init_var"} + %4:2 = hlfir.declare %3 {uniq_name = "_QFdo_concurrentElocal_init_var"} : (!fir.ref) -> (!fir.ref, !fir.ref) + %5 = fir.alloca i32 {bindc_name = "local_var", uniq_name = "_QFdo_concurrentElocal_var"} + %6:2 = hlfir.declare %5 {uniq_name = "_QFdo_concurrentElocal_var"} : (!fir.ref) -> (!fir.ref, !fir.ref) + %c1 = arith.constant 1 : index + %c10 = arith.constant 1 : index + fir.do_concurrent { + %9 = fir.alloca i32 {bindc_name = "i"} + %10:2 = hlfir.declare %9 {uniq_name = "_QFdo_concurrentEi"} : (!fir.ref) -> (!fir.ref, !fir.ref) + fir.do_concurrent.loop (%arg0) = (%c1) to (%c10) step (%c1) local(@local_localizer %6#0 -> %arg1, @local_init_localizer %4#0 -> %arg2 : !fir.ref, !fir.ref) { + %11 = fir.convert %arg0 : (index) -> i32 + fir.store %11 to %10#0 : !fir.ref + %13:2 = hlfir.declare %arg1 {uniq_name = "_QFdo_concurrentElocal_var"} : (!fir.ref) -> (!fir.ref, !fir.ref) + %15:2 = hlfir.declare %arg2 {uniq_name = "_QFdo_concurrentElocal_init_var"} : (!fir.ref) -> (!fir.ref, !fir.ref) + %17 = fir.load %10#0 : !fir.ref + %c5_i32 = arith.constant 5 : i32 + %18 = arith.cmpi slt, %17, %c5_i32 : i32 + fir.if %18 { + %c42_i32 = arith.constant 42 : i32 + hlfir.assign %c42_i32 to %13#0 : i32, !fir.ref + } else { + %c84_i32 = arith.constant 84 : i32 + hlfir.assign %c84_i32 to %15#0 : i32, !fir.ref + } + } + } + return +} + +// CHECK-LABEL: func.func @do_concurrent_locality_specs() { +// CHECK: %[[LOC_INIT_DECL:.*]]:2 = hlfir.declare %{{.*}} {uniq_name = "{{.*}}Elocal_init_var"} +// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} unordered { +// Verify localization of the `local` var. +// CHECK: %[[PRIV_LOC_ALLOC:.*]] = fir.alloca i32 + +// Verify localization of the `local_init` var. +// CHECK: %[[PRIV_LOC_INIT_ALLOC:.*]] = fir.alloca i32 +// CHECK: %[[LOC_INIT_VAL:.*]] = fir.load %[[LOC_INIT_DECL]]#0 : !fir.ref +// CHECK: fir.store %[[LOC_INIT_VAL]] to %[[PRIV_LOC_INIT_ALLOC]] : !fir.ref + +// CHECK: %[[VAL_15:.*]]:2 = hlfir.declare %[[PRIV_LOC_ALLOC]] +// CHECK: %[[VAL_16:.*]]:2 = hlfir.declare %[[PRIV_LOC_INIT_ALLOC]] + +// CHECK: hlfir.assign %{{.*}} to %[[VAL_15]]#0 : i32, !fir.ref +// CHECK: hlfir.assign %{{.*}} to %[[VAL_16]]#0 : i32, !fir.ref +// CHECK: } +// CHECK: return +// CHECK: }