[Mlir-commits] [mlir] 00bdfbf - [mlir][tensor] Preserve source encoding when folding insert_slice canonicalizers (#207239)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Aug 4 05:50:22 PDT 2026
Author: Dmitrii Makarenko
Date: 2026-08-04T12:50:17Z
New Revision: 00bdfbfffd2760a3382dddd36c2b61e0e04fafb5
URL: https://github.com/llvm/llvm-project/commit/00bdfbfffd2760a3382dddd36c2b61e0e04fafb5
DIFF: https://github.com/llvm/llvm-project/commit/00bdfbfffd2760a3382dddd36c2b61e0e04fafb5.diff
LOG: [mlir][tensor] Preserve source encoding when folding insert_slice canonicalizers (#207239)
`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`.
lit created with Claude Opus 4.7
Signed-off-by: Dmitrii Makarenko <dmitrii.makarenko at intel.com>
Added:
Modified:
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
mlir/test/Dialect/Tensor/canonicalize.mlir
Removed:
################################################################################
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