diff --git a/crates/backend-uzu/src/backends/common/kernel/mod.rs b/crates/backend-uzu/src/backends/common/kernel/mod.rs index 9d0d9c6d3..bf8c96082 100644 --- a/crates/backend-uzu/src/backends/common/kernel/mod.rs +++ b/crates/backend-uzu/src/backends/common/kernel/mod.rs @@ -4,6 +4,7 @@ pub mod attention_gemm; pub mod delta_net_chunked_prefill; pub mod delta_net_tree_verify; pub mod matmul; +pub mod radix_top_k_small; include!(concat!(env!("OUT_DIR"), "/traits.rs")); @@ -15,6 +16,7 @@ pub trait Kernels: Sized { type DeltaNetChunkedPrefill: delta_net_chunked_prefill::DeltaNetChunkedPrefill; type DeltaNetTreeVerify: delta_net_tree_verify::DeltaNetTreeVerify; type MatmulKernel: matmul::MatmulKernel; + type RadixTopKSmall: radix_top_k_small::RadixTopKSmall; } #[cfg(test)] diff --git a/crates/backend-uzu/src/backends/common/kernel/radix_top_k_small.rs b/crates/backend-uzu/src/backends/common/kernel/radix_top_k_small.rs new file mode 100644 index 000000000..c8767fa00 --- /dev/null +++ b/crates/backend-uzu/src/backends/common/kernel/radix_top_k_small.rs @@ -0,0 +1,20 @@ +use crate::backends::common::{Allocation, Backend, Encoder}; + +pub const MAX_K: u32 = 512; + +pub trait RadixTopKSmall: Sized { + fn new( + context: &B::Context, + columns: u32, + ) -> Result; + + fn encode( + &self, + input: &Allocation, + output_ids: &mut Allocation, + output_scores: &mut Allocation, + rows: u32, + k: u32, + encoder: &mut Encoder, + ) -> Result<(), B::Error>; +} diff --git a/crates/backend-uzu/src/backends/cpu/kernel/mod.rs b/crates/backend-uzu/src/backends/cpu/kernel/mod.rs index c140738d9..9675a849a 100644 --- a/crates/backend-uzu/src/backends/cpu/kernel/mod.rs +++ b/crates/backend-uzu/src/backends/cpu/kernel/mod.rs @@ -13,6 +13,7 @@ mod matmul; mod moe; mod normalization; mod pooling; +mod radix_top_k_small; mod sampling; mod short_conv; mod softmax; @@ -34,4 +35,5 @@ impl Kernels for CpuKernels { type DeltaNetChunkedPrefill = Infallible; type DeltaNetTreeVerify = Infallible; type MatmulKernel = matmul::MatmulCpuKernel; + type RadixTopKSmall = radix_top_k_small::CpuRadixTopKSmall; } diff --git a/crates/backend-uzu/src/backends/cpu/kernel/radix_top_k_small.rs b/crates/backend-uzu/src/backends/cpu/kernel/radix_top_k_small.rs new file mode 100644 index 000000000..eb3e2bf73 --- /dev/null +++ b/crates/backend-uzu/src/backends/cpu/kernel/radix_top_k_small.rs @@ -0,0 +1,72 @@ +use crate::{ + backends::{ + common::{ + Allocation, AsBufferRangeMut, AsBufferRangeRef, Encoder, + kernel::radix_top_k_small::{MAX_K, RadixTopKSmall}, + }, + cpu::{Cpu, context::CpuContext, error::CpuError}, + }, + utils::pointers::{SendPtr, SendPtrMut}, +}; + +pub struct CpuRadixTopKSmall { + columns: usize, +} + +impl RadixTopKSmall for CpuRadixTopKSmall { + fn new( + _context: &CpuContext, + columns: u32, + ) -> Result { + assert!(columns > 0); + Ok(Self { + columns: columns as usize, + }) + } + + fn encode( + &self, + input: &Allocation, + output_ids: &mut Allocation, + output_scores: &mut Allocation, + rows: u32, + k: u32, + encoder: &mut Encoder, + ) -> Result<(), CpuError> { + let rows = rows as usize; + let columns = self.columns; + let k = k as usize; + assert!(rows > 0 && k > 0 && k <= MAX_K as usize && k <= columns); + let input = input.as_buffer_range_ref(); + let input = SendPtr(unsafe { (&*input.buffer().get()).as_ptr().add(input.range().start).cast::() }); + let output_ids = output_ids.as_buffer_range_mut(); + let output_ids = SendPtrMut(unsafe { + (&mut *output_ids.buffer().get()).as_mut_ptr().add(output_ids.range().start).cast::() + }); + let output_scores = output_scores.as_buffer_range_mut(); + let output_scores = SendPtrMut(unsafe { + (&mut *output_scores.buffer().get()).as_mut_ptr().add(output_scores.range().start).cast::() + }); + encoder.as_command_buffer_mut().push_command(move || { + let values = unsafe { std::slice::from_raw_parts(input.as_ptr(), rows * columns) }; + let output_ids = unsafe { std::slice::from_raw_parts_mut(output_ids.as_ptr(), rows * k) }; + let output_scores = unsafe { std::slice::from_raw_parts_mut(output_scores.as_ptr(), rows * k) }; + for (row, values) in values.chunks_exact(columns).enumerate() { + let mut indices = (0..columns).collect::>(); + let compare = |&left: &usize, &right: &usize| { + values[right].total_cmp(&values[left]).then_with(|| left.cmp(&right)) + }; + if k < columns { + indices.select_nth_unstable_by(k, compare); + indices.truncate(k); + } + indices.sort_unstable_by(compare); + for (rank, index) in indices.into_iter().enumerate() { + output_ids[row * k + rank] = index as u32; + output_scores[row * k + rank] = values[index]; + } + } + }); + Ok(()) + } +} diff --git a/crates/backend-uzu/src/backends/metal/kernel/common/top_k.h b/crates/backend-uzu/src/backends/metal/kernel/common/top_k.h new file mode 100644 index 000000000..c17434d58 --- /dev/null +++ b/crates/backend-uzu/src/backends/metal/kernel/common/top_k.h @@ -0,0 +1,22 @@ +#pragma once + +#include +using namespace metal; + +constant uint TOP_K_SIGN_BIT = 0x80000000u; + +static inline uint top_k_score_key(float score) { + const uint bits = as_type(score); + return (bits & TOP_K_SIGN_BIT) ? ~bits : (bits ^ TOP_K_SIGN_BIT); +} + +static inline float top_k_score_from_key(uint key) { + const uint bits = (key & TOP_K_SIGN_BIT) ? (key ^ TOP_K_SIGN_BIT) : ~key; + return as_type(bits); +} + +// Larger keys follow the f32 total order; equal scores prefer smaller indices. +static inline ulong top_k_ordered_key(float score, uint index, uint index_bits) { + const ulong index_mask = (1ul << index_bits) - 1ul; + return (ulong(top_k_score_key(score)) << index_bits) | (index_mask - ulong(index)); +} diff --git a/crates/backend-uzu/src/backends/metal/kernel/mod.rs b/crates/backend-uzu/src/backends/metal/kernel/mod.rs index e3d1126b3..86353cfb9 100644 --- a/crates/backend-uzu/src/backends/metal/kernel/mod.rs +++ b/crates/backend-uzu/src/backends/metal/kernel/mod.rs @@ -9,6 +9,7 @@ use crate::backends::{ pub mod attention; pub mod gdn; pub mod matmul; +mod radix_top_k_small; pub const MTLB: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/default.metallib")); @@ -24,4 +25,5 @@ impl Kernels for MetalKernels { type DeltaNetChunkedPrefill = gdn::chunked::MetalDeltaNetChunkedPrefill; type DeltaNetTreeVerify = gdn::tree_verify::MetalDeltaNetTreeVerify; type MatmulKernel = matmul::MatmulMetalKernel; + type RadixTopKSmall = radix_top_k_small::MetalRadixTopKSmall; } diff --git a/crates/backend-uzu/src/backends/metal/kernel/radix_top_k_small.metal b/crates/backend-uzu/src/backends/metal/kernel/radix_top_k_small.metal new file mode 100644 index 000000000..55b059d3e --- /dev/null +++ b/crates/backend-uzu/src/backends/metal/kernel/radix_top_k_small.metal @@ -0,0 +1,216 @@ +#include +#include +#include "common/defines.h" +#include "common/dsl.h" +#include "common/top_k.h" + +using namespace metal; + +#define THREADS_PER_TG 256 +constant uint MAX_K = 512; +constant uint RADIX_BITS = 10; +constant uint BUCKETS = 1 << RADIX_BITS; +constant uint BUCKETS_PER_THREAD = BUCKETS / THREADS_PER_TG; +constant uint VECTOR_WIDTH = 4; + +METAL_FUNC bool arrive_last(device atomic_uint* count, uint expected) { + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + const bool last = atomic_fetch_add_explicit(count, 1u, memory_order_relaxed) == expected - 1; + if (last) + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + return last; +} + +METAL_FUNC void reset_arrival(device atomic_uint* count) { + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + atomic_store_explicit(count, 0u, memory_order_relaxed); +} + +template +METAL_FUNC void visit_partition_keys( + const device float* input, + uint columns, + uint index_bits, + uint partition, + uint partitions, + uint lid, + Visitor visit +) { + if (columns % VECTOR_WIDTH != 0) { + const uint begin = columns * partition / partitions; + const uint end = columns * (partition + 1u) / partitions; + for (uint column = begin + lid; column < end; column += THREADS_PER_TG) + visit(top_k_ordered_key(input[column], column, index_bits)); + return; + } + const device float4* input4 = reinterpret_cast(input); + const uint vector_columns = columns / VECTOR_WIDTH; + const uint vector_begin = vector_columns * partition / partitions; + const uint vector_end = vector_columns * (partition + 1u) / partitions; + for (uint vector = vector_begin + lid; vector < vector_end; vector += THREADS_PER_TG) { + const float4 values = input4[vector]; + for (uint lane = 0; lane < VECTOR_WIDTH; ++lane) + visit(top_k_ordered_key(values[lane], vector * VECTOR_WIDTH + lane, index_bits)); + } + if (partition + 1 == partitions) { + for (uint column = vector_columns * VECTOR_WIDTH + lid; column < columns; column += THREADS_PER_TG) + visit(top_k_ordered_key(input[column], column, index_bits)); + } +} + +KERNEL(RadixTopKSmallPass)( + const device float* input, + device atomic_uint* partial_histograms, + device ulong* prefixes, + device ulong* prefix_masks, + device uint* ranks, + device atomic_uint* done_counts, + constant uint& rows, + constant uint& k, + constant uint& pass, + const uint columns SPECIALIZE, + const uint partitions SPECIALIZE, + threadgroup uint histogram[BUCKETS], + threadgroup uint simd_prefixes[THREADS_PER_TG / METAL_SIMD_SIZE], + threadgroup uint& is_last, + const uint group GROUPS(rows* partitions), + const uint lid THREADS(THREADS_PER_TG) +) { + const uint row = group / partitions; + const uint partition = group % partitions; + const uint index_bits = columns <= 1 ? 1u : 32u - clz(columns - 1u); + const uint passes = (32u + index_bits + RADIX_BITS - 1u) / RADIX_BITS; + const uint shift = (passes - pass - 1u) * RADIX_BITS; + const device float* row_input = input + ulong(row) * columns; + const ulong prefix = prefixes[row]; + const ulong mask = prefix_masks[row]; + + // Build this partition's histogram. + device atomic_uint* partition_histogram = partial_histograms + (row * partitions + partition) * BUCKETS; + for (uint bucket = lid; bucket < BUCKETS; bucket += THREADS_PER_TG) + atomic_store_explicit(&partition_histogram[bucket], 0u, memory_order_relaxed); + // Order this partition's clears before its histogram updates. + threadgroup_barrier(mem_flags::mem_device); + visit_partition_keys(row_input, columns, index_bits, partition, partitions, lid, [&](ulong key) { + if ((key & mask) == prefix) + atomic_fetch_add_explicit(&partition_histogram[(key >> shift) & (BUCKETS - 1)], 1u, memory_order_relaxed); + }); + threadgroup_barrier(mem_flags::mem_device); + + // The last partition merges the row. + if (lid == 0) + is_last = arrive_last(&done_counts[row], partitions); + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + if (!is_last) + return; + + for (uint bucket = lid; bucket < BUCKETS; bucket += THREADS_PER_TG) { + uint count = 0; + for (uint p = 0; p < partitions; ++p) + count += + atomic_load_explicit(&partial_histograms[(row * partitions + p) * BUCKETS + bucket], memory_order_relaxed); + histogram[bucket] = count; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint sum = 0; + const uint begin = BUCKETS - (lid + 1) * BUCKETS_PER_THREAD; + for (uint offset = 0; offset < BUCKETS_PER_THREAD; ++offset) + sum += histogram[begin + offset]; + uint higher = simd_prefix_exclusive_sum(sum); + if ((lid + 1) % METAL_SIMD_SIZE == 0) + simd_prefixes[lid / METAL_SIMD_SIZE] = higher + sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (lid < METAL_SIMD_SIZE) { + const uint simdgroups = THREADS_PER_TG / METAL_SIMD_SIZE; + const uint simd_sum = lid < simdgroups ? simd_prefixes[lid] : 0; + const uint simd_prefix = simd_prefix_exclusive_sum(simd_sum); + if (lid < simdgroups) + simd_prefixes[lid] = simd_prefix; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + higher += simd_prefixes[lid / METAL_SIMD_SIZE]; + + uint rank = pass == 0 ? k - 1 : ranks[row]; + if (higher <= rank && rank < higher + sum) { + rank -= higher; + const uint group_begin = BUCKETS - (lid + 1) * BUCKETS_PER_THREAD; + for (uint offset = 0; offset < BUCKETS_PER_THREAD; ++offset) { + const uint bucket = group_begin + BUCKETS_PER_THREAD - offset - 1; + const uint count = histogram[bucket]; + if (rank < count) { + prefixes[row] |= ulong(bucket) << shift; + prefix_masks[row] |= ulong(BUCKETS - 1) << shift; + break; + } + rank -= count; + } + ranks[row] = rank; + } + threadgroup_barrier(mem_flags::mem_device); + if (lid == 0) + reset_arrival(&done_counts[row]); +} + +KERNEL(RadixTopKSmallCollect)( + const device float* input, + device uint* output_ids, + device float* output_scores, + const device ulong* prefixes, + device atomic_uint* selected_counts, + device atomic_uint* done_counts, + device ulong* selected_keys, + constant uint& rows, + constant uint& k, + const uint columns SPECIALIZE, + const uint partitions SPECIALIZE, + threadgroup ulong sorted_keys[MAX_K], + threadgroup uint& is_last, + const uint group GROUPS(rows* partitions), + const uint lid THREADS(THREADS_PER_TG) +) { + const uint row = group / partitions; + const uint partition = group % partitions; + const uint index_bits = columns <= 1 ? 1u : 32u - clz(columns - 1u); + const device float* row_input = input + ulong(row) * columns; + const ulong prefix = prefixes[row]; + + visit_partition_keys(row_input, columns, index_bits, partition, partitions, lid, [&](ulong key) { + if (key >= prefix) { + const uint index = atomic_fetch_add_explicit(&selected_counts[row], 1u, memory_order_relaxed); + // The final radix prefix selects exactly k unique keys. + selected_keys[row * k + index] = key; + } + }); + threadgroup_barrier(mem_flags::mem_device); + if (lid == 0) + is_last = arrive_last(&done_counts[row], partitions); + threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device); + if (!is_last) + return; + + for (uint index = lid; index < MAX_K; index += THREADS_PER_TG) + sorted_keys[index] = index < k ? selected_keys[row * k + index] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint size = 2; size <= MAX_K; size <<= 1) { + for (uint stride = size >> 1; stride; stride >>= 1) { + for (uint index = lid; index < MAX_K; index += THREADS_PER_TG) { + const uint other = index ^ stride; + if (other > index) { + const bool descending = (index & size) == 0; + if (descending ? sorted_keys[index] < sorted_keys[other] : sorted_keys[index] > sorted_keys[other]) { + const ulong key = sorted_keys[index]; + sorted_keys[index] = sorted_keys[other]; + sorted_keys[other] = key; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + for (uint index = lid; index < k; index += THREADS_PER_TG) { + const ulong key = sorted_keys[index]; + const ulong mask = (1ul << index_bits) - 1; + output_ids[row * k + index] = uint(mask - (key & mask)); + output_scores[row * k + index] = top_k_score_from_key(uint(key >> index_bits)); + } +} diff --git a/crates/backend-uzu/src/backends/metal/kernel/radix_top_k_small.rs b/crates/backend-uzu/src/backends/metal/kernel/radix_top_k_small.rs new file mode 100644 index 000000000..f93a8441a --- /dev/null +++ b/crates/backend-uzu/src/backends/metal/kernel/radix_top_k_small.rs @@ -0,0 +1,93 @@ +use std::mem::size_of; + +use super::{RadixTopKSmallCollectMetalKernel, RadixTopKSmallPassMetalKernel}; +use crate::backends::{ + common::{ + Allocation, Encoder, + kernel::radix_top_k_small::{MAX_K, RadixTopKSmall}, + }, + metal::{Metal, context::MetalContext, error::MetalError}, +}; + +const PARTITIONS: u32 = 4; +const RADIX_BITS: u32 = 10; +const RADIX_BUCKETS: usize = 1 << RADIX_BITS; + +pub struct MetalRadixTopKSmall { + pass: RadixTopKSmallPassMetalKernel, + collect: RadixTopKSmallCollectMetalKernel, + columns: u32, + passes: u32, +} + +impl RadixTopKSmall for MetalRadixTopKSmall { + fn new( + context: &MetalContext, + columns: u32, + ) -> Result { + assert!(columns > 0); + let index_bits = if columns <= 1 { + 1 + } else { + u32::BITS - (columns - 1).leading_zeros() + }; + Ok(Self { + pass: RadixTopKSmallPassMetalKernel::new(context, columns, PARTITIONS)?, + collect: RadixTopKSmallCollectMetalKernel::new(context, columns, PARTITIONS)?, + columns, + passes: (u32::BITS + index_bits).div_ceil(RADIX_BITS), + }) + } + + fn encode( + &self, + input: &Allocation, + output_ids: &mut Allocation, + output_scores: &mut Allocation, + rows: u32, + k: u32, + encoder: &mut Encoder, + ) -> Result<(), MetalError> { + assert!(rows > 0 && k > 0 && k <= MAX_K && k <= self.columns); + let rows = rows as usize; + let mut histograms = encoder.allocate_scratch(rows * PARTITIONS as usize * RADIX_BUCKETS * size_of::())?; + let mut prefixes = encoder.allocate_scratch(rows * size_of::())?; + let mut masks = encoder.allocate_scratch(rows * size_of::())?; + let mut ranks = encoder.allocate_scratch(rows * size_of::())?; + let mut counts = encoder.allocate_scratch(rows * size_of::())?; + let mut done = encoder.allocate_scratch(rows * size_of::())?; + let mut keys = encoder.allocate_scratch(rows * k as usize * size_of::())?; + encoder.encode_fill(&mut prefixes, 0); + encoder.encode_fill(&mut masks, 0); + encoder.encode_fill(&mut counts, 0); + encoder.encode_fill(&mut done, 0); + + for pass in 0..self.passes { + self.pass.encode( + input, + &mut histograms, + &mut prefixes, + &mut masks, + &mut ranks, + &mut done, + rows as u32, + k, + pass, + encoder, + ); + } + self.collect.encode( + input, + output_ids, + output_scores, + &prefixes, + &mut counts, + &mut done, + &mut keys, + rows as u32, + k, + encoder, + ); + Ok(()) + } +} diff --git a/crates/backend-uzu/unit/backends/common/kernel/mod.rs b/crates/backend-uzu/unit/backends/common/kernel/mod.rs index 5ac3ed84b..89b8930d1 100644 --- a/crates/backend-uzu/unit/backends/common/kernel/mod.rs +++ b/crates/backend-uzu/unit/backends/common/kernel/mod.rs @@ -10,6 +10,7 @@ mod matmul; mod moe; mod pooling; mod qkv_norm_test; +mod radix_top_k_small_test; mod short_conv; mod ssm; mod tensor_add_bias_test; diff --git a/crates/backend-uzu/unit/backends/common/kernel/radix_top_k_small_test.rs b/crates/backend-uzu/unit/backends/common/kernel/radix_top_k_small_test.rs new file mode 100644 index 000000000..dbc0e5a45 --- /dev/null +++ b/crates/backend-uzu/unit/backends/common/kernel/radix_top_k_small_test.rs @@ -0,0 +1,142 @@ +#[cfg(metal_backend)] +use std::time::Instant; + +use proc_macros::uzu_test; +use test_runner::for_each_non_cpu_backend; + +#[cfg(metal_backend)] +use crate::backends::metal::Metal; +use crate::{ + backends::{ + common::{Backend, Encoder, Kernels, kernel::radix_top_k_small::RadixTopKSmall}, + cpu::Cpu, + }, + tests::helpers::{alloc_allocation, alloc_allocation_with_data, allocation_to_vec, create_context}, +}; + +const TARGET_COLUMNS: usize = 248_320; +const TARGET_K: usize = 512; + +fn values( + rows: usize, + columns: usize, +) -> Vec { + (0..rows * columns).map(|index| ((index * 37 % 101) as f32).sin()).collect() +} + +fn radix_top_k_small( + input: &[f32], + rows: usize, + columns: usize, + k: usize, +) -> (Vec, Vec) { + let context = create_context::(); + let input = alloc_allocation_with_data::(&context, input); + let mut ids = alloc_allocation::(&context, rows * k); + let mut scores = alloc_allocation::(&context, rows * k); + let kernel = ::RadixTopKSmall::new(&context, columns as u32).unwrap(); + let mut encoder = Encoder::new(context.as_ref()).unwrap(); + kernel.encode(&input, &mut ids, &mut scores, rows as u32, k as u32, &mut encoder).unwrap(); + encoder.end_encoding().submit().wait_until_completed().unwrap(); + (allocation_to_vec(&ids), allocation_to_vec(&scores)) +} + +fn reference( + input: &[f32], + rows: usize, + columns: usize, + k: usize, +) -> (Vec, Vec) { + let mut ids = Vec::with_capacity(rows * k); + let mut scores = Vec::with_capacity(rows * k); + for values in input.chunks_exact(columns) { + let mut row = (0..columns).collect::>(); + row.sort_unstable_by(|&left, &right| values[right].total_cmp(&values[left]).then_with(|| left.cmp(&right))); + for index in row.into_iter().take(k) { + ids.push(index as u32); + scores.push(values[index]); + } + } + (ids, scores) +} + +fn assert_output( + actual: &(Vec, Vec), + expected: &(Vec, Vec), + shape: (usize, usize, usize), +) { + assert_eq!(actual.0, expected.0, "shape={shape:?}"); + assert!( + actual.1.iter().map(|value| value.to_bits()).eq(expected.1.iter().map(|value| value.to_bits())), + "shape={shape:?}", + ); +} + +#[uzu_test] +fn radix_top_k_small_matches_cpu() { + for shape @ (rows, columns, k) in [ + (1, 1, 1), + (1, 8, 8), + (3, 1025, 1), + (3, 1025, TARGET_K - 1), + (3, 1025, TARGET_K), + (15, TARGET_COLUMNS, TARGET_K), + ] { + let mut input = values(rows, columns); + let special = [ + f32::INFINITY, + f32::NEG_INFINITY, + -0.0, + 0.0, + 1.0, + 1.0, + f32::from_bits(0x7fc0_0001), + f32::from_bits(0xffc0_0001), + ]; + let special_count = special.len().min(input.len()); + input[..special_count].copy_from_slice(&special[..special_count]); + let expected = reference(&input, rows, columns, k); + assert_output(&radix_top_k_small::(&input, rows, columns, k), &expected, shape); + for_each_non_cpu_backend!(|B| { + let actual = radix_top_k_small::(&input, rows, columns, k); + assert_output(&actual, &expected, shape); + if columns == TARGET_COLUMNS { + for _ in 0..2 { + assert_output(&radix_top_k_small::(&input, rows, columns, k), &expected, shape); + } + } + }); + } +} + +#[cfg(metal_backend)] +#[uzu_test] +#[ignore = "benchmark"] +fn benchmark_radix_top_k_small() { + const ROWS: usize = 15; + const SAMPLES: usize = 50; + const BATCH: u32 = 16; + + let context = create_context::(); + let input = alloc_allocation_with_data::(&context, &values(ROWS, TARGET_COLUMNS)); + let mut ids = alloc_allocation::(&context, ROWS * TARGET_K); + let mut scores = alloc_allocation::(&context, ROWS * TARGET_K); + let kernel = + <::Kernels as Kernels>::RadixTopKSmall::new(&context, TARGET_COLUMNS as u32).unwrap(); + let mut run = || { + let start = Instant::now(); + let mut encoder = Encoder::new(context.as_ref()).unwrap(); + for _ in 0..BATCH { + kernel.encode(&input, &mut ids, &mut scores, ROWS as u32, TARGET_K as u32, &mut encoder).unwrap(); + } + let completed = encoder.end_encoding().submit().wait_until_completed().unwrap(); + (completed.gpu_execution_time().div_f64(BATCH as f64), start.elapsed().div_f64(BATCH as f64)) + }; + for _ in 0..5 { + run(); + } + let (mut gpu, mut wall): (Vec<_>, Vec<_>) = (0..SAMPLES).map(|_| run()).unzip(); + gpu.sort_unstable(); + wall.sort_unstable(); + eprintln!("radix_top_k_small gpu={:?} wall={:?}", gpu[SAMPLES / 2], wall[SAMPLES / 2]); +}