1//===- SerializeToLLVMBitcode.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/IR/BuiltinDialect.h"
11#include "mlir/IR/BuiltinOps.h"
12#include "mlir/IR/MLIRContext.h"
13#include "mlir/Parser/Parser.h"
14#include "mlir/Target/LLVM/ModuleToObject.h"
15#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
16#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
17
18#include "llvm/IRReader/IRReader.h"
19#include "llvm/Support/MemoryBufferRef.h"
20#include "llvm/Support/TargetSelect.h"
21#include "llvm/Support/raw_ostream.h"
22#include "llvm/TargetParser/Host.h"
23
24#include "gmock/gmock.h"
25
26using namespace mlir;
27
28// Skip the test if the native target was not built.
29#if LLVM_NATIVE_TARGET_TEST_ENABLED == 0
30#define SKIP_WITHOUT_NATIVE(x) DISABLED_##x
31#else
32#define SKIP_WITHOUT_NATIVE(x) x
33#endif
34
35namespace {
36// Dummy interface for testing.
37class TargetAttrImpl
38 : public gpu::TargetAttrInterface::FallbackModel<TargetAttrImpl> {
39public:
40 std::optional<SmallVector<char, 0>>
41 serializeToObject(Attribute attribute, Operation *module,
42 const gpu::TargetOptions &options) const;
43
44 Attribute createObject(Attribute attribute, Operation *module,
45 const SmallVector<char, 0> &object,
46 const gpu::TargetOptions &options) const;
47};
48} // namespace
49
50class MLIRTargetLLVM : public ::testing::Test {
51protected:
52 void SetUp() override {
53 llvm::InitializeNativeTarget();
54 llvm::InitializeNativeTargetAsmPrinter();
55 registry.addExtension(extensionFn: +[](MLIRContext *ctx, BuiltinDialect *dialect) {
56 IntegerAttr::attachInterface<TargetAttrImpl>(*ctx);
57 });
58 registerBuiltinDialectTranslation(registry);
59 registerLLVMDialectTranslation(registry);
60 registry.insert<gpu::GPUDialect>();
61 }
62
63 // Dialect registry.
64 DialectRegistry registry;
65
66 // MLIR module used for the tests.
67 std::string moduleStr = R"mlir(
68 llvm.func @foo(%arg0 : i32) {
69 llvm.return
70 }
71 )mlir";
72};
73
74TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(SerializeToLLVMBitcode)) {
75 MLIRContext context(registry);
76
77 OwningOpRef<ModuleOp> module =
78 parseSourceString<ModuleOp>(moduleStr, &context);
79 ASSERT_TRUE(!!module);
80
81 // Serialize the module.
82 std::string targetTriple = llvm::sys::getProcessTriple();
83 LLVM::ModuleToObject serializer(*(module->getOperation()), targetTriple, "",
84 "");
85 std::optional<SmallVector<char, 0>> serializedModule = serializer.run();
86 ASSERT_TRUE(!!serializedModule);
87 ASSERT_TRUE(!serializedModule->empty());
88
89 // Read the serialized module.
90 llvm::MemoryBufferRef buffer(
91 StringRef(serializedModule->data(), serializedModule->size()), "module");
92 llvm::LLVMContext llvmContext;
93 llvm::Expected<std::unique_ptr<llvm::Module>> llvmModule =
94 llvm::getLazyBitcodeModule(Buffer: buffer, Context&: llvmContext);
95 ASSERT_TRUE(!!llvmModule);
96 ASSERT_TRUE(!!*llvmModule);
97
98 // Check that it has a function named `foo`.
99 ASSERT_TRUE((*llvmModule)->getFunction("foo") != nullptr);
100}
101
102std::optional<SmallVector<char, 0>>
103TargetAttrImpl::serializeToObject(Attribute attribute, Operation *module,
104 const gpu::TargetOptions &options) const {
105 // Set a dummy attr to be retrieved by `createObject`.
106 module->setAttr("serialize_attr", UnitAttr::get(module->getContext()));
107 std::string targetTriple = llvm::sys::getProcessTriple();
108 LLVM::ModuleToObject serializer(
109 *module, targetTriple, "", "", 3, options.getInitialLlvmIRCallback(),
110 options.getLinkedLlvmIRCallback(), options.getOptimizedLlvmIRCallback());
111 return serializer.run();
112}
113
114Attribute
115TargetAttrImpl::createObject(Attribute attribute, Operation *module,
116 const SmallVector<char, 0> &object,
117 const gpu::TargetOptions &options) const {
118 // Create a GPU object with the GPU module dictionary as the object
119 // properties.
120 return gpu::ObjectAttr::get(
121 module->getContext(), attribute, gpu::CompilationTarget::Offload,
122 StringAttr::get(module->getContext(),
123 StringRef(object.data(), object.size())),
124 module->getAttrDictionary(), /*kernels=*/nullptr);
125}
126
127// This test checks the correct functioning of `TargetAttrInterface` as an API.
128// In particular, it shows how `TargetAttrInterface::createObject` can leverage
129// the `module` operation argument to retrieve information from the module.
130TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(TargetAttrAPI)) {
131 MLIRContext context(registry);
132 context.loadAllAvailableDialects();
133
134 OwningOpRef<ModuleOp> module =
135 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
136 ASSERT_TRUE(!!module);
137 Builder builder(&context);
138 IntegerAttr target = builder.getI32IntegerAttr(0);
139 auto targetAttr = dyn_cast<gpu::TargetAttrInterface>(target);
140 // Check the attribute holds the interface.
141 ASSERT_TRUE(!!targetAttr);
142 gpu::TargetOptions opts;
143 std::optional<SmallVector<char, 0>> serializedBinary =
144 targetAttr.serializeToObject(*module, opts);
145 // Check the serialized string.
146 ASSERT_TRUE(!!serializedBinary);
147 ASSERT_TRUE(!serializedBinary->empty());
148 // Create the object attribute.
149 auto object = cast<gpu::ObjectAttr>(
150 targetAttr.createObject(*module, *serializedBinary, opts));
151 // Check the object has properties.
152 DictionaryAttr properties = object.getProperties();
153 ASSERT_TRUE(!!properties);
154 // Check that it contains the attribute added to the module in
155 // `serializeToObject`.
156 ASSERT_TRUE(properties.contains("serialize_attr"));
157}
158
159// Test callback function invoked with initial LLVM IR
160TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(CallbackInvokedWithInitialLLVMIR)) {
161 MLIRContext context(registry);
162
163 OwningOpRef<ModuleOp> module =
164 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
165 ASSERT_TRUE(!!module);
166 Builder builder(&context);
167 IntegerAttr target = builder.getI32IntegerAttr(0);
168 auto targetAttr = dyn_cast<gpu::TargetAttrInterface>(target);
169
170 std::string initialLLVMIR;
171 auto initialCallback = [&initialLLVMIR](llvm::Module &module) {
172 llvm::raw_string_ostream ros(initialLLVMIR);
173 module.print(OS&: ros, AAW: nullptr);
174 };
175
176 gpu::TargetOptions opts(
177 {}, {}, {}, {}, mlir::gpu::TargetOptions::getDefaultCompilationTarget(),
178 {}, initialCallback);
179 std::optional<SmallVector<char, 0>> serializedBinary =
180 targetAttr.serializeToObject(*module, opts);
181
182 ASSERT_TRUE(serializedBinary != std::nullopt);
183 ASSERT_TRUE(!serializedBinary->empty());
184 ASSERT_TRUE(!initialLLVMIR.empty());
185}
186
187// Test callback function invoked with linked LLVM IR
188TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(CallbackInvokedWithLinkedLLVMIR)) {
189 MLIRContext context(registry);
190
191 OwningOpRef<ModuleOp> module =
192 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
193 ASSERT_TRUE(!!module);
194 Builder builder(&context);
195 IntegerAttr target = builder.getI32IntegerAttr(0);
196 auto targetAttr = dyn_cast<gpu::TargetAttrInterface>(target);
197
198 std::string linkedLLVMIR;
199 auto linkedCallback = [&linkedLLVMIR](llvm::Module &module) {
200 llvm::raw_string_ostream ros(linkedLLVMIR);
201 module.print(OS&: ros, AAW: nullptr);
202 };
203
204 gpu::TargetOptions opts(
205 {}, {}, {}, {}, mlir::gpu::TargetOptions::getDefaultCompilationTarget(),
206 {}, {}, linkedCallback);
207 std::optional<SmallVector<char, 0>> serializedBinary =
208 targetAttr.serializeToObject(*module, opts);
209
210 ASSERT_TRUE(serializedBinary != std::nullopt);
211 ASSERT_TRUE(!serializedBinary->empty());
212 ASSERT_TRUE(!linkedLLVMIR.empty());
213}
214
215// Test callback function invoked with optimized LLVM IR
216TEST_F(MLIRTargetLLVM,
217 SKIP_WITHOUT_NATIVE(CallbackInvokedWithOptimizedLLVMIR)) {
218 MLIRContext context(registry);
219
220 OwningOpRef<ModuleOp> module =
221 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
222 ASSERT_TRUE(!!module);
223 Builder builder(&context);
224 IntegerAttr target = builder.getI32IntegerAttr(0);
225 auto targetAttr = dyn_cast<gpu::TargetAttrInterface>(target);
226
227 std::string optimizedLLVMIR;
228 auto optimizedCallback = [&optimizedLLVMIR](llvm::Module &module) {
229 llvm::raw_string_ostream ros(optimizedLLVMIR);
230 module.print(OS&: ros, AAW: nullptr);
231 };
232
233 gpu::TargetOptions opts(
234 {}, {}, {}, {}, mlir::gpu::TargetOptions::getDefaultCompilationTarget(),
235 {}, {}, {}, optimizedCallback);
236 std::optional<SmallVector<char, 0>> serializedBinary =
237 targetAttr.serializeToObject(*module, opts);
238
239 ASSERT_TRUE(serializedBinary != std::nullopt);
240 ASSERT_TRUE(!serializedBinary->empty());
241 ASSERT_TRUE(!optimizedLLVMIR.empty());
242}
243

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