[mlir] Fix bytecode reader/writer on big-endian platforms

[uweigand]

Summary

This makes the bytecode reader/writer work on big-endian platforms. The only problem was related to encoding of multi-byte integers, where both reader and writer code make implicit assumptions about endianness of the host platform.

This fixes the current test failures on s390x, and in addition allows to remove the UNSUPPORTED markers from all other bytecode-related test cases - they now also all pass on s390x.

Note that there is still one other, unrelated, endian problem regarding decoding of external resources. This is not fixed by this patch, and still causes two tests to fail:

• test/IR/elements-attr-interface.mlir (already marked as XFAIL on s390x)
• Bytecode/./MLIRBytecodeTests/Bytecode/MultiModuleWithResource

The latter unit test currently fails very early due to the bytecode encoding bug fixed here. After this patch, it will still fail, but later on when verifying the contents of the resource. I've now added a GFAIL_SKIP here to get the build bot green on s390x again.

[mehdi_amini]

LGTM, but please wait for @rriddle to have a look as well!

[jpienaar]

Its really nice how minimal this change needed to be.

[GleasonK]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

8–9

This can probably be deleted now?

[rriddle]

Amazing!

[uweigand]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

8–9

Good point, I'm done this now.

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

 //===- BytecodeReader.cpp - MLIR Bytecode Reader --------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-// TODO: Support for big-endian architectures.
-
 #include "mlir/Bytecode/BytecodeReader.h"
 #include "mlir/AsmParser/AsmParser.h"
 #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"
 #include "mlir/IR/Verifier.h"
 #include "mlir/IR/Visitors.h"
 #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>
 #include <numeric>
 #include <optional>
