[Mlir-commits] [mlir] [mlir][vector] Use inferRankReducedResultType for subview type inference. (PR #84395)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Mar 7 14:58:36 PST 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Han-Chung Wang (hanhanW)
<details>
<summary>Changes</summary>
Fixes https://github.com/openxla/iree/issues/16475
---
Full diff: https://github.com/llvm/llvm-project/pull/84395.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp (+8-40)
- (modified) mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir (+26-2)
``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index a2d4e216633181..6f6b6dcdad2006 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1255,42 +1255,6 @@ getTransferFoldableInnerUnitDims(MemRefType srcType, VectorType vectorType) {
return result;
}
-/// Returns a MemRef type that drops inner `dimsToDrop` dimensions from
-/// `srcType`. E.g., if `srcType` is memref<512x16x1x1xf32> and `dimsToDrop` is
-/// two, it returns memref<512x16x16> type.
-static MemRefType getMemRefTypeWithDroppingInnerDims(OpBuilder &builder,
- MemRefType srcType,
- size_t dimsToDrop) {
- MemRefLayoutAttrInterface layout = srcType.getLayout();
- if (isa<AffineMapAttr>(layout) && layout.isIdentity()) {
- return MemRefType::get(srcType.getShape().drop_back(dimsToDrop),
- srcType.getElementType(), nullptr,
- srcType.getMemorySpace());
- }
- MemRefLayoutAttrInterface updatedLayout;
- if (auto strided = dyn_cast<StridedLayoutAttr>(layout)) {
- auto strides = llvm::to_vector(strided.getStrides().drop_back(dimsToDrop));
- updatedLayout = StridedLayoutAttr::get(strided.getContext(),
- strided.getOffset(), strides);
- return MemRefType::get(srcType.getShape().drop_back(dimsToDrop),
- srcType.getElementType(), updatedLayout,
- srcType.getMemorySpace());
- }
-
- // Non-strided layout case.
- AffineMap map = srcType.getLayout().getAffineMap();
- int numSymbols = map.getNumSymbols();
- for (size_t i = 0; i < dimsToDrop; ++i) {
- int dim = srcType.getRank() - i - 1;
- map = map.replace(builder.getAffineDimExpr(dim),
- builder.getAffineConstantExpr(0), map.getNumDims() - 1,
- numSymbols);
- }
- return MemRefType::get(srcType.getShape().drop_back(dimsToDrop),
- srcType.getElementType(), updatedLayout,
- srcType.getMemorySpace());
-}
-
/// Drop inner most contiguous unit dimensions from transfer_read operand.
class DropInnerMostUnitDimsTransferRead
: public OpRewritePattern<vector::TransferReadOp> {
@@ -1337,8 +1301,10 @@ class DropInnerMostUnitDimsTransferRead
rewriter.getIndexAttr(0));
SmallVector<OpFoldResult> strides(srcType.getRank(),
rewriter.getIndexAttr(1));
- MemRefType resultMemrefType =
- getMemRefTypeWithDroppingInnerDims(rewriter, srcType, dimsToDrop);
+ auto resultMemrefType =
+ cast<MemRefType>(memref::SubViewOp::inferRankReducedResultType(
+ srcType.getShape().drop_back(dimsToDrop), srcType, offsets, sizes,
+ strides));
ArrayAttr inBoundsAttr =
readOp.getInBounds()
? rewriter.getArrayAttr(
@@ -1421,8 +1387,10 @@ class DropInnerMostUnitDimsTransferWrite
rewriter.getIndexAttr(0));
SmallVector<OpFoldResult> strides(srcType.getRank(),
rewriter.getIndexAttr(1));
- MemRefType resultMemrefType =
- getMemRefTypeWithDroppingInnerDims(rewriter, srcType, dimsToDrop);
+ auto resultMemrefType =
+ cast<MemRefType>(memref::SubViewOp::inferRankReducedResultType(
+ srcType.getShape().drop_back(dimsToDrop), srcType, offsets, sizes,
+ strides));
ArrayAttr inBoundsAttr =
writeOp.getInBounds()
? rewriter.getArrayAttr(
diff --git a/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir b/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir
index 3984f17f9e8cdb..76760e3ee639bb 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir
@@ -43,7 +43,7 @@ func.func @contiguous_inner_most_dim(%A: memref<16x1xf32>, %i:index, %j:index) -
}
// CHECK: func @contiguous_inner_most_dim(%[[SRC:.+]]: memref<16x1xf32>, %[[I:.+]]: index, %[[J:.+]]: index) -> vector<8x1xf32>
// CHECK: %[[SRC_0:.+]] = memref.subview %[[SRC]]
-// CHECK-SAME: memref<16x1xf32> to memref<16xf32>
+// CHECK-SAME: memref<16x1xf32> to memref<16xf32, strided<[1]>>
// CHECK: %[[V:.+]] = vector.transfer_read %[[SRC_0]]
// CHECK: %[[RESULT]] = vector.shape_cast %[[V]] : vector<8xf32> to vector<8x1xf32>
// CHECK: return %[[RESULT]]
@@ -111,7 +111,7 @@ func.func @drop_two_inner_most_dim_for_transfer_write(%arg0: memref<1x512x16x1x1
// CHECK-SAME: %[[IDX:[a-zA-Z0-9]+]]
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[DEST]]
-// CHECK-SAME: memref<1x512x16x1x1xf32> to memref<1x512x16xf32>
+// CHECK-SAME: memref<1x512x16x1x1xf32> to memref<1x512x16xf32, strided<[8192, 16, 1]>>
// CHECK: %[[CAST:.+]] = vector.shape_cast %[[VEC]] : vector<1x16x16x1x1xf32> to vector<1x16x16xf32>
// CHECK: vector.transfer_write %[[CAST]], %[[SUBVIEW]]
// CHECK-SAME: [%[[C0]], %[[IDX]], %[[C0]]]
@@ -169,3 +169,27 @@ func.func @non_unit_strides(%arg0: memref<512x16x1xf32, strided<[8192, 16, 4], o
// The inner most unit dims can not be dropped if the strides are not ones.
// CHECK: func.func @non_unit_strides
// CHECK-NOT: memref.subview
+
+// -----
+
+func.func @contiguous_outer_dyn_inner_most_view2(%a: index, %b: index, %memref: memref<?x?x8x1xf32>) -> vector<8x1xf32> {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ %v = vector.transfer_read %memref[%a, %b, %c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?x8x1xf32>, vector<8x1xf32>
+ return %v : vector<8x1xf32>
+}
+// CHECK: func.func @contiguous_outer_dyn_inner_most_view2(
+// CHECK-SAME: %[[IDX0:[a-zA-Z0-9]+]]
+// CHECK-SAME: %[[IDX1:[a-zA-Z0-9]+]]
+// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[PAD:.+]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[D0:.+]] = memref.dim %[[SRC]], %[[C0]]
+// CHECK: %[[D1:.+]] = memref.dim %[[SRC]], %[[C1]]
+// CHECK: %[[VIEW:.+]] = memref.subview %[[SRC]][0, 0, 0, 0] [%[[D0]], %[[D1]], 8, 1] [1, 1, 1, 1]
+// CHECK-SAME: memref<?x?x8x1xf32> to memref<?x?x8xf32, strided<[?, 8, 1], offset: ?>>
+// CHECK: %[[VEC:.+]] = vector.transfer_read %[[VIEW]]
+// CHECK-SAME: memref<?x?x8xf32, strided<[?, 8, 1], offset: ?>>, vector<8xf32>
+// CHECK: %[[RESULT:.+]] = vector.shape_cast %[[VEC]]
+// CHECK: return %[[RESULT]]
``````````
</details>
https://github.com/llvm/llvm-project/pull/84395
More information about the Mlir-commits
mailing list