9a5092b358
- Fold an unpack(pack(x)) to x. - Rewrite a `tensor.pack` to an `tensor.expand_shape` if only one dimension is packed. Reviewed By: tyb0807, hanchung, mravishankar Differential Revision: https://reviews.llvm.org/D141123
25 lines
1.1 KiB
MLIR
25 lines
1.1 KiB
MLIR
// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns="test-simplify-pack-patterns" %s | FileCheck %s
|
|
|
|
// CHECK: func.func @single_dim_packing(
|
|
// CHECK-SAME: %[[ARG0:.+]]: tensor<256xf32>)
|
|
// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1]] : tensor<256xf32> into tensor<8x32xf32>
|
|
// CHECK: return %[[EXPANDED]] : tensor<8x32xf32>
|
|
func.func @single_dim_packing(%arg0: tensor<256xf32>) -> tensor<8x32xf32> {
|
|
%empty = tensor.empty() : tensor<8x32xf32>
|
|
%0 = tensor.pack %arg0 inner_dims_pos = [0] inner_tiles = [32] into %empty : tensor<256xf32> -> tensor<8x32xf32>
|
|
return %0 : tensor<8x32xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
// CHECK: func.func @single_dim_packing_with_padding(
|
|
// CHECK-SAME: %[[ARG0:.+]]: tensor<255xf32>)
|
|
// CHECK-NOT: tensor.expand_shape
|
|
// CHECK: tensor.pack
|
|
func.func @single_dim_packing_with_padding(%arg0: tensor<255xf32>) -> tensor<8x32xf32> {
|
|
%empty = tensor.empty() : tensor<8x32xf32>
|
|
%cst = arith.constant 0.000000e+00 : f32
|
|
%0 = tensor.pack %arg0 padding_value(%cst : f32) inner_dims_pos = [0] inner_tiles = [32] into %empty : tensor<255xf32> -> tensor<8x32xf32>
|
|
return %0 : tensor<8x32xf32>
|
|
}
|