ncnn

1
// Tencent is pleased to support the open source community by making ncnn available.
2
//
3
// Copyright (C) 2023 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
#endif
20
#if NCNN_fp16_arithmetic
21
#extension GL_EXT_shader_explicit_arithmetic_types_float16: require
22
#endif
23

24
layout (constant_id = 0) const int M = 0;
25
layout (constant_id = 1) const int N = 0;
26
layout (constant_id = 2) const int K = 0;
27
layout (constant_id = 3) const int B = 0;
28

29
#if NCNN_image_shader
30
layout (binding = 0) uniform unfp sampler3D qkcross_blob_3d;
31
layout (binding = 1) uniform unfp sampler3D v_blob_3d;
32
layout (binding = 2, imfmtc1) writeonly uniform unfp image3D qkvcross_blob_3d;
33
#else
34
layout (binding = 0) readonly buffer qkcross_blob { sfpvec4 qkcross_blob_data[]; };
35
layout (binding = 1) readonly buffer v_blob { sfp v_blob_data[]; };
36
layout (binding = 2) writeonly buffer qkvcross_blob { sfp qkvcross_blob_data[]; };
37
#endif
38

39
layout (push_constant) uniform parameter
40
{
41
    int M;
42
    int N;
43
    int K;
44
    int B;
45
} p;
46

47
void main()
48
{
49
    int gx = int(gl_GlobalInvocationID.x);
50
    int gy = int(gl_GlobalInvocationID.y);
51
    int gz = int(gl_GlobalInvocationID.z);
52

53
    if (gx >= psc(N) || gy >= psc(M) || gz >= psc(B))
54
        return;
55

56
    afpvec4 sum = afpvec4(0.f);
57

58
    for (int k = 0; k < psc(K); k++)
59
    {
60
#if NCNN_image_shader
61
        afpvec4 qk0 = image3d_ld4(qkcross_blob_3d, ivec3(k, gz * psc(M) + gy, 0));
62
        afp v0 = image3d_ld1(v_blob_3d, ivec3(k, gz * psc(N) + gx, 0));
63
#else
64
        const int ai = gz * psc(M) * psc(K) + gy * psc(K) + k;
65
        afpvec4 qk0 = buffer_ld4(qkcross_blob_data, ai);
66

67
        const int bi = gz * psc(N) * psc(K) + gx * psc(K) + k;
68
        afp v0 = buffer_ld1(v_blob_data, bi);
69
#endif
70

71
        sum += qk0 * v0;
72
    }
73

74
#if NCNN_image_shader
75
    image3d_st1(qkvcross_blob_3d, ivec3(gy * 4, gz * psc(N) + gx, 0), sum.r);
76
    image3d_st1(qkvcross_blob_3d, ivec3(gy * 4 + 1, gz * psc(N) + gx, 0), sum.g);
77
    image3d_st1(qkvcross_blob_3d, ivec3(gy * 4 + 2, gz * psc(N) + gx, 0), sum.b);
78
    image3d_st1(qkvcross_blob_3d, ivec3(gy * 4 + 3, gz * psc(N) + gx, 0), sum.a);
79
#else
80
    const int gi = (gz * psc(M) * psc(N) + gx * psc(M) + gy) * 4;
81

82
    buffer_st1(qkvcross_blob_data, gi, sum.r);
83
    buffer_st1(qkvcross_blob_data, gi + 1, sum.g);
84
    buffer_st1(qkvcross_blob_data, gi + 2, sum.b);
85
    buffer_st1(qkvcross_blob_data, gi + 3, sum.a);
86
#endif
87
}
88