From 6273e66516cbb67aa38b1c41edb587d71f865023 Mon Sep 17 00:00:00 2001
From: Ryan Kim <chokobole33@gmail.com>
Date: Mon, 30 Sep 2024 19:41:51 +0900
Subject: [PATCH] test(crypto): add mmcs gpu unittest

---
 .../merkle_tree/field_merkle_tree/BUILD.bazel |  16 +++
 .../field_merkle_tree_gpu_unittest.cc         | 101 ++++++++++++++++++
 2 files changed, 117 insertions(+)
 create mode 100644 tachyon/crypto/commitments/merkle_tree/field_merkle_tree/field_merkle_tree_gpu_unittest.cc

diff --git a/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/BUILD.bazel b/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/BUILD.bazel
index bc7ae4152..97c80f9f8 100644
--- a/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/BUILD.bazel
+++ b/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/BUILD.bazel
@@ -1,3 +1,4 @@
+load("//bazel:tachyon.bzl", "if_gpu_is_configured")
 load("//bazel:tachyon_cc.bzl", "tachyon_cc_library", "tachyon_cc_unittest")
 
 package(default_visibility = ["//visibility:public"])
@@ -60,3 +61,18 @@ tachyon_cc_unittest(
         "//tachyon/math/finite_fields/test:finite_field_test",
     ],
 )
+
+tachyon_cc_unittest(
+    name = "field_merkle_tree_gpu_unittests",
+    srcs = if_gpu_is_configured(["field_merkle_tree_gpu_unittest.cc"]),
+    deps = [
+        ":field_merkle_tree_mmcs",
+        "//tachyon/crypto/hashes/sponge:padding_free_sponge",
+        "//tachyon/crypto/hashes/sponge:truncated_permutation",
+        "//tachyon/crypto/hashes/sponge/poseidon2",
+        "//tachyon/crypto/hashes/sponge/poseidon2:poseidon2_params",
+        "//tachyon/crypto/hashes/sponge/poseidon2/icicle:icicle_poseidon2_holder",
+        "//tachyon/crypto/hashes/sponge/poseidon2/param_traits:poseidon2_baby_bear",
+        "//tachyon/math/finite_fields/test:finite_field_test",
+    ],
+)
diff --git a/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/field_merkle_tree_gpu_unittest.cc b/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/field_merkle_tree_gpu_unittest.cc
new file mode 100644
index 000000000..e6502539e
--- /dev/null
+++ b/tachyon/crypto/commitments/merkle_tree/field_merkle_tree/field_merkle_tree_gpu_unittest.cc
@@ -0,0 +1,101 @@
+#include "gtest/gtest.h"
+
+#include "tachyon/crypto/commitments/merkle_tree/field_merkle_tree/field_merkle_tree.h"
+#include "tachyon/crypto/hashes/sponge/padding_free_sponge.h"
+#include "tachyon/crypto/hashes/sponge/poseidon2/param_traits/poseidon2_baby_bear.h"
+#include "tachyon/crypto/hashes/sponge/poseidon2/poseidon2.h"
+#include "tachyon/crypto/hashes/sponge/poseidon2/poseidon2_params.h"
+#include "tachyon/crypto/hashes/sponge/truncated_permutation.h"
+#include "tachyon/math/finite_fields/test/finite_field_test.h"
+#include "tachyon/math/geometry/dimensions.h"
+
+#if TACHYON_CUDA
+#include "tachyon/crypto/hashes/sponge/poseidon2/icicle/icicle_poseidon2_holder.h"
+#endif
+
+namespace tachyon::crypto {
+
+constexpr size_t kRate = 8;
+constexpr size_t kChunk = 8;
+constexpr size_t kN = 2;
+
+using F = math::BabyBear;
+using PackedF = math::PackedBabyBear;
+using Params = Poseidon2Params<Poseidon2Vendor::kPlonky3,
+                               Poseidon2Vendor::kPlonky3, F, 15, 7>;
+using PackedParams = Poseidon2Params<Poseidon2Vendor::kPlonky3,
+                                     Poseidon2Vendor::kPlonky3, PackedF, 15, 7>;
+using Poseidon2 = Poseidon2Sponge<Params>;
+using PackedPoseidon2 = Poseidon2Sponge<PackedParams>;
+using MyHasher = PaddingFreeSponge<Poseidon2, kRate, kChunk>;
+using MyPackedHasher = PaddingFreeSponge<PackedPoseidon2, kRate, kChunk>;
+using MyCompressor = TruncatedPermutation<Poseidon2, kChunk, kN>;
+using MyPackedCompressor = TruncatedPermutation<PackedPoseidon2, kChunk, kN>;
+using Tree = FieldMerkleTree<F, kChunk>;
+
+namespace {
+
+class FieldMerkleTreeGpuTest : public math::FiniteFieldTest<PackedF> {
+ protected:
+  MyHasher hasher_;
+  MyCompressor compressor_;
+  MyPackedHasher packed_hasher_;
+  MyPackedCompressor packed_compressor_;
+};
+
+}  // namespace
+
+TEST_F(FieldMerkleTreeGpuTest, CommitGpu) {
+  struct Config {
+    size_t num;
+    math::Dimensions dimensions;
+  };
+
+  struct {
+    std::vector<Config> configs;
+  } tests[] = {
+      // NOTE(chokobole): Icicle supports only rows that are powers of 2.
+      {{{4, {8, 1024}}, {5, {8, 128}}, {6, {8, 8}}}},
+      {
+          {{1, {1, 32}},
+           {1, {2, 32}},
+           {1, {3, 32}},
+           {1, {4, 32}},
+           {1, {5, 32}},
+           {1, {6, 32}},
+           {1, {7, 32}},
+           {1, {8, 32}},
+           {1, {9, 32}},
+           {1, {10, 32}}},
+      },
+  };
+
+  IciclePoseidon2Holder<F> poseidon2_holder =
+      IciclePoseidon2Holder<F>::Create<kRate>(hasher_.derived().config);
+  IcicleMMCSHolder<F> mmcs_holder = IcicleMMCSHolder<F>::Create<kRate>(
+      poseidon2_holder.get()->impl(), poseidon2_holder.get()->impl());
+
+  for (const auto& test : tests) {
+    std::vector<math::RowMajorMatrix<F>> matrices;
+    for (size_t i = 0; i < test.configs.size(); ++i) {
+      math::Dimensions dimensions = test.configs[i].dimensions;
+      for (size_t j = 0; j < test.configs[i].num; ++j) {
+        matrices.push_back(math::RowMajorMatrix<F>::Random(dimensions.height,
+                                                           dimensions.width));
+      }
+    }
+
+    std::vector<math::RowMajorMatrix<F>> matrices_tmp = matrices;
+    Tree tree_cpu = Tree::MaybeBuildOwnedGpu(hasher_, packed_hasher_,
+                                             compressor_, packed_compressor_,
+                                             std::move(matrices_tmp), nullptr);
+    Tree tree_gpu = Tree::MaybeBuildOwnedGpu(hasher_, packed_hasher_,
+                                             compressor_, packed_compressor_,
+                                             std::move(matrices), &mmcs_holder);
+    EXPECT_EQ(tree_cpu.GetRoot(), tree_gpu.GetRoot());
+    EXPECT_EQ(tree_cpu.leaves(), tree_gpu.leaves());
+    EXPECT_EQ(tree_cpu.digest_layers(), tree_gpu.digest_layers());
+  }
+}
+
+}  // namespace tachyon::crypto