feat(crypto): add IcicleMMCS

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",
+    ],
+)

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 <memory>
+#include <vector>
+
+#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 <class F>
+struct IsIcicleMMCSSupportedImpl {
+  constexpr static bool value = false;
+};
+
+template <>
+struct IsIcicleMMCSSupportedImpl<math::BabyBear> {
+  constexpr static bool value = true;
+};
+
+template <>
+struct IsIcicleMMCSSupportedImpl<math::bn254::Fr> {
+  constexpr static bool value = true;
+};
+
+template <typename F>
+constexpr bool IsIcicleMMCSSupported = IsIcicleMMCSSupportedImpl<F>::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 <typename F>
+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<Eigen::Map<const math::RowMajorMatrix<F>>>&& matrices,
+      std::vector<std::vector<std::vector<F>>>* 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<math::BabyBear>::Commit(
+    std::vector<Eigen::Map<const math::RowMajorMatrix<math::BabyBear>>>&&
+        matrices,
+    std::vector<std::vector<std::vector<math::BabyBear>>>* outputs);
+
+template <>
+TACHYON_EXPORT bool IcicleMMCS<math::bn254::Fr>::Commit(
+    std::vector<Eigen::Map<const math::RowMajorMatrix<math::bn254::Fr>>>&&
+        matrices,
+    std::vector<std::vector<std::vector<math::bn254::Fr>>>* outputs);
+
+}  // namespace tachyon::crypto
+
+#endif  // TACHYON_CRYPTO_COMMITMENTS_MERKLE_TREE_FIELD_MERKLE_TREE_ICICLE_ICICLE_MMCS_H_

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 <algorithm>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#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<math::BabyBear>::Commit(
+    std::vector<Eigen::Map<const math::RowMajorMatrix<math::BabyBear>>>&&
+        matrices,
+    std::vector<std::vector<std::vector<math::BabyBear>>>* outputs) {
+#if FIELD_ID != BABY_BEAR
+#error Only BABY_BEAR is supported
+#endif
+  std::vector<std::vector<::babybear::scalar_t>> 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<const ::babybear::scalar_t*>(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<size_t>(matrices[i].cols()),
+            static_cast<size_t>(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<size_t>(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<std::vector<math::BabyBear>> 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<math::BabyBear> 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<math::BabyBear*>(&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

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_