naga/back/spv/
mesh_shader.rs

1use alloc::vec::Vec;
2use spirv::Word;
3
4use crate::{
5    back::spv::{
6        helpers::BindingDecorations, writer::FunctionInterface, Block, EntryPointContext, Error,
7        Instruction, WriterFlags,
8    },
9    non_max_u32::NonMaxU32,
10    Handle,
11};
12
13#[derive(Clone, Debug)]
14pub struct MeshReturnMember {
15    pub ty_id: u32,
16    pub binding: crate::Binding,
17}
18
19#[derive(Debug)]
20struct PerOutputTypeMeshReturnInfo {
21    max_length_constant: Word,
22    array_type_id: Word,
23    struct_members: Vec<MeshReturnMember>,
24
25    // * Most builtins must be in the same block.
26    // * All bindings must be in their own unique block.
27    // * The primitive indices builtin family needs its own block.
28    // * Cull primitive doesn't care about having its own block, but
29    //   some older validation layers didn't respect this.
30    builtin_block: Option<Word>,
31    bindings: Vec<Word>,
32}
33
34#[derive(Debug)]
35pub struct MeshReturnInfo {
36    /// Id of the workgroup variable containing the data to be output
37    out_variable_id: Word,
38    /// All members of the output variable struct type
39    out_members: Vec<MeshReturnMember>,
40    /// Id of the input variable for local invocation id
41    local_invocation_index_var_id: Word,
42    /// Total workgroup size (product)
43    workgroup_size: u32,
44
45    /// Vertex-specific info
46    vertex_info: PerOutputTypeMeshReturnInfo,
47    /// Primitive-specific info
48    primitive_info: PerOutputTypeMeshReturnInfo,
49    /// Array variable for the primitive indices builtin
50    primitive_indices: Option<Word>,
51}
52
53impl super::Writer {
54    /// Sets up an output variable that will handle part of the mesh shader output
55    pub(super) fn write_mesh_return_global_variable(
56        &mut self,
57        ty: u32,
58        array_size_id: u32,
59    ) -> Result<Word, Error> {
60        let array_ty = self.id_gen.next();
61        Instruction::type_array(array_ty, ty, array_size_id)
62            .to_words(&mut self.logical_layout.declarations);
63        let ptr_ty = self.get_pointer_type_id(array_ty, spirv::StorageClass::Output);
64        let var_id = self.id_gen.next();
65        Instruction::variable(ptr_ty, var_id, spirv::StorageClass::Output, None)
66            .to_words(&mut self.logical_layout.declarations);
67        Ok(var_id)
68    }
69
70    /// This does various setup things to allow mesh shader entry points
71    /// to be properly written, such as creating the output variables
72    pub(super) fn write_entry_point_mesh_shader_info(
73        &mut self,
74        iface: &mut FunctionInterface,
75        local_invocation_index_var_id: Option<Word>,
76        ir_module: &crate::Module,
77        ep_context: &mut EntryPointContext,
78    ) -> Result<(), Error> {
79        let Some(ref mesh_info) = iface.mesh_info else {
80            return Ok(());
81        };
82        // Collect the members in the output structs
83        let out_members: Vec<MeshReturnMember> =
84            match &ir_module.types[ir_module.global_variables[mesh_info.output_variable].ty] {
85                &crate::Type {
86                    inner: crate::TypeInner::Struct { ref members, .. },
87                    ..
88                } => members
89                    .iter()
90                    .map(|a| MeshReturnMember {
91                        ty_id: self.get_handle_type_id(a.ty),
92                        binding: a.binding.clone().unwrap(),
93                    })
94                    .collect(),
95                _ => unreachable!(),
96            };
97        let vertex_array_type_id = out_members
98            .iter()
99            .find(|a| a.binding == crate::Binding::BuiltIn(crate::BuiltIn::Vertices))
100            .unwrap()
101            .ty_id;
102        let primitive_array_type_id = out_members
103            .iter()
104            .find(|a| a.binding == crate::Binding::BuiltIn(crate::BuiltIn::Primitives))
105            .unwrap()
106            .ty_id;
107        let vertex_members = match &ir_module.types[mesh_info.vertex_output_type] {
108            &crate::Type {
109                inner: crate::TypeInner::Struct { ref members, .. },
110                ..
111            } => members
112                .iter()
113                .map(|a| MeshReturnMember {
114                    ty_id: self.get_handle_type_id(a.ty),
115                    binding: a.binding.clone().unwrap(),
116                })
117                .collect(),
118            _ => unreachable!(),
119        };
120        let primitive_members = match &ir_module.types[mesh_info.primitive_output_type] {
121            &crate::Type {
122                inner: crate::TypeInner::Struct { ref members, .. },
123                ..
124            } => members
125                .iter()
126                .map(|a| MeshReturnMember {
127                    ty_id: self.get_handle_type_id(a.ty),
128                    binding: a.binding.clone().unwrap(),
129                })
130                .collect(),
131            _ => unreachable!(),
132        };
133        // In the final return, we do a giant memcpy, for which this is helpful
134        let local_invocation_index_var_id = local_invocation_index_var_id
135            .expect("mesh entry points always declare LocalInvocationIndex");
136        // This is the information that is passed to the function writer
137        // so that it can write the final return logic
138        let mut mesh_return_info = MeshReturnInfo {
139            out_variable_id: self.global_variables[mesh_info.output_variable].var_id,
140            out_members,
141            local_invocation_index_var_id,
142            workgroup_size: self
143                .get_constant_scalar(crate::Literal::U32(iface.workgroup_size.iter().product())),
144
145            vertex_info: PerOutputTypeMeshReturnInfo {
146                array_type_id: vertex_array_type_id,
147                struct_members: vertex_members,
148                max_length_constant: self
149                    .get_constant_scalar(crate::Literal::U32(mesh_info.max_vertices)),
150                bindings: Vec::new(),
151                builtin_block: None,
152            },
153            primitive_info: PerOutputTypeMeshReturnInfo {
154                array_type_id: primitive_array_type_id,
155                struct_members: primitive_members,
156                max_length_constant: self
157                    .get_constant_scalar(crate::Literal::U32(mesh_info.max_primitives)),
158                bindings: Vec::new(),
159                builtin_block: None,
160            },
161            primitive_indices: None,
162        };
163        let vert_array_size_id =
164            self.get_constant_scalar(crate::Literal::U32(mesh_info.max_vertices));
165        let prim_array_size_id =
166            self.get_constant_scalar(crate::Literal::U32(mesh_info.max_primitives));
167
168        // Create the actual output variables and types.
169        // According to SPIR-V,
170        // * All builtins must be in the same output `Block` (except builtins for different output types like vertex/primitive)
171        // * Each member with `location` must be in its own `Block` decorated `struct`
172        // * Some builtins like CullPrimitiveEXT don't care as much (older validation layers don't know this! Wonderful!)
173        // * Some builtins like the indices ones need to be in their own output variable without a struct wrapper
174
175        // Write vertex builtin block
176        if mesh_return_info
177            .vertex_info
178            .struct_members
179            .iter()
180            .any(|a| matches!(a.binding, crate::Binding::BuiltIn(..)))
181        {
182            let builtin_block_ty_id = self.id_gen.next();
183            let mut ins = Instruction::type_struct(builtin_block_ty_id, &[]);
184            let mut bi_index = 0;
185            let mut decorations = Vec::new();
186            for member in &mesh_return_info.vertex_info.struct_members {
187                if let crate::Binding::BuiltIn(_) = member.binding {
188                    ins.add_operand(member.ty_id);
189                    let binding = self.map_binding(
190                        ir_module,
191                        iface.stage,
192                        spirv::StorageClass::Output,
193                        // Unused except in fragment shaders with other conditions, so we can pass null
194                        Handle::new(NonMaxU32::new(0).unwrap()),
195                        &member.binding,
196                    )?;
197                    match binding {
198                        BindingDecorations::BuiltIn(bi, others) => {
199                            decorations.push(Instruction::member_decorate(
200                                builtin_block_ty_id,
201                                bi_index,
202                                spirv::Decoration::BuiltIn,
203                                &[bi as Word],
204                            ));
205                            for other in others {
206                                decorations.push(Instruction::member_decorate(
207                                    builtin_block_ty_id,
208                                    bi_index,
209                                    other,
210                                    &[],
211                                ));
212                            }
213                        }
214                        _ => unreachable!(),
215                    }
216                    bi_index += 1;
217                }
218            }
219            ins.to_words(&mut self.logical_layout.declarations);
220            decorations.push(Instruction::decorate(
221                builtin_block_ty_id,
222                spirv::Decoration::Block,
223                &[],
224            ));
225            for dec in decorations {
226                dec.to_words(&mut self.logical_layout.annotations);
227            }
228            let v =
229                self.write_mesh_return_global_variable(builtin_block_ty_id, vert_array_size_id)?;
230            iface.varying_ids.push(v);
231            if self.flags.contains(WriterFlags::DEBUG) {
232                self.debugs
233                    .push(Instruction::name(v, "naga_vertex_builtin_outputs"));
234            }
235            mesh_return_info.vertex_info.builtin_block = Some(v);
236        }
237        // Write primitive builtin block
238        if mesh_return_info
239            .primitive_info
240            .struct_members
241            .iter()
242            .any(|a| {
243                !matches!(
244                    a.binding,
245                    crate::Binding::BuiltIn(
246                        crate::BuiltIn::PointIndex
247                            | crate::BuiltIn::LineIndices
248                            | crate::BuiltIn::TriangleIndices
249                    ) | crate::Binding::Location { .. }
250                )
251            })
252        {
253            let builtin_block_ty_id = self.id_gen.next();
254            let mut ins = Instruction::type_struct(builtin_block_ty_id, &[]);
255            let mut bi_index = 0;
256            let mut decorations = Vec::new();
257            for member in &mesh_return_info.primitive_info.struct_members {
258                if let crate::Binding::BuiltIn(bi) = member.binding {
259                    // These need to be in their own block, unlike other builtins
260                    if matches!(
261                        bi,
262                        crate::BuiltIn::PointIndex
263                            | crate::BuiltIn::LineIndices
264                            | crate::BuiltIn::TriangleIndices,
265                    ) {
266                        continue;
267                    }
268                    ins.add_operand(member.ty_id);
269                    let binding = self.map_binding(
270                        ir_module,
271                        iface.stage,
272                        spirv::StorageClass::Output,
273                        // Unused except in fragment shaders with other conditions, so we can pass null
274                        Handle::new(NonMaxU32::new(0).unwrap()),
275                        &member.binding,
276                    )?;
277                    match binding {
278                        BindingDecorations::BuiltIn(bi, others) => {
279                            decorations.push(Instruction::member_decorate(
280                                builtin_block_ty_id,
281                                bi_index,
282                                spirv::Decoration::BuiltIn,
283                                &[bi as Word],
284                            ));
285                            for other in others {
286                                decorations.push(Instruction::member_decorate(
287                                    builtin_block_ty_id,
288                                    bi_index,
289                                    other,
290                                    &[],
291                                ));
292                            }
293                        }
294                        _ => unreachable!(),
295                    }
296                    bi_index += 1;
297                }
298            }
299            ins.to_words(&mut self.logical_layout.declarations);
300            decorations.push(Instruction::decorate(
301                builtin_block_ty_id,
302                spirv::Decoration::Block,
303                &[],
304            ));
305            for dec in decorations {
306                dec.to_words(&mut self.logical_layout.annotations);
307            }
308            let v =
309                self.write_mesh_return_global_variable(builtin_block_ty_id, prim_array_size_id)?;
310            Instruction::decorate(v, spirv::Decoration::PerPrimitiveEXT, &[])
311                .to_words(&mut self.logical_layout.annotations);
312            iface.varying_ids.push(v);
313            if self.flags.contains(WriterFlags::DEBUG) {
314                self.debugs
315                    .push(Instruction::name(v, "naga_primitive_builtin_outputs"));
316            }
317            mesh_return_info.primitive_info.builtin_block = Some(v);
318        }
319
320        // Write vertex binding output blocks (1 array per output struct member)
321        for member in &mesh_return_info.vertex_info.struct_members {
322            match member.binding {
323                crate::Binding::Location { location, .. } => {
324                    // Create variable
325                    let v =
326                        self.write_mesh_return_global_variable(member.ty_id, vert_array_size_id)?;
327                    // Decorate the variable with Location
328                    Instruction::decorate(v, spirv::Decoration::Location, &[location])
329                        .to_words(&mut self.logical_layout.annotations);
330                    iface.varying_ids.push(v);
331                    mesh_return_info.vertex_info.bindings.push(v);
332                }
333                crate::Binding::BuiltIn(_) => (),
334            }
335        }
336        // Write primitive binding output blocks (1 array per output struct member)
337        // Also write indices output block
338        for member in &mesh_return_info.primitive_info.struct_members {
339            match member.binding {
340                crate::Binding::BuiltIn(
341                    crate::BuiltIn::PointIndex
342                    | crate::BuiltIn::LineIndices
343                    | crate::BuiltIn::TriangleIndices,
344                ) => {
345                    // This is written here instead of as part of the builtin block
346                    let v =
347                        self.write_mesh_return_global_variable(member.ty_id, prim_array_size_id)?;
348                    // This shouldn't be marked as PerPrimitiveEXT
349                    Instruction::decorate(
350                        v,
351                        spirv::Decoration::BuiltIn,
352                        &[match member.binding.to_built_in().unwrap() {
353                            crate::BuiltIn::PointIndex => spirv::BuiltIn::PrimitivePointIndicesEXT,
354                            crate::BuiltIn::LineIndices => spirv::BuiltIn::PrimitiveLineIndicesEXT,
355                            crate::BuiltIn::TriangleIndices => {
356                                spirv::BuiltIn::PrimitiveTriangleIndicesEXT
357                            }
358                            _ => unreachable!(),
359                        } as Word],
360                    )
361                    .to_words(&mut self.logical_layout.annotations);
362                    iface.varying_ids.push(v);
363                    if self.flags.contains(WriterFlags::DEBUG) {
364                        self.debugs
365                            .push(Instruction::name(v, "naga_primitive_indices_outputs"));
366                    }
367                    mesh_return_info.primitive_indices = Some(v);
368                }
369                crate::Binding::Location { location, .. } => {
370                    // Create variable
371                    let v =
372                        self.write_mesh_return_global_variable(member.ty_id, prim_array_size_id)?;
373                    // Decorate the variable with Location
374                    Instruction::decorate(v, spirv::Decoration::Location, &[location])
375                        .to_words(&mut self.logical_layout.annotations);
376                    // Decorate it with PerPrimitiveEXT
377                    Instruction::decorate(v, spirv::Decoration::PerPrimitiveEXT, &[])
378                        .to_words(&mut self.logical_layout.annotations);
379                    iface.varying_ids.push(v);
380
381                    mesh_return_info.primitive_info.bindings.push(v);
382                }
383                crate::Binding::BuiltIn(_) => (),
384            }
385        }
386
387        // Store this where it can be read later during function write
388        ep_context.mesh_state = Some(mesh_return_info);
389
390        Ok(())
391    }
392
393    pub(super) fn write_entry_point_task_return(
394        &mut self,
395        value_id: Word,
396        body: &mut Vec<Instruction>,
397        task_payload: Word,
398    ) -> Result<Instruction, Error> {
399        // OpEmitMeshTasksEXT must be called right before exiting (after setting other
400        // output variables if there are any)
401
402        // Extract the vec3<u32> into 3 u32's
403        let values = [self.id_gen.next(), self.id_gen.next(), self.id_gen.next()];
404        for (i, &value) in values.iter().enumerate() {
405            let instruction = Instruction::composite_extract(
406                self.get_u32_type_id(),
407                value,
408                value_id,
409                &[i as Word],
410            );
411            body.push(instruction);
412        }
413        let mut instruction = Instruction::new(spirv::Op::EmitMeshTasksEXT);
414        for id in values {
415            instruction.add_operand(id);
416        }
417        // We have to include the task payload in our call
418        instruction.add_operand(task_payload);
419        Ok(instruction)
420    }
421
422    /// This writes the actual loop
423    #[allow(clippy::too_many_arguments)]
424    fn write_mesh_copy_loop(
425        &mut self,
426        body: &mut Vec<Instruction>,
427        mut loop_body_block: Vec<Instruction>,
428        loop_header: u32,
429        loop_merge: u32,
430        count_id: u32,
431        index_var: u32,
432        return_info: &MeshReturnInfo,
433    ) {
434        let u32_id = self.get_u32_type_id();
435        let condition_check = self.id_gen.next();
436        let loop_continue = self.id_gen.next();
437        let loop_body = self.id_gen.next();
438
439        // Loop header
440        {
441            body.push(Instruction::label(loop_header));
442            body.push(Instruction::loop_merge(
443                loop_merge,
444                loop_continue,
445                spirv::SelectionControl::empty(),
446            ));
447            body.push(Instruction::branch(condition_check));
448        }
449        // Condition check - check if i is less than num vertices to copy
450        {
451            body.push(Instruction::label(condition_check));
452
453            let val_i = self.id_gen.next();
454            body.push(Instruction::load(u32_id, val_i, index_var, None));
455
456            let cond = self.id_gen.next();
457            body.push(Instruction::binary(
458                spirv::Op::ULessThan,
459                self.get_bool_type_id(),
460                cond,
461                val_i,
462                count_id,
463            ));
464            body.push(Instruction::branch_conditional(cond, loop_body, loop_merge));
465        }
466        // Loop body
467        {
468            body.push(Instruction::label(loop_body));
469            body.append(&mut loop_body_block);
470            body.push(Instruction::branch(loop_continue));
471        }
472        // Loop continue - increment i
473        {
474            body.push(Instruction::label(loop_continue));
475
476            let prev_val_i = self.id_gen.next();
477            body.push(Instruction::load(u32_id, prev_val_i, index_var, None));
478            let new_val_i = self.id_gen.next();
479            body.push(Instruction::binary(
480                spirv::Op::IAdd,
481                u32_id,
482                new_val_i,
483                prev_val_i,
484                return_info.workgroup_size,
485            ));
486            body.push(Instruction::store(index_var, new_val_i, None));
487
488            body.push(Instruction::branch(loop_header));
489        }
490    }
491
492    /// This generates the instructions used to copy all parts of a single output vertex/primitive
493    /// to their individual output locations
494    fn write_mesh_copy_body(
495        &mut self,
496        is_primitive: bool,
497        return_info: &MeshReturnInfo,
498        index_var: u32,
499        vert_array_ptr: u32,
500        prim_array_ptr: u32,
501    ) -> Vec<Instruction> {
502        let u32_type_id = self.get_u32_type_id();
503        let mut body = Vec::new();
504        // Current index to copy
505        let val_i = self.id_gen.next();
506        body.push(Instruction::load(u32_type_id, val_i, index_var, None));
507
508        let info = if is_primitive {
509            &return_info.primitive_info
510        } else {
511            &return_info.vertex_info
512        };
513        let array_ptr = if is_primitive {
514            prim_array_ptr
515        } else {
516            vert_array_ptr
517        };
518
519        let mut builtin_index = 0;
520        let mut binding_index = 0;
521        // Write individual members of the vertex
522        for (member_id, member) in info.struct_members.iter().enumerate() {
523            let val_to_copy_ptr = self.id_gen.next();
524            body.push(Instruction::access_chain(
525                self.get_pointer_type_id(member.ty_id, spirv::StorageClass::Workgroup),
526                val_to_copy_ptr,
527                array_ptr,
528                &[
529                    val_i,
530                    self.get_constant_scalar(crate::Literal::U32(member_id as u32)),
531                ],
532            ));
533            let val_to_copy = self.id_gen.next();
534            body.push(Instruction::load(
535                member.ty_id,
536                val_to_copy,
537                val_to_copy_ptr,
538                None,
539            ));
540            let mut needs_y_flip = false;
541            let ptr_to_copy_to = self.id_gen.next();
542            // Get a pointer to the struct member to copy
543            match member.binding {
544                crate::Binding::BuiltIn(
545                    crate::BuiltIn::PointIndex
546                    | crate::BuiltIn::LineIndices
547                    | crate::BuiltIn::TriangleIndices,
548                ) => {
549                    body.push(Instruction::access_chain(
550                        self.get_pointer_type_id(member.ty_id, spirv::StorageClass::Output),
551                        ptr_to_copy_to,
552                        return_info.primitive_indices.unwrap(),
553                        &[val_i],
554                    ));
555                }
556                crate::Binding::BuiltIn(bi) => {
557                    body.push(Instruction::access_chain(
558                        self.get_pointer_type_id(member.ty_id, spirv::StorageClass::Output),
559                        ptr_to_copy_to,
560                        info.builtin_block.unwrap(),
561                        &[
562                            val_i,
563                            self.get_constant_scalar(crate::Literal::U32(builtin_index)),
564                        ],
565                    ));
566                    needs_y_flip = matches!(bi, crate::BuiltIn::Position { .. })
567                        && self.flags.contains(WriterFlags::ADJUST_COORDINATE_SPACE);
568                    builtin_index += 1;
569                }
570                crate::Binding::Location { .. } => {
571                    body.push(Instruction::access_chain(
572                        self.get_pointer_type_id(member.ty_id, spirv::StorageClass::Output),
573                        ptr_to_copy_to,
574                        info.bindings[binding_index],
575                        &[val_i],
576                    ));
577                    binding_index += 1;
578                }
579            }
580            body.push(Instruction::store(ptr_to_copy_to, val_to_copy, None));
581            // Flip the vertex position y coordinate in some cases
582            // Can't use epilogue flip because can't read from this storage class
583            if needs_y_flip {
584                let prev_y = self.id_gen.next();
585                body.push(Instruction::composite_extract(
586                    self.get_f32_type_id(),
587                    prev_y,
588                    val_to_copy,
589                    &[1],
590                ));
591                let new_y = self.id_gen.next();
592                body.push(Instruction::unary(
593                    spirv::Op::FNegate,
594                    self.get_f32_type_id(),
595                    new_y,
596                    prev_y,
597                ));
598                let new_ptr_to_copy_to = self.id_gen.next();
599                body.push(Instruction::access_chain(
600                    self.get_f32_pointer_type_id(spirv::StorageClass::Output),
601                    new_ptr_to_copy_to,
602                    ptr_to_copy_to,
603                    &[self.get_constant_scalar(crate::Literal::U32(1))],
604                ));
605                body.push(Instruction::store(new_ptr_to_copy_to, new_y, None));
606            }
607        }
608        body
609    }
610
611    /// Writes the return call for a mesh shader, which involves copying previously
612    /// written vertices/primitives into the actual output location.
613    pub(super) fn write_mesh_shader_return(
614        &mut self,
615        return_info: &MeshReturnInfo,
616        block: &mut Block,
617        loop_counter_vertices: u32,
618        loop_counter_primitives: u32,
619        local_invocation_index_id: Word,
620    ) -> Result<(), Error> {
621        let u32_id = self.get_u32_type_id();
622
623        // Load the actual vertex and primitive counts
624        let mut load_u32_by_member_index =
625            |members: &[MeshReturnMember], bi: crate::BuiltIn, max: u32| {
626                let member_index = members
627                    .iter()
628                    .position(|a| a.binding == crate::Binding::BuiltIn(bi))
629                    .unwrap() as u32;
630                let ptr_id = self.id_gen.next();
631                block.body.push(Instruction::access_chain(
632                    self.get_pointer_type_id(u32_id, spirv::StorageClass::Workgroup),
633                    ptr_id,
634                    return_info.out_variable_id,
635                    &[self.get_constant_scalar(crate::Literal::U32(member_index))],
636                ));
637                let before_min_id = self.id_gen.next();
638                block
639                    .body
640                    .push(Instruction::load(u32_id, before_min_id, ptr_id, None));
641
642                // Clamp the values
643                let id = self.id_gen.next();
644                block.body.push(Instruction::ext_inst_gl_op(
645                    self.gl450_ext_inst_id,
646                    spirv::GlslStd450Op::UMin,
647                    u32_id,
648                    id,
649                    &[before_min_id, max],
650                ));
651                id
652            };
653        let vert_count_id = load_u32_by_member_index(
654            &return_info.out_members,
655            crate::BuiltIn::VertexCount,
656            return_info.vertex_info.max_length_constant,
657        );
658        let prim_count_id = load_u32_by_member_index(
659            &return_info.out_members,
660            crate::BuiltIn::PrimitiveCount,
661            return_info.primitive_info.max_length_constant,
662        );
663
664        // Get pointers to the arrays of data to extract
665        let mut get_array_ptr = |bi: crate::BuiltIn, array_type_id: u32| {
666            let id = self.id_gen.next();
667            block.body.push(Instruction::access_chain(
668                self.get_pointer_type_id(array_type_id, spirv::StorageClass::Workgroup),
669                id,
670                return_info.out_variable_id,
671                &[self.get_constant_scalar(crate::Literal::U32(
672                    return_info
673                        .out_members
674                        .iter()
675                        .position(|a| a.binding == crate::Binding::BuiltIn(bi))
676                        .unwrap() as u32,
677                ))],
678            ));
679            id
680        };
681        let vert_array_ptr = get_array_ptr(
682            crate::BuiltIn::Vertices,
683            return_info.vertex_info.array_type_id,
684        );
685        let prim_array_ptr = get_array_ptr(
686            crate::BuiltIn::Primitives,
687            return_info.primitive_info.array_type_id,
688        );
689
690        // This must be called exactly once before any other mesh outputs are written
691        {
692            let mut ins = Instruction::new(spirv::Op::SetMeshOutputsEXT);
693            ins.add_operand(vert_count_id);
694            ins.add_operand(prim_count_id);
695            block.body.push(ins);
696        }
697
698        // This is iterating over every returned vertex and splitting
699        // it out into the multiple per-output arrays.
700        let vertex_loop_header = self.id_gen.next();
701        let prim_loop_header = self.id_gen.next();
702        let in_between_loops = self.id_gen.next();
703        let func_end = self.id_gen.next();
704
705        block.body.push(Instruction::store(
706            loop_counter_vertices,
707            local_invocation_index_id,
708            None,
709        ));
710        block.body.push(Instruction::branch(vertex_loop_header));
711
712        let vertex_copy_body = self.write_mesh_copy_body(
713            false,
714            return_info,
715            loop_counter_vertices,
716            vert_array_ptr,
717            prim_array_ptr,
718        );
719        // Write vertex copy loop
720        self.write_mesh_copy_loop(
721            &mut block.body,
722            vertex_copy_body,
723            vertex_loop_header,
724            in_between_loops,
725            vert_count_id,
726            loop_counter_vertices,
727            return_info,
728        );
729
730        // In between loops, reset the initial index
731        {
732            block.body.push(Instruction::label(in_between_loops));
733
734            block.body.push(Instruction::store(
735                loop_counter_primitives,
736                local_invocation_index_id,
737                None,
738            ));
739
740            block.body.push(Instruction::branch(prim_loop_header));
741        }
742        let primitive_copy_body = self.write_mesh_copy_body(
743            true,
744            return_info,
745            loop_counter_primitives,
746            vert_array_ptr,
747            prim_array_ptr,
748        );
749        // Write primitive copy loop
750        self.write_mesh_copy_loop(
751            &mut block.body,
752            primitive_copy_body,
753            prim_loop_header,
754            func_end,
755            prim_count_id,
756            loop_counter_primitives,
757            return_info,
758        );
759
760        block.body.push(Instruction::label(func_end));
761        Ok(())
762    }
763
764    pub(super) fn write_mesh_shader_wrapper(
765        &mut self,
766        return_info: &MeshReturnInfo,
767        inner_id: u32,
768    ) -> Result<u32, Error> {
769        let out_id = self.id_gen.next();
770        let mut function = super::Function::default();
771        let lookup_function_type = super::LookupFunctionType {
772            parameter_type_ids: alloc::vec![],
773            return_type_id: self.void_type,
774        };
775        let function_type = self.get_function_type(lookup_function_type);
776        function.signature = Some(Instruction::function(
777            self.void_type,
778            out_id,
779            spirv::FunctionControl::empty(),
780            function_type,
781        ));
782        let u32_id = self.get_u32_type_id();
783        {
784            let mut block = Block::new(self.id_gen.next());
785            // A general function variable that we guarantee to allow in the final return. It must be
786            // declared at the top of the function. Currently it is used in the memcpy part to keep
787            // track of the current index to copy.
788            let loop_counter_vertices = self.id_gen.next();
789            let loop_counter_primitives = self.id_gen.next();
790            block.body.insert(
791                0,
792                Instruction::variable(
793                    self.get_pointer_type_id(u32_id, spirv::StorageClass::Function),
794                    loop_counter_vertices,
795                    spirv::StorageClass::Function,
796                    None,
797                ),
798            );
799            block.body.insert(
800                1,
801                Instruction::variable(
802                    self.get_pointer_type_id(u32_id, spirv::StorageClass::Function),
803                    loop_counter_primitives,
804                    spirv::StorageClass::Function,
805                    None,
806                ),
807            );
808            let local_invocation_index_id = self.id_gen.next();
809            block.body.push(Instruction::load(
810                u32_id,
811                local_invocation_index_id,
812                return_info.local_invocation_index_var_id,
813                None,
814            ));
815            block.body.push(Instruction::function_call(
816                self.void_type,
817                self.id_gen.next(),
818                inner_id,
819                &[],
820            ));
821            self.write_control_barrier(crate::Barrier::WORK_GROUP, &mut block.body);
822            self.write_mesh_shader_return(
823                return_info,
824                &mut block,
825                loop_counter_vertices,
826                loop_counter_primitives,
827                local_invocation_index_id,
828            )?;
829            function.consume(block, Instruction::return_void());
830        }
831        function.to_words(&mut self.logical_layout.function_definitions);
832        Ok(out_id)
833    }
834
835    pub(super) fn write_task_shader_wrapper(
836        &mut self,
837        task_payload: Word,
838        inner_id: u32,
839    ) -> Result<u32, Error> {
840        let out_id = self.id_gen.next();
841        let mut function = super::Function::default();
842        let lookup_function_type = super::LookupFunctionType {
843            parameter_type_ids: alloc::vec![],
844            return_type_id: self.void_type,
845        };
846        let function_type = self.get_function_type(lookup_function_type);
847        function.signature = Some(Instruction::function(
848            self.void_type,
849            out_id,
850            spirv::FunctionControl::empty(),
851            function_type,
852        ));
853
854        {
855            let mut block = Block::new(self.id_gen.next());
856            let result = self.id_gen.next();
857            block.body.push(Instruction::function_call(
858                self.get_vec3u_type_id(),
859                result,
860                inner_id,
861                &[],
862            ));
863            self.write_control_barrier(crate::Barrier::WORK_GROUP, &mut block.body);
864            let final_value = if let Some(task_limits) = self.task_dispatch_limits {
865                let zero_u32 = self.get_constant_scalar(crate::Literal::U32(0));
866                let max_per_dim = self.get_constant_scalar(crate::Literal::U32(
867                    task_limits.max_mesh_workgroups_per_dim,
868                ));
869                let max_total = self.get_constant_scalar(crate::Literal::U32(
870                    task_limits.max_mesh_workgroups_total,
871                ));
872                let combined_struct_type = self.get_tuple_of_u32s_ty_id();
873                let values = [self.id_gen.next(), self.id_gen.next(), self.id_gen.next()];
874                for (i, value) in values.into_iter().enumerate() {
875                    block.body.push(Instruction::composite_extract(
876                        self.get_u32_type_id(),
877                        value,
878                        result,
879                        &[i as u32],
880                    ));
881                }
882                let prod_1 = self.id_gen.next();
883                let overflows = [self.id_gen.next(), self.id_gen.next()];
884                {
885                    let struct_out = self.id_gen.next();
886                    block.body.push(Instruction::binary(
887                        spirv::Op::UMulExtended,
888                        combined_struct_type,
889                        struct_out,
890                        values[0],
891                        values[1],
892                    ));
893                    block.body.push(Instruction::composite_extract(
894                        self.get_u32_type_id(),
895                        prod_1,
896                        struct_out,
897                        &[0],
898                    ));
899                    block.body.push(Instruction::composite_extract(
900                        self.get_u32_type_id(),
901                        overflows[0],
902                        struct_out,
903                        &[1],
904                    ));
905                }
906                let prod_final = self.id_gen.next();
907                {
908                    let struct_out = self.id_gen.next();
909                    block.body.push(Instruction::binary(
910                        spirv::Op::UMulExtended,
911                        combined_struct_type,
912                        struct_out,
913                        prod_1,
914                        values[2],
915                    ));
916                    block.body.push(Instruction::composite_extract(
917                        self.get_u32_type_id(),
918                        prod_final,
919                        struct_out,
920                        &[0],
921                    ));
922                    block.body.push(Instruction::composite_extract(
923                        self.get_u32_type_id(),
924                        overflows[1],
925                        struct_out,
926                        &[1],
927                    ));
928                }
929                let total_too_large = self.id_gen.next();
930                block.body.push(Instruction::binary(
931                    spirv::Op::UGreaterThan,
932                    self.get_bool_type_id(),
933                    total_too_large,
934                    prod_final,
935                    max_total,
936                ));
937
938                let too_large = [self.id_gen.next(), self.id_gen.next(), self.id_gen.next()];
939                for (i, value) in values.into_iter().enumerate() {
940                    block.body.push(Instruction::binary(
941                        spirv::Op::UGreaterThan,
942                        self.get_bool_type_id(),
943                        too_large[i],
944                        value,
945                        max_per_dim,
946                    ));
947                }
948                let overflow_happens = [self.id_gen.next(), self.id_gen.next()];
949                for (i, value) in overflows.into_iter().enumerate() {
950                    block.body.push(Instruction::binary(
951                        spirv::Op::INotEqual,
952                        self.get_bool_type_id(),
953                        overflow_happens[i],
954                        value,
955                        zero_u32,
956                    ));
957                }
958                let mut current_violates_limits = total_too_large;
959                for is_too_large in too_large {
960                    let new = self.id_gen.next();
961                    block.body.push(Instruction::binary(
962                        spirv::Op::LogicalOr,
963                        self.get_bool_type_id(),
964                        new,
965                        current_violates_limits,
966                        is_too_large,
967                    ));
968                    current_violates_limits = new;
969                }
970                for overflow_happens in overflow_happens {
971                    let new = self.id_gen.next();
972                    block.body.push(Instruction::binary(
973                        spirv::Op::LogicalOr,
974                        self.get_bool_type_id(),
975                        new,
976                        current_violates_limits,
977                        overflow_happens,
978                    ));
979                    current_violates_limits = new;
980                }
981                let zero_vec3 = self.id_gen.next();
982                block.body.push(Instruction::composite_construct(
983                    self.get_vec3u_type_id(),
984                    zero_vec3,
985                    &[zero_u32, zero_u32, zero_u32],
986                ));
987                let final_result = self.id_gen.next();
988                block.body.push(Instruction::select(
989                    self.get_vec3u_type_id(),
990                    final_result,
991                    current_violates_limits,
992                    zero_vec3,
993                    result,
994                ));
995                final_result
996            } else {
997                result
998            };
999            let ins =
1000                self.write_entry_point_task_return(final_value, &mut block.body, task_payload)?;
1001            function.consume(block, ins);
1002        }
1003        function.to_words(&mut self.logical_layout.function_definitions);
1004        Ok(out_id)
1005    }
1006}
naga/back/spv/mesh_shader.rs

naga/back/spv/
mesh_shader.rs
