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]] 16func.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]] 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 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 57func.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 67func.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 77func.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 91func.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 101func.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 112func.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 123func.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} 127