xref: /aosp_15_r20/external/XNNPACK/src/f16-igemm/gen/3x16-minmax-avx2-broadcast.c (revision 4bdc94577ba0e567308109d787f7fec7b531ce36)
1 // Auto-generated file. Do not edit!
2 //   Template: src/f16-igemm/avx2-broadcast.c.in
3 //   Generator: tools/xngen
4 //
5 // Copyright 2022 Google LLC
6 //
7 // This source code is licensed under the BSD-style license found in the
8 // LICENSE file in the root directory of this source tree.
9 
10 #include <assert.h>
11 
12 #include <immintrin.h>
13 
14 #include <xnnpack/igemm.h>
15 #include <xnnpack/intrinsics-polyfill.h>
16 
17 
xnn_f16_igemm_minmax_ukernel_3x16__avx2_broadcast(size_t mr,size_t nc,size_t kc,size_t ks,const void ** restrict a,const void * restrict w,void * restrict c,size_t cm_stride,size_t cn_stride,size_t a_offset,const void * zero,const union xnn_f16_minmax_params params[restrict XNN_MIN_ELEMENTS (1)])18 void xnn_f16_igemm_minmax_ukernel_3x16__avx2_broadcast(
19     size_t mr,
20     size_t nc,
21     size_t kc,
22     size_t ks,
23     const void**restrict a,
24     const void*restrict w,
25     void*restrict c,
26     size_t cm_stride,
27     size_t cn_stride,
28     size_t a_offset,
29     const void* zero,
30     const union xnn_f16_minmax_params params[restrict XNN_MIN_ELEMENTS(1)])
31 {
32   assert(mr != 0);
33   assert(mr <= 3);
34   assert(nc != 0);
35   assert(kc != 0);
36   assert(kc % sizeof(uint16_t) == 0);
37   assert(ks != 0);
38   assert(ks % (3 * sizeof(void*)) == 0);
39   assert(a_offset % sizeof(uint16_t) == 0);
40   assert(a != NULL);
41   assert(w != NULL);
42   assert(c != NULL);
43 
44   uint16_t* c0 = c;
45   uint16_t* c1 = (uint16_t*) ((uintptr_t) c0 + cm_stride);
46   if XNN_UNPREDICTABLE(mr < 2) {
47     c1 = c0;
48   }
49   uint16_t* c2 = (uint16_t*) ((uintptr_t) c1 + cm_stride);
50   if XNN_UNPREDICTABLE(mr <= 2) {
51     c2 = c1;
52   }
53 
54   do {
55     __m256 vacc0x01234567 = _mm256_cvtph_ps(_mm_load_si128((const __m128i*) w));
56     __m256 vacc0x89ABCDEF = _mm256_cvtph_ps(_mm_load_si128((const __m128i*) ((const uint16_t*) w + 8)));
57     __m256 vacc1x01234567 = vacc0x01234567;
58     __m256 vacc1x89ABCDEF = vacc0x89ABCDEF;
59     __m256 vacc2x01234567 = vacc0x01234567;
60     __m256 vacc2x89ABCDEF = vacc0x89ABCDEF;
61     w = (const uint16_t*) w + 16;
62 
63     size_t p = ks;
64     do {
65       const uint16_t* restrict a0 = (const uint16_t*) a[0];
66       assert(a0 != NULL);
67       if XNN_UNPREDICTABLE(a0 != zero) {
68         a0 = (const uint16_t*) ((uintptr_t) a0 + a_offset);
69       }
70       const uint16_t* restrict a1 = (const uint16_t*) a[1];
71       assert(a1 != NULL);
72       if XNN_UNPREDICTABLE(a1 != zero) {
73         a1 = (const uint16_t*) ((uintptr_t) a1 + a_offset);
74       }
75       const uint16_t* restrict a2 = (const uint16_t*) a[2];
76       assert(a2 != NULL);
77       if XNN_UNPREDICTABLE(a2 != zero) {
78         a2 = (const uint16_t*) ((uintptr_t) a2 + a_offset);
79       }
80       a += 3;
81 
82       size_t k = kc;
83       do {
84         const __m256 vb01234567 = _mm256_cvtph_ps(_mm_load_si128((const __m128i*) w));
85         const __m256 vb89ABCDEF = _mm256_cvtph_ps(_mm_load_si128((const __m128i*) ((const uint16_t*) w + 8)));
86         w = (const uint16_t*) w + 16;
87 
88         const __m256 va0 = _mm256_cvtph_ps(_mm_set1_epi16((short) *a0));
89         a0 += 1;
90         const __m256 va1 = _mm256_cvtph_ps(_mm_set1_epi16((short) *a1));
91         a1 += 1;
92         const __m256 va2 = _mm256_cvtph_ps(_mm_set1_epi16((short) *a2));
93         a2 += 1;
94 
95         vacc0x01234567 = _mm256_cvtph_ps(_mm256_cvtps_ph(_mm256_fmadd_ps(va0, vb01234567, vacc0x01234567), _MM_FROUND_NO_EXC));
96         vacc0x89ABCDEF = _mm256_cvtph_ps(_mm256_cvtps_ph(_mm256_fmadd_ps(va0, vb89ABCDEF, vacc0x89ABCDEF), _MM_FROUND_NO_EXC));
97         vacc1x01234567 = _mm256_cvtph_ps(_mm256_cvtps_ph(_mm256_fmadd_ps(va1, vb01234567, vacc1x01234567), _MM_FROUND_NO_EXC));
98         vacc1x89ABCDEF = _mm256_cvtph_ps(_mm256_cvtps_ph(_mm256_fmadd_ps(va1, vb89ABCDEF, vacc1x89ABCDEF), _MM_FROUND_NO_EXC));
99         vacc2x01234567 = _mm256_cvtph_ps(_mm256_cvtps_ph(_mm256_fmadd_ps(va2, vb01234567, vacc2x01234567), _MM_FROUND_NO_EXC));
100         vacc2x89ABCDEF = _mm256_cvtph_ps(_mm256_cvtps_ph(_mm256_fmadd_ps(va2, vb89ABCDEF, vacc2x89ABCDEF), _MM_FROUND_NO_EXC));
101 
102         k -= sizeof(uint16_t);
103       } while (k != 0);
104       p -= 3 * sizeof(void*);
105     } while (p != 0);
106 
107     const __m256 vmin = _mm256_load_ps(params->avx.min);
108     vacc0x01234567 = _mm256_max_ps(vacc0x01234567, vmin);
109     vacc1x01234567 = _mm256_max_ps(vacc1x01234567, vmin);
110     vacc2x01234567 = _mm256_max_ps(vacc2x01234567, vmin);
111     vacc0x89ABCDEF = _mm256_max_ps(vacc0x89ABCDEF, vmin);
112     vacc1x89ABCDEF = _mm256_max_ps(vacc1x89ABCDEF, vmin);
113     vacc2x89ABCDEF = _mm256_max_ps(vacc2x89ABCDEF, vmin);
114 
115     const __m256 vmax = _mm256_load_ps(params->avx.max);
116     vacc0x01234567 = _mm256_min_ps(vacc0x01234567, vmax);
117     vacc1x01234567 = _mm256_min_ps(vacc1x01234567, vmax);
118     vacc2x01234567 = _mm256_min_ps(vacc2x01234567, vmax);
119     vacc0x89ABCDEF = _mm256_min_ps(vacc0x89ABCDEF, vmax);
120     vacc1x89ABCDEF = _mm256_min_ps(vacc1x89ABCDEF, vmax);
121     vacc2x89ABCDEF = _mm256_min_ps(vacc2x89ABCDEF, vmax);
122 
123     if XNN_LIKELY(nc >= 16) {
124       _mm_storeu_si128((__m128i*) c2, _mm256_cvtps_ph(vacc2x01234567, _MM_FROUND_NO_EXC));
125       _mm_storeu_si128((__m128i*) (c2 + 8), _mm256_cvtps_ph(vacc2x89ABCDEF, _MM_FROUND_NO_EXC));
126       c2 = (uint16_t*) ((uintptr_t) c2 + cn_stride);
127       _mm_storeu_si128((__m128i*) c1, _mm256_cvtps_ph(vacc1x01234567, _MM_FROUND_NO_EXC));
128       _mm_storeu_si128((__m128i*) (c1 + 8), _mm256_cvtps_ph(vacc1x89ABCDEF, _MM_FROUND_NO_EXC));
129       c1 = (uint16_t*) ((uintptr_t) c1 + cn_stride);
130       _mm_storeu_si128((__m128i*) c0, _mm256_cvtps_ph(vacc0x01234567, _MM_FROUND_NO_EXC));
131       _mm_storeu_si128((__m128i*) (c0 + 8), _mm256_cvtps_ph(vacc0x89ABCDEF, _MM_FROUND_NO_EXC));
132       c0 = (uint16_t*) ((uintptr_t) c0 + cn_stride);
133 
134       a = (const void**restrict) ((uintptr_t) a - ks);
135       nc -= 16;
136     } else {
137       __m128i vh2x01234567 = _mm256_cvtps_ph(vacc2x01234567, _MM_FROUND_NO_EXC);
138       __m128i vh1x01234567 = _mm256_cvtps_ph(vacc1x01234567, _MM_FROUND_NO_EXC);
139       __m128i vh0x01234567 = _mm256_cvtps_ph(vacc0x01234567, _MM_FROUND_NO_EXC);
140       if (nc & 8) {
141         _mm_storeu_si128((__m128i*) c2, vh2x01234567);
142         _mm_storeu_si128((__m128i*) c1, vh1x01234567);
143         _mm_storeu_si128((__m128i*) c0, vh0x01234567);
144 
145         vh2x01234567 = _mm256_cvtps_ph(vacc2x89ABCDEF, _MM_FROUND_NO_EXC);
146         vh1x01234567 = _mm256_cvtps_ph(vacc1x89ABCDEF, _MM_FROUND_NO_EXC);
147         vh0x01234567 = _mm256_cvtps_ph(vacc0x89ABCDEF, _MM_FROUND_NO_EXC);
148 
149         c2 += 8;
150         c1 += 8;
151         c0 += 8;
152       }
153       if (nc & 4) {
154         _mm_storel_epi64((__m128i*) c2, vh2x01234567);
155         _mm_storel_epi64((__m128i*) c1, vh1x01234567);
156         _mm_storel_epi64((__m128i*) c0, vh0x01234567);
157 
158         vh2x01234567 = _mm_unpackhi_epi64(vh2x01234567, vh2x01234567);
159         vh1x01234567 = _mm_unpackhi_epi64(vh1x01234567, vh1x01234567);
160         vh0x01234567 = _mm_unpackhi_epi64(vh0x01234567, vh0x01234567);
161 
162         c2 += 4;
163         c1 += 4;
164         c0 += 4;
165       }
166       if (nc & 2) {
167         _mm_storeu_si32(c2, vh2x01234567);
168         _mm_storeu_si32(c1, vh1x01234567);
169         _mm_storeu_si32(c0, vh0x01234567);
170 
171         vh2x01234567 = _mm_srli_epi64(vh2x01234567, 32);
172         vh1x01234567 = _mm_srli_epi64(vh1x01234567, 32);
173         vh0x01234567 = _mm_srli_epi64(vh0x01234567, 32);
174 
175         c2 += 2;
176         c1 += 2;
177         c0 += 2;
178       }
179       if (nc & 1) {
180         *c2 = _mm_extract_epi16(vh2x01234567, 0);
181         *c1 = _mm_extract_epi16(vh1x01234567, 0);
182         *c0 = _mm_extract_epi16(vh0x01234567, 0);
183       }
184 
185       nc = 0;
186     }
187   } while (nc != 0);
188 }
189