109 lines · plain
1// RUN: mlir-opt %s -transform-interpreter | FileCheck %s2 3#packed_maps = [4 affine_map<(d0, d1, d2) -> (d0, d2)>,5 affine_map<(d0, d1, d2) -> (d1, d2)>,6 affine_map<(d0, d1, d2) -> (d0, d1)>7]8 9// CHECK-LABEL: @test_vector_contract_to_usmmla_rev10 11// CHECK-SAME: %[[LHS:arg0]]: vector<4x8xi8>12// CHECK-SAME: %[[RHS:arg1]]: vector<[4]x8xi8>13// CHECK-SAME: %[[ACC:arg2]]: vector<4x[4]xi32>14 15// CHECK: [[P0:[0-9]+]] = ub.poison : vector<[8]xi32>16// CHECK-NEXT: [[P1:[0-9]+]] = ub.poison : vector<4x[4]xi32>17// CHECK-NEXT: [[P2:[0-9]+]] = ub.poison : vector<[16]xi8>18 19// Extract LHS rows 0 and 1, concatenate, turn into scalable vector20// CHECK: %[[T3:[0-9]+]] = vector.extract %[[LHS]][0] : vector<8xi8> from vector<4x8xi8>21// CHECK-NEXT: %[[T4:[0-9]+]] = vector.extract %[[LHS]][1] : vector<8xi8> from vector<4x8xi8>22// CHECK-NEXT: %[[T5:[0-9]+]] = vector.shuffle %[[T3]], %[[T4]] [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] : vector<8xi8>, vector<8xi8>23// CHECK-NEXT: %[[T6:[0-9]+]] = vector.scalable.insert %[[T5]], %[[P2]][0] : vector<16xi8> into vector<[16]xi8>24 25// Replicate across the entire length of the scalable vector26// CHECK-NEXT: %[[LHS_0:[0-9]+]] = arm_sve.dupq_lane %[[T6]][0] : vector<[16]xi8>27 28// Same for LHS rows 2 and 329// CHECK-NEXT: %[[T8:[0-9]+]] = vector.extract %[[LHS]][2] : vector<8xi8> from vector<4x8xi8>30// CHECK-NEXT: %[[T9:[0-9]+]] = vector.extract %[[LHS]][3] : vector<8xi8> from vector<4x8xi8>31// CHECK-NEXT: %[[T10:[0-9]+]] = vector.shuffle %[[T8]], %[[T9]] [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] : vector<8xi8>, vector<8xi8>32// CHECK-NEXT: %[[T11:[0-9]+]] = vector.scalable.insert %[[T10]], %[[P2]][0] : vector<16xi8> into vector<[16]xi8>33// CHECK-NEXT: %[[LHS_1:[0-9]+]] = arm_sve.dupq_lane %[[T11]][0] : vector<[16]xi8>34 35// Extract sub-tiles from the RHS36// CHECK-NEXT: %[[T13:[0-9]+]] = vector.shape_cast %[[RHS]] : vector<[4]x8xi8> to vector<[32]xi8>37// CHECK-NEXT: %[[RHS_0:[0-9]+]] = vector.scalable.extract %[[T13]][0] : vector<[16]xi8> from vector<[32]xi8>38// CHECK-NEXT: %[[RHS_1:[0-9]+]] = vector.scalable.extract %[[T13]][16] : vector<[16]xi8> from vector<[32]xi8>39 40// Extract accumulator rows 0 and 1 and pack (into "registers")41// Note the lack of bitcasts (i.e. interleave by pairs). This has the effect of transposing the tile.42// CHECK-NEXT: %[[T16:[0-9]+]] = vector.extract %[[ACC]][0] : vector<[4]xi32> from vector<4x[4]xi32>43// CHECK-NEXT: %[[T17:[0-9]+]] = vector.extract %[[ACC]][1] : vector<[4]xi32> from vector<4x[4]xi32>44// CHECK-NEXT: %[[T20:[0-9]+]] = vector.interleave %[[T16]], %[[T17]] : vector<[4]xi32> -> vector<[8]xi32>45// CHECK-NEXT: %[[ACC_0:[0-9]+]] = vector.scalable.extract %[[T20]][0] : vector<[4]xi32> from vector<[8]xi32>46// CHECK-NEXT: %[[ACC_1:[0-9]+]] = vector.scalable.extract %[[T20]][4] : vector<[4]xi32> from vector<[8]xi32>47 48// Same for accumulator rows 2 and 349// CHECK-NEXT: %[[T24:[0-9]+]] = vector.extract %[[ACC]][2] : vector<[4]xi32> from vector<4x[4]xi32>50// CHECK-NEXT: %[[T25:[0-9]+]] = vector.extract %[[ACC]][3] : vector<[4]xi32> from vector<4x[4]xi32>51// CHECK-NEXT: %[[T28:[0-9]+]] = vector.interleave %[[T24]], %[[T25]] : vector<[4]xi32> -> vector<[8]xi32>52// CHECK-NEXT: %[[ACC_2:[0-9]+]] = vector.scalable.extract %[[T28]][0] : vector<[4]xi32> from vector<[8]xi32>53// CHECK-NEXT: %[[ACC_3:[0-9]+]] = vector.scalable.extract %[[T28]][4] : vector<[4]xi32> from vector<[8]xi32>54 55// Do the sub-tile matrix multiplications56// The accumulator is transposed, the operands are in the opposite order,57// thus the result is obtained transposed too.58//59// CHECK-NEXT: %[[PACK_RES_00:[0-9]+]] = arm_sve.usmmla %[[ACC_0]], %[[RHS_0]], %[[LHS_0]] : vector<[16]xi8> to vector<[4]xi32>60// CHECK-NEXT: %[[PACK_RES_01:[0-9]+]] = arm_sve.usmmla %[[ACC_1]], %[[RHS_1]], %[[LHS_0]] : vector<[16]xi8> to vector<[4]xi32>61// CHECK-NEXT: %[[PACK_RES_10:[0-9]+]] = arm_sve.usmmla %[[ACC_2]], %[[RHS_0]], %[[LHS_1]] : vector<[16]xi8> to vector<[4]xi32>62// CHECK-NEXT: %[[PACK_RES_11:[0-9]+]] = arm_sve.usmmla %[[ACC_3]], %[[RHS_1]], %[[LHS_1]] : vector<[16]xi8> to vector<[4]xi32>63 64// Unpack (from "registers") and insert in the output result rows 0 and 165// Again, no bitcast/interleave-by-pairs, i.e. the transposed result of sub-tile66// multiplications is unpacked into the correct output layout.67//68// CHECK-NEXT: %[[T36:[0-9]+]] = vector.scalable.insert %[[PACK_RES_00]], %[[P0]][0] : vector<[4]xi32> into vector<[8]xi32>69// CHECK-NEXT: %[[T37:[0-9]+]] = vector.scalable.insert %[[PACK_RES_01]], %[[T36]][4] : vector<[4]xi32> into vector<[8]xi32>70// CHECK-NEXT: %[[UNPACK_RES_0:[_a-z0-9]+]], %[[UNPACK_RES_1:[_a-z0-9]+]] = vector.deinterleave %[[T37]] : vector<[8]xi32> -> vector<[4]xi32>71// CHECK-NEXT: %[[TMP_OUT_0:[0-9]+]] = vector.insert %[[UNPACK_RES_0]], %[[P1]] [0] : vector<[4]xi32> into vector<4x[4]xi32>72// CHECK-NEXT: %[[TMP_OUT_1:[0-9]+]] = vector.insert %[[UNPACK_RES_1]], %[[TMP_OUT_0]] [1] : vector<[4]xi32> into vector<4x[4]xi32>73 74// Same for result rows 2 and 375// CHECK-NEXT: %[[T43:[0-9]+]] = vector.scalable.insert %[[PACK_RES_10]], %[[P0]][0] : vector<[4]xi32> into vector<[8]xi32>76// CHECK-NEXT: %[[T44:[0-9]+]] = vector.scalable.insert %[[PACK_RES_11]], %[[T43]][4] : vector<[4]xi32> into vector<[8]xi32>77// CHECK-NEXT: %[[UNPACK_RES_2:[_a-z0-9]+]], %[[UNPACK_RES_3:[_a-z0-9]+]] = vector.deinterleave %[[T44]] : vector<[8]xi32> -> vector<[4]xi32>78// CHECK-NEXT: %[[TMP_OUT_2:[0-9]+]] = vector.insert %[[UNPACK_RES_2]], %[[TMP_OUT_1]] [2] : vector<[4]xi32> into vector<4x[4]xi32>79// CHECK-NEXT: %[[OUT:[0-9]+]] = vector.insert %[[UNPACK_RES_3]], %[[TMP_OUT_2]] [3] : vector<[4]xi32> into vector<4x[4]xi32>80 81// CHECK-NEXT: return %[[OUT]] : vector<4x[4]xi32>82 83func.func @test_vector_contract_to_usmmla_rev(84 %lhs: vector<4x8xi8>,85 %rhs: vector<[4]x8xi8>,86 %acc: vector<4x[4]xi32>) -> vector<4x[4]xi32> {87 88 %0 = arith.extsi %lhs : vector<4x8xi8> to vector<4x8xi32>89 %1 = arith.extui %rhs : vector<[4]x8xi8> to vector<[4]x8xi32>90 %2 = vector.contract {indexing_maps = #packed_maps,91 iterator_types = ["parallel", "parallel", "reduction"],92 kind = #vector.kind<add>} %0, %1, %acc93 : vector<4x8xi32>, vector<[4]x8xi32> into vector<4x[4]xi32>94 95 return %2 : vector<4x[4]xi32>96}97 98module attributes {transform.with_named_sequence} {99 transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {100 %func = transform.structured.match ops{["func.func"]} in %module : (!transform.any_op) -> !transform.op<"func.func">101 102 transform.apply_patterns to %func {103 transform.apply_patterns.arm_sve.vector_contract_to_i8mm104 } : !transform.op<"func.func">105 106 transform.yield107 }108}109