GPU-accelerated UI toolkit (Vulkan)
git clone https://git.lucas.co/cce-ui.git
feat: the path tracer draws a user image standing in the scene
set_rt_scene_with_image (and RtOffscreen::set_scene_with_image, which is
handed pixels) adds the image to the traced scene as a quad of two
triangles under a textured material, so both tiers meet it as they meet any
triangle and a scene that is an image alone is traced. Its colour is the
surface's albedo, sampled at the mip level a pixel's footprint asks for,
and a ray goes through where it is clear, by chance in proportion to the
alpha. A shared image is looked up every frame and this frame's bindings
pointed at it, since the 2D pass may replace it; a 1x1 stand-in is bound
while there is none. The denoiser's features are taken where a path first
lands, not on its first bounce.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
src/vk/image.rs | 7 +
src/vk/mod.rs | 2 +-
src/vk/renderer.rs | 23 ++-
src/vk/rt.rs | 472 +++++++++++++++++++++++++++++++++++++++++++++++++-
src/vk/rt_common.wgsl | 72 +++++++-
5 files changed, 563 insertions(+), 13 deletions(-)
diff --git a/src/vk/image.rs b/src/vk/image.rs
index 13a1f27..d228362 100644
--- a/src/vk/image.rs
+++ b/src/vk/image.rs
@@ -958,6 +958,13 @@ impl ImageStage {
self.images.get(&image_id).map(|gpu| gpu.descriptor_set)
}
+ /// An uploaded image's view and size, for a pass that samples it under
+ /// bindings of its own (the path tracer). None while its upload has not
+ /// landed.
+ pub(crate) fn view_and_size(&self, image_id: u32) -> Option<(vk::ImageView, u32, u32)> {
+ self.images.get(&image_id).map(|gpu| (gpu.view, gpu.width, gpu.height))
+ }
+
/// Record one image quad (index `i` of this frame's list). The caller
/// restores its own pipeline/scissor state afterwards. Returns false if the
/// image hasn't finished uploading (draw skipped).
diff --git a/src/vk/mod.rs b/src/vk/mod.rs
index a0242b6..3766eba 100644
--- a/src/vk/mod.rs
+++ b/src/vk/mod.rs
@@ -53,7 +53,7 @@ pub use image::{
upload_rgba_mipmapped, ImageQuad, PixelFormat,
};
pub use renderer::{Batch2D, Frame2D, PlatePush, VkRenderer, MAX_PLATE_FEATURES};
-pub use rt::{RtCamera, RtMaterial, RtOffscreen, RtTriangle};
+pub use rt::{RtCamera, RtImage, RtImagePixels, RtMaterial, RtOffscreen, RtTriangle};
pub use scene::{MeshId, SceneDraw, SceneImage, Vertex3D};
pub use text::TextSpan;
diff --git a/src/vk/renderer.rs b/src/vk/renderer.rs
index 7fd9a37..51f0f81 100644
--- a/src/vk/renderer.rs
+++ b/src/vk/renderer.rs
@@ -19,7 +19,7 @@ use crate::engine::Vertex;
use super::core::SurfaceLost;
use super::image::{ImageQuad, ImageStage};
-use super::rt::{RtCamera, RtMaterial, RtStage, RtTriangle};
+use super::rt::{RtCamera, RtImage, RtImageSource, RtMaterial, RtStage, RtTriangle};
use super::scene::{MeshId, SceneDraw, SceneImage, SceneStage, Vertex3D};
use super::text::{TextSpan, TextStage};
@@ -1495,6 +1495,19 @@ impl VkRenderer {
/// (geometry rebuilds), matching `update_mesh`. The first call compiles
/// the compute pipeline.
pub fn set_rt_scene(&mut self, triangles: &[RtTriangle], materials: &[RtMaterial]) {
+ self.set_rt_scene_with_image(triangles, materials, None);
+ }
+
+ /// [`set_rt_scene`](Self::set_rt_scene), with a user image standing in
+ /// the scene: the picture the raster pass draws as a `SceneImage`,
+ /// traced. The image's pixels changing is a change of scene like any
+ /// other — set it again, which restarts the accumulation.
+ pub fn set_rt_scene_with_image(
+ &mut self,
+ triangles: &[RtTriangle],
+ materials: &[RtMaterial],
+ image: Option<RtImage>,
+ ) {
unsafe {
let _ = self.core.device.device_wait_idle();
}
@@ -1508,6 +1521,8 @@ impl VkRenderer {
core.accel_loader.as_ref(),
core.as_scratch_align,
core.min_uniform_align,
+ core.queue,
+ core.command_pool,
)
});
rt.set_scene(
@@ -1517,6 +1532,7 @@ impl VkRenderer {
core.command_pool,
triangles,
materials,
+ image.map(|i| (RtImageSource::Shared(i.image), i.corners, i.opacity)),
);
}
@@ -1813,7 +1829,10 @@ impl VkRenderer {
clip_radius,
);
if let Some(rt) = self.rt.as_mut() {
- rt.write_frame_uniforms(frame_index);
+ let images = &self.image;
+ rt.write_frame_uniforms(&self.core.device, frame_index, &|id| {
+ images.view_and_size(id)
+ });
}
// Record.
diff --git a/src/vk/rt.rs b/src/vk/rt.rs
index f80b38e..bdb163c 100644
--- a/src/vk/rt.rs
+++ b/src/vk/rt.rs
@@ -39,6 +39,42 @@ pub struct RtMaterial {
pub emission: [f32; 3],
}
+/// An image standing in the traced scene: a quad whose surface is the
+/// image's colour, lit as any other surface is, and which lets a ray through
+/// where the image is clear. The tracer adds the quad to the scene itself.
+///
+/// The image is one the 2D pass already holds (an id from `upload_rgba`),
+/// so a picture shown in the raster viewport costs the tracer nothing more.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct RtImage {
+ pub image: u32,
+ /// The quad's corners in the scene's space, in the image's own order:
+ /// top-left, top-right, bottom-right, bottom-left.
+ pub corners: [[f32; 3]; 4],
+ /// Alpha multiplier over the image's own.
+ pub opacity: f32,
+}
+
+/// [`RtImage`] for the headless tracer, which has no 2D pass to share an
+/// image with and is handed the pixels: tightly packed sRGB RGBA8.
+#[derive(Debug, Clone, Copy)]
+pub struct RtImagePixels<'a> {
+ pub pixels: &'a [u8],
+ pub width: u32,
+ pub height: u32,
+ pub corners: [[f32; 3]; 4],
+ pub opacity: f32,
+}
+
+/// Where the stage's image comes from.
+pub(crate) enum RtImageSource<'a> {
+ /// The 2D pass's image of this id, looked up each frame: its upload may
+ /// not have landed when the scene is set, and it may be replaced after.
+ Shared(u32),
+ /// Pixels for the stage to upload and own.
+ Pixels { pixels: &'a [u8], width: u32, height: u32 },
+}
+
/// The full camera: the inverse of the raster path's `proj * view * model`.
/// Rays are unprojected from NDC through it, so any matrix stack that renders
/// the raster viewport drives the tracer unchanged.
@@ -84,6 +120,9 @@ struct RtParams {
max_bounces: u32,
spp: u32,
_pad: [u32; 3],
+ img_origin: [f32; 4],
+ img_u: [f32; 4],
+ img_v: [f32; 4],
}
// --- BVH construction (binned SAH) ---
@@ -300,6 +339,185 @@ struct RtFrame {
descriptor_set: vk::DescriptorSet,
}
+/// A texture the stage made and owns: the 1x1 stand-in bound while the
+/// scene has no image, and the headless tracer's image.
+struct OwnedTexture {
+ image: vk::Image,
+ view: vk::ImageView,
+ allocation: Option<Allocation>,
+ size: (u32, u32),
+}
+
+impl OwnedTexture {
+ /// Upload sRGB RGBA8 pixels as a one-level texture, with a blocking
+ /// one-time submit. One level: the headless tracer renders a still, and
+ /// its samples average what a mip chain would have.
+ fn new(
+ device: &ash::Device,
+ allocator: &mut Allocator,
+ queue: vk::Queue,
+ command_pool: vk::CommandPool,
+ pixels: &[u8],
+ width: u32,
+ height: u32,
+ ) -> Self {
+ assert_eq!(pixels.len(), (width * height * 4) as usize, "8888 size mismatch");
+ let range = vk::ImageSubresourceRange::default()
+ .aspect_mask(vk::ImageAspectFlags::COLOR)
+ .level_count(1)
+ .layer_count(1);
+ unsafe {
+ let image = device
+ .create_image(
+ &vk::ImageCreateInfo::default()
+ .image_type(vk::ImageType::TYPE_2D)
+ .format(vk::Format::R8G8B8A8_SRGB)
+ .extent(vk::Extent3D { width, height, depth: 1 })
+ .mip_levels(1)
+ .array_layers(1)
+ .samples(vk::SampleCountFlags::TYPE_1)
+ .tiling(vk::ImageTiling::OPTIMAL)
+ .usage(vk::ImageUsageFlags::SAMPLED | vk::ImageUsageFlags::TRANSFER_DST)
+ .initial_layout(vk::ImageLayout::UNDEFINED),
+ None,
+ )
+ .expect("Failed to create RT texture");
+ let requirements = device.get_image_memory_requirements(image);
+ let allocation = allocator
+ .allocate(&AllocationCreateDesc {
+ name: "rt-texture",
+ requirements,
+ location: MemoryLocation::GpuOnly,
+ linear: false,
+ allocation_scheme: AllocationScheme::GpuAllocatorManaged,
+ })
+ .expect("Failed to allocate RT texture memory");
+ device
+ .bind_image_memory(image, allocation.memory(), allocation.offset())
+ .expect("Failed to bind RT texture memory");
+ let mut staging = create_cpu_buffer(
+ device,
+ allocator,
+ pixels.len() as vk::DeviceSize,
+ vk::BufferUsageFlags::TRANSFER_SRC,
+ "rt-texture-staging",
+ );
+ staging.allocation.as_mut().unwrap().mapped_slice_mut().unwrap()[..pixels.len()]
+ .copy_from_slice(pixels);
+
+ let cmd = device
+ .allocate_command_buffers(
+ &vk::CommandBufferAllocateInfo::default()
+ .command_pool(command_pool)
+ .level(vk::CommandBufferLevel::PRIMARY)
+ .command_buffer_count(1),
+ )
+ .expect("Failed to allocate RT texture command buffer")[0];
+ device
+ .begin_command_buffer(
+ cmd,
+ &vk::CommandBufferBeginInfo::default()
+ .flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT),
+ )
+ .unwrap();
+ let barrier = |from_access, to_access, from_layout, to_layout, from_stage, to_stage| {
+ device.cmd_pipeline_barrier(
+ cmd,
+ from_stage,
+ to_stage,
+ vk::DependencyFlags::empty(),
+ &[],
+ &[],
+ &[vk::ImageMemoryBarrier::default()
+ .src_access_mask(from_access)
+ .dst_access_mask(to_access)
+ .old_layout(from_layout)
+ .new_layout(to_layout)
+ .src_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
+ .dst_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
+ .image(image)
+ .subresource_range(range)],
+ );
+ };
+ barrier(
+ vk::AccessFlags::empty(),
+ vk::AccessFlags::TRANSFER_WRITE,
+ vk::ImageLayout::UNDEFINED,
+ vk::ImageLayout::TRANSFER_DST_OPTIMAL,
+ vk::PipelineStageFlags::TOP_OF_PIPE,
+ vk::PipelineStageFlags::TRANSFER,
+ );
+ device.cmd_copy_buffer_to_image(
+ cmd,
+ staging.buffer,
+ image,
+ vk::ImageLayout::TRANSFER_DST_OPTIMAL,
+ &[vk::BufferImageCopy::default()
+ .buffer_row_length(width)
+ .buffer_image_height(height)
+ .image_subresource(
+ vk::ImageSubresourceLayers::default()
+ .aspect_mask(vk::ImageAspectFlags::COLOR)
+ .layer_count(1),
+ )
+ .image_extent(vk::Extent3D { width, height, depth: 1 })],
+ );
+ barrier(
+ vk::AccessFlags::TRANSFER_WRITE,
+ vk::AccessFlags::SHADER_READ,
+ vk::ImageLayout::TRANSFER_DST_OPTIMAL,
+ vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL,
+ vk::PipelineStageFlags::TRANSFER,
+ vk::PipelineStageFlags::COMPUTE_SHADER,
+ );
+ device.end_command_buffer(cmd).unwrap();
+ let cmds = [cmd];
+ device
+ .queue_submit(
+ queue,
+ &[vk::SubmitInfo::default().command_buffers(&cmds)],
+ vk::Fence::null(),
+ )
+ .expect("RT texture upload submit failed");
+ device.queue_wait_idle(queue).expect("RT texture upload wait failed");
+ device.free_command_buffers(command_pool, &cmds);
+ destroy_cpu_buffer(device, allocator, &mut staging);
+
+ let view = device
+ .create_image_view(
+ &vk::ImageViewCreateInfo::default()
+ .image(image)
+ .view_type(vk::ImageViewType::TYPE_2D)
+ .format(vk::Format::R8G8B8A8_SRGB)
+ .subresource_range(range),
+ None,
+ )
+ .expect("Failed to create RT texture view");
+ OwnedTexture { image, view, allocation: Some(allocation), size: (width, height) }
+ }
+ }
+
+ /// Caller must have the device idle.
+ fn destroy(&mut self, device: &ash::Device, allocator: &mut Allocator) {
+ unsafe {
+ device.destroy_image_view(self.view, None);
+ device.destroy_image(self.image, None);
+ }
+ if let Some(a) = self.allocation.take() {
+ let _ = allocator.free(a);
+ }
+ }
+}
+
+/// The scene's image as the stage holds it.
+struct StagedImage {
+ /// The 2D pass's image of this id, or None for one the stage owns.
+ shared: Option<u32>,
+ owned: Option<OwnedTexture>,
+ corners: [[f32; 3]; 4],
+ opacity: f32,
+}
+
#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct DenoiseParams {
@@ -620,6 +838,12 @@ pub(crate) struct RtStage {
materials: AllocatedBuffer,
tri_count: u32,
+ /// Bound at the image bindings while the scene has no image, or has one
+ /// whose upload has not landed: a shader's bindings are never empty.
+ stand_in: OwnedTexture,
+ image_sampler: vk::Sampler,
+ image: Option<StagedImage>,
+
accum: AllocatedBuffer,
/// Primary-hit features (2 vec4 per pixel) written by the tracer, read
/// by the denoiser.
@@ -653,6 +877,8 @@ impl RtStage {
accel_loader: Option<&ash::khr::acceleration_structure::Device>,
as_scratch_align: vk::DeviceSize,
min_uniform_align: vk::DeviceSize,
+ queue: vk::Queue,
+ command_pool: vk::CommandPool,
) -> Self {
let force_compute = std::env::var("CCE_VK_RT").is_ok_and(|v| v == "compute");
let denoise_on = !std::env::var("CCE_VK_RT_DENOISE")
@@ -710,6 +936,16 @@ impl RtStage {
.descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
.descriptor_count(1)
.stage_flags(vk::ShaderStageFlags::COMPUTE),
+ vk::DescriptorSetLayoutBinding::default()
+ .binding(7)
+ .descriptor_type(vk::DescriptorType::SAMPLED_IMAGE)
+ .descriptor_count(1)
+ .stage_flags(vk::ShaderStageFlags::COMPUTE),
+ vk::DescriptorSetLayoutBinding::default()
+ .binding(8)
+ .descriptor_type(vk::DescriptorType::SAMPLER)
+ .descriptor_count(1)
+ .stage_flags(vk::ShaderStageFlags::COMPUTE),
];
let descriptor_set_layout = device
.create_descriptor_set_layout(
@@ -766,6 +1002,12 @@ impl RtStage {
vk::DescriptorPoolSize::default()
.ty(vk::DescriptorType::STORAGE_IMAGE)
.descriptor_count(n),
+ vk::DescriptorPoolSize::default()
+ .ty(vk::DescriptorType::SAMPLED_IMAGE)
+ .descriptor_count(n),
+ vk::DescriptorPoolSize::default()
+ .ty(vk::DescriptorType::SAMPLER)
+ .descriptor_count(n),
];
if tier == RtTier::RayQuery {
pool_sizes.push(
@@ -820,6 +1062,34 @@ impl RtStage {
let denoiser = denoise_on
.then(|| Denoiser::new(device, allocator, frames_in_flight, min_uniform_align));
+ let stand_in =
+ OwnedTexture::new(device, allocator, queue, command_pool, &[255; 4], 1, 1);
+ // Linear within a level and between levels; the shader names
+ // the level, since a compute shader has no derivatives to
+ // choose one by.
+ let image_sampler = device
+ .create_sampler(
+ &vk::SamplerCreateInfo::default()
+ .mag_filter(vk::Filter::LINEAR)
+ .min_filter(vk::Filter::LINEAR)
+ .mipmap_mode(vk::SamplerMipmapMode::LINEAR)
+ .min_lod(0.0)
+ .max_lod(vk::LOD_CLAMP_NONE)
+ .address_mode_u(vk::SamplerAddressMode::CLAMP_TO_EDGE)
+ .address_mode_v(vk::SamplerAddressMode::CLAMP_TO_EDGE)
+ .address_mode_w(vk::SamplerAddressMode::CLAMP_TO_EDGE),
+ None,
+ )
+ .expect("Failed to create RT image sampler");
+ for frame in &frames {
+ Self::write_image_descriptor(
+ device,
+ frame.descriptor_set,
+ stand_in.view,
+ image_sampler,
+ );
+ }
+
RtStage {
tier,
accel_loader: accel_loader.cloned(),
@@ -835,6 +1105,9 @@ impl RtStage {
tris: AllocatedBuffer::null(),
materials: AllocatedBuffer::null(),
tri_count: 0,
+ stand_in,
+ image_sampler,
+ image: None,
accum: AllocatedBuffer::null(),
features: AllocatedBuffer::null(),
ping: AllocatedBuffer::null(),
@@ -855,9 +1128,43 @@ impl RtStage {
}
}
+ fn write_image_descriptor(
+ device: &ash::Device,
+ set: vk::DescriptorSet,
+ view: vk::ImageView,
+ sampler: vk::Sampler,
+ ) {
+ let image_infos = [vk::DescriptorImageInfo::default()
+ .image_view(view)
+ .image_layout(vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL)];
+ let sampler_infos = [vk::DescriptorImageInfo::default().sampler(sampler)];
+ unsafe {
+ device.update_descriptor_sets(
+ &[
+ vk::WriteDescriptorSet::default()
+ .dst_set(set)
+ .dst_binding(7)
+ .descriptor_type(vk::DescriptorType::SAMPLED_IMAGE)
+ .image_info(&image_infos),
+ vk::WriteDescriptorSet::default()
+ .dst_set(set)
+ .dst_binding(8)
+ .descriptor_type(vk::DescriptorType::SAMPLER)
+ .image_info(&sampler_infos),
+ ],
+ &[],
+ );
+ }
+ }
+
/// Replace the scene. Tier 1 builds the BVH on the CPU (reordering a copy
/// of the triangles); tier 2 builds driver acceleration structures on the
/// given queue instead. Caller must have the device idle.
+ ///
+ /// An image joins the scene as a quad of two triangles under a material
+ /// of its own, marked textured — so both tiers meet it as they meet any
+ /// triangle, and a scene that is an image alone is not an empty one.
+ #[allow(clippy::too_many_arguments)]
pub(crate) fn set_scene(
&mut self,
device: &ash::Device,
@@ -866,8 +1173,38 @@ impl RtStage {
command_pool: vk::CommandPool,
triangles: &[RtTriangle],
materials: &[RtMaterial],
+ image: Option<(RtImageSource, [[f32; 3]; 4], f32)>,
) {
let mut tris: Vec<RtTriangle> = triangles.to_vec();
+ // The index the image's material will have, past the scene's own
+ // (or past the one stood in for a scene that names none).
+ let image_material = materials.len().max(1) as u32;
+ if let Some((_, [tl, tr, br, bl], _)) = &image {
+ tris.push(RtTriangle { p0: *tl, p1: *bl, p2: *tr, material: image_material });
+ tris.push(RtTriangle { p0: *tr, p1: *bl, p2: *br, material: image_material });
+ }
+ if let Some(mut old) = self.image.take().and_then(|i| i.owned) {
+ old.destroy(device, allocator);
+ }
+ self.image = image.map(|(source, corners, opacity)| match source {
+ RtImageSource::Shared(id) => {
+ StagedImage { shared: Some(id), owned: None, corners, opacity }
+ }
+ RtImageSource::Pixels { pixels, width, height } => StagedImage {
+ shared: None,
+ owned: Some(OwnedTexture::new(
+ device,
+ allocator,
+ queue,
+ command_pool,
+ pixels,
+ width,
+ height,
+ )),
+ corners,
+ opacity,
+ },
+ });
let nodes = match self.tier {
RtTier::Compute => build_bvh(&mut tris),
RtTier::RayQuery => Vec::new(),
@@ -880,7 +1217,7 @@ impl RtStage {
p2: [t.p2[0], t.p2[1], t.p2[2], 0.0],
})
.collect();
- let gpu_mats: Vec<GpuMaterial> = if materials.is_empty() {
+ let mut gpu_mats: Vec<GpuMaterial> = if materials.is_empty() {
vec![GpuMaterial { albedo: [0.8, 0.8, 0.8, 0.0], emission: [0.0; 4] }]
} else {
materials
@@ -891,6 +1228,11 @@ impl RtStage {
})
.collect()
};
+ if self.image.is_some() {
+ // albedo.w marks it textured: the shader takes the colour from
+ // the image.
+ gpu_mats.push(GpuMaterial { albedo: [1.0, 1.0, 1.0, 1.0], emission: [0.0; 4] });
+ }
self.destroy_accel(device, allocator);
for buf in [&mut self.nodes, &mut self.tris, &mut self.materials] {
@@ -1444,12 +1786,49 @@ impl RtStage {
self.tri_count > 0 && self.sample_index < MAX_SAMPLES
}
- /// After the frame fence: write this frame's params.
- pub(crate) fn write_frame_uniforms(&mut self, frame_index: usize) {
+ /// After the frame fence: write this frame's params, and point this
+ /// frame's image bindings at the scene's image as it is NOW.
+ ///
+ /// Each frame, because a shared image is the 2D pass's to replace: its
+ /// upload may land after the scene was set, and freeing it destroys the
+ /// view. This frame's set is past its fence, so it is safe to rewrite,
+ /// and it is rewritten before every use — a view that is gone is never
+ /// one a dispatch reads. `shared` looks an id up: its view and size.
+ pub(crate) fn write_frame_uniforms(
+ &mut self,
+ device: &ash::Device,
+ frame_index: usize,
+ shared: &dyn Fn(u32) -> Option<(vk::ImageView, u32, u32)>,
+ ) {
if !self.staged || self.tri_count == 0 {
return;
}
let Some(camera) = self.camera else { return };
+ let bound = self.image.as_ref().and_then(|image| {
+ let (view, w, h) = match (&image.owned, image.shared) {
+ (Some(owned), _) => (owned.view, owned.size.0, owned.size.1),
+ (None, Some(id)) => shared(id)?,
+ (None, None) => return None,
+ };
+ Some((view, w, h, image.corners, image.opacity))
+ });
+ let (view, img_origin, img_u, img_v) = match bound {
+ Some((view, w, h, [tl, tr, _, bl], opacity)) => (
+ view,
+ [tl[0], tl[1], tl[2], opacity.clamp(0.0, 1.0)],
+ [tr[0] - tl[0], tr[1] - tl[1], tr[2] - tl[2], w as f32],
+ [bl[0] - tl[0], bl[1] - tl[1], bl[2] - tl[2], h as f32],
+ ),
+ // No image, or one not there yet: opacity 0 lets every ray
+ // through the quad, and the stand-in is never sampled for it.
+ None => (self.stand_in.view, [0.0; 4], [1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 0.0, 1.0]),
+ };
+ Self::write_image_descriptor(
+ device,
+ self.frames[frame_index].descriptor_set,
+ view,
+ self.image_sampler,
+ );
let params = RtParams {
inv_mvp: camera.inv_mvp,
width: self.output_size.0,
@@ -1458,6 +1837,9 @@ impl RtStage {
max_bounces: MAX_BOUNCES,
spp: self.spp,
_pad: [0; 3],
+ img_origin,
+ img_u,
+ img_v,
};
let frame = &mut self.frames[frame_index];
frame.uniforms.allocation.as_mut().unwrap().mapped_slice_mut().unwrap()
@@ -1685,7 +2067,12 @@ impl RtStage {
for buf in [&mut self.nodes, &mut self.tris, &mut self.materials] {
destroy_cpu_buffer(device, allocator, buf);
}
+ if let Some(mut owned) = self.image.take().and_then(|i| i.owned) {
+ owned.destroy(device, allocator);
+ }
+ self.stand_in.destroy(device, allocator);
unsafe {
+ device.destroy_sampler(self.image_sampler, None);
for frame in &mut self.frames {
let mut uniforms = std::mem::replace(&mut frame.uniforms, AllocatedBuffer::null());
destroy_cpu_buffer(device, allocator, &mut uniforms);
@@ -1732,6 +2119,7 @@ impl RtOffscreen {
let accel_loader = core.accel_loader.clone();
let as_scratch_align = core.as_scratch_align;
let min_uniform_align = core.min_uniform_align;
+ let (queue, command_pool) = (core.queue, core.command_pool);
let allocator = core.allocator.as_mut().unwrap();
let stage = RtStage::new(
&device,
@@ -1740,6 +2128,8 @@ impl RtOffscreen {
accel_loader.as_ref(),
as_scratch_align,
min_uniform_align,
+ queue,
+ command_pool,
);
unsafe {
let cmd = device
@@ -1768,6 +2158,16 @@ impl RtOffscreen {
/// Replace the scene (same schema as `VkRenderer::set_rt_scene`).
pub fn set_scene(&mut self, triangles: &[RtTriangle], materials: &[RtMaterial]) {
+ self.set_scene_with_image(triangles, materials, None);
+ }
+
+ /// [`set_scene`](Self::set_scene), with an image standing in the scene.
+ pub fn set_scene_with_image(
+ &mut self,
+ triangles: &[RtTriangle],
+ materials: &[RtMaterial],
+ image: Option<RtImagePixels>,
+ ) {
unsafe {
let _ = self.core.device.device_wait_idle();
}
@@ -1781,6 +2181,13 @@ impl RtOffscreen {
command_pool,
triangles,
materials,
+ image.map(|i| {
+ (
+ RtImageSource::Pixels { pixels: i.pixels, width: i.width, height: i.height },
+ i.corners,
+ i.opacity,
+ )
+ }),
);
}
@@ -1815,7 +2222,7 @@ impl RtOffscreen {
// stage() resets sample_index when the camera or size changed —
// keep our resume point, not the reset, after the first chunk.
self.stage.sample_index = done;
- self.stage.write_frame_uniforms(0);
+ self.stage.write_frame_uniforms(&device, 0, &|_| None);
unsafe {
device
.begin_command_buffer(self.cmd, &vk::CommandBufferBeginInfo::default())
@@ -2238,6 +2645,63 @@ mod tests {
assert!(!denoise.is_empty());
}
+ /// An image in the traced scene, end to end: a quad half red and half
+ /// clear, in front of a green wall. The red half shows red, the clear
+ /// half shows the wall behind it, and beside the quad is the wall too.
+ /// Run with: cargo test --lib vk::rt -- --ignored
+ #[test]
+ #[ignore = "requires a Vulkan device"]
+ fn test_offscreen_renders_an_image() {
+ let mut off = RtOffscreen::new();
+ // 2x1: a red texel, a clear one.
+ let pixels = [255u8, 0, 0, 255, 0, 0, 0, 0];
+ let wall = |p0, p1, p2| RtTriangle { p0, p1, p2, material: 0 };
+ off.set_scene_with_image(
+ &[
+ wall([-9.0, -9.0, -1.0], [9.0, -9.0, -1.0], [9.0, 9.0, -1.0]),
+ wall([-9.0, -9.0, -1.0], [9.0, 9.0, -1.0], [-9.0, 9.0, -1.0]),
+ ],
+ &[RtMaterial { albedo: [0.1, 0.9, 0.1], emission: [0.0; 3] }],
+ Some(RtImagePixels {
+ pixels: &pixels,
+ width: 2,
+ height: 1,
+ corners: [[-1.0, 0.5, 0.0], [1.0, 0.5, 0.0], [1.0, -0.5, 0.0], [-1.0, -0.5, 0.0]],
+ opacity: 1.0,
+ }),
+ );
+ let proj = glam::Mat4::perspective_rh(0.9, 1.0, 0.1, 100.0);
+ let view = glam::Mat4::look_at_rh(
+ glam::Vec3::new(0.0, 0.0, 3.0),
+ glam::Vec3::ZERO,
+ glam::Vec3::Y,
+ );
+ let camera = RtCamera { inv_mvp: (proj * view).inverse().to_cols_array_2d() };
+ let (w, h) = (64u32, 64u32);
+ let px = off.render(camera, w, h, 64);
+ let at = |x: u32, y: u32| {
+ let i = ((y * w + x) * 4) as usize;
+ (px[i] as i32, px[i + 1] as i32, px[i + 2] as i32)
+ };
+ // The quad spans x in -1..1 of a view about 2.9 wide at z = 0: the
+ // pane's columns 10 to 54, and rows 21 to 43.
+ let (r, g, _) = at(16, 32);
+ assert!(r > g + 40, "the image's red half is not red: r={r} g={g}");
+ let (r, g, _) = at(48, 32);
+ assert!(g > r + 40, "the image's clear half hides the wall: r={r} g={g}");
+ let (r, g, _) = at(32, 6);
+ assert!(g > r + 40, "beside the image is not the wall: r={r} g={g}");
+
+ // Without its image the scene is the wall alone.
+ off.set_scene(
+ &[wall([-9.0, -9.0, -1.0], [9.0, -9.0, -1.0], [9.0, 9.0, -1.0])],
+ &[RtMaterial { albedo: [0.1, 0.9, 0.1], emission: [0.0; 3] }],
+ );
+ let px = off.render(camera, w, h, 16);
+ let i = ((32 * w + 40) * 4) as usize;
+ assert!(px[i + 1] > px[i], "the image outlived its scene");
+ }
+
/// End-to-end GPU test — needs a Vulkan device, so ignored by default.
/// Run with: cargo test --lib vk::rt -- --ignored
#[test]
diff --git a/src/vk/rt_common.wgsl b/src/vk/rt_common.wgsl
index 4f79b50..2b153da 100644
--- a/src/vk/rt_common.wgsl
+++ b/src/vk/rt_common.wgsl
@@ -28,6 +28,16 @@ struct Params {
_pad0: u32,
_pad1: u32,
_pad2: u32,
+ // The scene's image, a quad of two triangles whose material is marked
+ // textured (albedo.w). xyz = its top-left corner; w = its opacity, 0
+ // when there is no image to sample (a textured hit then lets the ray
+ // through).
+ img_origin: vec4<f32>,
+ // xyz = the top edge, corner to corner; w = the texture's width in
+ // texels.
+ img_u: vec4<f32>,
+ // xyz = the left edge, top to bottom; w = the texture's height.
+ img_v: vec4<f32>,
}
@group(0) @binding(0) var<uniform> params: Params;
@@ -58,6 +68,11 @@ struct Material {
// pixel: [2i] = (shading normal, hit t — 1e30 for sky), [2i+1] = (albedo, 0).
@group(0) @binding(6) var<storage, read_write> features: array<vec4<f32>>;
+// The scene's image and its sampler. Always bound: to a 1x1 stand-in while
+// the scene has no image.
+@group(0) @binding(7) var img: texture_2d<f32>;
+@group(0) @binding(8) var img_sampler: sampler;
+
// PCG (O'Neill) — one u32 of state per path, advanced per draw.
fn rand(state: ptr<function, u32>) -> f32 {
var s = *state * 747796405u + 2891336453u;
@@ -117,13 +132,27 @@ fn cs_main(@builtin(global_invocation_id) gid: vec3<u32>) {
let p_far = params.inv_mvp * vec4<f32>(ndc_x, ndc_y, 1.0, 1.0);
var ro = p_near.xyz / p_near.w;
var rd = normalize(p_far.xyz / p_far.w - ro);
+ // The angle one pixel subtends, for the image's mip level: the
+ // ray through the next pixel along, against this one.
+ let p_next = params.inv_mvp
+ * vec4<f32>(ndc_x + 2.0 / f32(params.width), ndc_y, 1.0, 1.0);
+ let pixel_angle = length(normalize(p_next.xyz / p_next.w - ro) - rd);
+ let eye = ro;
var radiance = vec3<f32>(0.0);
var throughput = vec3<f32>(1.0);
+ // Until the path first lands on something or leaves for the sky:
+ // what it lands on is the pixel's feature for the denoiser. Not
+ // "the first bounce" — a ray let through the image's clear texels
+ // has used one and landed on nothing.
+ var primary = s == 0u;
+ // Until the path first scatters it is the camera's own ray, and a
+ // pixel's footprint on what it hits is known.
+ var straight = true;
for (var bounce: u32 = 0u; bounce < params.max_bounces; bounce = bounce + 1u) {
let hit = intersect_scene(ro, rd);
if hit.t >= 1e30 {
- if s == 0u && bounce == 0u {
+ if primary {
features[2u * idx] = vec4<f32>(0.0, 0.0, 0.0, 1e30);
features[2u * idx + 1u] = vec4<f32>(1.0, 1.0, 1.0, 0.0);
}
@@ -132,17 +161,48 @@ fn cs_main(@builtin(global_invocation_id) gid: vec3<u32>) {
}
let tri = tris[hit.tri];
let mat = materials[bitcast<u32>(tri.p0.w)];
+ let at = ro + rd * hit.t;
+ var albedo = mat.albedo.rgb;
+ if mat.albedo.w > 0.5 {
+ // The image: its colour is the surface's, and where it is
+ // clear the ray goes on as if nothing were there — by
+ // chance, in proportion, which over the samples is the
+ // image's own alpha.
+ let rel = at - params.img_origin.xyz;
+ let u = params.img_u.xyz;
+ let v = params.img_v.xyz;
+ let uv = vec2<f32>(dot(rel, u) / dot(u, u), dot(rel, v) / dot(v, v));
+ // The level whose texel is a pixel's footprint wide. By the
+ // footprint's SHORT axis: seen at a slant the long one is
+ // averaged by the samples, where a level chosen for it
+ // would blur both. A scattered ray has no footprint and
+ // takes a coarse level.
+ var lod = 3.0;
+ if straight {
+ let footprint = length(at - eye) * pixel_angle;
+ let texel = length(u) / max(params.img_u.w, 1.0);
+ lod = max(log2(footprint / max(texel, 1e-12)), 0.0);
+ }
+ let texel = textureSampleLevel(img, img_sampler, uv, lod);
+ if rand(&rng) >= texel.a * params.img_origin.w {
+ ro = at + rd * (1e-4 * max(1.0, hit.t));
+ continue;
+ }
+ albedo = texel.rgb;
+ }
radiance = radiance + throughput * mat.emission.rgb;
var n = normalize(cross(tri.p1.xyz - tri.p0.xyz, tri.p2.xyz - tri.p0.xyz));
if dot(n, rd) > 0.0 {
n = -n;
}
- if s == 0u && bounce == 0u {
- features[2u * idx] = vec4<f32>(n, hit.t);
- features[2u * idx + 1u] = vec4<f32>(mat.albedo.rgb, 0.0);
+ if primary {
+ features[2u * idx] = vec4<f32>(n, length(at - eye));
+ features[2u * idx + 1u] = vec4<f32>(albedo, 0.0);
}
- throughput = throughput * mat.albedo.rgb;
- ro = ro + rd * hit.t + n * 1e-4;
+ primary = false;
+ straight = false;
+ throughput = throughput * albedo;
+ ro = at + n * 1e-4;
rd = cosine_dir(n, rand(&rng), rand(&rng));
}
// Firefly clamp: rare sun-spike paths otherwise leave speckles the