yum/gpu_fft

git clone https://git.yummers.dev/yum/gpu_fft

yumAdd Pass constructcfe94cd2026-09-14

use num::Complex;
use std::borrow::Cow;
use std::num::NonZeroU64;
use std::sync::mpsc::{Receiver, Sender};
use wgpu::naga;

pub struct GpuShit {
    adapter: wgpu::Adapter,
    device: wgpu::Device,
    queue: wgpu::Queue,
}

impl GpuShit {
    pub fn get() -> anyhow::Result<GpuShit> {
        env_logger::init();
        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());
        let adapter =
            pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default()))?;
        println!("Got adapter: {:#?}", adapter.get_info());

        let (device, queue) =
            pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
                label: None,
                required_features: wgpu::Features::empty(),
                required_limits: wgpu::Limits::downlevel_defaults(),
                #[cfg(target_arch = "wasm32")]
                default_queue: wgpu::QueueDescriptor { label: None },
                experimental_features: wgpu::ExperimentalFeatures::disabled(),
                memory_hints: wgpu::MemoryHints::MemoryUsage,
                trace: wgpu::Trace::Off,
            }))?;

        Ok(GpuShit {
            adapter,
            device,
            queue,
        })
    }
}

fn get_entry_point_name(module: &naga::Module) -> anyhow::Result<String> {
    if module.entry_points.len() != 1 {
        anyhow::bail!(
            "Expected exactly one entry point, found {}",
            module.entry_points.len()
        );
    }
    for entry in module.entry_points.iter() {
        return Ok(entry.name.clone());
    }
    anyhow::bail!("Could not find entry point");
}

fn get_entry_point_workgroup_size(module: &naga::Module, name: &str) -> anyhow::Result<[u32; 3]> {
    for entry_point in module.entry_points.iter() {
        if entry_point.name != name {
            continue;
        }
        return Ok(entry_point.workgroup_size);
    }
    anyhow::bail!("Could not find entry point named {}", name);
}

fn get_bind_group_layout(
    gpu: &GpuShit,
    module: &naga::Module,
    module_info: &naga::valid::ModuleInfo,
    module_layout: &naga::proc::Layouter,
) -> anyhow::Result<wgpu::BindGroupLayout> {
    let mut entries = Vec::<wgpu::BindGroupLayoutEntry>::new();
    for (global_handle, global) in module.global_variables.iter() {
        let Some(binding) = &global.binding else {
            continue;
        };

        let can_write = match global.space {
            naga::AddressSpace::Storage { access } => {
                access.contains(naga::StorageAccess::STORE)
                    || access.contains(naga::StorageAccess::ATOMIC)
            }
            _ => {
                eprintln!(
                    "Skipping global {:?} with address space {:?}",
                    global.name, global.space
                );
                continue;
            }
        };

        let mut stages = wgpu::ShaderStages::empty();
        for (entry_idx, entry) in module.entry_points.iter().enumerate() {
            let usage = module_info.get_entry_point(entry_idx)[global_handle];
            if usage.is_empty() {
                continue;
            }

            stages |= match entry.stage {
                naga::ShaderStage::Vertex => wgpu::ShaderStages::VERTEX,
                naga::ShaderStage::Fragment => wgpu::ShaderStages::FRAGMENT,
                naga::ShaderStage::Compute => wgpu::ShaderStages::COMPUTE,
                _ => continue,
            }
        }
        if stages.is_empty() {
            eprintln!(
                "Skipping bound global {:?} which is not referenced by any stage",
                global.name
            );
            continue;
        }

        let global_ty = &module.types[global.ty];
        let entry = match module.types[global.ty].inner {
            naga::ir::TypeInner::Array { .. } => {
                let min_size = module_layout[global.ty].size;
                wgpu::BindGroupLayoutEntry {
                    binding: binding.binding,
                    visibility: stages,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage {
                            read_only: !can_write,
                        },
                        min_binding_size: NonZeroU64::new(min_size as u64),
                        has_dynamic_offset: false,
                    },
                    count: None,
                }
            }
            _ => {
                anyhow::bail!(format!(
                    "Unsupported global type: {:?} (name={:?})",
                    global_ty.name, global_ty.inner
                ))
            }
        };
        entries.push(entry);
    }
    let layout = wgpu::BindGroupLayoutDescriptor {
        label: None,
        entries: &entries,
    };
    Ok(gpu.device.create_bind_group_layout(&layout))
}

