[Mlir-commits] [mlir] [mlir][tensor] Preserve source encoding when folding insert_slice canonicalizers (PR #207239)

[Dmitrii Makarenko]

https://github.com/Devjiu created https://github.com/llvm/llvm-project/pull/207239


InsertSliceOpConstantArgumentFolder currently re-derives the refined source type via ExtractSliceOp::inferCanonicalRankReducedResultType, which copies the encoding of the passed-in "source template" — which in this pattern is insertSliceOp.getDestType(). For a static, encoding-less destination this silently drops any encoding the actual source carried. Downstream dialects that stash mandatory metadata in the encoding attribute (upper bounds, layout, sparsity descriptors) lose that metadata during --canonicalize.

Preserve the original source's encoding by rebuilding the refined RankedTensorType with srcType.getEncoding() explicitly. Shape refinement (the `?` -> static direction the pattern already performs) is unchanged and still satisfies preservesStaticInformation.

lit created with Claude Opus 4.7

>From 58b03acba8abd581f08554c4600b03d94031330a 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 currently re-derives the refined source
type via ExtractSliceOp::inferCanonicalRankReducedResultType, which copies
the encoding of the passed-in "source template" — which in this pattern is
insertSliceOp.getDestType(). For a static, encoding-less destination this
silently drops any encoding the actual source carried. Downstream dialects
that stash mandatory metadata in the encoding attribute (upper bounds,
layout, sparsity descriptors) lose that metadata during --canonicalize.

Preserve the original source's encoding by rebuilding the refined
RankedTensorType with srcType.getEncoding() explicitly. Shape refinement
(the `?` -> static direction the pattern already performs) is unchanged
and still satisfies preservesStaticInformation.

Signed-off-by: Dmitrii Makarenko <dmitrii.makarenko at intel.com>
---
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp   | 12 +++++-
 mlir/test/Dialect/Tensor/canonicalize.mlir | 48 ++++++++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 637366a289ac9..8dadc83fdc548 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -3014,10 +3014,18 @@ 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 and must be preserved: insert_slice does not convert between
+    // encodings, so the source's encoding is what the produced cast/op must
+    // carry (dropping it would silently discard downstream metadata such as
+    // bounds, layout, or sparsity descriptors).
+    auto sourceTypeBase = ExtractSliceOp::inferCanonicalRankReducedResultType(
         insertSliceOp.getSourceType().getRank(), insertSliceOp.getDestType(),
         mixedSizes);
+    auto sourceType = RankedTensorType::get(
+        sourceTypeBase.getShape(), sourceTypeBase.getElementType(),
+        insertSliceOp.getSourceType().getEncoding());
     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 67b7ab99c5d18..df92d3ccfd82b 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -937,6 +937,54 @@ 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>
+}
+
+// -----
+
 // 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>