117 lines · plain
1// RUN: mlir-opt %s --test-vector-chained-reduction-folding-patterns | FileCheck %s2 3// CHECK-LABEL: func.func @reduce_1x_fp32(4// CHECK-SAME: %[[ARG0:.+]]: vector<8xf32>) -> f32 {5// CHECK-DAG: %[[CST:.+]] = arith.constant 0.06// CHECK-NEXT: %[[RES:.+]] = vector.reduction <add>, %[[ARG0]], %[[CST]] : vector<8xf32> into f327// CHECK-NEXT: return %[[RES]] : f328func.func @reduce_1x_fp32(%arg0: vector<8xf32>) -> f32 {9 %cst0 = arith.constant 0.0 : f3210 %0 = vector.reduction <add>, %arg0, %cst0 : vector<8xf32> into f3211 return %0 : f3212}13 14// CHECK-LABEL: func.func @reduce_2x_fp32(15// CHECK-SAME: %[[ARG0:.+]]: vector<8xf32>, %[[ARG1:.+]]: vector<8xf32>) -> f32 {16// CHECK-DAG: %[[CST:.+]] = arith.constant 0.017// CHECK-DAG: %[[ADD:.+]] = arith.addf %[[ARG0]], %[[ARG1]] : vector<8xf32>18// CHECK-NEXT: %[[RES:.+]] = vector.reduction <add>, %[[ADD]], %[[CST]] : vector<8xf32> into f3219// CHECK-NEXT: return %[[RES]] : f3220func.func @reduce_2x_fp32(%arg0: vector<8xf32>, %arg1: vector<8xf32>) -> f32 {21 %cst0 = arith.constant 0.0 : f3222 %0 = vector.reduction <add>, %arg0, %cst0 : vector<8xf32> into f3223 %1 = vector.reduction <add>, %arg1, %0 : vector<8xf32> into f3224 return %1 : f3225}26 27// CHECK-LABEL: func.func @reduce_2x_no_acc_fp32(28// CHECK-SAME: %[[ARG0:.+]]: vector<8xf32>, %[[ARG1:.+]]: vector<8xf32>) -> f32 {29// CHECK: %[[ADD:.+]] = arith.addf %[[ARG0]], %[[ARG1]] : vector<8xf32>30// CHECK-NEXT: %[[RES:.+]] = vector.reduction <add>, %[[ADD]] : vector<8xf32> into f3231// CHECK-NEXT: return %[[RES]] : f3232func.func @reduce_2x_no_acc_fp32(%arg0: vector<8xf32>, %arg1: vector<8xf32>) -> f32 {33 %0 = vector.reduction <add>, %arg0 : vector<8xf32> into f3234 %1 = vector.reduction <add>, %arg1, %0 : vector<8xf32> into f3235 return %1 : f3236}37 38// CHECK-LABEL: func.func @reduce_2x_zero_add_fp32(39// CHECK-SAME: %[[ARG0:.+]]: vector<8xf32>, %[[ARG1:.+]]: vector<8xf32>) -> f32 {40// CHECK: %[[ADD:.+]] = arith.addf %[[ARG0]], %[[ARG1]] : vector<8xf32>41// CHECK-NEXT: %[[RES:.+]] = vector.reduction <add>, %[[ADD]] : vector<8xf32> into f3242// CHECK-NEXT: return %[[RES]] : f3243func.func @reduce_2x_zero_add_fp32(%arg0: vector<8xf32>, %arg1: vector<8xf32>) -> f32 {44 %cst0 = arith.constant dense<0.0> : vector<8xf32>45 %x = arith.addf %arg0, %cst0 : vector<8xf32>46 %0 = vector.reduction <add>, %x : vector<8xf32> into f3247 %1 = vector.reduction <add>, %arg1, %0 : vector<8xf32> into f3248 return %1 : f3249}50 51// CHECK-LABEL: func.func @reduce_3x_fp32(52// CHECK-SAME: %[[ARG0:.+]]: vector<8xf32>, %[[ARG1:.+]]: vector<8xf32>,53// CHECK-SAME: %[[ARG2:.+]]: vector<8xf32>) -> f32 {54// CHECK-DAG: %[[CST:.+]] = arith.constant 0.055// CHECK-DAG: %[[ADD0:.+]] = arith.addf %[[ARG1]], %[[ARG2]] : vector<8xf32>56// CHECK-DAG: %[[ADD1:.+]] = arith.addf %[[ARG0]], %[[ADD0]] : vector<8xf32>57// CHECK-NEXT: %[[RES:.+]] = vector.reduction <add>, %[[ADD1]], %[[CST]] : vector<8xf32> into f3258// CHECK-NEXT: return %[[RES]] : f3259func.func @reduce_3x_fp32(%arg0: vector<8xf32>, %arg1: vector<8xf32>,60 %arg2: vector<8xf32>) -> f32 {61 %cst0 = arith.constant 0.0 : f3262 %0 = vector.reduction <add>, %arg0, %cst0 : vector<8xf32> into f3263 %1 = vector.reduction <add>, %arg1, %0 : vector<8xf32> into f3264 %2 = vector.reduction <add>, %arg2, %1 : vector<8xf32> into f3265 return %2 : f3266}67 68// CHECK-LABEL: func.func @reduce_1x_i32(69// CHECK-SAME: %[[ARG0:.+]]: vector<8xi32>) -> i32 {70// CHECK-DAG: %[[CST:.+]] = arith.constant 071// CHECK-NEXT: %[[RES:.+]] = vector.reduction <add>, %[[ARG0]], %[[CST]] : vector<8xi32> into i3272// CHECK-NEXT: return %[[RES]] : i3273func.func @reduce_1x_i32(%arg0: vector<8xi32>) -> i32 {74 %cst0 = arith.constant 0 : i3275 %0 = vector.reduction <add>, %arg0, %cst0 : vector<8xi32> into i3276 return %0 : i3277}78 79// CHECK-LABEL: func.func @reduce_2x_i32(80// CHECK-SAME: %[[ARG0:.+]]: vector<8xi32>, %[[ARG1:.+]]: vector<8xi32>) -> i32 {81// CHECK-DAG: %[[CST:.+]] = arith.constant 082// CHECK-DAG: %[[ADD:.+]] = arith.addi %[[ARG0]], %[[ARG1]] : vector<8xi32>83// CHECK-NEXT: %[[RES:.+]] = vector.reduction <add>, %[[ADD]], %[[CST]] : vector<8xi32> into i3284// CHECK-NEXT: return %[[RES]] : i3285func.func @reduce_2x_i32(%arg0: vector<8xi32>, %arg1: vector<8xi32>) -> i32 {86 %cst0 = arith.constant 0 : i3287 %0 = vector.reduction <add>, %arg0, %cst0 : vector<8xi32> into i3288 %1 = vector.reduction <add>, %arg1, %0 : vector<8xi32> into i3289 return %1 : i3290}91 92// CHECK-LABEL: func.func @reduce_2x_no_acc_i32(93// CHECK-SAME: %[[ARG0:.+]]: vector<8xi32>, %[[ARG1:.+]]: vector<8xi32>) -> i32 {94// CHECK: %[[ADD:.+]] = arith.addi %[[ARG0]], %[[ARG1]] : vector<8xi32>95// CHECK-NEXT: %[[RES:.+]] = vector.reduction <add>, %[[ADD]] : vector<8xi32> into i3296// CHECK-NEXT: return %[[RES]] : i3297func.func @reduce_2x_no_acc_i32(%arg0: vector<8xi32>, %arg1: vector<8xi32>) -> i32 {98 %0 = vector.reduction <add>, %arg0 : vector<8xi32> into i3299 %1 = vector.reduction <add>, %arg1, %0 : vector<8xi32> into i32100 return %1 : i32101}102 103// CHECK-LABEL: func.func @reduce_2x_zero_add_i32(104// CHECK-SAME: %[[ARG0:.+]]: vector<8xi32>, %[[ARG1:.+]]: vector<8xi32>) -> i32 {105// CHECK-DAG: %[[CST:.+]] = arith.constant 0106// CHECK-DAG: %[[ADD:.+]] = arith.addi %[[ARG0]], %[[ARG1]] : vector<8xi32>107// CHECK-NEXT: %[[RES:.+]] = vector.reduction <add>, %[[ADD]], %[[CST]] : vector<8xi32> into i32108// CHECK-NEXT: return %[[RES]] : i32109func.func @reduce_2x_zero_add_i32(%arg0: vector<8xi32>, %arg1: vector<8xi32>) -> i32 {110 %cst0 = arith.constant 0 : i32111 %cstV = arith.constant dense<0> : vector<8xi32>112 %x = arith.addi %arg0, %cstV : vector<8xi32>113 %0 = vector.reduction <add>, %x, %cst0 : vector<8xi32> into i32114 %1 = vector.reduction <add>, %arg1, %0 : vector<8xi32> into i32115 return %1 : i32116}117