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