1 // Copyright 2015 The Gemmlowp Authors. All Rights Reserved.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 // pack_neon.h: optimized NEON specializations of the templates in pack.h.
16 
17 #ifndef GEMMLOWP_INTERNAL_PACK_NEON_H_
18 #define GEMMLOWP_INTERNAL_PACK_NEON_H_
19 
20 #include "pack.h"
21 
22 #include <arm_neon.h>
23 
24 namespace gemmlowp {
25 
26 typedef SideMap<const std::uint8_t, SideMapOrder::WidthMajor>
27     WidthMajorUint8SideMap;
28 
29 typedef SideMap<const std::int8_t, SideMapOrder::WidthMajor>
30     WidthMajorInt8SideMap;
31 
32 template <int Cells>
33 using DepthMajorSideFormatNCells4x2 = KernelSideFormat<CellFormat<4, 2>, Cells>;
34 
35 template <int Cells>
36 class PackingRegisterBlock<
37     WidthMajorUint8SideMap,
38     PackedSideBlock<DepthMajorSideFormatNCells4x2<Cells>>>
39     : public PackingRegisterBlockBase<
40           WidthMajorUint8SideMap,
41           PackedSideBlock<DepthMajorSideFormatNCells4x2<Cells>>> {
42  public:
43   typedef DepthMajorSideFormatNCells4x2<Cells> KernelSideFormat;
44   typedef typename KernelSideFormat::Cell CellFormat;
45   static const int kCells = KernelSideFormat::kCells;
46   static const int kCellWidth = CellFormat::kWidth;
47   static const int kKernelWidth = CellFormat::kWidth * kCells;
48   static const int kCellDepth = CellFormat::kDepth;
49   static const int kCellSize = CellFormat::kSize;
50 
Pack(PackedSideBlock<KernelSideFormat> * dst,int start_width)51   void Pack(PackedSideBlock<KernelSideFormat>* dst, int start_width) {
52     std::uint8_t* dst_ptr = dst->current_data();
53     const std::uint8_t* const src_ptr = this->complete_src_.data();
54     const int stride = this->complete_src_.stride();
55     // Load source WidthMajor data
56     uint8x16_t src_lines[4 * kCells];
57     for (int i = 0; i < 4 * kCells; i++) {
58       src_lines[i] = vld1q_u8(src_ptr + i * stride);
59     }
60     // Reorder the data within registers to make DepthMajor 4x2 cells
61     uint8x16x2_t src_lines_intertwined_2x[2 * kCells];
62     for (int i = 0; i < kCells; i++) {
63       src_lines_intertwined_2x[2 * i] =
64           vzipq_u8(src_lines[4 * i], src_lines[4 * i + 2]);
65       src_lines_intertwined_2x[2 * i + 1] =
66           vzipq_u8(src_lines[4 * i + 1], src_lines[4 * i + 3]);
67     }
68     uint8x16x2_t src_lines_intertwined_4x[2 * kCells];
69     for (int i = 0; i < kCells; i++) {
70       src_lines_intertwined_4x[2 * i] =
71           vzipq_u8(src_lines_intertwined_2x[2 * i].val[0],
72                    src_lines_intertwined_2x[2 * i + 1].val[0]);
73       src_lines_intertwined_4x[2 * i + 1] =
74           vzipq_u8(src_lines_intertwined_2x[2 * i].val[1],
75                    src_lines_intertwined_2x[2 * i + 1].val[1]);
76     }
77     // Store the resulting DepthMajor 4x2 cells in the destination packed block
78     for (int outer = 0; outer < 2; outer++) {
79       for (int inner = 0; inner < 2; inner++) {
80         for (int cell = 0; cell < kCells; cell++) {
81           uint8x8_t value = vget_low_u8(
82               src_lines_intertwined_4x[2 * cell + outer].val[inner]);
83           vst1_u8(dst_ptr, value);
84           dst_ptr += 8;
85         }
86         for (int cell = 0; cell < kCells; cell++) {
87           uint8x8_t value = vget_high_u8(
88               src_lines_intertwined_4x[2 * cell + outer].val[inner]);
89           vst1_u8(dst_ptr, value);
90           dst_ptr += 8;
91         }
92       }
93     }
94     // Compute sums across the depth dimension
95     uint16x8_t sums_of_2_cells[kCells][4];
96     for (int outer = 0; outer < 2; outer++) {
97       for (int inner = 0; inner < 2; inner++) {
98         int i = 2 * outer + inner;
99         for (int cell = 0; cell < kCells; cell++) {
100           sums_of_2_cells[cell][i] = vaddl_u8(
101               vget_low_u8(
102                   src_lines_intertwined_4x[2 * cell + outer].val[inner]),
103               vget_high_u8(
104                   src_lines_intertwined_4x[2 * cell + outer].val[inner]));
105         }
106       }
107     }
108     int32x4_t sums_of_4_cells[kCells][4];
109     for (int i = 0; i < 4; i++) {
110       for (int cell = 0; cell < kCells; cell++) {
111         sums_of_4_cells[cell][i] = vreinterpretq_s32_u32(
112             vaddl_u16(vget_low_u16(sums_of_2_cells[cell][i]),
113                       vget_high_u16(sums_of_2_cells[cell][i])));
114       }
115     }
116     // Update the sums_of_each_slice vector
117     for (int cell = 0; cell < kCells; cell++) {
118       int32x4_t s01 =
119           vaddq_s32(sums_of_4_cells[cell][0], sums_of_4_cells[cell][1]);
120       int32x4_t s23 =
121           vaddq_s32(sums_of_4_cells[cell][2], sums_of_4_cells[cell][3]);
122       int32x4_t s = vaddq_s32(s01, s23);
123       std::int32_t* sums_of_each_slice_ptr =
124           dst->sums_of_each_slice() + start_width + 4 * cell;
125       vst1q_s32(sums_of_each_slice_ptr,
126                 vaddq_s32(s, vld1q_s32(sums_of_each_slice_ptr)));
127     }
128     dst->seek_forward_n_cells(kCells * kRegisterSize / kCellDepth);
129   }
130 };
131 
132 template <int Cells>
133 using WidthMajorSideFormatNCells4x2 =
134     KernelSideFormat<CellFormat<4, 2, CellOrder::WidthMajor>, Cells>;
135 
136 template <int Cells>
137 class PackingRegisterBlock<
138     WidthMajorUint8SideMap,
139     PackedSideBlock<WidthMajorSideFormatNCells4x2<Cells>>>
140     : public PackingRegisterBlockBase<
141           WidthMajorUint8SideMap,
142           PackedSideBlock<WidthMajorSideFormatNCells4x2<Cells>>> {
143  public:
144   typedef WidthMajorSideFormatNCells4x2<Cells> KernelSideFormat;
145   typedef typename KernelSideFormat::Cell CellFormat;
146   static const int kCells = KernelSideFormat::kCells;
147   static const int kCellWidth = CellFormat::kWidth;
148   static const int kKernelWidth = CellFormat::kWidth * kCells;
149   static const int kCellDepth = CellFormat::kDepth;
150   static const int kCellSize = CellFormat::kSize;
151 
Pack(PackedSideBlock<KernelSideFormat> * dst,int start_width)152   void Pack(PackedSideBlock<KernelSideFormat>* dst, int start_width) {
153     std::uint8_t* dst_ptr = dst->current_data();
154     const std::uint8_t* src_ptr = this->complete_src_.data();
155     const int stride = this->complete_src_.stride();
156     // Load source WidthMajor data
157     uint16x8_t src_lines[kCells * 4];
158     for (int i = 0; i < kCells; i++) {
159       // This packing path is used with our current
160       // less-than-8-bit kernel, and the partial unrolling of this loop
161       // results in substantially faster code (thanks to better
162       // register allocation) on Nexus 5.
163 
164 #define GEMMLOWP_UNROLLED_LOOP_ITER(k)                            \
165   src_lines[4 * i + k] = vreinterpretq_u16_u8(vld1q_u8(src_ptr)); \
166   src_ptr += stride;
167 
168       GEMMLOWP_UNROLLED_LOOP_ITER(0)
169       GEMMLOWP_UNROLLED_LOOP_ITER(1)
170       GEMMLOWP_UNROLLED_LOOP_ITER(2)
171       GEMMLOWP_UNROLLED_LOOP_ITER(3)
172 
173 #undef GEMMLOWP_UNROLLED_LOOP_ITER
174     }
175     // Reorder the data within registers to make WidthMajor 4x2 cells
176     uint16x8x2_t src_lines_intertwined_2x[2 * kCells];
177     for (int i = 0; i < kCells; i++) {
178       src_lines_intertwined_2x[2 * i] =
179           vzipq_u16(src_lines[4 * i], src_lines[4 * i + 2]);
180       src_lines_intertwined_2x[2 * i + 1] =
181           vzipq_u16(src_lines[4 * i + 1], src_lines[4 * i + 3]);
182     }
183     uint16x8x2_t src_lines_intertwined_4x[2 * kCells];
184     for (int i = 0; i < kCells; i++) {
185       src_lines_intertwined_4x[2 * i] =
186           vzipq_u16(src_lines_intertwined_2x[2 * i].val[0],
187                     src_lines_intertwined_2x[2 * i + 1].val[0]);
188       src_lines_intertwined_4x[2 * i + 1] =
189           vzipq_u16(src_lines_intertwined_2x[2 * i].val[1],
190                     src_lines_intertwined_2x[2 * i + 1].val[1]);
191     }
192     // Store the resulting WidthMajor 4x2 cells in the destination packed block
193     for (int outer = 0; outer < 2; outer++) {
194       for (int inner = 0; inner < 2; inner++) {
195         for (int cell = 0; cell < kCells; cell++) {
196           uint8x8_t value = vreinterpret_u8_u16(vget_low_u16(
197               src_lines_intertwined_4x[2 * cell + outer].val[inner]));
198           vst1_u8(dst_ptr, value);
199           dst_ptr += 8;
200         }
201         for (int cell = 0; cell < kCells; cell++) {
202           uint8x8_t value = vreinterpret_u8_u16(vget_high_u16(
203               src_lines_intertwined_4x[2 * cell + outer].val[inner]));
204           vst1_u8(dst_ptr, value);
205           dst_ptr += 8;
206         }
207       }
208     }
209     // Compute sums across the depth dimension
210     uint16x8_t sums_of_2[kCells][4];
211     for (int outer = 0; outer < 2; outer++) {
212       for (int inner = 0; inner < 2; inner++) {
213         int i = 2 * outer + inner;
214         for (int cell = 0; cell < kCells; cell++) {
215           sums_of_2[cell][i] = vpaddlq_u8(vreinterpretq_u8_u16(
216               src_lines_intertwined_4x[2 * cell + outer].val[inner]));
217         }
218       }
219     }
220     uint16x8_t sums_of_4[kCells][2];
221     for (int i = 0; i < 2; i++) {
222       for (int cell = 0; cell < kCells; cell++) {
223         sums_of_4[cell][i] =
224             vaddq_u16(sums_of_2[cell][2 * i], sums_of_2[cell][2 * i + 1]);
225       }
226     }
227     uint16x8_t sums_of_8[kCells];
228     for (int cell = 0; cell < kCells; cell++) {
229       sums_of_8[cell] = vaddq_u16(sums_of_4[cell][0], sums_of_4[cell][1]);
230     }
231 
232     uint16x4_t sums_of_16[kCells];
233     for (int cell = 0; cell < kCells; cell++) {
234       sums_of_16[cell] = vadd_u16(vget_low_u16(sums_of_8[cell]),
235                                   vget_high_u16(sums_of_8[cell]));
236     }
237     // Update the sums_of_each_slice vector
238     for (int cell = 0; cell < kCells; cell++) {
239       int32x4_t s = vreinterpretq_s32_u32(vmovl_u16(sums_of_16[cell]));
240       std::int32_t* sums_of_each_slice_ptr =
241           dst->sums_of_each_slice() + start_width + 4 * cell;
242       vst1q_s32(sums_of_each_slice_ptr,
243                 vaddq_s32(s, vld1q_s32(sums_of_each_slice_ptr)));
244     }
245     dst->seek_forward_n_cells(kCells * kRegisterSize / kCellDepth);
246   }
247 };
248 
249 #ifdef GEMMLOWP_NEON_32
vpaddq_s16(int16x8_t a,int16x8_t b)250 inline int16x8_t vpaddq_s16(int16x8_t a, int16x8_t b) {
251   const int16x4_t c = vpadd_s16(vget_low_s16(a), vget_high_s16(a));
252   const int16x4_t d = vpadd_s16(vget_low_s16(b), vget_high_s16(b));
253   return vcombine_s16(c, d);
254 }
255 #endif
256 
257 template <int Width>
258 using Int8FastKernelFormat =
259     KernelSideFormatInt8<CellFormat<Width, 16, CellOrder::WidthMajor>, 1>;
260 
261 template <int Width>
262 class PackingRegisterBlock<WidthMajorUint8SideMap,
263                            PackedSideBlock<Int8FastKernelFormat<Width>>>
264     : public PackingRegisterBlockBase<
265           WidthMajorUint8SideMap,
266           PackedSideBlock<Int8FastKernelFormat<Width>>> {
267  public:
268   static_assert(Width == 2 || Width == 4, "");
269   typedef Int8FastKernelFormat<Width> KernelSideFormat;
270   typedef typename KernelSideFormat::Cell CellFormat;
271   static const int kCells = KernelSideFormat::kCells;
272   static const int kCellWidth = CellFormat::kWidth;
273   static const int kKernelWidth = CellFormat::kWidth * kCells;
274   static const int kCellDepth = CellFormat::kDepth;
275   static const int kCellSize = CellFormat::kSize;
276 
Pack(PackedSideBlock<KernelSideFormat> * dst,int start_width)277   void Pack(PackedSideBlock<KernelSideFormat>* dst, int start_width) {
278     std::int32_t* sums_ptr = dst->sums_of_each_slice() + start_width;
279     std::uint8_t* dst_ptr = dst->current_data();
280     const std::uint8_t* const src_ptr = this->complete_src_.data();
281     const int stride = this->complete_src_.stride();
282     // Load source WidthMajor data
283     uint8x16_t src_lines[Width];
284     for (int i = 0; i < Width; i++) {
285       src_lines[i] = vld1q_u8(src_ptr + i * stride);
286     }
287     const uint8x16_t sign_bit_dup = vdupq_n_u8(0x80);
288     for (int i = 0; i < Width; i++) {
289       src_lines[i] = veorq_u8(src_lines[i], sign_bit_dup);
290     }
291     for (int i = 0; i < Width; i++) {
292       vst1q_u8(dst_ptr + 16 * i, src_lines[i]);
293     }
294     int16x8_t sums2[Width];
295     for (int i = 0; i < Width; i++) {
296       const int8x8_t lo = vreinterpret_s8_u8(vget_low_u8(src_lines[i]));
297       const int8x8_t hi = vreinterpret_s8_u8(vget_high_u8(src_lines[i]));
298       sums2[i] = vaddl_s8(lo, hi);
299     }
300     int16x8_t sums4[Width / 2];
301     for (int i = 0; i < Width / 2; i++) {
302       sums4[i] = vpaddq_s16(sums2[2 * i], sums2[2 * i + 1]);
303     }
304     if (Width == 4) {
305       int32x4_t sum = vld1q_s32(sums_ptr);
306       int16x8_t sums8 = vpaddq_s16(sums4[0], sums4[1]);
307       sum = vpadalq_s16(sum, sums8);
308       vst1q_s32(sums_ptr, sum);
309     } else {
310       assert(Width == 2);
311       int32x2_t sum = vld1_s32(sums_ptr);
312       int16x4_t sums8 =
313           vpadd_s16(vget_low_s16(sums4[0]), vget_high_s16(sums4[0]));
314       sum = vpadal_s16(sum, sums8);
315       vst1_s32(sums_ptr, sum);
316     }
317     dst->seek_forward_n_cells(1);
318   }
319 };
320 
321 template <int Width>
322 using Int8InputsFastKernelFormat =
323     KernelSideFormatInt8Inputs<CellFormat<Width, 16, CellOrder::WidthMajor>, 1>;
324 
325 // Same as above, but for int8 inputs, avoiding the uint8 -> int8 conversion.
326 template <int Width>
327 class PackingRegisterBlock<WidthMajorInt8SideMap,
328                            PackedSideBlock<Int8InputsFastKernelFormat<Width>>>
329     : public PackingRegisterBlockBase<
330           WidthMajorInt8SideMap,
331           PackedSideBlock<Int8InputsFastKernelFormat<Width>>> {
332  public:
333   static_assert(Width == 2 || Width == 4, "");
334   typedef Int8InputsFastKernelFormat<Width> KernelSideFormat;
335   typedef typename KernelSideFormat::Cell CellFormat;
336   static const int kCells = KernelSideFormat::kCells;
337   static const int kCellWidth = CellFormat::kWidth;
338   static const int kKernelWidth = CellFormat::kWidth * kCells;
339   static const int kCellDepth = CellFormat::kDepth;
340   static const int kCellSize = CellFormat::kSize;
341 
Pack(PackedSideBlock<KernelSideFormat> * dst,int start_width)342   void Pack(PackedSideBlock<KernelSideFormat>* dst, int start_width) {
343     std::int32_t* sums_ptr = dst->sums_of_each_slice() + start_width;
344     std::int8_t* dst_ptr = reinterpret_cast<std::int8_t*>(dst->current_data());
345     const std::int8_t* const src_ptr = this->complete_src_.data();
346     const int stride = this->complete_src_.stride();
347     // Load source WidthMajor data
348     int8x16_t src_lines[Width];
349     for (int i = 0; i < Width; i++) {
350       src_lines[i] = vld1q_s8(src_ptr + i * stride);
351     }
352     for (int i = 0; i < Width; i++) {
353       vst1q_s8(dst_ptr + 16 * i, src_lines[i]);
354     }
355     int16x8_t sums2[Width];
356     for (int i = 0; i < Width; i++) {
357       const int8x8_t lo = vget_low_s8(src_lines[i]);
358       const int8x8_t hi = vget_high_s8(src_lines[i]);
359       sums2[i] = vaddl_s8(lo, hi);
360     }
361     int16x8_t sums4[Width / 2];
362     for (int i = 0; i < Width / 2; i++) {
363       sums4[i] = vpaddq_s16(sums2[2 * i], sums2[2 * i + 1]);
364     }
365     if (Width == 4) {
366       int32x4_t sum = vld1q_s32(sums_ptr);
367       int16x8_t sums8 = vpaddq_s16(sums4[0], sums4[1]);
368       sum = vpadalq_s16(sum, sums8);
369       vst1q_s32(sums_ptr, sum);
370     } else {
371       assert(Width == 2);
372       int32x2_t sum = vld1_s32(sums_ptr);
373       int16x4_t sums8 =
374           vpadd_s16(vget_low_s16(sums4[0]), vget_high_s16(sums4[0]));
375       sum = vpadal_s16(sum, sums8);
376       vst1_s32(sums_ptr, sum);
377     }
378     dst->seek_forward_n_cells(1);
379   }
380 };
381 
382 }  // namespace gemmlowp
383 
384 #endif  // GEMMLOWP_INTERNAL_PACK_NEON_H_
385