[wgsl in and IR] naga ray tracing pipelines

docs/api-specs/ray_tracing.md

@@ -93,10 +93,12 @@ fn render(/*whatever args you need to render*/) {
 
 ## `naga`'s raytracing API:
 
-`naga` supports ray queries (also known as inline raytracing) only. To enable basic ray query functions you must add
+`naga` supports ray queries (also known as inline raytracing). To enable basic ray query functions you must add
 `enable wgpu_ray_query` to the shader, ray queries and acceleration structures also support tags which require extra
-`enable` extensions (see Acceleration structure tags for more info). Ray tracing pipelines
-are currently unsupported. Naming is mostly taken from vulkan.
+`enable` extensions (see Acceleration structure tags for more info). Ray tracing pipelines are currently in
+development. Naming is mostly taken from vulkan.
+
+### Ray Queries
 
 ```wgsl
 // - Initializes the `ray_query` to check where (if anywhere) the ray defined by `ray_desc` hits in `acceleration_structure`
@@ -147,7 +149,7 @@ getCandidateHitVertexPositions(rq: ptr<function, ray_query<vertex_return>>) -> a
 
 > [!CAUTION]
 >
-> ### ⚠️Undefined behavior ⚠️:
+> #### ⚠️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
@@ -271,6 +273,39 @@ const RAY_QUERY_INTERSECTION_GENERATED = 2;
 const RAY_QUERY_INTERSECTION_AABB = 3;
 ```
 
+### Ray Tracing Pipelines
+
+Functions
+```wgsl
+// Begins to check where (if anywhere) the ray defined by `ray_desc` hits in `acceleration_structure` calling through the `any_hit` shaders and `closest_hit` shader if something was hit or the `miss` shader if no hit was found
+traceRay<T>(acceleration_structure: acceleration_structure, ray_desc: RayDesc, payload: ptr<ray_payload, T>)
+```
+
+> [!CAUTION]
+>
+> #### ⚠️Undefined behavior ⚠️:
+> Calling `traceRay` inside another `traceRay` more than `max_recursion_depth` times
+>
+> *this is only known undefined behaviour, and will be worked around in the future.
+
+New shader stages
+```wgsl
+// First stage to be called, allowed to call `traceRay`
+@ray_generation
+fn rg() {}
+
+// Stage called on any hit that is not opaque, not allowed to call `traceRay`
+@any_hit
+fn ah() {}
+
+// Stage called on the closest hit, allowed to call `traceRay`
+@closest_hit
+fn ch() {}
+
+// Stage call if there was never a hit, allowed to call `traceRay`
+@miss
+fn miss() {}
+```
 ### Acceleration structure tags
 
 These are tags that can be added to a acceleration structure (`acceleration_structure` ->

naga/src/back/dot/mod.rs

@@ -403,6 +403,22 @@ impl StatementGraph {
                         },
                     }
                 }
+                S::RayPipelineFunction(func) => match func {
+                    crate::RayPipelineFunction::TraceRay {
+                        acceleration_structure,
+                        descriptor,
+                        payload,
+                    } => {
+                        self.dependencies.push((
+                            id,
+                            acceleration_structure,
+                            "acceleration_structure",
+                        ));
+                        self.dependencies.push((id, descriptor, "descriptor"));
+                        self.dependencies.push((id, payload, "payload"));
+                        "TraceRay"
+                    }
+                },
             };
             // Set the last node to the merge node
             last_node = merge_id;

naga/src/back/glsl/mod.rs

@@ -140,7 +140,10 @@ impl crate::AddressSpace {
             | crate::AddressSpace::Storage { .. }
             | crate::AddressSpace::Handle
             | crate::AddressSpace::PushConstant
-            | crate::AddressSpace::TaskPayload => false,
+            | crate::AddressSpace::TaskPayload
+            // just a default impl, not really supported
+            | crate::AddressSpace::RayPayload
+            | crate::AddressSpace::IncomingRayPayload => false,
         }
     }
 }
@@ -504,7 +507,15 @@ impl fmt::Display for VaryingName<'_> {
                     (ShaderStage::Vertex, true) | (ShaderStage::Fragment, false) => "vs2fs",
                     // fragment to pipeline
                     (ShaderStage::Fragment, true) => "fs2p",
-                    (ShaderStage::Task | ShaderStage::Mesh, _) => unreachable!(),
+                    (
+                        ShaderStage::Task
+                        | ShaderStage::Mesh
+                        | ShaderStage::RayGeneration
+                        | ShaderStage::AnyHit
+                        | ShaderStage::ClosestHit
+                        | ShaderStage::Miss,
+                        _,
+                    ) => unreachable!(),
                 };
                 write!(f, "_{prefix}_location{location}",)
             }
