[Mlir-commits] [mlir] [mlir][tensor] Preserve source encoding when folding insert_slice canonicalizers (PR #207239)
Dmitrii Makarenko
llvmlistbot at llvm.org
Fri Jul 3 08:01:09 PDT 2026
https://github.com/Devjiu updated https://github.com/llvm/llvm-project/pull/207239
>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 1/2] [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>
>From 4dcd8019b046cc8fc61a9f2e6bbc3a41ac61b102 Mon Sep 17 00:00:00 2001
From: Dmitrii Makarenko <dmitrii.makarenko at intel.com>
Date: Fri, 3 Jul 2026 14:59:09 +0000
Subject: [PATCH 2/2] Align with `VerifiableTensorEncoding` verification
This commit applies same approach with encoding verification
---
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 34 ++++++++++++++++++----
mlir/test/Dialect/Tensor/canonicalize.mlir | 26 +++++++++++++++++
2 files changed, 55 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 8dadc83fdc548..8e7fa47d0b531 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"
@@ -44,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,
@@ -3016,16 +3036,20 @@ class InsertSliceOpConstantArgumentFolder final
// 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).
+ // 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(),
- insertSliceOp.getSourceType().getEncoding());
+ 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 df92d3ccfd82b..cac1b16cd87f7 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -985,6 +985,32 @@ func.func @preserve_source_encoding_on_parallel_insert_slice_folding(
// -----
+// `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