naga/back/
mod.rs

1/*!
2Backend functions that export shader [`Module`](super::Module)s into binary and text formats.
3*/
4#![cfg_attr(
5    not(any(dot_out, glsl_out, hlsl_out, msl_out, spv_out, wgsl_out)),
6    allow(
7        dead_code,
8        reason = "shared helpers can be dead if none of the enabled backends need it"
9    )
10)]
11
12use alloc::string::String;
13
14#[cfg(dot_out)]
15pub mod dot;
16#[cfg(glsl_out)]
17pub mod glsl;
18#[cfg(hlsl_out)]
19pub mod hlsl;
20#[cfg(msl_out)]
21pub mod msl;
22#[cfg(spv_out)]
23pub mod spv;
24#[cfg(wgsl_out)]
25pub mod wgsl;
26
27#[cfg(any(hlsl_out, msl_out, spv_out, glsl_out))]
28pub mod pipeline_constants;
29
30#[cfg(any(hlsl_out, glsl_out))]
31mod continue_forward;
32
33pub use nt::TaskDispatchLimits;
34
35/// Names of vector components.
36pub const COMPONENTS: &[char] = &['x', 'y', 'z', 'w'];
37/// Indent for backends.
38pub const INDENT: &str = "    ";
39
40/// Expressions that need baking.
41pub type NeedBakeExpressions = crate::FastHashSet<crate::Handle<crate::Expression>>;
42
43/// A type for displaying expression handles as baking identifiers.
44///
45/// Given an [`Expression`] [`Handle`] `h`, `Baked(h)` implements
46/// [`core::fmt::Display`], showing the handle's index prefixed by
47/// `_e`.
48///
49/// [`Expression`]: crate::Expression
50/// [`Handle`]: crate::Handle
51#[cfg_attr(
52    not(any(glsl_out, hlsl_out, msl_out, wgsl_out)),
53    allow(
54        dead_code,
55        reason = "shared helpers can be dead if none of the enabled backends need it"
56    )
57)]
58struct Baked(crate::Handle<crate::Expression>);
59
60impl core::fmt::Display for Baked {
61    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
62        self.0.write_prefixed(f, "_e")
63    }
64}
65
66bitflags::bitflags! {
67    /// How far through a ray query are we
68    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
69    #[cfg_attr(
70        not(any(hlsl_out, spv_out)),
71        allow(
72            dead_code,
73            reason = "shared helpers can be dead if none of the enabled backends need it"
74        )
75    )]
76    pub(super) struct RayQueryPoint: u32 {
77        /// Ray query has been successfully initialized.
78        const INITIALIZED = 1 << 0;
79        /// Proceed has been called on ray query.
80        const PROCEED = 1 << 1;
81        /// Proceed has returned false (have finished traversal).
82        const FINISHED_TRAVERSAL = 1 << 2;
83    }
84}
85
86/// Specifies the values of pipeline-overridable constants in the shader module.
87///
88/// If an `@id` attribute was specified on the declaration,
89/// the key must be the pipeline constant ID as a decimal ASCII number; if not,
90/// the key must be the constant's identifier name.
91///
92/// The value may represent any of WGSL's concrete scalar types.
93pub type PipelineConstants = hashbrown::HashMap<String, f64>;
94
95/// Indentation level.
96#[derive(Clone, Copy, Debug)]
97pub struct Level(pub usize);
98
99impl Level {
100    pub const fn next(&self) -> Self {
101        Level(self.0 + 1)
102    }
103}
104
105impl core::fmt::Display for Level {
106    fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> {
107        (0..self.0).try_for_each(|_| formatter.write_str(INDENT))
108    }
109}
110
111/// Locate the entry point(s) to write.
112///
113/// If `entry_point` is given, and the specified entry point exists, returns a
114/// length-1 `Range` containing the index of that entry point.  If no
115/// `entry_point` is given, returns the complete range of entry point indices.
116/// If `entry_point` is given but does not exist, returns an error.
117#[cfg(any(hlsl_out, msl_out))]
118fn get_entry_points(
119    module: &crate::ir::Module,
120    entry_point: Option<&(crate::ir::ShaderStage, String)>,
121) -> Result<core::ops::Range<usize>, (crate::ir::ShaderStage, String)> {
122    use alloc::borrow::ToOwned;
123
124    if let Some(&(stage, ref name)) = entry_point {
125        let Some(ep_index) = module
126            .entry_points
127            .iter()
128            .position(|ep| ep.stage == stage && ep.name == *name)
129        else {
130            return Err((stage, name.to_owned()));
131        };
132        Ok(ep_index..ep_index + 1)
133    } else {
134        Ok(0..module.entry_points.len())
135    }
136}
137
138/// Whether we're generating an entry point or a regular function.
139///
140/// Backend languages often require different code for a [`Function`]
141/// depending on whether it represents an [`EntryPoint`] or not.
142/// Backends can pass common code one of these values to select the
143/// right behavior.
144///
145/// These values also carry enough information to find the `Function`
146/// in the [`Module`]: the `Handle` for a regular function, or the
147/// index into [`Module::entry_points`] for an entry point.
148///
149/// [`Function`]: crate::Function
150/// [`EntryPoint`]: crate::EntryPoint
151/// [`Module`]: crate::Module
152/// [`Module::entry_points`]: crate::Module::entry_points
153#[derive(Clone, Copy, Debug)]
154pub enum FunctionType {
155    /// A regular function.
156    Function(crate::Handle<crate::Function>),
157    /// An [`EntryPoint`], and its index in [`Module::entry_points`].
158    ///
159    /// [`EntryPoint`]: crate::EntryPoint
160    /// [`Module::entry_points`]: crate::Module::entry_points
161    EntryPoint(crate::proc::EntryPointIndex),
162}
163
164impl FunctionType {
165    /// Returns true if the function is an entry point for a compute-like shader.
166    pub fn is_compute_like_entry_point(&self, module: &crate::Module) -> bool {
167        match *self {
168            FunctionType::EntryPoint(index) => {
169                module.entry_points[index as usize].stage.compute_like()
170            }
171            FunctionType::Function(_) => false,
172        }
173    }
174}
175
176/// Helper structure that stores data needed when writing the function
177#[derive(Debug)]
178pub struct FunctionCtx<'a> {
179    /// The current function being written
180    pub ty: FunctionType,
181    /// Analysis about the function
182    pub info: &'a crate::valid::FunctionInfo,
183    /// The expression arena of the current function being written
184    pub expressions: &'a crate::Arena<crate::Expression>,
185    /// Map of expressions that have associated variable names
186    pub named_expressions: &'a crate::NamedExpressions,
187}
188
189impl FunctionCtx<'_> {
190    /// Helper method that resolves a type of a given expression.
191    pub fn resolve_type<'a>(
192        &'a self,
193        handle: crate::Handle<crate::Expression>,
194        types: &'a crate::UniqueArena<crate::Type>,
195    ) -> &'a crate::TypeInner {
196        self.info[handle].ty.inner_with(types)
197    }
198
199    /// Helper method that generates a [`NameKey`](crate::proc::NameKey) for a local in the current function
200    pub const fn name_key(
201        &self,
202        local: crate::Handle<crate::LocalVariable>,
203    ) -> crate::proc::NameKey {
204        match self.ty {
205            FunctionType::Function(handle) => crate::proc::NameKey::FunctionLocal(handle, local),
206            FunctionType::EntryPoint(idx) => crate::proc::NameKey::EntryPointLocal(idx, local),
207        }
208    }
209
210    /// Helper method that generates a [`NameKey`](crate::proc::NameKey) for a function argument.
211    ///
212    /// # Panics
213    /// - If the function arguments are less or equal to `arg`
214    pub const fn argument_key(&self, arg: u32) -> crate::proc::NameKey {
215        match self.ty {
216            FunctionType::Function(handle) => crate::proc::NameKey::FunctionArgument(handle, arg),
217            FunctionType::EntryPoint(ep_index) => {
218                crate::proc::NameKey::EntryPointArgument(ep_index, arg)
219            }
220        }
221    }
222
223    /// Helper method that generates a [`NameKey`](crate::proc::NameKey) for an external texture
224    /// function argument.
225    ///
226    /// # Panics
227    /// - If the function arguments are less or equal to `arg`
228    /// - If `self.ty` is not `FunctionType::Function`.
229    pub const fn external_texture_argument_key(
230        &self,
231        arg: u32,
232        external_texture_key: crate::proc::ExternalTextureNameKey,
233    ) -> crate::proc::NameKey {
234        match self.ty {
235            FunctionType::Function(handle) => {
236                crate::proc::NameKey::ExternalTextureFunctionArgument(
237                    handle,
238                    arg,
239                    external_texture_key,
240                )
241            }
242            // This is a const function, which _sometimes_ gets called,
243            // so this lint is _sometimes_ triggered, depending on feature set.
244            #[expect(clippy::allow_attributes)]
245            #[allow(clippy::panic)]
246            FunctionType::EntryPoint(_) => {
247                panic!("External textures cannot be used as arguments to entry points")
248            }
249        }
250    }
251
252    /// Returns true if the given expression points to a fixed-function pipeline input.
253    pub fn is_fixed_function_input(
254        &self,
255        mut expression: crate::Handle<crate::Expression>,
256        module: &crate::Module,
257    ) -> Option<crate::BuiltIn> {
258        let ep_function = match self.ty {
259            FunctionType::Function(_) => return None,
260            FunctionType::EntryPoint(ep_index) => &module.entry_points[ep_index as usize].function,
261        };
262        let mut built_in = None;
263        loop {
264            match self.expressions[expression] {
265                crate::Expression::FunctionArgument(arg_index) => {
266                    return match ep_function.arguments[arg_index as usize].binding {
267                        Some(crate::Binding::BuiltIn(bi)) => Some(bi),
268                        _ => built_in,
269                    };
270                }
271                crate::Expression::AccessIndex { base, index } => {
272                    match *self.resolve_type(base, &module.types) {
273                        crate::TypeInner::Struct { ref members, .. } => {
274                            if let Some(crate::Binding::BuiltIn(bi)) =
275                                members[index as usize].binding
276                            {
277                                built_in = Some(bi);
278                            }
279                        }
280                        _ => return None,
281                    }
282                    expression = base;
283                }
284                _ => return None,
285            }
286        }
287    }
288}
289
290impl crate::Expression {
291    /// Returns the ref count, upon reaching which this expression
292    /// should be considered for baking.
293    ///
294    /// Note: we have to cache any expressions that depend on the control flow,
295    /// or otherwise they may be moved into a non-uniform control flow, accidentally.
296    /// See the [module-level documentation][emit] for details.
297    ///
298    /// [emit]: index.html#expression-evaluation-time
299    pub const fn bake_ref_count(&self) -> usize {
300        match *self {
301            // accesses are never cached, only loads are
302            crate::Expression::Access { .. } | crate::Expression::AccessIndex { .. } => usize::MAX,
303            // sampling may use the control flow, and image ops look better by themselves
304            crate::Expression::ImageSample { .. } | crate::Expression::ImageLoad { .. } => 1,
305            // derivatives use the control flow
306            crate::Expression::Derivative { .. } => 1,
307            // TODO: We need a better fix for named `Load` expressions
308            // More info - https://github.com/gfx-rs/naga/pull/914
309            // And https://github.com/gfx-rs/naga/issues/910
310            crate::Expression::Load { .. } => 1,
311            // cache expressions that are referenced multiple times
312            _ => 2,
313        }
314    }
315}
316
317/// Helper function that returns the string corresponding to the [`BinaryOperator`](crate::BinaryOperator)
318pub const fn binary_operation_str(op: crate::BinaryOperator) -> &'static str {
319    use crate::BinaryOperator as Bo;
320    match op {
321        Bo::Add => "+",
322        Bo::Subtract => "-",
323        Bo::Multiply => "*",
324        Bo::Divide => "/",
325        Bo::Modulo => "%",
326        Bo::Equal => "==",
327        Bo::NotEqual => "!=",
328        Bo::Less => "<",
329        Bo::LessEqual => "<=",
330        Bo::Greater => ">",
331        Bo::GreaterEqual => ">=",
332        Bo::And => "&",
333        Bo::ExclusiveOr => "^",
334        Bo::InclusiveOr => "|",
335        Bo::LogicalAnd => "&&",
336        Bo::LogicalOr => "||",
337        Bo::ShiftLeft => "<<",
338        Bo::ShiftRight => ">>",
339    }
340}
341
342impl crate::TypeInner {
343    /// Returns true if a variable of this type is a handle.
344    pub const fn is_handle(&self) -> bool {
345        match *self {
346            Self::Image { .. } | Self::Sampler { .. } | Self::AccelerationStructure { .. } => true,
347            _ => false,
348        }
349    }
350}
351
352impl crate::Statement {
353    /// Returns true if the statement directly terminates the current block.
354    ///
355    /// Used to decide whether case blocks require a explicit `break`.
356    pub const fn is_terminator(&self) -> bool {
357        match *self {
358            crate::Statement::Break
359            | crate::Statement::Continue
360            | crate::Statement::Return { .. }
361            | crate::Statement::Kill => true,
362            _ => false,
363        }
364    }
365}
366
367bitflags::bitflags! {
368    /// Ray flags, for a [`RayDesc`]'s `flags` field.
369    ///
370    /// Note that these exactly correspond to the SPIR-V "Ray Flags" mask, and
371    /// the SPIR-V backend passes them directly through to the
372    /// [`OpRayQueryInitializeKHR`][op] instruction. (We have to choose something, so
373    /// we might as well make one back end's life easier.)
374    ///
375    /// [`RayDesc`]: crate::Module::generate_ray_desc_type
376    /// [op]: https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpRayQueryInitializeKHR
377    #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
378    pub struct RayFlag: u32 {
379        const OPAQUE = 0x01;
380        const NO_OPAQUE = 0x02;
381        const TERMINATE_ON_FIRST_HIT = 0x04;
382        const SKIP_CLOSEST_HIT_SHADER = 0x08;
383        const CULL_BACK_FACING = 0x10;
384        const CULL_FRONT_FACING = 0x20;
385        const CULL_OPAQUE = 0x40;
386        const CULL_NO_OPAQUE = 0x80;
387        const SKIP_TRIANGLES = 0x100;
388        const SKIP_AABBS = 0x200;
389    }
390}
391
392/// The intersection test to use for ray queries.
393#[derive(Debug)]
394#[repr(u32)]
395pub enum RayIntersectionType {
396    Triangle = 1,
397    BoundingBox = 4,
398}