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 } 35 } 36 #-} 37 )"; 38 39 TEST(Bytecode, MultiModuleWithResource) { 40 MLIRContext context; 41 Builder builder(&context); 42 ParserConfig parseConfig(&context); 43 OwningOpRef<Operation *> module = 44 parseSourceString<Operation *>(IRWithResources, parseConfig); 45 ASSERT_TRUE(module); 46 47 // Write the module to bytecode 48 std::string buffer; 49 llvm::raw_string_ostream ostream(buffer); 50 ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), ostream))); 51 ostream.flush(); 52 53 // Create copy of buffer which is aligned to requested resource alignment. 54 constexpr size_t kAlignment = 0x20; 55 size_t buffer_size = buffer.size(); 56 buffer.reserve(buffer_size + kAlignment - 1); 57 size_t pad = ~(uintptr_t)buffer.data() + 1 & kAlignment - 1; 58 buffer.insert(0, pad, ' '); 59 StringRef aligned_buffer(buffer.data() + pad, buffer_size); 60 61 // Parse it back 62 OwningOpRef<Operation *> roundTripModule = 63 parseSourceString<Operation *>(aligned_buffer, parseConfig); 64 ASSERT_TRUE(roundTripModule); 65 66 // FIXME: Parsing external resources does not work on big-endian 67 // platforms currently. 68 if (llvm::support::endian::system_endianness() == 69 llvm::support::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 TEST(Bytecode, InsufficientAlignmentFailure) { 93 MLIRContext context; 94 Builder builder(&context); 95 ParserConfig parseConfig(&context); 96 OwningOpRef<Operation *> module = 97 parseSourceString<Operation *>(IRWithResources, parseConfig); 98 ASSERT_TRUE(module); 99 100 // Write the module to bytecode 101 std::string buffer; 102 llvm::raw_string_ostream ostream(buffer); 103 ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), ostream))); 104 ostream.flush(); 105 106 // Create copy of buffer which is insufficiently aligned. 107 constexpr size_t kAlignment = 0x20; 108 size_t buffer_size = buffer.size(); 109 buffer.reserve(buffer_size + kAlignment - 1); 110 size_t pad = ~(uintptr_t)buffer.data() + kAlignment / 2 + 1 & kAlignment - 1; 111 buffer.insert(0, pad, ' '); 112 StringRef misaligned_buffer(buffer.data() + pad, buffer_size); 113 114 std::unique_ptr<Diagnostic> diagnostic; 115 context.getDiagEngine().registerHandler([&](Diagnostic &diag) { 116 diagnostic = std::make_unique<Diagnostic>(std::move(diag)); 117 }); 118 119 // Try to parse it back and check for alignment error. 120 OwningOpRef<Operation *> roundTripModule = 121 parseSourceString<Operation *>(misaligned_buffer, parseConfig); 122 EXPECT_FALSE(roundTripModule); 123 ASSERT_TRUE(diagnostic); 124 EXPECT_THAT(diagnostic->str(), 125 StartsWith("expected bytecode buffer to be aligned to 32")); 126 } 127