[MLIR][Bytecode] Followup 8106c81

mlir/docs/BytecodeFormat.md

@@ -125,7 +125,9 @@ lazy-loading, and more. Each section contains a Section ID, whose high bit
 indicates if the section has alignment requirements, a length (which allows for
 skipping over the section), and an optional alignment. When an alignment is
 present, a variable number of padding bytes (0xCB) may appear before the section
-data. The alignment of a section must be a power of 2. The input bytecode buffer must satisfy the same alignment requirements as those of every section.
+data. The alignment of a section must be a power of 2. 
+The input bytecode buffer must satisfy the same alignment requirements as 
+those of every section.
 
 ## MLIR Encoding
 

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

@@ -22,8 +22,6 @@
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Endian.h"
-#include "llvm/Support/Format.h"
-#include "llvm/Support/LogicalResult.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/Support/SourceMgr.h"
 
@@ -296,12 +294,38 @@ class EncodingReader {
       if (failed(parseVarInt(alignment)))
         return failure();
 
-      // Check that the requested alignment is less than or equal to the
-      // alignment of the root buffer. If it is not, we cannot safely guarantee
-      // that the specified alignment is globally correct.
+      // Check that the requested alignment must not exceed the alignment of
+      // the root buffer itself. Otherwise we cannot guarantee that pointers
+      // derived from this buffer will actually satisfy the requested alignment
+      // globally.
       //
-      // E.g. if the buffer is 8k aligned and the section is 16k aligned,
-      // we could end up at an offset of 24k, which is not globally 16k aligned.
+      // Consider a bytecode buffer that is guaranteed to be 8k aligned, but not
+      // 16k aligned (e.g. absolute address 40960. If a section inside this
+      // buffer declares a 16k alignment requirement, two problems can arise:
+      //
+      //   (a) If we "align forward" the current pointer to the next
+      //       16k boundary, the amount of padding we skip depends on the
+      //       buffer's starting address. For example:
+      //
+      //         buffer_start = 40960
+      //         next 16k boundary = 49152
+      //         bytes skipped = 49152 - 40960 = 8192
+      //
+      //       This leaves behind variable padding that could be misinterpreted
+      //       as part of the next section.
+      //
+      //   (b) If we align relative to the buffer start, we may
+      //       obtain addresses that are multiples of "buffer_start +
+      //       section_alignment" rather than truly globally aligned
+      //       addresses. For example:
+      //
+      //         buffer_start = 40960 (5×8k, 8k aligned but not 16k)
+      //         offset       = 16384  (first multiple of 16k)
+      //         section_ptr  = 40960 + 16384 = 57344
+      //
+      //       57344 is 8k aligned but not 16k aligned.
+      //       Any consumer expecting true 16k alignment would see this as a
+      //       violation.
       if (failed(alignmentValidator(alignment)))
         return emitError("failed to align section ID: ", unsigned(sectionID));
 

mlir/unittests/Bytecode/BytecodeTest.cpp

@@ -72,6 +72,8 @@ TEST(Bytecode, MultiModuleWithResource) {
   ASSERT_TRUE(module);
 
   // Write the module to bytecode.
+  // Ensure that reserveExtraSpace is called with the size needed to write the
+  // bytecode buffer.
   MockOstream ostream;
   EXPECT_CALL(ostream, reserveExtraSpace).WillOnce([&](uint64_t space) {
     ostream.buffer = std::make_unique<std::byte[]>(space);
@@ -128,31 +130,28 @@ TEST(Bytecode, AlignmentFailure) {
   ASSERT_TRUE(module);
 
   // Write the module to bytecode.
-  MockOstream ostream;
-  EXPECT_CALL(ostream, reserveExtraSpace).WillOnce([&](uint64_t space) {
-    ostream.buffer = std::make_unique<std::byte[]>(space);
-    ostream.size = space;
-  });
+  std::string serializedBytecode;
+  llvm::raw_string_ostream ostream(serializedBytecode);
   ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), ostream)));
 
   // Create copy of buffer which is not aligned to requested resource alignment.
-  std::string buffer((char *)ostream.buffer.get(),
-                     (char *)ostream.buffer.get() + ostream.size);
+  std::string buffer(serializedBytecode);
   size_t bufferSize = buffer.size();
 
-  // Increment into the buffer until we get to a power of 2 alignment that is
-  // not 32 bit aligned.
+  // Increment into the buffer until we get to an address that is 2 byte aligned
+  // but not 32 byte aligned.
   size_t pad = 0;
   while (true) {
-    if (llvm::isAddrAligned(Align(2), &buffer[pad]) &&
-        !llvm::isAddrAligned(Align(32), &buffer[pad]))
+    if (llvm::isAddrAligned(Align(2), buffer.data() + pad) &&
+        !llvm::isAddrAligned(Align(32), buffer.data() + pad))
       break;
 
     pad++;
-    buffer.reserve(bufferSize + pad);
+    // Pad the beginning of the buffer to push the start point to an unaligned
+    // value.
+    buffer.insert(0, 1, ' ');
   }
 
-  buffer.insert(0, pad, ' ');
   StringRef alignedBuffer(buffer.data() + pad, bufferSize);
 
   // Attach a diagnostic handler to get the error message.