Skip to content

Commit 5e68a2f

Browse files
nihuinihui
committed
wip
1 parent ff795d8 commit 5e68a2f

137 files changed

Lines changed: 39210 additions & 1578 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎src/layer/loongarch/batchnorm_bf16s.h‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
// Tencent is pleased to support the open source community by making ncnn available.
2+
//
3+
// https://opensource.org/licenses/BSD-3-Clause
4+
//
5+
// Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
6+
//
7+
// SPDX-License-Identifier: BSD-3-Clause
8+
9+
#ifndef BATCHNORM_LOONGARCH_BF16S_H
10+
#define BATCHNORM_LOONGARCH_BF16S_H
11+
12+
static void batchnorm_bf16s_sse(unsigned short* ptr, const float* a, const float* b, int size, int elempack)
13+
{
14+
#if __loongarch_sx
15+
__m128 _a128 = (elempack == 4) ? (__m128)__lsx_vld(a, 0) : (__m128)__lsx_vreplfr2vr_s(a[0]);
16+
__m128 _b128 = (elempack == 4) ? (__m128)__lsx_vld(b, 0) : (__m128)__lsx_vreplfr2vr_s(b[0]);
17+
#if __loongarch_asx
18+
__m256 _a256 = (elempack == 8) ? (__m256)__lasx_xvld(a, 0) : combine4x2_ps(_a128, _a128);
19+
__m256 _b256 = (elempack == 8) ? (__m256)__lasx_xvld(b, 0) : combine4x2_ps(_b128, _b128);
20+
#endif
21+
#endif
22+
float sa = a[0];
23+
float sb = b[0];
24+
25+
int i = 0;
26+
#if __loongarch_sx
27+
#if __loongarch_asx
28+
for (; i + 7 < size; i += 8)
29+
{
30+
__m256 _p = bfloat2float_avx((__m128i)__lsx_vld(ptr, 0));
31+
_p = __lasx_xvfmadd_s(_p, _b256, _a256);
32+
__lsx_vst(float2bfloat_avx(_p), ptr, 0);
33+
ptr += 8;
34+
}
35+
#endif // __loongarch_asx
36+
for (; i + 3 < size; i += 4)
37+
{
38+
__m128 _p = bfloat2float_sse((__m128i)__lsx_vld(ptr, 0));
39+
_p = __lsx_vfmadd_s(_p, _b128, _a128);
40+
__lsx_vstelm_d(float2bfloat_sse(_p, _p), ptr, 0, 0);
41+
ptr += 4;
42+
}
43+
#endif // __loongarch_sx
44+
for (; i < size; i++)
45+
{
46+
*ptr = float32_to_bfloat16(sb * bfloat16_to_float32(*ptr) + sa);
47+
ptr++;
48+
}
49+
}
50+
51+
static void batchnorm_bf16s_per_element_sse(unsigned short* ptr, const float* a, const float* b, int size, int num_threads)
52+
{
53+
int nn_size = 0;
54+
int remain_size_start = 0;
55+
#if __loongarch_sx
56+
#if __loongarch_asx
57+
nn_size = (size - remain_size_start) / 8;
58+
#pragma omp parallel for num_threads(num_threads)
59+
for (int ii = 0; ii < nn_size; ii++)
60+
{
61+
int i = remain_size_start + ii * 8;
62+
__m256 _p = bfloat2float_avx((__m128i)__lsx_vld(ptr + i, 0));
63+
__m256 _a0 = (__m256)__lasx_xvld(a + i, 0);
64+
__m256 _b0 = (__m256)__lasx_xvld(b + i, 0);
65+
_p = __lasx_xvfmadd_s(_p, _b0, _a0);
66+
__lsx_vst(float2bfloat_avx(_p), ptr + i, 0);
67+
}
68+
remain_size_start += nn_size * 8;
69+
#endif // __loongarch_asx
70+
nn_size = (size - remain_size_start) / 4;
71+
#pragma omp parallel for num_threads(num_threads)
72+
for (int ii = 0; ii < nn_size; ii++)
73+
{
74+
int i = remain_size_start + ii * 4;
75+
__m128 _p = bfloat2float_sse((__m128i)__lsx_vld(ptr + i, 0));
76+
__m128 _a0 = (__m128)__lsx_vld(a + i, 0);
77+
__m128 _b0 = (__m128)__lsx_vld(b + i, 0);
78+
_p = __lsx_vfmadd_s(_p, _b0, _a0);
79+
__lsx_vstelm_d(float2bfloat_sse(_p, _p), ptr + i, 0, 0);
80+
}
81+
remain_size_start += nn_size * 4;
82+
#endif // __loongarch_sx
83+
#pragma omp parallel for num_threads(num_threads)
84+
for (int i = remain_size_start; i < size; i++)
85+
{
86+
ptr[i] = float32_to_bfloat16(b[i] * bfloat16_to_float32(ptr[i]) + a[i]);
87+
}
88+
}
89+
90+
#endif // BATCHNORM_LOONGARCH_BF16S_H

