Skip to content

Commit 670979c

Browse files
committed
Reapply "[mlir][vector] Drop inner unit dims for transfer ops on dynamic shapes." (llvm#80712)
This reverts commit b4c7152. Regression has been fixed.
1 parent 351f94d commit 670979c

File tree

2 files changed

+27
-6
lines changed

2 files changed

+27
-6
lines changed

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1236,7 +1236,7 @@ class DropInnerMostUnitDimsTransferRead
12361236
return failure();
12371237

12381238
auto srcType = dyn_cast<MemRefType>(readOp.getSource().getType());
1239-
if (!srcType || !srcType.hasStaticShape())
1239+
if (!srcType)
12401240
return failure();
12411241

12421242
if (!readOp.getPermutationMap().isMinorIdentity())
@@ -1260,19 +1260,21 @@ class DropInnerMostUnitDimsTransferRead
12601260
targetType.getElementType());
12611261

12621262
auto loc = readOp.getLoc();
1263+
SmallVector<OpFoldResult> sizes =
1264+
memref::getMixedSizes(rewriter, loc, readOp.getSource());
1265+
SmallVector<OpFoldResult> offsets(srcType.getRank(),
1266+
rewriter.getIndexAttr(0));
1267+
SmallVector<OpFoldResult> strides(srcType.getRank(),
1268+
rewriter.getIndexAttr(1));
12631269
MemRefType resultMemrefType =
12641270
getMemRefTypeWithDroppingInnerDims(rewriter, srcType, dimsToDrop);
1265-
SmallVector<int64_t> offsets(srcType.getRank(), 0);
1266-
SmallVector<int64_t> strides(srcType.getRank(), 1);
1267-
12681271
ArrayAttr inBoundsAttr =
12691272
readOp.getInBounds()
12701273
? rewriter.getArrayAttr(
12711274
readOp.getInBoundsAttr().getValue().drop_back(dimsToDrop))
12721275
: ArrayAttr();
12731276
Value rankedReducedView = rewriter.create<memref::SubViewOp>(
1274-
loc, resultMemrefType, readOp.getSource(), offsets, srcType.getShape(),
1275-
strides);
1277+
loc, resultMemrefType, readOp.getSource(), offsets, sizes, strides);
12761278
auto permMap = getTransferMinorIdentityMap(
12771279
cast<ShapedType>(rankedReducedView.getType()), resultTargetVecType);
12781280
Value result = rewriter.create<vector::TransferReadOp>(

mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,25 @@ func.func @contiguous_inner_most_view(%in: memref<1x1x8x1xf32, strided<[3072, 8,
1616

1717
// -----
1818

19+
func.func @contiguous_outer_dyn_inner_most_view(%in: memref<?x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>) -> vector<1x8x1xf32>{
20+
%c0 = arith.constant 0 : index
21+
%cst = arith.constant 0.0 : f32
22+
%0 = vector.transfer_read %in[%c0, %c0, %c0, %c0], %cst {in_bounds = [true, true, true]} : memref<?x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>, vector<1x8x1xf32>
23+
return %0 : vector<1x8x1xf32>
24+
}
25+
// CHECK: func @contiguous_outer_dyn_inner_most_view(
26+
// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]
27+
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
28+
// CHECK-DAG: %[[D0:.+]] = memref.dim %[[SRC]], %[[C0]]
29+
// CHECK: %[[SRC_0:.+]] = memref.subview %[[SRC]][0, 0, 0, 0] [%[[D0]], 1, 8, 1] [1, 1, 1, 1]
30+
// CHECK-SAME: memref<?x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>> to memref<?x1x8xf32, strided<[3072, 8, 1], offset: ?>>
31+
// CHECK: %[[VEC:.+]] = vector.transfer_read %[[SRC_0]]
32+
// CHECK-SAME: memref<?x1x8xf32, strided<[3072, 8, 1], offset: ?>>, vector<1x8xf32>
33+
// CHECK: %[[RESULT:.+]] = vector.shape_cast %[[VEC]]
34+
// CHECK: return %[[RESULT]]
35+
36+
// -----
37+
1938
func.func @contiguous_inner_most_dim(%A: memref<16x1xf32>, %i:index, %j:index) -> (vector<8x1xf32>) {
2039
%c0 = arith.constant 0 : index
2140
%f0 = arith.constant 0.0 : f32

0 commit comments

Comments
 (0)