[Mlir-commits] [mlir] [mlir][tensor] Preserve encoding in more canonicalizers (pad, convertToStaticExpandShape) (PR #207241)
Dmitrii Makarenko
llvmlistbot at llvm.org
Tue Aug 4 04:21:07 PDT 2026
https://github.com/Devjiu updated https://github.com/llvm/llvm-project/pull/207241
>From fad6a15ffec811be579573378cbde07071d823e1 Mon Sep 17 00:00:00 2001
From: Dmitrii Makarenko <dmitrii.makarenko at intel.com>
Date: Thu, 2 Jul 2026 16:32:59 +0000
Subject: [PATCH] [mlir][tensor] Preserve encoding in more canonicalizers (pad,
convertToStaticExpandShape)
Fixes encoding drop in `tensor.*` canonicalizers.
Patterns that only refine a tensor's shape (never merge/combine data) now propagate
the encoding: `ConvertToStaticExpandShape`, `PadOp::inferResultType` (and its callers
`FoldSourceTensorCast`, `FoldStaticPadding`).
An encoding implementing `VerifiableTensorEncoding` is re-verified against the refined
shape and dropped if invalid (e.g. sparse); an opaque encoding (no interface) is
propagated as-is.
Patterns that merge/combine tensors (`ConcatOp::inferResultType`, `InferConcatOperandTypes`,
`CollapseShapeOp::inferCollapsedType`) keep the existing drop-encoding behavior - there's
no static way to verify an arbitrary encoding survives a merge or rank change when
dynamic dims are involved.
Documents this contract on `VerifiableTensorEncoding` in `TensorEncoding.td`.
co-authored Claude Opus 4.7
Signed-off-by: Dmitrii Makarenko <dmitrii.makarenko at intel.com>
---
mlir/include/mlir/IR/TensorEncoding.td | 28 ++++
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 47 ++++++-
mlir/test/Dialect/Tensor/canonicalize.mlir | 145 ++++++++++++++++++++-
3 files changed, 213 insertions(+), 7 deletions(-)
diff --git a/mlir/include/mlir/IR/TensorEncoding.td b/mlir/include/mlir/IR/TensorEncoding.td
index d8ccd1f0c548e..a9f93a966599f 100644
--- a/mlir/include/mlir/IR/TensorEncoding.td
+++ b/mlir/include/mlir/IR/TensorEncoding.td
@@ -23,6 +23,34 @@ def VerifiableTensorEncoding : AttrInterface<"VerifiableTensorEncoding"> {
let cppNamespace = "::mlir";
let description = [{
Verifies an encoding attribute for a tensor.
+
+ This interface also doubles as the contract some `tensor.*`
+ canonicalization patterns rely on when refining a tensor's shape to be
+ more static (e.g. folding a `tensor.cast` into a consumer, or turning a
+ constant `Value` operand into a static dimension). When such a pattern
+ needs to decide whether an existing encoding still applies to the
+ refined shape:
+
+ - if the encoding implements this interface, the pattern re-runs
+ `verifyEncoding` against the refined shape and keeps the encoding
+ only if it still holds; this is how a rank- or shape-dependent
+ encoding (e.g. a sparse tensor encoding, which encodes a per-dimension
+ layout) gets dropped instead of ending up attached to a
+ `RankedTensorType` it is not valid for.
+ - if the encoding does not implement this interface, it is treated as
+ opaque and shape-agnostic, and is propagated onto the refined type
+ unconditionally. Not implementing the interface is thus an implicit
+ opt-out: the attribute's author is asserting that its validity does
+ not depend on the tensor's shape, and takes on responsibility for
+ that invariant.
+
+ Patterns that instead merge or otherwise combine multiple tensors with
+ possibly-differing encodings (e.g. `tensor.concat`, or collapsing
+ dimensions via `tensor.collapse_shape`) do not attempt to verify or
+ propagate an encoding at all, since there is no way to statically prove
+ an arbitrary encoding remains valid once tensors are combined or a
+ dynamic dimension folds two static ones together — the encoding is
+ dropped from the result in that case.
}];
let methods = [
InterfaceMethod<
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 637366a289ac9..680b6c3927ec2 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -24,6 +24,7 @@
#include "mlir/IR/Matchers.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/TensorEncoding.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
#include "mlir/Interfaces/InferIntRangeInterface.h"
@@ -44,6 +45,26 @@
using namespace mlir;
using namespace mlir::tensor;
+/// Implements the `VerifiableTensorEncoding` contract documented in
+/// TensorEncoding.td for patterns that refine a tensor's shape to be more
+/// static: verifiable encodings are re-checked against the refined shape and
+/// dropped if they no longer hold; opaque encodings (not implementing the
+/// interface) are propagated unconditionally.
+static Attribute propagateEncoding(Attribute encoding, ArrayRef<int64_t> shape,
+ Type elementType) {
+ auto verifiable = dyn_cast_or_null<VerifiableTensorEncoding>(encoding);
+ if (!verifiable)
+ return encoding;
+
+ MLIRContext *ctx = encoding.getContext();
+ // to avoid user's error stream
+ ScopedDiagnosticHandler swallow(ctx, [](Diagnostic &) { return success(); });
+ auto emit = [ctx]() { return mlir::emitError(UnknownLoc::get(ctx)); };
+ return succeeded(verifiable.verifyEncoding(shape, elementType, emit))
+ ? encoding
+ : Attribute{};
+}
+
/// Materialize a single constant operation from a given attribute value with
/// the desired resultant type.
Operation *TensorDialect::materializeConstant(OpBuilder &builder,
@@ -2273,10 +2294,18 @@ struct ConvertToStaticExpandShape : public OpRewritePattern<ExpandShapeOp> {
SmallVector<OpFoldResult> outputOfr =
getMixedValues(newOutputShape, dynamicOutputShape, rewriter);
+ // The refined types are still applied to the same src/result values, so
+ // propagate their encodings, letting each encoding self-decide whether it
+ // still holds on the more-static shape.
+ Type elementType = expandOp.getSrcType().getElementType();
auto inputType = RankedTensorType::get(
- newInputShape, expandOp.getSrcType().getElementType());
+ newInputShape, elementType,
+ propagateEncoding(expandOp.getSrcType().getEncoding(), newInputShape,
+ elementType));
auto outputType = RankedTensorType::get(
- newOutputShape, expandOp.getSrcType().getElementType());
+ newOutputShape, elementType,
+ propagateEncoding(expandOp.getResultType().getEncoding(),
+ newOutputShape, elementType));
auto inputCast = CastOp::create(rewriter, expandOp.getLoc(), inputType,
expandOp.getSrc());
auto newExpand = ExpandShapeOp::create(
@@ -2385,7 +2414,8 @@ RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType(
if (!dimsToProject.test(pos))
projectedShape.push_back(shape[pos]);
inferredType =
- RankedTensorType::get(projectedShape, inferredType.getElementType());
+ RankedTensorType::get(projectedShape, inferredType.getElementType(),
+ inferredType.getEncoding());
}
return inferredType;
}
@@ -3323,7 +3353,10 @@ RankedTensorType PadOp::inferResultType(RankedTensorType sourceType,
}
}
- return RankedTensorType::get(inferredShape, sourceType.getElementType());
+ Type elementType = sourceType.getElementType();
+ return RankedTensorType::get(
+ inferredShape, elementType,
+ propagateEncoding(sourceType.getEncoding(), inferredShape, elementType));
}
void PadOp::build(OpBuilder &b, OperationState &result, Type resultType,
@@ -3730,9 +3763,11 @@ struct FoldStaticPadding : public OpRewritePattern<PadOp> {
[&](int64_t x) { return x == ShapedType::kDynamic; }))
return failure();
- // Rewrite the op using the new static type.
+ Type elementType = padTensorOp.getType().getElementType();
auto newResultType = RankedTensorType::get(
- newOutDims, padTensorOp.getType().getElementType());
+ newOutDims, elementType,
+ propagateEncoding(padTensorOp.getType().getEncoding(), newOutDims,
+ elementType));
auto newOp = PadOp::create(
rewriter, padTensorOp->getLoc(), newResultType, input, staticLow,
staticHigh, newLows, newHighs, padTensorOp.getNofold(),
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 67b7ab99c5d18..26c40863e70dc 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -162,6 +162,26 @@ func.func @infer_concat_return_type(%arg0: tensor<5x12xi32>, %arg1: tensor<?x12x
// -----
+// ConcatOp genuinely merges data from multiple operands (unlike a plain
+// shape-refining cast), so `ConcatOp::inferResultType` / `InferConcatOperandTypes`
+// intentionally do NOT try to propagate an operand's encoding onto the
+// refined types
+// CHECK-LABEL: concat_drops_encoding
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<3x?xi32, "abc">
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xi32, "abc">
+// CHECK: %[[CAST0:.+]] = tensor.cast %[[ARG0]] : tensor<3x?xi32, "abc"> to tensor<3x?xi32>
+// CHECK: %[[CAST1:.+]] = tensor.cast %[[ARG1]] : tensor<?x?xi32, "abc"> to tensor<3x?xi32>
+// CHECK: %[[CONCAT:.+]] = tensor.concat dim(1) %[[CAST0]], %[[CAST1]] : (tensor<3x?xi32>, tensor<3x?xi32>) -> tensor<3x?xi32>
+// CHECK: tensor.cast %[[CONCAT]] : tensor<3x?xi32> to tensor<3x?xi32, "abc">
+func.func @concat_drops_encoding(
+ %a: tensor<3x?xi32, "abc">, %b: tensor<?x?xi32, "abc">) -> tensor<3x?xi32, "abc"> {
+ %r = tensor.concat dim(1) %a, %b
+ : (tensor<3x?xi32, "abc">, tensor<?x?xi32, "abc">) -> tensor<3x?xi32, "abc">
+ return %r : tensor<3x?xi32, "abc">
+}
+
+// -----
+
// CHECK-LABEL: func @fold_extract
func.func @fold_extract(%arg0 : index) -> (f32, f16, f16, i32, complex<f32>, i32) {
%const_0 = arith.constant 0 : index
@@ -1092,6 +1112,23 @@ func.func @collapse_of_cast(%t: tensor<8x12x32xf32>) -> tensor<?x32xf32> {
// -----
+// CollapseShapeOp::inferCollapsedType changes the tensor rank, and (like
+// concat) we can't statically verify an arbitrary encoding still holds on
+// the collapsed shape when dynamic dims are involved, so the encoding is
+// dropped
+// CHECK-LABEL: func.func @collapse_of_cast_drops_encoding(
+// CHECK-SAME: %[[IN:.*]]: tensor<8x12x32xf32, "abc">
+// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[IN]] {{\[}}[0, 1], [2]] : tensor<8x12x32xf32, "abc"> into tensor<96x32xf32>
+// CHECK: tensor.cast %[[COLLAPSE]] : tensor<96x32xf32> to tensor<?x32xf32>
+func.func @collapse_of_cast_drops_encoding(%t: tensor<8x12x32xf32, "abc">) -> tensor<?x32xf32> {
+ %0 = tensor.cast %t : tensor<8x12x32xf32, "abc"> to tensor<?x?x?xf32, "abc">
+ %1 = tensor.collapse_shape %0 [[0, 1], [2]] : tensor<?x?x?xf32, "abc"> into tensor<?x?xf32, "abc">
+ %2 = tensor.cast %1 : tensor<?x?xf32, "abc"> to tensor<?x32xf32>
+ return %2 : tensor<?x32xf32>
+}
+
+// -----
+
func.func @fold_collapse_of_expand(%arg0 : tensor<12x4xf32>) -> tensor<12x4xf32> {
%0 = tensor.expand_shape %arg0 [[0, 1], [2]] output_shape [3, 4, 4]
: tensor<12x4xf32> into tensor<3x4x4xf32>
@@ -1773,6 +1810,52 @@ func.func @pad_nofold_same_static_shape(%arg0: tensor<5x6xf32>, %a: index)
// -----
+// FoldSourceTensorCast (via PadOp::inferResultType) must preserve the source
+// encoding onto the refined pad result and the inserted result cast.
+// CHECK-LABEL: func @pad_after_cast_preserves_encoding(
+// CHECK-SAME: %[[INPUT:.*]]: tensor<?x64x?x?xf32, "abc">
+// CHECK: %[[PADDED:.*]] = tensor.pad %[[INPUT]]
+// CHECK: : tensor<?x64x?x?xf32, "abc"> to tensor<?x64x?x?xf32, "abc">
+// CHECK: %[[CAST:.*]] = tensor.cast %[[PADDED]] : tensor<?x64x?x?xf32, "abc"> to tensor<?x?x?x?xf32, "abc">
+// CHECK: return %[[CAST]]
+func.func @pad_after_cast_preserves_encoding(
+ %arg0: tensor<?x64x?x?xf32, "abc">) -> tensor<?x?x?x?xf32, "abc"> {
+ %cst = arith.constant 0.000000e+00 : f32
+ %dynamic = tensor.cast %arg0
+ : tensor<?x64x?x?xf32, "abc"> to tensor<?x?x?x?xf32, "abc">
+ %padded = tensor.pad %dynamic low[0, 0, 1, 1] high[0, 0, 1, 1] {
+ ^bb0(%a: index, %b: index, %c: index, %d: index):
+ tensor.yield %cst: f32
+ } : tensor<?x?x?x?xf32, "abc"> to tensor<?x?x?x?xf32, "abc">
+ return %padded: tensor<?x?x?x?xf32, "abc">
+}
+
+// -----
+
+// `VerifiableTensorEncoding` case: pad preserves rank, so the sparse encoding
+// stays valid on the refined shape and PadOp::inferResultType keeps it.
+#sparse_dense4b = #sparse_tensor.encoding<{
+ map = (d0, d1, d2, d3) -> (d0 : dense, d1 : dense, d2 : dense, d3 : dense)
+}>
+// CHECK-LABEL: func @pad_after_cast_preserves_sparse_encoding(
+// CHECK-SAME: %[[INPUT:.*]]: tensor<?x64x?x?xf32, #{{[a-z_0-9]+}}>
+// CHECK: %[[PADDED:.*]] = tensor.pad %[[INPUT]]
+// CHECK: : tensor<?x64x?x?xf32, #{{[a-z_0-9]+}}> to tensor<?x64x?x?xf32, #{{[a-z_0-9]+}}>
+// CHECK: tensor.cast %[[PADDED]] : tensor<?x64x?x?xf32, #{{[a-z_0-9]+}}> to tensor<?x?x?x?xf32, #{{[a-z_0-9]+}}>
+func.func @pad_after_cast_preserves_sparse_encoding(
+ %arg0: tensor<?x64x?x?xf32, #sparse_dense4b>) -> tensor<?x?x?x?xf32, #sparse_dense4b> {
+ %cst = arith.constant 0.000000e+00 : f32
+ %dynamic = tensor.cast %arg0
+ : tensor<?x64x?x?xf32, #sparse_dense4b> to tensor<?x?x?x?xf32, #sparse_dense4b>
+ %padded = tensor.pad %dynamic low[0, 0, 1, 1] high[0, 0, 1, 1] {
+ ^bb0(%a: index, %b: index, %c: index, %d: index):
+ tensor.yield %cst: f32
+ } : tensor<?x?x?x?xf32, #sparse_dense4b> to tensor<?x?x?x?xf32, #sparse_dense4b>
+ return %padded: tensor<?x?x?x?xf32, #sparse_dense4b>
+}
+
+// -----
+
// CHECK-LABEL: func @pad_after_cast_different_shape(
// CHECK-SAME: %[[INPUT:.*]]: tensor<?x64x?x?xf32>) -> tensor<?x?x?x?xf32> {
// CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
@@ -1921,6 +2004,26 @@ func.func @pad_static_zero_cast(%arg0: tensor<?x?x?xf32>, %pad_value: f32) -> te
// -----
+// FoldStaticPadding must preserve the pad's result encoding on the refined
+// (more-static) pad and the inserted result cast.
+// CHECK-LABEL: func @fold_static_padding_preserves_encoding(
+// CHECK-SAME: %[[SRC:.*]]: tensor<8x?xf32, "abc">
+// CHECK: %[[PADDED:.*]] = tensor.pad %[[SRC]] low[1, 2] high[1, 2]
+// CHECK: : tensor<8x?xf32, "abc"> to tensor<10x?xf32, "abc">
+// CHECK: tensor.cast %[[PADDED]] : tensor<10x?xf32, "abc"> to tensor<?x?xf32, "abc">
+func.func @fold_static_padding_preserves_encoding(
+ %arg0: tensor<8x?xf32, "abc">, %pv: f32) -> tensor<?x?xf32, "abc"> {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %r = tensor.pad %arg0 low[%c1, %c2] high[%c1, %c2] {
+ ^bb0(%a: index, %b: index):
+ tensor.yield %pv: f32
+ } : tensor<8x?xf32, "abc"> to tensor<?x?xf32, "abc">
+ return %r : tensor<?x?xf32, "abc">
+}
+
+// -----
+
// CHECK-LABEL: func @pad_nofold_static_zero(
// CHECK-SAME: %[[ARG0:.*]]: tensor<?x?x?xf32>
// CHECK: %[[PAD:.*]] = tensor.pad
@@ -2641,8 +2744,48 @@ func.func @partial_sink_expand_of_cast(%arg0 : tensor<10x10xf32>, %arg1 : index,
// CHECK-LABEL: func.func @partial_sink_expand_of_cast
// CHECK: %[[CAST:.+]] = tensor.cast
// CHECK-SAME: tensor<10x10xf32> to tensor<?x10xf32>
-// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2]]
+// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2]]
// CHECK-SAME: output_shape [%{{.*}}, %{{.*}}, 10]
// CHECK: %[[RES:.+]] = tensor.cast %[[EXPAND]]
// CHECK-SAME: tensor<?x?x10xf32> to tensor<?x?x?xf32>
// CHECK: return %[[RES]]
+
+// -----
+
+// ConvertToStaticExpandShape must carry the source's encoding onto the refined
+// cast-target and the (source-side) encoding of the new ExpandShapeOp result.
+// CHECK-LABEL: func.func @sink_expand_of_cast_preserves_encoding
+// CHECK: %[[EXPAND:.+]] = tensor.expand_shape
+// CHECK-SAME: tensor<64xf32, "abc"> into tensor<8x8xf32, "abc">
+// CHECK: tensor.cast %[[EXPAND]] : tensor<8x8xf32, "abc"> to tensor<?x?xf32, "abc">
+func.func @sink_expand_of_cast_preserves_encoding(%t: tensor<64xf32, "abc">) -> tensor<?x?xf32, "abc"> {
+ %c = tensor.cast %t : tensor<64xf32, "abc"> to tensor<?xf32, "abc">
+ %c8 = arith.constant 8 : index
+ %c8b = arith.constant 8 : index
+ %e = tensor.expand_shape %c [[0, 1]] output_shape [%c8, %c8b]
+ : tensor<?xf32, "abc"> into tensor<?x?xf32, "abc">
+ return %e : tensor<?x?xf32, "abc">
+}
+
+// -----
+
+// `VerifiableTensorEncoding` case: expand_shape refined shapes keep the src's
+// and result's ranks respectively, so both sparse encodings pass their
+// (rank-based) `verifyEncoding` and are preserved.
+#sparse_v = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
+#sparse_m = #sparse_tensor.encoding<{
+ map = (d0, d1) -> (d0 : compressed, d1 : compressed)
+}>
+// CHECK-LABEL: func.func @sink_expand_of_cast_preserves_sparse_encoding
+// CHECK: %[[EXPAND:.+]] = tensor.expand_shape
+// CHECK-SAME: tensor<64xf32, #{{[a-z_0-9]+}}> into tensor<8x8xf32, #{{[a-z_0-9]+}}>
+// CHECK: tensor.cast %[[EXPAND]] : tensor<8x8xf32, #{{[a-z_0-9]+}}> to tensor<?x?xf32, #{{[a-z_0-9]+}}>
+func.func @sink_expand_of_cast_preserves_sparse_encoding(
+ %t: tensor<64xf32, #sparse_v>) -> tensor<?x?xf32, #sparse_m> {
+ %c = tensor.cast %t : tensor<64xf32, #sparse_v> to tensor<?xf32, #sparse_v>
+ %c8a = arith.constant 8 : index
+ %c8b = arith.constant 8 : index
+ %e = tensor.expand_shape %c [[0, 1]] output_shape [%c8a, %c8b]
+ : tensor<?xf32, #sparse_v> into tensor<?x?xf32, #sparse_m>
+ return %e : tensor<?x?xf32, #sparse_m>
+}
More information about the Mlir-commits
mailing list