[wgsl-out] Ray tracing pipelines (gfx-rs#8970)

* wgsl-out rt pipelines

* Format.

* Enable the ray tracing pipeline extension.

* Add more checks for ray tracing pipelines.

* Changelog.

* Remove ir target.

Author: Vecvec

CHANGELOG.md

@@ -76,6 +76,7 @@ By @cwfitzgerald in [#8999](https://github.com/gfx-rs/wgpu/pull/8999).
 #### naga
 
 - Initial wgsl-in ray tracing pipelines. By @Vecvec in [#8570](https://github.com/gfx-rs/wgpu/pull/8570).
+- wgsl-out ray tracing pipelines. By @Vecvec in [#8970](https://github.com/gfx-rs/wgpu/pull/8970).
 - Allow parsing shaders which make use of `SPV_KHR_non_semantic_info` for debug info. Also removes `naga::front::spv::SUPPORTED_EXT_SETS`. By @inner-daemons in #8827.
 
 #### GLES

naga/src/back/wgsl/writer.rs

@@ -36,6 +36,7 @@ enum Attribute {
     MeshStage(String),
     TaskPayload(String),
     PerPrimitive,
+    IncomingRayPayload(String),
 }
 
 /// The WGSL form that `write_expr_with_indirection` should use to render a Naga
@@ -243,10 +244,17 @@ impl<W: Write> Writer<W> {
                         Attribute::WorkGroupSize(ep.workgroup_size),
                     ]
                 }
-                ShaderStage::RayGeneration
-                | ShaderStage::AnyHit
-                | ShaderStage::ClosestHit
-                | ShaderStage::Miss => unreachable!(),
+                ShaderStage::RayGeneration => vec![Attribute::Stage(ShaderStage::RayGeneration)],
+                ShaderStage::AnyHit | ShaderStage::ClosestHit | ShaderStage::Miss => {
+                    let payload_name = module.global_variables[ep.incoming_ray_payload.unwrap()]
+                        .name
+                        .clone()
+                        .unwrap();
+                    vec![
+                        Attribute::Stage(ep.stage),
+                        Attribute::IncomingRayPayload(payload_name),
+                    ]
+                }
             };
             self.write_attributes(&attributes)?;
             // Add a newline after attribute
@@ -287,6 +295,7 @@ impl<W: Write> Writer<W> {
             primitive_index: bool,
             cooperative_matrix: bool,
             draw_index: bool,
+            ray_tracing_pipeline: bool,
         }
         let mut needed = RequiredEnabled {
             mesh_shaders: module.uses_mesh_shaders(),
@@ -313,6 +322,22 @@ impl<W: Write> Writer<W> {
                 needed.mesh_shaders = true;
             }
             crate::Binding::BuiltIn(crate::BuiltIn::DrawIndex) => needed.draw_index = true,
+            crate::Binding::BuiltIn(
+                crate::BuiltIn::RayInvocationId
+                | crate::BuiltIn::NumRayInvocations
+                | crate::BuiltIn::InstanceCustomData
+                | crate::BuiltIn::GeometryIndex
+                | crate::BuiltIn::WorldRayOrigin
+                | crate::BuiltIn::WorldRayDirection
+                | crate::BuiltIn::ObjectRayOrigin
+                | crate::BuiltIn::ObjectRayDirection
+                | crate::BuiltIn::RayTmin
+                | crate::BuiltIn::RayTCurrentMax
+                | crate::BuiltIn::ObjectToWorld
+                | crate::BuiltIn::WorldToObject,
+            ) => {
+                needed.ray_tracing_pipeline = true;
+            }
             _ => {}
         };
 
@@ -332,6 +357,9 @@ impl<W: Write> Writer<W> {
                 TypeInner::CooperativeMatrix { .. } => {
                     needed.cooperative_matrix = true;
                 }
+                TypeInner::AccelerationStructure { .. } => {
+                    needed.ray_tracing_pipeline = true;
+                }
                 _ => {}
             }
         }
@@ -350,6 +378,44 @@ impl<W: Write> Writer<W> {
             }
         }
 
