Skip to content

Commit

Permalink
Merge pull request #3010 from autonomys/improve-pos-performance
Browse files Browse the repository at this point in the history 
Improve PoS performance
  • Loading branch information
@nazar-pc
nazar-pc authored Sep 11, 2024
2 parents 5c75e11 + 58f2b87 commit 9f67e67
Show file tree
Hide file tree
Showing 5 changed files with 94 additions and 46 deletions.
1 change: 1 addition & 0 deletions Cargo.lock
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions crates/subspace-proof-of-space/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ seq-macro = "0.3.5"
sha2 = { version = "0.10.7", default-features = false }
# Replacement for `parking_lot` in `no_std` environment
spin = "0.9.7"
static_assertions = "1.1.0"
subspace-core-primitives = { version = "0.1.0", path = "../subspace-core-primitives", default-features = false }

[dev-dependencies]
Expand Down
117 changes: 73 additions & 44 deletions crates/subspace-proof-of-space/src/chiapos/table.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,12 @@ use rayon::prelude::*;
use seq_macro::seq;
#[cfg(all(not(feature = "std"), any(feature = "parallel", test)))]
use spin::Mutex;
use static_assertions::const_assert;
use subspace_core_primitives::crypto::{blake3_hash, blake3_hash_list};

pub(super) const COMPUTE_F1_SIMD_FACTOR: usize = 8;
pub(super) const FIND_MATCHES_AND_COMPUTE_UNROLL_FACTOR: usize = 8;
pub(super) const HAS_MATCH_UNROLL_FACTOR: usize = 8;

/// Compute the size of `y` in bits
pub(super) const fn y_size_bits(k: u8) -> usize {
Expand Down Expand Up @@ -148,7 +151,7 @@ impl<const K: u8> Default for TablesCache<K> {
}

#[derive(Debug)]
pub(super) struct Match {
struct Match {
left_position: Position,
left_y: Y,
right_position: Position,
Expand Down Expand Up @@ -267,22 +270,31 @@ pub(super) fn compute_f1_simd<const K: u8>(
/// For verification purposes use [`has_match`] instead.
///
/// Returns `None` if either of buckets is empty.
fn find_matches<'a>(
left_bucket_ys: &'a [Y],
#[allow(clippy::too_many_arguments)]
fn find_matches<T, Map>(
left_bucket_ys: &[Y],
left_bucket_start_position: Position,
right_bucket_ys: &'a [Y],
right_bucket_ys: &[Y],
right_bucket_start_position: Position,
rmap_scratch: &'a mut Vec<RmapItem>,
left_targets: &'a LeftTargets,
) -> Option<impl Iterator<Item = Match> + 'a> {
rmap_scratch: &mut Vec<RmapItem>,
left_targets: &LeftTargets,
map: Map,
output: &mut Vec<T>,
) where
Map: Fn(Match) -> T,
{
// Clear and set to correct size with zero values
rmap_scratch.clear();
rmap_scratch.resize_with(usize::from(PARAM_BC), RmapItem::default);
let rmap = rmap_scratch;

// Both left and right buckets can be empty
let first_left_bucket_y = *left_bucket_ys.first()?;
let first_right_bucket_y = *right_bucket_ys.first()?;
let Some(&first_left_bucket_y) = left_bucket_ys.first() else {
return;
};
let Some(&first_right_bucket_y) = right_bucket_ys.first() else {
return;
};
// Since all entries in a bucket are obtained after division by `PARAM_BC`, we can compute
// quotient more efficiently by subtracting base value rather than computing remainder of
// division
Expand Down Expand Up @@ -315,31 +327,41 @@ fn find_matches<'a>(
}
};

Some(
left_bucket_ys
.iter()
.zip(left_bucket_start_position..)
.flat_map(move |(&y, left_position)| {
let r = usize::from(y) - base;
let left_targets_r = left_targets_parity
.chunks_exact(left_targets_parity.len() / usize::from(PARAM_BC))
.nth(r)
.expect("r is valid");

(0..usize::from(PARAM_M)).flat_map(move |m| {
let r_target = left_targets_r[m];
let rmap_item = rmap[usize::from(r_target)];

(rmap_item.start_position..rmap_item.start_position + rmap_item.count).map(
move |right_position| Match {
for (&y, left_position) in left_bucket_ys.iter().zip(left_bucket_start_position..) {
let r = usize::from(y) - base;
let left_targets_r = left_targets_parity
.chunks_exact(left_targets_parity.len() / usize::from(PARAM_BC))
.nth(r)
.expect("r is valid");

const_assert!(PARAM_M as usize % FIND_MATCHES_AND_COMPUTE_UNROLL_FACTOR == 0);

for r_targets in left_targets_r
.array_chunks::<{ FIND_MATCHES_AND_COMPUTE_UNROLL_FACTOR }>()
.take(usize::from(PARAM_M) / FIND_MATCHES_AND_COMPUTE_UNROLL_FACTOR)
{
let _: [(); FIND_MATCHES_AND_COMPUTE_UNROLL_FACTOR] = seq!(N in 0..8 {
[
#(
{
let rmap_item = rmap[usize::from(r_targets[N])];

for right_position in
rmap_item.start_position..rmap_item.start_position + rmap_item.count
{
let m = Match {
left_position,
left_y: y,
right_position,
},
)
})
}),
)
};
output.push(map(m));
}
},
)*
]
});
}
}
}

