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 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::endianness::native == llvm::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 namespace { 93 /// A custom operation for the purpose of showcasing how discardable attributes 94 /// are handled in absence of properties. 95 class OpWithoutProperties : public Op<OpWithoutProperties> { 96 public: 97 // Begin boilerplate. 98 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpWithoutProperties) 99 using Op::Op; 100 static ArrayRef<StringRef> getAttributeNames() { 101 static StringRef attributeNames[] = {StringRef("inherent_attr")}; 102 return ArrayRef(attributeNames); 103 }; 104 static StringRef getOperationName() { 105 return "test_op_properties.op_without_properties"; 106 } 107 // End boilerplate. 108 }; 109 110 // A trivial supporting dialect to register the above operation. 111 class TestOpPropertiesDialect : public Dialect { 112 public: 113 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestOpPropertiesDialect) 114 static constexpr StringLiteral getDialectNamespace() { 115 return StringLiteral("test_op_properties"); 116 } 117 explicit TestOpPropertiesDialect(MLIRContext *context) 118 : Dialect(getDialectNamespace(), context, 119 TypeID::get<TestOpPropertiesDialect>()) { 120 addOperations<OpWithoutProperties>(); 121 } 122 }; 123 } // namespace 124 125 constexpr StringLiteral withoutPropertiesAttrsSrc = R"mlir( 126 "test_op_properties.op_without_properties"() 127 {inherent_attr = 42, other_attr = 56} : () -> () 128 )mlir"; 129 130 TEST(Bytecode, OpWithoutProperties) { 131 MLIRContext context; 132 context.getOrLoadDialect<TestOpPropertiesDialect>(); 133 ParserConfig config(&context); 134 OwningOpRef<Operation *> op = 135 parseSourceString(withoutPropertiesAttrsSrc, config); 136 137 std::string bytecode; 138 llvm::raw_string_ostream os(bytecode); 139 ASSERT_TRUE(succeeded(writeBytecodeToFile(op.get(), os))); 140 std::unique_ptr<Block> block = std::make_unique<Block>(); 141 ASSERT_TRUE(succeeded(readBytecodeFile( 142 llvm::MemoryBufferRef(os.str(), "string-buffer"), block.get(), config))); 143 Operation *roundtripped = &block->front(); 144 EXPECT_EQ(roundtripped->getAttrs().size(), 2u); 145 EXPECT_TRUE(roundtripped->getInherentAttr("inherent_attr") != std::nullopt); 146 EXPECT_TRUE(roundtripped->getDiscardableAttr("other_attr") != Attribute()); 147 148 EXPECT_TRUE(OperationEquivalence::computeHash(op.get()) == 149 OperationEquivalence::computeHash(roundtripped)); 150 } 151