dynamis_gpu/
compute.rs

1use crate::buffer::GpuBuffer;
2use wgpu::{
3    BindGroup, BindGroupEntry, BindGroupLayout, BindGroupLayoutDescriptor, BindGroupLayoutEntry,
4    BindingType, BufferBindingType, CommandEncoder, ComputePassDescriptor,
5    ComputePipeline as WgpuComputePipeline, ComputePipelineDescriptor, Device,
6    PipelineLayoutDescriptor, ShaderModuleDescriptor, ShaderSource, ShaderStages,
7};
8
9#[derive(Clone, Copy)]
10pub enum BindingKind {
11    Uniform,
12    ReadOnlyStorage,
13    ReadWriteStorage,
14}
15
16pub struct BindingSpec {
17    pub binding: u32,
18    pub kind: BindingKind,
19}
20
21pub struct ComputePipeline {
22    pipeline: WgpuComputePipeline,
23    bind_group_layout: BindGroupLayout,
24    workgroup_size: u32,
25}
26
27impl ComputePipeline {
28    pub fn new(
29        device: &Device,
30        label: &str,
31        shader: &str,
32        entry: &str,
33        bindings: &[BindingSpec],
34        workgroup_size: u32,
35    ) -> Self {
36        let module = device.create_shader_module(ShaderModuleDescriptor {
37            label: Some(label),
38            source: ShaderSource::Wgsl(shader.into()),
39        });
40        let entries: Vec<BindGroupLayoutEntry> = bindings
41            .iter()
42            .map(|spec| BindGroupLayoutEntry {
43                binding: spec.binding,
44                visibility: ShaderStages::COMPUTE,
45                ty: match spec.kind {
46                    BindingKind::Uniform => BindingType::Buffer {
47                        ty: BufferBindingType::Uniform,
48                        has_dynamic_offset: false,
49                        min_binding_size: None,
50                    },
51                    BindingKind::ReadOnlyStorage => BindingType::Buffer {
52                        ty: BufferBindingType::Storage { read_only: true },
53                        has_dynamic_offset: false,
54                        min_binding_size: None,
55                    },
56                    BindingKind::ReadWriteStorage => BindingType::Buffer {
57                        ty: BufferBindingType::Storage { read_only: false },
58                        has_dynamic_offset: false,
59                        min_binding_size: None,
60                    },
61                },
62                count: None,
63            })
64            .collect();
65        let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
66            label: Some(label),
67            entries: &entries,
68        });
69        let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
70            label: Some(label),
71            bind_group_layouts: &[Some(&bind_group_layout)],
72            immediate_size: 0,
73        });
74        let pipeline = device.create_compute_pipeline(&ComputePipelineDescriptor {
75            label: Some(label),
76            layout: Some(&pipeline_layout),
77            module: &module,
78            entry_point: Some(entry),
79            compilation_options: Default::default(),
80            cache: None,
81        });
82        Self {
83            pipeline,
84            bind_group_layout,
85            workgroup_size,
86        }
87    }
88
89    pub fn bind_group_layout(&self) -> &BindGroupLayout {
90        &self.bind_group_layout
91    }
92
93    pub fn pipeline(&self) -> &WgpuComputePipeline {
94        &self.pipeline
95    }
96
97    pub fn workgroup_size(&self) -> u32 {
98        self.workgroup_size
99    }
100
101    pub fn create_bind_group(&self, device: &Device, entries: &[BindGroupEntry<'_>]) -> BindGroup {
102        device.create_bind_group(&wgpu::BindGroupDescriptor {
103            label: None,
104            layout: &self.bind_group_layout,
105            entries,
106        })
107    }
108
109    pub fn workgroup_count(&self, elements: u32) -> u32 {
110        elements.div_ceil(self.workgroup_size)
111    }
112
113    pub fn dispatch(&self, encoder: &mut CommandEncoder, bind_group: &BindGroup, count: u32) {
114        let mut pass = encoder.begin_compute_pass(&ComputePassDescriptor {
115            label: None,
116            timestamp_writes: None,
117        });
118        pass.set_pipeline(&self.pipeline);
119        pass.set_bind_group(0, bind_group, &[]);
120        pass.dispatch_workgroups(count, 1, 1);
121    }
122
123    pub fn dispatch_indirect(
124        &self,
125        encoder: &mut CommandEncoder,
126        bind_group: &BindGroup,
127        args: &GpuBuffer,
128    ) {
129        let mut pass = encoder.begin_compute_pass(&ComputePassDescriptor {
130            label: None,
131            timestamp_writes: None,
132        });
133        pass.set_pipeline(&self.pipeline);
134        pass.set_bind_group(0, bind_group, &[]);
135        pass.dispatch_workgroups_indirect(args.as_indirect_args(), 0);
136    }
137
138    pub fn dispatch_indirect_workgrouped(
139        &self,
140        encoder: &mut CommandEncoder,
141        bind_group: &BindGroup,
142        args: &GpuBuffer,
143    ) {
144        let mut pass = encoder.begin_compute_pass(&ComputePassDescriptor {
145            label: None,
146            timestamp_writes: None,
147        });
148        pass.set_pipeline(&self.pipeline);
149        pass.set_bind_group(0, bind_group, &[]);
150        pass.dispatch_workgroups_indirect(args.as_indirect_args(), 0);
151    }
152}
dynamis_gpu/compute.rs

dynamis_gpu/
compute.rs
