feat(gpu_prover): implement partial tree caching (#185)

## What ❔

This PR implements partial tree caching and recomputation.

## Why ❔

Partial tree caching and recomputation allows to only store bottom part
of the tree and recompute only small slice on demand when performing
queries. This means that the memory requirement is much lower vs full
tree caching while having almost no negative performance impact.

## Is this a breaking change?
- [ ] Yes
- [x] No

## Checklist

- [x] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [x] Tests for the changes have been added / updated.
- [ ] Documentation comments have been added / updated.
- [x] Code has been formatted.

Author: robik75

gpu_prover/native/blake2s.cu

@@ -9,6 +9,12 @@ typedef uint32_t u32;
 typedef uint64_t u64;
 typedef base_field bf;
 
+#define LOG_WARP_SIZE 5
+constexpr unsigned WARP_SIZE = 1 << LOG_WARP_SIZE;
+constexpr unsigned WARP_MASK = WARP_SIZE - 1;
+
+#define FULL_MASK 0xffffffff
+
 #define ROTR32(x, y) (((x) >> (y)) ^ ((x) << (32 - (y))))
 
 #define G(a, b, c, d, x, y)                                                                                                                                    \
@@ -159,6 +165,88 @@ EXTERN __global__ void ab_gather_merkle_paths_kernel(const unsigned *indexes, co
   results[dst_index] = values[src_index];
 }
 
+EXTERN __global__ void ab_gather_rows_and_merkle_paths_kernel(const unsigned *indexes, const unsigned indexes_count, const bool bit_reverse_indexes,
+                                                              const bf *values, const unsigned log_rows_per_leaf, const unsigned cols_count,
+                                                              const unsigned log_total_leaves_count, matrix_setter<bf, st_modifier::cs> leaf_values,
+                                                              const u32 *tree_bottom, const unsigned layers_count, u32 *merkle_paths) {
+  const unsigned lane_idx = threadIdx.x;
+  const unsigned idx = blockIdx.x;
+  const unsigned index_warp = indexes[idx];
+  const unsigned index_lane = index_warp & ~WARP_MASK | lane_idx;
+  const bool is_output_lane = index_warp == index_lane;
+  const unsigned leaf_index = bit_reverse_indexes ? __brev(index_lane) >> (32 - log_total_leaves_count) : index_lane;
+  values += leaf_index << log_rows_per_leaf;
+  leaf_values.add_row(idx);
+  merkle_paths += idx * STATE_SIZE;
+  const unsigned row_mask = (1u << log_rows_per_leaf) - 1;
+  auto read = [=](const unsigned offset) {
+    const unsigned row = offset & row_mask;
+    const unsigned col = offset >> log_rows_per_leaf;
+    const auto address = values + row + (col << (log_rows_per_leaf + log_total_leaves_count));
+    return col < cols_count ? bf::into_canonical_u32(load_cs(address)) : 0;
+  };
+  u32 state[STATE_SIZE];
+  u32 block[BLOCK_SIZE];
+  initialize(state);
+  u32 t = 0;
+  const unsigned values_count = cols_count << log_rows_per_leaf;
+  unsigned offset = 0;
+  while (offset < values_count) {
+    const unsigned remaining = values_count - offset;
+    const bool is_final_block = remaining <= BLOCK_SIZE;
+#pragma unroll
+    for (unsigned i = 0; i < BLOCK_SIZE; i++, offset++) {
+      const u32 value = read(offset);
+      block[i] = value;
+      if (offset >= values_count)
+        continue;
+      if (is_output_lane)
+        leaf_values.set(bf(value));
+      leaf_values.inc_col();
+    }
+    if (is_final_block)
+      compress<true>(state, t, block, remaining);
+    else
+      compress<false>(state, t, block, BLOCK_SIZE);
+  }
+#pragma unroll
+  for (unsigned layer = 0; layer < LOG_WARP_SIZE; layer++) {
+    u32 other_state[STATE_SIZE];
+    const bool take_other_first = lane_idx >> layer & 1;
+#pragma unroll
+    for (unsigned i = 0; i < STATE_SIZE; i++) {
+      other_state[i] = __shfl_xor_sync(FULL_MASK, state[i], 1 << layer);
+      if (is_output_lane)
+        merkle_paths[i] = other_state[i];
+      if (take_other_first) {
+        block[i] = other_state[i];
+        block[i + STATE_SIZE] = state[i];
+      } else {
+        block[i] = state[i];
+        block[i + STATE_SIZE] = other_state[i];
+      }
+    }
+    initialize(state);
+    t = 0;
+    compress<true>(state, t, block, BLOCK_SIZE);
+    merkle_paths += indexes_count * STATE_SIZE;
+  }
+  if (lane_idx >= STATE_SIZE)
+    return;
+  unsigned digest_index = index_warp >> LOG_WARP_SIZE;
+  unsigned log_digests_count = log_total_leaves_count - LOG_WARP_SIZE;
+  tree_bottom += lane_idx;
+  merkle_paths += lane_idx;
+  for (unsigned layer = LOG_WARP_SIZE; layer < layers_count; layer++) {
+    const unsigned other_index = digest_index ^ 1;
+    *merkle_paths = *(tree_bottom + other_index * STATE_SIZE);
+    digest_index >>= 1;
+    tree_bottom += (1u << log_digests_count) * STATE_SIZE;
+    log_digests_count--;
+    merkle_paths += indexes_count * STATE_SIZE;
+  }
+}
+
 EXTERN __global__ void ab_blake2s_pow_kernel(const u64 *seed, const u32 bits_count, const u64 max_nonce, volatile u64 *result) {
   const uint32_t digest_mask = 0xffffffff << 32 - bits_count;
   __align__(8) u32 m_u32[BLOCK_SIZE] = {};

gpu_prover/src/blake2s.rs

@@ -198,7 +198,7 @@ cuda_kernel!(
     )
 );
 
-pub fn gather_merkle_paths_device(
+pub fn gather_merkle_paths(
     indexes: &DeviceSlice<u32>,
     values: &DeviceSlice<Digest>,
     results: &mut DeviceSlice<Digest>,
@@ -228,41 +228,68 @@ pub fn gather_merkle_paths_device(
     GatherMerklePathsFunction::default().launch(&config, &args)
 }
 
-pub fn gather_merkle_paths_host(
-    indexes: &[u32],
-    values: &[Digest],
-    results: &mut [Digest],
+cuda_kernel!(
+    GatherRowsAndMerklePaths,
+    ab_gather_rows_and_merkle_paths_kernel(
+        indexes: *const u32,
+        indexes_count: u32,
+        bit_reverse_indexes: bool,
+        values: *const BF,
+        log_rows_per_leaf: u32,
+        cols_count: u32,
+        log_total_leaves_count: u32,
+        leaf_values: MutPtrAndStride<BF>,
+        tree_bottom: *const Digest,
+        layers_count: u32,
+        merkle_paths: *mut Digest,
+    )
+);
+
+pub fn gather_rows_and_merkle_paths(
+    indexes: &DeviceSlice<u32>,
+    bit_reverse_indexes: bool,
+    values: &DeviceSlice<BF>,
+    log_rows_per_index: u32,
+    leaf_values: &mut (impl DeviceMatrixChunkMutImpl<BF> + ?Sized),
+    tree_bottom: &DeviceSlice<Digest>,
+    merkle_paths: &mut DeviceSlice<Digest>,
     layers_count: u32,
-) {
-    assert!(indexes.len() <= u32::MAX as usize);
-    let indexes_count = indexes.len() as u32;
-    let values_count = values.len();
-    assert!(values_count.is_power_of_two());
-    let log_values_count = values_count.trailing_zeros();
-    assert_ne!(log_values_count, 0);
-    let log_leaves_count = log_values_count - 1;
-    assert!(layers_count < log_leaves_count);
-    assert_eq!(indexes.len() * layers_count as usize, results.len());
-    for layer_index in 0..layers_count {
-        let layer_offset =
-            (1 << (log_leaves_count + 1)) - (1 << (log_leaves_count + 1 - layer_index));
-        for (idx, &leaf_index) in indexes.iter().enumerate() {
-            let hash_offset = ((leaf_index >> layer_index) ^ 1) as usize;
-            let dst_index = idx + (layer_index * indexes_count) as usize;
-            let src_index = layer_offset + hash_offset;
-            results[dst_index] = values[src_index];
-        }
-    }
-    /*
-     const unsigned leaf_index = indexes[idx];
-     const unsigned layer_index = blockIdx.y;
-     const unsigned layer_offset = ((1u << log_leaves_count + 1) - (1u << log_leaves_count + 1 - layer_index)) * STATE_SIZE;
-     const unsigned hash_offset = (leaf_index >> layer_index ^ 1) * STATE_SIZE;
-     const unsigned element_offset = threadIdx.x;
-     const unsigned src_index = layer_offset + hash_offset + element_offset;
-     const unsigned dst_index = layer_index * indexes_count * STATE_SIZE + idx * STATE_SIZE + element_offset;
-     results[dst_index] = values[src_index];
-    */
+    stream: &CudaStream,
+) -> CudaResult<()> {
+    let indexes_len = indexes.len();
+    let values_len = values.len();
+    let cols_count = leaf_values.cols();
+    assert_eq!(values_len % cols_count, 0);
+    let log_rows_count = (values_len / cols_count).trailing_zeros();
+    assert_eq!(leaf_values.rows(), indexes_len << log_rows_per_index);
+    assert!(indexes_len <= u32::MAX as usize);
+    let indexes_count = indexes_len as u32;
+    assert!(layers_count >= LOG_WARP_SIZE);
+    assert_eq!(indexes_len * layers_count as usize, merkle_paths.len());
+    let cols_count = cols_count as u32;
+    let log_total_leaves_count = log_rows_count - log_rows_per_index;
+    let grid_dim = indexes_count;
+    let block_dim = WARP_SIZE;
+    let config = CudaLaunchConfig::basic(grid_dim, block_dim, stream);
+    let indexes = indexes.as_ptr();
+    let values = values.as_ptr();
+    let leaf_values = leaf_values.as_mut_ptr_and_stride();
+    let tree_bottom = tree_bottom.as_ptr();
+    let merkle_paths = merkle_paths.as_mut_ptr();
+    let args = GatherRowsAndMerklePathsArguments::new(
+        indexes,
+        indexes_count,
+        bit_reverse_indexes,
+        values,
+        log_rows_per_index,
+        cols_count,
+        log_total_leaves_count,
+        leaf_values,
+        tree_bottom,
+        layers_count,
+        merkle_paths,
+    );
+    GatherRowsAndMerklePathsFunction::default().launch(&config, &args)
 }
 
 pub fn merkle_tree_cap(
@@ -546,7 +573,7 @@ mod tests {
         let mut results_device = DeviceAllocation::alloc(results_host.len()).unwrap();
         memory_copy_async(&mut indexes_device, &indexes_host, &stream).unwrap();
         memory_copy_async(&mut values_device, &values_host, &stream).unwrap();
-        gather_merkle_paths_device(
+        super::gather_merkle_paths(
             &indexes_device,
             &values_device,
             &mut results_device,

gpu_prover/src/execution/gpu_worker.rs

@@ -62,7 +62,7 @@ fn get_trees_cache_mode(_circuit_type: CircuitType, _context: &ProverContext) ->
     //     },
     //     _ => TreesCacheMode::CacheFull,
     // }
-    TreesCacheMode::CacheNone
+    TreesCacheMode::CachePatrial
 }
 
 fn gpu_worker(

gpu_prover/src/prover/memory.rs

@@ -196,6 +196,7 @@ pub(crate) fn commit_memory<'a, A: GoodAllocator>(
             }
         },
     }
+    drop(evaluations);
     memory_holder.make_evaluations_sum_to_zero_extend_and_commit(context)?;
     let src_tree_cap_accessors = memory_holder.get_tree_caps_accessors();
     let mut tree_caps = Box::new(None);