Commit 512da70b authored by Nicolas Vasilache's avatar Nicolas Vasilache
Browse files

[mlir][Vector] Degrade masking information when forwarding linalg.copy to vector.transfer 

Summary:
linalg.copy + linalg.fill can be used to create a padded local buffer.
The `masked` attribute is only valid on this padded buffer.
When forwarding to vector.transfer ops, the attribute must be reset
conservatively.

Differential Revision: https://reviews.llvm.org/D83782
parent 694ded37

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

+10 −4
Original line number Diff line number Diff line
@@ -260,10 +260,13 @@ LogicalResult LinalgCopyVTRForwardingPattern::matchAndRewrite( 
  // `in` is the subview that linalg.copy reads. Replace it.

  Value in = copyOp.getInput(0);



  // linalg.copy + linalg.fill can be used to create a padded local buffer.

  // The `masked` attribute is only valid on this padded buffer.

  // When forwarding to vector.transfer_read, the attribute must be reset

  // conservatively.

  Value res = rewriter.create<vector::TransferReadOp>(

      xferOp.getLoc(), xferOp.getVectorType(), in, xferOp.indices(),

      xferOp.permutation_map(), xferOp.padding(),

      xferOp.masked() ? *xferOp.masked() : ArrayAttr());

      xferOp.permutation_map(), xferOp.padding(), ArrayAttr());



  if (maybeFillOp)

    rewriter.eraseOp(maybeFillOp);
@@ -308,10 +311,13 @@ LogicalResult LinalgCopyVTWForwardingPattern::matchAndRewrite( 
  Value out = copyOp.getOutputBuffer(0);



  // Forward vector.transfer into copy.

  // linalg.copy + linalg.fill can be used to create a padded local buffer.

  // The `masked` attribute is only valid on this padded buffer.

  // When forwarding to vector.transfer_write, the attribute must be reset

  // conservatively.

  rewriter.create<vector::TransferWriteOp>(

      xferOp.getLoc(), xferOp.vector(), out, xferOp.indices(),

      xferOp.permutation_map(),

      xferOp.masked() ? *xferOp.masked() : ArrayAttr());

      xferOp.permutation_map(), ArrayAttr());



  rewriter.eraseOp(copyOp);

  rewriter.eraseOp(xferOp);

mlir/test/Dialect/Linalg/forward-vector-transfers.mlir

+12 −6
Original line number Diff line number Diff line
@@ -6,13 +6,14 @@ 
//   CHECK-NOT: linalg.copy

//       CHECK: %[[ALLOC:.*]] = alloc

//       CHECK: vector.transfer_read %[[ARG0]]

//   CHECK-NOT: masked