/// Simplified version of [`find_matches`] for verification purposes.
Expand All @@ -348,11 +370,22 @@ pub(super) fn has_match(left_y: Y, right_y: Y) -> bool {
let parity = (u32::from(left_y) / u32::from(PARAM_BC)) % 2;
let left_r = u32::from(left_y) % u32::from(PARAM_BC);

for m in 0..u32::from(PARAM_M) {
let r_target = calculate_left_target_on_demand(parity, left_r, m);
if r_target == right_r {
return true;
}
const_assert!(PARAM_M as usize % HAS_MATCH_UNROLL_FACTOR == 0);

for m in 0..u32::from(PARAM_M) / HAS_MATCH_UNROLL_FACTOR as u32 {
let _: [(); HAS_MATCH_UNROLL_FACTOR] = seq!(N in 0..8 {
[
#(
{
#[allow(clippy::identity_op)]
let r_target = calculate_left_target_on_demand(parity, left_r, m * HAS_MATCH_UNROLL_FACTOR as u32 + N);
if r_target == right_r {
return true;
}
},
)*
]
});
}

false
Expand Down Expand Up @@ -478,7 +511,7 @@ fn match_and_compute_fn<'a, const K: u8, const TABLE_NUMBER: u8, const PARENT_TA
EvaluatableUsize<{ metadata_size_bytes(K, PARENT_TABLE_NUMBER) }>: Sized,
EvaluatableUsize<{ metadata_size_bytes(K, TABLE_NUMBER) }>: Sized,
{
let Some(matches) = find_matches(
find_matches(
&last_table.ys()[usize::from(left_bucket.start_position)..]
[..usize::from(left_bucket.size)],
left_bucket.start_position,
Expand All @@ -487,13 +520,9 @@ fn match_and_compute_fn<'a, const K: u8, const TABLE_NUMBER: u8, const PARENT_TA
right_bucket.start_position,
rmap_scratch,
left_targets,
) else {
return;
};

matches.for_each(|m| {
results_table.push(match_to_result(last_table, m));
});
|m| match_to_result(last_table, m),
results_table,
)
}

#[derive(Debug)]
Expand Down
7 changes: 5 additions & 2 deletions crates/subspace-proof-of-space/src/chiapos/table/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -139,15 +139,18 @@ fn test_matches() {
right_bucket_ys.sort_unstable();
right_bucket_ys.reverse();

let matches = find_matches(
let mut matches = Vec::new();
find_matches(
&left_bucket_ys,
Position::ZERO,
&right_bucket_ys,
Position::ZERO,
&mut rmap_scratch,
&left_targets,
|m| m,
&mut matches,
);
for m in matches.unwrap() {
for m in matches {
let yl = usize::from(*left_bucket_ys.get(usize::from(m.left_position)).unwrap());
let yr = usize::from(*right_bucket_ys.get(usize::from(m.right_position)).unwrap());

Expand Down
14 changes: 14 additions & 0 deletions crates/subspace-proof-of-space/src/chiapos/table/types.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,32 +12,38 @@ use derive_more::{Add, AddAssign, From, Into};
pub(in super::super) struct X(u32);

impl Step for X {
#[inline(always)]
fn steps_between(start: &Self, end: &Self) -> Option<usize> {
u32::steps_between(&start.0, &end.0)
}

#[inline(always)]
fn forward_checked(start: Self, count: usize) -> Option<Self> {
u32::forward_checked(start.0, count).map(Self)
}

#[inline(always)]
fn backward_checked(start: Self, count: usize) -> Option<Self> {
u32::backward_checked(start.0, count).map(Self)
}
}

impl From<X> for u64 {
#[inline(always)]
fn from(value: X) -> Self {
Self::from(value.0)
}
}

impl From<X> for u128 {
#[inline(always)]
fn from(value: X) -> Self {
Self::from(value.0)
}
}

impl From<X> for usize {
#[inline(always)]
fn from(value: X) -> Self {
value.0 as Self
}
Expand All @@ -56,12 +62,14 @@ impl X {
pub(in super::super) struct Y(u32);

impl From<Y> for u128 {
#[inline(always)]
fn from(value: Y) -> Self {
Self::from(value.0)
}
}

impl From<Y> for usize {
#[inline(always)]
fn from(value: Y) -> Self {
value.0 as Self
}
Expand All @@ -80,20 +88,24 @@ impl Y {
pub(in super::super) struct Position(u32);

impl Step for Position {
#[inline(always)]
fn steps_between(start: &Self, end: &Self) -> Option<usize> {
u32::steps_between(&start.0, &end.0)
}

#[inline(always)]
fn forward_checked(start: Self, count: usize) -> Option<Self> {
u32::forward_checked(start.0, count).map(Self)
}

#[inline(always)]
fn backward_checked(start: Self, count: usize) -> Option<Self> {
u32::backward_checked(start.0, count).map(Self)
}
}

impl From<Position> for usize {
#[inline(always)]
fn from(value: Position) -> Self {
value.0 as Self
}
Expand All @@ -117,6 +129,7 @@ impl<const K: u8, const TABLE_NUMBER: u8> Default for Metadata<K, TABLE_NUMBER>
where
EvaluatableUsize<{ metadata_size_bytes(K, TABLE_NUMBER) }>: Sized,
{
#[inline(always)]
fn default() -> Self {
Self([0; metadata_size_bytes(K, TABLE_NUMBER)])
}
Expand Down Expand Up @@ -154,6 +167,7 @@ impl<const K: u8, const TABLE_NUMBER: u8> From<X> for Metadata<K, TABLE_NUMBER>
where
EvaluatableUsize<{ metadata_size_bytes(K, TABLE_NUMBER) }>: Sized,
{
#[inline(always)]
fn from(value: X) -> Self {
Self::from(u128::from(value))
}
Expand Down

0 comments on commit 9f67e67

Please sign in to comment.