[mlir][Aarch64] Improve i8mm instruction sequence for `vector.contract`

[dcaballe]

The i8mm lowering for some vector.contract ops is currently functionally correct. However, performance wise there is some room for improvement. Looking at the generated asm for an mmt4d with 2x2x8 innermost tile sizes, we get:

    1470: 6e180483      mov     v3.d[1], v4.d[0]                                                                                                                                                                           
    1474: 4e006204      tbl     v4.16b, { v16.16b, v17.16b, v18.16b, v19.16b }, v0.16b                                                                                                                                     
    1478: 4e84a462      smmla   v2.4s, v3.16b, v4.16b                                                                                                                                                                      
    147c: 6e024041      ext     v1.16b, v2.16b, v2.16b, #0x8 

It calls my attention the mov instruction, esp. the indexing from 1 to 0, the tbl and the ext instructions. This may not seem a big deal but the problem is really exacerbated when using larger tile sizes. We observed large sequences of mov and ext instructions all over the place.

We should investigate what is going on and try to fix the problem. My suspicion is that this zero initialization and insertion for vecmat cases might be behind some of these instructions. We should try if using llvm.undef fixes part of the problem.

[llvmbot]

@llvm/issue-subscribers-mlir

Author: Diego Caballero (dcaballe)

The i8mm lowering for some `vector.contract` ops is currently functionally correct. However, performance wise there is some room for improvement. Looking at the generated asm for an mmt4d with 2x2x8 innermost tile sizes, we get:

    1470: 6e180483      mov     v3.d[1], v4.d[0]                                                                                                                                                                           
    1474: 4e006204      tbl     v4.16b, { v16.16b, v17.16b, v18.16b, v19.16b }, v0.16b                                                                                                                                     
    1478: 4e84a462      smmla   v2.4s, v3.16b, v4.16b                                                                                                                                                                      
    147c: 6e024041      ext     v1.16b, v2.16b, v2.16b, #<!-- -->0x8 

It calls my attention the mov instruction, esp. the indexing from 1 to 0, the tbl and the ext instructions. This may not seem a big deal but the problem is really exacerbated when using larger tile sizes. We observed large sequences of mov and ext instructions all over the place.

We should investigate what is going on and try to fix the problem. My suspicion is that this zero initialization and insertion for vecmat cases might be behind some of these instructions. We should try if using llvm.undef fixes part of the problem.

[momchil-velikov]

I had a look and got code like

	ldp	q7, q18, [x4, #-64]
	mov	v2.16b, v5.16b
	mov	v1.d[1], v4.d[0]
	zip2	v4.2d, v7.2d, v18.2d
	mov	v2.d[1], v3.d[0]
	zip2	v3.2d, v5.2d, v3.2d
	mov	v5.16b, v7.16b
	mov	v5.d[1], v18.d[0]
	smmla	v4.4s, v16.16b, v1.16b
	smmla	v2.4s, v6.16b, v0.16b
	smmla	v3.4s, v6.16b, v1.16b
	smmla	v5.4s, v16.16b, v0.16b
	mov	v0.16b, v2.16b
	zip2	v1.2d, v2.2d, v3.2d
	mov	v0.d[1], v3.d[0]
	mov	v2.16b, v5.16b
	zip2	v3.2d, v5.2d, v4.2d
	mov	v2.d[1], v4.d[0]

The mov/zip2 look like a reasonable amount of instructions needed to pack a discontiguous 2x2 tile into a SIMD register (and respectively unpack a SIMD register after the smmla. This example is from tiling and vectorising

  %out = linalg.mmt4d ins(%lhs, %rhs: tensor<2x2x4x8xi8>, tensor<2x2x4x8xi8>) outs(%acc: tensor<2x2x4x4xi32>) -> tensor<2x2x4x4xi32>

using

    // Step 1: Tile
    // Tile parallel dims
    %tiled_linalg_op_p, %loops:4 = transform.structured.tile_using_for %mmt4d tile_sizes [1, 1, 0, 4, 4, 0]
     : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
    // Tile reduction dims
    %tiled_linalg_op_r, %loops2:2 = transform.structured.tile_using_for %tiled_linalg_op_p tile_sizes [0, 0, 1, 0, 0, 8]
     : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)

    // Step 2: Vectorize
    transform.structured.vectorize %tiled_linalg_op_r vector_sizes [1, 1, 1, 4, 4, 8] : !transform.any_op