223 lines · plain
1// RUN: mlir-opt -split-input-file \2// RUN: -transform-preload-library='transform-library-paths=%p/td/decompose-unpack.mlir' \3// RUN: -transform-interpreter=entry-point=decompose_unpack %s | FileCheck %s4 5func.func @simple_KCRSsr_to_KCRS(%arg0: tensor<1x1x1x1x8x32xf32>, %arg1: tensor<1x1x32x8xf32>) -> tensor<1x1x32x8xf32> {6 %0 = linalg.unpack %arg0 inner_dims_pos = [3, 2] inner_tiles = [8, 32] into %arg1 : tensor<1x1x1x1x8x32xf32> -> tensor<1x1x32x8xf32>7 return %0 : tensor<1x1x32x8xf32>8}9// CHECK-LABEL: func.func @simple_KCRSsr_to_KCRS10// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]11// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]12// CHECK: %[[TILE:.+]] = tensor.extract_slice %[[SRC]][0, 0, 0, 0, 0, 0] [1, 1, 1, 1, 8, 32] [1, 1, 1, 1, 1, 1]13// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<32x8xf32>14// CHECK: %[[TRANSP:.+]] = linalg.transpose15// CHECK-SAME: ins(%[[TILE]] : tensor<8x32xf32>)16// CHECK-SAME: outs(%[[EMPTY]] : tensor<32x8xf32>)17// CHECK-SAME: permutation = [1, 0]18// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[TRANSP]] into %[[DEST]]19// CHECK-SAME: [0, 0, 0, 0] [1, 1, 32, 8] [1, 1, 1, 1]20// CHECK: return %[[INSERT]]21 22// -----23 24func.func @simple_unpack_static_tiles(%input: tensor<1x1x8x2xf32>, %output: tensor<5x1xf32>) -> tensor<5x1xf32> {25 %0 = linalg.unpack %input inner_dims_pos = [0, 1] inner_tiles = [8, 2] into %output : tensor<1x1x8x2xf32> -> tensor<5x1xf32>26 return %0 : tensor<5x1xf32>27}28// CHECK-LABEL: func.func @simple_unpack_static_tiles29// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]30// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]31// CHECK: %[[TILE:.+]] = tensor.extract_slice %[[SRC]][0, 0, 0, 0] [1, 1, 8, 2] [1, 1, 1, 1]32// CHECK-NOT: linalg.transpose33// They have the same type, so the insert_slice op is folded34// away.35// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[TILE]][0, 0] [5, 1] [1, 1]36// CHECK: return %[[SLICE]]37 38/// Same as example above, but with 1 dynamic tile size.39 40func.func @simple_unpack_dynamic_tile(%input: tensor<1x1x?x2xf32>, %output: tensor<5x1xf32>, %tile_dim: index) -> tensor<5x1xf32> {41 %0 = linalg.unpack %input inner_dims_pos = [0, 1] inner_tiles = [%tile_dim, 2] into %output : tensor<1x1x?x2xf32> -> tensor<5x1xf32>42 return %0 : tensor<5x1xf32>43}44// CHECK-LABEL: func.func @simple_unpack_dynamic_tile45// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]46// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]47// CHECK-SAME: %[[TILE_DIM:[a-zA-Z0-9]+]]48// CHECK: %[[TILE:.+]] = tensor.extract_slice %[[SRC]][0, 0, 0, 0] [1, 1, %[[TILE_DIM]], 2] [1, 1, 1, 1]49// CHECK-NOT: linalg.transpose50// They have the same type, so the insert_slice op is folded51// away.52// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[TILE]][0, 0] [5, 1] [1, 1]53// CHECK: return %[[SLICE]]54 55/// Same as example above, but with 1 dynamic tile size and a trasnpose56 57func.func @simple_unpack_dynamic_tile_transpose(%src: tensor<1x1x2x?xf32>, %dest: tensor<5x1xf32>, %tile_dim: index) -> tensor<5x1xf32> {58 %0 = linalg.unpack %src inner_dims_pos = [1, 0] inner_tiles = [2, %tile_dim] into %dest : tensor<1x1x2x?xf32> -> tensor<5x1xf32>59 return %0 : tensor<5x1xf32>60}61// CHECK-LABEL: func.func @simple_unpack_dynamic_tile_transpose62// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]63// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]64// CHECK-SAME: %[[TILE_DIM:[a-zA-Z0-9]+]]65// CHECK: %[[TILE:.*]] = tensor.extract_slice %[[SRC]][0, 0, 0, 0] [1, 1, 2, %[[TILE_DIM]]] [1, 1, 1, 1] : tensor<1x1x2x?xf32> to tensor<2x?xf32>66// CHECK: %[[EMPTY:.*]] = tensor.empty(%[[TILE_DIM]]) : tensor<?x2xf32>67// CHECK: %[[TRANSP:.*]] = linalg.transpose68// CHECK-SAME: ins(%[[TILE]] : tensor<2x?xf32>)69// CHECK-SAME: outs(%[[EMPTY]] : tensor<?x2xf32>)70// CHECK-SAME: permutation = [1, 0]71// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[TRANSP]][0, 0] [5, 1] [1, 1] : tensor<?x2xf32> to tensor<5x1xf32>72// CHECK: return %[[SLICE]] : tensor<5x1xf32>73 74 75/// Same as example above, but with 1 scalable tile size.76 77func.func @simple_unpack_scalable_tile(%input: tensor<1x1x?x2xf32>, %output: tensor<5x1xf32>) -> tensor<5x1xf32> {78 %c8 = arith.constant 8 : index79 %vscale = vector.vscale80 %c8_vscale = arith.muli %vscale, %c8 : index81 %0 = linalg.unpack %input inner_dims_pos = [0, 1] inner_tiles = [%c8_vscale, 2] into %output : tensor<1x1x?x2xf32> -> tensor<5x1xf32>82 return %0 : tensor<5x1xf32>83}84// CHECK-LABEL: func.func @simple_unpack_scalable_tile85// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]86// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]87// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index88// CHECK-DAG: %[[VS:.+]] = vector.vscale89// CHECK: %[[C8_VS:.+]] = arith.muli %[[VS]], %[[C8]] : index90// CHECK: %[[TILE:.+]] = tensor.extract_slice %[[SRC]][0, 0, 0, 0] [1, 1, %[[C8_VS]], 2] [1, 1, 1, 1]91// CHECK-NOT: linalg.transpose92// They have the same type, so the insert_slice op is folded93// away.94// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[TILE]][0, 0] [5, 1] [1, 1]95// CHECK: return %[[SLICE]]96 97// -----98 99func.func @simple_CNnc_to_NC(%arg0: tensor<1x1x32x8xf32>, %arg1: tensor<32x8xf32>) -> tensor<32x8xf32>{100 %0 = linalg.unpack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 8] into %arg1 : tensor<1x1x32x8xf32> -> tensor<32x8xf32>101 return %0 : tensor<32x8xf32>102}103// CHECK-LABEL: func.func @simple_CNnc_to_NC104// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]105// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]106// CHECK: %[[TILE:.+]] = tensor.extract_slice %[[SRC]][0, 0, 0, 0] [1, 1, 32, 8] [1, 1, 1, 1]107// CHECK-NOT: linalg.transpose108// They have the same type, so the insert_slice op is folded109// away.110// CHECK: return %[[TILE]]111 112// -----113 114func.func @simple_NCHWc_to_NCHW(%arg0: tensor<2x1x16x8x32xf32>, %arg1: tensor<2x32x16x8xf32>) -> tensor<2x32x16x8xf32> {115 %0 = linalg.unpack %arg0 inner_dims_pos = [1] inner_tiles = [32] into %arg1 : tensor<2x1x16x8x32xf32> -> tensor<2x32x16x8xf32>116 return %0 : tensor<2x32x16x8xf32>117}118// CHECK-LABEL: func.func @simple_NCHWc_to_NCHW119// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]120// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]121// CHECK: %[[TILE:.+]] = tensor.extract_slice %[[SRC]][0, 0, 0, 0, 0] [2, 1, 16, 8, 32] [1, 1, 1, 1, 1]122// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<2x32x16x8xf32>123// CHECK: %[[TRANSP:.+]] = linalg.transpose124// CHECK-SAME: ins(%[[TILE]] : tensor<2x16x8x32xf32>)125// CHECK-SAME: outs(%[[EMPTY]] : tensor<2x32x16x8xf32>)126// CHECK-SAME: permutation = [0, 3, 1, 2]127// They have the same type, so the insert_slice op is folded128// away.129// CHECK: return %[[TRANSP]]130 131// -----132 133func.func @simple_NHWC_to_NCHW(%arg0: tensor<1x16x8x32xf32>, %arg1: tensor<1x32x16x8xf32>) -> tensor<1x32x16x8xf32> {134 %0 = linalg.unpack %arg0 outer_dims_perm = [0, 2, 3, 1] inner_dims_pos = [] inner_tiles = [] into %arg1 : tensor<1x16x8x32xf32> -> tensor<1x32x16x8xf32>135 return %0 : tensor<1x32x16x8xf32>136}137// CHECK-LABEL: func.func @simple_NHWC_to_NCHW138// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]139// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]140// CHECK: %[[TILE:.+]] = tensor.extract_slice %[[SRC]][0, 0, 0, 0] [1, 16, 8, 32] [1, 1, 1, 1]141// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<32x16x8xf32>142// CHECK: %[[TRANSP:.+]] = linalg.transpose143// CHECK-SAME: ins(%[[TILE]] : tensor<16x8x32xf32>)144// CHECK-SAME: outs(%[[EMPTY]] : tensor<32x16x8xf32>)145// CHECK-SAME: permutation = [2, 0, 1]146// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[TRANSP]] into %[[DEST]]147// CHECK-SAME: [0, 0, 0, 0] [1, 32, 16, 8] [1, 1, 1, 1]148// CHECK: return %[[INSERT]]149 150// -----151 152func.func @unpack_with_dynamic_dims(%arg0: tensor<?x1x1x1x8x32xf32>, %arg1: tensor<?x1x32x8xf32>) -> tensor<?x1x32x8xf32> {153 %0 = linalg.unpack %arg0 inner_dims_pos = [3, 2] inner_tiles = [8, 32] into %arg1 : tensor<?x1x1x1x8x32xf32> -> tensor<?x1x32x8xf32>154 return %0 : tensor<?x1x32x8xf32>155}156// CHECK-LABEL: func.func @unpack_with_dynamic_dims157// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]158// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]159// CHECK: %[[C0:.+]] = arith.constant 0 : index160// CHECK: %[[DIM0_SRC:.+]] = tensor.dim %[[SRC]], %[[C0]] : tensor<?x1x1x1x8x32xf32>161// CHECK: %[[TILE:.+]] = tensor.extract_slice %[[SRC]][0, 0, 0, 0, 0, 0] [%[[DIM0_SRC]], 1, 1, 1, 8, 32] [1, 1, 1, 1, 1, 1]162// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM0_SRC]]) : tensor<?x32x8xf32>163// CHECK: %[[TRANSP:.+]] = linalg.transpose164// CHECK-SAME: ins(%[[TILE]] : tensor<?x8x32xf32>)165// CHECK-SAME: outs(%[[EMPTY]] : tensor<?x32x8xf32>)166// CHECK-SAME: permutation = [0, 2, 1]167// CHECK: %[[DIM0_DEST:.+]] = tensor.dim %[[DEST]], %[[C0]] : tensor<?x1x32x8xf32>168// CHECK: %[[EXTRACT_SLICE:.+]] = tensor.extract_slice %[[TRANSP]][0, 0, 0] [%[[DIM0_DEST]], 32, 8] [1, 1, 1] : tensor<?x32x8xf32> to tensor<?x32x8xf32>169// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[EXTRACT_SLICE]] into %[[DEST]]170// CHECK-SAME: [0, 0, 0, 0] [%[[DIM0_DEST]], 1, 32, 8] [1, 1, 1, 1]171// CHECK: return %[[INSERT]]172 173// -----174 175func.func @unpack_with_non_adjacent_inner_dims_pos_and_unit_outer(%arg0: tensor<1x1x1x4x1xf32>, %arg1: tensor<1x1x4xf32>) -> tensor<1x1x4xf32> {176 %0 = linalg.unpack %arg0 outer_dims_perm = [1, 2, 0] inner_dims_pos = [2, 0] inner_tiles = [4, 1] into %arg1 : tensor<1x1x1x4x1xf32> -> tensor<1x1x4xf32>177 return %0 : tensor<1x1x4xf32>178}179// CHECK-LABEL: func.func @unpack_with_non_adjacent_inner_dims_pos_and_unit_outer180// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]181// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]182// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[SRC]][0, 0, 0, 0, 0] [1, 1, 1, 4, 1] [1, 1, 1, 1, 1] : tensor<1x1x1x4x1xf32> to tensor<4x1xf32>183// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x4xf32>184// CHECK: %[[TRANSP:.+]] = linalg.transpose185// CHECK-SAME: ins(%[[SLICE]] : tensor<4x1xf32>)186// CHECK-SAME: outs(%[[EMPTY]] : tensor<1x4xf32>) permutation = [1, 0]187// CHECK: %[[INSERT:.+]] = tensor.insert_slice %transposed into %[[DEST]][0, 0, 0] [1, 1, 4] [1, 1, 1] : tensor<1x4xf32> into tensor<1x1x4xf32>188// CHECK: return %[[INSERT]]189 190// -----191 192func.func @unpack_with_non_trailing_dimensions_in_inner_dims(%arg0: tensor<1x1x1x4x1xf32>, %arg1: tensor<1x1x4xf32>) -> tensor<1x1x4xf32> {193 %pack = linalg.unpack %arg0 outer_dims_perm = [1, 2, 0] inner_dims_pos = [2, 1] inner_tiles = [4, 1] into %arg1 : tensor<1x1x1x4x1xf32> -> tensor<1x1x4xf32>194 return %pack : tensor<1x1x4xf32>195}196// CHECK-LABEL: func.func @unpack_with_non_trailing_dimensions_in_inner_dims197// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]198// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]199// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[SRC]][0, 0, 0, 0, 0] [1, 1, 1, 4, 1] [1, 1, 1, 1, 1] : tensor<1x1x1x4x1xf32> to tensor<4x1xf32>200// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x4xf32>201// CHECK: %[[TRANSP:.+]] = linalg.transpose202// CHECK-SAME: ins(%[[SLICE]] : tensor<4x1xf32>)203// CHECK-SAME: outs(%[[EMPTY]] : tensor<1x4xf32>) permutation = [1, 0]204// CHECK: %[[INSERT:.+]] = tensor.insert_slice %transposed into %[[DEST]][0, 0, 0] [1, 1, 4] [1, 1, 1] : tensor<1x4xf32> into tensor<1x1x4xf32>205// CHECK: return %[[INSERT]]206 207// -----208 209/// Note "126", which is a non-unit tile-outer-dim. This is not supported.210 211func.func @negative_non_unit_tiled_outer_dim(%src: tensor<1x126x1x1x8xf32>, %dest: tensor<1x1x1x1001xf32>) -> tensor<1x1x1x1001xf32> {212 %unpack = linalg.unpack %src213 outer_dims_perm = [0, 3, 2, 1]214 inner_dims_pos = [3]215 inner_tiles = [8]216 into %dest : tensor<1x126x1x1x8xf32>217 -> tensor<1x1x1x1001xf32>218 219 return %unpack : tensor<1x1x1x1001xf32>220}221// CHECK-LABEL: @negative_non_unit_tiled_outer_dim(222// CHECK: linalg.unpack223