vulkan: Implement top-k (#17418)
* vulkan: Implement top-k Each pass launches workgroups that each sort 2^N elements (where N is usually 7-10) and discards all but the top K. Repeat until only K are left. And there's a fast path when K==1 to just find the max value rather than sorting. * fix pipeline selection * vulkan: Add N-ary search algorithm for topk * microoptimizations
This commit is contained in:
Jeff Bolz
GitHub
@@ -0,0 +1,113 @@
|
||||
#version 450
|
||||
#extension GL_EXT_control_flow_attributes : enable
|
||||
|
||||
#include "types.glsl"
|
||||
|
||||
layout(constant_id = 0) const int BLOCK_SIZE = 1024;
|
||||
layout(constant_id = 1) const int NCOLS_PADDED_LOG2 = 10;
|
||||
|
||||
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
|
||||
|
||||
// Input can either be the source (A) or intermediate values (S).
|
||||
// Similarly, output can be either destination (D) or intermediate values (S).
|
||||
layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
|
||||
layout (binding = 0) readonly buffer S {ivec2 data_s[];};
|
||||
layout (binding = 1) writeonly buffer D {int data_d[];};
|
||||
layout (binding = 1) writeonly buffer T {ivec2 data_t[];};
|
||||
|
||||
layout (push_constant) uniform parameter {
|
||||
uint orig_ncols;
|
||||
uint ncols_input;
|
||||
uint ncols_output;
|
||||
uint nrows;
|
||||
uint first_pass;
|
||||
uint last_pass;
|
||||
} p;
|
||||
|
||||
// pairs of (gid, value)
|
||||
shared ivec2 dst_row[BLOCK_SIZE];
|
||||
|
||||
void topk(bool needs_bounds_check, const uint row) {
|
||||
const int col = int(gl_LocalInvocationID.x);
|
||||
|
||||
// initialize indices
|
||||
if (gl_GlobalInvocationID.x < p.ncols_input) {
|
||||
if (p.first_pass != 0) {
|
||||
const uint row_offset = row * p.ncols_input;
|
||||
dst_row[col] = ivec2(gl_GlobalInvocationID.x, floatBitsToInt(data_a[row_offset + gl_GlobalInvocationID.x]));
|
||||
} else {
|
||||
const uint row_offset = row * p.orig_ncols;
|
||||
dst_row[col] = data_s[row_offset + gl_GlobalInvocationID.x];
|
||||
}
|
||||
} else {
|
||||
dst_row[col] = ivec2(p.orig_ncols, 0);
|
||||
}
|
||||
barrier();
|
||||
|
||||
if (p.ncols_output == 1) {
|
||||
// Fast path for single output - just do a max reduction
|
||||
[[unroll]] for (int s = BLOCK_SIZE / 2; s >= 1; s /= 2) {
|
||||
if (col < s) {
|
||||
ivec2 a = dst_row[col];
|
||||
ivec2 b = dst_row[col + s];
|
||||
if (a.x >= p.orig_ncols ||
|
||||
b.x < p.orig_ncols && b.y > a.y) {
|
||||
dst_row[col] = b;
|
||||
}
|
||||
}
|
||||
barrier();
|
||||
}
|
||||
} else {
|
||||
// bitonic sort on this group of elements
|
||||
uint num_outer_loop_iters = NCOLS_PADDED_LOG2;
|
||||
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;
|
||||
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;
|
||||
|
||||
ivec2 sh_idx_0 = dst_row[idx_0];
|
||||
ivec2 sh_idx_1 = dst_row[idx_1];
|
||||
bool idx_0_oob = needs_bounds_check ? sh_idx_0.x >= p.orig_ncols : false;
|
||||
bool idx_1_oob = needs_bounds_check ? sh_idx_1.x >= p.orig_ncols : false;
|
||||
|
||||
if ((idx_0_oob ||
|
||||
(!idx_1_oob && intBitsToFloat(sh_idx_0.y) < intBitsToFloat(sh_idx_1.y))) && (ixj > col)) {
|
||||
dst_row[idx_0] = sh_idx_1;
|
||||
dst_row[idx_1] = sh_idx_0;
|
||||
}
|
||||
|
||||
barrier();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (col < p.ncols_output && gl_GlobalInvocationID.x < p.orig_ncols) {
|
||||
if (p.last_pass != 0) {
|
||||
const uint row_offset = row * p.ncols_output;
|
||||
data_d[row_offset + col] = dst_row[col].x;
|
||||
} else {
|
||||
const uint row_offset = row * p.orig_ncols + gl_WorkGroupID.x * p.ncols_output;
|
||||
data_t[row_offset + col] = dst_row[col];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void main() {
|
||||
// Fast path for fully occupied workgroups
|
||||
if ((p.ncols_input % BLOCK_SIZE) == 0) {
|
||||
uint row = gl_WorkGroupID.y;
|
||||
while (row < p.nrows) {
|
||||
topk(false, row);
|
||||
row += gl_WorkGroupSize.y * gl_NumWorkGroups.y;
|
||||
}
|
||||
} else {
|
||||
uint row = gl_WorkGroupID.y;
|
||||
while (row < p.nrows) {
|
||||
topk(true, row);
|
||||
row += gl_WorkGroupSize.y * gl_NumWorkGroups.y;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
#version 450
|
||||
#extension GL_EXT_control_flow_attributes : enable
|
||||
#extension GL_EXT_debug_printf : enable
|
||||
#extension GL_KHR_shader_subgroup_basic : enable
|
||||
#extension GL_KHR_shader_subgroup_ballot : enable
|
||||
#extension GL_KHR_shader_subgroup_arithmetic : enable
|
||||
#extension GL_KHR_shader_subgroup_shuffle : enable
|
||||
|
||||
#include "types.glsl"
|
||||
|
||||
layout(constant_id = 0) const int BLOCK_SIZE = 1024;
|
||||
layout(constant_id = 1) const int SUBGROUP_SIZE = 32;
|
||||
layout(constant_id = 2) const int SUBGROUP_SIZE_LOG2 = 5;
|
||||
|
||||
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
|
||||
|
||||
// Input can either be the source (A) or intermediate values (S).
|
||||
// Similarly, output can be either destination (D) or intermediate values (S).
|
||||
layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
|
||||
layout (binding = 0) readonly buffer S {ivec2 data_s[];};
|
||||
layout (binding = 1) writeonly buffer D {int data_d[];};
|
||||
layout (binding = 1) writeonly buffer T {ivec2 data_t[];};
|
||||
|
||||
layout (push_constant) uniform parameter {
|
||||
uint orig_ncols;
|
||||
uint ncols_input;
|
||||
uint ncols_output;
|
||||
uint nrows;
|
||||
uint first_pass;
|
||||
uint last_pass;
|
||||
} p;
|
||||
|
||||
// pairs of (gid, value)
|
||||
shared ivec2 dst_row[BLOCK_SIZE];
|
||||
|
||||
shared int counts[SUBGROUP_SIZE];
|
||||
shared int sh_min_idx;
|
||||
shared uint sh_total;
|
||||
shared uint offset_partials[BLOCK_SIZE / SUBGROUP_SIZE];
|
||||
|
||||
// Map float values to uint such that comparisons still work.
|
||||
// Positive values set the high bit, negative values are inverted.
|
||||
// +0.0 -> 0x80000000, -0.0 -> 0x7FFFFFFF are in the correct places.
|
||||
uint f2ui(float x) {
|
||||
uint y = floatBitsToUint(x);
|
||||
if ((y & 0x80000000) != 0) {
|
||||
y ^= ~0;
|
||||
} else {
|
||||
y |= 0x80000000;
|
||||
}
|
||||
return y;
|
||||
}
|
||||
|
||||
void topk(const uint row) {
|
||||
const int tid = int(gl_LocalInvocationID.x);
|
||||
|
||||
// initialize indices
|
||||
if (gl_GlobalInvocationID.x < p.ncols_input) {
|
||||
if (p.first_pass != 0) {
|
||||
const uint row_offset = row * p.ncols_input;
|
||||
dst_row[tid] = ivec2(gl_GlobalInvocationID.x, floatBitsToInt(data_a[row_offset + gl_GlobalInvocationID.x]));
|
||||
} else {
|
||||
const uint row_offset = row * p.orig_ncols;
|
||||
dst_row[tid] = data_s[row_offset + gl_GlobalInvocationID.x];
|
||||
}
|
||||
} else {
|
||||
dst_row[tid] = ivec2(p.orig_ncols, 0xFF800000); // -inf
|
||||
}
|
||||
barrier();
|
||||
|
||||
if (p.ncols_output == 1) {
|
||||
// Fast path for single output - just do a max reduction
|
||||
[[unroll]] for (int s = BLOCK_SIZE / 2; s >= 1; s /= 2) {
|
||||
if (tid < s) {
|
||||
ivec2 a = dst_row[tid];
|
||||
ivec2 b = dst_row[tid + s];
|
||||
if (a.x >= p.orig_ncols ||
|
||||
b.x < p.orig_ncols && b.y > a.y) {
|
||||
dst_row[tid] = b;
|
||||
}
|
||||
}
|
||||
barrier();
|
||||
}
|
||||
} else {
|
||||
// Do an N-ary search to find the K-th largest value.
|
||||
// We remap the float values to be comparable as unsigned integers,
|
||||
// and split the range into 2^N smaller ranges where N is the
|
||||
// subgroup size. Count how many values are in each range, if the K-th
|
||||
// largest value is in the middle of one of thee ranges then repeat
|
||||
// and split again.
|
||||
|
||||
// Mask is the current set of bits we're searching. Shift is the LSB index.
|
||||
int shift = 32 - SUBGROUP_SIZE_LOG2;
|
||||
uint mask = ((1 << SUBGROUP_SIZE_LOG2) - 1) << shift;
|
||||
|
||||
// The current range.
|
||||
uint range_min = 0;
|
||||
uint range_max = 0xFF800000;
|
||||
// How many are above the current range, and how many we need to find.
|
||||
uint total = 0;
|
||||
uint limit = min(p.ncols_output, p.ncols_input - gl_WorkGroupID.x * BLOCK_SIZE);
|
||||
|
||||
while (mask != 0) {
|
||||
barrier();
|
||||
// Initialize bucket counts to zero.
|
||||
if (tid < SUBGROUP_SIZE) {
|
||||
counts[tid] = 0;
|
||||
}
|
||||
barrier();
|
||||
// Count how many values are in each bucket.
|
||||
if (tid < p.ncols_input) {
|
||||
float y = intBitsToFloat(dst_row[tid].y);
|
||||
uint fy = f2ui(y);
|
||||
if (fy >= range_min && fy < range_max) {
|
||||
uint bucket = (fy & mask) >> shift;
|
||||
atomicAdd(counts[bucket], 1);
|
||||
}
|
||||
}
|
||||
barrier();
|
||||
|
||||
// On the first subgroup, do a scan to count (from the top down) how
|
||||
// many elements are in the top N buckets. Find the index of the first
|
||||
// that is over the limit. Copy it to the other invocations through
|
||||
// shared memory.
|
||||
if (tid < SUBGROUP_SIZE) {
|
||||
uint partial_sum = counts[SUBGROUP_SIZE - 1 - tid];
|
||||
partial_sum = subgroupInclusiveAdd(partial_sum) + total;
|
||||
uint t = subgroupBallotFindLSB(subgroupBallot(partial_sum >= limit));
|
||||
if (tid == t) {
|
||||
sh_min_idx = int(SUBGROUP_SIZE - 1 - t);
|
||||
sh_total = partial_sum;
|
||||
}
|
||||
}
|
||||
barrier();
|
||||
int min_idx = sh_min_idx;
|
||||
total = sh_total;
|
||||
|
||||
// Update the range, and break if we've found the K-th largest.
|
||||
range_max = range_min + ((min_idx + 1) << shift);
|
||||
range_min = range_min + (min_idx << shift);
|
||||
|
||||
if (total == p.ncols_output) {
|
||||
break;
|
||||
}
|
||||
total -= counts[min_idx];
|
||||
mask >>= SUBGROUP_SIZE_LOG2;
|
||||
shift -= SUBGROUP_SIZE_LOG2;
|
||||
if (shift < 0) {
|
||||
shift = 0;
|
||||
}
|
||||
}
|
||||
|
||||
ivec2 v = dst_row[tid];
|
||||
|
||||
// We need to compact these values to the start of the dst_row array.
|
||||
// Have each subgroup count how many items it'll store, so other
|
||||
// subgroups can compute their base offset.
|
||||
bool top = f2ui(intBitsToFloat(v.y)) >= range_min;
|
||||
uvec4 b = subgroupBallot(top);
|
||||
uint bit_count = subgroupBallotBitCount(b);
|
||||
if ((tid % SUBGROUP_SIZE) == 0) {
|
||||
offset_partials[tid / SUBGROUP_SIZE] = bit_count;
|
||||
}
|
||||
barrier();
|
||||
|
||||
uint out_idx = 0;
|
||||
[[unroll]] for (int i = 0; i < BLOCK_SIZE / SUBGROUP_SIZE; ++i) {
|
||||
if (i < tid / SUBGROUP_SIZE) {
|
||||
out_idx += offset_partials[i];
|
||||
}
|
||||
}
|
||||
|
||||
uint bit_count_ex = subgroupBallotExclusiveBitCount(b);
|
||||
if (top) {
|
||||
// TODO: Copy directly to the output?
|
||||
dst_row[out_idx + bit_count_ex] = v;
|
||||
}
|
||||
|
||||
barrier();
|
||||
}
|
||||
|
||||
if (tid < p.ncols_output && gl_GlobalInvocationID.x < p.orig_ncols) {
|
||||
if (p.last_pass != 0) {
|
||||
const uint row_offset = row * p.ncols_output;
|
||||
data_d[row_offset + tid] = dst_row[tid].x;
|
||||
} else {
|
||||
const uint row_offset = row * p.orig_ncols + gl_WorkGroupID.x * p.ncols_output;
|
||||
data_t[row_offset + tid] = dst_row[tid];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void main() {
|
||||
uint row = gl_WorkGroupID.y;
|
||||
while (row < p.nrows) {
|
||||
topk(row);
|
||||
row += gl_WorkGroupSize.y * gl_NumWorkGroups.y;
|
||||
}
|
||||
}
|
||||
@@ -913,6 +913,9 @@ void process_shaders() {
|
||||
string_to_spv("argsort_f32", "argsort.comp", {{"A_TYPE", "float"}});
|
||||
string_to_spv("argsort_large_f32", "argsort_large.comp", {{"A_TYPE", "float"}});
|
||||
|
||||
string_to_spv("topk_argsort_f32", "topk_argsort.comp", {{"A_TYPE", "float"}});
|
||||
string_to_spv("topk_nary_search_f32", "topk_nary_search.comp", {{"A_TYPE", "float"}});
|
||||
|
||||
string_to_spv("argmax_f32", "argmax.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "int"}}));
|
||||
string_to_spv("sum_rows_f32", "sum_rows.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
|
||||
string_to_spv("count_equal_i32", "count_equal.comp", merge_maps(base_dict, {{"A_TYPE", "int"}, {"B_TYPE", "int"}, {"D_TYPE", "int"}}));
|
||||
|
||||
Reference in New Issue
Block a user