Track GPU use of compute and render pipelines (#534)

wgpu-core/src/command/compute.rs

@@ -142,7 +142,10 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
                     }
                 }
                 ComputeCommand::SetPipeline(pipeline_id) => {
-                    let pipeline = &pipeline_guard[pipeline_id];
+                    let pipeline = cmb.trackers
+                        .compute_pipes
+                        .use_extend(&*pipeline_guard, pipeline_id, (), ())
+                        .unwrap();
 
                     unsafe {
                         raw.bind_compute_pipeline(&pipeline.raw);

wgpu-core/src/command/mod.rs

@@ -183,6 +183,8 @@ impl<B: GfxBackend> CommandBuffer<B> {
         base.views.merge_extend(&head.views).unwrap();
         base.bind_groups.merge_extend(&head.bind_groups).unwrap();
         base.samplers.merge_extend(&head.samplers).unwrap();
+        base.compute_pipes.merge_extend(&head.compute_pipes).unwrap();
+        base.render_pipes.merge_extend(&head.render_pipes).unwrap();
 
         let stages = all_buffer_stages() | all_image_stages();
         unsafe {

wgpu-core/src/command/render.rs

@@ -851,7 +851,10 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
                     };
                 }
                 RenderCommand::SetPipeline(pipeline_id) => {
-                    let pipeline = &pipeline_guard[pipeline_id];
+                    let pipeline = trackers
+                        .render_pipes
+                        .use_extend(&*pipeline_guard, pipeline_id, (), ())
+                        .unwrap();
 
                     assert!(
                         context.compatible(&pipeline.pass_context),

wgpu-core/src/device/life.rs

@@ -34,6 +34,8 @@ pub struct SuspectedResources {
     pub(crate) texture_views: Vec<id::TextureViewId>,
     pub(crate) samplers: Vec<id::SamplerId>,
     pub(crate) bind_groups: Vec<id::BindGroupId>,
+    pub(crate) compute_pipelines: Vec<id::ComputePipelineId>,
+    pub(crate) render_pipelines: Vec<id::RenderPipelineId>,
 }
 
 impl SuspectedResources {
@@ -43,6 +45,8 @@ impl SuspectedResources {
         self.texture_views.clear();
         self.samplers.clear();
         self.bind_groups.clear();
+        self.compute_pipelines.clear();
+        self.render_pipelines.clear();
     }
 
     pub fn extend(&mut self, other: &Self) {
@@ -51,6 +55,8 @@ impl SuspectedResources {
         self.texture_views.extend_from_slice(&other.texture_views);
         self.samplers.extend_from_slice(&other.samplers);
         self.bind_groups.extend_from_slice(&other.bind_groups);
+        self.compute_pipelines.extend_from_slice(&other.compute_pipelines);
+        self.render_pipelines.extend_from_slice(&other.render_pipelines);
     }
 }
 
@@ -65,6 +71,8 @@ struct NonReferencedResources<B: hal::Backend> {
     samplers: Vec<B::Sampler>,
     framebuffers: Vec<B::Framebuffer>,
     desc_sets: Vec<DescriptorSet<B>>,
+    compute_pipes: Vec<B::ComputePipeline>,
+    graphics_pipes: Vec<B::GraphicsPipeline>,
 }
 
 impl<B: hal::Backend> NonReferencedResources<B> {
@@ -76,6 +84,8 @@ impl<B: hal::Backend> NonReferencedResources<B> {
             samplers: Vec::new(),
             framebuffers: Vec::new(),
             desc_sets: Vec::new(),
+            compute_pipes: Vec::new(),
+            graphics_pipes: Vec::new(),
         }
     }
 
@@ -86,6 +96,8 @@ impl<B: hal::Backend> NonReferencedResources<B> {
         self.samplers.extend(other.samplers);
         self.framebuffers.extend(other.framebuffers);
         self.desc_sets.extend(other.desc_sets);
+        self.compute_pipes.extend(other.compute_pipes);
+        self.graphics_pipes.extend(other.graphics_pipes);
     }
 
     unsafe fn clean(
@@ -124,6 +136,13 @@ impl<B: hal::Backend> NonReferencedResources<B> {
                 .lock()
                 .free(self.desc_sets.drain(..));
         }
+
+        for raw in self.compute_pipes.drain(..) {
+            device.destroy_compute_pipeline(raw);
+        }
+        for raw in self.graphics_pipes.drain(..) {
+            device.destroy_graphics_pipeline(raw);
+        }
     }
 }
 
@@ -386,6 +405,44 @@ impl<B: GfxBackend> LifetimeTracker<B> {
                 }
             }
         }
+
+        if !self.suspected_resources.compute_pipelines.is_empty() {
+            let mut trackers = trackers.lock();
+            let (mut guard, _) = hub.compute_pipelines.write(token);
+
+            for id in self.suspected_resources.compute_pipelines.drain(..) {
+                if trackers.compute_pipes.remove_abandoned(id) {
+                    hub.compute_pipelines.free_id(id);
+                    let res = guard.remove(id).unwrap();
+
+                    let submit_index = res.life_guard.submission_index.load(Ordering::Acquire);
+                    self.active
+                        .iter_mut()
+                        .find(|a| a.index == submit_index)
+                        .map_or(&mut self.free_resources, |a| &mut a.last_resources)
+                        .compute_pipes.push(res.raw);
+                }
+            }
+        }
+
+        if !self.suspected_resources.render_pipelines.is_empty() {
+            let mut trackers = trackers.lock();
+            let (mut guard, _) = hub.render_pipelines.write(token);
+
+            for id in self.suspected_resources.render_pipelines.drain(..) {
+                if trackers.render_pipes.remove_abandoned(id) {
+                    hub.render_pipelines.free_id(id);
+                    let res = guard.remove(id).unwrap();
+
+                    let submit_index = res.life_guard.submission_index.load(Ordering::Acquire);
+                    self.active
+                        .iter_mut()
+                        .find(|a| a.index == submit_index)
+                        .map_or(&mut self.free_resources, |a| &mut a.last_resources)
+                        .graphics_pipes.push(res.raw);
+                }
+            }
+        }
     }
 
     pub(crate) fn triage_mapped<G: GlobalIdentityHandlerFactory>(

wgpu-core/src/device/mod.rs

@@ -1292,6 +1292,8 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
         // that it references for destruction in the next GC pass.
         {
             let (bind_group_guard, mut token) = hub.bind_groups.read(&mut token);
+            let (compute_pipe_guard, mut token) = hub.compute_pipelines.read(&mut token);
+            let (render_pipe_guard, mut token) = hub.render_pipelines.read(&mut token);
             let (buffer_guard, mut token) = hub.buffers.read(&mut token);
             let (texture_guard, mut token) = hub.textures.read(&mut token);
             let (texture_view_guard, mut token) = hub.texture_views.read(&mut token);
@@ -1322,6 +1324,16 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
                     device.temp_suspected.samplers.push(id);
                 }
             }
+            for id in comb.trackers.compute_pipes.used() {
+                if compute_pipe_guard[id].life_guard.ref_count.is_none() {
+                    device.temp_suspected.compute_pipelines.push(id);
+                }
+            }
+            for id in comb.trackers.render_pipes.used() {
+                if render_pipe_guard[id].life_guard.ref_count.is_none() {
+                    device.temp_suspected.render_pipelines.push(id);
+                }
+            }
         }
 
         device
