1 | //===- GPUToLLVMIRTranslation.cpp - Translate GPU dialect to LLVM IR ------===// |
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 | // This file implements a translation between the MLIR GPU dialect and LLVM IR. |
10 | // |
11 | //===----------------------------------------------------------------------===// |
12 | #include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h" |
13 | #include "mlir/Dialect/GPU/IR/GPUDialect.h" |
14 | #include "mlir/Target/LLVMIR/LLVMTranslationInterface.h" |
15 | #include "llvm/ADT/TypeSwitch.h" |
16 | |
17 | using namespace mlir; |
18 | |
19 | namespace { |
20 | LogicalResult launchKernel(gpu::LaunchFuncOp launchOp, |
21 | llvm::IRBuilderBase &builder, |
22 | LLVM::ModuleTranslation &moduleTranslation) { |
23 | auto kernelBinary = SymbolTable::lookupNearestSymbolFrom<gpu::BinaryOp>( |
24 | launchOp, launchOp.getKernelModuleName()); |
25 | if (!kernelBinary) { |
26 | launchOp.emitError("Couldn't find the binary holding the kernel: " ) |
27 | << launchOp.getKernelModuleName(); |
28 | return failure(); |
29 | } |
30 | auto offloadingHandler = |
31 | dyn_cast<gpu::OffloadingLLVMTranslationAttrInterface>( |
32 | kernelBinary.getOffloadingHandlerAttr()); |
33 | assert(offloadingHandler && "Invalid offloading handler." ); |
34 | return offloadingHandler.launchKernel(launchOp, kernelBinary, builder, |
35 | moduleTranslation); |
36 | } |
37 | |
38 | class GPUDialectLLVMIRTranslationInterface |
39 | : public LLVMTranslationDialectInterface { |
40 | public: |
41 | using LLVMTranslationDialectInterface::LLVMTranslationDialectInterface; |
42 | |
43 | LogicalResult |
44 | convertOperation(Operation *operation, llvm::IRBuilderBase &builder, |
45 | LLVM::ModuleTranslation &moduleTranslation) const override { |
46 | return llvm::TypeSwitch<Operation *, LogicalResult>(operation) |
47 | .Case(caseFn: [&](gpu::GPUModuleOp) { return success(); }) |
48 | .Case(caseFn: [&](gpu::BinaryOp op) { |
49 | auto offloadingHandler = |
50 | dyn_cast<gpu::OffloadingLLVMTranslationAttrInterface>( |
51 | op.getOffloadingHandlerAttr()); |
52 | assert(offloadingHandler && "Invalid offloading handler." ); |
53 | return offloadingHandler.embedBinary(op, builder, moduleTranslation); |
54 | }) |
55 | .Case(caseFn: [&](gpu::LaunchFuncOp op) { |
56 | return launchKernel(op, builder, moduleTranslation); |
57 | }) |
58 | .Default(defaultFn: [&](Operation *op) { |
59 | return op->emitError(message: "unsupported GPU operation: " ) << op->getName(); |
60 | }); |
61 | } |
62 | }; |
63 | |
64 | } // namespace |
65 | |
66 | void mlir::registerGPUDialectTranslation(DialectRegistry ®istry) { |
67 | registry.insert<gpu::GPUDialect>(); |
68 | registry.addExtension(extensionFn: +[](MLIRContext *ctx, gpu::GPUDialect *dialect) { |
69 | dialect->addInterfaces<GPUDialectLLVMIRTranslationInterface>(); |
70 | }); |
71 | } |
72 | |
73 | void mlir::registerGPUDialectTranslation(MLIRContext &context) { |
74 | DialectRegistry registry; |
75 | registerGPUDialectTranslation(registry); |
76 | context.appendDialectRegistry(registry); |
77 | } |
78 | |