runmat_plot/gpu/shaders/
stem.rs

1pub const F32: &str = r#"
2const WORKGROUP_SIZE: u32 = {{WORKGROUP_SIZE}}u;
3
4struct VertexRaw {
5    data: array<f32, 12u>,
6};
7
8struct StemParams {
9    color: vec4<f32>,
10    baseline_color: vec4<f32>,
11    baseline: f32,
12    min_x: f32,
13    max_x: f32,
14    point_count: u32,
15    line_style: u32,
16    baseline_visible: u32,
17};
18
19@group(0) @binding(0)
20var<storage, read> buf_x: array<f32>;
21
22@group(0) @binding(1)
23var<storage, read> buf_y: array<f32>;
24
25@group(0) @binding(2)
26var<storage, read_write> out_vertices: array<VertexRaw>;
27
28@group(0) @binding(3)
29var<uniform> params: StemParams;
30
31fn write_vertex(index: u32, pos: vec3<f32>, color: vec4<f32>) {
32    var v: VertexRaw;
33    v.data[0u] = pos.x;
34    v.data[1u] = pos.y;
35    v.data[2u] = 0.0;
36    v.data[3u] = color.x;
37    v.data[4u] = color.y;
38    v.data[5u] = color.z;
39    v.data[6u] = color.w;
40    v.data[7u] = 0.0;
41    v.data[8u] = 0.0;
42    v.data[9u] = 1.0;
43    v.data[10u] = 0.0;
44    v.data[11u] = 0.0;
45    out_vertices[index] = v;
46}
47
48fn include_segment(i: u32, style: u32) -> bool {
49    switch (style) {
50        case 0u: { return true; }
51        case 1u: { return (i % 4u) < 2u; }
52        case 2u: { return (i % 4u) == 0u; }
53        case 3u: {
54            let m = i % 6u;
55            return (m < 2u) || (m == 3u);
56        }
57        default: { return true; }
58    }
59}
60
61@compute @workgroup_size(WORKGROUP_SIZE)
62fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
63    let i = gid.x;
64    if (i >= params.point_count) {
65        return;
66    }
67
68    if (i == 0u && params.baseline_visible != 0u) {
69        write_vertex(0u, vec3<f32>(params.min_x, params.baseline, 0.0), params.baseline_color);
70        write_vertex(1u, vec3<f32>(params.max_x, params.baseline, 0.0), params.baseline_color);
71    }
72
73    if (!include_segment(i, params.line_style)) {
74        return;
75    }
76
77    let base = 2u + i * 2u;
78    write_vertex(base, vec3<f32>(buf_x[i], params.baseline, 0.0), params.color);
79    write_vertex(base + 1u, vec3<f32>(buf_x[i], buf_y[i], 0.0), params.color);
80}
81"#;
82
83pub const F64: &str = r#"
84const WORKGROUP_SIZE: u32 = {{WORKGROUP_SIZE}}u;
85
86struct VertexRaw {
87    data: array<f32, 12u>,
88};
89
90struct StemParams {
91    color: vec4<f32>,
92    baseline_color: vec4<f32>,
93    baseline: f32,
94    min_x: f32,
95    max_x: f32,
96    point_count: u32,
97    line_style: u32,
98    baseline_visible: u32,
99};
100
101@group(0) @binding(0)
102var<storage, read> buf_x: array<f64>;
103
104@group(0) @binding(1)
105var<storage, read> buf_y: array<f64>;
106
107@group(0) @binding(2)
108var<storage, read_write> out_vertices: array<VertexRaw>;
109
110@group(0) @binding(3)
111var<uniform> params: StemParams;
112
113fn write_vertex(index: u32, pos: vec3<f32>, color: vec4<f32>) {
114    var v: VertexRaw;
115    v.data[0u] = pos.x;
116    v.data[1u] = pos.y;
117    v.data[2u] = 0.0;
118    v.data[3u] = color.x;
119    v.data[4u] = color.y;
120    v.data[5u] = color.z;
121    v.data[6u] = color.w;
122    v.data[7u] = 0.0;
123    v.data[8u] = 0.0;
124    v.data[9u] = 1.0;
125    v.data[10u] = 0.0;
126    v.data[11u] = 0.0;
127    out_vertices[index] = v;
128}
129
130fn include_segment(i: u32, style: u32) -> bool {
131    switch (style) {
132        case 0u: { return true; }
133        case 1u: { return (i % 4u) < 2u; }
134        case 2u: { return (i % 4u) == 0u; }
135        case 3u: {
136            let m = i % 6u;
137            return (m < 2u) || (m == 3u);
138        }
139        default: { return true; }
140    }
141}
142
143@compute @workgroup_size(WORKGROUP_SIZE)
144fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
145    let i = gid.x;
146    if (i >= params.point_count) {
147        return;
148    }
149
150    if (i == 0u && params.baseline_visible != 0u) {
151        write_vertex(0u, vec3<f32>(params.min_x, params.baseline, 0.0), params.baseline_color);
152        write_vertex(1u, vec3<f32>(params.max_x, params.baseline, 0.0), params.baseline_color);
153    }
154
155    if (!include_segment(i, params.line_style)) {
156        return;
157    }
158
159    let base = 2u + i * 2u;
160    write_vertex(base, vec3<f32>(f32(buf_x[i]), params.baseline, 0.0), params.color);
161    write_vertex(base + 1u, vec3<f32>(f32(buf_x[i]), f32(buf_y[i]), 0.0), params.color);
162}
163"#;
runmat_plot/gpu/shaders/stem.rs

runmat_plot/gpu/shaders/
stem.rs
