1//===- SerializeROCDLTarget.cpp ---------------------------------*- C++ -*-===//
2//
3// This file is licensed 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/IR/GPUDialect.h"
10#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
11#include "mlir/IR/MLIRContext.h"
12#include "mlir/InitAllDialects.h"
13#include "mlir/Parser/Parser.h"
14#include "mlir/Target/LLVM/ROCDL/Target.h"
15#include "mlir/Target/LLVM/ROCDL/Utils.h"
16#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
17#include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
18#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
19#include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"
20
21#include "llvm/IRReader/IRReader.h"
22#include "llvm/Support/FileSystem.h"
23#include "llvm/Support/MemoryBufferRef.h"
24#include "llvm/Support/Path.h"
25#include "llvm/Support/TargetSelect.h"
26#include "llvm/Support/raw_ostream.h"
27#include "llvm/TargetParser/Host.h"
28
29#include "gmock/gmock.h"
30
31using namespace mlir;
32
33// Skip the test if the AMDGPU target was not built.
34#if MLIR_ENABLE_ROCM_CONVERSIONS
35#define SKIP_WITHOUT_AMDGPU(x) x
36#else
37#define SKIP_WITHOUT_AMDGPU(x) DISABLED_##x
38#endif
39
40class MLIRTargetLLVMROCDL : public ::testing::Test {
41protected:
42 void SetUp() override {
43 registerBuiltinDialectTranslation(registry);
44 registerLLVMDialectTranslation(registry);
45 registerGPUDialectTranslation(registry);
46 registerROCDLDialectTranslation(registry);
47 ROCDL::registerROCDLTargetInterfaceExternalModels(registry);
48 }
49
50 // Checks if a ROCm installation is available.
51 bool hasROCMTools() {
52 StringRef rocmPath = ROCDL::getROCMPath();
53 if (rocmPath.empty())
54 return false;
55 llvm::SmallString<128> lldPath(rocmPath);
56 llvm::sys::path::append(path&: lldPath, a: "llvm", b: "bin", c: "ld.lld");
57 return llvm::sys::fs::can_execute(Path: lldPath);
58 }
59
60 // Dialect registry.
61 DialectRegistry registry;
62
63 // MLIR module used for the tests.
64 const std::string moduleStr = R"mlir(
65 gpu.module @rocdl_test {
66 llvm.func @rocdl_kernel(%arg0: f32) attributes {gpu.kernel, rocdl.kernel} {
67 llvm.return
68 }
69 })mlir";
70};
71
72// Test ROCDL serialization to LLVM.
73TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(SerializeROCDLMToLLVM)) {
74 MLIRContext context(registry);
75
76 OwningOpRef<ModuleOp> module =
77 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
78 ASSERT_TRUE(!!module);
79
80 // Create a ROCDL target.
81 ROCDL::ROCDLTargetAttr target = ROCDL::ROCDLTargetAttr::get(&context);
82
83 // Serialize the module.
84 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
85 ASSERT_TRUE(!!serializer);
86 gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Offload);
87 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
88 std::optional<SmallVector<char, 0>> object =
89 serializer.serializeToObject(gpuModule, options);
90 // Check that the serializer was successful.
91 ASSERT_TRUE(object != std::nullopt);
92 ASSERT_TRUE(!object->empty());
93
94 // Read the serialized module.
95 llvm::MemoryBufferRef buffer(StringRef(object->data(), object->size()),
96 "module");
97 llvm::LLVMContext llvmContext;
98 llvm::Expected<std::unique_ptr<llvm::Module>> llvmModule =
99 llvm::getLazyBitcodeModule(buffer, llvmContext);
100 ASSERT_TRUE(!!llvmModule);
101 ASSERT_TRUE(!!*llvmModule);
102
103 // Check that it has a function named `foo`.
104 ASSERT_TRUE((*llvmModule)->getFunction("rocdl_kernel") != nullptr);
105 }
106}
107
108// Test ROCDL serialization to PTX.
109TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(SerializeROCDLToPTX)) {
110 MLIRContext context(registry);
111
112 OwningOpRef<ModuleOp> module =
113 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
114 ASSERT_TRUE(!!module);
115
116 // Create a ROCDL target.
117 ROCDL::ROCDLTargetAttr target = ROCDL::ROCDLTargetAttr::get(&context);
118
119 // Serialize the module.
120 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
121 ASSERT_TRUE(!!serializer);
122 gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Assembly);
123 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
124 std::optional<SmallVector<char, 0>> object =
125 serializer.serializeToObject(gpuModule, options);
126 // Check that the serializer was successful.
127 ASSERT_TRUE(object != std::nullopt);
128 ASSERT_TRUE(!object->empty());
129
130 ASSERT_TRUE(
131 StringRef(object->data(), object->size()).contains("rocdl_kernel"));
132 }
133}
134
135// Test ROCDL serialization to Binary.
136TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(SerializeROCDLToBinary)) {
137 if (!hasROCMTools())
138 GTEST_SKIP() << "ROCm installation not found, skipping test.";
139
140 MLIRContext context(registry);
141
142 OwningOpRef<ModuleOp> module =
143 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
144 ASSERT_TRUE(!!module);
145
146 // Create a ROCDL target.
147 ROCDL::ROCDLTargetAttr target = ROCDL::ROCDLTargetAttr::get(&context);
148
149 // Serialize the module.
150 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
151 ASSERT_TRUE(!!serializer);
152 gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Binary);
153 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
154 std::optional<SmallVector<char, 0>> object =
155 serializer.serializeToObject(gpuModule, options);
156 // Check that the serializer was successful.
157 ASSERT_TRUE(object != std::nullopt);
158 ASSERT_FALSE(object->empty());
159 }
160}
161

source code of mlir/unittests/Target/LLVM/SerializeROCDLTarget.cpp