avx2 计算一个uint8_t数组的平方加速

一 接到一个朋友需求，计算像素值的平方

for (int i = 0; i < height; i++)

    for (int j = 0; j < width; j++)

        result[i * stride + j] = pix[i * stride + j] * pix[i * stride + j]

二 avx加速

;void sqr_cal(uint8_t *in, uint16_t *out, int size)
;rdi rsi rdx
global sqr_cal
sqr_cal:     
     sub rdx, 16 ;每次计算16个像素
     movupd xmm0, [rdi] 非对齐加载16字节
     vpmovzxbw ymm1, xmm0 无符号padding uint8到uint16
     vpmullw ymm2, ymm1, ymm1 两个无符号uint16相乘，把低16位存入ymm2, 
     vmovdqa [rsi], ymm2 吧ymm2 写入到内存中
     add rdi, 16 移动数据指针
     add rsi, 32 移动数据指针 
     cmp rdx, 0 剩余计算数据和0比较
     jl .end 不到16字节了
     je .dret 刚好结束
     jmp sqr_cal 继续循环
.end:
     movupd xmm0, [rdi] 剩余数据载入
     vpmovsxbw ymm1, xmm0
     vpmullw ymm2, ymm1, ymm1
     vmovdqa [rsi], ymm2
.dret:
     ret

三 C语言使用范例

#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
extern "C" void sqr_cal(uint8_t *in, uint16_t *out, int size); 
  
int main()
{
 
    uint8_t in[32] = {1,1,234,2,2,2,3,3,3,4,4,4,5,5,5,6,6,6,7,7,7,8,8,8,9,9,9,10,10,10,11,11};
    uint16_t out[32] __attribute__((aligned(32))); //地址32对齐，使用avx对齐指令更快
    sqr_cal(in, out, 32);
    for (int i  = 0; i < 32; i++) {
        printf("%d ", out[i]);
    }
    printf("\n");
    return 0;
}

./main

1 1 54756 4 4 4 9 9 9 16 16 16 25 25 25 36 36 36 49 49 49 64 64 64 81 81 81 100 100 100 121 121