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