[mlir:Bytecode] Add encoding support for a majority of the builtin attributes

[rriddle]

Summary

This adds support for the non-location, non-elements, non-affine
builtin attributes.

[mehdi_amini]

Nice!

mlir/docs/BytecodeFormat.md

78

zigzag encoding for 64-bits integers:

mlir/include/mlir/Bytecode/BytecodeImplementation.h

175–176

What about this TODO?
We could support overloads with int64_t with zigzag encoding for example, without paying the price of APInt.

181

Worth mentioning that this is using ZigZag encoding for symmetry between negative and positive numbers.

Right, now if I have a uintX_t, it's not clear to me why one would pick this method instead of the one above directly for example.

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

137

This bit manipulation is magic :)

(Translation: I'm too lazy to try to figure it out right now, but hey if it works...)

[rriddle]

mlir/include/mlir/Bytecode/BytecodeImplementation.h

175–176

Yeah, I was deferring it until I had a method that needed it, but let me just add now (it just forwards to the underlying method anyways).

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

137

I added a comment of the actual formula, which is easier to read IMO. This right here though, is the same thing that protobufs use for zigzag decoding.

[jpienaar]

Nice, reads without surprises :-)

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

563

I would have thought

read<What><How>

would result in better method sorting order. E.g., if I want to read an APLfloat , autocomplete would give me the variants just by me saying readAPFloat (

mlir/lib/IR/BuiltinDialectBytecode.cpp

83

Should we add comment why we don't have or want DenseElementsAttr here?

300–301

Keep these in alphabetical order?

[rriddle]

mlir/lib/IR/BuiltinDialectBytecode.cpp

83

I'm leaving that out for now, because I just haven't gotten to all of the attributes (so I haven't had time to figure out what we want to do there). We'll need to have a conversation on some of the attributes though, to figure out if we are planning on deleting things. We could also just add support for now, but then just delete it whenever we do drop the attributes. This all factors into versioning a bit as well (I don't want to version until we figure out what we want to change and how).

mlir/docs/BytecodeFormat.md

 xxxx1000: 28 value bits, the encoding uses 4 bytes
 xxx10000: 35 value bits, the encoding uses 5 bytes
 xx100000: 42 value bits, the encoding uses 6 bytes
 x1000000: 49 value bits, the encoding uses 7 bytes
 10000000: 56 value bits, the encoding uses 8 bytes
 00000000: 64 value bits, the encoding uses 9 bytes
 ```
 
+##### Signed Variable-Width Integers
+
+Signed variable width integer values are encoded in a similar fashion to
+[varints](#variable-width-integers), but employ
+[zigzag encoding](https://en.wikipedia.org/wiki/Variable-length_quantity#Zigzag_encoding).
+This encoding uses the low bit of the value to indicate the sign, which allows
+for more efficiently encoding negative numbers. If a negative value were encoded
+using a normal [varint](#variable-width-integers), it would be treated as an
+extremely large unsigned value. Using zigzag encoding allows for a smaller
+number of active bits in the value, leading to a smaller encoding. Below is the
+basic computation for generating a zigzag encoding:
+
+```
+(value << 1) ^ (value >> 63)
+```
+
 #### Strings
 
 Strings are blobs of characters with an associated length.
 
 ### Sections
 
 ```
 section {

mlir/include/mlir/Bytecode/BytecodeImplementation.h

public:
 
   /// Read a reference to the given attribute.
   virtual LogicalResult readAttribute(Attribute &result) = 0;
   template <typename T>
   LogicalResult readAttributes(SmallVectorImpl<T> &attrs) {
     return readList(attrs, [this](T &attr) { return readAttribute(attr); });
   }
   template <typename T>
-  LogicalResult parseAttribute(T &result) {
+  LogicalResult readAttribute(T &result) {
     Attribute baseResult;
-    if (failed(parseAttribute(baseResult)))
+    if (failed(readAttribute(baseResult)))
       return failure();
     if ((result = baseResult.dyn_cast<T>()))
       return success();
-    return emitError() << "expected attribute of type: "
-                       << llvm::getTypeName<T>() << ", but got: " << baseResult;
+    return emitError() << "expected " << llvm::getTypeName<T>()
+                       << ", but got: " << baseResult;
   }
 
   /// Read a reference to the given type.
   virtual LogicalResult readType(Type &result) = 0;
   template <typename T>
   LogicalResult readTypes(SmallVectorImpl<T> &types) {
     return readList(types, [this](T &type) { return readType(type); });
   }
+  template <typename T>
+  LogicalResult readType(T &result) {
+    Type baseResult;
+    if (failed(readType(baseResult)))
+      return failure();
+    if ((result = baseResult.dyn_cast<T>()))
+      return success();
+    return emitError() << "expected " << llvm::getTypeName<T>()
+                       << ", but got: " << baseResult;
+  }
 
   //===--------------------------------------------------------------------===//
   // Primitives
   //===--------------------------------------------------------------------===//
 
   /// Read a variable width integer.
-  // TODO: Add a signed variant when necessary.
   virtual LogicalResult readVarInt(uint64_t &result) = 0;
 
+  /// Read a signed variable width integer.
+  virtual LogicalResult readSignedVarInt(int64_t &result) = 0;
+
+  /// Read an APInt that is known to have been encoded with the given width.
+  virtual FailureOr<APInt> readAPIntWithKnownWidth(unsigned bitWidth) = 0;
+
+  /// 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;
 };
 
 //===----------------------------------------------------------------------===//
 // DialectBytecodeWriter
 //===----------------------------------------------------------------------===//
 
void writeTypes(ArrayRef<T> types) {
     writeList(types, [this](T type) { writeType(type); });
   }
 
   //===--------------------------------------------------------------------===//
   // Primitives
   //===--------------------------------------------------------------------===//
 
   /// Write a variable width integer to the output stream. This should be the
   /// preferred method for emitting integers whenever possible.
-  // TODO: Add a signed variant when necessary.
   virtual void writeVarInt(uint64_t value) = 0;
 
+  /// Write a signed variable width integer to the output stream. This should be
+  /// the preferred method for emitting signed integers whenever possible.
+  virtual void writeSignedVarInt(int64_t value) = 0;
+
+  /// 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
+  /// method should generally only be invoked if you need an APInt, otherwise
+  /// use the varint methods above. APInt values are generally encoded using
+  /// zigzag encoding, to enable more efficient encodings for negative values.
+  virtual void writeAPIntWithKnownWidth(const APInt &value) = 0;
+
+  /// Write an APFloat to the bytecode stream whose semantics will be known
+  /// externally at read time. This method is useful for encoding APFloat values
+  /// 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;
 };
 
 //===----------------------------------------------------------------------===//

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

LogicalResult parseVarInt(uint64_t &result) {
     // 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));
     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)))
+      return failure();
+    // Essentially (but using unsigned): (x >> 1) ^ -(x & 1)
+    result = (result >> 1) ^ (~(result & 1) + 1);
+    return success();
+  }
+
   /// Parse a variable length encoded integer whose low bit is used to encode an
   /// unrelated flag, i.e: `(integerValue << 1) | (flag ? 1 : 0)`.
   LogicalResult parseVarIntWithFlag(uint64_t &result, bool &flag) {
     if (failed(parseVarInt(result)))
       return failure();
     flag = result & 1;
     result >>= 1;
     return success();
public:
   //===--------------------------------------------------------------------===//
   // Primitives
   //===--------------------------------------------------------------------===//
 
   LogicalResult readVarInt(uint64_t &result) override {
     return reader.parseVarInt(result);
   }
 
+  LogicalResult readSignedVarInt(int64_t &result) override {
+    uint64_t unsignedResult;
+    if (failed(reader.parseSignedVarInt(unsignedResult)))
+      return failure();
+    result = static_cast<int64_t>(unsignedResult);
+    return success();
+  }
+
+  FailureOr<APInt> readAPIntWithKnownWidth(unsigned bitWidth) override {
+    // Small values are encoded using a single byte.
+    if (bitWidth <= 8) {
+      uint8_t value;
+      if (failed(reader.parseByte(value)))
+        return failure();
+      return APInt(bitWidth, value);
+    }
+
+    // Large values up to 64 bits are encoded using a single varint.
+    if (bitWidth <= 64) {
+      uint64_t value;
+      if (failed(reader.parseSignedVarInt(value)))
+        return failure();
+      return APInt(bitWidth, value);
+    }
+
+    // Otherwise, for really big values we encode the array of active words in
+    // the value.
+    uint64_t numActiveWords;
+    if (failed(reader.parseVarInt(numActiveWords)))
+      return failure();
+    SmallVector<uint64_t, 4> words(numActiveWords);
+    for (uint64_t i = 0; i < numActiveWords; ++i)
+      if (failed(reader.parseSignedVarInt(words[i])))
+        return failure();
+    return APInt(bitWidth, words);
+  }
+
+  FailureOr<APFloat>
+  readAPFloatWithKnownSemantics(const llvm::fltSemantics &semantics) override {
+    FailureOr<APInt> intVal =
+        readAPIntWithKnownWidth(APFloat::getSizeInBits(semantics));
+    if (failed(intVal))
+      return failure();
+    return APFloat(semantics, *intVal);
+  }
+
   LogicalResult readString(StringRef &result) override {
     return stringReader.parseString(reader, result);
   }
 
 private:
   AttrTypeReader &attrTypeReader;
   StringSectionReader &stringReader;
   EncodingReader &reader;

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

public:
   void emitVarInt(uint64_t value) {
     // In the most common case, the value can be represented in a single byte.
     // Given how hot this case is, explicitly handle that here.
     if ((value >> 7) == 0)
       return emitByte((value << 1) | 0x1);
     emitMultiByteVarInt(value);
   }
 
+  /// Emit a signed variable length integer. Signed varints are encoded using
+  /// a varint with zigzag encoding, meaning that we use the low bit of the
+  /// value to indicate the sign of the value. This allows for more efficient
+  /// encoding of negative values by limiting the number of active bits
+  void emitSignedVarInt(uint64_t value) {
+    emitVarInt((value << 1) ^ (uint64_t)((int64_t)value >> 63));
+  }
+
   /// Emit a variable length integer whose low bit is used to encode the
   /// provided flag, i.e. encoded as: (value << 1) | (flag ? 1 : 0).
   void emitVarIntWithFlag(uint64_t value, bool flag) {
     emitVarInt((value << 1) | (flag ? 1 : 0));
   }
 
   //===--------------------------------------------------------------------===//
   // String Emission
public:
   }
 
   //===--------------------------------------------------------------------===//
   // Primitives
   //===--------------------------------------------------------------------===//
 
   void writeVarInt(uint64_t value) override { emitter.emitVarInt(value); }
 
