mlir/lib/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.cpp

   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([&](gpu::GPUModuleOp) { return success(); })
  48         .Case([&](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([&](gpu::LaunchFuncOp op) {
  56           return launchKernel(op, builder, moduleTranslation);
  57         })
  58         .Default([&](Operation *op) {
  59           return op->emitError("unsupported GPU operation: ") << op->getName();
  60         });
  61   }
  62 };
  63
  64 } // namespace
  65
  66 void mlir::registerGPUDialectTranslation(DialectRegistry &registry) {
  67   registry.insert<gpu::GPUDialect>();
  68   registry.addExtension(+[](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 }