143 lines · plain
1// RUN: mlir-opt -test-spirv-vector-unrolling -split-input-file %s | FileCheck %s2 3// CHECK-LABEL: @vaddi4// CHECK-SAME: (%[[ARG0:.+]]: vector<3xi32>, %[[ARG1:.+]]: vector<3xi32>, %[[ARG2:.+]]: vector<3xi32>, %[[ARG3:.+]]: vector<3xi32>)5func.func @vaddi(%arg0 : vector<6xi32>, %arg1 : vector<6xi32>) -> (vector<6xi32>) {6 // CHECK: %[[ADD0:.*]] = arith.addi %[[ARG0]], %[[ARG2]] : vector<3xi32>7 // CHECK: %[[ADD1:.*]] = arith.addi %[[ARG1]], %[[ARG3]] : vector<3xi32>8 // CHECK: return %[[ADD0]], %[[ADD1]] : vector<3xi32>, vector<3xi32>9 %0 = arith.addi %arg0, %arg1 : vector<6xi32>10 return %0 : vector<6xi32>11}12 13// CHECK-LABEL: @vaddi_2d14// CHECK-SAME: (%[[ARG0:.+]]: vector<2xi32>, %[[ARG1:.+]]: vector<2xi32>, %[[ARG2:.+]]: vector<2xi32>, %[[ARG3:.+]]: vector<2xi32>)15func.func @vaddi_2d(%arg0 : vector<2x2xi32>, %arg1 : vector<2x2xi32>) -> (vector<2x2xi32>) {16 // CHECK: %[[ADD0:.*]] = arith.addi %[[ARG0]], %[[ARG2]] : vector<2xi32>17 // CHECK: %[[ADD1:.*]] = arith.addi %[[ARG1]], %[[ARG3]] : vector<2xi32>18 // CHECK: return %[[ADD0]], %[[ADD1]] : vector<2xi32>, vector<2xi32>19 %0 = arith.addi %arg0, %arg1 : vector<2x2xi32>20 return %0 : vector<2x2xi32>21}22 23// CHECK-LABEL: @vaddi_2d_824// CHECK-SAME: (%[[ARG0:.+]]: vector<4xi32>, %[[ARG1:.+]]: vector<4xi32>, %[[ARG2:.+]]: vector<4xi32>, %[[ARG3:.+]]: vector<4xi32>, %[[ARG4:.+]]: vector<4xi32>, %[[ARG5:.+]]: vector<4xi32>, %[[ARG6:.+]]: vector<4xi32>, %[[ARG7:.+]]: vector<4xi32>)25func.func @vaddi_2d_8(%arg0 : vector<2x8xi32>, %arg1 : vector<2x8xi32>) -> (vector<2x8xi32>) {26 // CHECK: %[[ADD0:.*]] = arith.addi %[[ARG0]], %[[ARG4]] : vector<4xi32>27 // CHECK: %[[ADD1:.*]] = arith.addi %[[ARG1]], %[[ARG5]] : vector<4xi32>28 // CHECK: %[[ADD2:.*]] = arith.addi %[[ARG2]], %[[ARG6]] : vector<4xi32>29 // CHECK: %[[ADD3:.*]] = arith.addi %[[ARG3]], %[[ARG7]] : vector<4xi32>30 // CHECK: return %[[ADD0]], %[[ADD1]], %[[ADD2]], %[[ADD3]] : vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32>31 %0 = arith.addi %arg0, %arg1 : vector<2x8xi32>32 return %0 : vector<2x8xi32>33}34 35// -----36 37// CHECK-LABEL: @reduction_538// CHECK-SAME: (%[[ARG0:.+]]: vector<1xi32>, %[[ARG1:.+]]: vector<1xi32>, %[[ARG2:.+]]: vector<1xi32>, %[[ARG3:.+]]: vector<1xi32>, %[[ARG4:.+]]: vector<1xi32>)39func.func @reduction_5(%arg0 : vector<5xi32>) -> (i32) {40 // CHECK: %[[EXTRACT0:.*]] = vector.extract %[[ARG0]][0] : i32 from vector<1xi32>41 // CHECK: %[[EXTRACT1:.*]] = vector.extract %[[ARG1]][0] : i32 from vector<1xi32>42 // CHECK: %[[ADD0:.*]] = arith.addi %[[EXTRACT0]], %[[EXTRACT1]] : i3243 // CHECK: %[[EXTRACT2:.*]] = vector.extract %[[ARG2]][0] : i32 from vector<1xi32>44 // CHECK: %[[ADD1:.*]] = arith.addi %[[ADD0]], %[[EXTRACT2]] : i3245 // CHECK: %[[EXTRACT3:.*]] = vector.extract %[[ARG3]][0] : i32 from vector<1xi32>46 // CHECK: %[[ADD2:.*]] = arith.addi %[[ADD1]], %[[EXTRACT3]] : i3247 // CHECK: %[[EXTRACT4:.*]] = vector.extract %[[ARG4]][0] : i32 from vector<1xi32>48 // CHECK: %[[ADD3:.*]] = arith.addi %[[ADD2]], %[[EXTRACT4]] : i3249 // CHECK: return %[[ADD3]] : i3250 %0 = vector.reduction <add>, %arg0 : vector<5xi32> into i3251 return %0 : i3252}53 54// CHECK-LABEL: @reduction_855// CHECK-SAME: (%[[ARG0:.+]]: vector<4xi32>, %[[ARG1:.+]]: vector<4xi32>)56func.func @reduction_8(%arg0 : vector<8xi32>) -> (i32) {57 // CHECK: %[[REDUCTION0:.*]] = vector.reduction <add>, %[[ARG0]] : vector<4xi32> into i3258 // CHECK: %[[REDUCTION1:.*]] = vector.reduction <add>, %[[ARG1]] : vector<4xi32> into i3259 // CHECK: %[[ADD:.*]] = arith.addi %[[REDUCTION0]], %[[REDUCTION1]] : i3260 // CHECK: return %[[ADD]] : i3261 %0 = vector.reduction <add>, %arg0 : vector<8xi32> into i3262 return %0 : i3263}64 65// -----66 67// CHECK-LABEL: @vaddi_reduction68// CHECK-SAME: (%[[ARG0:.+]]: vector<4xi32>, %[[ARG1:.+]]: vector<4xi32>, %[[ARG2:.+]]: vector<4xi32>, %[[ARG3:.+]]: vector<4xi32>)69func.func @vaddi_reduction(%arg0 : vector<8xi32>, %arg1 : vector<8xi32>) -> (i32) {70 // CHECK: %[[ADD0:.*]] = arith.addi %[[ARG0]], %[[ARG2]] : vector<4xi32>71 // CHECK: %[[ADD1:.*]] = arith.addi %[[ARG1]], %[[ARG3]] : vector<4xi32>72 // CHECK: %[[REDUCTION0:.*]] = vector.reduction <add>, %[[ADD0]] : vector<4xi32> into i3273 // CHECK: %[[REDUCTION1:.*]] = vector.reduction <add>, %[[ADD1]] : vector<4xi32> into i3274 // CHECK: %[[ADD2:.*]] = arith.addi %[[REDUCTION0]], %[[REDUCTION1]] : i3275 // CHECK: return %[[ADD2]] : i3276 %0 = arith.addi %arg0, %arg1 : vector<8xi32>77 %1 = vector.reduction <add>, %0 : vector<8xi32> into i3278 return %1 : i3279}80 81// -----82 83// CHECK-LABEL: @transpose84// CHECK-SAME: (%[[ARG0:.+]]: vector<3xi32>, %[[ARG1:.+]]: vector<3xi32>)85func.func @transpose(%arg0 : vector<2x3xi32>) -> (vector<3x2xi32>) {86 // CHECK: %[[EXTRACT0:.*]] = vector.extract %[[ARG0]][0] : i32 from vector<3xi32>87 // CHECK: %[[EXTRACT1:.*]] = vector.extract %[[ARG1]][0] : i32 from vector<3xi32>88 // CHECK: %[[FROM_ELEMENTS0:.*]] = vector.from_elements %[[EXTRACT0]], %[[EXTRACT1]] : vector<2xi32>89 // CHECK: %[[EXTRACT2:.*]] = vector.extract %[[ARG0]][1] : i32 from vector<3xi32>90 // CHECK: %[[EXTRACT3:.*]] = vector.extract %[[ARG1]][1] : i32 from vector<3xi32>91 // CHECK: %[[FROM_ELEMENTS1:.*]] = vector.from_elements %[[EXTRACT2]], %[[EXTRACT3]] : vector<2xi32>92 // CHECK: %[[EXTRACT4:.*]] = vector.extract %[[ARG0]][2] : i32 from vector<3xi32>93 // CHECK: %[[EXTRACT5:.*]] = vector.extract %[[ARG1]][2] : i32 from vector<3xi32>94 // CHECK: %[[FROM_ELEMENTS2:.*]] = vector.from_elements %[[EXTRACT4]], %[[EXTRACT5]] : vector<2xi32>95 // CHECK: return %[[FROM_ELEMENTS0]], %[[FROM_ELEMENTS1]], %[[FROM_ELEMENTS2]] : vector<2xi32>, vector<2xi32>, vector<2xi32>96 %0 = vector.transpose %arg0, [1, 0] : vector<2x3xi32> to vector<3x2xi32>97 return %0 : vector<3x2xi32>98}99 100// -----101 102// In order to verify that the pattern is applied,103// we need to make sure that the the 2d vector does not104// come from the parameters. Otherwise, the pattern105// in unrollVectorsInSignatures which splits the 2d vector106// parameter will take precedent. Similarly, let's avoid107// returning a vector as another pattern would take precendence.108 109// CHECK-LABEL: @unroll_to_elements_2d110func.func @unroll_to_elements_2d() -> (f32, f32, f32, f32) {111 %1 = "test.op"() : () -> (vector<2x2xf32>)112 // CHECK: %[[VEC2D:.+]] = "test.op"113 // CHECK: %[[VEC0:.+]] = vector.extract %[[VEC2D]][0] : vector<2xf32> from vector<2x2xf32>114 // CHECK: %[[VEC1:.+]] = vector.extract %[[VEC2D]][1] : vector<2xf32> from vector<2x2xf32>115 // CHECK: %[[RES0:.+]]:2 = vector.to_elements %[[VEC0]]116 // CHECK: %[[RES1:.+]]:2 = vector.to_elements %[[VEC1]]117 %2:4 = vector.to_elements %1 : vector<2x2xf32>118 return %2#0, %2#1, %2#2, %2#3 : f32, f32, f32, f32119}120 121// -----122 123// In order to verify that the pattern is applied,124// we need to make sure that the the 2d vector is used125// by an operation and that extracts are not folded away.126// In other words we can't use "test.op" nor return the127// value `%0 = vector.from_elements`128 129// CHECK-LABEL: @unroll_from_elements_2d130// CHECK-SAME: (%[[ARG0:.+]]: f32, %[[ARG1:.+]]: f32, %[[ARG2:.+]]: f32, %[[ARG3:.+]]: f32)131func.func @unroll_from_elements_2d(%arg0: f32, %arg1: f32, %arg2: f32, %arg3: f32) -> (vector<2x2xf32>) {132 // CHECK: %[[VEC0:.+]] = vector.from_elements %[[ARG0]], %[[ARG1]] : vector<2xf32>133 // CHECK: %[[VEC1:.+]] = vector.from_elements %[[ARG2]], %[[ARG3]] : vector<2xf32>134 %0 = vector.from_elements %arg0, %arg1, %arg2, %arg3 : vector<2x2xf32>135 136 // CHECK: %[[RES0:.+]] = arith.addf %[[VEC0]], %[[VEC0]]137 // CHECK: %[[RES1:.+]] = arith.addf %[[VEC1]], %[[VEC1]]138 %1 = arith.addf %0, %0 : vector<2x2xf32>139 140 // return %[[RES0]], %%[[RES1]] : vector<2xf32>, vector<2xf32>141 return %1 : vector<2x2xf32>142}143