+  void writeSignedVarInt(int64_t value) override {
+    emitter.emitSignedVarInt(value);
+  }
+
+  void writeAPIntWithKnownWidth(const APInt &value) override {
+    size_t bitWidth = value.getBitWidth();
+
+    // If the value is a single byte, just emit it directly without going
+    // through a varint.
+    if (bitWidth <= 8)
+      return emitter.emitByte(value.getLimitedValue());
+
+    // If the value fits within a single varint, emit it directly.
+    if (bitWidth <= 64)
+      return emitter.emitSignedVarInt(value.getLimitedValue());
+
+    // Otherwise, we need to encode a variable number of active words. We use
+    // active words instead of the number of total words under the observation
+    // that smaller values will be more common.
+    unsigned numActiveWords = value.getActiveWords();
+    emitter.emitVarInt(numActiveWords);
+
+    const uint64_t *rawValueData = value.getRawData();
+    for (unsigned i = 0; i < numActiveWords; ++i)
+      emitter.emitSignedVarInt(rawValueData[i]);
+  }
+
+  void writeAPFloatWithKnownSemantics(const APFloat &value) override {
+    writeAPIntWithKnownWidth(value.bitcastToAPInt());
+  }
+
   void writeOwnedString(StringRef str) override {
     emitter.emitVarInt(stringSection.insert(str));
   }
 
 private:
   EncodingEmitter &emitter;
   IRNumberingState &numberingState;
   StringSectionBuilder &stringSection;

