naga/front/
atomic_upgrade.rs

1//! Upgrade the types of scalars observed to be accessed as atomics to [`Atomic`] types.
2//!
3//! In SPIR-V, atomic operations can be applied to any scalar value, but in Naga
4//! IR atomic operations can only be applied to values of type [`Atomic`]. Naga
5//! IR's restriction matches Metal Shading Language and WGSL, so we don't want
6//! to relax that. Instead, when the SPIR-V front end observes a value being
7//! accessed using atomic instructions, it promotes the value's type from
8//! [`Scalar`] to [`Atomic`]. This module implements `Module::upgrade_atomics`,
9//! the function that makes that change.
10//!
11//! Atomics can only appear in global variables in the [`Storage`] and
12//! [`Workgroup`] address spaces. These variables can either have `Atomic` types
13//! themselves, or be [`Array`]s of such, or be [`Struct`]s containing such.
14//! So we only need to change the types of globals and struct fields.
15//!
16//! Naga IR [`Load`] expressions and [`Store`] statements can operate directly
17//! on [`Atomic`] values, retrieving and depositing ordinary [`Scalar`] values,
18//! so changing the types doesn't have much effect on the code that operates on
19//! those values.
20//!
21//! Future work:
22//!
23//! - The GLSL front end could use this transformation as well.
24//!
25//! [`Atomic`]: TypeInner::Atomic
26//! [`Scalar`]: TypeInner::Scalar
27//! [`Storage`]: crate::AddressSpace::Storage
28//! [`WorkGroup`]: crate::AddressSpace::WorkGroup
29//! [`Array`]: TypeInner::Array
30//! [`Struct`]: TypeInner::Struct
31//! [`Load`]: crate::Expression::Load
32//! [`Store`]: crate::Statement::Store
33
34use alloc::{format, sync::Arc};
35use core::sync::atomic::AtomicUsize;
36
37use crate::{GlobalVariable, Handle, Module, Type, TypeInner};
38
39#[derive(Clone, Debug, thiserror::Error)]
40pub enum Error {
41    #[error("encountered an unsupported expression")]
42    Unsupported,
43    #[error("unexpected end of struct field access indices")]
44    UnexpectedEndOfIndices,
45    #[error("encountered unsupported global initializer in an atomic variable")]
46    GlobalInitUnsupported,
47    #[error("expected to find a global variable")]
48    GlobalVariableMissing,
49    #[error("atomic compare exchange requires a scalar base type")]
50    CompareExchangeNonScalarBaseType,
51}
52
53#[derive(Clone, Default)]
54struct Padding(Arc<AtomicUsize>);
55
56impl core::fmt::Display for Padding {
57    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
58        for _ in 0..self.0.load(core::sync::atomic::Ordering::Relaxed) {
59            f.write_str("  ")?;
60        }
61        Ok(())
62    }
63}
64
65impl Drop for Padding {
66    fn drop(&mut self) {
67        let _ = self.0.fetch_sub(1, core::sync::atomic::Ordering::Relaxed);
68    }
69}
70
71impl Padding {
72    fn trace(&self, msg: impl core::fmt::Display, t: impl core::fmt::Debug) {
73        format!("{msg} {t:#?}")
74            .split('\n')
75            .for_each(|ln| log::trace!("{self}{ln}"));
76    }
77
78    fn debug(&self, msg: impl core::fmt::Display, t: impl core::fmt::Debug) {
79        format!("{msg} {t:#?}")
80            .split('\n')
81            .for_each(|ln| log::debug!("{self}{ln}"));
82    }
83
84    fn inc_padding(&self) -> Padding {
85        let _ = self.0.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
86        self.clone()
87    }
88}
89
90#[derive(Debug, Default)]
91pub struct Upgrades {
92    /// Global variables that we've accessed using atomic operations.
93    ///
94    /// This includes globals with composite types (arrays, structs) where we've
95    /// only accessed some components (elements, fields) atomically.
96    globals: crate::arena::HandleSet<GlobalVariable>,
97
98    /// Struct fields that we've accessed using atomic operations.
99    ///
100    /// Each key refers to some [`Struct`] type, and each value is a set of
101    /// the indices of the fields in that struct that have been accessed
102    /// atomically.
103    ///
104    /// This includes fields with composite types (arrays, structs)
105    /// of which we've only accessed some components (elements, fields)
106    /// atomically.
107    ///
108    /// [`Struct`]: crate::TypeInner::Struct
109    fields: crate::FastHashMap<Handle<Type>, bit_set::BitSet>,
110}
111
112impl Upgrades {
113    pub fn insert_global(&mut self, global: Handle<GlobalVariable>) {
114        self.globals.insert(global);
115    }
116
117    pub fn insert_field(&mut self, struct_type: Handle<Type>, field: usize) {
118        self.fields.entry(struct_type).or_default().insert(field);
119    }
120
121    pub fn is_empty(&self) -> bool {
122        self.globals.is_empty()
123    }
124}
125
126struct UpgradeState<'a> {
127    padding: Padding,
128    module: &'a mut Module,
129
130    /// A map from old types to their upgraded versions.
131    ///
132    /// This ensures we never try to rebuild a type more than once.
133    upgraded_types: crate::FastHashMap<Handle<Type>, Handle<Type>>,
134}
135
136impl UpgradeState<'_> {
137    fn inc_padding(&self) -> Padding {
138        self.padding.inc_padding()
139    }
140
141    /// Get a type equivalent to `ty`, but with [`Scalar`] leaves upgraded to [`Atomic`] scalars.
142    ///
143    /// If such a type already exists in `self.module.types`, return its handle.
144    /// Otherwise, construct a new one and return that handle.
145    ///
146    /// If `ty` is a [`Pointer`], [`Array`], [`BindingArray`], recurse into the
147    /// type and upgrade its leaf types.
148    ///
149    /// If `ty` is a [`Struct`], recurse into it and upgrade only those fields
150    /// whose indices appear in `field_indices`.
151    ///
152    /// The existing type is not affected.
153    ///
154    /// [`Scalar`]: crate::TypeInner::Scalar
155    /// [`Atomic`]: crate::TypeInner::Atomic
156    /// [`Pointer`]: crate::TypeInner::Pointer
157    /// [`Array`]: crate::TypeInner::Array
158    /// [`Struct`]: crate::TypeInner::Struct
159    /// [`BindingArray`]: crate::TypeInner::BindingArray
160    fn upgrade_type(
161        &mut self,
162        ty: Handle<Type>,
163        upgrades: &Upgrades,
164    ) -> Result<Handle<Type>, Error> {
165        let padding = self.inc_padding();
166        padding.trace("visiting type: ", ty);
167
168        // If we've already upgraded this type, return the handle we produced at
169        // the time.
170        if let Some(&new) = self.upgraded_types.get(&ty) {
171            return Ok(new);
172        }
173
174        let inner = match self.module.types[ty].inner {
175            TypeInner::Scalar(scalar) => {
176                log::trace!("{padding}hit the scalar leaf, replacing with an atomic");
177                TypeInner::Atomic(scalar)
178            }
179            TypeInner::Pointer { base, space } => TypeInner::Pointer {
180                base: self.upgrade_type(base, upgrades)?,
181                space,
182            },
183            TypeInner::Array { base, size, stride } => TypeInner::Array {
184                base: self.upgrade_type(base, upgrades)?,
185                size,
186                stride,
187            },
188            TypeInner::Struct { ref members, span } => {
189                // If no field or subfield of this struct was ever accessed
190                // atomically, no change is needed. We should never have arrived here.
191                let Some(fields) = upgrades.fields.get(&ty) else {
192                    unreachable!("global or field incorrectly flagged as atomically accessed");
193                };
194
195                let mut new_members = members.clone();
196                for field in fields {
197                    new_members[field].ty = self.upgrade_type(new_members[field].ty, upgrades)?;
198                }
199
200                TypeInner::Struct {
201                    members: new_members,
202                    span,
203                }
204            }
205            TypeInner::BindingArray { base, size } => TypeInner::BindingArray {
206                base: self.upgrade_type(base, upgrades)?,
207                size,
208            },
209            _ => return Ok(ty),
210        };
211
212        // At this point, we have a `TypeInner` that is the upgraded version of
213        // `ty`. Find a suitable `Type` for this, creating a new one if
214        // necessary, and return its handle.
215        let r#type = &self.module.types[ty];
216        let span = self.module.types.get_span(ty);
217        let new_type = Type {
218            name: r#type.name.clone(),
219            inner,
220        };
221        padding.debug("ty: ", ty);
222        padding.debug("from: ", r#type);
223        padding.debug("to:   ", &new_type);
224        let new_handle = self.module.types.insert(new_type, span);
225        self.upgraded_types.insert(ty, new_handle);
226        Ok(new_handle)
227    }
228
229    fn upgrade_all(&mut self, upgrades: &Upgrades) -> Result<(), Error> {
230        for handle in upgrades.globals.iter() {
231            let padding = self.inc_padding();
232
233            let global = &self.module.global_variables[handle];
234            padding.trace("visiting global variable: ", handle);
235            padding.trace("var: ", global);
236
237            if global.init.is_some() {
238                return Err(Error::GlobalInitUnsupported);
239            }
240
241            let var_ty = global.ty;
242            let new_ty = self.upgrade_type(var_ty, upgrades)?;
243            if new_ty != var_ty {
244                padding.debug("upgrading global variable: ", handle);
245                padding.debug("from ty: ", var_ty);
246                padding.debug("to ty:   ", new_ty);
247                self.module.global_variables[handle].ty = new_ty;
248            }
249        }
250
251        Ok(())
252    }
253}
254
255impl Module {
256    /// Upgrade `global_var_handles` to have [`Atomic`] leaf types.
257    ///
258    /// [`Atomic`]: TypeInner::Atomic
259    pub(crate) fn upgrade_atomics(&mut self, upgrades: &Upgrades) -> Result<(), Error> {
260        let mut state = UpgradeState {
261            padding: Default::default(),
262            module: self,
263            upgraded_types: crate::FastHashMap::with_capacity_and_hasher(
264                upgrades.fields.len(),
265                Default::default(),
266            ),
267        };
268
269        state.upgrade_all(upgrades)?;
270
271        Ok(())
272    }
273}
naga/front/atomic_upgrade.rs

naga/front/
atomic_upgrade.rs