@@ -1357,6 +1369,8 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
             let (mut swap_chain_guard, mut token) = hub.swap_chains.write(&mut token);
             let (mut command_buffer_guard, mut token) = hub.command_buffers.write(&mut token);
             let (bind_group_guard, mut token) = hub.bind_groups.read(&mut token);
+            let (compute_pipe_guard, mut token) = hub.compute_pipelines.read(&mut token);
+            let (render_pipe_guard, mut token) = hub.render_pipelines.read(&mut token);
             let (buffer_guard, mut token) = hub.buffers.read(&mut token);
             let (texture_guard, mut token) = hub.textures.read(&mut token);
             let (texture_view_guard, mut token) = hub.texture_views.read(&mut token);
@@ -1415,6 +1429,16 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
                         device.temp_suspected.samplers.push(id);
                     }
                 }
+                for id in comb.trackers.compute_pipes.used() {
+                    if !compute_pipe_guard[id].life_guard.use_at(submit_index) {
+                        device.temp_suspected.compute_pipelines.push(id);
+                    }
+                }
+                for id in comb.trackers.render_pipes.used() {
+                    if !render_pipe_guard[id].life_guard.use_at(submit_index) {
+                        device.temp_suspected.render_pipelines.push(id);
+                    }
+                }
 
                 // execute resource transitions
                 let mut transit = device.com_allocator.extend(comb);
