Implement ray query candidate intersection generation and confirmation

CHANGELOG.md

@@ -45,6 +45,7 @@ Bottom level categories:
 #### Naga
 
 - Support @must_use attribute on function declarations. By @turbocrime in [#6801](https://github.com/gfx-rs/wgpu/pull/6801).
+- Support for generating the candidate intersections from AABB geometry, and confirming the hits. By @kvark in [#7047](https://github.com/gfx-rs/wgpu/pull/7047).
 
 ### Changes
 

etc/specs/ray_tracing.md

@@ -2,23 +2,23 @@
 
 🧪Experimental🧪
 
-`wgpu` supports an experimental version of ray tracing which is subject to change. The extensions allow for acceleration structures to be created and built (with 
+`wgpu` supports an experimental version of ray tracing which is subject to change. The extensions allow for acceleration structures to be created and built (with
 `Features::EXPERIMENTAL_RAY_TRACING_ACCELERATION_STRUCTURE` enabled) and interacted with in shaders. Currently `naga` only supports ray queries
 (accessible with `Features::EXPERIMENTAL_RAY_QUERY` enabled in wgpu).
 
 **Note**: The features documented here may have major bugs in them and are expected to be subject
 to breaking changes, suggestions for the API exposed by this should be posted on [the ray-tracing issue](https://github.com/gfx-rs/wgpu/issues/1040).
 Large changes may mean that this documentation may be out of date.
 
-***This is not*** an introduction to raytracing, and assumes basic prior knowledge, to look at the fundamentals look at 
+***This is not*** an introduction to raytracing, and assumes basic prior knowledge, to look at the fundamentals look at
 an [introduction](https://developer.nvidia.com/blog/introduction-nvidia-rtx-directx-ray-tracing/).
 
 ## `wgpu`'s raytracing API:
 
 The documentation and specific details of the functions and structures provided
-can be found with their definitions.  
+can be found with their definitions.
 
-A [`Blas`] can be created with [`Device::create_blas`].  
+A [`Blas`] can be created with [`Device::create_blas`].
 A [`Tlas`] can be created with [`Device::create_tlas`].
 
 Unless one is planning on using the unsafe building API (not recommended for beginners) a [`Tlas`] should be put inside
@@ -57,12 +57,21 @@ rayQueryInitialize(rq: ptr<function, ray_query>, acceleration_structure: acceler
 // - The hit is considered `Candidate` if this function returns true, and the hit is considered `Committed` if
 //   this function returns false.
 // - A `Candidate` intersection interrupts the ray traversal.
-// - A `Candidate` intersection may happen anywhere along the ray, it should not be relied on to give the closest hit. A 
+// - A `Candidate` intersection may happen anywhere along the ray, it should not be relied on to give the closest hit. A
 //   `Candidate` intersection is to allow the user themselves to decide if that intersection is valid*. If one wants to get
 //   the closest hit a `Committed` intersection should be used.
 // - Calling this function multiple times will cause the ray traversal to continue if it was interrupted by a `Candidate`
 //   intersection.
-rayQueryProceed(rq: ptr<function, ray_query>) -> bool`
+rayQueryProceed(rq: ptr<function, ray_query>) -> bool
+
+// - Generates a hit from procedural geometry at a particular distance.
+rayQueryGenerateIntersection(hit_t: f32)
+
+// - Commits a hit from triangular non-opaque geometry.
+rayQueryConfirmIntersection()
+
+// - Aborts the query.
+rayQueryTerminate()
 
 // - Returns intersection details about a hit considered `Committed`.
 rayQueryGetCommittedIntersection(rq: ptr<function, ray_query>) -> RayIntersection
@@ -71,19 +80,24 @@ rayQueryGetCommittedIntersection(rq: ptr<function, ray_query>) -> RayIntersectio
 rayQueryGetCandidateIntersection(rq: ptr<function, ray_query>) -> RayIntersection
 ```
 
-*The API to commit a candidate intersection is not yet implemented but would be possible to be user implemented.
-
 > [!CAUTION]
-> 
+>
 > ### ⚠️Undefined behavior ⚠️:
 > - Calling `rayQueryGetCommittedIntersection` or `rayQueryGetCandidateIntersection` when `rayQueryProceed` has not been
 > called on this ray query since it was initialized (or if the ray query has not been previously initialized).
 > - Calling `rayQueryGetCommittedIntersection` when `rayQueryProceed`'s latest return on this ray query is considered
->   `Candidate`. 
+>   `Candidate`.
 > - Calling `rayQueryGetCandidateIntersection` when `rayQueryProceed`'s latest return on this ray query is considered
 >   `Committed`.
 > - Calling `rayQueryProceed` when `rayQueryInitialize` has not previously been called on this ray query
-> 
+> - Calling `rayQueryGenerateIntersection` on a query with last intersection kind not being
+>   `RAY_QUERY_INTERSECTION_AABB`,
+> - Calling `rayQueryGenerateIntersection` with `hit_t` outside of `RayDesc::t_min .. RayDesc::t_max` range.
+>   or when `rayQueryProceed`'s latest return on this ray query is not considered `Candidate`.
+> - Calling `rayQueryConfirmIntersection` on a query with last intersection kind not being
+>   `RAY_QUERY_INTERSECTION_TRIANGLE`,
+>   or when `rayQueryProceed`'s latest return on this ray query is not considered `Candidate`.
+>
 > *this is only known undefined behaviour, and will be worked around in the future.
 
 ```wgsl

naga/src/back/dot/mod.rs

@@ -293,6 +293,13 @@ impl StatementGraph {
                             self.emits.push((id, result));
                             "RayQueryProceed"
                         }
+                        crate::RayQueryFunction::GenerateIntersection { hit_t } => {
+                            self.dependencies.push((id, hit_t, "hit_t"));
+                            "RayQueryGenerateIntersection"
+                        }
+                        crate::RayQueryFunction::ConfirmIntersection => {
+                            "RayQueryConfirmIntersection"
+                        }
                         crate::RayQueryFunction::Terminate => "RayQueryTerminate",
                     }
                 }

naga/src/back/hlsl/writer.rs

@@ -2452,7 +2452,20 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                     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();")?;
                 }

naga/src/back/msl/writer.rs

@@ -30,6 +30,7 @@ const RT_NAMESPACE: &str = "metal::raytracing";
 const RAY_QUERY_TYPE: &str = "_RayQuery";
 const RAY_QUERY_FIELD_INTERSECTOR: &str = "intersector";
 const RAY_QUERY_FIELD_INTERSECTION: &str = "intersection";
+const RAY_QUERY_MODERN_SUPPORT: bool = false; //TODO
 const RAY_QUERY_FIELD_READY: &str = "ready";
 const RAY_QUERY_FUN_MAP_INTERSECTION: &str = "_map_intersection_type";
 
@@ -3506,16 +3507,41 @@ impl<W: Write> Writer<W> {
                             self.named_expressions.insert(result, name);
                             self.put_expression(query, &context.expression, true)?;
                             writeln!(self.out, ".{RAY_QUERY_FIELD_READY};")?;
-                            //TODO: actually proceed?
-
-                            write!(self.out, "{level}")?;
-                            self.put_expression(query, &context.expression, true)?;
-                            writeln!(self.out, ".{RAY_QUERY_FIELD_READY} = false;")?;
+                            if RAY_QUERY_MODERN_SUPPORT {
+                                write!(self.out, "{level}")?;
+                                self.put_expression(query, &context.expression, true)?;
+                                writeln!(self.out, ".?.next();")?;
+                            }
+                        }
+                        crate::RayQueryFunction::GenerateIntersection { hit_t } => {
+                            if RAY_QUERY_MODERN_SUPPORT {
+                                write!(self.out, "{level}")?;
+                                self.put_expression(query, &context.expression, true)?;
+                                write!(self.out, ".?.commit_bounding_box_intersection(")?;
+                                self.put_expression(hit_t, &context.expression, true)?;
+                                writeln!(self.out, ");")?;
+                            } else {
+                                log::warn!("Ray Query GenerateIntersection is not yet supported");
+                            }
+                        }
+                        crate::RayQueryFunction::ConfirmIntersection => {
+                            if RAY_QUERY_MODERN_SUPPORT {
+                                write!(self.out, "{level}")?;
+                                self.put_expression(query, &context.expression, true)?;
+                                writeln!(self.out, ".?.commit_triangle_intersection();")?;
+                            } else {
+                                log::warn!("Ray Query ConfirmIntersection is not yet supported");
+                            }
                         }
                         crate::RayQueryFunction::Terminate => {
+                            if RAY_QUERY_MODERN_SUPPORT {
+                                write!(self.out, "{level}")?;
+                                self.put_expression(query, &context.expression, true)?;
+                                writeln!(self.out, ".?.abort();")?;
+                            }
                             write!(self.out, "{level}")?;
                             self.put_expression(query, &context.expression, true)?;
-                            writeln!(self.out, ".{RAY_QUERY_FIELD_INTERSECTION}.abort();")?;
+                            writeln!(self.out, ".{RAY_QUERY_FIELD_READY} = false;")?;
                         }
                     }
                 }

naga/src/back/pipeline_constants.rs

@@ -821,6 +821,10 @@ fn adjust_stmt(new_pos: &HandleVec<Expression, Handle<Expression>>, stmt: &mut S
                 crate::RayQueryFunction::Proceed { ref mut result } => {
                     adjust(result);
                 }
+                crate::RayQueryFunction::GenerateIntersection { ref mut hit_t } => {
+                    adjust(hit_t);
+                }
+                crate::RayQueryFunction::ConfirmIntersection => {}
                 crate::RayQueryFunction::Terminate => {}
             }
         }

naga/src/back/spv/instructions.rs

@@ -779,6 +779,19 @@ impl super::Instruction {
         instruction
     }
 
+    pub(super) fn ray_query_generate_intersection(query: Word, hit: Word) -> Self {
+        let mut instruction = Self::new(Op::RayQueryGenerateIntersectionKHR);
+        instruction.add_operand(query);
+        instruction.add_operand(hit);
+        instruction
+    }
+
+    pub(super) fn ray_query_confirm_intersection(query: Word) -> Self {
+        let mut instruction = Self::new(Op::RayQueryConfirmIntersectionKHR);
+        instruction.add_operand(query);
+        instruction
+    }
+
     pub(super) fn ray_query_get_intersection(
         op: Op,
         result_type_id: Word,

naga/src/back/spv/ray.rs

@@ -609,6 +609,19 @@ impl BlockContext<'_> {
                     .body
                     .push(Instruction::ray_query_proceed(result_type_id, id, query_id));
             }
+            crate::RayQueryFunction::GenerateIntersection { hit_t } => {
+                let hit_id = self.cached[hit_t];
+                block
+                    .body
+                    .push(Instruction::ray_query_generate_intersection(
+                        query_id, hit_id,
+                    ));
+            }
+            crate::RayQueryFunction::ConfirmIntersection => {
+                block
+                    .body
+                    .push(Instruction::ray_query_confirm_intersection(query_id));
+            }
             crate::RayQueryFunction::Terminate => {}
         }
     }

naga/src/compact/statements.rs

@@ -190,6 +190,10 @@ impl FunctionTracer<'_> {
             Qf::Proceed { result } => {
                 self.expressions_used.insert(result);
             }
+            Qf::GenerateIntersection { hit_t } => {
+                self.expressions_used.insert(hit_t);
+            }
+            Qf::ConfirmIntersection => {}
             Qf::Terminate => {}
         }
     }
@@ -393,6 +397,10 @@ impl FunctionMap {
             Qf::Proceed { ref mut result } => {
                 self.expressions.adjust(result);
             }
+            Qf::GenerateIntersection { ref mut hit_t } => {
+                self.expressions.adjust(hit_t);
+            }
+            Qf::ConfirmIntersection => {}
             Qf::Terminate => {}
         }
     }

naga/src/front/wgsl/lower/mod.rs

@@ -2692,6 +2692,40 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                                 .push(crate::Statement::RayQuery { query, fun }, span);
                             return Ok(Some(result));
                         }
+                        "rayQueryGenerateIntersection" => {
+                            let mut args = ctx.prepare_args(arguments, 2, span);
+                            let query = self.ray_query_pointer(args.next()?, ctx)?;
+                            let hit_t = self.expression(args.next()?, ctx)?;
+                            args.finish()?;
+
+                            let fun = crate::RayQueryFunction::GenerateIntersection { hit_t };
+                            let rctx = ctx.runtime_expression_ctx(span)?;
+                            rctx.block
+                                .push(crate::Statement::RayQuery { query, fun }, span);
+                            return Ok(None);
+                        }
+                        "rayQueryConfirmIntersection" => {
+                            let mut args = ctx.prepare_args(arguments, 1, span);
+                            let query = self.ray_query_pointer(args.next()?, ctx)?;
+                            args.finish()?;
+
+                            let fun = crate::RayQueryFunction::ConfirmIntersection;
+                            let rctx = ctx.runtime_expression_ctx(span)?;
+                            rctx.block
+                                .push(crate::Statement::RayQuery { query, fun }, span);
+                            return Ok(None);
+                        }
+                        "rayQueryTerminate" => {
+                            let mut args = ctx.prepare_args(arguments, 1, span);
+                            let query = self.ray_query_pointer(args.next()?, ctx)?;
+                            args.finish()?;
+
+                            let fun = crate::RayQueryFunction::Terminate;
+                            let rctx = ctx.runtime_expression_ctx(span)?;
+                            rctx.block
+                                .push(crate::Statement::RayQuery { query, fun }, span);
+                            return Ok(None);
+                        }
                         "rayQueryGetCommittedIntersection" => {
                             let mut args = ctx.prepare_args(arguments, 1, span);
                             let query = self.ray_query_pointer(args.next()?, ctx)?;

naga/src/lib.rs

@@ -1779,6 +1779,16 @@ pub enum RayQueryFunction {
         result: Handle<Expression>,
     },
 
+    /// Add a candidate generated intersection to be included
+    /// in the determination of the closest hit for a ray query.
+    GenerateIntersection {
+        hit_t: Handle<Expression>,
+    },
+
+    /// Confirm a triangle intersection to be included in the determination of
+    /// the closest hit for a ray query.
+    ConfirmIntersection,
+
     Terminate,
 }
 

naga/src/valid/analyzer.rs

@@ -1080,13 +1080,20 @@ impl FunctionInfo {
                 }
                 S::RayQuery { query, ref fun } => {
                     let _ = self.add_ref(query);
-                    if let crate::RayQueryFunction::Initialize {
-                        acceleration_structure,
-                        descriptor,
-                    } = *fun
-                    {
-                        let _ = self.add_ref(acceleration_structure);
-                        let _ = self.add_ref(descriptor);
+                    match *fun {
+                        crate::RayQueryFunction::Initialize {
+                            acceleration_structure,
+                            descriptor,
+                        } => {
+                            let _ = self.add_ref(acceleration_structure);
+                            let _ = self.add_ref(descriptor);
+                        }
+                        crate::RayQueryFunction::Proceed { result: _ } => {}
+                        crate::RayQueryFunction::GenerateIntersection { hit_t } => {
+                            let _ = self.add_ref(hit_t);
+                        }
+                        crate::RayQueryFunction::ConfirmIntersection => {}
+                        crate::RayQueryFunction::Terminate => {}
                     }
                     FunctionUniformity::new()
                 }

naga/src/valid/function.rs

@@ -177,6 +177,8 @@ pub enum FunctionError {
     InvalidRayDescriptor(Handle<crate::Expression>),
     #[error("Ray Query {0:?} does not have a matching type")]
     InvalidRayQueryType(Handle<crate::Type>),
+    #[error("Hit distance {0:?} must be an f32")]
+    InvalidHitDistanceType(Handle<crate::Expression>),
     #[error("Shader requires capability {0:?}")]
     MissingCapability(super::Capabilities),
     #[error(
@@ -1487,6 +1489,19 @@ impl super::Validator {
                         crate::RayQueryFunction::Proceed { result } => {
                             self.emit_expression(result, context)?;
                         }
+                        crate::RayQueryFunction::GenerateIntersection { hit_t } => {
+                            match *context.resolve_type(hit_t, &self.valid_expression_set)? {
+                                Ti::Scalar(crate::Scalar {
+                                    kind: crate::ScalarKind::Float,
+                                    width: _,
+                                }) => {}
+                                _ => {
+                                    return Err(FunctionError::InvalidHitDistanceType(hit_t)
+                                        .with_span_static(span, "invalid hit_t"))
+                                }
+                            }
+                        }
+                        crate::RayQueryFunction::ConfirmIntersection => {}
                         crate::RayQueryFunction::Terminate => {}
                     }
                 }

naga/src/valid/handles.rs

@@ -707,6 +707,10 @@ impl super::Validator {
                     crate::RayQueryFunction::Proceed { result } => {
                         validate_expr(result)?;
                     }
+                    crate::RayQueryFunction::GenerateIntersection { hit_t } => {
+                        validate_expr(hit_t)?;
+                    }
+                    crate::RayQueryFunction::ConfirmIntersection => {}
                     crate::RayQueryFunction::Terminate => {}
                 }
                 Ok(())

naga/tests/in/ray-query.wgsl

@@ -87,5 +87,11 @@ fn main_candidate() {
     var rq: ray_query;
     rayQueryInitialize(&rq, acc_struct, RayDesc(RAY_FLAG_TERMINATE_ON_FIRST_HIT, 0xFFu, 0.1, 100.0, pos, dir));
     let intersection = rayQueryGetCandidateIntersection(&rq);
-    output.visible = u32(intersection.kind == RAY_QUERY_INTERSECTION_AABB);
+    if (intersection.kind == RAY_QUERY_INTERSECTION_AABB) {
+        rayQueryGenerateIntersection(&rq, 10.0);
+    } else if (intersection.kind == RAY_QUERY_INTERSECTION_TRIANGLE) {
+        rayQueryConfirmIntersection(&rq);
+    } else {
+        rayQueryTerminate(&rq);
+    }
 }