加速霓虹灯的矩阵速度

如何解决加速霓虹灯的矩阵速度

我尝试以 1k x 1k 运行以下代码，所用时间为 1.4 秒。

有没有办法提高速度，代码在raspberry pi-4上测试过。

相同大小的基于 numpy 的乘法需要 0.14 秒来执行代码。

A 和 B 矩阵的大小相同，均为 1000 x 1000。

代码在编译时自动矢量化。

void matrix_multiply_neon(float32_t  *A,float32_t  *B,float32_t *C,uint32_t n,uint32_t m,uint32_t k) {
    int A_idx;
    int B_idx;
    int C_idx;

    float32x4_t A0;
    float32x4_t A1;
    float32x4_t A2;
    float32x4_t A3;
  
    float32x4_t B0;
    float32x4_t B1;
    float32x4_t B2;
    float32x4_t B3;
    
    float32x4_t C0;
    float32x4_t C1;
    float32x4_t C2;
    float32x4_t C3;
    for (int i_idx=0; i_idx<n; i_idx+=4) {
            for (int j_idx=0; j_idx<m; j_idx+=4) {
                    // Zero accumulators before matrix op
                    C0 = vmovq_n_f32(0);
                    C1 = vmovq_n_f32(0);
                    C2 = vmovq_n_f32(0);
                    C3 = vmovq_n_f32(0);
                    for (int k_idx=0; k_idx<k; k_idx+=4) {
                           
                            A_idx = i_idx + n*k_idx;
                            B_idx = k*j_idx + k_idx;
                            
                            A0 = vld1q_f32(A+A_idx);
                            A1 = vld1q_f32(A+A_idx+n);
                            A2 = vld1q_f32(A+A_idx+2*n);
                            A3 = vld1q_f32(A+A_idx+3*n);
                            
                            B0 = vld1q_f32(B+B_idx);
                            C0 = vfmaq_laneq_f32(C0,A0,B0,0);
                            C0 = vfmaq_laneq_f32(C0,A1,1);
                            C0 = vfmaq_laneq_f32(C0,A2,2);
                            C0 = vfmaq_laneq_f32(C0,A3,3);
                            B1 = vld1q_f32(B+B_idx+k);
                            C1 = vfmaq_laneq_f32(C1,B1,0);
                            C1 = vfmaq_laneq_f32(C1,1);
                            C1 = vfmaq_laneq_f32(C1,2);
                            C1 = vfmaq_laneq_f32(C1,3);
                            B2 = vld1q_f32(B+B_idx+2*k);
                            C2 = vfmaq_laneq_f32(C2,B2,0);
                            C2 = vfmaq_laneq_f32(C2,1);
                            C2 = vfmaq_laneq_f32(C2,2);
                            C2 = vfmaq_laneq_f32(C2,3);
                            B3 = vld1q_f32(B+B_idx+3*k);
                            C3 = vfmaq_laneq_f32(C3,B3,0);
                            C3 = vfmaq_laneq_f32(C3,1);
                            C3 = vfmaq_laneq_f32(C3,2);
                            C3 = vfmaq_laneq_f32(C3,3);
                    }
                    // Compute base index for stores
                    C_idx = n*j_idx + i_idx;
                    vst1q_f32(C+C_idx,C0);
                    vst1q_f32(C+C_idx+n,C1);
                    vst1q_f32(C+C_idx+2*n,C2);
                    vst1q_f32(C+C_idx+3*n,C3);
            }
    }
}