This is an archive of the discontinued LLVM Phabricator instance.

Differential D154082

[mlir][linalg] Add test case: vectorize tensor.pad and bufferize to allocation
ClosedPublic

Authored by springerm on Jun 29 2023, 6:57 AM.

Details

Reviewers
nicolasvasilache
Commits
rGe941e16de8ff: [mlir][linalg] Add test case: vectorize tensor.pad and bufferize to allocation
Summary

Add a test case that first vectorizes a tensor.pad op, then bufferizes it to a new allocation with a specified memory space.

Depends On: D154081

Diff Detail

Event Timeline

springerm created this revision.Jun 29 2023, 6:57 AM
Herald added a project: Restricted Project. · View Herald TranscriptJun 29 2023, 6:57 AM
Herald added subscribers: bviyer, Moerafaat, bzcheeseman and 24 others. · View Herald Transcript
springerm requested review of this revision.Jun 29 2023, 6:57 AM
Herald added a project: Restricted Project. · View Herald TranscriptJun 29 2023, 6:57 AM
Herald added subscribers: limo1996, stephenneuendorffer. · View Herald Transcript
springerm added a parent revision: D154081: [mlir][linalg] BufferizeToAllocationOp: Do not copy uninitialized buffers.Jun 29 2023, 6:57 AM
Harbormaster completed remote builds in B242074: Diff 535773.Jun 29 2023, 7:14 AM
nicolasvasilache accepted this revision.Jul 4 2023, 4:36 AM
This revision is now accepted and ready to land.Jul 4 2023, 4:36 AM
Herald added a subscriber: wangpc. · View Herald TranscriptJul 4 2023, 4:36 AM
This revision was landed with ongoing or failed builds.Jul 4 2023, 7:50 AM
Closed by commit rGe941e16de8ff: [mlir][linalg] Add test case: vectorize tensor.pad and bufferize to allocation (authored by springerm). · Explain Why
This revision was automatically updated to reflect the committed changes.
springerm added a commit: rGe941e16de8ff: [mlir][linalg] Add test case: vectorize tensor.pad and bufferize to allocation.

Revision Contents

PathSize
mlir/
test/
Dialect/
Linalg/
59 lines

Diff 537101

mlir/test/Dialect/Linalg/pad-to-specific-memory-space.mlir

Show First 20 Lines • Show All 52 Lines • ▼ Show 20 Lines %padded, %pad = transform.structured.pad %0 {
padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32], padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
padding_dimensions=[0, 1, 2], padding_dimensions=[0, 1, 2],
pack_paddings=[1, 1, 1] pack_paddings=[1, 1, 1]
} : (!transform.any_op) -> (!transform.any_op, !transform.any_op) } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
%buffer = transform.structured.bufferize_to_allocation %pad {memory_space = 3} : !transform.any_op %buffer = transform.structured.bufferize_to_allocation %pad {memory_space = 3} : !transform.any_op
%2 = transform.bufferization.one_shot_bufferize %arg1 {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op %2 = transform.bufferization.one_shot_bufferize %arg1 {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op
} }
// -----
#map = affine_map<()[s0] -> (-s0 + 12, 7)>
// CHECK-LABEL: func @vectorize_and_bufferize_pad(
// CHECK-SAME: %[[arg0:.*]]: memref<24x12xf32, strided<[?, ?], offset: ?>>,
// CHECK-SAME: %[[arg1:.*]]: memref<12x25xf32, strided<[?, ?], offset: ?>>,
// CHECK-SAME: %[[arg2:.*]]: memref<24x25xf32, strided<[?, ?], offset: ?>>,
func.func @vectorize_and_bufferize_pad(%arg0: tensor<24x12xf32>,
%arg1: tensor<12x25xf32>,
%arg2: tensor<24x25xf32>,
%iv0 : index, %iv1 : index,
%iv2 : index) -> tensor<24x25xf32> {
%0 = affine.min #map()[%iv2]
// CHECK: %[[s0:.*]] = memref.subview %[[arg0]]
%1 = tensor.extract_slice %arg0[%iv0, %iv2] [4, %0] [1, 1] : tensor<24x12xf32> to tensor<4x?xf32>
// CHECK: %[[s1:.*]] = memref.subview %[[arg1]]
%2 = tensor.extract_slice %arg1[%iv2, %iv1] [%0, 5] [1, 1] : tensor<12x25xf32> to tensor<?x5xf32>
// CHECK: %[[s2:.*]] = memref.subview %[[arg2]]
%3 = tensor.extract_slice %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<24x25xf32> to tensor<4x5xf32>
// CHECK: %[[v0:.*]] = vector.mask {{.*}} { vector.transfer_read %[[s0]]
// CHECK: %[[alloc0:.*]] = memref.alloc() : memref<4x7xf32, 3>
// CHECK: vector.mask {{.*}} { vector.transfer_write %[[v0]], %[[alloc0]]
// CHECK: %[[v1:.*]] = vector.mask {{.*}} { vector.transfer_read %[[s1]]
// CHECK: %[[alloc1:.*]] = memref.alloc() : memref<7x5xf32, 3>
// CHECK: vector.mask {{.*}} { vector.transfer_write %[[v1]], %[[alloc1]]
// CHECK: %[[v2:.*]] = vector.mask {{.*}} { vector.transfer_read %[[s2]]
// CHECK: %[[alloc2:.*]] = memref.alloc() : memref<4x5xf32, 3>
// CHECK: vector.mask {{.*}} { vector.transfer_write %[[v2]], %[[alloc0]]
// CHECK: linalg.matmul ins(%[[alloc0]], %[[alloc1]] : {{.*}}) outs(%[[alloc2]] : {{.*}})
// Copy back result.
// CHECK: memref.copy %[[alloc2]], %[[s2]]
%4 = linalg.matmul ins(%1, %2 : tensor<4x?xf32>, tensor<?x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>
// insert_slice bufferizes to a no-op.
%5 = tensor.insert_slice %4 into %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
func.return %5 : tensor<24x25xf32>
}
transform.sequence failures(propagate) {
^bb1(%arg1: !transform.any_op):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%padded, %pad = transform.structured.pad %0 {
padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
padding_dimensions=[0, 1, 2],
pack_paddings=[1, 1, 1]
} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.structured.masked_vectorize %pad vector_sizes [10, 12] : !transform.any_op
%vector_write = transform.structured.match ops{["vector.transfer_write"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%mask_op = transform.get_parent_op %vector_write {op_name = "vector.mask"} : (!transform.any_op) -> !transform.any_op
%buffer = transform.structured.bufferize_to_allocation %mask_op {memory_space = 3} : !transform.any_op
%2 = transform.bufferization.one_shot_bufferize %arg1 {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op
}