[Mlir-commits] [mlir] [mlirbc] Add AffineMap serialization support (PR #191970)
Jacques Pienaar
llvmlistbot at llvm.org
Fri May 29 02:27:56 PDT 2026
https://github.com/jpienaar updated https://github.com/llvm/llvm-project/pull/191970
>From d225ed9f73101b9daf023008d5cd5ba43c1d4970 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Fri, 20 Feb 2026 09:31:21 +0200
Subject: [PATCH] [mlirbc] Add AffineMap serialization support
Add binary bytecode encoding for AffineMapAttr, replacing the textual fallback.
AffineMap is encoded as numDims, numSymbols, numResults, followed by the result
expressions. Where each expression, AffineExpr, is encoded as a recursive tree
with a VarInt kind tag followed by kind-specific data.
Updated the checked in mlirbc file as memref has a default affinemap, so
updating it pre snap.
---
.../include/mlir/IR/BuiltinDialectBytecode.td | 10 ++
mlir/lib/IR/BuiltinDialectBytecode.cpp | 154 +++++++++++++++++-
.../Builtin/Bytecode/builtin_fixed_0.mlirbc | Bin 4471 -> 4423 bytes
3 files changed, 163 insertions(+), 1 deletion(-)
diff --git a/mlir/include/mlir/IR/BuiltinDialectBytecode.td b/mlir/include/mlir/IR/BuiltinDialectBytecode.td
index 207b99164d0d6..8792c791f061d 100644
--- a/mlir/include/mlir/IR/BuiltinDialectBytecode.td
+++ b/mlir/include/mlir/IR/BuiltinDialectBytecode.td
@@ -205,6 +205,15 @@ def DistinctAttr : DialectAttribute<(attr
Attribute:$referencedAttr
)>;
+// Make easy to disable until version number gets added.
+class EnableAffineMapPrintingJune2026<dag d> : DialectType<d>;
+
+def AffineMapAttr : EnableAffineMapPrintingJune2026<(attr
+ WithParser<"succeeded(readAffineMap($_reader, context, $_var))",
+ WithPrinter<"writeAffineMap($_writer, $_name)",
+ WithType<"AffineMap">>>:$value
+)>;
+
// Types
// -----
@@ -376,6 +385,7 @@ def BuiltinDialectAttributes : DialectAttributes<"Builtin"> {
SparseElementsAttr,
DistinctAttr,
FileLineColRange,
+ AffineMapAttr,
];
}
diff --git a/mlir/lib/IR/BuiltinDialectBytecode.cpp b/mlir/lib/IR/BuiltinDialectBytecode.cpp
index 14dc665184099..c690c9e433507 100644
--- a/mlir/lib/IR/BuiltinDialectBytecode.cpp
+++ b/mlir/lib/IR/BuiltinDialectBytecode.cpp
@@ -9,6 +9,8 @@
#include "BuiltinDialectBytecode.h"
#include "AttributeDetail.h"
#include "mlir/Bytecode/BytecodeImplementation.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinDialect.h"
#include "mlir/IR/BuiltinTypes.h"
@@ -33,7 +35,7 @@ namespace {
// TODO: Move these to separate file.
-// Returns the bitwidth if known, else return std::nullopt.
+// Returns the bitwidth if known, else return 0.
static std::optional<unsigned> getIntegerBitWidth(DialectBytecodeReader &reader,
Type type) {
if (auto intType = dyn_cast<IntegerType>(type))
@@ -88,6 +90,156 @@ static void writePotentiallySplatString(DialectBytecodeWriter &writer,
writer.writeOwnedString(str);
}
+//===----------------------------------------------------------------------===//
+// AffineExpr / AffineMap bytecode helpers
+//===----------------------------------------------------------------------===//
+
+// AffineExpr kind encoding:
+// Extra kinds may be appended here but the existing ones and their ordering
+// should not be changed.
+enum class AffineExprBytecodeKind : uint64_t {
+ DimId = 0,
+ SymbolId = 1,
+ Constant = 2,
+ Add = 3,
+ Mul = 4,
+ Mod = 5,
+ FloorDiv = 6,
+ CeilDiv = 7
+};
+
+static FailureOr<AffineExpr> readAffineExpr(DialectBytecodeReader &reader,
+ MLIRContext *context) {
+ uint64_t kind;
+ if (failed(reader.readVarInt(kind)))
+ return failure();
+
+ switch (static_cast<AffineExprBytecodeKind>(kind)) {
+ case AffineExprBytecodeKind::DimId: {
+ uint64_t position;
+ if (failed(reader.readVarInt(position)))
+ return failure();
+ return getAffineDimExpr(position, context);
+ }
+ case AffineExprBytecodeKind::SymbolId: {
+ uint64_t position;
+ if (failed(reader.readVarInt(position)))
+ return failure();
+ return getAffineSymbolExpr(position, context);
+ }
+ case AffineExprBytecodeKind::Constant: {
+ int64_t value;
+ if (failed(reader.readSignedVarInt(value)))
+ return failure();
+ return getAffineConstantExpr(value, context);
+ }
+ case AffineExprBytecodeKind::Add:
+ case AffineExprBytecodeKind::Mul:
+ case AffineExprBytecodeKind::Mod:
+ case AffineExprBytecodeKind::FloorDiv:
+ case AffineExprBytecodeKind::CeilDiv: { // Binary ops
+ auto lhs = readAffineExpr(reader, context);
+ if (failed(lhs))
+ return failure();
+ auto rhs = readAffineExpr(reader, context);
+ if (failed(rhs))
+ return failure();
+ AffineExprKind exprKind;
+ switch (static_cast<AffineExprBytecodeKind>(kind)) {
+ case AffineExprBytecodeKind::Add:
+ exprKind = AffineExprKind::Add;
+ break;
+ case AffineExprBytecodeKind::Mul:
+ exprKind = AffineExprKind::Mul;
+ break;
+ case AffineExprBytecodeKind::Mod:
+ exprKind = AffineExprKind::Mod;
+ break;
+ case AffineExprBytecodeKind::FloorDiv:
+ exprKind = AffineExprKind::FloorDiv;
+ break;
+ case AffineExprBytecodeKind::CeilDiv:
+ exprKind = AffineExprKind::CeilDiv;
+ break;
+ default:
+ llvm_unreachable("unhandled affine expr kind");
+ }
+ return getAffineBinaryOpExpr(exprKind, *lhs, *rhs);
+ }
+ }
+}
+
+static void writeAffineExpr(DialectBytecodeWriter &writer, AffineExpr expr) {
+ switch (expr.getKind()) {
+ case AffineExprKind::DimId:
+ writer.writeVarInt(static_cast<uint64_t>(AffineExprBytecodeKind::DimId));
+ writer.writeVarInt(cast<AffineDimExpr>(expr).getPosition());
+ break;
+ case AffineExprKind::SymbolId:
+ writer.writeVarInt(static_cast<uint64_t>(AffineExprBytecodeKind::SymbolId));
+ writer.writeVarInt(cast<AffineSymbolExpr>(expr).getPosition());
+ break;
+ case AffineExprKind::Constant:
+ writer.writeVarInt(static_cast<uint64_t>(AffineExprBytecodeKind::Constant));
+ writer.writeSignedVarInt(cast<AffineConstantExpr>(expr).getValue());
+ break;
+ case AffineExprKind::Add:
+ writer.writeVarInt(static_cast<uint64_t>(AffineExprBytecodeKind::Add));
+ writeAffineExpr(writer, cast<AffineBinaryOpExpr>(expr).getLHS());
+ writeAffineExpr(writer, cast<AffineBinaryOpExpr>(expr).getRHS());
+ break;
+ case AffineExprKind::Mul:
+ writer.writeVarInt(static_cast<uint64_t>(AffineExprBytecodeKind::Mul));
+ writeAffineExpr(writer, cast<AffineBinaryOpExpr>(expr).getLHS());
+ writeAffineExpr(writer, cast<AffineBinaryOpExpr>(expr).getRHS());
+ break;
+ case AffineExprKind::Mod:
+ writer.writeVarInt(static_cast<uint64_t>(AffineExprBytecodeKind::Mod));
+ writeAffineExpr(writer, cast<AffineBinaryOpExpr>(expr).getLHS());
+ writeAffineExpr(writer, cast<AffineBinaryOpExpr>(expr).getRHS());
+ break;
+ case AffineExprKind::FloorDiv:
+ writer.writeVarInt(static_cast<uint64_t>(AffineExprBytecodeKind::FloorDiv));
+ writeAffineExpr(writer, cast<AffineBinaryOpExpr>(expr).getLHS());
+ writeAffineExpr(writer, cast<AffineBinaryOpExpr>(expr).getRHS());
+ break;
+ case AffineExprKind::CeilDiv:
+ writer.writeVarInt(static_cast<uint64_t>(AffineExprBytecodeKind::CeilDiv));
+ writeAffineExpr(writer, cast<AffineBinaryOpExpr>(expr).getLHS());
+ writeAffineExpr(writer, cast<AffineBinaryOpExpr>(expr).getRHS());
+ break;
+ }
+}
+
+static LogicalResult readAffineMap(DialectBytecodeReader &reader,
+ MLIRContext *context, AffineMap &map) {
+ uint64_t numDims, numSymbols, numResults;
+ if (failed(reader.readVarInt(numDims)) ||
+ failed(reader.readVarInt(numSymbols)) ||
+ failed(reader.readVarInt(numResults)))
+ return failure();
+
+ SmallVector<AffineExpr> results;
+ results.reserve(numResults);
+ for (uint64_t i = 0; i < numResults; ++i) {
+ auto expr = readAffineExpr(reader, context);
+ if (failed(expr))
+ return failure();
+ results.push_back(*expr);
+ }
+ map = AffineMap::get(numDims, numSymbols, results, context);
+ return success();
+}
+
+static void writeAffineMap(DialectBytecodeWriter &writer, AffineMapAttr attr) {
+ AffineMap map = attr.getValue();
+ writer.writeVarInt(map.getNumDims());
+ writer.writeVarInt(map.getNumSymbols());
+ writer.writeVarInt(map.getNumResults());
+ for (AffineExpr expr : map.getResults())
+ writeAffineExpr(writer, expr);
+}
+
static FileLineColRange getFileLineColRange(MLIRContext *context,
StringAttr filename,
ArrayRef<uint64_t> lineCols) {
diff --git a/mlir/test/Dialect/Builtin/Bytecode/builtin_fixed_0.mlirbc b/mlir/test/Dialect/Builtin/Bytecode/builtin_fixed_0.mlirbc
index e2f58a42751a1cddc0dc2845d06735faafdc8b32..319b8985836bc0fdd4b947c864d29e4079732d64 100644
GIT binary patch
delta 56
zcmeyabX;kII-~MLjZ2Ku6MrW&icOAZ6y?xoWn^V!WMtkvmGL$+n?5rmGb7{XPS%4=
MOstNZ131(<0l|$8#{d8T
delta 104
zcmX at E^j&F!I%DfZjZ2KF6MrW&9-AD`D5{v4mX?{98lRh3V55;@preptsHvc9rvPKy
jF>J13yv;0$stl$?6DS?adXR~&)lu;D>C>APIMg`-B&Hz*
More information about the Mlir-commits
mailing list