[spirv] Make ray queries safer

CHANGELOG.md

@@ -90,6 +90,10 @@ SamplerDescriptor {
 - Validation errors from `CommandEncoder::finish()` will report the label of the invalid encoder. By @kpreid in [#8449](https://github.com/gfx-rs/wgpu/pull/8449).
 - Corrected documentation of the minimum alignment of the *end* of a mapped range of a buffer (it is 4, not 8). By @kpreid in [#8450](https://github.com/gfx-rs/wgpu/pull/8450).
 
+#### Naga
+
+- Prevent UB with invalid ray query calls on spirv. By @Vecvec in [#8390](https://github.com/gfx-rs/wgpu/pull/8390). 
+
 ### Bug Fixes
 
 #### naga

naga-test/src/lib.rs

@@ -114,6 +114,7 @@ pub struct SpirvOutParameters {
     pub separate_entry_points: bool,
     #[serde(deserialize_with = "deserialize_binding_map")]
     pub binding_map: naga::back::spv::BindingMap,
+    pub ray_query_initialization_tracking: bool,
     pub use_storage_input_output_16: bool,
 }
 impl Default for SpirvOutParameters {
@@ -126,6 +127,7 @@ impl Default for SpirvOutParameters {
             force_point_size: false,
             clamp_frag_depth: false,
             separate_entry_points: false,
+            ray_query_initialization_tracking: true,
             use_storage_input_output_16: true,
             binding_map: naga::back::spv::BindingMap::default(),
         }
@@ -159,6 +161,7 @@ impl SpirvOutParameters {
             binding_map: self.binding_map.clone(),
             zero_initialize_workgroup_memory: spv::ZeroInitializeWorkgroupMemoryMode::Polyfill,
             force_loop_bounding: true,
+            ray_query_initialization_tracking: true,
             debug_info,
             use_storage_input_output_16: self.use_storage_input_output_16,
         }

naga/src/back/spv/block.rs

@@ -203,7 +203,7 @@ impl Writer {
         ));
 
         let clamp_id = self.id_gen.next();
-        body.push(Instruction::ext_inst(
+        body.push(Instruction::ext_inst_gl_op(
             self.gl450_ext_inst_id,
             spirv::GLOp::FClamp,
             float_type_id,
@@ -1026,15 +1026,15 @@ impl BlockContext<'_> {
                             };
 
                             let max_id = self.gen_id();
-                            block.body.push(Instruction::ext_inst(
+                            block.body.push(Instruction::ext_inst_gl_op(
                                 self.writer.gl450_ext_inst_id,
                                 max_op,
                                 result_type_id,
                                 max_id,
                                 &[arg0_id, arg1_id],
                             ));
 
-                            MathOp::Custom(Instruction::ext_inst(
+                            MathOp::Custom(Instruction::ext_inst_gl_op(
                                 self.writer.gl450_ext_inst_id,
                                 min_op,
                                 result_type_id,
@@ -1068,7 +1068,7 @@ impl BlockContext<'_> {
                             arg2_id = self.writer.get_constant_composite(ty, &self.temp_list);
                         }
 
-                        MathOp::Custom(Instruction::ext_inst(
+                        MathOp::Custom(Instruction::ext_inst_gl_op(
                             self.writer.gl450_ext_inst_id,
                             spirv::GLOp::FClamp,
                             result_type_id,
@@ -1282,7 +1282,7 @@ impl BlockContext<'_> {
                                     &self.temp_list,
                                 ));
 
-                                MathOp::Custom(Instruction::ext_inst(
+                                MathOp::Custom(Instruction::ext_inst_gl_op(
                                     self.writer.gl450_ext_inst_id,
                                     spirv::GLOp::FMix,
                                     result_type_id,
@@ -1339,15 +1339,15 @@ impl BlockContext<'_> {
                         };
 
                         let lsb_id = self.gen_id();
-                        block.body.push(Instruction::ext_inst(
+                        block.body.push(Instruction::ext_inst_gl_op(
                             self.writer.gl450_ext_inst_id,
                             spirv::GLOp::FindILsb,
                             result_type_id,
                             lsb_id,
                             &[arg0_id],
                         ));
 
-                        MathOp::Custom(Instruction::ext_inst(
+                        MathOp::Custom(Instruction::ext_inst_gl_op(
                             self.writer.gl450_ext_inst_id,
                             spirv::GLOp::UMin,
                             result_type_id,
@@ -1388,7 +1388,7 @@ impl BlockContext<'_> {
                         };
 
                         let msb_id = self.gen_id();
-                        block.body.push(Instruction::ext_inst(
+                        block.body.push(Instruction::ext_inst_gl_op(
                             self.writer.gl450_ext_inst_id,
                             if width != 4 {
                                 spirv::GLOp::FindILsb
@@ -1445,7 +1445,7 @@ impl BlockContext<'_> {
 
                         // o = min(offset, w)
                         let offset_id = self.gen_id();
-                        block.body.push(Instruction::ext_inst(
+                        block.body.push(Instruction::ext_inst_gl_op(
                             self.writer.gl450_ext_inst_id,
                             spirv::GLOp::UMin,
                             u32_type,
@@ -1465,7 +1465,7 @@ impl BlockContext<'_> {
 
                         // c = min(count, tmp)
                         let count_id = self.gen_id();
-                        block.body.push(Instruction::ext_inst(
+                        block.body.push(Instruction::ext_inst_gl_op(
                             self.writer.gl450_ext_inst_id,
                             spirv::GLOp::UMin,
                             u32_type,
@@ -1495,7 +1495,7 @@ impl BlockContext<'_> {
 
                         // o = min(offset, w)
                         let offset_id = self.gen_id();
-                        block.body.push(Instruction::ext_inst(
+                        block.body.push(Instruction::ext_inst_gl_op(
                             self.writer.gl450_ext_inst_id,
                             spirv::GLOp::UMin,
                             u32_type,
@@ -1515,7 +1515,7 @@ impl BlockContext<'_> {
 
                         // c = min(count, tmp)
                         let count_id = self.gen_id();
-                        block.body.push(Instruction::ext_inst(
+                        block.body.push(Instruction::ext_inst_gl_op(
                             self.writer.gl450_ext_inst_id,
                             spirv::GLOp::UMin,
                             u32_type,
@@ -1610,7 +1610,7 @@ impl BlockContext<'_> {
                 };
 
                 block.body.push(match math_op {
-                    MathOp::Ext(op) => Instruction::ext_inst(
+                    MathOp::Ext(op) => Instruction::ext_inst_gl_op(
                         self.writer.gl450_ext_inst_id,
                         op,
                         result_type_id,
@@ -1621,7 +1621,27 @@ impl BlockContext<'_> {
                 });
                 id
             }
-            crate::Expression::LocalVariable(variable) => self.function.variables[&variable].id,
+            crate::Expression::LocalVariable(variable) => {
+                if let Some(rq_tracker) = self
+                    .function
+                    .ray_query_initialization_tracker_variables
+                    .get(&variable)
+                {
+                    self.ray_query_tracker_expr.insert(
+                        expr_handle,
+                        super::RayQueryTrackers {
+                            initialized_tracker: rq_tracker.id,
+                            t_max_tracker: self
+                                .function
+                                .ray_query_t_max_tracker_variables
+                                .get(&variable)
+                                .expect("Both trackers are set at the same time.")
+                                .id,
+                        },
+                    );
+                }
+                self.function.variables[&variable].id
+            }
             crate::Expression::Load { pointer } => {
                 self.write_checked_load(pointer, block, AccessTypeAdjustment::None, result_type_id)?
             }
@@ -1772,6 +1792,10 @@ impl BlockContext<'_> {
             crate::Expression::ArrayLength(expr) => self.write_runtime_array_length(expr, block)?,
             crate::Expression::RayQueryGetIntersection { query, committed } => {
                 let query_id = self.cached[query];
+                let init_tracker_id = *self
+                    .ray_query_tracker_expr
+                    .get(&query)
+                    .expect("not a cached ray query");
                 let func_id = self
                     .writer
                     .write_ray_query_get_intersection_function(committed, self.ir_module);
@@ -1782,7 +1806,7 @@ impl BlockContext<'_> {
                     intersection_type_id,
                     id,
                     func_id,
-                    &[query_id],
+                    &[query_id, init_tracker_id.initialized_tracker],
                 ));
                 id
             }
@@ -2008,7 +2032,7 @@ impl BlockContext<'_> {
                 let max_const_id = maybe_splat_const(self.writer, max_const_id);
 
                 let clamp_id = self.gen_id();
-                block.body.push(Instruction::ext_inst(
+                block.body.push(Instruction::ext_inst_gl_op(
                     self.writer.gl450_ext_inst_id,
                     spirv::GLOp::FClamp,
                     expr_type_id,
@@ -2671,7 +2695,7 @@ impl BlockContext<'_> {
             });
 
             let clamp_id = self.gen_id();
-            block.body.push(Instruction::ext_inst(
+            block.body.push(Instruction::ext_inst_gl_op(
                 self.writer.gl450_ext_inst_id,
                 clamp_op,
                 wide_vector_type_id,
@@ -2765,7 +2789,7 @@ impl BlockContext<'_> {
                 let [min, max] = [min, max].map(|lit| self.writer.get_constant_scalar(lit));
 
                 let clamp_id = self.gen_id();
-                block.body.push(Instruction::ext_inst(
+                block.body.push(Instruction::ext_inst_gl_op(
                     self.writer.gl450_ext_inst_id,
                     clamp_op,
                     result_type_id,

naga/src/back/spv/image.rs

@@ -446,7 +446,7 @@ impl BlockContext<'_> {
         // and negative values in a single instruction: negative values of
         // `input_id` get treated as very large positive values.
         let restricted_id = self.gen_id();
-        block.body.push(Instruction::ext_inst(
+        block.body.push(Instruction::ext_inst_gl_op(
             self.writer.gl450_ext_inst_id,
             spirv::GLOp::UMin,
             type_id,
@@ -580,7 +580,7 @@ impl BlockContext<'_> {
         // and negative values in a single instruction: negative values of
         // `coordinates` get treated as very large positive values.
         let restricted_coordinates_id = self.gen_id();
-        block.body.push(Instruction::ext_inst(
+        block.body.push(Instruction::ext_inst_gl_op(
             self.writer.gl450_ext_inst_id,
             spirv::GLOp::UMin,
             coordinates.type_id,
@@ -923,7 +923,7 @@ impl BlockContext<'_> {
 
             // Clamp the coords to the calculated margins
             let clamped_coords_id = self.gen_id();
-            block.body.push(Instruction::ext_inst(
+            block.body.push(Instruction::ext_inst_gl_op(
                 self.writer.gl450_ext_inst_id,
                 spirv::GLOp::NClamp,
                 vec2f_type_id,

naga/src/back/spv/index.rs

@@ -366,7 +366,7 @@ impl BlockContext<'_> {
         // One or the other of the index or length is dynamic, so emit code for
         // BoundsCheckPolicy::Restrict.
         let restricted_index_id = self.gen_id();
-        block.body.push(Instruction::ext_inst(
+        block.body.push(Instruction::ext_inst_gl_op(
             self.writer.gl450_ext_inst_id,
             spirv::GLOp::UMin,
             self.writer.get_u32_type_id(),

naga/src/back/spv/instructions.rs

@@ -156,18 +156,28 @@ impl super::Instruction {
         instruction
     }
 
-    pub(super) fn ext_inst(
+    pub(super) fn ext_inst_gl_op(
         set_id: Word,
         op: spirv::GLOp,
         result_type_id: Word,
         id: Word,
         operands: &[Word],
+    ) -> Self {
+        Self::ext_inst(set_id, op as u32, result_type_id, id, operands)
+    }
+
+    pub(super) fn ext_inst(
+        set_id: Word,
+        op: u32,
+        result_type_id: Word,
+        id: Word,
+        operands: &[Word],
     ) -> Self {
         let mut instruction = Self::new(Op::ExtInst);
         instruction.set_type(result_type_id);
         instruction.set_result(id);
         instruction.add_operand(set_id);
-        instruction.add_operand(op as u32);
+        instruction.add_operand(op);
         for operand in operands {
             instruction.add_operand(*operand)
         }
@@ -824,6 +834,14 @@ impl super::Instruction {
         instruction
     }
 
+    pub(super) fn ray_query_get_t_min(result_type_id: Word, id: Word, query: Word) -> Self {
+        let mut instruction = Self::new(Op::RayQueryGetRayTMinKHR);
+        instruction.set_type(result_type_id);
+        instruction.set_result(id);
+        instruction.add_operand(query);
+        instruction
+    }
+
     //
     //  Conversion Instructions
     //