yum/vfft

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

yumMetal optimization: use 4 batches instead of 22f018df

master
branchesmasterall refs →
60.3 KiB1861 linesraw
1use naga::GlobalVariable;
2use num::Complex;
3use std::borrow::Cow;
4use std::num::NonZeroU64;
5use std::sync::mpsc::{Receiver, Sender};
6use wgpu::naga;
7
8#[cfg(feature = "gpu")]
9#[macro_export]
10macro_rules! include_wgsl {
11    ($filename:literal) => {
12        include_str!(concat!(env!("OUT_DIR"), "/shaders/", $filename))
13    };
14}
15
16// Magic to allow converting a primitive type, a Vec, or a slice into a slice of bytes.
17pub trait Buffer<T> {
18    fn as_bytes(&self) -> &[T];
19}
20impl<T: bytemuck::Pod> Buffer<T> for T {
21    fn as_bytes(&self) -> &[T] {
22        std::slice::from_ref(self)
23    }
24}
25impl<T> Buffer<T> for Vec<T> {
26    fn as_bytes(&self) -> &[T] {
27        self.as_slice()
28    }
29}
30impl<T> Buffer<T> for &[T] {
31    fn as_bytes(&self) -> &[T] {
32        self
33    }
34}
35
36pub struct BufferBuilder<'a> {
37    label: Option<String>,
38    size: u64,
39    usage: wgpu::BufferUsages,
40    mapped_at_creation: bool,
41    gpu: &'a Gpu,
42}
43
44impl<'a> BufferBuilder<'a> {
45    fn new<T: bytemuck::Pod>(data: &impl Buffer<T>, gpu: &'a Gpu) -> Self {
46        Self::new_by_sz(std::mem::size_of_val(data.as_bytes()), gpu)
47    }
48
49    fn new_by_sz(sz: usize, gpu: &'a Gpu) -> Self {
50        Self {
51            label: None,
52            size: wgpu::util::align_to(sz, 16) as u64,
53            usage: wgpu::BufferUsages::empty(),
54            mapped_at_creation: false,
55            gpu: gpu,
56        }
57    }
58
59    pub fn copy_dst(&mut self) -> &mut Self {
60        self.usage |= wgpu::BufferUsages::COPY_DST;
61        self
62    }
63
64    pub fn copy_src(&mut self) -> &mut Self {
65        self.usage |= wgpu::BufferUsages::COPY_SRC;
66        self
67    }
68
69    pub fn map_read(&mut self) -> &mut Self {
70        self.usage |= wgpu::BufferUsages::MAP_READ;
71        self
72    }
73
74    pub fn query_resolve(&mut self) -> &mut Self {
75        self.usage |= wgpu::BufferUsages::QUERY_RESOLVE;
76        self
77    }
78
79    pub fn storage(&mut self) -> &mut Self {
80        self.usage |= wgpu::BufferUsages::STORAGE;
81        self
82    }
83
84    pub fn uniform(&mut self) -> &mut Self {
85        self.usage |= wgpu::BufferUsages::UNIFORM;
86        self
87    }
88
89    pub fn build(&self) -> wgpu::Buffer {
90        self.gpu.device.create_buffer(&wgpu::BufferDescriptor {
91            label: self.label.as_deref(),
92            size: self.size,
93            usage: self.usage,
94            mapped_at_creation: self.mapped_at_creation,
95        })
96    }
97}
98
99pub struct GpuBuilder {
100    required_features: wgpu::Features,
101}
102
103impl GpuBuilder {
104    pub fn new() -> Self {
105        Self {
106            required_features: wgpu::Features::empty(),
107        }
108    }
109
110    pub fn enable_profiling(&mut self) -> &mut Self {
111        self.required_features |= wgpu::Features::TIMESTAMP_QUERY;
112        self
113    }
114
115    pub fn build(&self) -> anyhow::Result<Gpu> {
116        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());
117        let adapter =
118            pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default()))?;
119        println!("Got adapter: {:#?}", adapter.get_info());
120
121        let supported = adapter.limits();
122        let mut limits = wgpu::Limits::downlevel_defaults();
123        limits.max_compute_invocations_per_workgroup =
124            supported.max_compute_invocations_per_workgroup;
125        limits.max_compute_workgroup_storage_size = supported.max_compute_workgroup_storage_size;
126        println!(
127            "Max compute invocations per workgroup: {}",
128            limits.max_compute_invocations_per_workgroup
129        );
130
131        let (device, queue) =
132            pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
133                label: None,
134                required_features: self.required_features,
135                required_limits: limits,
136                #[cfg(target_arch = "wasm32")]
137                default_queue: wgpu::QueueDescriptor { label: None },
138                experimental_features: wgpu::ExperimentalFeatures::disabled(),
139                memory_hints: wgpu::MemoryHints::MemoryUsage,
140                trace: wgpu::Trace::Off,
141            }))?;
142
143        let mut timestamps = None;
144        if self
145            .required_features
146            .contains(wgpu::Features::TIMESTAMP_QUERY)
147        {
148            timestamps = Some(device.create_query_set(&wgpu::QuerySetDescriptor {
149                label: None,
150                ty: wgpu::QueryType::Timestamp,
151                count: 2,
152            }));
153        }
154
155        Ok(Gpu {
156            adapter,
157            device,
158            queue,
159            timestamps,
160        })
161    }
162}
163
164pub struct Gpu {
165    adapter: wgpu::Adapter,
166    device: wgpu::Device,
167    queue: wgpu::Queue,
168    timestamps: Option<wgpu::QuerySet>,
169}
170
171impl Gpu {
172    pub fn buffer<T: bytemuck::Pod>(&self, data: &impl Buffer<T>) -> BufferBuilder<'_> {
173        BufferBuilder::new(data, self)
174    }
175
176    pub fn buffer_by_sz(&self, sz: usize) -> BufferBuilder<'_> {
177        BufferBuilder::new_by_sz(sz, self)
178    }
179}
180
181fn get_entry_point_name(module: &naga::Module) -> anyhow::Result<String> {
182    if module.entry_points.len() != 1 {
183        anyhow::bail!(
184            "Expected exactly one entry point, found {}",
185            module.entry_points.len()
186        );
187    }
188    for entry in module.entry_points.iter() {
189        return Ok(entry.name.clone());
190    }
191    anyhow::bail!("Could not find entry point");
192}
193
194fn get_entry_point_workgroup_size(module: &naga::Module, name: &str) -> anyhow::Result<[u32; 3]> {
195    for entry_point in module.entry_points.iter() {
196        if entry_point.name != name {
197            continue;
198        }
199        return Ok(entry_point.workgroup_size);
200    }
201    anyhow::bail!("Could not find entry point named {}", name);
202}
203
204fn for_each_global_variable<T>(
205    module: &naga::Module,
206    module_info: &naga::valid::ModuleInfo,
207    mut f: impl FnMut(
208        naga::Handle<GlobalVariable>,
209        &naga::GlobalVariable,
210        &naga::ResourceBinding,
211        &wgpu::ShaderStages,
212        bool,
213    ) -> anyhow::Result<Option<T>>,
214) -> anyhow::Result<Vec<T>> {
215    let mut result = Vec::new();
216    for (global_handle, global) in module.global_variables.iter() {
217        let Some(binding) = &global.binding else {
218            continue;
219        };
220
221        let can_write = match global.space {
222            naga::AddressSpace::Storage { access } => {
223                access.contains(naga::StorageAccess::STORE)
224                    || access.contains(naga::StorageAccess::ATOMIC)
225            }
226            naga::AddressSpace::Uniform => false,
227            _ => {
228                eprintln!(
229                    "Skipping global {:?} with address space {:?}",
230                    global.name, global.space
231                );
232                continue;
233            }
234        };
235        let mut stages = wgpu::ShaderStages::empty();
236        for (entry_idx, entry) in module.entry_points.iter().enumerate() {
237            let usage = module_info.get_entry_point(entry_idx)[global_handle];
238            if usage.is_empty() {
239                continue;
240            }
241
242            stages |= match entry.stage {
243                naga::ShaderStage::Vertex => wgpu::ShaderStages::VERTEX,
244                naga::ShaderStage::Fragment => wgpu::ShaderStages::FRAGMENT,
245                naga::ShaderStage::Compute => wgpu::ShaderStages::COMPUTE,
246                _ => continue,
247            }
248        }
249        if stages.is_empty() {
250            eprintln!(
251                "Skipping bound global {:?} which is not referenced by any stage",
252                global.name
253            );
254            continue;
255        }
256        if let Some(value) = f(global_handle, global, binding, &stages, can_write)? {
257            result.push(value);
258        }
259    }
260    Ok(result)
261}
262
263fn get_bind_group_layout(
264    gpu: &Gpu,
265    module: &naga::Module,
266    module_info: &naga::valid::ModuleInfo,
267    module_layout: &naga::proc::Layouter,
268) -> anyhow::Result<wgpu::BindGroupLayout> {
269    let mut entries = Vec::<wgpu::BindGroupLayoutEntry>::new();
270    for (global_handle, global) in module.global_variables.iter() {
271        let Some(binding) = &global.binding else {
272            continue;
273        };
274
275        let can_write = match global.space {
276            naga::AddressSpace::Storage { access } => {
277                access.contains(naga::StorageAccess::STORE)
278                    || access.contains(naga::StorageAccess::ATOMIC)
279            }
280            naga::AddressSpace::Uniform => false,
281            _ => {
282                eprintln!(
283                    "Skipping global {:?} with address space {:?}",
284                    global.name, global.space
285                );
286                continue;
287            }
288        };
289
290        let mut stages = wgpu::ShaderStages::empty();
291        for (entry_idx, entry) in module.entry_points.iter().enumerate() {
292            let usage = module_info.get_entry_point(entry_idx)[global_handle];
293            if usage.is_empty() {
294                continue;
295            }
296
297            stages |= match entry.stage {
298                naga::ShaderStage::Vertex => wgpu::ShaderStages::VERTEX,
299                naga::ShaderStage::Fragment => wgpu::ShaderStages::FRAGMENT,
300                naga::ShaderStage::Compute => wgpu::ShaderStages::COMPUTE,
301                _ => continue,
302            }
303        }
304        if stages.is_empty() {
305            eprintln!(
306                "Skipping bound global {:?} which is not referenced by any stage",
307                global.name
308            );
309            continue;
310        }
311
312        let global_ty = &module.types[global.ty];
313        let entry = match &module.types[global.ty].inner {
314            naga::ir::TypeInner::Array { .. } => {
315                let min_size = module_layout[global.ty].size;
316                wgpu::BindGroupLayoutEntry {
317                    binding: binding.binding,
318                    visibility: stages,
319                    ty: wgpu::BindingType::Buffer {
320                        ty: wgpu::BufferBindingType::Storage {
321                            read_only: !can_write,
322                        },
323                        min_binding_size: NonZeroU64::new(min_size as u64),
324                        has_dynamic_offset: false,
325                    },
326                    count: None,
327                }
328            }
329            naga::ir::TypeInner::Struct { .. } => {
330                let min_size = module_layout[global.ty].size;
331                wgpu::BindGroupLayoutEntry {
332                    binding: binding.binding,
333                    visibility: stages,
334                    ty: wgpu::BindingType::Buffer {
335                        ty: wgpu::BufferBindingType::Uniform,
336                        min_binding_size: NonZeroU64::new(min_size as u64),
337                        has_dynamic_offset: false,
338                    },
339                    count: None,
340                }
341            }
342            _ => {
343                anyhow::bail!(format!(
344                    "Unsupported global type: {:?} (name={:?})",
345                    global_ty.name, global_ty.inner
346                ))
347            }
348        };
349        entries.push(entry);
350    }
351    let layout = wgpu::BindGroupLayoutDescriptor {
352        label: None,
353        entries: &entries,
354    };
355    Ok(gpu.device.create_bind_group_layout(&layout))
356}
357
358pub struct Shader {
359    bind_group_layout: wgpu::BindGroupLayout,
360    shader_module: wgpu::ShaderModule,
361    entry_point_name: String,
362    entry_point_workgroup_size: [u32; 3],
363}
364
365impl Shader {
366    pub fn new(gpu: &Gpu, shader: &str) -> anyhow::Result<Self> {
367        let module = naga::front::wgsl::parse_str(shader)?;
368        let module_info = naga::valid::Validator::new(
369            naga::valid::ValidationFlags::all(),
370            naga::valid::Capabilities::all(),
371        )
372        .validate(&module)?;
373        let mut module_layout = naga::proc::Layouter::default();
374        module_layout.update(module.to_ctx())?;
375
376        Self::ensure_compute_compatibility(gpu, &module)?;
377        let bind_group_layout = get_bind_group_layout(gpu, &module, &module_info, &module_layout)?;
378        let entry_point_name = get_entry_point_name(&module)?;
379        let entry_point_workgroup_size =
380            get_entry_point_workgroup_size(&module, &entry_point_name)?;
381
382        let module_descriptor = wgpu::ShaderModuleDescriptor {
383            label: None,
384            source: wgpu::ShaderSource::Naga(Cow::Owned(module.clone())),
385        };
386        let shader_module = gpu.device.create_shader_module(module_descriptor.clone());
387
388        Ok(Self {
389            bind_group_layout,
390            // naga_module: module,
391            // naga_module_info: module_info,
392            shader_module,
393            entry_point_name,
394            entry_point_workgroup_size,
395            // module_layout,
396        })
397    }
398
399    fn ensure_compute_compatibility(gpu: &Gpu, module: &naga::Module) -> anyhow::Result<()> {
400        let mut has_compute_stage = false;
401        for entry_point in module.entry_points.iter() {
402            if entry_point.stage == naga::ShaderStage::Compute {
403                has_compute_stage = true;
404            }
405        }
406        if !has_compute_stage {
407            return Ok(());
408        }
409        let downlevel_caps = gpu.adapter.get_downlevel_capabilities();
410        if !downlevel_caps
411            .flags
412            .contains(wgpu::DownlevelFlags::COMPUTE_SHADERS)
413        {
414            anyhow::bail!(
415                "GPU {} does not support compute shaders.",
416                gpu.adapter.get_info().name
417            )
418        }
419        Ok(())
420    }
421}
422
423pub struct MaterialPass {
424    buffers: Vec<wgpu::Buffer>,
425    bind_group: wgpu::BindGroup,
426}
427
428pub struct MaterialPassBuilder {
429    buffers: Vec<wgpu::Buffer>,
430    bind_group: Option<wgpu::BindGroup>,
431}
432
433impl MaterialPassBuilder {
434    pub fn new() -> Self {
435        Self {
436            buffers: Vec::new(),
437            bind_group: None,
438        }
439    }
440
441    pub fn buffers(mut self, buffers: Vec<wgpu::Buffer>) -> Self {
442        self.buffers = buffers;
443        self
444    }
445
446    // Must be run after buffers(), etc.
447    pub fn bind_group(mut self, gpu: &Gpu, shader: &Shader) -> Self {
448        let mut descriptor_entries = Vec::<wgpu::BindGroupEntry>::new();
449        let entries =
450            self.buffers
451                .iter()
452                .enumerate()
453                .fold(Vec::new(), |mut entries, (index, buf)| {
454                    entries.push(wgpu::BindGroupEntry {
455                        binding: index as u32,
456                        resource: buf.as_entire_binding(),
457                    });
458                    entries
459                });
460        descriptor_entries.extend(entries);
461        let descriptor = wgpu::BindGroupDescriptor {
462            label: None,
463            layout: &shader.bind_group_layout,
464            entries: &descriptor_entries,
465        };
466        let bind_group = gpu.device.create_bind_group(&descriptor);
467        self.bind_group = Some(bind_group);
468        self
469    }
470
471    pub fn build(self) -> anyhow::Result<MaterialPass> {
472        if self.bind_group.is_none() {
473            anyhow::bail!("bind_group not set");
474        }
475        Ok(MaterialPass {
476            buffers: self.buffers,
477            bind_group: self.bind_group.unwrap(),
478        })
479    }
480}
481
482pub struct MaterialInstance {
483    passes: Vec<MaterialPass>,
484    readback: Option<wgpu::Buffer>,
485}
486
487impl MaterialInstance {
488    pub fn new(passes: Vec<MaterialPass>, readback: Option<wgpu::Buffer>) -> Self {
489        Self {
490            passes: passes,
491            readback: readback,
492        }
493    }
494}
495
496// A batch of Material work to be submitted together. This allows us to send a
497// batch of work while another is still processing on the GPU.
498pub struct MaterialBatch {
499    pub instances: Vec<MaterialInstance>,
500    ready_tx: Sender<(usize, Result<(), wgpu::BufferAsyncError>)>,
501    ready_rx: Receiver<(usize, Result<(), wgpu::BufferAsyncError>)>,
502    // Optional timestamp (ts) buffers for profiling.
503    timing: Option<MaterialTiming>,
504    // Whether the batch is currently processing. This is set true as soon as a batch is reserved with get_batch().
505    // The caller is responsible for setting this false
506    pub busy: bool,
507    // How many instances in the batch are currently processing. Set to
508    // batch_size (n_instancces) when the batch is reserved The caller may use
509    // this to track completion as each instance in the batch completes.
510    pub n_busy: usize,
511}
512
513impl MaterialBatch {
514    // Use MaterialBuilder instead. It handles creating batches.
515    fn new(
516        gpu: &Gpu,
517        mut passes_generator: impl FnMut() -> (Vec<MaterialPass>, Option<wgpu::Buffer>),
518        n_instances: usize,
519        enable_timing: bool,
520    ) -> Self {
521        let mut instances = Vec::<MaterialInstance>::new();
522        instances.reserve(n_instances);
523        for _ in 0..n_instances {
524            let (passes, readback) = passes_generator();
525            instances.push(MaterialInstance::new(passes, readback));
526        }
527        let (ready_tx, ready_rx) = std::sync::mpsc::channel();
528        let mut timing = None;
529        if enable_timing {
530            std::hint::cold_path();
531            timing = Some(MaterialTiming::new(gpu));
532        }
533        Self {
534            instances: instances,
535            ready_tx: ready_tx,
536            ready_rx: ready_rx,
537            timing: timing,
538            busy: false,
539            n_busy: 0,
540        }
541    }
542
543    pub fn enable_profiling(mut self, gpu: &Gpu) -> Self {
544        self.timing = Some(MaterialTiming::new(gpu));
545        self
546    }
547}
548
549pub struct MaterialTiming {
550    ts_buf: wgpu::Buffer,
551    ts_readback: wgpu::Buffer,
552    ready_tx: Sender<Result<(), wgpu::BufferAsyncError>>,
553    ready_rx: Receiver<Result<(), wgpu::BufferAsyncError>>,
554    busy: bool,
555}
556
557impl MaterialTiming {
558    pub fn new(gpu: &Gpu) -> Self {
559        let ts_buf = gpu.buffer_by_sz(16).query_resolve().copy_src().build();
560        let ts_readback = gpu.buffer_by_sz(16).copy_dst().map_read().build();
561        let (ready_tx, ready_rx) = std::sync::mpsc::channel();
562        Self {
563            ts_buf: ts_buf,
564            ts_readback: ts_readback,
565            ready_tx: ready_tx,
566            ready_rx: ready_rx,
567            busy: false,
568        }
569    }
570}
571
572pub struct Material {
573    pub batches: Vec<MaterialBatch>,
574    pub batch_size: usize,
575}
576
577impl Material {
578    pub fn get_batch(&mut self) -> Option<&mut MaterialBatch> {
579        let batch = self.batches.iter_mut().find(|batch| {
580            if batch.busy {
581                return false;
582            }
583            if let Some(timing) = &batch.timing {
584                std::hint::cold_path();
585                if timing.busy {
586                    return false;
587                }
588            }
589            true
590        })?;
591        // Found a batch.
592        batch.busy = true;
593        if let Some(timing) = &mut batch.timing {
594            std::hint::cold_path();
595            timing.busy = true;
596        }
597        batch.n_busy = self.batch_size;
598        Some(batch)
599    }
600}
601
602pub struct MaterialBuilder {
603    batch_size: usize, // instances per batch
604    n_batches: usize,
605    enable_timing: bool,
606}
607
608impl MaterialBuilder {
609    pub fn new() -> Self {
610        Self {
611            batch_size: 1,
612            n_batches: 1,
613            enable_timing: false,
614        }
615    }
616
617    pub fn batch_size(mut self, batch_size: usize) -> Self {
618        self.batch_size = batch_size;
619        self
620    }
621
622    pub fn n_batches(mut self, n_batches: usize) -> Self {
623        self.n_batches = n_batches;
624        self
625    }
626
627    pub fn enable_timing(mut self) -> Self {
628        self.enable_timing = true;
629        self
630    }
631
632    pub fn build(
633        self,
634        gpu: &Gpu,
635        mut passes_generator: impl FnMut() -> (Vec<MaterialPass>, Option<wgpu::Buffer>),
636    ) -> Material {
637        let mut batches = Vec::<MaterialBatch>::new();
638        batches.reserve(self.n_batches);
639        for _ in 0..self.n_batches {
640            batches.push(MaterialBatch::new(
641                gpu,
642                &mut passes_generator,
643                self.batch_size,
644                self.enable_timing,
645            ));
646        }
647        Material {
648            batches: batches,
649            batch_size: self.batch_size,
650        }
651    }
652}
653
654pub fn get_pipeline(gpu: &Gpu, shader: &Shader) -> wgpu::ComputePipeline {
655    let pipeline_layout = gpu
656        .device
657        .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
658            label: None,
659            bind_group_layouts: &[Some(&shader.bind_group_layout)],
660            immediate_size: 0,
661        });
662    // TODO we should use a builder. Don't want to skip
663    // zero-initializing memory in general case.
664    let compilation_options = wgpu::PipelineCompilationOptions {
665        zero_initialize_workgroup_memory: false,
666        ..Default::default()
667    };
668    gpu.device
669        .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
670            label: None,
671            layout: Some(&pipeline_layout),
672            module: &shader.shader_module,
673            entry_point: Some(&shader.entry_point_name),
674            compilation_options: compilation_options,
675            cache: None,
676        })
677}
678
679pub fn submit_dft_batch(
680    gpu: &Gpu,
681    shader: &Shader,
682    material: &mut Material,
683    pipeline: &wgpu::ComputePipeline,
684    data_batch: &[Vec<Complex<f32>>],
685) -> anyhow::Result<usize> {
686    assert_eq!(data_batch.len(), material.batch_size);
687
688    let Some(batch) = material.get_batch() else {
689        return Ok(0);
690    };
691
692    // Submit data to the GPU for each input in the batch.
693    for (data, instance) in data_batch.iter().zip(batch.instances.iter()) {
694        gpu.queue.write_buffer(
695            &instance.passes[0].buffers[0],
696            0,
697            bytemuck::cast_slice(data),
698        );
699    }
700
701    let mut encoder = gpu
702        .device
703        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
704    let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
705        label: None,
706        timestamp_writes: None,
707    });
708    compute_pass.set_pipeline(pipeline);
709    let shader_local_size = shader
710        .entry_point_workgroup_size
711        .iter()
712        .fold(1, |acc, &x| acc * x) as usize;
713    for idx in 0..data_batch.len() {
714        compute_pass.set_bind_group(0, &batch.instances[idx].passes[0].bind_group, &[]);
715        compute_pass.dispatch_workgroups(
716            data_batch[idx].len().div_ceil(shader_local_size) as u32,
717            1,
718            1,
719        );
720    }
721    drop(compute_pass);
722
723    for idx in 0..data_batch.len() {
724        let slot = &batch.instances[idx];
725        encoder.copy_buffer_to_buffer(
726            &slot.passes[0].buffers[1],
727            0,
728            &slot.readback.as_ref().unwrap(),
729            0,
730            slot.passes[0].buffers[1].size(),
731        );
732        let ready_tx = batch.ready_tx.clone();
733        encoder.map_buffer_on_submit(
734            &slot.readback.as_ref().unwrap(),
735            wgpu::MapMode::Read,
736            ..,
737            move |result| {
738                ready_tx.send((idx, result)).ok();
739            },
740        );
741    }
742
743    gpu.queue.submit([encoder.finish()]);
744    Ok(data_batch.len())
745}
746
747#[allow(dead_code)]
748fn drain_dft(
749    gpu: &Gpu,
750    material: &mut Material,
751) -> anyhow::Result<Vec<(usize, Vec<Complex<f32>>)>> {
752    gpu.device.poll(wgpu::PollType::Poll)?;
753    let mut results = Vec::new();
754    for batch in &mut material.batches {
755        while let Ok((instance_id, map_result)) = batch.ready_rx.try_recv() {
756            let slot = &mut batch.instances[instance_id];
757            let data = slot.readback.as_ref().unwrap().get_mapped_range(..)?;
758            let result = bytemuck::cast_slice(&data).to_vec();
759            drop(data);
760
761            slot.readback.as_ref().unwrap().unmap();
762            batch.n_busy -= 1;
763            if batch.n_busy == 0 {
764                batch.busy = false;
765            }
766
767            map_result?;
768            results.push((instance_id, result));
769        }
770    }
771    Ok(results)
772}
773
774#[allow(dead_code)]
775fn submit_fft_v0_batch(
776    gpu: &Gpu,
777    shader: &Shader,
778    material: &mut Material,
779    pipeline: &wgpu::ComputePipeline,
780    data_batch: &[Vec<Complex<f32>>],
781) -> anyhow::Result<usize> {
782    assert_eq!(data_batch.len(), material.batch_size);
783
784    let Some(batch) = material.get_batch() else {
785        return Ok(0);
786    };
787
788    // Submit data to the GPU for each input in the batch.
789    for (data, instance) in data_batch.iter().zip(batch.instances.iter()) {
790        gpu.queue.write_buffer(
791            &instance.passes[0].buffers[0],
792            0,
793            bytemuck::cast_slice(data),
794        );
795    }
796
797    let mut encoder = gpu
798        .device
799        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
800    let timestamp_writes = if let Some(timestamps) = &gpu.timestamps {
801        std::hint::cold_path();
802        Some(wgpu::ComputePassTimestampWrites {
803            query_set: &timestamps,
804            beginning_of_pass_write_index: Some(0),
805            end_of_pass_write_index: Some(1),
806        })
807    } else {
808        None
809    };
810
811    let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
812        label: None,
813        timestamp_writes: timestamp_writes,
814    });
815    compute_pass.set_pipeline(pipeline);
816    let shader_local_size = shader
817        .entry_point_workgroup_size
818        .iter()
819        .fold(1, |acc, &x| acc * x) as usize;
820    for (data, instance) in data_batch.iter().zip(batch.instances.iter()) {
821        for pass in &instance.passes {
822            compute_pass.set_bind_group(0, &pass.bind_group, &[]);
823            // Each instance computes two outputs, so we only need to dispatch len/2 instances.
824            compute_pass.dispatch_workgroups(
825                (data.len() / 2).div_ceil(shader_local_size) as u32,
826                1,
827                1,
828            );
829        }
830    }
831    drop(compute_pass);
832
833    for (instance_id, instance) in batch.instances.iter().enumerate() {
834        encoder.copy_buffer_to_buffer(
835            &instance.passes.last().unwrap().buffers[1],
836            0,
837            &instance.readback.as_ref().unwrap(),
838            0,
839            instance.passes.last().unwrap().buffers[1].size(),
840        );
841        let ready_tx = batch.ready_tx.clone();
842        encoder.map_buffer_on_submit(
843            &instance.readback.as_ref().unwrap(),
844            wgpu::MapMode::Read,
845            ..,
846            move |result| {
847                ready_tx.send((instance_id, result)).ok();
848            },
849        );
850    }
851
852    if let Some(timing) = &batch.timing {
853        std::hint::cold_path();
854        encoder.resolve_query_set(gpu.timestamps.as_ref().unwrap(), 0..2, &timing.ts_buf, 0);
855        encoder.copy_buffer_to_buffer(&timing.ts_buf, 0, &timing.ts_readback, 0, 16);
856        let ready_tx = timing.ready_tx.clone();
857        encoder.map_buffer_on_submit(
858            &timing.ts_readback,
859            wgpu::MapMode::Read,
860            ..,
861            move |result| {
862                ready_tx.send(result).ok();
863            },
864        );
865    }
866
867    gpu.queue.submit([encoder.finish()]);
868    Ok(data_batch.len())
869}
870
871#[allow(dead_code)]
872fn drain_fft_v0(
873    gpu: &Gpu,
874    material: &mut Material,
875    mut cb: impl FnMut(usize, &[Complex<f32>]) -> anyhow::Result<()>,
876) -> anyhow::Result<()> {
877    gpu.device.poll(wgpu::PollType::Poll)?;
878    for batch in &mut material.batches {
879        if !batch.busy {
880            continue;
881        }
882        while let Ok((instance_id, map_result)) = batch.ready_rx.try_recv() {
883            let instance = &batch.instances[instance_id];
884            map_result?;
885            let data = instance.readback.as_ref().unwrap().get_mapped_range(..)?;
886            let cb_result = cb(instance_id, bytemuck::cast_slice(&data));
887            drop(data);
888
889            instance.readback.as_ref().unwrap().unmap();
890            batch.n_busy -= 1;
891            if batch.n_busy == 0 {
892                batch.busy = false;
893            }
894
895            cb_result?;
896        }
897    }
898    Ok(())
899}
900
901#[allow(dead_code)]
902fn submit_fft_v1_batch(
903    gpu: &Gpu,
904    shader: &Shader,
905    material: &mut Material,
906    pipeline: &wgpu::ComputePipeline,
907    data_len: usize,                // length of 1 input
908    data_batch: &Vec<Complex<f32>>, // batch of N inputs
909) -> anyhow::Result<usize> {
910    assert!(data_len > 0);
911    assert_eq!(data_batch.len() % data_len, 0);
912
913    let Some(batch) = material.get_batch() else {
914        return Ok(0);
915    };
916    // We do not use instances since the input is pre-batched.
917    assert_eq!(batch.instances.len(), 1);
918    let instance = &batch.instances[0];
919
920    // Submit data
921    gpu.queue.write_buffer(
922        &instance.passes[0].buffers[0],
923        0,
924        bytemuck::cast_slice(&data_batch),
925    );
926
927    let mut encoder = gpu
928        .device
929        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
930    let timestamp_writes = if let Some(timestamps) = &gpu.timestamps {
931        std::hint::cold_path();
932        Some(wgpu::ComputePassTimestampWrites {
933            query_set: &timestamps,
934            beginning_of_pass_write_index: Some(0),
935            end_of_pass_write_index: Some(1),
936        })
937    } else {
938        None
939    };
940
941    let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
942        label: None,
943        timestamp_writes: timestamp_writes,
944    });
945    compute_pass.set_pipeline(pipeline);
946    let local_x = shader.entry_point_workgroup_size[0];
947    for pass in &instance.passes {
948        compute_pass.set_bind_group(0, &pass.bind_group, &[]);
949        // Each instance computes two outputs, so we only need to dispatch len/2 instances.
950        compute_pass.dispatch_workgroups(
951            (data_len / 2).div_ceil(local_x as usize) as u32,
952            (data_batch.len() / data_len) as u32,
953            1,
954        );
955    }
956    drop(compute_pass);
957
958    encoder.copy_buffer_to_buffer(
959        &instance.passes.last().unwrap().buffers[1],
960        0,
961        &instance.readback.as_ref().unwrap(),
962        0,
963        instance.passes.last().unwrap().buffers[1].size(),
964    );
965    let ready_tx = batch.ready_tx.clone();
966    encoder.map_buffer_on_submit(
967        &instance.readback.as_ref().unwrap(),
968        wgpu::MapMode::Read,
969        ..,
970        move |result| {
971            ready_tx.send((0, result)).ok();
972        },
973    );
974
975    if let Some(timing) = &batch.timing {
976        std::hint::cold_path();
977        encoder.resolve_query_set(gpu.timestamps.as_ref().unwrap(), 0..2, &timing.ts_buf, 0);
978        encoder.copy_buffer_to_buffer(&timing.ts_buf, 0, &timing.ts_readback, 0, 16);
979        let ready_tx = timing.ready_tx.clone();
980        encoder.map_buffer_on_submit(
981            &timing.ts_readback,
982            wgpu::MapMode::Read,
983            ..,
984            move |result| {
985                ready_tx.send(result).ok();
986            },
987        );
988    }
989
990    gpu.queue.submit([encoder.finish()]);
991    Ok(data_batch.len())
992}
993
994#[allow(dead_code)]
995fn submit_fft_v2_batch(
996    gpu: &Gpu,
997    shader: &Shader,
998    material: &mut Material,
999    pipeline: &wgpu::ComputePipeline,
1000    data_len: usize,                // length of 1 input
1001    data_batch: &Vec<Complex<f32>>, // batch of N inputs
1002) -> anyhow::Result<usize> {
1003    assert!(data_len > 0);
1004    assert_eq!(data_batch.len() % data_len, 0);
1005
1006    let Some(batch) = material.get_batch() else {
1007        return Ok(0);
1008    };
1009    // We do not use instances since the input is pre-batched.
1010    assert_eq!(batch.instances.len(), 1);
1011    let instance = &batch.instances[0];
1012
1013    // Submit data
1014    gpu.queue.write_buffer(
1015        &instance.passes[0].buffers[0],
1016        0,
1017        bytemuck::cast_slice(&data_batch),
1018    );
1019
1020    let mut encoder = gpu
1021        .device
1022        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
1023    let timestamp_writes = if let Some(timestamps) = &gpu.timestamps {
1024        std::hint::cold_path();
1025        Some(wgpu::ComputePassTimestampWrites {
1026            query_set: &timestamps,
1027            beginning_of_pass_write_index: Some(0),
1028            end_of_pass_write_index: Some(1),
1029        })
1030    } else {
1031        None
1032    };
1033
1034    let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
1035        label: None,
1036        timestamp_writes: timestamp_writes,
1037    });
1038    compute_pass.set_pipeline(pipeline);
1039    let local_x = shader.entry_point_workgroup_size[0];
1040    for pass in &instance.passes {
1041        compute_pass.set_bind_group(0, &pass.bind_group, &[]);
1042        // Each instance computes two outputs, so we only need to dispatch len/2 instances.
1043        compute_pass.dispatch_workgroups(1, 1, (data_batch.len() / data_len) as u32);
1044    }
1045    drop(compute_pass);
1046
1047    encoder.copy_buffer_to_buffer(
1048        &instance.passes.last().unwrap().buffers[1],
1049        0,
1050        &instance.readback.as_ref().unwrap(),
1051        0,
1052        instance.passes.last().unwrap().buffers[1].size(),
1053    );
1054    let ready_tx = batch.ready_tx.clone();
1055    encoder.map_buffer_on_submit(
1056        &instance.readback.as_ref().unwrap(),
1057        wgpu::MapMode::Read,
1058        ..,
1059        move |result| {
1060            ready_tx.send((0, result)).ok();
1061        },
1062    );
1063
1064    if let Some(timing) = &batch.timing {
1065        std::hint::cold_path();
1066        encoder.resolve_query_set(gpu.timestamps.as_ref().unwrap(), 0..2, &timing.ts_buf, 0);
1067        encoder.copy_buffer_to_buffer(&timing.ts_buf, 0, &timing.ts_readback, 0, 16);
1068        let ready_tx = timing.ready_tx.clone();
1069        encoder.map_buffer_on_submit(
1070            &timing.ts_readback,
1071            wgpu::MapMode::Read,
1072            ..,
1073            move |result| {
1074                ready_tx.send(result).ok();
1075            },
1076        );
1077    }
1078
1079    gpu.queue.submit([encoder.finish()]);
1080    Ok(data_batch.len())
1081}
1082
1083#[allow(dead_code)]
1084fn drain_fft_v0_timing(
1085    gpu: &Gpu,
1086    material: &mut Material,
1087    mut cb: impl FnMut(f64),
1088) -> anyhow::Result<()> {
1089    for batch in &mut material.batches {
1090        if let Some(timing) = &mut batch.timing {
1091            std::hint::cold_path();
1092
1093            if !timing.busy {
1094                continue;
1095            }
1096
1097            if let Ok(map_result) = timing.ready_rx.try_recv() {
1098                map_result?;
1099
1100                let data = timing.ts_readback.get_mapped_range(..)?;
1101                let ticks: &[u64] = bytemuck::cast_slice(&data);
1102                let elapsed_ns = (ticks[1].wrapping_sub(ticks[0]) as f64)
1103                    * (gpu.queue.get_timestamp_period() as f64);
1104
1105                drop(data);
1106                timing.ts_readback.unmap();
1107                cb(elapsed_ns);
1108                timing.busy = false;
1109            }
1110        }
1111    }
1112    Ok(())
1113}
1114
1115#[cfg(test)]
1116mod tests {
1117    use crate::util::*;
1118    use std::time::Duration;
1119    use std::time::Instant;
1120
1121    use num::Complex;
1122    use num::traits::Float;
1123    use rand::RngExt;
1124    use rand::SeedableRng;
1125    use rand::rngs::StdRng;
1126
1127    // Cast one float type to another, truncating if needed.
1128    fn t0_to_t1<T0: Float, T1: Float>(val: T0) -> T1 {
1129        return num::cast(val).unwrap();
1130    }
1131
1132    // Casts one Complex<T> to another.
1133    fn complex_to_t<T0: Float, T1: Float>(data: &[Complex<T0>]) -> Vec<Complex<T1>> {
1134        data.iter()
1135            .map(|c| Complex::new(t0_to_t1(c.re), t0_to_t1(c.im)))
1136            .collect()
1137    }
1138
1139    fn evaluate_results(
1140        result_ref_64: &Vec<Complex<f64>>,
1141        result_cur_64: &Vec<Complex<f64>>,
1142        duration_dft: &std::time::Duration,
1143        duration_cur: &std::time::Duration,
1144        algo_name: &str,
1145    ) {
1146        assert_eq!(result_ref_64.len(), result_cur_64.len());
1147
1148        println!("  Algorithm:  {}", algo_name);
1149        println!("    Duration:  {:?}", duration_cur);
1150        println!(
1151            "    Speedup:   {:?}",
1152            duration_dft.as_nanos() as f32 / duration_cur.as_nanos() as f32
1153        );
1154
1155        let mut max_err: f64 = 0.0;
1156        let mut sum_err: f64 = 0.0;
1157        // TODO: median, 90p, 99p
1158
1159        let result_64: Vec<Complex<f64>> = complex_to_t(&result_cur_64);
1160
1161        for i in 0..result_ref_64.len() {
1162            max_err = max_err.max((result_ref_64[i] - result_64[i]).norm());
1163            sum_err += (result_ref_64[i] - result_64[i]).norm();
1164        }
1165        println!("    Max err:   {}", max_err);
1166        println!("    Avg err:   {}", sum_err / (result_ref_64.len() as f64));
1167    }
1168
1169    fn evaluate_naive_dft(data: &[Complex<f32>], gpu: &super::Gpu) -> std::time::Duration {
1170        let shader = super::Shader::new(&gpu, include_wgsl!("naive_dft.wgsl")).unwrap();
1171        let mut result_dft = Vec::from(data);
1172
1173        const BATCH_SIZE: usize = 128;
1174        let inputs = vec![Vec::from(data); BATCH_SIZE];
1175
1176        let passes_gen = || -> (Vec<super::MaterialPass>, Option<wgpu::Buffer>) {
1177            let buffer_a = gpu.buffer(&data).storage().copy_dst().build();
1178            let buffer_b = gpu.buffer(&data).storage().copy_src().build();
1179            let buffers = vec![buffer_a, buffer_b];
1180
1181            let passes = vec![
1182                super::MaterialPassBuilder::new()
1183                    .buffers(buffers)
1184                    .bind_group(&gpu, &shader)
1185                    .build()
1186                    .unwrap(),
1187            ];
1188            let readback = gpu.buffer(&data).copy_dst().map_read().build();
1189            (passes, Some(readback))
1190        };
1191
1192        let mut material = super::MaterialBuilder::new()
1193            .batch_size(BATCH_SIZE)
1194            .n_batches(2)
1195            .build(gpu, passes_gen);
1196        let pipeline = super::get_pipeline(&gpu, &shader);
1197
1198        let warmup = Instant::now() + Duration::from_secs(3);
1199        let end = warmup + Duration::from_secs(5);
1200        let mut n_iter: u32 = 0;
1201        let mut now = Instant::now();
1202        while now < end {
1203            // Fill the pipeline.
1204            if super::submit_dft_batch(&gpu, &shader, &mut material, &pipeline, &inputs).unwrap()
1205                > 0
1206            {
1207                continue;
1208            }
1209            // Drain if pipeline is full.
1210            let results = super::drain_dft(&gpu, &mut material).unwrap();
1211            for (_, result) in &results {
1212                result_dft = result.clone();
1213            }
1214            if now > warmup {
1215                n_iter += results.len() as u32;
1216            }
1217            now = Instant::now();
1218        }
1219        let duration_dft = (now - warmup) / n_iter;
1220        println!("Duration: {:?}", duration_dft);
1221
1222        let mut result_ref: Vec<rustfft::num_complex::Complex<f64>> = complex_to_t(&data);
1223        let mut ref_planner = rustfft::FftPlanner::<f64>::new();
1224        let ref_fft = ref_planner.plan_fft_forward(result_ref.len());
1225        ref_fft.process(&mut result_ref);
1226        let result_ref_64: Vec<Complex<f64>> = complex_to_t(&result_ref);
1227
1228        let result_dft_64 = complex_to_t::<f32, f64>(&result_dft);
1229        evaluate_results(
1230            &result_ref_64,
1231            &result_dft_64,
1232            &duration_dft,
1233            &duration_dft,
1234            "Naive DFT",
1235        );
1236
1237        duration_dft
1238    }
1239
1240    fn evaluate_fft_v0(
1241        data: &[Complex<f32>],
1242        gpu: &super::Gpu,
1243        duration_dft: &std::time::Duration,
1244    ) {
1245        let shader = super::Shader::new(&gpu, include_wgsl!("fft_v0_basic.wgsl")).unwrap();
1246        let mut result = Vec::from(data);
1247        const BATCH_SIZE: usize = 8;
1248        let inputs = vec![Vec::from(data); BATCH_SIZE];
1249
1250        let twiddles = precompute_twiddles::<f32>(data.len());
1251        let twiddles_buffer = gpu.buffer(&twiddles).storage().copy_dst().build();
1252        gpu.queue.write_buffer(
1253            &twiddles_buffer,
1254            0,
1255            bytemuck::cast_slice(twiddles.as_slice()),
1256        );
1257
1258        let passes_gen = || -> (Vec<super::MaterialPass>, Option<wgpu::Buffer>) {
1259            let buffer_a = gpu.buffer(&data).storage().copy_src().copy_dst().build();
1260            let buffer_b = gpu.buffer(&data).storage().copy_src().copy_dst().build();
1261
1262            let mut stride = data.len() / 2;
1263            let mut big_n = 2;
1264            let num_passes = log_k_of::<2>(data.len());
1265            let mut passes = Vec::<super::MaterialPass>::new();
1266            passes.reserve(num_passes);
1267            for i in 0..num_passes {
1268                // TODO each pass needs its own uniforms, but each *instance* should share them
1269                #[repr(C)]
1270                #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
1271                struct Uniforms {
1272                    stride: u32,
1273                    big: u32,
1274                }
1275                let uniforms = Uniforms {
1276                    stride: stride as u32,
1277                    big: big_n as u32,
1278                };
1279                let uniforms_buffer = gpu.buffer(&uniforms).uniform().copy_dst().build();
1280                gpu.queue
1281                    .write_buffer(&uniforms_buffer, 0, bytemuck::bytes_of(&uniforms));
1282                stride /= 2;
1283                big_n *= 2;
1284
1285                let buffers = if i % 2 == 0 {
1286                    vec![
1287                        buffer_a.clone(),
1288                        buffer_b.clone(),
1289                        twiddles_buffer.clone(),
1290                        uniforms_buffer.clone(),
1291                    ]
1292                } else {
1293                    vec![
1294                        buffer_b.clone(),
1295                        buffer_a.clone(),
1296                        twiddles_buffer.clone(),
1297                        uniforms_buffer.clone(),
1298                    ]
1299                };
1300                passes.push(
1301                    super::MaterialPassBuilder::new()
1302                        .buffers(buffers)
1303                        .bind_group(&gpu, &shader)
1304                        .build()
1305                        .unwrap(),
1306                );
1307            }
1308            let readback = gpu.buffer(&data).copy_dst().map_read().build();
1309            (passes, Some(readback))
1310        };
1311
1312        // Allocate twice as many buffers as each batch uses.
1313        // This lets us send data over while the GPU is still working.
1314        let mut material = super::MaterialBuilder::new()
1315            .batch_size(BATCH_SIZE)
1316            .n_batches(2)
1317            .enable_timing()
1318            .build(gpu, passes_gen);
1319        let pipeline = super::get_pipeline(&gpu, &shader);
1320
1321        let warmup = Instant::now() + Duration::from_secs(3);
1322        let end = warmup + Duration::from_secs(5);
1323        let mut n_iter: u32 = 0;
1324        let mut sum_gpu_ns: f64 = 0.0;
1325        let mut now = Instant::now();
1326        while now < end {
1327            // Fill the pipeline.
1328            if super::submit_fft_v0_batch(&gpu, &shader, &mut material, &pipeline, &inputs).unwrap()
1329                > 0
1330            {
1331                continue;
1332            }
1333            // Drain if pipeline is full.
1334            let mut n_results = 0;
1335            super::drain_fft_v0(&gpu, &mut material, |_, data| {
1336                result.copy_from_slice(data);
1337                n_results += 1;
1338                Ok(())
1339            })
1340            .unwrap();
1341            let mut cur_sum_gpu_ns: f64 = 0.0;
1342            super::drain_fft_v0_timing(gpu, &mut material, |elapsed_ns| {
1343                cur_sum_gpu_ns += elapsed_ns;
1344            })
1345            .unwrap();
1346            if now > warmup {
1347                n_iter += n_results;
1348                sum_gpu_ns += cur_sum_gpu_ns;
1349            }
1350            now = Instant::now();
1351        }
1352        let duration = (now - warmup) / n_iter;
1353        println!("Duration (wall): {:?}", duration);
1354
1355        let gpu_duration = Duration::from_secs_f64(sum_gpu_ns / (n_iter as f64 * 1e9));
1356        println!("Duration (GPU):  {:?}", gpu_duration);
1357
1358        let mut result_ref: Vec<rustfft::num_complex::Complex<f64>> = complex_to_t(&data);
1359        let mut ref_planner = rustfft::FftPlanner::<f64>::new();
1360        let ref_fft = ref_planner.plan_fft_forward(result_ref.len());
1361        ref_fft.process(&mut result_ref);
1362        let result_ref_64: Vec<Complex<f64>> = complex_to_t(&result_ref);
1363
1364        let result_64 = complex_to_t::<f32, f64>(&result);
1365        evaluate_results(
1366            &result_ref_64,
1367            &result_64,
1368            &duration_dft,
1369            &duration,
1370            "FFT v0",
1371        );
1372    }
1373
1374    fn evaluate_fft_v1(
1375        data: &[Complex<f32>],
1376        gpu: &super::Gpu,
1377        duration_dft: &std::time::Duration,
1378    ) {
1379        let shader = super::Shader::new(&gpu, include_wgsl!("fft_v1_batched.wgsl")).unwrap();
1380        let mut result = Vec::from(data);
1381        const BATCH_SIZE: usize = 64;
1382
1383        let mut input_batch = Vec::<Complex<f32>>::new();
1384        input_batch.reserve(data.len() * BATCH_SIZE);
1385        for _ in 0..BATCH_SIZE {
1386            input_batch.extend_from_slice(data);
1387        }
1388
1389        let twiddles = precompute_twiddles::<f32>(data.len());
1390        let twiddles_buffer = gpu.buffer(&twiddles).storage().copy_dst().build();
1391        gpu.queue.write_buffer(
1392            &twiddles_buffer,
1393            0,
1394            bytemuck::cast_slice(twiddles.as_slice()),
1395        );
1396
1397        let passes_gen = || -> (Vec<super::MaterialPass>, Option<wgpu::Buffer>) {
1398            let buffer_a = gpu
1399                .buffer(&input_batch)
1400                .storage()
1401                .copy_src()
1402                .copy_dst()
1403                .build();
1404            let buffer_b = gpu
1405                .buffer(&input_batch)
1406                .storage()
1407                .copy_src()
1408                .copy_dst()
1409                .build();
1410
1411            let mut stride = data.len() / 2;
1412            let mut big_n = 2;
1413            let num_passes = log_k_of::<2>(data.len());
1414            let mut passes = Vec::<super::MaterialPass>::new();
1415            passes.reserve(num_passes);
1416            for i in 0..num_passes {
1417                // TODO each pass needs its own uniforms, but each *instance* should share them
1418                #[repr(C)]
1419                #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
1420                struct Uniforms {
1421                    stride: u32,
1422                    big_n: u32,
1423                    input_len: u32,
1424                }
1425                let uniforms = Uniforms {
1426                    stride: stride as u32,
1427                    big_n: big_n as u32,
1428                    input_len: data.len() as u32,
1429                };
1430                let uniforms_buffer = gpu.buffer(&uniforms).uniform().copy_dst().build();
1431                gpu.queue
1432                    .write_buffer(&uniforms_buffer, 0, bytemuck::bytes_of(&uniforms));
1433                stride /= 2;
1434                big_n *= 2;
1435
1436                let buffers = if i % 2 == 0 {
1437                    vec![
1438                        buffer_a.clone(),
1439                        buffer_b.clone(),
1440                        twiddles_buffer.clone(),
1441                        uniforms_buffer.clone(),
1442                    ]
1443                } else {
1444                    vec![
1445                        buffer_b.clone(),
1446                        buffer_a.clone(),
1447                        twiddles_buffer.clone(),
1448                        uniforms_buffer.clone(),
1449                    ]
1450                };
1451                passes.push(
1452                    super::MaterialPassBuilder::new()
1453                        .buffers(buffers)
1454                        .bind_group(&gpu, &shader)
1455                        .build()
1456                        .unwrap(),
1457                );
1458            }
1459            let readback = gpu.buffer(&input_batch).copy_dst().map_read().build();
1460            (passes, Some(readback))
1461        };
1462
1463        // Allocate twice as many buffers as each batch uses.
1464        // This lets us send data over while the GPU is still working.
1465        let mut material = super::MaterialBuilder::new()
1466            .batch_size(1) // use 1 instance per batch since we're pre-batching our inputs
1467            .n_batches(2)
1468            .enable_timing()
1469            .build(gpu, passes_gen);
1470        let pipeline = super::get_pipeline(&gpu, &shader);
1471
1472        let warmup = Instant::now() + Duration::from_secs(3);
1473        let end = warmup + Duration::from_secs(5);
1474        let mut n_iter: u32 = 0;
1475        let mut sum_gpu_ns: f64 = 0.0;
1476        let mut now = Instant::now();
1477        while now < end {
1478            // Fill the pipeline.
1479            if super::submit_fft_v1_batch(
1480                &gpu,
1481                &shader,
1482                &mut material,
1483                &pipeline,
1484                data.len(),
1485                &input_batch,
1486            )
1487            .unwrap()
1488                > 0
1489            {
1490                // Keep submitting as long as we're able to.
1491                continue;
1492            }
1493
1494            // Drain if pipeline is full.
1495            let mut n_results = 0;
1496            let data_len = data.len();
1497            super::drain_fft_v0(&gpu, &mut material, |_, data_batch| {
1498                for i in 0..BATCH_SIZE {
1499                    result.copy_from_slice(&data_batch[i * data_len..(i + 1) * data_len]);
1500                }
1501                n_results += BATCH_SIZE as u32;
1502                Ok(())
1503            })
1504            .unwrap();
1505            let mut cur_sum_gpu_ns: f64 = 0.0;
1506            super::drain_fft_v0_timing(gpu, &mut material, |elapsed_ns| {
1507                cur_sum_gpu_ns += elapsed_ns;
1508            })
1509            .unwrap();
1510            if now > warmup {
1511                n_iter += n_results;
1512                sum_gpu_ns += cur_sum_gpu_ns;
1513            }
1514            now = Instant::now();
1515        }
1516        let duration = (now - warmup) / n_iter;
1517        println!("Duration (wall): {:?}", duration);
1518
1519        let gpu_duration = Duration::from_secs_f64(sum_gpu_ns / (n_iter as f64 * 1e9));
1520        println!("Duration (GPU):  {:?}", gpu_duration);
1521
1522        let mut result_ref: Vec<rustfft::num_complex::Complex<f64>> = complex_to_t(&data);
1523        let mut ref_planner = rustfft::FftPlanner::<f64>::new();
1524        let ref_fft = ref_planner.plan_fft_forward(result_ref.len());
1525        ref_fft.process(&mut result_ref);
1526        let result_ref_64: Vec<Complex<f64>> = complex_to_t(&result_ref);
1527
1528        let result_64 = complex_to_t::<f32, f64>(&result);
1529        evaluate_results(
1530            &result_ref_64,
1531            &result_64,
1532            &duration_dft,
1533            &duration,
1534            "FFT v1",
1535        );
1536    }
1537
1538    fn evaluate_fft_v2(
1539        data: &[Complex<f32>],
1540        gpu: &super::Gpu,
1541        duration_dft: &std::time::Duration,
1542        shader: super::Shader,
1543    ) {
1544        let mut result = Vec::from(data);
1545        const BATCH_SIZE: usize = 1024;
1546
1547        let mut input_batch = Vec::<Complex<f32>>::new();
1548        input_batch.reserve(data.len() * BATCH_SIZE);
1549        for _ in 0..BATCH_SIZE {
1550            input_batch.extend_from_slice(data);
1551        }
1552
1553        let twiddles = precompute_twiddles::<f32>(data.len());
1554        let twiddles_buffer = gpu.buffer(&twiddles).storage().copy_dst().build();
1555        gpu.queue.write_buffer(
1556            &twiddles_buffer,
1557            0,
1558            bytemuck::cast_slice(twiddles.as_slice()),
1559        );
1560
1561        let passes_gen = || -> (Vec<super::MaterialPass>, Option<wgpu::Buffer>) {
1562            let buffer_a = gpu
1563                .buffer(&input_batch)
1564                .storage()
1565                .copy_src()
1566                .copy_dst()
1567                .build();
1568            let buffer_b = gpu
1569                .buffer(&input_batch)
1570                .storage()
1571                .copy_src()
1572                .copy_dst()
1573                .build();
1574
1575            let mut passes = Vec::<super::MaterialPass>::new();
1576
1577            // TODO each pass needs its own uniforms, but each *instance* should share them
1578            #[repr(C)]
1579            #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
1580            struct Uniforms {
1581                input_len: u32,
1582                n_iter: u32,
1583            }
1584            let uniforms = Uniforms {
1585                input_len: data.len() as u32,
1586                n_iter: log_k_of::<4>(data.len()) as u32,
1587            };
1588            let uniforms_buffer = gpu.buffer(&uniforms).uniform().copy_dst().build();
1589            gpu.queue
1590                .write_buffer(&uniforms_buffer, 0, bytemuck::bytes_of(&uniforms));
1591
1592            passes.push(
1593                super::MaterialPassBuilder::new()
1594                    .buffers(vec![
1595                        buffer_a,
1596                        buffer_b,
1597                        twiddles_buffer.clone(),
1598                        uniforms_buffer,
1599                    ])
1600                    .bind_group(&gpu, &shader)
1601                    .build()
1602                    .unwrap(),
1603            );
1604            let readback = gpu.buffer(&input_batch).copy_dst().map_read().build();
1605            (passes, Some(readback))
1606        };
1607
1608        // Allocate twice as many buffers as each batch uses.
1609        // This lets us send data over while the GPU is still working.
1610        let mut material = super::MaterialBuilder::new()
1611            .batch_size(1) // use 1 instance per batch since we're pre-batching our inputs
1612            .n_batches(2)
1613            .enable_timing()
1614            .build(gpu, passes_gen);
1615        let pipeline = super::get_pipeline(&gpu, &shader);
1616
1617        let warmup = Instant::now() + Duration::from_secs(3);
1618        let end = warmup + Duration::from_secs(5);
1619        let mut n_iter: u32 = 0;
1620        let mut sum_gpu_ns: f64 = 0.0;
1621        let mut now = Instant::now();
1622        while now < end {
1623            // Fill the pipeline.
1624            if super::submit_fft_v2_batch(
1625                &gpu,
1626                &shader,
1627                &mut material,
1628                &pipeline,
1629                data.len(),
1630                &input_batch,
1631            )
1632            .unwrap()
1633                > 0
1634            {
1635                // Keep submitting as long as we're able to.
1636                continue;
1637            }
1638
1639            // Drain if pipeline is full.
1640            let mut n_results = 0;
1641            let data_len = data.len();
1642            super::drain_fft_v0(&gpu, &mut material, |_, data_batch| {
1643                for i in 0..BATCH_SIZE {
1644                    result.copy_from_slice(&data_batch[i * data_len..(i + 1) * data_len]);
1645                }
1646                n_results += BATCH_SIZE as u32;
1647                Ok(())
1648            })
1649            .unwrap();
1650            let mut cur_sum_gpu_ns: f64 = 0.0;
1651            super::drain_fft_v0_timing(gpu, &mut material, |elapsed_ns| {
1652                cur_sum_gpu_ns += elapsed_ns;
1653            })
1654            .unwrap();
1655            if now > warmup {
1656                n_iter += n_results;
1657                sum_gpu_ns += cur_sum_gpu_ns;
1658            }
1659            now = Instant::now();
1660        }
1661        let duration = (now - warmup) / n_iter;
1662        println!("Duration (wall): {:?}", duration);
1663
1664        let gpu_duration = Duration::from_secs_f64(sum_gpu_ns / (n_iter as f64 * 1e9));
1665        println!("Duration (GPU):  {:?}", gpu_duration);
1666
1667        let mut result_ref: Vec<rustfft::num_complex::Complex<f64>> = complex_to_t(&data);
1668        let mut ref_planner = rustfft::FftPlanner::<f64>::new();
1669        let ref_fft = ref_planner.plan_fft_forward(result_ref.len());
1670        ref_fft.process(&mut result_ref);
1671        let result_ref_64: Vec<Complex<f64>> = complex_to_t(&result_ref);
1672
1673        let result_64 = complex_to_t::<f32, f64>(&result);
1674        evaluate_results(
1675            &result_ref_64,
1676            &result_64,
1677            &duration_dft,
1678            &duration,
1679            "FFT v2",
1680        );
1681    }
1682
1683    fn evaluate_fft_v4(
1684        data: &[Complex<f32>],
1685        gpu: &super::Gpu,
1686        duration_dft: &std::time::Duration,
1687        shader: super::Shader,
1688    ) {
1689        let mut result = Vec::from(data);
1690        const BATCH_SIZE: usize = 2048;
1691
1692        let mut input_batch = Vec::<Complex<f32>>::new();
1693        input_batch.reserve(data.len() * BATCH_SIZE);
1694        for _ in 0..BATCH_SIZE {
1695            input_batch.extend_from_slice(data);
1696        }
1697
1698        let twiddles = precompute_twiddles::<f32>(data.len());
1699        let twiddles_buffer = gpu.buffer(&twiddles).storage().copy_dst().build();
1700        gpu.queue.write_buffer(
1701            &twiddles_buffer,
1702            0,
1703            bytemuck::cast_slice(twiddles.as_slice()),
1704        );
1705
1706        let passes_gen = || -> (Vec<super::MaterialPass>, Option<wgpu::Buffer>) {
1707            let buffer_a = gpu
1708                .buffer(&input_batch)
1709                .storage()
1710                .copy_src()
1711                .copy_dst()
1712                .build();
1713            let buffer_b = gpu
1714                .buffer(&input_batch)
1715                .storage()
1716                .copy_src()
1717                .copy_dst()
1718                .build();
1719
1720            let mut passes = Vec::<super::MaterialPass>::new();
1721
1722            // TODO each pass needs its own uniforms, but each *instance* should share them
1723            #[repr(C)]
1724            #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
1725            struct Uniforms {
1726                input_len: u32,
1727                n_iter: u32,
1728            }
1729            let uniforms = Uniforms {
1730                input_len: data.len() as u32,
1731                n_iter: log_k_of::<8>(data.len()) as u32,
1732            };
1733            let uniforms_buffer = gpu.buffer(&uniforms).uniform().copy_dst().build();
1734            gpu.queue
1735                .write_buffer(&uniforms_buffer, 0, bytemuck::bytes_of(&uniforms));
1736
1737            passes.push(
1738                super::MaterialPassBuilder::new()
1739                    .buffers(vec![
1740                        buffer_a,
1741                        buffer_b,
1742                        twiddles_buffer.clone(),
1743                        uniforms_buffer,
1744                    ])
1745                    .bind_group(&gpu, &shader)
1746                    .build()
1747                    .unwrap(),
1748            );
1749            let readback = gpu.buffer(&input_batch).copy_dst().map_read().build();
1750            (passes, Some(readback))
1751        };
1752
1753        // Allocate twice as many buffers as each batch uses.
1754        // This lets us send data over while the GPU is still working.
1755        let mut material = super::MaterialBuilder::new()
1756            .batch_size(1) // use 1 instance per batch since we're pre-batching our inputs
1757            .n_batches(4)
1758            .enable_timing()
1759            .build(gpu, passes_gen);
1760        let pipeline = super::get_pipeline(&gpu, &shader);
1761
1762        let warmup = Instant::now() + Duration::from_secs(3);
1763        let end = warmup + Duration::from_secs(5);
1764        let mut n_iter: u32 = 0;
1765        let mut sum_gpu_ns: f64 = 0.0;
1766        let mut now = Instant::now();
1767        while now < end {
1768            // Fill the pipeline.
1769            if super::submit_fft_v2_batch(
1770                &gpu,
1771                &shader,
1772                &mut material,
1773                &pipeline,
1774                data.len(),
1775                &input_batch,
1776            )
1777            .unwrap()
1778                > 0
1779            {
1780                // Keep submitting as long as we're able to.
1781                continue;
1782            }
1783
1784            // Drain if pipeline is full.
1785            let mut n_results = 0;
1786            let data_len = data.len();
1787            super::drain_fft_v0(&gpu, &mut material, |_, data_batch| {
1788                for i in 0..BATCH_SIZE {
1789                    result.copy_from_slice(&data_batch[i * data_len..(i + 1) * data_len]);
1790                }
1791                n_results += BATCH_SIZE as u32;
1792                Ok(())
1793            })
1794            .unwrap();
1795            let mut cur_sum_gpu_ns: f64 = 0.0;
1796            super::drain_fft_v0_timing(gpu, &mut material, |elapsed_ns| {
1797                cur_sum_gpu_ns += elapsed_ns;
1798            })
1799            .unwrap();
1800            if now > warmup {
1801                n_iter += n_results;
1802                sum_gpu_ns += cur_sum_gpu_ns;
1803            }
1804            now = Instant::now();
1805        }
1806        let duration = (now - warmup) / n_iter;
1807        println!("Duration (wall): {:?}", duration);
1808
1809        let gpu_duration = Duration::from_secs_f64(sum_gpu_ns / (n_iter as f64 * 1e9));
1810        println!("Duration (GPU):  {:?}", gpu_duration);
1811
1812        let mut result_ref: Vec<rustfft::num_complex::Complex<f64>> = complex_to_t(&data);
1813        let mut ref_planner = rustfft::FftPlanner::<f64>::new();
1814        let ref_fft = ref_planner.plan_fft_forward(result_ref.len());
1815        ref_fft.process(&mut result_ref);
1816        let result_ref_64: Vec<Complex<f64>> = complex_to_t(&result_ref);
1817
1818        let result_64 = complex_to_t::<f32, f64>(&result);
1819        evaluate_results(
1820            &result_ref_64,
1821            &result_64,
1822            &duration_dft,
1823            &duration,
1824            "FFT v4",
1825        );
1826    }
1827
1828    #[test]
1829    fn run_it() {
1830        let mut data_f64: Vec<Complex<f64>> = Vec::new();
1831        data_f64.resize((2 as usize).pow(12), Complex::ZERO);
1832        // Randomize data.
1833        let mut rng = StdRng::seed_from_u64(42);
1834        for i in 0..data_f64.len() {
1835            data_f64[i] = Complex::new(rng.random(), rng.random());
1836        }
1837        let data_f32: Vec<Complex<f32>> = complex_to_t::<f64, f32>(&data_f64);
1838
1839        let gpu = super::GpuBuilder::new().enable_profiling().build().unwrap();
1840
1841        //let duration_dft = evaluate_naive_dft(&data_f32, &gpu);
1842        let duration_dft = Duration::from_secs(0);
1843        //evaluate_fft_v0(&data_f32, &gpu, &duration_dft);
1844        //evaluate_fft_v1(&data_f32, &gpu, &duration_dft);
1845        //evaluate_fft_v2(&data_f32, &gpu, &duration_dft, super::Shader::new(&gpu, include_wgsl!("fft_v2_barrier.wgsl")).unwrap());
1846        /*
1847        evaluate_fft_v2(
1848            &data_f32,
1849            &gpu,
1850            &duration_dft,
1851            super::Shader::new(&gpu, include_wgsl!("fft_v3_groupshared.wgsl")).unwrap(),
1852        );
1853        */
1854        evaluate_fft_v4(
1855            &data_f32,
1856            &gpu,
1857            &duration_dft,
1858            super::Shader::new(&gpu, include_wgsl!("fft_v4_radix_8.wgsl")).unwrap(),
1859        );
1860    }
1861}