127 lines · plain
1// RUN: mlir-opt %s --test-vector-break-down-reduction-patterns --cse | FileCheck %s2 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]] : f3210// CHECK-DAG: %[[R1:.+]] = arith.mulf %[[E0]], %[[E1]] : f3211// CHECK-DAG: %[[R2:.+]] = arith.minnumf %[[E0]], %[[E1]] : f3212// CHECK-DAG: %[[R3:.+]] = arith.maxnumf %[[E0]], %[[E1]] : f3213// CHECK-DAG: %[[R4:.+]] = arith.minimumf %[[E0]], %[[E1]] : f3214// CHECK-DAG: %[[R5:.+]] = arith.maximumf %[[E0]], %[[E1]] : f3215// CHECK: return %[[R0]], %[[R1]], %[[R2]], %[[R3]], %[[R4]], %[[R5]]16func.func @reduce_2x_f32(%arg0: vector<2xf32>) -> (f32, f32, f32, f32, f32, f32) {17 %0 = vector.reduction <add>, %arg0 : vector<2xf32> into f3218 %1 = vector.reduction <mul>, %arg0 : vector<2xf32> into f3219 %2 = vector.reduction <minnumf>, %arg0 : vector<2xf32> into f3220 %3 = vector.reduction <maxnumf>, %arg0 : vector<2xf32> into f3221 %4 = vector.reduction <minimumf>, %arg0 : vector<2xf32> into f3222 %5 = vector.reduction <maximumf>, %arg0 : vector<2xf32> into f3223 return %0, %1, %2, %3, %4, %5 : f32, f32, f32, f32, f32, f3224}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]] : i3231// CHECK-DAG: %[[R1:.+]] = arith.muli %[[E0]], %[[E1]] : i3232// CHECK-DAG: %[[R2:.+]] = arith.minsi %[[E0]], %[[E1]] : i3233// CHECK-DAG: %[[R3:.+]] = arith.maxsi %[[E0]], %[[E1]] : i3234// CHECK-DAG: %[[R4:.+]] = arith.minui %[[E0]], %[[E1]] : i3235// CHECK-DAG: %[[R5:.+]] = arith.maxui %[[E0]], %[[E1]] : i3236// CHECK-DAG: %[[R6:.+]] = arith.andi %[[E0]], %[[E1]] : i3237// CHECK-DAG: %[[R7:.+]] = arith.ori %[[E0]], %[[E1]] : i3238// CHECK-DAG: %[[R8:.+]] = arith.xori %[[E0]], %[[E1]] : i3239// CHECK: return %[[R0]], %[[R1]], %[[R2]], %[[R3]], %[[R4]], %[[R5]], %[[R6]], %[[R7]], %[[R8]]40func.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 i3242 %1 = vector.reduction <mul>, %arg0 : vector<2xi32> into i3243 %2 = vector.reduction <minsi>, %arg0 : vector<2xi32> into i3244 %3 = vector.reduction <maxsi>, %arg0 : vector<2xi32> into i3245 %4 = vector.reduction <minui>, %arg0 : vector<2xi32> into i3246 %5 = vector.reduction <maxui>, %arg0 : vector<2xi32> into i3247 %6 = vector.reduction <and>, %arg0 : vector<2xi32> into i3248 %7 = vector.reduction <or>, %arg0 : vector<2xi32> into i3249 %8 = vector.reduction <xor>, %arg0 : vector<2xi32> into i3250 return %0, %1, %2, %3, %4, %5, %6, %7, %8 : i32, i32, i32, i32, i32, i32, i32, i32, i3251}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]] : f3257func.func @reduce_1x_f32(%arg0: vector<1xf32>) -> f32 {58 %0 = vector.reduction <add>, %arg0 : vector<1xf32> into f3259 return %0 : f3260}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]] : f3266// CHECK-NEXT: return %[[R0]] : f3267func.func @reduce_1x_acc_f32(%arg0: vector<1xf32>, %arg1: f32) -> f32 {68 %0 = vector.reduction <add>, %arg0, %arg1 : vector<1xf32> into f3269 return %0 : f3270}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]] : i3276// CHECK-NEXT: return %[[R0]] : i3277func.func @reduce_1x_acc_i32(%arg0: vector<1xi32>, %arg1: i32) -> i32 {78 %0 = vector.reduction <add>, %arg0, %arg1 : vector<1xi32> into i3279 return %0 : i3280}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]] : f3287// CHECK: %[[R0:.+]] = arith.addf %[[A0]], %[[ARG1]] : f3288// CHECK: %[[M0:.+]] = arith.mulf %[[E0]], %[[E1]] fastmath<nnan> : f3289// CHECK: %[[R1:.+]] = arith.mulf %[[M0]], %[[ARG1]] fastmath<nnan> : f3290// CHECK-NEXT: return %[[R0]], %[[R1]] : f32, f3291func.func @reduce_2x_acc_f32(%arg0: vector<2xf32>, %arg1: f32) -> (f32, f32) {92 %0 = vector.reduction <add>, %arg0, %arg1 : vector<2xf32> into f3293 %1 = vector.reduction <mul>, %arg0, %arg1 fastmath<nnan> : vector<2xf32> into f3294 return %0, %1 : f32, f3295}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 f32100// CHECK-NEXT: return %[[R0]] : f32101func.func @reduce_3x_f32(%arg0: vector<3xf32>) -> f32 {102 %0 = vector.reduction <add>, %arg0 : vector<3xf32> into f32103 return %0 : f32104}105 106// Masking is not handled yet.107// CHECK-LABEL: func.func @reduce_mask_3x_f32108// CHECK-NEXT: %[[M:.+]] = vector.create_mask109// CHECK-NEXT: %[[R:.+]] = vector.mask %[[M]]110// CHECK-SAME: vector.reduction <add>111// CHECK-NEXT: return %[[R]] : f32112func.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> -> f32115 return %0 : f32116}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 f32122// CHECK-NEXT: return %[[R0]] : f32123func.func @reduce_scalable_f32(%arg0: vector<[1]xf32>) -> f32 {124 %0 = vector.reduction <add>, %arg0 : vector<[1]xf32> into f32125 return %0 : f32126}127