rendergraph: transitions

crates/renderer/Cargo.toml

@@ -22,6 +22,7 @@ bitflags.workspace = true
 petgraph.workspace = true
 itertools.workspace = true
 indexmap.workspace = true
+futures.workspace = true
 bytemuck = { version = "1.21.0", features = ["derive"] }
 
 

crates/renderer/src/images.rs

@@ -287,6 +287,14 @@ pub struct QueueOwnership {
     pub dst: u32,
 }
 
+pub const SUBRESOURCERANGE_ALL: vk::ImageSubresourceRange = vk::ImageSubresourceRange {
+    aspect_mask: vk::ImageAspectFlags::empty(),
+    base_mip_level: 0,
+    level_count: vk::REMAINING_MIP_LEVELS,
+    base_array_layer: 0,
+    layer_count: vk::REMAINING_ARRAY_LAYERS,
+};
+
 pub const SUBRESOURCERANGE_COLOR_ALL: vk::ImageSubresourceRange = vk::ImageSubresourceRange {
     aspect_mask: vk::ImageAspectFlags::COLOR,
     base_mip_level: 0,

crates/renderer/src/render_graph.rs

@@ -5,9 +5,9 @@ use std::{collections::BTreeMap, fmt::Debug, sync::Arc};
 use crate::{
     buffers::{Buffer, BufferDesc},
     commands, def_monotonic_id,
-    device::{self, Device},
+    device::{self, Device, DeviceOwned},
-    images::{Image, ImageDesc},
+    images::{self, Image, ImageDesc},
-    util::Rgba,
+    util::{self, Rgba},
     SwapchainFrame,
 };
 use ash::vk;
@@ -50,12 +50,31 @@ struct AttachmentInfo {
     store: StoreOp,
 }
 
-pub struct RenderContext<'a> {
+pub struct RenderContext {
     device: device::Device,
-    cmd: &'a commands::SingleUseCommand,
+    cmd: commands::SingleUseCommand,
 }
 
-pub trait Pass: Debug {
+#[derive(Debug, Clone, Copy, Default)]
+pub struct ResourceAccess {
+    stage: vk::PipelineStageFlags2,
+    mask: vk::AccessFlags2,
+    layout: Option<vk::ImageLayout>,
+}
+
+impl ResourceAccess {
+    fn undefined() -> Self {
+        Self {
+            stage: vk::PipelineStageFlags2::NONE,
+            mask: vk::AccessFlags2::empty(),
+            layout: Some(vk::ImageLayout::UNDEFINED),
+        }
+    }
+}
+
+pub trait Pass: Debug + Send {
+    fn get_read_resource_access(&self, id: RenderGraphResourceId) -> ResourceAccess;
+    fn get_write_resource_access(&self, id: RenderGraphResourceId) -> ResourceAccess;
     /// returns the layout the pass requires an image dependency to be in prior
     /// to the pass.
     fn get_layout_of_in_image_dependency(&self, id: RenderGraphResourceId) -> vk::ImageLayout;
@@ -65,10 +84,11 @@ pub trait Pass: Debug {
     /// mask of the queue capability requirements of this pass.
     fn get_queue_capability_requirements(&self) -> device::QueueFlags;
     /// returns an iterator over all (in) dependencies.
-    fn get_in_dependencies<'a>(&'a self) -> Box<dyn Iterator<Item = RenderGraphResourceId> + 'a>;
+    fn get_read_dependencies<'a>(&'a self) -> Box<dyn Iterator<Item = RenderGraphResourceId> + 'a>;
-    fn get_out_dependencies<'a>(&'a self) -> Box<dyn Iterator<Item = RenderGraphResourceId> + 'a>;
+    fn get_write_dependencies<'a>(&'a self)
+        -> Box<dyn Iterator<Item = RenderGraphResourceId> + 'a>;
 
-    fn record(self, ctx: &RenderContext) -> crate::Result<()>;
+    fn record(self: Box<Self>, ctx: &RenderContext) -> crate::Result<()>;
 }
 
 def_monotonic_id!(pub RenderGraphPassId);
@@ -82,6 +102,7 @@ def_monotonic_id!(pub RenderGraphPassId);
 pub struct RenderGraph {
     resource_descs: BTreeMap<RenderGraphResourceId, RenderGraphResourceDesc>,
     resources: BTreeMap<RenderGraphResourceId, RenderGraphResource>,
+    accesses: BTreeMap<RenderGraphResourceId, ResourceAccess>,
     passes: Vec<Box<dyn Pass>>,
     /// the rendergraph produces these resources. Any passes on which these
     /// outputs do not depend are pruned.
@@ -94,6 +115,7 @@ impl RenderGraph {
             resource_descs: BTreeMap::new(),
             resources: BTreeMap::new(),
             passes: Vec::new(),
+            accesses: BTreeMap::new(),
             outputs: Vec::new(),
         }
     }
