Merge branch 'dev' into oh_100_security_bits

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);