1 //===- TestConvertGPUKernelToCubin.cpp - Test gpu kernel cubin lowering ---===//
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 #include "mlir/Dialect/GPU/Transforms/Passes.h"
10 #include "mlir/Pass/Pass.h"
11 #include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
12 #include "mlir/Target/LLVMIR/Export.h"
13 #include "llvm/Support/TargetSelect.h"
14 
15 using namespace mlir;
16 
17 #if MLIR_CUDA_CONVERSIONS_ENABLED
18 namespace {
19 class TestSerializeToCubinPass
20     : public PassWrapper<TestSerializeToCubinPass, gpu::SerializeToBlobPass> {
21 public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestSerializeToCubinPass)22   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestSerializeToCubinPass)
23 
24   StringRef getArgument() const final { return "test-gpu-to-cubin"; }
getDescription() const25   StringRef getDescription() const final {
26     return "Lower GPU kernel function to CUBIN binary annotations";
27   }
28   TestSerializeToCubinPass();
29 
30 private:
31   void getDependentDialects(DialectRegistry &registry) const override;
32 
33   // Serializes PTX to CUBIN.
34   std::unique_ptr<std::vector<char>>
35   serializeISA(const std::string &isa) override;
36 };
37 } // namespace
38 
TestSerializeToCubinPass()39 TestSerializeToCubinPass::TestSerializeToCubinPass() {
40   this->triple = "nvptx64-nvidia-cuda";
41   this->chip = "sm_35";
42   this->features = "+ptx60";
43 }
44 
getDependentDialects(DialectRegistry & registry) const45 void TestSerializeToCubinPass::getDependentDialects(
46     DialectRegistry &registry) const {
47   registerNVVMDialectTranslation(registry);
48   gpu::SerializeToBlobPass::getDependentDialects(registry);
49 }
50 
51 std::unique_ptr<std::vector<char>>
serializeISA(const std::string &)52 TestSerializeToCubinPass::serializeISA(const std::string &) {
53   std::string data = "CUBIN";
54   return std::make_unique<std::vector<char>>(data.begin(), data.end());
55 }
56 
57 namespace mlir {
58 namespace test {
59 // Register test pass to serialize GPU module to a CUBIN binary annotation.
registerTestGpuSerializeToCubinPass()60 void registerTestGpuSerializeToCubinPass() {
61   PassRegistration<TestSerializeToCubinPass>([] {
62     // Initialize LLVM NVPTX backend.
63     LLVMInitializeNVPTXTarget();
64     LLVMInitializeNVPTXTargetInfo();
65     LLVMInitializeNVPTXTargetMC();
66     LLVMInitializeNVPTXAsmPrinter();
67 
68     return std::make_unique<TestSerializeToCubinPass>();
69   });
70 }
71 } // namespace test
72 } // namespace mlir
73 #endif // MLIR_CUDA_CONVERSIONS_ENABLED
74