diff --git a/.gitignore b/.gitignore index 3361cd802..0836e69f0 100644 --- a/.gitignore +++ b/.gitignore @@ -2,4 +2,6 @@ zig-out/ -test/bun/node_modules/ \ No newline at end of file +test/bun/node_modules/ + +test/bun/benchmark_data/ \ No newline at end of file diff --git a/src/compute_indices.zig b/src/committee_indices.zig similarity index 87% rename from src/compute_indices.zig rename to src/committee_indices.zig index 25057544b..b84d9dfd8 100644 --- a/src/compute_indices.zig +++ b/src/committee_indices.zig @@ -5,19 +5,16 @@ const native_endian = builtin.target.cpu.arch.endian(); const Allocator = std.mem.Allocator; pub const SEED_SIZE = 32; -const U32U32HashMap = std.AutoHashMap(u32, u32); -const U8SliceByU32 = std.AutoHashMap(u32, []const u8); -// note that AutoHashMap always copy data in put() api -// so value should be a pointer instead of U8SliceByU32 so that it can be freed -const U8SliceByU8ByU32 = std.AutoHashMap(u32, *U8SliceByU32); +// model this as []?[32]u8 does not show better performance +const U8ArrayByU32 = std.AutoHashMap(u32, [32]u8); /// a Zig implementation of https://github.com/ChainSafe/swap-or-not-shuffle/pull/5 pub const ComputeShuffledIndex = struct { // this ComputeShuffledIndex is always init() and deinit() inside consumer's function so use arena allocator here // to improve performance and implify deinit() arena: std.heap.ArenaAllocator, - pivot_by_index: U32U32HashMap, - source_by_position_by_index: U8SliceByU8ByU32, + pivot_by_index: []?u32, + source_by_position_by_index: []?*U8ArrayByU32, // 32 bytes seed + 1 byte i pivot_buffer: [33]u8, // 32 bytes seed + 1 byte i + 4 bytes positionDiv @@ -38,10 +35,12 @@ pub const ComputeShuffledIndex = struct { return error.InvalidRounds; } - const arena = std.heap.ArenaAllocator.init(parent_allocator); + var arena = std.heap.ArenaAllocator.init(parent_allocator); - const pivot_by_index = U32U32HashMap.init(parent_allocator); - const source_by_position_by_index = U8SliceByU8ByU32.init(parent_allocator); + const pivot_by_index = try arena.allocator().alloc(?u32, @intCast(rounds)); + @memset(pivot_by_index, null); + const source_by_position_by_index = try arena.allocator().alloc(?*U8ArrayByU32, @intCast(rounds)); + @memset(source_by_position_by_index, null); var pivot_buffer: [33]u8 = [_]u8{0} ** 33; var source_buffer: [37]u8 = [_]u8{0} ** 37; @@ -60,18 +59,16 @@ pub const ComputeShuffledIndex = struct { } pub fn deinit(self: *ComputeShuffledIndex) void { - self.pivot_by_index.deinit(); - - var it = self.source_by_position_by_index.iterator(); - while (it.next()) |entry| { - var source_by_position = entry.value_ptr.*; - // no need to loop through values and free the sources inside source_by_position thanks to arena - source_by_position.deinit(); - // we create() source_by_position in the below get() api - // but no need to destroy() it thanks to arena + // pivot_by_index is deinit() by arena allocator + + for (0..self.rounds) |i| { + const source_by_position = self.source_by_position_by_index[@intCast(i)]; + if (source_by_position) |item| { + item.deinit(); + } } - self.source_by_position_by_index.deinit(); + // source_by_position_by_index is deinit() by arena allocator // this needs to be the last step self.arena.deinit(); @@ -82,7 +79,7 @@ pub const ComputeShuffledIndex = struct { const allocator = self.arena.allocator(); for (0..self.rounds) |i| { - var pivot = self.pivot_by_index.get(@intCast(i)); + var pivot = self.pivot_by_index[@intCast(i)]; if (pivot == null) { self.pivot_buffer[SEED_SIZE] = @intCast(i % 256); var digest = [_]u8{0} ** 32; @@ -90,18 +87,21 @@ pub const ComputeShuffledIndex = struct { const u64Slice = std.mem.bytesAsSlice(u64, digest[0..8]); const u64_value = u64Slice[0]; const le_value = if (native_endian == .big) @byteSwap(u64_value) else u64_value; - pivot = @intCast(le_value % self.index_count); + const _pivot: u32 = @intCast(le_value % self.index_count); + self.pivot_by_index[@intCast(i)] = _pivot; + pivot = _pivot; } const flip = (pivot.? + self.index_count - permuted) % self.index_count; const position = @max(permuted, flip); const position_div: u32 = position / 256; - var source_by_position = self.source_by_position_by_index.get(@intCast(i)); + var source_by_position = self.source_by_position_by_index[@intCast(i)]; if (source_by_position == null) { - const _source_by_position = try allocator.create(U8SliceByU32); - _source_by_position.* = U8SliceByU32.init(allocator); - try self.source_by_position_by_index.put(@intCast(i), _source_by_position); + const _source_by_position = try allocator.create(U8ArrayByU32); + _source_by_position.* = U8ArrayByU32.init(allocator); + try _source_by_position.*.ensureTotalCapacity(256); + self.source_by_position_by_index[@intCast(i)] = _source_by_position; source_by_position = _source_by_position; } @@ -111,10 +111,8 @@ pub const ComputeShuffledIndex = struct { const u32Slice = std.mem.bytesAsSlice(u32, self.source_buffer[SEED_SIZE + 1 ..]); u32Slice[0] = if (native_endian == .big) @byteSwap(position_div) else position_div; - const _source = try allocator.alloc(u8, 32); - var hash = [_]u8{0} ** 32; - Sha256.hash(self.source_buffer[0..], &hash, .{}); - @memcpy(_source, hash[0..]); + var _source: [32]u8 = undefined; + Sha256.hash(self.source_buffer[0..], &_source, .{}); try source_by_position.?.put(position_div, _source); source = _source; } @@ -180,8 +178,9 @@ fn getCommitteeIndices(allocator: Allocator, seed: []const u8, active_indices: [ var compute_shuffled_index = try ComputeShuffledIndex.init(allocator, seed, @intCast(active_indices.len), rounds); defer compute_shuffled_index.deinit(); - var shuffled_result = U32U32HashMap.init(allocator); - defer shuffled_result.deinit(); + var shuffled_result = try allocator.alloc(?u32, @intCast(active_indices.len)); + defer allocator.free(shuffled_result); + @memset(shuffled_result, null); var i: u32 = 0; var cached_hash_input = [_]u8{0} ** (32 + 8); @@ -192,10 +191,10 @@ fn getCommitteeIndices(allocator: Allocator, seed: []const u8, active_indices: [ while (next_committee_index < out.len) { const index: u32 = @intCast(i % active_indices.len); - var shuffled_index = shuffled_result.get(index); + var shuffled_index = shuffled_result[index]; if (shuffled_index == null) { const _shuffled_index = try compute_shuffled_index.get(index); - try shuffled_result.put(index, _shuffled_index); + shuffled_result[index] = _shuffled_index; shuffled_index = _shuffled_index; } const candidate_index = active_indices[@intCast(shuffled_index.?)]; diff --git a/src/root.zig b/src/root.zig index c6fb263de..43a236b53 100644 --- a/src/root.zig +++ b/src/root.zig @@ -1,8 +1,8 @@ const std = @import("std"); const testing = std.testing; pub const PubkeyIndexMap = @import("./pubkey_index_map.zig").PubkeyIndexMap; -pub const ComputeIndices = @import("./compute_indices.zig"); -pub const ComputeShuffledIndex = ComputeIndices.ComputeShuffledIndex; +pub const CommitteeIndices = @import("./committee_indices.zig"); +pub const ComputeShuffledIndex = CommitteeIndices.ComputeShuffledIndex; export fn add(a: i32, b: i32) i32 { return a + b; diff --git a/src/root_c_abi.zig b/src/root_c_abi.zig index a4596f4f6..ad463f45d 100644 --- a/src/root_c_abi.zig +++ b/src/root_c_abi.zig @@ -4,7 +4,7 @@ pub const PubkeyIndexMap = @import("pubkey_index_map.zig").PubkeyIndexMap; const PUBKEY_INDEX_MAP_KEY_SIZE = @import("pubkey_index_map.zig").PUBKEY_INDEX_MAP_KEY_SIZE; const innerShuffleList = @import("shuffle.zig").innerShuffleList; const SEED_SIZE = @import("shuffle.zig").SEED_SIZE; -const compute_indices = @import("compute_indices.zig"); +const committee_indices = @import("committee_indices.zig"); const ErrorCode = @import("error.zig").ErrorCode; const NOT_FOUND_INDEX = @import("error.zig").NOT_FOUND_INDEX; const ERROR_INDEX = @import("error.zig").ERROR_INDEX; @@ -264,26 +264,26 @@ export fn doShuffleList(active_indices: [*c]u32, len: usize, seed: [*c]u8, seed_ export fn computeProposerIndexElectra(seed: [*c]u8, seed_len: usize, active_indices: [*c]u32, active_indices_len: usize, effective_balance_increments: [*c]u16, effective_balance_increments_len: usize, max_effective_balance_electra: u64, effective_balance_increment: u32, rounds: u32) u32 { const allocator = gpa.allocator(); // TODO: is it better to define a Result struct with code and value - const proposer_index = compute_indices.computeProposerIndexElectra(allocator, seed[0..seed_len], active_indices[0..active_indices_len], effective_balance_increments[0..effective_balance_increments_len], max_effective_balance_electra, effective_balance_increment, rounds) catch return ERROR_INDEX; + const proposer_index = committee_indices.computeProposerIndexElectra(allocator, seed[0..seed_len], active_indices[0..active_indices_len], effective_balance_increments[0..effective_balance_increments_len], max_effective_balance_electra, effective_balance_increment, rounds) catch return ERROR_INDEX; return proposer_index; } -export fn computeProposerIndex(seed: [*c]u8, seed_len: usize, active_indices: [*c]u32, active_indices_len: usize, effective_balance_increments: [*c]u16, effective_balance_increments_len: usize, rand_byte_count: compute_indices.ByteCount, max_effective_balance: u64, effective_balance_increment: u32, rounds: u32) u32 { +export fn computeProposerIndex(seed: [*c]u8, seed_len: usize, active_indices: [*c]u32, active_indices_len: usize, effective_balance_increments: [*c]u16, effective_balance_increments_len: usize, rand_byte_count: committee_indices.ByteCount, max_effective_balance: u64, effective_balance_increment: u32, rounds: u32) u32 { const allocator = gpa.allocator(); // TODO: is it better to define a Result struct with code and value - const proposer_index = compute_indices.computeProposerIndex(allocator, seed[0..seed_len], active_indices[0..active_indices_len], effective_balance_increments[0..effective_balance_increments_len], rand_byte_count, max_effective_balance, effective_balance_increment, rounds) catch return ERROR_INDEX; + const proposer_index = committee_indices.computeProposerIndex(allocator, seed[0..seed_len], active_indices[0..active_indices_len], effective_balance_increments[0..effective_balance_increments_len], rand_byte_count, max_effective_balance, effective_balance_increment, rounds) catch return ERROR_INDEX; return proposer_index; } export fn computeSyncCommitteeIndicesElectra(seed: [*c]u8, seed_len: usize, active_indices: [*c]u32, active_indices_len: usize, effective_balance_increments: [*c]u16, effective_balance_increments_len: usize, max_effective_balance_electra: u64, effective_balance_increment: u32, rounds: u32, out: [*c]u32, out_len: usize) c_uint { const allocator = gpa.allocator(); - compute_indices.computeSyncCommitteeIndicesElectra(allocator, seed[0..seed_len], active_indices[0..active_indices_len], effective_balance_increments[0..effective_balance_increments_len], max_effective_balance_electra, effective_balance_increment, rounds, out[0..out_len]) catch return ErrorCode.Error; + committee_indices.computeSyncCommitteeIndicesElectra(allocator, seed[0..seed_len], active_indices[0..active_indices_len], effective_balance_increments[0..effective_balance_increments_len], max_effective_balance_electra, effective_balance_increment, rounds, out[0..out_len]) catch return ErrorCode.Error; return ErrorCode.Success; } -export fn computeSyncCommitteeIndices(seed: [*c]u8, seed_len: usize, active_indices: [*c]u32, active_indices_len: usize, effective_balance_increments: [*c]u16, effective_balance_increments_len: usize, rand_byte_count: compute_indices.ByteCount, max_effective_balance: u64, effective_balance_increment: u32, rounds: u32, out: [*c]u32, out_len: usize) c_uint { +export fn computeSyncCommitteeIndices(seed: [*c]u8, seed_len: usize, active_indices: [*c]u32, active_indices_len: usize, effective_balance_increments: [*c]u16, effective_balance_increments_len: usize, rand_byte_count: committee_indices.ByteCount, max_effective_balance: u64, effective_balance_increment: u32, rounds: u32, out: [*c]u32, out_len: usize) c_uint { const allocator = gpa.allocator(); - compute_indices.computeSyncCommitteeIndices(allocator, seed[0..seed_len], active_indices[0..active_indices_len], effective_balance_increments[0..effective_balance_increments_len], rand_byte_count, max_effective_balance, effective_balance_increment, rounds, out[0..out_len]) catch return ErrorCode.Error; + committee_indices.computeSyncCommitteeIndices(allocator, seed[0..seed_len], active_indices[0..active_indices_len], effective_balance_increments[0..effective_balance_increments_len], rand_byte_count, max_effective_balance, effective_balance_increment, rounds, out[0..out_len]) catch return ErrorCode.Error; return ErrorCode.Success; } diff --git a/test/bun/package.json b/test/bun/package.json index e0ccefc89..10a684e08 100644 --- a/test/bun/package.json +++ b/test/bun/package.json @@ -20,7 +20,7 @@ "scripts": { "test:unit": "bun test test/unit", "lint": "biome check", - "benchmark": "bun benchmark:files 'test/perf/index.test.ts'", + "benchmark": "bun benchmark:files 'test/perf/*.test.ts'", "benchmark:files": "bun ./node_modules/.bin/benchmark --config .benchrc.yaml --defaultBranch main", "lint:fix": "yarn lint --write" } diff --git a/test/bun/test/perf/committeeIndices.test.ts b/test/bun/test/perf/committeeIndices.test.ts new file mode 100644 index 000000000..3dd3cdae6 --- /dev/null +++ b/test/bun/test/perf/committeeIndices.test.ts @@ -0,0 +1,57 @@ +import { randomBytes } from "node:crypto"; +import { bench, describe } from "@chainsafe/benchmark"; +import { + EFFECTIVE_BALANCE_INCREMENT, + MAX_EFFECTIVE_BALANCE_ELECTRA, + SHUFFLE_ROUND_COUNT, + SYNC_COMMITTEE_SIZE, +} from "@lodestar/params"; +import { computeSyncCommitteeIndicesElectra } from "../../src/index.js"; +import { naiveComputeSyncCommitteeIndicesElectra } from "../referenceImplementation.js"; + +describe("computeIndices", () => { + for (const listSize of [ + 16384, 250_000, + 1_000_000, + // Don't run 4_000_000 since it's very slow and not testnet has gotten there yet + // 4e6, + ]) { + // don't want to generate random seed to investigate performance + // each seed may lead to different cached items hence different performance + const seed = new Uint8Array(32).fill(1); + const vc = listSize; + const activeIndices = new Uint32Array( + Array.from({ length: vc }, (_, i) => i), + ); + const effectiveBalanceIncrements = new Uint16Array(vc); + for (let i = 0; i < vc; i++) { + effectiveBalanceIncrements[i] = 32 + 32 * (i % 64); + } + + bench({ + id: `JS - computeSyncCommitteeIndices - ${listSize} indices`, + fn: () => { + naiveComputeSyncCommitteeIndicesElectra( + seed, + activeIndices, + effectiveBalanceIncrements, + ); + }, + }); + + bench({ + id: `Zig - computeSyncCommitteeIndices - ${listSize} indices`, + fn: () => { + computeSyncCommitteeIndicesElectra( + seed, + activeIndices, + effectiveBalanceIncrements, + SYNC_COMMITTEE_SIZE, + MAX_EFFECTIVE_BALANCE_ELECTRA, + EFFECTIVE_BALANCE_INCREMENT, + SHUFFLE_ROUND_COUNT, + ); + }, + }); + } +}); diff --git a/test/bun/test/perf/index.test.ts b/test/bun/test/perf/pubkeyIndexMap.test.ts similarity index 100% rename from test/bun/test/perf/index.test.ts rename to test/bun/test/perf/pubkeyIndexMap.test.ts diff --git a/test/bun/test/perf/shuffle.test.ts b/test/bun/test/perf/shuffle.test.ts new file mode 100644 index 000000000..4a0076a3a --- /dev/null +++ b/test/bun/test/perf/shuffle.test.ts @@ -0,0 +1,50 @@ +import { bench, describe } from "@chainsafe/benchmark"; +import { unshuffleList } from "../../src/index.js"; +import * as referenceImplementation from "../referenceImplementation.js"; + +const SHUFFLE_ROUNDS_MINIMAL = 10; +// TODO: below is Rust implementation, double check with Zig +// Lighthouse Lodestar +// 512 254.04 us 1.6034 ms (x6) +// 16384 6.2046 ms 18.272 ms (x3) +// 4000000 1.5617 s 4.9690 s (x3) + +for (const listSize of [ + 16384, 250_000, + 1_000_000, + // Don't run 4_000_000 since it's very slow and not testnet has gotten there yet + // 4e6, +]) { + describe(`shuffle list - ${listSize} indices`, () => { + const seed = Buffer.alloc(32, 0xac); + const input: number[] = []; + for (let i = 0; i < listSize; i++) input[i] = i; + const indices = new Uint32Array(input); + + bench({ + id: `JS - unshuffleList - ${listSize} indices`, + fn: () => { + referenceImplementation.unshuffleList( + indices, + seed, + SHUFFLE_ROUNDS_MINIMAL, + ); + }, + }); + + bench({ + id: `Zig - unshuffleList - ${listSize} indices`, + fn: () => { + unshuffleList(indices, seed, SHUFFLE_ROUNDS_MINIMAL); + }, + }); + + // skip this as Bun use polling solution + // bench({ + // id: `Zig - asyncUnshuffleList - ${listSize} indices`, + // fn: async () => { + // await asyncUnshuffleList(indices, seed, SHUFFLE_ROUNDS_MINIMAL); + // }, + // }); + }); +}