Skip to main content

iridia_exec/
run.rs

1use std::collections::BTreeMap;
2
3use super::flat::{
4    AtomOp, BinOp, CvtOp, Executable, GlobalKind, LevelOp, Loc, MathOp, Member, Op, Pick, Ptr, Reg, Res, St,
5    UnOp, WaveGather, WaveKind, WaveOp, WaveScan,
6};
7use super::mask::{Mask, Wide};
8use super::rayquery::{Built, Query, Reported, Scenes};
9use super::request::{
10    Backend, Element, ImageAtomic, ImageLoad, ImageLod, ImageQuery, ImageSample, ImageStore,
11};
12use iridia::ResourceId;
13
14/// The most operations and loop turns one invocation may execute: a program that does not end is refused.
15pub const STEP_LIMIT: u64 = 1 << 24;
16
17/// Where a global's words are during a run.
18#[derive(Debug)]
19enum Words<'a> {
20    Owned(Vec<u32>),
21    /// A constant buffer as the caller holds it: rows of four lanes.
22    Rows(&'a [[u32; 4]]),
23    Slice(&'a [u32]),
24    /// The elements of an array of buffers; one the caller gave none of is `None`.
25    Elements(BTreeMap<u32, &'a [u32]>),
26    /// Independently bound descriptors of a constant-buffer array.
27    ConstantElements(BTreeMap<u32, &'a [[u32; 4]]>),
28    /// A buffer the backend answers reads of.
29    Backend,
30}
31
32/// The buffers a program reads, ready for runs: built once per program and reused for every pixel.
33#[derive(Debug)]
34pub struct Bound<'a> {
35    data: Vec<Words<'a>>,
36    scenes: Scenes,
37}
38
39impl Bound<'_> {
40    /// Give the acceleration structure of resource `id`: what the program's ray queries search.
41    pub fn with_scene(mut self, id: ResourceId, scene: std::sync::Arc<Built>) -> Self {
42        self.scenes.insert(id, scene);
43        self
44    }
45}
46
47/// What a run reads besides its inputs: constant buffers (rows of four lanes) and buffers (32-bit words), by resource.
48#[derive(Clone, Debug, Default)]
49pub struct Resources {
50    /// Constant buffers by interface resource: rows of four 32-bit words.
51    pub constants: BTreeMap<ResourceId, Vec<[u32; 4]>>,
52    /// Constant-buffer array descriptors, keyed by resource and zero-based descriptor element.
53    pub constant_elements: BTreeMap<(ResourceId, u32), Vec<[u32; 4]>>,
54    /// Buffers by interface resource: 32-bit words.
55    pub buffers: BTreeMap<ResourceId, Vec<u32>>,
56    /// The elements of arrays of buffers, by resource and element.
57    pub elements: BTreeMap<(ResourceId, u32), Vec<u32>>,
58    /// The acceleration structures, by resource.
59    pub scenes: BTreeMap<ResourceId, super::rayquery::Scene>,
60}
61
62impl Executable {
63    /// Bind the buffers the program reads, as `constants` and `buffers` find them by resource. A constant buffer the
64    /// program reads and the caller does not give is an error; a buffer not given is one the backend answers reads of
65    /// ([`Backend::read_buffer`]). Nothing is copied: the caller's buffers are read in place.
66    pub fn bind<'a>(
67        &self,
68        constants: impl Fn(ResourceId) -> Option<&'a [[u32; 4]]>,
69        buffers: impl Fn(ResourceId) -> Option<&'a [u32]>,
70    ) -> Result<Bound<'a>, String> {
71        self.bind_elements(constants, buffers, |_, _| None)
72    }
73
74    /// [`Executable::bind`] for a program that also reads arrays of buffers: `elements` finds the words of element
75    /// `i` of the array resource `id`. An element not given errors when a run reads it.
76    pub fn bind_elements<'a>(
77        &self,
78        constants: impl Fn(ResourceId) -> Option<&'a [[u32; 4]]>,
79        buffers: impl Fn(ResourceId) -> Option<&'a [u32]>,
80        elements: impl Fn(ResourceId, u32) -> Option<&'a [u32]>,
81    ) -> Result<Bound<'a>, String> {
82        if self
83            .globals
84            .iter()
85            .any(|g| matches!(g.kind, GlobalKind::StorageArray(_, iridia::Resource::UNBOUNDED)))
86        {
87            return Err("an unbounded buffer array needs bind_resources or bind_resource_arrays with explicit elements".into());
88        }
89        self.bind_resource_arrays(
90            constants,
91            buffers,
92            |id| {
93                let count = self
94                    .globals
95                    .iter()
96                    .find_map(|g| match g.kind {
97                        GlobalKind::StorageArray(r, count) if r == id => Some(count),
98                        _ => None,
99                    })
100                    .unwrap_or(0);
101                (0..count)
102                    .filter_map(|k| elements(id, k).map(|words| (k, words)))
103                    .collect()
104            },
105            |_| BTreeMap::new(),
106        )
107    }
108
109    /// Bind descriptor arrays by their supplied elements, including arrays of no stated size.
110    /// Keys are zero-based descriptor indices, independent of D3D register numbers. Missing
111    /// elements error only when read; a key past a bounded declaration's end errors at binding.
112    /// The maps are retained by the binding; the buffers themselves are borrowed without copying.
113    pub fn bind_resource_arrays<'a>(
114        &self,
115        constants: impl Fn(ResourceId) -> Option<&'a [[u32; 4]]>,
116        buffers: impl Fn(ResourceId) -> Option<&'a [u32]>,
117        elements: impl Fn(ResourceId) -> BTreeMap<u32, &'a [u32]>,
118        constant_elements: impl Fn(ResourceId) -> BTreeMap<u32, &'a [[u32; 4]]>,
119    ) -> Result<Bound<'a>, String> {
120        let mut data = Vec::with_capacity(self.globals.len());
121        for g in &self.globals {
122            data.push(match &g.kind {
123                GlobalKind::Constants(id) => Words::Rows(
124                    constants(*id).ok_or_else(|| format!("no constant buffer for resource {}", id.0))?,
125                ),
126                GlobalKind::ConstantsArray(id, count) => {
127                    let values = constant_elements(*id);
128                    if values.keys().any(|k| *k >= *count) {
129                        return Err(format!(
130                            "constant buffer {}: element past array count {count}",
131                            id.0
132                        ));
133                    }
134                    Words::ConstantElements(values)
135                }
136                GlobalKind::Private(words) => Words::Owned(words.clone()),
137                GlobalKind::Storage(id) => buffers(*id).map_or(Words::Backend, Words::Slice),
138                GlobalKind::StorageArray(id, count) => {
139                    let values = elements(*id);
140                    if values.keys().any(|k| *k >= *count) {
141                        return Err(format!("buffer {}: element past array count {count}", id.0));
142                    }
143                    Words::Elements(values)
144                }
145                // thread-group memory is the run's own
146                GlobalKind::Workgroup(_) => Words::Owned(Vec::new()),
147            });
148        }
149        Ok(Bound {
150            data,
151            scenes: BTreeMap::new(),
152        })
153    }
154
155    /// [`Executable::bind`] over a [`Resources`].
156    pub fn bind_resources<'a>(&self, resources: &'a Resources) -> Result<Bound<'a>, String> {
157        let mut bound = self.bind_resource_arrays(
158            |id| resources.constants.get(&id).map(Vec::as_slice),
159            |id| resources.buffers.get(&id).map(Vec::as_slice),
160            |id| {
161                resources
162                    .elements
163                    .iter()
164                    .filter_map(|(&(r, k), v)| (r == id).then_some((k, v.as_slice())))
165                    .collect()
166            },
167            |id| {
168                resources
169                    .constant_elements
170                    .iter()
171                    .filter_map(|(&(r, k), v)| (r == id).then_some((k, v.as_slice())))
172                    .collect()
173            },
174        )?;
175        for (id, scene) in &resources.scenes {
176            bound.scenes.insert(*id, std::sync::Arc::new(Built::new(scene)?));
177        }
178        Ok(bound)
179    }
180}
181
182/// What a ray tracing program decided about the hit it was offered, besides its payload.
183#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
184pub enum Decision {
185    /// Nothing: the hit stands.
186    #[default]
187    None,
188    /// `IgnoreHit`: an any-hit program rejects the candidate.
189    Ignore,
190    /// `AcceptHitAndEndSearch`: an any-hit program accepts it and ends the search.
191    AcceptAndEnd,
192}
193
194/// What one invocation produced.
195#[derive(Clone, Debug)]
196pub struct Output {
197    /// The words of each pointer argument (a ray tracing program's payload, attributes or parameter) as the program
198    /// left them, in argument order.
199    pub pointees: Vec<Vec<u32>>,
200    /// What the program decided about the hit it was offered; `Decision::None` for any other program.
201    pub decision: Decision,
202    /// One entry per member of the result, lanes beyond the member's width 0.
203    pub values: Vec<[u32; 4]>,
204    /// Whether the invocation discarded.
205    pub discarded: bool,
206    /// Operations and loop turns executed.
207    pub steps: u64,
208    /// Sample requests this lane made.
209    pub samples: u64,
210}
211
212/// The registers and local memory of a [`Runner`], kept from one run of a program to the next so that the literals are
213/// written once and nothing is allocated per pixel.
214#[derive(Debug)]
215pub struct Scratch<const N: usize> {
216    regs: Vec<Reg>,
217    mem: Vec<u32>,
218    /// Each lane's result members, their room kept.
219    outs: [Vec<Reg>; N],
220}
221impl<const N: usize> Scratch<N> {
222    /// Registers and local memory for `N` lanes of `exe`, with its literals already written.
223    pub fn new(exe: &Executable) -> Self {
224        let mut regs = vec![[0u32; 4]; exe.slots.saturating_mul(N)];
225        for (s, v) in &exe.inits {
226            for l in 0..N {
227                if let Some(register) = regs.get_mut(slot_at::<N>(*s).saturating_add(l)) {
228                    *register = *v;
229                }
230            }
231        }
232        Self {
233            regs,
234            mem: vec![0; exe.local_words.saturating_mul(N)],
235            outs: std::array::from_fn(|_| Vec::new()),
236        }
237    }
238}
239
240/// Runs an [`Executable`] for a 2 x 2 quad in lockstep (`N` = 4: lanes 0 and 1 the top row, 2 and 3 the bottom), as a
241/// GPU does: lanes that diverge are masked, a screen derivative is a difference inside the quad, a lane that is
242/// discarded goes on as a helper. Any other `N` runs that many invocations of their own (`N` = 1, one), in lockstep
243/// and masked where they diverge, each as it would run alone: its own steps and budget, a discard ending it, the
244/// derivatives of its inputs the ones the caller gives for every lane. A runner is reused for any number of runs.
245pub struct Runner<'a, const N: usize, M: Mask = u8> {
246    c: &'a Executable,
247    b: &'a Bound<'a>,
248    regs: Vec<Reg>,
249    mem: Vec<u32>,
250    sampler: &'a mut dyn Backend,
251    ftz: bool,
252    grouping: Grouping,
253    /// The steps of a quad or a thread group, counted once for all its lanes.
254    steps: u64,
255    /// The steps of each lane of independent invocations.
256    lane_steps: [u64; N],
257    samples: [u64; N],
258    outs: [Vec<Reg>; N],
259    /// What an any-hit program decided, per lane.
260    decisions: [Decision; N],
261    killed: M,
262    /// The memory a run changes, by global: thread-group memory and the buffers a compute run writes (empty for a
263    /// pixel run, which hands its writes to the backend).
264    store: Vec<Option<Vec<u32>>>,
265    /// The ray queries of each invocation, by query and lane.
266    queries: Vec<Vec<Option<Query>>>,
267    /// The same for the elements of arrays of buffers, by global and element.
268    elem_store: Vec<BTreeMap<u32, Vec<u32>>>,
269    /// The lanes of a thread group, and those of them that have left the program.
270    group: M,
271    retired: M,
272    /// The lanes in use and the width of a wave of them (0 outside a compute run).
273    lanes: usize,
274    wave: usize,
275    /// The derivatives of the arguments (scalar mode): per argument, d/dx then d/dy lanes.
276    derivs: Option<&'a [[Reg; 2]]>,
277}
278
279impl<const N: usize, M: Mask> std::fmt::Debug for Runner<'_, N, M> {
280    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
281        // independent invocations: the most steps a lane has taken
282        let steps = match self.grouping {
283            Grouping::Independent => self.lane_steps.iter().copied().max().unwrap_or(0),
284            Grouping::Quad | Grouping::Group => self.steps,
285        };
286        f.debug_struct("Runner")
287            .field("lanes", &N)
288            .field("grouping", &self.grouping)
289            .field("steps", &steps)
290            .finish_non_exhaustive()
291    }
292}
293
294/// What the lanes of a run are to each other.
295#[derive(Clone, Copy, Debug, PartialEq, Eq)]
296enum Grouping {
297    /// Invocations of their own.
298    Independent,
299    /// A 2 x 2 quad.
300    Quad,
301    /// A compute thread group.
302    Group,
303}
304
305#[derive(Clone, Copy)]
306struct Exits<M> {
307    brk: M,
308    cont: M,
309    ret: M,
310}
311impl<M: Mask> Default for Exits<M> {
312    fn default() -> Self {
313        Self {
314            brk: M::none(),
315            cont: M::none(),
316            ret: M::none(),
317        }
318    }
319}
320
321/// `$body` for each lane `$l` of `$active` (a runner's lanes): without a test per lane when every lane is active.
322macro_rules! each {
323    ($active:expr, $l:ident, $body:block) => {
324        if $active == M::first(N) {
325            for $l in 0..N $body
326        } else {
327            for $l in 0..N {
328                if $active.has($l) $body
329            }
330        }
331    };
332}
333
334/// A position in one of the tables of the compiled program or its bound data (a global, an operation). The compile pass
335/// numbers them, so a miss is an internal error, reported with the position instead of a panic.
336trait At<T> {
337    fn at(&self, i: usize) -> Result<&T, String>;
338}
339impl<T> At<T> for [T] {
340    fn at(&self, i: usize) -> Result<&T, String> {
341        self.get(i)
342            .ok_or_else(|| format!("internal error: position {i} of a table of {}", self.len()))
343    }
344}
345
346/// The word offset `o + extra`, an error where it does not fit.
347fn offset(o: u32, extra: u32) -> Result<usize, String> {
348    o.checked_add(extra)
349        .map(|word| word as usize)
350        .ok_or_else(|| format!("word offset {o} + {extra} overflows"))
351}
352
353/// The `n` (at most 4) words at `at` of `words`, as the leading words of a register; `what` names `words` in the error.
354fn take_words(words: &[u32], at: usize, n: usize, what: &str) -> Result<Reg, String> {
355    let from = at
356        .checked_add(n)
357        .and_then(|end| words.get(at..end))
358        .ok_or_else(|| format!("{what} index out of range"))?;
359    let mut r = [0u32; 4];
360    for (to, word) in r.iter_mut().zip(from) {
361        *to = *word;
362    }
363    Ok(r)
364}
365
366/// The first `n` words of the register `v` written at `at` of `words`.
367fn put_words(words: &mut [u32], at: usize, n: usize, v: &Reg, what: &str) -> Result<(), String> {
368    let to = at
369        .checked_add(n)
370        .and_then(|end| words.get_mut(at..end))
371        .ok_or_else(|| format!("{what} index out of range"))?;
372    for (to, word) in to.iter_mut().zip(v) {
373        *to = *word;
374    }
375    Ok(())
376}
377
378/// The first `n` words of a register as the coordinates of a texel (the rest 0).
379fn texel_of(r: &Reg, n: usize) -> [u32; 3] {
380    let [x, y, z, _] = *r;
381    let keep = |i: usize, w: u32| if i < n { w } else { 0 };
382    [keep(0, x), keep(1, y), keep(2, z)]
383}
384
385/// Component `axis` of a 3-vector of an id or a size (0 past its three).
386fn comp3(v: [u32; 3], axis: usize) -> u32 {
387    v.get(axis).copied().unwrap_or(0)
388}
389
390/// How many invocations a thread group of this size has.
391fn group_lanes(size: [u32; 3]) -> Result<u32, String> {
392    let [x, y, z] = size;
393    x.checked_mul(y)
394        .and_then(|xy| xy.checked_mul(z))
395        .ok_or_else(|| format!("a thread group of {x} by {y} by {z} invocations does not fit in 32 bits"))
396}
397
398/// Component `i` of a register; 0 past its four (a scalar operand repeated reads component 0).
399#[inline(always)]
400fn comp(r: &Reg, i: usize) -> u32 {
401    r.get(i).copied().unwrap_or(0)
402}
403
404/// HLSL's float to signed integer: saturating, a NaN 0, as the bits of the integer.
405#[inline(always)]
406#[allow(clippy::cast_possible_truncation)] // the saturating conversion is the operation
407fn float_to_int(x: f32) -> u32 {
408    (x as i32) as u32
409}
410
411/// HLSL's float to unsigned integer: saturating, a NaN 0.
412#[inline(always)]
413#[allow(clippy::cast_possible_truncation)] // the saturating conversion is the operation
414fn float_to_uint(x: f32) -> u32 {
415    x as u32
416}
417
418/// The register index of the first lane of slot `slot`: the slots outermost and the lanes innermost, as in `Scratch`.
419#[inline(always)]
420fn slot_at<const N: usize>(slot: u32) -> usize {
421    (slot as usize).saturating_mul(N)
422}
423
424/// A bug of the compile pass, as an error naming the register: the compile pass numbers a run's registers.
425fn no_register(index: usize, len: usize) -> String {
426    format!("internal error: register {index}, where the run has {len}")
427}
428
429#[inline(always)]
430fn fl(x: f32, on: bool) -> f32 {
431    if on && x.is_subnormal() {
432        f32::from_bits(x.to_bits() & 0x8000_0000)
433    } else {
434        x
435    }
436}
437#[inline(always)]
438fn f(x: u32) -> f32 {
439    f32::from_bits(x)
440}
441
442impl<'a, const N: usize, M: Mask> Runner<'a, N, M> {
443    /// A runner of `exe` over `bound`, answering requests through `backend`. Invocations of their own that read
444    /// derivatives of their inputs take them in `derivatives`, the same for every lane; a quad takes its own.
445    pub fn new(
446        exe: &'a Executable,
447        bound: &'a Bound<'a>,
448        backend: &'a mut dyn Backend,
449        derivatives: Option<&'a [[[u32; 4]; 2]]>,
450    ) -> Self {
451        Self::with_scratch(exe, bound, backend, derivatives, Scratch::new(exe))
452    }
453
454    /// [`Runner::new`] on registers kept from an earlier runner of the same program ([`Runner::into_scratch`]): the
455    /// literals are already in them.
456    pub fn with_scratch(
457        exe: &'a Executable,
458        bound: &'a Bound<'a>,
459        backend: &'a mut dyn Backend,
460        derivatives: Option<&'a [[[u32; 4]; 2]]>,
461        scratch: Scratch<N>,
462    ) -> Self {
463        let c = exe;
464        let (b, sampler, ftz, derivs) = (bound, backend, exe.flush_denormals, derivatives);
465        Self {
466            c,
467            b,
468            regs: scratch.regs,
469            mem: scratch.mem,
470            sampler,
471            ftz,
472            grouping: if N == 4 {
473                Grouping::Quad
474            } else {
475                Grouping::Independent
476            },
477            steps: 0,
478            lane_steps: [0; N],
479            samples: [0; N],
480            outs: scratch.outs,
481            decisions: [Decision::None; N],
482            killed: M::none(),
483            store: Vec::new(),
484            elem_store: Vec::new(),
485            queries: Vec::new(),
486            group: M::none(),
487            retired: M::none(),
488            lanes: 0,
489            wave: 0,
490            derivs,
491        }
492    }
493
494    /// Give the registers back, to start the next runner of this program with.
495    pub fn into_scratch(self) -> Scratch<N> {
496        Scratch {
497            regs: self.regs,
498            mem: self.mem,
499            outs: self.outs,
500        }
501    }
502
503    /// Register `lane` of the slot whose first lane is `base` (see [`slot_at`]).
504    #[inline(always)]
505    fn rd(&self, base: usize, lane: usize) -> Result<Reg, String> {
506        let index = base.saturating_add(lane);
507        self.regs
508            .get(index)
509            .copied()
510            .ok_or_else(|| no_register(index, self.regs.len()))
511    }
512    /// Set register `lane` of the slot whose first lane is `base` to `value`.
513    #[inline(always)]
514    fn wr(&mut self, base: usize, lane: usize, value: Reg) -> Result<(), String> {
515        let index = base.saturating_add(lane);
516        let len = self.regs.len();
517        *self.regs.get_mut(index).ok_or_else(|| no_register(index, len))? = value;
518        Ok(())
519    }
520
521    /// Count `n` steps the lanes of `active` take, refusing a run past its budget: each lane's own for independent
522    /// invocations, the quad's or the group's once.
523    #[inline(always)]
524    fn count(&mut self, active: M, n: u64) -> Result<(), String> {
525        let over = if self.grouping == Grouping::Independent {
526            let mut over = false;
527            for (l, steps) in self.lane_steps.iter_mut().enumerate() {
528                if active.has(l) {
529                    *steps = steps.saturating_add(n);
530                    over |= *steps > STEP_LIMIT;
531                }
532            }
533            over
534        } else {
535            self.steps = self.steps.saturating_add(n);
536            self.steps > STEP_LIMIT
537        };
538        if over {
539            return Err("execution budget exceeded".into());
540        }
541        Ok(())
542    }
543
544    /// Where local word `word` of lane `lane` is: the words outermost and the lanes innermost, as in the registers.
545    #[inline(always)]
546    fn local(word: usize, lane: usize) -> usize {
547        word.saturating_mul(N).saturating_add(lane)
548    }
549    /// The first local word of `n` from `word`, refusing words past the program's locals.
550    #[inline(always)]
551    fn locals(&self, word: usize, n: usize) -> Result<usize, String> {
552        if word.checked_add(n).is_none_or(|end| end > self.c.local_words) {
553            return Err("local index out of range".into());
554        }
555        Ok(word)
556    }
557    /// Local word `word` of lane `lane`.
558    #[inline(always)]
559    fn local_word(&self, word: usize, lane: usize) -> Result<u32, String> {
560        let at = Self::local(word, lane);
561        self.mem.get(at).copied().ok_or_else(|| {
562            format!(
563                "internal error: local word {at}, where the run has {}",
564                self.mem.len()
565            )
566        })
567    }
568    /// Set local word `word` of lane `lane`.
569    #[inline(always)]
570    fn set_local_word(&mut self, word: usize, lane: usize, value: u32) -> Result<(), String> {
571        let at = Self::local(word, lane);
572        let len = self.mem.len();
573        *self
574            .mem
575            .get_mut(at)
576            .ok_or_else(|| format!("internal error: local word {at}, where the run has {len}"))? = value;
577        Ok(())
578    }
579    /// Lane `lane`'s local words `at..at + n`.
580    fn read_locals(&self, lane: usize, at: usize, n: usize) -> Result<Vec<u32>, String> {
581        (at..at.saturating_add(n))
582            .map(|w| self.local_word(w, lane))
583            .collect()
584    }
585    /// Lane `lane`'s local words from `at` set to `words`.
586    fn write_locals(&mut self, lane: usize, at: usize, words: &[u32]) -> Result<(), String> {
587        for (k, w) in words.iter().enumerate() {
588            self.set_local_word(at.saturating_add(k), lane, *w)?;
589        }
590        Ok(())
591    }
592
593    fn mem_get(&mut self, loc: Loc, lane: usize, extra: u32, n8: u8) -> Result<Reg, String> {
594        let n = n8 as usize;
595        match loc {
596            Loc::Local(o) => {
597                let base = self.locals(offset(o, extra)?, n)?;
598                let mut r = [0u32; 4];
599                for (i, w) in r.iter_mut().take(n).enumerate() {
600                    *w = self.local_word(base.saturating_add(i), lane)?;
601                }
602                Ok(r)
603            }
604            Loc::Elem(g, o, pick) => {
605                let base = offset(o, extra)?;
606                let e = self.element_index(g, pick, lane)?;
607                if let Words::ConstantElements(values) = self.b.data.at(g as usize)? {
608                    let rows = values
609                        .get(&e)
610                        .ok_or_else(|| format!("no constant buffer given for descriptor {e}"))?;
611                    let mut result = [0; 4];
612                    for (k, value) in result.iter_mut().take(n).enumerate() {
613                        let word = base.checked_add(k).ok_or("constant word index overflow")?;
614                        *value = rows
615                            .get(word / 4)
616                            .and_then(|row| row.get(word % 4))
617                            .copied()
618                            .ok_or("constant buffer index out of range")?;
619                    }
620                    return Ok(result);
621                }
622                let d = self.element_words(g, e)?;
623                take_words(d, base, n, "buffer")
624            }
625            Loc::Global(g, o) => {
626                let base = offset(o, extra)?;
627                if let Some(Some(d)) = self.store.get(g as usize) {
628                    return take_words(d, base, n, "memory");
629                }
630                let bound = self.b;
631                match bound.data.at(g as usize)? {
632                    Words::Elements(_) | Words::ConstantElements(_) => {
633                        Err("an array of buffers read without an element".into())
634                    }
635                    Words::Owned(d) => take_words(d, base, n, "buffer"),
636                    Words::Slice(d) => take_words(d, base, n, "buffer"),
637                    Words::Rows(rows) => {
638                        let mut r = [0u32; 4];
639                        if n == 4 && base.is_multiple_of(4) {
640                            r = *rows.get(base / 4).ok_or("buffer index out of range")?;
641                        } else {
642                            for (i, w) in r.iter_mut().take(n).enumerate() {
643                                let word = base.saturating_add(i);
644                                *w = rows
645                                    .get(word / 4)
646                                    .and_then(|row| row.get(word % 4))
647                                    .copied()
648                                    .ok_or("buffer index out of range")?;
649                            }
650                        }
651                        Ok(r)
652                    }
653                    // a buffer the caller answers for word by word
654                    Words::Backend => {
655                        let GlobalKind::Storage(view) = self.c.globals.at(g as usize)?.kind else {
656                            return Err("a global that is neither bound nor a buffer".into());
657                        };
658                        let mut r = self.sampler.read_buffer(
659                            view,
660                            u32::try_from(base).map_err(|_| "buffer index out of range")?,
661                            n8,
662                        )?;
663                        for w in r.iter_mut().skip(n) {
664                            *w = 0;
665                        }
666                        Ok(r)
667                    }
668                }
669            }
670        }
671    }
672    fn mem_put(&mut self, loc: Loc, lane: usize, extra: u32, n: u8, v: Reg) -> Result<(), String> {
673        let count = n as usize;
674        match loc {
675            Loc::Local(o) => {
676                let base = self.locals(offset(o, extra)?, count)?;
677                for (i, w) in v.iter().take(count).enumerate() {
678                    self.set_local_word(base.saturating_add(i), lane, *w)?;
679                }
680                Ok(())
681            }
682            Loc::Elem(g, o, pick) => {
683                let e = self.element_index(g, pick, lane)?;
684                let at = offset(o, extra)?;
685                let Some(d) = self.elem_store.get_mut(g as usize).and_then(|v| v.get_mut(&e)) else {
686                    return Err("a store to an element of an array of buffers the run does not write".into());
687                };
688                put_words(d, at, count, &v, "memory")
689            }
690            Loc::Global(g, o) => {
691                if let Some(Some(d)) = self.store.get_mut(g as usize) {
692                    let at = offset(o, extra)?;
693                    return put_words(d, at, count, &v, "memory");
694                }
695                // a write to a storage buffer is the caller's effect (it reaches none of this evaluation's reads)
696                let GlobalKind::Storage(view) = self.c.globals.at(g as usize)?.kind else {
697                    return Err("store to a read-only global".into());
698                };
699                let first = o.checked_add(extra).ok_or("memory index out of range")?;
700                self.sampler.store_buffer(view, first, v, crate::lane_mask(n))
701            }
702        }
703    }
704    /// The element of an array of buffers a location names in `lane`.
705    fn element_index(&self, g: u32, pick: Pick, lane: usize) -> Result<u32, String> {
706        let index = match pick {
707            Pick::One => 0,
708            Pick::Fixed(i) => i,
709            Pick::Slot(s) => self.rd(slot_at::<N>(s), lane)?[0],
710        };
711        match &self.c.globals.at(g as usize)?.kind {
712            GlobalKind::StorageArray(id, count) | GlobalKind::ConstantsArray(id, count)
713                if index >= *count =>
714            {
715                Err(format!("buffer {}: element {index} of an array of {count}", id.0))
716            }
717            _ => Ok(index),
718        }
719    }
720    /// The words of element `e` of the array of buffers `g`: those the run writes, else those it was bound with.
721    fn element_words(&self, g: u32, e: u32) -> Result<&[u32], String> {
722        if let Some(d) = self.elem_store.get(g as usize).and_then(|v| v.get(&e)) {
723            return Ok(d);
724        }
725        match self.b.data.at(g as usize)? {
726            Words::Elements(v) => v
727                .get(&e)
728                .copied()
729                .ok_or_else(|| format!("no words given for element {e} of an array of buffers")),
730            _ => Err("a global that is not an array of buffers".into()),
731        }
732    }
733    fn dyn_extra(&self, p: &Ptr, lane: usize) -> Result<u32, String> {
734        match p.dynamic {
735            None => Ok(0),
736            Some((s, stride, len)) => {
737                let i = self.rd(slot_at::<N>(s), lane)?[0];
738                if i >= len {
739                    return Err(format!("dynamic index {i} past the end ({len})"));
740                }
741                i.checked_mul(stride)
742                    .ok_or_else(|| format!("dynamic index {i} of a stride of {stride} words overflows"))
743            }
744        }
745    }
746
747    fn ops(&mut self, lo: u32, hi: u32, active: M) -> Result<(), String> {
748        self.count(active, u64::from(hi.saturating_sub(lo)))?;
749        let c = self.c;
750        let ftz = self.ftz;
751        for k in lo..hi {
752            let (dst, op) = c.ops.at(k as usize)?;
753            let dst = slot_at::<N>(*dst);
754            match op {
755                // a local's words, at offsets the compiler checked: read straight from every lane's memory
756                Op::Load {
757                    loc: Loc::Local(o),
758                    n,
759                } => {
760                    let (base, n) = (self.locals(*o as usize, *n as usize)?, *n as usize);
761                    each!(active, l, {
762                        let mut r = [0u32; 4];
763                        for (i, w) in r.iter_mut().take(n).enumerate() {
764                            *w = self.local_word(base.saturating_add(i), l)?;
765                        }
766                        self.wr(dst, l, r)?;
767                    })
768                }
769                Op::LoadLanes {
770                    loc: Loc::Local(o),
771                    offs,
772                    n,
773                } => {
774                    let (o, n) = (*o as usize, *n as usize);
775                    for &k in offs.iter().take(n) {
776                        self.locals(o.saturating_add(k as usize), 1)?;
777                    }
778                    each!(active, l, {
779                        let mut r = [0u32; 4];
780                        for (w, &k) in r.iter_mut().zip(offs.iter().take(n)) {
781                            *w = self.local_word(o.saturating_add(k as usize), l)?;
782                        }
783                        self.wr(dst, l, r)?;
784                    })
785                }
786                Op::Load { loc, n } => each!(active, l, {
787                    let loaded = self.mem_get(*loc, l, 0, *n)?;
788                    self.wr(dst, l, loaded)?;
789                }),
790                Op::LoadLanes { loc, offs, n } => each!(active, l, {
791                    let mut r = [0u32; 4];
792                    for (w, &k) in r.iter_mut().zip(offs.iter().take(*n as usize)) {
793                        *w = self.mem_get(*loc, l, u32::from(k), 1)?[0];
794                    }
795                    self.wr(dst, l, r)?;
796                }),
797                Op::LoadDyn { ptr, n } => each!(active, l, {
798                    let x = self.dyn_extra(ptr, l)?;
799                    let loaded = self.mem_get(ptr.loc, l, x, *n)?;
800                    self.wr(dst, l, loaded)?;
801                }),
802                Op::Swizzle { src, pat, n } => {
803                    let src = slot_at::<N>(*src);
804                    each!(active, l, {
805                        let s = self.rd(src, l)?;
806                        let mut r = [0u32; 4];
807                        for (w, &p) in r.iter_mut().zip(pat.iter()).take(*n as usize) {
808                            *w = comp(&s, p as usize);
809                        }
810                        self.wr(dst, l, r)?;
811                    })
812                }
813                Op::Lane { src, i } => {
814                    let (src, i) = (slot_at::<N>(*src), *i as usize);
815                    each!(active, l, {
816                        let s = self.rd(src, l)?;
817                        self.wr(dst, l, [comp(&s, i); 4])?;
818                    })
819                }
820                Op::Splat { src, n } => {
821                    let _ = n;
822                    let src = slot_at::<N>(*src);
823                    each!(active, l, {
824                        let s = self.rd(src, l)?;
825                        self.wr(dst, l, [comp(&s, 0); 4])?;
826                    })
827                }
828                Op::Compose { parts, np } => each!(active, l, {
829                    let mut r = [0u32; 4];
830                    let mut at = 0usize;
831                    for p in parts.iter().take(*np as usize) {
832                        let s = self.rd(slot_at::<N>(p.0), l)?;
833                        for w in s.iter().take(p.1 as usize) {
834                            if let Some(slot) = r.get_mut(at) {
835                                *slot = *w;
836                            }
837                            at = at.saturating_add(1);
838                        }
839                    }
840                    self.wr(dst, l, r)?;
841                }),
842                Op::Un { op, src, n } => {
843                    let (src, n) = (slot_at::<N>(*src), *n as usize);
844                    // each operation's own loop over the lanes
845                    macro_rules! un {
846                        ($w:ident => $e:expr) => {
847                            each!(active, l, {
848                                let s = self.rd(src, l)?;
849                                let mut r = [0u32; 4];
850                                for (out, &$w) in r.iter_mut().zip(s.iter()).take(n) {
851                                    *out = $e;
852                                }
853                                self.wr(dst, l, r)?;
854                            })
855                        };
856                    }
857                    // each arm a constant operation, so its loop is specialised
858                    match op {
859                        UnOp::NegF => un!(w => unary(UnOp::NegF, w)),
860                        UnOp::NegI => un!(w => unary(UnOp::NegI, w)),
861                        UnOp::Not => un!(w => unary(UnOp::Not, w)),
862                        UnOp::BitNot => un!(w => unary(UnOp::BitNot, w)),
863                    }
864                }
865                Op::Bin { op, a, b, n, a1, b1 } => {
866                    let (a, b, n) = (slot_at::<N>(*a), slot_at::<N>(*b), *n as usize);
867                    // a scalar operand's lane 0 for every component
868                    let (sa, sb) = (usize::from(!*a1), usize::from(!*b1));
869                    // each operation's own loop over the lanes, `bin` of a constant operation inlined into it
870                    macro_rules! bin {
871                        ($($o:ident),*) => {
872                            match op {
873                                $(BinOp::$o => each!(active, l, {
874                                    let x = self.rd(a, l)?;
875                                    let y = self.rd(b, l)?;
876                                    let mut r = [0u32; 4];
877                                    for (i, out) in r.iter_mut().enumerate().take(n) {
878                                        *out = bin(BinOp::$o, comp(&x, i.saturating_mul(sa)), comp(&y, i.saturating_mul(sb)), ftz)?;
879                                    }
880                                    self.wr(dst, l, r)?;
881                                }),)*
882                            }
883                        };
884                    }
885                    bin!(
886                        AddF, SubF, MulF, DivF, ModF, AddU, SubU, MulU, DivU, ModU, And, Or, Xor, Shl, ShrU,
887                        ShrI, EqW, NeW, EqF, NeF, LtF, LeF, GtF, GeF, LtU, LeU, GtU, GeU, LtI, LeI, GtI, GeI,
888                        LAnd, LOr
889                    )
890                }
891                Op::Select { c: cs, a, b, n, c1 } => {
892                    let (cs, a, b, n) = (slot_at::<N>(*cs), slot_at::<N>(*a), slot_at::<N>(*b), *n as usize);
893                    let sc = usize::from(!*c1);
894                    each!(active, l, {
895                        let cc = self.rd(cs, l)?;
896                        let x = self.rd(a, l)?;
897                        let y = self.rd(b, l)?;
898                        let mut r = [0u32; 4];
899                        for (i, out) in r.iter_mut().enumerate().take(n) {
900                            *out = if comp(&cc, i.saturating_mul(sc)) != 0 {
901                                comp(&x, i)
902                            } else {
903                                comp(&y, i)
904                            };
905                        }
906                        self.wr(dst, l, r)?;
907                    })
908                }
909                Op::Cvt { op, src, n } => {
910                    let (src, n) = (slot_at::<N>(*src), *n as usize);
911                    macro_rules! cvt {
912                        ($w:ident => $e:expr) => {
913                            each!(active, l, {
914                                let s = self.rd(src, l)?;
915                                let mut r = [0u32; 4];
916                                for (out, &$w) in r.iter_mut().zip(s.iter()).take(n) {
917                                    *out = $e;
918                                }
919                                self.wr(dst, l, r)?;
920                            })
921                        };
922                    }
923                    // each arm a constant operation, so its loop is specialised
924                    match op {
925                        CvtOp::F2I => cvt!(w => convert(CvtOp::F2I, w)),
926                        CvtOp::F2U => cvt!(w => convert(CvtOp::F2U, w)),
927                        CvtOp::I2F => cvt!(w => convert(CvtOp::I2F, w)),
928                        CvtOp::U2F => cvt!(w => convert(CvtOp::U2F, w)),
929                        CvtOp::B2F => cvt!(w => convert(CvtOp::B2F, w)),
930                        CvtOp::B2U => cvt!(w => convert(CvtOp::B2U, w)),
931                    }
932                }
933                Op::Math {
934                    f: mf,
935                    a,
936                    b,
937                    c: cc,
938                    n,
939                    c1,
940                } => {
941                    let (a, b, cc) = (slot_at::<N>(*a), slot_at::<N>(*b), slot_at::<N>(*cc));
942                    each!(active, l, {
943                        let x = self.rd(a, l)?;
944                        let y = self.rd(b, l)?;
945                        let mut z = self.rd(cc, l)?;
946                        if *c1 {
947                            z = [z[0]; 4];
948                        }
949                        self.wr(dst, l, math(*mf, x, y, z, *n, ftz))?;
950                    })
951                }
952                Op::CopyWords { src, dst, words } => {
953                    for l in 0..N {
954                        if active.has(l) {
955                            for k in 0..*words {
956                                let w = self.mem_get(*src, l, k, 1)?[0];
957                                let at = self.locals(offset(*dst, k)?, 1)?;
958                                self.set_local_word(at, l, w)?;
959                            }
960                        }
961                    }
962                }
963                Op::RqHit { q, committed, first } => {
964                    for l in 0..N {
965                        if !active.has(l) {
966                            continue;
967                        }
968                        let b = self.b;
969                        let query = self
970                            .query_slot(*q, l)?
971                            .as_ref()
972                            .ok_or("a ray query read before it was initialized")?;
973                        let scene = b
974                            .scenes
975                            .get(&query.scene())
976                            .ok_or("a ray query without its scene")?;
977                        let hit = if *committed {
978                            &query.committed
979                        } else {
980                            &query.candidate
981                        };
982                        let r = hit
983                            .as_ref()
984                            .map_or_else(Reported::none, |h| Reported::of(h, scene));
985                        let first = *first as usize;
986                        let mut put = |k: usize, v: [u32; 4]| -> Result<(), String> {
987                            self.wr(first.saturating_add(k).saturating_mul(N), l, v)
988                        };
989                        put(0, [r.kind; 4])?;
990                        put(1, [r.t.to_bits(); 4])?;
991                        put(2, [r.custom; 4])?;
992                        put(3, [r.instance; 4])?;
993                        put(4, [r.sbt; 4])?;
994                        put(5, [r.geometry; 4])?;
995                        put(6, [r.primitive; 4])?;
996                        let [bx, by] = r.barycentrics;
997                        put(7, [bx.to_bits(), by.to_bits(), 0, 0])?;
998                        put(8, [u32::from(r.front); 4])?;
999                        for (c, (o, w)) in r.object_to_world.iter().zip(&r.world_to_object).enumerate() {
1000                            let [ox, oy, oz] = o.map(f32::to_bits);
1001                            put(9usize.saturating_add(c), [ox, oy, oz, 0])?;
1002                            let [wx, wy, wz] = w.map(f32::to_bits);
1003                            put(13usize.saturating_add(c), [wx, wy, wz, 0])?;
1004                        }
1005                    }
1006                }
1007                Op::ArrLen {
1008                    gid,
1009                    elem_words,
1010                    pick: Some(pick),
1011                } => {
1012                    for l in 0..N {
1013                        if active.has(l) {
1014                            let e = self.element_index(*gid, *pick, l)?;
1015                            let words = u32::try_from(self.element_words(*gid, e)?.len())
1016                                .map_err(|_| "an array of buffers over 4294967295 words")?;
1017                            self.wr(dst, l, [words.checked_div(*elem_words).unwrap_or(0); 4])?;
1018                        }
1019                    }
1020                }
1021                Op::ArrLen {
1022                    gid,
1023                    elem_words,
1024                    pick: None,
1025                } => {
1026                    let bound = self.b;
1027                    let count =
1028                        |n: usize| u32::try_from(n).map_err(|_| "a buffer over 4294967295 words".to_owned());
1029                    let words = match bound.data.at(*gid as usize)? {
1030                        _ if matches!(self.store.get(*gid as usize), Some(Some(_))) => count(
1031                            self.store
1032                                .get(*gid as usize)
1033                                .and_then(Option::as_ref)
1034                                .map_or(0, Vec::len),
1035                        )?,
1036                        Words::Elements(_) | Words::ConstantElements(_) => {
1037                            return Err("arrayLength of an array of buffers without an element".into());
1038                        }
1039                        Words::Owned(d) => count(d.len())?,
1040                        Words::Slice(d) => count(d.len())?,
1041                        Words::Rows(d) => count(d.len())?.saturating_mul(4),
1042                        Words::Backend => match self.c.globals.at(*gid as usize)?.kind {
1043                            GlobalKind::Storage(view) => self.sampler.buffer_words(view)?,
1044                            _ => return Err("arrayLength of a global that is not a buffer".into()),
1045                        },
1046                    };
1047                    let len = words.checked_div(*elem_words).unwrap_or(0);
1048                    for l in 0..N {
1049                        self.wr(dst, l, [len; 4])?;
1050                    }
1051                }
1052                Op::Deriv {
1053                    src,
1054                    axis,
1055                    fine,
1056                    n,
1057                    prov,
1058                } => {
1059                    if self.grouping == Grouping::Quad {
1060                        if active != M::first(4) {
1061                            return Err("a quad's lanes diverge at a derivative".into());
1062                        }
1063                        let mut v = [[0u32; 4]; 4];
1064                        for (l, lane) in v.iter_mut().enumerate() {
1065                            *lane = self.rd(slot_at::<N>(*src), l)?;
1066                        }
1067                        for l in 0..4 {
1068                            let (from, to) = match (*axis == 0, *fine) {
1069                                (true, false) => (0, 1),
1070                                (false, false) => (0, 2),
1071                                (true, true) => {
1072                                    if l < 2 {
1073                                        (0, 1)
1074                                    } else {
1075                                        (2, 3)
1076                                    }
1077                                }
1078                                (false, true) => {
1079                                    if l % 2 == 0 {
1080                                        (0, 2)
1081                                    } else {
1082                                        (1, 3)
1083                                    }
1084                                }
1085                            };
1086                            let mut r = [0u32; 4];
1087                            let (a, b) = (
1088                                v.get(to).copied().unwrap_or_default(),
1089                                v.get(from).copied().unwrap_or_default(),
1090                            );
1091                            for (out, (&x, &y)) in r.iter_mut().zip(a.iter().zip(&b)).take(*n as usize) {
1092                                *out = fl(fl(f(x), ftz) - fl(f(y), ftz), ftz).to_bits();
1093                            }
1094                            self.wr(dst, l, r)?;
1095                        }
1096                    } else {
1097                        let p = prov.as_ref().ok_or("a derivative of a computed value (a scalar evaluation has those of its inputs only)")?;
1098                        let d = self.derivs.ok_or("no derivatives given")?;
1099                        let mut r = [0u32; 4];
1100                        for (out, &(arg, lane)) in r.iter_mut().zip(p.iter()).take(*n as usize) {
1101                            *out = d
1102                                .get(arg as usize)
1103                                .ok_or("no derivatives given for the argument")?
1104                                .get(*axis as usize)
1105                                .and_then(|axis| axis.get(lane as usize))
1106                                .copied()
1107                                .ok_or("no derivative given for the axis or the lane")?;
1108                        }
1109                        for l in 0..N {
1110                            if active.has(l) {
1111                                self.wr(dst, l, r)?;
1112                            }
1113                        }
1114                    }
1115                }
1116                Op::Sample {
1117                    tex,
1118                    smp,
1119                    coord,
1120                    cn,
1121                    layer,
1122                    level,
1123                    dim,
1124                    offset,
1125                    gather,
1126                    depth_ref,
1127                    depth,
1128                } => {
1129                    let cn = *cn as usize;
1130                    let c3 = |me: &Self, l: usize, s: u32| -> Result<[f32; 3], String> {
1131                        let [x, y, z, _] = me.rd(slot_at::<N>(s), l)?;
1132                        let lane = |i: usize, w: u32| if i < cn { fl(f(w), ftz) } else { 0.0 };
1133                        Ok([lane(0, x), lane(1, y), lane(2, z)])
1134                    };
1135                    // a quad's implicit level of detail comes from the coordinates' differences
1136                    let footprint = if self.grouping == Grouping::Quad
1137                        && matches!(level, LevelOp::Auto | LevelOp::Bias(_))
1138                    {
1139                        if active != M::first(4) {
1140                            return Err("a quad's lanes diverge at an implicit-LOD sample".into());
1141                        }
1142                        let mut uv = [[0f32; 3]; 4];
1143                        for (l, corner) in uv.iter_mut().enumerate() {
1144                            *corner = c3(self, l, *coord)?;
1145                        }
1146                        let [origin, right, below, _] = uv;
1147                        let d = |a: [f32; 3]| -> [f32; 3] {
1148                            let [ax, ay, az] = a;
1149                            let [ox, oy, oz] = origin;
1150                            [ax - ox, ay - oy, az - oz]
1151                        };
1152                        Some([d(right), d(below)])
1153                    } else {
1154                        None
1155                    };
1156                    for l in 0..N {
1157                        if !active.has(l) {
1158                            continue;
1159                        }
1160                        let rd = |s: u32, me: &Self| me.rd(slot_at::<N>(s), l);
1161                        let lod = match (level, footprint) {
1162                            (LevelOp::Auto, None) => ImageLod::Implicit,
1163                            (LevelOp::Auto, Some(g)) => ImageLod::Gradients(g),
1164                            (LevelOp::Zero, _) => ImageLod::Level(0.0),
1165                            (LevelOp::Exact(s), _) => ImageLod::Level(fl(f(rd(*s, self)?[0]), ftz)),
1166                            (LevelOp::Bias(s), None) => ImageLod::Bias(fl(f(rd(*s, self)?[0]), ftz)),
1167                            (LevelOp::Bias(s), Some(g)) => {
1168                                ImageLod::GradientsBias(g, fl(f(rd(*s, self)?[0]), ftz))
1169                            }
1170                            (LevelOp::Gradient(x, y), _) => {
1171                                ImageLod::Gradients([c3(self, l, *x)?, c3(self, l, *y)?])
1172                            }
1173                        };
1174                        let req = ImageSample {
1175                            texture: self.element(tex, l)?,
1176                            sampler: self.element(smp, l)?,
1177                            dim: *dim,
1178                            coord: c3(self, l, *coord)?,
1179                            layer: layer.map(|s| rd(s, self).map(|r| r[0] as f32)).transpose()?,
1180                            lod,
1181                            offset: *offset,
1182                            gather: *gather,
1183                            depth_ref: depth_ref
1184                                .map(|s| rd(s, self).map(|r| fl(f(r[0]), ftz)))
1185                                .transpose()?,
1186                            depth: *depth,
1187                        };
1188                        let r = self.sampler.sample(req)?;
1189                        if let Some(count) = self.samples.get_mut(l) {
1190                            *count = count.saturating_add(1);
1191                        }
1192                        let [r, g, b, a] = r;
1193                        self.wr(dst, l, [r.to_bits(), g.to_bits(), b.to_bits(), a.to_bits()])?;
1194                    }
1195                }
1196                Op::Texel {
1197                    tex,
1198                    coord,
1199                    cn,
1200                    dim,
1201                    storage,
1202                    layer,
1203                    sample,
1204                    mip,
1205                } => {
1206                    for l in 0..N {
1207                        if !active.has(l) {
1208                            continue;
1209                        }
1210                        let cd = self.rd(slot_at::<N>(*coord), l)?;
1211                        let one = |s: &Option<u32>| -> Result<Option<i32>, String> {
1212                            s.map(|s| self.rd(slot_at::<N>(s), l).map(|r| r[0] as i32))
1213                                .transpose()
1214                        };
1215                        let req = ImageLoad {
1216                            texture: self.element(tex, l)?,
1217                            dim: *dim,
1218                            storage: *storage,
1219                            texel: texel_of(&cd, *cn as usize).map(|t| t as i32),
1220                            layer: one(layer)?,
1221                            sample: one(sample)?,
1222                            mip: one(mip)?,
1223                        };
1224                        let loaded = self.sampler.load(req)?;
1225                        self.wr(dst, l, loaded)?;
1226                    }
1227                }
1228                Op::Query {
1229                    tex,
1230                    dim,
1231                    kind,
1232                    level,
1233                } => {
1234                    for l in 0..N {
1235                        if !active.has(l) {
1236                            continue;
1237                        }
1238                        let level = level
1239                            .map(|s| self.rd(slot_at::<N>(s), l).map(|r| r[0]))
1240                            .transpose()?;
1241                        let texture = self.element(tex, l)?;
1242                        let answer = self.sampler.query(ImageQuery {
1243                            texture,
1244                            dim: *dim,
1245                            kind: *kind,
1246                            level,
1247                        })?;
1248                        self.wr(dst, l, answer)?;
1249                    }
1250                }
1251            }
1252        }
1253        Ok(())
1254    }
1255
1256    /// The lanes' states of ray query `q`.
1257    fn query_slots(&mut self, q: u32) -> Result<&mut Vec<Option<Query>>, String> {
1258        if self.queries.len() <= q as usize {
1259            self.queries.resize_with((q as usize).saturating_add(1), Vec::new);
1260        }
1261        let len = self.queries.len();
1262        let slots = self
1263            .queries
1264            .get_mut(q as usize)
1265            .ok_or_else(|| format!("internal error: ray query {q}, where the run has {len}"))?;
1266        if slots.len() < N {
1267            slots.resize(N, None);
1268        }
1269        Ok(slots)
1270    }
1271
1272    /// The state of ray query `q` in lane `lane`.
1273    fn query_slot(&mut self, q: u32, lane: usize) -> Result<&mut Option<Query>, String> {
1274        self.query_slots(q)?
1275            .get_mut(lane)
1276            .ok_or_else(|| format!("internal error: ray query lane {lane}"))
1277    }
1278
1279    /// The wave width and the lane count of a compute run; wave operations outside one are errors.
1280    fn wave_shape(&self) -> Result<(usize, usize), String> {
1281        if self.wave == 0 || self.wave > 128 {
1282            return Err("a wave operation outside a compute run (or of a wave wider than 128 lanes)".into());
1283        }
1284        Ok((self.wave, self.lanes))
1285    }
1286
1287    /// The resource an operation names for lane `l`: for an array, the element its index picks in that lane.
1288    fn element(&self, res: &Res, l: usize) -> Result<Element, String> {
1289        let index = match res.pick {
1290            Pick::One => 0,
1291            Pick::Fixed(i) => i,
1292            Pick::Slot(s) => self.rd(slot_at::<N>(s), l)?[0],
1293        };
1294        if !matches!(res.pick, Pick::One) && res.count != iridia::Resource::UNBOUNDED && index >= res.count {
1295            return Err(format!(
1296                "resource {}: element {index} of an array of {} (D3D12 leaves it undefined)",
1297                res.id.0, res.count
1298            ));
1299        }
1300        Ok(Element { id: res.id, index })
1301    }
1302
1303    fn block(&mut self, b: &'a [St], active: M) -> Result<Exits<M>, String> {
1304        let mut ex = Exits::<M>::default();
1305        let mut active = active;
1306        for st in b {
1307            if active.is_empty() {
1308                break;
1309            }
1310            match st {
1311                St::Run(lo, hi) => self.ops(*lo, *hi, active)?,
1312                St::Store {
1313                    loc: Loc::Local(o),
1314                    src,
1315                    n,
1316                } => {
1317                    let (base, n, src) = (
1318                        self.locals(*o as usize, *n as usize)?,
1319                        *n as usize,
1320                        slot_at::<N>(*src),
1321                    );
1322                    for l in 0..N {
1323                        if active.has(l) {
1324                            let v = self.rd(src, l)?;
1325                            for (i, w) in v.iter().take(n).enumerate() {
1326                                self.set_local_word(base.saturating_add(i), l, *w)?;
1327                            }
1328                        }
1329                    }
1330                }
1331                St::Store { loc, src, n } => {
1332                    for l in 0..N {
1333                        if active.has(l) {
1334                            let v = self.rd(slot_at::<N>(*src), l)?;
1335                            self.mem_put(*loc, l, 0, *n, v)?;
1336                        }
1337                    }
1338                }
1339                St::StoreDyn { ptr, src, n } => {
1340                    for l in 0..N {
1341                        if active.has(l) {
1342                            let x = self.dyn_extra(ptr, l)?;
1343                            let v = self.rd(slot_at::<N>(*src), l)?;
1344                            self.mem_put(ptr.loc, l, x, *n, v)?;
1345                        }
1346                    }
1347                }
1348                St::ImageStore {
1349                    tex,
1350                    dim,
1351                    coord,
1352                    cn,
1353                    layer,
1354                    value,
1355                    vn,
1356                } => {
1357                    for l in 0..N {
1358                        if active.has(l) {
1359                            let cd = self.rd(slot_at::<N>(*coord), l)?;
1360                            let element = self.element(tex, l)?;
1361                            self.sampler.store_image(ImageStore {
1362                                view: element.id,
1363                                index: element.index,
1364                                dim: *dim,
1365                                texel: texel_of(&cd, *cn as usize),
1366                                layer: layer
1367                                    .map(|s| self.rd(slot_at::<N>(s), l).map(|r| r[0]))
1368                                    .transpose()?,
1369                                value: self.rd(slot_at::<N>(*value), l)?,
1370                                lanes: *vn,
1371                            })?;
1372                        }
1373                    }
1374                }
1375                St::ImageAtomic {
1376                    tex,
1377                    dim,
1378                    coord,
1379                    cn,
1380                    layer,
1381                    op,
1382                    signed,
1383                    value,
1384                    result,
1385                } => {
1386                    for l in 0..N {
1387                        if active.has(l) {
1388                            let cd = self.rd(slot_at::<N>(*coord), l)?;
1389                            let element = self.element(tex, l)?;
1390                            let old = self.sampler.atomic_image(ImageAtomic {
1391                                view: element.id,
1392                                index: element.index,
1393                                dim: *dim,
1394                                texel: texel_of(&cd, *cn as usize),
1395                                layer: layer
1396                                    .map(|s| self.rd(slot_at::<N>(s), l).map(|r| r[0]))
1397                                    .transpose()?,
1398                                op: *op,
1399                                signed: *signed,
1400                                value: self.rd(slot_at::<N>(*value), l)?[0],
1401                            })?;
1402                            if let Some(r) = result {
1403                                self.wr(slot_at::<N>(*r), l, [old, 0, 0, 0])?;
1404                            }
1405                        }
1406                    }
1407                }
1408                St::Atomic {
1409                    ptr,
1410                    fun,
1411                    signed,
1412                    value,
1413                    compare,
1414                    result,
1415                    exchanged,
1416                } => {
1417                    for l in 0..N {
1418                        if active.has(l) {
1419                            let (g, o, elem) = match ptr.loc {
1420                                Loc::Global(g, o) => (g, o, None),
1421                                Loc::Elem(g, o, pick) => (g, o, Some(self.element_index(g, pick, l)?)),
1422                                Loc::Local(_) => return Err("atomic on a local".into()),
1423                            };
1424                            let word = offset(o, self.dyn_extra(ptr, l)?)?;
1425                            let v = self.rd(slot_at::<N>(*value), l)?[0];
1426                            let expected = match compare {
1427                                Some(c) => self.rd(slot_at::<N>(*c), l)?[0],
1428                                None => 0,
1429                            };
1430                            let memory = match elem {
1431                                Some(e) => self.elem_store.get_mut(g as usize).and_then(|v| v.get_mut(&e)),
1432                                None => self.store.get_mut(g as usize).and_then(Option::as_mut),
1433                            };
1434                            let (old, swapped) = if let Some(d) = memory {
1435                                let cell = d.get_mut(word).ok_or("atomic past the end of the memory")?;
1436                                let old = *cell;
1437                                let (new, swapped) = atomic(*fun, *signed, old, v, expected);
1438                                *cell = new;
1439                                (old, swapped)
1440                            } else {
1441                                let view = match self.c.globals.at(g as usize)?.kind {
1442                                    GlobalKind::Storage(view) => view,
1443                                    GlobalKind::StorageArray(..) => {
1444                                        return Err("atomic on an element of an array of buffers this run does not hold (only a compute run writes them; the backend is not asked for element atomics)".into());
1445                                    }
1446                                    _ => return Err("atomic on a read-only global".into()),
1447                                };
1448                                if *fun != AtomOp::Add {
1449                                    return Err(format!(
1450                                        "atomic {fun:?} on a buffer the backend holds (only add is passed on)"
1451                                    ));
1452                                }
1453                                (
1454                                    self.sampler.atomic_add(
1455                                        view,
1456                                        u32::try_from(word)
1457                                            .map_err(|_| "atomic past the end of the memory")?,
1458                                        v,
1459                                    )?,
1460                                    true,
1461                                )
1462                            };
1463                            if let Some(r) = result {
1464                                self.wr(slot_at::<N>(*r), l, [old; 4])?;
1465                            }
1466                            if let Some(x) = exchanged {
1467                                self.wr(slot_at::<N>(*x), l, [u32::from(swapped); 4])?;
1468                            }
1469                        }
1470                    }
1471                }
1472                St::WaveBallot { predicate, result } => {
1473                    let (w, lanes) = self.wave_shape()?;
1474                    for base in (0..lanes).step_by(w) {
1475                        let end = base.saturating_add(w).min(lanes);
1476                        let mut bits = [0u32; 4];
1477                        for l in base..end {
1478                            let on = match predicate {
1479                                Some(p) => self.rd(slot_at::<N>(*p), l)?[0] != 0,
1480                                None => true,
1481                            };
1482                            if active.has(l) && on {
1483                                let bit = l.saturating_sub(base);
1484                                if let Some(word) = bits.get_mut(bit / 32) {
1485                                    *word |= 1u32
1486                                        .checked_shl(u32::try_from(bit % 32).unwrap_or(0))
1487                                        .unwrap_or(0);
1488                                }
1489                            }
1490                        }
1491                        for l in base..end {
1492                            if active.has(l) {
1493                                self.wr(slot_at::<N>(*result), l, bits)?;
1494                            }
1495                        }
1496                    }
1497                }
1498                St::WaveGather {
1499                    mode,
1500                    value,
1501                    n,
1502                    result,
1503                } => {
1504                    let (w, lanes) = self.wave_shape()?;
1505                    for base in (0..lanes).step_by(w) {
1506                        let end = base.saturating_add(w).min(lanes);
1507                        let first = (base..end).find(|&l| active.has(l));
1508                        let mut taken: Vec<(usize, Reg)> = Vec::new();
1509                        for l in base..end {
1510                            if !active.has(l) {
1511                                continue;
1512                            }
1513                            let at = l.saturating_sub(base);
1514                            let arg = |s: u32| -> Result<usize, String> {
1515                                self.rd(slot_at::<N>(s), l).map(|r| r[0] as usize)
1516                            };
1517                            let from = match mode {
1518                                WaveGather::First => first.unwrap_or(l),
1519                                WaveGather::Lane(s) | WaveGather::Shuffle(s) => base.saturating_add(arg(*s)?),
1520                                WaveGather::Down(s) => base.saturating_add(at).saturating_add(arg(*s)?),
1521                                WaveGather::Up(s) => {
1522                                    let step = arg(*s)?;
1523                                    if step <= at {
1524                                        base.saturating_add(at.saturating_sub(step))
1525                                    } else {
1526                                        base.saturating_add(at)
1527                                    }
1528                                }
1529                                WaveGather::Xor(s) => base.saturating_add(at ^ arg(*s)?),
1530                                WaveGather::QuadLane(s) => {
1531                                    base.saturating_add(at & !3).saturating_add(arg(*s)? & 3)
1532                                }
1533                                WaveGather::QuadSwap(x) => base.saturating_add(at ^ *x as usize),
1534                            };
1535                            // a lane outside the wave is undefined in D3D: the lane's own value
1536                            let from = if from < end { from } else { l };
1537                            let mut r = self.rd(slot_at::<N>(*value), from)?;
1538                            for w in r.iter_mut().skip(*n as usize) {
1539                                *w = 0;
1540                            }
1541                            taken.push((l, r));
1542                        }
1543                        for (l, r) in taken {
1544                            self.wr(slot_at::<N>(*result), l, r)?;
1545                        }
1546                    }
1547                }
1548                St::WaveOp {
1549                    op,
1550                    scan,
1551                    value,
1552                    n,
1553                    kind,
1554                    result,
1555                } => {
1556                    let (w, lanes) = self.wave_shape()?;
1557                    let ftz = self.ftz;
1558                    for base in (0..lanes).step_by(w) {
1559                        let end = base.saturating_add(w).min(lanes);
1560                        let members: Vec<usize> = (base..end).filter(|&l| active.has(l)).collect();
1561                        for i in 0..*n as usize {
1562                            let values: Vec<u32> = members
1563                                .iter()
1564                                .map(|&l| {
1565                                    self.rd(slot_at::<N>(*value), l)
1566                                        .map(|r| r.get(i).copied().unwrap_or(0))
1567                                })
1568                                .collect::<Result<_, _>>()?;
1569                            let total = values
1570                                .iter()
1571                                .copied()
1572                                .reduce(|a, b| wave_combine(*op, *kind, a, b, ftz))
1573                                .unwrap_or(0);
1574                            let mut running = wave_identity(*op, *kind);
1575                            for (k, &l) in members.iter().enumerate() {
1576                                let value = values.get(k).copied().unwrap_or(0);
1577                                let out = match scan {
1578                                    WaveScan::Reduce => total,
1579                                    WaveScan::Inclusive => {
1580                                        running = if k == 0 {
1581                                            value
1582                                        } else {
1583                                            wave_combine(*op, *kind, running, value, ftz)
1584                                        };
1585                                        running
1586                                    }
1587                                    WaveScan::Exclusive => {
1588                                        let before = running;
1589                                        running = if k == 0 {
1590                                            value
1591                                        } else {
1592                                            wave_combine(*op, *kind, running, value, ftz)
1593                                        };
1594                                        before
1595                                    }
1596                                };
1597                                let mut r = self.rd(slot_at::<N>(*result), l)?;
1598                                if let Some(lane) = r.get_mut(i) {
1599                                    *lane = out;
1600                                }
1601                                self.wr(slot_at::<N>(*result), l, r)?;
1602                            }
1603                        }
1604                    }
1605                }
1606                St::StoreAgg { loc, members } => {
1607                    for l in 0..N {
1608                        if active.has(l) {
1609                            let mut at = 0;
1610                            for &(s, n) in members {
1611                                let v = self.rd(slot_at::<N>(s), l)?;
1612                                self.mem_put(*loc, l, at, n, v)?;
1613                                at = at.saturating_add(u32::from(n));
1614                            }
1615                        }
1616                    }
1617                }
1618                St::RqInit {
1619                    q,
1620                    scene,
1621                    flags,
1622                    mask,
1623                    tmin,
1624                    tmax,
1625                    origin,
1626                    direction,
1627                } => {
1628                    if !self.b.scenes.contains_key(scene) {
1629                        return Err(format!(
1630                            "no acceleration structure given for resource {}",
1631                            scene.0
1632                        ));
1633                    }
1634                    for l in 0..N {
1635                        if active.has(l) {
1636                            let reg = |s: u32| self.rd(slot_at::<N>(s), l);
1637                            let float = |s: u32| -> Result<f32, String> { Ok(fl(f(reg(s)?[0]), self.ftz)) };
1638                            let vec3 = |s: u32| -> Result<[f32; 3], String> {
1639                                let [x, y, z, _] = reg(s)?;
1640                                Ok([fl(f(x), self.ftz), fl(f(y), self.ftz), fl(f(z), self.ftz)])
1641                            };
1642                            let query = Query::new(
1643                                *scene,
1644                                reg(*flags)?[0],
1645                                reg(*mask)?[0],
1646                                float(*tmin)?,
1647                                float(*tmax)?,
1648                                vec3(*origin)?,
1649                                vec3(*direction)?,
1650                            );
1651                            *self.query_slot(*q, l)? = Some(query);
1652                        }
1653                    }
1654                }
1655                St::RqProceed { q, result } => {
1656                    for l in 0..N {
1657                        if active.has(l) {
1658                            let b = self.b;
1659                            let query = self
1660                                .query_slot(*q, l)?
1661                                .as_mut()
1662                                .ok_or("a ray query proceeded before it was initialized")?;
1663                            let scene = b
1664                                .scenes
1665                                .get(&query.scene())
1666                                .ok_or("a ray query without its scene")?;
1667                            let more = query.proceed(scene);
1668                            self.wr(slot_at::<N>(*result), l, [u32::from(more); 4])?;
1669                        }
1670                    }
1671                }
1672                St::RqConfirm { q } => {
1673                    for l in 0..N {
1674                        if active.has(l) {
1675                            self.query_slot(*q, l)?
1676                                .as_mut()
1677                                .ok_or("a ray query confirmed before it was initialized")?
1678                                .confirm()?;
1679                        }
1680                    }
1681                }
1682                St::RtDecide(decision) => {
1683                    for l in 0..N {
1684                        if active.has(l)
1685                            && let Some(slot) = self.decisions.get_mut(l)
1686                        {
1687                            *slot = *decision;
1688                        }
1689                    }
1690                }
1691                St::RtTrace {
1692                    scene,
1693                    operands,
1694                    payload,
1695                } => {
1696                    for l in 0..N {
1697                        if !active.has(l) {
1698                            continue;
1699                        }
1700                        let operand = |k: usize| -> Result<Reg, String> {
1701                            let slot = operands.at(k)?;
1702                            self.rd(slot_at::<N>(*slot), l)
1703                        };
1704                        let word = |k: usize| -> Result<u32, String> { Ok(operand(k)?[0]) };
1705                        let vec3 = |k: usize| -> Result<[f32; 3], String> {
1706                            let [x, y, z, _] = operand(k)?;
1707                            Ok([f(x), f(y), f(z)])
1708                        };
1709                        let (flags, mask, hit_offset, stride, miss) =
1710                            (word(0)?, word(1)?, word(2)?, word(3)?, word(4)?);
1711                        let (origin, direction) = (vec3(5)?, vec3(7)?);
1712                        let (t_min, t_max) = (f(word(6)?), f(word(8)?));
1713                        let at = payload.0 as usize;
1714                        let mut words = self.read_locals(l, at, payload.1 as usize)?;
1715                        self.sampler.trace_ray(super::request::TraceRay {
1716                            scene: *scene,
1717                            flags,
1718                            mask,
1719                            hit_group_offset: hit_offset,
1720                            hit_group_stride: stride,
1721                            miss_index: miss,
1722                            origin,
1723                            t_min,
1724                            direction,
1725                            t_max,
1726                            payload: &mut words,
1727                        })?;
1728                        self.write_locals(l, at, &words)?;
1729                    }
1730                }
1731                St::RtReport {
1732                    t,
1733                    kind,
1734                    attributes,
1735                    result,
1736                } => {
1737                    for l in 0..N {
1738                        if !active.has(l) {
1739                            continue;
1740                        }
1741                        let at = attributes.0 as usize;
1742                        let mut words = self.read_locals(l, at, attributes.1 as usize)?;
1743                        let accepted = self.sampler.report_hit(super::request::ReportHit {
1744                            t: f(self.rd(slot_at::<N>(*t), l)?[0]),
1745                            hit_kind: self.rd(slot_at::<N>(*kind), l)?[0],
1746                            attributes: &mut words,
1747                        })?;
1748                        self.write_locals(l, at, &words)?;
1749                        if let Some(r) = result {
1750                            self.wr(slot_at::<N>(*r), l, [u32::from(accepted); 4])?;
1751                        }
1752                    }
1753                }
1754                St::RtCall { index, parameter } => {
1755                    for l in 0..N {
1756                        if !active.has(l) {
1757                            continue;
1758                        }
1759                        let at = parameter.0 as usize;
1760                        let mut words = self.read_locals(l, at, parameter.1 as usize)?;
1761                        self.sampler.call_shader(super::request::CallShader {
1762                            index: self.rd(slot_at::<N>(*index), l)?[0],
1763                            parameter: &mut words,
1764                        })?;
1765                        self.write_locals(l, at, &words)?;
1766                    }
1767                }
1768                St::RqGenerate { q, t } => {
1769                    for l in 0..N {
1770                        if active.has(l) {
1771                            let t = fl(f(self.rd(slot_at::<N>(*t), l)?[0]), self.ftz);
1772                            self.query_slot(*q, l)?
1773                                .as_mut()
1774                                .ok_or("an intersection generated before the query was initialized")?
1775                                .generate(t)?;
1776                        }
1777                    }
1778                }
1779                St::RqTerminate { q } => {
1780                    for l in 0..N {
1781                        if active.has(l) {
1782                            self.query_slot(*q, l)?
1783                                .as_mut()
1784                                .ok_or("a ray query ended before it was initialized")?
1785                                .terminate();
1786                        }
1787                    }
1788                }
1789                St::Barrier => {
1790                    // lockstep makes every write before the barrier visible after it; what a barrier needs is that
1791                    // every invocation still in the program is at it (D3D leaves a divergent barrier undefined)
1792                    if self.store.is_empty() {
1793                        return Err("a barrier outside a compute run".into());
1794                    }
1795                    if active != self.group.minus(self.retired) {
1796                        return Err("a barrier that not every invocation of the group reaches".into());
1797                    }
1798                }
1799                St::If { c, a, r } => {
1800                    let mut take = M::none();
1801                    for l in 0..N {
1802                        if active.has(l) && self.rd(slot_at::<N>(*c), l)?[0] != 0 {
1803                            take.insert(l);
1804                        }
1805                    }
1806                    let e1 = self.block(a, take)?;
1807                    let e2 = self.block(r, active.minus(take))?;
1808                    let (b, cn, rt) = (e1.brk.or(e2.brk), e1.cont.or(e2.cont), e1.ret.or(e2.ret));
1809                    ex.brk = ex.brk.or(b);
1810                    ex.cont = ex.cont.or(cn);
1811                    ex.ret = ex.ret.or(rt);
1812                    active = active.minus(b.or(cn).or(rt));
1813                }
1814                St::Loop { body, cont, break_if } => {
1815                    let mut lm = active;
1816                    loop {
1817                        self.count(lm, 1)?;
1818                        let e = self.block(body, lm)?;
1819                        ex.ret = ex.ret.or(e.ret);
1820                        lm = lm.minus(e.ret.or(e.brk));
1821                        if lm.is_empty() {
1822                            break;
1823                        }
1824                        let e2 = self.block(cont, lm)?;
1825                        ex.ret = ex.ret.or(e2.ret);
1826                        lm = lm.minus(e2.ret.or(e2.brk));
1827                        if let Some(bi) = break_if {
1828                            for l in 0..N {
1829                                if lm.has(l) && self.rd(slot_at::<N>(*bi), l)?[0] != 0 {
1830                                    lm.remove(l);
1831                                }
1832                            }
1833                        }
1834                        if lm.is_empty() {
1835                            break;
1836                        }
1837                    }
1838                    active = active.minus(ex.ret);
1839                }
1840                St::Switch { sel, cases } => {
1841                    let mut chosen = [usize::MAX; N];
1842                    for (l, slot) in chosen.iter_mut().enumerate() {
1843                        if active.has(l) {
1844                            let v = self.rd(slot_at::<N>(*sel), l)?[0];
1845                            *slot = cases
1846                                .iter()
1847                                .position(|c| c.0 == Some(v))
1848                                .or_else(|| cases.iter().position(|c| c.0.is_none()))
1849                                .unwrap_or(usize::MAX);
1850                        }
1851                    }
1852                    let mut carry = M::none();
1853                    for (i, (_, fall, body)) in cases.iter().enumerate() {
1854                        let mut enter = carry;
1855                        for (l, slot) in chosen.iter().enumerate() {
1856                            if *slot == i {
1857                                enter.insert(l);
1858                            }
1859                        }
1860                        carry = M::none();
1861                        if enter.is_empty() {
1862                            continue;
1863                        }
1864                        let e = self.block(body, enter)?;
1865                        ex.cont = ex.cont.or(e.cont);
1866                        ex.ret = ex.ret.or(e.ret);
1867                        let done = e.brk.or(e.cont).or(e.ret);
1868                        if *fall {
1869                            carry = enter.minus(done);
1870                        }
1871                    }
1872                    active = active.minus(ex.cont.or(ex.ret));
1873                }
1874                St::Break => {
1875                    ex.brk = ex.brk.or(active);
1876                    active = M::none();
1877                }
1878                St::Continue => {
1879                    ex.cont = ex.cont.or(active);
1880                    active = M::none();
1881                }
1882                St::Return(result) => {
1883                    self.retired = self.retired.or(active);
1884                    for l in 0..N {
1885                        if active.has(l) {
1886                            // into the lane's own room, kept from run to run
1887                            let mut members = std::mem::take(
1888                                self.outs
1889                                    .get_mut(l)
1890                                    .ok_or("internal error: a result for a lane the run does not have")?,
1891                            );
1892                            members.clear();
1893                            for member in result {
1894                                match *member {
1895                                    Member::Reg(s, n) => {
1896                                        let mut r = [0u32; 4];
1897                                        let v = self.rd(slot_at::<N>(s), l)?;
1898                                        for (to, from) in r.iter_mut().zip(v.iter()).take(n as usize) {
1899                                            *to = *from;
1900                                        }
1901                                        members.push(r);
1902                                    }
1903                                    Member::Mem {
1904                                        word, count, width, ..
1905                                    } => {
1906                                        for k in 0..count as usize {
1907                                            let at = (word as usize)
1908                                                .saturating_add(k.saturating_mul(width as usize));
1909                                            let mut r = [0u32; 4];
1910                                            for (i, w) in r.iter_mut().take(width as usize).enumerate() {
1911                                                *w = self.local_word(at.saturating_add(i), l)?;
1912                                            }
1913                                            members.push(r);
1914                                        }
1915                                    }
1916                                }
1917                            }
1918                            if let Some(room) = self.outs.get_mut(l) {
1919                                *room = members;
1920                            }
1921                        }
1922                    }
1923                    ex.ret = ex.ret.or(active);
1924                    active = M::none();
1925                }
1926                St::Kill => {
1927                    self.killed = self.killed.or(active);
1928                    self.retired = self.retired.or(active);
1929                    if self.grouping == Grouping::Independent {
1930                        ex.ret = ex.ret.or(active);
1931                        active = M::none();
1932                    }
1933                }
1934            }
1935        }
1936        Ok(ex)
1937    }
1938
1939    /// Run once, each lane with its arguments (one register per entry argument, the lanes of its value first).
1940    pub fn run(&mut self, args: [&[[u32; 4]]; N]) -> Result<[Output; N], String> {
1941        self.run_first(args, N)
1942    }
1943
1944    /// [`Runner::run`] with lanes `0..count` only; the outputs of the others are empty (independent invocations).
1945    pub fn run_first(&mut self, args: [&[[u32; 4]]; N], count: usize) -> Result<[Output; N], String> {
1946        self.run_in_place(args, count)?;
1947        let c: &'a Executable = self.c;
1948        let mut outputs = Vec::with_capacity(N);
1949        for l in 0..N {
1950            let (values, discarded, steps, samples) = self.lane(l);
1951            outputs.push(Output {
1952                pointees: c
1953                    .pointer_args
1954                    .iter()
1955                    .map(|&(_, off, words)| self.read_locals(l, off as usize, words as usize))
1956                    .collect::<Result<_, _>>()?,
1957                decision: self.decisions.get(l).copied().unwrap_or_default(),
1958                values: values.to_vec(),
1959                discarded,
1960                steps,
1961                samples,
1962            });
1963        }
1964        outputs
1965            .try_into()
1966            .map_err(|_| "internal error: a run of N lanes made another number of outputs".to_owned())
1967    }
1968
1969    /// [`Runner::run_first`] with the results left in the runner, for [`Runner::lane`] to read: a run that
1970    /// allocates nothing once the runner's room is grown.
1971    pub(super) fn run_in_place(&mut self, args: [&[[u32; 4]]; N], count: usize) -> Result<(), String> {
1972        if count > N || (count < N && self.grouping != Grouping::Independent) {
1973            return Err(format!(
1974                "{count} of {N} lanes run: only independent invocations run fewer than all"
1975            ));
1976        }
1977        self.steps = 0;
1978        self.lane_steps = [0; N];
1979        self.mem.fill(0);
1980        self.samples = [0; N];
1981        self.killed = M::none();
1982        self.queries.clear();
1983        self.decisions = [Decision::None; N];
1984        for o in self.outs.iter_mut() {
1985            o.clear();
1986        }
1987        for (l, given) in args.iter().enumerate() {
1988            if given.len() < self.c.arg_len {
1989                return Err(format!(
1990                    "{} argument registers given, the function takes {}",
1991                    given.len(),
1992                    self.c.arg_len
1993                ));
1994            }
1995            for (i, (s, n)) in self.c.args.iter().enumerate() {
1996                if *n != 0 {
1997                    let argument = given.at(*self.c.arg_at.at(i)?)?;
1998                    self.wr(slot_at::<N>(*s), l, *argument)?;
1999                }
2000            }
2001            for &(at, off, count, width) in &self.c.array_args {
2002                for k in 0..count as usize {
2003                    let base = (off as usize).saturating_add(k.saturating_mul(width as usize));
2004                    let element = given.at(at.saturating_add(k))?;
2005                    for (i, w) in element.iter().take(width as usize).enumerate() {
2006                        self.set_local_word(base.saturating_add(i), l, *w)?;
2007                    }
2008                }
2009            }
2010            for &(at, off, words) in &self.c.pointer_args {
2011                for w in 0..words as usize {
2012                    let register = given.at(at.saturating_add(w / 4))?;
2013                    self.set_local_word((off as usize).saturating_add(w), l, comp(register, w % 4))?;
2014                }
2015            }
2016        }
2017        let c: &'a Executable = self.c;
2018        self.block(&c.body, M::first(count))?;
2019        Ok(())
2020    }
2021
2022    /// Lane `l` of the last run: its result members, whether it was discarded, its steps and its samples.
2023    pub(super) fn lane(&self, l: usize) -> (&[Reg], bool, u64, u64) {
2024        let steps = match self.grouping {
2025            Grouping::Independent => self.lane_steps.get(l).copied().unwrap_or(0),
2026            Grouping::Quad | Grouping::Group => self.steps,
2027        };
2028        (
2029            self.outs.get(l).map_or(&[], Vec::as_slice),
2030            self.killed.has(l),
2031            steps,
2032            self.samples.get(l).copied().unwrap_or(0),
2033        )
2034    }
2035}
2036
2037/// Two words combined by a wave operation (booleans are 0 and 1).
2038fn wave_combine(op: WaveOp, kind: WaveKind, a: u32, b: u32, ftz: bool) -> u32 {
2039    let g = |v: u32| fl(f(v), ftz);
2040    match (op, kind) {
2041        (WaveOp::All, _) => u32::from(a != 0 && b != 0),
2042        (WaveOp::Any, _) => u32::from(a != 0 || b != 0),
2043        (WaveOp::And, _) => a & b,
2044        (WaveOp::Or, _) => a | b,
2045        (WaveOp::Xor, _) => a ^ b,
2046        (WaveOp::Add, WaveKind::F) => fl(g(a) + g(b), ftz).to_bits(),
2047        (WaveOp::Mul, WaveKind::F) => fl(g(a) * g(b), ftz).to_bits(),
2048        (WaveOp::Min, WaveKind::F) => g(a).min(g(b)).to_bits(),
2049        (WaveOp::Max, WaveKind::F) => g(a).max(g(b)).to_bits(),
2050        (WaveOp::Add, _) => a.wrapping_add(b),
2051        (WaveOp::Mul, _) => a.wrapping_mul(b),
2052        (WaveOp::Min, WaveKind::I) => (a as i32).min(b as i32) as u32,
2053        (WaveOp::Max, WaveKind::I) => (a as i32).max(b as i32) as u32,
2054        (WaveOp::Min, WaveKind::U) => a.min(b),
2055        (WaveOp::Max, WaveKind::U) => a.max(b),
2056    }
2057}
2058/// What a wave operation over no lanes gives, and an exclusive scan's first lane.
2059fn wave_identity(op: WaveOp, kind: WaveKind) -> u32 {
2060    match (op, kind) {
2061        (WaveOp::All, _) => 1,
2062        (WaveOp::Any | WaveOp::Or | WaveOp::Xor, _) => 0,
2063        (WaveOp::And, _) => u32::MAX,
2064        (WaveOp::Add, WaveKind::F) => 0,
2065        (WaveOp::Add, _) => 0,
2066        (WaveOp::Mul, WaveKind::F) => 1.0f32.to_bits(),
2067        (WaveOp::Mul, _) => 1,
2068        (WaveOp::Min, WaveKind::F) => f32::INFINITY.to_bits(),
2069        (WaveOp::Max, WaveKind::F) => f32::NEG_INFINITY.to_bits(),
2070        (WaveOp::Min, WaveKind::I) => i32::MAX as u32,
2071        (WaveOp::Max, WaveKind::I) => i32::MIN as u32,
2072        (WaveOp::Min, WaveKind::U) => u32::MAX,
2073        (WaveOp::Max, WaveKind::U) => 0,
2074    }
2075}
2076
2077/// The new word and whether a compare-exchange swapped, of an atomic operation on `old`.
2078fn atomic(fun: AtomOp, signed: bool, old: u32, v: u32, expected: u32) -> (u32, bool) {
2079    let new = match fun {
2080        AtomOp::Add => old.wrapping_add(v),
2081        AtomOp::Sub => old.wrapping_sub(v),
2082        AtomOp::And => old & v,
2083        AtomOp::Or => old | v,
2084        AtomOp::Xor => old ^ v,
2085        AtomOp::Min if signed => (old as i32).min(v as i32) as u32,
2086        AtomOp::Max if signed => (old as i32).max(v as i32) as u32,
2087        AtomOp::Min => old.min(v),
2088        AtomOp::Max => old.max(v),
2089        AtomOp::Exchange => v,
2090        AtomOp::CompareExchange => {
2091            return if old == expected { (v, true) } else { (old, false) };
2092        }
2093    };
2094    (new, true)
2095}
2096
2097/// Unary operation `op` on the word `w`.
2098#[inline(always)]
2099pub(super) fn unary(op: UnOp, w: u32) -> u32 {
2100    match op {
2101        UnOp::NegF => w ^ 0x8000_0000,
2102        UnOp::NegI => (w as i32).wrapping_neg() as u32,
2103        UnOp::Not => (w == 0) as u32,
2104        UnOp::BitNot => !w,
2105    }
2106}
2107
2108/// Conversion `op` of the word `w`.
2109#[inline(always)]
2110pub(super) fn convert(op: CvtOp, w: u32) -> u32 {
2111    match op {
2112        CvtOp::F2I => float_to_int(f(w)),
2113        CvtOp::F2U => float_to_uint(f(w)),
2114        CvtOp::I2F => (w as i32 as f32).to_bits(),
2115        CvtOp::U2F => (w as f32).to_bits(),
2116        CvtOp::B2F => (if w != 0 { 1.0f32 } else { 0.0 }).to_bits(),
2117        CvtOp::B2U => u32::from(w != 0),
2118    }
2119}
2120
2121/// Binary operation `op` on the words `x` and `y`, subnormal floats flushed to zero where `ftz`.
2122#[inline(always)]
2123pub(super) fn bin(op: BinOp, x: u32, y: u32, ftz: bool) -> Result<u32, String> {
2124    use BinOp::*;
2125    let fa = || fl(f(x), ftz);
2126    let fb = || fl(f(y), ftz);
2127    Ok(match op {
2128        AddF => fl(fa() + fb(), ftz).to_bits(),
2129        SubF => fl(fa() - fb(), ftz).to_bits(),
2130        MulF => fl(fa() * fb(), ftz).to_bits(),
2131        DivF => fl(fa() / fb(), ftz).to_bits(),
2132        // C fmod: the remainder of a truncating division, the sign of the dividend
2133        ModF => fl(fa() % fb(), ftz).to_bits(),
2134        AddU => x.wrapping_add(y),
2135        SubU => x.wrapping_sub(y),
2136        MulU => x.wrapping_mul(y),
2137        DivU => x.checked_div(y).ok_or("integer divide by zero in the program")?,
2138        ModU => x.checked_rem(y).ok_or("integer modulo by zero in the program")?,
2139        And => x & y,
2140        Or => x | y,
2141        Xor => x ^ y,
2142        Shl => x.wrapping_shl(y),
2143        ShrU => x.wrapping_shr(y),
2144        ShrI => (x as i32).wrapping_shr(y) as u32,
2145        EqW => (x == y) as u32,
2146        NeW => (x != y) as u32,
2147        EqF => (fa() == fb()) as u32,
2148        NeF => (fa() != fb()) as u32,
2149        LtF => (fa() < fb()) as u32,
2150        LeF => (fa() <= fb()) as u32,
2151        GtF => (fa() > fb()) as u32,
2152        GeF => (fa() >= fb()) as u32,
2153        LtU => (x < y) as u32,
2154        LeU => (x <= y) as u32,
2155        GtU => (x > y) as u32,
2156        GeU => (x >= y) as u32,
2157        LtI => ((x as i32) < (y as i32)) as u32,
2158        LeI => ((x as i32) <= (y as i32)) as u32,
2159        GtI => ((x as i32) > (y as i32)) as u32,
2160        GeI => ((x as i32) >= (y as i32)) as u32,
2161        LAnd => (x != 0 && y != 0) as u32,
2162        LOr => (x != 0 || y != 0) as u32,
2163    })
2164}
2165
2166/// Math operation `mf` on the first `n` lanes of `a`, `b` and `c`, subnormal floats flushed to zero where `z`.
2167pub(super) fn math(mf: MathOp, a: Reg, b: Reg, c: Reg, n: u8, z: bool) -> Reg {
2168    let mut r = [0u32; 4];
2169    if mf == MathOp::Dot {
2170        let mut acc = 0.0f32;
2171        for i in 0..n as usize {
2172            let prod = fl(fl(f(comp(&a, i)), z) * fl(f(comp(&b, i)), z), z);
2173            acc = if i == 0 { prod } else { fl(acc + prod, z) };
2174        }
2175        r[0] = acc.to_bits();
2176        return r;
2177    }
2178    match mf {
2179        MathOp::All => {
2180            r[0] = u32::from(a.iter().take(n as usize).all(|&w| w != 0));
2181            return r;
2182        }
2183        MathOp::Any => {
2184            r[0] = u32::from(a.iter().take(n as usize).any(|&w| w != 0));
2185            return r;
2186        }
2187        MathOp::Normalize => {
2188            let lanes = a.get(..n as usize).unwrap_or_default();
2189            let mut sum = 0.0f32;
2190            for (i, &lane) in lanes.iter().enumerate() {
2191                let sq = fl(fl(f(lane), z) * fl(f(lane), z), z);
2192                sum = if i == 0 { sq } else { fl(sum + sq, z) };
2193            }
2194            let length = fl(crate::narrow(f64::from(sum).sqrt()), z);
2195            for (out, &lane) in r.iter_mut().zip(lanes) {
2196                *out = fl(fl(f(lane), z) / length, z).to_bits();
2197            }
2198            return r;
2199        }
2200        _ => {}
2201    }
2202    if mf == MathOp::PackHalf {
2203        r[0] = half::f16::from_f32(f(a[0])).to_bits() as u32
2204            | ((half::f16::from_f32(f(a[1])).to_bits() as u32) << 16);
2205        return r;
2206    }
2207    if mf == MathOp::UnpackHalf {
2208        let packed = comp(&a, 0);
2209        let half = |bits: u32| {
2210            half::f16::from_bits(u16::try_from(bits & 0xffff).unwrap_or(0))
2211                .to_f32()
2212                .to_bits()
2213        };
2214        r[0] = half(packed);
2215        r[1] = half(packed >> 16);
2216        return r;
2217    }
2218    for i in 0..n as usize {
2219        let (x, y, w) = (comp(&a, i), comp(&b, i), comp(&c, i));
2220        let g = |v: f32| fl(v, z);
2221        let one = |h: &dyn Fn(f32) -> f32| g(h(g(f(x)))).to_bits();
2222        let out = match mf {
2223            MathOp::AbsF => (f(x).abs()).to_bits(),
2224            MathOp::AbsI => (x as i32).wrapping_abs() as u32,
2225            MathOp::MinF => g(g(f(x)).min(g(f(y)))).to_bits(),
2226            MathOp::MaxF => g(g(f(x)).max(g(f(y)))).to_bits(),
2227            MathOp::MinU => x.min(y),
2228            MathOp::MaxU => x.max(y),
2229            MathOp::MinI => (x as i32).min(y as i32) as u32,
2230            MathOp::MaxI => (x as i32).max(y as i32) as u32,
2231            MathOp::ClampF => g(g(f(x)).max(g(f(y))).min(g(f(w)))).to_bits(),
2232            MathOp::SatF => one(&|v| if v.is_nan() { 0.0 } else { v.clamp(0.0, 1.0) }),
2233            MathOp::Floor => one(&|v| v.floor()),
2234            MathOp::Ceil => one(&|v| v.ceil()),
2235            MathOp::Trunc => one(&|v| v.trunc()),
2236            MathOp::Round => one(&|v| v.round_ties_even()),
2237            MathOp::Fract => one(&|v| v - v.floor()),
2238            MathOp::Exp2 => one(&|v| crate::narrow(f64::from(v).exp2())),
2239            MathOp::Log2 => one(&|v| crate::narrow(f64::from(v).log2())),
2240            MathOp::Sin => one(&|v| crate::narrow(f64::from(v).sin())),
2241            MathOp::Cos => one(&|v| crate::narrow(f64::from(v).cos())),
2242            // to the nearest half, ties to even, and back
2243            MathOp::QuantizeF16 => one(&|v| half::f16::from_f32(v).to_f32()),
2244            MathOp::Tan => one(&|v| crate::narrow(f64::from(v).tan())),
2245            MathOp::Asin => one(&|v| crate::narrow(f64::from(v).asin())),
2246            MathOp::Acos => one(&|v| crate::narrow(f64::from(v).acos())),
2247            MathOp::Atan => one(&|v| crate::narrow(f64::from(v).atan())),
2248            MathOp::Sinh => one(&|v| crate::narrow(f64::from(v).sinh())),
2249            MathOp::Cosh => one(&|v| crate::narrow(f64::from(v).cosh())),
2250            MathOp::Tanh => one(&|v| crate::narrow(f64::from(v).tanh())),
2251            // one rounding, as a fused multiply-add
2252            MathOp::Fma => g(g(f(x)).mul_add(g(f(y)), g(f(w)))).to_bits(),
2253            MathOp::Sqrt => one(&|v| crate::narrow(f64::from(v).sqrt())),
2254            MathOp::Rsqrt => one(&|v| crate::narrow(1.0 / f64::from(v).sqrt())),
2255            MathOp::Popcnt => x.count_ones(),
2256            MathOp::RevBits => x.reverse_bits(),
2257            MathOp::Clz => x.leading_zeros(),
2258            MathOp::Ctz => {
2259                if x == 0 {
2260                    u32::MAX
2261                } else {
2262                    x.trailing_zeros()
2263                }
2264            }
2265            // x * (1 - a) + y * a, each product and the sum rounded
2266            MathOp::Mix => g(g(g(f(x)) * g(1.0 - g(f(w)))) + g(g(f(y)) * g(f(w)))).to_bits(),
2267            MathOp::SignF => {
2268                let v = f(x);
2269                (if v > 0.0 {
2270                    1.0f32
2271                } else if v < 0.0 {
2272                    -1.0
2273                } else {
2274                    0.0
2275                })
2276                .to_bits()
2277            }
2278            MathOp::Pow => g(crate::narrow(f64::from(g(f(x))).powf(f64::from(g(f(y)))))).to_bits(),
2279            MathOp::SignI => (x as i32).signum() as u32,
2280            MathOp::IsNan => u32::from(f(x).is_nan()),
2281            MathOp::IsInf => u32::from(f(x).is_infinite()),
2282            MathOp::Dot
2283            | MathOp::PackHalf
2284            | MathOp::UnpackHalf
2285            | MathOp::All
2286            | MathOp::Any
2287            | MathOp::Normalize => 0,
2288        };
2289        if let Some(slot) = r.get_mut(i) {
2290            *slot = out;
2291        }
2292    }
2293    r
2294}
2295
2296// ---------------------------------------------------------------------------------------------------- compute
2297
2298/// What a run of a mesh program does with the memory of a thread group's output variable when the group has finished.
2299pub(super) type GroupHook<'a> = dyn FnMut([u32; 3], &[u32]) -> Result<(), String> + 'a;
2300
2301/// What a compute run did.
2302#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
2303pub struct Dispatched {
2304    /// Thread groups run.
2305    pub groups: u64,
2306    /// Operations and loop turns executed, summed over the groups (a group's lanes run each together).
2307    pub steps: u64,
2308}
2309
2310impl Executable {
2311    /// The thread group of a compute program (`numthreads`), or `None` for any other.
2312    pub fn workgroup_size(&self) -> Option<[u32; 3]> {
2313        self.workgroup
2314    }
2315
2316    /// Run `groups` thread groups of a compute program. Each group's invocations run in lockstep (a barrier is the
2317    /// point where every invocation still in the program has arrived), thread-group memory starts at zero in every
2318    /// group, and the groups run one after another in x, y, z order. `buffers` holds the words of every read-write
2319    /// buffer by resource: writes and atomics change them and later reads see the changes, as they do between
2320    /// groups. A buffer not given is one the backend answers reads of and takes writes to ([`Backend::read_buffer`],
2321    /// [`Backend::store_buffer`]). `wave` is the wave width the wave operations and lane built-ins report; it must
2322    /// divide the thread group's invocations into whole waves only for programs that care.
2323    pub fn dispatch(
2324        &self,
2325        bound: &Bound<'_>,
2326        backend: &mut dyn Backend,
2327        buffers: &mut BTreeMap<ResourceId, Vec<u32>>,
2328        groups: [u32; 3],
2329        wave: u32,
2330    ) -> Result<Dispatched, String> {
2331        self.dispatch_elements(bound, backend, buffers, &mut BTreeMap::new(), groups, wave)
2332    }
2333
2334    /// Run `groups` thread groups of a mesh program, handing the memory of the output variable (its words, as
2335    /// [`Executable::mesh_output`] says) to `after` when each group has finished.
2336    pub(super) fn dispatch_mesh(
2337        &self,
2338        bound: &Bound<'_>,
2339        backend: &mut dyn Backend,
2340        groups: [u32; 3],
2341        wave: u32,
2342        after: &mut GroupHook<'_>,
2343    ) -> Result<Dispatched, String> {
2344        let (global, _) = self.mesh_output.ok_or("not a mesh program")?;
2345        let size = self.workgroup.ok_or("not a mesh program")?;
2346        if wave == 0 {
2347            return Err("a wave of no lanes".into());
2348        }
2349        let ids = (0..groups[2])
2350            .flat_map(move |z| (0..groups[1]).flat_map(move |y| (0..groups[0]).map(move |x| [x, y, z])));
2351        let mut buffers = BTreeMap::new();
2352        let mut elements = BTreeMap::new();
2353        match group_lanes(size)? {
2354            0 => Err("a thread group of no invocations".into()),
2355            1..=64 => self.groups::<64>(
2356                bound,
2357                backend,
2358                &mut buffers,
2359                &mut elements,
2360                ids,
2361                size,
2362                wave,
2363                Some((global, after)),
2364            ),
2365            65..=256 => self.groups::<256>(
2366                bound,
2367                backend,
2368                &mut buffers,
2369                &mut elements,
2370                ids,
2371                size,
2372                wave,
2373                Some((global, after)),
2374            ),
2375            257..=1024 => self.groups::<1024>(
2376                bound,
2377                backend,
2378                &mut buffers,
2379                &mut elements,
2380                ids,
2381                size,
2382                wave,
2383                Some((global, after)),
2384            ),
2385            _ => Err("a thread group of more than 1024 invocations".into()),
2386        }
2387    }
2388
2389    /// [`Executable::dispatch`] for a program that also writes elements of arrays of buffers, `elements` holding
2390    /// their words by resource and element.
2391    pub fn dispatch_elements(
2392        &self,
2393        bound: &Bound<'_>,
2394        backend: &mut dyn Backend,
2395        buffers: &mut BTreeMap<ResourceId, Vec<u32>>,
2396        elements: &mut BTreeMap<(ResourceId, u32), Vec<u32>>,
2397        groups: [u32; 3],
2398        wave: u32,
2399    ) -> Result<Dispatched, String> {
2400        let size = self.workgroup.ok_or("not a compute program")?;
2401        if wave == 0 {
2402            return Err("a wave of no lanes".into());
2403        }
2404        let ids = (0..groups[2])
2405            .flat_map(move |z| (0..groups[1]).flat_map(move |y| (0..groups[0]).map(move |x| [x, y, z])));
2406        match group_lanes(size)? {
2407            0 => Err("a thread group of no invocations".into()),
2408            1..=64 => self.groups::<64>(bound, backend, buffers, elements, ids, size, wave, None),
2409            65..=256 => self.groups::<256>(bound, backend, buffers, elements, ids, size, wave, None),
2410            257..=1024 => self.groups::<1024>(bound, backend, buffers, elements, ids, size, wave, None),
2411            _ => Err("a thread group of more than 1024 invocations".into()),
2412        }
2413    }
2414
2415    #[allow(clippy::too_many_arguments)]
2416    fn groups<const N: usize>(
2417        &self,
2418        bound: &Bound<'_>,
2419        backend: &mut dyn Backend,
2420        buffers: &mut BTreeMap<ResourceId, Vec<u32>>,
2421        elements: &mut BTreeMap<(ResourceId, u32), Vec<u32>>,
2422        ids: impl Iterator<Item = [u32; 3]>,
2423        size: [u32; 3],
2424        wave: u32,
2425        mut after: Option<(usize, &mut GroupHook<'_>)>,
2426    ) -> Result<Dispatched, String> {
2427        let lanes = group_lanes(size)? as usize;
2428        // the buffers move into the run, and back whatever happens
2429        let store: Vec<Option<Vec<u32>>> = self
2430            .globals
2431            .iter()
2432            .map(|g| match g.kind {
2433                GlobalKind::Storage(id) => buffers.remove(&id),
2434                _ => None,
2435            })
2436            .collect();
2437        let elem_store: Vec<BTreeMap<u32, Vec<u32>>> = self
2438            .globals
2439            .iter()
2440            .map(|g| {
2441                let GlobalKind::StorageArray(id, count) = g.kind else {
2442                    return BTreeMap::new();
2443                };
2444                let keys: Vec<_> = elements
2445                    .keys()
2446                    .filter(|(r, k)| *r == id && *k < count)
2447                    .copied()
2448                    .collect();
2449                keys.into_iter()
2450                    .filter_map(|key| elements.remove(&key).map(|words| (key.1, words)))
2451                    .collect()
2452            })
2453            .collect();
2454        let scratch = Scratch::<N>::new(self);
2455        let mut runner = Runner::<N, Wide> {
2456            c: self,
2457            b: bound,
2458            regs: scratch.regs,
2459            mem: scratch.mem,
2460            sampler: backend,
2461            ftz: self.flush_denormals,
2462            grouping: Grouping::Group,
2463            steps: 0,
2464            lane_steps: [0; N],
2465            samples: [0; N],
2466            outs: scratch.outs,
2467            decisions: [Decision::None; N],
2468            killed: Wide::none(),
2469            store,
2470            elem_store,
2471            queries: Vec::new(),
2472            group: Wide::first(lanes),
2473            retired: Wide::none(),
2474            lanes,
2475            wave: wave as usize,
2476            derivs: None,
2477        };
2478        let mut done = Dispatched::default();
2479        let mut result = Ok(());
2480        for id in ids {
2481            match runner.run_group(id, size, wave) {
2482                Ok(steps) => {
2483                    done.groups = done.groups.saturating_add(1);
2484                    done.steps = done.steps.saturating_add(steps);
2485                    if let Some((global, hook)) = after.as_mut() {
2486                        let words = runner
2487                            .store
2488                            .get(*global)
2489                            .and_then(|memory| memory.as_deref())
2490                            .unwrap_or(&[]);
2491                        if let Err(e) = hook(id, words) {
2492                            result = Err(format!("group {id:?}: {e}"));
2493                            break;
2494                        }
2495                    }
2496                }
2497                Err(e) => {
2498                    result = Err(format!("group {id:?}: {e}"));
2499                    break;
2500                }
2501            }
2502        }
2503        for (g, memory) in self.globals.iter().zip(runner.store) {
2504            if let (GlobalKind::Storage(id), Some(words)) = (&g.kind, memory) {
2505                buffers.insert(*id, words);
2506            }
2507        }
2508        for (g, held) in self.globals.iter().zip(runner.elem_store) {
2509            if let GlobalKind::StorageArray(id, _) = g.kind {
2510                for (k, words) in held {
2511                    elements.insert((id, k), words);
2512                }
2513            }
2514        }
2515        result.map(|()| done)
2516    }
2517}
2518
2519impl<'a, const N: usize> Runner<'a, N, Wide> {
2520    /// Run one thread group: `id` its position in the dispatch, `size` its `numthreads`. Returns the steps taken.
2521    fn run_group(&mut self, id: [u32; 3], size: [u32; 3], wave: u32) -> Result<u64, String> {
2522        let c: &'a Executable = self.c;
2523        let lanes = group_lanes(size)? as usize;
2524        self.steps = 0;
2525        self.killed = Wide::none();
2526        self.retired = Wide::none();
2527        self.queries.clear();
2528        for (i, g) in c.globals.iter().enumerate() {
2529            if let GlobalKind::Workgroup(words) = g.kind {
2530                match self.store.get_mut(i) {
2531                    Some(Some(v)) => v.iter_mut().for_each(|w| *w = 0),
2532                    Some(slot) => *slot = Some(vec![0; words as usize]),
2533                    None => return Err(format!("internal error: global {i} has no memory in the run")),
2534                }
2535            }
2536        }
2537        for l in 0..lanes {
2538            let [width, height, _] = size;
2539            let index = u32::try_from(l).map_err(|_| "a thread group of over 4294967295 invocations")?;
2540            let local = [
2541                index.checked_rem(width).unwrap_or(0),
2542                index
2543                    .checked_div(width)
2544                    .unwrap_or(0)
2545                    .checked_rem(height)
2546                    .unwrap_or(0),
2547                index.checked_div(width.saturating_mul(height)).unwrap_or(0),
2548            ];
2549            for (i, (slot, n)) in c.args.iter().enumerate() {
2550                if *n == 0 {
2551                    continue;
2552                }
2553                let value = match c.builtins.get(i).copied().flatten() {
2554                    Some(naga::BuiltIn::WorkGroupId) => [id[0], id[1], id[2], 0],
2555                    Some(naga::BuiltIn::LocalInvocationId) => [local[0], local[1], local[2], 0],
2556                    Some(naga::BuiltIn::GlobalInvocationId) => {
2557                        let global = |axis: usize| {
2558                            comp3(id, axis)
2559                                .saturating_mul(comp3(size, axis))
2560                                .saturating_add(comp3(local, axis))
2561                        };
2562                        [global(0), global(1), global(2), 0]
2563                    }
2564                    Some(naga::BuiltIn::LocalInvocationIndex) => [index, 0, 0, 0],
2565                    Some(naga::BuiltIn::SubgroupInvocationId) => {
2566                        [index.checked_rem(wave).unwrap_or(0), 0, 0, 0]
2567                    }
2568                    Some(naga::BuiltIn::SubgroupSize) => [wave, 0, 0, 0],
2569                    Some(naga::BuiltIn::SubgroupId) => [index.checked_div(wave).unwrap_or(0), 0, 0, 0],
2570                    Some(naga::BuiltIn::NumSubgroups) => {
2571                        [u32::try_from(lanes).unwrap_or(u32::MAX).div_ceil(wave), 0, 0, 0]
2572                    }
2573                    other => return Err(format!("compute entry argument {i} carries {other:?}")),
2574                };
2575                self.wr(slot_at::<N>(*slot), l, value)?;
2576            }
2577            for w in 0..c.local_words {
2578                self.set_local_word(w, l, 0)?;
2579            }
2580        }
2581        self.block(&c.body, self.group)?;
2582        Ok(self.steps)
2583    }
2584}