[spv-out] handle rayQueryTerminate correctly.

CHANGELOG.md

@@ -162,6 +162,7 @@ By @cwfitzgerald in [#8579](https://github.com/gfx-rs/wgpu/pull/8579).
 #### naga
 
 - Fix a bug that resulted in the Metal error `program scope variable must reside in constant address space` in some cases. By @teoxoy in [#8311](https://github.com/gfx-rs/wgpu/pull/8311).
+- Handle `rayQueryTerminate` in spv-out instead of ignoring it. By @Vecvec in [#8581](https://github.com/gfx-rs/wgpu/pull/8581).
 
 #### General
 

docs/api-specs/ray_tracing.md

@@ -125,8 +125,9 @@ rayQueryGenerateIntersection(hit_t: f32)
 // - Commits a hit from triangular non-opaque geometry.
 rayQueryConfirmIntersection()
 
-// - Aborts the query.
-rayQueryTerminate()
+// Aborts the query which is in progress, that is, the next `rayQueryProceed` is guaranteed to return `false`
+// and any call to `rayQueryGetCommittedIntersection` will return the closest committed result so far.
+rayQueryTerminate(rq: ptr<function, ray_query>)
 
 // - Returns intersection details about a hit considered `Committed`.
 rayQueryGetCommittedIntersection(rq: ptr<function, ray_query>) -> RayIntersection

naga/src/back/spv/instructions.rs

@@ -842,6 +842,12 @@ impl super::Instruction {
         instruction
     }
 
+    pub(super) fn ray_query_terminate(query: Word) -> Self {
+        let mut instruction = Self::new(Op::RayQueryTerminateKHR);
+        instruction.add_operand(query);
+        instruction
+    }
+
     //
     //  Conversion Instructions
     //

naga/src/back/spv/mod.rs

@@ -459,6 +459,7 @@ enum LookupRayQueryFunction {
     ConfirmIntersection,
     GetVertexPositions { committed: bool },
     GetIntersection { committed: bool },
+    Terminate,
 }
 
 #[derive(Debug)]

naga/src/back/spv/ray.rs

@@ -1781,6 +1781,98 @@ impl Writer {
 
         func_id
     }
+
+    fn write_ray_query_terminate(&mut self) -> spirv::Word {
+        if let Some(&word) = self
+            .ray_query_functions
+            .get(&LookupRayQueryFunction::Terminate)
+        {
+            return word;
+        }
+
+        let ray_query_type_id = self.get_ray_query_pointer_id();
+
+        let u32_ty = self.get_u32_type_id();
+        let u32_ptr_ty = self.get_pointer_type_id(u32_ty, spirv::StorageClass::Function);
+
+        let bool_type_id = self.get_bool_type_id();
+
+        let (func_id, mut function, arg_ids) =
+            self.write_function_signature(&[ray_query_type_id, u32_ptr_ty], self.void_type);
+
+        let query_id = arg_ids[0];
+        let init_tracker_id = arg_ids[1];
+
+        let block_id = self.id_gen.next();
+        let mut block = Block::new(block_id);
+
+        let initialized_tracker_id = self.id_gen.next();
+        block.body.push(Instruction::load(
+            u32_ty,
+            initialized_tracker_id,
+            init_tracker_id,
+            None,
+        ));
+
+        let merge_id = self.id_gen.next();
+        let merge_block = Block::new(merge_id);
+
+        let valid_block_id = self.id_gen.next();
+        let mut valid_block = Block::new(valid_block_id);
+
+        let instruction = if self.ray_query_initialization_tracking {
+            let has_proceeded = write_ray_flags_contains_flags(
+                self,
+                &mut block,
+                initialized_tracker_id,
+                super::RayQueryPoint::PROCEED.bits(),
+            );
+
+            let finished_proceed_id = write_ray_flags_contains_flags(
+                self,
+                &mut block,
+                initialized_tracker_id,
+                super::RayQueryPoint::FINISHED_TRAVERSAL.bits(),
+            );
+
+            // Can't find anything to suggest double calling this function is invalid.
+
+            let not_finished_id = self.id_gen.next();
+            block.body.push(Instruction::unary(
+                spirv::Op::LogicalNot,
+                bool_type_id,
+                not_finished_id,
+                finished_proceed_id,
+            ));
+
+            let valid_call = self.write_logical_and(&mut block, not_finished_id, has_proceeded);
+
+            block.body.push(Instruction::selection_merge(
+                merge_id,
+                spirv::SelectionControl::NONE,
+            ));
+
+            Instruction::branch_conditional(valid_call, valid_block_id, merge_id)
+        } else {
+            Instruction::branch(valid_block_id)
+        };
+
+        function.consume(block, instruction);
+
+        valid_block
+            .body
+            .push(Instruction::ray_query_terminate(query_id));
+
+        function.consume(valid_block, Instruction::branch(merge_id));
+
+        function.consume(merge_block, Instruction::return_void());
+
+        function.to_words(&mut self.logical_layout.function_definitions);
+
+        self.ray_query_functions
+            .insert(LookupRayQueryFunction::Proceed, func_id);
+        func_id
+    }
 }
 
 impl BlockContext<'_> {
@@ -1863,7 +1955,17 @@ impl BlockContext<'_> {
                     &[query_id, tracker_ids.initialized_tracker],
                 ));
             }
