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(SerializeROCDLToLLVM)) {
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// Test ROCDL serialization to ISA with default code object version.
108TEST_F(MLIRTargetLLVMROCDL,
109 SKIP_WITHOUT_AMDGPU(SerializeROCDLToISAWithDefaultCOV)) {
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 EXPECT_TRUE(StringRef(object->data(), object->size())
128 .contains(".amdhsa_code_object_version 6"));
129 }
130}
131
132// Test ROCDL serialization to ISA with non-default code object version.
133TEST_F(MLIRTargetLLVMROCDL,
134 SKIP_WITHOUT_AMDGPU(SerializeROCDLToISAWithNonDefaultCOV)) {
135 MLIRContext context(registry);
136
137 OwningOpRef<ModuleOp> module =
138 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
139 ASSERT_TRUE(!!module);
140
141 // Create a ROCDL target.
142 ROCDL::ROCDLTargetAttr target = ROCDL::ROCDLTargetAttr::get(
143 &context, 2, "amdgcn-amd-amdhsa", "gfx900", "", "400");
144
145 // Serialize the module.
146 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
147 ASSERT_TRUE(!!serializer);
148 gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Assembly);
149 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
150 std::optional<SmallVector<char, 0>> object =
151 serializer.serializeToObject(gpuModule, options);
152 // Check that the serializer was successful.
153 EXPECT_TRUE(StringRef(object->data(), object->size())
154 .contains(".amdhsa_code_object_version 4"));
155 }
156}
157
158// Test ROCDL serialization to PTX.
159TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(SerializeROCDLToPTX)) {
160 MLIRContext context(registry);
161
162 OwningOpRef<ModuleOp> module =
163 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
164 ASSERT_TRUE(!!module);
165
166 // Create a ROCDL target.
167 ROCDL::ROCDLTargetAttr target = ROCDL::ROCDLTargetAttr::get(&context);
168
169 // Serialize the module.
170 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
171 ASSERT_TRUE(!!serializer);
172 gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Assembly);
173 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
174 std::optional<SmallVector<char, 0>> object =
175 serializer.serializeToObject(gpuModule, options);
176 // Check that the serializer was successful.
177 ASSERT_TRUE(object != std::nullopt);
178 ASSERT_TRUE(!object->empty());
179
180 ASSERT_TRUE(
181 StringRef(object->data(), object->size()).contains("rocdl_kernel"));
182 }
183}
184
185// Test ROCDL serialization to Binary.
186TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(SerializeROCDLToBinary)) {
187 if (!hasROCMTools())
188 GTEST_SKIP() << "ROCm installation not found, skipping test.";
189
190 MLIRContext context(registry);
191
192 OwningOpRef<ModuleOp> module =
193 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
194 ASSERT_TRUE(!!module);
195
196 // Create a ROCDL target.
197 ROCDL::ROCDLTargetAttr target = ROCDL::ROCDLTargetAttr::get(&context);
198
199 // Serialize the module.
200 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
201 ASSERT_TRUE(!!serializer);
202 gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Binary);
203 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
204 std::optional<SmallVector<char, 0>> object =
205 serializer.serializeToObject(gpuModule, options);
206 // Check that the serializer was successful.
207 ASSERT_TRUE(object != std::nullopt);
208 ASSERT_FALSE(object->empty());
209 }
210}
211
212// Test ROCDL metadata.
213TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(GetELFMetadata)) {
214 if (!hasROCMTools())
215 GTEST_SKIP() << "ROCm installation not found, skipping test.";
216
217 MLIRContext context(registry);
218
219 // MLIR module used for the tests.
220 const std::string moduleStr = R"mlir(
221 gpu.module @rocdl_test {
222 llvm.func @rocdl_kernel_1(%arg0: f32) attributes {gpu.kernel, rocdl.kernel} {
223 llvm.return
224 }
225 llvm.func @rocdl_kernel_0(%arg0: f32) attributes {gpu.kernel, rocdl.kernel} {
226 llvm.return
227 }
228 llvm.func @rocdl_kernel_2(%arg0: f32) attributes {gpu.kernel, rocdl.kernel} {
229 llvm.return
230 }
231 llvm.func @a_kernel(%arg0: f32) attributes {gpu.kernel, rocdl.kernel} {
232 llvm.return
233 }
234 })mlir";
235
236 OwningOpRef<ModuleOp> module =
237 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
238 ASSERT_TRUE(!!module);
239
240 // Create a ROCDL target.
241 ROCDL::ROCDLTargetAttr target = ROCDL::ROCDLTargetAttr::get(&context);
242
243 // Serialize the module.
244 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
245 ASSERT_TRUE(!!serializer);
246 gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Binary);
247 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
248 std::optional<SmallVector<char, 0>> object =
249 serializer.serializeToObject(gpuModule, options);
250 // Check that the serializer was successful.
251 ASSERT_TRUE(object != std::nullopt);
252 ASSERT_FALSE(object->empty());
253 if (!object)
254 continue;
255 // Get the metadata.
256 gpu::KernelTableAttr metadata =
257 ROCDL::getKernelMetadata(gpuModule, *object);
258 ASSERT_TRUE(metadata != nullptr);
259 // There should be 4 kernels.
260 ASSERT_TRUE(metadata.size() == 4);
261 // Check that the lookup method returns finds the kernel.
262 ASSERT_TRUE(metadata.lookup("a_kernel") != nullptr);
263 ASSERT_TRUE(metadata.lookup("rocdl_kernel_0") != nullptr);
264 // Check that the kernel doesn't exist.
265 ASSERT_TRUE(metadata.lookup("not_existent_kernel") == nullptr);
266 // Test the `KernelMetadataAttr` iterators.
267 for (gpu::KernelMetadataAttr kernel : metadata) {
268 // Check that the ELF metadata is present.
269 ASSERT_TRUE(kernel.getMetadata() != nullptr);
270 // Verify that `sgpr_count` is present and it is an integer attribute.
271 ASSERT_TRUE(kernel.getAttr<IntegerAttr>("sgpr_count") != nullptr);
272 // Verify that `vgpr_count` is present and it is an integer attribute.
273 ASSERT_TRUE(kernel.getAttr<IntegerAttr>("vgpr_count") != nullptr);
274 }
275 }
276}
277

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