[Mlir-commits] [mlir] e625da3 - [mlir][Tensor] Preserve correct rank expansions in InsertSlice canonicalizer (#217361)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Aug 25 01:13:06 PDT 2026
Author: Artem Gindinson
Date: 2026-08-25T10:12:59+02:00
New Revision: e625da37ad1f548bf07fee9ef525c8dea39a704a
URL: https://github.com/llvm/llvm-project/commit/e625da37ad1f548bf07fee9ef525c8dea39a704a
DIFF: https://github.com/llvm/llvm-project/commit/e625da37ad1f548bf07fee9ef525c8dea39a704a.diff
LOG: [mlir][Tensor] Preserve correct rank expansions in InsertSlice canonicalizer (#217361)
Inferring the source type of an `insert_slice` with constant-folded
arguments "canonically", i.e. by dropping off unit dimensions from the
front of the constant-folded destination shape, is not guaranteed to
match the actual rank transformation of the original `insert_slice` op.
I've encountered this as a loop peeling bug, where materializing
constant slice sizes for the peeled iteration would create the following
sequence:
```
// Before fold
%12 = arith.constant 2 : index
%13 = arith.constant 2 : index // Materialized peel iteration sizes
...
%inserted_slice_13 = tensor.insert_slice %conv into %extracted_slice_9[0, 0, 0, 0] [1, %12, 1, %13] [1, 1, 1, 1] : tensor<1x?x?xf32> into tensor<1x?x1x?xf32>
%inserted_slice_14 = tensor.insert_slice %inserted_slice_13 into %arg13[0, 0, 0, 0] [1, %12, 1, %13] [1, 1, 1, 1] : tensor<1x?x1x?xf32> into tensor<1x?x1x?xf32>
// After a series of incorrect folds
%cast = tensor.cast %conv : tensor<1x2x2xf32> to tensor<1x?x2xf32>
%cast_8 = tensor.cast %cast : tensor<1x?x2xf32> to tensor<?x1x2xf32> // already invalid IR
%cast_9 = tensor.cast %cast_8 : tensor<?x1x2xf32> to tensor<2x1x2xf32>
%inserted_slice_10 = tensor.insert_slice %cast_9 into %arg13[0, 0, 0, 0] [1, 2, 1, 2] [1, 1, 1, 1]
: tensor<2x1x2xf32> into tensor<1x2x1x2xf32> // misinterpreted slice dimensions
```
Instead. we should apply the original rank transformation from the
unfolded op when constructing the source type for the folded result.
Factors out the pre-existing logic from `extract_slice` canonicalization
into a more general helper, so that it can be employed for the
`insert_slice` constant folder too. Additionally, the only remaining
`inferCanonicalRankReducedResultType` use has been replaced with a
specialized call of the new helper.
Assisted-by: Codex
---------
Signed-off-by: Artem Gindinson <gindinson at roofline.ai>
Added:
Modified:
mlir/include/mlir/Dialect/Tensor/IR/Tensor.h
mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
mlir/test/Dialect/Tensor/canonicalize.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h b/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h
index e8e1342ef36fd..005400cddc6d9 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h
+++ b/mlir/include/mlir/Dialect/Tensor/IR/Tensor.h
@@ -26,6 +26,10 @@
#include "mlir/Interfaces/TilingInterface.h"
#include "mlir/Interfaces/ViewLikeInterface.h"
+namespace llvm {
+class SmallBitVector;
+} // namespace llvm
+
//===----------------------------------------------------------------------===//
// Tensor Dialect Helpers
//===----------------------------------------------------------------------===//
@@ -134,6 +138,20 @@ OpFoldResult getMixedSize(OpBuilder &builder, Location loc, Value value,
SmallVector<OpFoldResult> getMixedSizes(OpBuilder &builder, Location loc,
Value value);
+/// Infer a slice type for the given sizes and exact dropped-dimension mask. The
+/// result shape omits the sizes whose corresponding bits are set in
+/// `droppedDims`. The encoding of `sourceTensorType` is propagated to the
+/// inferred result type.
+RankedTensorType inferSliceType(RankedTensorType sourceTensorType,
+ ArrayRef<int64_t> staticSizes,
+ const llvm::SmallBitVector &droppedDims);
+/// SSA-valued sizes resolve to dynamic dimensions in the inferred type. Only
+/// static unit dimensions may be dropped from the source type to produce the
+/// result slice type.
+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/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..c02413bd05d08 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"
@@ -29,6 +30,7 @@
#include "mlir/Transforms/FoldUtils.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/WalkPatternRewriteDriver.h"
+#include "llvm/ADT/SmallBitVector.h"
#include "llvm/Support/Debug.h"
namespace mlir {
@@ -760,9 +762,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 2241acdcbb960..fb14c1953fd22 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2375,7 +2375,6 @@ ExtractSliceOp::inferResultType(RankedTensorType sourceTensorType,
sourceTensorType.getEncoding());
}
-// TODO: This uses neither offsets nor strides!
RankedTensorType
ExtractSliceOp::inferResultType(RankedTensorType sourceTensorType,
ArrayRef<OpFoldResult> sizes) {
@@ -2389,45 +2388,32 @@ ExtractSliceOp::inferResultType(RankedTensorType sourceTensorType,
sourceTensorType.getEncoding());
}
-/// 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
+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 ExtractSliceOp::inferCanonicalRankReducedResultType(
- unsigned desiredResultRank, RankedTensorType sourceRankedTensorType,
- ArrayRef<OpFoldResult> sizes) {
+RankedTensorType
+mlir::tensor::inferSliceType(RankedTensorType sourceTensorType,
+ ArrayRef<OpFoldResult> sizes,
+ const llvm::SmallBitVector &droppedDims) {
SmallVector<int64_t> staticSizes;
- SmallVector<Value> dynamicSizes;
- dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes);
- return ExtractSliceOp::inferCanonicalRankReducedResultType(
- desiredResultRank, sourceRankedTensorType, staticSizes);
+ std::tie(staticSizes, std::ignore) = decomposeMixedValues(sizes);
+ return inferSliceType(sourceTensorType, staticSizes, droppedDims);
}
/// Build an ExtractSliceOp with mixed static and dynamic entries and custom
@@ -2755,29 +2741,15 @@ void mlir::tensor::populateFoldConstantExtractSlicePatterns(
}
/// Return the canonical type of the result of an extract_slice op.
+/// Note: offsets and strides are not needed to determine the result type of
+/// an extract_slice. The operator arguments are just there for interface
+/// compatibility.
struct SliceReturnTypeCanonicalizer {
RankedTensorType operator()(ExtractSliceOp op,
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 +3016,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..d8db9d502fb99 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>
{
@@ -959,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>
@@ -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> {
More information about the Mlir-commits
mailing list