83 lines · plain
1// RUN: mlir-opt -split-input-file -transform-interpreter %s | FileCheck %s2 3module attributes {transform.with_named_sequence} {4 transform.named_sequence @__transform_main(%root : !transform.any_op {transform.readonly}) {5 %func_op = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.op<"func.func">6 transform.apply_patterns to %func_op {7 transform.apply_patterns.linalg.fold_pack_unpack_into_empty8 } : !transform.op<"func.func">9 transform.yield10 }11}12 13func.func @pack_empty(%arg0: tensor<8x8x32x32xf32>) -> tensor<8x8x32x32xf32> {14 %empty_unpacked = tensor.empty() : tensor<256x256xf32>15 %packed = linalg.pack %empty_unpacked16 inner_dims_pos = [0, 1] inner_tiles = [32, 32]17 into %arg0 : tensor<256x256xf32> -> tensor<8x8x32x32xf32>18 return %packed : tensor<8x8x32x32xf32>19}20 21// CHECK-LABEL: func.func @pack_empty(22// CHECK-SAME: %[[T:.+]]: tensor<8x8x32x32xf32>23// CHECK-NOT: linalg.pack24// CHECK: return %[[T]] : tensor<8x8x32x32xf32>25 26func.func @pack_empty_dynamic(%arg0: tensor<?x?x32x32xf32>, %dim0: index, %dim1: index) -> tensor<?x?x32x32xf32> {27 %empty_unpacked = tensor.empty(%dim0, %dim1) : tensor<?x?xf32>28 %packed = linalg.pack %empty_unpacked29 inner_dims_pos = [0, 1] inner_tiles = [32, 32]30 into %arg0 : tensor<?x?xf32> -> tensor<?x?x32x32xf32>31 return %packed : tensor<?x?x32x32xf32>32}33 34// CHECK-LABEL: func.func @pack_empty_dynamic(35// CHECK-SAME: %[[T:.+]]: tensor<?x?x32x32xf32>,36// CHECK-SAME: %[[DIM0:[a-zA-Z0-9_]+]]: index,37// CHECK-SAME: %[[DIM1:[a-zA-Z0-9_]+]]: index38// CHECK-NOT: linalg.pack39// CHECK: return %[[T]] : tensor<?x?x32x32xf32>40 41func.func @unpack_empty(%arg0: tensor<256x256xf32>) -> tensor<256x256xf32> {42 %empty_packed = tensor.empty() : tensor<8x8x32x32xf32>43 %unpacked = linalg.unpack %empty_packed44 inner_dims_pos = [0, 1] inner_tiles = [32, 32]45 into %arg0 : tensor<8x8x32x32xf32> -> tensor<256x256xf32>46 return %unpacked : tensor<256x256xf32>47}48 49// CHECK-LABEL: func.func @unpack_empty(50// CHECK-SAME: %[[T:.+]]: tensor<256x256xf32>51// CHECK-NOT: linalg.unpack52// CHECK: return %[[T]] : tensor<256x256xf32>53 54func.func @unpack_empty_dynamic(%arg0: tensor<?x?xf32>, %dim0: index, %dim1: index) -> tensor<?x?xf32> {55 %empty_packed = tensor.empty(%dim0, %dim1) : tensor<?x?x32x32xf32>56 %unpacked = linalg.unpack %empty_packed57 inner_dims_pos = [0, 1] inner_tiles = [32, 32]58 into %arg0 : tensor<?x?x32x32xf32> -> tensor<?x?xf32>59 return %unpacked : tensor<?x?xf32>60}61 62// CHECK-LABEL: func.func @unpack_empty_dynamic(63// CHECK-SAME: %[[T:.+]]: tensor<?x?xf32>,64// CHECK-SAME: %[[DIM0:[a-zA-Z0-9_]+]]: index,65// CHECK-SAME: %[[DIM1:[a-zA-Z0-9_]+]]: index66// CHECK-NOT: linalg.unpack67// CHECK: return %[[T]] : tensor<?x?xf32>68 69func.func @pack_padded_empty(%arg0: tensor<8x8x32x32xf32>) -> tensor<8x8x32x32xf32> {70 %pad = arith.constant 1.0 : f3271 %empty_unpacked = tensor.empty() : tensor<256x256xf32>72 %packed = linalg.pack %empty_unpacked73 padding_value(%pad : f32)74 inner_dims_pos = [0, 1] inner_tiles = [32, 32]75 into %arg0 : tensor<256x256xf32> -> tensor<8x8x32x32xf32>76 return %packed : tensor<8x8x32x32xf32>77}78 79// CHECK-LABEL: func.func @pack_padded_empty(80// CHECK-SAME: %[[T:.+]]: tensor<8x8x32x32xf32>81// CHECK: %[[PACK:.+]] = linalg.pack82// CHECK: return %[[PACK]] : tensor<8x8x32x32xf32>83