[Vectorization] Support integer truncation operations

[yzhang93]

An example of elementwise op with trunci is as

%3 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
                         iterator_types = ["parallel", "parallel", "parallel", "parallel"]
                        } ins(%2 : tensor<256x16x64x64xi32>) outs(%i8out : tensor<256x16x64x64xi8>) {
      ^bb0(%in: i32, %out: i8):
        %4 = arith.shrsi %in, %cst_shift : i32
        %5 = arith.trunci %4 : i32 to i8
        linalg.yield %5 : i8
    } -> tensor<256x16x64x64xi8>

Motivation:

We observed significant performance improvements when using a ukernel for the operation above. However, since the vectorization path with Peano doesn't support this operation, the code falls back to scalar execution, resulting in poor performance as below:

-----------------------------------------------------------------------------------------------------------
Benchmark                                                 Time             CPU   Iterations UserCounters...
-----------------------------------------------------------------------------------------------------------
BM_matmul4d_trunci/process_time/real_time              8719 us          150 us          100 items_per_second=114.69/s
BM_matmul4d_trunci/process_time/real_time              8719 us          159 us          100 items_per_second=114.697/s
BM_matmul4d_trunci/process_time/real_time              8688 us          155 us          100 items_per_second=115.098/s
BM_matmul4d_trunci/process_time/real_time              8689 us          159 us          100 items_per_second=115.094/s
BM_matmul4d_trunci/process_time/real_time              8709 us          159 us          100 items_per_second=114.83/s
BM_matmul4d_trunci/process_time/real_time_mean         8705 us          156 us            5 items_per_second=114.882/s
BM_matmul4d_trunci/process_time/real_time_median       8709 us          159 us            5 items_per_second=114.83/s
BM_matmul4d_trunci/process_time/real_time_stddev       15.4 us         4.13 us            5 items_per_second=0.20339/s
BM_matmul4d_trunci/process_time/real_time_cv           0.18 %          2.64 %             5 items_per_second=0.18%

Below is the performance with ukernel path:

-----------------------------------------------------------------------------------------------------------
Benchmark                                                 Time             CPU   Iterations UserCounters...
-----------------------------------------------------------------------------------------------------------
BM_matmul4d_trunci/process_time/real_time              1146 us         51.5 us          700 items_per_second=872.337/s
BM_matmul4d_trunci/process_time/real_time              1149 us         64.8 us          700 items_per_second=870.528/s
BM_matmul4d_trunci/process_time/real_time              1149 us         65.5 us          700 items_per_second=869.999/s
BM_matmul4d_trunci/process_time/real_time              1149 us         65.9 us          700 items_per_second=870.46/s
BM_matmul4d_trunci/process_time/real_time              1150 us         66.9 us          700 items_per_second=869.819/s
BM_matmul4d_trunci/process_time/real_time_mean         1149 us         62.9 us            5 items_per_second=870.629/s
BM_matmul4d_trunci/process_time/real_time_median       1149 us         65.5 us            5 items_per_second=870.46/s
BM_matmul4d_trunci/process_time/real_time_stddev       1.32 us         6.42 us            5 items_per_second=1.00111/s
BM_matmul4d_trunci/process_time/real_time_cv           0.11 %         10.20 %             5 items_per_second=0.11%

To reproduce the numbers, for ukernel run

python cpu_comparison/run.py results /**iree-installed-path**/ --peano_dir=/**llvm-aie-installed-path**/ --xrt_lite_n_core_rows=$XRT_LITE_N_CORE_ROWS --xrt_lite_n_core_cols=$XRT_LITE_N_CORE_COLS --target_device="npu4" -v --vitis_dir=/opt/Xilinx/Vitis/2024.2 --tests=matmul4d_scale_trunci_1024_16384_512_i8_i32_O1_npu4_outline_4_level_tiling_ukernel_peano_ctrlpkt_benchmark

For vectorization path, modify "use_ukernel": False in

iree-amd-aie/build_tools/ci/cpu_comparison/run.py

Line 2215 in b8ad10e

"use_ukernel": True,

[newling]

Can we run the above linalg.generic by itself in iree-amd-aie, or does it need to be preceded by a matmul for our compiler to handle it?

[newling]

end-to-end instructions to reproduce these numbers would be very useful!

[newling]

So I guess we need vector path for arith.muli, arith.shrsi, and arith.trunci

[yzhang93]

Can we run the above linalg.generic by itself in iree-amd-aie, or does it need to be preceded by a matmul for our compiler to handle it?

No, there's not a pipeline for elementwise op only, it has to go with matmul/matmul4d op.

I also added the command to reproduce the numbers.

[newling]

An example of a microkernel which seems to do what we want is:

iree-amd-aie/compiler/plugins/target/AMD-AIE/iree-amd-aie/Target/mm_npu4_peano.cc

Line 9 in 9cedcdd

template<int M, int N, int r>

First step might be to compile this peano and see what the LLVMIR generated is. There should be some AIE intrinsic instruction visible, which we can then target when lowering.

@jtuyls wrote this ukernel so might have some additional comments.

One thing I'm not 100% sure about: the is still looping over the i32 elements, so there isn't a 'vectorized' instruction to truncate/shift multiple i32s simultaneously?

Some other I found doing some digging

shrsi lowering:

using ShRSIOpLowering =
VectorConvertToLLVMPattern<arith::ShRSIOp, LLVM::AShrOp>;

ashr is the arithmetic shift right which is a bit more complex that

Signed bit shift:

arith.shrsi lowers to LLVM::AShrOp see here

whereas for unsigned bit shift

arith.shrui lowers to LLVM::LShrOp see here

The documentation for the signed and unsigned bit shifts in LLVM is here:

