78 lines · plain
1// RUN: mlir-opt %s --convert-vector-to-llvm='vector-contract-lowering=llvmintr' | FileCheck %s2 3#matmat_accesses = [4 affine_map<(i, j, k) -> (i, k)>,5 affine_map<(i, j, k) -> (k, j)>,6 affine_map<(i, j, k) -> (i, j)>7]8#matmat_trait = {9 indexing_maps = #matmat_accesses,10 iterator_types = ["parallel", "parallel", "reduction"]11}12 13// CHECK-LABEL: func.func @matmul(14// CHECK-SAME: %[[ARG0:.*]]: vector<2x4xf32>,15// CHECK-SAME: %[[ARG1:.*]]: vector<4x3xf32>,16// CHECK-SAME: %[[ARG2:.*]]: vector<2x3xf32>) -> vector<2x3xf32> {17// CHECK: %[[VAL_0:.*]] = builtin.unrealized_conversion_cast %[[ARG1]] : vector<4x3xf32> to !llvm.array<4 x vector<3xf32>>18// CHECK: %[[VAL_1:.*]] = builtin.unrealized_conversion_cast %[[ARG0]] : vector<2x4xf32> to !llvm.array<2 x vector<4xf32>>19// CHECK: %[[VAL_2:.*]] = ub.poison : vector<2x3xf32>20// CHECK: %[[VAL_3:.*]] = builtin.unrealized_conversion_cast %[[VAL_2]] : vector<2x3xf32> to !llvm.array<2 x vector<3xf32>>21// CHECK: %[[POISON_RHS:.*]] = ub.poison : vector<12xf32>22// CHECK: %[[POISON_LHS:.*]] = ub.poison : vector<8xf32>23 24// ===> Extract LHS25// | ROW_1 |26// | ----- | --> | ROW_1 | ROW_2 |27// | ROW_2 |28//29// CHECK: %[[LHS_ROW_1:.*]] = llvm.extractvalue %[[VAL_1]][0] : !llvm.array<2 x vector<4xf32>>30// CHECK: %[[TP_1:.*]] = llvm.shufflevector %[[LHS_ROW_1]], %[[LHS_ROW_1]] [0, 1, 2, 3, 0, 0, 0, 0] : vector<4xf32>31// CHECK: %[[TP_2:.*]] = llvm.shufflevector %[[TP_1]], %[[POISON_LHS]] [0, 1, 2, 3, 12, 13, 14, 15] : vector<8xf32>32// CHECK: %[[LHS_ROW_2:.*]] = llvm.extractvalue %[[VAL_1]][1] : !llvm.array<2 x vector<4xf32>>33// CHECK: %[[TP_3:.*]] = llvm.shufflevector %[[LHS_ROW_2]], %[[LHS_ROW_2]] [0, 1, 2, 3, 0, 0, 0, 0] : vector<4xf32>34// CHECK: %[[LHS:.*]] = llvm.shufflevector %[[TP_3]], %[[TP_2]] [8, 9, 10, 11, 0, 1, 2, 3] : vector<8xf32>35 36// == Extract RHS37// | ROW_1 |38// | ----- |39// | ROW_2 |40// | ----- | --> | ROW_1 | ROW_2 | ROW_3 | ROW_4 |41// | ROW_3 |42// | ----- |43// | ROW_4 |44// CHECK: %[[RHS_ROW_1:.*]] = llvm.extractvalue %[[VAL_0]][0] : !llvm.array<4 x vector<3xf32>>45// CHECK: %[[TP_4:.*]] = llvm.shufflevector %[[RHS_ROW_1]], %[[RHS_ROW_1]] [0, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0] : vector<3xf32>46// CHECK: %[[TP_5:.*]] = llvm.shufflevector %[[TP_4]], %[[POISON_RHS]] [0, 1, 2, 15, 16, 17, 18, 19, 20, 21, 22, 23] : vector<12xf32>47// CHECK: %[[RHS_ROW_2:.*]] = llvm.extractvalue %[[VAL_0]][1] : !llvm.array<4 x vector<3xf32>>48// CHECK: %[[TP_6:.*]] = llvm.shufflevector %[[RHS_ROW_2]], %[[RHS_ROW_2]] [0, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0] : vector<3xf32>49// CHECK: %[[TP_7:.*]] = llvm.shufflevector %[[TP_6]], %[[TP_5]] [12, 13, 14, 0, 1, 2, 18, 19, 20, 21, 22, 23] : vector<12xf32>50// CHECK: %[[RHS_ROW_3:.*]] = llvm.extractvalue %[[VAL_0]][2] : !llvm.array<4 x vector<3xf32>>51// CHECK: %[[TP_8:.*]] = llvm.shufflevector %[[RHS_ROW_3]], %[[RHS_ROW_3]] [0, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0] : vector<3xf32>52// CHECK: %[[TP_9:.*]] = llvm.shufflevector %[[TP_8]], %[[TP_7]] [12, 13, 14, 15, 16, 17, 0, 1, 2, 21, 22, 23] : vector<12xf32>53// CHECK: %[[RHS_ROW_4:.*]] = llvm.extractvalue %[[VAL_0]][3] : !llvm.array<4 x vector<3xf32>>54// CHECK: %[[TP_10:.*]] = llvm.shufflevector %[[RHS_ROW_4]], %[[RHS_ROW_4]] [0, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0] : vector<3xf32>55// CHECK: %[[RHS:.*]] = llvm.shufflevector %[[TP_10]], %[[TP_9]] [12, 13, 14, 15, 16, 17, 18, 19, 20, 0, 1, 2] : vector<12xf32>56 57// ===> Matrix multiply58// CHECK: %[[MM:.*]] = llvm.intr.matrix.multiply %[[LHS]], %[[RHS]] {lhs_columns = 4 : i32, lhs_rows = 2 : i32, rhs_columns = 3 : i32} : (vector<8xf32>, vector<12xf32>) -> vector<6xf32>59// CHECK: %[[RES:.*]] = arith.addf %[[ARG2]], %{{.*}} : vector<2x3xf32>60// CHECK: return %[[RES]] : vector<2x3xf32>61func.func @matmul(%arg0: vector<2x4xf32>,62 %arg1: vector<4x3xf32>,63 %arg2: vector<2x3xf32>) -> vector<2x3xf32> {64 %0 = vector.contract #matmat_trait %arg0, %arg1, %arg265 : vector<2x4xf32>, vector<4x3xf32> into vector<2x3xf32>66 return %0 : vector<2x3xf32>67}68 69// CHECK-LABEL: func @matmul_scalable70// CHECK-NOT: llvm.intr.matrix.multiply71func.func @matmul_scalable(%arg0: vector<2x4xf32>,72 %arg1: vector<4x[3]xf32>,73 %arg2: vector<2x[3]xf32>) -> vector<2x[3]xf32> {74 %0 = vector.contract #matmat_trait %arg0, %arg1, %arg275 : vector<2x4xf32>, vector<4x[3]xf32> into vector<2x[3]xf32>76 return %0 : vector<2x[3]xf32>77}78