[Mlir-commits] [mlir] [mlir][vector] Take dim sizes into account in DropInnerMostUnitDims. (PR #71752)

[Han-Chung Wang]

https://github.com/hanhanW updated https://github.com/llvm/llvm-project/pull/71752

>From 1095ab219bbc95080202c227c7c2274d7ca8206e Mon Sep 17 00:00:00 2001
From: hanhanW <hanhan0912 at gmail.com>
Date: Wed, 8 Nov 2023 17:04:09 -0800
Subject: [PATCH 1/3] [mlir][vector] Take vector sizes into account in
 DropInnerMostUnitDims.

The `stride == 1` does not imply that we can drop it. Because it could
be out of bounds. We should also take vector sizes into account.
---
 .../Dialect/Vector/Transforms/VectorTransforms.cpp |  9 ++++++---
 .../vector-transfer-collapse-inner-most-dims.mlir  | 14 ++++++++++++++
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index 16f9d7e4d57eef7..a96d75c5804f555 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1185,9 +1185,12 @@ class DropInnerMostUnitDims : public OpRewritePattern<vector::TransferReadOp> {
       return failure();
 
     size_t dimsToDrop = 0;
-    for (size_t i = 1; i < srcStrides.size(); ++i) {
-      int dim = srcType.getRank() - i - 1;
-      if (srcStrides[dim] == 1) {
+    int rankDiff = srcType.getRank() - readOp.getVectorType().getRank();
+    for (int64_t i = 0; i < targetType.getRank(); ++i) {
+      int dim = targetType.getRank() - i - 1;
+      if (srcStrides[dim + rankDiff] == 1 &&
+          srcType.getDimSize(dim + rankDiff) == 1 &&
+          targetType.getDimSize(dim) == 1) {
         dimsToDrop++;
       } else {
         break;
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 ef0bd9ddf8abf4b..1e97c482dc97ee7 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
@@ -60,3 +60,17 @@ func.func @contiguous_inner_most_dim_bounds_2d(%A: memref<1000x1x1xf32>, %i:inde
 //      CHECK:   %[[V:.+]] = vector.transfer_read %[[SRC_1]]
 // CHECK-SAME:       {in_bounds = [true]}
 // CHECK-SAME:       vector<4xf32>
+
+// -----
+
+func.func @contiguous_inner_most_dim_out_of_bounds_2d(%arg0: memref<1x1xf32>) -> vector<4x8xf32> {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = vector.transfer_read %arg0[%c0, %c0], %cst : memref<1x1xf32>, vector<4x8xf32>
+  return %0 : vector<4x8xf32>
+}
+//      CHECK: func.func @contiguous_inner_most_dim_out_of_bounds_2d
+// CHECK-SAME:   %[[SRC:[a-zA-Z0-9]+]]
+//  CHECK-NOT:   memref.subview
+//      CHECK:   %[[READ:.+]] = vector.transfer_read %[[SRC]]
+//      CHECK:   return %[[READ]] : vector<4x8xf32>

>From e7bfc2acec6db8b5e2c4a7831a04e459a7051823 Mon Sep 17 00:00:00 2001
From: hanhanW <hanhan0912 at gmail.com>
Date: Wed, 8 Nov 2023 17:21:39 -0800
Subject: [PATCH 2/3] use existing variables

---
 mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index a96d75c5804f555..e440babb5619e08 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1185,7 +1185,7 @@ class DropInnerMostUnitDims : public OpRewritePattern<vector::TransferReadOp> {
       return failure();
 
     size_t dimsToDrop = 0;
-    int rankDiff = srcType.getRank() - readOp.getVectorType().getRank();
+    int rankDiff = srcType.getRank() - targetType.getRank();
     for (int64_t i = 0; i < targetType.getRank(); ++i) {
       int dim = targetType.getRank() - i - 1;
       if (srcStrides[dim + rankDiff] == 1 &&

>From 58b4540911b91cbaa75838a5e101acb36ac2f3b8 Mon Sep 17 00:00:00 2001
From: hanhanW <hanhan0912 at gmail.com>
Date: Thu, 9 Nov 2023 16:23:23 -0800
Subject: [PATCH 3/3] add more comments

---
 mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp   | 8 +++++++-
 .../Vector/vector-transfer-collapse-inner-most-dims.mlir  | 2 ++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index e440babb5619e08..713f9cb72c82cec 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1184,9 +1184,15 @@ class DropInnerMostUnitDims : public OpRewritePattern<vector::TransferReadOp> {
     if (failed(getStridesAndOffset(srcType, srcStrides, srcOffset)))
       return failure();
 
+    // According to vector.transfer_read semantics, the result can be a slice.
+    // It pads the indices with `1` starting from beginning. Thus, we have to
+    // offset the check index with `rankDiff` in `srcStrides` and source dim
+    // sizes.
     size_t dimsToDrop = 0;
     int rankDiff = srcType.getRank() - targetType.getRank();
-    for (int64_t i = 0; i < targetType.getRank(); ++i) {
+    for (int64_t i = 0, e = targetType.getRank(); i < e; ++i) {
+      // Check that the inner dim size is 1 for both memref/tensor type and
+      // vector slice. It can be folded only if they are 1 and the stride is 1.
       int dim = targetType.getRank() - i - 1;
       if (srcStrides[dim + rankDiff] == 1 &&
           srcType.getDimSize(dim + rankDiff) == 1 &&
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 1e97c482dc97ee7..0d2743b9fe2e7f5 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
@@ -69,6 +69,8 @@ func.func @contiguous_inner_most_dim_out_of_bounds_2d(%arg0: memref<1x1xf32>) ->
   %0 = vector.transfer_read %arg0[%c0, %c0], %cst : memref<1x1xf32>, vector<4x8xf32>
   return %0 : vector<4x8xf32>
 }
+// The inner most unit dim can not be dropped. In this context, we do not
+// generate rank-reduced memref.subview ops.
 //      CHECK: func.func @contiguous_inner_most_dim_out_of_bounds_2d
 // CHECK-SAME:   %[[SRC:[a-zA-Z0-9]+]]
 //  CHECK-NOT:   memref.subview