‎src/layer/loongarch/binaryop_bf16s.h‎

Lines changed: 260 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,260 @@
1+
// Tencent is pleased to support the open source community by making ncnn available.
2+
//
3+
// https://opensource.org/licenses/BSD-3-Clause
4+
//
5+
// Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
6+
//
7+
// SPDX-License-Identifier: BSD-3-Clause
8+
9+
template<typename Op>
10+
static void binary_op_vector_no_broadcast_bf16s(const unsigned short* ptr, const unsigned short* ptr1, unsigned short* outptr, int size)
11+
{
12+
const Op op;
13+
14+
int i = 0;
15+
#if __loongarch_sx
16+
#if __loongarch_asx
17+
for (; i + 7 < size; i += 8)
18+
{
19+
__m256 _p = bfloat2float_avx((__m128i)__lsx_vld(ptr, 0));
20+
__m256 _b = bfloat2float_avx((__m128i)__lsx_vld(ptr1, 0));
21+
__m256 _outp = op(_p, _b);
22+
__lsx_vst(float2bfloat_avx(_outp), outptr, 0);
23+
ptr += 8;
24+
ptr1 += 8;
25+
outptr += 8;
26+
}
27+
#endif // __loongarch_asx
28+
for (; i + 3 < size; i += 4)
29+
{
30+
__m128 _p = bfloat2float_sse((__m128i)__lsx_vld(ptr, 0));
31+
__m128 _b = bfloat2float_sse((__m128i)__lsx_vld(ptr1, 0));
32+
__m128 _outp = op(_p, _b);
33+
__lsx_vstelm_d(float2bfloat_sse(_outp, _outp), outptr, 0, 0);
34+
ptr += 4;
35+
ptr1 += 4;
36+
outptr += 4;
37+
}
38+
#endif // __loongarch_sx
39+
for (; i < size; i++)
40+
{
41+
*outptr++ = float32_to_bfloat16(op(bfloat16_to_float32(*ptr++), bfloat16_to_float32(*ptr1++)));
42+
}
43+
}
44+
45+
template<typename Op>
46+
static void binary_op_vector_broadcast_b_bf16s(const unsigned short* ptr, const unsigned short* ptr1, unsigned short* outptr, int size, int elempack)
47+
{
48+
const Op op;
49+
50+
const float b = bfloat16_to_float32(*ptr1);
51+
52+
int i = 0;
53+
#if __loongarch_sx
54+
__m128 _b_128 = (elempack == 4) ? bfloat2float_sse((__m128i)__lsx_vld(ptr1, 0)) : (__m128)__lsx_vreplfr2vr_s(b);
55+
#if __loongarch_asx
56+
__m256 _b_256 = (elempack == 8) ? bfloat2float_avx((__m128i)__lsx_vld(ptr1, 0)) : combine4x2_ps(_b_128, _b_128);
57+
for (; i + 7 < size; i += 8)
58+
{
59+
__m256 _p = bfloat2float_avx((__m128i)__lsx_vld(ptr, 0));
60+
__m256 _outp = op(_p, _b_256);
61+
__lsx_vst(float2bfloat_avx(_outp), outptr, 0);
62+
ptr += 8;
63+
outptr += 8;
64+
}
65+
#endif // __loongarch_asx
66+
for (; i + 3 < size; i += 4)
67+
{
68+
__m128 _p = bfloat2float_sse((__m128i)__lsx_vld(ptr, 0));
69+
__m128 _outp = op(_p, _b_128);
70+
__lsx_vstelm_d(float2bfloat_sse(_outp, _outp), outptr, 0, 0);
71+
ptr += 4;
72+
outptr += 4;
73+
}
74+
#endif // __loongarch_sx
75+
for (; i < size; i++)
76+
{
77+
*outptr++ = float32_to_bfloat16(op(bfloat16_to_float32(*ptr++), b));
78+
}
79+
}
80+
81+
template<typename Op>
82+
static void binary_op_vector_broadcast_a_bf16s(const unsigned short* ptr, const unsigned short* ptr1, unsigned short* outptr, int size, int elempack)
83+
{
84+
const Op op;
85+
86+
const float a = bfloat16_to_float32(*ptr);
87+
88+
int i = 0;
89+
#if __loongarch_sx
90+
__m128 _a_128 = (elempack == 4) ? bfloat2float_sse((__m128i)__lsx_vld(ptr, 0)) : (__m128)__lsx_vreplfr2vr_s(a);
91+
#if __loongarch_asx
92+
__m256 _a_256 = (elempack == 8) ? bfloat2float_avx((__m128i)__lsx_vld(ptr, 0)) : combine4x2_ps(_a_128, _a_128);
93+
for (; i + 7 < size; i += 8)
94+
{
95+
__m256 _b = bfloat2float_avx((__m128i)__lsx_vld(ptr1, 0));
96+
__m256 _outp = op(_a_256, _b);
97+
__lsx_vst(float2bfloat_avx(_outp), outptr, 0);
98+
ptr1 += 8;
99+
outptr += 8;
100+
}
101+
#endif // __loongarch_asx
102+
for (; i + 3 < size; i += 4)
103+
{
104+
__m128 _b = bfloat2float_sse((__m128i)__lsx_vld(ptr1, 0));
105+
__m128 _outp = op(_a_128, _b);
106+
__lsx_vstelm_d(float2bfloat_sse(_outp, _outp), outptr, 0, 0);
107+
ptr1 += 4;
108+
outptr += 4;
109+
}
110+
#endif // __loongarch_sx
111+
for (; i < size; i++)
112+
{
113+
*outptr++ = float32_to_bfloat16(op(a, bfloat16_to_float32(*ptr1++)));
114+
}
115+
}
116+
117+
template<typename Op>
118+
static void binary_op_vector_broadcast_pb_bf16s(const unsigned short* ptr, const unsigned short* ptr1, unsigned short* outptr, int w, int elempack)
119+
{
120+
const Op op;
121+
122+
#if __loongarch_sx
123+
#if __loongarch_asx
124+
if (elempack == 8)
125+
{
126+
for (int i = 0; i < w; i++)
127+
{
128+
__m256 _p = bfloat2float_avx((__m128i)__lsx_vld(ptr, 0));
129+
__m256 _b = __lasx_xvreplfr2vr_s(bfloat16_to_float32(*ptr1));
130+
__m256 _outp = op(_p, _b);
131+
__lsx_vst(float2bfloat_avx(_outp), outptr, 0);
132+
ptr += 8;
133+
ptr1 += 1;
134+
outptr += 8;
135+
}
136+
return;
137+
}
138+
#endif // __loongarch_asx
139+
if (elempack == 4)
140+
{
141+
for (int i = 0; i < w; i++)
142+
{
143+
__m128 _p = bfloat2float_sse((__m128i)__lsx_vld(ptr, 0));
144+
__m128 _b = (__m128)__lsx_vreplfr2vr_s(bfloat16_to_float32(*ptr1));
145+
__m128 _outp = op(_p, _b);
146+
__lsx_vstelm_d(float2bfloat_sse(_outp, _outp), outptr, 0, 0);
147+
ptr += 4;
148+
ptr1 += 1;
149+
outptr += 4;
150+
}
151+
return;
152+
}
153+
#endif // __loongarch_sx
154+
}
155+
156+
template<typename Op>
157+
static void binary_op_vector_broadcast_pb_b_bf16s(const unsigned short* ptr, const unsigned short* ptr1, unsigned short* outptr, int w, int elempack)
158+
{
159+
const Op op;
160+
161+
const int size = w * elempack;
162+
const float b = bfloat16_to_float32(*ptr1);
163+
164+
int i = 0;
165+
#if __loongarch_sx
166+
#if __loongarch_asx
167+
{
168+
__m256 _b_256 = __lasx_xvreplfr2vr_s(b);
169+
for (; i + 7 < size; i += 8)
170+
{
171+
__m256 _p = bfloat2float_avx((__m128i)__lsx_vld(ptr, 0));
172+
__m256 _outp = op(_p, _b_256);
173+
__lsx_vst(float2bfloat_avx(_outp), outptr, 0);
174+
ptr += 8;
175+
outptr += 8;
176+
}
177+
}
178+
#endif // __loongarch_asx
179+
{
180+
__m128 _b_128 = (__m128)__lsx_vreplfr2vr_s(b);
181+
for (; i + 3 < size; i += 4)
182+
{
183+
__m128 _p = bfloat2float_sse((__m128i)__lsx_vld(ptr, 0));
184+
__m128 _outp = op(_p, _b_128);
185+
__lsx_vstelm_d(float2bfloat_sse(_outp, _outp), outptr, 0, 0);
186+
ptr += 4;
187+
outptr += 4;
188+
}
189+
}
190+
#endif // __loongarch_sx
191+
for (; i < size; i++)
192+
{
193+
*outptr++ = float32_to_bfloat16(op(bfloat16_to_float32(*ptr++), b));
194+
}
195+
}
196+
197+
template<typename Op>
198+
static void binary_op_vector_broadcast_pb_a_bf16s(const unsigned short* ptr, const unsigned short* ptr1, unsigned short* outptr, int w, int elempack)
199+
{
200+
const Op op;
201+
202+
#if __loongarch_sx
203+
#if __loongarch_asx
204+
if (elempack == 8)
205+
{
206+
__m256 _p = bfloat2float_avx((__m128i)__lsx_vld(ptr, 0));
207+
for (int i = 0; i < w; i++)
208+
{
209+
__m256 _b = __lasx_xvreplfr2vr_s(bfloat16_to_float32(*ptr1));
210+
__m256 _outp = op(_p, _b);
211+
__lsx_vst(float2bfloat_avx(_outp), outptr, 0);
212+
ptr1 += 1;
213+
outptr += 8;
214+
}
215+
return;
216+
}
217+
#endif // __loongarch_asx
218+
if (elempack == 4)
219+
{
220+
__m128 _p = bfloat2float_sse((__m128i)__lsx_vld(ptr, 0));
221+
for (int i = 0; i < w; i++)
222+
{
223+
__m128 _b = (__m128)__lsx_vreplfr2vr_s(bfloat16_to_float32(*ptr1));
224+
__m128 _outp = op(_p, _b);
225+
__lsx_vstelm_d(float2bfloat_sse(_outp, _outp), outptr, 0, 0);
226+
ptr1 += 1;
227+
outptr += 4;
228+
}
229+
return;
230+
}
231+
#endif // __loongarch_sx
232+
}
233+
234+
template<typename Op>
235+
static void binary_op_vector_bf16s(const unsigned short* ptr, const unsigned short* ptr1, unsigned short* outptr, int aw, int bw, int ap, int bp)
236+
{
237+
const int w = std::max(aw, bw);
238+
const int elempack = std::max(ap, bp);
239+
const int size = w * elempack;
240+
241+
if (ap == bp)
242+
{
243+
if (aw == bw)
244+
return binary_op_vector_no_broadcast_bf16s<Op>(ptr, ptr1, outptr, size);
245+
if (bw == 1)
246+
return binary_op_vector_broadcast_b_bf16s<Op>(ptr, ptr1, outptr, size, elempack);
247+
if (aw == 1)
248+
return binary_op_vector_broadcast_a_bf16s<Op>(ptr, ptr1, outptr, size, elempack);
249+
}
250+
251+
if (bp == 1)
252+
{
253+
if (aw == bw)
254+
return binary_op_vector_broadcast_pb_bf16s<Op>(ptr, ptr1, outptr, w, elempack);
255+
if (bw == 1)
256+
return binary_op_vector_broadcast_pb_b_bf16s<Op>(ptr, ptr1, outptr, w, elempack);
257+
if (aw == 1)
258+
return binary_op_vector_broadcast_pb_a_bf16s<Op>(ptr, ptr1, outptr, w, elempack);
259+
}
260+
}

0 commit comments

Comments
 (0)