@@ -1745,6 +1769,7 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
             index_format: desc.vertex_state.index_format,
             vertex_strides,
             sample_count: sc,
+            life_guard: LifeGuard::new(),
         };
 
         hub.render_pipelines
@@ -1754,9 +1779,18 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
     pub fn render_pipeline_destroy<B: GfxBackend>(&self, render_pipeline_id: id::RenderPipelineId) {
         let hub = B::hub(self);
         let mut token = Token::root();
-        let (_, mut token) = hub.devices.read(&mut token);
-        //TODO: track usage by GPU
-        hub.render_pipelines.unregister(render_pipeline_id, &mut token);
+        let (device_guard, mut token) = hub.devices.read(&mut token);
+
+        let device_id = {
+            let (mut pipeline_guard, _) = hub.render_pipelines.write(&mut token);
+            let pipeline = &mut pipeline_guard[render_pipeline_id];
+            pipeline.life_guard.ref_count.take();
+            pipeline.device_id.value
+        };
+
+        device_guard[device_id]
+            .lock_life(&mut token)
+            .suspected_resources.render_pipelines.push(render_pipeline_id);
     }
 
     pub fn device_create_compute_pipeline<B: GfxBackend>(
@@ -1812,6 +1846,7 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
                 value: device_id,
                 ref_count: device.life_guard.add_ref(),
             },
+            life_guard: LifeGuard::new(),
         };
         hub.compute_pipelines
             .register_identity(id_in, pipeline, &mut token)
