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
  53   // Create copy of buffer which is aligned to requested resource alignment.
  54   constexpr size_t kAlignment = 0x20;
  55   size_t bufferSize = buffer.size();
  56   buffer.reserve(bufferSize + kAlignment - 1);
  57   size_t pad = (~(uintptr_t)buffer.data() + 1) & (kAlignment - 1);
  58   buffer.insert(0, pad, ' ');
  59   StringRef alignedBuffer(buffer.data() + pad, bufferSize);
  60
  61   // Parse it back
  62   OwningOpRef<Operation *> roundTripModule =
  63       parseSourceString<Operation *>(alignedBuffer, parseConfig);
  64   ASSERT_TRUE(roundTripModule);
  65
  66   // FIXME: Parsing external resources does not work on big-endian
  67   // platforms currently.
  68   if (llvm::endianness::native == llvm::endianness::big)
  69     GTEST_SKIP();
  70
  71   // Try to see if we have a valid resource in the parsed module.
  72   auto checkResourceAttribute = [&](Operation *op) {
  73     Attribute attr = roundTripModule->getDiscardableAttr("bytecode.test");
  74     ASSERT_TRUE(attr);
  75     auto denseResourceAttr = dyn_cast<DenseI32ResourceElementsAttr>(attr);
  76     ASSERT_TRUE(denseResourceAttr);
  77     std::optional<ArrayRef<int32_t>> attrData =
  78         denseResourceAttr.tryGetAsArrayRef();
  79     ASSERT_TRUE(attrData.has_value());
  80     ASSERT_EQ(attrData->size(), static_cast<size_t>(4));
  81     EXPECT_EQ((*attrData)[0], 1);
  82     EXPECT_EQ((*attrData)[1], 2);
  83     EXPECT_EQ((*attrData)[2], 3);
  84     EXPECT_EQ((*attrData)[3], 4);
  85   };
  86
  87   checkResourceAttribute(*module);
  88   checkResourceAttribute(*roundTripModule);
  89 }
  90
  91 namespace {
  92 /// A custom operation for the purpose of showcasing how discardable attributes
  93 /// are handled in absence of properties.
  94 class OpWithoutProperties : public Op<OpWithoutProperties> {
  95 public:
  96   // Begin boilerplate.
  97   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpWithoutProperties)
  98   using Op::Op;
  99   static ArrayRef<StringRef> getAttributeNames() {
 100     static StringRef attributeNames[] = {StringRef("inherent_attr")};
 101     return ArrayRef(attributeNames);
 102   };
 103   static StringRef getOperationName() {
 104     return "test_op_properties.op_without_properties";
 105   }
 106   // End boilerplate.
 107 };
 108
 109 // A trivial supporting dialect to register the above operation.
 110 class TestOpPropertiesDialect : public Dialect {
 111 public:
 112   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestOpPropertiesDialect)
 113   static constexpr StringLiteral getDialectNamespace() {
 114     return StringLiteral("test_op_properties");
 115   }
 116   explicit TestOpPropertiesDialect(MLIRContext *context)
 117       : Dialect(getDialectNamespace(), context,
 118                 TypeID::get<TestOpPropertiesDialect>()) {
 119     addOperations<OpWithoutProperties>();
 120   }
 121 };
 122 } // namespace
 123
 124 constexpr StringLiteral withoutPropertiesAttrsSrc = R"mlir(
 125     "test_op_properties.op_without_properties"()
 126       {inherent_attr = 42, other_attr = 56} : () -> ()
 127 )mlir";
 128
 129 TEST(Bytecode, OpWithoutProperties) {
 130   MLIRContext context;
 131   context.getOrLoadDialect<TestOpPropertiesDialect>();
 132   ParserConfig config(&context);
 133   OwningOpRef<Operation *> op =
 134       parseSourceString(withoutPropertiesAttrsSrc, config);
 135
 136   std::string bytecode;
 137   llvm::raw_string_ostream os(bytecode);
 138   ASSERT_TRUE(succeeded(writeBytecodeToFile(op.get(), os)));
 139   std::unique_ptr<Block> block = std::make_unique<Block>();
 140   ASSERT_TRUE(succeeded(readBytecodeFile(
 141       llvm::MemoryBufferRef(bytecode, "string-buffer"), block.get(), config)));
 142   Operation *roundtripped = &block->front();
 143   EXPECT_EQ(roundtripped->getAttrs().size(), 2u);
 144   EXPECT_TRUE(roundtripped->getInherentAttr("inherent_attr") != std::nullopt);
 145   EXPECT_TRUE(roundtripped->getDiscardableAttr("other_attr") != Attribute());
 146
 147   EXPECT_TRUE(OperationEquivalence::computeHash(op.get()) ==
 148               OperationEquivalence::computeHash(roundtripped));
 149 }