Spaces:
Build error
Build error
File size: 3,203 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 |
#version 450
#include "dequant_head.comp"
layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
void main() {
[[unroll]] for (uint wgy = 0; wgy < 256; wgy++) {
const uint ib = gl_WorkGroupID.x * 256 + wgy;
if (ib >= p.M * p.K / QUANT_K) {
return;
}
const uint tid = gl_LocalInvocationID.x;
const uint il = tid / 16;
const uint ir = tid % 16;
const uint is = 2 * il;
const FLOAT_TYPE dall = FLOAT_TYPE(data_a[ib].d.x);
const FLOAT_TYPE dmin = FLOAT_TYPE(data_a[ib].d.y);
const uint y_idx = ib * QUANT_K + 64 * il + 2 * ir;
const uint qs_idx = 32*il + 2 * ir;
const uint qh_idx = 2 * ir;
uint scidx0 = (is < 4) ? is : (is + 4);
uint scidx1 = (is < 4) ? is : (is - 4);
uint scidxmask1 = (is < 4) ? 0x30 : 0xC0;
uint scidxshift1 = (is < 4) ? 0 : 2;
uint mbidx0 = is + 4;
uint mbidx1 = (is < 4) ? is + 4 : is;
uint mbidxmask0 = (is < 4) ? 0xF : 0xF0;
uint mbidxshift0 = (is < 4) ? 0 : 4;
uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0;
uint mbidxshift1 = (is < 4) ? 0 : 2;
uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1));
uint8_t mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1));
const FLOAT_TYPE d1 = dall * sc;
const FLOAT_TYPE m1 = dmin * mbyte;
scidx0 = (is < 4) ? is + 1 : (is + 5);
scidx1 = (is < 4) ? is + 1 : (is - 3);
scidxmask1 = (is < 4) ? 0x30 : 0xC0;
scidxshift1 = (is < 4) ? 0 : 2;
mbidx0 = is + 5;
mbidx1 = (is < 4) ? is + 5 : is + 1;
mbidxmask0 = (is < 4) ? 0xF : 0xF0;
mbidxshift0 = (is < 4) ? 0 : 4;
mbidxmask1 = (is < 4) ? 0x30 : 0xC0;
mbidxshift1 = (is < 4) ? 0 : 2;
sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1));
mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1));
const FLOAT_TYPE d2 = dall * sc;
const FLOAT_TYPE m2 = dmin * mbyte;
const uint8_t hm1 = uint8_t(1 << (2 * il ));
const uint8_t hm2 = uint8_t(1 << (2 * il + 1));
data_b[y_idx ] = D_TYPE(d1 * FLOAT_TYPE((data_a[ib].qs[qs_idx ] & 0xF) + (((data_a[ib].qh[qh_idx ] & hm1) != 0) ? 16 : 0)) - m1);
data_b[y_idx + 1] = D_TYPE(d1 * FLOAT_TYPE((data_a[ib].qs[qs_idx + 1] & 0xF) + (((data_a[ib].qh[qh_idx + 1] & hm1) != 0) ? 16 : 0)) - m1);
data_b[y_idx + 32] = D_TYPE(d2 * FLOAT_TYPE((data_a[ib].qs[qs_idx ] >> 4) + (((data_a[ib].qh[qh_idx ] & hm2) != 0) ? 16 : 0)) - m2);
data_b[y_idx + 33] = D_TYPE(d2 * FLOAT_TYPE((data_a[ib].qs[qs_idx + 1] >> 4) + (((data_a[ib].qh[qh_idx + 1] & hm2) != 0) ? 16 : 0)) - m2);
}
}
|