Hugging Face's logo

Spaces:
YZ-TAN
/
flask-llama
Build error

App Files Community
flask-llama
File size: 2,352 Bytes 
5a29263
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
 
#version 450

#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_spirv_intrinsics: enable
#extension GL_EXT_control_flow_attributes : require

#if RTE16
spirv_execution_mode(capabilities = [4467], 4462, 16); // RoundingModeRTE, 16 bits
#endif

layout (push_constant) uniform parameter
{
    uint batch_offset; uint offset_delta;
    uint IC;
    uint IW; uint IH;
    uint OW; uint OH;
    uint KW; uint KH;
    uint pelements;
    uint CHW;
    int s0; int s1;
    int p0; int p1;
    int d0; int d1;
} p;

#include "types.comp"

layout(constant_id = 0) const uint BLOCK_SIZE = 32;

const uint NUM_ITER = 512 / BLOCK_SIZE;

layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;

layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};

void main() {
    const uint gidx = gl_GlobalInvocationID.x;

    const uint oh = gl_GlobalInvocationID.y;
    const uint batch = gl_GlobalInvocationID.z / p.IC;
    const uint ic = gl_GlobalInvocationID.z % p.IC;

    A_TYPE values[NUM_ITER];
    uint offset_dst[NUM_ITER];
    [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) {
        values[idx] = A_TYPE(0);
    }

    [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) {

        const uint i = gidx * NUM_ITER + idx;

        const uint ksize = p.OW * (p.KH > 1 ? p.KW : 1);
        const uint kx = i / ksize;
        const uint kd = kx * ksize;
        const uint ky = (i - kd) / p.OW;
        const uint ix = i % p.OW;

        const uint iiw = ix * p.s0 + kx * p.d0 - p.p0;
        const uint iih = oh * p.s1 + ky * p.d1 - p.p1;

        offset_dst[idx] =
            ((batch * p.OH + oh) * p.OW + ix) * p.CHW +
            (ic * (p.KW * p.KH) + ky * p.KW + kx);

        if (i >= p.pelements) {
            continue;
        }

        if (iih < p.IH && iiw < p.IW) {
            const uint offset_src = ic * p.offset_delta + batch * p.batch_offset;
            values[idx] = data_a[offset_src + iih * p.IW + iiw];
        }
    }

    [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) {

        const uint i = gidx * NUM_ITER + idx;

        if (i >= p.pelements) {
            continue;
        }

        data_d[offset_dst[idx]] = D_TYPE(values[idx]);
    }

}