[Mlir-commits] [mlir] [mlir][tensor] Preserve source encoding when folding insert_slice canonicalizers (PR #207239)
Dmitrii Makarenko
llvmlistbot at llvm.org
Tue Aug 4 05:36:55 PDT 2026
https://github.com/Devjiu updated https://github.com/llvm/llvm-project/pull/207239
>From c06836e8e768a3b5c6a406de6d62ea19d4a82b9e Mon Sep 17 00:00:00 2001
From: Dmitrii Makarenko <dmitrii.makarenko at intel.com>
Date: Thu, 2 Jul 2026 16:28:11 +0000
Subject: [PATCH] [mlir][tensor] Preserve source encoding when folding
insert_slice canonicalizers
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
`InsertSliceOpConstantArgumentFolder` re-derives the refined source type via
`ExtractSliceOp::inferCanonicalRankReducedResultType`, which copies the encoding
of the passed-in "source template" — in this pattern, `insertSliceOp.getDestType()`.
For a static, encoding-less destination this silently drops any encoding the actual source carried,
which downstream dialects can use for mandatory metadata (upper bounds, layout, sparsity descriptors) lost during `--canonicalize`.
Rebuild the refined source type carrying the original source's encoding.
If the encoding implements `VerifiableTensorEncoding`, re-verify it against the refined
shape and drop it if it no longer holds (e.g. a rank-dependent sparse encoding);
otherwise treat it as opaque/shape-agnostic and propagate it unconditionally.
Shape refinement (the `?` -> static direction) is unchanged and still satisfies `preservesStaticInformation`.
Signed-off-by: Dmitrii Makarenko <dmitrii.makarenko at intel.com>
---
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 16 ++++-
mlir/test/Dialect/Tensor/canonicalize.mlir | 74 ++++++++++++++++++++++
2 files changed, 88 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 680b6c3927ec2..2241acdcbb960 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -3044,10 +3044,22 @@ class InsertSliceOpConstantArgumentFolder final
if (!sliceResult.isValid)
return failure();
- // Create the new op in canonical form.
- auto sourceType = ExtractSliceOp::inferCanonicalRankReducedResultType(
+ // Create the new op in canonical form. The refined shape is inferred from
+ // the destination type, but the encoding is a per-value property of the
+ // source: insert_slice does not convert between encodings, so the
+ // produced cast/op must carry the source's encoding (dropping it would
+ // silently discard downstream metadata such as bounds, layout, or
+ // sparsity descriptors). If the source's encoding no longer holds on the
+ // refined shape (e.g. a `VerifiableTensorEncoding` that self-invalidates),
+ // it is dropped in accordance with the encoding's own contract.
+ auto sourceTypeBase = ExtractSliceOp::inferCanonicalRankReducedResultType(
insertSliceOp.getSourceType().getRank(), insertSliceOp.getDestType(),
mixedSizes);
+ auto sourceType = RankedTensorType::get(
+ sourceTypeBase.getShape(), sourceTypeBase.getElementType(),
+ propagateEncoding(insertSliceOp.getSourceType().getEncoding(),
+ sourceTypeBase.getShape(),
+ sourceTypeBase.getElementType()));
Value toInsert = insertSliceOp.getSource();
if (sourceType != insertSliceOp.getSourceType()) {
OpBuilder::InsertionGuard g(rewriter);
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 26c40863e70dc..dab8c6e6ed2ec 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -957,6 +957,80 @@ func.func @insert_slice_cast_no_fold(%arg0 : tensor<1x?xf32>, %arg1 : tensor<?x?
// -----
+// Verify that the constant-argument folder for insert_slice preserves the
+// source's encoding on the inserted cast, rather than silently picking up the
+// destination's encoding (which is `none` here) via the shape template used by
+// ExtractSliceOp::inferCanonicalRankReducedResultType.
+// CHECK-LABEL: func @preserve_source_encoding_on_insert_slice_folding
+// CHECK-SAME: %[[SRC:[a-zA-Z0-9_]+]]: tensor<1x?x?x32xf16, "abc">
+// CHECK-SAME: %[[DST:[a-zA-Z0-9_]+]]: tensor<1x1280x32x32xf16>
+// CHECK-NOT: tensor.cast %{{.*}} : tensor<{{.*}}, "abc"> to tensor<{{[0-9x?]+}}xf16>
+// CHECK: %[[C:.+]] = tensor.cast %[[SRC]] : tensor<1x?x?x32xf16, "abc"> to tensor<1x48x16x32xf16, "abc">
+// CHECK: tensor.insert_slice %[[C]] into %[[DST]]
+func.func @preserve_source_encoding_on_insert_slice_folding(
+ %src: tensor<1x?x?x32xf16, "abc">,
+ %dst: tensor<1x1280x32x32xf16>) -> tensor<1x1280x32x32xf16> {
+ %c16 = arith.constant 16 : index
+ %sz1 = arith.constant 48 : index
+ %ivC = arith.constant 0 : index
+ %ivH = arith.constant 0 : index
+ %r = tensor.insert_slice %src into %dst[0, %ivC, %ivH, 0] [1, %sz1, %c16, 32] [1, 1, 1, 1]
+ : tensor<1x?x?x32xf16, "abc"> into tensor<1x1280x32x32xf16>
+ return %r : tensor<1x1280x32x32xf16>
+}
+
+// -----
+
+// Same invariant for the parallel_insert_slice variant.
+// CHECK-LABEL: func @preserve_source_encoding_on_parallel_insert_slice_folding
+// CHECK-SAME: %[[SRC:[a-zA-Z0-9_]+]]: tensor<1x?x?x32xf16, "abc">
+// CHECK-SAME: %[[DST:[a-zA-Z0-9_]+]]: tensor<1x1280x32x32xf16>
+// CHECK-NOT: tensor.cast %{{.*}} : tensor<{{.*}}, "abc"> to tensor<{{[0-9x?]+}}xf16>
+// CHECK: %[[C:.+]] = tensor.cast %[[SRC]] : tensor<1x?x?x32xf16, "abc"> to tensor<1x48x16x32xf16, "abc">
+// CHECK: tensor.parallel_insert_slice %[[C]] into
+func.func @preserve_source_encoding_on_parallel_insert_slice_folding(
+ %src: tensor<1x?x?x32xf16, "abc">,
+ %dst: tensor<1x1280x32x32xf16>,
+ %num_threads: index) -> tensor<1x1280x32x32xf16> {
+ %c16 = arith.constant 16 : index
+ %sz1 = arith.constant 48 : index
+ %r = scf.forall (%tid) in (%num_threads) shared_outs(%o = %dst) -> (tensor<1x1280x32x32xf16>) {
+ scf.forall.in_parallel {
+ tensor.parallel_insert_slice %src into %o[0, 0, 0, 0] [1, %sz1, %c16, 32] [1, 1, 1, 1]
+ : tensor<1x?x?x32xf16, "abc"> into tensor<1x1280x32x32xf16>
+ }
+ }
+ return %r : tensor<1x1280x32x32xf16>
+}
+
+// -----
+
+// `VerifiableTensorEncoding` case: insert_slice only refines the source shape
+// (dyn -> static), the source rank is unchanged. The sparse encoding's rank
+// invariant still holds on the refined shape, so `propagateEncoding` accepts
+// it and the inserted cast carries the sparse encoding through the folded op.
+#sparse_dense4_is = #sparse_tensor.encoding<{
+ map = (d0, d1, d2, d3) -> (d0 : dense, d1 : dense, d2 : dense, d3 : dense)
+}>
+// CHECK-LABEL: func @preserve_source_sparse_encoding_on_insert_slice_folding
+// CHECK-SAME: %[[SRC:[a-zA-Z0-9_]+]]: tensor<1x?x?x32xf16, #{{[a-z_0-9]+}}>
+// CHECK-SAME: %[[DST:[a-zA-Z0-9_]+]]: tensor<1x1280x32x32xf16>
+// CHECK: %[[C:.+]] = tensor.cast %[[SRC]] : tensor<1x?x?x32xf16, #{{[a-z_0-9]+}}> to tensor<1x48x16x32xf16, #{{[a-z_0-9]+}}>
+// CHECK: tensor.insert_slice %[[C]] into %[[DST]]
+func.func @preserve_source_sparse_encoding_on_insert_slice_folding(
+ %src: tensor<1x?x?x32xf16, #sparse_dense4_is>,
+ %dst: tensor<1x1280x32x32xf16>) -> tensor<1x1280x32x32xf16> {
+ %c16 = arith.constant 16 : index
+ %sz1 = arith.constant 48 : index
+ %ivC = arith.constant 0 : index
+ %ivH = arith.constant 0 : index
+ %r = tensor.insert_slice %src into %dst[0, %ivC, %ivH, 0] [1, %sz1, %c16, 32] [1, 1, 1, 1]
+ : tensor<1x?x?x32xf16, #sparse_dense4_is> into tensor<1x1280x32x32xf16>
+ return %r : tensor<1x1280x32x32xf16>
+}
+
+// -----
+
// CHECK-LABEL: func @insert_tensor_cast_on_insert_slice_src(
// CHECK-SAME: %[[arg0:.*]]: tensor<?x5x?xf32>, %[[arg1:.*]]: tensor<?x?x?xf32>
// CHECK: %[[cast:.*]] = tensor.cast %[[arg0]] : tensor<?x5x?xf32> to tensor<64x5x64xf32>
More information about the Mlir-commits
mailing list