@@ -521,7 +532,12 @@ impl ShaderStage {
             ShaderStage::Compute => "cs",
             ShaderStage::Fragment => "fs",
             ShaderStage::Vertex => "vs",
-            ShaderStage::Task | ShaderStage::Mesh => unreachable!(),
+            ShaderStage::Task
+            | ShaderStage::Mesh
+            | ShaderStage::RayGeneration
+            | ShaderStage::AnyHit
+            | ShaderStage::ClosestHit
+            | ShaderStage::Miss => unreachable!(),
         }
     }
 }
@@ -1309,6 +1325,10 @@ impl<'a, W: Write> Writer<'a, W> {
             crate::AddressSpace::Function => unreachable!(),
             // Textures and samplers are handled directly in `Writer::write`.
             crate::AddressSpace::Handle => unreachable!(),
+            // ray tracing pipelines unsupported
+            crate::AddressSpace::RayPayload | crate::AddressSpace::IncomingRayPayload => {
+                unreachable!()
+            }
         }
 
         Ok(())
@@ -1675,7 +1695,12 @@ impl<'a, W: Write> Writer<'a, W> {
             ShaderStage::Vertex => output,
             ShaderStage::Fragment => !output,
             ShaderStage::Compute => false,
-            ShaderStage::Task | ShaderStage::Mesh => unreachable!(),
+            ShaderStage::Task
+            | ShaderStage::Mesh
+            | ShaderStage::RayGeneration
+            | ShaderStage::AnyHit
+            | ShaderStage::ClosestHit
+            | ShaderStage::Miss => unreachable!(),
         };
 
         // Write the I/O locations, if allowed
@@ -2810,6 +2835,7 @@ impl<'a, W: Write> Writer<'a, W> {
                 }
                 writeln!(self.out, ");")?;
             }
+            Statement::RayPipelineFunction(_) => unimplemented!(),
         }
 
         Ok(())
@@ -5269,7 +5295,20 @@ const fn glsl_built_in(built_in: crate::BuiltIn, options: VaryingOptions) -> &'s
         | Bi::VertexCount
         | Bi::PrimitiveCount
         | Bi::Vertices
-        | Bi::Primitives => {
+        | Bi::Primitives
+        | 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 => {
             unimplemented!()
         }
     }
@@ -5287,7 +5326,7 @@ const fn glsl_storage_qualifier(space: crate::AddressSpace) -> Option<&'static s
         As::Handle => Some("uniform"),
         As::WorkGroup => Some("shared"),
         As::PushConstant => Some("uniform"),
-        As::TaskPayload => unreachable!(),
+        As::TaskPayload | As::RayPayload | As::IncomingRayPayload => unreachable!(),
     }
 }
 

naga/src/back/hlsl/conv.rs

@@ -192,6 +192,19 @@ impl crate::BuiltIn {
             | Self::PrimitiveCount
             | Self::Vertices
             | Self::Primitives => unreachable!(),
+            Self::RayInvocationId
+            | Self::NumRayInvocations
+            | Self::InstanceCustomData
+            | Self::GeometryIndex
+            | Self::WorldRayOrigin
+            | Self::WorldRayDirection
+            | Self::ObjectRayOrigin
+            | Self::ObjectRayDirection
+            | Self::RayTmin
+            | Self::RayTCurrentMax
+            | Self::ObjectToWorld
+            | Self::WorldToObject
+            | Self::HitKind => unreachable!(),
         })
     }
 }

naga/src/back/hlsl/mod.rs

@@ -285,6 +285,7 @@ impl crate::ShaderStage {
             Self::Compute => "cs",
             Self::Task => "as",
             Self::Mesh => "ms",
+            Self::RayGeneration | Self::AnyHit | Self::ClosestHit | Self::Miss => "lib",
         }
     }
 }

naga/src/back/hlsl/writer.rs

@@ -1008,6 +1008,9 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 write!(self.out, "ConstantBuffer<")?;
                 "b"
             }
+            crate::AddressSpace::RayPayload | crate::AddressSpace::IncomingRayPayload => {
+                unimplemented!()
+            }
         };
 
         // If the global is a push constant write the type now because it will be a
@@ -2756,6 +2759,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 }
                 writeln!(self.out, ");")?;
             }
+            Statement::RayPipelineFunction(_) => unreachable!(),
         }
 
         Ok(())
@@ -3086,7 +3090,9 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                                 | crate::AddressSpace::Private
                                 | crate::AddressSpace::WorkGroup
                                 | crate::AddressSpace::PushConstant
