[flang][OpenMP] - Add MapInfoOp instances for target private variables when needed
#109862
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,156 @@ | ||
| //===- MapsForPrivatizedSymbols.cpp | ||
| //-----------------------------------------===// | ||
| // | ||
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. | ||
| // See https://llvm.org/LICENSE.txt for license information. | ||
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | ||
| // | ||
| //===----------------------------------------------------------------------===// | ||
|
|
||
| //===----------------------------------------------------------------------===// | ||
| /// \file | ||
| /// An OpenMP dialect related pass for FIR/HLFIR which creates MapInfoOp | ||
| /// instances for certain privatized symbols. | ||
| /// For example, if an allocatable variable is used in a private clause attached | ||
| /// to a omp.target op, then the allocatable variable's descriptor will be | ||
| /// needed on the device (e.g. GPU). This descriptor needs to be separately | ||
| /// mapped onto the device. This pass creates the necessary omp.map.info ops for | ||
| /// this. | ||
| //===----------------------------------------------------------------------===// | ||
| // TODO: | ||
| // 1. Before adding omp.map.info, check if in case we already have an | ||
| // omp.map.info for the variable in question. | ||
| // 2. Generalize this for more than just omp.target ops. | ||
bhandarkar-pranav marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| //===----------------------------------------------------------------------===// | ||
|
|
||
| #include "flang/Optimizer/Builder/FIRBuilder.h" | ||
| #include "flang/Optimizer/Dialect/FIRType.h" | ||
| #include "flang/Optimizer/Dialect/Support/KindMapping.h" | ||
| #include "flang/Optimizer/HLFIR/HLFIROps.h" | ||
| #include "flang/Optimizer/OpenMP/Passes.h" | ||
| #include "mlir/Dialect/Func/IR/FuncOps.h" | ||
| #include "mlir/Dialect/OpenMP/OpenMPDialect.h" | ||
| #include "mlir/IR/BuiltinAttributes.h" | ||
| #include "mlir/IR/SymbolTable.h" | ||
| #include "mlir/Pass/Pass.h" | ||
| #include "llvm/Frontend/OpenMP/OMPConstants.h" | ||
| #include "llvm/Support/Debug.h" | ||
| #include <type_traits> | ||
|
|
||
| #define DEBUG_TYPE "omp-maps-for-privatized-symbols" | ||
|
|
||
| namespace flangomp { | ||
| #define GEN_PASS_DEF_MAPSFORPRIVATIZEDSYMBOLSPASS | ||
| #include "flang/Optimizer/OpenMP/Passes.h.inc" | ||
| } // namespace flangomp | ||
| using namespace mlir; | ||
| namespace { | ||
| class MapsForPrivatizedSymbolsPass | ||
| : public flangomp::impl::MapsForPrivatizedSymbolsPassBase< | ||
| MapsForPrivatizedSymbolsPass> { | ||
|
|
||
| bool privatizerNeedsMap(omp::PrivateClauseOp &privatizer) { | ||
| Region &allocRegion = privatizer.getAllocRegion(); | ||
| Value blockArg0 = allocRegion.getArgument(0); | ||
| if (blockArg0.use_empty()) | ||
| return false; | ||
| return true; | ||
| } | ||
| omp::MapInfoOp createMapInfo(Location loc, Value var, | ||
| fir::FirOpBuilder &builder) { | ||
| uint64_t mapTypeTo = static_cast< | ||
| std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>( | ||
| llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO); | ||
| Operation *definingOp = var.getDefiningOp(); | ||
| auto declOp = llvm::dyn_cast_or_null<hlfir::DeclareOp>(definingOp); | ||
| assert(declOp && | ||
| "Expected defining Op of privatized var to be hlfir.declare"); | ||
|
|
||
| // We want the first result of the hlfir.declare op because our goal | ||
| // is to map the descriptor (fir.box or fir.boxchar) and the first | ||
| // result for hlfir.declare is the descriptor if a the symbol being | ||
| // decalred needs a descriptor. | ||
| Value varPtr = declOp.getBase(); | ||
|
|
||
| // If we do not have a reference to descritor, but the descriptor itself | ||
| // then we need to store that on the stack so that we can map the | ||
| // address of the descriptor. | ||
| if (mlir::isa<fir::BaseBoxType>(varPtr.getType()) || | ||
| mlir::isa<fir::BoxCharType>(varPtr.getType())) { | ||
| OpBuilder::InsertPoint savedInsPoint = builder.saveInsertionPoint(); | ||
| mlir::Block *allocaBlock = builder.getAllocaBlock(); | ||
| assert(allocaBlock && "No allocablock found for a funcOp"); | ||
| builder.setInsertionPointToStart(allocaBlock); | ||
| auto alloca = builder.create<fir::AllocaOp>(loc, varPtr.getType()); | ||
|
Contributor
There was a problem hiding this comment. It may be important in the future to keep track of the allocas in scope made for descriptors and re-use them where there's an alloca in scope for that variable that we've already assigned and stored too, as it lead to a couple of issues in larger examples for at least the map case, I don't know if it'd be an issue here, so it's just something to keep in mind incase you run into any issues rather than a change suggestions :-) |
||
| builder.restoreInsertionPoint(savedInsPoint); | ||
| builder.create<fir::StoreOp>(loc, varPtr, alloca); | ||
| varPtr = alloca; | ||
| } | ||
| return builder.create<omp::MapInfoOp>( | ||
| loc, varPtr.getType(), varPtr, | ||
| TypeAttr::get(llvm::cast<omp::PointerLikeType>(varPtr.getType()) | ||
| .getElementType()), | ||
| /*varPtrPtr=*/Value{}, | ||
| /*members=*/SmallVector<Value>{}, | ||
| /*member_index=*/DenseIntElementsAttr{}, | ||
| /*bounds=*/ValueRange{}, | ||
| builder.getIntegerAttr(builder.getIntegerType(64, /*isSigned=*/false), | ||
| mapTypeTo), | ||
| builder.getAttr<omp::VariableCaptureKindAttr>( | ||
| omp::VariableCaptureKind::ByRef), | ||
| StringAttr(), builder.getBoolAttr(false)); | ||
| } | ||
| void addMapInfoOp(omp::TargetOp targetOp, omp::MapInfoOp mapInfoOp) { | ||
| auto argIface = llvm::cast<omp::BlockArgOpenMPOpInterface>(*targetOp); | ||
| unsigned insertIndex = | ||
| argIface.getMapBlockArgsStart() + argIface.numMapBlockArgs(); | ||
| targetOp.getMapVarsMutable().append(ValueRange{mapInfoOp}); | ||
| targetOp.getRegion().insertArgument(insertIndex, mapInfoOp.getType(), | ||
| mapInfoOp.getLoc()); | ||
| } | ||
| void addMapInfoOps(omp::TargetOp targetOp, | ||
| llvm::SmallVectorImpl<omp::MapInfoOp> &mapInfoOps) { | ||
| for (auto mapInfoOp : mapInfoOps) | ||
| addMapInfoOp(targetOp, mapInfoOp); | ||
| } | ||
| void runOnOperation() override { | ||
| ModuleOp module = getOperation()->getParentOfType<ModuleOp>(); | ||
| fir::KindMapping kindMap = fir::getKindMapping(module); | ||
| fir::FirOpBuilder builder{module, std::move(kindMap)}; | ||
| llvm::DenseMap<Operation *, llvm::SmallVector<omp::MapInfoOp, 4>> | ||
| mapInfoOpsForTarget; | ||
|
|
||
| getOperation()->walk([&](omp::TargetOp targetOp) { | ||
| if (targetOp.getPrivateVars().empty()) | ||
| return; | ||
| OperandRange privVars = targetOp.getPrivateVars(); | ||
| std::optional<ArrayAttr> privSyms = targetOp.getPrivateSyms(); | ||
| SmallVector<omp::MapInfoOp, 4> mapInfoOps; | ||
| for (auto [privVar, privSym] : llvm::zip_equal(privVars, *privSyms)) { | ||
bhandarkar-pranav marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| SymbolRefAttr privatizerName = llvm::cast<SymbolRefAttr>(privSym); | ||
| omp::PrivateClauseOp privatizer = | ||
| SymbolTable::lookupNearestSymbolFrom<omp::PrivateClauseOp>( | ||
| targetOp, privatizerName); | ||
| if (!privatizerNeedsMap(privatizer)) { | ||
| continue; | ||
| } | ||
| builder.setInsertionPoint(targetOp); | ||
| Location loc = targetOp.getLoc(); | ||
| omp::MapInfoOp mapInfoOp = createMapInfo(loc, privVar, builder); | ||
| mapInfoOps.push_back(mapInfoOp); | ||
| LLVM_DEBUG(llvm::dbgs() << "MapsForPrivatizedSymbolsPass created ->\n"); | ||
| LLVM_DEBUG(mapInfoOp.dump()); | ||
| } | ||
| if (!mapInfoOps.empty()) { | ||
| mapInfoOpsForTarget.insert({targetOp.getOperation(), mapInfoOps}); | ||
| } | ||
| }); | ||
| if (!mapInfoOpsForTarget.empty()) { | ||
| for (auto &[targetOp, mapInfoOps] : mapInfoOpsForTarget) { | ||
| addMapInfoOps(static_cast<omp::TargetOp>(targetOp), mapInfoOps); | ||
| } | ||
| } | ||
| } | ||
| }; | ||
| } // namespace | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,48 @@ | ||
| // RUN: fir-opt --split-input-file --omp-maps-for-privatized-symbols %s | FileCheck %s | ||
| module attributes {omp.is_target_device = false} { | ||
| omp.private {type = private} @_QFtarget_simpleEsimple_var_private_ref_box_heap_i32 : !fir.ref<!fir.box<!fir.heap<i32>>> alloc { | ||
| ^bb0(%arg0: !fir.ref<!fir.box<!fir.heap<i32>>>): | ||
| %0 = fir.alloca !fir.box<!fir.heap<i32>> {bindc_name = "simple_var", pinned, uniq_name = "_QFtarget_simpleEsimple_var"} | ||
| %1 = fir.load %arg0 : !fir.ref<!fir.box<!fir.heap<i32>>> | ||
| %5:2 = hlfir.declare %0 {fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFtarget_simpleEsimple_var"} : (!fir.ref<!fir.box<!fir.heap<i32>>>) -> (!fir.ref<!fir.box<!fir.heap<i32>>>, !fir.ref<!fir.box<!fir.heap<i32>>>) | ||
| omp.yield(%5#0 : !fir.ref<!fir.box<!fir.heap<i32>>>) | ||
| } | ||
| func.func @_QPtarget_simple() { | ||
| %0 = fir.alloca i32 {bindc_name = "a", uniq_name = "_QFtarget_simpleEa"} | ||
| %1:2 = hlfir.declare %0 {uniq_name = "_QFtarget_simpleEa"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>) | ||
| %2 = fir.alloca !fir.box<!fir.heap<i32>> {bindc_name = "simple_var", uniq_name = "_QFtarget_simpleEsimple_var"} | ||
| %3 = fir.zero_bits !fir.heap<i32> | ||
| %4 = fir.embox %3 : (!fir.heap<i32>) -> !fir.box<!fir.heap<i32>> | ||
| fir.store %4 to %2 : !fir.ref<!fir.box<!fir.heap<i32>>> | ||
| %5:2 = hlfir.declare %2 {fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFtarget_simpleEsimple_var"} : (!fir.ref<!fir.box<!fir.heap<i32>>>) -> (!fir.ref<!fir.box<!fir.heap<i32>>>, !fir.ref<!fir.box<!fir.heap<i32>>>) | ||
| %c2_i32 = arith.constant 2 : i32 | ||
| hlfir.assign %c2_i32 to %1#0 : i32, !fir.ref<i32> | ||
| %6 = omp.map.info var_ptr(%1#1 : !fir.ref<i32>, i32) map_clauses(to) capture(ByRef) -> !fir.ref<i32> {name = "a"} | ||
| omp.target map_entries(%6 -> %arg0 : !fir.ref<i32>) private(@_QFtarget_simpleEsimple_var_private_ref_box_heap_i32 %5#0 -> %arg1 : !fir.ref<!fir.box<!fir.heap<i32>>>) { | ||
| %11:2 = hlfir.declare %arg0 {uniq_name = "_QFtarget_simpleEa"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>) | ||
| %12:2 = hlfir.declare %arg1 {fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFtarget_simpleEsimple_var"} : (!fir.ref<!fir.box<!fir.heap<i32>>>) -> (!fir.ref<!fir.box<!fir.heap<i32>>>, !fir.ref<!fir.box<!fir.heap<i32>>>) | ||
| %c10_i32 = arith.constant 10 : i32 | ||
| %13 = fir.load %11#0 : !fir.ref<i32> | ||
| %14 = arith.addi %c10_i32, %13 : i32 | ||
| hlfir.assign %14 to %12#0 realloc : i32, !fir.ref<!fir.box<!fir.heap<i32>>> | ||
| omp.terminator | ||
| } | ||
| %7 = fir.load %5#1 : !fir.ref<!fir.box<!fir.heap<i32>>> | ||
| %8 = fir.box_addr %7 : (!fir.box<!fir.heap<i32>>) -> !fir.heap<i32> | ||
| %9 = fir.convert %8 : (!fir.heap<i32>) -> i64 | ||
| %c0_i64 = arith.constant 0 : i64 | ||
| %10 = arith.cmpi ne, %9, %c0_i64 : i64 | ||
| fir.if %10 { | ||
| %11 = fir.load %5#1 : !fir.ref<!fir.box<!fir.heap<i32>>> | ||
| %12 = fir.box_addr %11 : (!fir.box<!fir.heap<i32>>) -> !fir.heap<i32> | ||
| fir.freemem %12 : !fir.heap<i32> | ||
| %13 = fir.zero_bits !fir.heap<i32> | ||
| %14 = fir.embox %13 : (!fir.heap<i32>) -> !fir.box<!fir.heap<i32>> | ||
| fir.store %14 to %5#1 : !fir.ref<!fir.box<!fir.heap<i32>>> | ||
| } | ||
| return | ||
| } | ||
| } | ||
| // CHECK: %[[MAP0:.*]] = omp.map.info var_ptr({{.*}} : !fir.ref<i32>, i32) map_clauses(to) capture(ByRef) -> !fir.ref<i32> {name = "a"} | ||
| // CHECK: %[[MAP1:.*]] = omp.map.info var_ptr({{.*}} : !fir.ref<!fir.box<!fir.heap<i32>>>, !fir.box<!fir.heap<i32>>) map_clauses(to) capture(ByRef) -> !fir.ref<!fir.box<!fir.heap<i32>>> | ||
| // CHECK: omp.target map_entries(%[[MAP0]] -> %arg0, %[[MAP1]] -> %arg1 : !fir.ref<i32>, !fir.ref<!fir.box<!fir.heap<i32>>>) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.