wangwei
/
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/fp16/strategy.cpp

2558 lines
91 kB
Raw Blame History 

  1. #include "src/aarch64/matrix_mul/fp16/strategy.h"

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

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

  4. #include "src/common/utils.h"

  5. 

  6. #if __ARM_FEATURE_FP16_VECTOR_ARITHMETIC

  7. using namespace megdnn;

  8. using namespace aarch64;

  9. using namespace aarch64::matmul;

  10. 

  11. namespace {

  12. 

  13. void interleave_8x1(

  14.         __fp16* out, const __fp16* in, int ldin, int y0, int ymax, int k0, int kmax) {

  15.     __fp16* outptr = out;

  16.     const __fp16* inptr = in;

  17.     __fp16 zerobuff[24];

  18.     std::memset(zerobuff, 0, sizeof(__fp16) * 24);

  19. 

  20.     int y = y0;

  21.     for (; y + 8 <= ymax; y += 8) {

  22.         const __fp16* inptr0 = inptr + y * ldin + k0;

  23.         const __fp16* inptr1 = inptr0 + ldin;

  24.         const __fp16* inptr2 = inptr1 + ldin;

  25.         const __fp16* inptr3 = inptr2 + ldin;

  26.         const __fp16* inptr4 = inptr3 + ldin;

  27.         const __fp16* inptr5 = inptr4 + ldin;

  28.         const __fp16* inptr6 = inptr5 + ldin;

  29.         const __fp16* inptr7 = inptr6 + ldin;

  30. 

  31.         prefetch_2x(inptr0);

  32.         prefetch_2x(inptr1);

  33.         prefetch_2x(inptr2);

  34.         prefetch_2x(inptr3);

  35.         prefetch_2x(inptr4);

  36.         prefetch_2x(inptr5);

  37.         prefetch_2x(inptr6);

  38.         prefetch_2x(inptr7);

  39. 

  40.         int x = (kmax - k0);

  41.         for (; x > 7; x -= 8) {

  42.             int skippf = (x & 31);

  43.             interleave_8x1_8_h(

  44.                     inptr0, inptr1, inptr2, inptr3, inptr4, inptr5, inptr6, inptr7,

  45.                     outptr, skippf);

  46.         }

  47. 

  48.         for (; x > 0; x--) {

  49.             *outptr++ = *inptr0++;

  50.             *outptr++ = *inptr1++;

  51.             *outptr++ = *inptr2++;

  52.             *outptr++ = *inptr3++;

  53.             *outptr++ = *inptr4++;

  54.             *outptr++ = *inptr5++;

  55.             *outptr++ = *inptr6++;

  56.             *outptr++ = *inptr7++;

  57.         }

  58.     }

  59. 

  60.     for (; y < ymax; y += 4) {

  61.         const __fp16* inptr0 = inptr + y * ldin + k0;

  62.         const __fp16* inptr1 = inptr0 + ldin;

  63.         const __fp16* inptr2 = inptr1 + ldin;

  64.         const __fp16* inptr3 = inptr2 + ldin;

  65. 

  66.         prefetch_2x(inptr0);

  67.         prefetch_2x(inptr1);

  68.         prefetch_2x(inptr2);

  69.         prefetch_2x(inptr3);

  70. 

  71.         /* Cope with ragged cases by copying from a buffer of zeroes instead

  72.          */

  73.         int x = (kmax - k0);

  74.         for (; x > 3; x -= 4) {

  75.             if ((y + 3) >= ymax) {

  76.                 switch ((y + 3) - ymax) {

  77.                     /* Everything falls through in here */

  78.                     case 2:

  79.                         inptr1 = zerobuff;

  80.                     case 1:

  81.                         inptr2 = zerobuff;

  82.                     case 0:

  83.                         inptr3 = zerobuff;

  84.                         break;

  85.                     default:

  86.                         megdnn_assert(0);

  87.                 }

  88.             }

  89. 

  90.             interleave_4x1_4_h(inptr0, inptr1, inptr2, inptr3, outptr);

  91.         }

  92. 

  93.         for (; x > 0; x--) {

  94.             *outptr++ = *inptr0++;

  95.             *outptr++ = *inptr1++;

  96.             *outptr++ = *inptr2++;

  97.             *outptr++ = *inptr3++;

  98.         }

  99.     }

  100. }

  101. 

  102. void interleave_24x1(

  103.         __fp16* out, const __fp16* in, const int ldin, const int y0, const int ymax,

  104.         const int k0, const int kmax) {

  105.     __fp16* outptr = out;

  106.     const __fp16* inptr = in;

  107.     __fp16 zerobuff[24];

  108.     std::memset(zerobuff, 0, sizeof(__fp16) * 24);

  109.     int K16 = 16 * (kmax - k0);

  110.     int K24 = 24 * (kmax - k0);

  111. 

  112.     int y = y0;

  113.     for (; y + 24 <= ymax; y += 24) {

  114.         int yi = y;

  115.         for (; yi < y + 24; yi += 8) {

  116.             const __fp16* inptr0 = inptr + yi * ldin + k0;

  117.             const __fp16* inptr1 = inptr0 + ldin;

  118.             const __fp16* inptr2 = inptr1 + ldin;

  119.             const __fp16* inptr3 = inptr2 + ldin;

  120.             const __fp16* inptr4 = inptr3 + ldin;

  121.             const __fp16* inptr5 = inptr4 + ldin;

  122.             const __fp16* inptr6 = inptr5 + ldin;

  123.             const __fp16* inptr7 = inptr6 + ldin;

  124.             __fp16* outptr_inner = outptr + yi - y;

  125. 

  126.             prefetch_2x(inptr0);

  127.             prefetch_2x(inptr1);

  128.             prefetch_2x(inptr2);

  129.             prefetch_2x(inptr3);

  130.             prefetch_2x(inptr4);

  131.             prefetch_2x(inptr5);

  132.             prefetch_2x(inptr6);

  133.             prefetch_2x(inptr7);

  134. 

  135.             int x = (kmax - k0);

  136.             for (; x > 7; x -= 8) {

  137.                 int skippf = (x & 31);

  138.                 interleave_24x1_8_h_helper(

  139.                         inptr0, inptr1, inptr2, inptr3, inptr4, inptr5, inptr6, inptr7,

  140.                         outptr_inner, skippf);

  141.             }

  142.             for (; x > 0; x--) {

  143.                 *outptr_inner++ = *inptr0++;

  144.                 *outptr_inner++ = *inptr1++;

  145.                 *outptr_inner++ = *inptr2++;

  146.                 *outptr_inner++ = *inptr3++;

  147.                 *outptr_inner++ = *inptr4++;

  148.                 *outptr_inner++ = *inptr5++;

  149.                 *outptr_inner++ = *inptr6++;

  150.                 *outptr_inner++ = *inptr7++;

  151.                 outptr_inner += 16;

  152.             }

  153.         }

  154.         outptr += K24;

  155.     }

  156. 

  157.     for (; y + 16 <= ymax; y += 16) {

  158.         int yi = y;

  159.         for (; yi < y + 16; yi += 8) {

  160.             const __fp16* inptr0 = inptr + yi * ldin + k0;

  161.             const __fp16* inptr1 = inptr0 + ldin;

  162.             const __fp16* inptr2 = inptr1 + ldin;

  163.             const __fp16* inptr3 = inptr2 + ldin;

  164.             const __fp16* inptr4 = inptr3 + ldin;

  165.             const __fp16* inptr5 = inptr4 + ldin;

  166.             const __fp16* inptr6 = inptr5 + ldin;

  167.             const __fp16* inptr7 = inptr6 + ldin;

  168.             __fp16* outptr_inner = outptr + yi - y;

  169. 

  170.             prefetch_2x(inptr0);

  171.             prefetch_2x(inptr1);

  172.             prefetch_2x(inptr2);

  173.             prefetch_2x(inptr3);

  174.             prefetch_2x(inptr4);

  175.             prefetch_2x(inptr5);

  176.             prefetch_2x(inptr6);

  177.             prefetch_2x(inptr7);

  178. 

  179.             int x = (kmax - k0);

  180.             for (; x > 7; x -= 8) {

  181.                 int skippf = (x & 31);

  182.                 interleave_16x1_8_h_helper(

  183.                         inptr0, inptr1, inptr2, inptr3, inptr4, inptr5, inptr6, inptr7,

  184.                         outptr_inner, skippf);

  185.             }

  186.             for (; x > 0; x--) {

  187.                 *outptr_inner++ = *inptr0++;

  188.                 *outptr_inner++ = *inptr1++;

  189.                 *outptr_inner++ = *inptr2++;

  190.                 *outptr_inner++ = *inptr3++;

  191.                 *outptr_inner++ = *inptr4++;

  192.                 *outptr_inner++ = *inptr5++;

  193.                 *outptr_inner++ = *inptr6++;

  194.                 *outptr_inner++ = *inptr7++;

  195.                 outptr_inner += 8;

  196.             }

  197.         }

  198.         outptr += K16;

  199.     }

  200. 

  201.     for (; y + 8 <= ymax; y += 8) {

  202.         const __fp16* inptr0 = inptr + y * ldin + k0;

  203.         const __fp16* inptr1 = inptr0 + ldin;

  204.         const __fp16* inptr2 = inptr1 + ldin;

  205.         const __fp16* inptr3 = inptr2 + ldin;

  206.         const __fp16* inptr4 = inptr3 + ldin;

  207.         const __fp16* inptr5 = inptr4 + ldin;

  208.         const __fp16* inptr6 = inptr5 + ldin;

  209.         const __fp16* inptr7 = inptr6 + ldin;

  210. 

  211.         prefetch_2x(inptr0);

  212.         prefetch_2x(inptr1);

  213.         prefetch_2x(inptr2);

  214.         prefetch_2x(inptr3);

  215.         prefetch_2x(inptr4);

  216.         prefetch_2x(inptr5);

  217.         prefetch_2x(inptr6);

  218.         prefetch_2x(inptr7);

  219. 

  220.         int x = (kmax - k0);

  221.         for (; x > 7; x -= 8) {

  222.             int skippf = (x & 31);

  223.             interleave_8x1_8_h(

  224.                     inptr0, inptr1, inptr2, inptr3, inptr4, inptr5, inptr6, inptr7,

  225.                     outptr, skippf);

  226.         }

  227. 

  228.         for (; x > 0; x--) {

  229.             *outptr++ = *inptr0++;

  230.             *outptr++ = *inptr1++;

  231.             *outptr++ = *inptr2++;

  232.             *outptr++ = *inptr3++;

  233.             *outptr++ = *inptr4++;

  234.             *outptr++ = *inptr5++;

  235.             *outptr++ = *inptr6++;

  236.             *outptr++ = *inptr7++;

  237.         }

  238.     }

  239. 

  240.     for (; y < ymax; y += 4) {

  241.         const __fp16* inptr0 = inptr + y * ldin + k0;

  242.         const __fp16* inptr1 = inptr0 + ldin;

  243.         const __fp16* inptr2 = inptr1 + ldin;

  244.         const __fp16* inptr3 = inptr2 + ldin;

  245. 

  246.         prefetch_2x(inptr0);

  247.         prefetch_2x(inptr1);

  248.         prefetch_2x(inptr2);

  249.         prefetch_2x(inptr3);

  250. 

  251.         /* Cope with ragged cases by copying from a buffer of zeroes instead

  252.          */

  253.         int x = (kmax - k0);

  254.         for (; x > 3; x -= 4) {

  255.             if ((y + 3) >= ymax) {

  256.                 switch ((y + 3) - ymax) {

  257.                     /* Everything falls through in here */

  258.                     case 2:

  259.                         inptr1 = zerobuff;

  260.                     case 1:

  261.                         inptr2 = zerobuff;

  262.                     case 0:

  263.                         inptr3 = zerobuff;

  264.                         break;

  265.                     default:

  266.                         megdnn_assert(0);

  267.                 }

  268.             }

  269. 

  270.             interleave_4x1_4_h(inptr0, inptr1, inptr2, inptr3, outptr);

  271.         }

  272. 

  273.         for (; x > 0; x--) {

  274.             *outptr++ = *inptr0++;

  275.             *outptr++ = *inptr1++;

  276.             *outptr++ = *inptr2++;

  277.             *outptr++ = *inptr3++;

  278.         }

  279.     }

  280. }

  281. 

  282. void transpose_1x8(

  283.         __fp16* out, const __fp16* in, int ldin, int x0, int xmax, int k0, int kmax) {

  284.     int ksize = kmax - k0;

  285.     int ksize8 = (ksize << 3);

  286.     int ksize4 = (ksize << 2);

  287.     int k = ksize;

  288.     __fp16* outptr_base8 = out;

  289.     __fp16* outptr_base4 = out;

  290. 

  291.     const __fp16* inptr_base = in + x0 + k0 * ldin;

  292. 

  293.     for (; k > 3; k -= 4) {

  294.         __fp16* outptr = outptr_base8;

  295. 

  296.         const __fp16* inptr = inptr_base;

  297.         const __fp16* inptr1 = inptr + ldin;

  298.         const __fp16* inptr2 = inptr1 + ldin;

  299.         const __fp16* inptr3 = inptr2 + ldin;

  300. 

  301.         prefetch_3x(inptr);

  302.         prefetch_3x(inptr1);

  303.         prefetch_3x(inptr2);

  304.         prefetch_3x(inptr3);

  305. 

  306.         int x = x0;

  307.         for (; x + 8 <= xmax; x += 8) {

  308.             transpose_8x4_1_h(inptr, inptr1, inptr2, inptr3, outptr);

  309.             outptr += ksize8;

  310.         }

  311.         outptr += outptr_base4 - outptr_base8;

  312.         for (; x < xmax; x += 4) {

  313.             for (int i = 0; i < 4; i++) {

  314.                 __fp16 val = (x + i < xmax) ? (*inptr++) : (__fp16)(0);

  315.                 *outptr++ = val;

  316.             }

  317.             for (int i = 0; i < 4; i++) {

  318.                 __fp16 val = (x + i < xmax) ? (*inptr1++) : (__fp16)(0);

  319.                 *outptr++ = val;

  320.             }

  321.             for (int i = 0; i < 4; i++) {

  322.                 __fp16 val = (x + i < xmax) ? (*inptr2++) : (__fp16)(0);

  323.                 *outptr++ = val;

  324.             }

  325.             for (int i = 0; i < 4; i++) {

  326.                 __fp16 val = (x + i < xmax) ? (*inptr3++) : (__fp16)(0);

  327.                 *outptr++ = val;

  328.             }

  329.             outptr -= 16;

  330.             outptr += ksize4;

  331.         }

  332. 

  333.         inptr_base += ldin * 4;

  334.         outptr_base8 += 8 * 4;

  335.         outptr_base4 += 4 * 4;

  336.     }

  337. 

  338.     if (k) {

  339.         __fp16* outptr = outptr_base8;

  340.         const __fp16* inptr = inptr_base;

  341.         const __fp16* inptr1 = inptr + ldin;

  342.         const __fp16* inptr2 = inptr1 + ldin;

  343. 

  344.         prefetch_3x(inptr);

  345.         prefetch_3x(inptr1);

  346.         prefetch_3x(inptr2);

  347. 

  348.         int x = x0;

  349.         for (; x + 8 <= xmax; x += 8) {

  350.             switch (k) {

  351.                 case 3:

  352.                     transpose_8x2_1_h(inptr, inptr1, outptr);

  353.                     transpose_8x1_1_h(inptr2, outptr + 8 * 2);

  354.                     break;

  355. 

  356.                 case 2:

  357.                     transpose_8x2_1_h(inptr, inptr1, outptr);

  358.                     break;

  359. 

  360.                 case 1:

  361.                     transpose_8x1_1_h(inptr, outptr);

  362.                     break;

  363. 

  364.                 default:

  365.                     megdnn_assert(0);

  366.             }

  367.             outptr += ksize8;

  368.         }

  369. 

  370.         outptr += outptr_base4 - outptr_base8;

  371.         for (; x < xmax; x += 4) {

  372.             switch (k) {

  373.                 case 3:

  374.                     for (int i = 0; i < 4; i++) {

  375.                         __fp16 val = (x + i < xmax) ? (*inptr++) : (__fp16)(0);

  376.                         *outptr++ = val;

  377.                     }

  378.                     for (int i = 0; i < 4; i++) {

  379.                         __fp16 val = (x + i < xmax) ? (*inptr1++) : (__fp16)(0);

  380.                         *outptr++ = val;

  381.                     }

  382.                     for (int i = 0; i < 4; i++) {

  383.                         __fp16 val = (x + i < xmax) ? (*inptr2++) : (__fp16)(0);

  384.                         *outptr++ = val;

  385.                     }

  386.                     outptr -= 12;

  387.                     break;

  388.                 case 2:

  389.                     for (int i = 0; i < 4; i++) {

  390.                         __fp16 val = (x + i < xmax) ? (*inptr++) : (__fp16)(0);

  391.                         *outptr++ = val;

  392.                     }

  393.                     for (int i = 0; i < 4; i++) {

  394.                         __fp16 val = (x + i < xmax) ? (*inptr1++) : (__fp16)(0);

  395.                         *outptr++ = val;

  396.                     }

  397.                     outptr -= 8;

  398.                     break;

  399. 

  400.                 case 1:

  401.                     for (int i = 0; i < 4; i++) {

  402.                         __fp16 val = (x + i < xmax) ? (*inptr++) : (__fp16)(0);

  403.                         *outptr++ = val;

  404.                     }

  405.                     outptr -= 4;

  406.                     break;

  407. 

  408.                 default:

  409.                     megdnn_assert(0);

  410.             }

  411.             outptr += ksize4;

  412.         }

  413.     }

  414. }

  415. 

  416. void transpose_1x24(

  417.         __fp16* out, const __fp16* in, const int ldin, const int x0, const int xmax,

  418.         const int k0, const int kmax) {

  419.     int ksize = kmax - k0;

  420.     int ksize24 = ksize * 24;

  421.     int ksize16 = (ksize << 4);

  422.     int ksize8 = (ksize << 3);

  423.     int ksize4 = (ksize << 2);

  424.     int k = ksize;

  425.     __fp16* outptr_base = out;

  426.     __fp16* outptr_base16 = out;

  427.     __fp16* outptr_base8 = out;

  428.     __fp16* outptr_base4 = out;

  429. 

  430.     const __fp16* inptr_base = in + x0 + k0 * ldin;

  431. 

  432.     for (; k > 3; k -= 4) {

  433.         __fp16* outptr = outptr_base;

  434. 

  435.         const __fp16* inptr = inptr_base;

  436.         const __fp16* inptr1 = inptr + ldin;

  437.         const __fp16* inptr2 = inptr1 + ldin;

  438.         const __fp16* inptr3 = inptr2 + ldin;

  439. 

  440.         prefetch_3x(inptr);

  441.         prefetch_3x(inptr1);

  442.         prefetch_3x(inptr2);

  443.         prefetch_3x(inptr3);

  444. 

  445.         int x = x0;

  446.         for (; x + 24 <= xmax; x += 24) {

  447.             transpose_24x4_1_h(inptr, inptr1, inptr2, inptr3, outptr);

  448.             outptr += ksize24;

  449.         }

  450.         outptr += outptr_base16 - outptr_base;

  451.         for (; x + 16 <= xmax; x += 16) {

  452.             transpose_16x4_1_h(inptr, inptr1, inptr2, inptr3, outptr);

  453.             outptr += ksize16;

  454.         }

  455.         outptr += outptr_base8 - outptr_base16;

  456.         for (; x + 8 <= xmax; x += 8) {

  457.             transpose_8x4_1_h(inptr, inptr1, inptr2, inptr3, outptr);

  458.             outptr += ksize8;

  459.         }

  460.         outptr += outptr_base4 - outptr_base8;

  461.         for (; x < xmax; x += 4) {

  462.             for (int i = 0; i < 4; i++) {

  463.                 __fp16 val = (x + i < xmax) ? (*inptr++) : (__fp16)(0);

  464.                 *outptr++ = val;

  465.             }

  466.             for (int i = 0; i < 4; i++) {

  467.                 __fp16 val = (x + i < xmax) ? (*inptr1++) : (__fp16)(0);

  468.                 *outptr++ = val;

  469.             }

  470.             for (int i = 0; i < 4; i++) {

  471.                 __fp16 val = (x + i < xmax) ? (*inptr2++) : (__fp16)(0);

  472.                 *outptr++ = val;

  473.             }

  474.             for (int i = 0; i < 4; i++) {

  475.                 __fp16 val = (x + i < xmax) ? (*inptr3++) : (__fp16)(0);

  476.                 *outptr++ = val;

  477.             }

  478.             outptr -= 16;

  479.             outptr += ksize4;

  480.         }

  481. 

  482.         inptr_base += ldin * 4;

  483.         outptr_base += 24 * 4;

  484.         outptr_base16 += 16 * 4;

  485.         outptr_base8 += 8 * 4;

  486.         outptr_base4 += 4 * 4;

  487.     }

  488. 

  489.     if (k) {

  490.         __fp16* outptr = outptr_base;

  491.         const __fp16* inptr = inptr_base;

  492.         const __fp16* inptr1 = inptr + ldin;

  493.         const __fp16* inptr2 = inptr1 + ldin;

  494. 

  495.         prefetch_3x(inptr);

  496.         prefetch_3x(inptr1);

  497.         prefetch_3x(inptr2);

  498. 

  499.         int x = x0;

  500.         for (; x + 24 <= xmax; x += 24) {

  501.             switch (k) {

  502.                 case 3:

  503.                     transpose_24x2_1_h(inptr, inptr1, outptr);

  504.                     transpose_24x1_1_h(inptr2, outptr + 24 * 2);

  505.                     break;

  506. 

  507.                 case 2:

  508.                     transpose_24x2_1_h(inptr, inptr1, outptr);

  509.                     break;

  510. 

  511.                 case 1:

  512.                     transpose_24x1_1_h(inptr, outptr);

  513.                     break;

  514. 

  515.                 default:

  516.                     megdnn_assert(0);

  517.             }

  518.             outptr += ksize24;

  519.         }

  520. 

  521.         outptr += outptr_base16 - outptr_base;

  522.         for (; x + 16 <= xmax; x += 16) {

  523.             switch (k) {

  524.                 case 3:

  525.                     transpose_16x2_1_h(inptr, inptr1, outptr);

  526.                     transpose_16x1_1_h(inptr2, outptr + 16 * 2);

  527.                     break;

  528. 

  529.                 case 2:

  530.                     transpose_16x2_1_h(inptr, inptr1, outptr);

  531.                     break;

  532. 

  533.                 case 1:

  534.                     transpose_16x1_1_h(inptr, outptr);

  535.                     break;

  536. 

  537.                 default:

  538.                     megdnn_assert(0);

  539.             }

  540.             outptr += ksize16;

  541.         }

  542. 

  543.         outptr += outptr_base8 - outptr_base16;

  544.         for (; x + 8 <= xmax; x += 8) {

  545.             switch (k) {

  546.                 case 3:

  547.                     transpose_8x2_1_h(inptr, inptr1, outptr);

  548.                     transpose_8x1_1_h(inptr2, outptr + 8 * 2);

  549.                     break;

  550. 

  551.                 case 2:

  552.                     transpose_8x2_1_h(inptr, inptr1, outptr);

  553.                     break;

  554. 

  555.                 case 1:

  556.                     transpose_8x1_1_h(inptr, outptr);

  557.                     break;

  558. 

  559.                 default:

  560.                     megdnn_assert(0);

  561.             }

  562.             outptr += ksize8;

  563.         }

  564. 

  565.         outptr += outptr_base4 - outptr_base8;

  566.         for (; x < xmax; x += 4) {

  567.             switch (k) {

  568.                 case 3:

  569.                     for (int i = 0; i < 4; i++) {

  570.                         __fp16 val = (x + i < xmax) ? (*inptr++) : (__fp16)(0);

  571.                         *outptr++ = val;

  572.                     }

  573.                     for (int i = 0; i < 4; i++) {

  574.                         __fp16 val = (x + i < xmax) ? (*inptr1++) : (__fp16)(0);

  575.                         *outptr++ = val;

  576.                     }

  577.                     for (int i = 0; i < 4; i++) {

  578.                         __fp16 val = (x + i < xmax) ? (*inptr2++) : (__fp16)(0);

  579.                         *outptr++ = val;

  580.                     }

  581.                     outptr -= 12;

  582.                     break;

  583.                 case 2:

  584.                     for (int i = 0; i < 4; i++) {

  585.                         __fp16 val = (x + i < xmax) ? (*inptr++) : (__fp16)(0);

  586.                         *outptr++ = val;

  587.                     }

  588.                     for (int i = 0; i < 4; i++) {

  589.                         __fp16 val = (x + i < xmax) ? (*inptr1++) : (__fp16)(0);

  590.                         *outptr++ = val;

  591.                     }

  592.                     outptr -= 8;

  593.                     break;

  594. 

  595.                 case 1:

  596.                     for (int i = 0; i < 4; i++) {

  597.                         __fp16 val = (x + i < xmax) ? (*inptr++) : (__fp16)(0);

  598.                         *outptr++ = val;

  599.                     }

  600.                     outptr -= 4;

  601.                     break;

  602. 

  603.                 default:

  604.                     megdnn_assert(0);

  605.             }

  606.             outptr += ksize4;

  607.         }

  608.     }

  609. }

  610. 

  611. // Overview of register layout:

  612. //

  613. // A 2x24 cell of Rhs is stored in 16bit in q2-q7.

  614. // A 8x2 cell of Lhs is stored in 16bit in q0-q1

  615. // A 8x24 block of accumulators is stored in 16bit in q8--q31.

  616. //

  617. //                   +--------+--------+--------+

  618. //                   | v2[0-7]| v3[0-7]| v4[0-7]|

  619. //              Rhs  +--------+--------+--------+

  620. //                   | v5[0-7]| v6[0-7]| v7[0-7]|

  621. //                   +--------+--------+--------+

  622. //

  623. //                   |        |        |        |

  624. //

  625. //    Lhs            |        |        |        |

  626. //

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

  628. //  |v0|v1|          | v8[0-7]|v16[0-7]|v24[0-7]|

  629. //  |v0|v1|          | v9[0-7]|v17[0-7]|v25[0-7]|

  630. //  |v0|v1|          |v10[0-7]|v18[0-7]|v26[0-7]|

  631. //  |v0|v1|          |v11[0-7]|v19[0-7]|v27[0-7]|

  632. //  |v0|v1|          |v12[0-7]|v20[0-7]|v28[0-7]|

  633. //  |v0|v1|          |v13[0-7]|v21[0-7]|v29[0-7]|

  634. //  |v0|v1|          |v14[0-7]|v22[0-7]|v30[0-7]|

  635. //  |v0|v1|          |v15[0-7]|v23[0-7]|v31[0-7]|

  636. //  +--+--+ - - - -  +--------+--------+--------+

  637. //

  638. //                            Accumulator

  639. 

  640. void aarch64_hgemm_assembly_kernel_24x8(

  641.         const __fp16* a_ptr, const __fp16*& b_ptr, int K, __fp16* outptr0, int ldout,

  642.         int type) {

  643.     int oddk = (K & 1);

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

  645. 

  646.     register float16x8_t a0 asm("v0");

  647.     register float16x8_t a0a asm("v1");

  648.     register float16x8_t b0 asm("v2");

  649.     register float16x8_t b1 asm("v3");

  650.     register float16x8_t b2 asm("v4");

  651.     register float16x8_t b0a asm("v5");

  652.     register float16x8_t b1a asm("v6");

  653.     register float16x8_t b2a asm("v7");

  654. 

  655.     __fp16* outptr1 = outptr0 + ldout;

  656.     __fp16* outptr2 = outptr1 + ldout;

  657.     __fp16* outptr3 = outptr2 + ldout;

  658.     __fp16* outptr4 = outptr3 + ldout;

  659.     __fp16* outptr5 = outptr4 + ldout;

  660.     __fp16* outptr6 = outptr5 + ldout;

  661.     __fp16* outptr7 = outptr6 + ldout;

  662. 

  663.     asm volatile(

  664.             ".arch armv8.2-a+fp16\n"

  665. 

  666.             // load accumulator C

  667.             "cmp %w[type], #0\n"

  668.             "beq 5f\n"

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  685.             "b 6f\n"

  686. 

  687.             "5:\n"

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

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

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

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

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

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

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

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

  696. 

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

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

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

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

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

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

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

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

  705. 

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

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

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

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

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

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

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

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

  714. 

  715.             "6:\n"

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

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

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

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

  720.             "ldr %q[b0a], [%[b_ptr], #48]\n"

  721.             "ldr %q[b1a], [%[b_ptr], #64]\n"

  722. 

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

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

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

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

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

  728. 

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

  730. 

  731.             "1:\n"

  732.             "fmla v8.8h , %[b0].8h, %[a0].h[0]\n"

  733.             "fmla v9.8h , %[b0].8h, %[a0].h[1]\n"

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

  735.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  736.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  737.             "ldr %q[b2a], [%[b_ptr], #80]\n"

  738.             "fmla v12.8h, %[b0].8h, %[a0].h[4]\n"

  739.             "fmla v13.8h, %[b0].8h, %[a0].h[5]\n"

  740.             "fmla v14.8h, %[b0].8h, %[a0].h[6]\n"

  741.             "fmla v15.8h, %[b0].8h, %[a0].h[7]\n"

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

  743. 

  744.             "fmla v16.8h, %[b1].8h, %[a0].h[0]\n"

  745.             "fmla v17.8h, %[b1].8h, %[a0].h[1]\n"

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

  747.             "fmla v18.8h, %[b1].8h, %[a0].h[2]\n"

  748.             "fmla v19.8h, %[b1].8h, %[a0].h[3]\n"

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

  750.             "fmla v20.8h, %[b1].8h, %[a0].h[4]\n"

  751.             "fmla v21.8h, %[b1].8h, %[a0].h[5]\n"

  752.             "fmla v22.8h, %[b1].8h, %[a0].h[6]\n"

  753.             "fmla v23.8h, %[b1].8h, %[a0].h[7]\n"

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

  755. 

  756.             "fmla v24.8h, %[b2].8h, %[a0].h[0]\n"

  757.             "fmla v25.8h, %[b2].8h, %[a0].h[1]\n"

  758.             ASM_PREFETCH("[%[b_ptr], #288]")

  759.             "fmla v26.8h, %[b2].8h, %[a0].h[2]\n"

  760.             "fmla v27.8h, %[b2].8h, %[a0].h[3]\n"

  761.             "fmla v28.8h, %[b2].8h, %[a0].h[4]\n"

  762.             "fmla v29.8h, %[b2].8h, %[a0].h[5]\n"

  763.             "fmla v30.8h, %[b2].8h, %[a0].h[6]\n"

  764.             "fmla v31.8h, %[b2].8h, %[a0].h[7]\n"

  765.             "ldr %q[a0], [%[a_ptr], #32]\n"

  766. 

  767.             "fmla v8.8h , %[b0a].8h, %[a0a].h[0]\n"

  768.             "fmla v9.8h , %[b0a].8h, %[a0a].h[1]\n"

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

  770. 

  771.             "fmla v10.8h, %[b0a].8h, %[a0a].h[2]\n"

  772.             "fmla v11.8h, %[b0a].8h, %[a0a].h[3]\n"

  773.             "fmla v12.8h, %[b0a].8h, %[a0a].h[4]\n"

  774.             "fmla v13.8h, %[b0a].8h, %[a0a].h[5]\n"

  775.             "fmla v14.8h, %[b0a].8h, %[a0a].h[6]\n"

  776.             "fmla v15.8h, %[b0a].8h, %[a0a].h[7]\n"

  777.             "ldr %q[b0a], [%[b_ptr], #48]\n"

  778. 

  779.             "fmla v16.8h, %[b1a].8h, %[a0a].h[0]\n"

  780.             "fmla v17.8h, %[b1a].8h, %[a0a].h[1]\n"

  781.             ASM_PREFETCH("[%[b_ptr], #352]")

  782.             "fmla v18.8h, %[b1a].8h, %[a0a].h[2]\n"

  783.             "fmla v19.8h, %[b1a].8h, %[a0a].h[3]\n"

  784.             "fmla v20.8h, %[b1a].8h, %[a0a].h[4]\n"

  785.             "fmla v21.8h, %[b1a].8h, %[a0a].h[5]\n"

  786.             "fmla v22.8h, %[b1a].8h, %[a0a].h[6]\n"

  787.             "fmla v23.8h, %[b1a].8h, %[a0a].h[7]\n"

  788.             "ldr %q[b1a], [%[b_ptr], #64]\n"

  789. 

  790.             "fmla v24.8h, %[b2a].8h, %[a0a].h[0]\n"

  791.             "fmla v25.8h, %[b2a].8h, %[a0a].h[1]\n"

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

  793.             "fmla v26.8h, %[b2a].8h, %[a0a].h[2]\n"

  794.             "fmla v27.8h, %[b2a].8h, %[a0a].h[3]\n"

  795.             "fmla v28.8h, %[b2a].8h, %[a0a].h[4]\n"

  796.             "fmla v29.8h, %[b2a].8h, %[a0a].h[5]\n"

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

  798.             "fmla v30.8h, %[b2a].8h, %[a0a].h[6]\n"

  799.             "fmla v31.8h, %[b2a].8h, %[a0a].h[7]\n"

  800. 

  801.             "bne 1b\n"

  802.             "4:\n"

  803.             // Jump to odd tail if necessary.

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

  805. 

  806.             // Even tail

  807.             "fmla v8.8h , %[b0].8h, %[a0].h[0]\n"

  808.             "fmla v9.8h , %[b0].8h, %[a0].h[1]\n"

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

  810.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  811.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  812.             "ldr %q[b2a], [%[b_ptr], #80]\n"

  813.             "fmla v12.8h, %[b0].8h, %[a0].h[4]\n"

  814.             "fmla v13.8h, %[b0].8h, %[a0].h[5]\n"

  815.             "fmla v14.8h, %[b0].8h, %[a0].h[6]\n"

  816.             "fmla v15.8h, %[b0].8h, %[a0].h[7]\n"

  817. 

  818.             "fmla v16.8h, %[b1].8h, %[a0].h[0]\n"

  819.             "fmla v17.8h, %[b1].8h, %[a0].h[1]\n"

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

  821.             "fmla v18.8h, %[b1].8h, %[a0].h[2]\n"

  822.             "fmla v19.8h, %[b1].8h, %[a0].h[3]\n"

  823.             "fmla v20.8h, %[b1].8h, %[a0].h[4]\n"

  824.             "fmla v21.8h, %[b1].8h, %[a0].h[5]\n"

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

  826.             "fmla v22.8h, %[b1].8h, %[a0].h[6]\n"

  827.             "fmla v23.8h, %[b1].8h, %[a0].h[7]\n"

  828. 

  829.             "fmla v24.8h, %[b2].8h, %[a0].h[0]\n"

  830.             "fmla v25.8h, %[b2].8h, %[a0].h[1]\n"

  831.             "fmla v26.8h, %[b2].8h, %[a0].h[2]\n"

  832.             "fmla v27.8h, %[b2].8h, %[a0].h[3]\n"

  833.             "fmla v28.8h, %[b2].8h, %[a0].h[4]\n"

  834.             "fmla v29.8h, %[b2].8h, %[a0].h[5]\n"

  835.             "fmla v30.8h, %[b2].8h, %[a0].h[6]\n"

  836.             "fmla v31.8h, %[b2].8h, %[a0].h[7]\n"

  837. 

  838.             "fmla v8.8h , %[b0a].8h, %[a0a].h[0]\n"

  839.             "fmla v16.8h, %[b1a].8h, %[a0a].h[0]\n"

  840.             "str q8, [%[outptr0]]\n"

  841.             "fmla v24.8h, %[b2a].8h, %[a0a].h[0]\n"

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

  843. 

  844.             "fmla v9.8h , %[b0a].8h, %[a0a].h[1]\n"

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

  846.             "fmla v17.8h, %[b1a].8h, %[a0a].h[1]\n"

  847.             "str q9, [%[outptr1]]\n"

  848.             "fmla v25.8h, %[b2a].8h, %[a0a].h[1]\n"

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

  850. 

  851.             "fmla v10.8h, %[b0a].8h, %[a0a].h[2]\n"

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

  853.             "fmla v18.8h, %[b1a].8h, %[a0a].h[2]\n"

  854.             "str q10, [%[outptr2]]\n"

  855.             "fmla v26.8h, %[b2a].8h, %[a0a].h[2]\n"

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

  857. 

  858.             "fmla v11.8h, %[b0a].8h, %[a0a].h[3]\n"

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

  860.             "fmla v19.8h, %[b1a].8h, %[a0a].h[3]\n"

  861.             "str q11, [%[outptr3]]\n"

  862.             "fmla v27.8h, %[b2a].8h, %[a0a].h[3]\n"

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

  864. 

  865.             "fmla v12.8h, %[b0a].8h, %[a0a].h[4]\n"

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

  867.             "fmla v20.8h, %[b1a].8h, %[a0a].h[4]\n"

  868.             "str q12, [%[outptr4]]\n"

  869.             "fmla v28.8h, %[b2a].8h, %[a0a].h[4]\n"

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

  871. 

  872.             "fmla v13.8h, %[b0a].8h, %[a0a].h[5]\n"

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

  874.             "fmla v21.8h, %[b1a].8h, %[a0a].h[5]\n"

  875.             "str q13, [%[outptr5]]\n"

  876.             "fmla v29.8h, %[b2a].8h, %[a0a].h[5]\n"

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

  878. 

  879.             "fmla v14.8h, %[b0a].8h, %[a0a].h[6]\n"

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

  881.             "fmla v22.8h, %[b1a].8h, %[a0a].h[6]\n"

  882.             "str q14, [%[outptr6]]\n"

  883.             "fmla v30.8h, %[b2a].8h, %[a0a].h[6]\n"

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

  885. 

  886.             "fmla v15.8h, %[b0a].8h, %[a0a].h[7]\n"

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

  888.             "fmla v23.8h, %[b1a].8h, %[a0a].h[7]\n"

  889.             "str q15, [%[outptr7]]\n"

  890.             "fmla v31.8h, %[b2a].8h, %[a0a].h[7]\n"

  891.             "b 3f\n"

  892. 

  893.             // Odd tail

  894.             "2:\n"

  895.             "add %[a_ptr], %[a_ptr], #16\n"

  896.             "fmla v8.8h , %[b0].8h, %[a0].h[0]\n"

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

  898.             "fmla v16.8h, %[b1].8h, %[a0].h[0]\n"

  899.             "str q8, [%[outptr0]]\n"

  900.             "fmla v24.8h, %[b2].8h, %[a0].h[0]\n"

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

  902. 

  903.             "fmla v9.8h , %[b0].8h, %[a0].h[1]\n"

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

  905.             "fmla v17.8h, %[b1].8h, %[a0].h[1]\n"

  906.             "str q9, [%[outptr1]]\n"

  907.             "fmla v25.8h, %[b2].8h, %[a0].h[1]\n"

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

  909. 

  910.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

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

  912.             "fmla v18.8h, %[b1].8h, %[a0].h[2]\n"

  913.             "str q10, [%[outptr2]]\n"

  914.             "fmla v26.8h, %[b2].8h, %[a0].h[2]\n"

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

  916. 

  917.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

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

  919.             "fmla v19.8h, %[b1].8h, %[a0].h[3]\n"

  920.             "str q11, [%[outptr3]]\n"

  921.             "fmla v27.8h, %[b2].8h, %[a0].h[3]\n"

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

  923. 

  924.             "fmla v12.8h, %[b0].8h, %[a0].h[4]\n"

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

  926.             "fmla v20.8h, %[b1].8h, %[a0].h[4]\n"

  927.             "str q12, [%[outptr4]]\n"

  928.             "fmla v28.8h, %[b2].8h, %[a0].h[4]\n"

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

  930. 

  931.             "fmla v13.8h, %[b0].8h, %[a0].h[5]\n"

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

  933.             "fmla v21.8h, %[b1].8h, %[a0].h[5]\n"

  934.             "str q13, [%[outptr5]]\n"

  935.             "fmla v29.8h, %[b2].8h, %[a0].h[5]\n"

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

  937. 

  938.             "fmla v14.8h, %[b0].8h, %[a0].h[6]\n"

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

  940.             "fmla v22.8h, %[b1].8h, %[a0].h[6]\n"

  941.             "str q14, [%[outptr6]]\n"

  942.             "fmla v30.8h, %[b2].8h, %[a0].h[6]\n"

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

  944. 

  945.             "fmla v15.8h, %[b0].8h, %[a0].h[7]\n"

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

  947.             "fmla v23.8h, %[b1].8h, %[a0].h[7]\n"

  948.             "str q15, [%[outptr7]]\n"

  949.             "fmla v31.8h, %[b2].8h, %[a0].h[7]\n"

  950. 

  951.             "3:\n"

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

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

  954.             : [a0] "+w"(a0), [a0a] "+w"(a0a), [b0] "+w"(b0), [b1] "+w"(b1),

  955.               [b2] "+w"(b2), [k] "+r"(k), [b0a] "+w"(b0a), [b1a] "+w"(b1a),

  956.               [b2a] "+w"(b2a), [a_ptr] "+r"(a_ptr), [b_ptr] "+r"(b_ptr),

  957.               [outptr0] "+r"(outptr0), [outptr1] "+r"(outptr1),

  958.               [outptr2] "+r"(outptr2), [outptr3] "+r"(outptr3),

  959.               [outptr4] "+r"(outptr4), [outptr5] "+r"(outptr5),

  960.               [outptr6] "+r"(outptr6), [outptr7] "+r"(outptr7)

  961.             : [oddk] "r"(oddk), [type] "r"(type)

  962.             : "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16",

  963.               "v17", "v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25",

  964.               "v26", "v27", "v28", "v29", "v30", "v31", "cc", "memory");

  965. }

  966. 

  967. // Overview of register layout:

  968. //

  969. // A 2x16 cell of Rhs is stored in 16bit in q2,q3,q5,q6.

  970. // A 8x2 cell of Lhs is stored in 16bit in q0-q1

  971. // A 8x16 block of accumulators is stored in 16bit in q8-q15, q16-q23.

  972. //

  973. //                   +--------+--------+

  974. //                   | v2[0-7]| v3[0-7]|

  975. //              Rhs  +--------+--------+

  976. //                   | v5[0-7]| v6[0-7]|

  977. //                   +--------+--------+

  978. //

  979. //                   |        |        |

  980. //

  981. //    Lhs            |        |        |

  982. //

  983. //  +--+--+ - - - -  +--------+--------+

  984. //  |v0|v1|          | v8[0-7]|v16[0-7]|

  985. //  |v0|v1|          | v9[0-7]|v17[0-7]|

  986. //  |v0|v1|          |v10[0-7]|v18[0-7]|

  987. //  |v0|v1|          |v11[0-7]|v19[0-7]|

  988. //  |v0|v1|          |v12[0-7]|v20[0-7]|

  989. //  |v0|v1|          |v13[0-7]|v21[0-7]|

  990. //  |v0|v1|          |v14[0-7]|v22[0-7]|

  991. //  |v0|v1|          |v15[0-7]|v23[0-7]|

  992. //  +--+--+ - - - -  +--------+--------+

  993. //

  994. //                        Accumulator

  995. void aarch64_hgemm_assembly_kernel_16x8(

  996.         const __fp16* a_ptr, const __fp16*& b_ptr, int K, __fp16* outptr0, int ldout,

  997.         int type) {

  998.     int oddk = (K & 1);

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

  1000. 

  1001.     register float16x8_t a0 asm("v0");

  1002.     register float16x8_t a0a asm("v1");

  1003.     register float16x8_t b0 asm("v2");

  1004.     register float16x8_t b1 asm("v3");

  1005.     register float16x8_t b0a asm("v5");

  1006.     register float16x8_t b1a asm("v6");

  1007. 

  1008.     __fp16* outptr1 = outptr0 + ldout;

  1009.     __fp16* outptr2 = outptr1 + ldout;

  1010.     __fp16* outptr3 = outptr2 + ldout;

  1011.     __fp16* outptr4 = outptr3 + ldout;

  1012.     __fp16* outptr5 = outptr4 + ldout;

  1013.     __fp16* outptr6 = outptr5 + ldout;

  1014.     __fp16* outptr7 = outptr6 + ldout;

  1015. 

  1016.     asm volatile(

  1017.             ".arch armv8.2-a+fp16\n"

  1018. 

  1019.             // load accumulator C

  1020.             "cmp %w[type], #0\n"

  1021.             "beq 5f\n"

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

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

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

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

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

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

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

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

  1030.             "b 6f\n"

  1031. 

  1032.             "5:\n"

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

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

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

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

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

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

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

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

  1041. 

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

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

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

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

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

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

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

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

  1050. 

  1051.             "6:\n"

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

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

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

  1055.             "ldr %q[b0a], [%[b_ptr], #32]\n"

  1056.             "ldr %q[b1a], [%[b_ptr], #48]\n"

  1057. 

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

  1059. 

  1060.             "1:\n"

  1061.             "fmla v8.8h , %[b0].8h, %[a0].h[0]\n"

  1062.             "fmla v9.8h , %[b0].8h, %[a0].h[1]\n"

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

  1064.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  1065.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  1066.             "fmla v12.8h, %[b0].8h, %[a0].h[4]\n"

  1067.             "fmla v13.8h, %[b0].8h, %[a0].h[5]\n"

  1068.             "fmla v14.8h, %[b0].8h, %[a0].h[6]\n"

  1069.             "fmla v15.8h, %[b0].8h, %[a0].h[7]\n"

  1070.             "ldr %q[b0], [%[b_ptr], #64]\n"

  1071. 

  1072.             "fmla v16.8h, %[b1].8h, %[a0].h[0]\n"

  1073.             "fmla v17.8h, %[b1].8h, %[a0].h[1]\n"

  1074.             "fmla v18.8h, %[b1].8h, %[a0].h[2]\n"

  1075.             "fmla v19.8h, %[b1].8h, %[a0].h[3]\n"

  1076.             "add %[b_ptr], %[b_ptr], #64\n"

  1077.             "fmla v20.8h, %[b1].8h, %[a0].h[4]\n"

  1078.             "fmla v21.8h, %[b1].8h, %[a0].h[5]\n"

  1079.             "fmla v22.8h, %[b1].8h, %[a0].h[6]\n"

  1080.             "fmla v23.8h, %[b1].8h, %[a0].h[7]\n"

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

  1082. 

  1083.             "ldr %q[a0], [%[a_ptr], #32]\n"

  1084. 

  1085.             "fmla v8.8h , %[b0a].8h, %[a0a].h[0]\n"

  1086.             "fmla v9.8h , %[b0a].8h, %[a0a].h[1]\n"

  1087.             "fmla v10.8h, %[b0a].8h, %[a0a].h[2]\n"

  1088.             "fmla v11.8h, %[b0a].8h, %[a0a].h[3]\n"

  1089.             "fmla v12.8h, %[b0a].8h, %[a0a].h[4]\n"

  1090.             "fmla v13.8h, %[b0a].8h, %[a0a].h[5]\n"

  1091.             "fmla v14.8h, %[b0a].8h, %[a0a].h[6]\n"

  1092.             "fmla v15.8h, %[b0a].8h, %[a0a].h[7]\n"

  1093.             "ldr %q[b0a], [%[b_ptr], #32]\n"

  1094. 

  1095.             "fmla v16.8h, %[b1a].8h, %[a0a].h[0]\n"

  1096.             "fmla v17.8h, %[b1a].8h, %[a0a].h[1]\n"

  1097.             "fmla v18.8h, %[b1a].8h, %[a0a].h[2]\n"

  1098.             "fmla v19.8h, %[b1a].8h, %[a0a].h[3]\n"

  1099.             "fmla v20.8h, %[b1a].8h, %[a0a].h[4]\n"

  1100.             "fmla v21.8h, %[b1a].8h, %[a0a].h[5]\n"

  1101.             "fmla v22.8h, %[b1a].8h, %[a0a].h[6]\n"

  1102.             "fmla v23.8h, %[b1a].8h, %[a0a].h[7]\n"

  1103.             "ldr %q[b1a], [%[b_ptr], #48]\n"

  1104. 

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

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

  1107. 

  1108.             "bne 1b\n"

  1109.             "4:\n"

  1110.             // Jump to odd tail if necessary.

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

  1112. 

  1113.             // Even tail

  1114.             "fmla v8.8h , %[b0].8h, %[a0].h[0]\n"

  1115.             "fmla v9.8h , %[b0].8h, %[a0].h[1]\n"

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

  1117.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  1118.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  1119.             "fmla v12.8h, %[b0].8h, %[a0].h[4]\n"

  1120.             "fmla v13.8h, %[b0].8h, %[a0].h[5]\n"

  1121.             "fmla v14.8h, %[b0].8h, %[a0].h[6]\n"

  1122.             "fmla v15.8h, %[b0].8h, %[a0].h[7]\n"

  1123. 

  1124.             "fmla v16.8h, %[b1].8h, %[a0].h[0]\n"

  1125.             "fmla v17.8h, %[b1].8h, %[a0].h[1]\n"

  1126.             "add %[b_ptr], %[b_ptr], #64\n"

  1127.             "fmla v18.8h, %[b1].8h, %[a0].h[2]\n"

  1128.             "fmla v19.8h, %[b1].8h, %[a0].h[3]\n"

  1129.             "fmla v20.8h, %[b1].8h, %[a0].h[4]\n"

  1130.             "fmla v21.8h, %[b1].8h, %[a0].h[5]\n"

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

  1132.             "fmla v22.8h, %[b1].8h, %[a0].h[6]\n"

  1133.             "fmla v23.8h, %[b1].8h, %[a0].h[7]\n"

  1134. 

  1135.             "fmla v8.8h , %[b0a].8h, %[a0a].h[0]\n"

  1136.             "fmla v16.8h, %[b1a].8h, %[a0a].h[0]\n"

  1137.             "str q8, [%[outptr0]]\n"

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

  1139. 

  1140.             "fmla v9.8h , %[b0a].8h, %[a0a].h[1]\n"

  1141.             "fmla v17.8h, %[b1a].8h, %[a0a].h[1]\n"

  1142.             "str q9, [%[outptr1]]\n"

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

  1144. 

  1145.             "fmla v10.8h, %[b0a].8h, %[a0a].h[2]\n"

  1146.             "fmla v18.8h, %[b1a].8h, %[a0a].h[2]\n"

  1147.             "str q10, [%[outptr2]]\n"

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

  1149. 

  1150.             "fmla v11.8h, %[b0a].8h, %[a0a].h[3]\n"

  1151.             "fmla v19.8h, %[b1a].8h, %[a0a].h[3]\n"

  1152.             "str q11, [%[outptr3]]\n"

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

  1154. 

  1155.             "fmla v12.8h, %[b0a].8h, %[a0a].h[4]\n"

  1156.             "fmla v20.8h, %[b1a].8h, %[a0a].h[4]\n"

  1157.             "str q12, [%[outptr4]]\n"

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

  1159. 

  1160.             "fmla v13.8h, %[b0a].8h, %[a0a].h[5]\n"

  1161.             "fmla v21.8h, %[b1a].8h, %[a0a].h[5]\n"

  1162.             "str q13, [%[outptr5]]\n"

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

  1164. 

  1165.             "fmla v14.8h, %[b0a].8h, %[a0a].h[6]\n"

  1166.             "fmla v22.8h, %[b1a].8h, %[a0a].h[6]\n"

  1167.             "str q14, [%[outptr6]]\n"

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

  1169. 

  1170.             "fmla v15.8h, %[b0a].8h, %[a0a].h[7]\n"

  1171.             "fmla v23.8h, %[b1a].8h, %[a0a].h[7]\n"

  1172.             "str q15, [%[outptr7]]\n"

  1173.             "b 3f\n"

  1174. 

  1175.             // Odd tail

  1176.             "2:\n"

  1177.             "add %[a_ptr], %[a_ptr], #16\n"

  1178.             "fmla v8.8h , %[b0].8h, %[a0].h[0]\n"

  1179.             "add %[b_ptr], %[b_ptr], #32\n"

  1180.             "fmla v16.8h, %[b1].8h, %[a0].h[0]\n"

  1181.             "str q8, [%[outptr0]]\n"

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

  1183. 

  1184.             "fmla v9.8h , %[b0].8h, %[a0].h[1]\n"

  1185.             "fmla v17.8h, %[b1].8h, %[a0].h[1]\n"

  1186.             "str q9, [%[outptr1]]\n"

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

  1188. 

  1189.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  1190.             "fmla v18.8h, %[b1].8h, %[a0].h[2]\n"

  1191.             "str q10, [%[outptr2]]\n"

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

  1193. 

  1194.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  1195.             "fmla v19.8h, %[b1].8h, %[a0].h[3]\n"

  1196.             "str q11, [%[outptr3]]\n"

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

  1198. 

  1199.             "fmla v12.8h, %[b0].8h, %[a0].h[4]\n"

  1200.             "fmla v20.8h, %[b1].8h, %[a0].h[4]\n"

  1201.             "str q12, [%[outptr4]]\n"

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

  1203. 

  1204.             "fmla v13.8h, %[b0].8h, %[a0].h[5]\n"

  1205.             "fmla v21.8h, %[b1].8h, %[a0].h[5]\n"

  1206.             "str q13, [%[outptr5]]\n"

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

  1208. 

  1209.             "fmla v14.8h, %[b0].8h, %[a0].h[6]\n"

  1210.             "fmla v22.8h, %[b1].8h, %[a0].h[6]\n"

  1211.             "str q14, [%[outptr6]]\n"

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

  1213. 

  1214.             "fmla v15.8h, %[b0].8h, %[a0].h[7]\n"

  1215.             "fmla v23.8h, %[b1].8h, %[a0].h[7]\n"

  1216.             "str q15, [%[outptr7]]\n"

  1217. 

  1218.             "3:\n"

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

  1220.             : [a0] "+w"(a0), [a0a] "+w"(a0a), [b0] "+w"(b0), [b1] "+w"(b1), [k] "+r"(k),

  1221.               [b0a] "+w"(b0a), [b1a] "+w"(b1a), [a_ptr] "+r"(a_ptr),

  1222.               [b_ptr] "+r"(b_ptr), [outptr0] "+r"(outptr0), [outptr1] "+r"(outptr1),

  1223.               [outptr2] "+r"(outptr2), [outptr3] "+r"(outptr3), [outptr4] "+r"(outptr4),

  1224.               [outptr5] "+r"(outptr5), [outptr6] "+r"(outptr6), [outptr7] "+r"(outptr7)

  1225.             : [oddk] "r"(oddk), [type] "r"(type)

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

  1227.               "v19", "v20", "v21", "v22", "v23", "cc", "memory");

  1228. }

  1229. 

  1230. // Overview of register layout:

  1231. //

  1232. // A 2x8 cell of Rhs is stored in 16bit in q2,q5

  1233. // A 8x2 cell of Lhs is stored in 16bit in q0-q1

  1234. // A 8x8 block of accumulators is stored in 16bit in q8-q15.

  1235. //

  1236. //                   +--------+

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

  1238. //              Rhs  +--------+

  1239. //                   | v5[0-7]|

  1240. //                   +--------+

  1241. //

  1242. //                   |        |

  1243. //

  1244. //    Lhs            |        |

  1245. //

  1246. //  +--+--+ - - - -  +--------+

  1247. //  |v0|v1|          | v8[0-7]|

  1248. //  |v0|v1|          | v9[0-7]|

  1249. //  |v0|v1|          |v10[0-7]|

  1250. //  |v0|v1|          |v11[0-7]|

  1251. //  |v0|v1|          |v12[0-7]|

  1252. //  |v0|v1|          |v13[0-7]|

  1253. //  |v0|v1|          |v14[0-7]|

  1254. //  |v0|v1|          |v15[0-7]|

  1255. //  +--+--+ - - - -  +--------+

  1256. //

  1257. //                  Accumulator

  1258. void aarch64_hgemm_assembly_kernel_8x8(

  1259.         const __fp16* a_ptr, const __fp16*& b_ptr, int K, __fp16* outptr0, int ldout,

  1260.         int type) {

  1261.     int oddk = (K & 1);

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

  1263. 

  1264.     register float16x8_t a0 asm("v0");

  1265.     register float16x8_t a0a asm("v1");

  1266.     register float16x8_t b0 asm("v2");

  1267.     register float16x8_t b0a asm("v5");

  1268. 

  1269.     __fp16* outptr1 = outptr0 + ldout;

  1270.     __fp16* outptr2 = outptr1 + ldout;

  1271.     __fp16* outptr3 = outptr2 + ldout;

  1272.     __fp16* outptr4 = outptr3 + ldout;

  1273.     __fp16* outptr5 = outptr4 + ldout;

  1274.     __fp16* outptr6 = outptr5 + ldout;

  1275.     __fp16* outptr7 = outptr6 + ldout;

  1276. 

  1277.     asm volatile(

  1278.             ".arch armv8.2-a+fp16\n"

  1279. 

  1280.             // load accumulator C

  1281.             "cmp %w[type], #0\n"

  1282.             "beq 5f\n"

  1283.             "ldr q8, [%[outptr0]]\n"

  1284.             "ldr q9, [%[outptr1]]\n"

  1285.             "ldr q10, [%[outptr2]]\n"

  1286.             "ldr q11, [%[outptr3]]\n"

  1287.             "ldr q12, [%[outptr4]]\n"

  1288.             "ldr q13, [%[outptr5]]\n"

  1289.             "ldr q14, [%[outptr6]]\n"

  1290.             "ldr q15, [%[outptr7]]\n"

  1291.             "b 6f\n"

  1292. 

  1293.             "5:\n"

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

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

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

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

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

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

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

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

  1302. 

  1303.             "6:\n"

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

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

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

  1307. 

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

  1309. 

  1310.             "1:\n"

  1311.             "fmla v8.8h , %[b0].8h, %[a0].h[0]\n"

  1312.             "fmla v9.8h , %[b0].8h, %[a0].h[1]\n"

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

  1314.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  1315.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  1316.             "fmla v12.8h, %[b0].8h, %[a0].h[4]\n"

  1317.             "fmla v13.8h, %[b0].8h, %[a0].h[5]\n"

  1318.             "fmla v14.8h, %[b0].8h, %[a0].h[6]\n"

  1319.             "fmla v15.8h, %[b0].8h, %[a0].h[7]\n"

  1320.             "ldr %q[b0], [%[b_ptr], #32]\n"

  1321. 

  1322.             "add %[b_ptr], %[b_ptr], #32\n"

  1323.             "ldr %q[a0], [%[a_ptr], #32]\n"

  1324. 

  1325.             "fmla v8.8h , %[b0a].8h, %[a0a].h[0]\n"

  1326.             "fmla v9.8h , %[b0a].8h, %[a0a].h[1]\n"

  1327. 

  1328.             "fmla v10.8h, %[b0a].8h, %[a0a].h[2]\n"

  1329.             "fmla v11.8h, %[b0a].8h, %[a0a].h[3]\n"

  1330.             "fmla v12.8h, %[b0a].8h, %[a0a].h[4]\n"

  1331.             "fmla v13.8h, %[b0a].8h, %[a0a].h[5]\n"

  1332.             "fmla v14.8h, %[b0a].8h, %[a0a].h[6]\n"

  1333.             "fmla v15.8h, %[b0a].8h, %[a0a].h[7]\n"

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

  1335. 

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

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

  1338. 

  1339.             "bne 1b\n"

  1340.             "4:\n"

  1341.             // Jump to odd tail if necessary.

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

  1343. 

  1344.             // Even tail

  1345.             "fmla v8.8h , %[b0].8h, %[a0].h[0]\n"

  1346.             "fmla v9.8h , %[b0].8h, %[a0].h[1]\n"

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

  1348.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  1349.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  1350.             "fmla v12.8h, %[b0].8h, %[a0].h[4]\n"

  1351.             "fmla v13.8h, %[b0].8h, %[a0].h[5]\n"

  1352.             "fmla v14.8h, %[b0].8h, %[a0].h[6]\n"

  1353.             "fmla v15.8h, %[b0].8h, %[a0].h[7]\n"

  1354. 

  1355.             "add %[b_ptr], %[b_ptr], #32\n"

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

  1357. 

  1358.             "fmla v8.8h , %[b0a].8h, %[a0a].h[0]\n"

  1359.             "str q8, [%[outptr0]]\n"

  1360. 

  1361.             "fmla v9.8h , %[b0a].8h, %[a0a].h[1]\n"

  1362.             "str q9, [%[outptr1]]\n"

  1363. 

  1364.             "fmla v10.8h, %[b0a].8h, %[a0a].h[2]\n"

  1365.             "str q10, [%[outptr2]]\n"

  1366. 

  1367.             "fmla v11.8h, %[b0a].8h, %[a0a].h[3]\n"

  1368.             "str q11, [%[outptr3]]\n"

  1369. 

  1370.             "fmla v12.8h, %[b0a].8h, %[a0a].h[4]\n"

  1371.             "str q12, [%[outptr4]]\n"

  1372. 

  1373.             "fmla v13.8h, %[b0a].8h, %[a0a].h[5]\n"

  1374.             "str q13, [%[outptr5]]\n"

  1375. 

  1376.             "fmla v14.8h, %[b0a].8h, %[a0a].h[6]\n"

  1377.             "str q14, [%[outptr6]]\n"

  1378. 

  1379.             "fmla v15.8h, %[b0a].8h, %[a0a].h[7]\n"

  1380.             "str q15, [%[outptr7]]\n"

  1381.             "b 3f\n"

  1382. 

  1383.             // Odd tail

  1384.             "2:\n"

  1385.             "add %[a_ptr], %[a_ptr], #16\n"

  1386.             "fmla v8.8h , %[b0].8h, %[a0].h[0]\n"

  1387.             "add %[b_ptr], %[b_ptr], #16\n"

  1388.             "str q8, [%[outptr0]]\n"

  1389. 

  1390.             "fmla v9.8h , %[b0].8h, %[a0].h[1]\n"

  1391.             "str q9, [%[outptr1]]\n"

  1392. 

  1393.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  1394.             "str q10, [%[outptr2]]\n"

  1395. 

  1396.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  1397.             "str q11, [%[outptr3]]\n"

  1398. 

  1399.             "fmla v12.8h, %[b0].8h, %[a0].h[4]\n"

  1400.             "str q12, [%[outptr4]]\n"

  1401. 

  1402.             "fmla v13.8h, %[b0].8h, %[a0].h[5]\n"

  1403.             "str q13, [%[outptr5]]\n"

  1404. 

  1405.             "fmla v14.8h, %[b0].8h, %[a0].h[6]\n"

  1406.             "str q14, [%[outptr6]]\n"

  1407. 

  1408.             "fmla v15.8h, %[b0].8h, %[a0].h[7]\n"

  1409.             "str q15, [%[outptr7]]\n"

  1410. 

  1411.             "3:\n"

  1412.             : [a0] "+w"(a0), [a0a] "+w"(a0a), [b0] "+w"(b0), [k] "+r"(k),

  1413.               [b0a] "+w"(b0a), [a_ptr] "+r"(a_ptr), [b_ptr] "+r"(b_ptr),

  1414.               [outptr0] "+r"(outptr0), [outptr1] "+r"(outptr1), [outptr2] "+r"(outptr2),

  1415.               [outptr3] "+r"(outptr3), [outptr4] "+r"(outptr4), [outptr5] "+r"(outptr5),

  1416.               [outptr6] "+r"(outptr6), [outptr7] "+r"(outptr7)

  1417.             : [oddk] "r"(oddk), [type] "r"(type)

  1418.             : "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "cc", "memory");

  1419. }

  1420. 

  1421. // Overview of register layout:

  1422. //

  1423. // A 2x8 cell of Rhs is stored in 16bit in d2, d5

  1424. // A 8x2 cell of Lhs is stored in 16bit in q0 - q1

  1425. // A 8x8 block of accumulators is stored in 16bit in d8 - d15.

  1426. //

  1427. //                   +--------+

  1428. //                   | d2[0-3]|

  1429. //              Rhs  +--------+

  1430. //                   | d5[0-3]|

  1431. //                   +--------+

  1432. //

  1433. //                   |        |

  1434. //

  1435. //    Lhs            |        |

  1436. //

  1437. //  +--+--+ - - - -  +--------+

  1438. //  |v0|v1|          | d8[0-3]|

  1439. //  |v0|v1|          | d9[0-3]|

  1440. //  |v0|v1|          |d10[0-3]|

  1441. //  |v0|v1|          |d11[0-3]|

  1442. //  |v0|v1|          |d12[0-3]|

  1443. //  |v0|v1|          |d13[0-3]|

  1444. //  |v0|v1|          |d14[0-3]|

  1445. //  |v0|v1|          |d15[0-3]|

  1446. //  +--+--+ - - - -  +--------+

  1447. //

  1448. //                  Accumulator

  1449. void aarch64_hgemm_assembly_kernel_4x8(

  1450.         const __fp16* a_ptr, const __fp16*& b_ptr, int K, __fp16* outptr0, int ldout,

  1451.         int x_remain, int type) {

  1452.     int oddk = (K & 1);

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

  1454. 

  1455.     register float16x8_t a0 asm("v0");

  1456.     register float16x8_t a0a asm("v1");

  1457.     register float16x8_t b0 asm("v2");

  1458.     register float16x8_t b0a asm("v5");

  1459. 

  1460.     __fp16* outptr1 = outptr0 + ldout;

  1461.     __fp16* outptr2 = outptr1 + ldout;

  1462.     __fp16* outptr3 = outptr2 + ldout;

  1463.     __fp16* outptr4 = outptr3 + ldout;

  1464.     __fp16* outptr5 = outptr4 + ldout;

  1465.     __fp16* outptr6 = outptr5 + ldout;

  1466.     __fp16* outptr7 = outptr6 + ldout;

  1467. 

  1468. #define LOAD_LINE(reg_index, n)            \

  1469.     "mov x0, %[outptr" n                   \

  1470.     "]\n"                                  \

  1471.     "cmp %w[x_remain], #4\n"               \

  1472.     "b.lt REMAIN_LOAD_LINE_LESS_THAN_4_" n \

  1473.     "\n"                                   \

  1474.     "ldr d" reg_index                      \

  1475.     ", [x0]\n"                             \

  1476.     "b LOAD_LINE_END_" n                   \

  1477.     "\n"                                   \

  1478.                                            \

  1479.     "REMAIN_LOAD_LINE_LESS_THAN_4_" n      \

  1480.     ":\n"                                  \

  1481.     "cmp %w[x_remain], #0\n"               \

  1482.     "beq LOAD_LINE_END_" n                 \

  1483.     "\n"                                   \

  1484.     "ld1 {v" reg_index                     \

  1485.     ".h}[0], [x0], #2\n"                   \

  1486.     "cmp %w[x_remain], #1\n"               \

  1487.     "beq LOAD_LINE_END_" n                 \

  1488.     "\n"                                   \

  1489.     "ld1 {v" reg_index                     \

  1490.     ".h}[1], [x0], #2\n"                   \

  1491.     "cmp %w[x_remain], #2\n"               \

  1492.     "beq LOAD_LINE_END_" n                 \

  1493.     "\n"                                   \

  1494.     "ld1 {v" reg_index                     \

  1495.     ".h}[2], [x0], #2\n"                   \

  1496.     "LOAD_LINE_END_" n ":\n"

  1497. 

  1498. #define LOAD_C           \

  1499.     LOAD_LINE("8", "0")  \

  1500.     LOAD_LINE("9", "1")  \

  1501.     LOAD_LINE("10", "2") \

  1502.     LOAD_LINE("11", "3") \

  1503.     LOAD_LINE("12", "4") \

  1504.     LOAD_LINE("13", "5") \

  1505.     LOAD_LINE("14", "6") \

  1506.     LOAD_LINE("15", "7")

  1507. 

  1508. #define STORE_LINE(reg_index, n)            \

  1509.     "mov x0, %[outptr" n                    \

  1510.     "]\n"                                   \

  1511.     "cmp %w[x_remain], #4\n"                \

  1512.     "b.lt REMAIN_STORE_LINE_LESS_THAN_4_" n \

  1513.     "\n"                                    \

  1514.     "str d" reg_index                       \

  1515.     ", [x0]\n"                              \

  1516.     "b STORE_LINE_END_" n                   \

  1517.     "\n"                                    \

  1518.                                             \

  1519.     "REMAIN_STORE_LINE_LESS_THAN_4_" n      \

  1520.     ":\n"                                   \

  1521.     "cmp %w[x_remain], #0\n"                \

  1522.     "beq STORE_LINE_END_" n                 \

  1523.     "\n"                                    \

  1524.     "st1 {v" reg_index                      \

  1525.     ".h}[0], [x0], #2\n"                    \

  1526.     "cmp %w[x_remain], #1\n"                \

  1527.     "beq STORE_LINE_END_" n                 \

  1528.     "\n"                                    \

  1529.     "st1 {v" reg_index                      \

  1530.     ".h}[1], [x0], #2\n"                    \

  1531.     "cmp %w[x_remain], #2\n"                \

  1532.     "beq STORE_LINE_END_" n                 \

  1533.     "\n"                                    \

  1534.     "st1 {v" reg_index                      \

  1535.     ".h}[2], [x0], #2\n"                    \

  1536.     "STORE_LINE_END_" n ":\n"

  1537. 

  1538. #define STORE_C           \

  1539.     STORE_LINE("8", "0")  \

  1540.     STORE_LINE("9", "1")  \

  1541.     STORE_LINE("10", "2") \

  1542.     STORE_LINE("11", "3") \

  1543.     STORE_LINE("12", "4") \

  1544.     STORE_LINE("13", "5") \

  1545.     STORE_LINE("14", "6") \

  1546.     STORE_LINE("15", "7")

  1547. 

  1548.     asm volatile(

  1549.             ".arch armv8.2-a+fp16\n"

  1550. 

  1551.             // load accumulator C

  1552.             "cmp %w[type], #0\n"

  1553.             "beq 5f\n" LOAD_C

  1554.             "b 6f\n"

  1555. 

  1556.             "5:\n"

  1557.             "eor v8.8b,  v8.8b,  v8.8b\n"

  1558.             "eor v9.8b,  v9.8b,  v9.8b\n"

  1559.             "eor v10.8b, v10.8b, v10.8b\n"

  1560.             "eor v11.8b, v11.8b, v11.8b\n"

  1561.             "eor v12.8b, v12.8b, v12.8b\n"

  1562.             "eor v13.8b, v13.8b, v13.8b\n"

  1563.             "eor v14.8b, v14.8b, v14.8b\n"

  1564.             "eor v15.8b, v15.8b, v15.8b\n"

  1565. 

  1566.             "6:\n"

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

  1568. 

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

  1570. 

  1571.             "1:\n"

  1572.             "ldp %d[b0], %d[b0a], [%[b_ptr]]\n"

  1573.             "fmla v8.4h , %[b0].4h, %[a0].h[0]\n"

  1574.             "fmla v9.4h , %[b0].4h, %[a0].h[1]\n"

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

  1576.             "fmla v10.4h, %[b0].4h, %[a0].h[2]\n"

  1577.             "fmla v11.4h, %[b0].4h, %[a0].h[3]\n"

  1578.             "fmla v12.4h, %[b0].4h, %[a0].h[4]\n"

  1579.             "fmla v13.4h, %[b0].4h, %[a0].h[5]\n"

  1580.             "fmla v14.4h, %[b0].4h, %[a0].h[6]\n"

  1581.             "fmla v15.4h, %[b0].4h, %[a0].h[7]\n"

  1582. 

  1583.             "add %[b_ptr], %[b_ptr], #16\n"

  1584.             "ldr %q[a0], [%[a_ptr], #32]\n"

  1585. 

  1586.             "fmla v8.4h , %[b0a].4h, %[a0a].h[0]\n"

  1587.             "fmla v9.4h , %[b0a].4h, %[a0a].h[1]\n"

  1588.             "fmla v10.4h, %[b0a].4h, %[a0a].h[2]\n"

  1589.             "fmla v11.4h, %[b0a].4h, %[a0a].h[3]\n"

  1590.             "fmla v12.4h, %[b0a].4h, %[a0a].h[4]\n"

  1591.             "fmla v13.4h, %[b0a].4h, %[a0a].h[5]\n"

  1592.             "fmla v14.4h, %[b0a].4h, %[a0a].h[6]\n"

  1593.             "fmla v15.4h, %[b0a].4h, %[a0a].h[7]\n"

  1594. 

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

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

  1597. 

  1598.             "bne 1b\n"

  1599.             "4:\n"

  1600.             // Jump to odd tail if necessary.

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

  1602. 

  1603.             // Even tail

  1604.             "ldp %d[b0], %d[b0a], [%[b_ptr]]\n"

  1605.             "fmla v8.4h , %[b0].4h, %[a0].h[0]\n"

  1606.             "fmla v9.4h , %[b0].4h, %[a0].h[1]\n"

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

  1608.             "fmla v10.4h, %[b0].4h, %[a0].h[2]\n"

  1609.             "fmla v11.4h, %[b0].4h, %[a0].h[3]\n"

  1610.             "fmla v12.4h, %[b0].4h, %[a0].h[4]\n"

  1611.             "fmla v13.4h, %[b0].4h, %[a0].h[5]\n"

  1612.             "fmla v14.4h, %[b0].4h, %[a0].h[6]\n"

  1613.             "fmla v15.4h, %[b0].4h, %[a0].h[7]\n"

  1614. 

  1615.             "add %[b_ptr], %[b_ptr], #16\n"

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

  1617. 

  1618.             "fmla v8.4h , %[b0a].4h, %[a0a].h[0]\n"

  1619.             "fmla v9.4h , %[b0a].4h, %[a0a].h[1]\n"

  1620.             "fmla v10.4h, %[b0a].4h, %[a0a].h[2]\n"

  1621.             "fmla v11.4h, %[b0a].4h, %[a0a].h[3]\n"

  1622.             "fmla v12.4h, %[b0a].4h, %[a0a].h[4]\n"

  1623.             "fmla v13.4h, %[b0a].4h, %[a0a].h[5]\n"

  1624.             "fmla v14.4h, %[b0a].4h, %[a0a].h[6]\n"

  1625.             "fmla v15.4h, %[b0a].4h, %[a0a].h[7]\n"

  1626.             "b 3f\n"

  1627. 

  1628.             // Odd tail

  1629.             "2:\n"

  1630.             "ldr %d[b0], [%[b_ptr]]\n"

  1631.             "add %[a_ptr], %[a_ptr], #16\n"

  1632.             "fmla v8.4h , %[b0].4h, %[a0].h[0]\n"

  1633.             "add %[b_ptr], %[b_ptr], #8\n"

  1634.             "fmla v9.4h , %[b0].4h, %[a0].h[1]\n"

  1635.             "fmla v10.4h, %[b0].4h, %[a0].h[2]\n"

  1636.             "fmla v11.4h, %[b0].4h, %[a0].h[3]\n"

  1637.             "fmla v12.4h, %[b0].4h, %[a0].h[4]\n"

  1638.             "fmla v13.4h, %[b0].4h, %[a0].h[5]\n"

  1639.             "fmla v14.4h, %[b0].4h, %[a0].h[6]\n"

  1640.             "fmla v15.4h, %[b0].4h, %[a0].h[7]\n"

  1641. 

  1642.             "3:\n" STORE_C

  1643.             : [a0] "+w"(a0), [a0a] "+w"(a0a), [b0] "+w"(b0), [k] "+r"(k),

  1644.               [b0a] "+w"(b0a), [a_ptr] "+r"(a_ptr), [b_ptr] "+r"(b_ptr),

  1645.               [outptr0] "+r"(outptr0), [outptr1] "+r"(outptr1), [outptr2] "+r"(outptr2),

  1646.               [outptr3] "+r"(outptr3), [outptr4] "+r"(outptr4), [outptr5] "+r"(outptr5),

  1647.               [outptr6] "+r"(outptr6), [outptr7] "+r"(outptr7)

  1648.             : [oddk] "r"(oddk), [x_remain] "r"(x_remain), [type] "r"(type)

  1649.             : "x0", "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "cc",

  1650.               "memory");

  1651. #undef LOAD_LINE

  1652. #undef LOAD_C

  1653. #undef STORE_LINE

  1654. #undef STORE_C

  1655. }

  1656. 

  1657. // Overview of register layout:

  1658. //

  1659. // A 2x24 cell of Rhs is stored in 16bit in q2 - q7

  1660. // A 4x2 cell of Lhs is stored in 16bit in d0, d1

  1661. // A 4x24 block of accumulators is stored in 16bit in q8-q11, q16-q19, q24-q27.

  1662. //

  1663. //                   +--------+--------+--------+

  1664. //                   | v2[0-7]| v3[0-7]| v4[0-7]|

  1665. //              Rhs  +--------+--------+--------+

  1666. //                   | v5[0-7]| v6[0-7]| v7[0-7]|

  1667. //                   +--------+--------+--------+

  1668. //

  1669. //                   |        |        |        |

  1670. //

  1671. //    Lhs            |        |        |        |

  1672. //

  1673. //  +--+--+ - - - -  +--------+--------+--------+

  1674. //  |v0|v1|          | v8[0-7]|v16[0-7]|v24[0-7]|

  1675. //  |v0|v1|          | v9[0-7]|v17[0-7]|v25[0-7]|

  1676. //  |v0|v1|          |v10[0-7]|v18[0-7]|v26[0-7]|

  1677. //  |v0|v1|          |v11[0-7]|v19[0-7]|v27[0-7]|

  1678. //  +--+--+ - - - -  +--------+--------+--------+

  1679. //

  1680. //                            Accumulator

  1681. //! cannot load %[a0] and %[a0a] at same time!

  1682. void aarch64_hgemm_assembly_kernel_24x4(

  1683.         const __fp16* a_ptr, const __fp16*& b_ptr, int K, __fp16* outptr0, int ldout,

  1684.         int y_remain, int type) {

  1685.     int oddk = (K & 1);

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

  1687. 

  1688.     register float16x8_t a0 asm("v0");

  1689.     register float16x8_t a0a asm("v1");

  1690.     register float16x8_t b0 asm("v2");

  1691.     register float16x8_t b1 asm("v3");

  1692.     register float16x8_t b2 asm("v4");

  1693.     register float16x8_t b0a asm("v5");

  1694.     register float16x8_t b1a asm("v6");

  1695.     register float16x8_t b2a asm("v7");

  1696. 

  1697.     __fp16* outptr1 = outptr0 + ldout;

  1698.     __fp16* outptr2 = outptr1 + ldout;

  1699.     __fp16* outptr3 = outptr2 + ldout;

  1700. 

  1701. // clang-format off

  1702. #define LOAD_LINE(v1, v2, v3, n)        \

  1703.     "cbz w0, LOAD_24x4_C_END\n"         \

  1704.     "ldp q" v1 ", q" v2 ", [%[outptr" n \

  1705.     "]]\n"                              \

  1706.     "ldr q" v3 ", [%[outptr" n          \

  1707.     "], #32]\n"                         \

  1708.     "subs w0, w0, #1\n"

  1709. 

  1710. #define LOAD_C                       \

  1711.     "mov w0, %w[y_remain]\n"         \

  1712.     LOAD_LINE("8", "16", "24", "0")  \

  1713.     LOAD_LINE("9", "17", "25", "1")  \

  1714.     LOAD_LINE("10", "18", "26", "2") \

  1715.     LOAD_LINE("11", "19", "27", "3") \

  1716.     "LOAD_24x4_C_END:\n"

  1717. 

  1718. #define STORE_LINE(v1, v2, v3, n)       \

  1719.     "cbz w0, STORE_24x4_C_END\n"        \

  1720.     "stp q" v1 ", q" v2 ", [%[outptr" n \

  1721.     "]]\n"                              \

  1722.     "str q" v3 ", [%[outptr" n          \

  1723.     "], #32]\n"                         \

  1724.     "subs w0, w0, #1\n"

  1725. 

  1726. #define STORE_C "mov w0, %w[y_remain]\n"      \

  1727.             STORE_LINE("8", "16", "24", "0")  \

  1728.             STORE_LINE("9", "17", "25", "1")  \

  1729.             STORE_LINE("10", "18", "26", "2") \

  1730.             STORE_LINE("11", "19", "27", "3") \

  1731.             "STORE_24x4_C_END:\n"

  1732.     // clang-format on

  1733. 

  1734.     asm volatile(

  1735.             ".arch armv8.2-a+fp16\n"

  1736. 

  1737.             // load accumulator C

  1738.             "cmp %w[type], #0\n"

  1739.             "beq 5f\n" LOAD_C

  1740.             "b 6f\n"

  1741.             "5:\n"

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

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

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

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

  1746. 

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

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

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

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

  1751. 

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

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

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

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

  1756. 

  1757.             "6:\n"

  1758.             "ldr %d[a0], [%[a_ptr]]\n"

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

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

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

  1762.             "ldr %q[b0a], [%[b_ptr], #48]\n"

  1763.             "ldr %q[b1a], [%[b_ptr], #64]\n"

  1764. 

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

  1766. 

  1767.             "1:\n"

  1768.             "fmla v8.8h , %[b0].8h, %[a0].h[0]\n"

  1769.             "fmla v9.8h , %[b0].8h, %[a0].h[1]\n"

  1770.             "ldr %d[a0a], [%[a_ptr], #8]\n"

  1771.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  1772.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  1773.             "ldr %q[b2a], [%[b_ptr], #80]\n"

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

  1775. 

  1776.             "fmla v16.8h, %[b1].8h, %[a0].h[0]\n"

  1777.             "fmla v17.8h, %[b1].8h, %[a0].h[1]\n"

  1778.             "fmla v18.8h, %[b1].8h, %[a0].h[2]\n"

  1779.             "fmla v19.8h, %[b1].8h, %[a0].h[3]\n"

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

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

  1782. 

  1783.             "fmla v24.8h, %[b2].8h, %[a0].h[0]\n"

  1784.             "fmla v25.8h, %[b2].8h, %[a0].h[1]\n"

  1785.             "fmla v26.8h, %[b2].8h, %[a0].h[2]\n"

  1786.             "fmla v27.8h, %[b2].8h, %[a0].h[3]\n"

  1787.             "ldr %d[a0], [%[a_ptr], #16]\n"

  1788. 

  1789.             "fmla v8.8h , %[b0a].8h, %[a0a].h[0]\n"

  1790.             "fmla v9.8h , %[b0a].8h, %[a0a].h[1]\n"

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

  1792. 

  1793.             "fmla v10.8h, %[b0a].8h, %[a0a].h[2]\n"

  1794.             "fmla v11.8h, %[b0a].8h, %[a0a].h[3]\n"

  1795.             "ldr %q[b0a], [%[b_ptr], #48]\n"

  1796. 

  1797.             "fmla v16.8h, %[b1a].8h, %[a0a].h[0]\n"

  1798.             "fmla v17.8h, %[b1a].8h, %[a0a].h[1]\n"

  1799.             "fmla v18.8h, %[b1a].8h, %[a0a].h[2]\n"

  1800.             "fmla v19.8h, %[b1a].8h, %[a0a].h[3]\n"

  1801.             "ldr %q[b1a], [%[b_ptr], #64]\n"

  1802. 

  1803.             "fmla v24.8h, %[b2a].8h, %[a0a].h[0]\n"

  1804.             "fmla v25.8h, %[b2a].8h, %[a0a].h[1]\n"

  1805.             "add %[a_ptr], %[a_ptr], #16\n"

  1806.             "fmla v26.8h, %[b2a].8h, %[a0a].h[2]\n"

  1807.             "fmla v27.8h, %[b2a].8h, %[a0a].h[3]\n"

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

  1809. 

  1810.             "bne 1b\n"

  1811.             "4:\n"

  1812.             // Jump to odd tail if necessary.

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

  1814. 

  1815.             // Even tail

  1816.             "fmla v8.8h , %[b0].8h, %[a0].h[0]\n"

  1817.             "fmla v9.8h , %[b0].8h, %[a0].h[1]\n"

  1818.             "ldr %d[a0a], [%[a_ptr], #8]\n"

  1819.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  1820.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  1821.             "ldr %q[b2a], [%[b_ptr], #80]\n"

  1822. 

  1823.             "fmla v16.8h, %[b1].8h, %[a0].h[0]\n"

  1824.             "fmla v17.8h, %[b1].8h, %[a0].h[1]\n"

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

  1826.             "fmla v18.8h, %[b1].8h, %[a0].h[2]\n"

  1827.             "fmla v19.8h, %[b1].8h, %[a0].h[3]\n"

  1828.             "add %[a_ptr], %[a_ptr], #16\n"

  1829. 

  1830.             "fmla v24.8h, %[b2].8h, %[a0].h[0]\n"

  1831.             "fmla v25.8h, %[b2].8h, %[a0].h[1]\n"

  1832.             "fmla v26.8h, %[b2].8h, %[a0].h[2]\n"

  1833.             "fmla v27.8h, %[b2].8h, %[a0].h[3]\n"

  1834. 

  1835.             "fmla v8.8h, %[b0a].8h, %[a0a].h[0]\n"

  1836.             "fmla v9.8h, %[b0a].8h, %[a0a].h[1]\n"

  1837.             "fmla v10.8h, %[b0a].8h, %[a0a].h[2]\n"

  1838.             "fmla v11.8h, %[b0a].8h, %[a0a].h[3]\n"

  1839. 

  1840.             "fmla v16.8h, %[b1a].8h, %[a0a].h[0]\n"

  1841.             "fmla v17.8h, %[b1a].8h, %[a0a].h[1]\n"

  1842.             "fmla v18.8h, %[b1a].8h, %[a0a].h[2]\n"

  1843.             "fmla v19.8h, %[b1a].8h, %[a0a].h[3]\n"

  1844. 

  1845.             "fmla v24.8h, %[b2a].8h, %[a0a].h[0]\n"

  1846.             "fmla v25.8h, %[b2a].8h, %[a0a].h[1]\n"

  1847.             "fmla v26.8h, %[b2a].8h, %[a0a].h[2]\n"

  1848.             "fmla v27.8h, %[b2a].8h, %[a0a].h[3]\n"

  1849.             "b 3f\n"

  1850. 

  1851.             // Odd tail

  1852.             "2:\n"

  1853.             "add %[a_ptr], %[a_ptr], #8\n"

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

  1855. 

  1856.             "fmla v8.8h, %[b0].8h, %[a0].h[0]\n"

  1857.             "fmla v9.8h, %[b0].8h, %[a0].h[1]\n"

  1858.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  1859.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  1860. 

  1861.             "fmla v16.8h, %[b1].8h, %[a0].h[0]\n"

  1862.             "fmla v17.8h, %[b1].8h, %[a0].h[1]\n"

  1863.             "fmla v18.8h, %[b1].8h, %[a0].h[2]\n"

  1864.             "fmla v19.8h, %[b1].8h, %[a0].h[3]\n"

  1865. 

  1866.             "fmla v24.8h, %[b2].8h, %[a0].h[0]\n"

  1867.             "fmla v25.8h, %[b2].8h, %[a0].h[1]\n"

  1868.             "fmla v26.8h, %[b2].8h, %[a0].h[2]\n"

  1869.             "fmla v27.8h, %[b2].8h, %[a0].h[3]\n"

  1870. 

  1871.             "3:\n" STORE_C

  1872.             : [a0] "+w"(a0), [a0a] "+w"(a0a), [b0] "+w"(b0), [b1] "+w"(b1),

  1873.               [b2] "+w"(b2), [k] "+r"(k), [b0a] "+w"(b0a), [b1a] "+w"(b1a),

  1874.               [b2a] "+w"(b2a), [a_ptr] "+r"(a_ptr), [b_ptr] "+r"(b_ptr),

  1875.               [outptr0] "+r"(outptr0), [outptr1] "+r"(outptr1), [outptr2] "+r"(outptr2),

  1876.               [outptr3] "+r"(outptr3)

  1877.             : [oddk] "r"(oddk), [y_remain] "r"(y_remain), [type] "r"(type)

  1878.             : "w0", "v8", "v9", "v10", "v11", "v16", "v17", "v18", "v19", "v24", "v25",

  1879.               "v26", "v27", "cc", "memory");

  1880. #undef LOAD_LINE

  1881. #undef LOAD_C

  1882. #undef STORE_LINE

  1883. #undef STORE_C

  1884. }

  1885. 

  1886. // Overview of register layout:

  1887. //

  1888. // A 2x16 cell of Rhs is stored in 16bit in q2, q3, q5, q6

  1889. // A 4x2 cell of Lhs is stored in 16bit in d0, d1

  1890. // A 4x16 block of accumulators is stored in 16bit in q8-q11, q16-q19.

  1891. //

  1892. //                   +--------+--------+

  1893. //                   | v2[0-7]| v3[0-7]|

  1894. //              Rhs  +--------+--------+

  1895. //                   | v5[0-7]| v6[0-7]|

  1896. //                   +--------+--------+

  1897. //

  1898. //                   |        |        |

  1899. //

  1900. //    Lhs            |        |        |

  1901. //

  1902. //  +--+--+ - - - -  +--------+--------+

  1903. //  |v0|v1|          | v8[0-7]|v16[0-7]|

  1904. //  |v0|v1|          | v9[0-7]|v17[0-7]|

  1905. //  |v0|v1|          |v10[0-7]|v18[0-7]|

  1906. //  |v0|v1|          |v11[0-7]|v19[0-7]|

  1907. //  +--+--+ - - - -  +--------+--------+

  1908. //

  1909. //                       Accumulator

  1910. void aarch64_hgemm_assembly_kernel_16x4(

  1911.         const __fp16* a_ptr, const __fp16*& b_ptr, int K, __fp16* outptr0, int ldout,

  1912.         int y_remain, int type) {

  1913.     int oddk = (K & 1);

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

  1915. 

  1916.     register float16x8_t a0 asm("v0");

  1917.     register float16x8_t a0a asm("v1");

  1918.     register float16x8_t b0 asm("v2");

  1919.     register float16x8_t b1 asm("v3");

  1920.     register float16x8_t b0a asm("v5");

  1921.     register float16x8_t b1a asm("v6");

  1922. 

  1923.     __fp16* outptr1 = outptr0 + ldout;

  1924.     __fp16* outptr2 = outptr1 + ldout;

  1925.     __fp16* outptr3 = outptr2 + ldout;

  1926. 

  1927.     // clang-format off

  1928. 

  1929. #define LOAD_LINE(v1, v2, n)            \

  1930.     "cbz w0, LOAD_16x4_C_END\n"         \

  1931.     "ldp q" v1 ", q" v2 ", [%[outptr" n \

  1932.     "]]\n"                              \

  1933.     "subs w0, w0, #1\n"

  1934. 

  1935. #define LOAD_C "mov w0, %w[y_remain]\n" \

  1936.     LOAD_LINE("8", "16", "0")           \

  1937.     LOAD_LINE("9", "17", "1")           \

  1938.     LOAD_LINE("10", "18", "2")          \

  1939.     LOAD_LINE("11", "19", "3")          \

  1940.     "LOAD_16x4_C_END:\n"

  1941. 

  1942. #define STORE_LINE(v1, v2, n)           \

  1943.     "cbz w0, STORE_16x4_C_END\n"        \

  1944.     "stp q" v1 ", q" v2 ", [%[outptr" n \

  1945.     "]]\n"                              \

  1946.     "subs w0, w0, #1\n"

  1947. 

  1948. #define STORE_C "mov w0, %w[y_remain]\n" \

  1949.             STORE_LINE("8", "16", "0")   \

  1950.             STORE_LINE("9", "17", "1")   \

  1951.             STORE_LINE("10", "18", "2")  \

  1952.             STORE_LINE("11", "19", "3")  \

  1953.             "STORE_16x4_C_END:\n"

  1954. 

  1955.     // clang-format on

  1956. 

  1957.     asm volatile(

  1958.             ".arch armv8.2-a+fp16\n"

  1959. 

  1960.             // load accumulator C

  1961.             "cmp %w[type], #0\n"

  1962.             "beq 5f\n" LOAD_C

  1963.             "b 6f\n"

  1964. 

  1965.             "5:\n"

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

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

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

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

  1970. 

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

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

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

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

  1975. 

  1976.             "6:\n"

  1977.             "ldr %d[a0], [%[a_ptr]]\n"

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

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

  1980.             "ldr %q[b0a], [%[b_ptr], #32]\n"

  1981.             "ldr %q[b1a], [%[b_ptr], #48]\n"

  1982. 

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

  1984. 

  1985.             "1:\n"

  1986.             "fmla v8.8h , %[b0].8h, %[a0].h[0]\n"

  1987.             "fmla v9.8h , %[b0].8h, %[a0].h[1]\n"

  1988.             "ldr %d[a0a], [%[a_ptr], #8]\n"

  1989.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  1990.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  1991.             "ldr %q[b0], [%[b_ptr], #64]\n"

  1992. 

  1993.             "fmla v16.8h, %[b1].8h, %[a0].h[0]\n"

  1994.             "fmla v17.8h, %[b1].8h, %[a0].h[1]\n"

  1995.             "fmla v18.8h, %[b1].8h, %[a0].h[2]\n"

  1996.             "fmla v19.8h, %[b1].8h, %[a0].h[3]\n"

  1997.             "add %[b_ptr], %[b_ptr], #64\n"

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

  1999. 

  2000.             "fmla v8.8h , %[b0a].8h, %[a0a].h[0]\n"

  2001.             "fmla v9.8h , %[b0a].8h, %[a0a].h[1]\n"

  2002.             "ldr %d[a0], [%[a_ptr], #16]\n"

  2003.             "fmla v10.8h, %[b0a].8h, %[a0a].h[2]\n"

  2004.             "fmla v11.8h, %[b0a].8h, %[a0a].h[3]\n"

  2005.             "ldr %q[b0a], [%[b_ptr], #32]\n"

  2006. 

  2007.             "fmla v16.8h, %[b1a].8h, %[a0a].h[0]\n"

  2008.             "fmla v17.8h, %[b1a].8h, %[a0a].h[1]\n"

  2009.             "fmla v18.8h, %[b1a].8h, %[a0a].h[2]\n"

  2010.             "fmla v19.8h, %[b1a].8h, %[a0a].h[3]\n"

  2011.             "ldr %q[b1a], [%[b_ptr], #48]\n"

  2012. 

  2013.             "add %[a_ptr], %[a_ptr], #16\n"

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

  2015. 

  2016.             "bne 1b\n"

  2017.             "4:\n"

  2018.             // Jump to odd tail if necessary.

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

  2020. 

  2021.             // Even tail

  2022.             "fmla v8.8h , %[b0].8h, %[a0].h[0]\n"

  2023.             "fmla v9.8h , %[b0].8h, %[a0].h[1]\n"

  2024.             "ldr %d[a0a], [%[a_ptr], #8]\n"

  2025.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  2026.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  2027. 

  2028.             "fmla v16.8h, %[b1].8h, %[a0].h[0]\n"

  2029.             "fmla v17.8h, %[b1].8h, %[a0].h[1]\n"

  2030.             "add %[b_ptr], %[b_ptr], #64\n"

  2031.             "fmla v18.8h, %[b1].8h, %[a0].h[2]\n"

  2032.             "fmla v19.8h, %[b1].8h, %[a0].h[3]\n"

  2033.             "add %[a_ptr], %[a_ptr], #16\n"

  2034. 

  2035.             "fmla v8.8h, %[b0a].8h, %[a0a].h[0]\n"

  2036.             "fmla v9.8h, %[b0a].8h, %[a0a].h[1]\n"

  2037.             "fmla v10.8h, %[b0a].8h, %[a0a].h[2]\n"

  2038.             "fmla v11.8h, %[b0a].8h, %[a0a].h[3]\n"

  2039. 

  2040.             "fmla v16.8h, %[b1a].8h, %[a0a].h[0]\n"

  2041.             "fmla v17.8h, %[b1a].8h, %[a0a].h[1]\n"

  2042.             "fmla v18.8h, %[b1a].8h, %[a0a].h[2]\n"

  2043.             "fmla v19.8h, %[b1a].8h, %[a0a].h[3]\n"

  2044. 

  2045.             "b 3f\n"

  2046. 

  2047.             // Odd tail

  2048.             "2:\n"

  2049.             "add %[a_ptr], %[a_ptr], #8\n"

  2050.             "add %[b_ptr], %[b_ptr], #32\n"

  2051. 

  2052.             "fmla v8.8h, %[b0].8h, %[a0].h[0]\n"

  2053.             "fmla v9.8h, %[b0].8h, %[a0].h[1]\n"

  2054.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  2055.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  2056. 

  2057.             "fmla v16.8h, %[b1].8h, %[a0].h[0]\n"

  2058.             "fmla v17.8h, %[b1].8h, %[a0].h[1]\n"

  2059.             "fmla v18.8h, %[b1].8h, %[a0].h[2]\n"

  2060.             "fmla v19.8h, %[b1].8h, %[a0].h[3]\n"

  2061. 

  2062.             "3:\n" STORE_C

  2063.             : [a0] "+w"(a0), [a0a] "+w"(a0a), [b0] "+w"(b0), [b1] "+w"(b1), [k] "+r"(k),

  2064.               [b0a] "+w"(b0a), [b1a] "+w"(b1a), [a_ptr] "+r"(a_ptr),

  2065.               [b_ptr] "+r"(b_ptr), [outptr0] "+r"(outptr0), [outptr1] "+r"(outptr1),

  2066.               [outptr2] "+r"(outptr2), [outptr3] "+r"(outptr3)

  2067.             : [oddk] "r"(oddk), [y_remain] "r"(y_remain), [type] "r"(type)

  2068.             : "w0", "v8", "v9", "v10", "v11", "v16", "v17", "v18", "v19", "cc",

  2069.               "memory");

  2070. #undef LOAD_LINE

  2071. #undef LOAD_C

  2072. #undef STORE_LINE

  2073. #undef STORE_C

  2074. }

  2075. 

  2076. // Overview of register layout:

  2077. //

  2078. // A 2x8 cell of Rhs is stored in 16bit in q2, q5

  2079. // A 4x2 cell of Lhs is stored in 16bit in d0, d1

  2080. // A 4x8 block of accumulators is stored in 16bit in q8-q11.

  2081. //

  2082. //                   +--------+

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

  2084. //              Rhs  +--------+

  2085. //                   | v5[0-7]|

  2086. //                   +--------+

  2087. //

  2088. //                   |        |

  2089. //

  2090. //    Lhs            |        |

  2091. //

  2092. //  +--+--+ - - - -  +--------+

  2093. //  |v0|v1|          | v8[0-7]|

  2094. //  |v0|v1|          | v9[0-7]|

  2095. //  |v0|v1|          |v10[0-7]|

  2096. //  |v0|v1|          |v11[0-7]|

  2097. //  +--+--+ - - - -  +--------+

  2098. //

  2099. //                  Accumulator

  2100. void aarch64_hgemm_assembly_kernel_8x4(

  2101.         const __fp16* a_ptr, const __fp16*& b_ptr, int K, __fp16* outptr0, int ldout,

  2102.         int y_remain, int type) {

  2103.     int oddk = (K & 1);

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

  2105. 

  2106.     register float16x8_t a0 asm("v0");

  2107.     register float16x8_t a0a asm("v1");

  2108.     register float16x8_t b0 asm("v2");

  2109.     register float16x8_t b0a asm("v5");

  2110. 

  2111.     __fp16* outptr1 = outptr0 + ldout;

  2112.     __fp16* outptr2 = outptr1 + ldout;

  2113.     __fp16* outptr3 = outptr2 + ldout;

  2114. 

  2115. // clang-format off

  2116. #define LOAD_LINE(v1, n)       \

  2117.     "cbz w0, LOAD_8x4_C_END\n" \

  2118.     "ldr q" v1 ", [%[outptr" n \

  2119.     "]]\n"                     \

  2120.     "subs w0, w0, #1\n"

  2121. 

  2122. #define LOAD_C               \

  2123.     "mov w0, %w[y_remain]\n" \

  2124.     LOAD_LINE("8", "0")      \

  2125.     LOAD_LINE("9", "1")      \

  2126.     LOAD_LINE("10", "2")     \

  2127.     LOAD_LINE("11", "3")     \

  2128.     "LOAD_8x4_C_END:\n"

  2129. 

  2130. #define STORE_LINE(v1, n)       \

  2131.     "cbz w0, STORE_8x4_C_END\n" \

  2132.     "str q" v1 ", [%[outptr" n  \

  2133.     "]]\n"                      \

  2134.     "subs w0, w0, #1\n"

  2135. 

  2136. #define STORE_C              \

  2137.     "mov w0, %w[y_remain]\n" \

  2138.     STORE_LINE("8", "0")     \

  2139.     STORE_LINE("9", "1")     \

  2140.     STORE_LINE("10", "2")    \

  2141.     STORE_LINE("11", "3")    \

  2142.     "STORE_8x4_C_END:\n"

  2143.     // clang-format on

  2144. 

  2145.     asm volatile(

  2146.             ".arch armv8.2-a+fp16\n"

  2147. 

  2148.             // load accumulator C

  2149.             "cmp %w[type], #0\n"

  2150.             "beq 5f\n" LOAD_C

  2151.             "b 6f\n"

  2152.             "5:\n"

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

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

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

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

  2157. 

  2158.             "6:\n"

  2159.             "ldr %d[a0], [%[a_ptr]]\n"

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

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

  2162. 

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

  2164. 

  2165.             "1:\n"

  2166.             "fmla v8.8h , %[b0].8h, %[a0].h[0]\n"

  2167.             "fmla v9.8h , %[b0].8h, %[a0].h[1]\n"

  2168.             "ldr %d[a0a], [%[a_ptr], #8]\n"

  2169.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  2170.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  2171.             "ldr %q[b0], [%[b_ptr], #32]\n"

  2172. 

  2173.             "add %[b_ptr], %[b_ptr], #32\n"

  2174.             "ldr %d[a0], [%[a_ptr], #16]\n"

  2175. 

  2176.             "fmla v8.8h , %[b0a].8h, %[a0a].h[0]\n"

  2177.             "fmla v9.8h , %[b0a].8h, %[a0a].h[1]\n"

  2178.             "fmla v10.8h, %[b0a].8h, %[a0a].h[2]\n"

  2179.             "fmla v11.8h, %[b0a].8h, %[a0a].h[3]\n"

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

  2181. 

  2182.             "add %[a_ptr], %[a_ptr], #16\n"

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

  2184. 

  2185.             "bne 1b\n"

  2186.             "4:\n"

  2187.             // Jump to odd tail if necessary.

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

  2189. 

  2190.             // Even tail

  2191.             "fmla v8.8h , %[b0].8h, %[a0].h[0]\n"

  2192.             "fmla v9.8h , %[b0].8h, %[a0].h[1]\n"

  2193.             "ldr %d[a0a], [%[a_ptr], #8]\n"

  2194.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  2195.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  2196. 

  2197.             "add %[b_ptr], %[b_ptr], #32\n"

  2198.             "add %[a_ptr], %[a_ptr], #16\n"

  2199. 

  2200.             "fmla v8.8h, %[b0a].8h, %[a0a].h[0]\n"

  2201.             "fmla v9.8h, %[b0a].8h, %[a0a].h[1]\n"

  2202.             "fmla v10.8h, %[b0a].8h, %[a0a].h[2]\n"

  2203.             "fmla v11.8h, %[b0a].8h, %[a0a].h[3]\n"

  2204. 

  2205.             "b 3f\n"

  2206. 

  2207.             // Odd tail

  2208.             "2:\n"

  2209.             "add %[a_ptr], %[a_ptr], #8\n"

  2210.             "add %[b_ptr], %[b_ptr], #16\n"

  2211. 

  2212.             "fmla v8.8h, %[b0].8h, %[a0].h[0]\n"

  2213.             "fmla v9.8h, %[b0].8h, %[a0].h[1]\n"

  2214.             "fmla v10.8h, %[b0].8h, %[a0].h[2]\n"

  2215.             "fmla v11.8h, %[b0].8h, %[a0].h[3]\n"

  2216. 

  2217.             "3:\n" STORE_C

  2218.             : [a0] "+w"(a0), [a0a] "+w"(a0a), [b0] "+w"(b0), [k] "+r"(k),

  2219.               [b0a] "+w"(b0a), [a_ptr] "+r"(a_ptr), [b_ptr] "+r"(b_ptr),

  2220.               [outptr0] "+r"(outptr0), [outptr1] "+r"(outptr1), [outptr2] "+r"(outptr2),

  2221.               [outptr3] "+r"(outptr3)

  2222.             : [oddk] "r"(oddk), [y_remain] "r"(y_remain), [type] "r"(type)

  2223.             : "w0", "v8", "v9", "v10", "v11", "cc", "memory");

  2224. #undef LOAD_LINE

  2225. #undef LOAD_C

  2226. #undef STORE_LINE

  2227. #undef STORE_C

  2228. }

  2229. 

  2230. // Overview of register layout:

  2231. //

  2232. // A 2x8 cell of Rhs is stored in 16bit in d2, d5

  2233. // A 4x2 cell of Lhs is stored in 16bit in d0, d1

  2234. // A 4x8 block of accumulators is stored in 16bit in d8-d11.

  2235. //

  2236. //                   +--------+

  2237. //                   | d2[0-3]|

  2238. //              Rhs  +--------+

  2239. //                   | d5[0-3]|

  2240. //                   +--------+

  2241. //

  2242. //                   |        |

  2243. //

  2244. //    Lhs            |        |

  2245. //

  2246. //  +--+--+ - - - -  +--------+

  2247. //  |d0|d1|          | d8[0-3]|

  2248. //  |d0|d1|          | d9[0-3]|

  2249. //  |d0|d1|          |d10[0-3]|

  2250. //  |d0|d1|          |d11[0-3]|

  2251. //  +--+--+ - - - -  +--------+

  2252. //

  2253. //                  Accumulator

  2254. void aarch64_hgemm_assembly_kernel_4x4(

  2255.         const __fp16* a_ptr, const __fp16*& b_ptr, int K, __fp16* outptr0, int ldout,

  2256.         int x_remain, int y_remain, int type) {

  2257.     int oddk = (K & 1);

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

  2259. 

  2260.     register float16x8_t a0 asm("v0");

  2261.     register float16x8_t a0a asm("v1");

  2262.     register float16x8_t b0 asm("v2");

  2263.     register float16x8_t b0a asm("v5");

  2264. 

  2265.     __fp16* outptr1 = outptr0 + ldout;

  2266.     __fp16* outptr2 = outptr1 + ldout;

  2267.     __fp16* outptr3 = outptr2 + ldout;

  2268. 

  2269. #define LOAD_LINE(reg_index, n)                \

  2270.     "cbz w1, LOAD_4x4_C_END\n"                 \

  2271.     "mov x0, %[outptr" n                       \

  2272.     "]\n"                                      \

  2273.     "cmp %w[x_remain], #4\n"                   \

  2274.     "b.lt REMAIN_LOAD_4x4_LINE_LESS_THAN_4_" n \

  2275.     "\n"                                       \

  2276.     "ldr d" reg_index                          \

  2277.     ", [x0]\n"                                 \

  2278.     "b LOAD_4x4_LINE_END_" n                   \

  2279.     "\n"                                       \

  2280.                                                \

  2281.     "REMAIN_LOAD_4x4_LINE_LESS_THAN_4_" n      \

  2282.     ":\n"                                      \

  2283.     "cmp %w[x_remain], #0\n"                   \

  2284.     "beq LOAD_4x4_LINE_END_" n                 \

  2285.     "\n"                                       \

  2286.     "ld1 {v" reg_index                         \

  2287.     ".h}[0], [x0], #2\n"                       \

  2288.     "cmp %w[x_remain], #1\n"                   \

  2289.     "beq LOAD_4x4_LINE_END_" n                 \

  2290.     "\n"                                       \

  2291.     "ld1 {v" reg_index                         \

  2292.     ".h}[1], [x0], #2\n"                       \

  2293.     "cmp %w[x_remain], #2\n"                   \

  2294.     "beq LOAD_4x4_LINE_END_" n                 \

  2295.     "\n"                                       \

  2296.     "ld1 {v" reg_index                         \

  2297.     ".h}[2], [x0], #2\n"                       \

  2298.     "LOAD_4x4_LINE_END_" n                     \

  2299.     ":\n"                                      \

  2300.     "subs w1, w1, #1\n"

  2301. 

  2302. #define LOAD_C                                                       \

  2303.     "mov w1, %w[y_remain]\n" LOAD_LINE("8", "0") LOAD_LINE("9", "1") \

  2304.             LOAD_LINE("10", "2") LOAD_LINE("11", "3") "LOAD_4x4_C_END:\n"

  2305. 

  2306. #define STORE_LINE(reg_index, n)                \

  2307.     "cbz w1, STORE_4x4_C_END\n"                 \

  2308.     "mov x0, %[outptr" n                        \

  2309.     "]\n"                                       \

  2310.     "cmp %w[x_remain], #4\n"                    \

  2311.     "b.lt REMAIN_STORE_4x4_LINE_LESS_THAN_4_" n \

  2312.     "\n"                                        \

  2313.     "str d" reg_index                           \

  2314.     ", [x0]\n"                                  \

  2315.     "b STORE_4x4_LINE_END_" n                   \

  2316.     "\n"                                        \

  2317.                                                 \

  2318.     "REMAIN_STORE_4x4_LINE_LESS_THAN_4_" n      \

  2319.     ":\n"                                       \

  2320.     "cmp %w[x_remain], #0\n"                    \

  2321.     "beq STORE_4x4_LINE_END_" n                 \

  2322.     "\n"                                        \

  2323.     "st1 {v" reg_index                          \

  2324.     ".h}[0], [x0], #2\n"                        \

  2325.     "cmp %w[x_remain], #1\n"                    \

  2326.     "beq STORE_4x4_LINE_END_" n                 \

  2327.     "\n"                                        \

  2328.     "st1 {v" reg_index                          \

  2329.     ".h}[1], [x0], #2\n"                        \

  2330.     "cmp %w[x_remain], #2\n"                    \

  2331.     "beq STORE_4x4_LINE_END_" n                 \

  2332.     "\n"                                        \

  2333.     "st1 {v" reg_index                          \

  2334.     ".h}[2], [x0], #2\n"                        \

  2335.     "STORE_4x4_LINE_END_" n                     \

  2336.     ":\n"                                       \

  2337.     "subs w1, w1, #1\n"

  2338. 

  2339. #define STORE_C                                                        \

  2340.     "mov w1, %w[y_remain]\n" STORE_LINE("8", "0") STORE_LINE("9", "1") \

  2341.             STORE_LINE("10", "2") STORE_LINE("11", "3") "STORE_4x4_C_END:\n"

  2342. 

  2343.     asm volatile(

  2344.             ".arch armv8.2-a+fp16\n"

  2345. 

  2346.             // load accumulator C

  2347.             "cmp %w[type], #0\n"

  2348.             "beq 5f\n" LOAD_C

  2349.             "b 6f\n"

  2350. 

  2351.             "5:\n"

  2352.             "eor v8.8b,  v8.8b,  v8.8b\n"

  2353.             "eor v9.8b,  v9.8b,  v9.8b\n"

  2354.             "eor v10.8b, v10.8b, v10.8b\n"

  2355.             "eor v11.8b, v11.8b, v11.8b\n"

  2356. 

  2357.             "6:\n"

  2358.             "ldr %d[a0], [%[a_ptr]]\n"

  2359. 

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

  2361. 

  2362.             "1:\n"

  2363.             "ldp %d[b0], %d[b0a], [%[b_ptr]]\n"

  2364.             "fmla v8.4h , %[b0].4h, %[a0].h[0]\n"

  2365.             "fmla v9.4h , %[b0].4h, %[a0].h[1]\n"

  2366.             "ldr %d[a0a], [%[a_ptr], #8]\n"

  2367.             "fmla v10.4h, %[b0].4h, %[a0].h[2]\n"

  2368.             "fmla v11.4h, %[b0].4h, %[a0].h[3]\n"

  2369. 

  2370.             "add %[b_ptr], %[b_ptr], #16\n"

  2371.             "ldr %d[a0], [%[a_ptr], #16]\n"

  2372. 

  2373.             "fmla v8.4h , %[b0a].4h, %[a0a].h[0]\n"

  2374.             "fmla v9.4h , %[b0a].4h, %[a0a].h[1]\n"

  2375.             "fmla v10.4h, %[b0a].4h, %[a0a].h[2]\n"

  2376.             "fmla v11.4h, %[b0a].4h, %[a0a].h[3]\n"

  2377. 

  2378.             "add %[a_ptr], %[a_ptr], #16\n"

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

  2380. 

  2381.             "bne 1b\n"

  2382.             "4:\n"

  2383.             // Jump to odd tail if necessary.

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

  2385. 

  2386.             // Even tail

  2387.             "ldp %d[b0], %d[b0a], [%[b_ptr]]\n"

  2388.             "fmla v8.4h , %[b0].4h, %[a0].h[0]\n"

  2389.             "fmla v9.4h , %[b0].4h, %[a0].h[1]\n"

  2390.             "ldr %d[a0a], [%[a_ptr], #8]\n"

  2391.             "fmla v10.4h, %[b0].4h, %[a0].h[2]\n"

  2392.             "fmla v11.4h, %[b0].4h, %[a0].h[3]\n"

  2393. 

  2394.             "add %[b_ptr], %[b_ptr], #16\n"

  2395.             "add %[a_ptr], %[a_ptr], #16\n"

  2396. 

  2397.             "fmla v8.4h, %[b0a].4h, %[a0a].h[0]\n"

  2398.             "fmla v9.4h, %[b0a].4h, %[a0a].h[1]\n"

  2399.             "fmla v10.4h, %[b0a].4h, %[a0a].h[2]\n"

  2400.             "fmla v11.4h, %[b0a].4h, %[a0a].h[3]\n"

  2401.             "b 3f\n"

  2402. 

  2403.             // Odd tail

  2404.             "2:\n"

  2405.             "ldr %d[b0], [%[b_ptr]]\n"

  2406.             "add %[a_ptr], %[a_ptr], #8\n"

  2407.             "add %[b_ptr], %[b_ptr], #8\n"

  2408. 

  2409.             "fmla v8.4h, %[b0].4h, %[a0].h[0]\n"

  2410.             "fmla v9.4h, %[b0].4h, %[a0].h[1]\n"

  2411.             "fmla v10.4h, %[b0].4h, %[a0].h[2]\n"

  2412.             "fmla v11.4h, %[b0].4h, %[a0].h[3]\n"

  2413. 

  2414.             "3:\n" STORE_C

  2415.             : [a0] "+w"(a0), [a0a] "+w"(a0a), [b0] "+w"(b0), [k] "+r"(k),

  2416.               [b0a] "+w"(b0a), [a_ptr] "+r"(a_ptr), [b_ptr] "+r"(b_ptr),

  2417.               [outptr0] "+r"(outptr0), [outptr1] "+r"(outptr1), [outptr2] "+r"(outptr2),

  2418.               [outptr3] "+r"(outptr3)

  2419.             : [oddk] "r"(oddk), [x_remain] "r"(x_remain), [y_remain] "r"(y_remain),

  2420.               [type] "r"(type)

  2421.             : "x0", "w1", "v8", "v9", "v10", "v11", "cc", "memory");

  2422. 

  2423. #undef LOAD_LINE

  2424. #undef LOAD_C

  2425. #undef STORE_LINE

  2426. #undef STORE_C

  2427. }

  2428. 

  2429. void aarch64_hgemm_asimd_8x24(

  2430.         const __fp16* Apanel, const __fp16* Bpanel, __fp16* out, int ldout, int x0,

  2431.         int xmax, int y0, int ymax, int K, bool is_first_k) {

  2432.     const __fp16* a_ptr = Apanel;

  2433.     const int A_interleave = 8;

  2434.     const int B_transpose1xW = 24;

  2435.     const int K8 = (K << 3);

  2436.     const int K4 = (K << 2);

  2437.     int type = is_first_k ? 0 : 1;

  2438. 

  2439.     int y = y0;

  2440.     for (; y + A_interleave <= ymax; y += A_interleave) {

  2441.         const __fp16* a_ptr0 = a_ptr;

  2442.         const __fp16* b_ptr = Bpanel;

  2443. 

  2444.         __fp16* outptr0 = out + (y * ldout) + x0;

  2445. 

  2446.         int x = x0;

  2447. 

  2448.         for (; x + B_transpose1xW <= xmax; x += B_transpose1xW) {

  2449.             a_ptr = a_ptr0;

  2450.             aarch64_hgemm_assembly_kernel_24x8(a_ptr, b_ptr, K, outptr0, ldout, type);

  2451.             outptr0 += B_transpose1xW;

  2452.         }

  2453. 

  2454.         for (; x + 16 <= xmax; x += 16) {

  2455.             a_ptr = a_ptr0;

  2456.             aarch64_hgemm_assembly_kernel_16x8(a_ptr, b_ptr, K, outptr0, ldout, type);

  2457.             outptr0 += 16;

  2458.         }

  2459.         for (; x + 8 <= xmax; x += 8) {

  2460.             a_ptr = a_ptr0;

  2461.             aarch64_hgemm_assembly_kernel_8x8(a_ptr, b_ptr, K, outptr0, ldout, type);

  2462.             outptr0 += 8;

  2463.         }

  2464.         for (; x < xmax; x += 4) {

  2465.             int x_remain = xmax - x;

  2466.             a_ptr = a_ptr0;

  2467.             aarch64_hgemm_assembly_kernel_4x8(

  2468.                     a_ptr, b_ptr, K, outptr0, ldout, x_remain, type);

  2469.             outptr0 += 4;

  2470.         }

  2471.         a_ptr = a_ptr0 + K8;

  2472.     }

  2473. 

  2474.     for (; y < ymax; y += 4) {

  2475.         const __fp16* a_ptr0 = a_ptr;

  2476.         const __fp16* b_ptr = Bpanel;

  2477. 

  2478.         __fp16* outptr0 = out + (y * ldout) + x0;

  2479. 

  2480.         int x = x0;

  2481.         for (; x + B_transpose1xW <= xmax; x += B_transpose1xW) {

  2482.             a_ptr = a_ptr0;

  2483.             aarch64_hgemm_assembly_kernel_24x4(

  2484.                     a_ptr, b_ptr, K, outptr0, ldout, ymax - y, type);

  2485.             outptr0 += B_transpose1xW;

  2486.         }

  2487. 

  2488.         for (; x + 16 <= xmax; x += 16) {

  2489.             a_ptr = a_ptr0;

  2490.             aarch64_hgemm_assembly_kernel_16x4(

  2491.                     a_ptr, b_ptr, K, outptr0, ldout, ymax - y, type);

  2492.             outptr0 += 16;

  2493.         }

  2494.         for (; x + 8 <= xmax; x += 8) {

  2495.             a_ptr = a_ptr0;

  2496.             aarch64_hgemm_assembly_kernel_8x4(

  2497.                     a_ptr, b_ptr, K, outptr0, ldout, ymax - y, type);

  2498.             outptr0 += 8;

  2499.         }

  2500.         for (; x < xmax; x += 4) {

  2501.             a_ptr = a_ptr0;

  2502.             aarch64_hgemm_assembly_kernel_4x4(

  2503.                     a_ptr, b_ptr, K, outptr0, ldout, xmax - x, ymax - y, type);

  2504.             outptr0 += 4;

  2505.         }

  2506.         a_ptr = a_ptr0 + K4;

  2507.     }

  2508. }

  2509. }  // namespace

  2510. 

  2511. MEGDNN_REG_GEMM_STRATEGY_IMPL(hgemm_8x24);

  2512. 

  2513. void hgemm_8x24::pack_A(

  2514.         dt_float16* out, const dt_float16* in, int ldin, int y0, int ymax, int k0,

  2515.         int kmax, bool transpose_A) const {

  2516.     if (transpose_A) {

  2517.         transpose_1x8(

  2518.                 reinterpret_cast<__fp16*>(out), reinterpret_cast<const __fp16*>(in),

  2519.                 ldin, y0, ymax, k0, kmax);

  2520.     } else {

  2521.         interleave_8x1(

  2522.                 reinterpret_cast<__fp16*>(out), reinterpret_cast<const __fp16*>(in),

  2523.                 ldin, y0, ymax, k0, kmax);

  2524.     }

  2525. }

  2526. 

  2527. void hgemm_8x24::pack_B(

  2528.         dt_float16* out, const dt_float16* in, int ldin, int x0, int xmax, int k0,

  2529.         int kmax, bool transpose_B) const {

  2530.     if (transpose_B) {

  2531.         interleave_24x1(

  2532.                 reinterpret_cast<__fp16*>(out), reinterpret_cast<const __fp16*>(in),

  2533.                 ldin, x0, xmax, k0, kmax);

  2534.     } else {

  2535.         transpose_1x24(

  2536.                 reinterpret_cast<__fp16*>(out), reinterpret_cast<const __fp16*>(in),

  2537.                 ldin, x0, xmax, k0, kmax);

  2538.     }

  2539. }

  2540. 

  2541. void hgemm_8x24::kern(

  2542.         const dt_float16* packA, const dt_float16* packB, size_t M, size_t N, size_t K,

  2543.         dt_float16* C, size_t LDC, bool is_first_k, const dt_float16*,

  2544.         dt_float16*) const {

  2545.     megdnn_assert(

  2546.             A_dtype.enumv() == B_dtype.enumv() && A_dtype.enumv() == C_dtype.enumv() &&

  2547.             A_dtype.enumv() == DTypeEnum::Float16);

  2548.     MEGDNN_MARK_USED_VAR(A_dtype);

  2549.     MEGDNN_MARK_USED_VAR(B_dtype);

  2550.     MEGDNN_MARK_USED_VAR(C_dtype);

  2551.     aarch64_hgemm_asimd_8x24(

  2552.             reinterpret_cast<const __fp16*>(packA),

  2553.             reinterpret_cast<const __fp16*>(packB), reinterpret_cast<__fp16*>(C), LDC,

  2554.             0, N, 0, M, K, is_first_k);

  2555. }

  2556. #endif

  2557. // vim: syntax=cpp.doxygen