492 lines · plain
1// RUN: mlir-opt %s -transform-interpreter -canonicalize -cse -split-input-file | FileCheck %s2 3// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> (d0 * 32)>4// CHECK: func.func @NC_to_NCnc5// CHECK-SAME: %[[IN:.*]]: tensor<128x256xf32>,6// CHECK-SAME: %[[OUT:.*]]: tensor<4x8x32x32xf32>) -> tensor<4x8x32x32xf32> {7// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index8// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index9// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index10// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index11// CHECK: %[[RES0:.*]] = scf.for %[[N:.*]] = %[[C0]] to %[[C4]] step %[[C2]] iter_args(%[[ITER0:.*]] = %[[OUT]]) -> (tensor<4x8x32x32xf32>) {12// CHECK: %[[RES1:.+]] = scf.for %[[C:.*]] = %[[C0]] to %[[C8]] step %[[C4]] iter_args(%[[ITER1:.*]] = %[[ITER0]]) -> (tensor<4x8x32x32xf32>) {13// CHECK-DAG: %[[IN_N:.+]] = affine.apply #[[MAP0]](%[[N]])14// CHECK-DAG: %[[IN_C:.+]] = affine.apply #[[MAP0]](%[[C]])15// CHECK: %[[SUB_IN:.*]] = tensor.extract_slice %[[IN]][%[[IN_N]], %[[IN_C]]] [64, 128] [1, 1] : tensor<128x256xf32> to tensor<64x128xf32>16// CHECK: %[[SUB_OUT:.*]] = tensor.extract_slice %[[ITER1]][%[[N]], %[[C]], 0, 0] [2, 4, 32, 32] [1, 1, 1, 1] : tensor<4x8x32x32xf32> to tensor<2x4x32x32xf32>17// CHECK: %[[SUB_RES:.*]] = linalg.pack18// CHECK-SAME: %[[SUB_IN]] inner_dims_pos = [0, 1] inner_tiles = [32, 32] into %[[SUB_OUT]]19// CHECK: %[[INSERT:.*]] = tensor.insert_slice %[[SUB_RES]] into %[[ITER1]]20// CHECK: scf.yield %[[INSERT]] : tensor<4x8x32x32xf32>21// CHECK: }22// CHECK: scf.yield %[[RES1:.*]] : tensor<4x8x32x32xf32>23// CHECK: }24// CHECK: return %[[RES0:.*]] : tensor<4x8x32x32xf32>25// CHECK: }26func.func @NC_to_NCnc(%arg0: tensor<128x256xf32>, %arg1: tensor<4x8x32x32xf32>) -> tensor<4x8x32x32xf32> {27 %0 = linalg.pack %arg0 inner_dims_pos = [0, 1] inner_tiles = [32, 32] into %arg1 : tensor<128x256xf32> -> tensor<4x8x32x32xf32>28 return %0 : tensor<4x8x32x32xf32>29}30 31module attributes {transform.with_named_sequence} {32 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {33 %0 = transform.structured.match ops{["linalg.pack"]} in %arg1 : (!transform.any_op) -> !transform.any_op34 %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [2, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)35 transform.yield36 }37}38 39// -----40 41// CHECK: #[[MAP0:.+]] = affine_map<(d0) -> (d0 * 8)>42// CHECK: func.func @KC_to_CKkc43// CHECK-SAME: %[[IN:[A-Za-z0-9]+]]:44// CHECK-SAME: %[[OUT:[A-Za-z0-9]+]]:45// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index46// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index47// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index48// CHECK: scf.for %[[C:.+]] = %[[C0]] to %[[C32]] step %[[C2]]49// CHECK-DAG: %[[IN_C:.+]] = affine.apply #[[MAP0]](%[[C]])50// CHECK: %[[INPUT_SLICE:.+]] = tensor.extract_slice %[[IN]]51// CHECK-SAME: [0, %[[IN_C]]] [128, 16]52// CHECK: %[[OUTPUT_SLICE:.+]] = tensor.extract_slice %{{.+}}[%[[C]], 0, 0, 0] [2, 4, 32, 8]53// CHECK: linalg.pack54// CHECK-SAME: %[[INPUT_SLICE]] outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 8]55// CHECK-SAME: into %[[OUTPUT_SLICE]]56func.func @KC_to_CKkc(%arg0: tensor<128x256xf32>, %arg1: tensor<32x4x32x8xf32>) -> tensor<32x4x32x8xf32> {57 %0 = linalg.pack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 8] into %arg1 : tensor<128x256xf32> -> tensor<32x4x32x8xf32>58 return %0 : tensor<32x4x32x8xf32>59}60 61module attributes {transform.with_named_sequence} {62 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {63 %0 = transform.structured.match ops{["linalg.pack"]} in %arg1 : (!transform.any_op) -> !transform.any_op64 %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [2, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)65 transform.yield66 }67}68 69// -----70 71// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> (d0 * 2)>72// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0 * -2 + 15, 8)>73// CHECK: func.func @pad_and_pack_static(74// CHECK-SAME: %[[IN:.*]]: tensor<13x15xf32>,75// CHECK-SAME: %[[OUT:.*]]: tensor<2x8x8x2xf32>,76// CHECK-SAME: %[[PAD:.*]]: f32) -> tensor<2x8x8x2xf32> {77// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index78// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index79// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index80// CHECK-DAG: %[[RES0:.*]] = scf.for %[[J:.*]] = %[[C0]] to %[[C8]] step %[[C4]] iter_args(%[[ITER1:.*]] = %[[OUT]]) -> (tensor<2x8x8x2xf32>) {81// CHECK-DAG: %[[IN_J:.*]] = affine.apply #[[MAP0]](%[[J]])82// CHECK-DAG: %[[IN_J_SZ:.*]] = affine.min #[[MAP1]](%[[J]])83// CHECK: %[[SUB_IN:.*]] = tensor.extract_slice %[[IN]][0, %[[IN_J]]] [13, %[[IN_J_SZ]]] [1, 1]84// CHECK: %[[SUB_OUT:.*]] = tensor.extract_slice %[[ITER1]][0, %[[J]], 0, 0] [2, 4, 8, 2] [1, 1, 1, 1]85// CHECK: %[[SUB_RES:.*]] = linalg.pack86// CHECK-SAME: %[[SUB_IN]] padding_value(%[[PAD]] : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 2]87// CHECK-SAME: into %[[SUB_OUT]]88// CHECK: %[[INSERT:.*]] = tensor.insert_slice %[[SUB_RES]] into %[[ITER1]]89// CHECK: scf.yield %[[INSERT]] : tensor<2x8x8x2xf32>90// CHECK: }91// CHECK: return %[[RES0:.*]] : tensor<2x8x8x2xf32>92// CHECK: }93func.func @pad_and_pack_static(%input: tensor<13x15xf32>, %output: tensor<2x8x8x2xf32>, %pad: f32) -> tensor<2x8x8x2xf32> {94 %0 = linalg.pack %input padding_value(%pad : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 2] into %output : tensor<13x15xf32> -> tensor<2x8x8x2xf32>95 return %0 : tensor<2x8x8x2xf32>96}97 98module attributes {transform.with_named_sequence} {99 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {100 %0 = transform.structured.match ops{["linalg.pack"]} in %arg1 : (!transform.any_op) -> !transform.any_op101 %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [2, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)102 transform.yield103 }104}105 106// -----107 108// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0] -> (-d0 + s0, 2)>109// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0] -> (-d0 + s0, 4)>110// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> (d0 * 8)>111// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1)[s0] -> (d1 * -8 + s0, d0 * 8)>112// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0) -> (d0 * 2)>113// CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1)[s0] -> (d1 * -2 + s0, d0 * 2)>114// CHECK: func.func @pad_and_pack_partially_dynamic(115// CHECK-SAME: %[[IN:.*]]: tensor<?x?xf32>,116// CHECK-SAME: %[[OUT:.*]]: tensor<?x?x8x2xf32>,117// CHECK-SAME: %[[PAD:.*]]: f32) -> tensor<?x?x8x2xf32> {118// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index119// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index120// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index121// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index122// CHECK-DAG: %[[OUT_D0:.*]] = tensor.dim %[[OUT]], %[[C0]] : tensor<?x?x8x2xf32>123// CHECK-DAG: %[[OUT_D1:.*]] = tensor.dim %[[OUT]], %[[C1]] : tensor<?x?x8x2xf32>124// CHECK: %[[RES0:.*]] = scf.for %[[I:.*]] = %[[C0]] to %[[OUT_D0]] step %[[C2]] iter_args(%[[ITER0:.*]] = %[[OUT]]) -> (tensor<?x?x8x2xf32>) {125// CHECK: %[[RES1:.*]] = scf.for %[[J:.*]] = %[[C0]] to %[[OUT_D1]] step %[[C4]] iter_args(%[[ITER1:.*]] = %[[ITER0]]) -> (tensor<?x?x8x2xf32>) {126// CHECK-DAG: %[[OUT_I_SZ:.*]] = affine.min #[[MAP0]](%[[I]])[%[[OUT_D0]]]127// CHECK-DAG: %[[OUT_J_SZ:.*]] = affine.min #[[MAP1]](%[[J]])[%[[OUT_D1]]]128// CHECK-DAG: %[[IN_I:.*]] = affine.apply #[[MAP2]](%[[I]])129// CHECK-DAG: %[[IN_I_SZ:.*]] = affine.min #[[MAP3]]130// CHECK-DAG: %[[IN_J:.*]] = affine.apply #[[MAP4]](%[[J]])131// CHECK-DAG: %[[IN_J_SZ:.*]] = affine.min #[[MAP5]]132// CHECK: %[[SUB_IN:.*]] = tensor.extract_slice %[[IN]][%[[IN_I]], %[[IN_J]]] [%[[IN_I_SZ]], %[[IN_J_SZ]]] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>133// CHECK: %[[SUB_OUT:.*]] = tensor.extract_slice %[[ITER1]][%[[I]], %[[J]], 0, 0] [%[[OUT_I_SZ]], %[[OUT_J_SZ]], 8, 2] [1, 1, 1, 1] : tensor<?x?x8x2xf32> to tensor<?x?x8x2xf32>134// CHECK: %[[SUB_RES:.*]] = linalg.pack135// CHECK-SAME: %[[SUB_IN]] padding_value(%[[PAD]] : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 2]136// CHECK-SAME: into %[[SUB_OUT]]137// CHECK: %[[INSERT:.*]] = tensor.insert_slice %[[SUB_RES]] into %[[ITER1]]138// CHECK: scf.yield %[[INSERT]] : tensor<?x?x8x2xf32>139// CHECK: }140// CHECK: scf.yield %[[RES1:.*]] : tensor<?x?x8x2xf32>141// CHECK: }142// CHECK: return %[[VAL_34:.*]] : tensor<?x?x8x2xf32>143// CHECK: }144func.func @pad_and_pack_partially_dynamic(%input: tensor<?x?xf32>, %output: tensor<?x?x8x2xf32>, %pad: f32) -> tensor<?x?x8x2xf32> {145 %0 = linalg.pack %input padding_value(%pad : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 2] into %output : tensor<?x?xf32> -> tensor<?x?x8x2xf32>146 return %0 : tensor<?x?x8x2xf32>147}148 149module attributes {transform.with_named_sequence} {150 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {151 %0 = transform.structured.match ops{["linalg.pack"]} in %arg1 : (!transform.any_op) -> !transform.any_op152 %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [2, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)153 transform.yield154 }155}156 157// -----158 159// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0] -> (-d0 + s0, 2)>160// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0] -> (-d0 + s0, 4)>161// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0)[s0] -> (d0 * s0)>162// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s0, -(d1 * s0) + s1)>163// CHECK: func.func @pad_and_pack_fully_dynamic(164// CHECK-SAME: %[[IN:.*]]: tensor<?x?xf32>,165// CHECK-SAME: %[[OUT:.*]]: tensor<?x?x?x?xf32>,166// CHECK-SAME: %[[PAD:.*]]: f32,167// CHECK-SAME: %[[TILE_0:.*]]: index,168// CHECK-SAME: %[[TILE_1:.*]]: index) -> tensor<?x?x?x?xf32> {169// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index170// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index171// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index172// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index173// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index174// CHECK-DAG: %[[OUT_D0:.*]] = tensor.dim %[[OUT]], %[[C0]] : tensor<?x?x?x?xf32>175// CHECK-DAG: %[[OUT_D1:.*]] = tensor.dim %[[OUT]], %[[C1]] : tensor<?x?x?x?xf32>176// CHECK: %[[RES0:.*]] = scf.for %[[I:.*]] = %[[C0]] to %[[OUT_D0]] step %[[C2]] iter_args(%[[ITER0:.*]] = %[[OUT]]) -> (tensor<?x?x?x?xf32>) {177// CHECK: %[[RES1:.*]] = scf.for %[[J:.*]] = %[[C0]] to %[[OUT_D1]] step %[[C4]] iter_args(%[[ITER1:.*]] = %[[ITER0]]) -> (tensor<?x?x?x?xf32>) {178// CHECK-DAG: %[[OUT_I_SZ:.*]] = affine.min #[[MAP0]](%[[I]])[%[[OUT_D0]]]179// CHECK-DAG: %[[OUT_J_SZ:.*]] = affine.min #[[MAP1]](%[[J]])[%[[OUT_D1]]]180// CHECK-DAG: %[[IN_D0:.*]] = tensor.dim %[[IN]], %[[C0]]181// CHECK-DAG: %[[IN_D1:.*]] = tensor.dim %[[IN]], %[[C1]]182// CHECK: %[[IN_I:.*]] = affine.apply #[[MAP2]](%[[I]])[%[[TILE_0]]]183// CHECK: %[[IN_I_SZ:.*]] = affine.min #[[MAP3]](%[[OUT_I_SZ]], %[[I]])[%[[TILE_0]], %[[IN_D0]]]184// CHECK: %[[IN_J:.*]] = affine.apply #[[MAP2]](%[[J]])[%[[TILE_1]]]185// CHECK: %[[IN_J_SZ:.*]] = affine.min #[[MAP3]](%[[OUT_J_SZ]], %[[J]])[%[[TILE_1]], %[[IN_D1]]]186// CHECK: %[[SUB_IN:.*]] = tensor.extract_slice %[[IN]][%[[IN_I]], %[[IN_J]]] [%[[IN_I_SZ]], %[[IN_J_SZ]]] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>187// CHECK: %[[OUT_D2:.+]] = tensor.dim %[[ITER1]], %[[C2]]188// CHECK: %[[OUT_D3:.+]] = tensor.dim %[[ITER1]], %[[C3]]189// CHECK: %[[SUB_OUT:.*]] = tensor.extract_slice %[[ITER1]][%[[I]], %[[J]], 0, 0] [%[[OUT_I_SZ]], %[[OUT_J_SZ]], %[[OUT_D2]], %[[OUT_D3]]] [1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<?x?x?x?xf32>190// CHECK: %[[PACK:.*]] = linalg.pack191// CHECK-SAME: %[[SUB_IN]] padding_value(%[[PAD]] : f32) inner_dims_pos = [0, 1] inner_tiles = [%[[TILE_0]], %[[TILE_1]]]192// CHECK-SAME: into %[[SUB_OUT]]193// CHECK: %[[INSERT:.*]] = tensor.insert_slice %[[PACK]] into %[[ITER1]]194// CHECK: scf.yield %[[INSERT]] : tensor<?x?x?x?xf32>195// CHECK: }196// CHECK: scf.yield %[[RES1:.*]] : tensor<?x?x?x?xf32>197// CHECK: }198// CHECK: return %[[RES0:.*]] : tensor<?x?x?x?xf32>199// CHECK: }200func.func @pad_and_pack_fully_dynamic(%source: tensor<?x?xf32>, %dest: tensor<?x?x?x?xf32>, %pad: f32, %tile_n : index, %tile_m : index) -> tensor<?x?x?x?xf32> {201 %0 = linalg.pack %source padding_value(%pad : f32) inner_dims_pos = [0, 1] inner_tiles = [%tile_n, %tile_m] into %dest : tensor<?x?xf32> -> tensor<?x?x?x?xf32>202 return %0 : tensor<?x?x?x?xf32>203}204 205module attributes {transform.with_named_sequence} {206 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {207 %0 = transform.structured.match ops{["linalg.pack"]} in %arg1 : (!transform.any_op) -> !transform.any_op208 %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [2, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)209 transform.yield210 }211}212 213// -----214 215// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> (d0 floordiv 32)>216// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0 mod 32)>217// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> ((d0 + 1) floordiv 32 - d0 floordiv 32 + 1)>218// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0) -> (d0 floordiv 16)>219// CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0) -> (d0 mod 16)>220// CHECK-DAG: #[[MAP6:.+]] = affine_map<(d0) -> ((d0 + 3) floordiv 16 - d0 floordiv 16 + 1)>221// CHECK: func.func @NCnc_to_NC222// CHECK-SAME: %[[IN:[A-Za-z0-9]+]]:223// CHECK-SAME: %[[OUT:[A-Za-z0-9]+]]:224// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index225// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index226// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index227// CHECK-DAG: %[[C128:.*]] = arith.constant 128 : index228// CHECK-DAG: %[[C256:.*]] = arith.constant 256 : index229// CHECK: %{{.+}} = scf.for %[[I:.+]] = %[[C0]] to %[[C256]] step %[[C2]]230// CHECK: %{{.+}} = scf.for %[[J:.+]] = %[[C0]] to %[[C128]] step %[[C4]]231// CHECK-DAG: %[[IN_I:.+]] = affine.apply #[[MAP0]](%[[I]])232// CHECK-DAG: %[[OFFSET_I:.+]] = affine.apply #[[MAP1]](%[[I]])233// CHECK-DAG: %[[IN_I_SZ:.+]] = affine.apply #[[MAP2]](%[[I]])234// CHECK-DAG: %[[IN_J:.+]] = affine.apply #[[MAP4]](%[[J]])235// CHECK-DAG: %[[OFFSET_J:.+]] = affine.apply #[[MAP5]](%[[J]])236// CHECK-DAG: %[[IN_J_SZ:.+]] = affine.apply #[[MAP6]](%[[J]])237// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[IN]]238// CHECK-SAME: [%[[IN_I]], %[[IN_J]], 0, 0] [%[[IN_I_SZ]], %[[IN_J_SZ]], 32, 16]239// CHECK-SAME: : tensor<8x8x32x16xf32> to tensor<?x?x32x16xf32>240// CHECK: %[[EMPTY:.+]] = tensor.empty241// CHECK: %[[UNPACK:.+]] = linalg.unpack242// CHECK-SAME: %[[SLICE]] inner_dims_pos = [0, 1] inner_tiles = [32, 16]243// CHECK-SAME: into %[[EMPTY]]244// CHECK: %[[UNPACK_SLICE:.+]] = tensor.extract_slice %[[UNPACK]]245// CHECK-SAME: [%[[OFFSET_I]], %[[OFFSET_J]]] [2, 4]246// CHECK: %[[RES:.+]] = tensor.insert_slice %[[UNPACK_SLICE]]247// CHECK-SAME: into %{{.+}}[%[[I]], %[[J]]] [2, 4]248// CHECK: scf.yield %[[RES]]249func.func @NCnc_to_NC(%source: tensor<8x8x32x16xf32>, %dest: tensor<256x128xf32>) -> tensor<256x128xf32> {250 %0 = linalg.unpack %source inner_dims_pos = [0, 1] inner_tiles = [32, 16] into %dest : tensor<8x8x32x16xf32> -> tensor<256x128xf32>251 return %0 : tensor<256x128xf32>252}253 254module attributes {transform.with_named_sequence} {255 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {256 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op257 %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [2, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)258 transform.yield259 }260}261 262// -----263 264// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> (d0 floordiv 32)>265// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0 mod 32)>266// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> ((d0 + 1) floordiv 32 - d0 floordiv 32 + 1)>267// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0) -> (d0 floordiv 8)>268// CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0) -> (d0 mod 8)>269// CHECK-DAG: #[[MAP6:.+]] = affine_map<(d0) -> ((d0 + 3) floordiv 8 - d0 floordiv 8 + 1)>270// CHECK: func.func @CKkc_to_KC271// CHECK-SAME: %[[IN:[A-Za-z0-9]+]]:272// CHECK-SAME: %[[OUT:[A-Za-z0-9]+]]:273// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index274// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index275// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index276// CHECK-DAG: %[[C128:.*]] = arith.constant 128 : index277// CHECK-DAG: %[[C256:.*]] = arith.constant 256 : index278// CHECK: %{{.+}} = scf.for %[[K:.+]] = %[[C0]] to %[[C128]] step %[[C2]]279// CHECK: %{{.+}} = scf.for %[[C:.+]] = %[[C0]] to %[[C256]] step %[[C4]]280// CHECK-DAG: %[[IN_K:.+]] = affine.apply #[[MAP0]](%[[K]])281// CHECK-DAG: %[[OFFSET_K:.+]] = affine.apply #[[MAP1]](%[[K]])282// CHECK-DAG: %[[IN_K_SZ:.+]] = affine.apply #[[MAP2]](%[[K]])283// CHECK-DAG: %[[IN_C:.+]] = affine.apply #[[MAP4]](%[[C]])284// CHECK-DAG: %[[OFFSET_C:.+]] = affine.apply #[[MAP5]](%[[C]])285// CHECK-DAG: %[[IN_C_SZ:.+]] = affine.apply #[[MAP6]](%[[C]])286// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[IN]]287// CHECK: [%[[IN_C]], %[[IN_K]], 0, 0] [%[[IN_C_SZ]], %[[IN_K_SZ]], 32, 8]288// CHECK: %[[EMPTY:.+]] = tensor.empty289// CHECK: %[[UNPACK:.+]] = linalg.unpack290// CHECK-SAME: %[[IN_SLICE]] outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 8]291// CHECK-SAME: into %[[EMPTY]]292// CHECK: %[[UNPACK_SLICE:.+]] = tensor.extract_slice %[[UNPACK]]293// CHECK-SAME: [%[[OFFSET_K]], %[[OFFSET_C]]] [2, 4]294// CHECK: %[[RES:.+]] = tensor.insert_slice %[[UNPACK_SLICE]]295// CHECK-SAME: into %{{.+}}[%[[K]], %[[C]]] [2, 4]296// CHECK: scf.yield %[[RES]]297func.func @CKkc_to_KC(%source: tensor<32x4x32x8xf32>, %dest: tensor<128x256xf32>) -> tensor<128x256xf32> {298 %0 = linalg.unpack %source outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 8] into %dest : tensor<32x4x32x8xf32> -> tensor<128x256xf32>299 return %0 : tensor<128x256xf32>300}301 302module attributes {transform.with_named_sequence} {303 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {304 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op305 %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [2, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)306 transform.yield307 }308}309 310// -----311 312// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> (d0 floordiv 2)>313// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0 floordiv 4)>314// CHECK: func.func @perfect_CKkc_to_KC315// CHECK-SAME: %[[IN:[A-Za-z0-9]+]]:316// CHECK-SAME: %[[OUT:[A-Za-z0-9]+]]:317// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index318// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index319// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index320// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index321// CHECK-DAG: %[[C128:.*]] = arith.constant 128 : index322// CHECK: %{{.+}} = scf.for %[[K:.+]] = %[[C0]] to %[[C8]] step %[[C2]]323// CHECK: %{{.+}} = scf.for %[[C:.+]] = %[[C0]] to %[[C128]] step %[[C4]]324// CHECK-DAG: %[[IN_K:.+]] = affine.apply #[[MAP0]](%[[K]])325// CHECK-DAG: %[[IN_C:.+]] = affine.apply #[[MAP1]](%[[C]])326// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[IN]]327// CHECK: [%[[IN_C]], %[[IN_K]], 0, 0] [1, 1, 2, 4]328// CHECK: %[[ITER_SLICE:.+]] = tensor.extract_slice %{{.+}}[%[[K]], %[[C]]] [2, 4]329// CHECK: %[[UNPACK:.+]] = linalg.unpack330// CHECK-SAME: %[[IN_SLICE]] outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [2, 4]331// CHECK-SAME: into %[[ITER_SLICE]]332// CHECK: %[[RES:.+]] = tensor.insert_slice %[[UNPACK]]333// CHECK-SAME: into %{{.+}}[%[[K]], %[[C]]] [2, 4]334// CHECK: scf.yield %[[RES]]335func.func @perfect_CKkc_to_KC(%source: tensor<32x4x2x4xf32>, %dest: tensor<8x128xf32>) -> tensor<8x128xf32> {336 %0 = linalg.unpack %source outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [2, 4] into %dest : tensor<32x4x2x4xf32> -> tensor<8x128xf32>337 return %0 : tensor<8x128xf32>338}339 340module attributes {transform.with_named_sequence} {341 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {342 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op343 %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [2, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)344 transform.yield345 }346}347 348// -----349 350// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0] -> (-d0 + s0, 2)>351// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0] -> (-d0 + s0, 4)>352// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> (d0 floordiv 2)>353// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0) -> (d0 ceildiv 2)>354// CHECK: func.func @dynamic_perfect_CKkc_to_KC355// CHECK-SAME: %[[IN:[A-Za-z0-9]+]]:356// CHECK-SAME: %[[OUT:[A-Za-z0-9]+]]:357// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index358// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index359// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index360// CHECK-DAG: %[[DIM_0:.+]] = tensor.dim %[[OUT]], %[[C0]]361// CHECK-DAG: %[[DIM_1:.+]] = tensor.dim %[[OUT]], %[[C1]]362// CHECK: %{{.+}} = scf.for %[[K:.+]] = %[[C0]] to %[[DIM_0]] step %[[C2]]363// CHECK: %{{.+}} = scf.for %[[C:.+]] = %[[C0]] to %[[DIM_1]] step %[[C4]]364// CHECK-DAG: %[[OUT_K_SZ:.+]] = affine.min #[[MAP0]](%[[K]])[%[[DIM_0]]]365// CHECK-DAG: %[[OUT_C_SZ:.+]] = affine.min #[[MAP1]](%[[C]])[%[[DIM_1]]]366// CHECK-DAG: %[[IN_K:.+]] = affine.apply #[[MAP2]](%[[K]])367// CHECK-DAG: %[[IN_C:.+]] = affine.apply #[[MAP2]](%[[C]])368// CHECK-DAG: %[[IN_C_SZ:.+]] = affine.apply #[[MAP3]](%[[OUT_C_SZ]])369// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[IN]]370// CHECK: [%[[IN_C]], %[[IN_K]], 0, 0] [%[[IN_C_SZ]], 1, 2, 2]371// CHECK: %[[ITER_SLICE:.+]] = tensor.extract_slice %{{.+}}[%[[K]], %[[C]]] [%[[OUT_K_SZ]], %[[OUT_C_SZ]]]372// CHECK: %[[UNPACK:.+]] = linalg.unpack373// CHECK-SAME: %[[IN_SLICE]] outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [2, 2]374// CHECK-SAME: into %[[ITER_SLICE]]375// CHECK: %[[RES:.+]] = tensor.insert_slice %[[UNPACK]]376// CHECK-SAME: into %{{.+}}[%[[K]], %[[C]]] [%[[OUT_K_SZ]], %[[OUT_C_SZ]]]377// CHECK: scf.yield %[[RES]]378 379func.func @dynamic_perfect_CKkc_to_KC(%source: tensor<?x?x2x2xf32>, %dest: tensor<?x?xf32>) -> tensor<?x?xf32> {380 %0 = linalg.unpack %source outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [2, 2] into %dest : tensor<?x?x2x2xf32> -> tensor<?x?xf32>381 return %0 : tensor<?x?xf32>382}383 384module attributes {transform.with_named_sequence} {385 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {386 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op387 %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [2, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)388 transform.yield389 }390}391 392// -----393 394// CHECK: #[[MAP:.+]] = affine_map<(d0) -> (d0 floordiv 2)>395// CHECK: func.func @perfect_NKPQk_to_NPQK(396// CHECK-SAME: %[[SOURCE:.+]]: tensor<1x4x6x6x2xf32>,397// CHECK-SAME: %{{.+}}: tensor<1x6x6x8xf32>)398// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index399// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index400// CHECK-DAG: %[[C6:.*]] = arith.constant 6 : index401// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index402// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index403// CHECK: %{{.+}} = scf.for %[[P:.+]] = %[[C0]] to %[[C6]] step %[[C1]]404// CHECK: %{{.+}} = scf.for %[[Q:.+]] = %[[C0]] to %[[C6]] step %[[C1]]405// CHECK: %{{.+}} = scf.for %[[K:.+]] = %[[C0]] to %[[C8]] step %[[C4]]406// CHECK: %[[K_SZ:.+]] = affine.apply #[[MAP]](%[[K]])407// CHECK: %[[SLICE_SOURCE:.+]] = tensor.extract_slice %[[SOURCE]][0, %[[K_SZ]], %[[P]], %[[Q]], 0]408// CHECK: %[[SLICE_DEST:.+]] = tensor.extract_slice %{{.+}}[0, %[[P]], %[[Q]], %[[K]]]409// CHECK: %[[UNPACK:.+]] = linalg.unpack410// CHECK-SAME: %[[SLICE_SOURCE]] outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [2]411// CHECK-SAME: into %[[SLICE_DEST]]412// CHECK: %[[RES:.+]] = tensor.insert_slice %[[UNPACK]]413// CHECK-SAME: into %{{.+}}[0, %[[P]], %[[Q]], %[[K]]]414// CHECK: scf.yield %[[RES]]415 416func.func @perfect_NKPQk_to_NPQK(%source: tensor<1x4x6x6x2xf32>, %dest: tensor<1x6x6x8xf32>) -> tensor<1x6x6x8xf32> {417 %0 = linalg.unpack %source outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [2] into %dest : tensor<1x4x6x6x2xf32> -> tensor<1x6x6x8xf32>418 return %0 : tensor<1x6x6x8xf32>419}420 421module attributes {transform.with_named_sequence} {422 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {423 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op424 %1, %loops:4 = transform.structured.tile_using_for %0 tile_sizes [1, 1, 1, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)425 transform.yield426 }427}428 429// -----430 431func.func private @get_dynamic_tile_size() -> index432 433// CHECK-LABEL: func.func @fully_dynamic_unpack434// CHECK-SAME: %[[SRC:[0-9a-zA-Z]+]]435// CHECK-SAME: %[[DST:[0-9a-zA-Z]+]]436// CHECK: %[[INNER_TS:.+]] = call @get_dynamic_tile_size() : () -> index437// CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[DST]])438// CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]])439// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[SRC]]440// CHECK: %[[EMPTY:.+]] = tensor.empty441// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[SLICE]]442// CHECK-SAME: inner_dims_pos = [1, 0] inner_tiles = [%[[INNER_TS]], %[[INNER_TS]]] into %[[EMPTY]]443func.func @fully_dynamic_unpack(%source: tensor<?x?x?x?xf32>, %dest: tensor<?x?xf32>) -> tensor<?x?xf32> {444 %0 = func.call @get_dynamic_tile_size() : () -> index445 %1 = linalg.unpack %source inner_dims_pos = [1, 0] inner_tiles = [%0, %0] into %dest : tensor<?x?x?x?xf32> -> tensor<?x?xf32>446 return %1 : tensor<?x?xf32>447}448 449module attributes {transform.with_named_sequence} {450 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {451 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op452 %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [4, 8] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)453 transform.yield454 }455}456 457// -----458 459// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0 * 2)>460// CHECK: func.func @perfect_NPQK_to_NKPQk461// CHECK-SAME: %[[SOURCE:.+]]: tensor<1x6x6x8xf32>,462// CHECK-SAME: %{{.+}}: tensor<1x4x6x6x2xf32>)463// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index464// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index465// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index466// CHECK-DAG: %[[C6:.+]] = arith.constant 6 : index467// CHECK: %{{.+}} = scf.for %[[ARG2:.+]] = %[[C0]] to %[[C4]] step %[[C1]]468// CHECK: %{{.+}} = scf.for %[[ARG4:.+]] = %[[C0]] to %[[C6]] step %[[C1]]469// CHECK: %{{.+}} = scf.for %[[ARG6:.+]] = %[[C0]] to %[[C6]] step %[[C1]]470// CHECK: %[[APPLY:.+]] = affine.apply #[[MAP1]](%[[ARG2]])471// CHECK: %[[SLICE_SOURCE:.+]] = tensor.extract_slice %[[SOURCE]][0, %[[ARG4]], %[[ARG6]], %[[APPLY]]]472// CHECK: %[[SLICE_DEST:.+]] = tensor.extract_slice %{{.+}}[0, %[[ARG2]], %[[ARG4]], %[[ARG6]], 0]473// CHECK: %[[PACK:.+]] = linalg.pack474// CHECK-SAME: %[[SLICE_SOURCE]] outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [2]475// CHECK-SAME: into %[[SLICE_DEST]]476// CHECK: %[[RES:.+]] = tensor.insert_slice %[[PACK]]477// CHECK-SAME: into %{{.+}}[0, %[[ARG2]], %[[ARG4]], %[[ARG6]], 0]478// CHECK: scf.yield %[[RES]]479 480func.func @perfect_NPQK_to_NKPQk(%source: tensor<1x6x6x8xf32>, %dest: tensor<1x4x6x6x2xf32>) -> tensor<1x4x6x6x2xf32> {481 %0 = linalg.pack %source outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [2] into %dest : tensor<1x6x6x8xf32> -> tensor<1x4x6x6x2xf32>482 return %0 : tensor<1x4x6x6x2xf32>483}484 485module attributes {transform.with_named_sequence} {486 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {487 %0 = transform.structured.match ops{["linalg.pack"]} in %arg1 : (!transform.any_op) -> !transform.any_op488 %1, %loops:4 = transform.structured.tile_using_for %0 tile_sizes [1, 1, 1, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)489 transform.yield490 }491}492