[Mlir-commits] [mlir] [mlir][bytecode] Add option to elide locations during serialization (PR #201183)
Peter Hawkins
llvmlistbot at llvm.org
Tue Jun 2 14:51:11 PDT 2026
https://github.com/hawkinsp updated https://github.com/llvm/llvm-project/pull/201183
>From 64bb54d834b0e58415e2291e800c34f134f21a87 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins at google.com>
Date: Tue, 2 Jun 2026 21:40:12 +0000
Subject: [PATCH] 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
---
mlir/include/mlir/Bytecode/BytecodeWriter.h | 6 ++
mlir/lib/Bytecode/Writer/BytecodeWriter.cpp | 11 ++
mlir/lib/Bytecode/Writer/IRNumbering.cpp | 16 +++
mlir/lib/Bytecode/Writer/IRNumbering.h | 3 +
mlir/unittests/Bytecode/BytecodeTest.cpp | 109 ++++++++++++++++++++
5 files changed, 145 insertions(+)
diff --git a/mlir/include/mlir/Bytecode/BytecodeWriter.h b/mlir/include/mlir/Bytecode/BytecodeWriter.h
index c6cff0bc81314..8b883e32831a2 100644
--- a/mlir/include/mlir/Bytecode/BytecodeWriter.h
+++ b/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;
diff --git a/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp b/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp
index e32f8730f7d71..90a1de6691d99 100644
--- a/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp
+++ b/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;
}
diff --git a/mlir/lib/Bytecode/Writer/IRNumbering.cpp b/mlir/lib/Bytecode/Writer/IRNumbering.cpp
index d10f64494d22e..04625628fa5a6 100644
--- a/mlir/lib/Bytecode/Writer/IRNumbering.cpp
+++ b/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) {
diff --git a/mlir/lib/Bytecode/Writer/IRNumbering.h b/mlir/lib/Bytecode/Writer/IRNumbering.h
index 9b7ac0d3688e3..f236e4a31c40d 100644
--- a/mlir/lib/Bytecode/Writer/IRNumbering.h
+++ b/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);
diff --git a/mlir/unittests/Bytecode/BytecodeTest.cpp b/mlir/unittests/Bytecode/BytecodeTest.cpp
index 51eebc488d7f8..d8f6510924698 100644
--- a/mlir/unittests/Bytecode/BytecodeTest.cpp
+++ b/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");
+}
More information about the Mlir-commits
mailing list