[mlir][bytecode] Check that bytecode source buffer is sufficiently aligned. (llvm#66380)

Before this change, the `ByteCode` test failed on CentOS 7 with
devtoolset-9, because strings happen to be only 8 byte aligned. In
general though, strings have no alignment guarantee.

Increase resource alignment in test to 32 bytes. 
Adjust test to sufficiently align buffer.
Add test to check error when buffer is insufficiently aligned.

Author: chsigg

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

@@ -11,7 +11,6 @@
 #include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/Bytecode/BytecodeOpInterface.h"
 #include "mlir/Bytecode/Encoding.h"
-#include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/OpImplementation.h"
@@ -20,15 +19,13 @@
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/ScopeExit.h"
-#include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/MemoryBufferRef.h"
-#include "llvm/Support/SaveAndRestore.h"
 #include "llvm/Support/SourceMgr.h"
+
 #include <cstddef>
 #include <list>
 #include <memory>
@@ -93,25 +90,36 @@ namespace {
 class EncodingReader {
 public:
   explicit EncodingReader(ArrayRef<uint8_t> contents, Location fileLoc)
-      : dataIt(contents.data()), dataEnd(contents.end()), fileLoc(fileLoc) {}
+      : buffer(contents), dataIt(buffer.begin()), fileLoc(fileLoc) {}
   explicit EncodingReader(StringRef contents, Location fileLoc)
       : EncodingReader({reinterpret_cast<const uint8_t *>(contents.data()),
                         contents.size()},
                        fileLoc) {}
 
   /// Returns true if the entire section has been read.
-  bool empty() const { return dataIt == dataEnd; }
+  bool empty() const { return dataIt == buffer.end(); }
 
   /// Returns the remaining size of the bytecode.
-  size_t size() const { return dataEnd - dataIt; }
+  size_t size() const { return buffer.end() - dataIt; }
 
   /// Align the current reader position to the specified alignment.
   LogicalResult alignTo(unsigned alignment) {
     if (!llvm::isPowerOf2_32(alignment))
       return emitError("expected alignment to be a power-of-two");
 
+    auto isUnaligned = [&](const uint8_t *ptr) {
+      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 (uintptr_t(dataIt) & (uintptr_t(alignment) - 1)) {
+    while (isUnaligned(dataIt)) {
       uint8_t padding;
       if (failed(parseByte(padding)))
         return failure();
@@ -123,7 +131,7 @@ class EncodingReader {
 
     // Ensure the data iterator is now aligned. This case is unlikely because we
     // *just* went through the effort to align the data iterator.
-    if (LLVM_UNLIKELY(!llvm::isAddrAligned(llvm::Align(alignment), dataIt))) {
+    if (LLVM_UNLIKELY(isUnaligned(dataIt))) {
       return emitError("expected data iterator aligned to ", alignment,
                        ", but got pointer: '0x" +
                            llvm::utohexstr((uintptr_t)dataIt) + "'");
@@ -320,8 +328,11 @@ class EncodingReader {
     return success();
   }
 
-  /// The current data iterator, and an iterator to the end of the buffer.
-  const uint8_t *dataIt, *dataEnd;
+  /// The bytecode buffer.
+  ArrayRef<uint8_t> buffer;
+
+  /// The current iterator within the 'buffer'.
+  const uint8_t *dataIt;
 
   /// A location for the bytecode used to report errors.
   Location fileLoc;

mlir/unittests/Bytecode/BytecodeTest.cpp

@@ -6,7 +6,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Bytecode/BytecodeReader.h"
 #include "mlir/Bytecode/BytecodeWriter.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinAttributes.h"
@@ -22,7 +21,7 @@
 using namespace llvm;
 using namespace mlir;
 
-using testing::ElementsAre;
+using ::testing::StartsWith;
 
 StringLiteral IRWithResources = R"(
 module @TestDialectResources attributes {
@@ -31,7 +30,7 @@ module @TestDialectResources attributes {
 {-#
   dialect_resources: {
     builtin: {
-      resource: "0x1000000001000000020000000300000004000000"
+      resource: "0x2000000001000000020000000300000004000000"
     }
   }
 #-}
@@ -49,10 +48,19 @@ TEST(Bytecode, MultiModuleWithResource) {
   std::string buffer;
   llvm::raw_string_ostream ostream(buffer);
   ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), ostream)));
+  ostream.flush();
+
+  // Create copy of buffer which is aligned to requested resource alignment.
+  constexpr size_t kAlignment = 0x20;
+  size_t buffer_size = buffer.size();
+  buffer.reserve(buffer_size + kAlignment - 1);
+  size_t pad = ~(uintptr_t)buffer.data() + 1 & kAlignment - 1;
+  buffer.insert(0, pad, ' ');
+  StringRef aligned_buffer(buffer.data() + pad, buffer_size);
 
   // Parse it back
   OwningOpRef<Operation *> roundTripModule =
-      parseSourceString<Operation *>(ostream.str(), parseConfig);
+      parseSourceString<Operation *>(aligned_buffer, parseConfig);
   ASSERT_TRUE(roundTripModule);
 
   // FIXME: Parsing external resources does not work on big-endian
@@ -80,3 +88,39 @@ 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"));
+}

utils/bazel/llvm-project-overlay/mlir/unittests/BUILD.bazel

@@ -359,6 +359,27 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "bytecode_tests",
+    size = "small",
+    srcs = glob([
+        "Bytecode/*.cpp",
+        "Bytecode/*.h",
+        "Bytecode/*/*.cpp",
+        "Bytecode/*/*.h",
+    ]),
+    deps = [
+        "//llvm:Support",
+        "//mlir:BytecodeReader",
+        "//mlir:BytecodeWriter",
+        "//mlir:IR",
+        "//mlir:Parser",
+        "//third-party/unittest:gmock",
+        "//third-party/unittest:gtest",
+        "//third-party/unittest:gtest_main",
+    ],
+)
+
 cc_test(
     name = "conversion_tests",
     size = "small",