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

[harsh]

Summary

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.

[harsh]

Added unit test

[ThomasRaoux]

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

376–378

Shouldn't the condition be vectorType.getRank() <= 1?

mlir/test/Dialect/Vector/vector-transfer-flatten.mlir

60

does the problem not happen if memref is a higher rank but vector is still rank 0.

[ThomasRaoux]

@Benoit could you take a look as I believe you added those patterns. Why doesn't the memref has to be rank 1 for the pattern to skip? Do you know why the unit tests above are disaled?

[Benoit]

In D119202#3303339, @ThomasRaoux wrote:

@Benoit could you take a look as I believe you added those patterns. Why doesn't the memref has to be rank 1 for the pattern to skip? Do you know why the unit tests above are disaled?

Better ask @nicolasvasilache ! In https://reviews.llvm.org/D114993 I had originally written a larger number of tests. Nicolas took over and disabled/removed tests.

[ThomasRaoux]

In D119202#3303347, @Benoit wrote:

In D119202#3303339, @ThomasRaoux wrote:

@Benoit could you take a look as I believe you added those patterns. Why doesn't the memref has to be rank 1 for the pattern to skip? Do you know why the unit tests above are disaled?

Better ask @nicolasvasilache ! In https://reviews.llvm.org/D114993 I had originally written a larger number of tests. Nicolas took over and disabled/removed tests.

ok, what about the failure condition, why do we have to care about the memref rank?

[harsh]

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

376–378

If the vector rank is 0 and the source rank is 1, then the baseline code (before this patch) works. This is only an issue when
the source rank is 0 and the vector rank is 0. In that case, the assertion I run into is

mlir::AffineExpr mlir::makeCanonicalStridedLayoutExpr(ArrayRef<int64_t>, ArrayRef<mlir::AffineExpr>, mlir::MLIRContext *): Assertion `!sizes.empty() && !exprs.empty() && "expected non-empty sizes and exprs"' failed.

mlir/test/Dialect/Vector/vector-transfer-flatten.mlir

60

yes, for example if the memref is 1d (memref<8xi8> vs memref<i8>), then it works even though the vector is rank 0 in both cases.

[harsh]

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

376–378

In the case that the source rank is 0, I guess we don't need to enforce the vector rank since vector.transfer_read/write semantics enforce that the vector rank will also have to be 0D. So do you want to just check for sourceType.getRank() == 0? Or did you have another idea?

[Benoit]

In D119202#3303350, @ThomasRaoux wrote:

In D119202#3303347, @Benoit wrote:

In D119202#3303339, @ThomasRaoux wrote:

@Benoit could you take a look as I believe you added those patterns. Why doesn't the memref has to be rank 1 for the pattern to skip? Do you know why the unit tests above are disaled?

Better ask @nicolasvasilache ! In https://reviews.llvm.org/D114993 I had originally written a larger number of tests. Nicolas took over and disabled/removed tests.

ok, what about the failure condition, why do we have to care about the memref rank?

Good question, I think I just wrote this condition without really thinking about it.
I made the edit, rebuilt and reran, and so I can confirm, that it does not make a difference for my own usage patterns. Specifically: this diff does not matter to me:

--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
@@ -373,7 +373,7 @@ class FlattenContiguousRowMajorTransferReadPattern
     // Contiguity check is valid on tensors only.
     if (!sourceType)
       return failure();
-    if (vectorType.getRank() == 1 && sourceType.getRank() == 1)
+    if (vectorType.getRank() == 1)

[ThomasRaoux]

In D119202#3305284, @Benoit wrote:

In D119202#3303350, @ThomasRaoux wrote:

In D119202#3303347, @Benoit wrote:

In D119202#3303339, @ThomasRaoux wrote:

@Benoit could you take a look as I believe you added those patterns. Why doesn't the memref has to be rank 1 for the pattern to skip? Do you know why the unit tests above are disaled?

Better ask @nicolasvasilache ! In https://reviews.llvm.org/D114993 I had originally written a larger number of tests. Nicolas took over and disabled/removed tests.

ok, what about the failure condition, why do we have to care about the memref rank?

Good question, I think I just wrote this condition without really thinking about it.
I made the edit, rebuilt and reran, and so I can confirm, that it does not make a difference for my own usage patterns. Specifically: this diff does not matter to me:

--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
@@ -373,7 +373,7 @@ class FlattenContiguousRowMajorTransferReadPattern
     // Contiguity check is valid on tensors only.
     if (!sourceType)
       return failure();
-    if (vectorType.getRank() == 1 && sourceType.getRank() == 1)
+    if (vectorType.getRank() == 1)

Ok then I think we should change the condition to vectorType.getRank() <=1 for simplicity

[harsh]

Sounds good. I will change the condition to vectorType.getRank() <=1

[harsh]

Updated based on Thomas' comments

[ThomasRaoux]

Thanks Harsh!

[harsh]

Thank you! :)

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

LogicalResult matchAndRewrite(vector::TransferReadOp transferReadOp,
     auto loc = transferReadOp.getLoc();
     Value vector = transferReadOp.vector();
     VectorType vectorType = vector.getType().cast<VectorType>();
     Value source = transferReadOp.source();
     MemRefType sourceType = source.getType().dyn_cast<MemRefType>();
     // 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();
     if (getReducedRank(sourceType.getShape()) != sourceType.getRank())
       // This pattern requires the source to already be rank-reduced.
       return failure();
     if (sourceType.getNumElements() != vectorType.getNumElements())
       return failure();
     // TODO: generalize this pattern, relax the requirements here.
LogicalResult matchAndRewrite(vector::TransferWriteOp transferWriteOp,
     auto loc = transferWriteOp.getLoc();
     Value vector = transferWriteOp.vector();
     VectorType vectorType = vector.getType().cast<VectorType>();
     Value source = transferWriteOp.source();
     MemRefType sourceType = source.getType().dyn_cast<MemRefType>();
     // 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();
     if (getReducedRank(sourceType.getShape()) != sourceType.getRank())
       // This pattern requires the source to already be rank-reduced.
       return failure();
     if (sourceType.getNumElements() != vectorType.getNumElements())
       return failure();

mlir/test/Dialect/Vector/vector-transfer-flatten.mlir

 
 // C-HECK-LABEL: func @transfer_write_flattenable_with_offset
 // C-HECK-SAME:      %[[ARG:[0-9a-zA-Z]+]]: memref<5x4x3x2xi8
 // C-HECK-SAME:      %[[VEC:[0-9a-zA-Z]+]]: vector<5x4x3x2xi8>
 // C-HECK-DAG:     %[[COLLAPSED:.+]] = memref.collapse_shape %[[ARG]] {{.}}[0, 1, 2, 3]{{.}} : memref<5x4x3x2xi8, {{.+}}> into memref<120xi8, {{.+}}>
 // 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]]