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
This commit is contained in:
harsh 2022-02-08 19:50:00 +00:00
parent 079b6d02d1
commit 4a876b13fb
2 changed files with 30 additions and 4 deletions

View File

@ -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();

View File

@ -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]]