xref: /aosp_15_r20/external/XNNPACK/src/f32-dwconv2d-chw/3x3p1-ssse3.c.in (revision 4bdc94577ba0e567308109d787f7fec7b531ce36)
1// Copyright 2020 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 ROW_TILE >= 1
7$assert ACCUMULATORS >= 1
8#include <assert.h>
9
10#include <tmmintrin.h>
11
12#include <xnnpack/dwconv.h>
13#include <xnnpack/math.h>
14
15
16void xnn_f32_dwconv2d_chw_ukernel_3x3p1__ssse3_${ROW_TILE}x4${"_acc%d" % ACCUMULATORS if ACCUMULATORS > 1 else ""}(
17    size_t input_height,
18    size_t input_width,
19    const float* input,
20    const float* weights,
21    const float* zero,
22    float* output,
23    uint32_t padding_top,
24    const union xnn_f32_chw_params params[restrict XNN_MIN_ELEMENTS(1)]) XNN_OOB_READS
25{
26  assert(input_height != 0);
27  assert(input_width != 0);
28  assert(input_width % sizeof(float) == 0);
29  assert(padding_top == 1);
30
31  const __m128 vmask = _mm_load_ps((const float*) params->sse.mask);
32  const __m128 vmax = _mm_load_ps(params->sse.max);
33  const __m128 vmin = _mm_load_ps(params->sse.min);
34
35  const __m128 vbias = _mm_load1_ps(weights);
36  const __m128 vk00 = _mm_load1_ps(weights + 1);
37  const __m128 vk01 = _mm_load1_ps(weights + 2);
38  const __m128 vk02 = _mm_load1_ps(weights + 3);
39  const __m128 vk10 = _mm_load1_ps(weights + 4);
40  const __m128 vk11 = _mm_load1_ps(weights + 5);
41  const __m128 vk12 = _mm_load1_ps(weights + 6);
42  const __m128 vk20 = _mm_load1_ps(weights + 7);
43  const __m128 vk21 = _mm_load1_ps(weights + 8);
44  const __m128 vk22 = _mm_load1_ps(weights + 9);
45
46  const size_t input_decrement = round_up_po2(input_width, 4 * sizeof(float));
47
48  const float* i0 = zero;
49  const float* i1 = input;
50  $for M in range(2, 2 + ROW_TILE):
51    const float* i${M} = (const float*) ((uintptr_t) i${M-1} + input_width);
52
53  float* o0 = output;
54  $for M in range(1, ROW_TILE):
55    float* o${M} = (float*) ((uintptr_t) o${M-1} + input_width);
56
57  size_t output_height = input_height;
58  do {
59    $for M in range(2, 2 + ROW_TILE):
60      if XNN_UNPREDICTABLE(output_height < ${M}) {
61        i${M} = zero;
62        $if M <= ROW_TILE:
63          o${M-1} = o${M-2};
64      }
65
66    $for M in range(2 + ROW_TILE):
67      __m128 vi${M}x0123 = _mm_setzero_ps();
68
69    $for M in range(2 + ROW_TILE):
70      __m128 vi${M}x4567 = _mm_loadu_ps(i${M});
71      i${M} += 4;
72
73    size_t w = input_width;
74    for (; w > 4 * sizeof(float); w -= 4 * sizeof(float)) {
75      $for M in range(2 + ROW_TILE):
76        const __m128 vi${M}x89AB = _mm_loadu_ps(i${M});
77        i${M} += 4;
78
79      $for K in range(3):
80        $for M in range(ROW_TILE):
81          $if K == 0:
82            __m128 vo${M}p0 = _mm_add_ps(vbias, _mm_mul_ps(vi${M+K}x4567, vk${K}1));
83          $elif K < ACCUMULATORS:
84            __m128 vo${M}p${K} = _mm_mul_ps(vi${M+K}x4567, vk${K}1);
85          $else:
86            vo${M}p${K % ACCUMULATORS} = _mm_add_ps(vo${M}p${K % ACCUMULATORS}, _mm_mul_ps(vi${M+K}x4567, vk${K}1));
87
88      $for M in range(2 + ROW_TILE):
89        const __m128 vi${M}x3456 = _mm_castsi128_ps(_mm_alignr_epi8(_mm_castps_si128(vi${M}x4567), _mm_castps_si128(vi${M}x0123), 12));
90
91      $for K in range(3):
92        $for M in range(ROW_TILE):
93          $if K+3 < ACCUMULATORS:
94            __m128 vo${M}p${K+3} = _mm_mul_ps(vi${M+K}x3456, vk${K}0);
95          $else:
96            vo${M}p${(K+3) % ACCUMULATORS} = _mm_add_ps(vo${M}p${(K+3) % ACCUMULATORS}, _mm_mul_ps(vi${M+K}x3456, vk${K}0));
97
98      $for M in range(2 + ROW_TILE):
99        vi${M}x0123 = vi${M}x4567;
100
101      $for M in range(2 + ROW_TILE):
102        const __m128 vi${M}x5678 = _mm_castsi128_ps(_mm_alignr_epi8(_mm_castps_si128(vi${M}x89AB), _mm_castps_si128(vi${M}x4567), 4));
103
104      $for K in range(3):
105        $for M in range(ROW_TILE):
106          vo${M}p${(K+6) % ACCUMULATORS} = _mm_add_ps(vo${M}p${(K+6) % ACCUMULATORS}, _mm_mul_ps(vi${M+K}x5678, vk${K}2));
107
108      $for M in range(2 + ROW_TILE):
109        vi${M}x4567 = vi${M}x89AB;
110
111      $if ACCUMULATORS > 1:
112        $ACC_SLICE = 1
113        $while ACC_SLICE < ACCUMULATORS:
114          $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
115            $if A + ACC_SLICE < ACCUMULATORS:
116              $for M in range(ROW_TILE):
117                vo${M}p${A} = _mm_add_ps(vo${M}p${A}, vo${M}p${A + ACC_SLICE});
118          $ACC_SLICE *= 2
119
120      $for M in range(ROW_TILE):
121        __m128 vo${M} = _mm_max_ps(vo${M}p0, vmin);
122
123      $for M in range(ROW_TILE):
124        vo${M} = _mm_min_ps(vo${M}, vmax);
125
126      $for M in reversed(range(ROW_TILE)):
127        _mm_storeu_ps(o${M}, vo${M});
128        o${M} += 4;
129    }
130    // Always process the last block of 1..4 pixels.
131    assert(w >= 1 * sizeof(float));
132    assert(w <= 4 * sizeof(float));
133    {
134      $for M in range(2 + ROW_TILE):
135        vi${M}x4567 = _mm_and_ps(vmask, vi${M}x4567);
136
137      $for K in range(3):
138        $for M in range(ROW_TILE):
139          $if K == 0:
140            __m128 vo${M}p0 = _mm_add_ps(vbias, _mm_mul_ps(vi${M+K}x4567, vk${K}1));
141          $elif K < ACCUMULATORS:
142            __m128 vo${M}p${K} = _mm_mul_ps(vi${M+K}x4567, vk${K}1);
143          $else:
144            vo${M}p${K % ACCUMULATORS} = _mm_add_ps(vo${M}p${K % ACCUMULATORS}, _mm_mul_ps(vi${M+K}x4567, vk${K}1));
145
146      $for M in range(2 + ROW_TILE):
147        const __m128 vi${M}x3456 = _mm_castsi128_ps(_mm_alignr_epi8(_mm_castps_si128(vi${M}x4567), _mm_castps_si128(vi${M}x0123), 12));
148
149      $for K in range(3):
150        $for M in range(ROW_TILE):
151          $if K+3 < ACCUMULATORS:
152            __m128 vo${M}p${K+3} = _mm_mul_ps(vi${M+K}x3456, vk${K}0);
153          $else:
154            vo${M}p${(K+3) % ACCUMULATORS} = _mm_add_ps(vo${M}p${(K+3) % ACCUMULATORS}, _mm_mul_ps(vi${M+K}x3456, vk${K}0));
155
156      const __m128i vzero = _mm_setzero_si128();
157      $for M in range(2 + ROW_TILE):
158        const __m128 vi${M}x5678 = _mm_castsi128_ps(_mm_alignr_epi8(vzero, _mm_castps_si128(vi${M}x4567), 4));
159
160      $for K in range(3):
161        $for M in range(ROW_TILE):
162          vo${M}p${(K+6) % ACCUMULATORS} = _mm_add_ps(vo${M}p${(K+6) % ACCUMULATORS}, _mm_mul_ps(vi${M+K}x5678, vk${K}2));
163
164      $if ACCUMULATORS > 1:
165        $ACC_SLICE = 1
166        $while ACC_SLICE < ACCUMULATORS:
167          $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
168            $if A + ACC_SLICE < ACCUMULATORS:
169              $for M in range(ROW_TILE):
170                vo${M}p${A} = _mm_add_ps(vo${M}p${A}, vo${M}p${A + ACC_SLICE});
171          $ACC_SLICE *= 2
172
173      $for M in range(ROW_TILE):
174        __m128 vo${M} = _mm_max_ps(vo${M}p0, vmin);
175
176      $for M in range(ROW_TILE):
177        vo${M} = _mm_min_ps(vo${M}, vmax);
178
179      if XNN_LIKELY(w == 4 * sizeof(float)) {
180        $for M in reversed(range(ROW_TILE)):
181          _mm_storeu_ps(o${M}, vo${M});
182          o${M} += 4;
183      } else {
184        if (w & (2 * sizeof(float))) {
185          $for M in reversed(range(ROW_TILE)):
186            _mm_storel_pi((__m64*) o${M}, vo${M});
187            o${M} += 2;
188
189          $for M in range(ROW_TILE):
190            vo${M} = _mm_movehl_ps(vo${M}, vo${M});
191        }
192        if (w & (1 * sizeof(float))) {
193          $for M in reversed(range(ROW_TILE)):
194            _mm_store_ss(o${M}, vo${M});
195            o${M} += 1;
196        }
197      }
198    }
199
200    i0 = (const float*) ((uintptr_t) i${ROW_TILE} - input_decrement);
201    i1 = (const float*) ((uintptr_t) i${ROW_TILE+1} - input_decrement);
202    $for M in range(2, 2 + ROW_TILE):
203      i${M} = (const float*) ((uintptr_t) i${M-1} + input_width);
204
205    $if ROW_TILE > 1:
206      o0 = o${ROW_TILE - 1};
207      $for M in range(1, ROW_TILE):
208        o${M} = (float*) ((uintptr_t) o${M-1} + input_width);
209
210    $if ROW_TILE > 1:
211      output_height = doz(output_height, ${ROW_TILE});
212  } while (${"--" if ROW_TILE == 1 else ""}output_height != 0);
213}
214