1 | //===- OMPDescriptorMapInfoGen.cpp |
2 | //---------------------------------------------------===// |
3 | // |
4 | // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
5 | // See https://llvm.org/LICENSE.txt for license information. |
6 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
7 | // |
8 | //===----------------------------------------------------------------------===// |
9 | |
10 | //===----------------------------------------------------------------------===// |
11 | /// \file |
12 | /// An OpenMP dialect related pass for FIR/HLFIR which expands MapInfoOp's |
13 | /// containing descriptor related types (fir::BoxType's) into multiple |
14 | /// MapInfoOp's containing the parent descriptor and pointer member components |
15 | /// for individual mapping, treating the descriptor type as a record type for |
16 | /// later lowering in the OpenMP dialect. |
17 | //===----------------------------------------------------------------------===// |
18 | |
19 | #include "flang/Optimizer/Builder/FIRBuilder.h" |
20 | #include "flang/Optimizer/Dialect/FIRType.h" |
21 | #include "flang/Optimizer/Dialect/Support/KindMapping.h" |
22 | #include "flang/Optimizer/Transforms/Passes.h" |
23 | #include "mlir/Dialect/Func/IR/FuncOps.h" |
24 | #include "mlir/Dialect/OpenMP/OpenMPDialect.h" |
25 | #include "mlir/IR/BuiltinDialect.h" |
26 | #include "mlir/IR/BuiltinOps.h" |
27 | #include "mlir/IR/Operation.h" |
28 | #include "mlir/IR/SymbolTable.h" |
29 | #include "mlir/Pass/Pass.h" |
30 | #include "mlir/Support/LLVM.h" |
31 | #include "llvm/ADT/SmallPtrSet.h" |
32 | #include <iterator> |
33 | |
34 | namespace fir { |
35 | #define GEN_PASS_DEF_OMPDESCRIPTORMAPINFOGENPASS |
36 | #include "flang/Optimizer/Transforms/Passes.h.inc" |
37 | } // namespace fir |
38 | |
39 | namespace { |
40 | class OMPDescriptorMapInfoGenPass |
41 | : public fir::impl::OMPDescriptorMapInfoGenPassBase< |
42 | OMPDescriptorMapInfoGenPass> { |
43 | |
44 | void genDescriptorMemberMaps(mlir::omp::MapInfoOp op, |
45 | fir::FirOpBuilder &builder, |
46 | mlir::Operation *target) { |
47 | mlir::Location loc = builder.getUnknownLoc(); |
48 | mlir::Value descriptor = op.getVarPtr(); |
49 | |
50 | // If we enter this function, but the mapped type itself is not the |
51 | // descriptor, then it's likely the address of the descriptor so we |
52 | // must retrieve the descriptor SSA. |
53 | if (!fir::isTypeWithDescriptor(op.getVarType())) { |
54 | if (auto addrOp = mlir::dyn_cast_if_present<fir::BoxAddrOp>( |
55 | op.getVarPtr().getDefiningOp())) { |
56 | descriptor = addrOp.getVal(); |
57 | } |
58 | } |
59 | |
60 | // The fir::BoxOffsetOp only works with !fir.ref<!fir.box<...>> types, as |
61 | // allowing it to access non-reference box operations can cause some |
62 | // problematic SSA IR. However, in the case of assumed shape's the type |
63 | // is not a !fir.ref, in these cases to retrieve the appropriate |
64 | // !fir.ref<!fir.box<...>> to access the data we need to map we must |
65 | // perform an alloca and then store to it and retrieve the data from the new |
66 | // alloca. |
67 | if (mlir::isa<fir::BaseBoxType>(descriptor.getType())) { |
68 | mlir::OpBuilder::InsertPoint insPt = builder.saveInsertionPoint(); |
69 | builder.setInsertionPointToStart(builder.getAllocaBlock()); |
70 | auto alloca = builder.create<fir::AllocaOp>(loc, descriptor.getType()); |
71 | builder.restoreInsertionPoint(insPt); |
72 | builder.create<fir::StoreOp>(loc, descriptor, alloca); |
73 | descriptor = alloca; |
74 | } |
75 | |
76 | mlir::Value baseAddrAddr = builder.create<fir::BoxOffsetOp>( |
77 | loc, descriptor, fir::BoxFieldAttr::base_addr); |
78 | |
79 | // Member of the descriptor pointing at the allocated data |
80 | mlir::Value baseAddr = builder.create<mlir::omp::MapInfoOp>( |
81 | loc, baseAddrAddr.getType(), descriptor, |
82 | llvm::cast<mlir::omp::PointerLikeType>( |
83 | fir::unwrapRefType(baseAddrAddr.getType())) |
84 | .getElementType(), |
85 | baseAddrAddr, mlir::SmallVector<mlir::Value>{}, op.getBounds(), |
86 | builder.getIntegerAttr(builder.getIntegerType(64, false), |
87 | op.getMapType().value()), |
88 | builder.getAttr<mlir::omp::VariableCaptureKindAttr>( |
89 | mlir::omp::VariableCaptureKind::ByRef), |
90 | builder.getStringAttr("" ) /*name*/); |
91 | |
92 | // TODO: map the addendum segment of the descriptor, similarly to the |
93 | // above base address/data pointer member. |
94 | |
95 | if (auto mapClauseOwner = |
96 | llvm::dyn_cast<mlir::omp::MapClauseOwningOpInterface>(target)) { |
97 | llvm::SmallVector<mlir::Value> newMapOps; |
98 | mlir::OperandRange mapOperandsArr = mapClauseOwner.getMapOperands(); |
99 | |
100 | for (size_t i = 0; i < mapOperandsArr.size(); ++i) { |
101 | if (mapOperandsArr[i] == op) { |
102 | // Push new implicit maps generated for the descriptor. |
103 | newMapOps.push_back(baseAddr); |
104 | |
105 | // for TargetOp's which have IsolatedFromAbove we must align the |
106 | // new additional map operand with an appropriate BlockArgument, |
107 | // as the printing and later processing currently requires a 1:1 |
108 | // mapping of BlockArgs to MapInfoOp's at the same placement in |
109 | // each array (BlockArgs and MapOperands). |
110 | if (auto targetOp = llvm::dyn_cast<mlir::omp::TargetOp>(target)) |
111 | targetOp.getRegion().insertArgument(i, baseAddr.getType(), loc); |
112 | } |
113 | newMapOps.push_back(mapOperandsArr[i]); |
114 | } |
115 | mapClauseOwner.getMapOperandsMutable().assign(newMapOps); |
116 | } |
117 | |
118 | mlir::Value newDescParentMapOp = builder.create<mlir::omp::MapInfoOp>( |
119 | op->getLoc(), op.getResult().getType(), descriptor, |
120 | fir::unwrapRefType(descriptor.getType()), mlir::Value{}, |
121 | mlir::SmallVector<mlir::Value>{baseAddr}, |
122 | mlir::SmallVector<mlir::Value>{}, |
123 | builder.getIntegerAttr(builder.getIntegerType(64, false), |
124 | op.getMapType().value()), |
125 | op.getMapCaptureTypeAttr(), op.getNameAttr()); |
126 | op.replaceAllUsesWith(newDescParentMapOp); |
127 | op->erase(); |
128 | } |
129 | |
130 | // This pass executes on mlir::ModuleOp's finding omp::MapInfoOp's containing |
131 | // descriptor based types (allocatables, pointers, assumed shape etc.) and |
132 | // expanding them into multiple omp::MapInfoOp's for each pointer member |
133 | // contained within the descriptor. |
134 | void runOnOperation() override { |
135 | mlir::func::FuncOp func = getOperation(); |
136 | mlir::ModuleOp module = func->getParentOfType<mlir::ModuleOp>(); |
137 | fir::KindMapping kindMap = fir::getKindMapping(module); |
138 | fir::FirOpBuilder builder{module, std::move(kindMap)}; |
139 | |
140 | func->walk([&](mlir::omp::MapInfoOp op) { |
141 | if (fir::isTypeWithDescriptor(op.getVarType()) || |
142 | mlir::isa_and_present<fir::BoxAddrOp>( |
143 | op.getVarPtr().getDefiningOp())) { |
144 | builder.setInsertionPoint(op); |
145 | // TODO: Currently only supports a single user for the MapInfoOp, this |
146 | // is fine for the moment as the Fortran Frontend will generate a |
147 | // new MapInfoOp per Target operation for the moment. However, when/if |
148 | // we optimise/cleanup the IR, it likely isn't too difficult to |
149 | // extend this function, it would require some modification to create a |
150 | // single new MapInfoOp per new MapInfoOp generated and share it across |
151 | // all users appropriately, making sure to only add a single member link |
152 | // per new generation for the original originating descriptor MapInfoOp. |
153 | assert(llvm::hasSingleElement(op->getUsers()) && |
154 | "OMPDescriptorMapInfoGen currently only supports single users " |
155 | "of a MapInfoOp" ); |
156 | genDescriptorMemberMaps(op, builder, *op->getUsers().begin()); |
157 | } |
158 | }); |
159 | } |
160 | }; |
161 | |
162 | } // namespace |
163 | |
164 | namespace fir { |
165 | std::unique_ptr<mlir::Pass> createOMPDescriptorMapInfoGenPass() { |
166 | return std::make_unique<OMPDescriptorMapInfoGenPass>(); |
167 | } |
168 | } // namespace fir |
169 | |