Add interactive viewer

ray-tracing-egui/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "ray-tracing-egui"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+vulkano = "0.34.1"
+vulkano-shaders = "0.34.0"
+winit = "0.28"
+egui = "0.24"
+egui_winit_vulkano = "0.27.0"
+ray-tracing-core = { path = "../ray-tracing-core" }
+ray-tracing-scene = { path = "../ray-tracing-scene" }
+ray-tracing-renderer = { path = "../ray-tracing-renderer" }
+rayon = "1.10.0"
+rand = { version = "0.8.5", features = ["small_rng"] }
+clap = { version = "4.5.19", features = ["derive"] }

ray-tracing-egui/src/main.rs

@@ -0,0 +1,306 @@
+use egui_winit_vulkano::{Gui, GuiConfig};
+use render::render_thread;
+use setup::{get_compute_pipeline, get_framebuffers};
+use std::sync::Arc;
+use vulkano::{
+    buffer::{Buffer, BufferCreateInfo, BufferUsage},
+    command_buffer::{allocator::StandardCommandBufferAllocator, AutoCommandBufferBuilder},
+    descriptor_set::{
+        allocator::StandardDescriptorSetAllocator, PersistentDescriptorSet, WriteDescriptorSet,
+    },
+    memory::allocator::{
+        AllocationCreateInfo, MemoryAllocator, MemoryTypeFilter, StandardMemoryAllocator,
+    },
+    pipeline::Pipeline,
+    swapchain::{self, SwapchainCreateInfo, SwapchainPresentInfo},
+    sync::{self, future::FenceSignalFuture, GpuFuture},
+    Validated, VulkanError,
+};
+use winit::{
+    event::{Event, WindowEvent},
+    event_loop::ControlFlow,
+};
+
+mod render;
+mod setup;
+
+mod cs {
+    vulkano_shaders::shader! {
+        ty: "compute",
+        src: "
+            #version 460
+            layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;
+            layout(set = 0, binding = 0) uniform writeonly image2D outputImage;
+
+            layout(set = 0, binding = 1) uniform GPUSettings {
+                uint samples;
+            } settings;
+
+            layout(set = 0, binding = 2) buffer Data {
+                float data[];
+            } buf;
+
+
+            void main() {
+                uvec2 size = imageSize(outputImage);
+                uvec2 pos = gl_GlobalInvocationID.xy;
+                if (pos.x < size.x && pos.y < size.y) {
+                    uint index = (pos.y * size.x + pos.x) * 3;
+                    vec3 color = vec3(buf.data[index], buf.data[index + 1], buf.data[index + 2]);
+
+                    vec4 to_write = vec4(color / settings.samples, 1.0);
+                    imageStore(outputImage, ivec2(pos), to_write);
+                }
+            }
+        "
+    }
+}
+
+fn main() {
+    let (event_loop, physical_device, device, queue, window, surface) = setup::setup();
+
+    let (mut swapchain, mut images) = setup::create_swapchain(
+        physical_device.clone(),
+        device.clone(),
+        surface.clone(),
+        window.inner_size(),
+    );
+
+    let mut gui = Gui::new(
+        &event_loop,
+        surface.clone(),
+        queue.clone(),
+        swapchain.image_format(),
+        GuiConfig {
+            is_overlay: true,
+            ..Default::default()
+        },
+    );
+
+    let memory_allocator = Arc::new(StandardMemoryAllocator::new_default(device.clone()));
+
+    let command_buffer_allocator =
+        StandardCommandBufferAllocator::new(device.clone(), Default::default());
+    let descriptor_set_allocator =
+        StandardDescriptorSetAllocator::new(device.clone(), Default::default());
+
+    let (control_tx, control_rx) = std::sync::mpsc::channel();
+    let (data_tx, data_rx) = std::sync::mpsc::channel();
+
+    let mut settings = render::RenderSettings {
+        width: window.inner_size().width,
+        height: window.inner_size().height,
+    };
+
+    let send_settings = settings.clone();
+    let send_memory_allocator = memory_allocator.clone();
+
+    std::thread::spawn(move || {
+        render_thread(send_settings, control_rx, data_tx, send_memory_allocator)
+    });
+
+    let render_pass = setup::get_render_pass(device.clone(), &swapchain);
+    let mut framebuffers = setup::get_framebuffers(&images, &render_pass);
+
+    let mut recreate_swapchain = false;
+
+    let mut fences: Vec<Option<Arc<FenceSignalFuture<_>>>> = vec![None; images.len()];
+    let mut previous_fence_i = 0;
+
+    let mut buffer = None;
+
+    let cs = cs::load(device.clone()).unwrap();
+
+    let cpu_buffers: Vec<_> = framebuffers
+        .iter()
+        .map(|_| {
+            Buffer::from_data(
+                memory_allocator.clone(),
+                BufferCreateInfo {
+                    usage: BufferUsage::UNIFORM_BUFFER,
+                    ..Default::default()
+                },
+                AllocationCreateInfo {
+                    memory_type_filter: MemoryTypeFilter::PREFER_DEVICE
+                        | MemoryTypeFilter::HOST_SEQUENTIAL_WRITE,
+                    ..Default::default()
+                },
+                cs::GPUSettings { samples: 1 },
+            )
+            .unwrap()
+        })
+        .collect();
+
+    let mut pipeline = get_compute_pipeline(device.clone(), cs.clone());
+
+    event_loop.run(move |event, _, control_flow| match event {
+        Event::WindowEvent { window_id, event } => {
+            gui.update(&event);
+            match event {
+                WindowEvent::Resized(_) => {
+                    recreate_swapchain = true;
+                    window.request_redraw();
+                    settings.width = window.inner_size().width;
+                    settings.height = window.inner_size().height;
+                    let _ = control_tx.send(settings.clone());
+                }
+                WindowEvent::CloseRequested => {
+                    *control_flow = ControlFlow::Exit;
+                }
+                _ => (),
+            }
+        }
+        Event::MainEventsCleared => {
+            if recreate_swapchain {
+                let (new_swapchain, new_images) = swapchain
+                    .recreate(SwapchainCreateInfo {
+                        image_extent: window.inner_size().into(),
+                        ..swapchain.create_info()
+                    })
+                    .expect("failed to recreate swapchain: {e}");
+                swapchain = new_swapchain;
+                framebuffers = get_framebuffers(&new_images, &render_pass);
+                pipeline = get_compute_pipeline(device.clone(), cs.clone());
+            }
+
+            gui.immediate_ui(|gui| {
+                let ctx = gui.context();
+                egui::Window::new("panel").show(&ctx, |ui| {
+                    ui.heading("fdjkslf");
+
+                    ui.collapsing("Render Settings:", |ui| {
+                        ui.label(format!("{:#?}", &settings));
+                    });
+                });
+            });
+
+            let (image_i, suboptimal, acquire_future) =
+                match swapchain::acquire_next_image(swapchain.clone(), None)
+                    .map_err(Validated::unwrap)
+                {
+                    Ok(r) => r,
+                    Err(e) => panic!("failed to acquire next image: {e}"),
+                };
+
+            let previous_future = match fences[previous_fence_i as usize].clone() {
+                // Create a NowFuture
+                None => {
+                    let mut now = sync::now(device.clone());
+                    now.cleanup_finished();
+
+                    now.boxed()
+                }
+                // Use the existing FenceSignalFuture
+                Some(fence) => fence.boxed(),
+            };
+
+            if suboptimal {
+                recreate_swapchain = true;
+            }
+
+            if let Some(image_fence) = &fences[image_i as usize] {
+                image_fence.wait(None).unwrap();
+            }
+
+            while let Ok(b) = data_rx.try_recv() {
+                buffer = Some(b);
+            }
+
+            if buffer.as_ref().is_some_and(|b| {
+                b.width != window.inner_size().width || b.height != window.inner_size().height
+            }) {
+                buffer = None;
+            }
+
+            let compute_future = match &buffer {
+                Some(b) => {
+                    match cpu_buffers[image_i as usize].write() {
+                        Ok(mut local) => {
+                            local.samples = b.samples;
+                        }
+                        Err(e) => {
+                            println!("{}", e);
+                        }
+                    }
+
+                    let mut builder = AutoCommandBufferBuilder::primary(
+                        &command_buffer_allocator,
+                        queue.queue_family_index(),
+                        vulkano::command_buffer::CommandBufferUsage::OneTimeSubmit,
+                    )
+                    .unwrap();
+
+                    let pipeline_layout = pipeline.layout();
+                    let descriptor_set_layouts = pipeline_layout.set_layouts();
+
+                    let descriptor_set_layout = descriptor_set_layouts.first().unwrap();
+
+                    let descriptor_set = PersistentDescriptorSet::new(
+                        &descriptor_set_allocator,
+                        descriptor_set_layout.clone(),
+                        [
+                            WriteDescriptorSet::image_view(
+                                0,
+                                framebuffers[image_i as usize].attachments()[0].clone(),
+                            ),
+                            WriteDescriptorSet::buffer(1, cpu_buffers[image_i as usize].clone()),
+                            WriteDescriptorSet::buffer(2, b.buffer.clone()),
+                        ],
+                        [],
+                    )
+                    .unwrap();
+
+                    builder
+                        .bind_pipeline_compute(pipeline.clone())
+                        .unwrap()
+                        .bind_descriptor_sets(
+                            vulkano::pipeline::PipelineBindPoint::Compute,
+                            pipeline_layout.clone(),
+                            0,
+                            descriptor_set,
+                        )
+                        .unwrap()
+                        .dispatch([b.width.div_ceil(8), b.height.div_ceil(8), 1])
+                        .unwrap();
+
+                    let command = builder.build().unwrap();
+
+                    acquire_future
+                        .join(previous_future)
+                        .then_execute(queue.clone(), command)
+                        .unwrap()
+                        .then_signal_fence()
+                        .boxed()
+                }
+                None => acquire_future.join(previous_future).boxed(),
+            };
+
+            let gui_future = gui.draw_on_image(
+                compute_future,
+                framebuffers[image_i as usize].attachments()[0].clone(),
+            );
+
+            let future = gui_future
+                .then_swapchain_present(
+                    queue.clone(),
+                    SwapchainPresentInfo::swapchain_image_index(swapchain.clone(), image_i),
+                )
+                .then_signal_fence_and_flush();
+
+            fences[image_i as usize] = match future.map_err(Validated::unwrap) {
+                Ok(value) => Some(Arc::new(value)),
+                Err(VulkanError::OutOfDate) => {
+                    recreate_swapchain = true;
+                    None
+                }
+                Err(e) => {
+                    println!("failed to flush future: {e}");
+                    None
+                }
+            };
+
+            previous_fence_i = image_i;
+        }
+        _ => (),
+    })
+}

