[Mlir-commits] [mlir] [mlirbc] Add IntegerSetAttr (PR #207482)
Jacques Pienaar
llvmlistbot at llvm.org
Fri Jul 3 21:49:29 PDT 2026
https://github.com/jpienaar created https://github.com/llvm/llvm-project/pull/207482
Serialize IntegerSetAttr using VarInt for numDims and numSymbols and arrays for contraints and eqFlags. Add ArrayWithKnownSize` an Array variant that omits the length prefix when the size is known from a previously serialized member. This avoids encoding redundant array lengths when, for example, two arrays are always the same size.
Assisted-By: Gemini
>From 47026d2a85a42d284230e368c61569ae8f5ed0a2 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Mon, 23 Feb 2026 11:33:39 +0200
Subject: [PATCH] [mlirbc] Add IntegerSetAttr
Serialize IntegerSetAttr using VarInt for numDims and numSymbols and arrays for
contraints and eqFlags. Add ArrayWithKnownSize` an Array variant that omits the
length prefix when the size is known from a previously serialized member. This
avoids encoding redundant array lengths when, for example, two arrays are
always the same size.
Assisted-By: Gemini
---
.../mlir/Bytecode/BytecodeImplementation.h | 20 +++++++++++++
.../include/mlir/IR/BuiltinDialectBytecode.td | 29 +++++++++++++++++++
mlir/include/mlir/IR/BytecodeBase.td | 10 ++++++-
mlir/lib/IR/BuiltinDialectBytecode.cpp | 19 ++++++------
mlir/test/Dialect/Builtin/Bytecode/attrs.mlir | 16 ++++++++++
mlir/tools/mlir-tblgen/BytecodeDialectGen.cpp | 14 +++++++--
6 files changed, 95 insertions(+), 13 deletions(-)
diff --git a/mlir/include/mlir/Bytecode/BytecodeImplementation.h b/mlir/include/mlir/Bytecode/BytecodeImplementation.h
index f1381785a2703..d8d0ca7b3b9f7 100644
--- a/mlir/include/mlir/Bytecode/BytecodeImplementation.h
+++ b/mlir/include/mlir/Bytecode/BytecodeImplementation.h
@@ -76,6 +76,18 @@ class DialectBytecodeReader {
uint64_t size;
if (failed(readVarInt(size)))
return failure();
+ return readListWithKnownSize(result, size,
+ std::forward<CallbackFn>(callback));
+ }
+
+ /// Read out a list of elements with a known size, invoking the provided
+ /// callback for each element. Unlike readList, this does not read a length
+ /// prefix. The callback function may be in any of the following forms:
+ /// * LogicalResult(T &)
+ /// * FailureOr<T>()
+ template <typename T, typename CallbackFn>
+ LogicalResult readListWithKnownSize(SmallVectorImpl<T> &result, uint64_t size,
+ CallbackFn &&callback) {
result.reserve(size);
for (uint64_t i = 0; i < size; ++i) {
@@ -294,6 +306,14 @@ class DialectBytecodeWriter {
callback(element);
}
+ /// Write out a list of elements without a length prefix, for cases where the
+ /// size is known from another field.
+ template <typename RangeT, typename CallbackFn>
+ void writeListWithKnownSize(RangeT &&range, CallbackFn &&callback) {
+ for (auto &element : range)
+ callback(element);
+ }
+
/// Write a reference to the given attribute.
virtual void writeAttribute(Attribute attr) = 0;
virtual void writeOptionalAttribute(Attribute attr) = 0;
diff --git a/mlir/include/mlir/IR/BuiltinDialectBytecode.td b/mlir/include/mlir/IR/BuiltinDialectBytecode.td
index a3a0511791613..f087e5cb23081 100644
--- a/mlir/include/mlir/IR/BuiltinDialectBytecode.td
+++ b/mlir/include/mlir/IR/BuiltinDialectBytecode.td
@@ -214,6 +214,34 @@ def AffineMapAttr : EnableAffineMapPrintingJuly2026<(attr
WithType<"AffineMap">>>:$value
)>;
+def AffineExpr :
+ WithParser<"succeeded(readAffineExpr($_reader, context, $_var))",
+ WithBuilder<"$_args",
+ WithPrinter<"writeAffineExpr($_writer, $_getter)",
+ WithType<"AffineExpr">>>>;
+def AffineExprList : List<AffineExpr>;
+
+// Similar to AffineMapAttr, IntegerSetAttr has everything going via IntegerSet
+// and getValue, so extra indirections.
+def IntegerSetConstraints : Array<AffineExprList> {
+ let cGetter = "$_attrType.getValue().getConstraints()";
+}
+
+def IntegerSetEqFlags : ArrayWithKnownSize<BoolList, "constraints.size()"> {
+ let cGetter = "$_attrType.getValue().getEqFlags()";
+}
+
+class EnableIntegerSetPrintingJuly2026<dag d> : DialectType<d>;
+
+def IntegerSetAttr : EnableIntegerSetPrintingJuly2026<(attr
+ WithGetter<"$_attrType.getValue().getNumDims()", VarInt>:$numDims,
+ WithGetter<"$_attrType.getValue().getNumSymbols()", VarInt>:$numSymbols,
+ IntegerSetConstraints:$constraints,
+ IntegerSetEqFlags:$eqFlags
+)> {
+ let cBuilder = "IntegerSetAttr::get(IntegerSet::get(numDims, numSymbols, constraints, eqFlags))";
+}
+
// Types
// -----
@@ -388,6 +416,7 @@ def BuiltinDialectAttributes : DialectAttributes<"Builtin"> {
DistinctAttr,
FileLineColRange,
AffineMapAttr,
+ IntegerSetAttr,
];
}
diff --git a/mlir/include/mlir/IR/BytecodeBase.td b/mlir/include/mlir/IR/BytecodeBase.td
index df60800ef639b..3eab62493ffef 100644
--- a/mlir/include/mlir/IR/BytecodeBase.td
+++ b/mlir/include/mlir/IR/BytecodeBase.td
@@ -125,6 +125,15 @@ class Array<Bytecode t> {
Bytecode elemT = t;
string cBuilder = "$_args";
+
+ // Optional custom getter expression for the array from the parent type.
+ string cGetter = "";
+}
+// - Array variant where the size is known from a previously serialized member.
+// This avoids encoding a redundant length prefix.
+class ArrayWithKnownSize<Bytecode t, string sizeRef> : Array<t> {
+ // Expression providing the element count (e.g., "otherMember.size()").
+ string knownSizeRef = sizeRef;
}
// - Array elements currently needs a different bytecode type to accommodate
// for the list print/parsing.
@@ -177,4 +186,3 @@ def none;
def ReservedOrDead : DialectAttrOrType<(none)>;
#endif // BYTECODE_BASE
-
diff --git a/mlir/lib/IR/BuiltinDialectBytecode.cpp b/mlir/lib/IR/BuiltinDialectBytecode.cpp
index 366b3ceb06410..e5476082cd6a0 100644
--- a/mlir/lib/IR/BuiltinDialectBytecode.cpp
+++ b/mlir/lib/IR/BuiltinDialectBytecode.cpp
@@ -16,6 +16,7 @@
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/DialectResourceBlobManager.h"
+#include "mlir/IR/IntegerSet.h"
#include "mlir/IR/Location.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/TypeSwitch.h"
@@ -157,8 +158,8 @@ static AffineExprKind fromBytecodeKind(uint64_t kind) {
/// is prefix order (operator and then children), which is self-delimiting.
/// Instead of C++ recursion the reader uses an explicit work stack, bounding
/// memory to O(depth) and eliminating stack-overflow risk on malicious input.
-static FailureOr<AffineExpr> readAffineExpr(DialectBytecodeReader &reader,
- MLIRContext *context) {
+static LogicalResult readAffineExpr(DialectBytecodeReader &reader,
+ MLIRContext *context, AffineExpr &expr) {
// A work-stack item is either ReadOperand (0) or a combine marker whose
// payload is an AffineExprKind.
struct WorkItem {
@@ -175,15 +176,14 @@ static FailureOr<AffineExpr> readAffineExpr(DialectBytecodeReader &reader,
while (!work.empty()) {
// Bound total iterations to catch malformed input.
if (work.size() > 128)
- return reader.emitError("AffineExpr work stack overflow"), failure();
+ return reader.emitError("AffineExpr work stack overflow");
WorkItem item = work.pop_back_val();
if (item.isCombine) {
// Pop two operands and combine.
if (operands.size() < 2)
- return reader.emitError("malformed AffineExpr: operand underflow"),
- failure();
+ return reader.emitError("malformed AffineExpr: operand underflow");
AffineExpr rhs = operands.pop_back_val();
AffineExpr lhs = operands.pop_back_val();
operands.push_back(getAffineBinaryOpExpr(item.combineKind, lhs, rhs));
@@ -240,7 +240,8 @@ static FailureOr<AffineExpr> readAffineExpr(DialectBytecodeReader &reader,
return reader.emitError("malformed AffineExpr: expected single result"),
failure();
- return operands.front();
+ expr = operands.front();
+ return success();
}
/// Write an AffineExpr in prefix order (operator first, then children).
@@ -337,10 +338,10 @@ static LogicalResult readAffineMap(DialectBytecodeReader &reader,
SmallVector<AffineExpr> results;
results.reserve(numResults);
for (uint64_t i = 0; i < numResults; ++i) {
- auto expr = readAffineExpr(reader, context);
- if (failed(expr))
+ AffineExpr expr;
+ if (failed(readAffineExpr(reader, context, expr)))
return failure();
- results.push_back(*expr);
+ results.push_back(expr);
}
map = AffineMap::get(numDims, numSymbols, results, context);
return success();
diff --git a/mlir/test/Dialect/Builtin/Bytecode/attrs.mlir b/mlir/test/Dialect/Builtin/Bytecode/attrs.mlir
index aac1c44248995..3364a8e9fce8e 100644
--- a/mlir/test/Dialect/Builtin/Bytecode/attrs.mlir
+++ b/mlir/test/Dialect/Builtin/Bytecode/attrs.mlir
@@ -255,3 +255,19 @@ module @TestAffineMap attributes {
bytecode.projperm_single = affine_map<(d0, d1, d2) -> (d2)>
} {}
+//===----------------------------------------------------------------------===//
+// IntegerSetAttr
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @TestIntegerSetAttr
+module @TestIntegerSetAttr attributes {
+ // CHECK-DAG: bytecode.eq = affine_set<(d0) : (d0 == 0)>
+ // CHECK-DAG: bytecode.ineq = affine_set<(d0) : (d0 >= 0)>
+ // CHECK-DAG: bytecode.multi = affine_set<(d0, d1)[s0] : (d0 >= 0, -d0 + s0 >= 0, d1 >= 0)>
+ // CHECK-DAG: bytecode.eq_ineq = affine_set<(d0, d1) : (d0 == 0, d1 >= 0)>
+ bytecode.eq = affine_set<(d0) : (d0 == 0)>,
+ bytecode.ineq = affine_set<(d0) : (d0 >= 0)>,
+ bytecode.multi = affine_set<(d0, d1)[s0] : (d0 >= 0, -d0 + s0 >= 0, d1 >= 0)>,
+ bytecode.eq_ineq = affine_set<(d0, d1) : (d0 == 0, d1 >= 0)>
+} {}
+
diff --git a/mlir/tools/mlir-tblgen/BytecodeDialectGen.cpp b/mlir/tools/mlir-tblgen/BytecodeDialectGen.cpp
index dd178b5e5d232..031fe6cc94eb2 100644
--- a/mlir/tools/mlir-tblgen/BytecodeDialectGen.cpp
+++ b/mlir/tools/mlir-tblgen/BytecodeDialectGen.cpp
@@ -183,7 +183,12 @@ static void printParseConditional(mlir::raw_indented_ostream &ios,
parser = "succeeded($_reader.readAttributes($_var))";
else if (!composite && def->isSubClassOf("TypeKind"))
parser = "succeeded($_reader.readTypes($_var))";
- else
+ else if (attr->isSubClassOf("ArrayWithKnownSize")) {
+ std::string sizeRef = attr->getValueAsString("knownSizeRef").str();
+ parser = ("succeeded($_reader.readListWithKnownSize($_var, " +
+ sizeRef + ", " + listHelperName(std::get<1>(it)) + "))")
+ .str();
+ } else
parser = ("succeeded($_reader.readList($_var, " +
listHelperName(std::get<1>(it)) + "))")
.str();
@@ -385,8 +390,11 @@ void Generator::emitPrintHelper(const Record *memberRec, StringRef kind,
}
std::string returnType = getCType(def);
std::string nestedName = kind.str();
- ios << "writer.writeList(" << getter << ", [&](" << returnType << " "
- << nestedName << ") ";
+ StringRef writeMethod = memberRec->isSubClassOf("ArrayWithKnownSize")
+ ? "writer.writeListWithKnownSize("
+ : "writer.writeList(";
+ ios << writeMethod << getter << ", [&](" << returnType << " " << nestedName
+ << ") ";
auto lambdaScope = ios.scope("{\n", "});\n");
return emitPrintHelper(def, kind, nestedName, nestedName, ios);
}
More information about the Mlir-commits
mailing list