[Mlir-commits] [mlir] [MLIR][Bytecode] Followup 8106c81 (PR #157136)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Nikhil Kalra (nikalra)

<details>
<summary>Changes</summary>

Addressed code review feedback:
- Fixed some issues in the unit test
- Adjusted line wrapping in the docs
- Clarified comments in the bytecode reader

---
Full diff: https://github.com/llvm/llvm-project/pull/157136.diff


3 Files Affected:

- (modified) mlir/docs/BytecodeFormat.md (+3-1) 
- (modified) mlir/lib/Bytecode/Reader/BytecodeReader.cpp (+13-4) 
- (modified) mlir/unittests/Bytecode/BytecodeTest.cpp (+12-13) 


``````````diff
diff --git a/mlir/docs/BytecodeFormat.md b/mlir/docs/BytecodeFormat.md
index 9846df8726295..f50ddeb25cc33 100644
--- a/mlir/docs/BytecodeFormat.md
+++ b/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
 
diff --git a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
index d29053a2b6e65..287d926d66105 100644
--- a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
+++ b/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"
 
@@ -300,8 +298,19 @@ class EncodingReader {
       // alignment of the root buffer. If it is not, we cannot safely guarantee
       // that the specified alignment is globally correct.
       //
-      // 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.
+      // E.g. if the buffer is 8k aligned and the section is marked to be 16k
+      // aligned:
+      // - (a) the alignTo call early returns when the pointer is 16k
+      // aligned but given the original 8k alignment we could offset into the
+      // padding by ~8k giving us 16k pointer alignment leaving another ~8k of
+      // padding in the bytecode file that will inadvertently be read when we
+      // attempt to parse the next section.
+      // - (b) we update alignTo to align relative to the start of the buffer,
+      // but given an 8k aligned buffer and section alignment of 16k, we could
+      // end up with a pointer that is 24k aligned (8k start alignment + 16k
+      // offset) instead of globally 16k aligned (versus 16k start alignment +
+      // 16k offset). This would result in incorrectly stated alignment for
+      // resources that reference data inside of the bytecode buffer.
       if (failed(alignmentValidator(alignment)))
         return emitError("failed to align section ID: ", unsigned(sectionID));
 
diff --git a/mlir/unittests/Bytecode/BytecodeTest.cpp b/mlir/unittests/Bytecode/BytecodeTest.cpp
index 9ea6560f712a1..d7b442f6832d0 100644
--- a/mlir/unittests/Bytecode/BytecodeTest.cpp
+++ b/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.

``````````

</details>


https://github.com/llvm/llvm-project/pull/157136