atagen

monolith

public
Star0
Watch0
43e0e8e33c3e08f47bcca32dfe53708e351b36ea
1
ssh://brewery@pls.lobotomise.me:33/atagen/monolith.git
crates/render-vk/src/pipeline.rsRust319 lines11.8 KB
1
use crate::Error;
2
use crate::device::Shared;
3
use crate::format::{SUPPORTED, vk_format};
4
use ash::vk;
5
use std::collections::HashMap;
6
use std::path::PathBuf;
7
use std::sync::Arc;
8
use std::sync::atomic::{AtomicU32, Ordering};
9
use std::time::Instant;
10
11
/// Push constants shared by every shader, see `shaders/push.glsl`.
12
/// `dst` is the quad in normalized device coordinates as (x, y, w, h),
13
/// `uv` the texture coordinate at each corner (TL, TR, BL, BR).
14
#[repr(C)]
15
#[derive(Clone, Copy, Debug, Default)]
16
pub(crate) struct Push {
17
    pub color: [f32; 4],
18
    pub dst: [f32; 4],
19
    pub uv: [[f32; 2]; 4],
20
    /// Transition progress and kind, read only by `shaders/effect.frag`.
21
    pub effect: [f32; 4],
22
    /// See `shaders/clip.glsl`. A zero width turns clipping off.
23
    pub clip: [f32; 4],
24
    pub radii: [f32; 4],
25
    pub ring: [f32; 4],
26
}
27
28
impl Push {
29
    pub fn bytes(&self) -> &[u8] {
30
        // SAFETY: `Push` is `repr(C)` and all f32, so it has no padding bytes.
31
        unsafe { std::slice::from_raw_parts((self as *const Self).cast::<u8>(), size_of::<Self>()) }
32
    }
33
}
34
35
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
36
pub(crate) enum Kind {
37
    Solid,
38
    Texture,
39
    Effect,
40
}
41
42
pub(crate) struct Pipelines {
43
    shared: Arc<Shared>,
44
    pub set_layout: vk::DescriptorSetLayout,
45
    pub layout: vk::PipelineLayout,
46
    vert: vk::ShaderModule,
47
    solid: vk::ShaderModule,
48
    texture: vk::ShaderModule,
49
    effect: vk::ShaderModule,
50
    cache: HashMap<(vk::Format, Kind), vk::Pipeline>,
51
    vk_cache: vk::PipelineCache,
52
    cache_path: Option<PathBuf>,
53
    loaded_len: usize,
54
}
55
56
const QUAD_VERT: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/quad.vert.spv"));
57
const SOLID_FRAG: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/solid.frag.spv"));
58
const TEXTURE_FRAG: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/texture.frag.spv"));
59
const EFFECT_FRAG: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/effect.frag.spv"));
60
61
impl Pipelines {
62
    pub fn new(shared: &Arc<Shared>) -> Result<Self, Error> {
63
        let d = &shared.device.device;
64
        let bindings = [vk::DescriptorSetLayoutBinding::default()
65
            .binding(0)
66
            .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
67
            .descriptor_count(1)
68
            .stage_flags(vk::ShaderStageFlags::FRAGMENT)];
69
        let set_layout = unsafe {
70
            d.create_descriptor_set_layout(
71
                &vk::DescriptorSetLayoutCreateInfo::default()
72
                    .flags(vk::DescriptorSetLayoutCreateFlags::PUSH_DESCRIPTOR_KHR)
73
                    .bindings(&bindings),
74
                None,
75
            )?
76
        };
77
        let ranges = [vk::PushConstantRange::default()
78
            .stage_flags(vk::ShaderStageFlags::VERTEX | vk::ShaderStageFlags::FRAGMENT)
79
            .size(size_of::<Push>() as u32)];
80
        let set_layouts = [set_layout];
81
        let layout = unsafe {
82
            d.create_pipeline_layout(
83
                &vk::PipelineLayoutCreateInfo::default()
84
                    .set_layouts(&set_layouts)
85
                    .push_constant_ranges(&ranges),
86
                None,
87
            )?
88
        };
89
        let module = |code: &[u8]| -> Result<vk::ShaderModule, Error> {
90
            let words = ash::util::read_spv(&mut std::io::Cursor::new(code))
91
                .map_err(|_| vk::Result::ERROR_INVALID_SHADER_NV)?;
92
            Ok(unsafe {
93
                d.create_shader_module(&vk::ShaderModuleCreateInfo::default().code(&words), None)?
94
            })
95
        };
96
        let cache_path = cache_path();
97
        let props = unsafe {
98
            shared
99
                .device
100
                .instance
101
                .get_physical_device_properties(shared.device.physical)
102
        };
103
        let loaded = cache_path
104
            .as_ref()
105
            .and_then(|p| std::fs::read(p).ok())
106
            .filter(|data| header_matches(data, &props))
107
            .unwrap_or_default();
108
        let create = |data: &[u8]| unsafe {
109
            d.create_pipeline_cache(
110
                &vk::PipelineCacheCreateInfo::default().initial_data(data),
111
                None,
112
            )
113
        };
114
        let vk_cache = create(&loaded).or_else(|_| create(&[])).unwrap_or_default();
115
        Ok(Self {
116
            shared: shared.clone(),
117
            set_layout,
118
            layout,
119
            vert: module(QUAD_VERT)?,
120
            solid: module(SOLID_FRAG)?,
121
            texture: module(TEXTURE_FRAG)?,
122
            effect: module(EFFECT_FRAG)?,
123
            cache: HashMap::new(),
124
            vk_cache,
125
            cache_path,
126
            loaded_len: loaded.len(),
127
        })
128
    }
129
130
    /// Compiles the pipelines a frame draws with for every render format, so
131
    /// the first spawn animation does not stall on the driver compiler.
132
    pub fn warm(&mut self) {
133
        let began = Instant::now();
134
        for fourcc in SUPPORTED {
135
            let Some(format) = vk_format(fourcc) else {
136
                continue;
137
            };
138
            for kind in [Kind::Solid, Kind::Texture, Kind::Effect] {
139
                if let Err(e) = self.get(format.format, kind) {
140
                    tracing::warn!(?kind, "cannot warm a pipeline: {e}");
141
                    return;
142
                }
143
            }
144
        }
145
        tracing::debug!(elapsed = ?began.elapsed(), "warmed pipelines");
146
    }
147
148
    fn save_cache(&self) {
149
        let Some(path) = &self.cache_path else {
150
            return;
151
        };
152
        if self.vk_cache == vk::PipelineCache::null() {
153
            return;
154
        }
155
        let Ok(data) = (unsafe {
156
            self.shared
157
                .device
158
                .device
159
                .get_pipeline_cache_data(self.vk_cache)
160
        }) else {
161
            return;
162
        };
163
        if data.len() == self.loaded_len {
164
            return;
165
        }
166
        static WRITERS: AtomicU32 = AtomicU32::new(0);
167
        let tmp = path.with_extension(format!(
168
            "{}.{}.tmp",
169
            std::process::id(),
170
            WRITERS.fetch_add(1, Ordering::Relaxed)
171
        ));
172
        let written = path
173
            .parent()
174
            .map_or(Ok(()), std::fs::create_dir_all)
175
            .and_then(|()| std::fs::write(&tmp, &data))
176
            .and_then(|()| std::fs::rename(&tmp, path));
177
        if let Err(e) = written {
178
            tracing::debug!("cannot save the pipeline cache: {e}");
179
            let _ = std::fs::remove_file(&tmp);
180
        }
181
    }
182
183
    pub fn get(&mut self, format: vk::Format, kind: Kind) -> Result<vk::Pipeline, Error> {
184
        if let Some(&p) = self.cache.get(&(format, kind)) {
185
            return Ok(p);
186
        }
187
        let d = &self.shared.device.device;
188
        let frag = match kind {
189
            Kind::Solid => self.solid,
190
            Kind::Texture => self.texture,
191
            Kind::Effect => self.effect,
192
        };
193
        let stages = [
194
            vk::PipelineShaderStageCreateInfo::default()
195
                .stage(vk::ShaderStageFlags::VERTEX)
196
                .module(self.vert)
197
                .name(c"main"),
198
            vk::PipelineShaderStageCreateInfo::default()
199
                .stage(vk::ShaderStageFlags::FRAGMENT)
200
                .module(frag)
201
                .name(c"main"),
202
        ];
203
        let vertex = vk::PipelineVertexInputStateCreateInfo::default();
204
        let assembly = vk::PipelineInputAssemblyStateCreateInfo::default()
205
            .topology(vk::PrimitiveTopology::TRIANGLE_STRIP);
206
        let viewport = vk::PipelineViewportStateCreateInfo::default()
207
            .viewport_count(1)
208
            .scissor_count(1);
209
        let raster = vk::PipelineRasterizationStateCreateInfo::default()
210
            .polygon_mode(vk::PolygonMode::FILL)
211
            .cull_mode(vk::CullModeFlags::NONE)
212
            .line_width(1.0);
213
        let multisample = vk::PipelineMultisampleStateCreateInfo::default()
214
            .rasterization_samples(vk::SampleCountFlags::TYPE_1);
215
        let attachments = [vk::PipelineColorBlendAttachmentState::default()
216
            .blend_enable(true)
217
            .src_color_blend_factor(vk::BlendFactor::ONE)
218
            .dst_color_blend_factor(vk::BlendFactor::ONE_MINUS_SRC_ALPHA)
219
            .color_blend_op(vk::BlendOp::ADD)
220
            .src_alpha_blend_factor(vk::BlendFactor::ONE)
221
            .dst_alpha_blend_factor(vk::BlendFactor::ONE_MINUS_SRC_ALPHA)
222
            .alpha_blend_op(vk::BlendOp::ADD)
223
            .color_write_mask(vk::ColorComponentFlags::RGBA)];
224
        let blend = vk::PipelineColorBlendStateCreateInfo::default().attachments(&attachments);
225
        let dynamic_states = [vk::DynamicState::VIEWPORT, vk::DynamicState::SCISSOR];
226
        let dynamic = vk::PipelineDynamicStateCreateInfo::default().dynamic_states(&dynamic_states);
227
        let formats = [format];
228
        let mut rendering =
229
            vk::PipelineRenderingCreateInfo::default().color_attachment_formats(&formats);
230
        let info = vk::GraphicsPipelineCreateInfo::default()
231
            .stages(&stages)
232
            .vertex_input_state(&vertex)
233
            .input_assembly_state(&assembly)
234
            .viewport_state(&viewport)
235
            .rasterization_state(&raster)
236
            .multisample_state(&multisample)
237
            .color_blend_state(&blend)
238
            .dynamic_state(&dynamic)
239
            .layout(self.layout)
240
            .push_next(&mut rendering);
241
        let pipeline = unsafe { d.create_graphics_pipelines(self.vk_cache, &[info], None) }
242
            .map_err(|(_, e)| e)?[0];
243
        self.cache.insert((format, kind), pipeline);
244
        Ok(pipeline)
245
    }
246
}
247
248
impl Drop for Pipelines {
249
    fn drop(&mut self) {
250
        self.save_cache();
251
        let d = &self.shared.device.device;
252
        unsafe {
253
            for (_, p) in self.cache.drain() {
254
                d.destroy_pipeline(p, None);
255
            }
256
            d.destroy_pipeline_cache(self.vk_cache, None);
257
            d.destroy_shader_module(self.vert, None);
258
            d.destroy_shader_module(self.solid, None);
259
            d.destroy_shader_module(self.texture, None);
260
            d.destroy_shader_module(self.effect, None);
261
            d.destroy_pipeline_layout(self.layout, None);
262
            d.destroy_descriptor_set_layout(self.set_layout, None);
263
        }
264
    }
265
}
266
267
fn cache_path() -> Option<PathBuf> {
268
    let base = std::env::var_os("XDG_CACHE_HOME")
269
        .filter(|v| !v.is_empty())
270
        .map(PathBuf::from)
271
        .or_else(|| std::env::var_os("HOME").map(|home| PathBuf::from(home).join(".cache")))?;
272
    Some(base.join("monolith").join("pipelines.bin"))
273
}
274
275
/// Whether `data` opens with a cache header written by this driver on this
276
/// device. The spec lets a driver reject a foreign cache but some crash on
277
/// one instead, so the header is checked first.
278
fn header_matches(data: &[u8], props: &vk::PhysicalDeviceProperties) -> bool {
279
    let word = |at: usize| {
280
        data.get(at..at + 4)
281
            .and_then(|b| b.try_into().ok())
282
            .map(u32::from_ne_bytes)
283
    };
284
    word(0).is_some_and(|size| size >= 32 && size as usize <= data.len())
285
        && word(4) == Some(vk::PipelineCacheHeaderVersion::ONE.as_raw() as u32)
286
        && word(8) == Some(props.vendor_id)
287
        && word(12) == Some(props.device_id)
288
        && data.get(16..32) == Some(&props.pipeline_cache_uuid[..])
289
}
290
291
#[cfg(test)]
292
mod tests {
293
    use super::*;
294
295
    fn props() -> vk::PhysicalDeviceProperties {
296
        vk::PhysicalDeviceProperties {
297
            vendor_id: 0x10de,
298
            device_id: 0x2487,
299
            pipeline_cache_uuid: [7; 16],
300
            ..Default::default()
301
        }
302
    }
303
304
    fn header(vendor: u32, uuid: u8) -> Vec<u8> {
305
        let words = [32u32, 1, vendor, 0x2487];
306
        let mut data: Vec<u8> = words.iter().flat_map(|w| w.to_ne_bytes()).collect();
307
        data.extend([uuid; 16]);
308
        data
309
    }
310
311
    #[test]
312
    fn cache_header_must_match_the_device_and_driver() {
313
        assert!(header_matches(&header(0x10de, 7), &props()));
314
        assert!(!header_matches(&header(0x1002, 7), &props()));
315
        assert!(!header_matches(&header(0x10de, 8), &props()));
316
        assert!(!header_matches(&header(0x10de, 7)[..20], &props()));
317
        assert!(!header_matches(&[], &props()));
318
    }
319
}