ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp

#version 450
#extension GL_EXT_control_flow_attributes : enable

#include "types.comp"

layout(constant_id = 0) const int BLOCK_SIZE = 1024;
layout(constant_id = 1) const int BLOCK_SIZE_LOG2 = 10;
#define ASC 0

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

layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
layout (binding = 1)          buffer D {int data_d[];};

layout (push_constant) uniform parameter {
    uint ncols;
    uint order;
} p;

shared int dst_row[BLOCK_SIZE];
shared A_TYPE a_sh[BLOCK_SIZE];

void swap(uint idx0, uint idx1) {
    int tmp = dst_row[idx0];
    dst_row[idx0] = dst_row[idx1];
    dst_row[idx1] = tmp;
}

void argsort(bool needs_bounds_check) {
    // bitonic sort
    const int col = int(gl_LocalInvocationID.x);
    const uint row = gl_WorkGroupID.y;

    const uint row_offset = row * p.ncols;

    // initialize indices
    dst_row[col] = col;
    a_sh[col] = data_a[row_offset + col];
    barrier();

    uint num_outer_loop_iters = BLOCK_SIZE_LOG2;
    [[unroll]] for (uint k = 2, outer_idx = 0; outer_idx < num_outer_loop_iters; k *= 2, outer_idx++) {
        uint num_inner_loop_iters = outer_idx + 1;
        [[unroll]] for (uint j = k / 2, inner_idx = 0; inner_idx < num_inner_loop_iters; j /= 2, inner_idx++) {
            const int ixj = int(col ^ j);

            int idx_0 = (col & k) == 0 ? col : ixj;
            int idx_1 = (col & k) == 0 ? ixj : col;

            int sh_idx_0 = dst_row[idx_0];
            int sh_idx_1 = dst_row[idx_1];
            bool idx_0_oob = needs_bounds_check ? sh_idx_0 >= p.ncols : false;
            bool idx_1_oob = needs_bounds_check ? sh_idx_1 >= p.ncols : false;

            if ((idx_0_oob ||
                (!idx_1_oob && a_sh[sh_idx_0] > a_sh[sh_idx_1])) && (ixj > col)) {
                swap(idx_0, idx_1);
            }

            barrier();
        }
    }

    if (col < p.ncols) {
        if (p.order == ASC) {
            data_d[row_offset + col] = dst_row[col];
        } else {
            data_d[row_offset + p.ncols - col - 1] = dst_row[col];
        }
    }
}

void main() {
    if (p.ncols == BLOCK_SIZE) {
        argsort(false);
    } else {
        argsort(true);
    }
}