-            crate::RayQueryFunction::Terminate => {}
+            crate::RayQueryFunction::Terminate => {
+                let id = self.gen_id();
+
+                let func_id = self.writer.write_ray_query_terminate();
+                block.body.push(Instruction::function_call(
+                    self.writer.void_type,
+                    id,
+                    func_id,
+                    &[query_id, tracker_ids.initialized_tracker],
+                ));
+            }
         }
     }
 

naga/tests/out/spv/wgsl-aliased-ray-query.spvasm

@@ -1,7 +1,7 @@
 ; SPIR-V
 ; Version: 1.4
 ; Generator: rspirv
-; Bound: 244
+; Bound: 258
 OpCapability Shader
 OpCapability RayQueryKHR
 OpExtension "SPV_KHR_ray_query"
@@ -308,6 +308,25 @@ OpBranch %231
 %231 = OpLabel
 OpReturn
 OpFunctionEnd
+%245 = OpFunction  %2  None %225
+%246 = OpFunctionParameter  %32
+%247 = OpFunctionParameter  %33
+%248 = OpLabel
+%249 = OpLoad  %7  %247
+%252 = OpBitwiseAnd  %7  %249 %90
+%253 = OpINotEqual  %10  %252 %35
+%254 = OpBitwiseAnd  %7  %249 %23
+%255 = OpINotEqual  %10  %254 %35
+%256 = OpLogicalNot  %10  %255
+%257 = OpLogicalAnd  %10  %256 %253
+OpSelectionMerge %250 None
+OpBranchConditional %257 %251 %250
+%251 = OpLabel
+OpRayQueryTerminateKHR %246
+OpBranch %250
+%250 = OpLabel
+OpReturn
+OpFunctionEnd
 %16 = OpFunction  %2  None %17
 %15 = OpLabel
 %31 = OpVariable  %32  Function
@@ -334,6 +353,7 @@ OpBranchConditional %221 %223 %224
 %243 = OpFunctionCall  %2  %226 %31 %34
 OpReturn
 %224 = OpLabel
+%244 = OpFunctionCall  %2  %245 %31 %34
 OpReturn
 %222 = OpLabel
 OpBranch %181

naga/tests/out/spv/wgsl-ray-query-no-init-tracking.spvasm

@@ -1,7 +1,7 @@
 ; SPIR-V
 ; Version: 1.4
 ; Generator: rspirv
-; Bound: 382
+; Bound: 396
 OpCapability Shader
 OpCapability RayQueryKHR
 OpExtension "SPV_KHR_ray_query"
