xref: /llvm-project/mlir/unittests/Bytecode/BytecodeTest.cpp (revision bb0bbed610d86ba155f9c066c23038f7f34e2dbb)
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 "gmock/gmock.h"
20 #include "gtest/gtest.h"
21 
22 using namespace llvm;
23 using namespace mlir;
24 
25 using testing::ElementsAre;
26 
27 StringLiteral IRWithResources = R"(
28 module @TestDialectResources attributes {
29   bytecode.test = dense_resource<resource> : tensor<4xi32>
30 } {}
31 {-#
32   dialect_resources: {
33     builtin: {
34       resource: "0x1000000001000000020000000300000004000000"
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   // Parse it back
54   OwningOpRef<Operation *> roundTripModule =
55       parseSourceString<Operation *>(ostream.str(), parseConfig);
56   ASSERT_TRUE(roundTripModule);
57 
58   // FIXME: Parsing external resources does not work on big-endian
59   // platforms currently.
60   if (llvm::support::endian::system_endianness() ==
61       llvm::support::endianness::big)
62     GTEST_SKIP();
63 
64   // Try to see if we have a valid resource in the parsed module.
65   auto checkResourceAttribute = [&](Operation *op) {
66     Attribute attr = roundTripModule->getAttr("bytecode.test");
67     ASSERT_TRUE(attr);
68     auto denseResourceAttr = dyn_cast<DenseI32ResourceElementsAttr>(attr);
69     ASSERT_TRUE(denseResourceAttr);
70     std::optional<ArrayRef<int32_t>> attrData =
71         denseResourceAttr.tryGetAsArrayRef();
72     ASSERT_TRUE(attrData.has_value());
73     ASSERT_EQ(attrData->size(), static_cast<size_t>(4));
74     EXPECT_EQ((*attrData)[0], 1);
75     EXPECT_EQ((*attrData)[1], 2);
76     EXPECT_EQ((*attrData)[2], 3);
77     EXPECT_EQ((*attrData)[3], 4);
78   };
79 
80   checkResourceAttribute(*module);
81   checkResourceAttribute(*roundTripModule);
82 }
83