[Mlir-commits] [mlir] 2b5134f - [mlir] Fix bytecode reading of resource sections

[llvmlistbot at llvm.org]

Author: Mogball
Date: 2023-09-29T18:39:56-07:00
New Revision: 2b5134f1b7f4d2ded14f138c84939d2037326c51

URL: https://github.com/llvm/llvm-project/commit/2b5134f1b7f4d2ded14f138c84939d2037326c51
DIFF: https://github.com/llvm/llvm-project/commit/2b5134f1b7f4d2ded14f138c84939d2037326c51.diff

LOG: [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"
    }
  }
```

Added: 
    

Modified: 
    mlir/lib/Bytecode/Reader/BytecodeReader.cpp
    mlir/unittests/Bytecode/BytecodeTest.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
index 98a080d23a1dce5..0bc2a2f676ba4bc 100644
--- a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
+++ b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
@@ -111,13 +111,6 @@ class EncodingReader {
       return ((uintptr_t)ptr & (alignment - 1)) != 0;
     };
 
-    // Ensure the data buffer was sufficiently aligned in the first place.
-    if (LLVM_UNLIKELY(isUnaligned(buffer.begin()))) {
-      return emitError("expected bytecode buffer to be aligned to ", alignment,
-                       ", but got pointer: '0x" +
-                           llvm::utohexstr((uintptr_t)buffer.begin()) + "'");
-    }
-
     // Shift the reader position to the next alignment boundary.
     while (isUnaligned(dataIt)) {
       uint8_t padding;
@@ -319,7 +312,8 @@ class EncodingReader {
 
     // Parse in the remaining bytes of the value.
     llvm::support::ulittle64_t resultLE(result);
-    if (failed(parseBytes(numBytes, reinterpret_cast<uint8_t *>(&resultLE) + 1)))
+    if (failed(
+            parseBytes(numBytes, reinterpret_cast<uint8_t *>(&resultLE) + 1)))
       return failure();
 
     // Shift out the low-order bits that were used to mark how the value was

diff  --git a/mlir/unittests/Bytecode/BytecodeTest.cpp b/mlir/unittests/Bytecode/BytecodeTest.cpp
index b5f8c09b617bdd2..4ec90101f48c856 100644
--- a/mlir/unittests/Bytecode/BytecodeTest.cpp
+++ b/mlir/unittests/Bytecode/BytecodeTest.cpp
@@ -30,7 +30,8 @@ module @TestDialectResources attributes {
 {-#
   dialect_resources: {
     builtin: {
-      resource: "0x2000000001000000020000000300000004000000"
+      resource: "0x2000000001000000020000000300000004000000",
+      resource_2: "0x2000000001000000020000000300000004000000"
     }
   }
 #-}
@@ -88,39 +89,3 @@ TEST(Bytecode, MultiModuleWithResource) {
   checkResourceAttribute(*module);
   checkResourceAttribute(*roundTripModule);
 }
-
-TEST(Bytecode, InsufficientAlignmentFailure) {
-  MLIRContext context;
-  Builder builder(&context);
-  ParserConfig parseConfig(&context);
-  OwningOpRef<Operation *> module =
-      parseSourceString<Operation *>(IRWithResources, parseConfig);
-  ASSERT_TRUE(module);
-
-  // Write the module to bytecode
-  std::string buffer;
-  llvm::raw_string_ostream ostream(buffer);
-  ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), ostream)));
-  ostream.flush();
-
-  // Create copy of buffer which is insufficiently aligned.
-  constexpr size_t kAlignment = 0x20;
-  size_t buffer_size = buffer.size();
-  buffer.reserve(buffer_size + kAlignment - 1);
-  size_t pad = ~(uintptr_t)buffer.data() + kAlignment / 2 + 1 & kAlignment - 1;
-  buffer.insert(0, pad, ' ');
-  StringRef misaligned_buffer(buffer.data() + pad, buffer_size);
-
-  std::unique_ptr<Diagnostic> diagnostic;
-  context.getDiagEngine().registerHandler([&](Diagnostic &diag) {
-    diagnostic = std::make_unique<Diagnostic>(std::move(diag));
-  });
-
-  // Try to parse it back and check for alignment error.
-  OwningOpRef<Operation *> roundTripModule =
-      parseSourceString<Operation *>(misaligned_buffer, parseConfig);
-  EXPECT_FALSE(roundTripModule);
-  ASSERT_TRUE(diagnostic);
-  EXPECT_THAT(diagnostic->str(),
-              StartsWith("expected bytecode buffer to be aligned to 32"));
-}