[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!(),
  }
 }