ncnn

1
// Tencent is pleased to support the open source community by making ncnn available.
2
//
3
// Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.
4
//
5
// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
6
// in compliance with the License. You may obtain a copy of the License at
7
//
8
// https://opensource.org/licenses/BSD-3-Clause
9
//
10
// Unless required by applicable law or agreed to in writing, software distributed
11
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
12
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
13
// specific language governing permissions and limitations under the License.
14

15
#version 450
16

17
#if NCNN_fp16_storage
18
#extension GL_EXT_shader_16bit_storage: require
19
struct sfpvec8 { f16vec4 abcd; f16vec4 efgh; };
20
#endif
21
#if NCNN_fp16_arithmetic
22
#extension GL_EXT_shader_explicit_arithmetic_types_float16: require
23
#endif
24

25
layout (constant_id = 0) const int bugihfa = 0;
26

27
#define shape_constant_id_offset 1
28
layout (constant_id = shape_constant_id_offset + 0) const int dims = 0;
29
layout (constant_id = shape_constant_id_offset + 1) const int w = 0;
30
layout (constant_id = shape_constant_id_offset + 2) const int h = 0;
31
layout (constant_id = shape_constant_id_offset + 3) const int d = 0;
32
layout (constant_id = shape_constant_id_offset + 4) const int c = 0;
33
layout (constant_id = shape_constant_id_offset + 5) const int cstep = 0;
34

35
layout (constant_id = shape_constant_id_offset + 6) const int outdims = 0;
36
layout (constant_id = shape_constant_id_offset + 7) const int outw = 0;
37
layout (constant_id = shape_constant_id_offset + 8) const int outh = 0;
38
layout (constant_id = shape_constant_id_offset + 9) const int outd = 0;
39
layout (constant_id = shape_constant_id_offset + 10) const int outc = 0;
40
layout (constant_id = shape_constant_id_offset + 11) const int outcstep = 0;
41

42
#if NCNN_image_shader
43
layout (binding = 0) uniform unfp sampler3D bottom_blob;
44
layout (binding = 1, imfmtc4) writeonly uniform unfp image3D top_blob;
45
#else
46
#if NCNN_fp16_packed
47
layout (binding = 0) readonly buffer bottom_blob { sfpvec2 bottom_blob_data[]; };
48
#else
49
layout (binding = 0) readonly buffer bottom_blob { sfp bottom_blob_data[]; };
50
#endif
51
layout (binding = 1) writeonly buffer top_blob { sfpvec4 top_blob_data[]; };
52
#endif
53

54
layout (push_constant) uniform parameter
55
{
56
    int dims;
57
    int w;
58
    int h;
59
    int d;
60
    int c;
61
    int cstep;
62

63
    int outdims;
64
    int outw;
65
    int outh;
66
    int outd;
67
    int outc;
68
    int outcstep;
69

70
    int woffset;
71
    int hoffset;
72
    int doffset;
73
    int coffset;
74
} p;
75

