[HLSL] Make ray queries safer.

CHANGELOG.md

@@ -63,6 +63,12 @@ Bottom level categories:
 
 - Expanded documentation of `QuerySet`, `QueryType`, and `resolve_query_set()` describing how to use queries. By @kpreid in [#8776](https://github.com/gfx-rs/wgpu/pull/8776).
 
+### Changes
+
+#### Naga
+
+- Prevent UB from incorrectly using ray queries on HLSL. By @Vecvec in [#8763](https://github.com/gfx-rs/wgpu/pull/8763).
+
 ## v28.0.0 (2025-12-17)
 
 ### Major Changes

naga/src/back/hlsl/keywords.rs

@@ -935,4 +935,6 @@ pub static RESERVED_CASE_INSENSITIVE_SET: RacyLock<CaseInsensitiveKeywordSet> =
 pub const RESERVED_PREFIXES: &[&str] = &[
     "__dynamic_buffer_offsets",
     super::help::IMAGE_STORAGE_LOAD_SCALAR_WRAPPER,
+    super::writer::RAY_QUERY_TRACKER_VARIABLE_PREFIX,
+    super::writer::INTERNAL_PREFIX,
 ];

naga/src/back/hlsl/mod.rs

@@ -543,6 +543,9 @@ pub struct Options {
     /// If set, loops will have code injected into them, forcing the compiler
     /// to think the number of iterations is bounded.
     pub force_loop_bounding: bool,
+    /// if set, ray queries will get a variable to track their state to prevent
+    /// misuse.
+    pub ray_query_initialization_tracking: bool,
 }
 
 impl Default for Options {
@@ -560,6 +563,7 @@ impl Default for Options {
             zero_initialize_workgroup_memory: true,
             restrict_indexing: true,
             force_loop_bounding: true,
+            ray_query_initialization_tracking: true,
         }
     }
 }

naga/src/back/hlsl/writer.rs

@@ -46,6 +46,9 @@ pub(crate) const F2U64_FUNCTION: &str = "naga_f2u64";
 pub(crate) const IMAGE_SAMPLE_BASE_CLAMP_TO_EDGE_FUNCTION: &str =
     "nagaTextureSampleBaseClampToEdge";
 pub(crate) const IMAGE_LOAD_EXTERNAL_FUNCTION: &str = "nagaTextureLoadExternal";
+pub(crate) const RAY_QUERY_TRACKER_VARIABLE_PREFIX: &str = "naga_query_init_tracker_for_";
+/// Prefix for variables in a naga statement
+pub(crate) const INTERNAL_PREFIX: &str = "naga_";
 
 enum Index {
     Expression(Handle<crate::Expression>),
@@ -1664,9 +1667,9 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 self.write_array_size(module, base, size)?;
             }
 
-            match module.types[local.ty].inner {
+            let is_ray_query = match module.types[local.ty].inner {
                 // from https://microsoft.github.io/DirectX-Specs/d3d/Raytracing.html#tracerayinline-example-1 it seems that ray queries shouldn't be zeroed
-                TypeInner::RayQuery { .. } => {}
+                TypeInner::RayQuery { .. } => true,
                 _ => {
                     write!(self.out, " = ")?;
                     // Write the local initializer if needed
@@ -1676,10 +1679,21 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                         // Zero initialize local variables
                         self.write_default_init(module, local.ty)?;
                     }
+                    false
                 }
-            }
+            };
             // Finish the local with `;` and add a newline (only for readability)
-            writeln!(self.out, ";")?
+            writeln!(self.out, ";")?;
+            // If it's a ray query, we also want a tracker variable
+            if is_ray_query {
+                write!(self.out, "{}", back::INDENT)?;
+                self.write_value_type(module, &TypeInner::Scalar(Scalar::U32))?;
+                writeln!(
+                    self.out,
+                    " {RAY_QUERY_TRACKER_VARIABLE_PREFIX}{} = 0;",
+                    self.names[&func_ctx.name_key(handle)]
+                )?;
+            }
         }
 
         if !func.local_variables.is_empty() {
@@ -2564,50 +2578,77 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
             } => {
                 self.write_switch(module, func_ctx, level, selector, cases)?;
             }
