110 lines · plain
1// RUN: mlir-opt -split-input-file -transform-interpreter --canonicalize \2// RUN: -transform-preload-library='transform-library-paths=%p/td/decompose-unpack.mlir' \3// RUN: -transform-interpreter=entry-point=decompose_unpack \4// RUN: -transform-interpreter %s | FileCheck %s5 6func.func @KCRSsr_to_KCRS(%arg0: tensor<1x1x4x8x8x32xf32>, %arg1: tensor<1x1x128x64xf32>) -> tensor<1x1x128x64xf32> {7 %0 = linalg.unpack %arg0 inner_dims_pos = [3, 2] inner_tiles = [8, 32] into %arg1 : tensor<1x1x4x8x8x32xf32> -> tensor<1x1x128x64xf32>8 return %0 : tensor<1x1x128x64xf32>9}10 11module attributes {transform.with_named_sequence} {12 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {13 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op14 %1, %loops:4 = transform.structured.tile_using_for %0 tile_sizes [1, 1, 32, 8] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)15 transform.yield16 }17}18// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> (d0 floordiv 32)>19// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0 floordiv 8)>20// CHECK: func.func @KCRSsr_to_KCRS21// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]22// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]23// CHECK: %{{.+}} = scf.for %[[R:[a-zA-Z0-9]+]] =24// CHECK: %{{.+}} = scf.for %[[S:[a-zA-Z0-9]+]] =25// CHECK: %[[IN_R:.+]] = affine.apply #[[MAP0]](%[[R]])26// CHECK: %[[IN_S:.+]] = affine.apply #[[MAP1]](%[[S]])27// CHECK: %[[SRC_SLICE:.+]] = tensor.extract_slice %[[SRC]]28// CHECK-SAME: [0, 0, %[[IN_R]], %[[IN_S]], 0, 0] [1, 1, 1, 1, 8, 32] [1, 1, 1, 1, 1, 1]29// CHECK: %[[TILE:.+]] = tensor.extract_slice %[[SRC_SLICE]]30// CHECK-SAME: [0, 0, 0, 0, 0, 0] [1, 1, 1, 1, 8, 32] [1, 1, 1, 1, 1, 1] : tensor<1x1x1x1x8x32xf32> to tensor<8x32xf32>31// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<32x8xf32>32// CHECK: %[[TRANSP:.+]] = linalg.transpose33// CHECK-SAME: ins(%[[TILE]]34// CHECK-SAME: outs(%[[EMPTY]]35// CHECK-SAME: permutation = [1, 0]36// CHECK: %{{.+}} = tensor.insert_slice %[[TRANSP]] into %{{.+}}37 38// -----39 40func.func @unpack_and_extract_slice(%arg0: tensor<2x8x8x2xf32>, %arg1: tensor<13x15xf32>) -> tensor<13x15xf32> {41 %0 = linalg.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 2] into %arg1 : tensor<2x8x8x2xf32> -> tensor<13x15xf32>42 return %0 : tensor<13x15xf32>43}44// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> (-d0 + 13, 8)>45// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (-d0 + 15, 2)>46// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> (d0 floordiv 8)>47// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0) -> (d0 floordiv 2)>48// CHECK: func.func @unpack_and_extract_slice49// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]50// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]51// CHECK: %{{.+}} = scf.for %[[I:[a-zA-Z0-9]+]] =52// CHECK: %{{.+}} = scf.for %[[J:[a-zA-Z0-9]+]] =53// CHECK-DAG: %[[OUT_I_SZ:.+]] = affine.min #[[MAP0]](%[[I]])54// CHECK-DAG: %[[OUT_J_SZ:.+]] = affine.min #[[MAP1]](%[[J]])55// CHECK-DAG: %[[IN_I:.+]] = affine.apply #[[MAP2]](%[[I]])56// CHECK-DAG: %[[IN_J:.+]] = affine.apply #[[MAP3]](%[[J]])57// CHECK: %[[SRC_SLICE:.+]] = tensor.extract_slice %[[SRC]]58// CHECK-SAME: [%[[IN_I]], %[[IN_J]], 0, 0] [1, 1, 8, 2] [1, 1, 1, 1]59// CHECK: %[[ITER_SLICE:.+]] = tensor.extract_slice %{{[a-zA-Z0-9]+}}60// CHECK-SAME: [%[[I]], %[[J]]] [%[[OUT_I_SZ]], %[[OUT_J_SZ]]]61// CHECK: %[[TILE:.+]] = tensor.extract_slice %[[SRC_SLICE]]62// CHECK-SAME: [0, 0, 0, 0] [1, 1, 8, 2] [1, 1, 1, 1] : tensor<1x1x8x2xf32> to tensor<8x2xf32>63// CHECK-NOT: linalg.transpose64// CHECK: %[[UNPACK_TILE:.+]] = tensor.extract_slice %[[TILE]]65// CHECK-SAME: [0, 0] [%[[OUT_I_SZ]], %[[OUT_J_SZ]]] [1, 1]66// CHECK: %[[INSERT1:.+]] = tensor.insert_slice %[[UNPACK_TILE]] into %[[ITER_SLICE]]67// CHECK-SAME: [0, 0] [%[[OUT_I_SZ]], %[[OUT_J_SZ]]] [1, 1]68// CHECK: %[[INSERT2:.+]] = tensor.insert_slice %[[INSERT1]] into %{{[a-zA-Z0-9]+}}69// CHECK-SAME: [%[[I]], %[[J]]] [%[[OUT_I_SZ]], %[[OUT_J_SZ]]] [1, 1]70 71module attributes {transform.with_named_sequence} {72 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {73 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op74 %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [8, 2] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)75 transform.yield76 }77}78 79// -----80 81func.func @CKkc_to_KC(%arg0: tensor<32x4x32x8xf32>, %arg1: tensor<128x256xf32>) -> tensor<128x256xf32> {82 %0 = linalg.unpack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 8] into %arg1 : tensor<32x4x32x8xf32> -> tensor<128x256xf32>83 return %0 : tensor<128x256xf32>84}85// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> (d0 floordiv 32)>86// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0 floordiv 8)>87// CHECK: func.func @CKkc_to_KC88// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]89// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]90// CHECK: %{{.+}} = scf.for %[[K:[a-zA-Z0-9]+]] =91// CHECK: %{{.+}} = scf.for %[[C:[a-zA-Z0-9]+]] =92// CHECK: %[[IN_K:.+]] = affine.apply #[[MAP0]](%[[K]])93// CHECK: %[[IN_C:.+]] = affine.apply #[[MAP1]](%[[C]])94// CHECK: %[[SRC_SLICE:.+]] = tensor.extract_slice %[[SRC]]95// CHECK-SAME: [%[[IN_C]], %[[IN_K]], 0, 0] [1, 1, 32, 8] [1, 1, 1, 1]96// CHECK: %[[TILE:.+]] = tensor.extract_slice %[[SRC_SLICE]]97// CHECK-SAME: [0, 0, 0, 0] [1, 1, 32, 8] [1, 1, 1, 1] : tensor<1x1x32x8xf32> to tensor<32x8xf32>98// CHECK-NOT: linalg.transpose99// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[TILE]] into %{{[a-zA-Z0-9]+}}100// CHECK-SAME: [%[[K]], %[[C]]] [32, 8] [1, 1]101 102 103module attributes {transform.with_named_sequence} {104 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {105 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op106 %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [32, 8] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)107 transform.yield108 }109}110