76
void main()
77
{
78
    int gx = int(gl_GlobalInvocationID.x);
79
    int gy = int(gl_GlobalInvocationID.y);
80
    int gz = int(gl_GlobalInvocationID.z);
81

82
    if (gx >= psc(outw) || gy >= psc(outh) * psc(outd) || gz >= psc(outc))
83
        return;
84

85
    if (psc(dims) == 1)
86
    {
87
        ivec4 x4 = gx * 4 + p.woffset + ivec4(0, 1, 2, 3);
88

89
#if NCNN_image_shader
90
        afpvec8 v0 = image3d_ld8(bottom_blob, ivec3(x4.r / 8, 0, 0));
91
        afpvec8 v1 = image3d_ld8(bottom_blob, ivec3(x4.g / 8, 0, 0));
92
        afpvec8 v2 = image3d_ld8(bottom_blob, ivec3(x4.b / 8, 0, 0));
93
        afpvec8 v3 = image3d_ld8(bottom_blob, ivec3(x4.a / 8, 0, 0));
94

95
        afpvec4 v;
96
#if NCNN_fp16_arithmetic
97
        if (bugihfa == 1)
98
        {
99
            ivec4 x4lane2 = (x4 % 8) / 4;
100
            ivec4 x4m4 = x4 % 4;
101

102
            if (x4m4.r == 0) v.r = v0[x4lane2.r].r;
103
            if (x4m4.r == 1) v.r = v0[x4lane2.r].g;
104
            if (x4m4.r == 2) v.r = v0[x4lane2.r].b;
105
            if (x4m4.r == 3) v.r = v0[x4lane2.r].a;
106
            if (x4m4.g == 0) v.g = v0[x4lane2.g].r;
107
            if (x4m4.g == 1) v.g = v0[x4lane2.g].g;
108
            if (x4m4.g == 2) v.g = v0[x4lane2.g].b;
109
            if (x4m4.g == 3) v.g = v0[x4lane2.g].a;
110
            if (x4m4.b == 0) v.b = v0[x4lane2.b].r;
111
            if (x4m4.b == 1) v.b = v0[x4lane2.b].g;
112
            if (x4m4.b == 2) v.b = v0[x4lane2.b].b;
113
            if (x4m4.b == 3) v.b = v0[x4lane2.b].a;
114
            if (x4m4.a == 0) v.a = v0[x4lane2.a].r;
115
            if (x4m4.a == 1) v.a = v0[x4lane2.a].g;
116
            if (x4m4.a == 2) v.a = v0[x4lane2.a].b;
117
            if (x4m4.a == 3) v.a = v0[x4lane2.a].a;
118
        }
119
        else
120
#endif
121
        {
122
            v.r = v0[(x4.r % 8) / 4][x4.r % 4];
123
            v.g = v1[(x4.g % 8) / 4][x4.g % 4];
124
            v.b = v2[(x4.b % 8) / 4][x4.b % 4];
125
            v.a = v3[(x4.a % 8) / 4][x4.a % 4];
126
        }
127

128
        image3d_st4(top_blob, ivec3(gx, 0, 0), v);
129
#else
130
#if NCNN_fp16_packed
131
        ivec4 v_offset = (x4 / 8) * 4 + (x4 % 8) / 2;
132
        ivec4 lane2 = x4 % 2;
133

134
        afpvec2 vr = buffer_ld2(bottom_blob_data, v_offset.r);
135
        afpvec2 vg = buffer_ld2(bottom_blob_data, v_offset.g);
136
        afpvec2 vb = buffer_ld2(bottom_blob_data, v_offset.b);
137
        afpvec2 va = buffer_ld2(bottom_blob_data, v_offset.a);
138

139
        afpvec4 v = afpvec4(vr[lane2.r], vg[lane2.g], vb[lane2.b], va[lane2.a]);
140

141
        buffer_st4(top_blob_data, gx, v);
142
#else
143
        ivec4 v_offset = (x4 / 8) * 8 + x4 % 8;
144

145
        buffer_cp1to4(top_blob_data, gx, bottom_blob_data, v_offset);
146
#endif
147
#endif
148
    }
149
    else if (psc(dims) == 2)
150
    {
151
        int x = gx + p.woffset;
152
        ivec4 y4 = gy * 4 + p.hoffset + ivec4(0, 1, 2, 3);
153

154
#if NCNN_image_shader
155
        afpvec8 v0 = image3d_ld8(bottom_blob, ivec3(x, y4.r / 8, 0));
156
        afpvec8 v1 = image3d_ld8(bottom_blob, ivec3(x, y4.g / 8, 0));
157
        afpvec8 v2 = image3d_ld8(bottom_blob, ivec3(x, y4.b / 8, 0));
158
        afpvec8 v3 = image3d_ld8(bottom_blob, ivec3(x, y4.a / 8, 0));
159

160
        afpvec4 v;
161
#if NCNN_fp16_arithmetic
162
        if (bugihfa == 1)
163
        {
164
            ivec4 y4lane2 = (y4 % 8) / 4;
165
            ivec4 y4m4 = y4 % 4;
166

167
            if (y4m4.r == 0) v.r = v0[y4lane2.r].r;
168
            if (y4m4.r == 1) v.r = v0[y4lane2.r].g;
169
            if (y4m4.r == 2) v.r = v0[y4lane2.r].b;
170
            if (y4m4.r == 3) v.r = v0[y4lane2.r].a;
171
            if (y4m4.g == 0) v.g = v0[y4lane2.g].r;
172
            if (y4m4.g == 1) v.g = v0[y4lane2.g].g;
173
            if (y4m4.g == 2) v.g = v0[y4lane2.g].b;
174
            if (y4m4.g == 3) v.g = v0[y4lane2.g].a;
175
            if (y4m4.b == 0) v.b = v0[y4lane2.b].r;
176
            if (y4m4.b == 1) v.b = v0[y4lane2.b].g;
177
            if (y4m4.b == 2) v.b = v0[y4lane2.b].b;
178
            if (y4m4.b == 3) v.b = v0[y4lane2.b].a;
179
            if (y4m4.a == 0) v.a = v0[y4lane2.a].r;
180
            if (y4m4.a == 1) v.a = v0[y4lane2.a].g;
181
            if (y4m4.a == 2) v.a = v0[y4lane2.a].b;
182
            if (y4m4.a == 3) v.a = v0[y4lane2.a].a;
183
        }
184
        else
185
#endif
186
        {
187
            v.r = v0[(y4.r % 8) / 4][y4.r % 4];
188
            v.g = v1[(y4.g % 8) / 4][y4.g % 4];
189
            v.b = v2[(y4.b % 8) / 4][y4.b % 4];
190
            v.a = v3[(y4.a % 8) / 4][y4.a % 4];
191
        }
192

193
        image3d_st4(top_blob, ivec3(gx, gy, 0), v);
194
#else
195
        int gi = gy * psc(outw) + gx;
196

197
#if NCNN_fp16_packed
198
        ivec4 v_offset = ((y4 / 8) * psc(w) + x) * 4 + (y4 % 8) / 2;
199
        ivec4 lane2 = y4 % 2;
200

201
        afpvec2 vr = buffer_ld2(bottom_blob_data, v_offset.r);
202
        afpvec2 vg = buffer_ld2(bottom_blob_data, v_offset.g);
203
        afpvec2 vb = buffer_ld2(bottom_blob_data, v_offset.b);
204
        afpvec2 va = buffer_ld2(bottom_blob_data, v_offset.a);
205

206
        afpvec4 v = afpvec4(vr[lane2.r], vg[lane2.g], vb[lane2.b], va[lane2.a]);
207

208
        buffer_st4(top_blob_data, gi, v);
209
#else
210
        ivec4 v_offset = ((y4 / 8) * psc(w) + x) * 8 + y4 % 8;
211

212
        buffer_cp1to4(top_blob_data, gi, bottom_blob_data, v_offset);
213
#endif
214
#endif
215
    }
216
    else if (psc(dims) == 3)
217
    {
218
        int x = gx + p.woffset;
219
        int y = gy + p.hoffset;
220
        ivec4 z4 = gz * 4 + p.coffset + ivec4(0, 1, 2, 3);
221

222
#if NCNN_image_shader
223
        afpvec8 v0 = image3d_ld8(bottom_blob, ivec3(x, y, z4.r / 8));
224
        afpvec8 v1 = image3d_ld8(bottom_blob, ivec3(x, y, z4.g / 8));
225
        afpvec8 v2 = image3d_ld8(bottom_blob, ivec3(x, y, z4.b / 8));
226
        afpvec8 v3 = image3d_ld8(bottom_blob, ivec3(x, y, z4.a / 8));
227

228
        afpvec4 v;
229
#if NCNN_fp16_arithmetic
230
        if (bugihfa == 1)
231
        {
232
            ivec4 z4lane2 = (z4 % 8) / 4;
233
            ivec4 z4m4 = z4 % 4;
234

235
            if (z4m4.r == 0) v.r = v0[z4lane2.r].r;
236
            if (z4m4.r == 1) v.r = v0[z4lane2.r].g;
237
            if (z4m4.r == 2) v.r = v0[z4lane2.r].b;
238
            if (z4m4.r == 3) v.r = v0[z4lane2.r].a;
239
            if (z4m4.g == 0) v.g = v0[z4lane2.g].r;
240
            if (z4m4.g == 1) v.g = v0[z4lane2.g].g;
241
            if (z4m4.g == 2) v.g = v0[z4lane2.g].b;
242
            if (z4m4.g == 3) v.g = v0[z4lane2.g].a;
243
            if (z4m4.b == 0) v.b = v0[z4lane2.b].r;
244
            if (z4m4.b == 1) v.b = v0[z4lane2.b].g;
245
            if (z4m4.b == 2) v.b = v0[z4lane2.b].b;
246
            if (z4m4.b == 3) v.b = v0[z4lane2.b].a;
247
            if (z4m4.a == 0) v.a = v0[z4lane2.a].r;
248
            if (z4m4.a == 1) v.a = v0[z4lane2.a].g;
249
            if (z4m4.a == 2) v.a = v0[z4lane2.a].b;
250
            if (z4m4.a == 3) v.a = v0[z4lane2.a].a;
251
        }
252
        else
253
#endif
254
        {
255
            v.r = v0[(z4.r % 8) / 4][z4.r % 4];
256
            v.g = v1[(z4.g % 8) / 4][z4.g % 4];
257
            v.b = v2[(z4.b % 8) / 4][z4.b % 4];
258
            v.a = v3[(z4.a % 8) / 4][z4.a % 4];
259
        }
260

261
        image3d_st4(top_blob, ivec3(gx, gy, gz), v);
262
#else
263
        int gi = gz * psc(outcstep) + gy * psc(outw) + gx;
264

265
#if NCNN_fp16_packed
266
        ivec4 v_offset = ((z4 / 8) * psc(cstep) + y * psc(w) + x) * 4 + (z4 % 8) / 2;
267
        ivec4 lane2 = z4 % 2;
268

269
        afpvec2 vr = buffer_ld2(bottom_blob_data, v_offset.r);
270
        afpvec2 vg = buffer_ld2(bottom_blob_data, v_offset.g);
271
        afpvec2 vb = buffer_ld2(bottom_blob_data, v_offset.b);
272
        afpvec2 va = buffer_ld2(bottom_blob_data, v_offset.a);
273

274
        afpvec4 v = afpvec4(vr[lane2.r], vg[lane2.g], vb[lane2.b], va[lane2.a]);
275

276
        buffer_st4(top_blob_data, gi, v);
277
#else
278
        ivec4 v_offset = ((z4 / 8) * psc(cstep) + y * psc(w) + x) * 8 + z4 % 8;
279

280
        buffer_cp1to4(top_blob_data, gi, bottom_blob_data, v_offset);
281
#endif
282
#endif
283
    }
284
    else // if (psc(dims) == 4)
285
    {
286
        int yd = gy / psc(outh);
287
        int yh = gy % psc(outh);
288

289
        int x = gx + p.woffset;
290
        int y = (yd + p.doffset) * psc(h) + (yh + p.hoffset);
291
        ivec4 z4 = gz * 4 + p.coffset + ivec4(0, 1, 2, 3);
292

293
#if NCNN_image_shader
294
        afpvec8 v0 = image3d_ld8(bottom_blob, ivec3(x, y, z4.r / 8));
295
        afpvec8 v1 = image3d_ld8(bottom_blob, ivec3(x, y, z4.g / 8));
296
        afpvec8 v2 = image3d_ld8(bottom_blob, ivec3(x, y, z4.b / 8));
297
        afpvec8 v3 = image3d_ld8(bottom_blob, ivec3(x, y, z4.a / 8));
298

299
        afpvec4 v;
300
#if NCNN_fp16_arithmetic
301
        if (bugihfa == 1)
302
        {
303
            ivec4 z4lane2 = (z4 % 8) / 4;
304
            ivec4 z4m4 = z4 % 4;
305

306
            if (z4m4.r == 0) v.r = v0[z4lane2.r].r;
307
            if (z4m4.r == 1) v.r = v0[z4lane2.r].g;
308
            if (z4m4.r == 2) v.r = v0[z4lane2.r].b;
309
            if (z4m4.r == 3) v.r = v0[z4lane2.r].a;
310
            if (z4m4.g == 0) v.g = v0[z4lane2.g].r;
311
            if (z4m4.g == 1) v.g = v0[z4lane2.g].g;
312
            if (z4m4.g == 2) v.g = v0[z4lane2.g].b;
313
            if (z4m4.g == 3) v.g = v0[z4lane2.g].a;
314
            if (z4m4.b == 0) v.b = v0[z4lane2.b].r;
315
            if (z4m4.b == 1) v.b = v0[z4lane2.b].g;
316
            if (z4m4.b == 2) v.b = v0[z4lane2.b].b;
317
            if (z4m4.b == 3) v.b = v0[z4lane2.b].a;
318
            if (z4m4.a == 0) v.a = v0[z4lane2.a].r;
319
            if (z4m4.a == 1) v.a = v0[z4lane2.a].g;
320
            if (z4m4.a == 2) v.a = v0[z4lane2.a].b;
321
            if (z4m4.a == 3) v.a = v0[z4lane2.a].a;
322
        }
323
        else
324
#endif
325
        {
326
            v.r = v0[(z4.r % 8) / 4][z4.r % 4];
327
            v.g = v1[(z4.g % 8) / 4][z4.g % 4];
328
            v.b = v2[(z4.b % 8) / 4][z4.b % 4];
329
            v.a = v3[(z4.a % 8) / 4][z4.a % 4];
330
        }
331

332
        image3d_st4(top_blob, ivec3(gx, gy, gz), v);
333
#else
334
        int gi = gz * psc(outcstep) + gy * psc(outw) + gx;
335

336
#if NCNN_fp16_packed
337
        ivec4 v_offset = ((z4 / 8) * psc(cstep) + y * psc(w) + x) * 4 + (z4 % 8) / 2;
338
        ivec4 lane2 = z4 % 2;
339

340
        afpvec2 vr = buffer_ld2(bottom_blob_data, v_offset.r);
341
        afpvec2 vg = buffer_ld2(bottom_blob_data, v_offset.g);
342
        afpvec2 vb = buffer_ld2(bottom_blob_data, v_offset.b);
343
        afpvec2 va = buffer_ld2(bottom_blob_data, v_offset.a);
344

345
        afpvec4 v = afpvec4(vr[lane2.r], vg[lane2.g], vb[lane2.b], va[lane2.a]);
346

347
        buffer_st4(top_blob_data, gi, v);
348
#else
349
        ivec4 v_offset = ((z4 / 8) * psc(cstep) + y * psc(w) + x) * 8 + z4 % 8;
350

351
        buffer_cp1to4(top_blob_data, gi, bottom_blob_data, v_offset);
352
#endif
353
#endif
354
    }
355
}
356