ray-tracing-egui/src/render.rs

@@ -0,0 +1,111 @@
+use std::sync::{
+    mpsc::{Receiver, Sender},
+    Arc,
+};
+
+use rand::{rngs::SmallRng, SeedableRng};
+use ray_tracing_core::{camera::BasicCamera, renderer::ClassicalRenderer};
+use ray_tracing_renderer::path_tracer_importance::PathTracerImportance;
+use ray_tracing_scene::examples::{self, example_scenes, ExampleScene};
+use rayon::{
+    iter::{IndexedParallelIterator, ParallelIterator},
+    slice::ParallelSliceMut,
+};
+use vulkano::{
+    buffer::{Buffer, BufferCreateInfo, BufferUsage, Subbuffer},
+    memory::allocator::{
+        AllocationCreateInfo, FreeListAllocator, GenericMemoryAllocator, MemoryAllocator,
+        MemoryTypeFilter,
+    },
+};
+
+#[derive(Debug, Clone)]
+pub struct RenderSettings {
+    pub width: u32,
+    pub height: u32,
+}
+
+pub enum ControlMessages {
+    SetSettings(RenderSettings),
+}
+
+pub struct Data {
+    pub width: u32,
+    pub height: u32,
+    pub samples: u32,
+    pub buffer: Subbuffer<[f32]>,
+}
+
+pub fn render_thread(
+    s: RenderSettings,
+    rx: Receiver<RenderSettings>,
+    tx: Sender<Data>,
+    allocator: Arc<impl MemoryAllocator>,
+) {
+    let mut settings = s;
+
+    let mut buffer = vec![0.0; settings.width as usize * settings.height as usize * 3];
+
+    let e = examples::basic_cornell();
+
+    let scene = e.scene;
+
+    let camera = BasicCamera::new(
+        settings.width,
+        settings.height,
+        e.camera_pos,
+        e.camera_dir,
+        e.camera_up,
+        e.horizontal_fov,
+    );
+
+    let mut renderer = PathTracerImportance::new(settings.width, settings.height);
+
+    let mut samples = 0;
+    loop {
+        while let Ok(s) = rx.try_recv() {
+            settings = s;
+            buffer = vec![0.0; settings.width as usize * settings.height as usize * 3];
+            renderer = PathTracerImportance::new(settings.width, settings.height);
+            samples = 0;
+        }
+
+        buffer.par_chunks_mut(3).enumerate().for_each(|(i, c)| {
+            let x = (i % settings.width as usize) as u32;
+            let y = (i / settings.height as usize) as u32;
+
+            let mut rng = SmallRng::seed_from_u64(
+                (x + y * settings.width) as u64
+                    + (settings.width as u64 * settings.height as u64 * samples as u64),
+            );
+            let r = renderer.render_pixel(&scene, &camera, x, y, &mut rng);
+            c[0] += r.r();
+            c[1] += r.g();
+            c[2] += r.b();
+        });
+
+        samples += 1;
+
+        let data = Data {
+            width: settings.width,
+            height: settings.height,
+            samples,
+            buffer: Buffer::from_iter(
+                allocator.clone(),
+                BufferCreateInfo {
+                    usage: BufferUsage::STORAGE_BUFFER,
+                    ..Default::default()
+                },
+                AllocationCreateInfo {
+                    memory_type_filter: MemoryTypeFilter::PREFER_DEVICE
+                        | MemoryTypeFilter::HOST_SEQUENTIAL_WRITE,
+                    ..Default::default()
+                },
+                buffer.iter().copied(),
+            )
+            .unwrap(),
+        };
+
+        let _ = tx.send(data);
+    }
+}

