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

runmat_plot/gpu/shaders/
stem.rs
