[Mlir-commits] [mlir] [mlir][tensor] Preserve encoding in more canonicalizers (concat, reshape, pad) (PR #207241)
Dmitrii Makarenko
llvmlistbot at llvm.org
Thu Jul 2 10:59:22 PDT 2026
https://github.com/Devjiu created https://github.com/llvm/llvm-project/pull/207241
This commit fixes encoding drop in similar with InsertSliceOpConstantArgumentFolder cases. The same anti-pattern — silently producing a `RankedTensorType` without an encoding while a source value with an encoding flows through the refined type — is present in several other canonicalizers/inferResultType implementations:
* ConcatOp::inferResultType built the result type with no encoding. When all inputs share an encoding, propagate it; when they differ, drop it (concat semantics are undefined across differing encodings, so we do not pick one).
* InferConcatOperandTypes built per-operand refined types with the result element type but no encoding. The cast targets an operand value, so it must carry that operand's own encoding.
* CollapseShapeOp::inferCollapsedType dropped the encoding. Since collapse changes the tensor rank, `VerifiableTensorEncoding` attrs (e.g. sparse tensor encoding) would become invalid on the new shape; those are dropped as before. Encodings that do not implement `VerifiableTensorEncoding` (opaque dict/string/bounds attrs used by downstream dialects) are rank-agnostic and are now preserved.
* ConvertToStaticExpandShape built the inserted `tensor.cast` target and the refined `expand_shape` result with no encoding. Both types preserve the rank of the value flowing through them, so they now carry the source's and result's encodings respectively.
* PadOp::inferResultType dropped the encoding. Pad preserves rank, so the source encoding stays valid. FoldSourceTensorCast and FoldStaticPadding inherit the fix through inferResultType; FoldStaticPadding additionally rebuilt the type directly and is updated to propagate the pad's result encoding.
>From 2d3a314f3d670b262df33319ce0cb6b7e223d469 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
(concat, reshape, pad)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
This commit fixes encoding drop in similar with InsertSliceOpConstantArgumentFolder cases.
The same anti-pattern — silently producing a `RankedTensorType` without an
encoding while a source value with an encoding flows through the
refined type — is present in several other canonicalizers/inferResultType
implementations:
* ConcatOp::inferResultType built the result type with no encoding. When all
inputs share an encoding, propagate it; when they differ, drop it (concat
semantics are undefined across differing encodings, so we do not pick one).
* InferConcatOperandTypes built per-operand refined types with the result
element type but no encoding. The cast targets an operand value, so it
must carry that operand's own encoding.
* CollapseShapeOp::inferCollapsedType dropped the encoding. Since collapse
changes the tensor rank, `VerifiableTensorEncoding` attrs (e.g. sparse
tensor encoding) would become invalid on the new shape; those are dropped
as before. Encodings that do not implement `VerifiableTensorEncoding`
(opaque dict/string/bounds attrs used by downstream dialects) are
rank-agnostic and are now preserved.
* ConvertToStaticExpandShape built the inserted `tensor.cast` target and
the refined `expand_shape` result with no encoding. Both types preserve
the rank of the value flowing through them, so they now carry the source's
and result's encodings respectively.
* PadOp::inferResultType dropped the encoding. Pad preserves rank, so the
source encoding stays valid. FoldSourceTensorCast and FoldStaticPadding
inherit the fix through inferResultType; FoldStaticPadding additionally
rebuilt the type directly and is updated to propagate the pad's result
encoding.
Signed-off-by: Dmitrii Makarenko <dmitrii.makarenko at intel.com>
---
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 45 ++++++++---
mlir/test/Dialect/Tensor/canonicalize.mlir | 93 +++++++++++++++++++++-
2 files changed, 125 insertions(+), 13 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 637366a289ac9..2406df1db81d0 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"
@@ -577,7 +578,17 @@ RankedTensorType ConcatOp::inferResultType(int64_t dim, TypeRange inputTypes) {
concatSize =
concatSize + SaturatedInteger::wrap(tensorType.getDimSize(dim));
sizes[dim] = concatSize.asInteger();
- return RankedTensorType::get(sizes, tensorTypes[0].getElementType());
+ // Preserve the encoding when all inputs share it; otherwise drop it (the
+ // semantics of concatenating tensors with differing encodings are undefined
+ // at this level, so we don't try to pick one).
+ Attribute encoding = tensorTypes[0].getEncoding();
+ for (auto tensorType : llvm::drop_begin(tensorTypes)) {
+ if (tensorType.getEncoding() != encoding) {
+ encoding = Attribute();
+ break;
+ }
+ }
+ return RankedTensorType::get(sizes, tensorTypes[0].getElementType(), encoding);
}
void ConcatOp::build(OpBuilder &builder, OperationState &result, int64_t dim,
@@ -811,11 +822,14 @@ struct InferConcatOperandTypes : public OpRewritePattern<ConcatOp> {
SmallVector<int64_t> inferredOperandShape(inferredResultType.getShape());
for (auto [operandIdx, operandType] :
llvm::enumerate(concatOp->getOperandTypes())) {
- // Compute inferred type for operand.
- inferredOperandShape[dim] =
- cast<RankedTensorType>(operandType).getDimSize(dim);
+ // Compute inferred type for operand. The refined type is applied to the
+ // operand itself, so it must carry the operand's own encoding rather
+ // than the (potentially different or missing) result encoding.
+ auto operandRankedType = cast<RankedTensorType>(operandType);
+ inferredOperandShape[dim] = operandRankedType.getDimSize(dim);
auto inferredOperandType = RankedTensorType::get(
- inferredOperandShape, inferredResultType.getElementType());
+ inferredOperandShape, inferredResultType.getElementType(),
+ operandRankedType.getEncoding());
// Check if inferred type is more static.
if (!preservesStaticInformation(inferredOperandType, operandType)) {
@@ -2023,7 +2037,10 @@ CollapseShapeOp::inferCollapsedType(RankedTensorType type,
currentDim += dim;
}
- return RankedTensorType::get(newShape, type.getElementType());
+ Attribute encoding = type.getEncoding();
+ if (llvm::isa_and_present<VerifiableTensorEncoding>(encoding))
+ encoding = {};
+ return RankedTensorType::get(newShape, type.getElementType(), encoding);
}
void CollapseShapeOp::build(OpBuilder &b, OperationState &result, Value src,
@@ -2273,10 +2290,13 @@ struct ConvertToStaticExpandShape : public OpRewritePattern<ExpandShapeOp> {
SmallVector<OpFoldResult> outputOfr =
getMixedValues(newOutputShape, dynamicOutputShape, rewriter);
+ // The refined types keep the ranks of the src / result respectively
auto inputType = RankedTensorType::get(
- newInputShape, expandOp.getSrcType().getElementType());
+ newInputShape, expandOp.getSrcType().getElementType(),
+ expandOp.getSrcType().getEncoding());
auto outputType = RankedTensorType::get(
- newOutputShape, expandOp.getSrcType().getElementType());
+ newOutputShape, expandOp.getSrcType().getElementType(),
+ expandOp.getResultType().getEncoding());
auto inputCast = CastOp::create(rewriter, expandOp.getLoc(), inputType,
expandOp.getSrc());
auto newExpand = ExpandShapeOp::create(
@@ -2385,7 +2405,7 @@ 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 +3343,8 @@ RankedTensorType PadOp::inferResultType(RankedTensorType sourceType,
}
}
- return RankedTensorType::get(inferredShape, sourceType.getElementType());
+ return RankedTensorType::get(inferredShape, sourceType.getElementType(),
+ sourceType.getEncoding());
}
void PadOp::build(OpBuilder &b, OperationState &result, Type resultType,
@@ -3730,9 +3751,9 @@ struct FoldStaticPadding : public OpRewritePattern<PadOp> {
[&](int64_t x) { return x == ShapedType::kDynamic; }))
return failure();
- // Rewrite the op using the new static type.
auto newResultType = RankedTensorType::get(
- newOutDims, padTensorOp.getType().getElementType());
+ newOutDims, padTensorOp.getType().getElementType(),
+ padTensorOp.getType().getEncoding());
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..c5fc71e6c7bd5 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -162,6 +162,22 @@ func.func @infer_concat_return_type(%arg0: tensor<5x12xi32>, %arg1: tensor<?x12x
// -----
+// ConcatOp::inferResultType must carry the (uniformly-shared) operand encoding
+// onto both the refined operand cast and the refined ConcatOp result.
+// CHECK-LABEL: concat_preserves_uniform_encoding
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<3x?xi32, "abc">
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xi32, "abc">
+// CHECK: %[[CAST:.+]] = tensor.cast %[[ARG1]] : tensor<?x?xi32, "abc"> to tensor<3x?xi32, "abc">
+// CHECK: tensor.concat dim(1) %[[ARG0]], %[[CAST]] : (tensor<3x?xi32, "abc">, tensor<3x?xi32, "abc">) -> tensor<3x?xi32, "abc">
+func.func @concat_preserves_uniform_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 +1108,22 @@ func.func @collapse_of_cast(%t: tensor<8x12x32xf32>) -> tensor<?x32xf32> {
// -----
+// A user-defined (non-VerifiableTensorEncoding) encoding must be preserved
+// through the collapse_of_cast folder; inferCollapsedType propagates it
+// alongside the refined shape.
+// CHECK-LABEL: func.func @collapse_of_cast_preserves_encoding(
+// CHECK-SAME: %[[IN:.*]]: tensor<8x12x32xf32, "abc">
+// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[IN]] {{\[}}[0, 1], [2]] : tensor<8x12x32xf32, "abc"> into tensor<96x32xf32, "abc">
+// CHECK: tensor.cast %[[COLLAPSE]] : tensor<96x32xf32, "abc"> to tensor<?x32xf32>
+func.func @collapse_of_cast_preserves_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 +1805,28 @@ 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">
+}
+
+// -----
+
// 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 +1975,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 +2715,25 @@ 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">
+}
More information about the Mlir-commits
mailing list