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 26 using 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 35 namespace { 36 // Dummy interface for testing. 37 class TargetAttrImpl 38 : public gpu::TargetAttrInterface::FallbackModel<TargetAttrImpl> { 39 public: 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 50 class MLIRTargetLLVM : public ::testing::Test { 51 protected: 52 void SetUp() override { 53 llvm::InitializeNativeTarget(); 54 llvm::InitializeNativeTargetAsmPrinter(); 55 registry.addExtension(+[](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 74 TEST_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, 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 102 std::optional<SmallVector<char, 0>> 103 TargetAttrImpl::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 114 Attribute 115 TargetAttrImpl::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. 130 TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(TargetAttrAPI)) { 131 MLIRContext context(registry); 132 context.loadAllAvailableDialects(); 133 134 OwningOpRef<ModuleOp> module = 135 parseSourceString<ModuleOp>(moduleStr, &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 160 TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(CallbackInvokedWithInitialLLVMIR)) { 161 MLIRContext context(registry); 162 163 OwningOpRef<ModuleOp> module = 164 parseSourceString<ModuleOp>(moduleStr, &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(ros, 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 188 TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(CallbackInvokedWithLinkedLLVMIR)) { 189 MLIRContext context(registry); 190 191 OwningOpRef<ModuleOp> module = 192 parseSourceString<ModuleOp>(moduleStr, &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(ros, 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 216 TEST_F(MLIRTargetLLVM, 217 SKIP_WITHOUT_NATIVE(CallbackInvokedWithOptimizedLLVMIR)) { 218 MLIRContext context(registry); 219 220 OwningOpRef<ModuleOp> module = 221 parseSourceString<ModuleOp>(moduleStr, &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(ros, 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 }