[Mlir-commits] [mlir] c83ebf1 - [mlirbc] Serialize dense elements attr i1 using packed (#182233)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Feb 19 00:57:37 PST 2026
Author: Jacques Pienaar
Date: 2026-02-19T08:57:33Z
New Revision: c83ebf19cf3e649ed4d2ccadddd8c7e447850363
URL: https://github.com/llvm/llvm-project/commit/c83ebf19cf3e649ed4d2ccadddd8c7e447850363
DIFF: https://github.com/llvm/llvm-project/commit/c83ebf19cf3e649ed4d2ccadddd8c7e447850363.diff
LOG: [mlirbc] Serialize dense elements attr i1 using packed (#182233)
Extra cost is in serialization layer localized while resulting in
smaller bytecode files, this also keeps the format compatible with what
was previously.
Added:
mlir/test/Bytecode/bytecode_callback_write_unowned_blob.mlir
mlir/test/Bytecode/i1_splat_roundtrip.mlir
Modified:
mlir/include/mlir/Bytecode/BytecodeImplementation.h
mlir/include/mlir/IR/BuiltinDialectBytecode.td
mlir/lib/Bytecode/Writer/BytecodeWriter.cpp
mlir/lib/Bytecode/Writer/IRNumbering.cpp
mlir/lib/IR/BuiltinDialectBytecode.cpp
mlir/test/lib/IR/TestBytecodeRoundtrip.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Bytecode/BytecodeImplementation.h b/mlir/include/mlir/Bytecode/BytecodeImplementation.h
index 0ddc531073e23..49c5b33bd8290 100644
--- a/mlir/include/mlir/Bytecode/BytecodeImplementation.h
+++ b/mlir/include/mlir/Bytecode/BytecodeImplementation.h
@@ -398,6 +398,10 @@ class DialectBytecodeWriter {
/// written as-is, with no additional compression or compaction.
virtual void writeOwnedBlob(ArrayRef<char> blob) = 0;
+ /// Write a blob to the bytecode, which is not owned by the caller. The blob
+ /// is copied into the bytecode, and need not strictly outlive the call.
+ virtual void writeUnownedBlob(ArrayRef<char> blob) = 0;
+
/// Write a bool to the output stream.
virtual void writeOwnedBool(bool value) = 0;
diff --git a/mlir/include/mlir/IR/BuiltinDialectBytecode.td b/mlir/include/mlir/IR/BuiltinDialectBytecode.td
index 0208e8cdbf293..b5ffa3eeb58c4 100644
--- a/mlir/include/mlir/IR/BuiltinDialectBytecode.td
+++ b/mlir/include/mlir/IR/BuiltinDialectBytecode.td
@@ -175,10 +175,11 @@ def DenseArrayAttr : DialectAttribute<(attr
def DenseElementsAttr : WithType<"DenseElementsAttr", Attribute>;
def DenseIntOrFPElementsAttr : DialectAttribute<(attr
ShapedType:$type,
- Blob:$rawData
-)> {
- let cBuilder = "cast<$_resultType>($_resultType::getFromRawBuffer($_args))";
-}
+ WithBuilder<"$_args",
+ WithType<"SmallVector<char>",
+ WithParser<"succeeded(readDenseIntOrFPElementsAttr($_reader, type, $_var))",
+ WithPrinter<"writeDenseIntOrFPElementsAttr($_writer, $_name)">>>>:$rawData
+)>;
def DenseStringElementsAttr : DialectAttribute<(attr
ShapedType:$type,
diff --git a/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp b/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp
index eacb936c763b7..a04e3c3e3f177 100644
--- a/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp
+++ b/mlir/lib/Bytecode/Writer/BytecodeWriter.cpp
@@ -465,6 +465,14 @@ class DialectWriter : public DialectBytecodeWriter {
"dialect blob");
}
+ void writeUnownedBlob(ArrayRef<char> blob) override {
+ emitter.emitVarInt(blob.size(), "dialect blob");
+ emitter.emitBytes(
+ ArrayRef<uint8_t>(reinterpret_cast<const uint8_t *>(blob.data()),
+ blob.size()),
+ "dialect blob");
+ }
+
void writeOwnedBool(bool value) override {
emitter.emitByte(value, "dialect bool");
}
diff --git a/mlir/lib/Bytecode/Writer/IRNumbering.cpp b/mlir/lib/Bytecode/Writer/IRNumbering.cpp
index 8e8e7148ee70e..d10f64494d22e 100644
--- a/mlir/lib/Bytecode/Writer/IRNumbering.cpp
+++ b/mlir/lib/Bytecode/Writer/IRNumbering.cpp
@@ -50,6 +50,7 @@ struct IRNumberingState::NumberingDialectWriter : public DialectBytecodeWriter {
}
void writeOwnedBlob(ArrayRef<char> blob) override {}
void writeOwnedBool(bool value) override {}
+ void writeUnownedBlob(ArrayRef<char> blob) override {}
int64_t getBytecodeVersion() const override {
return state.getDesiredBytecodeVersion();
diff --git a/mlir/lib/IR/BuiltinDialectBytecode.cpp b/mlir/lib/IR/BuiltinDialectBytecode.cpp
index 31aff47ceaa67..a87224a48c32a 100644
--- a/mlir/lib/IR/BuiltinDialectBytecode.cpp
+++ b/mlir/lib/IR/BuiltinDialectBytecode.cpp
@@ -148,6 +148,68 @@ static void writeFileLineColRangeLocs(DialectBytecodeWriter &writer,
writer.writeVarInt(range.getEndColumn());
}
+static LogicalResult
+readDenseIntOrFPElementsAttr(DialectBytecodeReader &reader, ShapedType type,
+ SmallVectorImpl<char> &rawData) {
+ ArrayRef<char> blob;
+ if (failed(reader.readBlob(blob)))
+ return failure();
+
+ // If the type is not i1, just copy the blob.
+ if (!type.getElementType().isInteger(1)) {
+ rawData.append(blob.begin(), blob.end());
+ return success();
+ }
+
+ // Check to see if this is using the packed format.
+ // Note: this could be asserted instead as this should be the case. But we
+ // did have period where the unpacked was being serialized, this enables
+ // consuming those still and the check for which case we are in is pretty
+ // cheap.
+ size_t numElements = type.getNumElements();
+ size_t packedSize = llvm::divideCeil(numElements, 8);
+ if (blob.size() == packedSize && blob.size() != numElements &&
+ blob.size() != 1) {
+ // Unpack the blob.
+ rawData.resize(numElements);
+ for (size_t i = 0; i < numElements; ++i)
+ rawData[i] = (blob[i / 8] & (1 << (i % 8))) ? 0xFF : 0x00;
+ return success();
+ }
+ // Otherwise, fallback to the default behavior.
+ rawData.append(blob.begin(), blob.end());
+ return success();
+}
+
+static void writeDenseIntOrFPElementsAttr(DialectBytecodeWriter &writer,
+ DenseIntOrFPElementsAttr attr) {
+ // Check to see if this is an i1 dense attribute.
+ if (attr.getElementType().isInteger(1)) {
+ // Pack the data.
+ SmallVector<char> data;
+ ArrayRef<char> rawData = attr.getRawData();
+
+ // If the attribute is a splat, we can just splat the value directly.
+ if (attr.isSplat()) {
+ data.resize(1);
+ data[0] = rawData[0] ? 0xFF : 0x00;
+ writer.writeUnownedBlob(data);
+ return;
+ }
+
+ size_t numElements = attr.getNumElements();
+ data.resize(llvm::divideCeil(numElements, 8));
+ // Otherwise, pack the data manually.
+ for (size_t i = 0; i < numElements; ++i)
+ if (rawData[i])
+ data[i / 8] |= (1 << (i % 8));
+ writer.writeUnownedBlob(data);
+ return;
+ }
+
+ writer.writeOwnedBlob(attr.getRawData());
+}
+
#include "mlir/IR/BuiltinDialectBytecode.cpp.inc"
/// This class implements the bytecode interface for the builtin dialect.
diff --git a/mlir/test/Bytecode/bytecode_callback_write_unowned_blob.mlir b/mlir/test/Bytecode/bytecode_callback_write_unowned_blob.mlir
new file mode 100644
index 0000000000000..6fa0d71db5452
--- /dev/null
+++ b/mlir/test/Bytecode/bytecode_callback_write_unowned_blob.mlir
@@ -0,0 +1,9 @@
+// RUN: mlir-opt %s -split-input-file --test-bytecode-roundtrip="test-kind=7" | FileCheck %s
+
+func.func @base_test(%arg0: !test.i32) {
+ return
+}
+
+// CHECK: Writing unowned blob...
+// CHECK: Successfully read the unowned blob.
+// CHECK: func.func @base_test([[ARG0:%.+]]: !test.i32) {
diff --git a/mlir/test/Bytecode/i1_splat_roundtrip.mlir b/mlir/test/Bytecode/i1_splat_roundtrip.mlir
new file mode 100644
index 0000000000000..dde6fe61934e6
--- /dev/null
+++ b/mlir/test/Bytecode/i1_splat_roundtrip.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt %s -emit-bytecode | mlir-opt | FileCheck %s
+
+func.func @test_i1_splat_true() -> tensor<100xi1> {
+ %0 = arith.constant dense<true> : tensor<100xi1>
+ return %0 : tensor<100xi1>
+}
+
+// CHECK-LABEL: func.func @test_i1_splat_true
+// CHECK: arith.constant dense<true> : tensor<100xi1>
+
+func.func @test_i1_splat_false() -> tensor<100xi1> {
+ %0 = arith.constant dense<false> : tensor<100xi1>
+ return %0 : tensor<100xi1>
+}
+
+// CHECK-LABEL: func.func @test_i1_splat_false
+// CHECK: arith.constant dense<false> : tensor<100xi1>
diff --git a/mlir/test/lib/IR/TestBytecodeRoundtrip.cpp b/mlir/test/lib/IR/TestBytecodeRoundtrip.cpp
index 4894ad5294990..e83a3432a0286 100644
--- a/mlir/test/lib/IR/TestBytecodeRoundtrip.cpp
+++ b/mlir/test/lib/IR/TestBytecodeRoundtrip.cpp
@@ -84,6 +84,8 @@ struct TestBytecodeRoundtripPass
// test-kind 6 is a plain roundtrip with downgrade/upgrade to/from
// `targetVersion`.
return runTest6(getOperation());
+ case (7):
+ return runTest7(getOperation());
default:
llvm_unreachable("unhandled test kind for TestBytecodeCallbacks pass");
}
@@ -412,6 +414,59 @@ struct TestBytecodeRoundtripPass
doRoundtripWithConfigs(op, writeConfig, parseConfig);
}
+ // Test7: When writing bytecode, we override the encoding of TestI32Type with
+ // the encoding of builtin IntegerType, but we also write an unowned blob.
+ // We can natively parse this without the use of a callback, relying on the
+ // existing builtin reader mechanism.
+ void runTest7(Operation *op) {
+ auto *builtin = op->getContext()->getLoadedDialect<mlir::BuiltinDialect>();
+ BytecodeDialectInterface *iface =
+ builtin->getRegisteredInterface<BytecodeDialectInterface>();
+ BytecodeWriterConfig writeConfig;
+ writeConfig.attachTypeCallback(
+ [&](Type entryValue, std::optional<StringRef> &dialectGroupName,
+ DialectBytecodeWriter &writer) -> LogicalResult {
+ // Emit TestIntegerType using the builtin dialect encoding.
+ if (llvm::isa<test::TestI32Type>(entryValue)) {
+ auto builtinI32Type =
+ IntegerType::get(op->getContext(), 32,
+ IntegerType::SignednessSemantics::Signless);
+ // Specify that this type will need to be written as part of the
+ // builtin group. This will override the default dialect group of
+ // the attribute (test).
+ dialectGroupName = StringLiteral("builtin");
+ if (succeeded(iface->writeType(builtinI32Type, writer))) {
+ char dummyBlob[] = "test_blob";
+ llvm::outs() << "Writing unowned blob...\n";
+ writer.writeUnownedBlob(ArrayRef<char>(dummyBlob, 9));
+ return success();
+ }
+ }
+ return failure();
+ });
+ ParserConfig parseConfig(op->getContext(), /*verifyAfterParse=*/true);
+ parseConfig.getBytecodeReaderConfig().attachTypeCallback(
+ [&](DialectBytecodeReader &reader, StringRef dialectName,
+ Type &entry) -> LogicalResult {
+ if (dialectName != StringLiteral("builtin"))
+ return failure();
+ Type builtinAttr = iface->readType(reader);
+ if (auto integerType =
+ llvm::dyn_cast_or_null<IntegerType>(builtinAttr)) {
+ if (integerType.getWidth() == 32 && integerType.isSignless()) {
+ ArrayRef<char> blob;
+ if (succeeded(reader.readBlob(blob)) &&
+ blob == ArrayRef<char>("test_blob", 9)) {
+ llvm::outs() << "Successfully read the unowned blob.\n";
+ entry = test::TestI32Type::get(reader.getContext());
+ }
+ }
+ }
+ return success();
+ });
+ doRoundtripWithConfigs(op, writeConfig, parseConfig);
+ }
+
test::TestDialect *testDialect;
};
} // namespace
More information about the Mlir-commits
mailing list