OpenI
/
MegEngine
Not watched
1
0
Fork 0
Code
Releases 31 Wiki Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: c2e9860feb
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'c2e9860feb'
${ noResults }
MegEngine/dnn/src/aarch64/matrix_mul/int8_dot/kernel_8x12x4.h

1547 lines
58 kB
Raw Blame History 

  1. #if MGB_ENABLE_DOT

  2. #include "src/aarch64/matrix_mul/asm/common.h"

  3. #include "src/arm_common/simd_macro/marm_neon.h"

  4. 

  5. namespace megdnn {

  6. namespace aarch64 {

  7. namespace matmul_8x12x4 {

  8. 

  9. // Overview of register layout:

  10. //

  11. // A 12x4 cell of Rhs is stored in 8bit in q2-q4.

  12. // A 8x4x2 cell of Lhs is stored in 8bit in q0-q1,q5-q6

  13. // A 8x12 block of accumulators is stored in 8bit in q8--q31.

  14. //

  15. //                            +--------+--------+--------+

  16. //                            |v2[0-16]|v3[0-16]|v4[0-16]|

  17. //                       Rhs  +--------+--------+--------+

  18. //

  19. //                            |        |        |        |

  20. //

  21. //    Lhs                     |        |        |        |

  22. //

  23. //  +-------+-------+ - - - - +--------+--------+--------+

  24. //  |v0[0-4]|v5[0-4]|         | v8[0-4]|v16[0-4]|v24[0-4]|

  25. //  |v0[0-4]|v5[0-4]|         | v9[0-4]|v17[0-4]|v25[0-4]|

  26. //  |v0[0-4]|v5[0-4]|         |v10[0-4]|v18[0-4]|v26[0-4]|

  27. //  |v0[0-4]|v5[0-4]|         |v11[0-4]|v19[0-4]|v27[0-4]|

  28. //  |v1[0-4]|v6[0-4]|         |v12[0-4]|v20[0-4]|v28[0-4]|

  29. //  |v1[0-4]|v6[0-4]|         |v13[0-4]|v21[0-4]|v29[0-4]|

  30. //  |v1[0-4]|v6[0-4]|         |v14[0-4]|v22[0-4]|v30[0-4]|

  31. //  |v1[0-4]|v6[0-4]|         |v15[0-4]|v23[0-4]|v31[0-4]|

  32. //  +-------+-------+ - - - - +--------+--------+--------+

  33. //

  34. //                            Accumulator

  35. 

  36. /**

  37.  * \note The performance of reorder instruction and use prefetch is almost the

  38.  * same, I test in kirin980 with small and big core, here i just keep both the

  39.  * implementation.

  40.  */

  41. 

  42. #if 1

  43. MEGDNN_ATTRIBUTE_TARGET("dotprod")

  44. static void kern_8x12(

  45.         const int8_t* packA, const int8_t* packB, int K, int32_t* output, int LDC,

  46.         bool is_first_k) {

  47.     K /= 4;

  48.     const int8_t* a_ptr = packA;

  49.     const int8_t* b_ptr = packB;

  50.     // Fix up for odd lengths - set a flag if K is odd, but make

  51.     // sure we round up the iteration count.

  52.     int oddk = (K & 1);

  53.     int k = ((K + 1) / 2) - 1;

  54.     int32x4_t a0;

  55.     int32x4_t a1;

  56.     int32x4_t b0;

  57.     int32x4_t b1;

  58.     int32x4_t b2;

  59.     int32x4_t a0a;

  60.     int32x4_t a1a;

  61.     LDC = LDC * sizeof(int32_t);

  62. 

  63.     int32_t* outptr0 = output;

  64.     int32_t* outptr1;

  65.     int32_t* outptr2;

  66.     int32_t* outptr3;

  67.     int32_t* outptr4;

  68.     int32_t* outptr5;

  69.     int32_t* outptr6;

  70.     int32_t* outptr7;

  71. 

  72.     asm volatile (

  73.             // load accumulator C

  74.             "add %[outptr1], %[outptr0], %x[LDC]\n"

  75.             "add %[outptr2], %[outptr1], %x[LDC]\n"

  76.             "add %[outptr3], %[outptr2], %x[LDC]\n"

  77.             "add %[outptr4], %[outptr3], %x[LDC]\n"

  78.             "add %[outptr5], %[outptr4], %x[LDC]\n"

  79.             "add %[outptr6], %[outptr5], %x[LDC]\n"

  80.             "add %[outptr7], %[outptr6], %x[LDC]\n"

  81.             "cmp %w[is_first_k], #1\n"

  82.             "beq 5f\n"

  83.             // we can not use ld1, as it can not encode {v8, v16, v24}

  84.             "ldp q8, q16, [%[outptr0]]\n"

  85.             "ldr q24, [%[outptr0], #32]\n"

  86.             "ldp q9, q17, [%[outptr1]]\n"

  87.             "ldr q25, [%[outptr1], #32]\n"

  88.             "ldp q10, q18, [%[outptr2]]\n"

  89.             "ldr q26, [%[outptr2], #32]\n"

  90.             "ldp q11, q19, [%[outptr3]]\n"

  91.             "ldr q27, [%[outptr3], #32]\n"

  92.             "ldp q12, q20, [%[outptr4]]\n"

  93.             "ldr q28, [%[outptr4], #32]\n"

  94.             "ldp q13, q21, [%[outptr5]]\n"

  95.             "ldr q29, [%[outptr5], #32]\n"

  96.             "ldp q14, q22, [%[outptr6]]\n"

  97.             "ldr q30, [%[outptr6], #32]\n"

  98.             "ldp q15, q23, [%[outptr7]]\n"

  99.             "ldr q31, [%[outptr7], #32]\n"

  100.             "b 6f\n"

  101. 

  102.             "5:\n"

  103.             "eor v8.16b,  v8.16b,  v8.16b\n"

  104.             "eor v9.16b,  v9.16b,  v9.16b\n"

  105.             "eor v10.16b, v10.16b, v10.16b\n"

  106.             "eor v11.16b, v11.16b, v11.16b\n"

  107.             "eor v12.16b, v12.16b, v12.16b\n"

  108.             "eor v13.16b, v13.16b, v13.16b\n"

  109.             "eor v14.16b, v14.16b, v14.16b\n"

  110.             "eor v15.16b, v15.16b, v15.16b\n"

  111. 

  112.             "eor v16.16b, v16.16b, v16.16b\n"

  113.             "eor v17.16b, v17.16b, v17.16b\n"

  114.             "eor v18.16b, v18.16b, v18.16b\n"

  115.             "eor v19.16b, v19.16b, v19.16b\n"

  116.             "eor v20.16b, v20.16b, v20.16b\n"

  117.             "eor v21.16b, v21.16b, v21.16b\n"

  118.             "eor v22.16b, v22.16b, v22.16b\n"

  119.             "eor v23.16b, v23.16b, v23.16b\n"

  120. 

  121.             "eor v24.16b, v24.16b, v24.16b\n"

  122.             "eor v25.16b, v25.16b, v25.16b\n"

  123.             "eor v26.16b, v26.16b, v26.16b\n"

  124.             "eor v27.16b, v27.16b, v27.16b\n"

  125.             "eor v28.16b, v28.16b, v28.16b\n"

  126.             "eor v29.16b, v29.16b, v29.16b\n"

  127.             "eor v30.16b, v30.16b, v30.16b\n"

  128.             "eor v31.16b, v31.16b, v31.16b\n"

  129. 

  130.             "6: \n"

  131.             // Initialize result registers, load initial operands, prime prefetches.

  132.             "ldr  %q[a0], [%[a_ptr]]\n"

  133.             "ldr  %q[b0], [%[b_ptr]]\n"

  134.             "ldr  %q[a1], [%[a_ptr], #16]\n"

  135.             "ldr  %q[b1], [%[b_ptr], #16]\n"

  136.             ASM_PREFETCH("[%[b_ptr], #64]")

  137.             ASM_PREFETCH("[%[a_ptr], #64]")

  138.             ASM_PREFETCH("[%[b_ptr], #128]")

  139.             ASM_PREFETCH("[%[a_ptr], #128]")

  140.             ASM_PREFETCH("[%[b_ptr], #192]")

  141.             ASM_PREFETCH("[%[b_ptr], #256]")

  142.             ASM_PREFETCH("[%[a_ptr], #192]")

  143.             ASM_PREFETCH("[%[b_ptr], #320]")

  144.             ASM_PREFETCH("[%[a_ptr], #256]")

  145.             ASM_PREFETCH("[%[b_ptr], #384]")

  146. 

  147.             // Skip loop if we are doing zero iterations of it.

  148.             "cbz  %w[k], 4f\n"

  149. 

  150.             // Loop proper

  151.             "1:\n"

  152.             "sdot  v8.4s , %[b0].16b, %[a0].4b[0]\n"

  153.             "sdot    v9.4s , %[b0].16b, %[a0].4b[1]\n"

  154. 

  155.             "ldr  %q[b2], [%[b_ptr], #32]\n"

  156.             "sdot  v10.4s, %[b0].16b, %[a0].4b[2]\n"

  157.             "sdot  v11.4s, %[b0].16b, %[a0].4b[3]\n"

  158.             "ldr  %q[a0a], [%[a_ptr], #32]\n"

  159.             "sdot   v12.4s, %[b0].16b, %[a1].4b[0]\n"

  160.             "sdot  v13.4s, %[b0].16b, %[a1].4b[1]\n"

  161.             "ldr  %q[a1a], [%[a_ptr], #48]\n"

  162.             "sdot  v14.4s, %[b0].16b, %[a1].4b[2]\n"

  163.             "sdot  v15.4s, %[b0].16b, %[a1].4b[3]\n"

  164.             "ldr  %q[b0], [%[b_ptr], #48]\n"

  165. 

  166.             "sdot  v16.4s, %[b1].16b, %[a0].4b[0]\n"

  167.             "sdot  v17.4s, %[b1].16b, %[a0].4b[1]\n"

  168.             ASM_PREFETCH("[%[a_ptr], #320]")

  169.             "sdot  v18.4s, %[b1].16b, %[a0].4b[2]\n"

  170.             "sdot  v19.4s, %[b1].16b, %[a0].4b[3]\n"

  171.             "sdot  v20.4s, %[b1].16b, %[a1].4b[0]\n"

  172.             "sdot  v21.4s, %[b1].16b, %[a1].4b[1]\n"

  173.             "sdot  v22.4s, %[b1].16b, %[a1].4b[2]\n"

  174.             "sdot  v23.4s, %[b1].16b, %[a1].4b[3]\n"

  175.             "ldr  %q[b1], [%[b_ptr], #64]\n"

  176. 

  177.             "sdot  v24.4s, %[b2].16b, %[a0].4b[0]\n"

  178.             "sdot  v25.4s, %[b2].16b, %[a0].4b[1]\n"

  179.             ASM_PREFETCH("[%[b_ptr], #448]")

  180.             "sdot  v26.4s, %[b2].16b, %[a0].4b[2]\n"

  181.             "sdot  v27.4s, %[b2].16b, %[a0].4b[3]\n"

  182.             "sdot  v28.4s, %[b2].16b, %[a1].4b[0]\n"

  183.             "sdot  v29.4s, %[b2].16b, %[a1].4b[1]\n"

  184.             "sdot  v30.4s, %[b2].16b, %[a1].4b[2]\n"

  185.             "sdot  v31.4s, %[b2].16b, %[a1].4b[3]\n"

  186.             "ldr  %q[b2], [%[b_ptr], #80]\n"

  187. 

  188.             "sdot  v8.4s , %[b0].16b, %[a0a].4b[0]\n"

  189.             "sdot  v9.4s , %[b0].16b, %[a0a].4b[1]\n"

  190.             "ldr  %q[a0], [%[a_ptr], #64]\n"

  191.             "sdot  v10.4s, %[b0].16b, %[a0a].4b[2]\n"

  192.             "sdot  v11.4s, %[b0].16b, %[a0a].4b[3]\n"

  193.             "sdot   v12.4s, %[b0].16b, %[a1a].4b[0]\n"

  194.             "ldr  %q[a1], [%[a_ptr], #80]\n"

  195.             "sdot   v13.4s, %[b0].16b, %[a1a].4b[1]\n"

  196.             "sdot  v14.4s, %[b0].16b, %[a1a].4b[2]\n"

  197.             "sdot  v15.4s, %[b0].16b, %[a1a].4b[3]\n"

  198.             "ldr  %q[b0], [%[b_ptr], #96]\n"

  199. 

  200.             "sdot  v16.4s, %[b1].16b, %[a0a].4b[0]\n"

  201.             "sdot  v17.4s, %[b1].16b, %[a0a].4b[1]\n"

  202.             ASM_PREFETCH("[%[b_ptr], #512]")

  203.             "sdot  v18.4s, %[b1].16b, %[a0a].4b[2]\n"

  204.             "sdot  v19.4s, %[b1].16b, %[a0a].4b[3]\n"

  205.             "sdot  v20.4s, %[b1].16b, %[a1a].4b[0]\n"

  206.             "sdot  v21.4s, %[b1].16b, %[a1a].4b[1]\n"

  207.             "sdot  v22.4s, %[b1].16b, %[a1a].4b[2]\n"

  208.             "sdot  v23.4s, %[b1].16b, %[a1a].4b[3]\n"

  209.             "ldr  %q[b1], [%[b_ptr], #112]\n"

  210. 

  211.             "sdot  v24.4s, %[b2].16b, %[a0a].4b[0]\n"

  212.             "sdot  v25.4s, %[b2].16b, %[a0a].4b[1]\n"

  213.             "add  %[a_ptr], %[a_ptr], #64\n"

  214.             "sdot  v26.4s, %[b2].16b, %[a0a].4b[2]\n"

  215.             "sdot  v27.4s, %[b2].16b, %[a0a].4b[3]\n"

  216.             "add  %[b_ptr], %[b_ptr], #96\n"

  217.             "sdot  v28.4s, %[b2].16b, %[a1a].4b[0]\n"

  218.             "sdot  v29.4s, %[b2].16b, %[a1a].4b[1]\n"

  219.             "subs  %w[k], %w[k], #1\n"

  220.             "sdot  v30.4s, %[b2].16b, %[a1a].4b[2]\n"

  221.             "sdot  v31.4s, %[b2].16b, %[a1a].4b[3]\n"

  222.             "bne  1b\n"

  223. 

  224.             // Target to use when K is 1 or 2 (i.e. zero iterations of main loop)

  225.             "4:\n"

  226. 

  227.             // Branch to alternative tail for odd K

  228.             "cbnz  %w[oddk], 2f\n"

  229. 

  230.             // Detached final iteration (even K)

  231.             "sdot  v8.4s , %[b0].16b, %[a0].4b[0]\n"

  232.             "sdot   v9.4s , %[b0].16b, %[a0].4b[1]\n"

  233.             "ldr  %q[b2], [%[b_ptr], #32]\n"

  234.             "sdot  v10.4s, %[b0].16b, %[a0].4b[2]\n"

  235.             "sdot  v11.4s, %[b0].16b, %[a0].4b[3]\n"

  236.             "ldr  %q[a0a], [%[a_ptr], #32]\n"

  237.             "sdot   v12.4s, %[b0].16b, %[a1].4b[0]\n"

  238.             "sdot   v13.4s, %[b0].16b, %[a1].4b[1]\n"

  239.             "ldr  %q[a1a], [%[a_ptr], #48]\n"

  240.             "sdot  v14.4s, %[b0].16b, %[a1].4b[2]\n"

  241.             "sdot  v15.4s, %[b0].16b, %[a1].4b[3]\n"

  242.             "ldr  %q[b0], [%[b_ptr], #48]\n"

  243. 

  244.             "sdot  v16.4s, %[b1].16b, %[a0].4b[0]\n"

  245.             "sdot  v17.4s, %[b1].16b, %[a0].4b[1]\n"

  246.             "sdot  v18.4s, %[b1].16b, %[a0].4b[2]\n"

  247.             "sdot  v19.4s, %[b1].16b, %[a0].4b[3]\n"

  248.             "sdot  v20.4s, %[b1].16b, %[a1].4b[0]\n"

  249.             "sdot  v21.4s, %[b1].16b, %[a1].4b[1]\n"

  250.             "sdot  v22.4s, %[b1].16b, %[a1].4b[2]\n"

  251.             "sdot  v23.4s, %[b1].16b, %[a1].4b[3]\n"

  252.             "ldr  %q[b1], [%[b_ptr], #64]\n"

  253. 

  254.             "sdot  v24.4s, %[b2].16b, %[a0].4b[0]\n"

  255.             "sdot  v25.4s, %[b2].16b, %[a0].4b[1]\n"

  256.             "add  %[a_ptr], %[a_ptr], #64\n"

  257.             "sdot  v26.4s, %[b2].16b, %[a0].4b[2]\n"

  258.             "sdot  v27.4s, %[b2].16b, %[a0].4b[3]\n"

  259.             "sdot  v28.4s, %[b2].16b, %[a1].4b[0]\n"

  260.             "sdot  v29.4s, %[b2].16b, %[a1].4b[1]\n"

  261.             "sdot  v30.4s, %[b2].16b, %[a1].4b[2]\n"

  262.             "sdot  v31.4s, %[b2].16b, %[a1].4b[3]\n"

  263.             "ldr  %q[b2], [%[b_ptr], #80]\n"

  264. 

  265.             "sdot  v8.4s , %[b0].16b, %[a0a].4b[0]\n"

  266. 

  267.             "sdot  v16.4s, %[b1].16b, %[a0a].4b[0]\n"

  268.             "add  %[b_ptr], %[b_ptr], #96\n"

  269.             "sdot   v9.4s , %[b0].16b, %[a0a].4b[1]\n"

  270.             "str  q8, [%[outptr0], #0]\n"

  271.             "sdot  v17.4s, %[b1].16b, %[a0a].4b[1]\n"

  272.             "str  q16, [%[outptr0], #16]\n"

  273.             "sdot  v24.4s, %[b2].16b, %[a0a].4b[0]\n"

  274.             "str  q24, [%[outptr0], #32]\n"

  275. 

  276.             "sdot  v25.4s, %[b2].16b, %[a0a].4b[1]\n"

  277.             "str  q9, [%[outptr1], #0]\n"

  278.             "sdot  v10.4s, %[b0].16b, %[a0a].4b[2]\n"

  279.             "str  q17, [%[outptr1], #16]\n"

  280.             "sdot  v18.4s, %[b1].16b, %[a0a].4b[2]\n"

  281.             "str  q25, [%[outptr1], #32]\n"

  282.             "sdot  v26.4s, %[b2].16b, %[a0a].4b[2]\n"

  283.             "str  q10, [%[outptr2], #0]\n"

  284. 

  285.             "sdot  v11.4s, %[b0].16b, %[a0a].4b[3]\n"

  286.             "str  q18, [%[outptr2], #16]\n"

  287.             "sdot  v19.4s, %[b1].16b, %[a0a].4b[3]\n"

  288.             "str  q26, [%[outptr2], #32]\n"

  289.             "sdot  v27.4s, %[b2].16b, %[a0a].4b[3]\n"

  290.             "str  q11, [%[outptr3], #0]\n"

  291. 

  292.             "sdot   v12.4s, %[b0].16b, %[a1a].4b[0]\n"

  293.             "str  q19, [%[outptr3], #16]\n"

  294.             "sdot  v20.4s, %[b1].16b, %[a1a].4b[0]\n"

  295.             "str  q27, [%[outptr3], #32]\n"

  296.             "sdot  v28.4s, %[b2].16b, %[a1a].4b[0]\n"

  297.             "str  q12, [%[outptr4], #0]\n"

  298. 

  299.             "sdot   v13.4s, %[b0].16b, %[a1a].4b[1]\n"

  300.             "str  q20, [%[outptr4], #16]\n"

  301.             "sdot  v21.4s, %[b1].16b, %[a1a].4b[1]\n"

  302.             "str  q28, [%[outptr4], #32]\n"

  303.             "sdot  v29.4s, %[b2].16b, %[a1a].4b[1]\n"

  304.             "str  q13, [%[outptr5], #0]\n"

  305. 

  306.             "sdot  v14.4s, %[b0].16b, %[a1a].4b[2]\n"

  307.             "str  q21, [%[outptr5], #16]\n"

  308.             "sdot  v22.4s, %[b1].16b, %[a1a].4b[2]\n"

  309.             "str  q29, [%[outptr5], #32]\n"

  310.             "sdot  v30.4s, %[b2].16b, %[a1a].4b[2]\n"

  311.             "str  q14, [%[outptr6], #0]\n"

  312. 

  313.             "sdot  v15.4s, %[b0].16b, %[a1a].4b[3]\n"

  314.             "str  q22, [%[outptr6], #16]\n"

  315.             "sdot  v23.4s, %[b1].16b, %[a1a].4b[3]\n"

  316.             "str  q30, [%[outptr6], #32]\n"

  317.             "sdot  v31.4s, %[b2].16b, %[a1a].4b[3]\n"

  318.             "str  q15, [%[outptr7], #0]\n"

  319. 

  320.             "b  3f\n"

  321. 

  322.             // Detached final iteration (odd K)

  323.             "2:\n"

  324.             "sdot  v8.4s , %[b0].16b, %[a0].4b[0]\n"

  325.             "ldr  %q[b2], [%[b_ptr], #32]\n"

  326.             "sdot  v16.4s, %[b1].16b, %[a0].4b[0]\n"

  327.             "sdot   v9.4s , %[b0].16b, %[a0].4b[1]\n"

  328.             "str  q8, [%[outptr0], #0]\n"

  329.             "sdot  v17.4s, %[b1].16b, %[a0].4b[1]\n"

  330.             "str  q16, [%[outptr0], #16]\n"

  331.             "sdot  v24.4s, %[b2].16b, %[a0].4b[0]\n"

  332.             "add  %[b_ptr], %[b_ptr], #48\n"

  333.             "add  %[a_ptr], %[a_ptr], #32\n"

  334.             "str  q24, [%[outptr0], #32]\n"

  335.             "sdot  v25.4s, %[b2].16b, %[a0].4b[1]\n"

  336.             "str  q9, [%[outptr1], #0]\n"

  337. 

  338.             "sdot  v10.4s, %[b0].16b, %[a0].4b[2]\n"

  339.             "str  q17, [%[outptr1], #16]\n"

  340.             "sdot  v18.4s, %[b1].16b, %[a0].4b[2]\n"

  341.             "str  q25, [%[outptr1], #32]\n"

  342.             "sdot  v26.4s, %[b2].16b, %[a0].4b[2]\n"

  343.             "str  q10, [%[outptr2], #0]\n"

  344. 

  345.             "sdot  v11.4s, %[b0].16b, %[a0].4b[3]\n"

  346.             "str  q18, [%[outptr2], #16]\n"

  347.             "sdot  v19.4s, %[b1].16b, %[a0].4b[3]\n"

  348.             "str  q26, [%[outptr2], #32]\n"

  349.             "sdot  v27.4s, %[b2].16b, %[a0].4b[3]\n"

  350.             "str  q11, [%[outptr3], #0]\n"

  351. 

  352.             "sdot   v12.4s, %[b0].16b, %[a1].4b[0]\n"

  353.             "str  q19, [%[outptr3], #16]\n"

  354.             "sdot  v20.4s, %[b1].16b, %[a1].4b[0]\n"

  355.             "str  q27, [%[outptr3], #32]\n"

  356.             "sdot  v28.4s, %[b2].16b, %[a1].4b[0]\n"

  357.             "str  q12, [%[outptr4], #0]\n"

  358. 

  359.             "sdot   v13.4s, %[b0].16b, %[a1].4b[1]\n"

  360.             "str  q20, [%[outptr4], #16]\n"

  361.             "sdot  v21.4s, %[b1].16b, %[a1].4b[1]\n"

  362.             "str  q28, [%[outptr4], #32]\n"

  363.             "sdot  v29.4s, %[b2].16b, %[a1].4b[1]\n"

  364.             "str  q13, [%[outptr5], #0]\n"

  365. 

  366.             "sdot  v14.4s, %[b0].16b, %[a1].4b[2]\n"

  367.             "str  q21, [%[outptr5], #16]\n"

  368.             "sdot  v22.4s, %[b1].16b, %[a1].4b[2]\n"

  369.             "str  q29, [%[outptr5], #32]\n"

  370.             "sdot  v30.4s, %[b2].16b, %[a1].4b[2]\n"

  371.             "str  q14, [%[outptr6], #0]\n"

  372. 

  373.             "sdot  v15.4s, %[b0].16b, %[a1].4b[3]\n"

  374.             "str  q22, [%[outptr6], #16]\n"

  375.             "sdot  v23.4s, %[b1].16b, %[a1].4b[3]\n"

  376.             "str  q30, [%[outptr6], #32]\n"

  377.             "sdot  v31.4s, %[b2].16b, %[a1].4b[3]\n"

  378.             "str  q15, [%[outptr7], #0]\n"

  379. 

  380. 

  381.             // Common tail

  382.             "3:\n"

  383.             "str  q23, [%[outptr7], #16]\n"

  384.             "str  q31, [%[outptr7], #32]\n"

  385.             :

  386.               [a_ptr] "+r" (a_ptr), [b_ptr] "+r" (b_ptr),[oddk] "+r" (oddk),

  387.               [is_first_k] "+r" (is_first_k), [k] "+r" (k), [LDC] "+r" (LDC),

  388.               [a0] "=w" (a0), [a1] "=w" (a1), [a0a] "=w" (a0a), [a1a] "=w" (a1a),

  389.               [b0] "=w" (b0), [b1] "=w" (b1), [b2] "=w" (b2),

  390.               [outptr0] "+r"(outptr0), [outptr1] "=r"(outptr1),

  391.               [outptr2] "=r"(outptr2), [outptr3] "=r"(outptr3),

  392.               [outptr4] "=r"(outptr4), [outptr5] "=r"(outptr5),

  393.               [outptr6] "=r"(outptr6), [outptr7] "=r"(outptr7)

  394.             :

  395.             : "v8", "v9", "v10", "v11", "v12", "v13", "v14",

  396.               "v15", "v16", "v17", "v18", "v19", "v20", "v21", "v22", "v23",

  397.               "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31", "cc",

  398.               "memory"

  399.             );

  400. }

  401. #else

  402. MEGDNN_ATTRIBUTE_TARGET("dotprod")

  403. static void kern_8x12(

  404.         const int8_t* packA, const int8_t* packB, int K, int32_t* output, int LDC,

  405.         bool is_first_k) {

  406.     K /= 4;

  407.     const int8_t* a_ptr = packA;

  408.     const int8_t* b_ptr = packB;

  409.     // Fix up for odd lengths - set a flag if K is odd, but make

  410.     // sure we round up the iteration count.

  411.     int oddk = (K & 1);

  412.     int k = K / 2;

  413. 

  414.     int32x4_t a0;

  415.     int32x4_t a1;

  416.     int32x4_t b0;

  417.     int32x4_t b1;

  418.     int32x4_t b2;

  419.     int32x4_t a0a;

  420.     int32x4_t a1a;

  421.     LDC = LDC * sizeof(int32_t);

  422. 

  423.     int32_t* outptr0 = output;

  424.     int32_t* outptr1;

  425.     int32_t* outptr2;

  426.     int32_t* outptr3;

  427.     int32_t* outptr4;

  428.     int32_t* outptr5;

  429.     int32_t* outptr6;

  430.     int32_t* outptr7;

  431. 

  432.     asm volatile(

  433.             // load accumulator C

  434.             "add %[outptr1], %[outptr0], %x[LDC]\n"

  435.             "add %[outptr2], %[outptr1], %x[LDC]\n"

  436.             "add %[outptr3], %[outptr2], %x[LDC]\n"

  437.             "add %[outptr4], %[outptr3], %x[LDC]\n"

  438.             "add %[outptr5], %[outptr4], %x[LDC]\n"

  439.             "add %[outptr6], %[outptr5], %x[LDC]\n"

  440.             "add %[outptr7], %[outptr6], %x[LDC]\n"

  441.             "cmp %w[is_first_k], #1\n"

  442.             "beq 1f\n"

  443.             // we can not use ld1, as it can not encode {v8, v16, v24}

  444.             "ldp q8, q16, [%[outptr0]]\n"

  445.             "ldr q24, [%[outptr0], #32]\n"

  446.             "ldp q9, q17, [%[outptr1]]\n"

  447.             "ldr q25, [%[outptr1], #32]\n"

  448.             "ldp q10, q18, [%[outptr2]]\n"

  449.             "ldr q26, [%[outptr2], #32]\n"

  450.             "ldp q11, q19, [%[outptr3]]\n"

  451.             "ldr q27, [%[outptr3], #32]\n"

  452.             "ldp q12, q20, [%[outptr4]]\n"

  453.             "ldr q28, [%[outptr4], #32]\n"

  454.             "ldp q13, q21, [%[outptr5]]\n"

  455.             "ldr q29, [%[outptr5], #32]\n"

  456.             "ldp q14, q22, [%[outptr6]]\n"

  457.             "ldr q30, [%[outptr6], #32]\n"

  458.             "ldp q15, q23, [%[outptr7]]\n"

  459.             "ldr q31, [%[outptr7], #32]\n"

  460.             "b 2f\n"

  461. 

  462.             "1:\n"

  463.             "eor v8.16b,  v8.16b,  v8.16b\n"

  464.             "eor v9.16b,  v9.16b,  v9.16b\n"

  465.             "eor v10.16b, v10.16b, v10.16b\n"

  466.             "eor v11.16b, v11.16b, v11.16b\n"

  467.             "eor v12.16b, v12.16b, v12.16b\n"

  468.             "eor v13.16b, v13.16b, v13.16b\n"

  469.             "eor v14.16b, v14.16b, v14.16b\n"

  470.             "eor v15.16b, v15.16b, v15.16b\n"

  471. 

  472.             "eor v16.16b, v16.16b, v16.16b\n"

  473.             "eor v17.16b, v17.16b, v17.16b\n"

  474.             "eor v18.16b, v18.16b, v18.16b\n"

  475.             "eor v19.16b, v19.16b, v19.16b\n"

  476.             "eor v20.16b, v20.16b, v20.16b\n"

  477.             "eor v21.16b, v21.16b, v21.16b\n"

  478.             "eor v22.16b, v22.16b, v22.16b\n"

  479.             "eor v23.16b, v23.16b, v23.16b\n"

  480. 

  481.             "eor v24.16b, v24.16b, v24.16b\n"

  482.             "eor v25.16b, v25.16b, v25.16b\n"

  483.             "eor v26.16b, v26.16b, v26.16b\n"

  484.             "eor v27.16b, v27.16b, v27.16b\n"

  485.             "eor v28.16b, v28.16b, v28.16b\n"

  486.             "eor v29.16b, v29.16b, v29.16b\n"

  487.             "eor v30.16b, v30.16b, v30.16b\n"

  488.             "eor v31.16b, v31.16b, v31.16b\n"

  489. 

  490.             "2: \n"

  491.             "cbz  %w[oddk], 3f\n"

  492.             // parse the oddk

  493.             "ldr  %q[a0], [%[a_ptr]], #16\n"

  494.             "ldr  %q[a1], [%[a_ptr]], #16\n"

  495.             "ldr  %q[b0], [%[b_ptr]], #16\n"

  496.             "ldr  %q[b1], [%[b_ptr]], #16\n"

  497.             "ldr  %q[b2], [%[b_ptr]], #16\n"

  498.             "sdot  v8.4s, %[b0].16b, %[a0].4b[0]\n"

  499.             "sdot  v9.4s, %[b0].16b, %[a0].4b[1]\n"

  500.             "sdot  v10.4s, %[b0].16b, %[a0].4b[2]\n"

  501.             "sdot  v11.4s, %[b0].16b, %[a0].4b[3]\n"

  502.             "sdot  v12.4s, %[b0].16b, %[a1].4b[0]\n"

  503.             "sdot  v13.4s, %[b0].16b, %[a1].4b[1]\n"

  504.             "sdot  v14.4s, %[b0].16b, %[a1].4b[2]\n"

  505.             "sdot  v15.4s, %[b0].16b, %[a1].4b[3]\n"

  506.             "sdot  v16.4s, %[b1].16b, %[a0].4b[0]\n"

  507.             "sdot  v17.4s, %[b1].16b, %[a0].4b[1]\n"

  508.             "sdot  v18.4s, %[b1].16b, %[a0].4b[2]\n"

  509.             "sdot  v19.4s, %[b1].16b, %[a0].4b[3]\n"

  510.             "sdot  v20.4s, %[b1].16b, %[a1].4b[0]\n"

  511.             "sdot  v21.4s, %[b1].16b, %[a1].4b[1]\n"

  512.             "sdot  v22.4s, %[b1].16b, %[a1].4b[2]\n"

  513.             "sdot  v23.4s, %[b1].16b, %[a1].4b[3]\n"

  514.             "sdot  v24.4s, %[b2].16b, %[a0].4b[0]\n"

  515.             "sdot  v25.4s, %[b2].16b, %[a0].4b[1]\n"

  516.             "sdot  v26.4s, %[b2].16b, %[a0].4b[2]\n"

  517.             "sdot  v27.4s, %[b2].16b, %[a0].4b[3]\n"

  518.             "sdot  v28.4s, %[b2].16b, %[a1].4b[0]\n"

  519.             "sdot  v29.4s, %[b2].16b, %[a1].4b[1]\n"

  520.             "sdot  v30.4s, %[b2].16b, %[a1].4b[2]\n"

  521.             "sdot  v31.4s, %[b2].16b, %[a1].4b[3]\n"

  522. 

  523.             "cbz %w[k], 4f\n"

  524.             // Loop proper

  525.             "3:\n"

  526.             "ldr  %q[a0], [%[a_ptr]], #16\n"

  527.             "ldr  %q[a1], [%[a_ptr]], #16\n"

  528.             "ldr  %q[a0a], [%[a_ptr]], #16\n"

  529.             "ldr  %q[a1a], [%[a_ptr]], #16\n"

  530.             "ldr  %q[b0], [%[b_ptr]], #16\n"

  531.             "ldr  %q[b1], [%[b_ptr]], #16\n"

  532.             "ldr  %q[b2], [%[b_ptr]], #16\n"

  533.             "sdot  v8.4s, %[b0].16b, %[a0].4b[0]\n"

  534.             "sdot  v9.4s, %[b0].16b, %[a0].4b[1]\n"

  535.             "sdot  v10.4s, %[b0].16b, %[a0].4b[2]\n"

  536.             "sdot  v11.4s, %[b0].16b, %[a0].4b[3]\n"

  537.             "sdot  v12.4s, %[b0].16b, %[a1].4b[0]\n"

  538.             "sdot  v13.4s, %[b0].16b, %[a1].4b[1]\n"

  539.             "sdot  v14.4s, %[b0].16b, %[a1].4b[2]\n"

  540.             "sdot  v15.4s, %[b0].16b, %[a1].4b[3]\n"

  541.             "sdot  v16.4s, %[b1].16b, %[a0].4b[0]\n"

  542.             "sdot  v17.4s, %[b1].16b, %[a0].4b[1]\n"

  543.             "sdot  v18.4s, %[b1].16b, %[a0].4b[2]\n"

  544.             "sdot  v19.4s, %[b1].16b, %[a0].4b[3]\n"

  545.             "sdot  v20.4s, %[b1].16b, %[a1].4b[0]\n"

  546.             "sdot  v21.4s, %[b1].16b, %[a1].4b[1]\n"

  547.             "sdot  v22.4s, %[b1].16b, %[a1].4b[2]\n"

  548.             "sdot  v23.4s, %[b1].16b, %[a1].4b[3]\n"

  549.             "sdot  v24.4s, %[b2].16b, %[a0].4b[0]\n"

  550.             "sdot  v25.4s, %[b2].16b, %[a0].4b[1]\n"

  551.             "sdot  v26.4s, %[b2].16b, %[a0].4b[2]\n"

  552.             "sdot  v27.4s, %[b2].16b, %[a0].4b[3]\n"

  553.             "sdot  v28.4s, %[b2].16b, %[a1].4b[0]\n"

  554.             "sdot  v29.4s, %[b2].16b, %[a1].4b[1]\n"

  555.             "sdot  v30.4s, %[b2].16b, %[a1].4b[2]\n"

  556.             "sdot  v31.4s, %[b2].16b, %[a1].4b[3]\n"

  557.             "ldr  %q[b0], [%[b_ptr]], #16\n"

  558.             "ldr  %q[b1], [%[b_ptr]], #16\n"

  559.             "ldr  %q[b2], [%[b_ptr]], #16\n"

  560.             "sdot  v8.4s, %[b0].16b, %[a0a].4b[0]\n"

  561.             "sdot  v9.4s, %[b0].16b, %[a0a].4b[1]\n"

  562.             "sdot  v10.4s, %[b0].16b, %[a0a].4b[2]\n"

  563.             "sdot  v11.4s, %[b0].16b, %[a0a].4b[3]\n"

  564.             "sdot  v12.4s, %[b0].16b, %[a1a].4b[0]\n"

  565.             "sdot  v13.4s, %[b0].16b, %[a1a].4b[1]\n"

  566.             "sdot  v14.4s, %[b0].16b, %[a1a].4b[2]\n"

  567.             "sdot  v15.4s, %[b0].16b, %[a1a].4b[3]\n"

  568.             "sdot  v16.4s, %[b1].16b, %[a0a].4b[0]\n"

  569.             "sdot  v17.4s, %[b1].16b, %[a0a].4b[1]\n"

  570.             "sdot  v18.4s, %[b1].16b, %[a0a].4b[2]\n"

  571.             "sdot  v19.4s, %[b1].16b, %[a0a].4b[3]\n"

  572.             "sdot  v20.4s, %[b1].16b, %[a1a].4b[0]\n"

  573.             "sdot  v21.4s, %[b1].16b, %[a1a].4b[1]\n"

  574.             "sdot  v22.4s, %[b1].16b, %[a1a].4b[2]\n"

  575.             "sdot  v23.4s, %[b1].16b, %[a1a].4b[3]\n"

  576.             "sdot  v24.4s, %[b2].16b, %[a0a].4b[0]\n"

  577.             "sdot  v25.4s, %[b2].16b, %[a0a].4b[1]\n"

  578.             "sdot  v26.4s, %[b2].16b, %[a0a].4b[2]\n"

  579.             "sdot  v27.4s, %[b2].16b, %[a0a].4b[3]\n"

  580.             "sdot  v28.4s, %[b2].16b, %[a1a].4b[0]\n"

  581.             "sdot  v29.4s, %[b2].16b, %[a1a].4b[1]\n"

  582.             "sdot  v30.4s, %[b2].16b, %[a1a].4b[2]\n"

  583.             "sdot  v31.4s, %[b2].16b, %[a1a].4b[3]\n"

  584. 

  585.             "subs %w[k], %w[k], #1\n"

  586.             "bne 3b\n"

  587. 

  588.             "4:\n"

  589.             "stp q8, q16, [%[outptr0]]\n"

  590.             "str q24, [%[outptr0], #32]\n"

  591.             "stp q9, q17, [%[outptr1]]\n"

  592.             "str q25, [%[outptr1], #32]\n"

  593.             "stp q10, q18, [%[outptr2]]\n"

  594.             "str q26, [%[outptr2], #32]\n"

  595.             "stp q11, q19, [%[outptr3]]\n"

  596.             "str q27, [%[outptr3], #32]\n"

  597.             "stp q12, q20, [%[outptr4]]\n"

  598.             "str q28, [%[outptr4], #32]\n"

  599.             "stp q13, q21, [%[outptr5]]\n"

  600.             "str q29, [%[outptr5], #32]\n"

  601.             "stp q14, q22, [%[outptr6]]\n"

  602.             "str q30, [%[outptr6], #32]\n"

  603.             "stp q15, q23, [%[outptr7]]\n"

  604.             "str q31, [%[outptr7], #32]\n"

  605. 

  606.             : [a_ptr] "+r"(a_ptr), [b_ptr] "+r"(b_ptr), [a0] "+w"(a0), [a1] "+w"(a1),

  607.               [a0a] "+w"(a0a), [a1a] "+w"(a1a), [b0] "+w"(b0), [b1] "+w"(b1),

  608.               [b2] "+w"(b2), [k] "+r"(k), [LDC] "+r"(LDC), [oddk] "+r"(oddk),

  609.               [is_first_k] "+r"(is_first_k), [outptr0] "+r"(outptr0),

  610.               [outptr1] "=r"(outptr1), [outptr2] "=r"(outptr2), [outptr3] "=r"(outptr3),

  611.               [outptr4] "=r"(outptr4), [outptr5] "=r"(outptr5), [outptr6] "=r"(outptr6),

  612.               [outptr7] "=r"(outptr7)

  613.             :

  614.             : "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18",

  615.               "v19", "v20", "v21", "v22", "v23", "v24", "v25", "v26", "v27", "v28",

  616.               "v29", "v30", "v31", "cc", "memory");

  617. }

  618. 

  619. #endif

  620. 

  621. // Overview of register layout:

  622. //

  623. // A 12x4 cell of Rhs is stored in 8bit in q2-q4.

  624. // A 8x4x2 cell of Lhs is stored in 8bit in q0-q1,q5-q6

  625. // A 8x12 block of accumulators is stored in 8bit in q8--q31.

  626. //

  627. //                            +--------+--------+--------+

  628. //                            |v1[0-16]|v2[0-16]|v3[0-16]|

  629. //                       Rhs  +--------+--------+--------+

  630. //                            |v5[0-16]|v6[0-16]|v7[0-16]|

  631. //                            +--------+--------+--------+

  632. //

  633. //                            |        |        |        |

  634. //

  635. //    Lhs                     |        |        |        |

  636. //

  637. //  +-------+-------+ - - - - +--------+--------+--------+

  638. //  |v0[0-4]|v4[0-4]|         | v8[0-4]|v12[0-4]|v16[0-4]|

  639. //  |v0[0-4]|v4[0-4]|         | v9[0-4]|v13[0-4]|v17[0-4]|

  640. //  |v0[0-4]|v4[0-4]|         |v10[0-4]|v14[0-4]|v18[0-4]|

  641. //  |v0[0-4]|v4[0-4]|         |v11[0-4]|v15[0-4]|v19[0-4]|

  642. //  +-------+-------+ - - - - +--------+--------+--------+

  643. //

  644. //                            Accumulator

  645. MEGDNN_ATTRIBUTE_TARGET("dotprod")

  646. static void kern_4x12(

  647.         const int8_t* packA, const int8_t* packB, int K, int32_t* output, int LDC,

  648.         bool is_first_k, int m_remain) {

  649.     K /= 4;

  650.     const int8_t* a_ptr = packA;

  651.     const int8_t* b_ptr = packB;

  652.     // Fix up for odd lengths - set a flag if K is odd, but make

  653.     // sure we round up the iteration count.

  654.     int oddk = (K & 1);

  655.     int k = K / 2;

  656.     int32x4_t a0;

  657.     int32x4_t b0;

  658.     int32x4_t b1;

  659.     int32x4_t b2;

  660.     int32x4_t a0a;

  661.     int32x4_t b0a;

  662.     int32x4_t b1a;

  663.     int32x4_t b2a;

  664. 

  665.     LDC = LDC * sizeof(int32_t);

  666.     int32_t* outptr0 = output;

  667.     int32_t* outptr1;

  668.     int32_t* outptr2;

  669.     int32_t* outptr3;

  670.     size_t x0;

  671. 

  672. // clang-format off

  673. #define LOAD_LINE(v1, v2, v3, m)            \

  674.     "cbz %[x0], 100f\n"                        \

  675.     "ldp " v1 "," v2 ", [%[outptr" m "]]\n" \

  676.     "ldr " v3 ",  [%[outptr" m "], #32]\n"  \

  677.     "subs %[x0], %[x0], #1\n"

  678. 

  679. #define LOAD_C \

  680.     "mov %[x0], %x[m_remain]\n"            \

  681.     LOAD_LINE("q8", "q12", "q16", "0")  \

  682.     LOAD_LINE("q9", "q13", "q17", "1")  \

  683.     LOAD_LINE("q10", "q14", "q18", "2") \

  684.     LOAD_LINE("q11", "q15", "q19", "3") \

  685.     "100:\n"

  686. 

  687. #define STORE_LINE(v1, v2, v3, m)          \

  688.     "cbz %[x0], 101f\n"                       \

  689.     "stp " v1 "," v2", [%[outptr" m "]]\n" \

  690.     "str " v3 ", [%[outptr" m "], #32]\n"  \

  691.     "subs %[x0], %[x0], #1\n"

  692. 

  693. #define STORE_C \

  694.     "mov %[x0], %x[m_remain]\n"             \

  695.     STORE_LINE("q8", "q12", "q16", "0")  \

  696.     STORE_LINE("q9", "q13", "q17", "1")  \

  697.     STORE_LINE("q10", "q14", "q18", "2") \

  698.     STORE_LINE("q11", "q15", "q19", "3") \

  699.     "101:\n"

  700. 

  701.     // clang-format on

  702. 

  703.     asm volatile(

  704.             // load accumulator C

  705.             "add %[outptr1], %[outptr0], %x[LDC]\n"

  706.             "add %[outptr2], %[outptr1], %x[LDC]\n"

  707.             "add %[outptr3], %[outptr2], %x[LDC]\n"

  708.             "cmp %w[is_first_k], #1\n"

  709.             "beq 1f\n" LOAD_C

  710. 

  711.             "b 2f\n"

  712. 

  713.             "1:\n"

  714.             "eor v8.16b,  v8.16b,  v8.16b\n"

  715.             "eor v9.16b,  v9.16b,  v9.16b\n"

  716.             "eor v10.16b, v10.16b, v10.16b\n"

  717.             "eor v11.16b, v11.16b, v11.16b\n"

  718.             "eor v12.16b, v12.16b, v12.16b\n"

  719.             "eor v13.16b, v13.16b, v13.16b\n"

  720.             "eor v14.16b, v14.16b, v14.16b\n"

  721.             "eor v15.16b, v15.16b, v15.16b\n"

  722.             "eor v16.16b, v16.16b, v16.16b\n"

  723.             "eor v17.16b, v17.16b, v17.16b\n"

  724.             "eor v18.16b, v18.16b, v18.16b\n"

  725.             "eor v19.16b, v19.16b, v19.16b\n"

  726. 

  727.             "2: \n"

  728.             "cbz  %w[oddk], 3f\n"

  729. 

  730.             // parse the oddk

  731.             "ldr  %q[a0], [%[a_ptr]], #16\n"

  732.             "ldr  %q[b0], [%[b_ptr]], #16\n"

  733.             "ldr  %q[b1], [%[b_ptr]], #16\n"

  734.             "ldr  %q[b2], [%[b_ptr]], #16\n"

  735.             "sdot  v8.4s, %[b0].16b, %[a0].4b[0]\n"

  736.             "sdot    v9.4s, %[b0].16b, %[a0].4b[1]\n"

  737.             "sdot  v10.4s, %[b0].16b, %[a0].4b[2]\n"

  738.             "sdot  v11.4s, %[b0].16b, %[a0].4b[3]\n"

  739.             "sdot  v12.4s, %[b1].16b, %[a0].4b[0]\n"

  740.             "sdot    v13.4s, %[b1].16b, %[a0].4b[1]\n"

  741.             "sdot  v14.4s, %[b1].16b, %[a0].4b[2]\n"

  742.             "sdot  v15.4s, %[b1].16b, %[a0].4b[3]\n"

  743.             "sdot  v16.4s, %[b2].16b, %[a0].4b[0]\n"

  744.             "sdot    v17.4s, %[b2].16b, %[a0].4b[1]\n"

  745.             "sdot  v18.4s, %[b2].16b, %[a0].4b[2]\n"

  746.             "sdot  v19.4s, %[b2].16b, %[a0].4b[3]\n"

  747. 

  748.             "cbz %w[k], 4f\n"

  749.             // Loop proper

  750.             "3:\n"

  751.             "ldr  %q[a0], [%[a_ptr]], #16\n"

  752.             "ldr  %q[b0], [%[b_ptr]], #16\n"

  753.             "ldr  %q[b1], [%[b_ptr]], #16\n"

  754.             "ldr  %q[b2], [%[b_ptr]], #16\n"

  755.             "ldr  %q[a0a], [%[a_ptr]], #16\n"

  756.             "ldr  %q[b0a], [%[b_ptr]], #16\n"

  757.             "ldr  %q[b1a], [%[b_ptr]], #16\n"

  758.             "ldr  %q[b2a], [%[b_ptr]], #16\n"

  759. 

  760.             "sdot  v8.4s, %[b0].16b, %[a0].4b[0]\n"

  761.             "sdot    v9.4s, %[b0].16b, %[a0].4b[1]\n"

  762.             "sdot  v10.4s, %[b0].16b, %[a0].4b[2]\n"

  763.             "sdot  v11.4s, %[b0].16b, %[a0].4b[3]\n"

  764.             "sdot  v12.4s, %[b1].16b, %[a0].4b[0]\n"

  765.             "sdot    v13.4s, %[b1].16b, %[a0].4b[1]\n"

  766.             "sdot  v14.4s, %[b1].16b, %[a0].4b[2]\n"

  767.             "sdot  v15.4s, %[b1].16b, %[a0].4b[3]\n"

  768.             "sdot  v16.4s, %[b2].16b, %[a0].4b[0]\n"

  769.             "sdot    v17.4s, %[b2].16b, %[a0].4b[1]\n"

  770.             "sdot  v18.4s, %[b2].16b, %[a0].4b[2]\n"

  771.             "sdot  v19.4s, %[b2].16b, %[a0].4b[3]\n"

  772.             "sdot  v8.4s , %[b0a].16b, %[a0a].4b[0]\n"

  773.             "sdot    v9.4s , %[b0a].16b, %[a0a].4b[1]\n"

  774.             "sdot  v10.4s, %[b0a].16b, %[a0a].4b[2]\n"

  775.             "sdot  v11.4s, %[b0a].16b, %[a0a].4b[3]\n"

  776.             "sdot  v12.4s, %[b1a].16b, %[a0a].4b[0]\n"

  777.             "sdot    v13.4s, %[b1a].16b, %[a0a].4b[1]\n"

  778.             "sdot  v14.4s, %[b1a].16b, %[a0a].4b[2]\n"

  779.             "sdot  v15.4s, %[b1a].16b, %[a0a].4b[3]\n"

  780.             "sdot  v16.4s, %[b2a].16b, %[a0a].4b[0]\n"

  781.             "sdot    v17.4s, %[b2a].16b, %[a0a].4b[1]\n"

  782.             "sdot  v18.4s, %[b2a].16b, %[a0a].4b[2]\n"

  783.             "sdot  v19.4s, %[b2a].16b, %[a0a].4b[3]\n"

  784. 

  785.             "subs %w[k], %w[k], #1\n"

  786.             "bne 3b\n"

  787. 

  788.             "4:\n" STORE_C

  789.             : [a_ptr] "+r"(a_ptr), [b_ptr] "+r"(b_ptr), [k] "+r"(k),

  790.               [outptr0] "+r"(outptr0), [oddk] "+r"(oddk), [is_first_k] "+r"(is_first_k),

  791.               [m_remain] "+r"(m_remain), [LDC] "+r"(LDC), [a0] "=w"(a0),

  792.               [a0a] "=w"(a0a), [b0] "=w"(b0), [b1] "=w"(b1), [b2] "=w"(b2),

  793.               [b0a] "=w"(b0a), [b1a] "=w"(b1a), [b2a] "=w"(b2a),

  794.               [outptr1] "=r"(outptr1), [outptr2] "=r"(outptr2), [outptr3] "=r"(outptr3),

  795.               [x0] "=r"(x0)

  796.             :

  797.             : "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18",

  798.               "v19", "memory", "cc");

  799. 

  800. #undef LOAD_LINE

  801. #undef LOAD_C

  802. #undef STORE_LINE

  803. #undef STORE_C

  804. }

  805. 

  806. // Overview of register layout:

  807. //

  808. // A (4x4)x2 cell of Rhs is stored in 8bit in q2-q3.

  809. // A 4x4x2 cell of Lhs is stored in 8bit in q0-q1, q4-a5

  810. // A 8x4 block of accumulators is stored in 8bit in q4--q7.

  811. //

  812. //                            +--------+

  813. //                            |v2[0-16]|

  814. //                       Rhs  +--------+

  815. //                            |v3[0-16]|

  816. //                            +--------+

  817. //                            |        |

  818. //

  819. //    Lhs                     |        |

  820. //

  821. //  +-------+-------+ - - - - +--------+

  822. //  |v0[0-4]|v4[0-4]|         | v6[0-4]|

  823. //  |v0[0-4]|v4[0-4]|         | v7[0-4]|

  824. //  |v0[0-4]|v4[0-4]|         | v8[0-4]|

  825. //  |v0[0-4]|v4[0-4]|         | v9[0-4]|

  826. //  |v1[0-4]|v5[0-4]|         |v10[0-4]|

  827. //  |v1[0-4]|v5[0-4]|         |v11[0-4]|

  828. //  |v1[0-4]|v5[0-4]|         |v12[0-4]|

  829. //  |v1[0-4]|v5[0-4]|         |v13[0-4]|

  830. //  +-------+-------+ - - - - +---------+

  831. //

  832. //                            Accumulator

  833. MEGDNN_ATTRIBUTE_TARGET("dotprod")

  834. static void kern_8x4(

  835.         const int8_t* packA, const int8_t* packB, int K, int32_t* output, int LDC,

  836.         bool is_first_k, int n_remain) {

  837.     K /= 4;

  838.     const int8_t* a_ptr = packA;

  839.     const int8_t* b_ptr = packB;

  840.     // Fix up for odd lengths - set a flag if K is odd, but make

  841.     // sure we round up the iteration count.

  842.     int oddk = (K & 1);

  843.     int k = K / 2;

  844.     int32x4_t a0;

  845.     int32x4_t a1;

  846.     int32x4_t b0;

  847.     int32x4_t b0a;

  848.     int32x4_t a0a;

  849.     int32x4_t a1a;

  850. 

  851.     LDC = LDC * sizeof(int32_t);

  852.     int32_t* outptr0 = output;

  853.     int32_t* outptr1;

  854.     int32_t* outptr2;

  855.     int32_t* outptr3;

  856.     int32_t* outptr4;

  857.     int32_t* outptr5;

  858.     int32_t* outptr6;

  859.     int32_t* outptr7;

  860. 

  861.     size_t x0;

  862. 

  863. // clang-format off

  864. #define LOAD_LINE(reg_index, n)             \

  865.     "mov %[x0], %[outptr" n "]\n"              \

  866.     "cmp %w[n_remain], #4\n"                \

  867.     "blt 100" n "f\n"                       \

  868.     "ldr q" reg_index ", [%[x0]] \n"           \

  869.     "b 101" n "f\n"                         \

  870.     "100" n ":\n"                           \

  871.     "cmp %w[n_remain], #0\n"                \

  872.     "beq 101" n "f\n"                       \

  873.     "ld1 {v" reg_index ".s}[0], [%[x0]], #4\n" \

  874.     "cmp %w[n_remain], #1\n"                \

  875.     "beq 101" n "f\n"                       \

  876.     "ld1 {v" reg_index ".s}[1], [%[x0]], #4\n" \

  877.     "cmp %w[n_remain], #2\n"                \

  878.     "beq 101" n "f\n"                       \

  879.     "ld1 {v" reg_index ".s}[2], [%[x0]], #4\n" \

  880.     "101" n ":\n"

  881. 

  882. 

  883. #define LOAD_C           \

  884.     LOAD_LINE("6", "0")  \

  885.     LOAD_LINE("7", "1")  \

  886.     LOAD_LINE("8", "2")  \

  887.     LOAD_LINE("9", "3")  \

  888.     LOAD_LINE("10", "4") \

  889.     LOAD_LINE("11", "5") \

  890.     LOAD_LINE("12", "6") \

  891.     LOAD_LINE("13", "7")

  892. 

  893. #define STORE_LINE(reg_index, n)            \

  894.     "mov %[x0], %[outptr" n "]\n"              \

  895.     "cmp %w[n_remain], #4\n"                \

  896.     "blt 102" n "f\n"                       \

  897.     "str q" reg_index ", [%[x0]]\n"            \

  898.     "b 103" n "f\n"                         \

  899.     "102" n ":\n"                           \

  900.     "cmp %w[n_remain], #0\n"                \

  901.     "beq 103" n "f\n"                       \

  902.     "st1 {v" reg_index ".s}[0], [%[x0]], #4\n" \

  903.     "cmp %w[n_remain], #1\n"                \

  904.     "beq 103" n "f\n"                       \

  905.     "st1 {v" reg_index ".s}[1], [%[x0]], #4\n" \

  906.     "cmp %w[n_remain], #2\n"                \

  907.     "beq 103" n "f\n"                       \

  908.     "st1 {v" reg_index ".s}[2], [%[x0]], #4\n" \

  909.     "103" n ":\n"

  910. 

  911. #define STORE_C           \

  912.     STORE_LINE("6", "0")  \

  913.     STORE_LINE("7", "1")  \

  914.     STORE_LINE("8", "2")  \

  915.     STORE_LINE("9", "3")  \

  916.     STORE_LINE("10", "4") \

  917.     STORE_LINE("11", "5") \

  918.     STORE_LINE("12", "6") \

  919.     STORE_LINE("13", "7")

  920. 

  921.     // clang-format on

  922. 

  923.     asm volatile(

  924.             // load accumulator C

  925.             "add %[outptr1], %[outptr0], %x[LDC]\n"

  926.             "add %[outptr2], %[outptr1], %x[LDC]\n"

  927.             "add %[outptr3], %[outptr2], %x[LDC]\n"

  928.             "add %[outptr4], %[outptr3], %x[LDC]\n"

  929.             "add %[outptr5], %[outptr4], %x[LDC]\n"

  930.             "add %[outptr6], %[outptr5], %x[LDC]\n"

  931.             "add %[outptr7], %[outptr6], %x[LDC]\n"

  932.             "cmp %w[is_first_k], #1\n"

  933.             "beq 1f\n" LOAD_C

  934. 

  935.             "b 2f\n"

  936. 

  937.             "1:\n"

  938.             "eor v6.16b, v6.16b, v6.16b\n"

  939.             "eor v7.16b, v7.16b, v7.16b\n"

  940.             "eor v8.16b,  v8.16b,  v8.16b\n"

  941.             "eor v9.16b,  v9.16b,  v9.16b\n"

  942.             "eor v10.16b, v10.16b, v10.16b\n"

  943.             "eor v11.16b, v11.16b, v11.16b\n"

  944.             "eor v12.16b, v12.16b, v12.16b\n"

  945.             "eor v13.16b, v13.16b, v13.16b\n"

  946. 

  947.             "2: \n"

  948.             "cbz  %w[oddk], 3f\n"

  949. 

  950.             // parse the oddk

  951.             "ldr  %q[a0], [%[a_ptr]], #16\n"

  952.             "ldr  %q[b0], [%[b_ptr]], #16\n"

  953.             "ldr  %q[a1], [%[a_ptr]], #16\n"

  954.             "sdot  v6.4s , %[b0].16b, %[a0].4b[0]\n"

  955.             "sdot  v7.4s , %[b0].16b, %[a0].4b[1]\n"

  956.             "sdot  v8.4s, %[b0].16b, %[a0].4b[2]\n"

  957.             "sdot  v9.4s, %[b0].16b, %[a0].4b[3]\n"

  958.             "sdot  v10.4s, %[b0].16b, %[a1].4b[0]\n"

  959.             "sdot  v11.4s, %[b0].16b, %[a1].4b[1]\n"

  960.             "sdot  v12.4s, %[b0].16b, %[a1].4b[2]\n"

  961.             "sdot  v13.4s, %[b0].16b, %[a1].4b[3]\n"

  962. 

  963.             "cbz %w[k], 4f\n"

  964.             // Loop proper

  965.             "3:\n"

  966.             "ldr  %q[a0], [%[a_ptr]], #16\n"

  967.             "ldr  %q[b0], [%[b_ptr]], #16\n"

  968.             "ldr  %q[a1], [%[a_ptr]], #16\n"

  969.             "ldr  %q[a0a], [%[a_ptr]], #16\n"

  970.             "ldr  %q[a1a], [%[a_ptr]], #16\n"

  971.             "ldr  %q[b0a], [%[b_ptr]], #16\n"

  972.             "sdot  v6.4s , %[b0].16b, %[a0].4b[0]\n"

  973.             "sdot  v7.4s , %[b0].16b, %[a0].4b[1]\n"

  974.             "sdot  v8.4s, %[b0].16b, %[a0].4b[2]\n"

  975.             "sdot  v9.4s, %[b0].16b, %[a0].4b[3]\n"

  976.             "sdot  v10.4s, %[b0].16b, %[a1].4b[0]\n"

  977.             "sdot  v11.4s, %[b0].16b, %[a1].4b[1]\n"

  978.             "sdot  v12.4s, %[b0].16b, %[a1].4b[2]\n"

  979.             "sdot  v13.4s, %[b0].16b, %[a1].4b[3]\n"

  980. 

  981.             "sdot  v6.4s , %[b0a].16b, %[a0a].4b[0]\n"

  982.             "sdot  v7.4s , %[b0a].16b, %[a0a].4b[1]\n"

  983.             "sdot  v8.4s, %[b0a].16b, %[a0a].4b[2]\n"

  984.             "sdot  v9.4s, %[b0a].16b, %[a0a].4b[3]\n"

  985.             "sdot  v10.4s, %[b0a].16b, %[a1a].4b[0]\n"

  986.             "sdot  v11.4s, %[b0a].16b, %[a1a].4b[1]\n"

  987.             "sdot  v12.4s, %[b0a].16b, %[a1a].4b[2]\n"

  988.             "sdot  v13.4s, %[b0a].16b, %[a1a].4b[3]\n"

  989. 

  990.             "subs %w[k], %w[k], #1\n"

  991.             "bne 3b\n"

  992. 

  993.             "4:\n" STORE_C

  994.             : [a_ptr] "+r"(a_ptr), [b_ptr] "+r"(b_ptr), [LDC] "+r"(LDC),

  995.               [oddk] "+r"(oddk), [is_first_k] "+r"(is_first_k),

  996.               [n_remain] "+r"(n_remain), [k] "+r"(k), [outptr0] "+r"(outptr0),

  997.               [a0] "=w"(a0), [a1] "=w"(a1), [a0a] "=w"(a0a), [a1a] "=w"(a1a),

  998.               [b0] "=w"(b0), [b0a] "=w"(b0a), [outptr1] "=r"(outptr1),

  999.               [outptr2] "=r"(outptr2), [outptr3] "=r"(outptr3), [outptr4] "=r"(outptr4),

  1000.               [outptr5] "=r"(outptr5), [outptr6] "=r"(outptr6), [outptr7] "=r"(outptr7),

  1001.               [x0] "=r"(x0)

  1002.             :

  1003.             : "v6", "v7", "v8", "v9", "v10", "v11", "v12", "v13", "memory", "cc");

  1004. 

  1005. #undef LOAD_LINE

  1006. #undef LOAD_C

  1007. #undef STORE_LINE

  1008. #undef STORE_C

  1009. }

  1010. 

  1011. // Overview of register layout:

  1012. //

  1013. // A 4x4x2 cell of Rhs is stored in 8bit in q2-q3.

  1014. // A 4x4x2 cell of Lhs is stored in 8bit in q0-q1

  1015. // A 4x4x2 block of accumulators is stored in 8bit in q4--q7.

  1016. //

  1017. //                            +--------+

  1018. //                            | v2[0-7]|

  1019. //                       Rhs  +--------+

  1020. //                            | v3[0-7]|

  1021. //                            +--------+

  1022. //                            |        |

  1023. //

  1024. //    Lhs                     |        |

  1025. //

  1026. //  +-------+-------+ - - - - +--------+

  1027. //  |v0[0-4]|v1[0-4]|         | v4[0-7]|

  1028. //  |v0[0-4]|v1[0-4]|         | v5[0-7]|

  1029. //  |v0[0-4]|v1[0-4]|         | v6[0-7]|

  1030. //  |v0[0-4]|v1[0-4]|         | v7[0-7]|

  1031. //  +-------+-------+ - - - - +--------+

  1032. //

  1033. //                            Accumulator

  1034. MEGDNN_ATTRIBUTE_TARGET("dotprod")

  1035. static void kern_4x4(

  1036.         const int8_t* packA, const int8_t* packB, int K, int32_t* output, int LDC,

  1037.         bool is_first_k, int m_remain, int n_remain) {

  1038.     K /= 4;

  1039.     const int32_t* a_ptr = reinterpret_cast<const int32_t*>(packA);

  1040.     const int32_t* b_ptr = reinterpret_cast<const int32_t*>(packB);

  1041.     // Fix up for odd lengths - set a flag if K is odd, but make

  1042.     // sure we round up the iteration count.

  1043.     int oddk = (K & 1);

  1044.     int k = K / 2;

  1045.     int32x4_t a0;

  1046.     int32x4_t a0a;

  1047.     int32x4_t b0;

  1048.     int32x4_t b0a;

  1049.     LDC = LDC * sizeof(int32_t);

  1050. 

  1051.     int32_t* outptr0 = output;

  1052.     int32_t* outptr1;

  1053.     int32_t* outptr2;

  1054.     int32_t* outptr3;

  1055.     size_t x0, x1;

  1056. 

  1057. // clang-format off

  1058. #define LOAD_LINE(reg_index, n)                          \

  1059.     "cbz %[x1], 102f\n"                                     \

  1060.     "mov %[x0], %[outptr" n "]\n"                           \

  1061.     "cmp %w[n_remain], #4\n"                             \

  1062.     "blt 100" n "f\n"                                    \

  1063.     "ldr q" reg_index ", [%[x0]]\n"                         \

  1064.     "b 101" n "f\n"                                      \

  1065.     "100" n ":\n"                                        \

  1066.     "cmp %w[n_remain], #0\n"                             \

  1067.     "beq 101" n "f\n"                                    \

  1068.     "ld1 {v" reg_index ".s}[0], [%[x0]], #4\n"              \

  1069.     "cmp %w[n_remain], #1\n"                             \

  1070.     "beq 101" n "f\n"                                    \

  1071.     "ld1 {v" reg_index ".s}[1], [%[x0]], #4\n"              \

  1072.     "cmp %w[n_remain], #2\n"                             \

  1073.     "beq 101" n "f\n"                                    \

  1074.     "ld1 {v" reg_index ".s}[2], [%[x0]], #4\n"              \

  1075.     "101" n ":\n" \

  1076.     "subs %[x1], %[x1], #1\n"

  1077. 

  1078. #define LOAD_C               \

  1079.     "mov %[x1], %x[m_remain]\n" \

  1080.     LOAD_LINE("4", "0")      \

  1081.     LOAD_LINE("5", "1")      \

  1082.     LOAD_LINE("6", "2")      \

  1083.     LOAD_LINE("7", "3")      \

  1084.     "102:\n"

  1085. 

  1086. #define STORE_LINE(reg_index, n)            \

  1087.     "cbz %[x1], 105f\n"                        \

  1088.     "mov %[x0], %[outptr" n "]\n"              \

  1089.     "cmp %w[n_remain], #4\n"                \

  1090.     "blt 103" n "f\n"                       \

  1091.     "str q" reg_index ", [%[x0]]\n"            \

  1092.     "b 104" n "f\n"                         \

  1093.     "103" n ":\n"                           \

  1094.     "cmp %w[n_remain], #0\n"                \

  1095.     "beq 104" n "f\n"                       \

  1096.     "st1 {v" reg_index ".s}[0], [%[x0]], #4\n" \

  1097.     "cmp %w[n_remain], #1\n"                \

  1098.     "beq 104" n "f\n"                       \

  1099.     "st1 {v" reg_index ".s}[1], [%[x0]], #4\n" \

  1100.     "cmp %w[n_remain], #2\n"                \

  1101.     "beq 104" n "f\n"                       \

  1102.     "st1 {v" reg_index ".s}[2], [%[x0]], #4\n" \

  1103.     "104" n ":\n" \

  1104.     "subs %[x1], %[x1], #1\n"

  1105. 

  1106. #define STORE_C              \

  1107.     "mov %[x1], %x[m_remain]\n" \

  1108.     STORE_LINE("4", "0")     \

  1109.     STORE_LINE("5", "1")     \

  1110.     STORE_LINE("6", "2")     \

  1111.     STORE_LINE("7", "3")     \

  1112.     "105:\n"

  1113. 

  1114.     // clang-format on

  1115. 

  1116.     asm volatile(

  1117.             // load accumulator C

  1118.             "add %[outptr1], %[outptr0], %x[LDC]\n"

  1119.             "add %[outptr2], %[outptr1], %x[LDC]\n"

  1120.             "add %[outptr3], %[outptr2], %x[LDC]\n"

  1121.             "cmp %w[is_first_k], #1\n"

  1122.             "beq 1f\n"  //

  1123.             LOAD_C      //

  1124. 

  1125.             "b 2f\n"

  1126. 

  1127.             "1:\n"

  1128.             "eor v4.16b,  v4.16b,  v4.16b\n"

  1129.             "eor v5.16b,  v5.16b,  v5.16b\n"

  1130.             "eor v6.16b, v6.16b, v6.16b\n"

  1131.             "eor v7.16b, v7.16b, v7.16b\n"

  1132. 

  1133.             "2: \n"

  1134.             "cbz  %w[oddk], 3f\n"

  1135. 

  1136.             // parse the oddk

  1137.             "ldr  %q[a0], [%[a_ptr]], #16\n"

  1138.             "ldr  %q[b0], [%[b_ptr]], #16\n"

  1139.             "sdot  v4.4s , %[b0].16b, %[a0].4b[0]\n"

  1140.             "sdot  v5.4s , %[b0].16b, %[a0].4b[1]\n"

  1141.             "sdot  v6.4s, %[b0].16b, %[a0].4b[2]\n"

  1142.             "sdot  v7.4s, %[b0].16b, %[a0].4b[3]\n"

  1143. 

  1144.             "cbz %w[k], 4f\n"

  1145.             // Loop proper

  1146.             "3:\n"

  1147.             "ldr  %q[a0], [%[a_ptr]], #16\n"

  1148.             "ldr  %q[b0], [%[b_ptr]], #16\n"

  1149.             "ldr  %q[a0a], [%[a_ptr]], #16\n"

  1150.             "ldr  %q[b0a], [%[b_ptr]], #16\n"

  1151.             "sdot  v4.4s , %[b0].16b, %[a0].4b[0]\n"

  1152.             "sdot  v5.4s , %[b0].16b, %[a0].4b[1]\n"

  1153.             "sdot  v6.4s, %[b0].16b, %[a0].4b[2]\n"

  1154.             "sdot  v7.4s, %[b0].16b, %[a0].4b[3]\n"

  1155.             "sdot  v4.4s , %[b0a].16b, %[a0a].4b[0]\n"

  1156.             "sdot  v5.4s , %[b0a].16b, %[a0a].4b[1]\n"

  1157.             "sdot  v6.4s, %[b0a].16b, %[a0a].4b[2]\n"

  1158.             "sdot  v7.4s, %[b0a].16b, %[a0a].4b[3]\n"

  1159. 

  1160.             "subs %w[k], %w[k], #1\n"

  1161.             "bne 3b\n"

  1162. 

  1163.             "4:\n" STORE_C

  1164.             : [a_ptr] "+r"(a_ptr), [b_ptr] "+r"(b_ptr), [oddk] "+r"(oddk),

  1165.               [is_first_k] "+r"(is_first_k), [n_remain] "+r"(n_remain),

  1166.               [m_remain] "+r"(m_remain), [LDC] "+r"(LDC), [outptr0] "+r"(outptr0),

  1167.               [k] "+r"(k), [a0] "=w"(a0), [a0a] "=w"(a0a), [b0] "=w"(b0),

  1168.               [b0a] "=w"(b0a), [outptr1] "=r"(outptr1), [outptr2] "=r"(outptr2),

  1169.               [outptr3] "=r"(outptr3), [x0] "=r"(x0), [x1] "=r"(x1)

  1170.             :

  1171.             : "v4", "v5", "v6", "v7", "memory", "cc");

  1172. 

  1173. #undef LOAD_LINE

  1174. #undef LOAD_C

  1175. #undef STORE_LINE

  1176. #undef STORE_C

  1177. }

  1178. 

  1179. static void gemm_s8_8x12_pack_A_n(

  1180.         dt_int8* outptr, const dt_int8* inptr, int ldin, int y0, int ymax, int k0,

  1181.         int kmax) {

  1182.     int8_t zerobuff[16];

  1183.     std::memset(zerobuff, 0, sizeof(int8_t) * 16);

  1184. 

  1185.     int y = y0;

  1186.     for (; y + 7 < ymax; y += 8) {

  1187.         const int8_t* inptr0 = inptr + y * ldin + k0;

  1188.         const int8_t* inptr1 = inptr0 + ldin;

  1189.         const int8_t* inptr2 = inptr1 + ldin;

  1190.         const int8_t* inptr3 = inptr2 + ldin;

  1191.         const int8_t* inptr4 = inptr3 + ldin;

  1192.         const int8_t* inptr5 = inptr4 + ldin;

  1193.         const int8_t* inptr6 = inptr5 + ldin;

  1194.         const int8_t* inptr7 = inptr6 + ldin;

  1195. 

  1196.         prefetch_2x(inptr0);

  1197.         prefetch_2x(inptr1);

  1198.         prefetch_2x(inptr2);

  1199.         prefetch_2x(inptr3);

  1200.         prefetch_2x(inptr4);

  1201.         prefetch_2x(inptr5);

  1202.         prefetch_2x(inptr6);

  1203.         prefetch_2x(inptr7);

  1204. 

  1205.         int K = kmax - k0;

  1206.         //! read 8 * 4 in each row

  1207.         for (; K > 15; K -= 16) {

  1208.             interleave_8x4_4_b(

  1209.                     inptr0, inptr1, inptr2, inptr3, inptr4, inptr5, inptr6, inptr7,

  1210.                     outptr);

  1211.         }

  1212. 

  1213.         if (K > 0) {

  1214.             interleave_8(

  1215.                     inptr0, inptr1, inptr2, inptr3, inptr4, inptr5, inptr6, inptr7,

  1216.                     outptr, 4, K);

  1217.         }

  1218.     }

  1219.     for (; y < ymax; y += 4) {

  1220.         const int8_t* inptr0 = inptr + y * ldin + k0;

  1221.         const int8_t* inptr1 = inptr0 + ldin;

  1222.         const int8_t* inptr2 = inptr1 + ldin;

  1223.         const int8_t* inptr3 = inptr2 + ldin;

  1224. 

  1225.         prefetch_2x(inptr0);

  1226.         prefetch_2x(inptr1);

  1227.         prefetch_2x(inptr2);

  1228.         prefetch_2x(inptr3);

  1229. 

  1230.         int K = kmax - k0;

  1231.         //! read 4 * 4 in each row

  1232.         for (; K > 15; K -= 16) {

  1233.             if (y + 3 >= ymax) {

  1234.                 switch (y + 3 - ymax) {

  1235.                     case 2:

  1236.                         inptr1 = zerobuff;

  1237.                     case 1:

  1238.                         inptr2 = zerobuff;

  1239.                     case 0:

  1240.                         inptr3 = zerobuff;

  1241.                         break;

  1242.                     default:

  1243.                         megdnn_assert(0);

  1244.                 }

  1245.             }

  1246. 

  1247.             interleave_4x4_4_b(inptr0, inptr1, inptr2, inptr3, outptr);

  1248.         }

  1249. 

  1250.         if (K > 0) {

  1251.             if (y + 3 >= ymax) {

  1252.                 switch (y + 3 - ymax) {

  1253.                     case 2:

  1254.                         inptr1 = zerobuff;

  1255.                     case 1:

  1256.                         inptr2 = zerobuff;

  1257.                     case 0:

  1258.                         inptr3 = zerobuff;

  1259.                         break;

  1260.                     default:

  1261.                         megdnn_assert(0);

  1262.                 }

  1263.             }

  1264.             interleave_4(inptr0, inptr1, inptr2, inptr3, outptr, 4, K);

  1265.         }

  1266.     }

  1267. }

  1268. 

  1269. static void gemm_s8_8x12_pack_A_t(

  1270.         dt_int8* out, const dt_int8* in, int ldin, int x0, int xmax, int k0, int kmax) {

  1271.     int8_t zerobuff[16];

  1272.     std::memset(zerobuff, 0, sizeof(int8_t) * 16);

  1273.     const int ksize = kmax - k0;

  1274.     const int ksize8 = round_up<int>(ksize, 4) * 8;

  1275.     const int ksize4 = round_up(ksize, 4) * 4;

  1276.     int8_t* outptr = out;

  1277.     int8_t* outptr_base = out;

  1278.     //! 4x4 block output start pos

  1279.     int8_t* outptr_base4 = out + ((xmax - x0) / 8) * ksize8;

  1280. 

  1281.     int k = k0;

  1282.     for (; k < kmax; k += 4) {

  1283.         const int8_t* inptr0 = in + k * ldin + x0;

  1284.         const int8_t* inptr1 = inptr0 + ldin;

  1285.         const int8_t* inptr2 = inptr1 + ldin;

  1286.         const int8_t* inptr3 = inptr2 + ldin;

  1287.         prefetch_2x(inptr0);

  1288.         prefetch_2x(inptr1);

  1289.         prefetch_2x(inptr2);

  1290.         prefetch_2x(inptr3);

  1291. 

  1292.         int x = x0;

  1293.         outptr = outptr_base;

  1294.         for (; x + 7 < xmax; x += 8) {

  1295.             if (k + 3 >= kmax) {

  1296.                 switch (k + 3 - kmax) {

  1297.                     case 2:

  1298.                         inptr1 = zerobuff;

  1299.                     case 1:

  1300.                         inptr2 = zerobuff;

  1301.                     case 0:

  1302.                         inptr3 = zerobuff;

  1303.                         break;

  1304.                     default:

  1305.                         megdnn_assert(0);

  1306.                 }

  1307.             }

  1308. 

  1309.             transpose_8x4_1_b(inptr0, inptr1, inptr2, inptr3, outptr);

  1310.             outptr += ksize8;

  1311.         }

  1312. 

  1313.         outptr = outptr_base4;

  1314.         for (; x + 3 < xmax; x += 4) {

  1315.             if (k + 3 >= kmax) {

  1316.                 switch (k + 3 - kmax) {

  1317.                     case 2:

  1318.                         inptr1 = zerobuff;

  1319.                     case 1:

  1320.                         inptr2 = zerobuff;

  1321.                     case 0:

  1322.                         inptr3 = zerobuff;

  1323.                         break;

  1324.                     default:

  1325.                         megdnn_assert(0);

  1326.                 }

  1327.             }

  1328. 

  1329.             transpose_4(inptr0, inptr1, inptr2, inptr3, outptr, 4, 4);

  1330.             outptr += ksize4;

  1331.         }

  1332. 

  1333.         if (x < xmax) {

  1334.             if (k + 3 >= kmax) {

  1335.                 switch (k + 3 - kmax) {

  1336.                     case 2:

  1337.                         inptr1 = zerobuff;

  1338.                     case 1:

  1339.                         inptr2 = zerobuff;

  1340.                     case 0:

  1341.                         inptr3 = zerobuff;

  1342.                         break;

  1343.                     default:

  1344.                         megdnn_assert(0);

  1345.                 }

  1346.             }

  1347. 

  1348.             transpose_4(inptr0, inptr1, inptr2, inptr3, outptr, 4, xmax - x);

  1349.         }

  1350. 

  1351.         outptr_base += 8 * 4;

  1352.         outptr_base4 += 4 * 4;

  1353.     }

  1354. }

  1355. 

  1356. static void gemm_s8_8x12_pack_B_n(

  1357.         dt_int8* out, const dt_int8* in, int ldin, int x0, int xmax, int k0, int kmax) {

  1358.     int8_t zerobuff[16];

  1359.     std::memset(zerobuff, 0, sizeof(int8_t) * 16);

  1360.     const int ksize = kmax - k0;

  1361.     const int ksize12 = round_up<int>(ksize, 4) * 12;

  1362.     const int ksize4 = round_up(ksize, 4) * 4;

  1363.     int8_t* outptr = out;

  1364.     int8_t* outptr_base = out;

  1365.     //! 4x4 block output start pos

  1366.     int8_t* outptr_base4 = out + ((xmax - x0) / 12) * ksize12;

  1367. 

  1368.     int k = k0;

  1369.     for (; k < kmax; k += 4) {

  1370.         const int8_t* inptr0 = in + k * ldin + x0;

  1371.         const int8_t* inptr1 = inptr0 + ldin;

  1372.         const int8_t* inptr2 = inptr1 + ldin;

  1373.         const int8_t* inptr3 = inptr2 + ldin;

  1374.         prefetch_2x(inptr0);

  1375.         prefetch_2x(inptr1);

  1376.         prefetch_2x(inptr2);

  1377.         prefetch_2x(inptr3);

  1378. 

  1379.         int x = x0;

  1380.         outptr = outptr_base;

  1381.         for (; x + 11 < xmax; x += 12) {

  1382.             if (k + 3 >= kmax) {

  1383.                 switch (k + 3 - kmax) {

  1384.                     case 2:

  1385.                         inptr1 = zerobuff;

  1386.                     case 1:

  1387.                         inptr2 = zerobuff;

  1388.                     case 0:

  1389.                         inptr3 = zerobuff;

  1390.                         break;

  1391.                     default:

  1392.                         megdnn_assert(0);

  1393.                 }

  1394.             }

  1395. 

  1396.             transpose_12x4_1_b(inptr0, inptr1, inptr2, inptr3, outptr);

  1397.             outptr += ksize12;

  1398.         }

  1399. 

  1400.         outptr = outptr_base4;

  1401.         for (; x + 3 < xmax; x += 4) {

  1402.             if (k + 3 >= kmax) {

  1403.                 switch (k + 3 - kmax) {

  1404.                     case 2:

  1405.                         inptr1 = zerobuff;

  1406.                     case 1:

  1407.                         inptr2 = zerobuff;

  1408.                     case 0:

  1409.                         inptr3 = zerobuff;

  1410.                         break;

  1411.                     default:

  1412.                         megdnn_assert(0);

  1413.                 }

  1414.             }

  1415. 

  1416.             transpose_4(inptr0, inptr1, inptr2, inptr3, outptr, 4, 4);

  1417.             outptr += ksize4;

  1418.         }

  1419. 

  1420.         if (x < xmax) {

  1421.             if (k + 3 >= kmax) {

  1422.                 switch (k + 3 - kmax) {

  1423.                     case 2:

  1424.                         inptr1 = zerobuff;

  1425.                     case 1:

  1426.                         inptr2 = zerobuff;

  1427.                     case 0:

  1428.                         inptr3 = zerobuff;

  1429.                         break;

  1430.                     default:

  1431.                         megdnn_assert(0);

  1432.                 }

  1433.             }

  1434. 

  1435.             transpose_4(inptr0, inptr1, inptr2, inptr3, outptr, 4, xmax - x);

  1436.         }

  1437. 

  1438.         outptr_base += 12 * 4;

  1439.         outptr_base4 += 4 * 4;

  1440.     }

  1441. }

  1442. 

  1443. static void gemm_s8_8x12_pack_B_t(

  1444.         dt_int8* outptr, const dt_int8* inptr, int ldin, int y0, int ymax, int k0,

  1445.         int kmax) {

  1446.     int8_t zerobuff[16];

  1447.     std::memset(zerobuff, 0, sizeof(int8_t) * 16);

  1448. 

  1449.     int y = y0;

  1450.     for (; y + 11 < ymax; y += 12) {

  1451.         const int8_t* inptr0 = inptr + y * ldin + k0;

  1452.         const int8_t* inptr1 = inptr0 + ldin;

  1453.         const int8_t* inptr2 = inptr1 + ldin;

  1454.         const int8_t* inptr3 = inptr2 + ldin;

  1455.         const int8_t* inptr4 = inptr3 + ldin;

  1456.         const int8_t* inptr5 = inptr4 + ldin;

  1457.         const int8_t* inptr6 = inptr5 + ldin;

  1458.         const int8_t* inptr7 = inptr6 + ldin;

  1459.         const int8_t* inptr8 = inptr7 + ldin;

  1460.         const int8_t* inptr9 = inptr8 + ldin;

  1461.         const int8_t* inptr10 = inptr9 + ldin;

  1462.         const int8_t* inptr11 = inptr10 + ldin;

  1463. 

  1464.         prefetch_2x(inptr0);

  1465.         prefetch_2x(inptr1);

  1466.         prefetch_2x(inptr2);

  1467.         prefetch_2x(inptr3);

  1468.         prefetch_2x(inptr4);

  1469.         prefetch_2x(inptr5);

  1470.         prefetch_2x(inptr6);

  1471.         prefetch_2x(inptr7);

  1472.         prefetch_2x(inptr8);

  1473.         prefetch_2x(inptr9);

  1474.         prefetch_2x(inptr10);

  1475.         prefetch_2x(inptr11);

  1476. 

  1477.         int K = kmax - k0;

  1478.         //! read 12 * 4 in each row

  1479.         for (; K > 15; K -= 16) {

  1480.             interleave_12x4_4_b(

  1481.                     inptr0, inptr1, inptr2, inptr3, inptr4, inptr5, inptr6, inptr7,

  1482.                     inptr8, inptr9, inptr10, inptr11, outptr);

  1483.         }

  1484. 

  1485.         if (K > 0) {

  1486.             interleave_12(

  1487.                     inptr0, inptr1, inptr2, inptr3, inptr4, inptr5, inptr6, inptr7,

  1488.                     inptr8, inptr9, inptr10, inptr11, outptr, 4, K);

  1489.         }

  1490.     }

  1491.     for (; y < ymax; y += 4) {

  1492.         const int8_t* inptr0 = inptr + y * ldin + k0;

  1493.         const int8_t* inptr1 = inptr0 + ldin;

  1494.         const int8_t* inptr2 = inptr1 + ldin;

  1495.         const int8_t* inptr3 = inptr2 + ldin;

  1496. 

  1497.         prefetch_2x(inptr0);

  1498.         prefetch_2x(inptr1);

  1499.         prefetch_2x(inptr2);

  1500.         prefetch_2x(inptr3);

  1501. 

  1502.         int K = kmax - k0;

  1503.         //! read 4 * 4 in each row

  1504.         for (; K > 15; K -= 16) {

  1505.             if (y + 3 >= ymax) {

  1506.                 switch (y + 3 - ymax) {

  1507.                     case 2:

  1508.                         inptr1 = zerobuff;

  1509.                     case 1:

  1510.                         inptr2 = zerobuff;

  1511.                     case 0:

  1512.                         inptr3 = zerobuff;

  1513.                         break;

  1514.                     default:

  1515.                         megdnn_assert(0);

  1516.                 }

  1517.             }

  1518. 

  1519.             interleave_4x4_4_b(inptr0, inptr1, inptr2, inptr3, outptr);

  1520.         }

  1521. 

  1522.         if (K > 0) {

  1523.             if (y + 3 >= ymax) {

  1524.                 switch (y + 3 - ymax) {

  1525.                     case 2:

  1526.                         inptr1 = zerobuff;

  1527.                     case 1:

  1528.                         inptr2 = zerobuff;

  1529.                     case 0:

  1530.                         inptr3 = zerobuff;

  1531.                         break;

  1532.                     default:

  1533.                         megdnn_assert(0);

  1534.                 }

  1535.             }

  1536.             interleave_4(inptr0, inptr1, inptr2, inptr3, outptr, 4, K);

  1537.         }

  1538.     }

  1539. }

  1540. 

  1541. }  // namespace matmul_8x12x4

  1542. }  // namespace aarch64

  1543. }  // namespace megdnn

  1544. 

  1545. #endif

  1546. // vim: syntax=cpp.doxygen