naga/back/spv/
mod.rs

1/*!
2Backend for [SPIR-V][spv] (Standard Portable Intermediate Representation).
3
4[spv]: https://www.khronos.org/registry/SPIR-V/
5*/
6
7mod block;
8mod f16_polyfill;
9mod helpers;
10mod image;
11mod index;
12mod instructions;
13mod layout;
14mod ray;
15mod recyclable;
16mod selection;
17mod subgroup;
18mod writer;
19
20pub use spirv::{Capability, SourceLanguage};
21
22use alloc::{string::String, vec::Vec};
23use core::ops;
24
25use spirv::Word;
26use thiserror::Error;
27
28use crate::arena::{Handle, HandleVec};
29use crate::proc::{BoundsCheckPolicies, TypeResolution};
30
31#[derive(Clone)]
32struct PhysicalLayout {
33    magic_number: Word,
34    version: Word,
35    generator: Word,
36    bound: Word,
37    instruction_schema: Word,
38}
39
40#[derive(Default)]
41struct LogicalLayout {
42    capabilities: Vec<Word>,
43    extensions: Vec<Word>,
44    ext_inst_imports: Vec<Word>,
45    memory_model: Vec<Word>,
46    entry_points: Vec<Word>,
47    execution_modes: Vec<Word>,
48    debugs: Vec<Word>,
49    annotations: Vec<Word>,
50    declarations: Vec<Word>,
51    function_declarations: Vec<Word>,
52    function_definitions: Vec<Word>,
53}
54
55struct Instruction {
56    op: spirv::Op,
57    wc: u32,
58    type_id: Option<Word>,
59    result_id: Option<Word>,
60    operands: Vec<Word>,
61}
62
63const BITS_PER_BYTE: crate::Bytes = 8;
64
65#[derive(Clone, Debug, Error)]
66pub enum Error {
67    #[error("The requested entry point couldn't be found")]
68    EntryPointNotFound,
69    #[error("target SPIRV-{0}.{1} is not supported")]
70    UnsupportedVersion(u8, u8),
71    #[error("using {0} requires at least one of the capabilities {1:?}, but none are available")]
72    MissingCapabilities(&'static str, Vec<Capability>),
73    #[error("unimplemented {0}")]
74    FeatureNotImplemented(&'static str),
75    #[error("module is not validated properly: {0}")]
76    Validation(&'static str),
77    #[error("overrides should not be present at this stage")]
78    Override,
79    #[error(transparent)]
80    ResolveArraySizeError(#[from] crate::proc::ResolveArraySizeError),
81    #[error("mapping of {0:?} is missing")]
82    MissingBinding(crate::ResourceBinding),
83}
84
85#[derive(Default)]
86struct IdGenerator(Word);
87
88impl IdGenerator {
89    fn next(&mut self) -> Word {
90        self.0 += 1;
91        self.0
92    }
93}
94
95#[derive(Debug, Clone)]
96pub struct DebugInfo<'a> {
97    pub source_code: &'a str,
98    pub file_name: &'a str,
99    pub language: SourceLanguage,
100}
101
102/// A SPIR-V block to which we are still adding instructions.
103///
104/// A `Block` represents a SPIR-V block that does not yet have a termination
105/// instruction like `OpBranch` or `OpReturn`.
106///
107/// The `OpLabel` that starts the block is implicit. It will be emitted based on
108/// `label_id` when we write the block to a `LogicalLayout`.
109///
110/// To terminate a `Block`, pass the block and the termination instruction to
111/// `Function::consume`. This takes ownership of the `Block` and transforms it
112/// into a `TerminatedBlock`.
113struct Block {
114    label_id: Word,
115    body: Vec<Instruction>,
116}
117
118/// A SPIR-V block that ends with a termination instruction.
119struct TerminatedBlock {
120    label_id: Word,
121    body: Vec<Instruction>,
122}
123
124impl Block {
125    const fn new(label_id: Word) -> Self {
126        Block {
127            label_id,
128            body: Vec::new(),
129        }
130    }
131}
132
133struct LocalVariable {
134    id: Word,
135    instruction: Instruction,
136}
137
138struct ResultMember {
139    id: Word,
140    type_id: Word,
141    built_in: Option<crate::BuiltIn>,
142}
143
144struct EntryPointContext {
145    argument_ids: Vec<Word>,
146    results: Vec<ResultMember>,
147}
148
149#[derive(Default)]
150struct Function {
151    signature: Option<Instruction>,
152    parameters: Vec<FunctionArgument>,
153    variables: crate::FastHashMap<Handle<crate::LocalVariable>, LocalVariable>,
154    /// Map from a local variable that is a ray query to its u32 tracker.
155    ray_query_initialization_tracker_variables:
156        crate::FastHashMap<Handle<crate::LocalVariable>, LocalVariable>,
157    /// Map from a local variable that is a ray query to its tracker for the t max.
158    ray_query_t_max_tracker_variables:
159        crate::FastHashMap<Handle<crate::LocalVariable>, LocalVariable>,
160    /// List of local variables used as a counters to ensure that all loops are bounded.
161    force_loop_bounding_vars: Vec<LocalVariable>,
162
163    /// A map from a Naga expression to the temporary SPIR-V variable we have
164    /// spilled its value to, if any.
165    ///
166    /// Naga IR lets us apply [`Access`] expressions to expressions whose value
167    /// is an array or matrix---not a pointer to such---but SPIR-V doesn't have
168    /// instructions that can do the same. So when we encounter such code, we
169    /// spill the expression's value to a generated temporary variable. That, we
170    /// can obtain a pointer to, and then use an `OpAccessChain` instruction to
171    /// do whatever series of [`Access`] and [`AccessIndex`] operations we need
172    /// (with bounds checks). Finally, we generate an `OpLoad` to get the final
173    /// value.
174    ///
175    /// [`Access`]: crate::Expression::Access
176    /// [`AccessIndex`]: crate::Expression::AccessIndex
177    spilled_composites: crate::FastIndexMap<Handle<crate::Expression>, LocalVariable>,
178
179    /// A set of expressions that are either in [`spilled_composites`] or refer
180    /// to some component/element of such.
181    ///
182    /// [`spilled_composites`]: Function::spilled_composites
183    spilled_accesses: crate::arena::HandleSet<crate::Expression>,
184
185    /// A map taking each expression to the number of [`Access`] and
186    /// [`AccessIndex`] expressions that uses it as a base value. If an
187    /// expression has no entry, its count is zero: it is never used as a
188    /// [`Access`] or [`AccessIndex`] base.
189    ///
190    /// We use this, together with [`ExpressionInfo::ref_count`], to recognize
191    /// the tips of chains of [`Access`] and [`AccessIndex`] expressions that
192    /// access spilled values --- expressions in [`spilled_composites`]. We
193    /// defer generating code for the chain until we reach its tip, so we can
194    /// handle it with a single instruction.
195    ///
196    /// [`Access`]: crate::Expression::Access
197    /// [`AccessIndex`]: crate::Expression::AccessIndex
198    /// [`ExpressionInfo::ref_count`]: crate::valid::ExpressionInfo
199    /// [`spilled_composites`]: Function::spilled_composites
200    access_uses: crate::FastHashMap<Handle<crate::Expression>, usize>,
201
202    blocks: Vec<TerminatedBlock>,
203    entry_point_context: Option<EntryPointContext>,
204}
205
206impl Function {
207    fn consume(&mut self, mut block: Block, termination: Instruction) {
208        block.body.push(termination);
209        self.blocks.push(TerminatedBlock {
210            label_id: block.label_id,
211            body: block.body,
212        })
213    }
214
215    fn parameter_id(&self, index: u32) -> Word {
216        match self.entry_point_context {
217            Some(ref context) => context.argument_ids[index as usize],
218            None => self.parameters[index as usize]
219                .instruction
220                .result_id
221                .unwrap(),
222        }
223    }
224}
225
226/// Characteristics of a SPIR-V `OpTypeImage` type.
227///
228/// SPIR-V requires non-composite types to be unique, including images. Since we
229/// use `LocalType` for this deduplication, it's essential that `LocalImageType`
230/// be equal whenever the corresponding `OpTypeImage`s would be. To reduce the
231/// likelihood of mistakes, we use fields that correspond exactly to the
232/// operands of an `OpTypeImage` instruction, using the actual SPIR-V types
233/// where practical.
234#[derive(Debug, PartialEq, Hash, Eq, Copy, Clone)]
235struct LocalImageType {
236    sampled_type: crate::Scalar,
237    dim: spirv::Dim,
238    flags: ImageTypeFlags,
239    image_format: spirv::ImageFormat,
240}
241
242bitflags::bitflags! {
243    /// Flags corresponding to the boolean(-ish) parameters to OpTypeImage.
244    #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
245    pub struct ImageTypeFlags: u8 {
246        const DEPTH = 0x1;
247        const ARRAYED = 0x2;
248        const MULTISAMPLED = 0x4;
249        const SAMPLED = 0x8;
250    }
251}
252
253impl LocalImageType {
254    /// Construct a `LocalImageType` from the fields of a `TypeInner::Image`.
255    fn from_inner(dim: crate::ImageDimension, arrayed: bool, class: crate::ImageClass) -> Self {
256        let make_flags = |multi: bool, other: ImageTypeFlags| -> ImageTypeFlags {
257            let mut flags = other;
258            flags.set(ImageTypeFlags::ARRAYED, arrayed);
259            flags.set(ImageTypeFlags::MULTISAMPLED, multi);
260            flags
261        };
262
263        let dim = spirv::Dim::from(dim);
264
265        match class {
266            crate::ImageClass::Sampled { kind, multi } => LocalImageType {
267                sampled_type: crate::Scalar { kind, width: 4 },
268                dim,
269                flags: make_flags(multi, ImageTypeFlags::SAMPLED),
270                image_format: spirv::ImageFormat::Unknown,
271            },
272            crate::ImageClass::Depth { multi } => LocalImageType {
273                sampled_type: crate::Scalar {
274                    kind: crate::ScalarKind::Float,
275                    width: 4,
276                },
277                dim,
278                flags: make_flags(multi, ImageTypeFlags::DEPTH | ImageTypeFlags::SAMPLED),
279                image_format: spirv::ImageFormat::Unknown,
280            },
281            crate::ImageClass::Storage { format, access: _ } => LocalImageType {
282                sampled_type: format.into(),
283                dim,
284                flags: make_flags(false, ImageTypeFlags::empty()),
285                image_format: format.into(),
286            },
287            crate::ImageClass::External => unimplemented!(),
288        }
289    }
290}
291
292/// A numeric type, for use in [`LocalType`].
293#[derive(Debug, PartialEq, Hash, Eq, Copy, Clone)]
294enum NumericType {
295    Scalar(crate::Scalar),
296    Vector {
297        size: crate::VectorSize,
298        scalar: crate::Scalar,
299    },
300    Matrix {
301        columns: crate::VectorSize,
302        rows: crate::VectorSize,
303        scalar: crate::Scalar,
304    },
305}
306
307impl NumericType {
308    const fn from_inner(inner: &crate::TypeInner) -> Option<Self> {
309        match *inner {
310            crate::TypeInner::Scalar(scalar) | crate::TypeInner::Atomic(scalar) => {
311                Some(NumericType::Scalar(scalar))
312            }
313            crate::TypeInner::Vector { size, scalar } => Some(NumericType::Vector { size, scalar }),
314            crate::TypeInner::Matrix {
315                columns,
316                rows,
317                scalar,
318            } => Some(NumericType::Matrix {
319                columns,
320                rows,
321                scalar,
322            }),
323            _ => None,
324        }
325    }
326
327    const fn scalar(self) -> crate::Scalar {
328        match self {
329            NumericType::Scalar(scalar)
330            | NumericType::Vector { scalar, .. }
331            | NumericType::Matrix { scalar, .. } => scalar,
332        }
333    }
334
335    const fn with_scalar(self, scalar: crate::Scalar) -> Self {
336        match self {
337            NumericType::Scalar(_) => NumericType::Scalar(scalar),
338            NumericType::Vector { size, .. } => NumericType::Vector { size, scalar },
339            NumericType::Matrix { columns, rows, .. } => NumericType::Matrix {
340                columns,
341                rows,
342                scalar,
343            },
344        }
345    }
346}
347
348/// A SPIR-V type constructed during code generation.
349///
350/// This is the variant of [`LookupType`] used to represent types that might not
351/// be available in the arena. Variants are present here for one of two reasons:
352///
353/// -   They represent types synthesized during code generation, as explained
354///     in the documentation for [`LookupType`].
355///
356/// -   They represent types for which SPIR-V forbids duplicate `OpType...`
357///     instructions, requiring deduplication.
358///
359/// This is not a complete copy of [`TypeInner`]: for example, SPIR-V generation
360/// never synthesizes new struct types, so `LocalType` has nothing for that.
361///
362/// Each `LocalType` variant should be handled identically to its analogous
363/// `TypeInner` variant. You can use the [`Writer::localtype_from_inner`]
364/// function to help with this, by converting everything possible to a
365/// `LocalType` before inspecting it.
366///
367/// ## `LocalType` equality and SPIR-V `OpType` uniqueness
368///
369/// The definition of `Eq` on `LocalType` is carefully chosen to help us follow
370/// certain SPIR-V rules. SPIR-V ยง2.8 requires some classes of `OpType...`
371/// instructions to be unique; for example, you can't have two `OpTypeInt 32 1`
372/// instructions in the same module. All 32-bit signed integers must use the
373/// same type id.
374///
375/// All SPIR-V types that must be unique can be represented as a `LocalType`,
376/// and two `LocalType`s are always `Eq` if SPIR-V would require them to use the
377/// same `OpType...` instruction. This lets us avoid duplicates by recording the
378/// ids of the type instructions we've already generated in a hash table,
379/// [`Writer::lookup_type`], keyed by `LocalType`.
380///
381/// As another example, [`LocalImageType`], stored in the `LocalType::Image`
382/// variant, is designed to help us deduplicate `OpTypeImage` instructions. See
383/// its documentation for details.
384///
385/// SPIR-V does not require pointer types to be unique - but different
386/// SPIR-V ids are considered to be distinct pointer types. Since Naga
387/// uses structural type equality, we need to represent each Naga
388/// equivalence class with a single SPIR-V `OpTypePointer`.
389///
390/// As it always must, the `Hash` implementation respects the `Eq` relation.
391///
392/// [`TypeInner`]: crate::TypeInner
393#[derive(Debug, PartialEq, Hash, Eq, Copy, Clone)]
394enum LocalType {
395    /// A numeric type.
396    Numeric(NumericType),
397    Pointer {
398        base: Word,
399        class: spirv::StorageClass,
400    },
401    Image(LocalImageType),
402    SampledImage {
403        image_type_id: Word,
404    },
405    Sampler,
406    BindingArray {
407        base: Handle<crate::Type>,
408        size: u32,
409    },
410    AccelerationStructure,
411    RayQuery,
412}
413
414/// A type encountered during SPIR-V generation.
415///
416/// In the process of writing SPIR-V, we need to synthesize various types for
417/// intermediate results and such: pointer types, vector/matrix component types,
418/// or even booleans, which usually appear in SPIR-V code even when they're not
419/// used by the module source.
420///
421/// However, we can't use `crate::Type` or `crate::TypeInner` for these, as the
422/// type arena may not contain what we need (it only contains types used
423/// directly by other parts of the IR), and the IR module is immutable, so we
424/// can't add anything to it.
425///
426/// So for local use in the SPIR-V writer, we use this type, which holds either
427/// a handle into the arena, or a [`LocalType`] containing something synthesized
428/// locally.
429///
430/// This is very similar to the [`proc::TypeResolution`] enum, with `LocalType`
431/// playing the role of `TypeInner`. However, `LocalType` also has other
432/// properties needed for SPIR-V generation; see the description of
433/// [`LocalType`] for details.
434///
435/// [`proc::TypeResolution`]: crate::proc::TypeResolution
436#[derive(Debug, PartialEq, Hash, Eq, Copy, Clone)]
437enum LookupType {
438    Handle(Handle<crate::Type>),
439    Local(LocalType),
440}
441
442impl From<LocalType> for LookupType {
443    fn from(local: LocalType) -> Self {
444        Self::Local(local)
445    }
446}
447
448#[derive(Debug, PartialEq, Clone, Hash, Eq)]
449struct LookupFunctionType {
450    parameter_type_ids: Vec<Word>,
451    return_type_id: Word,
452}
453
454#[derive(Debug, PartialEq, Clone, Hash, Eq)]
455enum LookupRayQueryFunction {
456    Initialize,
457    Proceed,
458    GenerateIntersection,
459    ConfirmIntersection,
460    GetVertexPositions { committed: bool },
461    GetIntersection { committed: bool },
462}
463
464#[derive(Debug)]
465enum Dimension {
466    Scalar,
467    Vector,
468    Matrix,
469}
470
471/// Key used to look up an operation which we have wrapped in a helper
472/// function, which should be called instead of directly emitting code
473/// for the expression. See [`Writer::wrapped_functions`].
474#[derive(Debug, Eq, PartialEq, Hash)]
475enum WrappedFunction {
476    BinaryOp {
477        op: crate::BinaryOperator,
478        left_type_id: Word,
479        right_type_id: Word,
480    },
481}
482
483/// A map from evaluated [`Expression`](crate::Expression)s to their SPIR-V ids.
484///
485/// When we emit code to evaluate a given `Expression`, we record the
486/// SPIR-V id of its value here, under its `Handle<Expression>` index.
487///
488/// A `CachedExpressions` value can be indexed by a `Handle<Expression>` value.
489///
490/// [emit]: index.html#expression-evaluation-time-and-scope
491#[derive(Default)]
492struct CachedExpressions {
493    ids: HandleVec<crate::Expression, Word>,
494}
495impl CachedExpressions {
496    fn reset(&mut self, length: usize) {
497        self.ids.clear();
498        self.ids.resize(length, 0);
499    }
500}
501impl ops::Index<Handle<crate::Expression>> for CachedExpressions {
502    type Output = Word;
503    fn index(&self, h: Handle<crate::Expression>) -> &Word {
504        let id = &self.ids[h];
505        if *id == 0 {
506            unreachable!("Expression {:?} is not cached!", h);
507        }
508        id
509    }
510}
511impl ops::IndexMut<Handle<crate::Expression>> for CachedExpressions {
512    fn index_mut(&mut self, h: Handle<crate::Expression>) -> &mut Word {
513        let id = &mut self.ids[h];
514        if *id != 0 {
515            unreachable!("Expression {:?} is already cached!", h);
516        }
517        id
518    }
519}
520impl recyclable::Recyclable for CachedExpressions {
521    fn recycle(self) -> Self {
522        CachedExpressions {
523            ids: self.ids.recycle(),
524        }
525    }
526}
527
528#[derive(Eq, Hash, PartialEq)]
529enum CachedConstant {
530    Literal(crate::proc::HashableLiteral),
531    Composite {
532        ty: LookupType,
533        constituent_ids: Vec<Word>,
534    },
535    ZeroValue(Word),
536}
537
538/// The SPIR-V representation of a [`crate::GlobalVariable`].
539///
540/// In the Vulkan spec 1.3.296, the section [Descriptor Set Interface][dsi] says:
541///
542/// > Variables identified with the `Uniform` storage class are used to access
543/// > transparent buffer backed resources. Such variables *must* be:
544/// >
545/// > -   typed as `OpTypeStruct`, or an array of this type,
546/// >
547/// > -   identified with a `Block` or `BufferBlock` decoration, and
548/// >
549/// > -   laid out explicitly using the `Offset`, `ArrayStride`, and `MatrixStride`
550/// >     decorations as specified in "Offset and Stride Assignment".
551///
552/// This is followed by identical language for the `StorageBuffer`,
553/// except that a `BufferBlock` decoration is not allowed.
554///
555/// When we encounter a global variable in the [`Storage`] or [`Uniform`]
556/// address spaces whose type is not already [`Struct`], this backend implicitly
557/// wraps the global variable in a struct: we generate a SPIR-V global variable
558/// holding an `OpTypeStruct` with a single member, whose type is what the Naga
559/// global's type would suggest, decorated as required above.
560///
561/// The [`helpers::global_needs_wrapper`] function determines whether a given
562/// [`crate::GlobalVariable`] needs to be wrapped.
563///
564/// [dsi]: https://registry.khronos.org/vulkan/specs/1.3-extensions/html/vkspec.html#interfaces-resources-descset
565/// [`Storage`]: crate::AddressSpace::Storage
566/// [`Uniform`]: crate::AddressSpace::Uniform
567/// [`Struct`]: crate::TypeInner::Struct
568#[derive(Clone)]
569struct GlobalVariable {
570    /// The SPIR-V id of the `OpVariable` that declares the global.
571    ///
572    /// If this global has been implicitly wrapped in an `OpTypeStruct`, this id
573    /// refers to the wrapper, not the original Naga value it contains. If you
574    /// need the Naga value, use [`access_id`] instead of this field.
575    ///
576    /// If this global is not implicitly wrapped, this is the same as
577    /// [`access_id`].
578    ///
579    /// This is used to compute the `access_id` pointer in function prologues,
580    /// and used for `ArrayLength` expressions, which need to pass the wrapper
581    /// struct.
582    ///
583    /// [`access_id`]: GlobalVariable::access_id
584    var_id: Word,
585
586    /// The loaded value of a `AddressSpace::Handle` global variable.
587    ///
588    /// If the current function uses this global variable, this is the id of an
589    /// `OpLoad` instruction in the function's prologue that loads its value.
590    /// (This value is assigned as we write the prologue code of each function.)
591    /// It is then used for all operations on the global, such as `OpImageSample`.
592    handle_id: Word,
593
594    /// The SPIR-V id of a pointer to this variable's Naga IR value.
595    ///
596    /// If the current function uses this global variable, and it has been
597    /// implicitly wrapped in an `OpTypeStruct`, this is the id of an
598    /// `OpAccessChain` instruction in the function's prologue that refers to
599    /// the wrapped value inside the struct. (This value is assigned as we write
600    /// the prologue code of each function.) If you need the wrapper struct
601    /// itself, use [`var_id`] instead of this field.
602    ///
603    /// If this global is not implicitly wrapped, this is the same as
604    /// [`var_id`].
605    ///
606    /// [`var_id`]: GlobalVariable::var_id
607    access_id: Word,
608}
609
610impl GlobalVariable {
611    const fn dummy() -> Self {
612        Self {
613            var_id: 0,
614            handle_id: 0,
615            access_id: 0,
616        }
617    }
618
619    const fn new(id: Word) -> Self {
620        Self {
621            var_id: id,
622            handle_id: 0,
623            access_id: 0,
624        }
625    }
626
627    /// Prepare `self` for use within a single function.
628    fn reset_for_function(&mut self) {
629        self.handle_id = 0;
630        self.access_id = 0;
631    }
632}
633
634struct FunctionArgument {
635    /// Actual instruction of the argument.
636    instruction: Instruction,
637    handle_id: Word,
638}
639
640/// Tracks the expressions for which the backend emits the following instructions:
641/// - OpConstantTrue
642/// - OpConstantFalse
643/// - OpConstant
644/// - OpConstantComposite
645/// - OpConstantNull
646struct ExpressionConstnessTracker {
647    inner: crate::arena::HandleSet<crate::Expression>,
648}
649
650impl ExpressionConstnessTracker {
651    fn from_arena(arena: &crate::Arena<crate::Expression>) -> Self {
652        let mut inner = crate::arena::HandleSet::for_arena(arena);
653        for (handle, expr) in arena.iter() {
654            let insert = match *expr {
655                crate::Expression::Literal(_)
656                | crate::Expression::ZeroValue(_)
657                | crate::Expression::Constant(_) => true,
658                crate::Expression::Compose { ref components, .. } => {
659                    components.iter().all(|&h| inner.contains(h))
660                }
661                crate::Expression::Splat { value, .. } => inner.contains(value),
662                _ => false,
663            };
664            if insert {
665                inner.insert(handle);
666            }
667        }
668        Self { inner }
669    }
670
671    fn is_const(&self, value: Handle<crate::Expression>) -> bool {
672        self.inner.contains(value)
673    }
674}
675
676/// General information needed to emit SPIR-V for Naga statements.
677struct BlockContext<'w> {
678    /// The writer handling the module to which this code belongs.
679    writer: &'w mut Writer,
680
681    /// The [`Module`](crate::Module) for which we're generating code.
682    ir_module: &'w crate::Module,
683
684    /// The [`Function`](crate::Function) for which we're generating code.
685    ir_function: &'w crate::Function,
686
687    /// Information module validation produced about
688    /// [`ir_function`](BlockContext::ir_function).
689    fun_info: &'w crate::valid::FunctionInfo,
690
691    /// The [`spv::Function`](Function) to which we are contributing SPIR-V instructions.
692    function: &'w mut Function,
693
694    /// SPIR-V ids for expressions we've evaluated.
695    cached: CachedExpressions,
696
697    /// The `Writer`'s temporary vector, for convenience.
698    temp_list: Vec<Word>,
699
700    /// Tracks the constness of `Expression`s residing in `self.ir_function.expressions`
701    expression_constness: ExpressionConstnessTracker,
702
703    force_loop_bounding: bool,
704
705    /// Hash from an expression whose type is a ray query / pointer to a ray query to its tracker.
706    /// Note: this is sparse, so can't be a handle vec
707    ray_query_tracker_expr: crate::FastHashMap<Handle<crate::Expression>, RayQueryTrackers>,
708}
709
710#[derive(Clone, Copy)]
711struct RayQueryTrackers {
712    // Initialization tracker
713    initialized_tracker: Word,
714    // Tracks the t max from ray query initialize.
715    // Unlike HLSL, spir-v's equivalent getter for the current committed t has UB (instead of just
716    // returning t_max) if there was no previous hit (though in some places it treats the behaviour as
717    // defined), therefore we must track the tmax inputted into ray query initialize.
718    t_max_tracker: Word,
719}
720
721impl BlockContext<'_> {
722    fn gen_id(&mut self) -> Word {
723        self.writer.id_gen.next()
724    }
725
726    fn get_type_id(&mut self, lookup_type: LookupType) -> Word {
727        self.writer.get_type_id(lookup_type)
728    }
729
730    fn get_handle_type_id(&mut self, handle: Handle<crate::Type>) -> Word {
731        self.writer.get_handle_type_id(handle)
732    }
733
734    fn get_expression_type_id(&mut self, tr: &TypeResolution) -> Word {
735        self.writer.get_expression_type_id(tr)
736    }
737
738    fn get_index_constant(&mut self, index: Word) -> Word {
739        self.writer.get_constant_scalar(crate::Literal::U32(index))
740    }
741
742    fn get_scope_constant(&mut self, scope: Word) -> Word {
743        self.writer
744            .get_constant_scalar(crate::Literal::I32(scope as _))
745    }
746
747    fn get_pointer_type_id(&mut self, base: Word, class: spirv::StorageClass) -> Word {
748        self.writer.get_pointer_type_id(base, class)
749    }
750
751    fn get_numeric_type_id(&mut self, numeric: NumericType) -> Word {
752        self.writer.get_numeric_type_id(numeric)
753    }
754}
755
756pub struct Writer {
757    physical_layout: PhysicalLayout,
758    logical_layout: LogicalLayout,
759    id_gen: IdGenerator,
760
761    /// The set of capabilities modules are permitted to use.
762    ///
763    /// This is initialized from `Options::capabilities`.
764    capabilities_available: Option<crate::FastHashSet<Capability>>,
765
766    /// The set of capabilities used by this module.
767    ///
768    /// If `capabilities_available` is `Some`, then this is always a subset of
769    /// that.
770    capabilities_used: crate::FastIndexSet<Capability>,
771
772    /// The set of spirv extensions used.
773    extensions_used: crate::FastIndexSet<&'static str>,
774
775    debug_strings: Vec<Instruction>,
776    debugs: Vec<Instruction>,
777    annotations: Vec<Instruction>,
778    flags: WriterFlags,
779    bounds_check_policies: BoundsCheckPolicies,
780    zero_initialize_workgroup_memory: ZeroInitializeWorkgroupMemoryMode,
781    force_loop_bounding: bool,
782    use_storage_input_output_16: bool,
783    void_type: Word,
784    //TODO: convert most of these into vectors, addressable by handle indices
785    lookup_type: crate::FastHashMap<LookupType, Word>,
786    lookup_function: crate::FastHashMap<Handle<crate::Function>, Word>,
787    lookup_function_type: crate::FastHashMap<LookupFunctionType, Word>,
788    /// Operations which have been wrapped in a helper function. The value is
789    /// the ID of the function, which should be called instead of emitting code
790    /// for the operation directly.
791    wrapped_functions: crate::FastHashMap<WrappedFunction, Word>,
792    /// Indexed by const-expression handle indexes
793    constant_ids: HandleVec<crate::Expression, Word>,
794    cached_constants: crate::FastHashMap<CachedConstant, Word>,
795    global_variables: HandleVec<crate::GlobalVariable, GlobalVariable>,
796    fake_missing_bindings: bool,
797    binding_map: BindingMap,
798
799    // Cached expressions are only meaningful within a BlockContext, but we
800    // retain the table here between functions to save heap allocations.
801    saved_cached: CachedExpressions,
802
803    gl450_ext_inst_id: Word,
804
805    // Just a temporary list of SPIR-V ids
806    temp_list: Vec<Word>,
807
808    ray_query_functions: crate::FastHashMap<LookupRayQueryFunction, Word>,
809
810    /// F16 I/O polyfill manager for handling `f16` input/output variables
811    /// when `StorageInputOutput16` capability is not available.
812    io_f16_polyfills: f16_polyfill::F16IoPolyfill,
813
814    /// Non semantic debug printf extension `OpExtInstImport`
815    debug_printf: Option<Word>,
816    pub(crate) ray_query_initialization_tracking: bool,
817}
818
819bitflags::bitflags! {
820    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
821    pub struct WriterFlags: u32 {
822        /// Include debug labels for everything.
823        const DEBUG = 0x1;
824
825        /// Flip Y coordinate of [`BuiltIn::Position`] output.
826        ///
827        /// [`BuiltIn::Position`]: crate::BuiltIn::Position
828        const ADJUST_COORDINATE_SPACE = 0x2;
829
830        /// Emit [`OpName`][op] for input/output locations.
831        ///
832        /// Contrary to spec, some drivers treat it as semantic, not allowing
833        /// any conflicts.
834        ///
835        /// [op]: https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpName
836        const LABEL_VARYINGS = 0x4;
837
838        /// Emit [`PointSize`] output builtin to vertex shaders, which is
839        /// required for drawing with `PointList` topology.
840        ///
841        /// [`PointSize`]: crate::BuiltIn::PointSize
842        const FORCE_POINT_SIZE = 0x8;
843
844        /// Clamp [`BuiltIn::FragDepth`] output between 0 and 1.
845        ///
846        /// [`BuiltIn::FragDepth`]: crate::BuiltIn::FragDepth
847        const CLAMP_FRAG_DEPTH = 0x10;
848
849        /// Instead of silently failing if the arguments to generate a ray query are
850        /// invalid, uses debug printf extension to print to the command line
851        ///
852        /// Note: VK_KHR_shader_non_semantic_info must be enabled. This will have no
853        /// effect if `options.ray_query_initialization_tracking` is set to false.
854        const PRINT_ON_RAY_QUERY_INITIALIZATION_FAIL = 0x20;
855    }
856}
857
858bitflags::bitflags! {
859    /// How far through a ray query are we
860    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
861    pub(super) struct RayQueryPoint: u32 {
862        /// Ray query has been successfully initialized.
863        const INITIALIZED = 1 << 0;
864        /// Proceed has been called on ray query.
865        const PROCEED = 1 << 1;
866        /// Proceed has returned false (have finished traversal).
867        const FINISHED_TRAVERSAL = 1 << 2;
868    }
869}
870
871#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash)]
872#[cfg_attr(feature = "serialize", derive(serde::Serialize))]
873#[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
874pub struct BindingInfo {
875    pub descriptor_set: u32,
876    pub binding: u32,
877    /// If the binding is an unsized binding array, this overrides the size.
878    pub binding_array_size: Option<u32>,
879}
880
881// Using `BTreeMap` instead of `HashMap` so that we can hash itself.
882pub type BindingMap = alloc::collections::BTreeMap<crate::ResourceBinding, BindingInfo>;
883
884#[derive(Clone, Copy, Debug, PartialEq, Eq)]
885pub enum ZeroInitializeWorkgroupMemoryMode {
886    /// Via `VK_KHR_zero_initialize_workgroup_memory` or Vulkan 1.3
887    Native,
888    /// Via assignments + barrier
889    Polyfill,
890    None,
891}
892
893#[derive(Debug, Clone)]
894pub struct Options<'a> {
895    /// (Major, Minor) target version of the SPIR-V.
896    pub lang_version: (u8, u8),
897
898    /// Configuration flags for the writer.
899    pub flags: WriterFlags,
900
901    /// Don't panic on missing bindings. Instead use fake values for `Binding`
902    /// and `DescriptorSet` decorations. This may result in invalid SPIR-V.
903    pub fake_missing_bindings: bool,
904
905    /// Map of resources to information about the binding.
906    pub binding_map: BindingMap,
907
908    /// If given, the set of capabilities modules are allowed to use. Code that
909    /// requires capabilities beyond these is rejected with an error.
910    ///
911    /// If this is `None`, all capabilities are permitted.
912    pub capabilities: Option<crate::FastHashSet<Capability>>,
913
914    /// How should generate code handle array, vector, matrix, or image texel
915    /// indices that are out of range?
916    pub bounds_check_policies: BoundsCheckPolicies,
917
918    /// Dictates the way workgroup variables should be zero initialized
919    pub zero_initialize_workgroup_memory: ZeroInitializeWorkgroupMemoryMode,
920
921    /// If set, loops will have code injected into them, forcing the compiler
922    /// to think the number of iterations is bounded.
923    pub force_loop_bounding: bool,
924
925    /// if set, ray queries will get a variable to track their state to prevent
926    /// misuse.
927    pub ray_query_initialization_tracking: bool,
928
929    /// Whether to use the `StorageInputOutput16` capability for `f16` shader I/O.
930    /// When false, `f16` I/O is polyfilled using `f32` types with conversions.
931    pub use_storage_input_output_16: bool,
932
933    pub debug_info: Option<DebugInfo<'a>>,
934}
935
936impl Default for Options<'_> {
937    fn default() -> Self {
938        let mut flags = WriterFlags::ADJUST_COORDINATE_SPACE
939            | WriterFlags::LABEL_VARYINGS
940            | WriterFlags::CLAMP_FRAG_DEPTH;
941        if cfg!(debug_assertions) {
942            flags |= WriterFlags::DEBUG;
943        }
944        Options {
945            lang_version: (1, 0),
946            flags,
947            fake_missing_bindings: true,
948            binding_map: BindingMap::default(),
949            capabilities: None,
950            bounds_check_policies: BoundsCheckPolicies::default(),
951            zero_initialize_workgroup_memory: ZeroInitializeWorkgroupMemoryMode::Polyfill,
952            force_loop_bounding: true,
953            ray_query_initialization_tracking: true,
954            use_storage_input_output_16: true,
955            debug_info: None,
956        }
957    }
958}
959
960// A subset of options meant to be changed per pipeline.
961#[derive(Debug, Clone)]
962#[cfg_attr(feature = "serialize", derive(serde::Serialize))]
963#[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
964pub struct PipelineOptions {
965    /// The stage of the entry point.
966    pub shader_stage: crate::ShaderStage,
967    /// The name of the entry point.
968    ///
969    /// If no entry point that matches is found while creating a [`Writer`], a error will be thrown.
970    pub entry_point: String,
971}
972
973pub fn write_vec(
974    module: &crate::Module,
975    info: &crate::valid::ModuleInfo,
976    options: &Options,
977    pipeline_options: Option<&PipelineOptions>,
978) -> Result<Vec<u32>, Error> {
979    let mut words: Vec<u32> = Vec::new();
980    let mut w = Writer::new(options)?;
981
982    w.write(
983        module,
984        info,
985        pipeline_options,
986        &options.debug_info,
987        &mut words,
988    )?;
989    Ok(words)
990}