logical shift : https://llvm.org/docs/LangRef.html#lshr-instruction
arithmetic shift : https://llvm.org/docs/LangRef.html#ashr-instruction

[newling]

@Sameeranjoshi ?

[newling]

The initial problem here is that the linalg.generic is never being vectorized because

arith::ShRSIOp is missing from https://github.com/nod-ai/iree-amd-aie/blame/4d14c00de4184c772d51d94dcf0010253e1fdaa6/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEVectorization.cpp#L85

Adding that as a one line change and recompiling iree-amd-aie, then arith.shrsiop gets vectorized to

%12 = arith.shrsi %11, %cst : vector<1x1x1x1x4x8xi32>

where the number of bits to shift is

%cst = arith.constant dense<5> : vector<1x1x1x1x4x8xi32>

So that's good. Continuing on from there, it gets flattened somewhere to

 %30 = arith.shrsi %29, %cst_64 : vector<32xi32>

and then in convert-arith-to-llvm it becomes

%58 = llvm.ashr %57, %29 : vector<32xi32>

i.e. the llvm arithmetic shift operator. Still looking good. This gets translated to LLVM IR, and then peano in llc fails with

LLVM ERROR: unable to legalize instruction: %373:_(<32 x s32>) = G_ASHR %372:_, %19:_(<32 x s32>) (in function: core_1_3)

Basically peano doesn't know what to do with this vectorized arithmetic shift instruction. Why not?

A couple of possible issues,

1. the shape isn't right
2. in isolation the op cannot be matched, more context is needed (what are the surrounding ops).
3. peano just hasn't implemented this legalization
4. peano can fundamentally not support arithmetic shift right because of HW constrains, can only support logical shift right

If 1, easy fix on our end.
If 2 or 3, peano needs better matching. It looks like there is supports for ASHL: https://github.com/Xilinx/llvm-aie/blob/77e395c8db5c08f230884b9980ec2647eb182695/llvm/lib/Target/AIE/aie2p/AIE2PInstrPatterns.td#L187 which seems very promising, and a TODO comment mentioning ASHR!
If 4, we might have to reconsider how quantization works (or have scalar performance)

UPDATE: code changes on branch main...newling:iree-amd-aie:shrsi_peano_fail_repro

[newling]

In

https://github.com/nod-ai/iree-amd-aie/blob/main/build_tools/ci/cpu_comparison/matmul_template/matmul_trunci_scaling_MxK_KxN.mlir

the sequence of ops is

    %4 = arith.extsi %in : ${TYPE2} to ${TYPE_MUL_RESULT}
     %5 = arith.muli %4, %cst_mul : ${TYPE_MUL_RESULT}
     %6 = arith.shrsi %5, %cst_shift : ${TYPE_MUL_RESULT}
     %7 = arith.trunci %6 : ${TYPE_MUL_RESULT} to ${TYPE1}

In LLVM IR this is

sext
mul
ashr
trunc

I think this should lower as

sext -> UPS
mul -> ? surely something
ashr, trunc -> SRS

i.e. the combination of ashr+trunc should be srs.

[newling]

Update on the lowerings of sext, mul, and ashr + trunc:

mul: see also Xilinx/llvm-aie#488

Following LLVM IR, to be compiled (function by function) with peano's llc:

@buff_src = external global [128 x i32]
@buff_dst = external global [128 x i32]

; See https://github.com/Xilinx/llvm-aie/issues/488
define void @mul() {
  %1 = load <32 x i16>, ptr @buff_src, align 256
  %2 = load <32 x i16>, ptr @buff_src, align 256
  ; add lowers, but not mul. 
  ;             ===========
  mul <32 x i16> %1, %2
  store <32 x i16> %3, ptr @buff_dst, align 256
  ret void
}

; sign extend
; MLIR:
;   arith.extsi
; Can target UPS I think, support sign bit replication 
define void @sext() {
  %1 = load <16 x i16>, ptr @buff_src, align 256
  %2 = sext <16 x i16> %1 to <16 x i32>
  store <16 x i32> %2, ptr @buff_dst, align 256
  ret void
}
;  
; signed shift right (arithmetic shift right) followed by truncation
; MLIR:
;   arith.shrsi
;   arith.trunci
; I think this can be supported by ups
define void @ashr_trunc() {
  %1 = load  <16 x i32>, ptr @buff_src, align 256
  %2 = ashr <16 x i32> %1, splat (i32 3)
  %3 = trunc <16 x i32> %2 to <16 x i16>
  store <16 x i16> %3, ptr @buff_dst, align 256
  ret void
}

; truncate
; MLIR:
;   arith.trunci
; remove top half of bits
; SUPPORTED for aie2p: 
; See https://github.com/Xilinx/llvm-aie/issues/482#issuecomment-2957149011
; with shuffle
define void @trunc() { 
  %1 = load <16 x i32>, ptr @buff_src, align 256
  %2 = trunc <16 x i32> %1 to <16 x i16>
  store <16 x i16> %2, ptr @buff_dst, align 256
  ret void
}

[newling]

Another finding:

The shrsi op is intercepted in mlir-aie and converted.

shrsi(x) becomes srs(ups(x))

https://github.com/Xilinx/mlir-aie/blob/0ce127c3835a6777538f37c397385a7f1ecd8108/lib/Dialect/AIEVec/Transforms/VectorToAIEVecConversions.cpp#L2881

Suppose it's x is v16i32 (vector of 16 i32 values).
ups casts this to a vector of 16 i64 values
then srs casts it to a vector of 16 i32s, but does a bit shift in the process (srs is shift round scale)

Now we could just pull that mlir-aie code back into the codebase, but I'm not sure we should.