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