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

naga/back/
mod.rs
