xref: /aosp_15_r20/external/XNNPACK/src/f16-vbinary/vopc-f16c.c.in (revision 4bdc94577ba0e567308109d787f7fec7b531ce36)
1*4bdc9457SAndroid Build Coastguard Worker// Copyright 2022 Google LLC
2*4bdc9457SAndroid Build Coastguard Worker//
3*4bdc9457SAndroid Build Coastguard Worker// This source code is licensed under the BSD-style license found in the
4*4bdc9457SAndroid Build Coastguard Worker// LICENSE file in the root directory of this source tree.
5*4bdc9457SAndroid Build Coastguard Worker
6*4bdc9457SAndroid Build Coastguard Worker$assert BATCH_TILE % 8 == 0
7*4bdc9457SAndroid Build Coastguard Worker$assert BATCH_TILE >= 8
8*4bdc9457SAndroid Build Coastguard Worker$ABC = "01234567456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
9*4bdc9457SAndroid Build Coastguard Worker$assert OP in ["ADD", "DIV", "RDIV", "MAX", "MIN", "MUL", "SUB", "RSUB", "SQRDIFF"]
10*4bdc9457SAndroid Build Coastguard Worker$assert ACTIVATION in ["LINEAR", "MINMAX"]
11*4bdc9457SAndroid Build Coastguard Worker#include <assert.h>
12*4bdc9457SAndroid Build Coastguard Worker
13*4bdc9457SAndroid Build Coastguard Worker#include <immintrin.h>
14*4bdc9457SAndroid Build Coastguard Worker
15*4bdc9457SAndroid Build Coastguard Worker#include <xnnpack/common.h>
16*4bdc9457SAndroid Build Coastguard Worker#include <xnnpack/intrinsics-polyfill.h>
17*4bdc9457SAndroid Build Coastguard Worker#include <xnnpack/vbinary.h>
18*4bdc9457SAndroid Build Coastguard Worker
19*4bdc9457SAndroid Build Coastguard Worker
20*4bdc9457SAndroid Build Coastguard Worker$_MM256_OP_PS = {
21*4bdc9457SAndroid Build Coastguard Worker$  "ADD": lambda x: "_mm256_add_ps(%s, vb)" % x,
22*4bdc9457SAndroid Build Coastguard Worker$  "DIV": lambda x: "_mm256_div_ps(%s, vb)" % x,
23*4bdc9457SAndroid Build Coastguard Worker$  "RDIV": lambda x: "_mm256_div_ps(vb, %s)" % x,
24*4bdc9457SAndroid Build Coastguard Worker$  "MAX": lambda x: "_mm256_max_ps(%s, vb)" % x,
25*4bdc9457SAndroid Build Coastguard Worker$  "MIN": lambda x: "_mm256_min_ps(%s, vb)" % x,
26*4bdc9457SAndroid Build Coastguard Worker$  "MUL": lambda x: "_mm256_mul_ps(%s, vb)" % x,
27*4bdc9457SAndroid Build Coastguard Worker$  "SUB": lambda x: "_mm256_sub_ps(%s, vb)" % x,
28*4bdc9457SAndroid Build Coastguard Worker$  "RSUB": lambda x: "_mm256_sub_ps(vb, %s)" % x,
29*4bdc9457SAndroid Build Coastguard Worker$  "SQRDIFF": lambda x: "_mm256_sub_ps(%s, vb)" % x,
30*4bdc9457SAndroid Build Coastguard Worker$}[OP]
31*4bdc9457SAndroid Build Coastguard Worker$SUFFIX = {"LINEAR": "", "MINMAX": "_minmax"}[ACTIVATION]
32*4bdc9457SAndroid Build Coastguard Worker$PARAMS = {"LINEAR": "xnn_f16_default_params", "MINMAX": "xnn_f16_minmax_params"}[ACTIVATION]
33*4bdc9457SAndroid Build Coastguard Workervoid xnn_f16_v${OP.lower()}c${SUFFIX}_ukernel__f16c_x${BATCH_TILE}(
34*4bdc9457SAndroid Build Coastguard Worker    size_t n,
35*4bdc9457SAndroid Build Coastguard Worker    const void* restrict a_ptr,
36*4bdc9457SAndroid Build Coastguard Worker    const void* restrict b_ptr,
37*4bdc9457SAndroid Build Coastguard Worker    void* restrict y_ptr,
38*4bdc9457SAndroid Build Coastguard Worker    const union ${PARAMS} params[restrict XNN_MIN_ELEMENTS(1)]) XNN_OOB_READS
39*4bdc9457SAndroid Build Coastguard Worker{
40*4bdc9457SAndroid Build Coastguard Worker  assert(n != 0);
41*4bdc9457SAndroid Build Coastguard Worker  assert(n % sizeof(uint16_t) == 0);
42*4bdc9457SAndroid Build Coastguard Worker  assert(a_ptr != NULL);
43*4bdc9457SAndroid Build Coastguard Worker  assert(b_ptr != NULL);
44*4bdc9457SAndroid Build Coastguard Worker  assert(y_ptr != NULL);
45*4bdc9457SAndroid Build Coastguard Worker
46*4bdc9457SAndroid Build Coastguard Worker  const uint16_t* a = (const uint16_t*) a_ptr;
47*4bdc9457SAndroid Build Coastguard Worker  const uint16_t* b = (const uint16_t*) b_ptr;
48*4bdc9457SAndroid Build Coastguard Worker  uint16_t* y = (uint16_t*) y_ptr;
49*4bdc9457SAndroid Build Coastguard Worker
50*4bdc9457SAndroid Build Coastguard Worker  $if ACTIVATION == "MINMAX":
51*4bdc9457SAndroid Build Coastguard Worker    const __m256 vy_min = _mm256_load_ps(params->avx.min);
52*4bdc9457SAndroid Build Coastguard Worker    const __m256 vy_max = _mm256_load_ps(params->avx.max);
53*4bdc9457SAndroid Build Coastguard Worker
54*4bdc9457SAndroid Build Coastguard Worker  const __m256 vb = _mm256_cvtph_ps(_mm_set1_epi16((short) *b));
55*4bdc9457SAndroid Build Coastguard Worker  $if BATCH_TILE > 8:
56*4bdc9457SAndroid Build Coastguard Worker    for (; n >= ${BATCH_TILE} * sizeof(uint16_t); n -= ${BATCH_TILE} * sizeof(uint16_t)) {
57*4bdc9457SAndroid Build Coastguard Worker      const __m256 va${ABC[0:8]} = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) a));
58*4bdc9457SAndroid Build Coastguard Worker      $for N in range(8, BATCH_TILE, 8):
59*4bdc9457SAndroid Build Coastguard Worker        const __m256 va${ABC[N:N+8]} = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) (a + ${N})));
60*4bdc9457SAndroid Build Coastguard Worker      a += ${BATCH_TILE};
61*4bdc9457SAndroid Build Coastguard Worker
62*4bdc9457SAndroid Build Coastguard Worker      $for N in range(0, BATCH_TILE, 8):
63*4bdc9457SAndroid Build Coastguard Worker        __m256 vy${ABC[N:N+8]} = _mm256_cvtph_ps(_mm256_cvtps_ph(${_MM256_OP_PS("va" + ABC[N:N+8])}, _MM_FROUND_NO_EXC));
64*4bdc9457SAndroid Build Coastguard Worker
65*4bdc9457SAndroid Build Coastguard Worker      $if OP == "SQRDIFF":
66*4bdc9457SAndroid Build Coastguard Worker        $for N in range(0, BATCH_TILE, 8):
67*4bdc9457SAndroid Build Coastguard Worker          vy${ABC[N:N+8]} = _mm256_cvtph_ps(_mm256_cvtps_ph(_mm256_mul_ps(vy${ABC[N:N+8]}, vy${ABC[N:N+8]}), _MM_FROUND_NO_EXC));
68*4bdc9457SAndroid Build Coastguard Worker
69*4bdc9457SAndroid Build Coastguard Worker      $if ACTIVATION == "MINMAX":
70*4bdc9457SAndroid Build Coastguard Worker        $for N in range(0, BATCH_TILE, 8):
71*4bdc9457SAndroid Build Coastguard Worker          vy${ABC[N:N+8]} = _mm256_max_ps(vy${ABC[N:N+8]}, vy_min);
72*4bdc9457SAndroid Build Coastguard Worker
73*4bdc9457SAndroid Build Coastguard Worker        $for N in range(0, BATCH_TILE, 8):
74*4bdc9457SAndroid Build Coastguard Worker          vy${ABC[N:N+8]} = _mm256_min_ps(vy${ABC[N:N+8]}, vy_max);
75*4bdc9457SAndroid Build Coastguard Worker
76*4bdc9457SAndroid Build Coastguard Worker      _mm_storeu_si128((__m128i*) y, _mm256_cvtps_ph(vy${ABC[0:8]}, _MM_FROUND_NO_EXC));
77*4bdc9457SAndroid Build Coastguard Worker      $for N in range(8, BATCH_TILE, 8):
78*4bdc9457SAndroid Build Coastguard Worker        _mm_storeu_si128((__m128i*) (y + ${N}), _mm256_cvtps_ph(vy${ABC[N:N+8]}, _MM_FROUND_NO_EXC));
79*4bdc9457SAndroid Build Coastguard Worker      y += ${BATCH_TILE};
80*4bdc9457SAndroid Build Coastguard Worker    }
81*4bdc9457SAndroid Build Coastguard Worker  for (; n >= 8 * sizeof(uint16_t); n -= 8 * sizeof(uint16_t)) {
82*4bdc9457SAndroid Build Coastguard Worker    const __m256 va = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) a));
83*4bdc9457SAndroid Build Coastguard Worker    a += 8;
84*4bdc9457SAndroid Build Coastguard Worker
85*4bdc9457SAndroid Build Coastguard Worker    __m256 vy = _mm256_cvtph_ps(_mm256_cvtps_ph(${_MM256_OP_PS("va")}, _MM_FROUND_NO_EXC));
86*4bdc9457SAndroid Build Coastguard Worker    $if OP == "SQRDIFF":
87*4bdc9457SAndroid Build Coastguard Worker      vy = _mm256_cvtph_ps(_mm256_cvtps_ph(_mm256_mul_ps(vy, vy), _MM_FROUND_NO_EXC));
88*4bdc9457SAndroid Build Coastguard Worker
89*4bdc9457SAndroid Build Coastguard Worker    $if ACTIVATION == "MINMAX":
90*4bdc9457SAndroid Build Coastguard Worker      vy = _mm256_max_ps(vy, vy_min);
91*4bdc9457SAndroid Build Coastguard Worker      vy = _mm256_min_ps(vy, vy_max);
92*4bdc9457SAndroid Build Coastguard Worker
93*4bdc9457SAndroid Build Coastguard Worker    _mm_storeu_si128((__m128i*) y, _mm256_cvtps_ph(vy, _MM_FROUND_NO_EXC));
94*4bdc9457SAndroid Build Coastguard Worker    y += 8;
95*4bdc9457SAndroid Build Coastguard Worker  }
96*4bdc9457SAndroid Build Coastguard Worker  if XNN_UNLIKELY(n != 0) {
97*4bdc9457SAndroid Build Coastguard Worker    const __m256 va = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) a));
98*4bdc9457SAndroid Build Coastguard Worker
99*4bdc9457SAndroid Build Coastguard Worker    __m256 vy = _mm256_cvtph_ps(_mm256_cvtps_ph(${_MM256_OP_PS("va")}, _MM_FROUND_NO_EXC));
100*4bdc9457SAndroid Build Coastguard Worker    $if OP == "SQRDIFF":
101*4bdc9457SAndroid Build Coastguard Worker      vy = _mm256_cvtph_ps(_mm256_cvtps_ph(_mm256_mul_ps(vy, vy), _MM_FROUND_NO_EXC));
102*4bdc9457SAndroid Build Coastguard Worker
103*4bdc9457SAndroid Build Coastguard Worker    $if ACTIVATION == "MINMAX":
104*4bdc9457SAndroid Build Coastguard Worker      vy = _mm256_max_ps(vy, vy_min);
105*4bdc9457SAndroid Build Coastguard Worker      vy = _mm256_min_ps(vy, vy_max);
106*4bdc9457SAndroid Build Coastguard Worker
107*4bdc9457SAndroid Build Coastguard Worker    __m128i vh = _mm256_cvtps_ph(vy, _MM_FROUND_NO_EXC);
108*4bdc9457SAndroid Build Coastguard Worker    if (n & (4 * sizeof(uint16_t))) {
109*4bdc9457SAndroid Build Coastguard Worker      _mm_storel_epi64((__m128i*) y, vh);
110*4bdc9457SAndroid Build Coastguard Worker      vh = _mm_unpackhi_epi64(vh, vh);
111*4bdc9457SAndroid Build Coastguard Worker      y += 4;
112*4bdc9457SAndroid Build Coastguard Worker    }
113*4bdc9457SAndroid Build Coastguard Worker    if (n & (2 * sizeof(uint16_t))) {
114*4bdc9457SAndroid Build Coastguard Worker      _mm_storeu_si32(y, vh);
115*4bdc9457SAndroid Build Coastguard Worker      vh = _mm_srli_epi64(vh, 32);
116*4bdc9457SAndroid Build Coastguard Worker      y += 2;
117*4bdc9457SAndroid Build Coastguard Worker    }
118*4bdc9457SAndroid Build Coastguard Worker    if (n & (1 * sizeof(uint16_t))) {
119*4bdc9457SAndroid Build Coastguard Worker      *y = (uint16_t) _mm_extract_epi16(vh, 0);
120*4bdc9457SAndroid Build Coastguard Worker    }
121*4bdc9457SAndroid Build Coastguard Worker  }
122*4bdc9457SAndroid Build Coastguard Worker}
123