-                                | crate::AddressSpace::TaskPayload,
+                                | crate::AddressSpace::TaskPayload
+                                | crate::AddressSpace::RayPayload
+                                | crate::AddressSpace::IncomingRayPayload,
                             )
                             | None => true,
                             Some(crate::AddressSpace::Uniform) => {

naga/src/back/msl/mod.rs

@@ -718,7 +718,20 @@ impl ResolvedBinding {
                     | Bi::VertexCount
                     | Bi::PrimitiveCount
                     | Bi::Vertices
-                    | Bi::Primitives => unreachable!(),
+                    | Bi::Primitives
+                    | 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 => unreachable!(),
                 };
                 write!(out, "{name}")?;
             }

naga/src/back/msl/writer.rs

@@ -596,7 +596,9 @@ impl crate::AddressSpace {
             | Self::WorkGroup
             | Self::PushConstant
             | Self::Handle
-            | Self::TaskPayload => true,
+            | Self::TaskPayload
+            | Self::RayPayload
+            | Self::IncomingRayPayload => true,
             Self::Function => false,
         }
     }
@@ -609,7 +611,7 @@ impl crate::AddressSpace {
             // may end up with "const" even if the binding is read-write,
             // and that should be OK.
             Self::Storage { .. } => true,
-            Self::TaskPayload => unimplemented!(),
+            Self::TaskPayload | Self::RayPayload | Self::IncomingRayPayload => unimplemented!(),
             // These should always be read-write.
             Self::Private | Self::WorkGroup => false,
             // These translate to `constant` address space, no need for qualifiers.
@@ -624,9 +626,13 @@ impl crate::AddressSpace {
             Self::Handle => None,
             Self::Uniform | Self::PushConstant => Some("constant"),
             Self::Storage { .. } => Some("device"),
-            Self::Private | Self::Function => Some("thread"),
+            // note for `RayPayload`, this probably needs to be emulated as a
+            // private variable, as metal has essentially an inout input
+            // for where it is passed.
+            Self::Private | Self::Function | Self::RayPayload => Some("thread"),
             Self::WorkGroup => Some("threadgroup"),
             Self::TaskPayload => Some("object_data"),
+            Self::IncomingRayPayload => Some("ray_data"),
         }
     }
 }
@@ -4194,6 +4200,7 @@ impl<W: Write> Writer<W> {
                     }
                     writeln!(self.out, ");")?;
                 }
+                crate::Statement::RayPipelineFunction(_) => unreachable!(),
             }
         }
 
@@ -6672,6 +6679,10 @@ template <typename A>
                     false,
                 ),
                 crate::ShaderStage::Task | crate::ShaderStage::Mesh => unimplemented!(),
+                crate::ShaderStage::RayGeneration
+                | crate::ShaderStage::AnyHit
+                | crate::ShaderStage::ClosestHit
+                | crate::ShaderStage::Miss => unimplemented!(),
             };
 
             // Should this entry point be modified to do vertex pulling?
@@ -6744,6 +6755,8 @@ template <typename A>
                         crate::AddressSpace::Function
                         | crate::AddressSpace::Private
                         | crate::AddressSpace::WorkGroup => {}
+                        crate::AddressSpace::RayPayload
+                        | crate::AddressSpace::IncomingRayPayload => unimplemented!(),
                     }
                 }
                 if needs_buffer_sizes {

naga/src/back/pipeline_constants.rs

@@ -860,6 +860,17 @@ fn adjust_stmt(new_pos: &HandleVec<Expression, Handle<Expression>>, stmt: &mut S
                 crate::RayQueryFunction::Terminate => {}
             }
         }
+        Statement::RayPipelineFunction(ref mut func) => match *func {
+            crate::RayPipelineFunction::TraceRay {
+                ref mut acceleration_structure,
+                ref mut descriptor,
+                ref mut payload,
+            } => {
+                adjust(acceleration_structure);
+                adjust(descriptor);
+                adjust(payload);
+            }
+        },
         Statement::Break
         | Statement::Continue
         | Statement::Kill

naga/src/back/spv/block.rs

@@ -3679,6 +3679,7 @@ impl BlockContext<'_> {
                 } => {
                     self.write_subgroup_gather(mode, argument, result, &mut block)?;
                 }
+                Statement::RayPipelineFunction(_) => unreachable!(),
             }
         }
 

naga/src/back/spv/helpers.rs

@@ -55,6 +55,7 @@ pub(super) const fn map_storage_class(space: crate::AddressSpace) -> spirv::Stor
         crate::AddressSpace::WorkGroup => spirv::StorageClass::Workgroup,
         crate::AddressSpace::PushConstant => spirv::StorageClass::PushConstant,
         crate::AddressSpace::TaskPayload => unreachable!(),
+        crate::AddressSpace::IncomingRayPayload | crate::AddressSpace::RayPayload => unreachable!(),
     }
 }