xref: /aosp_15_r20/external/XNNPACK/src/f16-vmulcaddc/fma3.c.in (revision 4bdc94577ba0e567308109d787f7fec7b531ce36)
1// Copyright 2022 Google LLC
2//
3// This source code is licensed under the BSD-style license found in the
4// LICENSE file in the root directory of this source tree.
5
6$assert CHANNEL_TILE % 8 == 0
7$assert CHANNEL_TILE >= 8
8$assert ROW_TILE >= 1
9$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
10#include <assert.h>
11
12#include <immintrin.h>
13
14#include <xnnpack/intrinsics-polyfill.h>
15#include <xnnpack/math.h>
16#include <xnnpack/vmulcaddc.h>
17
18
19void xnn_f16_vmulcaddc_minmax_ukernel_c${CHANNEL_TILE}__fma3_${ROW_TILE}x(
20    size_t rows,
21    size_t channels,
22    const void*restrict input,
23    size_t input_stride,
24    const void*restrict weights,
25    void*restrict output,
26    size_t output_stride,
27    const union xnn_f16_minmax_params params[restrict XNN_MIN_ELEMENTS(1)]) XNN_OOB_READS
28{
29  assert(rows != 0);
30  assert(channels != 0);
31  assert(channels % sizeof(uint16_t) == 0);
32
33  const uint16_t* i0 = (const uint16_t*) input;
34  uint16_t* o0 = (uint16_t*) output;
35  $for M in range(1, ROW_TILE):
36    const uint16_t* i${M} = (const uint16_t*) ((uintptr_t) i${M-1} + input_stride);
37    uint16_t* o${M} = (uint16_t*) ((uintptr_t) o${M-1} + output_stride);
38
39  const size_t input_increment = input_stride * ${ROW_TILE} - channels;
40  const size_t output_increment = output_stride * ${ROW_TILE} - channels;
41
42  const __m256 vmin = _mm256_load_ps(params->avx.min);
43  const __m256 vmax = _mm256_load_ps(params->avx.max);
44  do {
45    $for M in range(1, ROW_TILE):
46      $if M % 2 == 0:
47        if XNN_UNPREDICTABLE(rows <= ${M}) {
48          i${M} = i${M-1};
49          o${M} = o${M-1};
50        }
51      $else:
52        if XNN_UNPREDICTABLE(rows < ${M+1}) {
53          i${M} = i${M-1};
54          o${M} = o${M-1};
55        }
56
57    const uint16_t* w = (const uint16_t*) weights;
58    size_t c = channels;
59    $if CHANNEL_TILE > 8:
60      for (; c >= ${CHANNEL_TILE} * sizeof(uint16_t); c -= ${CHANNEL_TILE} * sizeof(uint16_t)) {
61        const __m256 vscale${ABC[0:8]} = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) w));
62        $for C in range(8, CHANNEL_TILE, 8):
63          const __m256 vscale${ABC[C:C+8]} = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) (w + ${C})));
64
65        $for M in range(ROW_TILE):
66          __m256 vacc${M}x${ABC[0:8]} = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) i${M}));
67          $for C in range(8, CHANNEL_TILE, 8):
68            __m256 vacc${M}x${ABC[C:C+8]} = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) (i${M} + ${C})));
69          i${M} += ${CHANNEL_TILE};
70
71        $for C in range(0, CHANNEL_TILE, 8):
72          const __m256 vbias${ABC[C:C+8]} = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) (w + ${CHANNEL_TILE + C})));
73        w += ${2 * CHANNEL_TILE};
74
75        $for M in range(ROW_TILE):
76          $for C in range(0, CHANNEL_TILE, 8):
77            vacc${M}x${ABC[C:C+8]} = _mm256_fmadd_ps(vacc${M}x${ABC[C:C+8]}, vscale${ABC[C:C+8]}, vbias${ABC[C:C+8]});
78
79        $for M in range(ROW_TILE):
80          $for C in range(0, CHANNEL_TILE, 8):
81            vacc${M}x${ABC[C:C+8]} = _mm256_max_ps(vacc${M}x${ABC[C:C+8]}, vmin);
82
83        $for M in range(ROW_TILE):
84          $for C in range(0, CHANNEL_TILE, 8):
85            vacc${M}x${ABC[C:C+8]} = _mm256_min_ps(vacc${M}x${ABC[C:C+8]}, vmax);
86
87        $for M in range(ROW_TILE):
88          _mm_storeu_si128((__m128i*) o${M}, _mm256_cvtps_ph(vacc${M}x${ABC[0:8]}, _MM_FROUND_NO_EXC));
89          $for C in range(8, CHANNEL_TILE, 8):
90            _mm_storeu_si128((__m128i*) (o${M} + ${C}), _mm256_cvtps_ph(vacc${M}x${ABC[C:C+8]}, _MM_FROUND_NO_EXC));
91          o${M} += ${CHANNEL_TILE};
92      }
93    for (; c >= 8 * sizeof(uint16_t); c -= 8 * sizeof(uint16_t)) {
94      const __m256 vscale = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) w));
95
96      $for M in range(ROW_TILE):
97        __m256 vacc${M} = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) i${M}));
98        i${M} += 8;
99
100      const __m256 vbias = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) (w + ${CHANNEL_TILE})));
101      w += ${8 if CHANNEL_TILE > 8 else CHANNEL_TILE * 2};
102
103      $for M in range(ROW_TILE):
104        vacc${M} = _mm256_fmadd_ps(vacc${M}, vscale, vbias);
105
106      $for M in range(ROW_TILE):
107        vacc${M} = _mm256_max_ps(vacc${M}, vmin);
108
109      $for M in range(ROW_TILE):
110        vacc${M} = _mm256_min_ps(vacc${M}, vmax);
111
112      $for M in range(ROW_TILE):
113        _mm_storeu_si128((__m128i*) o${M}, _mm256_cvtps_ph(vacc${M}, _MM_FROUND_NO_EXC));
114        o${M} += 8;
115    }
116    if XNN_UNLIKELY(c != 0) {
117      const __m256 vscale = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) w));
118
119      $for M in range(ROW_TILE):
120        __m256 vacc${M} = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) i${M}));
121        i${M} = (const uint16_t*) ((uintptr_t) i${M} + c);
122
123      const __m256 vbias = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) (w + ${CHANNEL_TILE})));
124
125      $for M in range(ROW_TILE):
126        vacc${M} = _mm256_fmadd_ps(vacc${M}, vscale, vbias);
127
128      $for M in range(ROW_TILE):
129        vacc${M} = _mm256_max_ps(vacc${M}, vmin);
130
131      $for M in range(ROW_TILE):
132        vacc${M} = _mm256_min_ps(vacc${M}, vmax);
133
134      $for M in range(ROW_TILE):
135        __m128i vh${M} = _mm256_cvtps_ph(vacc${M}, _MM_FROUND_NO_EXC);
136
137      if (c & (4 * sizeof(uint16_t))) {
138        $for M in range(ROW_TILE):
139          _mm_storel_epi64((__m128i*) o${M}, vh${M});
140
141        $for M in range(ROW_TILE):
142          vh${M} = _mm_unpackhi_epi64(vh${M}, vh${M});
143
144        $for M in range(ROW_TILE):
145          o${M} += 4;
146      }
147      if (c & (2 * sizeof(uint16_t))) {
148        $for M in range(ROW_TILE):
149          _mm_storeu_si32(o${M}, vh${M});
150
151        $for M in range(ROW_TILE):
152          vh${M} = _mm_srli_epi64(vh${M}, 32);
153
154        $for M in range(ROW_TILE):
155          o${M} += 2;
156      }
157      if (c & (1 * sizeof(uint16_t))) {
158        $for M in range(ROW_TILE):
159          *o${M} = (uint16_t) _mm_extract_epi16(vh${M}, 0);
160
161        $for M in range(ROW_TILE):
162          o${M} += 1;
163      }
164    }
165    $for M in range(ROW_TILE):
166      i${M} = (const uint16_t*) ((uintptr_t) i${M} + input_increment);
167      o${M} = (uint16_t*) ((uintptr_t) o${M} + output_increment);
168    rows = doz(rows, ${ROW_TILE});
169  } while (rows != 0);
170}
171