mlir/test/Dialect/Vector/fold-arith-extf-into-vector-contract.mlir

   1 // RUN: mlir-opt -split-input-file -test-fold-arith-extf-into-vector-contract-patterns %s | FileCheck %s
   2
   3
   4 // CHECK-DAG: #[[$map0:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
   5 // CHECK-DAG: #[[$map1:.*]] = affine_map<(d0, d1, d2) -> (d2, d1)>
   6 // CHECK-DAG: #[[$map2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
   7 // CHECK-LABEL: func.func @fold_arith_extf_into_contract
   8 //  CHECK-SAME: (%[[ARG0:.*]]: vector<64x64xf16>, %[[ARG1:.*]]: vector<64x64xf16>, %[[ARG2:.*]]: vector<64x64xf32>)
   9 //  CHECK-NEXT:   %[[R:.+]] = vector.contract {indexing_maps = [#[[$map0]], #[[$map1]], #[[$map2]]],
  10 //  CHECK-SAME:   iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
  11 //  CHECK-SAME:   %[[ARG0]], %[[ARG1]], %[[ARG2]] : vector<64x64xf16>, vector<64x64xf16> into vector<64x64xf32>
  12 //  CHECK-NEXT:   return %[[R]] : vector<64x64xf32>
  13 func.func @fold_arith_extf_into_contract(
  14   %arg0: vector<64x64xf16>,
  15   %arg1: vector<64x64xf16>,
  16   %arg2: vector<64x64xf32>) -> vector<64x64xf32> {
  17     %lhs_f32 = arith.extf %arg0 : vector<64x64xf16> to vector<64x64xf32>
  18     %rhs_f32 = arith.extf %arg1 : vector<64x64xf16> to vector<64x64xf32>
  19     %result = vector.contract {
  20       indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>],
  21       iterator_types = ["parallel", "parallel", "reduction"],
  22       kind = #vector.kind<add>}
  23       %lhs_f32, %rhs_f32, %arg2 : vector<64x64xf32>, vector<64x64xf32> into vector<64x64xf32>
  24     return %result : vector<64x64xf32>
  25 }
  26
  27 // -----
  28
  29 // CHECK-DAG: #[[$map0:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
  30 // CHECK-DAG: #[[$map1:.*]] = affine_map<(d0, d1, d2) -> (d2, d1)>
  31 // CHECK-DAG: #[[$map2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
  32 // CHECK-LABEL: func.func @fold_arith_extf_into_contract_scalable
  33 //  CHECK-SAME: (%[[ARG0:.*]]: vector<[64]x64xf16>, %[[ARG1:.*]]: vector<64x64xf16>, %[[ARG2:.*]]: vector<[64]x64xf32>)
  34 //  CHECK-NEXT:   %[[R:.+]] = vector.contract {indexing_maps = [#[[$map0]], #[[$map1]], #[[$map2]]],
  35 //  CHECK-SAME:   iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
  36 //  CHECK-SAME:   %[[ARG0]], %[[ARG1]], %[[ARG2]] : vector<[64]x64xf16>, vector<64x64xf16> into vector<[64]x64xf32>
  37 //  CHECK-NEXT:   return %[[R]] : vector<[64]x64xf32>
  38 func.func @fold_arith_extf_into_contract_scalable(
  39   %arg0: vector<[64]x64xf16>,
  40   %arg1: vector<64x64xf16>,
  41   %arg2: vector<[64]x64xf32>) -> vector<[64]x64xf32> {
  42     %lhs_f32 = arith.extf %arg0 : vector<[64]x64xf16> to vector<[64]x64xf32>
  43     %rhs_f32 = arith.extf %arg1 : vector<64x64xf16> to vector<64x64xf32>
  44     %result = vector.contract {
  45       indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>],
  46       iterator_types = ["parallel", "parallel", "reduction"],
  47       kind = #vector.kind<add>}
  48       %lhs_f32, %rhs_f32, %arg2 : vector<[64]x64xf32>, vector<64x64xf32> into vector<[64]x64xf32>
  49     return %result : vector<[64]x64xf32>
  50 }
  51
  52 // -----
  53
  54 // CHECK-LABEL: func.func @fold_arith_extsi_into_contract
  55 //  CHECK-SAME: (%[[ARG0:.*]]: vector<64x64xi8>, %[[ARG1:.*]]: vector<64x64xi8>, %[[ARG2:.*]]: vector<64x64xi32>)
  56 //  CHECK-NEXT:   %[[R:.+]] = vector.contract
  57 //  CHECK-SAME:   iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
  58 //  CHECK-SAME:   %[[ARG0]], %[[ARG1]], %[[ARG2]] : vector<64x64xi8>, vector<64x64xi8> into vector<64x64xi32>
  59 //  CHECK-NEXT:   return %[[R]] : vector<64x64xi32>
  60 func.func @fold_arith_extsi_into_contract(
  61   %arg0: vector<64x64xi8>,
  62   %arg1: vector<64x64xi8>,
  63   %arg2: vector<64x64xi32>) -> vector<64x64xi32> {
  64     %lhs_i32 = arith.extsi %arg0 : vector<64x64xi8> to vector<64x64xi32>
  65     %rhs_i32 = arith.extsi %arg1 : vector<64x64xi8> to vector<64x64xi32>
  66     %result = vector.contract {
  67       indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>],
  68       iterator_types = ["parallel", "parallel", "reduction"],
  69       kind = #vector.kind<add>}
  70       %lhs_i32, %rhs_i32, %arg2 : vector<64x64xi32>, vector<64x64xi32> into vector<64x64xi32>
  71     return %result : vector<64x64xi32>
  72 }