pub struct ShaderShit {
    bind_group_layout: wgpu::BindGroupLayout,
    shader_module: wgpu::ShaderModule,
    entry_point_name: String,
    entry_point_workgroup_size: [u32; 3],
}

impl ShaderShit {
    pub fn new(gpu: &GpuShit, shader: &str) -> anyhow::Result<Self> {
        let module = naga::front::wgsl::parse_str(shader)?;
        let module_info = naga::valid::Validator::new(
            naga::valid::ValidationFlags::all(),
            naga::valid::Capabilities::all(),
        )
        .validate(&module)?;
        let mut module_layout = naga::proc::Layouter::default();
        module_layout.update(module.to_ctx())?;

        Self::ensure_compute_compatibility(gpu, &module)?;
        let bind_group_layout = get_bind_group_layout(gpu, &module, &module_info, &module_layout)?;
        let entry_point_name = get_entry_point_name(&module)?;
        let entry_point_workgroup_size =
            get_entry_point_workgroup_size(&module, &entry_point_name)?;

        let module_descriptor = wgpu::ShaderModuleDescriptor {
            label: None,
            source: wgpu::ShaderSource::Naga(Cow::Owned(module)),
        };
        let shader_module = gpu.device.create_shader_module(module_descriptor);

        Ok(Self {
            bind_group_layout,
            shader_module,
            entry_point_name,
            entry_point_workgroup_size,
        })
    }

    fn ensure_compute_compatibility(gpu: &GpuShit, module: &naga::Module) -> anyhow::Result<()> {
        let mut has_compute_stage = false;
        for entry_point in module.entry_points.iter() {
            if entry_point.stage == naga::ShaderStage::Compute {
                has_compute_stage = true;
            }
        }
        if !has_compute_stage {
            return Ok(());
        }
        let downlevel_caps = gpu.adapter.get_downlevel_capabilities();
        if !downlevel_caps
            .flags
            .contains(wgpu::DownlevelFlags::COMPUTE_SHADERS)
        {
            anyhow::bail!(
                "GPU {} does not support compute shaders.",
                gpu.adapter.get_info().name
            )
        }
        Ok(())
    }
}

#[derive(Clone)]
pub struct MaterialPass {
    buffers: Vec<wgpu::Buffer>,
    bind_group: wgpu::BindGroup,
}

impl MaterialPass {
    pub fn new(gpu: &GpuShit, shader: &ShaderShit, data: &[Complex<f32>]) -> Self {
        let input_data_buffer = gpu.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: std::mem::size_of_val(data) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let output_data_buffer = gpu.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: input_data_buffer.size(),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        let download_buffer = gpu.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: input_data_buffer.size(),
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
            mapped_at_creation: false,
        });
        let bind_group = gpu.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &shader.bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: input_data_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: output_data_buffer.as_entire_binding(),
                },
            ],
        });
        Self {
            buffers: vec![input_data_buffer, output_data_buffer, download_buffer],
            bind_group: bind_group,
        }
    }

    pub fn from_buffers(gpu: &GpuShit, shader: &ShaderShit, buffers: Vec<wgpu::Buffer>) -> Self {
        let descriptor_entries =
            buffers
                .iter()
                .enumerate()
                .fold(Vec::new(), |mut entries, (index, buf)| {
                    entries.push(wgpu::BindGroupEntry {
                        binding: index as u32,
                        resource: buf.as_entire_binding(),
                    });
                    entries
                });
        let descriptor_seed = wgpu::BindGroupDescriptor {
            label: None,
            layout: &shader.bind_group_layout,
            entries: &descriptor_entries,
        };
        let bind_group = gpu.device.create_bind_group(&descriptor_seed);
        Self {
            buffers: buffers,
            bind_group: bind_group,
        }
    }
}

struct MaterialInstance {
    job_id: u64,
    busy: bool,
    passes: Vec<MaterialPass>,
}

impl MaterialInstance {
    pub fn new(passes: Vec<MaterialPass>) -> Self {
        Self {
            job_id: 0, // Placeholder. Job ID is assigned when data is submitted.
            busy: false,
            passes: passes,
        }
    }
}

struct MaterialShit {
    slots: Vec<MaterialInstance>,
    next_job_id: u64,
    ready_tx: Sender<(usize, u64, Result<(), wgpu::BufferAsyncError>)>,
    ready_rx: Receiver<(usize, u64, Result<(), wgpu::BufferAsyncError>)>,
}

