[mlir][Vector] Don't fully unroll transfer_writes of n-D scalable vectors (#71924)

It is not possible to unroll a scalable vector at compile time. This
currently prevents transfer_writes from being lowered to
arm_sme.tile_writes (downstream).
This commit is contained in:
Benjamin Maxwell 2023-11-13 11:22:55 +00:00 committed by GitHub
parent a604c4b562
commit ced97ffd08
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 15 additions and 0 deletions

View File

@ -1218,6 +1218,11 @@ struct UnrollTransferWriteConversion
auto vec = getDataVector(xferOp);
auto xferVecType = xferOp.getVectorType();
if (xferVecType.getScalableDims()[0]) {
// Cannot unroll a scalable dimension at compile time.
return failure();
}
int64_t dimSize = xferVecType.getShape()[0];
Value source = xferOp.getSource(); // memref or tensor to be written to.
auto sourceType = isTensorOp(xferOp) ? xferOp.getShapedType() : Type();

View File

@ -737,4 +737,14 @@ func.func @cannot_lower_transfer_read_with_leading_scalable(%arg0: memref<?x4xf3
// CHECK-SAME: %[[MEMREF:.*]]: memref<?x4xf32>)
// CHECK: %{{.*}} = vector.transfer_read %[[MEMREF]][%{{.*}}, %{{.*}}], %{{.*}}, %{{.*}} {in_bounds = [true, true]} : memref<?x4xf32>, vector<[4]x4xf32>
// -----
// FULL-UNROLL-LABEL: @cannot_fully_unroll_transfer_write_of_nd_scalable_vector
func.func @cannot_fully_unroll_transfer_write_of_nd_scalable_vector(%vec: vector<[4]x[4]xf32>, %memref: memref<?x?xf32>) {
// FULL-UNROLL-NOT: vector.extract
// FULL-UNROLL: vector.transfer_write {{.*}} : vector<[4]x[4]xf32>, memref<?x?xf32>
// FULL-UNROLL-NOT: vector.extract
%c0 = arith.constant 0 : index
vector.transfer_write %vec, %memref[%c0, %c0] {in_bounds = [true, true]} : vector<[4]x[4]xf32>, memref<?x?xf32>
return
}