+        if module.global_variables.iter().any(|gv| {
+            gv.1.space == crate::AddressSpace::IncomingRayPayload
+                || gv.1.space == crate::AddressSpace::RayPayload
+        }) {
+            needed.ray_tracing_pipeline = true;
+        }
+
+        if module.entry_points.iter().any(|ep| {
+            matches!(
+                ep.stage,
+                ShaderStage::RayGeneration
+                    | ShaderStage::AnyHit
+                    | ShaderStage::ClosestHit
+                    | ShaderStage::Miss
+            )
+        }) {
+            needed.ray_tracing_pipeline = true;
+        }
+
+        if module.global_variables.iter().any(|gv| {
+            gv.1.space == crate::AddressSpace::IncomingRayPayload
+                || gv.1.space == crate::AddressSpace::RayPayload
+        }) {
+            needed.ray_tracing_pipeline = true;
+        }
+
+        if module.entry_points.iter().any(|ep| {
+            matches!(
+                ep.stage,
+                ShaderStage::RayGeneration
+                    | ShaderStage::AnyHit
+                    | ShaderStage::ClosestHit
+                    | ShaderStage::Miss
+            )
+        }) {
+            needed.ray_tracing_pipeline = true;
+        }
+
         // Write required declarations
         let mut any_written = false;
         if needed.f16 {
@@ -380,6 +446,10 @@ impl<W: Write> Writer<W> {
             writeln!(self.out, "enable wgpu_cooperative_matrix;")?;
             any_written = true;
         }
+        if needed.ray_tracing_pipeline {
+            writeln!(self.out, "enable wgpu_ray_tracing_pipeline;")?;
+            any_written = true;
+        }
         if any_written {
             // Empty line for readability
             writeln!(self.out)?;
@@ -501,10 +571,10 @@ impl<W: Write> Writer<W> {
                         ShaderStage::Task => "task",
                         //Handled by another variant in the Attribute enum, so this code should never be hit.
                         ShaderStage::Mesh => unreachable!(),
-                        ShaderStage::RayGeneration
-                        | ShaderStage::AnyHit
-                        | ShaderStage::ClosestHit
-                        | ShaderStage::Miss => unreachable!(),
+                        ShaderStage::RayGeneration => "ray_generation",
+                        ShaderStage::AnyHit => "any_hit",
+                        ShaderStage::ClosestHit => "closest_hit",
+                        ShaderStage::Miss => "miss",
                     };
 
                     write!(self.out, "@{stage_str} ")?;
@@ -542,6 +612,9 @@ impl<W: Write> Writer<W> {
                     write!(self.out, "@payload({payload_name}) ")?;
                 }
                 Attribute::PerPrimitive => write!(self.out, "@per_primitive ")?,
+                Attribute::IncomingRayPayload(ref payload_name) => {
+                    write!(self.out, "@incoming_payload({payload_name}) ")?;
+                }
             };
         }
         Ok(())
@@ -1103,7 +1176,21 @@ impl<W: Write> Writer<W> {
                 self.write_expr(module, data.stride, func_ctx)?;
                 writeln!(self.out, ");")?
             }
-            Statement::RayPipelineFunction(_) => unreachable!(),
+            Statement::RayPipelineFunction(fun) => match fun {
+                crate::RayPipelineFunction::TraceRay {
+                    acceleration_structure,
+                    descriptor,
+                    payload,
+                } => {
+                    write!(self.out, "{level}traceRay(")?;
+                    self.write_expr(module, acceleration_structure, func_ctx)?;
+                    write!(self.out, ", ")?;
+                    self.write_expr(module, descriptor, func_ctx)?;
+                    write!(self.out, ", ")?;
+                    self.write_expr(module, payload, func_ctx)?;
+                    writeln!(self.out, ");")?
+                }
+            },
         }
 
         Ok(())

naga/src/common/wgsl/to_wgsl.rs

@@ -195,25 +195,26 @@ impl TryToWgsl for crate::BuiltIn {
             Bi::PrimitiveCount => "primitive_count",
             Bi::CullPrimitive => "cull_primitive",
 
+            Bi::RayInvocationId => "ray_invocation_id",
+            Bi::NumRayInvocations => "num_ray_invocations",
+            Bi::InstanceCustomData => "instance_custom_data",
+            Bi::GeometryIndex => "geometry_index",
+            Bi::WorldRayOrigin => "world_ray_origin",
+            Bi::WorldRayDirection => "world_ray_direction",
+            Bi::ObjectRayOrigin => "object_ray_origin",
+            Bi::ObjectRayDirection => "object_ray_direction",
+            Bi::RayTmin => "ray_t_min",
+            Bi::RayTCurrentMax => "ray_t_current_max",
+            Bi::ObjectToWorld => "object_to_world",
+            Bi::WorldToObject => "world_to_object",
+            Bi::HitKind => "hit_kind",
+
             Bi::BaseInstance
             | Bi::BaseVertex
             | Bi::CullDistance
             | Bi::PointSize
             | Bi::PointCoord
-            | Bi::WorkGroupSize
-            | Bi::RayInvocationId
-            | Bi::NumRayInvocations
-            | Bi::InstanceCustomData
-            | Bi::GeometryIndex
-            | Bi::WorldRayOrigin
-            | Bi::WorldRayDirection
-            | Bi::ObjectRayOrigin
-            | Bi::ObjectRayDirection
-            | Bi::RayTmin
-            | Bi::RayTCurrentMax
-            | Bi::ObjectToWorld
-            | Bi::WorldToObject
-            | Bi::HitKind => return None,
+            | Bi::WorkGroupSize => return None,
         })
     }
 }
@@ -387,7 +388,8 @@ pub const fn address_space_str(
             As::Handle => return (None, None),
             As::Function => "function",
             As::TaskPayload => "task_payload",
-            As::IncomingRayPayload | As::RayPayload => return (None, None),
+            As::IncomingRayPayload => "incoming_ray_payload",
+            As::RayPayload => "ray_payload",
         }),
         None,
     )

naga/tests/in/wgsl/ray-tracing-pipeline.toml

@@ -1,2 +1,2 @@
 capabilities = "RAY_TRACING_PIPELINE"
-targets = "IR"
+targets = "WGSL"