[mlir:Bytecode] Add encoding support for builtin location attributes

[rriddle]

Summary

This provides a significantly more efficient encoding for locations.

[mehdi_amini]

mlir/lib/IR/BuiltinDialectBytecode.cpp

405

I wonder why you split the writer/reader?
I felt it may be easier to have for each attribute/type/loc the reader and writer method side-by-side: after all they are supposed to be symmetric.

[rriddle]

mlir/lib/IR/BuiltinDialectBytecode.cpp

405

It was just easier for me to reason about in the beginning. I don't have a huge preference and could move them next to each other, let me know if you want that I'll move them in a follow up.

[mehdi_amini]

mlir/lib/IR/BuiltinDialectBytecode.cpp

405

Take it as a suggestion for a future revision :)

[jpienaar]

mlir/lib/IR/BuiltinDialectBytecode.cpp

379

Should there be an error emitted about failure here? (I don't know where the errors are reported from or how we ensure at least an error is reported)

[rriddle]

mlir/lib/IR/BuiltinDialectBytecode.cpp

379

All of the API provided by the DialectBytecodeReader/DialectBytecodeWriter call into the main bytecode API, which will emit errors on failure (so we don't emit here to avoid double-emit).

[rriddle]

mlir/lib/IR/BuiltinDialectBytecode.cpp

405

Done in D132583

mlir/lib/IR/BuiltinDialectBytecode.cpp

enum AttributeCode {
   ///   }
   kIntegerAttr = 8,
 
   ///   FloatAttr {
   ///     type: FloatType
   ///     value: APFloat
   ///   }
   kFloatAttr = 9,
+
+  ///   CallSiteLoc {
+  ///    callee: LocationAttr,
+  ///    caller: LocationAttr
+  ///   }
+  kCallSiteLoc = 10,
+
+  ///   FileLineColLoc {
+  ///     file: StringAttr,
+  ///     line: varint,
+  ///     column: varint
+  ///   }
+  kFileLineColLoc = 11,
+
+  ///   FusedLoc {
+  ///     locations: LocationAttr[]
+  ///   }
+  kFusedLoc = 12,
+
+  ///   FusedLocWithMetadata {
+  ///     locations: LocationAttr[],
+  ///     metadata: Attribute
+  ///   }
+  /// A variant of FusedLoc with metadata.
+  kFusedLocWithMetadata = 13,
+
+  ///   NameLoc {
+  ///     name: StringAttr,
+  ///     childLoc: LocationAttr
+  ///   }
+  kNameLoc = 14,
+
+  ///   UnknownLoc {
+  ///   }
+  kUnknownLoc = 15,
 };
 
 /// 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 {
   DictionaryAttr readDictionaryAttr(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;
 
+  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(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;
 
+  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;
+  LogicalResult write(OpaqueLoc attr, DialectBytecodeWriter &writer) const;
+
   //===--------------------------------------------------------------------===//
   // Types
 
   Type readType(DialectBytecodeReader &reader) const override;
   IntegerType readIntegerType(DialectBytecodeReader &reader) const;
   FunctionType readFunctionType(DialectBytecodeReader &reader) const;
 
   LogicalResult writeType(Type type,
Attribute BuiltinDialectBytecodeInterface::readAttribute(
   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);
+  case builtin_encoding::kCallSiteLoc:
+    return readCallSiteLoc(reader);
+  case builtin_encoding::kFileLineColLoc:
+    return readFileLineColLoc(reader);
+  case builtin_encoding::kFusedLoc:
+    return readFusedLoc(reader, /*hasMetadata=*/false);
+  case builtin_encoding::kFusedLocWithMetadata:
+    return readFusedLoc(reader, /*hasMetadata=*/true);
+  case builtin_encoding::kNameLoc:
+    return readNameLoc(reader);
+  case builtin_encoding::kUnknownLoc:
+    return UnknownLoc::get(getContext());
   default:
     reader.emitError() << "unknown builtin attribute code: " << code;
     return Attribute();
   }
 }
 
 ArrayAttr BuiltinDialectBytecodeInterface::readArrayAttr(
     DialectBytecodeReader &reader) const {

 TypeAttr BuiltinDialectBytecodeInterface::readTypeAttr(
     DialectBytecodeReader &reader) const {
   Type type;
   if (failed(reader.readType(type)))
     return TypeAttr();
   return TypeAttr::get(type);
 }
 
+LocationAttr BuiltinDialectBytecodeInterface::readCallSiteLoc(
+    DialectBytecodeReader &reader) const {
+  LocationAttr callee, caller;
+  if (failed(reader.readAttribute(callee)) ||
+      failed(reader.readAttribute(caller)))
+    return LocationAttr();
+  return CallSiteLoc::get(callee, caller);
+}
+
+LocationAttr BuiltinDialectBytecodeInterface::readFileLineColLoc(
+    DialectBytecodeReader &reader) const {
+  StringAttr filename;
+  uint64_t line, column;
+  if (failed(reader.readAttribute(filename)) ||
+      failed(reader.readVarInt(line)) || failed(reader.readVarInt(column)))
+    return LocationAttr();
+  return FileLineColLoc::get(filename, line, column);
+}
+
+LocationAttr
+BuiltinDialectBytecodeInterface::readFusedLoc(DialectBytecodeReader &reader,
+                                              bool hasMetadata) const {
+  // Parse the child locations.
+  auto readLoc = [&]() -> FailureOr<Location> {
+    LocationAttr locAttr;
+    if (failed(reader.readAttribute(locAttr)))
+      return failure();
+    return Location(locAttr);
+  };
+  SmallVector<Location> locations;
+  if (failed(reader.readList(locations, readLoc)))
+    return LocationAttr();
+
+  // Parse the metadata if present.
+  Attribute metadata;
+  if (hasMetadata && failed(reader.readAttribute(metadata)))
+    return LocationAttr();
+
+  return FusedLoc::get(locations, metadata, getContext());
+}
+
+LocationAttr BuiltinDialectBytecodeInterface::readNameLoc(
+    DialectBytecodeReader &reader) const {
+  StringAttr name;
+  LocationAttr childLoc;
+  if (failed(reader.readAttribute(name)) ||
+      failed(reader.readAttribute(childLoc)))
+    return LocationAttr();
+  return NameLoc::get(name, childLoc);
+}
+
 //===----------------------------------------------------------------------===//
 // Attributes: Writer
 
 LogicalResult BuiltinDialectBytecodeInterface::writeAttribute(
     Attribute attr, DialectBytecodeWriter &writer) const {
   return TypeSwitch<Attribute, LogicalResult>(attr)
       .Case<ArrayAttr, DictionaryAttr, IntegerAttr, FloatAttr, StringAttr,
-            SymbolRefAttr, TypeAttr>([&](auto attr) {
+            SymbolRefAttr, TypeAttr, CallSiteLoc, FileLineColLoc, FusedLoc,
+            NameLoc>([&](auto attr) {
         write(attr, writer);
         return success();
       })
+      .Case([&](OpaqueLoc attr) { return write(attr, writer); })
       .Case([&](UnitAttr) {
         writer.writeVarInt(builtin_encoding::kUnitAttr);
         return success();
       })
+      .Case([&](UnknownLoc) {
+        writer.writeVarInt(builtin_encoding::kUnknownLoc);
+        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(
     TypeAttr attr, DialectBytecodeWriter &writer) const {
   writer.writeVarInt(builtin_encoding::kTypeAttr);
   writer.writeType(attr.getValue());
 }
 
+void BuiltinDialectBytecodeInterface::write(
+    CallSiteLoc attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kCallSiteLoc);
+  writer.writeAttribute(attr.getCallee());
+  writer.writeAttribute(attr.getCaller());
+}
+
+void BuiltinDialectBytecodeInterface::write(
+    FileLineColLoc attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kFileLineColLoc);
+  writer.writeAttribute(attr.getFilename());
+  writer.writeVarInt(attr.getLine());
+  writer.writeVarInt(attr.getColumn());
+}
+
+void BuiltinDialectBytecodeInterface::write(
+    FusedLoc attr, DialectBytecodeWriter &writer) const {
+  if (Attribute metadata = attr.getMetadata()) {
+    writer.writeVarInt(builtin_encoding::kFusedLocWithMetadata);
+    writer.writeAttributes(attr.getLocations());
+    writer.writeAttribute(metadata);
+  } else {
+    writer.writeVarInt(builtin_encoding::kFusedLoc);
+    writer.writeAttributes(attr.getLocations());
+  }
+}
+
+void BuiltinDialectBytecodeInterface::write(
+    NameLoc attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kNameLoc);
+  writer.writeAttribute(attr.getName());
+  writer.writeAttribute(attr.getChildLoc());
+}
+
+LogicalResult
+BuiltinDialectBytecodeInterface::write(OpaqueLoc attr,
+                                       DialectBytecodeWriter &writer) const {
+  // We can't encode an OpaqueLoc directly given that it is in-memory only, so
+  // encode the fallback instead.
+  return writeAttribute(attr.getFallbackLocation(), writer);
+}
+
 //===----------------------------------------------------------------------===//
 // 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
+// RUN: mlir-opt -emit-bytecode %s | mlir-opt -mlir-print-local-scope | FileCheck %s
 
 // Bytecode currently does not support big-endian platforms
 // UNSUPPORTED: s390x-
 
 // CHECK-LABEL: @TestArray
 module @TestArray attributes {
   // CHECK: bytecode.array = [unit]
   bytecode.array = [unit]
module @TestInt attributes {
   // 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
 } {}
 
+// CHECK-LABEL: @TestLocCallSite
+module @TestLocCallSite attributes {
+  // CHECK: bytecode.loc = loc(callsite("foo" at "mysource.cc":10:8))
+  bytecode.loc = loc(callsite("foo" at "mysource.cc":10:8))
+} {}
+
+// CHECK-LABEL: @TestLocFileLineCol
+module @TestLocFileLineCol attributes {
+  // CHECK: bytecode.loc = loc("mysource.cc":10:8)
+  bytecode.loc = loc("mysource.cc":10:8)
+} {}
+
+// CHECK-LABEL: @TestLocFused
+module @TestLocFused attributes {
+  // CHECK: bytecode.loc = loc(fused["foo", "mysource.cc":10:8])
+  // CHECK: bytecode.loc2 = loc(fused<"myPass">["foo", "foo2"])
+  bytecode.loc = loc(fused["foo", "mysource.cc":10:8]),
+  bytecode.loc2 = loc(fused<"myPass">["foo", "foo2"])
+} {}
+
+// CHECK-LABEL: @TestLocName
+module @TestLocName attributes {
+  // CHECK: bytecode.loc = loc("foo")
+  // CHECK: bytecode.loc2 = loc("foo"("mysource.cc":10:8))
+  bytecode.loc = loc("foo"),
+  bytecode.loc2 = loc("foo"("mysource.cc":10:8))
+} {}
+
+// CHECK-LABEL: @TestLocUnknown
+module @TestLocUnknown attributes {
+  // CHECK: bytecode.loc = loc(unknown)
+  bytecode.loc = loc(unknown)
+} {}
+
 // CHECK-LABEL: @TestString
 module @TestString attributes {
   // CHECK: bytecode.string = "hello"
   // CHECK: bytecode.string2 = "hello" : i32
   bytecode.string = "hello",
   bytecode.string2 = "hello" : i32
 } {}