mlir/unittests/Target/LLVM/SerializeToLLVMBitcode.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-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 }