mlir/lib/Bytecode/Writer/IRNumbering.cpp

 struct IRNumberingState::NumberingDialectWriter : public DialectBytecodeWriter {
   NumberingDialectWriter(IRNumberingState &state) : state(state) {}
 
   void writeAttribute(Attribute attr) override { state.number(attr); }
   void writeType(Type type) override { state.number(type); }
 
   /// Stubbed out methods that are not used for numbering.
   void writeVarInt(uint64_t) override {}
+  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.
   }
 
   /// The parent numbering state that is populated by this writer.
   IRNumberingState &state;

mlir/lib/IR/BuiltinDialectBytecode.cpp

enum AttributeCode {
   kArrayAttr = 0,
 
   ///   DictionaryAttr {
   ///     attrs: <StringAttr, Attribute>[]
   ///   }
   kDictionaryAttr = 1,
 
   ///   StringAttr {
-  ///     string
+  ///     value: string
   ///   }
   kStringAttr = 2,
+
+  ///   StringAttrWithType {
+  ///     value: string,
+  ///     type: Type
+  ///   }
+  /// A variant of StringAttr with a type.
+  kStringAttrWithType = 3,
+
+  ///   FlatSymbolRefAttr {
+  ///     rootReference: StringAttr
+  ///   }
+  /// A variant of SymbolRefAttr with no leaf references.
+  kFlatSymbolRefAttr = 4,
+
+  ///   SymbolRefAttr {
+  ///     rootReference: StringAttr,
+  ///     leafReferences: FlatSymbolRefAttr[]
+  ///   }
+  kSymbolRefAttr = 5,
+
+  ///   TypeAttr {
+  ///     value: Type
+  ///   }
+  kTypeAttr = 6,
+
+  ///   UnitAttr {
+  ///   }
+  kUnitAttr = 7,
+
+  ///   IntegerAttr {
+  ///     type: Type
+  ///     value: APInt,
+  ///   }
+  kIntegerAttr = 8,
+
+  ///   FloatAttr {
+  ///     type: FloatType
+  ///     value: APFloat
+  ///   }
+  kFloatAttr = 9,
 };
 
 /// 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 {
