spirv_webgpu_transform/
splitcombined.rs

1use super::*;
2
3mod function_call;
4mod function_parameter;
5mod load;
6mod type_function;
7mod type_pointer;
8mod variable;
9
10use function_call::*;
11use function_parameter::*;
12use load::*;
13use type_function::*;
14use type_pointer::*;
15use variable::*;
16
17/// Perform the operation on a `Vec<u32>`.
18/// Use [u8_slice_to_u32_vec] to convert a `&[u8]` into a `Vec<u32>`
19pub fn combimgsampsplitter(
20    in_spv: &[u32],
21    corrections: &mut CorrectionMap,
22) -> Result<Vec<u32>, ()> {
23    let spv = in_spv.to_owned();
24
25    let mut instruction_bound = spv[SPV_HEADER_INSTRUCTION_BOUND_OFFSET];
26    let magic_number = spv[SPV_HEADER_MAGIC_NUM_OFFSET];
27
28    let spv_header = spv[0..SPV_HEADER_LENGTH].to_owned();
29
30    assert_eq!(magic_number, SPV_HEADER_MAGIC);
31
32    let mut instruction_inserts = vec![];
33    let mut word_inserts = vec![];
34
35    let spv = spv.into_iter().skip(SPV_HEADER_LENGTH).collect::<Vec<_>>();
36    let mut new_spv = spv.clone();
37
38    let mut op_type_sampler_idx = None;
39    let mut first_op_deocrate_idx = None;
40    let mut first_op_type_void_idx = None;
41
42    let mut op_type_image_idxs = vec![];
43    let mut op_type_sampled_image_idxs = vec![];
44    let mut op_type_pointer_idxs = vec![];
45    let mut op_variables_idxs = vec![];
46    let mut op_loads_idxs = vec![];
47    let mut op_decorate_idxs = vec![];
48    let mut op_type_function_idxs = vec![];
49    let mut op_function_parameter_idxs = vec![];
50    let mut op_function_call_idxs = vec![];
51
52    // 1. Find locations instructions we need
53    let mut spv_idx = 0;
54    while spv_idx < spv.len() {
55        let op = spv[spv_idx];
56        let word_count = hiword(op);
57        let instruction = loword(op);
58
59        match instruction {
60            SPV_INSTRUCTION_OP_TYPE_VOID => {
61                first_op_type_void_idx = Some(spv_idx);
62            }
63            SPV_INSTRUCTION_OP_TYPE_SAMPLER => {
64                op_type_sampler_idx = Some(spv_idx);
65                new_spv[spv_idx] = encode_word(word_count, SPV_INSTRUCTION_OP_NOP);
66            }
67            SPV_INSTRUCTION_OP_TYPE_IMAGE => {
68                op_type_image_idxs.push(spv_idx);
69            }
70            SPV_INSTRUCTION_OP_TYPE_SAMPLED_IMAGE => op_type_sampled_image_idxs.push(spv_idx),
71            SPV_INSTRUCTION_OP_TYPE_POINTER => {
72                // This should probably go elsewhere.
73                #[allow(clippy::collapsible_match)]
74                if spv[spv_idx + 2] == SPV_STORAGE_CLASS_UNIFORM_CONSTANT {
75                    op_type_pointer_idxs.push(spv_idx);
76                }
77            }
78            SPV_INSTRUCTION_OP_VARIABLE => op_variables_idxs.push(spv_idx),
79            SPV_INSTRUCTION_OP_LOAD => op_loads_idxs.push(spv_idx),
80            SPV_INSTRUCTION_OP_DECORATE => {
81                op_decorate_idxs.push(spv_idx);
82                first_op_deocrate_idx.get_or_insert(spv_idx);
83            }
84            SPV_INSTRUCTION_OP_TYPE_FUNCTION => op_type_function_idxs.push(spv_idx),
85            SPV_INSTRUCTION_OP_FUNCTION_PARAMETER => op_function_parameter_idxs.push(spv_idx),
86            SPV_INSTRUCTION_OP_FUNCTION_CALL => op_function_call_idxs.push(spv_idx),
87
88            _ => {}
89        }
90
91        spv_idx += word_count as usize;
92    }
93
94    // 2. Insert OpTypeSampler and respective OpTypePointer if neccessary
95
96    // - If there has been no OpTypeImage, there will be nothing to do
97    if op_type_image_idxs.is_empty() {
98        return Ok(in_spv.to_vec());
99    };
100
101    let op_type_sampler_res_id = if let Some(idx) = op_type_sampler_idx {
102        spv[idx + 1]
103    } else {
104        let op_type_sampler_res_id = instruction_bound;
105        instruction_bound += 1;
106        op_type_sampler_res_id
107    };
108
109    let op_type_pointer_sampler_res_id = instruction_bound;
110    instruction_bound += 1;
111    instruction_inserts.push(InstructionInsert {
112        // Let's avoid trouble and just insert after OpTypeVoid.
113        // previous_spv_idx: op_type_image_idx,
114        previous_spv_idx: first_op_type_void_idx.unwrap(),
115        instruction: vec![
116            encode_word(2, SPV_INSTRUCTION_OP_TYPE_SAMPLER),
117            op_type_sampler_res_id,
118            encode_word(4, SPV_INSTRUCTION_OP_TYPE_POINTER),
119            op_type_pointer_sampler_res_id,
120            SPV_STORAGE_CLASS_UNIFORM_CONSTANT,
121            op_type_sampler_res_id,
122        ],
123    });
124
125    // 3. OpTypePointer
126    let tp_res = type_pointer(TypePointerIn {
127        spv: &spv,
128        new_spv: &mut new_spv,
129
130        op_type_pointer_idxs: &op_type_pointer_idxs,
131        op_type_sampled_image_idxs: &op_type_sampled_image_idxs,
132    });
133
134    // 4. OpVariable
135    let v_res = variable(VariableIn {
136        spv: &spv,
137        instruction_bound: &mut instruction_bound,
138        instruction_inserts: &mut instruction_inserts,
139        op_type_pointer_sampler_res_id,
140        op_variables_idxs: &op_variables_idxs,
141        tp_res: &tp_res,
142    });
143
144    // 5. OpTypeFunction
145    type_function(TypeFunctionIn {
146        spv: &spv,
147        word_inserts: &mut word_inserts,
148        op_type_pointer_sampler_res_id,
149        op_type_function_idxs: &op_type_function_idxs,
150        tp_res: &tp_res,
151    });
152
153    // 6. OpFunctionParameter
154    let parameter_res = function_parameter(FunctionParameterIn {
155        spv: &spv,
156        instruction_bound: &mut instruction_bound,
157        instruction_inserts: &mut instruction_inserts,
158        op_type_pointer_sampler_res_id,
159        op_function_parameter_idxs: &op_function_parameter_idxs,
160        tp_res: &tp_res,
161    });
162
163    // 7. OpFunctionCall
164    function_call(FunctionCallIn {
165        spv: &spv,
166        word_inserts: &mut word_inserts,
167        op_function_call_idxs: &op_function_call_idxs,
168        v_res: &v_res,
169        parameter_res: &parameter_res,
170    });
171
172    // 8. OpLoad
173    load(LoadIn {
174        spv: &spv,
175        new_spv: &mut new_spv,
176        instruction_bound: &mut instruction_bound,
177        instruction_inserts: &mut instruction_inserts,
178        op_type_sampler_res_id,
179        op_loads_idxs: &op_loads_idxs,
180        v_res: &v_res,
181        parameter_res: &parameter_res,
182    });
183
184    // 9. OpDecorate
185    let DecorateOut {
186        descriptor_sets_to_correct,
187    } = util::decorate(DecorateIn {
188        spv: &spv,
189        instruction_inserts: &mut instruction_inserts,
190        first_op_deocrate_idx,
191        op_decorate_idxs: &op_decorate_idxs,
192        affected_decorations: &v_res
193            .iter()
194            .map(
195                |VariableOut {
196                     v_res_id,
197                     new_sampler_v_res_id,
198                     ..
199                 }| {
200                    AffectedDecoration {
201                        original_res_id: *v_res_id,
202                        new_res_ids: vec![*new_sampler_v_res_id],
203                        correction_type: CorrectionType::SplitCombined,
204                    }
205                },
206            )
207            .collect::<Vec<_>>(),
208        corrections,
209    });
210
211    // 10. Insert New Instructions
212    insert_new_instructions(&spv, &mut new_spv, &word_inserts, &instruction_inserts);
213
214    // 11. Correct OpDecorate Bindings
215    util::correct_decorate(CorrectDecorateIn {
216        new_spv: &mut new_spv,
217        descriptor_sets_to_correct,
218    });
219
220    // 12. Remove Instructions that have been Whited Out.
221    prune_noops(&mut new_spv);
222
223    // 13. Write New Header and New Code
224    Ok(fuse_final(spv_header, new_spv, instruction_bound))
225}
spirv_webgpu_transform/splitcombined.rs

spirv_webgpu_transform/
splitcombined.rs
