[mlir] Add bytecode encodings for the builtin ElementsAttr attributes

[rriddle]

Summary

This adds bytecode support for DenseArrayAttr, DenseIntOrFpElementsAttr,
DenseStringElementsAttr, and SparseElementsAttr.

Depends on D133743

[jpienaar]

Nice!

mlir/include/mlir/Bytecode/BytecodeImplementation.h

233

Same lifetime requirements as String above?

[abeakkas]

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

92

There's a small typo here. It should be "CHECK{LITERAL}" instead.

mlir/include/mlir/Bytecode/BytecodeImplementation.h

public:
   /// Read an APFloat that is known to have been encoded with the given
   /// semantics.
   virtual FailureOr<APFloat>
   readAPFloatWithKnownSemantics(const llvm::fltSemantics &semantics) = 0;
 
   /// Read a string from the bytecode.
   virtual LogicalResult readString(StringRef &result) = 0;
 
+  /// Read a blob from the bytecode.
+  virtual LogicalResult readBlob(ArrayRef<char> &result) = 0;
+
 private:
   /// Read a handle to a dialect resource.
   virtual FailureOr<AsmDialectResourceHandle> readResourceHandle() = 0;
 };
 
 //===----------------------------------------------------------------------===//
 // DialectBytecodeWriter
 //===----------------------------------------------------------------------===//
public:
   /// when the semantics are known via external means, such as via a type.
   virtual void writeAPFloatWithKnownSemantics(const APFloat &value) = 0;
 
   /// Write a string to the bytecode, which is owned by the caller and is
   /// guaranteed to not die before the end of the bytecode process. This should
   /// only be called if such a guarantee can be made, such as when the string is
   /// owned by an attribute or type.
   virtual void writeOwnedString(StringRef str) = 0;
+
+  /// Write a blob to the bytecode, which is owned by the caller and is
+  /// guaranteed to not die before the end of the bytecode process. The blob is
+  /// written as-is, with no additional compression or compaction.
+  virtual void writeOwnedBlob(ArrayRef<char> blob) = 0;
 };
 
 //===----------------------------------------------------------------------===//
 // BytecodeDialectInterface
 //===----------------------------------------------------------------------===//
 
 class BytecodeDialectInterface
     : public DialectInterface::Base<BytecodeDialectInterface> {

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

if (failed(intVal))
       return failure();
     return APFloat(semantics, *intVal);
   }
 
   LogicalResult readString(StringRef &result) override {
     return stringReader.parseString(reader, result);
   }
 
+  LogicalResult readBlob(ArrayRef<char> &result) override {
+    uint64_t dataSize;
+    ArrayRef<uint8_t> data;
+    if (failed(reader.parseVarInt(dataSize)) ||
+        failed(reader.parseBytes(dataSize, data)))
+      return failure();
+    result = llvm::makeArrayRef(reinterpret_cast<const char *>(data.data()),
+                                data.size());
+    return success();
+  }
+
 private:
   AttrTypeReader &attrTypeReader;
   StringSectionReader &stringReader;
   ResourceSectionReader &resourceReader;
   EncodingReader &reader;
 };
 } // namespace
 

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

public:
   void writeAPFloatWithKnownSemantics(const APFloat &value) override {
     writeAPIntWithKnownWidth(value.bitcastToAPInt());
   }
 
   void writeOwnedString(StringRef str) override {
     emitter.emitVarInt(stringSection.insert(str));
   }
 