@@ -101,22 +123,33 @@ impl RenderGraph {
     pub fn add_resource(&mut self, desc: RenderGraphResourceDesc) -> RenderGraphResourceId {
         let id = RenderGraphResourceId::new();
         self.resource_descs.insert(id, desc);
+        self.accesses.insert(id, ResourceAccess::undefined());
         id
     }
     pub fn mark_as_output(&mut self, id: RenderGraphResourceId) {
         // TODO: dedup
         self.outputs.push(id);
     }
-    pub fn import_image(&mut self, image: Arc<Image>) -> RenderGraphResourceId {
+    pub fn import_image(
+        &mut self,
+        image: Arc<Image>,
+        access: ResourceAccess,
+    ) -> RenderGraphResourceId {
         let id = RenderGraphResourceId::new();
         self.resources
             .insert(id, RenderGraphResource::ImportedImage(image));
+        self.accesses.insert(id, access);
         id
     }
-    pub fn import_buffer(&mut self, buffer: Arc<Buffer>) -> RenderGraphResourceId {
+    pub fn import_buffer(
+        &mut self,
+        buffer: Arc<Buffer>,
+        access: ResourceAccess,
+    ) -> RenderGraphResourceId {
         let id = RenderGraphResourceId::new();
         self.resources
             .insert(id, RenderGraphResource::ImportedBuffer(buffer));
+        self.accesses.insert(id, access);
         id
     }
     pub fn import_framebuffer(&mut self, frame: Arc<SwapchainFrame>) -> RenderGraphResourceId {
@@ -132,7 +165,7 @@ impl RenderGraph {
     // https://blog.traverseresearch.nl/render-graph-101-f42646255636
     // https://github.com/EmbarkStudios/kajiya/blob/main/crates/lib/kajiya-rg/src/graph.rs
     // https://themaister.net/blog/2017/08/15/render-graphs-and-vulkan-a-deep-dive/
-    pub fn resolve(mut self) {
+    pub fn resolve(mut self, device: device::Device) -> crate::Result<()> {
         eprintln!("{:#?}", &self);
         let mut last_write = BTreeMap::new();
 
@@ -141,13 +174,13 @@ impl RenderGraph {
         // insert edges between write->read edges of 2 passes
         for (i, pass) in self.passes.iter().enumerate() {
             let node = dag.add_node(i);
-            for dep in pass.get_in_dependencies() {
+            for dep in pass.get_read_dependencies() {
                 if let Some(&other) = last_write.get(&dep) {
                     dag.add_edge(other, node, dep);
                 }
             }
 
-            for dep in pass.get_out_dependencies() {
+            for dep in pass.get_write_dependencies() {
                 // keep track of which pass last wrote to a resource
                 // this is the node to build an edge from next time this resource is read
                 last_write.insert(dep, node);
@@ -158,8 +191,8 @@ impl RenderGraph {
         let output = dag.add_node(!0);
         for (id, node) in self
             .outputs
-            .into_iter()
+            .iter()
-            .filter_map(|id| last_write.get(&id).cloned().map(|node| (id, node)))
+            .filter_map(|&id| last_write.get(&id).cloned().map(|node| (id, node)))
         {
             dag.add_edge(node, output, id);
         }
@@ -193,12 +226,18 @@ impl RenderGraph {
         )
         .expect("writing render_graph repr");
 
+        struct RenderGraphStage {
+            passes: Vec<Box<dyn Pass>>,
+            accesses: Vec<(RenderGraphResourceId, ResourceAccess)>,
+        }
+
         let mut topological_map = Vec::new();
+        let mut top_dag = dag.clone();
         loop {
-            let (sources, indices): (Vec<_>, Vec<_>) = dag
+            let (sources, indices): (Vec<_>, Vec<_>) = top_dag
                 .externals(petgraph::Direction::Incoming)
                 .filter(|&id| id != output)
-                .filter_map(|id| dag.node_weight(id).cloned().map(|idx| (id, idx)))
+                .filter_map(|id| top_dag.node_weight(id).cloned().map(|idx| (id, idx)))
                 .unzip();
 
             if sources.is_empty() {
@@ -206,20 +245,203 @@ impl RenderGraph {
             }
 
             for &source in &sources {
-                dag.remove_node(source);
+                top_dag.remove_node(source);
             }
 
-            topological_map.push(indices);
+            let mut accesses = BTreeMap::<RenderGraphResourceId, ResourceAccess>::new();
+
+            // TODO: distinguish between make-available and make-visible
+            for &idx in &indices {
+                let pass = &self.passes[idx];
+                for id in pass.get_read_dependencies() {
+                    let access = pass.get_read_resource_access(id);
+                    accesses
+                        .entry(id)
+                        .and_modify(|entry| {
+                            entry.stage = entry.stage.max(access.stage);
+                            entry.mask |= access.mask;
+                            // TODO: layout
+                        })
+                        .or_insert(access);
+                }
+            }
+
+            topological_map.push((indices, accesses));
         }
 
         // I don't think this can currently happen with the way passes are added.
-        dag.remove_node(output);
+        top_dag.remove_node(output);
         if dag.node_count() > 0 {
             panic!("dag is cyclic!");
         }
 
         eprintln!("topology: {:?}", topological_map);
+
+        let pool =
+            commands::SingleUseCommandPool::new(device.clone(), device.graphics_queue().clone())?;
+
+        let tasks = topological_map
+            .iter()
+            .map(|(set, accesses)| {
+                let pool = pool.clone();
+                let device = device.clone();
+                let passes = set
+                    .into_iter()
+                    .map(|i| self.passes.remove(*i))
+                    .collect::<Vec<_>>();
+
+                let cmd = pool.alloc()?;
+
+                // transitions
+                for (&id, &access) in accesses.iter() {
+                    self.transition_resource_to(device.dev(), unsafe { &cmd.buffer() }, id, access);
+                }
+
+                let task = smol::spawn(async move {
+                    let ctx = RenderContext { device, cmd };
+
+                    for pass in passes {
+                        pass.record(&ctx)?;
+                    }
+
+                    crate::Result::Ok(ctx.cmd)
+                });
+
+                crate::Result::Ok(task)
+            })
+            .collect::<crate::Result<Vec<_>>>()?;
+
+        let commands = smol::block_on(futures::future::join_all(tasks))
+            .into_iter()
+            .collect::<crate::Result<Vec<_>>>()?;
+
+        Ok(())
     }
+
+    fn transition_resource_to(
+        &mut self,
+        dev: &ash::Device,
+        cmd: &vk::CommandBuffer,
+        id: RenderGraphResourceId,
+        access: ResourceAccess,
+    ) {
+        let old_access = self.accesses.get(&id);
+        let res = self.resources.get(&id);
+
+        if let (Some(&old_access), Some(res)) = (old_access, res) {
+            let barrier: Barrier = match res {
+                RenderGraphResource::Framebuffer(arc) => {
+                    image_barrier(arc.image, arc.format, old_access, access, None).into()
+                }
+                RenderGraphResource::ImportedImage(arc) => {
+                    image_barrier(arc.handle(), arc.format(), old_access, access, None).into()
+                }
+                RenderGraphResource::ImportedBuffer(arc) => {
+                    buffer_barrier(arc.handle(), 0, arc.len(), old_access, access, None).into()
+                }
+                RenderGraphResource::Image(image) => {
+                    image_barrier(image.handle(), image.format(), old_access, access, None).into()
+                }
+                RenderGraphResource::Buffer(buffer) => {
+                    buffer_barrier(buffer.handle(), 0, buffer.len(), old_access, access, None)
+                        .into()
+                }
+            };
+
+            self.accesses.insert(id, access);
+
+            unsafe {
+                dev.cmd_pipeline_barrier2(*cmd, &((&barrier).into()));
+            }
+        }
+    }
+}
+
+enum Barrier {
+    Image(vk::ImageMemoryBarrier2<'static>),
+    Buffer(vk::BufferMemoryBarrier2<'static>),
+}
+
+impl<'a> From<&'a Barrier> for vk::DependencyInfo<'a> {
+    fn from(value: &'a Barrier) -> Self {
+        let info = vk::DependencyInfo::default();
+        let info = match value {
+            Barrier::Image(barrier) => info.image_memory_barriers(core::slice::from_ref(barrier)),
+            Barrier::Buffer(barrier) => info.buffer_memory_barriers(core::slice::from_ref(barrier)),
+        };
+
+        info
+    }
+}
+
+impl From<vk::ImageMemoryBarrier2<'static>> for Barrier {
+    fn from(value: vk::ImageMemoryBarrier2<'static>) -> Self {
+        Self::Image(value)
+    }
+}
+impl From<vk::BufferMemoryBarrier2<'static>> for Barrier {
+    fn from(value: vk::BufferMemoryBarrier2<'static>) -> Self {
+        Self::Buffer(value)
+    }
+}
+
+pub fn buffer_barrier(
+    buffer: vk::Buffer,
+    offset: u64,
+    size: u64,
+    before: ResourceAccess,
+    after: ResourceAccess,
+    queue_families: Option<(u32, u32)>,
+) -> vk::BufferMemoryBarrier2<'static> {
+    vk::BufferMemoryBarrier2::default()
+        .buffer(buffer)
+        .offset(offset)
+        .size(size)
+        .src_access_mask(before.mask)
+        .src_stage_mask(before.stage)
+        .dst_access_mask(after.mask)
+        .dst_stage_mask(after.stage)
+        .src_queue_family_index(
+            queue_families
+                .map(|(src, _)| src)
+                .unwrap_or(vk::QUEUE_FAMILY_IGNORED),
+        )
+        .dst_queue_family_index(
+            queue_families
+                .map(|(_, dst)| dst)
+                .unwrap_or(vk::QUEUE_FAMILY_IGNORED),
+        )
+}
+
+pub fn image_barrier(
+    image: vk::Image,
+    format: vk::Format,
+    before_access: ResourceAccess,
+    after_access: ResourceAccess,
+    queue_families: Option<(u32, u32)>,
+) -> vk::ImageMemoryBarrier2<'static> {
+    vk::ImageMemoryBarrier2::default()
+        .src_access_mask(before_access.mask)
+        .src_stage_mask(before_access.stage)
+        .dst_access_mask(after_access.mask)
+        .dst_stage_mask(after_access.stage)
+        .image(image)
+        .old_layout(before_access.layout.unwrap_or_default())
+        .new_layout(after_access.layout.unwrap_or_default())
+        .subresource_range(vk::ImageSubresourceRange {
+            aspect_mask: util::image_aspect_from_format(format),
+            ..images::SUBRESOURCERANGE_ALL
+        })
+        .src_queue_family_index(
+            queue_families
+                .map(|(src, _)| src)
+                .unwrap_or(vk::QUEUE_FAMILY_IGNORED),
+        )
+        .dst_queue_family_index(
+            queue_families
+                .map(|(_, dst)| dst)
+                .unwrap_or(vk::QUEUE_FAMILY_IGNORED),
+        )
 }
 
 #[cfg(test)]
@@ -231,6 +453,12 @@ mod tests {
             #[derive(Debug, Clone)]
             struct $name(Vec<RenderGraphResourceId>, Vec<RenderGraphResourceId>);
             impl Pass for $name {
+                fn get_read_resource_access(&self, _id: RenderGraphResourceId) -> ResourceAccess {
+                    ResourceAccess::default()
+                }
+                fn get_write_resource_access(&self, _id: RenderGraphResourceId) -> ResourceAccess {
+                    ResourceAccess::default()
+                }
                 fn get_layout_of_in_image_dependency(
                     &self,
                     _id: RenderGraphResourceId,
@@ -249,19 +477,19 @@ mod tests {
                     $queue
                 }
 
-                fn get_in_dependencies<'a>(
+                fn get_read_dependencies<'a>(
                     &'a self,
                 ) -> Box<dyn Iterator<Item = RenderGraphResourceId> + 'a> {
                     Box::new(self.0.iter().cloned())
                 }
 
-                fn get_out_dependencies<'a>(
+                fn get_write_dependencies<'a>(
                     &'a self,
                 ) -> Box<dyn Iterator<Item = RenderGraphResourceId> + 'a> {
                     Box::new(self.1.iter().cloned())
                 }
 
-                fn record(self, _ctx: &RenderContext) -> crate::Result<()> {
+                fn record(self: Box<Self>, _ctx: &RenderContext) -> crate::Result<()> {
                     Ok(())
                 }
             }
@@ -323,6 +551,6 @@ mod tests {
         graph.mark_as_output(gbuffer);
         graph.mark_as_output(depth_image);
 
-        graph.resolve();
+        // graph.resolve();
     }
 }

crates/renderer/src/util.rs

@@ -359,3 +359,18 @@ impl Rgba {
         self.0.map(|f| (f.clamp(0.0, 1.0) * 255.0) as i32 - 128)
     }
 }
+
+#[allow(dead_code)]
+pub fn image_aspect_from_format(format: vk::Format) -> vk::ImageAspectFlags {
+    use vk::{Format, ImageAspectFlags};
+    match format {
+        Format::D32_SFLOAT | Format::X8_D24_UNORM_PACK32 | Format::D16_UNORM => {
+            ImageAspectFlags::DEPTH
+        }
+        Format::S8_UINT => ImageAspectFlags::STENCIL,
+        Format::D32_SFLOAT_S8_UINT | Format::D16_UNORM_S8_UINT | Format::D24_UNORM_S8_UINT => {
+            ImageAspectFlags::DEPTH | ImageAspectFlags::STENCIL
+        }
+        _ => ImageAspectFlags::COLOR,
+    }
+}