BuiltinDialectBytecodeInterface(Dialect *dialect)
       : BytecodeDialectInterface(dialect) {}
 
   //===--------------------------------------------------------------------===//
   // Attributes
 
   Attribute readAttribute(DialectBytecodeReader &reader) const override;
   ArrayAttr readArrayAttr(DialectBytecodeReader &reader) const;
   DictionaryAttr readDictionaryAttr(DialectBytecodeReader &reader) const;
-  StringAttr readStringAttr(DialectBytecodeReader &reader) const;
+  FloatAttr readFloatAttr(DialectBytecodeReader &reader) const;
+  IntegerAttr readIntegerAttr(DialectBytecodeReader &reader) const;
+  StringAttr readStringAttr(DialectBytecodeReader &reader, bool hasType) const;
+  SymbolRefAttr readSymbolRefAttr(DialectBytecodeReader &reader,
+                                  bool hasNestedRefs) const;
+  TypeAttr readTypeAttr(DialectBytecodeReader &reader) const;
 
   LogicalResult writeAttribute(Attribute attr,
                                DialectBytecodeWriter &writer) const override;
   void write(ArrayAttr 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(StringAttr attr, DialectBytecodeWriter &writer) const;
+  void write(SymbolRefAttr attr, DialectBytecodeWriter &writer) const;
+  void write(TypeAttr attr, DialectBytecodeWriter &writer) const;
 
   //===--------------------------------------------------------------------===//
   // Types
 
   Type readType(DialectBytecodeReader &reader) const override;
   IntegerType readIntegerType(DialectBytecodeReader &reader) const;
   FunctionType readFunctionType(DialectBytecodeReader &reader) const;
 
Attribute BuiltinDialectBytecodeInterface::readAttribute(
   if (failed(reader.readVarInt(code)))
     return Attribute();
   switch (code) {
   case builtin_encoding::kArrayAttr:
     return readArrayAttr(reader);
   case builtin_encoding::kDictionaryAttr:
     return readDictionaryAttr(reader);
   case builtin_encoding::kStringAttr:
-    return readStringAttr(reader);
+    return readStringAttr(reader, /*hasType=*/false);
+  case builtin_encoding::kStringAttrWithType:
+    return readStringAttr(reader, /*hasType=*/true);
+  case builtin_encoding::kFlatSymbolRefAttr:
+    return readSymbolRefAttr(reader, /*hasNestedRefs=*/false);
+  case builtin_encoding::kSymbolRefAttr:
+    return readSymbolRefAttr(reader, /*hasNestedRefs=*/true);
+  case builtin_encoding::kTypeAttr:
+    return readTypeAttr(reader);
+  case builtin_encoding::kUnitAttr:
+    return UnitAttr::get(getContext());
+  case builtin_encoding::kIntegerAttr:
+    return readIntegerAttr(reader);
+  case builtin_encoding::kFloatAttr:
+    return readFloatAttr(reader);
   default:
     reader.emitError() << "unknown builtin attribute code: " << code;
     return Attribute();
   }
 }
 
 ArrayAttr BuiltinDialectBytecodeInterface::readArrayAttr(
     DialectBytecodeReader &reader) const {
auto readNamedAttr = [&]() -> FailureOr<NamedAttribute> {
     return NamedAttribute(name, value);
   };
   SmallVector<NamedAttribute> attrs;
   if (failed(reader.readList(attrs, readNamedAttr)))
     return DictionaryAttr();
   return DictionaryAttr::get(getContext(), attrs);
 }
 
-StringAttr BuiltinDialectBytecodeInterface::readStringAttr(
+FloatAttr BuiltinDialectBytecodeInterface::readFloatAttr(
+    DialectBytecodeReader &reader) const {
+  FloatType type;
+  if (failed(reader.readType(type)))
+    return FloatAttr();
+  FailureOr<APFloat> value =
+      reader.readAPFloatWithKnownSemantics(type.getFloatSemantics());
+  if (failed(value))
+    return FloatAttr();
+  return FloatAttr::get(type, *value);
+}
+
+IntegerAttr BuiltinDialectBytecodeInterface::readIntegerAttr(
     DialectBytecodeReader &reader) const {
+  Type type;
+  if (failed(reader.readType(type)))
+    return IntegerAttr();
+
+  // Extract the value storage width from the type.
+  unsigned bitWidth;
+  if (auto intType = type.dyn_cast<IntegerType>()) {
+    bitWidth = intType.getWidth();
+  } else if (type.isa<IndexType>()) {
+    bitWidth = IndexType::kInternalStorageBitWidth;
+  } else {
+    reader.emitError()
+        << "expected integer or index type for IntegerAttr, but got: " << type;
+    return IntegerAttr();
+  }
+
+  FailureOr<APInt> value = reader.readAPIntWithKnownWidth(bitWidth);
+  if (failed(value))
+    return IntegerAttr();
+  return IntegerAttr::get(type, *value);
+}
+
+StringAttr
+BuiltinDialectBytecodeInterface::readStringAttr(DialectBytecodeReader &reader,
+                                                bool hasType) const {
   StringRef string;
   if (failed(reader.readString(string)))
     return StringAttr();
-  return StringAttr::get(getContext(), string);
+
+  // Read the type if present.
+  Type type;
+  if (!hasType)
+    type = NoneType::get(getContext());
+  else if (failed(reader.readType(type)))
+    return StringAttr();
+  return StringAttr::get(string, type);
+}
+
+SymbolRefAttr BuiltinDialectBytecodeInterface::readSymbolRefAttr(
+    DialectBytecodeReader &reader, bool hasNestedRefs) const {
+  StringAttr rootReference;
+  if (failed(reader.readAttribute(rootReference)))
+    return SymbolRefAttr();
+  SmallVector<FlatSymbolRefAttr> nestedReferences;
+  if (hasNestedRefs && failed(reader.readAttributes(nestedReferences)))
+    return SymbolRefAttr();
+  return SymbolRefAttr::get(rootReference, nestedReferences);
+}
+
+TypeAttr BuiltinDialectBytecodeInterface::readTypeAttr(
+    DialectBytecodeReader &reader) const {
+  Type type;
+  if (failed(reader.readType(type)))
+    return TypeAttr();
+  return TypeAttr::get(type);
 }
 
 //===----------------------------------------------------------------------===//
 // Attributes: Writer
 
 LogicalResult BuiltinDialectBytecodeInterface::writeAttribute(
     Attribute attr, DialectBytecodeWriter &writer) const {
   return TypeSwitch<Attribute, LogicalResult>(attr)
-      .Case<ArrayAttr, DictionaryAttr, StringAttr>([&](auto attr) {
+      .Case<ArrayAttr, DictionaryAttr, FloatAttr, IntegerAttr, StringAttr,
+            SymbolRefAttr, TypeAttr>([&](auto attr) {
         write(attr, writer);
         return success();
       })
+      .Case([&](UnitAttr) {
+        writer.writeVarInt(builtin_encoding::kUnitAttr);
+        return success();
+      })
       .Default([&](Attribute) { return failure(); });
 }
 
 void BuiltinDialectBytecodeInterface::write(
     ArrayAttr attr, DialectBytecodeWriter &writer) const {
   writer.writeVarInt(builtin_encoding::kArrayAttr);
   writer.writeAttributes(attr.getValue());
 }
 
 void BuiltinDialectBytecodeInterface::write(
     DictionaryAttr attr, DialectBytecodeWriter &writer) const {
   writer.writeVarInt(builtin_encoding::kDictionaryAttr);
   writer.writeList(attr.getValue(), [&](NamedAttribute attr) {
     writer.writeAttribute(attr.getName());
     writer.writeAttribute(attr.getValue());
   });
 }
 
 void BuiltinDialectBytecodeInterface::write(
+    FloatAttr attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kFloatAttr);
+  writer.writeType(attr.getType());
+  writer.writeAPFloatWithKnownSemantics(attr.getValue());
+}
+
+void BuiltinDialectBytecodeInterface::write(
+    IntegerAttr attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kIntegerAttr);
+  writer.writeType(attr.getType());
+  writer.writeAPIntWithKnownWidth(attr.getValue());
+}
+
+void BuiltinDialectBytecodeInterface::write(
     StringAttr attr, DialectBytecodeWriter &writer) const {
+  // We only encode the type if it isn't NoneType, which is significantly less
+  // common.
+  Type type = attr.getType();
+  if (!type.isa<NoneType>()) {
+    writer.writeVarInt(builtin_encoding::kStringAttrWithType);
+    writer.writeOwnedString(attr.getValue());
+    writer.writeType(type);
+    return;
+  }
   writer.writeVarInt(builtin_encoding::kStringAttr);
   writer.writeOwnedString(attr.getValue());
 }
 
+void BuiltinDialectBytecodeInterface::write(
+    SymbolRefAttr attr, DialectBytecodeWriter &writer) const {
+  ArrayRef<FlatSymbolRefAttr> nestedRefs = attr.getNestedReferences();
+  writer.writeVarInt(nestedRefs.empty() ? builtin_encoding::kFlatSymbolRefAttr
+                                        : builtin_encoding::kSymbolRefAttr);
+
+  writer.writeAttribute(attr.getRootReference());
+  if (!nestedRefs.empty())
+    writer.writeAttributes(nestedRefs);
+}
+
+void BuiltinDialectBytecodeInterface::write(
+    TypeAttr attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kTypeAttr);
+  writer.writeType(attr.getValue());
+}
+
 //===----------------------------------------------------------------------===//
 // Types: Reader
 
 Type BuiltinDialectBytecodeInterface::readType(
     DialectBytecodeReader &reader) const {
   uint64_t code;
   if (failed(reader.readVarInt(code)))
     return Type();

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

 // RUN: mlir-opt -emit-bytecode %s | mlir-opt | 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]
 } {}
 
+//===----------------------------------------------------------------------===//
+// 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
+  // CHECK: bytecode.float3 = -5.000000e-01 : bf16
+  bytecode.float = 10.0 : f64,
+  bytecode.float1 = 0.1 : f80,
+  bytecode.float2 = 0.1 : f128,
+  bytecode.float3 = -0.5 : bf16
+} {}
+
+//===----------------------------------------------------------------------===//
+// IntegerAttr
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @TestInt
+module @TestInt attributes {
+  // CHECK: bytecode.int = false
+  // CHECK: bytecode.int1 = -1 : i8
+  // CHECK: bytecode.int2 = 800 : ui64
+  // CHECK: bytecode.int3 = 90000000000000000300000000000000000001 : i128
+  bytecode.int = false,
+  bytecode.int1 = -1 : i8,
+  bytecode.int2 = 800 : ui64,
+  bytecode.int3 = 90000000000000000300000000000000000001 : i128
+} {}
+
+//===----------------------------------------------------------------------===//
+// StringAttr
+//===----------------------------------------------------------------------===//
+
 // CHECK-LABEL: @TestString
 module @TestString attributes {
   // CHECK: bytecode.string = "hello"
-  bytecode.string = "hello"
+  // CHECK: bytecode.string2 = "hello" : i32
+  bytecode.string = "hello",
+  bytecode.string2 = "hello" : i32
+} {}
+
+//===----------------------------------------------------------------------===//
+// SymbolRefAttr
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @TestSymbolRef
+module @TestSymbolRef attributes {
+  // CHECK: bytecode.ref = @foo
+  // CHECK: bytecode.ref2 = @foo::@bar::@foo
+  bytecode.ref = @foo,
+  bytecode.ref2 = @foo::@bar::@foo
+} {}
+
+//===----------------------------------------------------------------------===//
+// TypeAttr
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @TestType
+module @TestType attributes {
+  // CHECK: bytecode.type = i178
+  bytecode.type = i178
 } {}