func @testAllocRead(%in: memref<? x f32>) -> vector<32 x f32> {

  %c0 = constant 0: index

  %f0 = constant 0.0: f32

  %alloc = alloc() : memref<32 x f32>

  %subview = subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>

  linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>

  %0 = vector.transfer_read %alloc[%c0], %f0: memref<32 x f32>, vector<32 x f32>

  %0 = vector.transfer_read %alloc[%c0], %f0 {masked = [false]} : memref<32 x f32>, vector<32 x f32>

  dealloc %alloc : memref<32 x f32>

  return %0: vector<32 x f32>

}
@@ -23,6 +24,7 @@ func @testAllocRead(%in: memref<? x f32>) -> vector<32 x f32> { 
//   CHECK-NOT: linalg.copy

//       CHECK: %[[ALLOC:.*]] = alloc

//       CHECK: vector.transfer_read %[[ARG0]]

//   CHECK-NOT: masked

func @testAllocFillRead(%in: memref<? x f32>) -> vector<32 x f32> {

  %c0 = constant 0: index

  %f0 = constant 0.0: f32
@@ -30,7 +32,7 @@ func @testAllocFillRead(%in: memref<? x f32>) -> vector<32 x f32> { 
  linalg.fill(%alloc, %f0): memref<32 x f32>, f32

  %subview = subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>

  linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>

  %0 = vector.transfer_read %alloc[%c0], %f0: memref<32 x f32>, vector<32 x f32>

  %0 = vector.transfer_read %alloc[%c0], %f0 {masked = [false]} : memref<32 x f32>, vector<32 x f32>

  dealloc %alloc : memref<32 x f32>

  return %0: vector<32 x f32>

}
@@ -41,6 +43,7 @@ func @testAllocFillRead(%in: memref<? x f32>) -> vector<32 x f32> { 
//   CHECK-NOT: linalg.copy

//       CHECK: %[[ALLOC:.*]] = alloc

//       CHECK: vector.transfer_read %[[ARG0]]

//   CHECK-NOT: masked

func @testViewRead(%in: memref<? x f32>) -> vector<32 x f32> {

  %c0 = constant 0: index

  %f0 = constant 0.0: f32
@@ -48,7 +51,7 @@ func @testViewRead(%in: memref<? x f32>) -> vector<32 x f32> { 
  %view = view %alloc[%c0][] : memref<128 x i8> to memref<32 x f32>

  %subview = subview %view[0][16][1] : memref<32 x f32> to memref<16 x f32>

  linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>

  %0 = vector.transfer_read %view[%c0], %f0: memref<32 x f32>, vector<32 x f32>

  %0 = vector.transfer_read %view[%c0], %f0 {masked = [false]} : memref<32 x f32>, vector<32 x f32>

  dealloc %alloc : memref<128 x i8>

  return %0: vector<32 x f32>

}
@@ -59,6 +62,7 @@ func @testViewRead(%in: memref<? x f32>) -> vector<32 x f32> { 
//   CHECK-NOT: linalg.copy

//       CHECK: %[[ALLOC:.*]] = alloc

//       CHECK: vector.transfer_read %[[ARG0]]

//   CHECK-NOT: masked

func @testViewFillRead(%in: memref<? x f32>) -> vector<32 x f32> {

  %c0 = constant 0: index

  %f0 = constant 0.0: f32
@@ -67,7 +71,7 @@ func @testViewFillRead(%in: memref<? x f32>) -> vector<32 x f32> { 
  %subview = subview %view[0][16][1] : memref<32 x f32> to memref<16 x f32>

  linalg.fill(%view, %f0): memref<32 x f32>, f32

  linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>

  %0 = vector.transfer_read %view[%c0], %f0: memref<32 x f32>, vector<32 x f32>

  %0 = vector.transfer_read %view[%c0], %f0 {masked = [false]} : memref<32 x f32>, vector<32 x f32>

  dealloc %alloc : memref<128 x i8>

  return %0: vector<32 x f32>

}
@@ -78,12 +82,13 @@ func @testViewFillRead(%in: memref<? x f32>) -> vector<32 x f32> { 
//   CHECK-NOT: linalg.copy

//       CHECK: %[[ALLOC:.*]] = alloc

//       CHECK: vector.transfer_write %[[ARG0]], %[[ARG1]]

//   CHECK-NOT: masked

func @testAllocWrite(%vec: vector<32 x f32>, %out: memref<? x f32>) {

  %c0 = constant 0: index

  %f0 = constant 0.0: f32

  %alloc = alloc() : memref<32 x f32>

  %subview = subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>

  vector.transfer_write %vec, %alloc[%c0] : vector<32 x f32>, memref<32 x f32>

  vector.transfer_write %vec, %alloc[%c0] {masked = [false]} : vector<32 x f32>, memref<32 x f32>

  linalg.copy(%subview, %out): memref<16 x f32>, memref<? x f32>

  dealloc %alloc : memref<32 x f32>

  return
@@ -95,13 +100,14 @@ func @testAllocWrite(%vec: vector<32 x f32>, %out: memref<? x f32>) { 
//   CHECK-NOT: linalg.copy

//       CHECK: %[[ALLOC:.*]] = alloc

//       CHECK: vector.transfer_write %[[ARG0]], %[[ARG1]]

//   CHECK-NOT: masked

func @testViewWrite(%vec: vector<32 x f32>, %out: memref<? x f32>) {

  %c0 = constant 0: index

  %f0 = constant 0.0: f32

  %alloc = alloc() : memref<128 x i8>

  %view = view %alloc[%c0][] : memref<128 x i8> to memref<32 x f32>

  %subview = subview %view[0][16][1] : memref<32 x f32> to memref<16 x f32>

  vector.transfer_write %vec, %view[%c0] : vector<32 x f32>, memref<32 x f32>

  vector.transfer_write %vec, %view[%c0] {masked = [false]} : vector<32 x f32>, memref<32 x f32>

  linalg.copy(%subview, %out): memref<16 x f32>, memref<? x f32>

  dealloc %alloc : memref<128 x i8>

  return