impl MaterialShit {
    #[allow(dead_code)]
    pub fn new(gpu: &GpuShit, shader: &ShaderShit, data: &[Complex<f32>], n_slots: usize) -> Self {
        let mut slots = Vec::<MaterialInstance>::new();
        for _ in 0..n_slots {
            let pass = MaterialPass::new(gpu, shader, data);
            slots.push(MaterialInstance::new(vec![pass.clone()]));
        }
        let (ready_tx, ready_rx) = std::sync::mpsc::channel();
        Self {
            slots: slots,
            next_job_id: 0,
            ready_tx: ready_tx,
            ready_rx: ready_rx,
        }
    }

    pub fn get_instance(&mut self) -> Option<(&mut MaterialInstance, usize, u64)> {
        let (slot_idx, slot) = self
            .slots
            .iter_mut()
            .enumerate()
            .find(|(_, slot)| !slot.busy)?;
        let job_id = self.next_job_id;
        slot.busy = true;
        slot.job_id = self.next_job_id;
        self.next_job_id += 1;
        Some((slot, slot_idx, job_id))
    }
}

#[allow(dead_code)]
fn get_pipeline(gpu: &GpuShit, shader: &ShaderShit) -> wgpu::ComputePipeline {
    let pipeline_layout = gpu
        .device
        .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: None,
            bind_group_layouts: &[Some(&shader.bind_group_layout)],
            immediate_size: 0,
        });

    gpu.device
        .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: None,
            layout: Some(&pipeline_layout),
            module: &shader.shader_module,
            entry_point: Some(&shader.entry_point_name),
            compilation_options: wgpu::PipelineCompilationOptions::default(),
            cache: None,
        })
}

#[allow(dead_code)]
fn submit_dft_batch(
    gpu: &GpuShit,
    shader: &ShaderShit,
    material: &mut MaterialShit,
    pipeline: &wgpu::ComputePipeline,
    data_batch: &[Vec<Complex<f32>>],
) -> anyhow::Result<usize> {
    // Submit data to the GPU for each input in the batch.
    let mut jobs = Vec::with_capacity(data_batch.len());
    for data in data_batch {
        let Some((slot, slot_idx, job_id)) = material.get_instance() else {
            break;
        };
        gpu.queue
            .write_buffer(&slot.passes[0].buffers[0], 0, bytemuck::cast_slice(data));
        jobs.push((slot_idx, job_id, data.len()));
    }
    if jobs.is_empty() {
        return Ok(0);
    }

    let mut encoder = gpu
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
    let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
        label: None,
        timestamp_writes: None,
    });
    compute_pass.set_pipeline(pipeline);
    let shader_local_size = shader
        .entry_point_workgroup_size
        .iter()
        .fold(1, |acc, &x| acc * x) as usize;
    for &(slot_idx, _, len) in &jobs {
        compute_pass.set_bind_group(0, &material.slots[slot_idx].passes[0].bind_group, &[]);
        compute_pass.dispatch_workgroups(len.div_ceil(shader_local_size) as u32, 1, 1);
    }
    drop(compute_pass);

    for &(slot_idx, job_id, _) in &jobs {
        let slot = &material.slots[slot_idx];
        encoder.copy_buffer_to_buffer(
            &slot.passes[0].buffers[1],
            0,
            &slot.passes[0].buffers[2],
            0,
            slot.passes[0].buffers[1].size(),
        );
        let ready_tx = material.ready_tx.clone();
        encoder.map_buffer_on_submit(
            &slot.passes[0].buffers[2],
            wgpu::MapMode::Read,
            ..,
            move |result| {
                ready_tx.send((slot_idx, job_id, result)).ok();
            },
        );
    }

    gpu.queue.submit([encoder.finish()]);
    Ok(jobs.len())
}

#[allow(dead_code)]
fn drain_dft(
    gpu: &GpuShit,
    material: &mut MaterialShit,
) -> anyhow::Result<Vec<(u64, Vec<Complex<f32>>)>> {
    gpu.device.poll(wgpu::PollType::Poll)?;
    let mut results = Vec::new();
    while let Ok((slot_index, job_id, map_result)) = material.ready_rx.try_recv() {
        let slot = &mut material.slots[slot_index];
        let data = slot.passes[0].buffers[2].get_mapped_range(..)?;
        let result = bytemuck::cast_slice(&data).to_vec();
        drop(data);

        slot.passes[0].buffers[2].unmap();
        slot.busy = false;

        map_result?;
        results.push((job_id, result));
    }
    Ok(results)
}

