ferrotensor_backend/engine/
misc.rs

1//! Top-k selection and depthwise 1-D convolution.
2
3use super::{groups64, uniform, ComputeEngine};
4use crate::{UnifiedTensor, WgpuContext};
5
6impl ComputeEngine {
7    /// Writes the indices of the `k` largest values of `input`, largest first,
8    /// as `u32`s into `out_idx`. Runs on a single GPU thread; meant for sampling
9    /// from one logits vector.
10    pub fn top_k(&self, ctx: &WgpuContext, input: &UnifiedTensor, out_idx: &UnifiedTensor, k: u32) {
11        let params = uniform(ctx, bytemuck::cast_slice(&[k]));
12        self.run(
13            ctx,
14            &self.top_k,
15            &[(0, input.buffer()), (1, out_idx.buffer()), (2, &params)],
16            (1, 1, 1),
17        );
18    }
19
20    /// Depthwise 1-D convolution with kernel size 3 and padding 1:
21    /// `out[d, s] = bias[d] + sum_k weight[d, k] * input[d, s + k - 1]`, with
22    /// `input`/`out` `[dim, seq_len]`, `weight` `[dim, 3]` and `bias` `[dim]`.
23    pub fn conv1d(
24        &self,
25        ctx: &WgpuContext,
26        input: &UnifiedTensor,
27        weight: &UnifiedTensor,
28        bias: &UnifiedTensor,
29        out: &UnifiedTensor,
30    ) {
31        let dim = weight.shape[0];
32        let seq_len = input.shape[input.shape.len() - 1];
33        let params = uniform(ctx, bytemuck::cast_slice(&[dim as u32, seq_len as u32]));
34        self.run(
35            ctx,
36            &self.conv1d,
37            &[
38                (0, input.buffer()),
39                (1, weight.buffer()),
40                (2, out.buffer()),
41                (3, bias.buffer()),
42                (4, &params),
43            ],
44            groups64(dim * seq_len),
45        );
46    }
47}
ferrotensor_backend/engine/misc.rs

ferrotensor_backend/engine/
misc.rs
