1  #version 450
 2  
 3  #ifdef FLOAT16
 4  #extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
 5  #endif
 6  
 7  #include "mul_mat_vec_base.comp"
 8  
 9  layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
10  
11  layout (constant_id = 0) const uint BLOCK_SIZE = 32;
12  
13  shared FLOAT_TYPE tmp[BLOCK_SIZE];
14  
15  void main() {
16      const uint row = gl_WorkGroupID.x;
17      const uint tid = gl_LocalInvocationID.x;
18  
19      uint a_offset, b_offset, d_offset;
20      get_offsets(a_offset, b_offset, d_offset);
21  
22      const uint y_offset = QUANT_R == 1 ? 1 : QUANT_K/2;
23  
24      tmp[tid] = FLOAT_TYPE(0.0f);
25  
26      [[unroll]] for (uint i = 0; i < p.ncols/BLOCK_SIZE; i += 2) {
27          const uint col = i*BLOCK_SIZE + 2*tid;
28          const uint ib = (row*p.ncols + col)/QUANT_K; // block index
29          const uint iqs = (col%QUANT_K)/QUANT_R; // quant index
30          const uint iybs = col - col%QUANT_K; // y block start index
31  
32          vec2 v = dequantize(ib, iqs, a_offset / QUANT_K);
33  
34          // matrix multiplication
35          tmp[tid] += FLOAT_TYPE(v.x) * FLOAT_TYPE(data_b[b_offset + iybs + iqs]) +
36                      FLOAT_TYPE(v.y) * FLOAT_TYPE(data_b[b_offset + iybs + iqs + y_offset]);
37      }
38  
39      // sum up partial sums and write back result
40      barrier();
41      [[unroll]] for (uint s = BLOCK_SIZE/2; s > 0; s >>= 1) {
42          if (tid < s) {
43              tmp[tid] += tmp[tid + s];
44          }
45          barrier();
46      }
47      if (tid == 0) {
48          data_d[d_offset + row] = D_TYPE(tmp[0]);
49      }
50  }