char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Ruben Ortlamvulkan: add fwht support for Intel with shmem reduction (#23964)e82beaa60

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
3.2 KiB115 linesraw
1#version 450
2
3#extension GL_EXT_control_flow_attributes : require
4#ifndef FWHT_SHMEM
5#extension GL_KHR_shader_subgroup_basic : enable
6#extension GL_KHR_shader_subgroup_shuffle : enable
7#endif
8
9layout(constant_id = 0) const uint BLOCK_SIZE = 32;
10layout(constant_id = 1) const uint N = 128;
11
12layout(local_size_x_id = 0, local_size_y = 4, local_size_z = 1) in;
13
14layout(push_constant) uniform parameter
15{
16    uint n_rows;
17    uint src_offset;
18    uint dst_offset;
19    float scale;
20};
21
22layout(binding = 0, std430) readonly buffer A { float data_a[]; };
23layout(binding = 1, std430) writeonly buffer D { float data_d[]; };
24
25const uint EL_W = N / BLOCK_SIZE;
26
27#ifdef FWHT_SHMEM
28shared float shmem[4 * N];
29#endif
30
31void main() {
32#ifdef FWHT_SHMEM
33    const uint tid = gl_LocalInvocationID.x;
34    const uint shmem_base = gl_LocalInvocationID.y * N;
35    const uint row_id = gl_LocalInvocationID.y;
36#else
37    const uint tid = gl_SubgroupInvocationID;
38    const uint row_id = gl_SubgroupID;
39#endif
40
41    for (uint base_row = gl_WorkGroupID.x * gl_WorkGroupSize.y;
42            base_row < n_rows;
43            base_row += gl_NumWorkGroups.x * gl_WorkGroupSize.y) {
44        const uint row = base_row + row_id;
45        const uint row_offset = row * N;
46
47#ifndef FWHT_SHMEM
48        if (row >= n_rows) {
49            continue;
50        }
51#endif
52
53        float reg[EL_W];
54
55        [[unroll]]
56        for (uint i = 0; i < EL_W; ++i) {
57            reg[i] = row < n_rows ? data_a[src_offset + row_offset + i * BLOCK_SIZE + tid] * scale : 0.0;
58        }
59
60#ifdef FWHT_SHMEM
61        [[unroll]]
62        for (uint h = 1; h < BLOCK_SIZE; h <<= 1) {
63            [[unroll]]
64            for (uint i = 0; i < EL_W; ++i) {
65                shmem[shmem_base + i * BLOCK_SIZE + tid] = reg[i];
66            }
67            barrier();
68            [[unroll]]
69            for (uint j = 0; j < EL_W; ++j) {
70                const float val = reg[j];
71                const float other = shmem[shmem_base + j * BLOCK_SIZE + (tid ^ h)];
72                reg[j] = (tid & h) == 0 ? val + other : other - val;
73            }
74            barrier();
75        }
76#else
77        [[unroll]]
78        for (uint h = 1; h < BLOCK_SIZE; h <<= 1) {
79            [[unroll]]
80            for (uint j = 0; j < EL_W; ++j) {
81                const float val = reg[j];
82                const float val2 = subgroupShuffleXor(val, h);
83                reg[j] = (tid & h) == 0 ? val + val2 : val2 - val;
84            }
85        }
86#endif
87
88        [[unroll]]
89        for (uint h = BLOCK_SIZE; h < N; h <<= 1) {
90            const uint step = h / BLOCK_SIZE;
91            [[unroll]]
92            for (uint j = 0; j < EL_W; j += 2 * step) {
93                [[unroll]]
94                for (uint k = 0; k < step; ++k) {
95                    const float x = reg[j + k];
96                    const float y = reg[j + k + step];
97                    reg[j + k] = x + y;
98                    reg[j + k + step] = x - y;
99                }
100            }
101        }
102
103#ifdef FWHT_SHMEM
104        if (row < n_rows) {
105#endif
106            [[unroll]]
107            for (uint i = 0; i < EL_W; ++i) {
108                data_d[dst_offset + row_offset + i * BLOCK_SIZE + tid] = reg[i];
109            }
110#ifdef FWHT_SHMEM
111        }
112        barrier();
113#endif
114    }
115}