-            Statement::RayQuery { query, ref fun } => match *fun {
-                RayQueryFunction::Initialize {
-                    acceleration_structure,
-                    descriptor,
-                } => {
-                    write!(self.out, "{level}")?;
-                    self.write_expr(module, query, func_ctx)?;
-                    write!(self.out, ".TraceRayInline(")?;
-                    self.write_expr(module, acceleration_structure, func_ctx)?;
-                    write!(self.out, ", ")?;
-                    self.write_expr(module, descriptor, func_ctx)?;
-                    write!(self.out, ".flags, ")?;
-                    self.write_expr(module, descriptor, func_ctx)?;
-                    write!(self.out, ".cull_mask, ")?;
-                    write!(self.out, "RayDescFromRayDesc_(")?;
-                    self.write_expr(module, descriptor, func_ctx)?;
-                    writeln!(self.out, "));")?;
-                }
-                RayQueryFunction::Proceed { result } => {
-                    write!(self.out, "{level}")?;
-                    let name = Baked(result).to_string();
-                    write!(self.out, "const bool {name} = ")?;
-                    self.named_expressions.insert(result, name);
-                    self.write_expr(module, query, func_ctx)?;
-                    writeln!(self.out, ".Proceed();")?;
-                }
-                RayQueryFunction::GenerateIntersection { hit_t } => {
-                    write!(self.out, "{level}")?;
-                    self.write_expr(module, query, func_ctx)?;
-                    write!(self.out, ".CommitProceduralPrimitiveHit(")?;
-                    self.write_expr(module, hit_t, func_ctx)?;
-                    writeln!(self.out, ");")?;
-                }
-                RayQueryFunction::ConfirmIntersection => {
-                    write!(self.out, "{level}")?;
-                    self.write_expr(module, query, func_ctx)?;
-                    writeln!(self.out, ".CommitNonOpaqueTriangleHit();")?;
-                }
-                RayQueryFunction::Terminate => {
-                    write!(self.out, "{level}")?;
-                    self.write_expr(module, query, func_ctx)?;
-                    writeln!(self.out, ".Abort();")?;
+            Statement::RayQuery { query, ref fun } => {
+                // There are three possibilities for a ptr to be:
+                // 1. A variable
+                // 2. A function argument
+                // 3. part of a struct
+                //
+                // 2 and 3 are not possible, a ray query (in naga IR)
+                // is not allowed to be passed into a function, and
+                // all languages disallow it in a struct (you get fun results if
+                // you try it :) ).
+                //
+                // Therefore, the ray query expression must be a variable.
+                let crate::Expression::LocalVariable(query_var) = func_ctx.expressions[query]
+                else {
+                    unreachable!()
+                };
+
+                let tracker_expr_name = format!(
+                    "{RAY_QUERY_TRACKER_VARIABLE_PREFIX}{}",
+                    self.names[&func_ctx.name_key(query_var)]
+                );
+
+                match *fun {
+                    RayQueryFunction::Initialize {
+                        acceleration_structure,
+                        descriptor,
+                    } => {
+                        self.write_initialize_function(
+                            module,
+                            level,
+                            query,
+                            acceleration_structure,
+                            descriptor,
+                            &tracker_expr_name,
+                            func_ctx,
+                        )?;
+                    }
+                    RayQueryFunction::Proceed { result } => {
+                        self.write_proceed(
+                            module,
+                            level,
+                            query,
+                            result,
+                            &tracker_expr_name,
+                            func_ctx,
+                        )?;
+                    }
+                    RayQueryFunction::GenerateIntersection { hit_t } => {
+                        self.write_generate_intersection(
+                            module,
+                            level,
+                            query,
+                            hit_t,
+                            &tracker_expr_name,
+                            func_ctx,
+                        )?;
+                    }
+                    RayQueryFunction::ConfirmIntersection => {
+                        self.write_confirm_intersection(
+                            module,
+                            level,
+                            query,
+                            &tracker_expr_name,
+                            func_ctx,
+                        )?;
+                    }
+                    RayQueryFunction::Terminate => {
+                        self.write_terminate(module, level, query, &tracker_expr_name, func_ctx)?;
+                    }
                 }
-            },
+            }
             Statement::SubgroupBallot { result, predicate } => {
                 write!(self.out, "{level}")?;
                 let name = Baked(result).to_string();
@@ -4275,14 +4316,24 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 write!(self.out, ")")?
             }
             Expression::RayQueryGetIntersection { query, committed } => {
+                // For reasoning, see write_stmt
+                let Expression::LocalVariable(query_var) = func_ctx.expressions[query] else {
+                    unreachable!()
+                };
+
+                let tracker_expr_name = format!(
+                    "{RAY_QUERY_TRACKER_VARIABLE_PREFIX}{}",
+                    self.names[&func_ctx.name_key(query_var)]
+                );
+
                 if committed {
                     write!(self.out, "GetCommittedIntersection(")?;
                     self.write_expr(module, query, func_ctx)?;
-                    write!(self.out, ")")?;
+                    write!(self.out, ", {tracker_expr_name})")?;
                 } else {
                     write!(self.out, "GetCandidateIntersection(")?;
                     self.write_expr(module, query, func_ctx)?;
-                    write!(self.out, ")")?;
+                    write!(self.out, ", {tracker_expr_name})")?;
                 }
             }
             // Not supported yet

naga/src/back/mod.rs

@@ -61,6 +61,26 @@ impl core::fmt::Display for Baked {
     }
 }
 
+bitflags::bitflags! {
+    /// How far through a ray query are we
+    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
+    #[cfg_attr(
+        not(any(hlsl_out, spv_out)),
+        allow(
+            dead_code,
+            reason = "shared helpers can be dead if none of the enabled backends need it"
+        )
+    )]
+    pub(super) struct RayQueryPoint: u32 {
+        /// Ray query has been successfully initialized.
+        const INITIALIZED = 1 << 0;
+        /// Proceed has been called on ray query.
+        const PROCEED = 1 << 1;
+        /// Proceed has returned false (have finished traversal).
+        const FINISHED_TRAVERSAL = 1 << 2;
+    }
+}
+
 /// Specifies the values of pipeline-overridable constants in the shader module.
 ///
 /// If an `@id` attribute was specified on the declaration,

naga/src/back/spv/mod.rs

@@ -895,19 +895,6 @@ bitflags::bitflags! {
     }
 }
 
-bitflags::bitflags! {
-    /// How far through a ray query are we
-    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
-    pub(super) struct RayQueryPoint: u32 {
-        /// Ray query has been successfully initialized.
-        const INITIALIZED = 1 << 0;
-        /// Proceed has been called on ray query.
-        const PROCEED = 1 << 1;
-        /// Proceed has returned false (have finished traversal).
-        const FINISHED_TRAVERSAL = 1 << 2;
-    }
-}
-
 #[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash)]
 #[cfg_attr(feature = "serialize", derive(serde::Serialize))]
 #[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]