vyre-driver-wgpu 0.7.2

wgpu backend for vyre IR - implements VyreBackend, owns GPU runtime, buffer pool, pipeline cache
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

use vyre_driver::BackendError;
use vyre_megakernel::{
    compile_selected_modules, Artifact, EmittedTargetModule, TargetCompileError, TargetCompiler,
    TargetPayload, TargetPayloadFormat, TargetProfile,
};

pub(crate) const WGPU_TARGET_FORMAT: &str = "wgsl";
pub(crate) const WGPU_TARGET_FORMAT_VERSION: u16 = 2;

pub(crate) const WGPU_TARGET_MODULE_SCHEMA_VERSION: u16 = 2;

#[derive(serde::Serialize, serde::Deserialize)]
pub(crate) struct WgpuTargetModule {
    pub(crate) schema_version: u16,
    pub(crate) wgsl: String,
}

pub(crate) struct WgpuTargetCompiler {
    format: TargetPayloadFormat,
    profile: TargetProfile,
}

impl TargetCompiler for WgpuTargetCompiler {
    fn format(&self) -> &TargetPayloadFormat {
        &self.format
    }

    fn profile(&self) -> &TargetProfile {
        &self.profile
    }

    fn compile(&self, artifact: &Artifact) -> Result<TargetPayload, TargetCompileError> {
        compile_selected_modules(
            artifact,
            self.format.clone(),
            self.profile.clone(),
            |selected, _profile| {
                let descriptor = selected.descriptor.clone();
                let module =
                    crate::emit::emit_naga_module_for_descriptor(&descriptor).map_err(|error| {
                        TargetCompileError::Emission(format!("WGSL emission failed: {error}"))
                    })?;
                let wgsl = crate::emit::write_wgsl(&module).map_err(|error| {
                    TargetCompileError::Emission(format!("WGSL writing failed: {error}"))
                })?;
                let target = WgpuTargetModule {
                    schema_version: WGPU_TARGET_MODULE_SCHEMA_VERSION,
                    wgsl,
                };
                let grid_size = vyre_driver::infer_dispatch_grid_for_count(
                    selected.logical_element_count,
                    selected.descriptor.dispatch.workgroup_size,
                )
                .map_err(|error| TargetCompileError::Emission(error.to_string()))?;
                serde_json::to_vec(&target)
                    .map(|bytes| EmittedTargetModule {
                        entry_point: "main".to_string(),
                        grid_size,
                        dynamic_shared_bytes: 0,
                        workgroup_size: selected.descriptor.dispatch.workgroup_size,
                        resource_bindings: selected.canonical_bindings.clone(),
                        bytes,
                    })
                    .map_err(|error| {
                        TargetCompileError::Emission(format!(
                            "WGSL target module serialization failed: {error}"
                        ))
                    })
            },
        )
    }
}

pub(crate) fn target_profile() -> Result<TargetProfile, BackendError> {
    TargetProfile::new(
        WGPU_TARGET_FORMAT,
        u64::from(WGPU_TARGET_FORMAT_VERSION),
        [256, 256, 64],
        256,
        16_384,
        0,
    )
    .map_err(|error| BackendError::KernelCompileFailed {
        backend: "wgpu".to_string(),
        compiler_message: format!(
            "WGSL target profile is invalid: {error}. Fix: repair the registered profile."
        ),
    })
}

pub(crate) fn target_compiler_factory() -> Result<Box<dyn TargetCompiler>, BackendError> {
    let format = TargetPayloadFormat::new(WGPU_TARGET_FORMAT, WGPU_TARGET_FORMAT_VERSION).map_err(|error| {
        BackendError::KernelCompileFailed {
            backend: "wgpu".to_string(),
            compiler_message: format!(
                "WGSL target format is invalid: {error}. Fix: repair the registered format identity."
            ),
        }
    })?;
    let profile = target_profile()?;
    Ok(Box::new(WgpuTargetCompiler { format, profile }))
}