[mlir][bytecode] Add option to elide locations during serialization (#201183)

Adds a setElideLocations option to BytecodeWriterConfig to elide
locations during bytecode serialization. When enabled, all LocationAttrs
are mapped to UnknownLoc during numbering and writing to produce
location-invariant bytecode (e.g., for stable fingerprinting).

Another way to achieve the same thing would be to apply the
strip-debuginfo pass,
but that requires mutating the module, which in turn requires cloning
the module if one still requires the unstripped original.

Assisted-by: Antigravity / Gemini

Author: hawkinsp

mlir/include/mlir/Bytecode/BytecodeWriter.h

@@ -177,6 +177,12 @@ class BytecodeWriterConfig {
       attachResourcePrinter(std::move(printer));
   }
 
+  /// Set a boolean flag to skip emission of unique locations into the bytecode
+  /// file. When enabled, all locations are mapped to UnknownLoc during
+  /// numbering.
+  void setElideLocations(bool shouldElideLocations = true);
+  bool shouldElideLocations() const;
+
 private:
   /// A pointer to allocated storage for the impl state.
   std::unique_ptr<Impl> impl;

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

@@ -44,6 +44,9 @@ struct BytecodeWriterConfig::Impl {
   /// file.
   bool shouldElideResourceData = false;
 
+  /// A flag specifying whether to elide emission of locations.
+  bool shouldElideLocations = false;
+
   /// A map containing dialect version information for each dialect to emit.
   llvm::StringMap<std::unique_ptr<DialectVersion>> dialectVersionMap;
 
@@ -102,6 +105,14 @@ void BytecodeWriterConfig::setElideResourceDataFlag(
   impl->shouldElideResourceData = shouldElideResourceData;
 }
 
+void BytecodeWriterConfig::setElideLocations(bool shouldElideLocations) {
+  impl->shouldElideLocations = shouldElideLocations;
+}
+
+bool BytecodeWriterConfig::shouldElideLocations() const {
+  return impl->shouldElideLocations;
+}
+
 void BytecodeWriterConfig::setDesiredBytecodeVersion(int64_t bytecodeVersion) {
   impl->bytecodeVersion = bytecodeVersion;
 }

mlir/lib/Bytecode/Writer/IRNumbering.cpp

@@ -13,6 +13,7 @@
 #include "mlir/Bytecode/Encoding.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Location.h"
 #include "mlir/IR/OpDefinition.h"
 
 using namespace mlir;
@@ -198,6 +199,13 @@ IRNumberingState::IRNumberingState(Operation *op,
   finalizeDialectResourceNumberings(op);
 }
 
+unsigned IRNumberingState::getNumber(Location loc) {
+  if (config.shouldElideLocations()) {
+    return getNumber(Attribute(UnknownLoc::get(loc.getContext())));
+  }
+  return getNumber(Attribute(loc));
+}
+
 void IRNumberingState::computeGlobalNumberingState(Operation *rootOp) {
   // A simple state struct tracking data used when walking operations.
   struct StackState {
@@ -308,6 +316,14 @@ void IRNumberingState::computeGlobalNumberingState(Operation *rootOp) {
   });
 }
 
+void IRNumberingState::number(Location loc) {
+  if (config.shouldElideLocations()) {
+    number(Attribute(UnknownLoc::get(loc.getContext())));
+  } else {
+    number(Attribute(loc));
+  }
+}
+
 void IRNumberingState::number(Attribute attr) {
   auto it = attrs.try_emplace(attr);
   if (!it.second) {

mlir/lib/Bytecode/Writer/IRNumbering.h

@@ -14,6 +14,7 @@
 #ifndef LIB_MLIR_BYTECODE_WRITER_IRNUMBERING_H
 #define LIB_MLIR_BYTECODE_WRITER_IRNUMBERING_H
 
+#include "mlir/IR/Location.h"
 #include "mlir/IR/OpImplementation.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SetVector.h"
@@ -165,6 +166,7 @@ class IRNumberingState {
     assert(attrs.count(attr) && "attribute not numbered");
     return attrs[attr]->number;
   }
+  unsigned getNumber(Location loc);
   unsigned getNumber(Block *block) {
     assert(blockIDs.count(block) && "block not numbered");
     return blockIDs[block];
@@ -221,6 +223,7 @@ class IRNumberingState {
 
   /// Number the given IR unit for bytecode emission.
   void number(Attribute attr);
+  void number(Location loc);
   void number(Block &block);
   DialectNumbering &numberDialect(Dialect *dialect);
   DialectNumbering &numberDialect(StringRef dialect);

mlir/unittests/Bytecode/BytecodeTest.cpp

@@ -292,3 +292,112 @@ TEST(Bytecode, EmptyFusedLocRoundtrip) {
 
   module.erase();
 }
+
+TEST(Bytecode, LocationElision) {
+  MLIRContext context;
+  context.allowUnregisteredDialects();
+  ParserConfig config(&context);
+
+  // Module 1: Reuses the same location "foo" everywhere.
+  StringRef ir1 = R"mlir(
+    module @Test {
+      "test.op"() : () -> () loc("foo")
+    } loc("foo")
+  )mlir";
+
+  // Module 2: Uses unique locations everywhere.
+  StringRef ir2 = R"mlir(
+    module @Test {
+      "test.op"() : () -> () loc("a")
+    } loc("b")
+  )mlir";
+
+  OwningOpRef<Operation *> op1 = parseSourceString(ir1, config);
+  OwningOpRef<Operation *> op2 = parseSourceString(ir2, config);
+  ASSERT_TRUE(op1);
+  ASSERT_TRUE(op2);
+
+  // Serialize both with location elision enabled.
+  BytecodeWriterConfig writerConfig;
+  writerConfig.setElideLocations(true);
+
+  std::string bytecode1;
+  {
+    llvm::raw_string_ostream os(bytecode1);
+    ASSERT_TRUE(succeeded(writeBytecodeToFile(op1.get(), os, writerConfig)));
+  }
+
+  std::string bytecode2;
+  {
+    llvm::raw_string_ostream os(bytecode2);
+    ASSERT_TRUE(succeeded(writeBytecodeToFile(op2.get(), os, writerConfig)));
+  }
+
+  // If location elision is working correctly, both modules must produce
+  // the EXACT same bytecode representation, because all locations (shared or
+  // unique) will have been collapsed into a single shared UnknownLoc.
+  EXPECT_EQ(bytecode1, bytecode2);
+}
+
+TEST(Bytecode, LocationElisionPreservesAttributes) {
+  MLIRContext context;
+  context.allowUnregisteredDialects();
+  ParserConfig config(&context);
+
+  // An operation with a debug location ("elide_me") AND a semantic attribute
+  // that is a LocationAttr ("preserve_me").
+  StringRef ir = R"mlir(
+    module @Test {
+      "test.op"() {some_loc_attr = loc("preserve_me")} : () -> () loc("elide_me")
+    } loc("elide_me")
+  )mlir";
+
+  OwningOpRef<Operation *> op = parseSourceString(ir, config);
+  ASSERT_TRUE(op);
+
+  // Serialize with location elision enabled.
+  BytecodeWriterConfig writerConfig;
+  writerConfig.setElideLocations(true);
+
+  std::string bytecode;
+  {
+    llvm::raw_string_ostream os(bytecode);
+    ASSERT_TRUE(succeeded(writeBytecodeToFile(op.get(), os, writerConfig)));
+  }
+
+  // Parse it back using the bytecode reader.
+  std::unique_ptr<Block> block = std::make_unique<Block>();
+  ASSERT_TRUE(succeeded(readBytecodeFile(
+      llvm::MemoryBufferRef(bytecode, "string-buffer"), block.get(), config)));
+
+  // Verify we got the roundtripped module.
+  ASSERT_FALSE(block->empty());
+  Operation *roundTrippedModule = &block->front();
+  ASSERT_TRUE(roundTrippedModule);
+
+  // Find the inner "test.op" operation.
+  Operation *innerOp = nullptr;
+  roundTrippedModule->walk([&](Operation *op) {
+    if (op->getName().getStringRef() == "test.op") {
+      innerOp = op;
+    }
+  });
+  ASSERT_TRUE(innerOp);
+
+  // 1. Verify that the debug location of "test.op" WAS elided (became
+  // UnknownLoc).
+  EXPECT_TRUE(isa<UnknownLoc>(innerOp->getLoc()));
+
+  // 2. Verify that the semantic location attribute WAS PRESERVED.
+  Attribute semanticLocAttr = innerOp->getAttr("some_loc_attr");
+  ASSERT_TRUE(semanticLocAttr);
+  auto locAttr = dyn_cast<LocationAttr>(semanticLocAttr);
+  ASSERT_TRUE(locAttr);
+
+  // It should still be loc("preserve_me"), not UnknownLoc.
+  EXPECT_FALSE(isa<UnknownLoc>(locAttr));
+
+  auto nameLoc = dyn_cast<NameLoc>(locAttr);
+  ASSERT_TRUE(nameLoc);
+  EXPECT_EQ(nameLoc.getName().getValue(), "preserve_me");
+}