1//===- SerializeNVVMTarget.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/Config/mlir-config.h"
10#include "mlir/Dialect/GPU/IR/GPUDialect.h"
11#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
12#include "mlir/IR/MLIRContext.h"
13#include "mlir/InitAllDialects.h"
14#include "mlir/Parser/Parser.h"
15#include "mlir/Target/LLVM/NVVM/Target.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/NVVM/NVVMToLLVMIRTranslation.h"
20
21#include "llvm/Bitcode/BitcodeWriter.h"
22#include "llvm/Config/Targets.h" // for LLVM_HAS_NVPTX_TARGET
23#include "llvm/IRReader/IRReader.h"
24#include "llvm/Support/MemoryBufferRef.h"
25#include "llvm/Support/Process.h"
26#include "llvm/Support/SourceMgr.h"
27#include "llvm/Support/TargetSelect.h"
28#include "llvm/Support/raw_ostream.h"
29#include "llvm/TargetParser/Host.h"
30
31#include "gmock/gmock.h"
32#include <cstdint>
33
34using namespace mlir;
35
36// Skip the test if the NVPTX target was not built.
37#if LLVM_HAS_NVPTX_TARGET
38#define SKIP_WITHOUT_NVPTX(x) x
39#else
40#define SKIP_WITHOUT_NVPTX(x) DISABLED_##x
41#endif
42
43class MLIRTargetLLVMNVVM : public ::testing::Test {
44protected:
45 void SetUp() override {
46 registerBuiltinDialectTranslation(registry);
47 registerLLVMDialectTranslation(registry);
48 registerGPUDialectTranslation(registry);
49 registerNVVMDialectTranslation(registry);
50 NVVM::registerNVVMTargetInterfaceExternalModels(registry);
51 }
52
53 // Checks if PTXAS is in PATH.
54 bool hasPtxas() {
55 // Find the `ptxas` compiler.
56 std::optional<std::string> ptxasCompiler =
57 llvm::sys::Process::FindInEnvPath(EnvName: "PATH", FileName: "ptxas");
58 return ptxasCompiler.has_value();
59 }
60
61 // Dialect registry.
62 DialectRegistry registry;
63
64 // MLIR module used for the tests.
65 const std::string moduleStr = R"mlir(
66 gpu.module @nvvm_test {
67 llvm.func @nvvm_kernel(%arg0: f32) attributes {gpu.kernel, nvvm.kernel} {
68 llvm.return
69 }
70 })mlir";
71};
72
73// Test NVVM serialization to LLVM.
74TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMMToLLVM)) {
75 MLIRContext context(registry);
76
77 OwningOpRef<ModuleOp> module =
78 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
79 ASSERT_TRUE(!!module);
80
81 // Create an NVVM target.
82 NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);
83
84 // Serialize the module.
85 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
86 ASSERT_TRUE(!!serializer);
87 gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Offload);
88 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
89 std::optional<SmallVector<char, 0>> object =
90 serializer.serializeToObject(gpuModule, options);
91 // Check that the serializer was successful.
92 ASSERT_TRUE(object != std::nullopt);
93 ASSERT_TRUE(!object->empty());
94
95 // Read the serialized module.
96 llvm::MemoryBufferRef buffer(StringRef(object->data(), object->size()),
97 "module");
98 llvm::LLVMContext llvmContext;
99 llvm::Expected<std::unique_ptr<llvm::Module>> llvmModule =
100 llvm::getLazyBitcodeModule(buffer, llvmContext);
101 ASSERT_TRUE(!!llvmModule);
102 ASSERT_TRUE(!!*llvmModule);
103
104 // Check that it has a function named `foo`.
105 ASSERT_TRUE((*llvmModule)->getFunction("nvvm_kernel") != nullptr);
106 }
107}
108
109// Test NVVM serialization to PTX.
110TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToPTX)) {
111 MLIRContext context(registry);
112
113 OwningOpRef<ModuleOp> module =
114 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
115 ASSERT_TRUE(!!module);
116
117 // Create an NVVM target.
118 NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);
119
120 // Serialize the module.
121 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
122 ASSERT_TRUE(!!serializer);
123 gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Assembly);
124 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
125 std::optional<SmallVector<char, 0>> object =
126 serializer.serializeToObject(gpuModule, options);
127 // Check that the serializer was successful.
128 ASSERT_TRUE(object != std::nullopt);
129 ASSERT_TRUE(!object->empty());
130
131 ASSERT_TRUE(
132 StringRef(object->data(), object->size()).contains("nvvm_kernel"));
133 ASSERT_TRUE(StringRef(object->data(), object->size()).count('\0') == 0);
134 }
135}
136
137// Test NVVM serialization to Binary.
138TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToBinary)) {
139 if (!hasPtxas())
140 GTEST_SKIP() << "PTXAS compiler not found, skipping test.";
141
142 MLIRContext context(registry);
143
144 OwningOpRef<ModuleOp> module =
145 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
146 ASSERT_TRUE(!!module);
147
148 // Create an NVVM target.
149 NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);
150
151 // Serialize the module.
152 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
153 ASSERT_TRUE(!!serializer);
154 gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Binary);
155 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
156 std::optional<SmallVector<char, 0>> object =
157 serializer.serializeToObject(gpuModule, options);
158 // Check that the serializer was successful.
159 ASSERT_TRUE(object != std::nullopt);
160 ASSERT_TRUE(!object->empty());
161 }
162}
163
164// Test callback functions invoked with LLVM IR and ISA.
165TEST_F(MLIRTargetLLVMNVVM,
166 SKIP_WITHOUT_NVPTX(CallbackInvokedWithLLVMIRAndISA)) {
167 MLIRContext context(registry);
168
169 OwningOpRef<ModuleOp> module =
170 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
171 ASSERT_TRUE(!!module);
172
173 NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);
174
175 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
176 ASSERT_TRUE(!!serializer);
177
178 std::string initialLLVMIR;
179 auto initialCallback = [&initialLLVMIR](llvm::Module &module) {
180 llvm::raw_string_ostream ros(initialLLVMIR);
181 module.print(OS&: ros, AAW: nullptr);
182 };
183
184 std::string linkedLLVMIR;
185 auto linkedCallback = [&linkedLLVMIR](llvm::Module &module) {
186 llvm::raw_string_ostream ros(linkedLLVMIR);
187 module.print(OS&: ros, AAW: nullptr);
188 };
189
190 std::string optimizedLLVMIR;
191 auto optimizedCallback = [&optimizedLLVMIR](llvm::Module &module) {
192 llvm::raw_string_ostream ros(optimizedLLVMIR);
193 module.print(OS&: ros, AAW: nullptr);
194 };
195
196 std::string isaResult;
197 auto isaCallback = [&isaResult](llvm::StringRef isa) {
198 isaResult = isa.str();
199 };
200
201 gpu::TargetOptions options({}, {}, {}, {}, gpu::CompilationTarget::Assembly,
202 {}, initialCallback, linkedCallback,
203 optimizedCallback, isaCallback);
204
205 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
206 std::optional<SmallVector<char, 0>> object =
207 serializer.serializeToObject(gpuModule, options);
208
209 ASSERT_TRUE(object != std::nullopt);
210 ASSERT_TRUE(!object->empty());
211 ASSERT_TRUE(!initialLLVMIR.empty());
212 ASSERT_TRUE(!linkedLLVMIR.empty());
213 ASSERT_TRUE(!optimizedLLVMIR.empty());
214 ASSERT_TRUE(!isaResult.empty());
215
216 initialLLVMIR.clear();
217 linkedLLVMIR.clear();
218 optimizedLLVMIR.clear();
219 isaResult.clear();
220 }
221}
222
223// Test linking LLVM IR from a resource attribute.
224TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(LinkedLLVMIRResource)) {
225 MLIRContext context(registry);
226 std::string moduleStr = R"mlir(
227 gpu.module @nvvm_test {
228 llvm.func @bar()
229 llvm.func @nvvm_kernel(%arg0: f32) attributes {gpu.kernel, nvvm.kernel} {
230 llvm.call @bar() : () -> ()
231 llvm.return
232 }
233 }
234 )mlir";
235 // Provide the library to link as a serialized bitcode blob.
236 SmallVector<char> bitcodeToLink;
237 {
238 std::string linkedLib = R"llvm(
239 define void @bar() {
240 ret void
241 }
242 )llvm";
243 llvm::SMDiagnostic err;
244 llvm::MemoryBufferRef buffer(linkedLib, "linkedLib");
245 llvm::LLVMContext llvmCtx;
246 std::unique_ptr<llvm::Module> module = llvm::parseIR(Buffer: buffer, Err&: err, Context&: llvmCtx);
247 ASSERT_TRUE(module) << " Can't parse IR: " << err.getMessage();
248 {
249 llvm::raw_svector_ostream os(bitcodeToLink);
250 WriteBitcodeToFile(M: *module, Out&: os);
251 }
252 }
253
254 OwningOpRef<ModuleOp> module =
255 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
256 ASSERT_TRUE(!!module);
257 Builder builder(&context);
258
259 NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);
260 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
261
262 // Hook to intercept the LLVM IR after linking external libs.
263 std::string linkedLLVMIR;
264 auto linkedCallback = [&linkedLLVMIR](llvm::Module &module) {
265 llvm::raw_string_ostream ros(linkedLLVMIR);
266 module.print(OS&: ros, AAW: nullptr);
267 };
268
269 // Store the bitcode as a DenseI8ArrayAttr.
270 SmallVector<Attribute> librariesToLink;
271 librariesToLink.push_back(DenseI8ArrayAttr::get(
272 &context,
273 ArrayRef<int8_t>((int8_t *)bitcodeToLink.data(), bitcodeToLink.size())));
274 gpu::TargetOptions options({}, librariesToLink, {}, {},
275 gpu::CompilationTarget::Assembly, {}, {},
276 linkedCallback);
277 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
278 std::optional<SmallVector<char, 0>> object =
279 serializer.serializeToObject(gpuModule, options);
280
281 // Verify that we correctly linked in the library: the external call is
282 // replaced by the definition.
283 ASSERT_TRUE(!linkedLLVMIR.empty());
284 {
285 llvm::SMDiagnostic err;
286 llvm::MemoryBufferRef buffer(linkedLLVMIR, "linkedLLVMIR");
287 llvm::LLVMContext llvmCtx;
288 std::unique_ptr<llvm::Module> module =
289 llvm::parseIR(buffer, err, llvmCtx);
290 ASSERT_TRUE(module) << " Can't parse linkedLLVMIR: " << err.getMessage()
291 << " IR: \n\b" << linkedLLVMIR;
292 llvm::Function *bar = module->getFunction("bar");
293 ASSERT_TRUE(bar);
294 ASSERT_FALSE(bar->empty());
295 }
296 ASSERT_TRUE(object != std::nullopt);
297 ASSERT_TRUE(!object->empty());
298 }
299}
300

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