169 lines · plain
1// RUN: mlir-opt %s -convert-vector-to-xegpu -split-input-file | FileCheck %s2 3#map = affine_map<(d0, d1, d2) -> (d0, d2)>4#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>5#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>6func.func @dpas_gemm_f16(%lhs: vector<8x16xf16>, %rhs: vector<16x16xf16>,7 %acc: vector<8x16xf32>) -> vector<8x16xf32> {8 %3 = vector.contract9 {indexing_maps = [#map, #map1, #map2],10 iterator_types = ["parallel", "parallel", "reduction"],11 kind = #vector.kind<add>} %lhs, %rhs, %acc12 : vector<8x16xf16>, vector<16x16xf16> into vector<8x16xf32>13 return %3 : vector<8x16xf32>14}15 16// CHECK-LABEL: @dpas_gemm_f16(17// CHECK-SAME: %[[LHS:.+]]: vector<8x16xf16>,18// CHECK-SAME: %[[RHS:.+]]: vector<16x16xf16>,19// CHECK-SAME: %[[ACC:.+]]: vector<8x16xf32>20// CHECK: %[[DPAS:.+]] = xegpu.dpas21// CHECK-SAME: %[[LHS]], %[[RHS]], %[[ACC]]22// CHECK-SAME: {{.*}}-> vector<8x16xf32>23// CHECK: return %[[DPAS]]24 25// -----26 27#map = affine_map<(d0, d1, d2) -> (d0, d2)>28#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>29#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>30func.func @dpas_gemm_i8(%lhs: vector<8x32xi8>, %rhs: vector<32x16xi8>,31 %acc: vector<8x16xi32>) -> vector<8x16xi32> {32 %3 = vector.contract33 {indexing_maps = [#map, #map1, #map2],34 iterator_types = ["parallel", "parallel", "reduction"],35 kind = #vector.kind<add>} %lhs, %rhs, %acc36 : vector<8x32xi8>, vector<32x16xi8> into vector<8x16xi32>37 return %3 : vector<8x16xi32>38}39 40// CHECK-LABEL: @dpas_gemm_i8(41// CHECK-SAME: %[[LHS:.+]]: vector<8x32xi8>,42// CHECK-SAME: %[[RHS:.+]]: vector<32x16xi8>,43// CHECK-SAME: %[[ACC:.+]]: vector<8x16xi32>44// CHECK: %[[DPAS:.+]] = xegpu.dpas45// CHECK-SAME: %[[LHS]], %[[RHS]], %[[ACC]]46// CHECK-SAME: {{.*}}-> vector<8x16xi32>47// CHECK: return %[[DPAS]]48 49// -----50 51// No restriction on vector sizes to allow capturing workgroup-sized operations.52// The operations can then be progressively resized through distribution down53// to hardware compatible sizes.54 55#map = affine_map<(d0, d1, d2) -> (d0, d2)>56#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>57#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>58func.func @dpas_large_dims(%lhs: vector<128x512xf16>, %rhs: vector<512x256xf16>,59 %acc: vector<128x256xf32>) -> vector<128x256xf32> {60 %3 = vector.contract61 {indexing_maps = [#map, #map1, #map2],62 iterator_types = ["parallel", "parallel", "reduction"],63 kind = #vector.kind<add>} %lhs, %rhs, %acc64 : vector<128x512xf16>, vector<512x256xf16> into vector<128x256xf32>65 return %3 : vector<128x256xf32>66}67 68// CHECK-LABEL: @dpas_large_dims(69// CHECK-SAME: %[[LHS:.+]]: vector<128x512xf16>,70// CHECK-SAME: %[[RHS:.+]]: vector<512x256xf16>,71// CHECK-SAME: %[[ACC:.+]]: vector<128x256xf32>72// CHECK: %[[DPAS:.+]] = xegpu.dpas73// CHECK-SAME: %[[LHS]], %[[RHS]], %[[ACC]]74// CHECK-SAME: {{.*}}-> vector<128x256xf32>75// CHECK: return %[[DPAS]]76 77// -----78 79// For simplicity, only plain data layouts are currently supported.80// VNNI packing is applied later as a separate lowering step.81 82#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>83#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>84#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>85func.func @negative_vnni_packed(%lhs: vector<8x8x2xf16>, %rhs: vector<8x16x2xf16>,86 %acc: vector<8x16xf32>) -> vector<8x16xf32> {87 %3 = vector.contract88 {indexing_maps = [#map, #map1, #map2],89 iterator_types = ["parallel", "parallel", "reduction", "reduction"],90 kind = #vector.kind<add>} %lhs, %rhs, %acc91 : vector<8x8x2xf16>, vector<8x16x2xf16> into vector<8x16xf32>92 return %3 : vector<8x16xf32>93}94 95// CHECK-LABEL: @negative_vnni_packed(96// CHECK: vector.contract97 98// -----99 100#map = affine_map<(d0, d1, d2) -> (d0, d2)>101#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>102#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>103func.func @negative_combining_kind(%lhs: vector<8x16xf16>, %rhs: vector<16x16xf16>,104 %acc: vector<8x16xf32>) -> vector<8x16xf32> {105 %3 = vector.contract106 {indexing_maps = [#map, #map1, #map2],107 iterator_types = ["parallel", "parallel", "reduction"],108 kind = #vector.kind<mul>} %lhs, %rhs, %acc109 : vector<8x16xf16>, vector<16x16xf16> into vector<8x16xf32>110 return %3 : vector<8x16xf32>111}112 113// CHECK-LABEL: @negative_combining_kind(114// CHECK: vector.contract115 116// -----117 118#map = affine_map<(d0, d1, d2) -> (d0, d2)>119#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>120#map2 = affine_map<(d0, d1, d2) -> ()>121func.func @negative_accumulator_shape(%lhs: vector<8x16xf16>, %rhs: vector<16x16xf16>,122 %acc: vector<f32>) -> vector<f32> {123 %3 = vector.contract124 {indexing_maps = [#map, #map1, #map2],125 iterator_types = ["reduction", "reduction", "reduction"],126 kind = #vector.kind<add>} %lhs, %rhs, %acc127 : vector<8x16xf16>, vector<16x16xf16> into vector<f32>128 return %3 : vector<f32>129}130 131// CHECK-LABEL: @negative_accumulator_shape(132// CHECK: vector.contract133 134// -----135 136#map = affine_map<(d0, d1, d2) -> (d2, d0)>137#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>138#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>139func.func @negative_gemm_transpose_a(%lhs: vector<16x8xf16>, %rhs: vector<16x16xf16>,140 %acc: vector<8x16xf32>) -> vector<8x16xf32> {141 %3 = vector.contract142 {indexing_maps = [#map, #map1, #map2],143 iterator_types = ["parallel", "parallel", "reduction"],144 kind = #vector.kind<add>} %lhs, %rhs, %acc145 : vector<16x8xf16>, vector<16x16xf16> into vector<8x16xf32>146 return %3 : vector<8x16xf32>147}148 149// CHECK-LABEL: @negative_gemm_transpose_a(150// CHECK: vector.contract151 152// -----153 154#map = affine_map<(d0, d1, d2) -> (d0, d2)>155#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>156#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>157func.func @negative_gemm_transpose_b(%lhs: vector<8x16xf16>, %rhs: vector<16x16xf16>,158 %acc: vector<8x16xf32>) -> vector<8x16xf32> {159 %3 = vector.contract160 {indexing_maps = [#map, #map1, #map2],161 iterator_types = ["parallel", "parallel", "reduction"],162 kind = #vector.kind<add>} %lhs, %rhs, %acc163 : vector<8x16xf16>, vector<16x16xf16> into vector<8x16xf32>164 return %3 : vector<8x16xf32>165}166 167// CHECK-LABEL: @negative_gemm_transpose_b(168// CHECK: vector.contract169