[Mlir-commits] [mlir] [mlir][tensor] Preserve encoding in more canonicalizers (concat, reshape, pad) (PR #207241)
Dmitrii Makarenko
llvmlistbot at llvm.org
Fri Jul 3 07:43:15 PDT 2026
https://github.com/Devjiu updated https://github.com/llvm/llvm-project/pull/207241
>From 65e00b4915e4894a77b7fd083c1d6aa06a964d2d 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 1/2] [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 | 47 ++++++++---
mlir/test/Dialect/Tensor/canonicalize.mlir | 93 +++++++++++++++++++++-
2 files changed, 127 insertions(+), 13 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 637366a289ac9..2c94a10a8e870 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,18 @@ 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 +823,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 +2038,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 +2291,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 +2406,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 +3345,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 +3753,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">
+}
>From 8e19f97a147696c5cd2a6f30a64f93f584359a23 Mon Sep 17 00:00:00 2001
From: Dmitrii Makarenko <dmitrii.makarenko at intel.com>
Date: Fri, 3 Jul 2026 14:40:16 +0000
Subject: [PATCH 2/2] Add VerifiableTensorEncoding check
This commit checks tensor encoding for `VerifiableTensorEncoding::verifyEncoding`
if possible before assigning it source type, updated by canonicalizer.
Signed-off-by: Dmitrii Makarenko <dmitrii.makarenko at intel.com>
---
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 73 +++++++++++++-----
mlir/test/Dialect/Tensor/canonicalize.mlir | 90 ++++++++++++++++++++++
2 files changed, 142 insertions(+), 21 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 2c94a10a8e870..903eb5f6c411d 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -45,6 +45,25 @@
using namespace mlir;
using namespace mlir::tensor;
+/// The intent is to keep canonicalizers from silently erasing downstream
+/// metadata: a `RankedTensorType` encoding is a per-value property, and the
+/// canonicalizer is not a place to guess whether it still applies — the
+/// encoding's own contract is.
+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,
@@ -578,9 +597,9 @@ RankedTensorType ConcatOp::inferResultType(int64_t dim, TypeRange inputTypes) {
concatSize =
concatSize + SaturatedInteger::wrap(tensorType.getDimSize(dim));
sizes[dim] = concatSize.asInteger();
- // 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).
+ // Only propagate an encoding when all inputs agree on it (concat semantics
+ // across differing encodings are undefined at this level). Let the encoding
+ // itself decide whether it still holds on the new shape.
Attribute encoding = tensorTypes[0].getEncoding();
for (auto tensorType : llvm::drop_begin(tensorTypes)) {
if (tensorType.getEncoding() != encoding) {
@@ -588,8 +607,9 @@ RankedTensorType ConcatOp::inferResultType(int64_t dim, TypeRange inputTypes) {
break;
}
}
- return RankedTensorType::get(sizes, tensorTypes[0].getElementType(),
- encoding);
+ Type elementType = tensorTypes[0].getElementType();
+ return RankedTensorType::get(sizes, elementType,
+ propagateEncoding(encoding, sizes, elementType));
}
void ConcatOp::build(OpBuilder &builder, OperationState &result, int64_t dim,
@@ -825,12 +845,15 @@ struct InferConcatOperandTypes : public OpRewritePattern<ConcatOp> {
llvm::enumerate(concatOp->getOperandTypes())) {
// 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.
+ // than the (potentially different or missing) result encoding, subject
+ // to the encoding still holding on the refined shape.
auto operandRankedType = cast<RankedTensorType>(operandType);
inferredOperandShape[dim] = operandRankedType.getDimSize(dim);
+ Type elementType = inferredResultType.getElementType();
auto inferredOperandType = RankedTensorType::get(
- inferredOperandShape, inferredResultType.getElementType(),
- operandRankedType.getEncoding());
+ inferredOperandShape, elementType,
+ propagateEncoding(operandRankedType.getEncoding(),
+ inferredOperandShape, elementType));
// Check if inferred type is more static.
if (!preservesStaticInformation(inferredOperandType, operandType)) {
@@ -2038,10 +2061,9 @@ CollapseShapeOp::inferCollapsedType(RankedTensorType type,
currentDim += dim;
}
- Attribute encoding = type.getEncoding();
- if (llvm::isa_and_present<VerifiableTensorEncoding>(encoding))
- encoding = {};
- return RankedTensorType::get(newShape, type.getElementType(), encoding);
+ return RankedTensorType::get(
+ newShape, type.getElementType(),
+ propagateEncoding(type.getEncoding(), newShape, type.getElementType()));
}
void CollapseShapeOp::build(OpBuilder &b, OperationState &result, Value src,
@@ -2291,13 +2313,18 @@ struct ConvertToStaticExpandShape : public OpRewritePattern<ExpandShapeOp> {
SmallVector<OpFoldResult> outputOfr =
getMixedValues(newOutputShape, dynamicOutputShape, rewriter);
- // The refined types keep the ranks of the src / result respectively
+ // 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(),
- expandOp.getSrcType().getEncoding());
+ newInputShape, elementType,
+ propagateEncoding(expandOp.getSrcType().getEncoding(), newInputShape,
+ elementType));
auto outputType = RankedTensorType::get(
- newOutputShape, expandOp.getSrcType().getElementType(),
- expandOp.getResultType().getEncoding());
+ newOutputShape, elementType,
+ propagateEncoding(expandOp.getResultType().getEncoding(),
+ newOutputShape, elementType));
auto inputCast = CastOp::create(rewriter, expandOp.getLoc(), inputType,
expandOp.getSrc());
auto newExpand = ExpandShapeOp::create(
@@ -3345,8 +3372,10 @@ RankedTensorType PadOp::inferResultType(RankedTensorType sourceType,
}
}
- return RankedTensorType::get(inferredShape, sourceType.getElementType(),
- sourceType.getEncoding());
+ Type elementType = sourceType.getElementType();
+ return RankedTensorType::get(
+ inferredShape, elementType,
+ propagateEncoding(sourceType.getEncoding(), inferredShape, elementType));
}
void PadOp::build(OpBuilder &b, OperationState &result, Type resultType,
@@ -3753,9 +3782,11 @@ struct FoldStaticPadding : public OpRewritePattern<PadOp> {
[&](int64_t x) { return x == ShapedType::kDynamic; }))
return failure();
+ Type elementType = padTensorOp.getType().getElementType();
auto newResultType = RankedTensorType::get(
- newOutDims, padTensorOp.getType().getElementType(),
- padTensorOp.getType().getEncoding());
+ 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 c5fc71e6c7bd5..a92c94f4cc652 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -178,6 +178,30 @@ func.func @concat_preserves_uniform_encoding(
// -----
+// `VerifiableTensorEncoding` case: concat preserves rank, so the sparse
+// encoding's rank invariant still holds on the refined shape and it is kept.
+// The operand refinement inserts a cast whose target must carry the same
+// sparse encoding (not `null`, which would need re-encoding back later).
+#sparse_dense4 = #sparse_tensor.encoding<{
+ map = (d0, d1, d2, d3) -> (d0 : dense, d1 : dense, d2 : dense, d3 : dense)
+}>
+// CHECK-LABEL: concat_preserves_sparse_encoding
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<3x?x8x8xf32, #{{[a-z_0-9]+}}>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x8x8xf32, #{{[a-z_0-9]+}}>
+// CHECK: %[[CAST:.+]] = tensor.cast %[[ARG1]] : tensor<?x?x8x8xf32, #{{[a-z_0-9]+}}> to tensor<3x?x8x8xf32, #{{[a-z_0-9]+}}>
+// CHECK: tensor.concat dim(1) %[[ARG0]], %[[CAST]]
+// CHECK-SAME: -> tensor<3x?x8x8xf32, #{{[a-z_0-9]+}}>
+func.func @concat_preserves_sparse_encoding(
+ %a: tensor<3x?x8x8xf32, #sparse_dense4>,
+ %b: tensor<?x?x8x8xf32, #sparse_dense4>) -> tensor<3x?x8x8xf32, #sparse_dense4> {
+ %r = tensor.concat dim(1) %a, %b
+ : (tensor<3x?x8x8xf32, #sparse_dense4>, tensor<?x?x8x8xf32, #sparse_dense4>)
+ -> tensor<3x?x8x8xf32, #sparse_dense4>
+ return %r : tensor<3x?x8x8xf32, #sparse_dense4>
+}
+
+// -----
+
// CHECK-LABEL: func @fold_extract
func.func @fold_extract(%arg0 : index) -> (f32, f16, f16, i32, complex<f32>, i32) {
%const_0 = arith.constant 0 : index
@@ -1124,6 +1148,25 @@ func.func @collapse_of_cast_preserves_encoding(%t: tensor<8x12x32xf32, "abc">) -
// -----
+// `VerifiableTensorEncoding` case: collapse changes the rank, so the sparse
+// encoding's rank invariant no longer holds; propagateEncoding asks the
+// encoding via `verifyEncoding` and drops it from the refined type.
+#sparse_dense3 = #sparse_tensor.encoding<{
+ map = (d0, d1, d2) -> (d0 : dense, d1 : dense, d2 : dense)
+}>
+// CHECK-LABEL: func.func @collapse_of_cast_drops_rank_invalid_encoding(
+// CHECK-NOT: collapse_shape {{.*}}#{{[a-z_0-9]+}}
+// CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0, 1], [2]] : tensor<8x12x32xf32, #{{[a-z_0-9]+}}> into tensor<96x32xf32>
+func.func @collapse_of_cast_drops_rank_invalid_encoding(
+ %t: tensor<8x12x32xf32, #sparse_dense3>) -> tensor<?x32xf32> {
+ %0 = tensor.cast %t : tensor<8x12x32xf32, #sparse_dense3> to tensor<?x?x?xf32, #sparse_dense3>
+ %1 = tensor.collapse_shape %0 [[0, 1], [2]] : tensor<?x?x?xf32, #sparse_dense3> into tensor<?x?xf32>
+ %2 = tensor.cast %1 : tensor<?x?xf32> 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>
@@ -1827,6 +1870,30 @@ func.func @pad_after_cast_preserves_encoding(
// -----
+// `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
@@ -2737,3 +2804,26 @@ func.func @sink_expand_of_cast_preserves_encoding(%t: tensor<64xf32, "abc">) ->
: 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