[Mlir-commits] [mlir] 3cbb24b - [mlir][vector] Update CastAway{Extract|Insert}StridedSliceLeadingOneDim (#210902)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jul 23 07:10:32 PDT 2026
Author: Andrzej Warzyński
Date: 2026-07-23T15:10:27+01:00
New Revision: 3cbb24bd5412a5ecea5a0e4d1603e0a38676432d
URL: https://github.com/llvm/llvm-project/commit/3cbb24bd5412a5ecea5a0e4d1603e0a38676432d
DIFF: https://github.com/llvm/llvm-project/commit/3cbb24bd5412a5ecea5a0e4d1603e0a38676432d.diff
LOG: [mlir][vector] Update CastAway{Extract|Insert}StridedSliceLeadingOneDim (#210902)
Update both:
* CastAwayExtractStridedSliceLeadingOneDim
* CastAwayInsertStridedSliceLeadingOneDim
to use vector.shape_cast, rather than vector.extract and
vector.broadcast, as the canonical form for stripping unit dimensions.
This change was originally implemented by @krzysz00 in #196206, but was
subsequently reverted in #199546. This PR intentionally restores only a
subset of #196206, making it easier to identify and triage any potential
regressions.
Co-authored-by: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Added:
Modified:
mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
index 26a702ef0f512..d86e9f224bf9d 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
@@ -78,8 +78,8 @@ struct CastAwayExtractStridedSliceLeadingOneDim
Location loc = extractOp.getLoc();
- Value newSrcVector = vector::ExtractOp::create(
- rewriter, loc, extractOp.getSource(), splatZero(dropCount));
+ Value newSrcVector = rewriter.createOrFold<ShapeCastOp>(
+ loc, newSrcType, extractOp.getSource());
// The offsets/sizes/strides attribute can have a less number of elements
// than the input vector's rank: it is meant for the leading dimensions.
@@ -94,7 +94,7 @@ struct CastAwayExtractStridedSliceLeadingOneDim
rewriter, loc, newDstType, newSrcVector, newOffsets, newSizes,
newStrides);
- rewriter.replaceOpWithNewOp<vector::BroadcastOp>(extractOp, oldDstType,
+ rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(extractOp, oldDstType,
newExtractOp);
return success();
@@ -122,10 +122,10 @@ struct CastAwayInsertStridedSliceLeadingOneDim
// Trim leading one dimensions from both operands.
Location loc = insertOp.getLoc();
- Value newSrcVector = vector::ExtractOp::create(
- rewriter, loc, insertOp.getValueToStore(), splatZero(srcDropCount));
- Value newDstVector = vector::ExtractOp::create(
- rewriter, loc, insertOp.getDest(), splatZero(dstDropCount));
+ Value newSrcVector = rewriter.createOrFold<vector::ShapeCastOp>(
+ loc, newSrcType, insertOp.getValueToStore());
+ Value newDstVector = rewriter.createOrFold<vector::ShapeCastOp>(
+ loc, newDstType, insertOp.getDest());
auto newOffsets = rewriter.getArrayAttr(
insertOp.getOffsets().getValue().take_back(newDstType.getRank()));
@@ -136,7 +136,7 @@ struct CastAwayInsertStridedSliceLeadingOneDim
rewriter, loc, newDstType, newSrcVector, newDstVector, newOffsets,
newStrides);
- rewriter.replaceOpWithNewOp<vector::BroadcastOp>(insertOp, oldDstType,
+ rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(insertOp, oldDstType,
newInsertOp);
return success();
diff --git a/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir b/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
index bf01c8a8589d9..41cd2846fd752 100644
--- a/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
@@ -270,44 +270,45 @@ func.func @cast_away_contraction_leading_one_dims_to_dot_product(%arg0: vector<6
}
// -----
+
// CHECK-LABEL: func @cast_away_extract_strided_slice_leading_one_dims
func.func @cast_away_extract_strided_slice_leading_one_dims(%arg0: vector<1x8x8xf16>) -> vector<1x1x8xf16> {
- // CHECK: %[[SRC:.+]] = vector.extract %{{.*}}[0] : vector<8x8xf16> from vector<1x8x8xf16>
+ // CHECK: %[[SRC:.+]] = vector.shape_cast %{{.*}} : vector<1x8x8xf16> to vector<8x8xf16>
// CHECK: %[[EXTRACT:.+]] = vector.extract_strided_slice %[[SRC]] {offsets = [4], sizes = [1], strides = [1]} : vector<8x8xf16> to vector<1x8xf16>
%0 = vector.extract_strided_slice %arg0 {offsets = [0, 4], sizes = [1, 1], strides = [1, 1]} : vector<1x8x8xf16> to vector<1x1x8xf16>
- // CHECK: %[[RET:.+]] = vector.broadcast %[[EXTRACT]] : vector<1x8xf16> to vector<1x1x8xf16>
+ // CHECK: %[[RET:.+]] = vector.shape_cast %[[EXTRACT]] : vector<1x8xf16> to vector<1x1x8xf16>
// CHECK: return %[[RET]]
return %0: vector<1x1x8xf16>
}
// CHECK-LABEL: func @cast_away_extract_strided_slice_leading_one_dims_scalable
func.func @cast_away_extract_strided_slice_leading_one_dims_scalable(%arg0: vector<1x8x[8]xf16>) -> vector<1x1x[8]xf16> {
- // CHECK: %[[SRC:.+]] = vector.extract %{{.*}}[0] : vector<8x[8]xf16> from vector<1x8x[8]xf16>
+ // CHECK: %[[SRC:.+]] = vector.shape_cast %{{.*}} : vector<1x8x[8]xf16> to vector<8x[8]xf16>
// CHECK: %[[EXTRACT:.+]] = vector.extract_strided_slice %[[SRC]] {offsets = [4], sizes = [1], strides = [1]} : vector<8x[8]xf16> to vector<1x[8]xf16>
%0 = vector.extract_strided_slice %arg0 {offsets = [0, 4], sizes = [1, 1], strides = [1, 1]} : vector<1x8x[8]xf16> to vector<1x1x[8]xf16>
- // CHECK: %[[RET:.+]] = vector.broadcast %[[EXTRACT]] : vector<1x[8]xf16> to vector<1x1x[8]xf16>
+ // CHECK: %[[RET:.+]] = vector.shape_cast %[[EXTRACT]] : vector<1x[8]xf16> to vector<1x1x[8]xf16>
// CHECK: return %[[RET]]
return %0: vector<1x1x[8]xf16>
}
// CHECK-LABEL: func @cast_away_insert_strided_slice_leading_one_dims
func.func @cast_away_insert_strided_slice_leading_one_dims(%arg0: vector<1x8xf16>, %arg1: vector<1x8x8xf16>) -> vector<1x8x8xf16> {
- // CHECK: %[[SRC:.+]] = vector.extract %{{.*}}[0] : vector<8xf16> from vector<1x8xf16>
- // CHECK: %[[DST:.+]] = vector.extract %{{.*}}[0] : vector<8x8xf16> from vector<1x8x8xf16>
+ // CHECK: %[[SRC:.+]] = vector.shape_cast %{{.*}} : vector<1x8xf16> to vector<8xf16>
+ // CHECK: %[[DST:.+]] = vector.shape_cast %{{.*}} : vector<1x8x8xf16> to vector<8x8xf16>
// CHECK: %[[INSERT:.+]] = vector.insert_strided_slice %[[SRC]], %[[DST]] {offsets = [0, 0], strides = [1]} : vector<8xf16> into vector<8x8xf16>
%0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [0, 0, 0], strides = [1, 1]} : vector<1x8xf16> into vector<1x8x8xf16>
- // CHECK: %[[RET:.+]] = vector.broadcast %[[INSERT]] : vector<8x8xf16> to vector<1x8x8xf16>
+ // CHECK: %[[RET:.+]] = vector.shape_cast %[[INSERT]] : vector<8x8xf16> to vector<1x8x8xf16>
// CHECK: return %[[RET]]
return %0: vector<1x8x8xf16>
}
// CHECK-LABEL: func @cast_away_insert_strided_slice_leading_one_dims_scalable
func.func @cast_away_insert_strided_slice_leading_one_dims_scalable(%arg0: vector<1x[8]xf16>, %arg1: vector<1x8x[8]xf16>) -> vector<1x8x[8]xf16> {
- // CHECK: %[[SRC:.+]] = vector.extract %{{.*}}[0] : vector<[8]xf16> from vector<1x[8]xf16>
- // CHECK: %[[DST:.+]] = vector.extract %{{.*}}[0] : vector<8x[8]xf16> from vector<1x8x[8]xf16>
+ // CHECK: %[[SRC:.+]] = vector.shape_cast %{{.*}} : vector<1x[8]xf16> to vector<[8]xf16>
+ // CHECK: %[[DST:.+]] = vector.shape_cast %{{.*}} : vector<1x8x[8]xf16> to vector<8x[8]xf16>
// CHECK: %[[INSERT:.+]] = vector.insert_strided_slice %[[SRC]], %[[DST]] {offsets = [0, 0], strides = [1]} : vector<[8]xf16> into vector<8x[8]xf16>
%0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [0, 0, 0], strides = [1, 1]} : vector<1x[8]xf16> into vector<1x8x[8]xf16>
- // CHECK: %[[RET:.+]] = vector.broadcast %[[INSERT]] : vector<8x[8]xf16> to vector<1x8x[8]xf16>
+ // CHECK: %[[RET:.+]] = vector.shape_cast %[[INSERT]] : vector<8x[8]xf16> to vector<1x8x[8]xf16>
// CHECK: return %[[RET]]
return %0: vector<1x8x[8]xf16>
}
@@ -315,20 +316,18 @@ func.func @cast_away_insert_strided_slice_leading_one_dims_scalable(%arg0: vecto
// CHECK-LABEL: func @cast_away_insert_strided_slice_leading_one_dims_one_element
// CHECK-SAME: %[[ARG0:.+]]: vector<1x1xf16>, %{{.+}}: vector<1x1x1xf16>
func.func @cast_away_insert_strided_slice_leading_one_dims_one_element(%arg0: vector<1x1xf16>, %arg1: vector<1x1x1xf16>) -> vector<1x1x1xf16> {
- // CHECK: %[[EXT:.+]] = vector.extract %{{.*}}[0] : vector<1xf16> from vector<1x1xf16>
- // CHECK: %[[B:.+]] = vector.broadcast %[[EXT]] : vector<1xf16> to vector<1x1x1xf16>
+ // CHECK: %[[RET:.+]] = vector.shape_cast %{{.*}} : vector<1x1xf16> to vector<1x1x1xf16>
%0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [0, 0, 0], strides = [1, 1]} : vector<1x1xf16> into vector<1x1x1xf16>
- // CHECK: return %[[B]]
+ // CHECK: return %[[RET]]
return %0: vector<1x1x1xf16>
}
// CHECK-LABEL: func @cast_away_insert_strided_slice_leading_one_dims_one_element_scalable
// CHECK-SAME: %[[ARG0:.+]]: vector<1x[1]xf16>, %{{.+}}: vector<1x1x[1]xf16>
func.func @cast_away_insert_strided_slice_leading_one_dims_one_element_scalable(%arg0: vector<1x[1]xf16>, %arg1: vector<1x1x[1]xf16>) -> vector<1x1x[1]xf16> {
- // CHECK: %[[EXT:.+]] = vector.extract %{{.*}}[0] : vector<[1]xf16> from vector<1x[1]xf16>
- // CHECK: %[[B:.+]] = vector.broadcast %[[EXT]] : vector<[1]xf16> to vector<1x1x[1]xf16>
+ // CHECK: %[[RET:.+]] = vector.shape_cast %{{.*}} : vector<1x[1]xf16> to vector<1x1x[1]xf16>
%0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [0, 0, 0], strides = [1, 1]} : vector<1x[1]xf16> into vector<1x1x[1]xf16>
- // CHECK: return %[[B]]
+ // CHECK: return %[[RET]]
return %0: vector<1x1x[1]xf16>
}
More information about the Mlir-commits
mailing list