LogicalResult parseVarInt(uint64_t &result) {
     if (LLVM_LIKELY(result & 1)) {
       result >>= 1;
       return success();
     }
 
     // Handle the overwhelming uncommon case where the value required all 8
     // bytes (i.e. a really really big number). In this case, the marker byte is
     // all zeros: `00000000`.
-    if (LLVM_UNLIKELY(result == 0))
-      return parseBytes(sizeof(result), reinterpret_cast<uint8_t *>(&result));
+    if (LLVM_UNLIKELY(result == 0)) {
+      llvm::support::ulittle64_t resultLE;
+      if (failed(parseBytes(sizeof(resultLE),
+                            reinterpret_cast<uint8_t *>(&resultLE))))
+        return failure();
+      result = resultLE;
+      return success();
+    }
     return parseMultiByteVarInt(result);
   }
 
   /// Parse a signed variable length encoded integer from the byte stream. A
   /// signed varint is encoded as a normal varint with zigzag encoding applied,
   /// i.e. the low bit of the value is used to indicate the sign.
   LogicalResult parseSignedVarInt(uint64_t &result) {
     if (failed(parseVarInt(result)))
LLVM_ATTRIBUTE_NOINLINE LogicalResult parseMultiByteVarInt(uint64_t &result) {
     // here because we only care about the first byte, and so that be actually
     // get ctz intrinsic calls when possible (the `uint8_t` overload uses a loop
     // implementation).
     uint32_t numBytes = llvm::countr_zero<uint32_t>(result);
     assert(numBytes > 0 && numBytes <= 7 &&
            "unexpected number of trailing zeros in varint encoding");
 
     // Parse in the remaining bytes of the value.
-    if (failed(parseBytes(numBytes, reinterpret_cast<uint8_t *>(&result) + 1)))
+    llvm::support::ulittle64_t resultLE(result);
+    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
     // encoded.
-    result >>= (numBytes + 1);
+    result = resultLE >> (numBytes + 1);
     return success();
   }
 
   /// The current data iterator, and an iterator to the end of the buffer.
   const uint8_t *dataIt, *dataEnd;
 
   /// A location for the bytecode used to report errors.
   Location fileLoc;

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/CachedHashString.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Support/Endian.h"
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
 #include <optional>
 #include <sys/types.h>
 
 #define DEBUG_TYPE "mlir-bytecode-writer"
 
void EncodingEmitter::emitMultiByteVarInt(uint64_t value) {
   // Compute the number of bytes needed to encode the value. Each byte can hold
   // up to 7-bits of data. We only check up to the number of bits we can encode
   // in the first byte (8).
   uint64_t it = value >> 7;
   for (size_t numBytes = 2; numBytes < 9; ++numBytes) {
     if (LLVM_LIKELY(it >>= 7) == 0) {
       uint64_t encodedValue = (value << 1) | 0x1;
       encodedValue <<= (numBytes - 1);
-      emitBytes({reinterpret_cast<uint8_t *>(&encodedValue), numBytes});
+      llvm::support::ulittle64_t encodedValueLE(encodedValue);
+      emitBytes({reinterpret_cast<uint8_t *>(&encodedValueLE), numBytes});
       return;
     }
   }
 
   // If the value is too large to encode in a single byte, emit a special all
   // zero marker byte and splat the value directly.
   emitByte(0);
-  emitBytes({reinterpret_cast<uint8_t *>(&value), sizeof(value)});
+  llvm::support::ulittle64_t valueLE(value);
+  emitBytes({reinterpret_cast<uint8_t *>(&valueLE), sizeof(valueLE)});
 }
 
 //===----------------------------------------------------------------------===//
 // Bytecode Writer
 //===----------------------------------------------------------------------===//
 
 namespace {
 class BytecodeWriter {

mlir/test/Bytecode/general.mlir

 // RUN: mlir-opt -allow-unregistered-dialect -emit-bytecode %s | mlir-opt -allow-unregistered-dialect | FileCheck %s
 
-// Bytecode currently does not support big-endian platforms
-// UNSUPPORTED: target=s390x-{{.*}}
-
 // CHECK-LABEL: "bytecode.test1"
 // CHECK-NEXT:    "unregistered.op"() {test_attr = #test.dynamic_singleton} : () -> ()
 // CHECK-NEXT:    "bytecode.empty"() : () -> ()
 // CHECK-NEXT:    "bytecode.attributes"() {attra = 10 : i64, attrb = #bytecode.attr} : () -> ()
 // CHECK-NEXT{LITERAL}: "bytecode.sparse"() {value = sparse<[[2, 1], [1, 1], [1, 2]], [1.
 // CHECK-NEXT:    test.graph_region {
 // CHECK-NEXT:      "bytecode.operands"(%[[RESULTS:.*]]#0, %[[RESULTS]]#1, %[[RESULTS]]#2) : (i32, i64, i32) -> ()
 // CHECK-NEXT:      %[[RESULTS]]:3 = "bytecode.results"() : () -> (i32, i64, i32)

mlir/test/Bytecode/invalid/invalid-dialect_section.mlir

 // This file contains various failure test cases related to the structure of
 // the dialect section.
 
-// Bytecode currently does not support big-endian platforms
-// UNSUPPORTED: target=s390x-{{.*}}
-
 //===--------------------------------------------------------------------===//
 // Dialect Name
 //===--------------------------------------------------------------------===//
 
 // RUN: not mlir-opt %S/invalid-dialect_section-dialect_string.mlirbc 2>&1 | FileCheck %s --check-prefix=DIALECT_STR
 // DIALECT_STR: invalid string index: 15
 
 //===--------------------------------------------------------------------===//
 // OpName
 //===--------------------------------------------------------------------===//
 
 // RUN: not mlir-opt %S/invalid-dialect_section-opname_dialect.mlirbc 2>&1 | FileCheck %s --check-prefix=OPNAME_DIALECT
 // OPNAME_DIALECT: invalid dialect index: 7
 
 // RUN: not mlir-opt %S/invalid-dialect_section-opname_string.mlirbc 2>&1 | FileCheck %s --check-prefix=OPNAME_STR
 // OPNAME_STR: invalid string index: 31

mlir/test/Bytecode/invalid/invalid-ir_section.mlir

 // This file contains various failure test cases related to the structure of
 // the IR section.
 
-// Bytecode currently does not support big-endian platforms
-// UNSUPPORTED: target=s390x-{{.*}}
-
 //===--------------------------------------------------------------------===//
 // Operations
 //===--------------------------------------------------------------------===//
 
 //===--------------------------------------------------------------------===//
 // Name
 
 // RUN: not mlir-opt %S/invalid-ir_section-opname.mlirbc -allow-unregistered-dialect 2>&1 | FileCheck %s --check-prefix=OP_NAME

mlir/test/Bytecode/invalid/invalid-string_section.mlir

 // This file contains various failure test cases related to the structure of
 // the string section.
 
-// Bytecode currently does not support big-endian platforms
-// UNSUPPORTED: target=s390x-{{.*}}
-
 //===--------------------------------------------------------------------===//
 // Count
 //===--------------------------------------------------------------------===//
 
 // RUN: not mlir-opt %S/invalid-string_section-count.mlirbc 2>&1 | FileCheck %s --check-prefix=COUNT
 // COUNT: attempting to parse a byte at the end of the bytecode
 
 //===--------------------------------------------------------------------===//

mlir/test/Bytecode/invalid/invalid-structure.mlir

 // This file contains various failure test cases related to the structure of
 // a bytecode file.
 
-// Bytecode currently does not support big-endian platforms
-// UNSUPPORTED: target=s390x-{{.*}}
-
 //===--------------------------------------------------------------------===//
 // Version
 //===--------------------------------------------------------------------===//
 
 // RUN: not mlir-opt %S/invalid-structure-version.mlirbc 2>&1 | FileCheck %s --check-prefix=VERSION
 // VERSION: bytecode version 127 is newer than the current version {{[0-9]+}}
 
 //===--------------------------------------------------------------------===//

mlir/test/Bytecode/invalid/invalid_attr_type_offset_section.mlir

 // This file contains various failure test cases related to the structure of
 // the attribute/type offset section.
 
-// Bytecode currently does not support big-endian platforms
-// UNSUPPORTED: target=s390x-{{.*}}
-
 //===--------------------------------------------------------------------===//
 // Offset
 //===--------------------------------------------------------------------===//
 
 // RUN: not mlir-opt %S/invalid-attr_type_offset_section-large_offset.mlirbc 2>&1 | FileCheck %s --check-prefix=LARGE_OFFSET
 // LARGE_OFFSET: Attribute or Type entry offset points past the end of section
 
 //===--------------------------------------------------------------------===//
 // Trailing Data
 //===--------------------------------------------------------------------===//
 
 // RUN: not mlir-opt %S/invalid-attr_type_offset_section-trailing_data.mlirbc 2>&1 | FileCheck %s --check-prefix=TRAILING_DATA
 // TRAILING_DATA: unexpected trailing data in the Attribute/Type offset section

mlir/test/Bytecode/invalid/invalid_attr_type_section.mlir

 // This file contains various failure test cases related to the structure of
 // the attribute/type offset section.
 
-// Bytecode currently does not support big-endian platforms
-// UNSUPPORTED: target=s390x-{{.*}}
-
 //===--------------------------------------------------------------------===//
 // Index
 //===--------------------------------------------------------------------===//
 
 // RUN: not mlir-opt %S/invalid-attr_type_section-index.mlirbc 2>&1 | FileCheck %s --check-prefix=INDEX
 // INDEX: invalid Attribute index: 3
 
 //===--------------------------------------------------------------------===//
 // Trailing Data
 //===--------------------------------------------------------------------===//
 
 // RUN: not mlir-opt %S/invalid-attr_type_section-trailing_data.mlirbc 2>&1 | FileCheck %s --check-prefix=TRAILING_DATA
 // TRAILING_DATA: trailing characters found after Attribute assembly format: trailing

mlir/test/Bytecode/resources.mlir

 // RUN: mlir-opt -emit-bytecode %s | mlir-opt | FileCheck %s
 
-// Bytecode currently does not support big-endian platforms
-// UNSUPPORTED: target=s390x-{{.*}}
-
 // CHECK-LABEL: @TestDialectResources
 module @TestDialectResources attributes {
   // CHECK: bytecode.test = dense_resource<decl_resource> : tensor<2xui32>
   // CHECK: bytecode.test2 = dense_resource<resource> : tensor<4xf64>
   // CHECK: bytecode.test3 = dense_resource<resource_2> : tensor<4xf64>
   bytecode.test = dense_resource<decl_resource> : tensor<2xui32>,
   bytecode.test2 = dense_resource<resource> : tensor<4xf64>,
   bytecode.test3 = dense_resource<resource_2> : tensor<4xf64>

mlir/test/Bytecode/versioning/versioned_attr.mlir

 // This file contains a test case representative of a dialect parsing an
 // attribute with versioned custom encoding.
 
-// Bytecode currently does not support big-endian platforms
-// UNSUPPORTED: target=s390x-{{.*}}
-
 //===--------------------------------------------------------------------===//
 // Test attribute upgrade
 //===--------------------------------------------------------------------===//
 
 // COM: bytecode contains
 // COM: module {
 // COM:   version: 1.12
 // COM:   "test.versionedB"() <{attribute = #test.attr_params<24, 42>}> : () -> ()

mlir/test/Bytecode/versioning/versioned_bytecode.mlir

 // This file contains test cases related to roundtripping.
 
-// Bytecode currently does not support big-endian platforms
-// UNSUPPORTED: target=s390x-{{.*}}
-
 //===--------------------------------------------------------------------===//
 // Test roundtrip
 //===--------------------------------------------------------------------===//
 
 // RUN: mlir-opt %S/versioned-op-1.12.mlirbc -emit-bytecode \
 // RUN:   -emit-bytecode-version=0 | mlir-opt -o %t.1 && \
 // RUN: mlir-opt %S/versioned-op-1.12.mlirbc -o %t.2 && \
 // RUN: diff %t.1 %t.2

mlir/test/Bytecode/versioning/versioned_op.mlir

 // This file contains test cases related to the dialect post-parsing upgrade
 // mechanism.
 
-// Bytecode currently does not support big-endian platforms
-// UNSUPPORTED: target=s390x-{{.*}}
-
 //===--------------------------------------------------------------------===//
 // Test generic
 //===--------------------------------------------------------------------===//
 
 // COM: bytecode contains
 // COM: module {
 // COM:   version: 2.0
 // COM:   "test.versionedA"() <{dims = 123 : i64, modifier = false}> : () -> ()

mlir/test/Dialect/Builtin/Bytecode/attrs.mlir

 // RUN: mlir-opt -emit-bytecode -allow-unregistered-dialect %s | mlir-opt -allow-unregistered-dialect -mlir-print-local-scope | FileCheck %s
 
-// Bytecode currently does not support big-endian platforms
-// UNSUPPORTED: target=s390x-{{.*}}
-
 //===----------------------------------------------------------------------===//
 // ArrayAttr
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @TestArray
 module @TestArray attributes {
   // CHECK: bytecode.array = [unit]
   bytecode.array = [unit]

mlir/test/Dialect/Builtin/Bytecode/types.mlir

 // RUN: mlir-opt -emit-bytecode %s | mlir-opt | FileCheck %s
 
-// Bytecode currently does not support big-endian platforms
-// UNSUPPORTED: target=s390x-{{.*}}
-
 //===----------------------------------------------------------------------===//
 // ComplexType
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @TestComplex
 module @TestComplex attributes {
   // CHECK: bytecode.test = complex<i32>
   bytecode.test = complex<i32>

mlir/test/Dialect/Quant/Bytecode/types.mlir

 // RUN: mlir-opt -emit-bytecode %s | mlir-opt | FileCheck %s
 
-// Bytecode currently does not support big-endian platforms
-// UNSUPPORTED: target=s390x-{{.*}}
-
 //===----------------------------------------------------------------------===//
 // AnyQuantized
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: parseAnyFullySpecified
 module @parseAnyFullySpecified attributes {
   // CHECK: bytecode.test = !quant.any<i8<-8:7>:f32>
   bytecode.test = !quant.any<i8<-8:7>:f32>

mlir/unittests/Bytecode/BytecodeTest.cpp

 #include "mlir/Bytecode/BytecodeWriter.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/OwningOpRef.h"
 #include "mlir/Parser/Parser.h"
 
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Endian.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 using namespace mlir;
 
 using testing::ElementsAre;
 
TEST(Bytecode, MultiModuleWithResource) {
   llvm::raw_string_ostream ostream(buffer);
   ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), ostream)));
 
   // Parse it back
   OwningOpRef<Operation *> roundTripModule =
       parseSourceString<Operation *>(ostream.str(), parseConfig);
   ASSERT_TRUE(roundTripModule);
 
+  // FIXME: Parsing external resources does not work on big-endian
+  // platforms currently.
+  if (llvm::support::endian::system_endianness() ==
+      llvm::support::endianness::big)
+    GTEST_SKIP();
+
   // Try to see if we have a valid resource in the parsed module.
   auto checkResourceAttribute = [&](Operation *op) {
     Attribute attr = roundTripModule->getAttr("bytecode.test");
     ASSERT_TRUE(attr);
     auto denseResourceAttr = dyn_cast<DenseI32ResourceElementsAttr>(attr);
     ASSERT_TRUE(denseResourceAttr);
     std::optional<ArrayRef<int32_t>> attrData =
         denseResourceAttr.tryGetAsArrayRef();