Skip to content

Commit 07d0d58

Browse files
netrunnerevenetrunnereve
committed
q3_k
1 parent b0e4ccb commit 07d0d58

File tree

1 file changed

+14
-22
lines changed

1 file changed

+14
-22
lines changed

‎ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_q3_k.comp

Lines changed: 14 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,9 @@
55

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

8+
shared uint32_t scu8[BLOCK_SIZE/16][12];
9+
shared FLOAT_TYPE sccache[BLOCK_SIZE/16][2][12];
10+
811
void compute_outputs(const uint32_t first_row, const uint32_t num_rows) {
912
uint a_offset, b_offset, d_offset;
1013
get_offsets(a_offset, b_offset, d_offset);
@@ -16,6 +19,7 @@ void compute_outputs(const uint32_t first_row, const uint32_t num_rows) {
1619
const uint tid = gl_LocalInvocationID.x;
1720
const uint itid = tid%16; // 0...15
1821
const uint ix = tid/16;
22+
const uint itid8 = itid%8;
1923

2024
const uint v_im = itid/8; // 0 or 1. 0 computes 0..., 1 computes 128...
2125
const uint v_in = itid - 8*v_im; // 0...15 or 0...7
@@ -34,27 +38,15 @@ void compute_outputs(const uint32_t first_row, const uint32_t num_rows) {
3438
}
3539
}
3640

37-
const uint s_shift = 4 * v_im;
38-
3941
[[unroll]] for (uint i = ix; i < num_blocks_per_row; i += it_size) {
4042
const uint y_idx = i * QUANT_K + y_offset;
4143

4244
[[unroll]] for (uint n = 0; n < num_rows; ++n) {
4345
const uint ib0 = a_offset / QUANT_K + (first_row+n)*num_blocks_per_row;
4446
const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib0 + i].d);
4547

46-
uint16_t s0_16 = data_a_packed16[ib0 + i].scales[0];
47-
uint16_t s2_16 = data_a_packed16[ib0 + i].scales[1];
48-
uint16_t s4_16 = data_a_packed16[ib0 + i].scales[2];
49-
uint16_t s6_16 = data_a_packed16[ib0 + i].scales[3];
50-
uint16_t s8_16 = data_a_packed16[ib0 + i].scales[4];
51-
uint16_t s10_16 = data_a_packed16[ib0 + i].scales[5];
52-
u8vec2 s0 = unpack8(s0_16);
53-
u8vec2 s2 = unpack8(s2_16);
54-
u8vec2 s4 = unpack8(s4_16);
55-
u8vec2 s6 = unpack8(s6_16);
56-
u8vec2 s8 = unpack8(s8_16);
57-
u8vec2 s10 = unpack8(s10_16);
48+
sccache[ix][v_im][itid8] = FLOAT_TYPE(int8_t(((int8_t(data_a[ib0+i].scales[itid8]) >> 4*v_im) & 0xF) | ((int8_t(data_a[ib0+i].scales[itid%4+8]) >> (4*v_im + 2*(itid8/4)) & 0x3) << 4)) - 32);
49+
barrier();
5850

5951
[[unroll]] for (uint j = 0; j < NUM_COLS; ++j) {
6052

@@ -69,14 +61,14 @@ void compute_outputs(const uint32_t first_row, const uint32_t num_rows) {
6961

7062
FLOAT_TYPE sum = FLOAT_TYPE(0.0);
7163
[[unroll]] for (int l = 0; l < 2; ++l) {
72-
sum = fma(FLOAT_TYPE(b0[l]) * FLOAT_TYPE(int8_t(((s0[0] >> s_shift) & 0xF) | ((s8[0] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] ) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 0)) != 0) ? 0 : 4)),
73-
fma(FLOAT_TYPE(b32[l]) * FLOAT_TYPE(int8_t(((s2[0] >> s_shift) & 0xF) | ((s10[0] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 2) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 1)) != 0) ? 0 : 4)),
74-
fma(FLOAT_TYPE(b64[l]) * FLOAT_TYPE(int8_t(((s4[0] >> s_shift) & 0xF) | ((s8[0] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 4) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 2)) != 0) ? 0 : 4)),
75-
fma(FLOAT_TYPE(b96[l]) * FLOAT_TYPE(int8_t(((s6[0] >> s_shift) & 0xF) | ((s10[0] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 6) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 3)) != 0) ? 0 : 4)),
76-
fma(FLOAT_TYPE(b16[l]) * FLOAT_TYPE(int8_t(((s0[1] >> s_shift) & 0xF) | ((s8[1] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] ) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 0)) != 0) ? 0 : 4)),
77-
fma(FLOAT_TYPE(b48[l]) * FLOAT_TYPE(int8_t(((s2[1] >> s_shift) & 0xF) | ((s10[1] >> (s_shift + 0) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 2) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 1)) != 0) ? 0 : 4)),
78-
fma(FLOAT_TYPE(b80[l]) * FLOAT_TYPE(int8_t(((s4[1] >> s_shift) & 0xF) | ((s8[1] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 4) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 2)) != 0) ? 0 : 4)),
79-
fma(FLOAT_TYPE(b112[l]) * FLOAT_TYPE(int8_t(((s6[1] >> s_shift) & 0xF) | ((s10[1] >> (s_shift + 2) & 0x3) << 4)) - 32), FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 6) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 3)) != 0) ? 0 : 4)), sum))))))));
64+
sum = fma(FLOAT_TYPE( b0[l]) * sccache[ix][v_im][0], FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] ) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m )) != 0) ? 0 : 4)),
65+
fma(FLOAT_TYPE( b16[l]) * sccache[ix][v_im][1], FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] ) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m )) != 0) ? 0 : 4)),
66+
fma(FLOAT_TYPE( b32[l]) * sccache[ix][v_im][2], FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 2) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 1)) != 0) ? 0 : 4)),
67+
fma(FLOAT_TYPE( b48[l]) * sccache[ix][v_im][3], FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 2) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 1)) != 0) ? 0 : 4)),
68+
fma(FLOAT_TYPE( b64[l]) * sccache[ix][v_im][4], FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 4) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 2)) != 0) ? 0 : 4)),
69+
fma(FLOAT_TYPE( b80[l]) * sccache[ix][v_im][5], FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 4) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 2)) != 0) ? 0 : 4)),
70+
fma(FLOAT_TYPE( b96[l]) * sccache[ix][v_im][6], FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 6) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 3)) != 0) ? 0 : 4)),
71+
fma(FLOAT_TYPE(b112[l]) * sccache[ix][v_im][7], FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 6) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 3)) != 0) ? 0 : 4)), sum))))))));
8072
}
8173
temp[j][n] = fma(d, sum, temp[j][n]);
8274
}

0 commit comments

Comments
 (0)