ray-tracing-egui/src/setup.rs

@@ -0,0 +1,206 @@
+use std::sync::Arc;
+
+use vulkano::{
+    device::{
+        physical::{PhysicalDevice, PhysicalDeviceType},
+        Device, DeviceCreateInfo, DeviceExtensions, Features, Queue, QueueCreateInfo, QueueFlags,
+    },
+    format::Format,
+    image::{view::ImageView, Image, ImageUsage},
+    instance::{Instance, InstanceCreateFlags, InstanceCreateInfo},
+    pipeline::{
+        compute::ComputePipelineCreateInfo, layout::PipelineDescriptorSetLayoutCreateInfo,
+        ComputePipeline, PipelineLayout, PipelineShaderStageCreateInfo,
+    },
+    render_pass::{Framebuffer, FramebufferCreateInfo, RenderPass},
+    shader::ShaderModule,
+    swapchain::{Surface, Swapchain, SwapchainCreateInfo},
+    VulkanLibrary,
+};
+use winit::{
+    dpi::PhysicalSize,
+    event_loop::{EventLoop, EventLoopBuilder},
+    window::{Window, WindowBuilder},
+};
+
+fn select_physical_device(
+    instance: &Arc<Instance>,
+    surface: &Arc<Surface>,
+    device_extensions: &DeviceExtensions,
+    device_features: &Features,
+) -> (Arc<PhysicalDevice>, u32) {
+    instance
+        .enumerate_physical_devices()
+        .expect("could not enumerate device")
+        .filter(|p| p.supported_extensions().contains(device_extensions))
+        .filter(|p| p.supported_features().contains(device_features))
+        .filter_map(|p| {
+            p.queue_family_properties()
+                .iter()
+                .enumerate()
+                .position(|(i, q)| {
+                    q.queue_flags.contains(QueueFlags::GRAPHICS)
+                        && p.surface_support(i as u32, surface).unwrap_or(false)
+                })
+                .map(|q| (p, q as u32))
+        })
+        .min_by_key(|(p, _)| match p.properties().device_type {
+            PhysicalDeviceType::DiscreteGpu => 0,
+            PhysicalDeviceType::IntegratedGpu => 1,
+            PhysicalDeviceType::VirtualGpu => 2,
+            PhysicalDeviceType::Cpu => 3,
+            _ => 4,
+        })
+        .expect("no device available")
+}
+
+pub fn setup() -> (
+    EventLoop<()>,
+    Arc<PhysicalDevice>,
+    Arc<Device>,
+    Arc<Queue>,
+    Arc<Window>,
+    Arc<Surface>,
+) {
+    let event_loop: EventLoop<()> = EventLoopBuilder::default().build();
+
+    let library = VulkanLibrary::new().unwrap();
+
+    let required_extensions = Surface::required_extensions(&event_loop);
+
+    let instance = Instance::new(
+        library,
+        InstanceCreateInfo {
+            flags: InstanceCreateFlags::ENUMERATE_PORTABILITY,
+            enabled_extensions: required_extensions,
+            ..Default::default()
+        },
+    )
+    .unwrap();
+
+    let window = Arc::new(WindowBuilder::new().build(&event_loop).unwrap());
+
+    let surface = Surface::from_window(instance.clone(), window.clone()).unwrap();
+
+    let device_extensions = DeviceExtensions {
+        khr_swapchain: true,
+        ..Default::default()
+    };
+
+    let device_features = Features {
+        shader_float64: true,
+        ..Features::empty()
+    };
+
+    let (physical_device, queue_family_index) =
+        select_physical_device(&instance, &surface, &device_extensions, &device_features);
+
+    let (device, mut queues) = Device::new(
+        physical_device.clone(),
+        DeviceCreateInfo {
+            queue_create_infos: vec![QueueCreateInfo {
+                queue_family_index,
+                ..Default::default()
+            }],
+            enabled_extensions: device_extensions,
+            enabled_features: device_features,
+            ..Default::default()
+        },
+    )
+    .unwrap();
+
+    let queue = queues.next().unwrap();
+
+    (event_loop, physical_device, device, queue, window, surface)
+}
+
+pub fn create_swapchain(
+    physical_device: Arc<PhysicalDevice>,
+    device: Arc<Device>,
+    surface: Arc<Surface>,
+    dimensions: PhysicalSize<u32>,
+) -> (Arc<Swapchain>, Vec<Arc<Image>>) {
+    let caps = physical_device
+        .surface_capabilities(&surface, Default::default())
+        .unwrap();
+    let image_format = physical_device
+        .surface_formats(&surface, Default::default())
+        .unwrap()[0]
+        .0;
+    let composite_alpha = caps.supported_composite_alpha.into_iter().next().unwrap();
+    Swapchain::new(
+        device.clone(),
+        surface.clone(),
+        SwapchainCreateInfo {
+            min_image_count: caps.min_image_count + 1,
+            image_format,
+            image_extent: dimensions.into(),
+            image_usage: ImageUsage::COLOR_ATTACHMENT.union(ImageUsage::STORAGE),
+            composite_alpha,
+            ..Default::default()
+        },
+    )
+    .unwrap()
+}
+
+pub fn get_render_pass(device: Arc<Device>, swapchain: &Arc<Swapchain>) -> Arc<RenderPass> {
+    vulkano::single_pass_renderpass!(
+        device,
+        attachments: {
+            color: {
+                format: swapchain.image_format(),
+                samples: 1,
+                load_op: Clear,
+                store_op: Store,
+            },
+        },
+        pass: {
+            color: [color],
+            depth_stencil: {},
+        },
+    )
+    .unwrap()
+}
+
+pub fn get_framebuffers(
+    images: &[Arc<Image>],
+    render_pass: &Arc<RenderPass>,
+) -> Vec<Arc<Framebuffer>> {
+    images
+        .iter()
+        .map(|image| {
+            let view = ImageView::new_default(image.clone()).unwrap();
+            Framebuffer::new(
+                render_pass.clone(),
+                FramebufferCreateInfo {
+                    attachments: vec![view],
+                    ..Default::default()
+                },
+            )
+            .unwrap()
+        })
+        .collect()
+}
+
+pub fn get_compute_pipeline(device: Arc<Device>, cs: Arc<ShaderModule>) -> Arc<ComputePipeline> {
+    let cs = cs.entry_point("main").unwrap();
+
+    let stage = PipelineShaderStageCreateInfo::new(cs);
+
+    let layout = PipelineLayout::new(
+        device.clone(),
+        PipelineDescriptorSetLayoutCreateInfo::from_stages([&stage])
+            .into_pipeline_layout_create_info(device.clone())
+            .unwrap(),
+    )
+    .unwrap();
+
+    ComputePipeline::new(
+        device,
+        None,
+        ComputePipelineCreateInfo {
+            ..ComputePipelineCreateInfo::stage_layout(stage, layout)
+        },
+    )
+    .unwrap()
+}