Skip to content

Commit 56496fd

Browse files
dsharletgxnnpack-bot
dsharletg
authored and
xnnpack-bot
committed
Add int32 sum kernels
These will be necessary to support "factoring" larger reductions into multiple stages. Benchmarks: ``` ------------------------------------------------------------------------------------------- Benchmark Time CPU Iterations UserCounters... ------------------------------------------------------------------------------------------- bench/sum_int32_avx512/real_time 3109 ns 3108 ns 47834 OP=21.0808G/s 256x1x1x256 bench/sum_int32_avx2/real_time 4023 ns 4021 ns 33691 OP=16.2923G/s 256x1x1x256 bench/sum_int32_sse2/real_time 6368 ns 6368 ns 21828 OP=10.2912G/s 256x1x1x256 bench/sum_int32/real_time 13319 ns 13316 ns 10717 OP=4.92059G/s 256x1x1x256 ``` PiperOrigin-RevId: 907644825
1 parent 6408104 commit 56496fd

5 files changed

Lines changed: 69 additions & 0 deletions

File tree

‎ynnpack/kernels/reduce/arm_neon.cc‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ using simd::bf16x8;
5050
using simd::f16x8;
5151
using simd::f32x4;
5252
using simd::f32x8;
53+
using simd::s32x4;
5354
using simd::s16x8;
5455
using simd::s8x16;
5556
using simd::u8x16;
@@ -124,6 +125,22 @@ void sum_fp32_neon(size_t n, size_t k3, size_t k2, size_t k1,
124125
}
125126
}
126127

