Update ODS variadic segments "magic" attributes to use native Properties

The operand_segment_sizes and result_segment_sizes Attributes are now inlined
in the operation as native propertie. We continue to support building an
Attribute on the fly for `getAttr("operand_segment_sizes")` and setting the
property from an attribute with `setAttr("operand_segment_sizes", attr)`.

A new bytecode version is introduced to support backward compatibility and
backdeployments.

Differential Revision: https://reviews.llvm.org/D155919

Author: joker-eph

mlir/include/mlir/Bytecode/BytecodeImplementation.h

@@ -20,6 +20,7 @@
 #include "mlir/IR/DialectInterface.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Twine.h"
 
 namespace mlir {
@@ -39,6 +40,9 @@ class DialectBytecodeReader {
   /// Emit an error to the reader.
   virtual InFlightDiagnostic emitError(const Twine &msg = {}) = 0;
 
+  /// Return the bytecode version being read.
+  virtual uint64_t getBytecodeVersion() const = 0;
+
   /// Read out a list of elements, invoking the provided callback for each
   /// element. The callback function may be in any of the following forms:
   ///   * LogicalResult(T &)
@@ -148,6 +152,76 @@ class DialectBytecodeReader {
                     [this](int64_t &value) { return readSignedVarInt(value); });
   }
 
+  /// Parse a variable length encoded integer whose low bit is used to encode an
+  /// unrelated flag, i.e: `(integerValue << 1) | (flag ? 1 : 0)`.
+  LogicalResult readVarIntWithFlag(uint64_t &result, bool &flag) {
+    if (failed(readVarInt(result)))
+      return failure();
+    flag = result & 1;
+    result >>= 1;
+    return success();
+  }
+
+  /// Read a "small" sparse array of integer <= 32 bits elements, where
+  /// index/value pairs can be compressed when the array is small.
+  /// Note that only some position of the array will be read and the ones
+  /// not stored in the bytecode are gonne be left untouched.
+  /// If the provided array is too small for the stored indices, an error
+  /// will be returned.
+  template <typename T>
+  LogicalResult readSparseArray(MutableArrayRef<T> array) {
+    static_assert(sizeof(T) < sizeof(uint64_t), "expect integer < 64 bits");
+    static_assert(std::is_integral<T>::value, "expects integer");
+    uint64_t nonZeroesCount;
+    bool useSparseEncoding;
+    if (failed(readVarIntWithFlag(nonZeroesCount, useSparseEncoding)))
+      return failure();
+    if (nonZeroesCount == 0)
+      return success();
+    if (!useSparseEncoding) {
+      // This is a simple dense array.
+      if (nonZeroesCount > array.size()) {
+        emitError("trying to read an array of ")
+            << nonZeroesCount << " but only " << array.size()
+            << " storage available.";
+        return failure();
+      }
+      for (int64_t index : llvm::seq<int64_t>(0, nonZeroesCount)) {
+        uint64_t value;
+        if (failed(readVarInt(value)))
+          return failure();
+        array[index] = value;
+      }
+      return success();
+    }
+    // Read sparse encoding
+    // This is the number of bits used for packing the index with the value.
+    uint64_t indexBitSize;
+    if (failed(readVarInt(indexBitSize)))
+      return failure();
+    constexpr uint64_t maxIndexBitSize = 8;
+    if (indexBitSize > maxIndexBitSize) {
+      emitError("reading sparse array with indexing above 8 bits: ")
+          << indexBitSize;
+      return failure();
+    }
+    for (uint32_t count : llvm::seq<uint32_t>(0, nonZeroesCount)) {
+      (void)count;
+      uint64_t indexValuePair;
+      if (failed(readVarInt(indexValuePair)))
+        return failure();
+      uint64_t index = indexValuePair & ~(uint64_t(-1) << (indexBitSize));
+      uint64_t value = indexValuePair >> indexBitSize;
+      if (index >= array.size()) {
+        emitError("reading a sparse array found index ")
+            << index << " but only " << array.size() << " storage available.";
+        return failure();
+      }
+      array[index] = value;
+    }
+    return success();
+  }
+
   /// Read an APInt that is known to have been encoded with the given width.
   virtual FailureOr<APInt> readAPIntWithKnownWidth(unsigned bitWidth) = 0;
 
@@ -230,6 +304,55 @@ class DialectBytecodeWriter {
     writeList(value, [this](int64_t value) { writeSignedVarInt(value); });
   }
 
+  /// Write a VarInt and a flag packed together.
+  void writeVarIntWithFlag(uint64_t value, bool flag) {
+    writeVarInt((value << 1) | (flag ? 1 : 0));
+  }
+
+  /// Write out a "small" sparse array of integer <= 32 bits elements, where
+  /// index/value pairs can be compressed when the array is small. This method
+  /// will scan the array multiple times and should not be used for large
+  /// arrays. The optional provided "zero" can be used to adjust for the
+  /// expected repeated value. We assume here that the array size fits in a 32
+  /// bits integer.
+  template <typename T>
+  void writeSparseArray(ArrayRef<T> array) {
+    static_assert(sizeof(T) < sizeof(uint64_t), "expect integer < 64 bits");
+    static_assert(std::is_integral<T>::value, "expects integer");
+    uint32_t size = array.size();
+    uint32_t nonZeroesCount = 0, lastIndex = 0;
+    for (uint32_t index : llvm::seq<uint32_t>(0, size)) {
+      if (!array[index])
+        continue;
+      nonZeroesCount++;
+      lastIndex = index;
+    }
+    // If the last position is too large, or the array isn't at least 50%
+    // sparse, emit it with a dense encoding.
+    if (lastIndex > 256 || nonZeroesCount > size / 2) {
+      // Emit the array size and a flag which indicates whether it is sparse.
+      writeVarIntWithFlag(size, false);
+      for (const T &elt : array)
+        writeVarInt(elt);
+      return;
+    }
+    // Emit sparse: first the number of elements we'll write and a flag
+    // indicating it is a sparse encoding.
+    writeVarIntWithFlag(nonZeroesCount, true);
+    if (nonZeroesCount == 0)
+      return;
+    // This is the number of bits used for packing the index with the value.
+    int indexBitSize = llvm::Log2_32_Ceil(lastIndex + 1);
+    writeVarInt(indexBitSize);
+    for (uint32_t index : llvm::seq<uint32_t>(0, lastIndex + 1)) {
+      T value = array[index];
+      if (!value)
+        continue;
+      uint64_t indexValuePair = (value << indexBitSize) | (index);
+      writeVarInt(indexValuePair);
+    }
+  }
+
   /// Write an APInt to the bytecode stream whose bitwidth will be known
   /// externally at read time. This method is useful for encoding APInt values
   /// when the width is known via external means, such as via a type. This

mlir/include/mlir/Bytecode/Encoding.h

@@ -45,8 +45,12 @@ enum BytecodeVersion {
   /// with the discardable attributes.
   kNativePropertiesEncoding = 5,
 
+  /// ODS emits operand/result segment_size as native properties instead of
+  /// an attribute.
+  kNativePropertiesODSSegmentSize = 6,
+
   /// The current bytecode version.
-  kVersion = 5,
+  kVersion = 6,
 
   /// An arbitrary value used to fill alignment padding.
   kAlignmentByte = 0xCB,

mlir/include/mlir/IR/ODSSupport.h

@@ -37,6 +37,13 @@ Attribute convertToAttribute(MLIRContext *ctx, int64_t storage);
 LogicalResult convertFromAttribute(MutableArrayRef<int64_t> storage,
                                    Attribute attr, InFlightDiagnostic *diag);
 
+/// Convert a DenseI32ArrayAttr to the provided storage. It is expected that the
+/// storage has the same size as the array. An error is returned if the
+/// attribute isn't a DenseI32ArrayAttr or it does not have the same size. If
+/// the optional diagnostic is provided an error message is also emitted.
+LogicalResult convertFromAttribute(MutableArrayRef<int32_t> storage,
+                                   Attribute attr, InFlightDiagnostic *diag);
+
 /// Convert the provided ArrayRef<int64_t> to a DenseI64ArrayAttr attribute.
 Attribute convertToAttribute(MLIRContext *ctx, ArrayRef<int64_t> storage);
 

mlir/include/mlir/IR/OpBase.td

@@ -1241,6 +1241,7 @@ class ArrayProperty<string storageTypeParam = "", int n, string desc = ""> :
   let interfaceType = "::llvm::ArrayRef<" # storageTypeParam # ">";
   let convertFromStorage = "$_storage";
   let assignToStorage = "::llvm::copy($_value, $_storage)";
+  let hashProperty = "llvm::hash_combine_range(std::begin($_storage), std::end($_storage));";
 }
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/OpDefinition.h

@@ -20,6 +20,7 @@
 #define MLIR_IR_OPDEFINITION_H
 
 #include "mlir/IR/Dialect.h"
+#include "mlir/IR/ODSSupport.h"
 #include "mlir/IR/Operation.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 

mlir/include/mlir/IR/OperationSupport.h

@@ -555,7 +555,8 @@ class RegisteredOperationName : public OperationName {
                                              StringRef name) final {
       if constexpr (hasProperties) {
         auto concreteOp = cast<ConcreteOp>(op);
-        return ConcreteOp::getInherentAttr(concreteOp.getProperties(), name);
+        return ConcreteOp::getInherentAttr(concreteOp.getContext(),
+                                           concreteOp.getProperties(), name);
       }
       // If the op does not have support for properties, we dispatch back to the
       // dictionnary of discardable attributes for now.
@@ -575,7 +576,8 @@ class RegisteredOperationName : public OperationName {
     void populateInherentAttrs(Operation *op, NamedAttrList &attrs) final {
       if constexpr (hasProperties) {
         auto concreteOp = cast<ConcreteOp>(op);
-        ConcreteOp::populateInherentAttrs(concreteOp.getProperties(), attrs);
+        ConcreteOp::populateInherentAttrs(concreteOp.getContext(),
+                                          concreteOp.getProperties(), attrs);
       }
     }
     LogicalResult

mlir/include/mlir/TableGen/Property.h

@@ -35,51 +35,76 @@ class Property {
 public:
   explicit Property(const llvm::Record *record);
   explicit Property(const llvm::DefInit *init);
+  Property(StringRef storageType, StringRef interfaceType,
+           StringRef convertFromStorageCall, StringRef assignToStorageCall,
+           StringRef convertToAttributeCall, StringRef convertFromAttributeCall,
+           StringRef readFromMlirBytecodeCall,
+           StringRef writeToMlirBytecodeCall, StringRef hashPropertyCall,
+           StringRef defaultValue);
 
   // Returns the storage type.
-  StringRef getStorageType() const;
+  StringRef getStorageType() const { return storageType; }
 
   // Returns the interface type for this property.
-  StringRef getInterfaceType() const;
+  StringRef getInterfaceType() const { return interfaceType; }
 
   // Returns the template getter method call which reads this property's
   // storage and returns the value as of the desired return type.
-  StringRef getConvertFromStorageCall() const;
+  StringRef getConvertFromStorageCall() const { return convertFromStorageCall; }
 
   // Returns the template setter method call which reads this property's
   // in the provided interface type and assign it to the storage.
-  StringRef getAssignToStorageCall() const;
+  StringRef getAssignToStorageCall() const { return assignToStorageCall; }
 
   // Returns the conversion method call which reads this property's
   // in the storage type and builds an attribute.
-  StringRef getConvertToAttributeCall() const;
+  StringRef getConvertToAttributeCall() const { return convertToAttributeCall; }
 
   // Returns the setter method call which reads this property's
   // in the provided interface type and assign it to the storage.
-  StringRef getConvertFromAttributeCall() const;
+  StringRef getConvertFromAttributeCall() const {
+    return convertFromAttributeCall;
+  }
 
   // Returns the method call which reads this property from
   // bytecode and assign it to the storage.
-  StringRef getReadFromMlirBytecodeCall() const;
+  StringRef getReadFromMlirBytecodeCall() const {
+    return readFromMlirBytecodeCall;
+  }
 
   // Returns the method call which write this property's
   // to the the bytecode.
-  StringRef getWriteToMlirBytecodeCall() const;
+  StringRef getWriteToMlirBytecodeCall() const {
+    return writeToMlirBytecodeCall;
+  }
 
   // Returns the code to compute the hash for this property.
-  StringRef getHashPropertyCall() const;
+  StringRef getHashPropertyCall() const { return hashPropertyCall; }
 
   // Returns whether this Property has a default value.
-  bool hasDefaultValue() const;
+  bool hasDefaultValue() const { return !defaultValue.empty(); }
+
   // Returns the default value for this Property.
-  StringRef getDefaultValue() const;
+  StringRef getDefaultValue() const { return defaultValue; }
 
   // Returns the TableGen definition this Property was constructed from.
-  const llvm::Record &getDef() const;
+  const llvm::Record &getDef() const { return *def; }
 
 private:
   // The TableGen definition of this constraint.
   const llvm::Record *def;
+
+  // Elements describing a Property, in general fetched from the record.
+  StringRef storageType;
+  StringRef interfaceType;
+  StringRef convertFromStorageCall;
+  StringRef assignToStorageCall;
+  StringRef convertToAttributeCall;
+  StringRef convertFromAttributeCall;
+  StringRef readFromMlirBytecodeCall;
+  StringRef writeToMlirBytecodeCall;
+  StringRef hashPropertyCall;
+  StringRef defaultValue;
 };
 
 // A struct wrapping an op property and its name together