@@ -1820,9 +1855,19 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
     pub fn compute_pipeline_destroy<B: GfxBackend>(&self, compute_pipeline_id: id::ComputePipelineId) {
         let hub = B::hub(self);
         let mut token = Token::root();
-        let (_, mut token) = hub.devices.read(&mut token);
-        //TODO: track usage by GPU
-        hub.compute_pipelines.unregister(compute_pipeline_id, &mut token);
+        let (device_guard, mut token) = hub.devices.read(&mut token);
+
+        let device_id = {
+            let (mut pipeline_guard, _) = hub.compute_pipelines.write(&mut token);
+            let pipeline = &mut pipeline_guard[compute_pipeline_id];
+            pipeline.life_guard.ref_count.take();
+            pipeline.device_id.value
+        };
+
+        device_guard[device_id]
+            .lock_life(&mut token)
+            .suspected_resources.compute_pipelines.push(compute_pipeline_id);
+
     }
 
     pub fn device_create_swap_chain<B: GfxBackend>(

wgpu-core/src/hub.rs

@@ -184,6 +184,7 @@ impl<B: hal::Backend> Access<ComputePipeline<B>> for Device<B> {}
 impl<B: hal::Backend> Access<ComputePipeline<B>> for BindGroup<B> {}
 impl<B: hal::Backend> Access<RenderPipeline<B>> for Device<B> {}
 impl<B: hal::Backend> Access<RenderPipeline<B>> for BindGroup<B> {}
+impl<B: hal::Backend> Access<RenderPipeline<B>> for ComputePipeline<B> {}
 impl<B: hal::Backend> Access<ShaderModule<B>> for Device<B> {}
 impl<B: hal::Backend> Access<ShaderModule<B>> for PipelineLayout<B> {}
 impl<B: hal::Backend> Access<Buffer<B>> for Root {}

wgpu-core/src/pipeline.rs

@@ -5,11 +5,14 @@
 use crate::{
     device::RenderPassContext,
     id::{DeviceId, PipelineLayoutId, ShaderModuleId},
+    LifeGuard,
     RawString,
+    RefCount,
     Stored,
     U32Array
 };
 use wgt::{BufferAddress, ColorStateDescriptor, DepthStencilStateDescriptor, IndexFormat, InputStepMode, PrimitiveTopology, RasterizationStateDescriptor, VertexAttributeDescriptor};
+use std::borrow::Borrow;
 
 #[repr(C)]
 #[derive(Debug)]
@@ -59,6 +62,13 @@ pub struct ComputePipeline<B: hal::Backend> {
     pub(crate) raw: B::ComputePipeline,
     pub(crate) layout_id: PipelineLayoutId,
     pub(crate) device_id: Stored<DeviceId>,
+    pub(crate) life_guard: LifeGuard,
+}
+
+impl<B: hal::Backend> Borrow<RefCount> for ComputePipeline<B> {
+    fn borrow(&self) -> &RefCount {
+        self.life_guard.ref_count.as_ref().unwrap()
+    }
 }
 
 #[repr(C)]
@@ -96,4 +106,11 @@ pub struct RenderPipeline<B: hal::Backend> {
     pub(crate) index_format: IndexFormat,
     pub(crate) sample_count: u8,
     pub(crate) vertex_strides: Vec<(BufferAddress, InputStepMode)>,
+    pub(crate) life_guard: LifeGuard,
+}
+
+impl<B: hal::Backend> Borrow<RefCount> for RenderPipeline<B> {
+    fn borrow(&self) -> &RefCount {
+        self.life_guard.ref_count.as_ref().unwrap()
+    }
 }

wgpu-core/src/track/mod.rs

@@ -9,7 +9,7 @@ mod texture;
 use crate::{
     conv,
     hub::Storage,
-    id::{BindGroupId, SamplerId, TextureViewId, TypedId},
+    id::{self, TypedId},
     resource,
     Epoch,
     FastHashMap,
@@ -464,9 +464,11 @@ pub const DUMMY_SELECTOR: () = ();
 pub struct TrackerSet {
     pub buffers: ResourceTracker<BufferState>,
     pub textures: ResourceTracker<TextureState>,
-    pub views: ResourceTracker<PhantomData<TextureViewId>>,
-    pub bind_groups: ResourceTracker<PhantomData<BindGroupId>>,
-    pub samplers: ResourceTracker<PhantomData<SamplerId>>,
+    pub views: ResourceTracker<PhantomData<id::TextureViewId>>,
+    pub bind_groups: ResourceTracker<PhantomData<id::BindGroupId>>,
+    pub samplers: ResourceTracker<PhantomData<id::SamplerId>>,
+    pub compute_pipes: ResourceTracker<PhantomData<id::ComputePipelineId>>,
+    pub render_pipes: ResourceTracker<PhantomData<id::RenderPipelineId>>,
 }
 
 impl TrackerSet {
@@ -478,6 +480,8 @@ impl TrackerSet {
             views: ResourceTracker::new(backend),
             bind_groups: ResourceTracker::new(backend),
             samplers: ResourceTracker::new(backend),
+            compute_pipes: ResourceTracker::new(backend),
+            render_pipes: ResourceTracker::new(backend),
         }
     }
 
@@ -488,6 +492,8 @@ impl TrackerSet {
         self.views.clear();
         self.bind_groups.clear();
         self.samplers.clear();
+        self.compute_pipes.clear();
+        self.render_pipes.clear();
     }
 
     /// Try to optimize the tracking representation.
@@ -497,6 +503,8 @@ impl TrackerSet {
         self.views.optimize();
         self.bind_groups.optimize();
         self.samplers.optimize();
+        self.compute_pipes.optimize();
+        self.render_pipes.optimize();
     }
 
     /// Merge all the trackers of another instance by extending
@@ -507,6 +515,8 @@ impl TrackerSet {
         self.views.merge_extend(&other.views).unwrap();
         self.bind_groups.merge_extend(&other.bind_groups).unwrap();
         self.samplers.merge_extend(&other.samplers).unwrap();
+        self.compute_pipes.merge_extend(&other.compute_pipes).unwrap();
+        self.render_pipes.merge_extend(&other.render_pipes).unwrap();
     }
 
     pub fn backend(&self) -> wgt::Backend {