@@ -508,6 +508,25 @@ OpBranch %369
 %369 = OpLabel
 OpReturn
 OpFunctionEnd
+%383 = OpFunction  %2  None %363
+%384 = OpFunctionParameter  %32
+%385 = OpFunctionParameter  %33
+%386 = OpLabel
+%387 = OpLoad  %6  %385
+%390 = OpBitwiseAnd  %6  %387 %93
+%391 = OpINotEqual  %8  %390 %35
+%392 = OpBitwiseAnd  %6  %387 %27
+%393 = OpINotEqual  %8  %392 %35
+%394 = OpLogicalNot  %8  %393
+%395 = OpLogicalAnd  %8  %394 %391
+OpSelectionMerge %388 None
+OpBranchConditional %395 %389 %388
+%389 = OpLabel
+OpRayQueryTerminateKHR %384
+OpBranch %388
+%388 = OpLabel
+OpReturn
+OpFunctionEnd
 %262 = OpFunction  %2  None %243
 %261 = OpLabel
 %266 = OpVariable  %32  Function
@@ -534,6 +553,7 @@ OpBranchConditional %359 %361 %362
 %381 = OpFunctionCall  %2  %364 %266 %267
 OpReturn
 %362 = OpLabel
+%382 = OpFunctionCall  %2  %383 %266 %267
 OpReturn
 %360 = OpLabel
 OpBranch %319

naga/tests/out/spv/wgsl-ray-query.spvasm

@@ -1,7 +1,7 @@
 ; SPIR-V
 ; Version: 1.4
 ; Generator: rspirv
-; Bound: 382
+; Bound: 396
 OpCapability Shader
 OpCapability RayQueryKHR
 OpExtension "SPV_KHR_ray_query"
@@ -508,6 +508,25 @@ OpBranch %369
 %369 = OpLabel
 OpReturn
 OpFunctionEnd
+%383 = OpFunction  %2  None %363
+%384 = OpFunctionParameter  %32
+%385 = OpFunctionParameter  %33
+%386 = OpLabel
+%387 = OpLoad  %6  %385
+%390 = OpBitwiseAnd  %6  %387 %93
+%391 = OpINotEqual  %8  %390 %35
+%392 = OpBitwiseAnd  %6  %387 %27
+%393 = OpINotEqual  %8  %392 %35
+%394 = OpLogicalNot  %8  %393
+%395 = OpLogicalAnd  %8  %394 %391
+OpSelectionMerge %388 None
+OpBranchConditional %395 %389 %388
+%389 = OpLabel
+OpRayQueryTerminateKHR %384
+OpBranch %388
+%388 = OpLabel
+OpReturn
+OpFunctionEnd
 %262 = OpFunction  %2  None %243
 %261 = OpLabel
 %266 = OpVariable  %32  Function
@@ -534,6 +553,7 @@ OpBranchConditional %359 %361 %362
 %381 = OpFunctionCall  %2  %364 %266 %267
 OpReturn
 %362 = OpLabel
+%382 = OpFunctionCall  %2  %383 %266 %267
 OpReturn
 %360 = OpLabel
 OpBranch %319

tests/tests/wgpu-gpu/ray_tracing/shader.wgsl

@@ -124,4 +124,11 @@ fn invalid_usages() {
         rayQueryProceed(&rq);
         let intersection = rayQueryGetCommittedIntersection(&rq);
     }
+    {
+        var rq: ray_query;
+        rayQueryInitialize(&rq, acc_struct, RayDesc(0u, 0xFFu, 0.001, 100000.0, vec3f(0.0, 0.0, 0.0), vec3f(0.0, 0.0, 1.0)));
+        rayQueryProceed(&rq);
+        // The acceleration structure has been set up to not generate an intersections, meaning terminate is invalid here.
+        rayQueryTerminate(&rq);
+    }
 }