ray query: validation, better test

Author: kvark

src/valid/expression.rs

@@ -35,6 +35,8 @@ pub enum ExpressionError {
     InvalidPointerType(Handle<crate::Expression>),
     #[error("Array length of {0:?} can't be done")]
     InvalidArrayType(Handle<crate::Expression>),
+    #[error("Get intersection of {0:?} can't be done")]
+    InvalidRayQueryType(Handle<crate::Expression>),
     #[error("Splatting {0:?} can't be done")]
     InvalidSplatType(Handle<crate::Expression>),
     #[error("Swizzling {0:?} can't be done")]
@@ -1427,7 +1429,26 @@ impl super::Validator {
                     return Err(ExpressionError::InvalidArrayType(expr));
                 }
             },
-            E::RayQueryProceedResult | E::RayQueryGetIntersection { .. } => ShaderStages::all(),
+            E::RayQueryProceedResult => ShaderStages::all(),
+            E::RayQueryGetIntersection {
+                query,
+                committed: _,
+            } => match resolver[query] {
+                Ti::Pointer {
+                    base,
+                    space: crate::AddressSpace::Function,
+                } => match resolver.types[base].inner {
+                    Ti::RayQuery => ShaderStages::all(),
+                    ref other => {
+                        log::error!("Intersection result of a pointer to {:?}", other);
+                        return Err(ExpressionError::InvalidRayQueryType(query));
+                    }
+                },
+                ref other => {
+                    log::error!("Intersection result of {:?}", other);
+                    return Err(ExpressionError::InvalidRayQueryType(query));
+                }
+            },
         };
         Ok(stages)
     }

src/valid/function.rs

@@ -47,8 +47,6 @@ pub enum AtomicError {
     InvalidPointer(Handle<crate::Expression>),
     #[error("Operand {0:?} has invalid type.")]
     InvalidOperand(Handle<crate::Expression>),
-    #[error("Result expression {0:?} has already been introduced earlier")]
-    ResultAlreadyInScope(Handle<crate::Expression>),
     #[error("Result type for {0:?} doesn't match the statement")]
     ResultTypeMismatch(Handle<crate::Expression>),
 }
@@ -131,6 +129,14 @@ pub enum FunctionError {
     },
     #[error("Atomic operation is invalid")]
     InvalidAtomic(#[from] AtomicError),
+    #[error("Ray Query {0:?} is not a local variable")]
+    InvalidRayQueryExpression(Handle<crate::Expression>),
+    #[error("Acceleration structure {0:?} is not a matching expression")]
+    InvalidAccelerationStructure(Handle<crate::Expression>),
+    #[error("Ray descriptor {0:?} is not a matching expression")]
+    InvalidRayDescriptor(Handle<crate::Expression>),
+    #[error("Ray Query {0:?} does not have a matching type")]
+    InvalidRayQueryType(Handle<crate::Type>),
     #[error(
         "Required uniformity of control flow for {0:?} in {1:?} is not fulfilled because of {2:?}"
     )]
@@ -169,8 +175,10 @@ struct BlockContext<'a> {
     info: &'a FunctionInfo,
     expressions: &'a Arena<crate::Expression>,
     types: &'a UniqueArena<crate::Type>,
+    local_vars: &'a Arena<crate::LocalVariable>,
     global_vars: &'a Arena<crate::GlobalVariable>,
     functions: &'a Arena<crate::Function>,
+    special_types: &'a crate::SpecialTypes,
     prev_infos: &'a [FunctionInfo],
     return_type: Option<Handle<crate::Type>>,
 }
@@ -188,8 +196,10 @@ impl<'a> BlockContext<'a> {
             info,
             expressions: &fun.expressions,
             types: &module.types,
+            local_vars: &fun.local_variables,
             global_vars: &module.global_variables,
             functions: &module.functions,
+            special_types: &module.special_types,
             prev_infos,
             return_type: fun.result.as_ref().map(|fr| fr.ty),
         }
@@ -299,6 +309,21 @@ impl super::Validator {
         Ok(callee_info.available_stages)
     }
 
+    #[cfg(feature = "validate")]
+    fn emit_expression(
+        &mut self,
+        handle: Handle<crate::Expression>,
+        context: &BlockContext,
+    ) -> Result<(), WithSpan<FunctionError>> {
+        if self.valid_expression_set.insert(handle.index()) {
+            self.valid_expression_list.push(handle);
+            Ok(())
+        } else {
+            Err(FunctionError::ExpressionAlreadyInScope(handle)
+                .with_span_handle(handle, context.expressions))
+        }
+    }
+
     #[cfg(feature = "validate")]
     fn validate_atomic(
         &mut self,
@@ -347,13 +372,7 @@ impl super::Validator {
             }
         }
 
