231 lines · cpp
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-exception6//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/Diagnostics.h"14#include "mlir/IR/OpImplementation.h"15#include "mlir/IR/OwningOpRef.h"16#include "mlir/Parser/Parser.h"17 18#include "llvm/ADT/StringRef.h"19#include "llvm/Support/Alignment.h"20#include "llvm/Support/Endian.h"21#include "llvm/Support/MemoryBufferRef.h"22#include "llvm/Support/raw_ostream.h"23#include "gmock/gmock.h"24#include "gtest/gtest.h"25 26using namespace llvm;27using namespace mlir;28 29StringLiteral irWithResources = R"(30module @TestDialectResources attributes {31 bytecode.test = dense_resource<resource> : tensor<4xi32>32} {}33{-#34 dialect_resources: {35 builtin: {36 resource: "0x2000000001000000020000000300000004000000",37 resource_2: "0x2000000001000000020000000300000004000000"38 }39 }40#-}41)";42 43struct MockOstream final : public raw_ostream {44 std::unique_ptr<std::byte[]> buffer;45 size_t size = 0;46 47 MOCK_METHOD(void, reserveExtraSpace, (uint64_t extraSpace), (override));48 49 MockOstream() : raw_ostream(true) {}50 uint64_t current_pos() const override { return pos; }51 52private:53 size_t pos = 0;54 55 void write_impl(const char *ptr, size_t length) override {56 if (pos + length <= size) {57 memcpy((void *)(buffer.get() + pos), ptr, length);58 pos += length;59 } else {60 report_fatal_error(61 "Attempted to write past the end of the fixed size buffer.");62 }63 }64};65 66TEST(Bytecode, MultiModuleWithResource) {67 MLIRContext context;68 Builder builder(&context);69 ParserConfig parseConfig(&context);70 OwningOpRef<Operation *> module =71 parseSourceString<Operation *>(irWithResources, parseConfig);72 ASSERT_TRUE(module);73 74 // Write the module to bytecode.75 // Ensure that reserveExtraSpace is called with the size needed to write the76 // bytecode buffer.77 MockOstream ostream;78 EXPECT_CALL(ostream, reserveExtraSpace).WillOnce([&](uint64_t space) {79 ostream.buffer = std::make_unique<std::byte[]>(space);80 ostream.size = space;81 });82 ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), ostream)));83 84 // Create copy of buffer which is aligned to requested resource alignment.85 std::string buffer((char *)ostream.buffer.get(),86 (char *)ostream.buffer.get() + ostream.size);87 constexpr size_t kAlignment = 0x20;88 size_t bufferSize = buffer.size();89 buffer.reserve(bufferSize + kAlignment - 1);90 size_t pad = (~(uintptr_t)buffer.data() + 1) & (kAlignment - 1);91 buffer.insert(0, pad, ' ');92 StringRef alignedBuffer(buffer.data() + pad, bufferSize);93 94 // Parse it back95 OwningOpRef<Operation *> roundTripModule =96 parseSourceString<Operation *>(alignedBuffer, parseConfig);97 ASSERT_TRUE(roundTripModule);98 99 // FIXME: Parsing external resources does not work on big-endian100 // platforms currently.101 if (llvm::endianness::native == llvm::endianness::big)102 GTEST_SKIP();103 104 // Try to see if we have a valid resource in the parsed module.105 auto checkResourceAttribute = [](Operation *parsedModule) {106 Attribute attr = parsedModule->getDiscardableAttr("bytecode.test");107 ASSERT_TRUE(attr);108 auto denseResourceAttr = dyn_cast<DenseI32ResourceElementsAttr>(attr);109 ASSERT_TRUE(denseResourceAttr);110 std::optional<ArrayRef<int32_t>> attrData =111 denseResourceAttr.tryGetAsArrayRef();112 ASSERT_TRUE(attrData.has_value());113 ASSERT_EQ(attrData->size(), static_cast<size_t>(4));114 EXPECT_EQ((*attrData)[0], 1);115 EXPECT_EQ((*attrData)[1], 2);116 EXPECT_EQ((*attrData)[2], 3);117 EXPECT_EQ((*attrData)[3], 4);118 };119 120 checkResourceAttribute(*module);121 checkResourceAttribute(*roundTripModule);122}123 124TEST(Bytecode, AlignmentFailure) {125 MLIRContext context;126 Builder builder(&context);127 ParserConfig parseConfig(&context);128 OwningOpRef<Operation *> module =129 parseSourceString<Operation *>(irWithResources, parseConfig);130 ASSERT_TRUE(module);131 132 // Write the module to bytecode.133 std::string serializedBytecode;134 llvm::raw_string_ostream ostream(serializedBytecode);135 ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), ostream)));136 137 // Create copy of buffer which is not aligned to requested resource alignment.138 std::string buffer(serializedBytecode);139 size_t bufferSize = buffer.size();140 141 // Increment into the buffer until we get to an address that is 2 byte aligned142 // but not 32 byte aligned.143 size_t pad = 0;144 while (true) {145 if (llvm::isAddrAligned(Align(2), buffer.data() + pad) &&146 !llvm::isAddrAligned(Align(32), buffer.data() + pad))147 break;148 149 pad++;150 // Pad the beginning of the buffer to push the start point to an unaligned151 // value.152 buffer.insert(0, 1, ' ');153 }154 155 StringRef alignedBuffer(buffer.data() + pad, bufferSize);156 157 // Attach a diagnostic handler to get the error message.158 llvm::SmallVector<std::string> msg;159 ScopedDiagnosticHandler handler(160 &context, [&msg](Diagnostic &diag) { msg.push_back(diag.str()); });161 162 // Parse it back163 OwningOpRef<Operation *> roundTripModule =164 parseSourceString<Operation *>(alignedBuffer, parseConfig);165 ASSERT_FALSE(roundTripModule);166 ASSERT_THAT(msg[0].data(), ::testing::StartsWith(167 "expected section alignment 32 but bytecode "168 "buffer"));169 ASSERT_STREQ(msg[1].data(), "failed to align section ID: 5");170}171 172namespace {173/// A custom operation for the purpose of showcasing how discardable attributes174/// are handled in absence of properties.175class OpWithoutProperties : public Op<OpWithoutProperties> {176public:177 // Begin boilerplate.178 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpWithoutProperties)179 using Op::Op;180 static ArrayRef<StringRef> getAttributeNames() {181 static StringRef attributeNames[] = {StringRef("inherent_attr")};182 return ArrayRef(attributeNames);183 };184 static StringRef getOperationName() {185 return "test_op_properties.op_without_properties";186 }187 // End boilerplate.188};189 190// A trivial supporting dialect to register the above operation.191class TestOpPropertiesDialect : public Dialect {192public:193 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestOpPropertiesDialect)194 static constexpr StringLiteral getDialectNamespace() {195 return StringLiteral("test_op_properties");196 }197 explicit TestOpPropertiesDialect(MLIRContext *context)198 : Dialect(getDialectNamespace(), context,199 TypeID::get<TestOpPropertiesDialect>()) {200 addOperations<OpWithoutProperties>();201 }202};203} // namespace204 205constexpr StringLiteral withoutPropertiesAttrsSrc = R"mlir(206 "test_op_properties.op_without_properties"()207 {inherent_attr = 42, other_attr = 56} : () -> ()208)mlir";209 210TEST(Bytecode, OpWithoutProperties) {211 MLIRContext context;212 context.getOrLoadDialect<TestOpPropertiesDialect>();213 ParserConfig config(&context);214 OwningOpRef<Operation *> op =215 parseSourceString(withoutPropertiesAttrsSrc, config);216 217 std::string bytecode;218 llvm::raw_string_ostream os(bytecode);219 ASSERT_TRUE(succeeded(writeBytecodeToFile(op.get(), os)));220 std::unique_ptr<Block> block = std::make_unique<Block>();221 ASSERT_TRUE(succeeded(readBytecodeFile(222 llvm::MemoryBufferRef(bytecode, "string-buffer"), block.get(), config)));223 Operation *roundtripped = &block->front();224 EXPECT_EQ(roundtripped->getAttrs().size(), 2u);225 EXPECT_TRUE(roundtripped->getInherentAttr("inherent_attr") != std::nullopt);226 EXPECT_TRUE(roundtripped->getDiscardableAttr("other_attr") != Attribute());227 228 EXPECT_TRUE(OperationEquivalence::computeHash(op.get()) ==229 OperationEquivalence::computeHash(roundtripped));230}231