1 // 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 #include <assert.h>
7
8 #include <arm_neon.h>
9
10 #include <xnnpack/packx.h>
11
12
xnn_x32_packx_ukernel_4x__neon_st4(size_t m,size_t k,const uint32_t * restrict x,size_t x_stride,uint32_t * restrict y)13 void xnn_x32_packx_ukernel_4x__neon_st4(
14 size_t m,
15 size_t k,
16 const uint32_t* restrict x,
17 size_t x_stride,
18 uint32_t* restrict y)
19 {
20 assert(m != 0);
21 assert(k != 0);
22
23 const uint32_t* x0 = x;
24 const uint32_t* x1 = (const uint32_t*) ((uintptr_t) x0 + x_stride);
25 if (m < 2) {
26 x1 = x0;
27 }
28 const uint32_t* x2 = (const uint32_t*) ((uintptr_t) x1 + x_stride);
29 if (m <= 2) {
30 x2 = x1;
31 }
32 const uint32_t* x3 = (const uint32_t*) ((uintptr_t) x2 + x_stride);
33 if (m != 4) {
34 x3 = x2;
35 }
36
37 for (; k >= 4; k -= 4) {
38 const uint32x4_t vx0 = vld1q_u32(x0); x0 += 4;
39 const uint32x4_t vx1 = vld1q_u32(x1); x1 += 4;
40 const uint32x4_t vx2 = vld1q_u32(x2); x2 += 4;
41 const uint32x4_t vx3 = vld1q_u32(x3); x3 += 4;
42
43 const uint32x4x4_t vy = { vx0, vx1, vx2, vx3 };
44 vst4q_u32(y, vy); y += 16;
45 }
46 if XNN_UNLIKELY(k != 0) {
47 do {
48 const uint32x2_t vx00 = vld1_dup_u32(x0); x0 += 1;
49 const uint32x2_t vx22 = vld1_dup_u32(x2); x2 += 1;
50 const uint32x2_t vx01 = vld1_lane_u32(x1, vx00, 1); x1 += 1;
51 const uint32x2_t vx23 = vld1_lane_u32(x3, vx22, 1); x3 += 1;
52 const uint32x4_t vy = vcombine_u32(vx01, vx23);
53 vst1q_u32(y, vy); y += 4;
54 } while (--k != 0);
55 }
56 }
57