243 lines · cpp
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-exception6//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 == 030#define SKIP_WITHOUT_NATIVE(x) DISABLED_##x31#else32#define SKIP_WITHOUT_NATIVE(x) x33#endif34 35namespace {36// Dummy interface for testing.37class TargetAttrImpl38 : 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} // namespace49 50class MLIRTargetLLVM : public ::testing::Test {51protected: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.return70 }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, 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 114Attribute115TargetAttrImpl::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 object119 // 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 leverage129// 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>(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 in155 // `serializeToObject`.156 ASSERT_TRUE(properties.contains("serialize_attr"));157}158 159// Test callback function invoked with initial LLVM IR160TEST_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 IR188TEST_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 IR216TEST_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}243