kroma-network / tachyon

Modular ZK(Zero Knowledge) backend accelerated by GPU
MIT License
7.77k stars 231 forks source link

feat: enable using rust allocated ldes #539

Closed chokobole closed 2 months ago

chokobole commented 2 months ago

Description

The commit to matrices are roughly performed in the order below:

baby_bear_poseidon2::TwoAdicFri::commit()tachyon_sp1_baby_bear_poseidon2_two_adic_fri_allocate_ldes() -> tachyon_sp1_baby_bear_poseidon2_two_adic_fri_commit()FieldMerkleTreeMMCS::Commit()Radix2EvaluationDomain::CosetLDEBatch()ExpandInPlaceWithZeroPad() (This generates the leaves of a merkle tree, which is a vector of the matrix(Low Degree Extension, a.k.a, LDE) which is stored in a ProverData).

Previously, LDE was owned by both the C++ and Rust side. In order to prevent the resulting double-free, we add an additional function std::mem::forget() on the Rust side. This approach forces us to modify Plonky3 code, but the problem is we can't find where to modify for 0.1.3-succinct version in Plonky3. A naive possible solution is to allocate memory on the C++ side and de-allocate it on the Rust side, but this doesn't work because of the different memory management system and instead causes memory corruption.

To prevent a double-free or memory corruption without modifications to Plonky3 code, the Rust side should own the leaves of the FieldMerkleTree. The FieldMerkleTreeMMCS::Commit() must accept Rust allocated leaves, while the existing version is renamed to FieldMerkleTreeMMCS::CommitOwned() for internal use, such as unit tests. Similarly, FieldMerkleTree::Build() becomes FieldMerkleTree::BuildOwned().

Radix2EvaluationDomain::CosetLDEBatch() now requires Rust allocated expanded matrices, internal memory allocation is removed, and ExpandInPlaceWithZeroPad() is renamed to ExpandWithZeroPad() .

tachyon_sp1_baby_bear_poseidon2_two_adic_fri_commit() is changed to accept Rust allocated leaves for the same reason.

TomTaehoonKim commented 2 months ago

@chokobole 4ee4d79 accept to -> to accept