From 5814118d3fb97e0a7703ba9b5db3a443bfbd6837 Mon Sep 17 00:00:00 2001 From: Janis Date: Sun, 5 Jan 2025 01:16:54 +0100 Subject: [PATCH] egui renderpass --- crates/renderer/src/egui.rs | 710 ++++++++++++++++++++++++++++++++++++ crates/renderer/src/lib.rs | 2 + 2 files changed, 712 insertions(+) create mode 100644 crates/renderer/src/egui.rs diff --git a/crates/renderer/src/egui.rs b/crates/renderer/src/egui.rs new file mode 100644 index 0000000..d9cfc1d --- /dev/null +++ b/crates/renderer/src/egui.rs @@ -0,0 +1,710 @@ +use std::{collections::BTreeMap, sync::Arc}; + +use ash::{prelude::VkResult, vk}; +use indexmap::IndexMap; + +use crate::{ + buffers::{Buffer, BufferDesc}, + device::{self, DeviceOwned}, + images::{Image, ImageDesc, ImageViewDesc}, + render_graph::{ + buffer_barrier, image_barrier, Access, Barrier, GraphResourceDesc, GraphResourceId, + PassDesc, RecordFn, RenderContext, RenderGraph, + }, + texture, + util::Rect2D, + EguiState, +}; + +pub fn egui_pre_pass( + dev: &device::Device, + rg: &mut RenderGraph, + textures: &mut crate::texture::TextureManager, + egui_state: &mut EguiState, + output: &egui::FullOutput, +) -> VkResult<()> { + // allocate resource ids for textures in tessellated list (imported from texture manager) + // define accesses for resource ids + + // create textures for new egui textures + for (egui_id, delta) in output + .textures_delta + .set + .iter() + .filter(|(_, image)| image.is_whole()) + { + let image = Image::new( + dev.clone(), + ImageDesc { + name: Some(format!("egui-texture-{egui_id:?}").into()), + format: vk::Format::R8G8B8A8_UNORM, + extent: vk::Extent3D { + width: delta.image.width() as u32, + height: delta.image.height() as u32, + depth: 1, + }, + usage: vk::ImageUsageFlags::SAMPLED | vk::ImageUsageFlags::TRANSFER_DST, + mem_usage: vk_mem::MemoryUsage::AutoPreferDevice, + ..Default::default() + }, + )?; + + let tid = textures.insert_image(Arc::new(image)); + if let Some(old) = egui_state.textures.insert( + *egui_id, + crate::EguiTextureInfo { + id: tid, + options: delta.options, + }, + ) { + textures.remove_texture(old.id); + } + } + + // calculate size for staging buffer. + // calculate size for staging image. + let (staging_size, image_size) = output.textures_delta.set.iter().fold( + (0usize, 0usize), + |(mut buffer, mut image), (_id, delta)| { + let bytes = delta.image.height() * delta.image.width() * delta.image.bytes_per_pixel(); + if !delta.is_whole() { + image = image.max(bytes); + } + buffer = buffer + bytes; + + (buffer, image) + }, + ); + + let mut staging_buffer = Buffer::new( + dev.clone(), + BufferDesc { + name: Some("egui-prepass-staging-buffer".into()), + size: staging_size as u64, + usage: vk::BufferUsageFlags::TRANSFER_SRC, + queue_families: device::QueueFlags::empty(), + mem_usage: vk_mem::MemoryUsage::AutoPreferHost, + alloc_flags: vk_mem::AllocationCreateFlags::MAPPED + | vk_mem::AllocationCreateFlags::HOST_ACCESS_SEQUENTIAL_WRITE + | vk_mem::AllocationCreateFlags::STRATEGY_FIRST_FIT, + ..Default::default() + }, + )?; + let staging_image = Arc::new(Image::new( + dev.clone(), + ImageDesc { + name: Some("egui-prepass-staging-buffer".into()), + format: vk::Format::R8G8B8A8_UNORM, + extent: vk::Extent3D { + width: (image_size / 2) as u32, + height: (image_size - (image_size / 2)) as u32, + depth: 1, + }, + usage: vk::ImageUsageFlags::TRANSFER_SRC | vk::ImageUsageFlags::TRANSFER_DST, + queue_families: device::QueueFlags::empty(), + mem_usage: vk_mem::MemoryUsage::AutoPreferDevice, + ..Default::default() + }, + )?); + + let aliased_images = { + let mut staging_map = staging_buffer.map()?; + let mut offset = 0; + + let aliased_images = output + .textures_delta + .set + .iter() + .filter_map(|(id, delta)| { + let bytes = + delta.image.height() * delta.image.width() * delta.image.bytes_per_pixel(); + + let mem = &mut staging_map[offset..offset + bytes]; + match &delta.image { + egui::ImageData::Color(arc) => { + let slice = unsafe { + core::slice::from_raw_parts( + arc.pixels.as_ptr().cast::(), + arc.pixels.len() * size_of::(), + ) + }; + mem[..slice.len()].copy_from_slice(slice); + } + egui::ImageData::Font(font_image) => { + for (i, c) in font_image.srgba_pixels(None).enumerate() { + let bytes = c.to_array(); + mem[i * 4..(i + 1) * 4].copy_from_slice(&bytes); + } + } + } + + let old_offset = offset; + offset += bytes; + + if !delta.is_whole() { + unsafe { + let alias = staging_image + .clone() + .get_alias(ImageDesc { + name: Some(format!("egui-prepass-staging-aliased-{id:?}").into()), + format: vk::Format::R8G8B8A8_UNORM, + extent: vk::Extent3D { + width: delta.image.width() as u32, + height: delta.image.height() as u32, + depth: 1, + }, + usage: vk::ImageUsageFlags::TRANSFER_SRC + | vk::ImageUsageFlags::TRANSFER_DST, + queue_families: device::QueueFlags::empty(), + ..Default::default() + }) + .unwrap(); + let pos = delta.pos.unwrap(); + let rect = Rect2D::new_from_size( + glam::ivec2(pos[0] as i32, pos[1] as i32), + glam::ivec2(delta.image.width() as i32, delta.image.height() as i32), + ); + Some((*id, (old_offset, bytes, rect, alias))) + } + } else { + None + } + }) + .collect::>(); + + // let tessellated = egui.tessellate(output.shapes, output.pixels_per_point); + + aliased_images + }; + + let textures = output + .textures_delta + .set + .iter() + .filter_map(|(egui_id, _)| { + egui_state + .lookup_texture(*egui_id) + .and_then(|tid| textures.get_texture(tid)) + .map(|img| (*egui_id, img)) + }) + .map(|(id, img)| { + ( + id, + rg.import_image( + img, + Access { + layout: Some(vk::ImageLayout::GENERAL), + ..Access::undefined() + }, + ), + ) + }) + .collect::>(); + + let staging_buffer = rg.import_buffer(Arc::new(staging_buffer), Access::undefined()); + let staging_image = rg.import_image(staging_image, Access::undefined()); + + let record = Box::new({ + let textures = textures.clone(); + move |ctx: &RenderContext| -> crate::Result<()> { + let staging_image = ctx.get_image(staging_image).unwrap().clone(); + let staging_buffer = ctx.get_buffer(staging_buffer).unwrap(); + + for (id, (offset, _, rect, _)) in aliased_images { + let alias = unsafe { + staging_image.get_alias(ImageDesc { + name: Some(format!("egui-prepass-staging-aliased-{id:?}v").into()), + format: vk::Format::R8G8B8A8_UNORM, + extent: vk::Extent3D { + width: rect.width() as u32, + height: rect.height() as u32, + depth: 1, + }, + usage: vk::ImageUsageFlags::TRANSFER_SRC + | vk::ImageUsageFlags::TRANSFER_DST, + queue_families: device::QueueFlags::empty(), + ..Default::default() + })? + }; + + let texture = textures.get(&id).and_then(|id| ctx.get_image(*id)).unwrap(); + + let image: Barrier = image_barrier( + alias.handle(), + alias.format(), + Access { + stage: vk::PipelineStageFlags2::NONE, + mask: vk::AccessFlags2::empty(), + layout: Some(vk::ImageLayout::UNDEFINED), + }, + Access { + stage: vk::PipelineStageFlags2::TRANSFER, + mask: vk::AccessFlags2::TRANSFER_WRITE, + layout: Some(vk::ImageLayout::TRANSFER_DST_OPTIMAL), + }, + None, + ) + .into(); + + unsafe { + ctx.device + .dev() + .cmd_pipeline_barrier2(ctx.cmd.buffer(), &((&image).into())); + } + + ctx.cmd.copy_buffer_to_image( + staging_buffer.handle(), + alias.handle(), + vk::ImageLayout::TRANSFER_DST_OPTIMAL, + &[vk::BufferImageCopy { + buffer_offset: offset as u64, + buffer_row_length: alias.width(), + buffer_image_height: alias.height(), + image_subresource: vk::ImageSubresourceLayers::default() + .aspect_mask(vk::ImageAspectFlags::COLOR) + .base_array_layer(0) + .mip_level(0) + .layer_count(1), + image_offset: vk::Offset3D { x: 0, y: 0, z: 0 }, + image_extent: alias.size(), + }], + ); + + let from_barrier = image_barrier( + alias.handle(), + alias.format(), + Access { + stage: vk::PipelineStageFlags2::TRANSFER, + mask: vk::AccessFlags2::TRANSFER_WRITE, + layout: Some(vk::ImageLayout::TRANSFER_DST_OPTIMAL), + }, + Access { + stage: vk::PipelineStageFlags2::TRANSFER, + mask: vk::AccessFlags2::TRANSFER_READ, + layout: Some(vk::ImageLayout::TRANSFER_SRC_OPTIMAL), + }, + None, + ); + let to_barrier = image_barrier( + texture.handle(), + texture.format(), + Access { + stage: vk::PipelineStageFlags2::NONE, + mask: vk::AccessFlags2::empty(), + layout: Some(vk::ImageLayout::GENERAL), + }, + Access { + stage: vk::PipelineStageFlags2::TRANSFER, + mask: vk::AccessFlags2::TRANSFER_WRITE, + layout: Some(vk::ImageLayout::TRANSFER_DST_OPTIMAL), + }, + None, + ); + + unsafe { + ctx.device.dev().cmd_pipeline_barrier2( + ctx.cmd.buffer(), + &vk::DependencyInfo::default() + .image_memory_barriers(&[from_barrier, to_barrier]), + ); + } + ctx.cmd.copy_images( + alias.handle(), + vk::ImageLayout::TRANSFER_SRC_OPTIMAL, + texture.handle(), + vk::ImageLayout::TRANSFER_DST_OPTIMAL, + &[vk::ImageCopy { + src_subresource: vk::ImageSubresourceLayers::default() + .aspect_mask(vk::ImageAspectFlags::COLOR) + .base_array_layer(0) + .mip_level(0) + .layer_count(1), + src_offset: vk::Offset3D { x: 0, y: 0, z: 0 }, + dst_subresource: vk::ImageSubresourceLayers::default() + .aspect_mask(vk::ImageAspectFlags::COLOR) + .base_array_layer(0) + .mip_level(0) + .layer_count(1), + dst_offset: vk::Offset3D { + x: rect.top_left().x, + y: rect.top_left().y, + z: 0, + }, + extent: alias.size(), + }], + ); + + let image: Barrier = image_barrier( + texture.handle(), + texture.format(), + Access { + stage: vk::PipelineStageFlags2::TRANSFER, + mask: vk::AccessFlags2::TRANSFER_WRITE, + layout: Some(vk::ImageLayout::TRANSFER_DST_OPTIMAL), + }, + Access { + stage: vk::PipelineStageFlags2::ALL_COMMANDS, + mask: vk::AccessFlags2::empty(), + layout: Some(vk::ImageLayout::GENERAL), + }, + None, + ) + .into(); + + unsafe { + ctx.device + .dev() + .cmd_pipeline_barrier2(ctx.cmd.buffer(), &((&image).into())); + } + } + Ok(()) + } + }); + + rg.add_pass(PassDesc { + reads: [ + ( + staging_buffer, + Access { + stage: vk::PipelineStageFlags2::TRANSFER, + mask: vk::AccessFlags2::TRANSFER_READ, + ..Access::undefined() + }, + ), + (staging_image, Access::undefined()), + ] + .to_vec(), + writes: textures + .iter() + .map(|(_, id)| { + ( + *id, + Access { + layout: Some(vk::ImageLayout::GENERAL), + ..Access::undefined() + }, + ) + }) + .collect(), + record, + }); + Ok(()) +} + +// fn egui_pass() +pub fn egui_pass( + dev: &device::Device, + rg: &mut RenderGraph, + texture_handler: &mut crate::texture::TextureManager, + samplers: &mut crate::SamplerCache, + egui_state: &mut EguiState, + egui: &egui::Context, + output: egui::FullOutput, + target: GraphResourceId, +) -> VkResult> { + let draw_data = egui.tessellate(output.shapes, output.pixels_per_point); + + #[repr(C)] + #[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)] + struct Vertex { + pos: glam::Vec2, + uv: glam::Vec2, + color: egui::epaint::Color32, + } + + #[repr(transparent)] + #[derive(Debug, Clone, Copy)] + struct DrawCall(vk::DrawIndexedIndirectCommand); + unsafe impl bytemuck::Zeroable for DrawCall {} + unsafe impl bytemuck::Pod for DrawCall {} + + let mut vertices = Vec::new(); + let mut indices = Vec::new(); + let mut draw_calls = Vec::new(); + let mut textures = IndexMap::new(); + let mut textures_indices = Vec::new(); + for draw in draw_data { + let egui::epaint::Primitive::Mesh(mesh) = draw.primitive else { + continue; + }; + + draw_calls.push(DrawCall(vk::DrawIndexedIndirectCommand { + index_count: mesh.indices.len() as u32, + instance_count: 1, + first_index: indices.len() as u32, + vertex_offset: vertices.len() as i32, + first_instance: 0, + })); + + vertices.extend(mesh.vertices.iter().map(|v| Vertex { + pos: glam::vec2(v.pos.x, v.pos.y), + uv: glam::vec2(v.uv.x, v.uv.y), + color: v.color, + })); + + indices.extend(mesh.indices); + let texture = egui_state.textures.get(&mesh.texture_id).cloned().unwrap(); + if !textures.contains_key(&texture.id) { + textures.insert(texture.id, texture); + } + let idx = textures.get_index_of(&texture.id).unwrap(); + textures_indices.push(idx as u32); + } + + let num_draw_calls = draw_calls.len(); + let vertices_size = vertices.len() * size_of::(); + let indices_size = indices.len() * size_of::(); + let draw_calls_size = draw_calls.len() * size_of::(); + + let (draw_staging, vertices, indices, draw_calls, texture_ids) = { + let staging_size = vertices_size + indices_size + draw_calls_size; + + let mut staging = Buffer::new( + dev.clone(), + BufferDesc { + name: Some("egui-draw-staging".into()), + size: staging_size as u64, + usage: vk::BufferUsageFlags::TRANSFER_SRC, + mem_usage: vk_mem::MemoryUsage::AutoPreferHost, + alloc_flags: vk_mem::AllocationCreateFlags::MAPPED + | vk_mem::AllocationCreateFlags::HOST_ACCESS_SEQUENTIAL_WRITE + | vk_mem::AllocationCreateFlags::STRATEGY_FIRST_FIT, + ..Default::default() + }, + )?; + + { + let mut map = staging.map()?; + + let (st_vertices, rest) = map.split_at_mut(vertices_size); + let (st_indices, st_drawcalls) = rest.split_at_mut(indices_size); + st_vertices.copy_from_slice(bytemuck::cast_slice(&vertices)); + st_indices.copy_from_slice(bytemuck::cast_slice(&indices)); + st_drawcalls.copy_from_slice(bytemuck::cast_slice(&draw_calls)); + } + let staging = rg.import_buffer(Arc::new(staging), Access::undefined()); + + let vertices = rg.add_resource(GraphResourceDesc::Buffer(BufferDesc { + name: Some("egui-draw-vertices".into()), + size: vertices_size as u64, + usage: vk::BufferUsageFlags::TRANSFER_DST | vk::BufferUsageFlags::VERTEX_BUFFER, + mem_usage: vk_mem::MemoryUsage::AutoPreferDevice, + ..Default::default() + })); + let indices = rg.add_resource(GraphResourceDesc::Buffer(BufferDesc { + name: Some("egui-draw-indices".into()), + size: indices_size as u64, + usage: vk::BufferUsageFlags::TRANSFER_DST | vk::BufferUsageFlags::INDEX_BUFFER, + mem_usage: vk_mem::MemoryUsage::AutoPreferDevice, + ..Default::default() + })); + let draw_calls = rg.add_resource(GraphResourceDesc::Buffer(BufferDesc { + name: Some("egui-draw-draw_calls".into()), + size: draw_calls_size as u64, + usage: vk::BufferUsageFlags::TRANSFER_DST | vk::BufferUsageFlags::INDIRECT_BUFFER, + mem_usage: vk_mem::MemoryUsage::AutoPreferDevice, + ..Default::default() + })); + + let mut texture_ids = Buffer::new( + dev.clone(), + BufferDesc { + name: Some("egui-draw-texture_ids".into()), + size: (textures_indices.len() * size_of::()) as u64, + usage: vk::BufferUsageFlags::STORAGE_BUFFER, + mem_usage: vk_mem::MemoryUsage::AutoPreferDevice, + alloc_flags: vk_mem::AllocationCreateFlags::HOST_ACCESS_SEQUENTIAL_WRITE, + ..Default::default() + }, + )?; + { + let mut map = texture_ids.map()?; + map.copy_from_slice(bytemuck::cast_slice(&textures_indices)); + } + + (staging, vertices, indices, draw_calls, texture_ids) + }; + + let descriptor_infos = textures + .values() + .map(|entry| { + let texture = texture_handler.get_texture(entry.id).unwrap(); + let info = vk::DescriptorImageInfo { + sampler: samplers.get_sampler(entry.into_sampler_desc()).unwrap(), + image_view: texture + .get_view(ImageViewDesc { + kind: vk::ImageViewType::TYPE_2D, + format: texture.format(), + aspect: vk::ImageAspectFlags::COLOR, + mip_range: (0..1).into(), + layer_range: (0..1).into(), + ..Default::default() + }) + .unwrap(), + image_layout: vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL, + }; + + info + }) + .collect::>(); + + let uniform_info = vk::DescriptorBufferInfo { + buffer: texture_ids.buffer(), + offset: 0, + range: texture_ids.len(), + }; + + let descriptor_writes = descriptor_infos + .iter() + .enumerate() + .map(|(i, info)| { + vk::WriteDescriptorSet::default() + .image_info(core::slice::from_ref(info)) + .descriptor_count(1) + .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER) + .dst_set(egui_state.descriptor_set) + .dst_binding(EguiState::TEXTURE_BINDING) + .dst_array_element(i as u32) + }) + .chain(core::iter::once({ + vk::WriteDescriptorSet::default() + .buffer_info(core::slice::from_ref(&uniform_info)) + .descriptor_count(1) + .dst_binding(EguiState::UNIFORM_BINDING) + .descriptor_type(vk::DescriptorType::STORAGE_BUFFER) + .dst_array_element(0) + .dst_set(egui_state.descriptor_set) + })) + .collect::>(); + + unsafe { + dev.dev().update_descriptor_sets(&descriptor_writes, &[]); + } + + let to_remove_tex_ids = output + .textures_delta + .free + .iter() + .filter_map(|id| egui_state.textures.get(id).cloned()) + .map(|entry| entry.id) + .collect::>(); + + let reads = textures + .keys() + .filter_map(|id| { + texture_handler + .get_texture(*id) + .map(|img| (rg.import_image(img, Access::general()), Access::general())) + }) + .chain([(target, Access::color_attachment_read_write())]) + .collect::>(); + let writes = [(target, Access::color_attachment_write_only())].to_vec(); + + let record: Box = Box::new({ + let pipeline = egui_state.pipeline.clone(); + let pipeline_layout = egui_state.pipeline_layout.clone(); + let descriptor_set = egui_state.descriptor_set; + + move |ctx: &RenderContext| -> crate::Result<()> { + let cmd = &ctx.cmd; + let staging = ctx.get_buffer(draw_staging).unwrap(); + let vertices = ctx.get_buffer(vertices).unwrap(); + let indices = ctx.get_buffer(indices).unwrap(); + let draw_calls = ctx.get_buffer(draw_calls).unwrap(); + let target = ctx.get_image(target).unwrap(); + + cmd.copy_buffers( + staging.buffer(), + vertices.buffer(), + &[vk::BufferCopy { + src_offset: 0, + dst_offset: 0, + size: vertices_size as u64, + }], + ); + cmd.copy_buffers( + staging.buffer(), + indices.buffer(), + &[vk::BufferCopy { + src_offset: vertices_size as u64, + dst_offset: 0, + size: indices_size as u64, + }], + ); + cmd.copy_buffers( + staging.buffer(), + draw_calls.buffer(), + &[vk::BufferCopy { + src_offset: (vertices_size + indices_size) as u64, + dst_offset: 0, + size: draw_calls_size as u64, + }], + ); + + let barriers = [ + buffer_barrier( + vertices.handle(), + 0, + vertices.len(), + Access::transfer_write(), + Access::vertex_read(), + None, + ), + buffer_barrier( + indices.handle(), + 0, + indices.len(), + Access::transfer_write(), + Access::index_read(), + None, + ), + buffer_barrier( + draw_calls.handle(), + 0, + draw_calls.len(), + Access::transfer_write(), + Access::indirect_read(), + None, + ), + ]; + unsafe { + ctx.device.dev().cmd_pipeline_barrier2( + cmd.buffer(), + &vk::DependencyInfo::default().buffer_memory_barriers(&barriers), + ); + } + + cmd.bind_pipeline(&pipeline); + cmd.bind_indices(indices.buffer(), 0, vk::IndexType::UINT32); + cmd.bind_vertices(vertices.buffer(), 0); + cmd.push_constants( + &pipeline_layout, + vk::ShaderStageFlags::VERTEX, + 0, + bytemuck::cast_slice( + &[target.width() as f32, target.height() as f32].map(|f| f.to_bits()), + ), + ); + cmd.bind_descriptor_sets( + &pipeline_layout, + vk::PipelineBindPoint::GRAPHICS, + &[descriptor_set], + ); + cmd.draw_indexed_indirect( + draw_calls.buffer(), + 0, + num_draw_calls as u32, + size_of::() as u32, + ); + Ok(()) + } + }); + + rg.add_pass(PassDesc { + reads, + writes, + record, + }); + + Ok(to_remove_tex_ids) +} diff --git a/crates/renderer/src/lib.rs b/crates/renderer/src/lib.rs index f50da19..163c9e4 100644 --- a/crates/renderer/src/lib.rs +++ b/crates/renderer/src/lib.rs @@ -38,6 +38,8 @@ use raw_window_handle::RawDisplayHandle; mod buffers; mod commands; mod device; +#[path = "egui.rs"] +mod egui_pass; mod images; mod pipeline; mod render_graph;