Matmul NEON 4x4 fully unrolled

matmul-neon.s

neon_matmul_4x4_full_unrolled_v3:
    /* r0 address of A
       r1 address of B
       r2 address of C
    */
    push {r4, r5, r6, r7, r8, lr} /* Keep integer registers */
    vpush {q4-q7}               /* Floating point registers starting from s16 must be preserved */

    /* Load A as row vectors */
	mov r4, r0
	add r5, r0, #16
	add r6, r0, #32
	add r7, r0, #48
    vld1.f32 {q12}, [r4]  /* q12 ← {a[0][0], a[0][1], a[0][2], a[0][3]} */
    vld1.f32 {q13}, [r5]  /* q13 ← {a[1][0], a[1][1], a[1][2], a[1][3]} */
    vld1.f32 {q14}, [r6]  /* q14 ← {a[2][0], a[2][1], a[2][2], a[2][3]} */
    vld1.f32 {q15}, [r7]  /* q15 ← {a[3][0], a[3][1], a[3][2], a[3][3]} */

    # Manually build the vector column
    mov r4, r1
    vldr.f32 s16, [r4]
    vldr.f32 s17, [r4, #16]
    vldr.f32 s18, [r4, #32]
    vldr.f32 s19, [r4, #48]
    /* q4 = { b[0][0], b[1][0], b[2][0], b[3][0] } */

    vldr.f32 s20, [r4, #4]
    vldr.f32 s21, [r4, #4+16]
    vldr.f32 s22, [r4, #4+32]
    vldr.f32 s23, [r4, #4+48]
    /* q5 = { b[0][1], b[1][1], b[2][1], b[3][1] } */

    vldr.f32 s24, [r4, #8]
    vldr.f32 s25, [r4, #8+16]
    vldr.f32 s26, [r4, #8+32]
    vldr.f32 s27, [r4, #8+48]
    /* q6 = { b[0][2], b[1][2], b[2][2], b[3][2] } */

    vldr.f32 s28, [r4, #12]
    vldr.f32 s29, [r4, #12+16]
    vldr.f32 s30, [r4, #12+32]
    vldr.f32 s31, [r4, #12+48]
    /* q7 = { b[0][3], b[1][3], b[2][3], b[3][3] } */
    
    /* Row c[0][:] */
    vmul.f32 q0, q12, q4
    /* q0 ← {a[0][0] * b[0][0], a[0][1] * b[1][0], a[0][2] * b[2][0], a[0][3] * b[3][0] } */
    vmul.f32 q1, q12, q5
    /* q1 ← {a[0][0] * b[0][1], a[0][1] * b[1][1], a[0][2] * b[2][1], a[0][3] * b[3][1] } */
    vmul.f32 q2, q12, q6
    /* q2 ← {a[0][0] * b[0][2], a[0][1] * b[1][2], a[0][2] * b[2][2], a[0][3] * b[3][2] } */
    vmul.f32 q3, q12, q7
    /* q3 ← {a[0][0] * b[0][3], a[0][1] * b[1][3], a[0][2] * b[2][3], a[0][3] * b[3][3] } */

    /* Sum */
    vpadd.f32 d0, d0, d1
    vpadd.f32 d2, d2, d3
    vpadd.f32 d4, d4, d5
    vpadd.f32 d6, d6, d7
    vpadd.f32 d0, d0, d0
    vpadd.f32 d2, d2, d2
    vpadd.f32 d4, d4, d4
    vpadd.f32 d6, d6, d6

    /* That block above is the same as
       the following but reordered to
       avoid stalls
  
    vpadd.f32 d0, d0, d1
    ## d0 ← { a[0][0] * b[0][0] + a[0][1] * b[1][0],
    ##        a[0][2] * b[2][0] + a[0][3] * b[3][0] }
    vpadd.f32 d0, d0, d0
    ## d0 ← { a[0][0] * b[0][0] + a[0][1] * b[1][0] + a[0][2] * b[2][0] + a[0][3] * b[3][0],
    ##        a[0][0] * b[0][0] + a[0][1] * b[1][0] + a[0][2] * b[2][0] + a[0][3] * b[3][0] }

    vpadd.f32 d2, d2, d3
    ## d2 ← { a[0][0] * b[0][1] + a[0][1] * b[1][1],
    ##        a[0][2] * b[2][1] + a[0][3] * b[3][1] }
    vpadd.f32 d2, d2, d2
    ## d2 ← { a[0][0] * b[0][1] + a[0][1] * b[1][1] + a[0][2] * b[2][1] + a[0][3] * b[3][1],
    ##        a[0][0] * b[0][1] + a[0][1] * b[1][1] + a[0][2] * b[2][1] + a[0][3] * b[3][1] }

    vpadd.f32 d4, d4, d5
    ## d4 ← { a[0][0] * b[0][2] + a[0][1] * b[1][2],
    ##        a[0][2] * b[2][2] + a[0][3] * b[3][2] }
    vpadd.f32 d4, d4, d4
    ## d4 ← { a[0][0] * b[0][2] + a[0][1] * b[1][2] + a[0][2] * b[2][2] + a[0][3] * b[3][2],
    ##        a[0][0] * b[0][2] + a[0][1] * b[1][2] + a[0][2] * b[2][2] + a[0][3] * b[3][2] }

    vpadd.f32 d6, d6, d7
    ## d6 ← { a[0][0] * b[0][3] + a[0][1] * b[1][3],
    ##        a[0][2] * b[2][3] + a[0][3] * b[3][3] }
    vpadd.f32 d6, d6, d6
    ## d6 ← { a[0][0] * b[0][3] + a[0][1] * b[1][3] + a[0][2] * b[2][3] + a[0][3] * b[3][3],
    ##        a[0][0] * b[0][3] + a[0][1] * b[1][3] + a[0][2] * b[2][3] + a[0][3] * b[3][3] }

    */

    /* Now build q0 so we can store it directly */
    vmov.f32 s1, s4
    vmov.f32 s2, s8
    vmov.f32 s3, s12

    /* Store into c[0][:] */
    vst1.f32 {q0}, [r2]!

    /* Row c[1][:] */
    vmul.f32 q0, q13, q4
    vmul.f32 q1, q13, q5
    vmul.f32 q2, q13, q6
    vmul.f32 q3, q13, q7

    vpadd.f32 d0, d0, d1
    vpadd.f32 d2, d2, d3
    vpadd.f32 d4, d4, d5
    vpadd.f32 d6, d6, d7
    vpadd.f32 d0, d0, d0
    vpadd.f32 d2, d2, d2
    vpadd.f32 d4, d4, d4
    vpadd.f32 d6, d6, d6

    vmov.f32 s1, s4
    vmov.f32 s2, s8
    vmov.f32 s3, s12

    vst1.f32 {q0}, [r2]!

    /* Row c[2][:] */
    vmul.f32 q0, q14, q4
    vmul.f32 q1, q14, q5
    vmul.f32 q2, q14, q6
    vmul.f32 q3, q14, q7

    vpadd.f32 d0, d0, d1
    vpadd.f32 d2, d2, d3
    vpadd.f32 d4, d4, d5
    vpadd.f32 d6, d6, d7
    vpadd.f32 d0, d0, d0
    vpadd.f32 d2, d2, d2
    vpadd.f32 d4, d4, d4
    vpadd.f32 d6, d6, d6

    vmov.f32 s1, s4
    vmov.f32 s2, s8
    vmov.f32 s3, s12

    vst1.f32 {q0}, [r2]!

    /* Row c[3][:] */
    vmul.f32 q0, q15, q4
    vmul.f32 q1, q15, q5
    vmul.f32 q2, q15, q6
    vmul.f32 q3, q15, q7

    vpadd.f32 d0, d0, d1
    vpadd.f32 d2, d2, d3
    vpadd.f32 d4, d4, d5
    vpadd.f32 d6, d6, d7
    vpadd.f32 d0, d0, d0
    vpadd.f32 d2, d2, d2
    vpadd.f32 d4, d4, d4
    vpadd.f32 d6, d6, d6

    vmov.f32 s1, s4
    vmov.f32 s2, s8
    vmov.f32 s3, s12

    vst1.f32 {q0}, [r2]!

    vpop {q4-q7}                /* Restore preserved floating registers */
    pop {r4, r5, r6, r7, r8, lr}  /* Restore integer registers */
    bx lr /* Leave function */