-        if self.valid_expression_set.insert(result.index()) {
-            self.valid_expression_list.push(result);
-        } else {
-            return Err(AtomicError::ResultAlreadyInScope(result)
-                .with_span_handle(result, context.expressions)
-                .into_other());
-        }
+        self.emit_expression(result, context)?;
         match context.expressions[result] {
             crate::Expression::AtomicResult { ty, comparison }
                 if {
@@ -401,12 +420,7 @@ impl super::Validator {
             match *statement {
                 S::Emit(ref range) => {
                     for handle in range.clone() {
-                        if self.valid_expression_set.insert(handle.index()) {
-                            self.valid_expression_list.push(handle);
-                        } else {
-                            return Err(FunctionError::ExpressionAlreadyInScope(handle)
-                                .with_span_handle(handle, context.expressions));
-                        }
+                        self.emit_expression(handle, context)?;
                     }
                 }
                 S::Block(ref block) => {
@@ -807,8 +821,55 @@ impl super::Validator {
                 } => {
                     self.validate_atomic(pointer, fun, value, result, context)?;
                 }
-                S::RayQuery { query: _, fun: _ } => {
-                    //TODO
+                S::RayQuery { query, ref fun } => {
+                    let query_var = match *context.get_expression(query) {
+                        crate::Expression::LocalVariable(var) => &context.local_vars[var],
+                        ref other => {
+                            log::error!("Unexpected ray query expression {other:?}");
+                            return Err(FunctionError::InvalidRayQueryExpression(query)
+                                .with_span_static(span, "invalid query expression"));
+                        }
+                    };
+                    match context.types[query_var.ty].inner {
+                        Ti::RayQuery => {}
+                        ref other => {
+                            log::error!("Unexpected ray query type {other:?}");
+                            return Err(FunctionError::InvalidRayQueryType(query_var.ty)
+                                .with_span_static(span, "invalid query type"));
+                        }
+                    }
+                    match *fun {
+                        crate::RayQueryFunction::Initialize {
+                            acceleration_structure,
+                            descriptor,
+                        } => {
+                            match *context
+                                .resolve_type(acceleration_structure, &self.valid_expression_set)?
+                            {
+                                Ti::AccelerationStructure => {}
+                                _ => {
+                                    return Err(FunctionError::InvalidAccelerationStructure(
+                                        acceleration_structure,
+                                    )
+                                    .with_span_static(span, "invalid acceleration structure"))
+                                }
+                            }
+                            let desc_ty_given =
+                                context.resolve_type(descriptor, &self.valid_expression_set)?;
+                            let desc_ty_expected = context
+                                .special_types
+                                .ray_desc
+                                .map(|handle| &context.types[handle].inner);
+                            if Some(desc_ty_given) != desc_ty_expected {
+                                return Err(FunctionError::InvalidRayDescriptor(descriptor)
+                                    .with_span_static(span, "invalid ray descriptor"));
+                            }
+                        }
+                        crate::RayQueryFunction::Proceed { result } => {
+                            self.emit_expression(result, context)?;
+                        }
+                        crate::RayQueryFunction::Terminate => {}
+                    }
                 }
             }
         }

tests/in/ray-query.wgsl

@@ -45,19 +45,29 @@ struct RayIntersection {
 
 struct Output {
     visible: u32,
+    normal: vec3<f32>,
 }
 
 @group(0) @binding(1)
 var<storage, read_write> output: Output;
 
+fn get_torus_normal(world_point: vec3<f32>, intersection: RayIntersection) -> vec3<f32> {
+    let local_point = intersection.world_to_object * vec4<f32>(world_point, 1.0);
+    let point_on_guiding_line = normalize(local_point.xy) * 2.4;
+    let world_point_on_guiding_line = intersection.object_to_world * vec4<f32>(point_on_guiding_line, 0.0, 1.0);
+    return normalize(world_point - world_point_on_guiding_line);
+}
+
 @compute @workgroup_size(1)
 fn main() {
     var rq: ray_query;
 
-    rayQueryInitialize(&rq, acc_struct, RayDesc(RAY_FLAG_TERMINATE_ON_FIRST_HIT, 0xFFu, 0.1, 100.0, vec3<f32>(0.0), vec3<f32>(0.0, 1.0, 0.0)));
+    let dir = vec3<f32>(0.0, 1.0, 0.0);
+    rayQueryInitialize(&rq, acc_struct, RayDesc(RAY_FLAG_TERMINATE_ON_FIRST_HIT, 0xFFu, 0.1, 100.0, vec3<f32>(0.0), dir));
 
-    rayQueryProceed(&rq);
+    while (rayQueryProceed(&rq)) {}
 
     let intersection = rayQueryGetCommittedIntersection(&rq);
     output.visible = u32(intersection.kind == RAY_QUERY_INTERSECTION_NONE);
+    output.normal = get_torus_normal(dir * intersection.t, intersection);
 }

tests/out/msl/ray-query.msl

@@ -15,14 +15,8 @@ constexpr metal::uint _map_intersection_type(const metal::raytracing::intersecti
 
 struct Output {
     uint visible_;
-};
-struct RayDesc {
-    uint flags;
-    uint cull_mask;
-    float tmin;
-    float tmax;
-    metal::float3 origin;
-    metal::float3 dir;
+    char _pad1[12];
+    metal::float3 normal;
 };
 struct RayIntersection {
     uint kind;
@@ -38,21 +32,48 @@ struct RayIntersection {
     metal::float4x3 object_to_world;
     metal::float4x3 world_to_object;
 };
+struct RayDesc {
+    uint flags;
+    uint cull_mask;
+    float tmin;
+    float tmax;
+    metal::float3 origin;
+    metal::float3 dir;
+};
+
+metal::float3 get_torus_normal(
+    metal::float3 world_point,
+    RayIntersection intersection
+) {
+    metal::float3 local_point = intersection.world_to_object * metal::float4(world_point, 1.0);
+    metal::float2 point_on_guiding_line = metal::normalize(local_point.xy) * 2.4000000953674316;
+    metal::float3 world_point_on_guiding_line = intersection.object_to_world * metal::float4(point_on_guiding_line, 0.0, 1.0);
+    return metal::normalize(world_point - world_point_on_guiding_line);
+}
 
 kernel void main_(
   metal::raytracing::instance_acceleration_structure acc_struct [[user(fake0)]]
 , device Output& output [[user(fake0)]]
 ) {
     _RayQuery rq = {};
-    RayDesc _e12 = RayDesc {4u, 255u, 0.10000000149011612, 100.0, metal::float3(0.0), metal::float3(0.0, 1.0, 0.0)};
+    metal::float3 dir = metal::float3(0.0, 1.0, 0.0);
+    RayDesc _e12 = RayDesc {4u, 255u, 0.10000000149011612, 100.0, metal::float3(0.0), dir};
     rq.intersector.assume_geometry_type(metal::raytracing::geometry_type::triangle);
     rq.intersector.set_opacity_cull_mode((_e12.flags & 64) != 0 ? metal::raytracing::opacity_cull_mode::opaque : (_e12.flags & 128) != 0 ? metal::raytracing::opacity_cull_mode::non_opaque : metal::raytracing::opacity_cull_mode::none);
     rq.intersector.force_opacity((_e12.flags & 1) != 0 ? metal::raytracing::forced_opacity::opaque : (_e12.flags & 2) != 0 ? metal::raytracing::forced_opacity::non_opaque : metal::raytracing::forced_opacity::none);
     rq.intersector.accept_any_intersection((_e12.flags & 4) != 0);
     rq.intersection = rq.intersector.intersect(metal::raytracing::ray(_e12.origin, _e12.dir, _e12.tmin, _e12.tmax), acc_struct, _e12.cull_mask);    rq.ready = true;
-    bool _e13 = rq.ready;
-    rq.ready = false;
-    RayIntersection intersection = RayIntersection {_map_intersection_type(rq.intersection.type), rq.intersection.distance, rq.intersection.user_instance_id, rq.intersection.instance_id, {}, rq.intersection.geometry_id, rq.intersection.primitive_id, rq.intersection.triangle_barycentric_coord, rq.intersection.triangle_front_facing, {}, rq.intersection.object_to_world_transform, rq.intersection.world_to_object_transform};
-    output.visible_ = static_cast<uint>(intersection.kind == 0u);
+    while(true) {
+        bool _e13 = rq.ready;
+        rq.ready = false;
+        if (_e13) {
+        } else {
+            break;
+        }
+    }
+    RayIntersection intersection_1 = RayIntersection {_map_intersection_type(rq.intersection.type), rq.intersection.distance, rq.intersection.user_instance_id, rq.intersection.instance_id, {}, rq.intersection.geometry_id, rq.intersection.primitive_id, rq.intersection.triangle_barycentric_coord, rq.intersection.triangle_front_facing, {}, rq.intersection.object_to_world_transform, rq.intersection.world_to_object_transform};
+    output.visible_ = static_cast<uint>(intersection_1.kind == 0u);
+    metal::float3 _e25 = get_torus_normal(dir * intersection_1.t, intersection_1);
+    output.normal = _e25;
     return;
 }