KBaba7/llama.cpp

1#version 4502 3#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require4 5#include "mul_mat_vec_base.comp"6 7layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;8 9#if !defined(DATA_A_F32) && !defined(DATA_A_F16)10#define K_PER_ITER 811#else12#define K_PER_ITER 213#endif14 15 16uint a_offset, b_offset, d_offset, y_offset;17 18void iter(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const uint num_rows, const uint tid, const uint i, bool lastiter)19{20    [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) {21        const uint col = i*BLOCK_SIZE + K_PER_ITER*tid;22        const uint iqs = (col%QUANT_K)/QUANT_R; // quant index23        const uint iybs = col - col%QUANT_K; // y block start index24 25#if K_PER_ITER == 826#if QUANT_R == 227        const vec4 bv02 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4]);28        const vec4 bv13 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs + y_offset) / 4]);29        const vec4 bv0 = vec4(bv02.x, bv13.x, bv02.y, bv13.y);30        const vec4 bv1 = vec4(bv02.z, bv13.z, bv02.w, bv13.w);31#else32        const vec4 bv0 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4]);33        const vec4 bv1 = vec4(data_b_v4[(j*p.batch_stride_b + b_offset + iybs + iqs) / 4 + 1]);34#endif35#else36        // Check if the second of the pair of elements is OOB, and don't fetch B or37        // accumulate it. We still fetch a pair of elements for A, which is fine for38        // quantized formats since they'll be within the same block. We should39        // probably skip fetching the second element for F16/F32, but as of now we40        // still do.41        const bool OOB = lastiter && (iybs + iqs + y_offset >= p.ncols);42 43        FLOAT_TYPE b0 = 0, b1 = 0;44        b0 = FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs]);45        if (!OOB) {46            b1 = FLOAT_TYPE(data_b[j*p.batch_stride_b + b_offset + iybs + iqs + y_offset]);47        }48#endif49        uint ibi = first_row*p.ncols;50        [[unroll]] for (uint n = 0; n < num_rows; ++n) {51            const uint ib = (ibi + col)/QUANT_K; // block index52            ibi += p.ncols;53 54#if K_PER_ITER == 855            vec4 v = dequantize4(ib, iqs, a_offset);56            vec4 v2 = dequantize4(ib, iqs+(4/QUANT_R), a_offset);57 58            const vec2 dm = get_dm(ib, a_offset);59            if (dm.y != 0) { // quant has min component60                v = v * dm.x + dm.y;61                v2 = v2 * dm.x + dm.y;62            }63 64            // matrix multiplication65            FLOAT_TYPE rowtmp = dot(bv0, v);66            rowtmp += dot(bv1, v2);67 68            if (dm.y == 0)69                rowtmp *= dm.x;70 71            temp[j][n] += rowtmp;72#else73            const vec2 v = dequantize(ib, iqs, a_offset);74 75            // matrix multiplication76            temp[j][n] = fma(FLOAT_TYPE(v.x), b0, temp[j][n]);77            if (!OOB) {78                temp[j][n] = fma(FLOAT_TYPE(v.y), b1, temp[j][n]);79            }80#endif81        }82    }83}84 85void compute_outputs(const uint32_t first_row, const uint32_t num_rows) {86    const uint tid = gl_LocalInvocationID.x;87 88    get_offsets(a_offset, b_offset, d_offset);89    a_offset /= QUANT_K;90 91    y_offset = QUANT_R == 1 ? 1 : QUANT_K/2;92 93    FLOAT_TYPE temp[NUM_COLS][NUM_ROWS];94 95    [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) {96        [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) {97            temp[j][i] = FLOAT_TYPE(0);98        }99    }100 101    uint num_iters = p.ncols / (K_PER_ITER * BLOCK_SIZE);102    if (num_iters * K_PER_ITER * BLOCK_SIZE + K_PER_ITER*tid < p.ncols) {103        num_iters++;104    }105    int unroll_count = 4;106    uint unrolled_iters = num_iters & ~(unroll_count - 1);107 108    uint i = 0;109    while (i < unrolled_iters) {110        // Manually partially unroll the loop111        [[unroll]] for (uint k = 0; k < unroll_count; ++k) {112            iter(temp, first_row, num_rows, tid, i*K_PER_ITER, false);113            i++;114        }115    }116    unroll_count = 2;117    unrolled_iters = num_iters & ~(unroll_count - 1);118    while (i < unrolled_iters) {119        // Manually partially unroll the loop120        [[unroll]] for (uint k = 0; k < unroll_count; ++k) {121            iter(temp, first_row, num_rows, tid, i*K_PER_ITER, false);122            i++;123        }124    }125    while (i < num_iters) {126        iter(temp, first_row, num_rows, tid, i*K_PER_ITER, true);127        i++;128    }129 130    reduce_result(temp, d_offset, first_row, num_rows, tid);131}132 133void main() {134    const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z);135 136#if defined(DATA_A_IQ2_XXS) || defined(DATA_A_IQ2_XS) || defined(DATA_A_IQ2_S) || defined(DATA_A_IQ3_XXS) || defined(DATA_A_IQ3_S) || defined(DATA_A_IQ4_XS) || defined(DATA_A_IQ4_NL)137    init_iq_shmem(gl_WorkGroupSize);138#endif139 140    // do NUM_ROWS at a time, unless there aren't enough remaining rows141    if (first_row + NUM_ROWS <= p.stride_d) {142        compute_outputs(first_row, NUM_ROWS);143    } else {144        if (first_row >= p.stride_d) {145            return;146        }147        compute_outputs(first_row, p.stride_d - first_row);148    }149}150