qualia_core_db/wgsl_forge/emit/
spirv.rs

1//! SPIR-V emission target.
2//!
3//! Unlike the other native targets (HLSL/MSL/PTX/CUDA-C), which print source in
4//! the foreign language, the SPIR-V target produces a *binary* SPIR-V module by
5//! reusing the deterministic WGSL the [`emit_wgsl`] path already generates,
6//! parsing it back into a `naga::Module`, validating it with the same
7//! capabilities as [`crate::wgsl_forge::validate::validate_wgsl`]
8//! (`RAY_QUERY | COOPERATIVE_MATRIX | SHADER_FLOAT16`), then lowering it through
9//! naga's `spv-out` backend.
10//!
11//! ## Word encoding in [`GeneratedShader::source`]
12//!
13//! `GeneratedShader::source` is a `String`, but SPIR-V is a sequence of 32-bit
14//! words. We serialize the `Vec<u32>` as **`;`-joined decimal words** (e.g.
15//! `"119734787;65536;..."`). Decimal (rather than hex) keeps it unambiguous and
16//! trivially round-trippable with `split(';')` + `u32::from_str`; the first word
17//! is always the SPIR-V magic number `0x07230203` = `119734787` decimal.
18
19use super::{emit_wgsl, GeneratedShader};
20use crate::wgsl_forge::{ForgeError, KernelSpec, Schedule};
21
22/// SPIR-V opcodes used by the workgroup-size patcher.
23///
24/// naga emits `OpExecutionMode` with `LocalSize` mode (not the deprecated
25/// `OpLocalSize` instruction) to set the workgroup size.
26const OP_EXECUTION_MODE: u32 = 16;
27const OP_EXECUTION_MODE_LOCAL_SIZE_WORD_COUNT: u32 = 6;
28const EXEC_MODE_LOCAL_SIZE: u32 = 17; // LocalSize mode value
29
30/// Patch the workgroup size in a SPIR-V binary to `(x, 1, 1)`.
31///
32/// This is the alternative to `OpSpecConstant` + `OpExecutionModeId LocalSizeId`:
33/// naga does not support specialization constants for `@workgroup_size`, and
34/// wgpu does not expose `VkSpecializationInfo`. Instead, we emit SPIR-V once
35/// with a base workgroup size, then binary-patch the `OpExecutionMode LocalSize`
36/// words to produce variants for different schedules — avoiding a full naga
37/// re-parse + validate + spv-out pass per variant.
38///
39/// Returns `Err` if `OpExecutionMode LocalSize` is not found (malformed module).
40pub fn patch_spirv_workgroup_size(
41    words: &mut [u32],
42    workgroup_size: u32,
43) -> Result<(), ForgeError> {
44    // Search for OpExecutionMode with LocalSize mode:
45    // word[0] = (6 << 16) | 16, word[2] = 17 (LocalSize),
46    // word[3..5] = x, y, z.
47    let target = (OP_EXECUTION_MODE_LOCAL_SIZE_WORD_COUNT << 16) | OP_EXECUTION_MODE;
48    for i in 0..words
49        .len()
50        .saturating_sub(OP_EXECUTION_MODE_LOCAL_SIZE_WORD_COUNT as usize)
51    {
52        if words[i] == target && words[i + 2] == EXEC_MODE_LOCAL_SIZE {
53            words[i + 3] = workgroup_size;
54            words[i + 4] = 1;
55            words[i + 5] = 1;
56            return Ok(());
57        }
58    }
59    Err(ForgeError::Emission(
60        "OpExecutionMode LocalSize not found in SPIR-V binary".to_string(),
61    ))
62}
63
64/// Encode `Vec<u32>` SPIR-V words as a `;`-joined decimal string.
65fn encode_spirv_words(words: &[u32]) -> String {
66    let mut source = String::with_capacity(words.len() * 6);
67    for (index, word) in words.iter().enumerate() {
68        if index > 0 {
69            source.push(SPIRV_WORD_SEPARATOR);
70        }
71        source.push_str(&word.to_string());
72    }
73    source
74}
75
76/// Separator between decimal SPIR-V words in [`GeneratedShader::source`].
77pub const SPIRV_WORD_SEPARATOR: char = ';';
78
79/// Emit a validated SPIR-V module for `kernel`/`schedule`.
80///
81/// The returned [`GeneratedShader::source`] holds the SPIR-V words as a
82/// `;`-joined decimal string (see module docs); all other fields mirror the
83/// WGSL record this was derived from, except `source_hash`, which is recomputed
84/// over the SPIR-V word string so it identifies the actual emitted artifact.
85pub fn emit_spirv(kernel: &KernelSpec, schedule: Schedule) -> Result<GeneratedShader, ForgeError> {
86    // 1. Generate the deterministic WGSL the forge already produces.
87    let wgsl = emit_wgsl(kernel, schedule)?;
88
89    // 2. Parse it back into a naga module.
90    let module = naga::front::wgsl::parse_str(&wgsl.source)
91        .map_err(|error| ForgeError::WgslParse(error.emit_to_string(&wgsl.source)))?;
92
93    // 3. Validate with the same capabilities validate_wgsl uses; the returned
94    //    ModuleInfo is exactly what the spv backend needs as its `info` arg.
95    let mut validator = naga::valid::Validator::new(
96        naga::valid::ValidationFlags::all(),
97        naga::valid::Capabilities::RAY_QUERY
98            | naga::valid::Capabilities::COOPERATIVE_MATRIX
99            | naga::valid::Capabilities::SHADER_FLOAT16,
100    );
101    let info = validator
102        .validate(&module)
103        .map_err(|error| ForgeError::WgslValidation(format!("{error:?}")))?;
104
105    // 4. Lower to SPIR-V words. Default Options targets SPIR-V 1.0 with portable
106    //    flags; passing `None` for pipeline options emits every entry point.
107    let words =
108        naga::back::spv::write_vec(&module, &info, &naga::back::spv::Options::default(), None)
109            .map_err(|error| ForgeError::Emission(format!("SPIR-V backend: {error}")))?;
110
111    if words.is_empty() {
112        return Err(ForgeError::Emission(
113            "SPIR-V backend produced an empty module".to_string(),
114        ));
115    }
116
117    // 5. Encode as `;`-joined decimal words and rehash over that artifact.
118    let source = encode_spirv_words(&words);
119    let source_hash = blake3::hash(source.as_bytes()).to_hex().to_string();
120
121    Ok(GeneratedShader {
122        kernel_id: wgsl.kernel_id,
123        semantic_hash: wgsl.semantic_hash,
124        source_hash,
125        schedule: wgsl.schedule,
126        source,
127    })
128}
129
130/// Decode a `;`-joined decimal SPIR-V word string back into `Vec<u32>`.
131///
132/// Provided so consumers (pipeline upload, on-disk caches) can recover the
133/// binary module from a [`GeneratedShader`] without re-deriving the encoding.
134pub fn decode_spirv_words(source: &str) -> Result<Vec<u32>, ForgeError> {
135    source
136        .split(SPIRV_WORD_SEPARATOR)
137        .map(|token| {
138            token.parse::<u32>().map_err(|error| {
139                ForgeError::Emission(format!("invalid SPIR-V word {token:?}: {error}"))
140            })
141        })
142        .collect()
143}
144
145/// Emit SPIR-V and patch the workgroup size to `schedule.workgroup_size`.
146///
147/// This is the specialization-constant alternative: emit once with naga's
148/// default, then binary-patch `OpLocalSize` to the desired size. Avoids
149/// re-running naga for each schedule variant.
150pub fn emit_spirv_patched(
151    kernel: &KernelSpec,
152    schedule: Schedule,
153) -> Result<GeneratedShader, ForgeError> {
154    let mut words = {
155        let base = emit_spirv(kernel, schedule)?;
156        decode_spirv_words(&base.source)?
157    };
158    patch_spirv_workgroup_size(&mut words, schedule.workgroup_size)?;
159    let source = encode_spirv_words(&words);
160    let source_hash = blake3::hash(source.as_bytes()).to_hex().to_string();
161    Ok(GeneratedShader {
162        kernel_id: kernel.id.clone(),
163        semantic_hash: kernel.semantic_hash()?,
164        source_hash,
165        schedule,
166        source,
167    })
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173    use crate::wgsl_forge::{BuiltinKernel, Schedule};
174
175    #[test]
176    fn affine_emits_non_empty_spirv_words() {
177        let kernel = BuiltinKernel::AffineF32.spec();
178        let generated = emit_spirv(&kernel, Schedule::default()).expect("spirv emission");
179        let words = decode_spirv_words(&generated.source).expect("decode words");
180        assert!(!words.is_empty(), "SPIR-V module must contain words");
181        // First word of any SPIR-V module is the magic number 0x07230203.
182        assert_eq!(words[0], 0x0723_0203, "SPIR-V magic number header");
183        assert_eq!(generated.kernel_id, kernel.id);
184    }
185
186    #[test]
187    fn patch_workgroup_size_modifies_op_local_size() {
188        let kernel = BuiltinKernel::AffineF32.spec();
189        let base = emit_spirv(&kernel, Schedule::default()).expect("spirv emission");
190        let mut words = decode_spirv_words(&base.source).expect("decode words");
191
192        // Find the original OpExecutionMode LocalSize and record its x value.
193        let target = (OP_EXECUTION_MODE_LOCAL_SIZE_WORD_COUNT << 16) | OP_EXECUTION_MODE;
194        let mut orig_x = 0;
195        for i in 0..words.len().saturating_sub(6) {
196            if words[i] == target && words[i + 2] == EXEC_MODE_LOCAL_SIZE {
197                orig_x = words[i + 3];
198                break;
199            }
200        }
201        assert!(
202            orig_x > 0,
203            "OpExecutionMode LocalSize should exist in emitted SPIR-V"
204        );
205
206        // Patch to 128.
207        patch_spirv_workgroup_size(&mut words, 128).expect("patch");
208
209        // Verify the patch took effect.
210        for i in 0..words.len().saturating_sub(6) {
211            if words[i] == target && words[i + 2] == EXEC_MODE_LOCAL_SIZE {
212                assert_eq!(words[i + 3], 128, "workgroup x patched to 128");
213                assert_eq!(words[i + 4], 1, "workgroup y = 1");
214                assert_eq!(words[i + 5], 1, "workgroup z = 1");
215                return;
216            }
217        }
218        panic!("OpExecutionMode LocalSize disappeared after patch");
219    }
220
221    #[test]
222    fn emit_spirv_patched_produces_valid_variant() {
223        let kernel = BuiltinKernel::AffineF32.spec();
224        let schedule = Schedule {
225            workgroup_size: 128,
226            ..Default::default()
227        };
228        let generated = emit_spirv_patched(&kernel, schedule).expect("patched spirv");
229        let words = decode_spirv_words(&generated.source).expect("decode");
230        assert!(!words.is_empty());
231        assert_eq!(words[0], 0x0723_0203, "magic number preserved");
232
233        // Verify workgroup size is 128 in the binary.
234        let target = (OP_EXECUTION_MODE_LOCAL_SIZE_WORD_COUNT << 16) | OP_EXECUTION_MODE;
235        for i in 0..words.len().saturating_sub(6) {
236            if words[i] == target && words[i + 2] == EXEC_MODE_LOCAL_SIZE {
237                assert_eq!(words[i + 3], 128);
238                return;
239            }
240        }
241        panic!("OpExecutionMode LocalSize not found in patched SPIR-V");
242    }
243}
qualia_core_db/wgsl_forge/emit/spirv.rs

qualia_core_db/wgsl_forge/emit/
spirv.rs
