Skip to main content

runmat_plot/gpu/
bar.rs

1use crate::core::renderer::Vertex;
2use crate::core::scene::GpuVertexBuffer;
3use crate::gpu::shaders;
4use crate::gpu::{tuning, ScalarType};
5use glam::Vec4;
6use std::sync::Arc;
7use wgpu::util::DeviceExt;
8const VERTICES_PER_BAR: u32 = 6;
9
10/// Inputs required to pack bar vertices directly on the GPU.
11#[derive(Clone, Debug)]
12pub struct BarGpuInputs {
13    pub values_buffer: Arc<wgpu::Buffer>,
14    pub row_count: u32,
15    pub scalar: ScalarType,
16}
17
18/// Parameters describing how the GPU vertices should be generated.
19#[derive(Clone, Copy, Debug)]
20pub struct BarGpuParams {
21    pub color: Vec4,
22    pub bar_width: f32,
23    pub series_index: u32,
24    pub series_count: u32,
25    pub source_row_count: u32,
26    pub transpose_source: bool,
27    pub group_index: u32,
28    pub group_count: u32,
29    pub orientation: BarOrientation,
30    pub layout: BarLayoutMode,
31}
32
33#[derive(Clone, Copy, Debug)]
34pub enum BarLayoutMode {
35    Grouped,
36    Stacked,
37}
38
39impl BarLayoutMode {
40    fn as_u32(self) -> u32 {
41        match self {
42            BarLayoutMode::Grouped => 0,
43            BarLayoutMode::Stacked => 1,
44        }
45    }
46}
47
48#[derive(Clone, Copy, Debug)]
49pub enum BarOrientation {
50    Vertical,
51    Horizontal,
52}
53
54impl BarOrientation {
55    fn as_u32(self) -> u32 {
56        match self {
57            BarOrientation::Vertical => 0,
58            BarOrientation::Horizontal => 1,
59        }
60    }
61}
62
63#[repr(C)]
64#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
65struct BarUniforms {
66    color: [f32; 4],
67    bar_width: f32,
68    row_count: u32,
69    series_index: u32,
70    series_count: u32,
71    source_row_count: u32,
72    transpose_source: u32,
73    group_index: u32,
74    group_count: u32,
75    orientation: u32,
76    layout: u32,
77    _pad: [u32; 2],
78}
79
80/// Builds a GPU-resident vertex buffer for bar charts directly from provider-owned value arrays.
81pub fn pack_vertices_from_values(
82    device: &Arc<wgpu::Device>,
83    queue: &Arc<wgpu::Queue>,
84    inputs: &BarGpuInputs,
85    params: &BarGpuParams,
86) -> Result<GpuVertexBuffer, String> {
87    if inputs.row_count == 0 {
88        return Err("bar: input cannot be empty".to_string());
89    }
90
91    let workgroup_size = tuning::effective_workgroup_size();
92    let shader = compile_shader(device, workgroup_size, inputs.scalar);
93
94    let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
95        label: Some("bar-pack-bind-layout"),
96        entries: &[
97            wgpu::BindGroupLayoutEntry {
98                binding: 0,
99                visibility: wgpu::ShaderStages::COMPUTE,
100                ty: wgpu::BindingType::Buffer {
101                    ty: wgpu::BufferBindingType::Storage { read_only: true },
102                    has_dynamic_offset: false,
103                    min_binding_size: None,
104                },
105                count: None,
106            },
107            wgpu::BindGroupLayoutEntry {
108                binding: 1,
109                visibility: wgpu::ShaderStages::COMPUTE,
110                ty: wgpu::BindingType::Buffer {
111                    ty: wgpu::BufferBindingType::Storage { read_only: false },
112                    has_dynamic_offset: false,
113                    min_binding_size: None,
114                },
115                count: None,
116            },
117            wgpu::BindGroupLayoutEntry {
118                binding: 2,
119                visibility: wgpu::ShaderStages::COMPUTE,
120                ty: wgpu::BindingType::Buffer {
121                    ty: wgpu::BufferBindingType::Uniform,
122                    has_dynamic_offset: false,
123                    min_binding_size: None,
124                },
125                count: None,
126            },
127        ],
128    });
129
130    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
131        label: Some("bar-pack-pipeline-layout"),
132        bind_group_layouts: &[&bind_group_layout],
133        push_constant_ranges: &[],
134    });
135
136    let pipeline =
137        device.create_compute_pipeline(&crate::wgpu_compat::wgpu_compute_pipeline_descriptor! {
138            label: Some("bar-pack-pipeline"),
139            layout: Some(&pipeline_layout),
140            module: &shader,
141            entry_point: "main",
142        });
143
144    let vertex_count = inputs.row_count as u64 * VERTICES_PER_BAR as u64;
145    let output_size = vertex_count * std::mem::size_of::<Vertex>() as u64;
146    let output_buffer = Arc::new(device.create_buffer(&wgpu::BufferDescriptor {
147        label: Some("bar-gpu-vertices"),
148        size: output_size,
149        usage: wgpu::BufferUsages::STORAGE
150            | wgpu::BufferUsages::VERTEX
151            | wgpu::BufferUsages::COPY_DST
152            | wgpu::BufferUsages::COPY_SRC,
153        mapped_at_creation: false,
154    }));
155
156    let uniforms = BarUniforms {
157        color: params.color.to_array(),
158        bar_width: params.bar_width,
159        row_count: inputs.row_count,
160        series_index: params.series_index,
161        series_count: params.series_count.max(1),
162        source_row_count: params.source_row_count.max(inputs.row_count),
163        transpose_source: u32::from(params.transpose_source),
164        group_index: params.group_index,
165        group_count: params.group_count.max(1),
166        orientation: params.orientation.as_u32(),
167        layout: params.layout.as_u32(),
168        _pad: [0, 0],
169    };
170    let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
171        label: Some("bar-pack-uniforms"),
172        contents: bytemuck::bytes_of(&uniforms),
173        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
174    });
175
176    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
177        label: Some("bar-pack-bind-group"),
178        layout: &bind_group_layout,
179        entries: &[
180            wgpu::BindGroupEntry {
181                binding: 0,
182                resource: inputs.values_buffer.as_entire_binding(),
183            },
184            wgpu::BindGroupEntry {
185                binding: 1,
186                resource: output_buffer.as_entire_binding(),
187            },
188            wgpu::BindGroupEntry {
189                binding: 2,
190                resource: uniform_buffer.as_entire_binding(),
191            },
192        ],
193    });
194
195    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
196        label: Some("bar-pack-encoder"),
197    });
198    {
199        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
200            label: Some("bar-pack-pass"),
201            timestamp_writes: None,
202        });
203        pass.set_pipeline(&pipeline);
204        pass.set_bind_group(0, &bind_group, &[]);
205        let workgroups = inputs.row_count.div_ceil(workgroup_size);
206        pass.dispatch_workgroups(workgroups, 1, 1);
207    }
208    queue.submit(Some(encoder.finish()));
209
210    Ok(GpuVertexBuffer::new(output_buffer, vertex_count as usize))
211}
212
213fn compile_shader(
214    device: &Arc<wgpu::Device>,
215    workgroup_size: u32,
216    scalar: ScalarType,
217) -> wgpu::ShaderModule {
218    let template = match scalar {
219        ScalarType::F32 => shaders::bar::F32,
220        ScalarType::F64 => shaders::bar::F64,
221    };
222    let source = template.replace("{{WORKGROUP_SIZE}}", &workgroup_size.to_string());
223    device.create_shader_module(wgpu::ShaderModuleDescriptor {
224        label: Some("bar-pack-shader"),
225        source: wgpu::ShaderSource::Wgsl(source.into()),
226    })
227}