xref: /aosp_15_r20/external/XNNPACK/src/xx-fill/neon-x64.c (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 #include <assert.h>
7 
8 #include <arm_neon.h>
9 
10 #include <xnnpack/fill.h>
11 
12 
xnn_xx_fill_ukernel__neon_x64(size_t rows,size_t channels,void * output,size_t output_stride,const uint32_t fill_pattern)13 void xnn_xx_fill_ukernel__neon_x64(
14     size_t rows,
15     size_t channels,
16     void* output,
17     size_t output_stride,
18     const uint32_t fill_pattern)
19 {
20   assert(rows != 0);
21   assert(channels != 0);
22 
23   const size_t output_increment = output_stride - channels;
24 
25   const uint8x16_t vfill_pattern = vreinterpretq_u8_u32(vdupq_n_u32(fill_pattern));
26   do {
27     size_t c = channels;
28     for (; c >= 64 * sizeof(uint8_t); c -= 64 * sizeof(uint8_t)) {
29       vst1q_u8(output, vfill_pattern); output = ((uint8_t*) output + 16);
30       vst1q_u8(output, vfill_pattern); output = ((uint8_t*) output + 16);
31       vst1q_u8(output, vfill_pattern); output = ((uint8_t*) output + 16);
32       vst1q_u8(output, vfill_pattern); output = ((uint8_t*) output + 16);
33     }
34     for (; c >= 16 * sizeof(uint8_t); c -= 16 * sizeof(uint8_t)) {
35       vst1q_u8(output, vfill_pattern); output = ((uint8_t*) output + 16);
36     }
37     if XNN_UNLIKELY(c != 0) {
38       if XNN_LIKELY(c & (8 * sizeof(uint8_t))) {
39         vst1_u8(output, vget_low_u8(vfill_pattern)); output = ((uint8_t*) output + 8);
40       }
41       if XNN_LIKELY(c & (4 * sizeof(uint8_t))) {
42         vst1q_lane_u32(output, vreinterpretq_u32_u8(vfill_pattern), 0); output = ((uint8_t*) output + 4);
43       }
44       uint8x8_t vfill_subpattern = vget_low_u8(vfill_pattern);
45       if XNN_LIKELY(c & (2 * sizeof(uint8_t))) {
46         vst1_lane_u16(output, vreinterpret_u16_u8(vfill_subpattern), 0); output = ((uint8_t*) output + 2);
47         vfill_subpattern = vext_u8(vfill_subpattern, vfill_subpattern, 2);
48       }
49       if XNN_LIKELY(c & (1 * sizeof(uint8_t))) {
50         vst1_lane_u8(output, vfill_subpattern, 0); output = ((uint8_t*) output + 1);
51       }
52     }
53     output = (void*) ((uintptr_t) output + output_increment);
54   } while (--rows != 0);
55 }
56