Skip to main content

runmat_plot/gpu/shaders/
bar.rs

1pub const F32: &str = r#"const WORKGROUP_SIZE: u32 = {{WORKGROUP_SIZE}}u;
2
3const VERTICES_PER_BAR: u32 = 6u;
4
5struct VertexRaw {
6    data: array<f32, 12u>,
7};
8
9struct BarParams {
10    color: vec4<f32>,
11    bar_width: f32,
12    row_count: u32,
13    series_index: u32,
14    series_count: u32,
15    source_row_count: u32,
16    transpose_source: u32,
17    group_index: u32,
18    group_count: u32,
19    orientation: u32,
20    layout: u32,
21};
22
23@group(0) @binding(0)
24var<storage, read> values: array<f32>;
25
26@group(0) @binding(1)
27var<storage, read_write> out_vertices: array<VertexRaw>;
28
29@group(0) @binding(2)
30var<uniform> params: BarParams;
31
32fn encode_vertex(position: vec3<f32>) -> VertexRaw {
33    var vertex: VertexRaw;
34    vertex.data[0u] = position.x;
35    vertex.data[1u] = position.y;
36    vertex.data[2u] = position.z;
37    vertex.data[3u] = params.color.x;
38    vertex.data[4u] = params.color.y;
39    vertex.data[5u] = params.color.z;
40    vertex.data[6u] = params.color.w;
41    vertex.data[7u] = 0.0;
42    vertex.data[8u] = 0.0;
43    vertex.data[9u] = 1.0;
44    vertex.data[10u] = 0.0;
45    vertex.data[11u] = 0.0;
46    return vertex;
47}
48
49fn write_vertices(base_index: u32, quad: array<vec3<f32>, 4u>) {
50    out_vertices[base_index + 0u] = encode_vertex(quad[0u]);
51    out_vertices[base_index + 1u] = encode_vertex(quad[1u]);
52    out_vertices[base_index + 2u] = encode_vertex(quad[2u]);
53    out_vertices[base_index + 3u] = encode_vertex(quad[0u]);
54    out_vertices[base_index + 4u] = encode_vertex(quad[2u]);
55    out_vertices[base_index + 5u] = encode_vertex(quad[3u]);
56}
57
58fn build_vertical_quad(idx: u32, start: f32, end: f32, per_group_width: f32, local_offset: f32) -> array<vec3<f32>, 4u> {
59    let center = (f32(idx) + 1.0) + local_offset;
60    let half = per_group_width * 0.5;
61    let left = center - half;
62    let right = center + half;
63    let bottom = min(start, end);
64    let top = max(start, end);
65
66    return array<vec3<f32>, 4u>(
67        vec3<f32>(left, bottom, 0.0),
68        vec3<f32>(right, bottom, 0.0),
69        vec3<f32>(right, top, 0.0),
70        vec3<f32>(left, top, 0.0),
71    );
72}
73
74fn build_horizontal_quad(idx: u32, start: f32, end: f32, per_group_width: f32, local_offset: f32) -> array<vec3<f32>, 4u> {
75    let center = (f32(idx) + 1.0) + local_offset;
76    let half = per_group_width * 0.5;
77    let bottom = center - half;
78    let top = center + half;
79    let left = min(start, end);
80    let right = max(start, end);
81
82    return array<vec3<f32>, 4u>(
83        vec3<f32>(left, bottom, 0.0),
84        vec3<f32>(right, bottom, 0.0),
85        vec3<f32>(right, top, 0.0),
86        vec3<f32>(left, top, 0.0),
87    );
88}
89
90@compute @workgroup_size(WORKGROUP_SIZE)
91fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
92    let idx = gid.x;
93    if (idx >= params.row_count) {
94        return;
95    }
96
97    let safe_group_count = max(params.group_count, 1u);
98    let per_group_width = max(params.bar_width / f32(safe_group_count), 0.01);
99    let group_offset_start = -params.bar_width * 0.5;
100    let local_offset = group_offset_start
101        + per_group_width * f32(min(params.group_index, safe_group_count - 1u))
102        + per_group_width * 0.5;
103    let source_stride = max(params.source_row_count, 1u);
104    let value_index = if (params.transpose_source == 1u) {
105        idx * source_stride + params.series_index
106    } else {
107        params.series_index * source_stride + idx
108    };
109    let value = values[value_index];
110
111    var base_pos = 0.0;
112    var base_neg = 0.0;
113    if (params.layout == 1u && params.series_index > 0u) {
114        var col: u32 = 0u;
115        loop {
116            if (col >= params.series_index) {
117                break;
118            }
119            let prev_index = if (params.transpose_source == 1u) {
120                idx * source_stride + col
121            } else {
122                col * source_stride + idx
123            };
124            let prev = values[prev_index];
125            if (isFinite(prev)) {
126                if (prev >= 0.0) {
127                    base_pos += prev;
128                } else {
129                    base_neg += prev;
130                }
131            }
132            col = col + 1u;
133        }
134    }
135
136    var start = 0.0;
137    var end = 0.0;
138    if (!isFinite(value)) {
139        start = 0.0;
140        end = 0.0;
141    } else if (params.layout == 1u) {
142        if (value >= 0.0) {
143            start = base_pos;
144            end = base_pos + value;
145        } else {
146            start = base_neg + value;
147            end = base_neg;
148        }
149    } else {
150        start = 0.0;
151        end = value;
152    }
153
154    let quad = if (params.orientation == 0u) {
155        build_vertical_quad(idx, start, end, per_group_width, local_offset)
156    } else {
157        build_horizontal_quad(idx, start, end, per_group_width, local_offset)
158    };
159
160    let base_index = idx * VERTICES_PER_BAR;
161    write_vertices(base_index, quad);
162}
163"#;
164
165pub const F64: &str = r#"const WORKGROUP_SIZE: u32 = {{WORKGROUP_SIZE}}u;
166
167const VERTICES_PER_BAR: u32 = 6u;
168
169struct VertexRaw {
170    data: array<f32, 12u>,
171};
172
173struct BarParams {
174    color: vec4<f32>,
175    bar_width: f32,
176    row_count: u32,
177    series_index: u32,
178    series_count: u32,
179    source_row_count: u32,
180    transpose_source: u32,
181    group_index: u32,
182    group_count: u32,
183    orientation: u32,
184    layout: u32,
185};
186
187@group(0) @binding(0)
188var<storage, read> values: array<f64>;
189
190@group(0) @binding(1)
191var<storage, read_write> out_vertices: array<VertexRaw>;
192
193@group(0) @binding(2)
194var<uniform> params: BarParams;
195
196fn encode_vertex(position: vec3<f32>) -> VertexRaw {
197    var vertex: VertexRaw;
198    vertex.data[0u] = position.x;
199    vertex.data[1u] = position.y;
200    vertex.data[2u] = position.z;
201    vertex.data[3u] = params.color.x;
202    vertex.data[4u] = params.color.y;
203    vertex.data[5u] = params.color.z;
204    vertex.data[6u] = params.color.w;
205    vertex.data[7u] = 0.0;
206    vertex.data[8u] = 0.0;
207    vertex.data[9u] = 1.0;
208    vertex.data[10u] = 0.0;
209    vertex.data[11u] = 0.0;
210    return vertex;
211}
212
213fn write_vertices(base_index: u32, quad: array<vec3<f32>, 4u>) {
214    out_vertices[base_index + 0u] = encode_vertex(quad[0u]);
215    out_vertices[base_index + 1u] = encode_vertex(quad[1u]);
216    out_vertices[base_index + 2u] = encode_vertex(quad[2u]);
217    out_vertices[base_index + 3u] = encode_vertex(quad[0u]);
218    out_vertices[base_index + 4u] = encode_vertex(quad[2u]);
219    out_vertices[base_index + 5u] = encode_vertex(quad[3u]);
220}
221
222fn build_vertical_quad(idx: u32, start: f32, end: f32, per_group_width: f32, local_offset: f32) -> array<vec3<f32>, 4u> {
223    let center = (f32(idx) + 1.0) + local_offset;
224    let half = per_group_width * 0.5;
225    let left = center - half;
226    let right = center + half;
227    let bottom = min(start, end);
228    let top = max(start, end);
229
230    return array<vec3<f32>, 4u>(
231        vec3<f32>(left, bottom, 0.0),
232        vec3<f32>(right, bottom, 0.0),
233        vec3<f32>(right, top, 0.0),
234        vec3<f32>(left, top, 0.0),
235    );
236}
237
238fn build_horizontal_quad(idx: u32, start: f32, end: f32, per_group_width: f32, local_offset: f32) -> array<vec3<f32>, 4u> {
239    let center = (f32(idx) + 1.0) + local_offset;
240    let half = per_group_width * 0.5;
241    let bottom = center - half;
242    let top = center + half;
243    let left = min(start, end);
244    let right = max(start, end);
245
246    return array<vec3<f32>, 4u>(
247        vec3<f32>(left, bottom, 0.0),
248        vec3<f32>(right, bottom, 0.0),
249        vec3<f32>(right, top, 0.0),
250        vec3<f32>(left, top, 0.0),
251    );
252}
253
254@compute @workgroup_size(WORKGROUP_SIZE)
255fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
256    let idx = gid.x;
257    if (idx >= params.row_count) {
258        return;
259    }
260
261    let safe_group_count = max(params.group_count, 1u);
262    let per_group_width = max(params.bar_width / f32(safe_group_count), 0.01);
263    let group_offset_start = -params.bar_width * 0.5;
264    let local_offset = group_offset_start
265        + per_group_width * f32(min(params.group_index, safe_group_count - 1u))
266        + per_group_width * 0.5;
267    let source_stride = max(params.source_row_count, 1u);
268    let value_index = if (params.transpose_source == 1u) {
269        idx * source_stride + params.series_index
270    } else {
271        params.series_index * source_stride + idx
272    };
273    let value = f32(values[value_index]);
274
275    var base_pos = 0.0;
276    var base_neg = 0.0;
277    if (params.layout == 1u && params.series_index > 0u) {
278        var col: u32 = 0u;
279        loop {
280            if (col >= params.series_index) {
281                break;
282            }
283            let prev_index = if (params.transpose_source == 1u) {
284                idx * source_stride + col
285            } else {
286                col * source_stride + idx
287            };
288            let prev = f32(values[prev_index]);
289            if (isFinite(prev)) {
290                if (prev >= 0.0) {
291                    base_pos += prev;
292                } else {
293                    base_neg += prev;
294                }
295            }
296            col = col + 1u;
297        }
298    }
299
300    var start = 0.0;
301    var end = 0.0;
302    if (!isFinite(value)) {
303        start = 0.0;
304        end = 0.0;
305    } else if (params.layout == 1u) {
306        if (value >= 0.0) {
307            start = base_pos;
308            end = base_pos + value;
309        } else {
310            start = base_neg + value;
311            end = base_neg;
312        }
313    } else {
314        start = 0.0;
315        end = value;
316    }
317
318    let quad = if (params.orientation == 0u) {
319        build_vertical_quad(idx, start, end, per_group_width, local_offset)
320    } else {
321        build_horizontal_quad(idx, start, end, per_group_width, local_offset)
322    };
323
324    let base_index = idx * VERTICES_PER_BAR;
325    write_vertices(base_index, quad);
326}
327"#;