perf(MerkleTre): Parallelize over subtrees

Previously, building a Merkle tree in parallel meant that each layer was
built in parallel, with all worker threads synchronizing after each
layer. This is nice to implement but, as it turns out, has more overhead
than expected. A more performant strategy is to take a number of
subtrees of the big Merkle tree and compute all subtrees in parallel.

f

Author: jan-ferdinand

twenty-first/src/util_types/merkle_tree.rs

@@ -9,6 +9,7 @@ use itertools::Itertools;
 use rayon::prelude::*;
 use thiserror::Error;
 
+use crate::error::U32_TO_USIZE_ERR;
 use crate::prelude::*;
 
 /// Indexes internal nodes of a [`MerkleTree`].
@@ -136,13 +137,9 @@ impl MerkleTree {
     /// - If the number of leafs is zero.
     /// - If the number of leafs is not a power of two.
     pub fn sequential_new(leafs: &[Digest]) -> Result<Self> {
-        let mut nodes = Self::initialize_merkle_tree_nodes(leafs)?;
-
-        for i in (ROOT_INDEX..leafs.len()).rev() {
-            nodes[i] = Tip5::hash_pair(nodes[i * 2], nodes[i * 2 + 1]);
-        }
-
-        Ok(MerkleTree { nodes })
+        let nodes = Self::initialize_merkle_tree_nodes(leafs)?;
+        let num_remaining_nodes = leafs.len();
+        Self::sequentially_fill_tree(nodes, num_remaining_nodes)
     }
 
     /// Build a MerkleTree with the given leafs.
@@ -159,27 +156,114 @@ impl MerkleTree {
         let mut nodes = Self::initialize_merkle_tree_nodes(leafs)?;
 
         // parallel
-        let mut num_nodes_on_this_level = leafs.len();
-        while num_nodes_on_this_level >= crate::config::merkle_tree_parallelization_cutoff() {
-            num_nodes_on_this_level /= 2;
-            let node_indices_on_this_level = num_nodes_on_this_level..2 * num_nodes_on_this_level;
-            let nodes_on_this_level = node_indices_on_this_level
-                .clone()
-                .into_par_iter()
-                .map(|i| Tip5::hash_pair(nodes[i * 2], nodes[i * 2 + 1]))
-                .collect::<Vec<_>>();
-            nodes[node_indices_on_this_level].copy_from_slice(&nodes_on_this_level);
+        let mut num_remaining_nodes = leafs.len();
+        let mut num_threads = Self::num_threads();
+        while num_remaining_nodes >= crate::config::merkle_tree_parallelization_cutoff() {
+            // If the number of threads is so large that the chunk size is 1
+            // (or even 0), each individual thread performs no work anymore.
+            // In such a case, the most reasonable course of action is to reduce
+            // the effective number of worker threads, which increases the chunk
+            // size.
+            //
+            // Since parallelization_cutoff >= 2, it follows that
+            // num_nodes_missing / 2 >= 1. Hence, the loop terminates at latest
+            // once num_threads equals 1.
+            while num_threads > num_remaining_nodes / 2 {
+                num_threads /= 2;
+            }
+
+            // re-slice to only include
+            // 1. the nodes that need to be computed and
+            // 2. exactly those nodes required to compute them.
+            let nodes = &mut nodes[..2 * num_remaining_nodes];
+            let subtrees = Self::subtrees_mut(nodes, num_threads);
+            subtrees.into_par_iter().for_each(|mut tree_layers| {
+                debug_assert!(tree_layers.len() > 1, "internal error: infinite iteration");
+                let mut previous_layer = tree_layers.pop().unwrap();
+                for next_layer in tree_layers.into_iter().rev() {
+                    for (node, (&left, &right)) in
+                        next_layer.iter_mut().zip(previous_layer.iter().tuples())
+                    {
+                        *node = Tip5::hash_pair(left, right);
+                    }
+                    previous_layer = next_layer;
+                }
+            });
+
+            // Update the number of remaining nodes by subtracting the number of
+            // freshly computed nodes. Equivalently, record that the tree has
+            // grown by subtree_height many layers.
+            let current_tree_height = num_remaining_nodes.ilog2();
+            let subtree_height = current_tree_height - num_threads.ilog2();
+            num_remaining_nodes >>= subtree_height;
         }
 
-        // sequential
-        let num_remaining_nodes = num_nodes_on_this_level;
+        Self::sequentially_fill_tree(nodes, num_remaining_nodes)
+    }
+
+    /// Internal helper function to de-duplicate code between
+    /// [`Self::sequential_new`] and [`Self::par_new`].
+    fn sequentially_fill_tree(mut nodes: Vec<Digest>, num_remaining_nodes: usize) -> Result<Self> {
         for i in (ROOT_INDEX..num_remaining_nodes).rev() {
             nodes[i] = Tip5::hash_pair(nodes[i * 2], nodes[i * 2 + 1]);
         }
 
         Ok(MerkleTree { nodes })
     }
 
+    /// Divides the given, contiguous slice into `num_trees` subtrees.
+    ///
+    /// The passed-in slice must represent a complete Merkle tree. Each subtree
+    /// is returned as a number of mutable slices, where each such slice
+    /// represents one layer in the subtree. The first slice contains only one
+    /// element, the subtree's root, the next slice contains two elements, the
+    /// root's children, and so on.
+    ///
+    /// In general, the top-most nodes of the complete tree will not be covered
+    /// by the returned subtrees. For example, when requesting 2 subtrees, the
+    /// complete tree's root will not be covered; when requesting 4 subtrees,
+    /// the root and its direct children will not be covered; and so on.
+    ///
+    /// The number of subtrees must be a power of two, and must not exceed the
+    /// number of leafs in the complete Merkle tree represented by the given
+    /// slice.
+    ///
+    /// Because this is an internal helper function (and only for that reason),
+    /// it's the caller's responsibility to ensure that the arguments are
+    /// integral. To recap:
+    /// - the number of `nodes` must be a power of 2
+    /// - the number of `num_trees` must be a power of 2
+    /// - the number of `nodes` must be at least `2 * num_trees`
+    fn subtrees_mut<T>(nodes: &mut [T], num_trees: usize) -> Vec<Vec<&mut [T]>> {
+        let num_leafs = nodes.len() / 2;
+        let total_tree_height = num_leafs.ilog2();
+        let sub_tree_height =
+            usize::try_from(total_tree_height - num_trees.ilog2()).expect(U32_TO_USIZE_ERR);
+
+        // a tree's “height” is the number of layers excluding the root,
+        // but we want to include the root
+        let num_layers = sub_tree_height + 1;
+        let mut subtrees = (0..num_trees)
+            .map(|_| Vec::with_capacity(num_layers))
+            .collect_vec();
+
+        // the number of nodes to skip includes the dummy node at index 0
+        let nodes_to_skip = num_trees;
+        let (_, mut nodes) = nodes.split_at_mut(nodes_to_skip);
+
+        for layer_idx in 0..num_layers {
+            let nodes_at_this_layer = 1 << layer_idx;
+            for tree in &mut subtrees {
+                let (layer, rest) = nodes.split_at_mut(nodes_at_this_layer);
+                tree.push(layer);
+                nodes = rest;
+            }
+        }
+        debug_assert!(nodes.is_empty());
+
+        subtrees
+    }
+
     /// Compute the Merkle root from the given leafs without recording any
     /// internal nodes.
     ///
@@ -240,17 +324,8 @@ impl MerkleTree {
             return Err(MerkleTreeError::IncorrectNumberOfLeafs);
         }
 
-        // To guarantee that all chunks correspond to trees of the same height,
-        // the number of threads must divide the number of leafs cleanly.
-        let num_threads = rayon::current_num_threads();
-        let num_threads = if num_threads.is_power_of_two() {
-            num_threads
-        } else {
-            num_threads.next_power_of_two() / 2 // previous power of 2
-        };
-        let mut num_threads = num_threads.max(1); // avoid division by 0
-
         // parallel
+        let mut num_threads = Self::num_threads();
         let mut leafs = Cow::Borrowed(leafs);
         while leafs.len() >= crate::config::merkle_tree_parallelization_cutoff() {
             // If the number of threads is so large that the chunk size is 1
@@ -278,6 +353,30 @@ impl MerkleTree {
         Self::sequential_frugal_root(&leafs)
     }
 
+    /// Internal helper function to determine the number of threads to use for
+    /// parallel Merkle tree construction.
+    ///
+    /// Can be used to figure out the number of chunks to split the work into,
+    /// but take care that each chunk contains at least 2 nodes, else no
+    /// meaningful work will be done and your iteration might run forever.
+    ///
+    /// Guaranteed to be a power of two.
+    ///
+    /// Respects the [`RAYON_NUM_THREADS`][rayon] environment variable, if set.
+    fn num_threads() -> usize {
+        // To guarantee that all chunks correspond to trees of the same height,
+        // the number of threads must divide the number of leafs cleanly.
+        let num_threads = rayon::current_num_threads();
+        let num_threads = if num_threads.is_power_of_two() {
+            num_threads
+        } else {
+            num_threads.next_power_of_two() / 2 // previous power of 2
+        };
+
+        // avoid division by 0
+        num_threads.max(1)
+    }
+
     /// Helps to kick off Merkle tree construction. Sets up the Merkle tree's
     /// internal nodes if (and only if) it is possible to construct a Merkle
     /// tree with the given leafs.
@@ -1360,4 +1459,62 @@ pub mod merkle_tree_test {
 
         assert_eq!(expected_paths, auth_paths);
     }
+
+    #[test]
+    fn merkle_subtrees_are_sliced_correctly() {
+        const TREE_HEIGHT: usize = 5;
+        const NUM_LEAFS: u32 = 1 << TREE_HEIGHT;
+        const NUM_NODES: u32 = 2 * NUM_LEAFS;
+        debug_assert_eq!(64, NUM_NODES);
+
+        let mut nodes = (0..NUM_NODES).collect_vec();
+
+        let all_nodes = MerkleTree::subtrees_mut(&mut nodes, 1);
+        assert_eq!(1, all_nodes.len());
+        let subtree = &all_nodes[0];
+        assert_eq!([1], subtree[0]);
+        assert_eq!([2, 3], subtree[1]);
+        assert_eq!((4..8).collect_vec().as_slice(), subtree[2]);
+        assert_eq!((8..16).collect_vec().as_slice(), subtree[3]);
+        assert_eq!((16..32).collect_vec().as_slice(), subtree[4]);
+        assert_eq!((32..64).collect_vec().as_slice(), subtree[5]);
+
+        let two_trees = MerkleTree::subtrees_mut(&mut nodes, 2);
+        assert_eq!(2, two_trees.len());
+        let left_tree = &two_trees[0];
+        assert_eq!([2], left_tree[0]);
+        assert_eq!([4, 5], left_tree[1]);
+        assert_eq!((8..12).collect_vec().as_slice(), left_tree[2]);
+        assert_eq!((16..24).collect_vec().as_slice(), left_tree[3]);
+        assert_eq!((32..48).collect_vec().as_slice(), left_tree[4]);
+        let right_tree = &two_trees[1];
+        assert_eq!([3], right_tree[0]);
+        assert_eq!([6, 7], right_tree[1]);
+        assert_eq!((12..16).collect_vec().as_slice(), right_tree[2]);
+        assert_eq!((24..32).collect_vec().as_slice(), right_tree[3]);
+        assert_eq!((48..64).collect_vec().as_slice(), right_tree[4]);
+
+        let four_trees = MerkleTree::subtrees_mut(&mut nodes, 4);
+        assert_eq!(4, four_trees.len());
+        let left_left_tree = &four_trees[0];
+        assert_eq!([4], left_left_tree[0]);
+        assert_eq!([8, 9], left_left_tree[1]);
+        assert_eq!((16..20).collect_vec().as_slice(), left_left_tree[2]);
+        assert_eq!((32..40).collect_vec().as_slice(), left_left_tree[3]);
+        let left_right_tree = &four_trees[1];
+        assert_eq!([5], left_right_tree[0]);
+        assert_eq!([10, 11], left_right_tree[1]);
+        assert_eq!((20..24).collect_vec().as_slice(), left_right_tree[2]);
+        assert_eq!((40..48).collect_vec().as_slice(), left_right_tree[3]);
+        let right_left_tree = &four_trees[2];
+        assert_eq!([6], right_left_tree[0]);
+        assert_eq!([12, 13], right_left_tree[1]);
+        assert_eq!((24..28).collect_vec().as_slice(), right_left_tree[2]);
+        assert_eq!((48..56).collect_vec().as_slice(), right_left_tree[3]);
+        let right_right_tree = &four_trees[3];
+        assert_eq!([7], right_right_tree[0]);
+        assert_eq!([14, 15], right_right_tree[1]);
+        assert_eq!((28..32).collect_vec().as_slice(), right_right_tree[2]);
+        assert_eq!((56..64).collect_vec().as_slice(), right_right_tree[3]);
+    }
 }