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
14pub const STEP_LIMIT: u64 = 1 << 24;
16
17#[derive(Debug)]
19enum Words<'a> {
20 Owned(Vec<u32>),
21 Rows(&'a [[u32; 4]]),
23 Slice(&'a [u32]),
24 Elements(BTreeMap<u32, &'a [u32]>),
26 ConstantElements(BTreeMap<u32, &'a [[u32; 4]]>),
28 Backend,
30}
31
32#[derive(Debug)]
34pub struct Bound<'a> {
35 data: Vec<Words<'a>>,
36 scenes: Scenes,
37}
38
39impl Bound<'_> {
40 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#[derive(Clone, Debug, Default)]
49pub struct Resources {
50 pub constants: BTreeMap<ResourceId, Vec<[u32; 4]>>,
52 pub constant_elements: BTreeMap<(ResourceId, u32), Vec<[u32; 4]>>,
54 pub buffers: BTreeMap<ResourceId, Vec<u32>>,
56 pub elements: BTreeMap<(ResourceId, u32), Vec<u32>>,
58 pub scenes: BTreeMap<ResourceId, super::rayquery::Scene>,
60}
61
62impl Executable {
63 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 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 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 GlobalKind::Workgroup(_) => Words::Owned(Vec::new()),
147 });
148 }
149 Ok(Bound {
150 data,
151 scenes: BTreeMap::new(),
152 })
153 }
154
155 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#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
184pub enum Decision {
185 #[default]
187 None,
188 Ignore,
190 AcceptAndEnd,
192}
193
194#[derive(Clone, Debug)]
196pub struct Output {
197 pub pointees: Vec<Vec<u32>>,
200 pub decision: Decision,
202 pub values: Vec<[u32; 4]>,
204 pub discarded: bool,
206 pub steps: u64,
208 pub samples: u64,
210}
211
212#[derive(Debug)]
215pub struct Scratch<const N: usize> {
216 regs: Vec<Reg>,
217 mem: Vec<u32>,
218 outs: [Vec<Reg>; N],
220}
221impl<const N: usize> Scratch<N> {
222 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
240pub 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 steps: u64,
255 lane_steps: [u64; N],
257 samples: [u64; N],
258 outs: [Vec<Reg>; N],
259 decisions: [Decision; N],
261 killed: M,
262 store: Vec<Option<Vec<u32>>>,
265 queries: Vec<Vec<Option<Query>>>,
267 elem_store: Vec<BTreeMap<u32, Vec<u32>>>,
269 group: M,
271 retired: M,
272 lanes: usize,
274 wave: usize,
275 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 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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
296enum Grouping {
297 Independent,
299 Quad,
301 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
321macro_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
334trait 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
346fn 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
353fn 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
366fn 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
378fn 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
385fn comp3(v: [u32; 3], axis: usize) -> u32 {
387 v.get(axis).copied().unwrap_or(0)
388}
389
390fn 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#[inline(always)]
400fn comp(r: &Reg, i: usize) -> u32 {
401 r.get(i).copied().unwrap_or(0)
402}
403
404#[inline(always)]
406#[allow(clippy::cast_possible_truncation)] fn float_to_int(x: f32) -> u32 {
408 (x as i32) as u32
409}
410
411#[inline(always)]
413#[allow(clippy::cast_possible_truncation)] fn float_to_uint(x: f32) -> u32 {
415 x as u32
416}
417
418#[inline(always)]
420fn slot_at<const N: usize>(slot: u32) -> usize {
421 (slot as usize).saturating_mul(N)
422}
423
424fn 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 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 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 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 #[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 #[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 #[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 #[inline(always)]
546 fn local(word: usize, lane: usize) -> usize {
547 word.saturating_mul(N).saturating_add(lane)
548 }
549 #[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 #[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 #[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 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 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 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 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 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 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 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 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 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 let (sa, sb) = (usize::from(!*a1), usize::from(!*b1));
869 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 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 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 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 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 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 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 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 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 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 pub fn run(&mut self, args: [&[[u32; 4]]; N]) -> Result<[Output; N], String> {
1941 self.run_first(args, N)
1942 }
1943
1944 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 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 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
2037fn 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}
2058fn 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
2077fn 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#[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#[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#[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 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
2166pub(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 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 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 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
2296pub(super) type GroupHook<'a> = dyn FnMut([u32; 3], &[u32]) -> Result<(), String> + 'a;
2300
2301#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
2303pub struct Dispatched {
2304 pub groups: u64,
2306 pub steps: u64,
2308}
2309
2310impl Executable {
2311 pub fn workgroup_size(&self) -> Option<[u32; 3]> {
2313 self.workgroup
2314 }
2315
2316 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 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 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 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 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}