#[cfg(test)]
mod tests {
    use std::time::Duration;
    use std::time::Instant;

    use num::Complex;
    use num::traits::Float;
    use rand::RngExt;
    use rand::SeedableRng;
    use rand::rngs::StdRng;

    // Cast one float type to another, truncating if needed.
    fn t0_to_t1<T0: Float, T1: Float>(val: T0) -> T1 {
        return num::cast(val).unwrap();
    }

    // Casts one Complex<T> to another.
    fn complex_to_t<T0: Float, T1: Float>(data: &[Complex<T0>]) -> Vec<Complex<T1>> {
        data.iter()
            .map(|c| Complex::new(t0_to_t1(c.re), t0_to_t1(c.im)))
            .collect()
    }

    fn evaluate_results(
        result_ref_64: &Vec<Complex<f64>>,
        result_cur_64: &Vec<Complex<f64>>,
        duration_dft: std::time::Duration,
        duration_cur: std::time::Duration,
        algo_name: &str,
    ) {
        assert_eq!(result_ref_64.len(), result_cur_64.len());

        println!("  Algorithm:  {}", algo_name);
        println!("    Duration:  {:?}", duration_cur);
        println!(
            "    Speedup:   {:?}",
            duration_dft.as_nanos() as f32 / duration_cur.as_nanos() as f32
        );

        let mut max_err: f64 = 0.0;
        let mut sum_err: f64 = 0.0;
        // TODO: median, 90p, 99p

        let result_64: Vec<Complex<f64>> = complex_to_t(&result_cur_64);

        for i in 0..result_ref_64.len() {
            max_err = max_err.max((result_ref_64[i] - result_64[i]).norm());
            sum_err += (result_ref_64[i] - result_64[i]).norm();
        }
        println!("    Max err:   {}", max_err);
        println!("    Avg err:   {}", sum_err / (result_ref_64.len() as f64));
    }

    fn evaluate_naive_dft(data: &[Complex<f32>], gpu: &super::GpuShit) {
        let shader =
            super::ShaderShit::new(&gpu, include_str!("generated/shaders/naive_dft.wgsl")).unwrap();
        let mut result_dft = Vec::from(data);

        const BATCH_SIZE: usize = 64;
        let inputs = vec![Vec::from(data); BATCH_SIZE];
        let mut material = super::MaterialShit::new(&gpu, &shader, &result_dft, BATCH_SIZE * 2);

        let pipeline = super::get_pipeline(&gpu, &shader);

        let warmup = Instant::now() + Duration::from_secs(3);
        let end = warmup + Duration::from_secs(5);
        let mut n_iter: u32 = 0;
        let mut now = Instant::now();
        while now < end {
            // Fill the pipeline.
            if super::submit_dft_batch(&gpu, &shader, &mut material, &pipeline, &inputs).unwrap()
                > 0
            {
                continue;
            }
            // Drain if pipeline is full.
            let results = super::drain_dft(&gpu, &mut material).unwrap();
            for (_, result) in &results {
                result_dft = result.clone();
            }
            if now > warmup {
                n_iter += results.len() as u32;
            }
            now = Instant::now();
        }
        let duration_dft = (now - warmup) / n_iter;
        println!("Duration: {:?}", duration_dft);

        let mut result_ref: Vec<rustfft::num_complex::Complex<f64>> = complex_to_t(&data);
        let mut ref_planner = rustfft::FftPlanner::<f64>::new();
        let ref_fft = ref_planner.plan_fft_forward(result_ref.len());
        ref_fft.process(&mut result_ref);
        let result_ref_64: Vec<Complex<f64>> = complex_to_t(&result_ref);

        let result_dft_64 = complex_to_t::<f32, f64>(&result_dft);
        evaluate_results(
            &result_ref_64,
            &result_dft_64,
            duration_dft,
            duration_dft,
            "Naive DFT",
        );
    }

    #[test]
    fn run_it() {
        let mut data_f64: Vec<Complex<f64>> = Vec::new();
        data_f64.resize((2 as usize).pow(12), Complex::ZERO);
        // Randomize data.
        let mut rng = StdRng::seed_from_u64(42);
        for i in 0..data_f64.len() {
            data_f64[i] = Complex::new(rng.random(), rng.random());
        }
        let data_f32: Vec<Complex<f32>> = complex_to_t::<f64, f32>(&data_f64);

        let gpu = super::GpuShit::get().unwrap();

        evaluate_naive_dft(&data_f32, &gpu);
    }
}