vtsampler 0.1.0

A pure Rust implementation similar to ID3D11VideoProcessor
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

use std::collections::HashMap;

use wgpu::{
    BindGroupLayout, BindGroupLayoutDescriptor, BindGroupLayoutEntry, BindingType,
    ComputePipeline, ComputePipelineDescriptor, Device, PipelineCompilationOptions,
    ShaderStages, StorageTextureAccess, TextureSampleType, TextureViewDimension,
};

use crate::{
    VTFormat,
    format::{VTColorSpace, VTSampleError, VTScaleFilter},
    shader::{compile_wgsl, create_shader_module},
};

#[derive(Hash, Eq, PartialEq, Clone, Copy, Debug)]
pub struct PipelineKey {
    pub input: VTFormat,
    pub output: VTFormat,
    pub need_scale: bool,
    pub color_space: VTColorSpace,
    pub scale_filter: VTScaleFilter,
}

pub(crate) struct PipelineSet {
    pub layout: BindGroupLayout,
    pub pipeline: ComputePipeline,
}

pub struct ShaderRegistry {
    pipelines: HashMap<PipelineKey, PipelineSet>,
}

impl ShaderRegistry {
    pub fn new() -> Self {
        Self {
            pipelines: HashMap::new(),
        }
    }

    pub fn get(
        &mut self,
        device: &Device,
        key: PipelineKey,
    ) -> Result<&PipelineSet, VTSampleError> {
        if let std::collections::hash_map::Entry::Vacant(e) = self.pipelines.entry(key) {
            e.insert(Self::build_pipeline(device, key)?);
        }
        Ok(self.pipelines.get(&key).unwrap())
    }

    fn build_pipeline(device: &Device, key: PipelineKey) -> Result<PipelineSet, VTSampleError> {
        let source = compile_wgsl(
            key.input,
            key.output,
            key.need_scale,
            key.color_space,
            key.scale_filter,
        )?;
        let module = create_shader_module(device, &source);
        let layout = Self::create_bind_group_layout(device, key.input, key.output);
        let pipeline = device.create_compute_pipeline(&ComputePipelineDescriptor {
            label: Some("vtsampler_process"),
            layout: Some(
                &device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                    label: Some("vtsampler_process_layout"),
                    bind_group_layouts: &[&layout],
                    push_constant_ranges: &[],
                }),
            ),
            module: &module,
            entry_point: Some("main"),
            compilation_options: PipelineCompilationOptions::default(),
            cache: None,
        });
        Ok(PipelineSet { layout, pipeline })
    }

    fn create_bind_group_layout(
        device: &Device,
        input: VTFormat,
        output: VTFormat,
    ) -> BindGroupLayout {
        let mut entries = Vec::new();
        for i in 0..input.plane_count() {
            entries.push(BindGroupLayoutEntry {
                binding: i as u32,
                visibility: ShaderStages::COMPUTE,
                ty: BindingType::Texture {
                    sample_type: TextureSampleType::Float { filterable: true },
                    view_dimension: TextureViewDimension::D2,
                    multisampled: false,
                },
                count: None,
            });
        }
        for (i, format) in output.plane_formats().iter().enumerate() {
            entries.push(BindGroupLayoutEntry {
                binding: (input.plane_count() + i) as u32,
                visibility: ShaderStages::COMPUTE,
                ty: BindingType::StorageTexture {
                    access: StorageTextureAccess::WriteOnly,
                    format: *format,
                    view_dimension: TextureViewDimension::D2,
                },
                count: None,
            });
        }
        device.create_bind_group_layout(&BindGroupLayoutDescriptor {
            label: Some("vtsampler_bind_group_layout"),
            entries: &entries,
        })
    }
}