Lock rank cleanup, part 1

lock-analyzer/src/main.rs

@@ -22,6 +22,10 @@ use std::{
 
 use anyhow::{Context, Result};
 
+// When printing the rank graph, include the source locations where the acquisitions were
+// observed. Setting this to `false` produces output that more closely matches `rank.rs`.
+const PRINT_ACQUISITION_LOCATIONS: bool = false;
+
 fn main() -> Result<()> {
     let mut ranks: BTreeMap<u32, Rank> = BTreeMap::default();
 
@@ -113,69 +117,108 @@ fn main() -> Result<()> {
         }
     }
 
-    for older_rank in ranks.values() {
+    // Perform a DFS traversal of the graph to generate a topological order for the ranks.
+    let mut order = Vec::with_capacity(ranks.len());
+    let mut visited = 0; // mask of ranks visited by the DFS
+    let mut connected = 0; // mask of ranks that are connected to the graph
+
+    fn visit(
+        order: &mut Vec<u32>,
+        visited: &mut u64,
+        connected: &mut u64,
+        ranks: &BTreeMap<u32, Rank>,
+        bit: u32,
+    ) {
+        if *visited & (1 << bit) != 0 {
+            return;
+        }
+        *visited |= 1 << bit;
+
+        for &next_rank_bit in ranks[&bit].acquisitions.keys() {
+            *connected |= 1 << next_rank_bit;
+            visit(order, visited, connected, ranks, next_rank_bit);
+        }
+
+        order.push(bit);
+    }
+
+    for &bit in ranks.keys() {
+        visit(&mut order, &mut visited, &mut connected, &ranks, bit);
+    }
+
+    order.reverse();
+
+    // Assemble the ranks into three groups:
+    // 1. Non-leaf ranks, in the topological order.
+    // 2. Leaf ranks that are connected to the graph, alphabetically.
+    // 3. Ranks that aren't connected to the graph at all, alphabetically.
+    let mut grouped_order = Vec::with_capacity(order.len());
+
+    grouped_order.extend(order.iter().filter(|bit| !ranks[bit].is_leaf()));
+
+    grouped_order.extend({
+        let mut tmp = order
+            .iter()
+            .copied()
+            .filter(|bit| ranks[bit].is_leaf() && connected & (1 << *bit) != 0)
+            .collect::<Vec<_>>();
+        tmp.sort_by_key(|bit| &ranks[bit].const_name);
+        tmp
+    });
+
+    grouped_order.extend({
+        let mut tmp = order
+            .iter()
+            .copied()
+            .filter(|bit| ranks[bit].is_leaf() && connected & (1 << *bit) == 0)
+            .collect::<Vec<_>>();
+        tmp.sort_by_key(|bit| &ranks[bit].const_name);
+        tmp
+    });
+
+    // Finally, compute indexes so we can sort the followers of each rank in the same order
+    // as the ranks themselves.
+    let mut indices = vec![0; ranks.len()];
+    for (i, &bit) in grouped_order.iter().enumerate() {
+        indices[bit as usize] = i;
+    }
+
+    for bit in &grouped_order {
+        let older_rank = &ranks[bit];
         if older_rank.is_leaf() {
-            // We'll print leaf locks separately, below.
+            println!(
+                "    rank {} {:?} followed by {{ }};",
+                older_rank.const_name, older_rank.member_name
+            );
             continue;
         }
         println!(
             "    rank {} {:?} followed by {{",
             older_rank.const_name, older_rank.member_name
         );
-        let mut acquired_any_leaf_locks = false;
-        let mut first_newer = true;
-        for (newer_rank, locations) in &older_rank.acquisitions {
-            // List acquisitions of leaf locks at the end.
-            if ranks[newer_rank].is_leaf() {
-                acquired_any_leaf_locks = true;
-                continue;
-            }
-            if !first_newer {
-                println!();
-            }
-            for (older_location, newer_locations) in locations {
-                if newer_locations.len() == 1 {
-                    for newer_loc in newer_locations {
-                        println!("        // holding {older_location} while locking {newer_loc}");
-                    }
-                } else {
-                    println!("        // holding {older_location} while locking:");
-                    for newer_loc in newer_locations {
-                        println!("        //     {newer_loc}");
+        let mut acquisitions = older_rank.acquisitions.iter().collect::<Vec<_>>();
+        acquisitions.sort_by_key(|(&newer_rank, _)| indices[newer_rank as usize]);
+        for (newer_rank, locations) in acquisitions {
+            if PRINT_ACQUISITION_LOCATIONS {
+                for (older_location, newer_locations) in locations {
+                    if newer_locations.len() == 1 {
+                        for newer_loc in newer_locations {
+                            println!(
+                                "        // holding {older_location} while locking {newer_loc}"
+                            );
+                        }
+                    } else {
+                        println!("        // holding {older_location} while locking:");
+                        for newer_loc in newer_locations {
+                            println!("        //     {newer_loc}");
+                        }
                     }
                 }
             }
             println!("        {},", ranks[newer_rank].const_name);
-            first_newer = false;
         }
 
-        if acquired_any_leaf_locks {
-            // We checked that older_rank isn't a leaf lock, so we
-            // must have printed something above.
-            if !first_newer {
-                println!();
-            }
-            println!("        // leaf lock acquisitions:");
-            for newer_rank in older_rank.acquisitions.keys() {
-                if !ranks[newer_rank].is_leaf() {
-                    continue;
-                }
-                println!("        {},", ranks[newer_rank].const_name);
-            }
-        }
-        println!("    }};");
-        println!();
-    }
-
-    for older_rank in ranks.values() {
-        if !older_rank.is_leaf() {
-            continue;
-        }
-
-        println!(
-            "    rank {} {:?} followed by {{ }};",
-            older_rank.const_name, older_rank.member_name
-        );
+        println!("    }}");
     }
 
     Ok(())

wgpu-core/src/command/bundle.rs

@@ -341,9 +341,9 @@ impl RenderBundleEncoder {
                 .map_pass_err(scope)?;
         };
 
+        let buffer_guard = hub.buffers.read();
         let bind_group_guard = hub.bind_groups.read();
         let pipeline_guard = hub.render_pipelines.read();
-        let buffer_guard = hub.buffers.read();
 
         let mut state = State {
             trackers: RenderBundleScope::new(),

wgpu-core/src/command/mod.rs

@@ -1226,9 +1226,9 @@ impl CommandEncoder {
         self: &Arc<Self>,
         desc: &wgt::CommandBufferDescriptor<Label>,
     ) -> (Arc<CommandBuffer>, Option<CommandEncoderError>) {
-        let mut cmd_enc_status = self.data.lock();
+        let status = self.data.lock().finish();
 
-        let res = match cmd_enc_status.finish() {
+        let res = match status {
             CommandEncoderStatus::Finished(mut cmd_buf_data) => {
                 match Self::encode_commands(&self.device, &mut cmd_buf_data) {
                     Ok(()) => Ok(cmd_buf_data),

wgpu-core/src/hub.rs

@@ -124,6 +124,7 @@ use crate::{
     command::{CommandBuffer, CommandEncoder, RenderBundle},
     device::{queue::Queue, Device},
     instance::Adapter,
+    lock::rank,
     pipeline::{ComputePipeline, PipelineCache, RenderPipeline, ShaderModule},
     registry::{Registry, RegistryReport},
     resource::{
@@ -209,29 +210,36 @@ pub struct Hub {
 
 impl Hub {
     pub(crate) fn new() -> Self {
+        // Unique lock ranks are required only for registries that are accessed concurrently.
+        // This happens in render pass/bundle encoding, and bind group creation. (Concurrent
+        // access could probably be avoided even in those cases, but acquiring all the locks
+        // at once simplifies the code.)
+        //
+        // The _first_ concurrently-held registry lock uses REGISTRY_STORAGE. Others have
+        // their own named lock rank.
         Self {
             adapters: Registry::new(),
             devices: Registry::new(),
             queues: Registry::new(),
             pipeline_layouts: Registry::new(),
             shader_modules: Registry::new(),
             bind_group_layouts: Registry::new(),
-            bind_groups: Registry::new(),
+            bind_groups: Registry::with_rank(rank::REGISTRY_BIND_GROUPS),
             command_encoders: Registry::new(),
             command_buffers: Registry::new(),
             render_bundles: Registry::new(),
-            render_pipelines: Registry::new(),
+            render_pipelines: Registry::with_rank(rank::REGISTRY_RENDER_PIPELINES),
             compute_pipelines: Registry::new(),
             pipeline_caches: Registry::new(),
             query_sets: Registry::new(),
             buffers: Registry::new(),
             staging_buffers: Registry::new(),
             textures: Registry::new(),
-            texture_views: Registry::new(),
-            external_textures: Registry::new(),
-            samplers: Registry::new(),
+            texture_views: Registry::with_rank(rank::REGISTRY_TEXTURE_VIEWS),
+            external_textures: Registry::with_rank(rank::REGISTRY_EXTERNAL_TEXTURES),
+            samplers: Registry::with_rank(rank::REGISTRY_SAMPLERS),
             blas_s: Registry::new(),
-            tlas_s: Registry::new(),
+            tlas_s: Registry::with_rank(rank::REGISTRY_TLAS),
         }
     }
 

wgpu-core/src/lock/rank.rs

@@ -178,6 +178,25 @@ define_lock_ranks! {
     rank COMMAND_ALLOCATOR_FREE_ENCODERS "CommandAllocator::free_encoders" followed by {
         SHARED_TRACKER_INDEX_ALLOCATOR_INNER,
     }
+    rank REGISTRY_STORAGE "Registry::storage" followed by {
+        REGISTRY_TEXTURE_VIEWS,
+        REGISTRY_BIND_GROUPS,
+    }
+    rank REGISTRY_BIND_GROUPS "Registry::bind_groups" followed by {
+        REGISTRY_RENDER_PIPELINES,
+    }
+    rank REGISTRY_RENDER_PIPELINES "Registry::render_pipelines" followed by {
+        SHARED_TRACKER_INDEX_ALLOCATOR_INNER,
+    }
+    rank REGISTRY_TEXTURE_VIEWS "Registry::texture_views" followed by {
+        REGISTRY_SAMPLERS,
+    }
+    rank REGISTRY_SAMPLERS "Registry::samplers" followed by {
+        REGISTRY_TLAS,
+    }
+    rank REGISTRY_TLAS "Registry::tlas" followed by {
+        REGISTRY_EXTERNAL_TEXTURES,
+    }
 
     // Leaf ranks reachable from the graph above, alphabetical.
     rank BLAS_COMPACTION_STATE "Blas::compaction_state" followed by { }
@@ -187,7 +206,7 @@ define_lock_ranks! {
     rank DEVICE_DEFERRED_DESTROY "Device::deferred_destroy" followed by { }
     rank DEVICE_TRACE "Device::trace" followed by { }
     rank DEVICE_USAGE_SCOPES "Device::usage_scopes" followed by { }
-    rank REGISTRY_STORAGE "Registry::storage" followed by { }
+    rank REGISTRY_EXTERNAL_TEXTURES "Registry::external_textures" followed by { }
     rank SHARED_TRACKER_INDEX_ALLOCATOR_INNER "SharedTrackerIndexAllocator::inner" followed by { }
     rank TEXTURE_BIND_GROUPS "Texture::bind_groups" followed by { }
     rank TEXTURE_CLEAR_MODE "Texture::clear_mode" followed by { }

wgpu-core/src/lock/ranked.rs

@@ -80,6 +80,7 @@ pub struct Mutex<T> {
 /// For details, see [the module documentation][self].
 pub struct MutexGuard<'a, T> {
     inner: parking_lot::MutexGuard<'a, T>,
+    #[cfg_attr(not(miri), expect(unused))] // but `Drop` has important side effects
     saved: LockStateGuard,
 }
 
@@ -163,17 +164,58 @@ fn acquire(new_rank: LockRank, location: &'static Location<'static>) -> LockStat
 /// Check that locks are being acquired in stacking order, and update the
 /// per-thread state accordingly.
 fn release(saved: LockState) {
+    let saved_info = saved
+        .last_acquired
+        .as_ref()
+        .map(|(rank, location)| (rank.bit.member_name(), location));
+
     let prior = LOCK_STATE.replace(saved);
 
+    let (prior_name, prior_location) = prior
+        .last_acquired
+        .as_ref()
+        .map(|(rank, location)| (rank.bit.member_name(), location))
+        .expect("Releasing a lock, but no acquisition recorded");
+
     // Although Rust allows mutex guards to be dropped in any
     // order, this analysis requires that locks be acquired and
     // released in stack order: the next lock to be released must be
     // the most recently acquired lock still held.
-    assert_eq!(
-        prior.depth,
-        saved.depth + 1,
-        "Lock not released in stacking order"
-    );
+
+    match (saved.depth, saved_info) {
+        (saved_depth @ 0, None) => {
+            assert_eq!(
+                prior.depth,
+                saved_depth + 1,
+                "Lock not released in stacking order\n\
+                released {:<35} locked at {:?}\n\
+                when not expecting any locks to be held\n",
+                prior_name,
+                prior_location,
+            );
+        }
+        (0, Some(_)) => {
+            panic!("Found previous lock acquisition information, but saved.depth = 0");
+        }
+        (saved_depth, Some((saved_name, saved_location))) => {
+            assert_eq!(
+                prior.depth,
+                saved_depth + 1,
+                "Lock not released in stacking order\n\
+                expecting release of {:<35} locked at {:?}\n\
+                but instead released {:<35} locked at {:?}\n",
+                saved_name,
+                saved_location,
+                prior_name,
+                prior_location,
+            );
+        }
+        (saved_depth, None) => {
+            panic!(
+                "Found saved.depth = {saved_depth}, but no previous lock acquisition information"
+            );
+        }
+    }
 }
 
 impl<T> Mutex<T> {