1 //===- SerializeToBlob.cpp - MLIR GPU lowering pass -----------------------===//
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 base class for a pass to serialize a gpu module
10 // into a binary blob that can be executed on a GPU. The binary blob is added
11 // as a string attribute to the gpu module.
12 //
13 //===----------------------------------------------------------------------===//
14 
15 #include "mlir/Dialect/GPU/Passes.h"
16 #include "mlir/Pass/Pass.h"
17 #include "llvm/IR/LegacyPassManager.h"
18 #include "llvm/Support/TargetRegistry.h"
19 #include "llvm/Support/TargetSelect.h"
20 #include "llvm/Target/TargetMachine.h"
21 
22 using namespace mlir;
23 
24 std::string gpu::getDefaultGpuBinaryAnnotation() { return "gpu.binary"; }
25 
26 gpu::SerializeToBlobPass::SerializeToBlobPass(TypeID passID)
27     : OperationPass<gpu::GPUModuleOp>(passID) {}
28 
29 gpu::SerializeToBlobPass::SerializeToBlobPass(const SerializeToBlobPass &other)
30     : OperationPass<gpu::GPUModuleOp>(other) {
31   // Pass::Option has no copy constructor, copy them manually.
32   triple = other.triple;
33   chip = other.chip;
34   features = other.features;
35   gpuBinaryAnnotation = other.gpuBinaryAnnotation;
36 }
37 
38 static std::string translateToISA(llvm::Module &llvmModule,
39                                   llvm::TargetMachine &targetMachine) {
40   llvmModule.setDataLayout(targetMachine.createDataLayout());
41 
42   std::string targetISA;
43   llvm::raw_string_ostream stream(targetISA);
44   llvm::buffer_ostream pstream(stream);
45   llvm::legacy::PassManager codegenPasses;
46   targetMachine.addPassesToEmitFile(codegenPasses, pstream, nullptr,
47                                     llvm::CGFT_AssemblyFile);
48   codegenPasses.run(llvmModule);
49   return targetISA;
50 }
51 
52 void gpu::SerializeToBlobPass::runOnOperation() {
53   // Lower the module to an LLVM IR module using a separate context to enable
54   // multi-threaded processing.
55   llvm::LLVMContext llvmContext;
56   std::unique_ptr<llvm::Module> llvmModule = translateToLLVMIR(llvmContext);
57   if (!llvmModule)
58     return signalPassFailure();
59 
60   // Lower the LLVM IR module to target ISA.
61   std::unique_ptr<llvm::TargetMachine> targetMachine = createTargetMachine();
62   if (!targetMachine)
63     return signalPassFailure();
64 
65   std::string targetISA = translateToISA(*llvmModule, *targetMachine);
66 
67   // Serialize the target ISA.
68   std::unique_ptr<std::vector<char>> blob = serializeISA(targetISA);
69   if (!blob)
70     return signalPassFailure();
71 
72   // Add the blob as module attribute.
73   auto attr = StringAttr::get(&getContext(), {blob->data(), blob->size()});
74   getOperation()->setAttr(gpuBinaryAnnotation, attr);
75 }
76 
77 std::unique_ptr<llvm::TargetMachine>
78 gpu::SerializeToBlobPass::createTargetMachine() {
79   Location loc = getOperation().getLoc();
80   std::string error;
81   const llvm::Target *target =
82       llvm::TargetRegistry::lookupTarget(triple, error);
83   if (!target) {
84     emitError(loc, Twine("failed to lookup target: ") + error);
85     return {};
86   }
87   llvm::TargetMachine *machine =
88       target->createTargetMachine(triple, chip, features, {}, {});
89   if (!machine) {
90     emitError(loc, "failed to create target machine");
91     return {};
92   }
93 
94   return std::unique_ptr<llvm::TargetMachine>{machine};
95 }
96