[Mlir-commits] [mlir] [mlir][Tensor] Preserve correct rank expansions in InsertSlice canonicalizer (PR #217361)
Artem Gindinson
llvmlistbot at llvm.org
Thu Aug 20 01:30:01 PDT 2026
https://github.com/AGindinson updated https://github.com/llvm/llvm-project/pull/217361
>From 6097d8c1d351b1f1ecec2ab8820252a1655af024 Mon Sep 17 00:00:00 2001
From: Artem Gindinson <gindinson at roofline.ai>
Date: Wed, 19 Aug 2026 13:56:05 +0000
Subject: [PATCH 1/2] [mlir][Tensor] Preserve correct rank expansions in
`InsertSliceOpConstantArgumentFolder`
For the source type of an `insert_slice` with constant-folded arguments
to be inferable "canonically", i.e. by dropping off unit dimensions from
the constant-folded destination type, the unit dimension positions actually
have to match the rank-reduced ones in the source slice. This is a wrong
assumption, so instead we should apply the original rank transformation
from the unfolded op when constructing the source type for the folded one.
Factors out the pre-existing logic which `extract_slice` canonicalization
uses into a more general helper, so that it can be employed for the main
change too.
Assisted-by: Codex
Signed-off-by: Artem Gindinson <gindinson at roofline.ai>
---
mlir/include/mlir/Dialect/Tensor/IR/Tensor.h | 11 +++
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 65 ++++++++---------
mlir/test/Dialect/Tensor/canonicalize.mlir | 76 ++++++++++++++++++++
3 files changed, 118 insertions(+), 34 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h b/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h
index e8e1342ef36fd..c383e0fb34b1b 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h
+++ b/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h
@@ -25,6 +25,7 @@
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Interfaces/TilingInterface.h"
#include "mlir/Interfaces/ViewLikeInterface.h"
+#include "llvm/ADT/SmallBitVector.h"
//===----------------------------------------------------------------------===//
// Tensor Dialect Helpers
@@ -134,6 +135,16 @@ OpFoldResult getMixedSize(OpBuilder &builder, Location loc, Value value,
SmallVector<OpFoldResult> getMixedSizes(OpBuilder &builder, Location loc,
Value value);
+/// Infer a slice type with an exact rank-reduction pattern. The source tensor
+/// type provides the element type and encoding; its rank may differ from the
+/// number of sizes when dimensions are dropped.
+RankedTensorType inferSliceType(RankedTensorType sourceTensorType,
+ ArrayRef<int64_t> staticSizes,
+ const llvm::SmallBitVector &droppedDims);
+RankedTensorType inferSliceType(RankedTensorType sourceTensorType,
+ ArrayRef<OpFoldResult> sizes,
+ const llvm::SmallBitVector &droppedDims);
+
/// Create a rank-reducing ExtractSliceOp @[0 .. 0] with strides [1 .. 1] and
/// appropriate sizes (i.e. `tensor.getSizes()`) to reduce the rank of `tensor`
/// to that of `targetType`.
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 2241acdcbb960..99ba0bc4d5e0d 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2389,6 +2389,34 @@ ExtractSliceOp::inferResultType(RankedTensorType sourceTensorType,
sourceTensorType.getEncoding());
}
+RankedTensorType
+mlir::tensor::inferSliceType(RankedTensorType sourceTensorType,
+ ArrayRef<int64_t> staticSizes,
+ const llvm::SmallBitVector &droppedDims) {
+ assert(staticSizes.size() == droppedDims.size() &&
+ "expected one dropped-dimension bit per size");
+
+ SmallVector<int64_t> resultShape;
+ resultShape.reserve(staticSizes.size() - droppedDims.count());
+ for (auto [idx, size] : llvm::enumerate(staticSizes))
+ if (!droppedDims.test(idx))
+ resultShape.push_back(size);
+
+ Type elementType = sourceTensorType.getElementType();
+ return RankedTensorType::get(resultShape, elementType,
+ propagateEncoding(sourceTensorType.getEncoding(),
+ resultShape, elementType));
+}
+
+RankedTensorType
+mlir::tensor::inferSliceType(RankedTensorType sourceTensorType,
+ ArrayRef<OpFoldResult> sizes,
+ const llvm::SmallBitVector &droppedDims) {
+ SmallVector<int64_t> staticSizes;
+ std::tie(staticSizes, std::ignore) = decomposeMixedValues(sizes);
+ return inferSliceType(sourceTensorType, staticSizes, droppedDims);
+}
+
/// If the rank is reduced (i.e. the desiredResultRank is smaller than the
/// number of sizes), drop as many size 1 as needed to produce an inferred
/// type with the desired rank.
@@ -2760,24 +2788,7 @@ struct SliceReturnTypeCanonicalizer {
ArrayRef<OpFoldResult> mixedOffsets,
ArrayRef<OpFoldResult> mixedSizes,
ArrayRef<OpFoldResult> mixedStrides) {
- // Infer a tensor type without taking into account any rank reductions.
- RankedTensorType nonReducedType =
- ExtractSliceOp::inferResultType(op.getSourceType(), mixedSizes);
-
- // Directly return the non-rank reduced type if there are no dropped
- // dims.
- llvm::SmallBitVector droppedDims = op.getDroppedDims();
- if (droppedDims.none())
- return nonReducedType;
-
- // Build the reduced shape, preserving the original rank reduction pattern.
- SmallVector<int64_t> targetShape;
- for (auto i : llvm::seq<int64_t>(mixedSizes.size()))
- if (!droppedDims.test(i))
- targetShape.push_back(nonReducedType.getDimSize(i));
-
- return RankedTensorType::get(targetShape, nonReducedType.getElementType(),
- nonReducedType.getEncoding());
+ return inferSliceType(op.getSourceType(), mixedSizes, op.getDroppedDims());
}
};
@@ -3044,22 +3055,8 @@ class InsertSliceOpConstantArgumentFolder final
if (!sliceResult.isValid)
return failure();
- // 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()));
+ auto sourceType = inferSliceType(insertSliceOp.getSourceType(), mixedSizes,
+ insertSliceOp.getDroppedDims());
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 dab8c6e6ed2ec..ee7a7650a43c9 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -826,6 +826,56 @@ func.func @rank_reducing_insert_slice_canonicalize(%arg0 : tensor<?x?xf32>, %arg
// -----
+// The trailing unit dimension is rank-reduced in the inserted op. When the
+// first dimension folds to unit, make sure that the rank-reduction pattern
+// is preserved.
+// CHECK-LABEL: func @rank_reducing_insert_slice_preserves_shapes
+// CHECK-SAME: %[[SRC:.+]]: tensor<?x2xf32>
+// CHECK-SAME: %[[DST:.+]]: tensor<?x2x1xf32>
+// CHECK: %[[CAST:.+]] = tensor.cast %[[SRC]] : tensor<?x2xf32> to tensor<1x2xf32>
+// CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[CAST]] into %[[DST]][0, 0, 0] [1, 2, 1] [1, 1, 1]
+// CHECK-SAME: : tensor<1x2xf32> into tensor<?x2x1xf32>
+// CHECK: return %[[RESULT]]
+func.func @rank_reducing_insert_slice_preserves_shapes(
+ %src: tensor<?x2xf32>, %dst: tensor<?x2x1xf32>) -> tensor<?x2x1xf32> {
+ %c1 = arith.constant 1 : index
+ %r = tensor.insert_slice %src into %dst[0, 0, 0] [%c1, 2, 1] [1, 1, 1]
+ : tensor<?x2xf32> into tensor<?x2x1xf32>
+ return %r : tensor<?x2x1xf32>
+}
+
+// -----
+
+// A non-leading unit dimension is rank-reduced in the source op being inserted.
+// Verify that constant sizes are correctly folded through consecutive insert_slice
+// ops, preserving cast compatibility between all shapes.
+// CHECK-LABEL: func @rank_reducing_insert_slice_preserves_shapes_consecutive
+// CHECK-SAME: %[[SRC:[a-zA-Z0-9_]+]]: tensor<1x2x2xf32>
+// CHECK-SAME: %[[INNER:[a-zA-Z0-9_]+]]: tensor<1x?x1x?xf32>
+// CHECK-SAME: %[[OUTER:[a-zA-Z0-9_]+]]: tensor<1x?x1x?xf32>
+// CHECK: %[[FIRST:.+]] = tensor.insert_slice %[[SRC]] into %[[INNER]][0, 0, 0, 0] [1, 2, 1, 2] [1, 1, 1, 1]
+// CHECK-SAME: : tensor<1x2x2xf32> into tensor<1x?x1x?xf32>
+// CHECK: %[[CAST:.+]] = tensor.cast %[[FIRST]] : tensor<1x?x1x?xf32> to tensor<1x2x1x2xf32>
+// CHECK: %[[SECOND:.+]] = tensor.insert_slice %[[CAST]] into %[[OUTER]][0, 0, 0, 0] [1, 2, 1, 2] [1, 1, 1, 1]
+// CHECK-SAME: : tensor<1x2x1x2xf32> into tensor<1x?x1x?xf32>
+// CHECK: return %[[SECOND]]
+func.func @rank_reducing_insert_slice_preserves_shapes_consecutive(
+ %src: tensor<1x2x2xf32>, %inner: tensor<1x?x1x?xf32>,
+ %outer: tensor<1x?x1x?xf32>) -> tensor<1x?x1x?xf32> {
+ %c2a = arith.constant 2 : index
+ %c2b = arith.constant 2 : index
+ %cast = tensor.cast %src : tensor<1x2x2xf32> to tensor<1x?x?xf32>
+ %first = tensor.insert_slice %cast into %inner[0, 0, 0, 0]
+ [1, %c2a, 1, %c2b] [1, 1, 1, 1]
+ : tensor<1x?x?xf32> into tensor<1x?x1x?xf32>
+ %second = tensor.insert_slice %first into %outer[0, 0, 0, 0]
+ [1, %c2a, 1, %c2b] [1, 1, 1, 1]
+ : tensor<1x?x1x?xf32> into tensor<1x?x1x?xf32>
+ return %second : tensor<1x?x1x?xf32>
+}
+
+// -----
+
func.func @rank_reducing_slice_to_insert_slice_canonicalize(%arg0 : tensor<?x?x?xf32>, %arg1 : index,
%arg2 : index, %arg3 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
{
@@ -2397,6 +2447,32 @@ func.func @canonicalize_parallel_insert_slice_indices(
// -----
+// The trailing unit dimension is rank-reduced in the inserted op. When the
+// first dimension folds to unit, make sure that the rank-reduction pattern
+// is preserved.
+// CHECK-LABEL: func @canonicalize_rank_reducing_parallel_insert_slice
+// CHECK-SAME: %[[SRC:.+]]: tensor<?x2xf32>
+// CHECK: scf.forall
+// CHECK: %[[CAST:.+]] = tensor.cast %[[SRC]] : tensor<?x2xf32> to tensor<1x2xf32>
+// CHECK-NEXT: scf.forall.in_parallel {
+// CHECK-NEXT: tensor.parallel_insert_slice %[[CAST]] into %{{.+}}[0, 0, 0] [1, 2, 1] [1, 1, 1]
+// CHECK-SAME: : tensor<1x2xf32> into tensor<?x2x1xf32>
+func.func @canonicalize_rank_reducing_parallel_insert_slice(
+ %src: tensor<?x2xf32>, %dst: tensor<?x2x1xf32>,
+ %num_threads: index) -> tensor<?x2x1xf32> {
+ %c1 = arith.constant 1 : index
+ %r = scf.forall (%tid) in (%num_threads) shared_outs(%o = %dst)
+ -> (tensor<?x2x1xf32>) {
+ scf.forall.in_parallel {
+ tensor.parallel_insert_slice %src into %o[0, 0, 0] [%c1, 2, 1] [1, 1, 1]
+ : tensor<?x2xf32> into tensor<?x2x1xf32>
+ }
+ }
+ return %r : tensor<?x2x1xf32>
+}
+
+// -----
+
// CHECK-LABEL: func.func @fold_insert_slice_after_extract_slice
// CHECK-SAME: (%[[INPUT:.+]]: tensor<1x2x2x4xf32>)
func.func @fold_insert_slice_after_extract_slice(%input: tensor<1x2x2x4xf32>) -> tensor<1x2x2x4xf32> {
>From 68a3de05a0658c2d5e18e03b90a5f69d416a55a6 Mon Sep 17 00:00:00 2001
From: Artem Gindinson <gindinson at roofline.ai>
Date: Thu, 20 Aug 2026 08:14:11 +0000
Subject: [PATCH 2/2] [NFC] Replace inferCanonicalRankReducedResultType
The new `inferSliceSizes` helper can be simply specified to drop off
unit dimensions at the leftmost positions. Employ that and drop the
original helper.
Signed-off-by: Artem Gindinson <gindinson at roofline.ai>
---
.../mlir/Dialect/Tensor/IR/TensorOps.td | 21 +---------
.../Linalg/Transforms/DropUnitDims.cpp | 10 +++--
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 41 -------------------
mlir/test/Dialect/Tensor/canonicalize.mlir | 4 +-
4 files changed, 11 insertions(+), 65 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
index c9b858519a592..a711402609775 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
+++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
@@ -402,8 +402,8 @@ def Tensor_ExtractSliceOp : Tensor_OpWithOffsetSizesAndStrides<"extract_slice",
Note that there may be multiple ways to infer a resulting rank-reduced type.
e.g. 1x6x1 could potentially rank-reduce to either 1x6 or 6x1 2-D shapes.
- To disambiguate, the inference helpers `inferCanonicalRankReducedResultType`
- only drop the first unit dimensions, in order:
+ To disambiguate, canonical rank-reduction inference drops only the first
+ unit dimensions, in order:
e.g. 1x6x1 rank-reduced to 2-D will infer the 6x1 2-D shape, but not 1x6.
Verification however has access to result type and does not need to infer.
@@ -503,23 +503,6 @@ def Tensor_ExtractSliceOp : Tensor_OpWithOffsetSizesAndStrides<"extract_slice",
RankedTensorType sourceTensorType,
ArrayRef<OpFoldResult> staticSizes);
- /// If the rank is reduced (i.e. the desiredResultRank is smaller than the
- /// number of sizes), drop as many size 1 as needed to produce an inferred type
- /// with the desired rank.
- ///
- /// Note that there may be multiple ways to compute this rank-reduced type:
- /// e.g. 1x6x1 can rank-reduce to either 1x6 or 6x1 2-D tensors.
- ///
- /// To disambiguate, this function always drops the first 1 sizes occurrences.
- static RankedTensorType inferCanonicalRankReducedResultType(
- unsigned resultRank,
- RankedTensorType sourceRankedTensorType,
- ArrayRef<int64_t> staticSizes);
- static RankedTensorType inferCanonicalRankReducedResultType(
- unsigned resultRank,
- RankedTensorType sourceRankedTensorType,
- ArrayRef<OpFoldResult> staticSizes);
-
/// Return the expected rank of each of the`static_offsets`, `static_sizes`
/// and `static_strides` attributes.
std::array<unsigned, 3> getArrayAttrMaxRanks() {
diff --git a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
index c3dca148b7f94..ee74b82f9d5c5 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
@@ -16,6 +16,7 @@
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
@@ -760,9 +761,12 @@ struct RankReducedExtractSliceOp
SmallVector<OpFoldResult> offsets = sliceOp.getMixedOffsets();
SmallVector<OpFoldResult> strides = sliceOp.getMixedStrides();
SmallVector<OpFoldResult> sizes = sliceOp.getMixedSizes();
- auto rankReducedType = cast<RankedTensorType>(
- tensor::ExtractSliceOp::inferCanonicalRankReducedResultType(
- reassociation->size(), sliceOp.getSourceType(), sizes));
+ SmallVector<int64_t> staticSizes;
+ std::tie(staticSizes, std::ignore) = decomposeMixedValues(sizes);
+ llvm::SmallBitVector droppedDims = getPositionsOfShapeOne(
+ sizes.size() - reassociation->size(), staticSizes);
+ RankedTensorType rankReducedType =
+ tensor::inferSliceType(sliceOp.getSourceType(), sizes, droppedDims);
Location loc = sliceOp.getLoc();
Value newSlice = tensor::ExtractSliceOp::create(
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 99ba0bc4d5e0d..43992e12f9385 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2417,47 +2417,6 @@ mlir::tensor::inferSliceType(RankedTensorType sourceTensorType,
return inferSliceType(sourceTensorType, staticSizes, droppedDims);
}
-/// If the rank is reduced (i.e. the desiredResultRank is smaller than the
-/// number of sizes), drop as many size 1 as needed to produce an inferred
-/// type with the desired rank.
-///
-/// Note that there may be multiple ways to compute this rank-reduced type:
-/// e.g. 1x6x1 can rank-reduce to either 1x6 or 6x1 2-D tensors.
-///
-/// To disambiguate, this function always drops the first 1 sizes occurrences.
-RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType(
- unsigned desiredResultRank, RankedTensorType sourceRankedTensorType,
- ArrayRef<int64_t> sizes) {
- // Type inferred in the absence of rank-reducing behavior.
- auto inferredType = llvm::cast<RankedTensorType>(
- inferResultType(sourceRankedTensorType, sizes));
- int rankDiff = inferredType.getRank() - desiredResultRank;
- if (rankDiff > 0) {
- auto shape = inferredType.getShape();
- llvm::SmallBitVector dimsToProject =
- getPositionsOfShapeOne(rankDiff, shape);
- SmallVector<int64_t> projectedShape;
- // Best effort rank-reducing: drop 1s in order.
- for (unsigned pos = 0, e = shape.size(); pos < e; ++pos)
- if (!dimsToProject.test(pos))
- projectedShape.push_back(shape[pos]);
- inferredType =
- RankedTensorType::get(projectedShape, inferredType.getElementType(),
- inferredType.getEncoding());
- }
- return inferredType;
-}
-
-RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType(
- unsigned desiredResultRank, RankedTensorType sourceRankedTensorType,
- ArrayRef<OpFoldResult> sizes) {
- SmallVector<int64_t> staticSizes;
- SmallVector<Value> dynamicSizes;
- dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes);
- return ExtractSliceOp::inferCanonicalRankReducedResultType(
- desiredResultRank, sourceRankedTensorType, staticSizes);
-}
-
/// Build an ExtractSliceOp with mixed static and dynamic entries and custom
/// result type. If the type passed is nullptr, it is inferred.
void ExtractSliceOp::build(OpBuilder &b, OperationState &result,
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index ee7a7650a43c9..d8db9d502fb99 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -1009,8 +1009,8 @@ 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.
+// destination's encoding (which is `none` here) while deriving the folded
+// source shape.
// 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>
More information about the Mlir-commits
mailing list