xref: /llvm-project/mlir/unittests/Bytecode/BytecodeTest.cpp (revision 3b34c117db175ae6a1404f7c07857c4aa6fc1ae3)
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/BytecodeWriter.h"
10 #include "mlir/IR/AsmState.h"
11 #include "mlir/IR/BuiltinAttributes.h"
12 #include "mlir/IR/OpImplementation.h"
13 #include "mlir/IR/OwningOpRef.h"
14 #include "mlir/Parser/Parser.h"
15 
16 #include "llvm/ADT/StringRef.h"
17 #include "llvm/Support/Endian.h"
18 #include "gmock/gmock.h"
19 #include "gtest/gtest.h"
20 
21 using namespace llvm;
22 using namespace mlir;
23 
24 StringLiteral IRWithResources = R"(
25 module @TestDialectResources attributes {
26   bytecode.test = dense_resource<resource> : tensor<4xi32>
27 } {}
28 {-#
29   dialect_resources: {
30     builtin: {
31       resource: "0x2000000001000000020000000300000004000000",
32       resource_2: "0x2000000001000000020000000300000004000000"
33     }
34   }
35 #-}
36 )";
37 
38 TEST(Bytecode, MultiModuleWithResource) {
39   MLIRContext context;
40   Builder builder(&context);
41   ParserConfig parseConfig(&context);
42   OwningOpRef<Operation *> module =
43       parseSourceString<Operation *>(IRWithResources, parseConfig);
44   ASSERT_TRUE(module);
45 
46   // Write the module to bytecode
47   std::string buffer;
48   llvm::raw_string_ostream ostream(buffer);
49   ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), ostream)));
50   ostream.flush();
51 
52   // Create copy of buffer which is aligned to requested resource alignment.
53   constexpr size_t kAlignment = 0x20;
54   size_t buffer_size = buffer.size();
55   buffer.reserve(buffer_size + kAlignment - 1);
56   size_t pad = ~(uintptr_t)buffer.data() + 1 & kAlignment - 1;
57   buffer.insert(0, pad, ' ');
58   StringRef aligned_buffer(buffer.data() + pad, buffer_size);
59 
60   // Parse it back
61   OwningOpRef<Operation *> roundTripModule =
62       parseSourceString<Operation *>(aligned_buffer, parseConfig);
63   ASSERT_TRUE(roundTripModule);
64 
65   // FIXME: Parsing external resources does not work on big-endian
66   // platforms currently.
67   if (llvm::support::endian::system_endianness() ==
68       llvm::support::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