+  void writeOwnedBlob(ArrayRef<char> blob) override {
+    emitter.emitVarInt(blob.size());
+    emitter.emitOwnedBlob(ArrayRef<uint8_t>(
+        reinterpret_cast<const uint8_t *>(blob.data()), blob.size()));
+  }
+
 private:
   EncodingEmitter &emitter;
   IRNumberingState &numberingState;
   StringSectionBuilder &stringSection;
 };
 } // namespace
 
 void BytecodeWriter::writeAttrTypeSection(EncodingEmitter &emitter) {

mlir/lib/Bytecode/Writer/IRNumbering.cpp

struct IRNumberingState::NumberingDialectWriter : public DialectBytecodeWriter {
   void writeSignedVarInt(int64_t value) override {}
   void writeAPIntWithKnownWidth(const APInt &value) override {}
   void writeAPFloatWithKnownSemantics(const APFloat &value) override {}
   void writeOwnedString(StringRef) override {
     // TODO: It might be nice to prenumber strings and sort by the number of
     // references. This could potentially be useful for optimizing things like
     // file locations.
   }
+  void writeOwnedBlob(ArrayRef<char> blob) override {}
 
   /// The parent numbering state that is populated by this writer.
   IRNumberingState &state;
 };
 
 //===----------------------------------------------------------------------===//
 // IR Numbering
 //===----------------------------------------------------------------------===//

mlir/lib/IR/BuiltinDialectBytecode.cpp

enum AttributeCode {
   ///   }
   kUnknownLoc = 15,
 
   ///   DenseResourceElementsAttr {
   ///     type: Type,
   ///     handle: ResourceHandle
   ///   }
   kDenseResourceElementsAttr = 16,
+
+  ///   DenseArrayAttr {
+  ///     type: RankedTensorType,
+  ///     data: blob
+  ///   }
+  kDenseArrayAttr = 17,
+
+  ///   DenseIntOrFPElementsAttr {
+  ///     type: ShapedType,
+  ///     data: blob
+  ///   }
+  kDenseIntOrFPElementsAttr = 18,
+
+  ///   DenseStringElementsAttr {
+  ///     type: ShapedType,
+  ///     isSplat: varint,
+  ///     data: string[]
+  ///   }
+  kDenseStringElementsAttr = 19,
+
+  ///   SparseElementsAttr {
+  ///     type: ShapedType,
+  ///     indices: DenseIntElementsAttr,
+  ///     values: DenseElementsAttr
+  ///   }
+  kSparseElementsAttr = 20,
 };
 
 /// This enum contains marker codes used to indicate which type is currently
 /// being decoded, and how it should be decoded. The order of these codes should
 /// generally be unchanged, as any changes will inevitably break compatibility
 /// with older bytecode.
 enum TypeCode {
   ///   IntegerType {
struct BuiltinDialectBytecodeInterface : public BytecodeDialectInterface {
   BuiltinDialectBytecodeInterface(Dialect *dialect)
       : BytecodeDialectInterface(dialect) {}
 
   //===--------------------------------------------------------------------===//
   // Attributes
 
   Attribute readAttribute(DialectBytecodeReader &reader) const override;
   ArrayAttr readArrayAttr(DialectBytecodeReader &reader) const;
+  DenseArrayAttr readDenseArrayAttr(DialectBytecodeReader &reader) const;
+  DenseElementsAttr
+  readDenseIntOrFPElementsAttr(DialectBytecodeReader &reader) const;
+  DenseStringElementsAttr
+  readDenseStringElementsAttr(DialectBytecodeReader &reader) const;
   DenseResourceElementsAttr
   readDenseResourceElementsAttr(DialectBytecodeReader &reader) const;
   DictionaryAttr readDictionaryAttr(DialectBytecodeReader &reader) const;
   FloatAttr readFloatAttr(DialectBytecodeReader &reader) const;
   IntegerAttr readIntegerAttr(DialectBytecodeReader &reader) const;
+  SparseElementsAttr
+  readSparseElementsAttr(DialectBytecodeReader &reader) const;
   StringAttr readStringAttr(DialectBytecodeReader &reader, bool hasType) const;
   SymbolRefAttr readSymbolRefAttr(DialectBytecodeReader &reader,
                                   bool hasNestedRefs) const;
   TypeAttr readTypeAttr(DialectBytecodeReader &reader) const;
 
   LocationAttr readCallSiteLoc(DialectBytecodeReader &reader) const;
   LocationAttr readFileLineColLoc(DialectBytecodeReader &reader) const;
   LocationAttr readFusedLoc(DialectBytecodeReader &reader,
                             bool hasMetadata) const;
   LocationAttr readNameLoc(DialectBytecodeReader &reader) const;
 
   LogicalResult writeAttribute(Attribute attr,
                                DialectBytecodeWriter &writer) const override;
   void write(ArrayAttr attr, DialectBytecodeWriter &writer) const;
+  void write(DenseArrayAttr attr, DialectBytecodeWriter &writer) const;
+  void write(DenseIntOrFPElementsAttr attr,
+             DialectBytecodeWriter &writer) const;
+  void write(DenseStringElementsAttr attr, DialectBytecodeWriter &writer) const;
   void write(DenseResourceElementsAttr attr,
              DialectBytecodeWriter &writer) const;
   void write(DictionaryAttr attr, DialectBytecodeWriter &writer) const;
   void write(IntegerAttr attr, DialectBytecodeWriter &writer) const;
   void write(FloatAttr attr, DialectBytecodeWriter &writer) const;
+  void write(SparseElementsAttr attr, DialectBytecodeWriter &writer) const;
   void write(StringAttr attr, DialectBytecodeWriter &writer) const;
   void write(SymbolRefAttr attr, DialectBytecodeWriter &writer) const;
   void write(TypeAttr attr, DialectBytecodeWriter &writer) const;
 
   void write(CallSiteLoc attr, DialectBytecodeWriter &writer) const;
   void write(FileLineColLoc attr, DialectBytecodeWriter &writer) const;
   void write(FusedLoc attr, DialectBytecodeWriter &writer) const;
   void write(NameLoc attr, DialectBytecodeWriter &writer) const;
Attribute BuiltinDialectBytecodeInterface::readAttribute(
   case builtin_encoding::kFusedLocWithMetadata:
     return readFusedLoc(reader, /*hasMetadata=*/true);
   case builtin_encoding::kNameLoc:
     return readNameLoc(reader);
   case builtin_encoding::kUnknownLoc:
     return UnknownLoc::get(getContext());
   case builtin_encoding::kDenseResourceElementsAttr:
     return readDenseResourceElementsAttr(reader);
+  case builtin_encoding::kDenseArrayAttr:
+    return readDenseArrayAttr(reader);
+  case builtin_encoding::kDenseIntOrFPElementsAttr:
+    return readDenseIntOrFPElementsAttr(reader);
+  case builtin_encoding::kDenseStringElementsAttr:
+    return readDenseStringElementsAttr(reader);
+  case builtin_encoding::kSparseElementsAttr:
+    return readSparseElementsAttr(reader);
   default:
     reader.emitError() << "unknown builtin attribute code: " << code;
     return Attribute();
   }
 }
 
 LogicalResult BuiltinDialectBytecodeInterface::writeAttribute(
     Attribute attr, DialectBytecodeWriter &writer) const {
   return TypeSwitch<Attribute, LogicalResult>(attr)
-      .Case<ArrayAttr, DenseResourceElementsAttr, DictionaryAttr, FloatAttr,
-            IntegerAttr, StringAttr, SymbolRefAttr, TypeAttr>([&](auto attr) {
+      .Case<ArrayAttr, DenseArrayAttr, DenseIntOrFPElementsAttr,
+            DenseStringElementsAttr, DenseResourceElementsAttr, DictionaryAttr,
+            FloatAttr, IntegerAttr, SparseElementsAttr, StringAttr,
+            SymbolRefAttr, TypeAttr>([&](auto attr) {
         write(attr, writer);
         return success();
       })
       .Case<CallSiteLoc, FileLineColLoc, FusedLoc, NameLoc>([&](auto attr) {
         write(attr, writer);
         return success();
       })
       .Case([&](OpaqueLoc attr) { return write(attr, writer); })

 
 void BuiltinDialectBytecodeInterface::write(
     ArrayAttr attr, DialectBytecodeWriter &writer) const {
   writer.writeVarInt(builtin_encoding::kArrayAttr);
   writer.writeAttributes(attr.getValue());
 }
 
 //===----------------------------------------------------------------------===//
+// DenseArrayAttr
+
+DenseArrayAttr BuiltinDialectBytecodeInterface::readDenseArrayAttr(
+    DialectBytecodeReader &reader) const {
+  RankedTensorType type;
+  ArrayRef<char> blob;
+  if (failed(reader.readType(type)) || failed(reader.readBlob(blob)))
+    return DenseArrayAttr();
+  return DenseArrayAttr::get(type, blob);
+}
+
+void BuiltinDialectBytecodeInterface::write(
+    DenseArrayAttr attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kDenseArrayAttr);
+  writer.writeType(attr.getType());
+  writer.writeOwnedBlob(attr.getRawData());
+}
+
+//===----------------------------------------------------------------------===//
+// DenseIntOrFPElementsAttr
+
+DenseElementsAttr BuiltinDialectBytecodeInterface::readDenseIntOrFPElementsAttr(
+    DialectBytecodeReader &reader) const {
+  ShapedType type;
+  ArrayRef<char> blob;
+  if (failed(reader.readType(type)) || failed(reader.readBlob(blob)))
+    return DenseIntOrFPElementsAttr();
+  return DenseIntOrFPElementsAttr::getFromRawBuffer(type, blob);
+}
+
+void BuiltinDialectBytecodeInterface::write(
+    DenseIntOrFPElementsAttr attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kDenseIntOrFPElementsAttr);
+  writer.writeType(attr.getType());
+  writer.writeOwnedBlob(attr.getRawData());
+}
+
+//===----------------------------------------------------------------------===//
+// DenseStringElementsAttr
+
+DenseStringElementsAttr
+BuiltinDialectBytecodeInterface::readDenseStringElementsAttr(
+    DialectBytecodeReader &reader) const {
+  ShapedType type;
+  uint64_t isSplat;
+  if (failed(reader.readType(type)) || failed(reader.readVarInt(isSplat)))
+    return DenseStringElementsAttr();
+
+  SmallVector<StringRef> values(isSplat ? 1 : type.getNumElements());
+  for (StringRef &value : values)
+    if (failed(reader.readString(value)))
+      return DenseStringElementsAttr();
+  return DenseStringElementsAttr::get(type, values);
+}
+
+void BuiltinDialectBytecodeInterface::write(
+    DenseStringElementsAttr attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kDenseStringElementsAttr);
+  writer.writeType(attr.getType());
+
+  bool isSplat = attr.isSplat();
+  writer.writeVarInt(isSplat);
+
+  // If the attribute is a splat, only write out the single value.
+  if (isSplat)
+    return writer.writeOwnedString(attr.getRawStringData().front());
+
+  for (StringRef str : attr.getRawStringData())
+    writer.writeOwnedString(str);
+}
+
+//===----------------------------------------------------------------------===//
 // DenseResourceElementsAttr
 
 DenseResourceElementsAttr
 BuiltinDialectBytecodeInterface::readDenseResourceElementsAttr(
     DialectBytecodeReader &reader) const {
   ShapedType type;
   if (failed(reader.readType(type)))
     return DenseResourceElementsAttr();

 void BuiltinDialectBytecodeInterface::write(
     IntegerAttr attr, DialectBytecodeWriter &writer) const {
   writer.writeVarInt(builtin_encoding::kIntegerAttr);
   writer.writeType(attr.getType());
   writer.writeAPIntWithKnownWidth(attr.getValue());
 }
 
 //===----------------------------------------------------------------------===//
+// SparseElementsAttr
+
+SparseElementsAttr BuiltinDialectBytecodeInterface::readSparseElementsAttr(
+    DialectBytecodeReader &reader) const {
+  ShapedType type;
+  DenseIntElementsAttr indices;
+  DenseElementsAttr values;
+  if (failed(reader.readType(type)) || failed(reader.readAttribute(indices)) ||
+      failed(reader.readAttribute(values)))
+    return SparseElementsAttr();
+  return SparseElementsAttr::get(type, indices, values);
+}
+
+void BuiltinDialectBytecodeInterface::write(
+    SparseElementsAttr attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kSparseElementsAttr);
+  writer.writeType(attr.getType());
+  writer.writeAttribute(attr.getIndices());
+  writer.writeAttribute(attr.getValues());
+}
+
+//===----------------------------------------------------------------------===//
 // StringAttr
 
 StringAttr
 BuiltinDialectBytecodeInterface::readStringAttr(DialectBytecodeReader &reader,
                                                 bool hasType) const {
   StringRef string;
   if (failed(reader.readString(string)))
     return StringAttr();

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

-// RUN: mlir-opt -emit-bytecode %s | mlir-opt -mlir-print-local-scope | FileCheck %s
+// 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: s390x-
 
 //===----------------------------------------------------------------------===//
 // ArrayAttr
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @TestArray
 module @TestArray attributes {
   // CHECK: bytecode.array = [unit]
   bytecode.array = [unit]
 } {}
 
 //===----------------------------------------------------------------------===//
+// DenseArrayAttr
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @TestDenseArray
+module @TestDenseArray attributes {
+  // CHECK: bytecode.test1 = array<i1: true, false, true, false, false>
+  // CHECK: bytecode.test2 = array<i8: 10, 32, -1>
+  // CHECK: bytecode.test3 = array<f64: 1.{{.*}}e+01, 3.2{{.*}}e+01, 1.809{{.*}}e+03
+  bytecode.test1 = array<i1: true, false, true, false, false>,
+  bytecode.test2 = array<i8: 10, 32, 255>,
+  bytecode.test3 = array<f64: 10.0, 32.0, 1809.0>
+} {}
+
+//===----------------------------------------------------------------------===//
+// DenseIntOfFPElementsAttr
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @TestDenseIntOrFPElements
+// CHECK: bytecode.test1 = dense<true> : tensor<256xi1>
+// CHECK: bytecode.test2 = dense<[10, 32, -1]> : tensor<3xi8>
+// CHECK: bytecode.test3 = dense<[1.{{.*}}e+01, 3.2{{.*}}e+01, 1.809{{.*}}e+03]> : tensor<3xf64>
+module @TestDenseIntOrFPElements attributes {
+  bytecode.test1 = dense<true> : tensor<256xi1>,
+  bytecode.test2 = dense<[10, 32, 255]> : tensor<3xi8>,
+  bytecode.test3 = dense<[10.0, 32.0, 1809.0]> : tensor<3xf64>
+} {}
+
+//===----------------------------------------------------------------------===//
+// DenseStringElementsAttr
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @TestDenseStringElementsAttr
+module @TestDenseStringElementsAttr attributes {
+  bytecode.test1 = dense<"splat"> : tensor<256x!bytecode.string>,
+  bytecode.test2 = dense<["foo", "bar", "baz"]> : tensor<3x!bytecode.string>
+} {}
+
+//===----------------------------------------------------------------------===//
 // FloatAttr
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @TestFloat
 module @TestFloat attributes {
   // CHECK: bytecode.float = 1.000000e+01 : f64
   // CHECK: bytecode.float1 = 0.10000{{.*}} : f80
   // CHECK: bytecode.float2 = 0.10000{{.*}} : f128
module @TestInt attributes {
   // CHECK: bytecode.int3 = 90000000000000000300000000000000000001 : i128
   bytecode.int = false,
   bytecode.int1 = -1 : i8,
   bytecode.int2 = 800 : ui64,
   bytecode.int3 = 90000000000000000300000000000000000001 : i128
 } {}
 
 //===----------------------------------------------------------------------===//
+// SparseElementsAttr
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @TestSparseElements
+module @TestSparseElements attributes {
+  // CHECK-LITERAL: bytecode.sparse = sparse<[[0, 0], [1, 2]], [1, 5]> : tensor<3x4xi32>
+  bytecode.sparse = sparse<[[0, 0], [1, 2]], [1, 5]> : tensor<3x4xi32>
+} {}
+
+
+//===----------------------------------------------------------------------===//
 // StringAttr
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @TestString
 module @TestString attributes {
   // CHECK: bytecode.string = "hello"
   // CHECK: bytecode.string2 = "hello" : i32
   bytecode.string = "hello",