193 lines · plain
1// RUN: mlir-opt --split-input-file --verify-diagnostics \2// RUN: --test-vector-reduction-to-spirv-dot-prod %s -o - | FileCheck %s3 4// Positive tests.5 6// CHECK-LABEL: func.func @to_sdot7// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>)8// CHECK-NEXT: [[DOT:%.+]] = spirv.SDot [[ARG0]], [[ARG1]] : vector<4xi8> -> i329// CHECK-NEXT: return [[DOT]] : i3210func.func @to_sdot(%arg0: vector<4xi8>, %arg1: vector<4xi8>) -> i32 {11 %lhs = arith.extsi %arg0 : vector<4xi8> to vector<4xi32>12 %rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi32>13 %mul = arith.muli %lhs, %rhs : vector<4xi32>14 %red = vector.reduction <add>, %mul : vector<4xi32> into i3215 return %red : i3216}17 18// CHECK-LABEL: func.func @to_sdot_acc19// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>, [[ACC:%.+]]: i32)20// CHECK-NEXT: [[DOT:%.+]] = spirv.SDotAccSat [[ARG0]], [[ARG1]], [[ACC]] : vector<4xi8> -> i3221// CHECK-NEXT: return [[DOT]] : i3222func.func @to_sdot_acc(%arg0: vector<4xi8>, %arg1: vector<4xi8>, %acc: i32) -> i32 {23 %lhs = arith.extsi %arg0 : vector<4xi8> to vector<4xi32>24 %rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi32>25 %mul = arith.muli %lhs, %rhs : vector<4xi32>26 %red = vector.reduction <add>, %mul, %acc : vector<4xi32> into i3227 return %red : i3228}29 30// CHECK-LABEL: func.func @to_sdot_i6431// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>)32// CHECK-NEXT: [[DOT:%.+]] = spirv.SDot [[ARG0]], [[ARG1]] : vector<4xi8> -> i6433// CHECK-NEXT: return [[DOT]] : i6434func.func @to_sdot_i64(%arg0: vector<4xi8>, %arg1: vector<4xi8>) -> i64 {35 %lhs = arith.extsi %arg0 : vector<4xi8> to vector<4xi64>36 %rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi64>37 %mul = arith.muli %lhs, %rhs : vector<4xi64>38 %red = vector.reduction <add>, %mul : vector<4xi64> into i6439 return %red : i6440}41 42// CHECK-LABEL: func.func @to_sdot_acc_i6443// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>, [[ACC:%.+]]: i64)44// CHECK-NEXT: [[DOT:%.+]] = spirv.SDotAccSat [[ARG0]], [[ARG1]], [[ACC]] : vector<4xi8> -> i6445// CHECK-NEXT: return [[DOT]] : i6446func.func @to_sdot_acc_i64(%arg0: vector<4xi8>, %arg1: vector<4xi8>, %acc: i64) -> i64 {47 %lhs = arith.extsi %arg0 : vector<4xi8> to vector<4xi64>48 %rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi64>49 %mul = arith.muli %lhs, %rhs : vector<4xi64>50 %red = vector.reduction <add>, %mul, %acc : vector<4xi64> into i6451 return %red : i6452}53 54// CHECK-LABEL: func.func @to_udot55// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>)56// CHECK-NEXT: [[DOT:%.+]] = spirv.UDot [[ARG0]], [[ARG1]] : vector<4xi8> -> i3257// CHECK-NEXT: return [[DOT]] : i3258func.func @to_udot(%arg0: vector<4xi8>, %arg1: vector<4xi8>) -> i32 {59 %lhs = arith.extui %arg0 : vector<4xi8> to vector<4xi32>60 %rhs = arith.extui %arg1 : vector<4xi8> to vector<4xi32>61 %mul = arith.muli %lhs, %rhs : vector<4xi32>62 %red = vector.reduction <add>, %mul : vector<4xi32> into i3263 return %red : i3264}65 66// CHECK-LABEL: func.func @to_udot_acc67// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>, [[ACC:%.+]]: i32)68// CHECK-NEXT: [[DOT:%.+]] = spirv.UDotAccSat [[ARG0]], [[ARG1]], [[ACC]] : vector<4xi8> -> i3269// CHECK-NEXT: return [[DOT]] : i3270func.func @to_udot_acc(%arg0: vector<4xi8>, %arg1: vector<4xi8>, %acc: i32) -> i32 {71 %lhs = arith.extui %arg0 : vector<4xi8> to vector<4xi32>72 %rhs = arith.extui %arg1 : vector<4xi8> to vector<4xi32>73 %mul = arith.muli %lhs, %rhs : vector<4xi32>74 %red = vector.reduction <add>, %mul, %acc : vector<4xi32> into i3275 return %red : i3276}77 78// CHECK-LABEL: func.func @to_signed_unsigned_dot79// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>)80// CHECK-NEXT: [[DOT:%.+]] = spirv.SUDot [[ARG0]], [[ARG1]] : vector<4xi8> -> i3281// CHECK-NEXT: return [[DOT]] : i3282func.func @to_signed_unsigned_dot(%arg0: vector<4xi8>, %arg1: vector<4xi8>) -> i32 {83 %lhs = arith.extsi %arg0 : vector<4xi8> to vector<4xi32>84 %rhs = arith.extui %arg1 : vector<4xi8> to vector<4xi32>85 %mul = arith.muli %lhs, %rhs : vector<4xi32>86 %red = vector.reduction <add>, %mul : vector<4xi32> into i3287 return %red : i3288}89 90// CHECK-LABEL: func.func @to_signed_unsigned_dot_acc91// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>, [[ACC:%.+]]: i32)92// CHECK-NEXT: [[DOT:%.+]] = spirv.SUDotAccSat [[ARG0]], [[ARG1]], [[ACC]] : vector<4xi8> -> i3293// CHECK-NEXT: return [[DOT]] : i3294func.func @to_signed_unsigned_dot_acc(%arg0: vector<4xi8>, %arg1: vector<4xi8>, %acc: i32) -> i32 {95 %lhs = arith.extsi %arg0 : vector<4xi8> to vector<4xi32>96 %rhs = arith.extui %arg1 : vector<4xi8> to vector<4xi32>97 %mul = arith.muli %lhs, %rhs : vector<4xi32>98 %red = vector.reduction <add>, %mul, %acc : vector<4xi32> into i3299 return %red : i32100}101 102// CHECK-LABEL: func.func @to_unsigned_signed_dot103// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>)104// CHECK-NEXT: [[DOT:%.+]] = spirv.SUDot [[ARG1]], [[ARG0]] : vector<4xi8> -> i32105// CHECK-NEXT: return [[DOT]] : i32106func.func @to_unsigned_signed_dot(%arg0: vector<4xi8>, %arg1: vector<4xi8>) -> i32 {107 %lhs = arith.extui %arg0 : vector<4xi8> to vector<4xi32>108 %rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi32>109 %mul = arith.muli %lhs, %rhs : vector<4xi32>110 %red = vector.reduction <add>, %mul : vector<4xi32> into i32111 return %red : i32112}113 114// CHECK-LABEL: func.func @to_unsigned_signed_dot_acc115// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>, [[ACC:%.+]]: i32)116// CHECK-NEXT: [[DOT:%.+]] = spirv.SUDotAccSat [[ARG1]], [[ARG0]], [[ACC]] : vector<4xi8> -> i32117// CHECK-NEXT: return [[DOT]] : i32118func.func @to_unsigned_signed_dot_acc(%arg0: vector<4xi8>, %arg1: vector<4xi8>, %acc: i32) -> i32 {119 %lhs = arith.extui %arg0 : vector<4xi8> to vector<4xi32>120 %rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi32>121 %mul = arith.muli %lhs, %rhs : vector<4xi32>122 %red = vector.reduction <add>, %mul, %acc : vector<4xi32> into i32123 return %red : i32124}125 126// CHECK-LABEL: func.func @to_sdot_vector3127// CHECK-SAME: (%[[ARG0:.+]]: vector<3xi8>, %[[ARG1:.+]]: vector<3xi8>)128// CHECK: %[[ZERO:.+]] = spirv.Constant 0 : i8129// CHECK: %[[LHS:.+]] = spirv.CompositeConstruct %[[ARG0]], %[[ZERO]] : (vector<3xi8>, i8) -> vector<4xi8>130// CHECK: %[[RHS:.+]] = spirv.CompositeConstruct %[[ARG1]], %[[ZERO]] : (vector<3xi8>, i8) -> vector<4xi8>131// CHECK: %[[SDOT:.+]] = spirv.SDot %[[LHS]], %[[RHS]] : vector<4xi8> -> i32132// CHECK: return %[[SDOT]]133func.func @to_sdot_vector3(%arg0: vector<3xi8>, %arg1: vector<3xi8>) -> i32 {134 %lhs = arith.extsi %arg0 : vector<3xi8> to vector<3xi32>135 %rhs = arith.extsi %arg1 : vector<3xi8> to vector<3xi32>136 %mul = arith.muli %lhs, %rhs : vector<3xi32>137 %red = vector.reduction <add>, %mul : vector<3xi32> into i32138 return %red : i32139}140 141// -----142 143// Negative tests.144 145// CHECK-LABEL: func.func @too_short146// CHECK-SAME: ([[ARG0:%.+]]: vector<2xi8>, [[ARG1:%.+]]: vector<2xi8>)147// CHECK: [[RED:%.+]] = vector.reduction148// CHECK-NEXT: return [[RED]] : i32149func.func @too_short(%arg0: vector<2xi8>, %arg1: vector<2xi8>) -> i32 {150 %lhs = arith.extsi %arg0 : vector<2xi8> to vector<2xi32>151 %rhs = arith.extsi %arg1 : vector<2xi8> to vector<2xi32>152 %mul = arith.muli %lhs, %rhs : vector<2xi32>153 %red = vector.reduction <add>, %mul : vector<2xi32> into i32154 return %red : i32155}156 157// CHECK-LABEL: func.func @too_long158// CHECK-SAME: ([[ARG0:%.+]]: vector<6xi8>, [[ARG1:%.+]]: vector<6xi8>)159// CHECK: [[RED:%.+]] = vector.reduction160// CHECK-NEXT: return [[RED]] : i32161func.func @too_long(%arg0: vector<6xi8>, %arg1: vector<6xi8>) -> i32 {162 %lhs = arith.extsi %arg0 : vector<6xi8> to vector<6xi32>163 %rhs = arith.extsi %arg1 : vector<6xi8> to vector<6xi32>164 %mul = arith.muli %lhs, %rhs : vector<6xi32>165 %red = vector.reduction <add>, %mul : vector<6xi32> into i32166 return %red : i32167}168 169// CHECK-LABEL: func.func @wrong_reduction_kind170// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>)171// CHECK: [[RED:%.+]] = vector.reduction <mul>172// CHECK-NEXT: return [[RED]] : i32173func.func @wrong_reduction_kind(%arg0: vector<4xi8>, %arg1: vector<4xi8>) -> i32 {174 %lhs = arith.extsi %arg0 : vector<4xi8> to vector<4xi32>175 %rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi32>176 %mul = arith.muli %lhs, %rhs : vector<4xi32>177 %red = vector.reduction <mul>, %mul : vector<4xi32> into i32178 return %red : i32179}180 181// CHECK-LABEL: func.func @wrong_arith_op182// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>)183// CHECK: [[ADD:%.+]] = arith.addi184// CHECK: [[RED:%.+]] = vector.reduction <mul>, [[ADD]]185// CHECK-NEXT: return [[RED]] : i32186func.func @wrong_arith_op(%arg0: vector<4xi8>, %arg1: vector<4xi8>) -> i32 {187 %lhs = arith.extsi %arg0 : vector<4xi8> to vector<4xi32>188 %rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi32>189 %add = arith.addi %lhs, %rhs : vector<4xi32>190 %red = vector.reduction <mul>, %add : vector<4xi32> into i32191 return %red : i32192}193