93 lines · plain
1// RUN: mlir-opt -split-input-file -transform-interpreter --canonicalize \2// RUN: -transform-preload-library='transform-library-paths=%p/td/decompose-pack.mlir' \3// RUN: -transform-interpreter=entry-point=decompose_pack \4// RUN: -transform-interpreter %s | FileCheck %s5 6func.func @KCRS_to_KCRSsr(%arg0: tensor<1x1x128x64xf32>, %arg1: tensor<1x1x4x8x8x32xf32>) -> tensor<1x1x4x8x8x32xf32> {7 %0 = linalg.pack %arg0 inner_dims_pos = [3, 2] inner_tiles = [8, 32] into %arg1 : tensor<1x1x128x64xf32> -> tensor<1x1x4x8x8x32xf32>8 return %0 : tensor<1x1x4x8x8x32xf32>9}10// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> (d0 * 32)>11// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> (d0 * 8)>12// CHECK: func.func @KCRS_to_KCRSsr13// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]14// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]15// CHECK: scf.for %[[R:[a-zA-Z0-9]+]] =16// CHECK: scf.for %[[S:[a-zA-Z0-9]+]] {{.*}} iter_args(%[[ITER_SLICE:.*]] =17// CHECK: %[[IN_R:.+]] = affine.apply #[[MAP0]](%[[R]])18// CHECK: %[[IN_S:.+]] = affine.apply #[[MAP2]](%[[S]])19// CHECK: %[[SRC_SLICE:.+]] = tensor.extract_slice %[[SRC]]20// CHECK-SAME: [0, 0, %[[IN_R]], %[[IN_S]]] [1, 1, 32, 8] [1, 1, 1, 1]21// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<1x1x8x32xf32>22// CHECK: %[[TRANSP:.*]] = linalg.transpose23// CHECK-SAME: ins(%[[SRC_SLICE]] : tensor<1x1x32x8xf32>)24// CHECK-SAME: outs(%[[EMPTY]] : tensor<1x1x8x32xf32>)25// CHECK-SAME: permutation = [0, 1, 3, 2]26// CHECK: %{{.+}} = tensor.insert_slice %[[TRANSP]] into %{{.+}}27 28module attributes {transform.with_named_sequence} {29 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {30 %0 = transform.structured.match ops{["linalg.pack"]} in %arg1 : (!transform.any_op) -> !transform.any_op31 %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)32 transform.yield33 }34}35 36// -----37 38func.func @pad_and_pack(%arg0: tensor<13x15xf32>, %arg1: tensor<2x8x8x2xf32>, %arg2: f32) -> tensor<2x8x8x2xf32> {39 %0 = linalg.pack %arg0 padding_value(%arg2 : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 2] into %arg1 : tensor<13x15xf32> -> tensor<2x8x8x2xf32>40 return %0 : tensor<2x8x8x2xf32>41}42// CHECK: func.func @pad_and_pack43// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]44// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]45// CHECK-SAME: %[[PAD_VAL:[a-zA-Z0-9]+]]46// CHECK: scf.for47// CHECK: scf.for48// CHECK: %[[SRC_SLICE]] = tensor.extract_slice %[[SRC]]49// CHECK: %[[PAD:.+]] = tensor.pad %[[SRC_SLICE]]50// CHECK: tensor.yield %[[PAD_VAL]]51// CHECK: } : tensor<?x?xf32> to tensor<8x2xf32>52// CHECK-NOT: linalg.transpose53// CHECK: %{{.+}} = tensor.insert_slice %[[PAD]] into %{{.+}}54 55module attributes {transform.with_named_sequence} {56 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {57 %0 = transform.structured.match ops{["linalg.pack"]} in %arg1 : (!transform.any_op) -> !transform.any_op58 %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [1, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)59 transform.yield60 }61}62 63// -----64 65 66func.func @KC_to_CKkc(%arg0: tensor<128x256xf32>, %arg1: tensor<32x4x32x8xf32>) -> tensor<32x4x32x8xf32> {67 %0 = linalg.pack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 8] into %arg1 : tensor<128x256xf32> -> tensor<32x4x32x8xf32>68 return %0 : tensor<32x4x32x8xf32>69}70// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> (d0 * 32)>71// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> (d0 * 8)>72// CHECK: func.func @KC_to_CKkc73// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]74// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]75// CHECK: %{{.+}} = scf.for %[[C:[a-zA-Z0-9]+]] =76// CHECK: %{{.+}} = scf.for %[[K:[a-zA-Z0-9]+]] =77// CHECK-DAG: %[[IN_K:.+]] = affine.apply #[[MAP0]](%[[K]])78// CHECK-DAG: %[[IN_C:.+]] = affine.apply #[[MAP2]](%[[C]])79// CHECK: %[[TILE:.+]] = tensor.extract_slice %[[SRC]]80// CHECK-SAME: [%[[IN_K]], %[[IN_C]]] [32, 8] [1, 1]81// CHECK-NOT: linalg.transpose82// CHECK: %[[SUB_ITER:.+]] = tensor.insert_slice %[[TILE]] into %{{[a-zA-Z0-9]+}}83// CHECK-SAME: [0, 0, 0, 0] [1, 1, 32, 8] [1, 1, 1, 1] : tensor<32x8xf32> into tensor<1x1x32x8xf32>84// CHECK: %{{.+}} = tensor.insert_slice %[[SUB_ITER]] into %{{[a-zA-Z0-9]+}}85// CHECK-SAME: [%[[C]], %[[K]], 0, 0] [1, 1, 32, 8] [1, 1, 1, 1] : tensor<1x1x32x8xf32> into tensor<32x4x32x8xf32>86module attributes {transform.with_named_sequence} {87 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {88 %0 = transform.structured.match ops{["linalg.pack"]} in %arg1 : (!transform.any_op) -> !transform.any_op89 %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [1, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)90 transform.yield91 }92}93