blob: fcde413804fa6dd51bbc17cc8eae0f9cf012ffb8 [file] [log] [blame]
Marat Dukhan9a88efe2019-12-10 15:54:24 -08001// Copyright 2019 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 BATCH_TILE % 16 == 0
7$assert BATCH_TILE >= 16
8$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
Marat Dukhan13bafb02020-06-05 00:43:11 -07009$assert OP in ["ADD", "DIV", "MAX", "MIN", "MUL", "SUB", "SQRDIFF"]
Marat Dukhan91cd2b72020-04-09 23:57:31 -070010$assert ACTIVATION in ["LINEAR", "MINMAX"]
Marat Dukhan9a88efe2019-12-10 15:54:24 -080011#include <assert.h>
12
13#include <immintrin.h>
14
15#include <xnnpack/common.h>
Marat Dukhan9e0e8ee2019-12-13 15:34:13 -080016#include <xnnpack/intrinsics-polyfill.h>
Marat Dukhan9a88efe2019-12-10 15:54:24 -080017#include <xnnpack/vbinary.h>
18
19
20$_MM512_OP_PS = {
21$ "ADD": lambda x, y: "_mm512_add_ps(%s, %s)" % (x, y),
22$ "DIV": lambda x, y: "_mm512_div_ps(%s, %s)" % (x, y),
23$ "MAX": lambda x, y: "_mm512_max_ps(%s, %s)" % (x, y),
24$ "MIN": lambda x, y: "_mm512_min_ps(%s, %s)" % (x, y),
25$ "MUL": lambda x, y: "_mm512_mul_ps(%s, %s)" % (x, y),
26$ "SUB": lambda x, y: "_mm512_sub_ps(%s, %s)" % (x, y),
Marat Dukhan13bafb02020-06-05 00:43:11 -070027$ "SQRDIFF": lambda x, y: "_mm512_sub_ps(%s, %s)" % (x, y),
Marat Dukhan9a88efe2019-12-10 15:54:24 -080028$}[OP]
Marat Dukhan91cd2b72020-04-09 23:57:31 -070029$SUFFIX = {"LINEAR": "", "MINMAX": "_minmax"}[ACTIVATION]
30$PARAMS = {"LINEAR": "xnn_f32_default_params", "MINMAX": "xnn_f32_minmax_params"}[ACTIVATION]
31void xnn_f32_v${OP.lower()}${SUFFIX}_ukernel__avx512f_x${BATCH_TILE}(
Marat Dukhan9a88efe2019-12-10 15:54:24 -080032 size_t n,
33 const float* a,
34 const float* b,
35 float* y,
Marat Dukhanf196d012020-04-15 11:50:03 -070036 const union ${PARAMS} params[restrict XNN_MIN_ELEMENTS(1)])
Marat Dukhan9a88efe2019-12-10 15:54:24 -080037{
38 assert(n != 0);
39 assert(n % sizeof(float) == 0);
40
Marat Dukhan91cd2b72020-04-09 23:57:31 -070041 $if ACTIVATION == "MINMAX":
42 const __m512 vy_min = _mm512_broadcast_f32x4(_mm_load_ps(params->sse.min));
43 const __m512 vy_max = _mm512_broadcast_f32x4(_mm_load_ps(params->sse.max));
Marat Dukhan9a88efe2019-12-10 15:54:24 -080044
45 for (; n >= ${BATCH_TILE} * sizeof(float); n -= ${BATCH_TILE} * sizeof(float)) {
46 const __m512 va${ABC[0:16]} = _mm512_loadu_ps(a);
47 $for N in range(16, BATCH_TILE, 16):
48 const __m512 va${ABC[N:N+16]} = _mm512_loadu_ps(a + ${N});
49 a += ${BATCH_TILE};
50
51 const __m512 vb${ABC[0:16]} = _mm512_loadu_ps(b);
52 $for N in range(16, BATCH_TILE, 16):
53 const __m512 vb${ABC[N:N+16]} = _mm512_loadu_ps(b + ${N});
54 b += ${BATCH_TILE};
55
56 $for N in range(0, BATCH_TILE, 16):
57 __m512 vy${ABC[N:N+16]} = ${_MM512_OP_PS("va" + ABC[N:N+16], "vb" + ABC[N:N+16])};
58
Marat Dukhan13bafb02020-06-05 00:43:11 -070059 $if OP == "SQRDIFF":
60 $for N in range(0, BATCH_TILE, 16):
61 vy${ABC[N:N+16]} = _mm512_mul_ps(vy${ABC[N:N+16]}, vy${ABC[N:N+16]});
62
Marat Dukhan91cd2b72020-04-09 23:57:31 -070063 $if ACTIVATION == "MINMAX":
64 $for N in range(0, BATCH_TILE, 16):
65 vy${ABC[N:N+16]} = _mm512_max_ps(vy${ABC[N:N+16]}, vy_min);
Marat Dukhan9a88efe2019-12-10 15:54:24 -080066
Marat Dukhan91cd2b72020-04-09 23:57:31 -070067 $for N in range(0, BATCH_TILE, 16):
68 vy${ABC[N:N+16]} = _mm512_min_ps(vy${ABC[N:N+16]}, vy_max);
Marat Dukhan9a88efe2019-12-10 15:54:24 -080069
70 _mm512_storeu_ps(y, vy${ABC[0:16]});
71 $for N in range(16, BATCH_TILE, 16):
72 _mm512_storeu_ps(y + ${N}, vy${ABC[N:N+16]});
73 y += ${BATCH_TILE};
74 }
Marat Dukhan0219fd52020-06-08 13:55:46 -070075 $if BATCH_TILE > 16:
Marat Dukhan9a88efe2019-12-10 15:54:24 -080076 for (; n >= 16 * sizeof(float); n -= 16 * sizeof(float)) {
77 const __m512 va = _mm512_loadu_ps(a);
78 a += 16;
79
80 const __m512 vb = _mm512_loadu_ps(b);
81 b += 16;
82
83 __m512 vy = ${_MM512_OP_PS("va", "vb")};
Marat Dukhan13bafb02020-06-05 00:43:11 -070084 $if OP == "SQRDIFF":
85 vy = _mm512_mul_ps(vy, vy);
Marat Dukhan91cd2b72020-04-09 23:57:31 -070086 $if ACTIVATION == "MINMAX":
87 vy = _mm512_max_ps(vy, vy_min);
88 vy = _mm512_min_ps(vy, vy_max);
Marat Dukhan9a88efe2019-12-10 15:54:24 -080089 _mm512_storeu_ps(y, vy);
90 y += 16;
91 }
92 if XNN_UNLIKELY(n != 0) {
93 assert(n >= 1 * sizeof(float));
94 assert(n <= 15 * sizeof(float));
95 // Prepare mask for valid 32-bit elements (depends on n).
96 n >>= 2 /* log2(sizeof(float)) */;
97 const __mmask16 vmask = _cvtu32_mask16((uint16_t) ((uint32_t) (UINT32_C(1) << n) - UINT32_C(1)));
98
99 const __m512 va = _mm512_maskz_loadu_ps(vmask, a);
100 const __m512 vb = _mm512_maskz_loadu_ps(vmask, b);
101
102 __m512 vy = ${_MM512_OP_PS("va", "vb")};
Marat Dukhan13bafb02020-06-05 00:43:11 -0700103 $if OP == "SQRDIFF":
104 vy = _mm512_mul_ps(vy, vy);
Marat Dukhan91cd2b72020-04-09 23:57:31 -0700105 $if ACTIVATION == "MINMAX":
106 vy = _mm512_max_ps(vy, vy_min);
107 vy = _mm512_min_ps(vy, vy_max);
Marat Dukhan9a88efe2019-12-10 15:54:24 -0800108 _mm512_mask_storeu_ps(y, vmask, vy);
109 }
110}