Skip to content

Commit 2b5134f

Browse files
author
Mogball
committed
[mlir] Fix bytecode reading of resource sections
This partially reverts #66380. The assertion that the underlying buffer of an EncodingReader is aligned to any required alignments for resource sections. Resources know their own alignment and pad their buffers accordingly, but the bytecode reader doesn't know that ahead of time. Consequently, it cannot give the resource EncodingReader a base buffer aligned to the maximum required alignment. A simple example from the test fails without this: ```mlir module @TestDialectResources attributes { bytecode.test = dense_resource<resource> : tensor<4xi32> } {} {-# dialect_resources: { builtin: { resource: "0x2000000001000000020000000300000004000000", resource_2: "0x2000000001000000020000000300000004000000" } } ```
1 parent 2048836 commit 2b5134f

File tree

2 files changed

+4
-45
lines changed

2 files changed

+4
-45
lines changed

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -111,13 +111,6 @@ class EncodingReader {
111111
return ((uintptr_t)ptr & (alignment - 1)) != 0;
112112
};
113113

114-
// Ensure the data buffer was sufficiently aligned in the first place.
115-
if (LLVM_UNLIKELY(isUnaligned(buffer.begin()))) {
116-
return emitError("expected bytecode buffer to be aligned to ", alignment,
117-
", but got pointer: '0x" +
118-
llvm::utohexstr((uintptr_t)buffer.begin()) + "'");
119-
}
120-
121114
// Shift the reader position to the next alignment boundary.
122115
while (isUnaligned(dataIt)) {
123116
uint8_t padding;
@@ -319,7 +312,8 @@ class EncodingReader {
319312

320313
// Parse in the remaining bytes of the value.
321314
llvm::support::ulittle64_t resultLE(result);
322-
if (failed(parseBytes(numBytes, reinterpret_cast<uint8_t *>(&resultLE) + 1)))
315+
if (failed(
316+
parseBytes(numBytes, reinterpret_cast<uint8_t *>(&resultLE) + 1)))
323317
return failure();
324318

325319
// Shift out the low-order bits that were used to mark how the value was

mlir/unittests/Bytecode/BytecodeTest.cpp

Lines changed: 2 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,8 @@ module @TestDialectResources attributes {
3030
{-#
3131
dialect_resources: {
3232
builtin: {
33-
resource: "0x2000000001000000020000000300000004000000"
33+
resource: "0x2000000001000000020000000300000004000000",
34+
resource_2: "0x2000000001000000020000000300000004000000"
3435
}
3536
}
3637
#-}
@@ -88,39 +89,3 @@ TEST(Bytecode, MultiModuleWithResource) {
8889
checkResourceAttribute(*module);
8990
checkResourceAttribute(*roundTripModule);
9091
}
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-
}

0 commit comments

Comments
 (0)