mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2025-11-20 12:07:33 +00:00
de5627910d
- Launch an appropriate number of invocations (next larger power of two). 32 invocations is common and the barrier is much cheaper there. - Specialize for "needs bounds checking" vs not. - Make the code less branchy and [[unroll]] the loops. In the final code, I see no branches inside the main loop (only predicated stores) when needs_bounds_check is false. - Always sort ascending, then apply the ascending vs descending option when doing the final stores to memory. - Copy the values into shared memory, makes them slightly cheaper to access.
80 lines
2.2 KiB
Plaintext
80 lines
2.2 KiB
Plaintext
#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);
|
|
}
|
|
}
|