diff --git a/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/BUILD.bazel b/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/BUILD.bazel new file mode 100644 index 000000000..8b6334b49 --- /dev/null +++ b/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/BUILD.bazel @@ -0,0 +1,57 @@ +load("@icicle//:build_defs.bzl", "BABY_BEAR", "BN254", "icicle_defines") +load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda") +load("//bazel:tachyon_cc.bzl", "tachyon_cc_library", "tachyon_cuda_library") + +tachyon_cuda_library( + name = "icicle_mmcs_baby_bear", + srcs = if_cuda(["icicle_mmcs_baby_bear.cc"]), + hdrs = ["icicle_mmcs_baby_bear.h"], + local_defines = icicle_defines(BABY_BEAR), + deps = [ + ":icicle_mmcs", + "//tachyon/base:bits", + "//tachyon/base:openmp_util", + "//tachyon/base/containers:container_util", + "//tachyon/device/gpu:gpu_enums", + "//tachyon/device/gpu:gpu_logging", + "@com_google_absl//absl/numeric:bits", + "@icicle//:merkle_tree_baby_bear", + ] + if_cuda([ + "@local_config_cuda//cuda:cudart_static", + ]), +) + +tachyon_cuda_library( + name = "icicle_mmcs_bn254", + srcs = if_cuda(["icicle_mmcs_bn254.cc"]), + hdrs = ["icicle_mmcs_bn254.h"], + local_defines = icicle_defines(BN254), + deps = [ + ":icicle_mmcs", + "//tachyon/base:bits", + "//tachyon/base:openmp_util", + "//tachyon/base/containers:container_util", + "//tachyon/device/gpu:gpu_enums", + "//tachyon/device/gpu:gpu_logging", + "@com_google_absl//absl/numeric:bits", + "@icicle//:merkle_tree_bn254", + ] + if_cuda([ + "@local_config_cuda//cuda:cudart_static", + ]), +) + +tachyon_cc_library( + name = "icicle_mmcs", + hdrs = ["icicle_mmcs.h"], + deps = [ + "//tachyon:export", + "//tachyon/base:logging", + "//tachyon/device/gpu:gpu_device_functions", + "//tachyon/math/elliptic_curves/bls12/bls12_381:fr", + "//tachyon/math/elliptic_curves/bn/bn254:fr", + "//tachyon/math/finite_fields/baby_bear", + "//tachyon/math/matrix:matrix_types", + "@com_google_absl//absl/types:span", + "@icicle//:hdrs", + ], +) diff --git a/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs.h b/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs.h new file mode 100644 index 000000000..1a45c2710 --- /dev/null +++ b/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs.h @@ -0,0 +1,101 @@ +#ifndef TACHYON_CRYPTO_COMMITMENTS_MERKLE_TREE_FIELD_MERKLE_TREE_ICICLE_ICICLE_MMCS_H_ +#define TACHYON_CRYPTO_COMMITMENTS_MERKLE_TREE_FIELD_MERKLE_TREE_ICICLE_ICICLE_MMCS_H_ + +#include +#include + +#include "absl/types/span.h" +#include "third_party/icicle/include/merkle-tree/merkle_tree_config.h" + +#include "tachyon/base/logging.h" +#include "tachyon/device/gpu/gpu_device_functions.h" +#include "tachyon/export.h" +#include "tachyon/math/elliptic_curves/bn/bn254/fr.h" +#include "tachyon/math/finite_fields/baby_bear/baby_bear.h" +#include "tachyon/math/matrix/matrix_types.h" + +namespace tachyon::crypto { + +template +struct IsIcicleMMCSSupportedImpl { + constexpr static bool value = false; +}; + +template <> +struct IsIcicleMMCSSupportedImpl { + constexpr static bool value = true; +}; + +template <> +struct IsIcicleMMCSSupportedImpl { + constexpr static bool value = true; +}; + +template +constexpr bool IsIcicleMMCSSupported = IsIcicleMMCSSupportedImpl::value; + +struct TACHYON_EXPORT IcicleMMCSOptions { + unsigned int arity = 2; + unsigned int keep_rows = 0; + unsigned int digest_elements = 1; + bool are_inputs_on_device = false; + bool are_outputs_on_device = false; + bool is_async = false; +}; + +template +class IcicleMMCS { + public: + IcicleMMCS(gpuMemPool_t mem_pool, gpuStream_t stream, const void* hasher, + const void* compressor, size_t rate, + const IcicleMMCSOptions& options = IcicleMMCSOptions()) + : mem_pool_(mem_pool), + stream_(stream), + hasher_(hasher), + compressor_(compressor), + rate_(rate) { + ::device_context::DeviceContext ctx{stream_, /*device_id=*/0, mem_pool_}; + config_.reset(new ::merkle_tree::TreeBuilderConfig{ + ctx, + options.arity, + options.keep_rows, + options.digest_elements, + options.are_inputs_on_device, + options.are_outputs_on_device, + options.is_async, + }); + VLOG(1) << "IcicleMMCS is created"; + } + IcicleMMCS(const IcicleMMCS& other) = delete; + IcicleMMCS& operator=(const IcicleMMCS& other) = delete; + + [[nodiscard]] bool Commit( + std::vector>>&& matrices, + std::vector>>* outputs); + + private: + gpuMemPool_t mem_pool_ = nullptr; + gpuStream_t stream_ = nullptr; + // not owned + const void* hasher_ = nullptr; + // not owned + const void* compressor_ = nullptr; + size_t rate_ = 0; + std::unique_ptr<::merkle_tree::TreeBuilderConfig> config_; +}; + +template <> +TACHYON_EXPORT bool IcicleMMCS::Commit( + std::vector>>&& + matrices, + std::vector>>* outputs); + +template <> +TACHYON_EXPORT bool IcicleMMCS::Commit( + std::vector>>&& + matrices, + std::vector>>* outputs); + +} // namespace tachyon::crypto + +#endif // TACHYON_CRYPTO_COMMITMENTS_MERKLE_TREE_FIELD_MERKLE_TREE_ICICLE_ICICLE_MMCS_H_ diff --git a/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs_baby_bear.cc b/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs_baby_bear.cc new file mode 100644 index 000000000..86e931c24 --- /dev/null +++ b/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs_baby_bear.cc @@ -0,0 +1,118 @@ +#include "tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs_baby_bear.h" + +#include +#include +#include +#include + +#include "absl/numeric/bits.h" +#include "third_party/icicle/include/poseidon2/poseidon2.cu.h" +#include "third_party/icicle/src/merkle-tree/merkle.cu.cc" // NOLINT(build/include) +#include "third_party/icicle/src/merkle-tree/mmcs.cu.cc" // NOLINT(build/include) + +#include "tachyon/base/bits.h" +#include "tachyon/base/containers/container_util.h" +#include "tachyon/base/openmp_util.h" +#include "tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs.h" +#include "tachyon/device/gpu/gpu_enums.h" +#include "tachyon/device/gpu/gpu_logging.h" + +gpuError_t tachyon_babybear_mmcs_commit_cuda( + const ::matrix::Matrix<::babybear::scalar_t>* leaves, + unsigned int number_of_inputs, ::babybear::scalar_t* digests, + const ::hash::Hasher<::babybear::scalar_t, ::babybear::scalar_t>* hasher, + const ::hash::Hasher<::babybear::scalar_t, ::babybear::scalar_t>* + compressor, + const ::merkle_tree::TreeBuilderConfig& tree_config) { + // NOTE(GideokKim): The internal logic of the icicle Merkle tree always + // assumes that the leaves exist in multiples of 2. + return ::merkle_tree::mmcs_commit<::babybear::scalar_t, ::babybear::scalar_t>( + leaves, number_of_inputs, digests, *hasher, *compressor, tree_config); +} + +namespace tachyon::crypto { + +template <> +bool IcicleMMCS::Commit( + std::vector>>&& + matrices, + std::vector>>* outputs) { +#if FIELD_ID != BABY_BEAR +#error Only BABY_BEAR is supported +#endif + std::vector> matrices_tmp(matrices.size()); + for (size_t i = 0; i < matrices.size(); ++i) { + matrices_tmp[i].resize(matrices[i].rows() * matrices[i].cols()); + OMP_PARALLEL_FOR(size_t j = 0; j < matrices[i].size(); ++j) { + matrices_tmp[i][j] = ::babybear::scalar_t::from_montgomery( + reinterpret_cast(matrices[i].data())[j]); + } + } + + std::vector<::matrix::Matrix<::babybear::scalar_t>> leaves = + base::CreateVector(matrices.size(), [&matrices_tmp, &matrices](size_t i) { + return ::matrix::Matrix<::babybear::scalar_t>{ + matrices_tmp[i].data(), + static_cast(matrices[i].cols()), + static_cast(matrices[i].rows()), + }; + }); + + size_t max_tree_height = 0; + for (const auto& matrix : matrices) { + size_t tree_height = base::bits::Log2Ceiling( + absl::bit_ceil(static_cast(matrix.rows()))); + max_tree_height = std::max(max_tree_height, tree_height); + } + config_->keep_rows = max_tree_height + 1; + config_->digest_elements = rate_; + size_t digests_len = ::merkle_tree::get_digests_len( + config_->keep_rows - 1, config_->arity, config_->digest_elements); + + std::unique_ptr<::babybear::scalar_t[]> icicle_digest( + new ::babybear::scalar_t[digests_len]); + + gpuError_t error = tachyon_babybear_mmcs_commit_cuda( + leaves.data(), leaves.size(), icicle_digest.get(), + reinterpret_cast< + const ::hash::Hasher<::babybear::scalar_t, ::babybear::scalar_t>*>( + hasher_), + reinterpret_cast< + const ::hash::Hasher<::babybear::scalar_t, ::babybear::scalar_t>*>( + compressor_), + *config_); + if (error != gpuSuccess) { + GPU_LOG(ERROR, error) << "Failed tachyon_babybear_mmcs_commit_cuda()"; + return false; + } + + // TODO(GideokKim): Optimize this. + outputs->reserve(config_->keep_rows); + size_t previous_number_of_element = 0; + for (size_t layer_idx = 0; layer_idx <= max_tree_height; ++layer_idx) { + std::vector> digest_layer; + size_t number_of_node = 1 << (max_tree_height - layer_idx); + digest_layer.reserve(number_of_node); + + for (size_t node_idx = 0; node_idx < number_of_node; ++node_idx) { + std::vector digest; + digest.reserve(config_->digest_elements); + + for (size_t element_idx = 0; element_idx < config_->digest_elements; + ++element_idx) { + size_t idx = previous_number_of_element + + config_->digest_elements * node_idx + element_idx; + icicle_digest[idx] = + ::babybear::scalar_t::to_montgomery(icicle_digest[idx]); + digest.emplace_back( + *reinterpret_cast(&icicle_digest[idx])); + } + digest_layer.emplace_back(std::move(digest)); + } + outputs->emplace_back(std::move(digest_layer)); + previous_number_of_element += number_of_node * config_->digest_elements; + } + return true; +} + +} // namespace tachyon::crypto diff --git a/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs_baby_bear.h b/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs_baby_bear.h new file mode 100644 index 000000000..dfb0a8145 --- /dev/null +++ b/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs_baby_bear.h @@ -0,0 +1,19 @@ +#ifndef TACHYON_CRYPTO_COMMITMENTS_MERKLE_TREE_FIELD_MERKLE_TREE_ICICLE_ICICLE_MMCS_BABY_BEAR_H_ +#define TACHYON_CRYPTO_COMMITMENTS_MERKLE_TREE_FIELD_MERKLE_TREE_ICICLE_ICICLE_MMCS_BABY_BEAR_H_ + +#include "third_party/icicle/include/fields/stark_fields/babybear.cu.h" +#include "third_party/icicle/include/hash/hash.cu.h" +#include "third_party/icicle/include/matrix/matrix.cu.h" +#include "third_party/icicle/include/merkle-tree/merkle_tree_config.h" + +#include "tachyon/device/gpu/gpu_device_functions.h" + +extern "C" gpuError_t tachyon_babybear_mmcs_commit_cuda( + const ::matrix::Matrix<::babybear::scalar_t>* leaves, + unsigned int number_of_inputs, ::babybear::scalar_t* digests, + const ::hash::Hasher<::babybear::scalar_t, ::babybear::scalar_t>* hasher, + const ::hash::Hasher<::babybear::scalar_t, ::babybear::scalar_t>* + compressor, + const ::merkle_tree::TreeBuilderConfig& tree_config); + +#endif // TACHYON_CRYPTO_COMMITMENTS_MERKLE_TREE_FIELD_MERKLE_TREE_ICICLE_ICICLE_MMCS_BABY_BEAR_H_ diff --git a/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs_bn254.cc b/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs_bn254.cc new file mode 100644 index 000000000..25f053253 --- /dev/null +++ b/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs_bn254.cc @@ -0,0 +1,116 @@ +#include "tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs_bn254.h" + +#include +#include +#include +#include + +#include "absl/numeric/bits.h" +#include "third_party/icicle/include/poseidon2/poseidon2.cu.h" +#include "third_party/icicle/src/merkle-tree/merkle.cu.cc" // NOLINT(build/include) +#include "third_party/icicle/src/merkle-tree/mmcs.cu.cc" // NOLINT(build/include) + +#include "tachyon/base/bits.h" +#include "tachyon/base/containers/container_util.h" +#include "tachyon/base/openmp_util.h" +#include "tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs.h" +#include "tachyon/device/gpu/gpu_enums.h" +#include "tachyon/device/gpu/gpu_logging.h" + +gpuError_t tachyon_bn254_mmcs_commit_cuda( + const ::matrix::Matrix<::bn254::scalar_t>* leaves, + unsigned int number_of_inputs, ::bn254::scalar_t* digests, + const ::hash::Hasher<::bn254::scalar_t, ::bn254::scalar_t>* hasher, + const ::hash::Hasher<::bn254::scalar_t, ::bn254::scalar_t>* compressor, + const ::merkle_tree::TreeBuilderConfig& tree_config) { + // NOTE(GideokKim): The internal logic of the icicle Merkle tree always + // assumes that the leaves exist in multiples of 2. + return ::merkle_tree::mmcs_commit<::bn254::scalar_t, ::bn254::scalar_t>( + leaves, number_of_inputs, digests, *hasher, *compressor, tree_config); +} + +namespace tachyon::crypto { + +template <> +bool IcicleMMCS::Commit( + std::vector>>&& + matrices, + std::vector>>* outputs) { +#if FIELD_ID != BN254 +#error Only BN254 is supported +#endif + std::vector> matrices_tmp(matrices.size()); + for (size_t i = 0; i < matrices.size(); ++i) { + matrices_tmp[i].resize(matrices[i].rows() * matrices[i].cols()); + OMP_PARALLEL_FOR(size_t j = 0; j < matrices[i].size(); ++j) { + matrices_tmp[i][j] = ::bn254::scalar_t::from_montgomery( + reinterpret_cast(matrices[i].data())[j]); + } + } + + std::vector<::matrix::Matrix<::bn254::scalar_t>> leaves = + base::CreateVector(matrices.size(), [&matrices_tmp, &matrices](size_t i) { + return ::matrix::Matrix<::bn254::scalar_t>{ + matrices_tmp[i].data(), + static_cast(matrices[i].cols()), + static_cast(matrices[i].rows()), + }; + }); + + size_t max_tree_height = 0; + for (const auto& matrix : matrices) { + size_t tree_height = base::bits::Log2Ceiling( + absl::bit_ceil(static_cast(matrix.rows()))); + max_tree_height = std::max(max_tree_height, tree_height); + } + config_->keep_rows = max_tree_height + 1; + config_->digest_elements = rate_; + size_t digests_len = ::merkle_tree::get_digests_len( + config_->keep_rows - 1, config_->arity, config_->digest_elements); + + std::unique_ptr<::bn254::scalar_t[]> icicle_digest( + new ::bn254::scalar_t[digests_len]); + + gpuError_t error = tachyon_bn254_mmcs_commit_cuda( + leaves.data(), leaves.size(), icicle_digest.get(), + reinterpret_cast< + const ::hash::Hasher<::bn254::scalar_t, ::bn254::scalar_t>*>(hasher_), + reinterpret_cast< + const ::hash::Hasher<::bn254::scalar_t, ::bn254::scalar_t>*>( + compressor_), + *config_); + if (error != gpuSuccess) { + GPU_LOG(ERROR, error) << "Failed tachyon_bn254_mmcs_commit_cuda()"; + return false; + } + + // TODO(GideokKim): Optimize this. + outputs->reserve(config_->keep_rows); + size_t previous_number_of_element = 0; + for (size_t layer_idx = 0; layer_idx <= max_tree_height; ++layer_idx) { + std::vector> digest_layer; + size_t number_of_node = 1 << (max_tree_height - layer_idx); + digest_layer.reserve(number_of_node); + + for (size_t node_idx = 0; node_idx < number_of_node; ++node_idx) { + std::vector digest; + digest.reserve(config_->digest_elements); + + for (size_t element_idx = 0; element_idx < config_->digest_elements; + ++element_idx) { + size_t idx = previous_number_of_element + + config_->digest_elements * node_idx + element_idx; + icicle_digest[idx] = + ::bn254::scalar_t::to_montgomery(icicle_digest[idx]); + digest.emplace_back( + *reinterpret_cast(&icicle_digest[idx])); + } + digest_layer.emplace_back(std::move(digest)); + } + outputs->emplace_back(std::move(digest_layer)); + previous_number_of_element += number_of_node * config_->digest_elements; + } + return true; +} + +} // namespace tachyon::crypto diff --git a/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs_bn254.h b/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs_bn254.h new file mode 100644 index 000000000..2fb6f6534 --- /dev/null +++ b/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/icicle/icicle_mmcs_bn254.h @@ -0,0 +1,18 @@ +#ifndef TACHYON_CRYPTO_COMMITMENTS_MERKLE_TREE_FIELD_MERKLE_TREE_ICICLE_ICICLE_MMCS_BN254_H_ +#define TACHYON_CRYPTO_COMMITMENTS_MERKLE_TREE_FIELD_MERKLE_TREE_ICICLE_ICICLE_MMCS_BN254_H_ + +#include "third_party/icicle/include/curves/params/bn254.cu.h" +#include "third_party/icicle/include/hash/hash.cu.h" +#include "third_party/icicle/include/matrix/matrix.cu.h" +#include "third_party/icicle/include/merkle-tree/merkle_tree_config.h" + +#include "tachyon/device/gpu/gpu_device_functions.h" + +extern "C" gpuError_t tachyon_bn254_mmcs_commit_cuda( + const ::matrix::Matrix<::bn254::scalar_t>* leaves, + unsigned int number_of_inputs, ::bn254::scalar_t* digests, + const ::hash::Hasher<::bn254::scalar_t, ::bn254::scalar_t>* hasher, + const ::hash::Hasher<::bn254::scalar_t, ::bn254::scalar_t>* compressor, + const ::merkle_tree::TreeBuilderConfig& tree_config); + +#endif // TACHYON_CRYPTO_COMMITMENTS_MERKLE_TREE_FIELD_MERKLE_TREE_ICICLE_ICICLE_MMCS_BN254_H_ diff --git a/third_party/icicle/change_sort_function.patch b/third_party/icicle/change_sort_function.patch new file mode 100644 index 000000000..8ee17c515 --- /dev/null +++ b/third_party/icicle/change_sort_function.patch @@ -0,0 +1,19 @@ +diff --git a/icicle/src/merkle-tree/mmcs.cu.cc b/icicle/src/merkle-tree/mmcs.cu.cc +index 1d695483..e0a0cf76 100644 +--- a/icicle/src/merkle-tree/mmcs.cu.cc ++++ b/icicle/src/merkle-tree/mmcs.cu.cc +@@ -203,10 +203,10 @@ namespace merkle_tree { + IcicleError_t::InvalidArgument, + "Hash max preimage length does not match merkle tree arity multiplied by digest elements"); + +- std::vector> sorted_inputs(number_of_inputs); +- std::partial_sort_copy( +- inputs, inputs + number_of_inputs, sorted_inputs.begin(), sorted_inputs.end(), +- [](const Matrix& left, const Matrix& right) { return left.height > right.height; }); ++ std::vector> sorted_inputs(inputs, inputs + number_of_inputs); ++ std::stable_sort(sorted_inputs.begin(), sorted_inputs.end(), [](const Matrix& left, const Matrix& right) { ++ return left.height > right.height; ++ }); + + // Check that the height of any two given matrices either rounds up + // to the same next power of two or otherwise equal diff --git a/third_party/icicle/fix_first_digest_copy_bug.patch b/third_party/icicle/fix_first_digest_copy_bug.patch new file mode 100644 index 000000000..2851e0b3d --- /dev/null +++ b/third_party/icicle/fix_first_digest_copy_bug.patch @@ -0,0 +1,13 @@ +diff --git a/icicle/src/merkle-tree/mmcs.cu.cc b/icicle/src/merkle-tree/mmcs.cu.cc +index 1d695483..1c123caa 100644 +--- a/icicle/src/merkle-tree/mmcs.cu.cc ++++ b/icicle/src/merkle-tree/mmcs.cu.cc +@@ -165,7 +165,7 @@ namespace merkle_tree { + d_leaves_info, params.number_of_leaves_to_inject, params.number_of_rows, states, params.digest_elements, + *params.hasher, *params.ctx)); + +- CHK_IF_RETURN(maybe_copy_digests(digests, big_tree_digests, params)); ++ CHK_IF_RETURN(maybe_copy_digests(states, big_tree_digests, params)); + + params.number_of_rows_padded /= params.arity; + params.segment_size /= params.arity; diff --git a/third_party/icicle/separate_merkle_tree_config.patch b/third_party/icicle/separate_merkle_tree_config.patch new file mode 100644 index 000000000..3a9677120 --- /dev/null +++ b/third_party/icicle/separate_merkle_tree_config.patch @@ -0,0 +1,76 @@ +diff --git a/icicle/include/merkle-tree/merkle.cu.h b/icicle/include/merkle-tree/merkle.cu.h +index d0a020db..2be4e4d0 100644 +--- a/icicle/include/merkle-tree/merkle.cu.h ++++ b/icicle/include/merkle-tree/merkle.cu.h +@@ -2,7 +2,7 @@ + #ifndef MERKLE_H + #define MERKLE_H + +-#include "gpu-utils/device_context.cu.h" ++#include "merkle-tree/merkle_tree_config.h" + #include "gpu-utils/error_handler.cu.h" + #include "utils/utils.h" + #include "hash/hash.cu.h" +@@ -57,27 +57,6 @@ namespace merkle_tree { + return height; + } + +- /** +- * @struct TreeBuilderConfig +- * Struct that encodes various Tree builder parameters. +- */ +- struct TreeBuilderConfig { +- device_context::DeviceContext ctx; /**< Details related to the device such as its id and stream id. */ +- unsigned int arity; +- unsigned int +- keep_rows; /**< How many rows of the Merkle tree rows should be written to output. '0' means all of them */ +- unsigned int +- digest_elements; /** @param digest_elements the size of output for each bottom layer hash and compression. +- * Will also be equal to the size of the root of the tree. Default value 1 */ +- bool are_inputs_on_device; /**< True if inputs are on device and false if they're on host. Default value: false. */ +- bool +- are_outputs_on_device; /**< True if outputs are on device and false if they're on host. Default value: false. */ +- bool is_async; /**< Whether to run the tree builder asynchronously. If set to `true`, the build_merkle_tree +- * function will be non-blocking and you'd need to synchronize it explicitly by running +- * `cudaStreamSynchronize` or `cudaDeviceSynchronize`. If set to false, the +- * function will block the current CPU thread. */ +- }; +- + static TreeBuilderConfig + default_merkle_config(const device_context::DeviceContext& ctx = device_context::get_default_device_context()) + { +diff --git a/icicle/include/merkle-tree/merkle_tree_config.h b/icicle/include/merkle-tree/merkle_tree_config.h +new file mode 100644 +index 00000000..bdccb892 +--- /dev/null ++++ b/icicle/include/merkle-tree/merkle_tree_config.h +@@ -0,0 +1,29 @@ ++#ifndef MERKLE_TREE_CONFIG_H ++#define MERKLE_TREE_CONFIG_H ++ ++#include "gpu-utils/device_context.cu.h" ++ ++namespace merkle_tree { ++ /** ++ * @struct TreeBuilderConfig ++ * Struct that encodes various Tree builder parameters. ++ */ ++ struct TreeBuilderConfig { ++ device_context::DeviceContext ctx; /**< Details related to the device such as its id and stream id. */ ++ unsigned int arity; ++ unsigned int ++ keep_rows; /**< How many rows of the Merkle tree rows should be written to output. '0' means all of them */ ++ unsigned int ++ digest_elements; /** @param digest_elements the size of output for each bottom layer hash and compression. ++ * Will also be equal to the size of the root of the tree. Default value 1 */ ++ bool are_inputs_on_device; /**< True if inputs are on device and false if they're on host. Default value: false. */ ++ bool ++ are_outputs_on_device; /**< True if outputs are on device and false if they're on host. Default value: false. */ ++ bool is_async; /**< Whether to run the tree builder asynchronously. If set to `true`, the build_merkle_tree ++ * function will be non-blocking and you'd need to synchronize it explicitly by running ++ * `cudaStreamSynchronize` or `cudaDeviceSynchronize`. If set to false, the ++ * function will block the current CPU thread. */ ++ }; ++} // namespace merkle_tree ++ ++#endif // MERKLE_TREE_CONFIG_H diff --git a/third_party/icicle/workspace.bzl b/third_party/icicle/workspace.bzl index fbba8b5ae..22eac1452 100644 --- a/third_party/icicle/workspace.bzl +++ b/third_party/icicle/workspace.bzl @@ -16,6 +16,9 @@ def repo(): "@kroma_network_tachyon//third_party/icicle:separate_msm_config.patch", "@kroma_network_tachyon//third_party/icicle:separate_ntt_algorithm.patch", "@kroma_network_tachyon//third_party/icicle:separate_hash_config.patch", + "@kroma_network_tachyon//third_party/icicle:separate_merkle_tree_config.patch", + "@kroma_network_tachyon//third_party/icicle:change_sort_function.patch", + "@kroma_network_tachyon//third_party/icicle:fix_first_digest_copy_bug.patch", ], build_file = "//third_party/icicle:icicle.BUILD", link_files = {"//third_party/icicle:build_defs.bzl": "build_defs.bzl"},