128+
void sum_int32_neon(size_t n, size_t k3, size_t k2, size_t k1,
129+
size_t a_stride_n, size_t a_stride_k3, size_t a_stride_k2,
130+
const void* a, size_t, void* c) {
131+
if (k1 == 1 && a_stride_n == sizeof(int32_t)) {
132+
stream_reduce<sum_accumulator_k1_1<s32x4>, int32_t, int32_t>(
133+
n, k3, k2, a_stride_k3, a_stride_k2,
134+
reinterpret_cast<const int32_t*>(a),
135+
/*C_stride_m=*/0, reinterpret_cast<int32_t*>(c));
136+
} else {
137+
tiled_reduce<sum_accumulator_x32<s32x4, 4>, int32_t, int32_t>(
138+
n, k3, k2, k1, a_stride_n, a_stride_k3, a_stride_k2,
139+
reinterpret_cast<const int32_t*>(a), /*C_stride_m=*/0,
140+
reinterpret_cast<int32_t*>(c));
141+
}
142+
}
143+
127144
void sum_squared_fp32_neon(size_t n, size_t k3, size_t k2, size_t k1,
128145
size_t a_stride_n, size_t a_stride_k3,
129146
size_t a_stride_k2, const void* a, size_t, void* c) {

‎ynnpack/kernels/reduce/sum.inc‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ YNN_UNARY_REDUCE_KERNEL(arch_flag::neondot, sum_uint8_int32_neondot, uint8_t, in
1212
#endif // YNN_ARCH_ARM_NEONDOT
1313
#ifdef YNN_ARCH_ARM_NEON
1414
YNN_UNARY_REDUCE_KERNEL(arch_flag::neon, sum_fp32_neon, float, float)
15+
YNN_UNARY_REDUCE_KERNEL(arch_flag::neon, sum_int32_neon, int32_t, int32_t)
1516
YNN_UNARY_REDUCE_KERNEL(arch_flag::neon, sum_bf16_fp32_neon, bfloat16, float)
1617
#ifdef YNN_ARCH_ARM64
1718
YNN_UNARY_REDUCE_KERNEL(arch_flag::neon, sum_fp64_neon, double, double)
@@ -35,13 +36,15 @@ YNN_UNARY_REDUCE_KERNEL(arch_flag::avx512, sum_bf16_fp32_avx512, bfloat16, float
3536
YNN_UNARY_REDUCE_KERNEL(arch_flag::avx512, sum_fp16_fp32_avx512, half, float)
3637
YNN_UNARY_REDUCE_KERNEL(arch_flag::avx512, sum_uint8_int32_avx512, uint8_t, int32_t)
3738
YNN_UNARY_REDUCE_KERNEL(arch_flag::avx512, sum_int8_int32_avx512, int8_t, int32_t)
39+
YNN_UNARY_REDUCE_KERNEL(arch_flag::avx512, sum_int32_avx512, int32_t, int32_t)
3840
YNN_UNARY_REDUCE_KERNEL(arch_flag::avx512, sum_fp32_avx512, float, float)
3941
YNN_UNARY_REDUCE_KERNEL(arch_flag::avx512, sum_fp64_avx512, double, double)
4042
#endif // YNN_ARCH_X86_AVX512
4143
#ifdef YNN_ARCH_X86_AVX2
4244
YNN_UNARY_REDUCE_KERNEL(arch_flag::avx2, sum_bf16_fp32_avx2, bfloat16, float)
4345
YNN_UNARY_REDUCE_KERNEL(arch_flag::avx2, sum_uint8_int32_avx2, uint8_t, int32_t)
4446
YNN_UNARY_REDUCE_KERNEL(arch_flag::avx2, sum_int8_int32_avx2, int8_t, int32_t)
47+
YNN_UNARY_REDUCE_KERNEL(arch_flag::avx2, sum_int32_avx2, int32_t, int32_t)
4548
#endif // YNN_ARCH_X86_AVX2
4649
#ifdef YNN_ARCH_X86_AVX
4750
YNN_UNARY_REDUCE_KERNEL(arch_flag::avx, sum_fp32_avx, float, float)
@@ -59,6 +62,7 @@ YNN_UNARY_REDUCE_KERNEL(arch_flag::sse2, sum_fp32_sse2, float, float)
5962
YNN_UNARY_REDUCE_KERNEL(arch_flag::sse2, sum_bf16_fp32_sse2, bfloat16, float)
6063
YNN_UNARY_REDUCE_KERNEL(arch_flag::sse2, sum_uint8_int32_sse2, uint8_t, int32_t)
6164
YNN_UNARY_REDUCE_KERNEL(arch_flag::sse2, sum_int8_int32_sse2, int8_t, int32_t)
65+
YNN_UNARY_REDUCE_KERNEL(arch_flag::sse2, sum_int32_sse2, int32_t, int32_t)
6266
#endif // YNN_ARCH_X86_SSE2
6367
#ifdef YNN_ARCH_X86_SSSE3
6468
YNN_UNARY_REDUCE_KERNEL(arch_flag::ssse3, sum_uint8_int32_ssse3, uint8_t, int32_t)

‎ynnpack/kernels/reduce/x86_avx2.cc‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,22 @@ void sum_squared_uint8_int32_avx2(size_t n, size_t k3, size_t k2, size_t k1,
197197
}
198198
}
199199

200+
void sum_int32_avx2(size_t n, size_t k3, size_t k2, size_t k1,
201+
size_t a_stride_n, size_t a_stride_k3, size_t a_stride_k2,
202+
const void* a, size_t, void* c) {
203+
if (k1 == 1 && a_stride_n == sizeof(int32_t)) {
204+
stream_reduce<sum_accumulator_k1_1<s32x8>, int32_t, int32_t>(
205+
n, k3, k2, a_stride_k3, a_stride_k2,
206+
reinterpret_cast<const int32_t*>(a),
207+
/*C_stride_m=*/0, reinterpret_cast<int32_t*>(c));
208+
} else {
209+
tiled_reduce<sum_accumulator_x32<s32x8, 8>, int32_t, int32_t>(
210+
n, k3, k2, k1, a_stride_n, a_stride_k3, a_stride_k2,
211+
reinterpret_cast<const int32_t*>(a), /*C_stride_m=*/0,
212+
reinterpret_cast<int32_t*>(c));
213+
}
214+
}
215+
200216
void sum_bf16_fp32_avx2(size_t n, size_t k3, size_t k2, size_t k1,
201217
size_t a_stride_n, size_t a_stride_k3,
202218
size_t a_stride_k2, const void* a, size_t, void* c) {

‎ynnpack/kernels/reduce/x86_avx512.cc‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,22 @@ void sum_squared_fp64_avx512(size_t n, size_t k3, size_t k2, size_t k1,
203203
}
204204
}
205205

206+
void sum_int32_avx512(size_t n, size_t k3, size_t k2, size_t k1,
207+
size_t a_stride_n, size_t a_stride_k3, size_t a_stride_k2,
208+
const void* a, size_t, void* c) {
209+
if (k1 == 1 && a_stride_n == sizeof(int32_t)) {
210+
stream_reduce<sum_accumulator_k1_1<s32x16>, int32_t, int32_t>(
211+
n, k3, k2, a_stride_k3, a_stride_k2,
212+
reinterpret_cast<const int32_t*>(a),
213+
/*C_stride_m=*/0, reinterpret_cast<int32_t*>(c));
214+
} else {
215+
tiled_reduce<sum_accumulator_x32<s32x16, 16>, int32_t, int32_t>(
216+
n, k3, k2, k1, a_stride_n, a_stride_k3, a_stride_k2,
217+
reinterpret_cast<const int32_t*>(a), /*C_stride_m=*/0,
218+
reinterpret_cast<int32_t*>(c));
219+
}
220+
}
221+
206222
void sum_fp32_avx512(size_t n, size_t k3, size_t k2, size_t k1,
207223
size_t a_stride_n, size_t a_stride_k3,
208224
size_t a_stride_k2, const void* a, size_t, void* c) {

‎ynnpack/kernels/reduce/x86_sse2.cc‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -237,6 +237,22 @@ void sum_fp32_sse2(size_t n, size_t k3, size_t k2, size_t k1,
237237
}
238238
}
239239

240+
void sum_int32_sse2(size_t n, size_t k3, size_t k2, size_t k1,
241+
size_t a_stride_n, size_t a_stride_k3, size_t a_stride_k2,
242+
const void* a, size_t, void* c) {
243+
if (k1 == 1 && a_stride_n == sizeof(int32_t)) {
244+
stream_reduce<sum_accumulator_k1_1<s32x4>, int32_t, int32_t>(
245+
n, k3, k2, a_stride_k3, a_stride_k2,
246+
reinterpret_cast<const int32_t*>(a),
247+
/*C_stride_m=*/0, reinterpret_cast<int32_t*>(c));
248+
} else {
249+
tiled_reduce<sum_accumulator_x32<s32x4, 4>, int32_t, int32_t>(
250+
n, k3, k2, k1, a_stride_n, a_stride_k3, a_stride_k2,
251+
reinterpret_cast<const int32_t*>(a), /*C_stride_m=*/0,
252+
reinterpret_cast<int32_t*>(c));
253+
}
254+
}
255+
240256
void sum_squared_fp32_sse2(size_t n, size_t k3, size_t k2, size_t k1,
241257
size_t a_stride_n, size_t a_stride_k3,
242258
size_t a_stride_k2, const void* a, size_t, void* c) {

0 commit comments

Comments
 (0)