rlx-wgpu 0.2.6

Cross-platform GPU backend for RLX via wgpu (Metal/Vulkan/DX12/WebGPU)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.

use crate::buffer::Arena;
use crate::kernels::{WelchPeaksGpuParams, welch_peaks_gpu_kernel};

fn dispatch_dims(n: u32, wg: u32) -> (u32, u32, u32) {
    (n.div_ceil(wg).max(1), 1, 1)
}

/// Dispatch native GPU WelchPeaks inside an existing compute pass.
pub fn dispatch_welch_peaks_gpu_in_pass(
    pass: &mut wgpu::ComputePass<'_>,
    device: &wgpu::Device,
    uniform: &wgpu::Buffer,
    bind_group: &wgpu::BindGroup,
    welch_batch: u32,
) {
    let k = welch_peaks_gpu_kernel(device);
    pass.set_pipeline(&k.pipeline);
    pass.set_bind_group(0, bind_group, &[]);
    let _ = uniform;
    let (gx, gy, gz) = dispatch_dims(welch_batch, 64);
    pass.dispatch_workgroups(gx, gy, gz);
}

/// Standalone WelchPeaks GPU dispatch (submit + poll).
pub fn run_welch_peaks_gpu(
    _arena: &Arena,
    device: &wgpu::Device,
    queue: &wgpu::Queue,
    uniform: &wgpu::Buffer,
    bind_group: &wgpu::BindGroup,
    welch_batch: u32,
) {
    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
        label: Some("rlx-wgpu welch_peaks_gpu"),
    });
    {
        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
            label: Some("welch_peaks_gpu"),
            timestamp_writes: None,
        });
        dispatch_welch_peaks_gpu_in_pass(&mut pass, device, uniform, bind_group, welch_batch);
    }
    queue.submit(std::iter::once(encoder.finish()));
}

pub fn welch_peaks_gpu_params(
    spec_off: u32,
    dst_off: u32,
    welch_batch: u32,
    n_fft: u32,
    n_segments: u32,
    k: u32,
    n_bins: u32,
) -> WelchPeaksGpuParams {
    WelchPeaksGpuParams {
        spec_off,
        dst_off,
        welch_batch,
        n_fft,
        n_segments,
        k,
        n_bins,
        _p0: 0,
        _p1: 0,
    }
}