[Mlir-commits] [mlir] 4a876b1 - Add case to handle 0-D vectors in FlattenContiguousRowMajorTransferWritePattern and FlattenContiguousRowMajorTransferReadPattern.

[llvmlistbot at llvm.org]

Author: harsh
Date: 2022-02-08T20:00:12Z
New Revision: 4a876b13fbba8ad9ad7375bd1fe22c7d71c4ad05

URL: https://github.com/llvm/llvm-project/commit/4a876b13fbba8ad9ad7375bd1fe22c7d71c4ad05
DIFF: https://github.com/llvm/llvm-project/commit/4a876b13fbba8ad9ad7375bd1fe22c7d71c4ad05.diff

LOG: Add case to handle 0-D vectors in FlattenContiguousRowMajorTransferWritePattern and FlattenContiguousRowMajorTransferReadPattern.

For 0-D as well as 1-D vectors, both these patterns should
return a failure as there is no need to collapse the shape
of the source. Currently, only 1-D vectors were handled. This
patch handles the 0-D case as well.

Reviewed By: Benoit, ThomasRaoux

Differential Revision: https://reviews.llvm.org/D119202

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
    mlir/test/Dialect/Vector/vector-transfer-flatten.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
index 40897ef53b053..2092d93f3ece1 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
@@ -373,8 +373,8 @@ class FlattenContiguousRowMajorTransferReadPattern
     // Contiguity check is valid on tensors only.
     if (!sourceType)
       return failure();
-    if (vectorType.getRank() == 1 && sourceType.getRank() == 1)
-      // Already 1D, nothing to do.
+    if (vectorType.getRank() <= 1)
+      // Already 0D/1D, nothing to do.
       return failure();
     if (!isStaticShapeAndContiguousRowMajor(sourceType))
       return failure();
@@ -425,8 +425,8 @@ class FlattenContiguousRowMajorTransferWritePattern
     // Contiguity check is valid on tensors only.
     if (!sourceType)
       return failure();
-    if (vectorType.getRank() == 1 && sourceType.getRank() == 1)
-      // Already 1D, nothing to do.
+    if (vectorType.getRank() <= 1)
+      // Already 0D/1D, nothing to do.
       return failure();
     if (!isStaticShapeAndContiguousRowMajor(sourceType))
       return failure();

diff  --git a/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir b/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
index 68a6779461d62..65025b4c9b9c6 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
@@ -33,3 +33,29 @@ func @transfer_write_flattenable_with_offset(
 // C-HECK-DAG:     %[[VEC1D:.+]] = vector.shape_cast %[[VEC]] : vector<5x4x3x2xi8> to vector<120xi8>
 // C-HECK:         vector.transfer_write %[[VEC1D]], %[[COLLAPSED]]
 
+// -----
+
+func @transfer_write_0d(%arg : memref<i8>, %vec : vector<i8>) {
+      vector.transfer_write %vec, %arg[] : vector<i8>, memref<i8>
+      return
+}
+
+// CHECK-LABEL: func @transfer_write_0d
+// CHECK-SAME:       %[[ARG:.+]]: memref<i8>
+// CHECK-SAME:       %[[VEC:.+]]: vector<i8>
+// CHECK:          vector.transfer_write %[[VEC]], %[[ARG]][] : vector<i8>, memref<i8>
+// CHECK:          return
+
+// -----
+
+func @transfer_read_0d(%arg : memref<i8>) -> vector<i8> {
+      %cst = arith.constant 0 : i8
+      %0 = vector.transfer_read %arg[], %cst : memref<i8>, vector<i8>
+      return %0 : vector<i8>
+}
+
+// CHECK-LABEL: func @transfer_read_0d
+// CHECK-SAME:       %[[ARG:.+]]: memref<i8>
+// CHECK:            %[[CST:.+]] = arith.constant 0 : i8
+// CHECK:            %[[READ:.+]] = vector.transfer_read %[[ARG]][], %[[CST]] : memref<i8>
+// CHECK:            return %[[READ]]