Unverified Commit ced97ffd authored by Benjamin Maxwell's avatar Benjamin Maxwell Committed by GitHub
Browse files

[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).
parent a604c4b5

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

+5 −0
Original line number Diff line number Diff line
@@ -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();

mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir

+10 −0
Original line number Diff line number Diff line
@@ -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

}