mlir/unittests/Bytecode/BytecodeTest.cpp

   1 //===- AdaptorTest.cpp - Adaptor unit tests -------------------------------===//
   2 //
   3 // Part of the LLVM Project, 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/Bytecode/BytecodeReader.h"
  10 #include "mlir/Bytecode/BytecodeWriter.h"
  11 #include "mlir/IR/AsmState.h"
  12 #include "mlir/IR/BuiltinAttributes.h"
  13 #include "mlir/IR/OpImplementation.h"
  14 #include "mlir/IR/OwningOpRef.h"
  15 #include "mlir/Parser/Parser.h"
  16
  17 #include "llvm/ADT/StringRef.h"
  18 #include "llvm/Support/Endian.h"
  19 #include "llvm/Support/MemoryBufferRef.h"
  20 #include "gmock/gmock.h"
  21 #include "gtest/gtest.h"
  22
  23 using namespace llvm;
  24 using namespace mlir;
  25
  26 StringLiteral IRWithResources = R"(
  27 module @TestDialectResources attributes {
  28   bytecode.test = dense_resource<resource> : tensor<4xi32>
  29 } {}
  30 {-#
  31   dialect_resources: {
  32     builtin: {
  33       resource: "0x2000000001000000020000000300000004000000",
  34       resource_2: "0x2000000001000000020000000300000004000000"
  35     }
  36   }
  37 #-}
  38 )";
  39
  40 TEST(Bytecode, MultiModuleWithResource) {
  41   MLIRContext context;
  42   Builder builder(&context);
  43   ParserConfig parseConfig(&context);
  44   OwningOpRef<Operation *> module =
  45       parseSourceString<Operation *>(IRWithResources, parseConfig);
  46   ASSERT_TRUE(module);
  47
  48   // Write the module to bytecode
  49   std::string buffer;
  50   llvm::raw_string_ostream ostream(buffer);
  51   ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), ostream)));
  52   ostream.flush();
  53
  54   // Create copy of buffer which is aligned to requested resource alignment.
  55   constexpr size_t kAlignment = 0x20;
  56   size_t buffer_size = buffer.size();
  57   buffer.reserve(buffer_size + kAlignment - 1);
  58   size_t pad = ~(uintptr_t)buffer.data() + 1 & kAlignment - 1;
  59   buffer.insert(0, pad, ' ');
  60   StringRef aligned_buffer(buffer.data() + pad, buffer_size);
  61
  62   // Parse it back
  63   OwningOpRef<Operation *> roundTripModule =
  64       parseSourceString<Operation *>(aligned_buffer, parseConfig);
  65   ASSERT_TRUE(roundTripModule);
  66
  67   // FIXME: Parsing external resources does not work on big-endian
  68   // platforms currently.
  69   if (llvm::endianness::native == llvm::endianness::big)
  70     GTEST_SKIP();
  71
  72   // Try to see if we have a valid resource in the parsed module.
  73   auto checkResourceAttribute = [&](Operation *op) {
  74     Attribute attr = roundTripModule->getDiscardableAttr("bytecode.test");
  75     ASSERT_TRUE(attr);
  76     auto denseResourceAttr = dyn_cast<DenseI32ResourceElementsAttr>(attr);
  77     ASSERT_TRUE(denseResourceAttr);
  78     std::optional<ArrayRef<int32_t>> attrData =
  79         denseResourceAttr.tryGetAsArrayRef();
  80     ASSERT_TRUE(attrData.has_value());
  81     ASSERT_EQ(attrData->size(), static_cast<size_t>(4));
  82     EXPECT_EQ((*attrData)[0], 1);
  83     EXPECT_EQ((*attrData)[1], 2);
  84     EXPECT_EQ((*attrData)[2], 3);
  85     EXPECT_EQ((*attrData)[3], 4);
  86   };
  87
  88   checkResourceAttribute(*module);
  89   checkResourceAttribute(*roundTripModule);
  90 }
  91
  92 namespace {
  93 /// A custom operation for the purpose of showcasing how discardable attributes
  94 /// are handled in absence of properties.
  95 class OpWithoutProperties : public Op<OpWithoutProperties> {
  96 public:
  97   // Begin boilerplate.
  98   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpWithoutProperties)
  99   using Op::Op;
 100   static ArrayRef<StringRef> getAttributeNames() {
 101     static StringRef attributeNames[] = {StringRef("inherent_attr")};
 102     return ArrayRef(attributeNames);
 103   };
 104   static StringRef getOperationName() {
 105     return "test_op_properties.op_without_properties";
 106   }
 107   // End boilerplate.
 108 };
 109
 110 // A trivial supporting dialect to register the above operation.
 111 class TestOpPropertiesDialect : public Dialect {
 112 public:
 113   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestOpPropertiesDialect)
 114   static constexpr StringLiteral getDialectNamespace() {
 115     return StringLiteral("test_op_properties");
 116   }
 117   explicit TestOpPropertiesDialect(MLIRContext *context)
 118       : Dialect(getDialectNamespace(), context,
 119                 TypeID::get<TestOpPropertiesDialect>()) {
 120     addOperations<OpWithoutProperties>();
 121   }
 122 };
 123 } // namespace
 124
 125 constexpr StringLiteral withoutPropertiesAttrsSrc = R"mlir(
 126     "test_op_properties.op_without_properties"()
 127       {inherent_attr = 42, other_attr = 56} : () -> ()
 128 )mlir";
 129
 130 TEST(Bytecode, OpWithoutProperties) {
 131   MLIRContext context;
 132   context.getOrLoadDialect<TestOpPropertiesDialect>();
 133   ParserConfig config(&context);
 134   OwningOpRef<Operation *> op =
 135       parseSourceString(withoutPropertiesAttrsSrc, config);
 136
 137   std::string bytecode;
 138   llvm::raw_string_ostream os(bytecode);
 139   ASSERT_TRUE(succeeded(writeBytecodeToFile(op.get(), os)));
 140   std::unique_ptr<Block> block = std::make_unique<Block>();
 141   ASSERT_TRUE(succeeded(readBytecodeFile(
 142       llvm::MemoryBufferRef(os.str(), "string-buffer"), block.get(), config)));
 143   Operation *roundtripped = &block->front();
 144   EXPECT_EQ(roundtripped->getAttrs().size(), 2u);
 145   EXPECT_TRUE(roundtripped->getInherentAttr("inherent_attr") != std::nullopt);
 146   EXPECT_TRUE(roundtripped->getDiscardableAttr("other_attr") != Attribute());
 147
 148   EXPECT_TRUE(OperationEquivalence::computeHash(op.get()) ==
 149               OperationEquivalence::computeHash(roundtripped));
 150 }