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 53 // Create copy of buffer which is aligned to requested resource alignment. 54 constexpr size_t kAlignment = 0x20; 55 size_t bufferSize = buffer.size(); 56 buffer.reserve(bufferSize + kAlignment - 1); 57 size_t pad = (~(uintptr_t)buffer.data() + 1) & (kAlignment - 1); 58 buffer.insert(0, pad, ' '); 59 StringRef alignedBuffer(buffer.data() + pad, bufferSize); 60 61 // Parse it back 62 OwningOpRef<Operation *> roundTripModule = 63 parseSourceString<Operation *>(alignedBuffer, parseConfig); 64 ASSERT_TRUE(roundTripModule); 65 66 // FIXME: Parsing external resources does not work on big-endian 67 // platforms currently. 68 if (llvm::endianness::native == llvm::endianness::big) 69 GTEST_SKIP(); 70 71 // Try to see if we have a valid resource in the parsed module. 72 auto checkResourceAttribute = [](Operation *parsedModule) { 73 Attribute attr = parsedModule->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 91 namespace { 92 /// A custom operation for the purpose of showcasing how discardable attributes 93 /// are handled in absence of properties. 94 class OpWithoutProperties : public Op<OpWithoutProperties> { 95 public: 96 // Begin boilerplate. 97 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpWithoutProperties) 98 using Op::Op; 99 static ArrayRef<StringRef> getAttributeNames() { 100 static StringRef attributeNames[] = {StringRef("inherent_attr")}; 101 return ArrayRef(attributeNames); 102 }; 103 static StringRef getOperationName() { 104 return "test_op_properties.op_without_properties"; 105 } 106 // End boilerplate. 107 }; 108 109 // A trivial supporting dialect to register the above operation. 110 class TestOpPropertiesDialect : public Dialect { 111 public: 112 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestOpPropertiesDialect) 113 static constexpr StringLiteral getDialectNamespace() { 114 return StringLiteral("test_op_properties"); 115 } 116 explicit TestOpPropertiesDialect(MLIRContext *context) 117 : Dialect(getDialectNamespace(), context, 118 TypeID::get<TestOpPropertiesDialect>()) { 119 addOperations<OpWithoutProperties>(); 120 } 121 }; 122 } // namespace 123 124 constexpr StringLiteral withoutPropertiesAttrsSrc = R"mlir( 125 "test_op_properties.op_without_properties"() 126 {inherent_attr = 42, other_attr = 56} : () -> () 127 )mlir"; 128 129 TEST(Bytecode, OpWithoutProperties) { 130 MLIRContext context; 131 context.getOrLoadDialect<TestOpPropertiesDialect>(); 132 ParserConfig config(&context); 133 OwningOpRef<Operation *> op = 134 parseSourceString(withoutPropertiesAttrsSrc, config); 135 136 std::string bytecode; 137 llvm::raw_string_ostream os(bytecode); 138 ASSERT_TRUE(succeeded(writeBytecodeToFile(op.get(), os))); 139 std::unique_ptr<Block> block = std::make_unique<Block>(); 140 ASSERT_TRUE(succeeded(readBytecodeFile( 141 llvm::MemoryBufferRef(bytecode, "string-buffer"), block.get(), config))); 142 Operation *roundtripped = &block->front(); 143 EXPECT_EQ(roundtripped->getAttrs().size(), 2u); 144 EXPECT_TRUE(roundtripped->getInherentAttr("inherent_attr") != std::nullopt); 145 EXPECT_TRUE(roundtripped->getDiscardableAttr("other_attr") != Attribute()); 146 147 EXPECT_TRUE(OperationEquivalence::computeHash(op.get()) == 148 OperationEquivalence::computeHash(roundtripped)); 149 } 150