naga/compact/
statements.rs

1use alloc::{vec, vec::Vec};
2
3use super::functions::FunctionTracer;
4use super::FunctionMap;
5use crate::arena::Handle;
6use crate::compact::handle_set_map::HandleMap;
7
8impl FunctionTracer<'_> {
9    pub fn trace_block(&mut self, block: &[crate::Statement]) {
10        let mut worklist: Vec<&[crate::Statement]> = vec![block];
11        while let Some(last) = worklist.pop() {
12            for stmt in last {
13                use crate::Statement as St;
14                match *stmt {
15                    St::Emit(ref _range) => {
16                        // If we come across a statement that actually uses an
17                        // expression in this range, it'll get traced from
18                        // there. But since evaluating expressions has no
19                        // effect, we don't need to assume that everything
20                        // emitted is live.
21                    }
22                    St::Block(ref block) => worklist.push(block),
23                    St::If {
24                        condition,
25                        ref accept,
26                        ref reject,
27                    } => {
28                        self.expressions_used.insert(condition);
29                        worklist.push(accept);
30                        worklist.push(reject);
31                    }
32                    St::Switch {
33                        selector,
34                        ref cases,
35                    } => {
36                        self.expressions_used.insert(selector);
37                        for case in cases {
38                            worklist.push(&case.body);
39                        }
40                    }
41                    St::Loop {
42                        ref body,
43                        ref continuing,
44                        break_if,
45                    } => {
46                        if let Some(break_if) = break_if {
47                            self.expressions_used.insert(break_if);
48                        }
49                        worklist.push(body);
50                        worklist.push(continuing);
51                    }
52                    St::Return { value: Some(value) } => {
53                        self.expressions_used.insert(value);
54                    }
55                    St::Store { pointer, value } => {
56                        self.expressions_used.insert(pointer);
57                        self.expressions_used.insert(value);
58                    }
59                    St::ImageStore {
60                        image,
61                        coordinate,
62                        array_index,
63                        value,
64                    } => {
65                        self.expressions_used.insert(image);
66                        self.expressions_used.insert(coordinate);
67                        if let Some(array_index) = array_index {
68                            self.expressions_used.insert(array_index);
69                        }
70                        self.expressions_used.insert(value);
71                    }
72                    St::Atomic {
73                        pointer,
74                        ref fun,
75                        value,
76                        result,
77                    } => {
78                        self.expressions_used.insert(pointer);
79                        self.trace_atomic_function(fun);
80                        self.expressions_used.insert(value);
81                        if let Some(result) = result {
82                            self.expressions_used.insert(result);
83                        }
84                    }
85                    St::ImageAtomic {
86                        image,
87                        coordinate,
88                        array_index,
89                        fun: _,
90                        value,
91                    } => {
92                        self.expressions_used.insert(image);
93                        self.expressions_used.insert(coordinate);
94                        if let Some(array_index) = array_index {
95                            self.expressions_used.insert(array_index);
96                        }
97                        self.expressions_used.insert(value);
98                    }
99                    St::WorkGroupUniformLoad { pointer, result } => {
100                        self.expressions_used.insert(pointer);
101                        self.expressions_used.insert(result);
102                    }
103                    St::Call {
104                        function,
105                        ref arguments,
106                        result,
107                    } => {
108                        self.trace_call(function);
109                        for expr in arguments {
110                            self.expressions_used.insert(*expr);
111                        }
112                        if let Some(result) = result {
113                            self.expressions_used.insert(result);
114                        }
115                    }
116                    St::RayQuery { query, ref fun } => {
117                        self.expressions_used.insert(query);
118                        self.trace_ray_query_function(fun);
119                    }
120                    St::SubgroupBallot { result, predicate } => {
121                        if let Some(predicate) = predicate {
122                            self.expressions_used.insert(predicate);
123                        }
124                        self.expressions_used.insert(result);
125                    }
126                    St::SubgroupCollectiveOperation {
127                        op: _,
128                        collective_op: _,
129                        argument,
130                        result,
131                    } => {
132                        self.expressions_used.insert(argument);
133                        self.expressions_used.insert(result);
134                    }
135                    St::SubgroupGather {
136                        mode,
137                        argument,
138                        result,
139                    } => {
140                        match mode {
141                            crate::GatherMode::BroadcastFirst => {}
142                            crate::GatherMode::Broadcast(index)
143                            | crate::GatherMode::Shuffle(index)
144                            | crate::GatherMode::ShuffleDown(index)
145                            | crate::GatherMode::ShuffleUp(index)
146                            | crate::GatherMode::ShuffleXor(index)
147                            | crate::GatherMode::QuadBroadcast(index) => {
148                                self.expressions_used.insert(index);
149                            }
150                            crate::GatherMode::QuadSwap(_) => {}
151                        }
152                        self.expressions_used.insert(argument);
153                        self.expressions_used.insert(result);
154                    }
155                    St::CooperativeStore { target, ref data } => {
156                        self.expressions_used.insert(target);
157                        self.expressions_used.insert(data.pointer);
158                        self.expressions_used.insert(data.stride);
159                    }
160                    St::RayPipelineFunction(func) => match func {
161                        crate::RayPipelineFunction::TraceRay {
162                            acceleration_structure,
163                            descriptor,
164                            payload,
165                        } => {
166                            self.expressions_used.insert(acceleration_structure);
167                            self.expressions_used.insert(descriptor);
168                            self.expressions_used.insert(payload);
169                        }
170                    },
171                    St::DebugPrintf {
172                        format: _,
173                        ref arguments,
174                    } => {
175                        for &expr in arguments {
176                            self.expressions_used.insert(expr);
177                        }
178                    }
179
180                    // Trivial statements.
181                    St::Break
182                    | St::Continue
183                    | St::Kill
184                    | St::ControlBarrier(_)
185                    | St::MemoryBarrier(_)
186                    | St::Return { value: None } => {}
187                }
188            }
189        }
190    }
191
192    fn trace_atomic_function(&mut self, fun: &crate::AtomicFunction) {
193        use crate::AtomicFunction as Af;
194        match *fun {
195            Af::Exchange {
196                compare: Some(expr),
197            } => {
198                self.expressions_used.insert(expr);
199            }
200            Af::Exchange { compare: None }
201            | Af::Add
202            | Af::Subtract
203            | Af::And
204            | Af::ExclusiveOr
205            | Af::InclusiveOr
206            | Af::Min
207            | Af::Max => {}
208        }
209    }
210
211    fn trace_ray_query_function(&mut self, fun: &crate::RayQueryFunction) {
212        use crate::RayQueryFunction as Qf;
213        match *fun {
214            Qf::Initialize {
215                acceleration_structure,
216                descriptor,
217            } => {
218                self.expressions_used.insert(acceleration_structure);
219                self.expressions_used.insert(descriptor);
220            }
221            Qf::Proceed { result } => {
222                self.expressions_used.insert(result);
223            }
224            Qf::GenerateIntersection { hit_t } => {
225                self.expressions_used.insert(hit_t);
226            }
227            Qf::ConfirmIntersection => {}
228            Qf::Terminate => {}
229            Qf::Begin => {}
230        }
231    }
232}
233
234impl FunctionMap {
235    /// Adjust statements in the body of `function`.
236    ///
237    /// Adjusts expressions using `self.expressions`, and adjusts calls to other
238    /// functions using `function_map`.
239    pub fn adjust_body(
240        &self,
241        function: &mut crate::Function,
242        function_map: &HandleMap<crate::Function>,
243    ) {
244        let block = &mut function.body;
245        let mut worklist: Vec<&mut [crate::Statement]> = vec![block];
246        let adjust = |handle: &mut Handle<crate::Expression>| {
247            self.expressions.adjust(handle);
248        };
249        while let Some(last) = worklist.pop() {
250            for stmt in last {
251                use crate::Statement as St;
252                match *stmt {
253                    St::Emit(ref mut range) => {
254                        self.expressions.adjust_range(range, &function.expressions);
255                    }
256                    St::Block(ref mut block) => worklist.push(block),
257                    St::If {
258                        ref mut condition,
259                        ref mut accept,
260                        ref mut reject,
261                    } => {
262                        adjust(condition);
263                        worklist.push(accept);
264                        worklist.push(reject);
265                    }
266                    St::Switch {
267                        ref mut selector,
268                        ref mut cases,
269                    } => {
270                        adjust(selector);
271                        for case in cases {
272                            worklist.push(&mut case.body);
273                        }
274                    }
275                    St::Loop {
276                        ref mut body,
277                        ref mut continuing,
278                        ref mut break_if,
279                    } => {
280                        if let Some(ref mut break_if) = *break_if {
281                            adjust(break_if);
282                        }
283                        worklist.push(body);
284                        worklist.push(continuing);
285                    }
286                    St::Return {
287                        value: Some(ref mut value),
288                    } => adjust(value),
289                    St::Store {
290                        ref mut pointer,
291                        ref mut value,
292                    } => {
293                        adjust(pointer);
294                        adjust(value);
295                    }
296                    St::ImageStore {
297                        ref mut image,
298                        ref mut coordinate,
299                        ref mut array_index,
300                        ref mut value,
301                    } => {
302                        adjust(image);
303                        adjust(coordinate);
304                        if let Some(ref mut array_index) = *array_index {
305                            adjust(array_index);
306                        }
307                        adjust(value);
308                    }
309                    St::Atomic {
310                        ref mut pointer,
311                        ref mut fun,
312                        ref mut value,
313                        ref mut result,
314                    } => {
315                        adjust(pointer);
316                        self.adjust_atomic_function(fun);
317                        adjust(value);
318                        if let Some(ref mut result) = *result {
319                            adjust(result);
320                        }
321                    }
322                    St::ImageAtomic {
323                        ref mut image,
324                        ref mut coordinate,
325                        ref mut array_index,
326                        fun: _,
327                        ref mut value,
328                    } => {
329                        adjust(image);
330                        adjust(coordinate);
331                        if let Some(ref mut array_index) = *array_index {
332                            adjust(array_index);
333                        }
334                        adjust(value);
335                    }
336                    St::WorkGroupUniformLoad {
337                        ref mut pointer,
338                        ref mut result,
339                    } => {
340                        adjust(pointer);
341                        adjust(result);
342                    }
343                    St::Call {
344                        ref mut function,
345                        ref mut arguments,
346                        ref mut result,
347                    } => {
348                        function_map.adjust(function);
349                        for expr in arguments {
350                            adjust(expr);
351                        }
352                        if let Some(ref mut result) = *result {
353                            adjust(result);
354                        }
355                    }
356                    St::RayQuery {
357                        ref mut query,
358                        ref mut fun,
359                    } => {
360                        adjust(query);
361                        self.adjust_ray_query_function(fun);
362                    }
363                    St::SubgroupBallot {
364                        ref mut result,
365                        ref mut predicate,
366                    } => {
367                        if let Some(ref mut predicate) = *predicate {
368                            adjust(predicate);
369                        }
370                        adjust(result);
371                    }
372                    St::SubgroupCollectiveOperation {
373                        op: _,
374                        collective_op: _,
375                        ref mut argument,
376                        ref mut result,
377                    } => {
378                        adjust(argument);
379                        adjust(result);
380                    }
381                    St::SubgroupGather {
382                        ref mut mode,
383                        ref mut argument,
384                        ref mut result,
385                    } => {
386                        match *mode {
387                            crate::GatherMode::BroadcastFirst => {}
388                            crate::GatherMode::Broadcast(ref mut index)
389                            | crate::GatherMode::Shuffle(ref mut index)
390                            | crate::GatherMode::ShuffleDown(ref mut index)
391                            | crate::GatherMode::ShuffleUp(ref mut index)
392                            | crate::GatherMode::ShuffleXor(ref mut index)
393                            | crate::GatherMode::QuadBroadcast(ref mut index) => adjust(index),
394                            crate::GatherMode::QuadSwap(_) => {}
395                        }
396                        adjust(argument);
397                        adjust(result);
398                    }
399                    St::CooperativeStore {
400                        ref mut target,
401                        ref mut data,
402                    } => {
403                        adjust(target);
404                        adjust(&mut data.pointer);
405                        adjust(&mut data.stride);
406                    }
407                    St::RayPipelineFunction(ref mut func) => match *func {
408                        crate::RayPipelineFunction::TraceRay {
409                            ref mut acceleration_structure,
410                            ref mut descriptor,
411                            ref mut payload,
412                        } => {
413                            adjust(acceleration_structure);
414                            adjust(descriptor);
415                            adjust(payload);
416                        }
417                    },
418                    St::DebugPrintf {
419                        format: _,
420                        ref mut arguments,
421                    } => {
422                        for expr in arguments {
423                            adjust(expr);
424                        }
425                    }
426
427                    // Trivial statements.
428                    St::Break
429                    | St::Continue
430                    | St::Kill
431                    | St::ControlBarrier(_)
432                    | St::MemoryBarrier(_)
433                    | St::Return { value: None } => {}
434                }
435            }
436        }
437    }
438
439    fn adjust_atomic_function(&self, fun: &mut crate::AtomicFunction) {
440        use crate::AtomicFunction as Af;
441        match *fun {
442            Af::Exchange {
443                compare: Some(ref mut expr),
444            } => {
445                self.expressions.adjust(expr);
446            }
447            Af::Exchange { compare: None }
448            | Af::Add
449            | Af::Subtract
450            | Af::And
451            | Af::ExclusiveOr
452            | Af::InclusiveOr
453            | Af::Min
454            | Af::Max => {}
455        }
456    }
457
458    fn adjust_ray_query_function(&self, fun: &mut crate::RayQueryFunction) {
459        use crate::RayQueryFunction as Qf;
460        match *fun {
461            Qf::Initialize {
462                ref mut acceleration_structure,
463                ref mut descriptor,
464            } => {
465                self.expressions.adjust(acceleration_structure);
466                self.expressions.adjust(descriptor);
467            }
468            Qf::Proceed { ref mut result } => {
469                self.expressions.adjust(result);
470            }
471            Qf::GenerateIntersection { ref mut hit_t } => {
472                self.expressions.adjust(hit_t);
473            }
474            Qf::ConfirmIntersection => {}
475            Qf::Terminate => {}
476            Qf::Begin => {}
477        }
478    }
479}
naga/compact/statements.rs

naga/compact/
statements.rs
