1 | //===-- TargetRewrite.cpp -------------------------------------------------===// |
2 | // |
3 | // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
4 | // See https://llvm.org/LICENSE.txt for license information. |
5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
6 | // |
7 | //===----------------------------------------------------------------------===// |
8 | // |
9 | // Target rewrite: rewriting of ops to make target-specific lowerings manifest. |
10 | // LLVM expects different lowering idioms to be used for distinct target |
11 | // triples. These distinctions are handled by this pass. |
12 | // |
13 | // Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/ |
14 | // |
15 | //===----------------------------------------------------------------------===// |
16 | |
17 | #include "flang/Optimizer/CodeGen/CodeGen.h" |
18 | |
19 | #include "flang/Optimizer/Builder/Character.h" |
20 | #include "flang/Optimizer/Builder/FIRBuilder.h" |
21 | #include "flang/Optimizer/Builder/Todo.h" |
22 | #include "flang/Optimizer/CodeGen/Target.h" |
23 | #include "flang/Optimizer/Dialect/FIRDialect.h" |
24 | #include "flang/Optimizer/Dialect/FIROps.h" |
25 | #include "flang/Optimizer/Dialect/FIROpsSupport.h" |
26 | #include "flang/Optimizer/Dialect/FIRType.h" |
27 | #include "flang/Optimizer/Dialect/Support/FIRContext.h" |
28 | #include "flang/Optimizer/Support/DataLayout.h" |
29 | #include "mlir/Dialect/DLTI/DLTI.h" |
30 | #include "mlir/Dialect/GPU/IR/GPUDialect.h" |
31 | #include "mlir/Dialect/LLVMIR/LLVMDialect.h" |
32 | #include "mlir/Transforms/DialectConversion.h" |
33 | #include "llvm/ADT/STLExtras.h" |
34 | #include "llvm/ADT/TypeSwitch.h" |
35 | #include "llvm/Support/Debug.h" |
36 | #include <optional> |
37 | |
38 | namespace fir { |
39 | #define GEN_PASS_DEF_TARGETREWRITEPASS |
40 | #include "flang/Optimizer/CodeGen/CGPasses.h.inc" |
41 | } // namespace fir |
42 | |
43 | #define DEBUG_TYPE "flang-target-rewrite" |
44 | |
45 | namespace { |
46 | |
47 | /// Fixups for updating a FuncOp's arguments and return values. |
48 | struct FixupTy { |
49 | enum class Codes { |
50 | ArgumentAsLoad, |
51 | ArgumentType, |
52 | CharPair, |
53 | ReturnAsStore, |
54 | ReturnType, |
55 | Split, |
56 | Trailing, |
57 | TrailingCharProc |
58 | }; |
59 | |
60 | FixupTy(Codes code, std::size_t index, std::size_t second = 0) |
61 | : code{code}, index{index}, second{second} {} |
62 | FixupTy(Codes code, std::size_t index, |
63 | std::function<void(mlir::func::FuncOp)> &&finalizer) |
64 | : code{code}, index{index}, finalizer{finalizer} {} |
65 | FixupTy(Codes code, std::size_t index, |
66 | std::function<void(mlir::gpu::GPUFuncOp)> &&finalizer) |
67 | : code{code}, index{index}, gpuFinalizer{finalizer} {} |
68 | FixupTy(Codes code, std::size_t index, std::size_t second, |
69 | std::function<void(mlir::func::FuncOp)> &&finalizer) |
70 | : code{code}, index{index}, second{second}, finalizer{finalizer} {} |
71 | FixupTy(Codes code, std::size_t index, std::size_t second, |
72 | std::function<void(mlir::gpu::GPUFuncOp)> &&finalizer) |
73 | : code{code}, index{index}, second{second}, gpuFinalizer{finalizer} {} |
74 | |
75 | Codes code; |
76 | std::size_t index; |
77 | std::size_t second{}; |
78 | std::optional<std::function<void(mlir::func::FuncOp)>> finalizer{}; |
79 | std::optional<std::function<void(mlir::gpu::GPUFuncOp)>> gpuFinalizer{}; |
80 | }; // namespace |
81 | |
82 | /// Target-specific rewriting of the FIR. This is a prerequisite pass to code |
83 | /// generation that traverses the FIR and modifies types and operations to a |
84 | /// form that is appropriate for the specific target. LLVM IR has specific |
85 | /// idioms that are used for distinct target processor and ABI combinations. |
86 | class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> { |
87 | public: |
88 | using TargetRewritePassBase<TargetRewrite>::TargetRewritePassBase; |
89 | |
90 | void runOnOperation() override final { |
91 | auto &context = getContext(); |
92 | mlir::OpBuilder rewriter(&context); |
93 | |
94 | auto mod = getModule(); |
95 | if (!forcedTargetTriple.empty()) |
96 | fir::setTargetTriple(mod, forcedTargetTriple); |
97 | |
98 | if (!forcedTargetCPU.empty()) |
99 | fir::setTargetCPU(mod, forcedTargetCPU); |
100 | |
101 | if (!forcedTuneCPU.empty()) |
102 | fir::setTuneCPU(mod, forcedTuneCPU); |
103 | |
104 | if (!forcedTargetFeatures.empty()) |
105 | fir::setTargetFeatures(mod, forcedTargetFeatures); |
106 | |
107 | // TargetRewrite will require querying the type storage sizes, if it was |
108 | // not set already, create a DataLayoutSpec for the ModuleOp now. |
109 | std::optional<mlir::DataLayout> dl = |
110 | fir::support::getOrSetMLIRDataLayout(mod, /*allowDefaultLayout=*/true); |
111 | if (!dl) { |
112 | mlir::emitError(mod.getLoc(), |
113 | "module operation must carry a data layout attribute " |
114 | "to perform target ABI rewrites on FIR" ); |
115 | signalPassFailure(); |
116 | return; |
117 | } |
118 | |
119 | auto specifics = fir::CodeGenSpecifics::get( |
120 | mod.getContext(), fir::getTargetTriple(mod), fir::getKindMapping(mod), |
121 | fir::getTargetCPU(mod), fir::getTargetFeatures(mod), *dl, |
122 | fir::getTuneCPU(mod)); |
123 | |
124 | setMembers(specifics.get(), &rewriter, &*dl); |
125 | |
126 | // Perform type conversion on signatures and call sites. |
127 | if (mlir::failed(convertTypes(mod))) { |
128 | mlir::emitError(mlir::UnknownLoc::get(&context), |
129 | "error in converting types to target abi" ); |
130 | signalPassFailure(); |
131 | } |
132 | |
133 | // Convert ops in target-specific patterns. |
134 | mod.walk([&](mlir::Operation *op) { |
135 | if (auto call = mlir::dyn_cast<fir::CallOp>(op)) { |
136 | if (!hasPortableSignature(call.getFunctionType(), op)) |
137 | convertCallOp(call, call.getFunctionType()); |
138 | } else if (auto dispatch = mlir::dyn_cast<fir::DispatchOp>(op)) { |
139 | if (!hasPortableSignature(dispatch.getFunctionType(), op)) |
140 | convertCallOp(dispatch, dispatch.getFunctionType()); |
141 | } else if (auto gpuLaunchFunc = |
142 | mlir::dyn_cast<mlir::gpu::LaunchFuncOp>(op)) { |
143 | llvm::SmallVector<mlir::Type> operandsTypes; |
144 | for (auto arg : gpuLaunchFunc.getKernelOperands()) |
145 | operandsTypes.push_back(arg.getType()); |
146 | auto fctTy = mlir::FunctionType::get(&context, operandsTypes, {}); |
147 | if (!hasPortableSignature(fctTy, op)) |
148 | convertCallOp(gpuLaunchFunc, fctTy); |
149 | } else if (auto addr = mlir::dyn_cast<fir::AddrOfOp>(op)) { |
150 | if (mlir::isa<mlir::FunctionType>(addr.getType()) && |
151 | !hasPortableSignature(addr.getType(), op)) |
152 | convertAddrOp(addr); |
153 | } |
154 | }); |
155 | |
156 | clearMembers(); |
157 | } |
158 | |
159 | mlir::ModuleOp getModule() { return getOperation(); } |
160 | |
161 | template <typename Ty, typename Callback> |
162 | std::optional<std::function<mlir::Value(mlir::Operation *)>> |
163 | rewriteCallResultType(mlir::Location loc, mlir::Type originalResTy, |
164 | Ty &newResTys, |
165 | fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs, |
166 | Callback &newOpers, mlir::Value &savedStackPtr, |
167 | fir::CodeGenSpecifics::Marshalling &m) { |
168 | // Currently, targets mandate COMPLEX or STRUCT is a single aggregate or |
169 | // packed scalar, including the sret case. |
170 | assert(m.size() == 1 && "return type not supported on this target" ); |
171 | auto resTy = std::get<mlir::Type>(m[0]); |
172 | auto attr = std::get<fir::CodeGenSpecifics::Attributes>(m[0]); |
173 | if (attr.isSRet()) { |
174 | assert(fir::isa_ref_type(resTy) && "must be a memory reference type" ); |
175 | // Save the stack pointer, if it has not been saved for this call yet. |
176 | // We will need to restore it after the call, because the alloca |
177 | // needs to be deallocated. |
178 | if (!savedStackPtr) |
179 | savedStackPtr = genStackSave(loc); |
180 | mlir::Value stack = |
181 | rewriter->create<fir::AllocaOp>(loc, fir::dyn_cast_ptrEleTy(resTy)); |
182 | newInTyAndAttrs.push_back(m[0]); |
183 | newOpers.push_back(stack); |
184 | return [=](mlir::Operation *) -> mlir::Value { |
185 | auto memTy = fir::ReferenceType::get(originalResTy); |
186 | auto cast = rewriter->create<fir::ConvertOp>(loc, memTy, stack); |
187 | return rewriter->create<fir::LoadOp>(loc, cast); |
188 | }; |
189 | } |
190 | newResTys.push_back(resTy); |
191 | return [=, &savedStackPtr](mlir::Operation *call) -> mlir::Value { |
192 | // We are going to generate an alloca, so save the stack pointer. |
193 | if (!savedStackPtr) |
194 | savedStackPtr = genStackSave(loc); |
195 | return this->convertValueInMemory(loc, call->getResult(0), originalResTy, |
196 | /*inputMayBeBigger=*/true); |
197 | }; |
198 | } |
199 | |
200 | template <typename Ty, typename Callback> |
201 | std::optional<std::function<mlir::Value(mlir::Operation *)>> |
202 | rewriteCallComplexResultType( |
203 | mlir::Location loc, mlir::ComplexType ty, Ty &newResTys, |
204 | fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs, Callback &newOpers, |
205 | mlir::Value &savedStackPtr) { |
206 | if (noComplexConversion) { |
207 | newResTys.push_back(ty); |
208 | return std::nullopt; |
209 | } |
210 | auto m = specifics->complexReturnType(loc, ty.getElementType()); |
211 | return rewriteCallResultType(loc, ty, newResTys, newInTyAndAttrs, newOpers, |
212 | savedStackPtr, m); |
213 | } |
214 | |
215 | template <typename Ty, typename Callback> |
216 | std::optional<std::function<mlir::Value(mlir::Operation *)>> |
217 | rewriteCallStructResultType( |
218 | mlir::Location loc, fir::RecordType recTy, Ty &newResTys, |
219 | fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs, Callback &newOpers, |
220 | mlir::Value &savedStackPtr) { |
221 | if (noStructConversion) { |
222 | newResTys.push_back(recTy); |
223 | return std::nullopt; |
224 | } |
225 | auto m = specifics->structReturnType(loc, recTy); |
226 | return rewriteCallResultType(loc, recTy, newResTys, newInTyAndAttrs, |
227 | newOpers, savedStackPtr, m); |
228 | } |
229 | |
230 | void passArgumentOnStackOrWithNewType( |
231 | mlir::Location loc, fir::CodeGenSpecifics::TypeAndAttr newTypeAndAttr, |
232 | mlir::Type oldType, mlir::Value oper, |
233 | llvm::SmallVectorImpl<mlir::Value> &newOpers, |
234 | mlir::Value &savedStackPtr) { |
235 | auto resTy = std::get<mlir::Type>(newTypeAndAttr); |
236 | auto attr = std::get<fir::CodeGenSpecifics::Attributes>(newTypeAndAttr); |
237 | // We are going to generate an alloca, so save the stack pointer. |
238 | if (!savedStackPtr) |
239 | savedStackPtr = genStackSave(loc); |
240 | if (attr.isByVal()) { |
241 | mlir::Value mem = rewriter->create<fir::AllocaOp>(loc, oldType); |
242 | rewriter->create<fir::StoreOp>(loc, oper, mem); |
243 | if (mem.getType() != resTy) |
244 | mem = rewriter->create<fir::ConvertOp>(loc, resTy, mem); |
245 | newOpers.push_back(mem); |
246 | } else { |
247 | mlir::Value bitcast = |
248 | convertValueInMemory(loc, oper, resTy, /*inputMayBeBigger=*/false); |
249 | newOpers.push_back(bitcast); |
250 | } |
251 | } |
252 | |
253 | // Do a bitcast (convert a value via its memory representation). |
254 | // The input and output types may have different storage sizes, |
255 | // "inputMayBeBigger" should be set to indicate which of the input or |
256 | // output type may be bigger in order for the load/store to be safe. |
257 | // The mismatch comes from the fact that the LLVM register used for passing |
258 | // may be bigger than the value being passed (e.g., passing |
259 | // a `!fir.type<t{fir.array<3xi8>}>` into an i32 LLVM register). |
260 | mlir::Value convertValueInMemory(mlir::Location loc, mlir::Value value, |
261 | mlir::Type newType, bool inputMayBeBigger) { |
262 | if (inputMayBeBigger) { |
263 | auto newRefTy = fir::ReferenceType::get(newType); |
264 | auto mem = rewriter->create<fir::AllocaOp>(loc, value.getType()); |
265 | rewriter->create<fir::StoreOp>(loc, value, mem); |
266 | auto cast = rewriter->create<fir::ConvertOp>(loc, newRefTy, mem); |
267 | return rewriter->create<fir::LoadOp>(loc, cast); |
268 | } else { |
269 | auto oldRefTy = fir::ReferenceType::get(value.getType()); |
270 | auto mem = rewriter->create<fir::AllocaOp>(loc, newType); |
271 | auto cast = rewriter->create<fir::ConvertOp>(loc, oldRefTy, mem); |
272 | rewriter->create<fir::StoreOp>(loc, value, cast); |
273 | return rewriter->create<fir::LoadOp>(loc, mem); |
274 | } |
275 | } |
276 | |
277 | void passSplitArgument(mlir::Location loc, |
278 | fir::CodeGenSpecifics::Marshalling splitArgs, |
279 | mlir::Type oldType, mlir::Value oper, |
280 | llvm::SmallVectorImpl<mlir::Value> &newOpers, |
281 | mlir::Value &savedStackPtr) { |
282 | // COMPLEX or struct argument split into separate arguments |
283 | if (!fir::isa_complex(oldType)) { |
284 | // Cast original operand to a tuple of the new arguments |
285 | // via memory. |
286 | llvm::SmallVector<mlir::Type> partTypes; |
287 | for (auto argPart : splitArgs) |
288 | partTypes.push_back(std::get<mlir::Type>(argPart)); |
289 | mlir::Type tupleType = |
290 | mlir::TupleType::get(oldType.getContext(), partTypes); |
291 | if (!savedStackPtr) |
292 | savedStackPtr = genStackSave(loc); |
293 | oper = convertValueInMemory(loc, oper, tupleType, |
294 | /*inputMayBeBigger=*/false); |
295 | } |
296 | auto iTy = rewriter->getIntegerType(32); |
297 | for (auto e : llvm::enumerate(splitArgs)) { |
298 | auto &tup = e.value(); |
299 | auto ty = std::get<mlir::Type>(tup); |
300 | auto index = e.index(); |
301 | auto idx = rewriter->getIntegerAttr(iTy, index); |
302 | auto val = rewriter->create<fir::ExtractValueOp>( |
303 | loc, ty, oper, rewriter->getArrayAttr(idx)); |
304 | newOpers.push_back(val); |
305 | } |
306 | } |
307 | |
308 | void rewriteCallOperands( |
309 | mlir::Location loc, fir::CodeGenSpecifics::Marshalling passArgAs, |
310 | mlir::Type originalArgTy, mlir::Value oper, |
311 | llvm::SmallVectorImpl<mlir::Value> &newOpers, mlir::Value &savedStackPtr, |
312 | fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs) { |
313 | if (passArgAs.size() == 1) { |
314 | // COMPLEX or derived type is passed as a single argument. |
315 | passArgumentOnStackOrWithNewType(loc, passArgAs[0], originalArgTy, oper, |
316 | newOpers, savedStackPtr); |
317 | } else { |
318 | // COMPLEX or derived type is split into separate arguments |
319 | passSplitArgument(loc, passArgAs, originalArgTy, oper, newOpers, |
320 | savedStackPtr); |
321 | } |
322 | newInTyAndAttrs.insert(newInTyAndAttrs.end(), passArgAs.begin(), |
323 | passArgAs.end()); |
324 | } |
325 | |
326 | template <typename CPLX> |
327 | void rewriteCallComplexInputType( |
328 | mlir::Location loc, CPLX ty, mlir::Value oper, |
329 | fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs, |
330 | llvm::SmallVectorImpl<mlir::Value> &newOpers, |
331 | mlir::Value &savedStackPtr) { |
332 | if (noComplexConversion) { |
333 | newInTyAndAttrs.push_back(fir::CodeGenSpecifics::getTypeAndAttr(ty)); |
334 | newOpers.push_back(oper); |
335 | return; |
336 | } |
337 | auto m = specifics->complexArgumentType(loc, ty.getElementType()); |
338 | rewriteCallOperands(loc, m, ty, oper, newOpers, savedStackPtr, |
339 | newInTyAndAttrs); |
340 | } |
341 | |
342 | void rewriteCallStructInputType( |
343 | mlir::Location loc, fir::RecordType recTy, mlir::Value oper, |
344 | fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs, |
345 | llvm::SmallVectorImpl<mlir::Value> &newOpers, |
346 | mlir::Value &savedStackPtr) { |
347 | if (noStructConversion) { |
348 | newInTyAndAttrs.push_back(fir::CodeGenSpecifics::getTypeAndAttr(recTy)); |
349 | newOpers.push_back(oper); |
350 | return; |
351 | } |
352 | auto structArgs = |
353 | specifics->structArgumentType(loc, recTy, newInTyAndAttrs); |
354 | rewriteCallOperands(loc, structArgs, recTy, oper, newOpers, savedStackPtr, |
355 | newInTyAndAttrs); |
356 | } |
357 | |
358 | static bool hasByValOrSRetArgs( |
359 | const fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs) { |
360 | return llvm::any_of(newInTyAndAttrs, [](auto arg) { |
361 | const auto &attr = std::get<fir::CodeGenSpecifics::Attributes>(arg); |
362 | return attr.isByVal() || attr.isSRet(); |
363 | }); |
364 | } |
365 | |
366 | // Convert fir.call and fir.dispatch Ops. |
367 | template <typename A> |
368 | void convertCallOp(A callOp, mlir::FunctionType fnTy) { |
369 | auto loc = callOp.getLoc(); |
370 | rewriter->setInsertionPoint(callOp); |
371 | llvm::SmallVector<mlir::Type> newResTys; |
372 | fir::CodeGenSpecifics::Marshalling newInTyAndAttrs; |
373 | llvm::SmallVector<mlir::Value> newOpers; |
374 | mlir::Value savedStackPtr = nullptr; |
375 | |
376 | // If the call is indirect, the first argument must still be the function |
377 | // to call. |
378 | int dropFront = 0; |
379 | if constexpr (std::is_same_v<std::decay_t<A>, fir::CallOp>) { |
380 | if (!callOp.getCallee()) { |
381 | newInTyAndAttrs.push_back( |
382 | fir::CodeGenSpecifics::getTypeAndAttr(fnTy.getInput(0))); |
383 | newOpers.push_back(callOp.getOperand(0)); |
384 | dropFront = 1; |
385 | } |
386 | } else if constexpr (std::is_same_v<std::decay_t<A>, fir::DispatchOp>) { |
387 | dropFront = 1; // First operand is the polymorphic object. |
388 | } |
389 | |
390 | // Determine the rewrite function, `wrap`, for the result value. |
391 | std::optional<std::function<mlir::Value(mlir::Operation *)>> wrap; |
392 | if (fnTy.getResults().size() == 1) { |
393 | mlir::Type ty = fnTy.getResult(0); |
394 | llvm::TypeSwitch<mlir::Type>(ty) |
395 | .template Case<mlir::ComplexType>([&](mlir::ComplexType cmplx) { |
396 | wrap = rewriteCallComplexResultType(loc, cmplx, newResTys, |
397 | newInTyAndAttrs, newOpers, |
398 | savedStackPtr); |
399 | }) |
400 | .template Case<fir::RecordType>([&](fir::RecordType recTy) { |
401 | wrap = rewriteCallStructResultType(loc, recTy, newResTys, |
402 | newInTyAndAttrs, newOpers, |
403 | savedStackPtr); |
404 | }) |
405 | .Default([&](mlir::Type ty) { newResTys.push_back(ty); }); |
406 | } else if (fnTy.getResults().size() > 1) { |
407 | TODO(loc, "multiple results not supported yet" ); |
408 | } |
409 | |
410 | llvm::SmallVector<mlir::Type> trailingInTys; |
411 | llvm::SmallVector<mlir::Value> trailingOpers; |
412 | llvm::SmallVector<mlir::Value> operands; |
413 | unsigned passArgShift = 0; |
414 | if constexpr (std::is_same_v<std::decay_t<A>, mlir::gpu::LaunchFuncOp>) |
415 | operands = callOp.getKernelOperands(); |
416 | else |
417 | operands = callOp.getOperands().drop_front(dropFront); |
418 | for (auto e : llvm::enumerate( |
419 | llvm::zip(fnTy.getInputs().drop_front(dropFront), operands))) { |
420 | mlir::Type ty = std::get<0>(e.value()); |
421 | mlir::Value oper = std::get<1>(e.value()); |
422 | unsigned index = e.index(); |
423 | llvm::TypeSwitch<mlir::Type>(ty) |
424 | .template Case<fir::BoxCharType>([&](fir::BoxCharType boxTy) { |
425 | if constexpr (std::is_same_v<std::decay_t<A>, fir::CallOp>) { |
426 | if (noCharacterConversion) { |
427 | newInTyAndAttrs.push_back( |
428 | fir::CodeGenSpecifics::getTypeAndAttr(boxTy)); |
429 | newOpers.push_back(oper); |
430 | return; |
431 | } |
432 | } else { |
433 | // TODO: dispatch case; it used to be a to-do because of sret, |
434 | // but is not tested and maybe should be removed. This pass is |
435 | // anyway ran after lowering fir.dispatch in flang, so maybe that |
436 | // should just be a requirement of the pass. |
437 | TODO(loc, "ABI of fir.dispatch with character arguments" ); |
438 | } |
439 | auto m = specifics->boxcharArgumentType(boxTy.getEleTy()); |
440 | auto unbox = rewriter->create<fir::UnboxCharOp>( |
441 | loc, std::get<mlir::Type>(m[0]), std::get<mlir::Type>(m[1]), |
442 | oper); |
443 | // unboxed CHARACTER arguments |
444 | for (auto e : llvm::enumerate(m)) { |
445 | unsigned idx = e.index(); |
446 | auto attr = |
447 | std::get<fir::CodeGenSpecifics::Attributes>(e.value()); |
448 | auto argTy = std::get<mlir::Type>(e.value()); |
449 | if (attr.isAppend()) { |
450 | trailingInTys.push_back(argTy); |
451 | trailingOpers.push_back(unbox.getResult(idx)); |
452 | } else { |
453 | newInTyAndAttrs.push_back(e.value()); |
454 | newOpers.push_back(unbox.getResult(idx)); |
455 | } |
456 | } |
457 | }) |
458 | .template Case<mlir::ComplexType>([&](mlir::ComplexType cmplx) { |
459 | rewriteCallComplexInputType(loc, cmplx, oper, newInTyAndAttrs, |
460 | newOpers, savedStackPtr); |
461 | }) |
462 | .template Case<fir::RecordType>([&](fir::RecordType recTy) { |
463 | rewriteCallStructInputType(loc, recTy, oper, newInTyAndAttrs, |
464 | newOpers, savedStackPtr); |
465 | }) |
466 | .template Case<mlir::TupleType>([&](mlir::TupleType tuple) { |
467 | if (fir::isCharacterProcedureTuple(tuple)) { |
468 | mlir::ModuleOp module = getModule(); |
469 | if constexpr (std::is_same_v<std::decay_t<A>, fir::CallOp>) { |
470 | if (callOp.getCallee()) { |
471 | llvm::StringRef charProcAttr = |
472 | fir::getCharacterProcedureDummyAttrName(); |
473 | // The charProcAttr attribute is only used as a safety to |
474 | // confirm that this is a dummy procedure and should be split. |
475 | // It cannot be used to match because attributes are not |
476 | // available in case of indirect calls. |
477 | auto funcOp = module.lookupSymbol<mlir::func::FuncOp>( |
478 | *callOp.getCallee()); |
479 | if (funcOp && |
480 | !funcOp.template getArgAttrOfType<mlir::UnitAttr>( |
481 | index, charProcAttr)) |
482 | mlir::emitError(loc, "tuple argument will be split even " |
483 | "though it does not have the `" + |
484 | charProcAttr + "` attribute" ); |
485 | } |
486 | } |
487 | mlir::Type funcPointerType = tuple.getType(0); |
488 | mlir::Type lenType = tuple.getType(1); |
489 | fir::FirOpBuilder builder(*rewriter, module); |
490 | auto [funcPointer, len] = |
491 | fir::factory::extractCharacterProcedureTuple(builder, loc, |
492 | oper); |
493 | newInTyAndAttrs.push_back( |
494 | fir::CodeGenSpecifics::getTypeAndAttr(funcPointerType)); |
495 | newOpers.push_back(funcPointer); |
496 | trailingInTys.push_back(lenType); |
497 | trailingOpers.push_back(len); |
498 | } else { |
499 | newInTyAndAttrs.push_back( |
500 | fir::CodeGenSpecifics::getTypeAndAttr(tuple)); |
501 | newOpers.push_back(oper); |
502 | } |
503 | }) |
504 | .Default([&](mlir::Type ty) { |
505 | if constexpr (std::is_same_v<std::decay_t<A>, fir::DispatchOp>) { |
506 | if (callOp.getPassArgPos() && *callOp.getPassArgPos() == index) |
507 | passArgShift = newOpers.size() - *callOp.getPassArgPos(); |
508 | } |
509 | newInTyAndAttrs.push_back( |
510 | fir::CodeGenSpecifics::getTypeAndAttr(ty)); |
511 | newOpers.push_back(oper); |
512 | }); |
513 | } |
514 | |
515 | llvm::SmallVector<mlir::Type> newInTypes = toTypeList(newInTyAndAttrs); |
516 | newInTypes.insert(newInTypes.end(), trailingInTys.begin(), |
517 | trailingInTys.end()); |
518 | newOpers.insert(newOpers.end(), trailingOpers.begin(), trailingOpers.end()); |
519 | |
520 | llvm::SmallVector<mlir::Value, 1> newCallResults; |
521 | // TODO propagate/update call argument and result attributes. |
522 | if constexpr (std::is_same_v<std::decay_t<A>, mlir::gpu::LaunchFuncOp>) { |
523 | auto newCall = rewriter->create<A>( |
524 | loc, callOp.getKernel(), callOp.getGridSizeOperandValues(), |
525 | callOp.getBlockSizeOperandValues(), |
526 | callOp.getDynamicSharedMemorySize(), newOpers); |
527 | if (callOp.getClusterSizeX()) |
528 | newCall.getClusterSizeXMutable().assign(callOp.getClusterSizeX()); |
529 | if (callOp.getClusterSizeY()) |
530 | newCall.getClusterSizeYMutable().assign(callOp.getClusterSizeY()); |
531 | if (callOp.getClusterSizeZ()) |
532 | newCall.getClusterSizeZMutable().assign(callOp.getClusterSizeZ()); |
533 | newCallResults.append(newCall.result_begin(), newCall.result_end()); |
534 | if (auto cudaProcAttr = |
535 | callOp->template getAttrOfType<cuf::ProcAttributeAttr>( |
536 | cuf::getProcAttrName())) { |
537 | newCall->setAttr(cuf::getProcAttrName(), cudaProcAttr); |
538 | } |
539 | } else if constexpr (std::is_same_v<std::decay_t<A>, fir::CallOp>) { |
540 | fir::CallOp newCall; |
541 | if (callOp.getCallee()) { |
542 | newCall = rewriter->create<fir::CallOp>(loc, *callOp.getCallee(), |
543 | newResTys, newOpers); |
544 | } else { |
545 | newOpers[0].setType(mlir::FunctionType::get( |
546 | callOp.getContext(), |
547 | mlir::TypeRange{newInTypes}.drop_front(dropFront), newResTys)); |
548 | newCall = rewriter->create<fir::CallOp>(loc, newResTys, newOpers); |
549 | } |
550 | newCall.setFastmathAttr(callOp.getFastmathAttr()); |
551 | // Always set ABI argument attributes on call operations, even when |
552 | // direct, as required by |
553 | // https://llvm.org/docs/LangRef.html#parameter-attributes. |
554 | if (hasByValOrSRetArgs(newInTyAndAttrs)) { |
555 | llvm::SmallVector<mlir::Attribute> argAttrsArray; |
556 | for (const auto &arg : |
557 | llvm::ArrayRef<fir::CodeGenSpecifics::TypeAndAttr>(newInTyAndAttrs) |
558 | .drop_front(dropFront)) { |
559 | mlir::NamedAttrList argAttrs; |
560 | const auto &attr = std::get<fir::CodeGenSpecifics::Attributes>(arg); |
561 | if (attr.isByVal()) { |
562 | mlir::Type elemType = |
563 | fir::dyn_cast_ptrOrBoxEleTy(std::get<mlir::Type>(arg)); |
564 | argAttrs.set(mlir::LLVM::LLVMDialect::getByValAttrName(), |
565 | mlir::TypeAttr::get(elemType)); |
566 | } else if (attr.isSRet()) { |
567 | mlir::Type elemType = |
568 | fir::dyn_cast_ptrOrBoxEleTy(std::get<mlir::Type>(arg)); |
569 | argAttrs.set(mlir::LLVM::LLVMDialect::getStructRetAttrName(), |
570 | mlir::TypeAttr::get(elemType)); |
571 | } |
572 | if (auto align = attr.getAlignment()) { |
573 | argAttrs.set( |
574 | mlir::LLVM::LLVMDialect::getAlignAttrName(), |
575 | rewriter->getIntegerAttr(rewriter->getIntegerType(32), align)); |
576 | } |
577 | argAttrsArray.emplace_back( |
578 | argAttrs.getDictionary(rewriter->getContext())); |
579 | } |
580 | newCall.setArgAttrsAttr(rewriter->getArrayAttr(argAttrsArray)); |
581 | } |
582 | LLVM_DEBUG(llvm::dbgs() << "replacing call with " << newCall << '\n'); |
583 | if (wrap) |
584 | newCallResults.push_back((*wrap)(newCall.getOperation())); |
585 | else |
586 | newCallResults.append(newCall.result_begin(), newCall.result_end()); |
587 | } else { |
588 | fir::DispatchOp dispatchOp = rewriter->create<A>( |
589 | loc, newResTys, rewriter->getStringAttr(callOp.getMethod()), |
590 | callOp.getOperands()[0], newOpers, |
591 | rewriter->getI32IntegerAttr(*callOp.getPassArgPos() + passArgShift), |
592 | /*arg_attrs=*/nullptr, /*res_attrs=*/nullptr, |
593 | callOp.getProcedureAttrsAttr()); |
594 | if (wrap) |
595 | newCallResults.push_back((*wrap)(dispatchOp.getOperation())); |
596 | else |
597 | newCallResults.append(dispatchOp.result_begin(), |
598 | dispatchOp.result_end()); |
599 | } |
600 | |
601 | if (newCallResults.size() <= 1) { |
602 | if (savedStackPtr) { |
603 | if (newCallResults.size() == 1) { |
604 | // We assume that all the allocas are inserted before |
605 | // the operation that defines the new call result. |
606 | rewriter->setInsertionPointAfterValue(newCallResults[0]); |
607 | } else { |
608 | // If the call does not have results, then insert |
609 | // stack restore after the original call operation. |
610 | rewriter->setInsertionPointAfter(callOp); |
611 | } |
612 | genStackRestore(loc, savedStackPtr); |
613 | } |
614 | replaceOp(callOp, newCallResults); |
615 | } else { |
616 | // The TODO is duplicated here to make sure this part |
617 | // handles the stackrestore insertion properly, if |
618 | // we add support for multiple call results. |
619 | TODO(loc, "multiple results not supported yet" ); |
620 | } |
621 | } |
622 | |
623 | // Result type fixup for ComplexType. |
624 | template <typename Ty> |
625 | void lowerComplexSignatureRes( |
626 | mlir::Location loc, mlir::ComplexType cmplx, Ty &newResTys, |
627 | fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs) { |
628 | if (noComplexConversion) { |
629 | newResTys.push_back(cmplx); |
630 | return; |
631 | } |
632 | for (auto &tup : |
633 | specifics->complexReturnType(loc, cmplx.getElementType())) { |
634 | auto argTy = std::get<mlir::Type>(tup); |
635 | if (std::get<fir::CodeGenSpecifics::Attributes>(tup).isSRet()) |
636 | newInTyAndAttrs.push_back(tup); |
637 | else |
638 | newResTys.push_back(argTy); |
639 | } |
640 | } |
641 | |
642 | // Argument type fixup for ComplexType. |
643 | void lowerComplexSignatureArg( |
644 | mlir::Location loc, mlir::ComplexType cmplx, |
645 | fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs) { |
646 | if (noComplexConversion) { |
647 | newInTyAndAttrs.push_back(fir::CodeGenSpecifics::getTypeAndAttr(cmplx)); |
648 | } else { |
649 | auto cplxArgs = |
650 | specifics->complexArgumentType(loc, cmplx.getElementType()); |
651 | newInTyAndAttrs.insert(newInTyAndAttrs.end(), cplxArgs.begin(), |
652 | cplxArgs.end()); |
653 | } |
654 | } |
655 | |
656 | template <typename Ty> |
657 | void |
658 | lowerStructSignatureRes(mlir::Location loc, fir::RecordType recTy, |
659 | Ty &newResTys, |
660 | fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs) { |
661 | if (noComplexConversion) { |
662 | newResTys.push_back(recTy); |
663 | return; |
664 | } else { |
665 | for (auto &tup : specifics->structReturnType(loc, recTy)) { |
666 | if (std::get<fir::CodeGenSpecifics::Attributes>(tup).isSRet()) |
667 | newInTyAndAttrs.push_back(tup); |
668 | else |
669 | newResTys.push_back(std::get<mlir::Type>(tup)); |
670 | } |
671 | } |
672 | } |
673 | |
674 | void |
675 | lowerStructSignatureArg(mlir::Location loc, fir::RecordType recTy, |
676 | fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs) { |
677 | if (noStructConversion) { |
678 | newInTyAndAttrs.push_back(fir::CodeGenSpecifics::getTypeAndAttr(recTy)); |
679 | return; |
680 | } |
681 | auto structArgs = |
682 | specifics->structArgumentType(loc, recTy, newInTyAndAttrs); |
683 | newInTyAndAttrs.insert(newInTyAndAttrs.end(), structArgs.begin(), |
684 | structArgs.end()); |
685 | } |
686 | |
687 | llvm::SmallVector<mlir::Type> |
688 | toTypeList(const fir::CodeGenSpecifics::Marshalling &marshalled) { |
689 | llvm::SmallVector<mlir::Type> typeList; |
690 | for (auto &typeAndAttr : marshalled) |
691 | typeList.emplace_back(std::get<mlir::Type>(typeAndAttr)); |
692 | return typeList; |
693 | } |
694 | |
695 | /// Taking the address of a function. Modify the signature as needed. |
696 | void convertAddrOp(fir::AddrOfOp addrOp) { |
697 | rewriter->setInsertionPoint(addrOp); |
698 | auto addrTy = mlir::cast<mlir::FunctionType>(addrOp.getType()); |
699 | fir::CodeGenSpecifics::Marshalling newInTyAndAttrs; |
700 | llvm::SmallVector<mlir::Type> newResTys; |
701 | auto loc = addrOp.getLoc(); |
702 | for (mlir::Type ty : addrTy.getResults()) { |
703 | llvm::TypeSwitch<mlir::Type>(ty) |
704 | .Case<mlir::ComplexType>([&](mlir::ComplexType ty) { |
705 | lowerComplexSignatureRes(loc, ty, newResTys, newInTyAndAttrs); |
706 | }) |
707 | .Case<fir::RecordType>([&](fir::RecordType ty) { |
708 | lowerStructSignatureRes(loc, ty, newResTys, newInTyAndAttrs); |
709 | }) |
710 | .Default([&](mlir::Type ty) { newResTys.push_back(ty); }); |
711 | } |
712 | llvm::SmallVector<mlir::Type> trailingInTys; |
713 | for (mlir::Type ty : addrTy.getInputs()) { |
714 | llvm::TypeSwitch<mlir::Type>(ty) |
715 | .Case<fir::BoxCharType>([&](auto box) { |
716 | if (noCharacterConversion) { |
717 | newInTyAndAttrs.push_back( |
718 | fir::CodeGenSpecifics::getTypeAndAttr(box)); |
719 | } else { |
720 | for (auto &tup : specifics->boxcharArgumentType(box.getEleTy())) { |
721 | auto attr = std::get<fir::CodeGenSpecifics::Attributes>(tup); |
722 | auto argTy = std::get<mlir::Type>(tup); |
723 | if (attr.isAppend()) |
724 | trailingInTys.push_back(argTy); |
725 | else |
726 | newInTyAndAttrs.push_back(tup); |
727 | } |
728 | } |
729 | }) |
730 | .Case<mlir::ComplexType>([&](mlir::ComplexType ty) { |
731 | lowerComplexSignatureArg(loc, ty, newInTyAndAttrs); |
732 | }) |
733 | .Case<mlir::TupleType>([&](mlir::TupleType tuple) { |
734 | if (fir::isCharacterProcedureTuple(tuple)) { |
735 | newInTyAndAttrs.push_back( |
736 | fir::CodeGenSpecifics::getTypeAndAttr(tuple.getType(0))); |
737 | trailingInTys.push_back(tuple.getType(1)); |
738 | } else { |
739 | newInTyAndAttrs.push_back( |
740 | fir::CodeGenSpecifics::getTypeAndAttr(ty)); |
741 | } |
742 | }) |
743 | .template Case<fir::RecordType>([&](fir::RecordType recTy) { |
744 | lowerStructSignatureArg(loc, recTy, newInTyAndAttrs); |
745 | }) |
746 | .Default([&](mlir::Type ty) { |
747 | newInTyAndAttrs.push_back( |
748 | fir::CodeGenSpecifics::getTypeAndAttr(ty)); |
749 | }); |
750 | } |
751 | llvm::SmallVector<mlir::Type> newInTypes = toTypeList(newInTyAndAttrs); |
752 | // append trailing input types |
753 | newInTypes.insert(newInTypes.end(), trailingInTys.begin(), |
754 | trailingInTys.end()); |
755 | // replace this op with a new one with the updated signature |
756 | auto newTy = rewriter->getFunctionType(newInTypes, newResTys); |
757 | auto newOp = rewriter->create<fir::AddrOfOp>(addrOp.getLoc(), newTy, |
758 | addrOp.getSymbol()); |
759 | replaceOp(addrOp, newOp.getResult()); |
760 | } |
761 | |
762 | /// Convert the type signatures on all the functions present in the module. |
763 | /// As the type signature is being changed, this must also update the |
764 | /// function itself to use any new arguments, etc. |
765 | llvm::LogicalResult convertTypes(mlir::ModuleOp mod) { |
766 | mlir::MLIRContext *ctx = mod->getContext(); |
767 | auto targetCPU = specifics->getTargetCPU(); |
768 | mlir::StringAttr targetCPUAttr = |
769 | targetCPU.empty() ? nullptr : mlir::StringAttr::get(ctx, targetCPU); |
770 | auto tuneCPU = specifics->getTuneCPU(); |
771 | mlir::StringAttr tuneCPUAttr = |
772 | tuneCPU.empty() ? nullptr : mlir::StringAttr::get(ctx, tuneCPU); |
773 | auto targetFeaturesAttr = specifics->getTargetFeatures(); |
774 | |
775 | for (auto fn : mod.getOps<mlir::func::FuncOp>()) { |
776 | if (targetCPUAttr) |
777 | fn->setAttr("target_cpu" , targetCPUAttr); |
778 | |
779 | if (tuneCPUAttr) |
780 | fn->setAttr("tune_cpu" , tuneCPUAttr); |
781 | |
782 | if (targetFeaturesAttr) |
783 | fn->setAttr("target_features" , targetFeaturesAttr); |
784 | |
785 | convertSignature<mlir::func::ReturnOp, mlir::func::FuncOp>(fn); |
786 | } |
787 | |
788 | for (auto gpuMod : mod.getOps<mlir::gpu::GPUModuleOp>()) { |
789 | for (auto fn : gpuMod.getOps<mlir::func::FuncOp>()) |
790 | convertSignature<mlir::func::ReturnOp, mlir::func::FuncOp>(fn); |
791 | for (auto fn : gpuMod.getOps<mlir::gpu::GPUFuncOp>()) |
792 | convertSignature<mlir::gpu::ReturnOp, mlir::gpu::GPUFuncOp>(fn); |
793 | } |
794 | |
795 | return mlir::success(); |
796 | } |
797 | |
798 | // Returns true if the function should be interoperable with C. |
799 | static bool isFuncWithCCallingConvention(mlir::Operation *op) { |
800 | auto funcOp = mlir::dyn_cast<mlir::func::FuncOp>(op); |
801 | if (!funcOp) |
802 | return false; |
803 | return op->hasAttrOfType<mlir::UnitAttr>( |
804 | fir::FIROpsDialect::getFirRuntimeAttrName()) || |
805 | op->hasAttrOfType<mlir::StringAttr>(fir::getSymbolAttrName()); |
806 | } |
807 | |
808 | /// If the signature does not need any special target-specific conversions, |
809 | /// then it is considered portable for any target, and this function will |
810 | /// return `true`. Otherwise, the signature is not portable and `false` is |
811 | /// returned. |
812 | bool hasPortableSignature(mlir::Type signature, mlir::Operation *op) { |
813 | assert(mlir::isa<mlir::FunctionType>(signature)); |
814 | auto func = mlir::dyn_cast<mlir::FunctionType>(signature); |
815 | bool hasCCallingConv = isFuncWithCCallingConvention(op); |
816 | for (auto ty : func.getResults()) |
817 | if ((mlir::isa<fir::BoxCharType>(ty) && !noCharacterConversion) || |
818 | (fir::isa_complex(ty) && !noComplexConversion) || |
819 | (mlir::isa<mlir::IntegerType>(ty) && hasCCallingConv) || |
820 | (mlir::isa<fir::RecordType>(ty) && !noStructConversion)) { |
821 | LLVM_DEBUG(llvm::dbgs() << "rewrite " << signature << " for target\n" ); |
822 | return false; |
823 | } |
824 | for (auto ty : func.getInputs()) |
825 | if (((mlir::isa<fir::BoxCharType>(ty) || |
826 | fir::isCharacterProcedureTuple(ty)) && |
827 | !noCharacterConversion) || |
828 | (fir::isa_complex(ty) && !noComplexConversion) || |
829 | (mlir::isa<mlir::IntegerType>(ty) && hasCCallingConv) || |
830 | (mlir::isa<fir::RecordType>(ty) && !noStructConversion)) { |
831 | LLVM_DEBUG(llvm::dbgs() << "rewrite " << signature << " for target\n" ); |
832 | return false; |
833 | } |
834 | return true; |
835 | } |
836 | |
837 | /// Determine if the signature has host associations. The host association |
838 | /// argument may need special target specific rewriting. |
839 | template <typename OpTy> |
840 | static bool hasHostAssociations(OpTy func) { |
841 | std::size_t end = func.getFunctionType().getInputs().size(); |
842 | for (std::size_t i = 0; i < end; ++i) |
843 | if (func.template getArgAttrOfType<mlir::UnitAttr>( |
844 | i, fir::getHostAssocAttrName())) |
845 | return true; |
846 | return false; |
847 | } |
848 | |
849 | /// Rewrite the signatures and body of the `FuncOp`s in the module for |
850 | /// the immediately subsequent target code gen. |
851 | template <typename ReturnOpTy, typename FuncOpTy> |
852 | void convertSignature(FuncOpTy func) { |
853 | auto funcTy = mlir::cast<mlir::FunctionType>(func.getFunctionType()); |
854 | if (hasPortableSignature(funcTy, func) && !hasHostAssociations(func)) |
855 | return; |
856 | llvm::SmallVector<mlir::Type> newResTys; |
857 | fir::CodeGenSpecifics::Marshalling newInTyAndAttrs; |
858 | llvm::SmallVector<std::pair<unsigned, mlir::NamedAttribute>> savedAttrs; |
859 | llvm::SmallVector<std::pair<unsigned, mlir::NamedAttribute>> extraAttrs; |
860 | llvm::SmallVector<FixupTy> fixups; |
861 | llvm::SmallVector<std::pair<unsigned, mlir::NamedAttrList>, 1> resultAttrs; |
862 | |
863 | // Save argument attributes in case there is a shift so we can replace them |
864 | // correctly. |
865 | for (auto e : llvm::enumerate(funcTy.getInputs())) { |
866 | unsigned index = e.index(); |
867 | llvm::ArrayRef<mlir::NamedAttribute> attrs = |
868 | mlir::function_interface_impl::getArgAttrs(func, index); |
869 | for (mlir::NamedAttribute attr : attrs) { |
870 | savedAttrs.push_back({index, attr}); |
871 | } |
872 | } |
873 | |
874 | // Convert return value(s) |
875 | for (auto ty : funcTy.getResults()) |
876 | llvm::TypeSwitch<mlir::Type>(ty) |
877 | .template Case<mlir::ComplexType>([&](mlir::ComplexType cmplx) { |
878 | if (noComplexConversion) |
879 | newResTys.push_back(cmplx); |
880 | else |
881 | doComplexReturn(func, cmplx, newResTys, newInTyAndAttrs, fixups); |
882 | }) |
883 | .template Case<mlir::IntegerType>([&](mlir::IntegerType intTy) { |
884 | auto m = specifics->integerArgumentType(func.getLoc(), intTy); |
885 | assert(m.size() == 1); |
886 | auto attr = std::get<fir::CodeGenSpecifics::Attributes>(m[0]); |
887 | auto retTy = std::get<mlir::Type>(m[0]); |
888 | std::size_t resId = newResTys.size(); |
889 | llvm::StringRef extensionAttrName = attr.getIntExtensionAttrName(); |
890 | if (!extensionAttrName.empty() && |
891 | isFuncWithCCallingConvention(func)) |
892 | resultAttrs.emplace_back( |
893 | resId, rewriter->getNamedAttr(extensionAttrName, |
894 | rewriter->getUnitAttr())); |
895 | newResTys.push_back(retTy); |
896 | }) |
897 | .template Case<fir::RecordType>([&](fir::RecordType recTy) { |
898 | doStructReturn(func, recTy, newResTys, newInTyAndAttrs, fixups); |
899 | }) |
900 | .Default([&](mlir::Type ty) { newResTys.push_back(ty); }); |
901 | |
902 | // Saved potential shift in argument. Handling of result can add arguments |
903 | // at the beginning of the function signature. |
904 | unsigned argumentShift = newInTyAndAttrs.size(); |
905 | |
906 | // Convert arguments |
907 | llvm::SmallVector<mlir::Type> trailingTys; |
908 | for (auto e : llvm::enumerate(funcTy.getInputs())) { |
909 | auto ty = e.value(); |
910 | unsigned index = e.index(); |
911 | llvm::TypeSwitch<mlir::Type>(ty) |
912 | .template Case<fir::BoxCharType>([&](fir::BoxCharType boxTy) { |
913 | if (noCharacterConversion) { |
914 | newInTyAndAttrs.push_back( |
915 | fir::CodeGenSpecifics::getTypeAndAttr(boxTy)); |
916 | } else { |
917 | // Convert a CHARACTER argument type. This can involve separating |
918 | // the pointer and the LEN into two arguments and moving the LEN |
919 | // argument to the end of the arg list. |
920 | for (auto &tup : |
921 | specifics->boxcharArgumentType(boxTy.getEleTy())) { |
922 | auto attr = std::get<fir::CodeGenSpecifics::Attributes>(tup); |
923 | auto argTy = std::get<mlir::Type>(tup); |
924 | if (attr.isAppend()) { |
925 | trailingTys.push_back(argTy); |
926 | } else { |
927 | fixups.emplace_back(FixupTy::Codes::Trailing, |
928 | newInTyAndAttrs.size(), |
929 | trailingTys.size()); |
930 | newInTyAndAttrs.push_back(tup); |
931 | } |
932 | } |
933 | } |
934 | }) |
935 | .template Case<mlir::ComplexType>([&](mlir::ComplexType cmplx) { |
936 | doComplexArg(func, cmplx, newInTyAndAttrs, fixups); |
937 | }) |
938 | .template Case<mlir::TupleType>([&](mlir::TupleType tuple) { |
939 | if (fir::isCharacterProcedureTuple(tuple)) { |
940 | fixups.emplace_back(FixupTy::Codes::TrailingCharProc, |
941 | newInTyAndAttrs.size(), trailingTys.size()); |
942 | newInTyAndAttrs.push_back( |
943 | fir::CodeGenSpecifics::getTypeAndAttr(tuple.getType(0))); |
944 | trailingTys.push_back(tuple.getType(1)); |
945 | } else { |
946 | newInTyAndAttrs.push_back( |
947 | fir::CodeGenSpecifics::getTypeAndAttr(ty)); |
948 | } |
949 | }) |
950 | .template Case<mlir::IntegerType>([&](mlir::IntegerType intTy) { |
951 | auto m = specifics->integerArgumentType(func.getLoc(), intTy); |
952 | assert(m.size() == 1); |
953 | auto attr = std::get<fir::CodeGenSpecifics::Attributes>(m[0]); |
954 | auto argNo = newInTyAndAttrs.size(); |
955 | llvm::StringRef extensionAttrName = attr.getIntExtensionAttrName(); |
956 | if (!extensionAttrName.empty() && |
957 | isFuncWithCCallingConvention(func)) |
958 | fixups.emplace_back(FixupTy::Codes::ArgumentType, argNo, |
959 | [=](FuncOpTy func) { |
960 | func.setArgAttr( |
961 | argNo, extensionAttrName, |
962 | mlir::UnitAttr::get(func.getContext())); |
963 | }); |
964 | |
965 | newInTyAndAttrs.push_back(m[0]); |
966 | }) |
967 | .template Case<fir::RecordType>([&](fir::RecordType recTy) { |
968 | doStructArg(func, recTy, newInTyAndAttrs, fixups); |
969 | }) |
970 | .Default([&](mlir::Type ty) { |
971 | newInTyAndAttrs.push_back( |
972 | fir::CodeGenSpecifics::getTypeAndAttr(ty)); |
973 | }); |
974 | |
975 | if (func.template getArgAttrOfType<mlir::UnitAttr>( |
976 | index, fir::getHostAssocAttrName())) { |
977 | extraAttrs.push_back( |
978 | {newInTyAndAttrs.size() - 1, |
979 | rewriter->getNamedAttr("llvm.nest" , rewriter->getUnitAttr())}); |
980 | } |
981 | } |
982 | |
983 | if (!func.empty()) { |
984 | // If the function has a body, then apply the fixups to the arguments and |
985 | // return ops as required. These fixups are done in place. |
986 | auto loc = func.getLoc(); |
987 | const auto fixupSize = fixups.size(); |
988 | const auto oldArgTys = func.getFunctionType().getInputs(); |
989 | int offset = 0; |
990 | for (std::remove_const_t<decltype(fixupSize)> i = 0; i < fixupSize; ++i) { |
991 | const auto &fixup = fixups[i]; |
992 | mlir::Type fixupType = |
993 | fixup.index < newInTyAndAttrs.size() |
994 | ? std::get<mlir::Type>(newInTyAndAttrs[fixup.index]) |
995 | : mlir::Type{}; |
996 | switch (fixup.code) { |
997 | case FixupTy::Codes::ArgumentAsLoad: { |
998 | // Argument was pass-by-value, but is now pass-by-reference and |
999 | // possibly with a different element type. |
1000 | auto newArg = |
1001 | func.front().insertArgument(fixup.index, fixupType, loc); |
1002 | rewriter->setInsertionPointToStart(&func.front()); |
1003 | auto oldArgTy = |
1004 | fir::ReferenceType::get(oldArgTys[fixup.index - offset]); |
1005 | auto cast = rewriter->create<fir::ConvertOp>(loc, oldArgTy, newArg); |
1006 | auto load = rewriter->create<fir::LoadOp>(loc, cast); |
1007 | func.getArgument(fixup.index + 1).replaceAllUsesWith(load); |
1008 | func.front().eraseArgument(fixup.index + 1); |
1009 | } break; |
1010 | case FixupTy::Codes::ArgumentType: { |
1011 | // Argument is pass-by-value, but its type has likely been modified to |
1012 | // suit the target ABI convention. |
1013 | auto oldArgTy = oldArgTys[fixup.index - offset]; |
1014 | // If type did not change, keep the original argument. |
1015 | if (fixupType == oldArgTy) |
1016 | break; |
1017 | |
1018 | auto newArg = |
1019 | func.front().insertArgument(fixup.index, fixupType, loc); |
1020 | rewriter->setInsertionPointToStart(&func.front()); |
1021 | mlir::Value bitcast = convertValueInMemory(loc, newArg, oldArgTy, |
1022 | /*inputMayBeBigger=*/true); |
1023 | func.getArgument(fixup.index + 1).replaceAllUsesWith(bitcast); |
1024 | func.front().eraseArgument(fixup.index + 1); |
1025 | LLVM_DEBUG(llvm::dbgs() |
1026 | << "old argument: " << oldArgTy << ", repl: " << bitcast |
1027 | << ", new argument: " |
1028 | << func.getArgument(fixup.index).getType() << '\n'); |
1029 | } break; |
1030 | case FixupTy::Codes::CharPair: { |
1031 | // The FIR boxchar argument has been split into a pair of distinct |
1032 | // arguments that are in juxtaposition to each other. |
1033 | auto newArg = |
1034 | func.front().insertArgument(fixup.index, fixupType, loc); |
1035 | if (fixup.second == 1) { |
1036 | rewriter->setInsertionPointToStart(&func.front()); |
1037 | auto boxTy = oldArgTys[fixup.index - offset - fixup.second]; |
1038 | auto box = rewriter->create<fir::EmboxCharOp>( |
1039 | loc, boxTy, func.front().getArgument(fixup.index - 1), newArg); |
1040 | func.getArgument(fixup.index + 1).replaceAllUsesWith(box); |
1041 | func.front().eraseArgument(fixup.index + 1); |
1042 | offset++; |
1043 | } |
1044 | } break; |
1045 | case FixupTy::Codes::ReturnAsStore: { |
1046 | // The value being returned is now being returned in memory (callee |
1047 | // stack space) through a hidden reference argument. |
1048 | auto newArg = |
1049 | func.front().insertArgument(fixup.index, fixupType, loc); |
1050 | offset++; |
1051 | func.walk([&](ReturnOpTy ret) { |
1052 | rewriter->setInsertionPoint(ret); |
1053 | auto oldOper = ret.getOperand(0); |
1054 | auto oldOperTy = fir::ReferenceType::get(oldOper.getType()); |
1055 | auto cast = |
1056 | rewriter->create<fir::ConvertOp>(loc, oldOperTy, newArg); |
1057 | rewriter->create<fir::StoreOp>(loc, oldOper, cast); |
1058 | rewriter->create<ReturnOpTy>(loc); |
1059 | ret.erase(); |
1060 | }); |
1061 | } break; |
1062 | case FixupTy::Codes::ReturnType: { |
1063 | // The function is still returning a value, but its type has likely |
1064 | // changed to suit the target ABI convention. |
1065 | func.walk([&](ReturnOpTy ret) { |
1066 | rewriter->setInsertionPoint(ret); |
1067 | auto oldOper = ret.getOperand(0); |
1068 | mlir::Value bitcast = |
1069 | convertValueInMemory(loc, oldOper, newResTys[fixup.index], |
1070 | /*inputMayBeBigger=*/false); |
1071 | rewriter->create<ReturnOpTy>(loc, bitcast); |
1072 | ret.erase(); |
1073 | }); |
1074 | } break; |
1075 | case FixupTy::Codes::Split: { |
1076 | // The FIR argument has been split into a pair of distinct arguments |
1077 | // that are in juxtaposition to each other. (For COMPLEX value or |
1078 | // derived type passed with VALUE in BIND(C) context). |
1079 | auto newArg = |
1080 | func.front().insertArgument(fixup.index, fixupType, loc); |
1081 | if (fixup.second == 1) { |
1082 | rewriter->setInsertionPointToStart(&func.front()); |
1083 | mlir::Value firstArg = func.front().getArgument(fixup.index - 1); |
1084 | mlir::Type originalTy = |
1085 | oldArgTys[fixup.index - offset - fixup.second]; |
1086 | mlir::Type pairTy = originalTy; |
1087 | if (!fir::isa_complex(originalTy)) { |
1088 | pairTy = mlir::TupleType::get( |
1089 | originalTy.getContext(), |
1090 | mlir::TypeRange{firstArg.getType(), newArg.getType()}); |
1091 | } |
1092 | auto undef = rewriter->create<fir::UndefOp>(loc, pairTy); |
1093 | auto iTy = rewriter->getIntegerType(32); |
1094 | auto zero = rewriter->getIntegerAttr(iTy, 0); |
1095 | auto one = rewriter->getIntegerAttr(iTy, 1); |
1096 | mlir::Value pair1 = rewriter->create<fir::InsertValueOp>( |
1097 | loc, pairTy, undef, firstArg, rewriter->getArrayAttr(zero)); |
1098 | mlir::Value pair = rewriter->create<fir::InsertValueOp>( |
1099 | loc, pairTy, pair1, newArg, rewriter->getArrayAttr(one)); |
1100 | // Cast local argument tuple to original type via memory if needed. |
1101 | if (pairTy != originalTy) |
1102 | pair = convertValueInMemory(loc, pair, originalTy, |
1103 | /*inputMayBeBigger=*/true); |
1104 | func.getArgument(fixup.index + 1).replaceAllUsesWith(pair); |
1105 | func.front().eraseArgument(fixup.index + 1); |
1106 | offset++; |
1107 | } |
1108 | } break; |
1109 | case FixupTy::Codes::Trailing: { |
1110 | // The FIR argument has been split into a pair of distinct arguments. |
1111 | // The first part of the pair appears in the original argument |
1112 | // position. The second part of the pair is appended after all the |
1113 | // original arguments. (Boxchar arguments.) |
1114 | auto newBufArg = |
1115 | func.front().insertArgument(fixup.index, fixupType, loc); |
1116 | auto newLenArg = |
1117 | func.front().addArgument(trailingTys[fixup.second], loc); |
1118 | auto boxTy = oldArgTys[fixup.index - offset]; |
1119 | rewriter->setInsertionPointToStart(&func.front()); |
1120 | auto box = rewriter->create<fir::EmboxCharOp>(loc, boxTy, newBufArg, |
1121 | newLenArg); |
1122 | func.getArgument(fixup.index + 1).replaceAllUsesWith(box); |
1123 | func.front().eraseArgument(fixup.index + 1); |
1124 | } break; |
1125 | case FixupTy::Codes::TrailingCharProc: { |
1126 | // The FIR character procedure argument tuple must be split into a |
1127 | // pair of distinct arguments. The first part of the pair appears in |
1128 | // the original argument position. The second part of the pair is |
1129 | // appended after all the original arguments. |
1130 | auto newProcPointerArg = |
1131 | func.front().insertArgument(fixup.index, fixupType, loc); |
1132 | auto newLenArg = |
1133 | func.front().addArgument(trailingTys[fixup.second], loc); |
1134 | auto tupleType = oldArgTys[fixup.index - offset]; |
1135 | rewriter->setInsertionPointToStart(&func.front()); |
1136 | fir::FirOpBuilder builder(*rewriter, getModule()); |
1137 | auto tuple = fir::factory::createCharacterProcedureTuple( |
1138 | builder, loc, tupleType, newProcPointerArg, newLenArg); |
1139 | func.getArgument(fixup.index + 1).replaceAllUsesWith(tuple); |
1140 | func.front().eraseArgument(fixup.index + 1); |
1141 | } break; |
1142 | } |
1143 | } |
1144 | } |
1145 | |
1146 | llvm::SmallVector<mlir::Type> newInTypes = toTypeList(newInTyAndAttrs); |
1147 | // Set the new type and finalize the arguments, etc. |
1148 | newInTypes.insert(newInTypes.end(), trailingTys.begin(), trailingTys.end()); |
1149 | auto newFuncTy = |
1150 | mlir::FunctionType::get(func.getContext(), newInTypes, newResTys); |
1151 | LLVM_DEBUG(llvm::dbgs() << "new func: " << newFuncTy << '\n'); |
1152 | func.setType(newFuncTy); |
1153 | |
1154 | for (std::pair<unsigned, mlir::NamedAttribute> extraAttr : extraAttrs) |
1155 | func.setArgAttr(extraAttr.first, extraAttr.second.getName(), |
1156 | extraAttr.second.getValue()); |
1157 | |
1158 | for (auto [resId, resAttrList] : resultAttrs) |
1159 | for (mlir::NamedAttribute resAttr : resAttrList) |
1160 | func.setResultAttr(resId, resAttr.getName(), resAttr.getValue()); |
1161 | |
1162 | // Replace attributes to the correct argument if there was an argument shift |
1163 | // to the right. |
1164 | if (argumentShift > 0) { |
1165 | for (std::pair<unsigned, mlir::NamedAttribute> savedAttr : savedAttrs) { |
1166 | func.removeArgAttr(savedAttr.first, savedAttr.second.getName()); |
1167 | func.setArgAttr(savedAttr.first + argumentShift, |
1168 | savedAttr.second.getName(), |
1169 | savedAttr.second.getValue()); |
1170 | } |
1171 | } |
1172 | |
1173 | for (auto &fixup : fixups) { |
1174 | if constexpr (std::is_same_v<FuncOpTy, mlir::func::FuncOp>) |
1175 | if (fixup.finalizer) |
1176 | (*fixup.finalizer)(func); |
1177 | if constexpr (std::is_same_v<FuncOpTy, mlir::gpu::GPUFuncOp>) |
1178 | if (fixup.gpuFinalizer) |
1179 | (*fixup.gpuFinalizer)(func); |
1180 | } |
1181 | } |
1182 | |
1183 | template <typename OpTy, typename Ty, typename FIXUPS> |
1184 | void doReturn(OpTy func, Ty &newResTys, |
1185 | fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs, |
1186 | FIXUPS &fixups, fir::CodeGenSpecifics::Marshalling &m) { |
1187 | assert(m.size() == 1 && |
1188 | "expect result to be turned into single argument or result so far" ); |
1189 | auto &tup = m[0]; |
1190 | auto attr = std::get<fir::CodeGenSpecifics::Attributes>(tup); |
1191 | auto argTy = std::get<mlir::Type>(tup); |
1192 | if (attr.isSRet()) { |
1193 | unsigned argNo = newInTyAndAttrs.size(); |
1194 | if (auto align = attr.getAlignment()) |
1195 | fixups.emplace_back( |
1196 | FixupTy::Codes::ReturnAsStore, argNo, [=](OpTy func) { |
1197 | auto elemType = fir::dyn_cast_ptrOrBoxEleTy( |
1198 | func.getFunctionType().getInput(argNo)); |
1199 | func.setArgAttr(argNo, "llvm.sret" , |
1200 | mlir::TypeAttr::get(elemType)); |
1201 | func.setArgAttr(argNo, "llvm.align" , |
1202 | rewriter->getIntegerAttr( |
1203 | rewriter->getIntegerType(32), align)); |
1204 | }); |
1205 | else |
1206 | fixups.emplace_back(FixupTy::Codes::ReturnAsStore, argNo, |
1207 | [=](OpTy func) { |
1208 | auto elemType = fir::dyn_cast_ptrOrBoxEleTy( |
1209 | func.getFunctionType().getInput(argNo)); |
1210 | func.setArgAttr(argNo, "llvm.sret" , |
1211 | mlir::TypeAttr::get(elemType)); |
1212 | }); |
1213 | newInTyAndAttrs.push_back(tup); |
1214 | return; |
1215 | } |
1216 | if (auto align = attr.getAlignment()) |
1217 | fixups.emplace_back( |
1218 | FixupTy::Codes::ReturnType, newResTys.size(), [=](OpTy func) { |
1219 | func.setArgAttr( |
1220 | newResTys.size(), "llvm.align" , |
1221 | rewriter->getIntegerAttr(rewriter->getIntegerType(32), align)); |
1222 | }); |
1223 | else |
1224 | fixups.emplace_back(FixupTy::Codes::ReturnType, newResTys.size()); |
1225 | newResTys.push_back(argTy); |
1226 | } |
1227 | |
1228 | /// Convert a complex return value. This can involve converting the return |
1229 | /// value to a "hidden" first argument or packing the complex into a wide |
1230 | /// GPR. |
1231 | template <typename OpTy, typename Ty, typename FIXUPS> |
1232 | void doComplexReturn(OpTy func, mlir::ComplexType cmplx, Ty &newResTys, |
1233 | fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs, |
1234 | FIXUPS &fixups) { |
1235 | if (noComplexConversion) { |
1236 | newResTys.push_back(cmplx); |
1237 | return; |
1238 | } |
1239 | auto m = |
1240 | specifics->complexReturnType(func.getLoc(), cmplx.getElementType()); |
1241 | doReturn(func, newResTys, newInTyAndAttrs, fixups, m); |
1242 | } |
1243 | |
1244 | template <typename OpTy, typename Ty, typename FIXUPS> |
1245 | void doStructReturn(OpTy func, fir::RecordType recTy, Ty &newResTys, |
1246 | fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs, |
1247 | FIXUPS &fixups) { |
1248 | if (noStructConversion) { |
1249 | newResTys.push_back(recTy); |
1250 | return; |
1251 | } |
1252 | auto m = specifics->structReturnType(func.getLoc(), recTy); |
1253 | doReturn(func, newResTys, newInTyAndAttrs, fixups, m); |
1254 | } |
1255 | |
1256 | template <typename OpTy, typename FIXUPS> |
1257 | void createFuncOpArgFixups( |
1258 | OpTy func, fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs, |
1259 | fir::CodeGenSpecifics::Marshalling &argsInTys, FIXUPS &fixups) { |
1260 | const auto fixupCode = argsInTys.size() > 1 ? FixupTy::Codes::Split |
1261 | : FixupTy::Codes::ArgumentType; |
1262 | for (auto e : llvm::enumerate(argsInTys)) { |
1263 | auto &tup = e.value(); |
1264 | auto index = e.index(); |
1265 | auto attr = std::get<fir::CodeGenSpecifics::Attributes>(tup); |
1266 | auto argNo = newInTyAndAttrs.size(); |
1267 | if (attr.isByVal()) { |
1268 | if (auto align = attr.getAlignment()) |
1269 | fixups.emplace_back(FixupTy::Codes::ArgumentAsLoad, argNo, |
1270 | [=](OpTy func) { |
1271 | auto elemType = fir::dyn_cast_ptrOrBoxEleTy( |
1272 | func.getFunctionType().getInput(argNo)); |
1273 | func.setArgAttr(argNo, "llvm.byval" , |
1274 | mlir::TypeAttr::get(elemType)); |
1275 | func.setArgAttr( |
1276 | argNo, "llvm.align" , |
1277 | rewriter->getIntegerAttr( |
1278 | rewriter->getIntegerType(32), align)); |
1279 | }); |
1280 | else |
1281 | fixups.emplace_back(FixupTy::Codes::ArgumentAsLoad, |
1282 | newInTyAndAttrs.size(), [=](OpTy func) { |
1283 | auto elemType = fir::dyn_cast_ptrOrBoxEleTy( |
1284 | func.getFunctionType().getInput(argNo)); |
1285 | func.setArgAttr(argNo, "llvm.byval" , |
1286 | mlir::TypeAttr::get(elemType)); |
1287 | }); |
1288 | } else { |
1289 | if (auto align = attr.getAlignment()) |
1290 | fixups.emplace_back( |
1291 | fixupCode, argNo, index, [=](OpTy func) { |
1292 | func.setArgAttr(argNo, "llvm.align" , |
1293 | rewriter->getIntegerAttr( |
1294 | rewriter->getIntegerType(32), align)); |
1295 | }); |
1296 | else |
1297 | fixups.emplace_back(fixupCode, argNo, index); |
1298 | } |
1299 | newInTyAndAttrs.push_back(tup); |
1300 | } |
1301 | } |
1302 | |
1303 | /// Convert a complex argument value. This can involve storing the value to |
1304 | /// a temporary memory location or factoring the value into two distinct |
1305 | /// arguments. |
1306 | template <typename OpTy, typename FIXUPS> |
1307 | void doComplexArg(OpTy func, mlir::ComplexType cmplx, |
1308 | fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs, |
1309 | FIXUPS &fixups) { |
1310 | if (noComplexConversion) { |
1311 | newInTyAndAttrs.push_back(fir::CodeGenSpecifics::getTypeAndAttr(cmplx)); |
1312 | return; |
1313 | } |
1314 | auto cplxArgs = |
1315 | specifics->complexArgumentType(func.getLoc(), cmplx.getElementType()); |
1316 | createFuncOpArgFixups(func, newInTyAndAttrs, cplxArgs, fixups); |
1317 | } |
1318 | |
1319 | template <typename OpTy, typename FIXUPS> |
1320 | void doStructArg(OpTy func, fir::RecordType recTy, |
1321 | fir::CodeGenSpecifics::Marshalling &newInTyAndAttrs, |
1322 | FIXUPS &fixups) { |
1323 | if (noStructConversion) { |
1324 | newInTyAndAttrs.push_back(fir::CodeGenSpecifics::getTypeAndAttr(recTy)); |
1325 | return; |
1326 | } |
1327 | auto structArgs = |
1328 | specifics->structArgumentType(func.getLoc(), recTy, newInTyAndAttrs); |
1329 | createFuncOpArgFixups(func, newInTyAndAttrs, structArgs, fixups); |
1330 | } |
1331 | |
1332 | private: |
1333 | // Replace `op` and remove it. |
1334 | void replaceOp(mlir::Operation *op, mlir::ValueRange newValues) { |
1335 | op->replaceAllUsesWith(newValues); |
1336 | op->dropAllReferences(); |
1337 | op->erase(); |
1338 | } |
1339 | |
1340 | inline void setMembers(fir::CodeGenSpecifics *s, mlir::OpBuilder *r, |
1341 | mlir::DataLayout *dl) { |
1342 | specifics = s; |
1343 | rewriter = r; |
1344 | dataLayout = dl; |
1345 | } |
1346 | |
1347 | inline void clearMembers() { setMembers(nullptr, nullptr, nullptr); } |
1348 | |
1349 | // Inserts a call to llvm.stacksave at the current insertion |
1350 | // point and the given location. Returns the call's result Value. |
1351 | inline mlir::Value genStackSave(mlir::Location loc) { |
1352 | fir::FirOpBuilder builder(*rewriter, getModule()); |
1353 | return builder.genStackSave(loc); |
1354 | } |
1355 | |
1356 | // Inserts a call to llvm.stackrestore at the current insertion |
1357 | // point and the given location and argument. |
1358 | inline void genStackRestore(mlir::Location loc, mlir::Value sp) { |
1359 | fir::FirOpBuilder builder(*rewriter, getModule()); |
1360 | return builder.genStackRestore(loc, sp); |
1361 | } |
1362 | |
1363 | fir::CodeGenSpecifics *specifics = nullptr; |
1364 | mlir::OpBuilder *rewriter = nullptr; |
1365 | mlir::DataLayout *dataLayout = nullptr; |
1366 | }; |
1367 | } // namespace |
1368 | |