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 using ::testing::StartsWith; 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::support::endian::system_endianness() == 70 llvm::support::endianness::big) 71 GTEST_SKIP(); 72 73 // Try to see if we have a valid resource in the parsed module. 74 auto checkResourceAttribute = [&](Operation *op) { 75 Attribute attr = roundTripModule->getDiscardableAttr("bytecode.test"); 76 ASSERT_TRUE(attr); 77 auto denseResourceAttr = dyn_cast<DenseI32ResourceElementsAttr>(attr); 78 ASSERT_TRUE(denseResourceAttr); 79 std::optional<ArrayRef<int32_t>> attrData = 80 denseResourceAttr.tryGetAsArrayRef(); 81 ASSERT_TRUE(attrData.has_value()); 82 ASSERT_EQ(attrData->size(), static_cast<size_t>(4)); 83 EXPECT_EQ((*attrData)[0], 1); 84 EXPECT_EQ((*attrData)[1], 2); 85 EXPECT_EQ((*attrData)[2], 3); 86 EXPECT_EQ((*attrData)[3], 4); 87 }; 88 89 checkResourceAttribute(*module); 90 checkResourceAttribute(*roundTripModule); 91 } 92