mlir/test/Dialect/Vector/break-down-vector-reduction.mlir

   1 // RUN: mlir-opt %s --test-vector-break-down-reduction-patterns --cse | FileCheck %s
   2
   3 // NOTE: This test pass is set break down vector reductions of size 2 or fewer.
   4
   5 // CHECK-LABEL:   func.func @reduce_2x_f32(
   6 // CHECK-SAME:     %[[ARG0:.+]]: vector<2xf32>) -> (f32, f32, f32, f32, f32, f32) {
   7 // CHECK-DAG:      %[[E0:.+]] = vector.extract %[[ARG0]][0] : f32 from vector<2xf32>
   8 // CHECK-DAG:      %[[E1:.+]] = vector.extract %[[ARG0]][1] : f32 from vector<2xf32>
   9 // CHECK-DAG:      %[[R0:.+]] = arith.addf %[[E0]], %[[E1]] : f32
  10 // CHECK-DAG:      %[[R1:.+]] = arith.mulf %[[E0]], %[[E1]] : f32
  11 // CHECK-DAG:      %[[R2:.+]] = arith.minnumf %[[E0]], %[[E1]] : f32
  12 // CHECK-DAG:      %[[R3:.+]] = arith.maxnumf %[[E0]], %[[E1]] : f32
  13 // CHECK-DAG:      %[[R4:.+]] = arith.minimumf %[[E0]], %[[E1]] : f32
  14 // CHECK-DAG:      %[[R5:.+]] = arith.maximumf %[[E0]], %[[E1]] : f32
  15 // CHECK:          return %[[R0]], %[[R1]], %[[R2]], %[[R3]], %[[R4]], %[[R5]]
  16 func.func @reduce_2x_f32(%arg0: vector<2xf32>) -> (f32, f32, f32, f32, f32, f32) {
  17   %0 = vector.reduction <add>, %arg0 : vector<2xf32> into f32
  18   %1 = vector.reduction <mul>, %arg0 : vector<2xf32> into f32
  19   %2 = vector.reduction <minnumf>, %arg0 : vector<2xf32> into f32
  20   %3 = vector.reduction <maxnumf>, %arg0 : vector<2xf32> into f32
  21   %4 = vector.reduction <minimumf>, %arg0 : vector<2xf32> into f32
  22   %5 = vector.reduction <maximumf>, %arg0 : vector<2xf32> into f32
  23   return %0, %1, %2, %3, %4, %5 : f32, f32, f32, f32, f32, f32
  24 }
  25
  26 // CHECK-LABEL:   func.func @reduce_2x_i32(
  27 // CHECK-SAME:     %[[ARG0:.+]]: vector<2xi32>) -> (i32, i32, i32, i32, i32, i32, i32, i32, i32) {
  28 // CHECK-DAG:      %[[E0:.+]] = vector.extract %[[ARG0]][0] : i32 from vector<2xi32>
  29 // CHECK-DAG:      %[[E1:.+]] = vector.extract %[[ARG0]][1] : i32 from vector<2xi32>
  30 // CHECK-DAG:      %[[R0:.+]] = arith.addi %[[E0]], %[[E1]] : i32
  31 // CHECK-DAG:      %[[R1:.+]] = arith.muli %[[E0]], %[[E1]] : i32
  32 // CHECK-DAG:      %[[R2:.+]] = arith.minsi %[[E0]], %[[E1]] : i32
  33 // CHECK-DAG:      %[[R3:.+]] = arith.maxsi %[[E0]], %[[E1]] : i32
  34 // CHECK-DAG:      %[[R4:.+]] = arith.minui %[[E0]], %[[E1]] : i32
  35 // CHECK-DAG:      %[[R5:.+]] = arith.maxui %[[E0]], %[[E1]] : i32
  36 // CHECK-DAG:      %[[R6:.+]] = arith.andi %[[E0]], %[[E1]] : i32
  37 // CHECK-DAG:      %[[R7:.+]] = arith.ori %[[E0]], %[[E1]] : i32
  38 // CHECK-DAG:      %[[R8:.+]] = arith.xori %[[E0]], %[[E1]] : i32
  39 // CHECK:          return %[[R0]], %[[R1]], %[[R2]], %[[R3]], %[[R4]], %[[R5]], %[[R6]], %[[R7]], %[[R8]]
  40 func.func @reduce_2x_i32(%arg0: vector<2xi32>) -> (i32, i32, i32, i32, i32, i32, i32, i32, i32) {
  41   %0 = vector.reduction <add>, %arg0 : vector<2xi32> into i32
  42   %1 = vector.reduction <mul>, %arg0 : vector<2xi32> into i32
  43   %2 = vector.reduction <minsi>, %arg0 : vector<2xi32> into i32
  44   %3 = vector.reduction <maxsi>, %arg0 : vector<2xi32> into i32
  45   %4 = vector.reduction <minui>, %arg0 : vector<2xi32> into i32
  46   %5 = vector.reduction <maxui>, %arg0 : vector<2xi32> into i32
  47   %6 = vector.reduction <and>, %arg0 : vector<2xi32> into i32
  48   %7 = vector.reduction <or>, %arg0 : vector<2xi32> into i32
  49   %8 = vector.reduction <xor>, %arg0 : vector<2xi32> into i32
  50   return %0, %1, %2, %3, %4, %5, %6, %7, %8 : i32, i32, i32, i32, i32, i32, i32, i32, i32
  51 }
  52
  53 // CHECK-LABEL:   func.func @reduce_1x_f32(
  54 // CHECK-SAME:     %[[ARG0:.+]]: vector<1xf32>) -> f32 {
  55 // CHECK-NEXT:     %[[E0:.+]] = vector.extract %[[ARG0]][0] : f32 from vector<1xf32>
  56 // CHECK-NEXT:     return %[[E0]] : f32
  57 func.func @reduce_1x_f32(%arg0: vector<1xf32>) -> f32 {
  58   %0 = vector.reduction <add>, %arg0 : vector<1xf32> into f32
  59   return %0 : f32
  60 }
  61
  62 // CHECK-LABEL:   func.func @reduce_1x_acc_f32(
  63 // CHECK-SAME:     %[[ARG0:.+]]: vector<1xf32>, %[[ARG1:.+]]: f32) -> f32 {
  64 // CHECK-NEXT:     %[[E0:.+]] = vector.extract %[[ARG0]][0] : f32 from vector<1xf32>
  65 // CHECK-NEXT:     %[[R0:.+]] = arith.addf %[[E0]], %[[ARG1]] : f32
  66 // CHECK-NEXT:     return %[[R0]] : f32
  67 func.func @reduce_1x_acc_f32(%arg0: vector<1xf32>, %arg1: f32) -> f32 {
  68   %0 = vector.reduction <add>, %arg0, %arg1 : vector<1xf32> into f32
  69   return %0 : f32
  70 }
  71
  72 // CHECK-LABEL:   func.func @reduce_1x_acc_i32(
  73 // CHECK-SAME:     %[[ARG0:.+]]: vector<1xi32>, %[[ARG1:.+]]: i32) -> i32 {
  74 // CHECK-NEXT:     %[[E0:.+]] = vector.extract %[[ARG0]][0] : i32 from vector<1xi32>
  75 // CHECK-NEXT:     %[[R0:.+]] = arith.addi %[[E0]], %[[ARG1]] : i32
  76 // CHECK-NEXT:     return %[[R0]] : i32
  77 func.func @reduce_1x_acc_i32(%arg0: vector<1xi32>, %arg1: i32) -> i32 {
  78   %0 = vector.reduction <add>, %arg0, %arg1 : vector<1xi32> into i32
  79   return %0 : i32
  80 }
  81
  82 // CHECK-LABEL:   func.func @reduce_2x_acc_f32(
  83 // CHECK-SAME:     %[[ARG0:.+]]: vector<2xf32>, %[[ARG1:.+]]: f32) -> (f32, f32) {
  84 // CHECK-DAG:      %[[E0:.+]] = vector.extract %[[ARG0]][0] : f32 from vector<2xf32>
  85 // CHECK-DAG:      %[[E1:.+]] = vector.extract %[[ARG0]][1] : f32 from vector<2xf32>
  86 // CHECK:          %[[A0:.+]] = arith.addf %[[E0]], %[[E1]] : f32
  87 // CHECK:          %[[R0:.+]] = arith.addf %[[A0]], %[[ARG1]] : f32
  88 // CHECK:          %[[M0:.+]] = arith.mulf %[[E0]], %[[E1]] fastmath<nnan> : f32
  89 // CHECK:          %[[R1:.+]] = arith.mulf %[[M0]], %[[ARG1]] fastmath<nnan> : f32
  90 // CHECK-NEXT:     return %[[R0]], %[[R1]] : f32, f32
  91 func.func @reduce_2x_acc_f32(%arg0: vector<2xf32>, %arg1: f32) -> (f32, f32) {
  92   %0 = vector.reduction <add>, %arg0, %arg1 : vector<2xf32> into f32
  93   %1 = vector.reduction <mul>, %arg0, %arg1 fastmath<nnan> : vector<2xf32> into f32
  94   return %0, %1 : f32, f32
  95 }
  96
  97 // CHECK-LABEL:   func.func @reduce_3x_f32(
  98 // CHECK-SAME:     %[[ARG0:.+]]: vector<3xf32>) -> f32 {
  99 // CHECK-NEXT:     %[[R0:.+]] = vector.reduction <add>, %[[ARG0]] : vector<3xf32> into f32
 100 // CHECK-NEXT:     return %[[R0]] : f32
 101 func.func @reduce_3x_f32(%arg0: vector<3xf32>) -> f32 {
 102   %0 = vector.reduction <add>, %arg0 : vector<3xf32> into f32
 103   return %0 : f32
 104 }
 105
 106 // Masking is not handled yet.
 107 // CHECK-LABEL:   func.func @reduce_mask_3x_f32
 108 // CHECK-NEXT:     %[[M:.+]] = vector.create_mask
 109 // CHECK-NEXT:     %[[R:.+]] = vector.mask %[[M]]
 110 // CHECK-SAME:       vector.reduction <add>
 111 // CHECK-NEXT:     return %[[R]] : f32
 112 func.func @reduce_mask_3x_f32(%arg0: vector<3xf32>, %arg1: index) -> f32 {
 113   %mask = vector.create_mask %arg1 : vector<3xi1>
 114   %0 = vector.mask %mask { vector.reduction <add>, %arg0 : vector<3xf32> into f32 } : vector<3xi1> -> f32
 115   return %0 : f32
 116 }
 117
 118 // Scalable vectors are not supported.
 119 // CHECK-LABEL:   func.func @reduce_scalable_f32(
 120 // CHECK-SAME:     %[[ARG0:.+]]: vector<[1]xf32>) -> f32 {
 121 // CHECK-NEXT:     %[[R0:.+]] = vector.reduction <add>, %[[ARG0]] : vector<[1]xf32> into f32
 122 // CHECK-NEXT:     return %[[R0]] : f32
 123 func.func @reduce_scalable_f32(%arg0: vector<[1]xf32>) -> f32 {
 124   %0 = vector.reduction <add>, %arg0 : vector<[1]xf32> into f32
 125   return %0 : f32
 126 }