diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp @@ -260,10 +260,13 @@ // `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( 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 @@ 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( 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); diff --git a/mlir/test/Dialect/Linalg/forward-vector-transfers.mlir b/mlir/test/Dialect/Linalg/forward-vector-transfers.mlir --- a/mlir/test/Dialect/Linalg/forward-vector-transfers.mlir +++ b/mlir/test/Dialect/Linalg/forward-vector-transfers.mlir @@ -6,13 +6,14 @@ // CHECK-NOT: linalg.copy // CHECK: %[[ALLOC:.*]] = alloc // CHECK: vector.transfer_read %[[ARG0]] +// CHECK-NOT: masked func @testAllocRead(%in: memref) -> 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, 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 @@ // CHECK-NOT: linalg.copy // CHECK: %[[ALLOC:.*]] = alloc // CHECK: vector.transfer_read %[[ARG0]] +// CHECK-NOT: masked func @testAllocFillRead(%in: memref) -> vector<32 x f32> { %c0 = constant 0: index %f0 = constant 0.0: f32 @@ -30,7 +32,7 @@ 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, 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 @@ // CHECK-NOT: linalg.copy // CHECK: %[[ALLOC:.*]] = alloc // CHECK: vector.transfer_read %[[ARG0]] +// CHECK-NOT: masked func @testViewRead(%in: memref) -> vector<32 x f32> { %c0 = constant 0: index %f0 = constant 0.0: f32 @@ -48,7 +51,7 @@ %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, 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 @@ // CHECK-NOT: linalg.copy // CHECK: %[[ALLOC:.*]] = alloc // CHECK: vector.transfer_read %[[ARG0]] +// CHECK-NOT: masked func @testViewFillRead(%in: memref) -> vector<32 x f32> { %c0 = constant 0: index %f0 = constant 0.0: f32 @@ -67,7 +71,7 @@ %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, 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 @@ // 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) { %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 dealloc %alloc : memref<32 x f32> return @@ -95,13 +100,14 @@ // 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) { %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 dealloc %alloc : memref<128 x i8> return