Feat/roman/tree builder (#525)

# Updates:

## Hashing

 - Added SpongeHasher class
 - Can be used to accept any hash function as an argument
 - Absorb and squeeze are now separated
- Memory management is now mostly done by SpongeHasher class, each hash
function only describes permutation kernels

## Tree builder

 - Tree builder is now hash-agnostic. 
 - Tree builder now supports 2D input (matrices)
- Tree builder can now use two different hash functions for layer 0 and
compression layers

## Poseidon1

 - Interface changed to classes
 - Now allows for any alpha
 - Now allows passing constants not in a single vector
 - Now allows for any domain tag
 - Constants are now released upon going out of scope
 - Rust wrappers changed to Poseidon struct
 
 ## Poseidon2
 
 - Interface changed to classes
 - Constants are now released upon going out of scope
 - Rust wrappers changed to Poseidon2 struct
 
## Keccak

 - Added Keccak class which inherits SpongeHasher
 - Now doesn't use gpu registers for storing states
 
 To do:
- [x] Update poseidon1 golang bindings
- [x] Update poseidon1 examples
- [x] Fix poseidon2 cuda test
- [x] Fix poseidon2 merkle tree builder test
- [x] Update keccak class with new design
- [x] Update keccak test
- [x] Check keccak correctness
- [x] Update tree builder rust wrappers
- [x] Leave doc comments

Future work:  
- [ ] Add keccak merkle tree builder externs
- [ ] Add keccak rust tree builder wrappers
- [ ] Write docs
- [ ] Add example
- [ ] Fix device output for tree builder

---------

Co-authored-by: Jeremy Felder <jeremy.felder1@gmail.com>
Co-authored-by: nonam3e <71525212+nonam3e@users.noreply.github.com>

examples/c++/multi-gpu-poseidon/example.cu

@@ -6,6 +6,9 @@
 #include "api/bn254.h"
 #include "gpu-utils/error_handler.cuh"
 
+#include "poseidon/poseidon.cuh"
+#include "hash/hash.cuh"
+
 using namespace poseidon;
 using namespace bn254;
 
@@ -20,31 +23,20 @@ void checkCudaError(cudaError_t error)
 // these global constants go into template calls
 const int size_col = 11;
 
-// this function executes the Poseidon thread
 void threadPoseidon(
   device_context::DeviceContext ctx,
   unsigned size_partition,
   scalar_t* layers,
   scalar_t* column_hashes,
-  PoseidonConstants<scalar_t>* constants)
+  Poseidon<scalar_t> * poseidon)
 {
   cudaError_t err_result = CHK_STICKY(cudaSetDevice(ctx.device_id));
   if (err_result != cudaSuccess) {
     std::cerr << "CUDA error: " << cudaGetErrorString(err_result) << std::endl;
     return;
   }
-  // CHK_IF_RETURN(); I can't use it in a standard thread function
-  PoseidonConfig column_config = {
-    ctx,   // ctx
-    false, // are_inputes_on_device
-    false, // are_outputs_on_device
-    false, // input_is_a_state
-    false, // aligned
-    false, // loop_state
-    false, // is_async
-  };
-  cudaError_t err =
-    bn254_poseidon_hash_cuda(layers, column_hashes, (size_t)size_partition, size_col, *constants, column_config);
+  SpongeConfig column_config = default_sponge_config(ctx);
+  cudaError_t err = poseidon->hash_many(layers, column_hashes, (size_t) size_partition, size_col, 1, column_config);
   checkCudaError(err);
 }
 
@@ -59,6 +51,11 @@ using FpMilliseconds = std::chrono::duration<float, std::chrono::milliseconds::p
     exit(EXIT_FAILURE);                                                                                                \
   }
 
+#define CHECK_ALLOC(ptr) if ((ptr) == nullptr) { \
+    std::cerr << "Memory allocation for '" #ptr "' failed." << std::endl; \
+    exit(EXIT_FAILURE); \
+}
+
 int main()
 {
   const unsigned size_row = (1 << 30);
@@ -116,19 +113,18 @@ int main()
   scalar_t* column_hash1 = static_cast<scalar_t*>(malloc(size_partition * sizeof(scalar_t)));
   CHECK_ALLOC(column_hash1);
 
-  PoseidonConstants<scalar_t> column_constants0, column_constants1;
-  bn254_init_optimized_poseidon_constants_cuda(size_col, ctx0, &column_constants0);
-  cudaError_t err_result = CHK_STICKY(cudaSetDevice(ctx1.device_id));
-  if (err_result != cudaSuccess) {
-    std::cerr << "CUDA error: " << cudaGetErrorString(err_result) << std::endl;
-    return;
-  }
-  bn254_init_optimized_poseidon_constants_cuda(size_col, ctx1, &column_constants1);
+    Poseidon<scalar_t> column_poseidon0(size_col, ctx0);
+    cudaError_t err_result =  CHK_STICKY(cudaSetDevice(ctx1.device_id));
+    if (err_result != cudaSuccess) {
+        std::cerr << "CUDA error: " << cudaGetErrorString(err_result) << std::endl;
+        return; 
+    }
+    Poseidon<scalar_t> column_poseidon1(size_col, ctx1);
 
   std::cout << "Parallel execution of Poseidon threads" << std::endl;
   START_TIMER(parallel);
-  std::thread thread0(threadPoseidon, ctx0, size_partition, layers0, column_hash0, &column_constants0);
-  std::thread thread1(threadPoseidon, ctx1, size_partition, layers1, column_hash1, &column_constants1);
+  std::thread thread0(threadPoseidon, ctx0, size_partition, layers0, column_hash0, &column_poseidon0);
+  std::thread thread1(threadPoseidon, ctx1, size_partition, layers1, column_hash1, &column_poseidon1);
 
   // Wait for the threads to finish
   thread0.join();
@@ -141,9 +137,9 @@ int main()
 
   std::cout << "Sequential execution of Poseidon threads" << std::endl;
   START_TIMER(sequential);
-  std::thread thread2(threadPoseidon, ctx0, size_partition, layers0, column_hash0, &column_constants0);
+  std::thread thread2(threadPoseidon, ctx0, size_partition, layers0, column_hash0, &column_poseidon0);
   thread2.join();
-  std::thread thread3(threadPoseidon, ctx0, size_partition, layers1, column_hash1, &column_constants0);
+  std::thread thread3(threadPoseidon, ctx0, size_partition, layers1, column_hash1, &column_poseidon0);
   thread3.join();
   END_TIMER(sequential, "1 GPU");
   std::cout << "Output Data from Thread 2: ";

examples/c++/polynomial-api/example.cu

@@ -3,13 +3,11 @@
 #include "polynomials/polynomials.h"
 #include "polynomials/cuda_backend/polynomial_cuda_backend.cuh"
 #include "ntt/ntt.cuh"
-#include "poseidon/tree/merkle.cuh"
+
 #include "api/bn254.h"
 #include <chrono>
 
-// using namespace field_config;
 using namespace polynomials;
-using namespace merkle;
 using namespace bn254;
 
 // define the polynomial type

examples/c++/poseidon/example.cu

@@ -4,6 +4,8 @@
 
 #include "api/bn254.h"
 #include "curves/params/bn254.cuh"
+#include "poseidon/poseidon.cuh"
+#include "hash/hash.cuh"
 using namespace poseidon;
 using namespace bn254;
 
@@ -14,13 +16,12 @@ inline uint32_t tree_index(uint32_t level, uint32_t offset) { return (1 << level
 
 // We assume the tree has leaves already set, compute all other levels
 void build_tree(
-  const uint32_t tree_height, scalar_t* tree, PoseidonConstants<scalar_t>* constants, PoseidonConfig config)
+  const uint32_t tree_height, scalar_t* tree, Poseidon<scalar_t> &poseidon, SpongeConfig &config)
 {
   for (uint32_t level = tree_height - 1; level > 0; level--) {
     const uint32_t next_level = level - 1;
     const uint32_t next_level_width = 1 << next_level;
-    bn254_poseidon_hash_cuda(
-      &tree[tree_index(level, 0)], &tree[tree_index(next_level, 0)], next_level_width, 2, *constants, config);
+    poseidon.hash_many(&tree[tree_index(level, 0)], &tree[tree_index(next_level, 0)], next_level_width, 2, 1, config);
   }
 }
 
@@ -65,8 +66,8 @@ uint32_t validate_proof(
   const uint32_t tree_height,
   const uint32_t* proof_lr,
   const scalar_t* proof_hash,
-  PoseidonConstants<scalar_t>* constants,
-  PoseidonConfig config)
+  Poseidon<scalar_t> &poseidon,
+  SpongeConfig &config)
 {
   scalar_t hashes_in[2], hash_out[1], level_hash;
   level_hash = hash;
@@ -79,7 +80,7 @@ uint32_t validate_proof(
       hashes_in[1] = level_hash;
     }
     // next level hash
-    bn254_poseidon_hash_cuda(hashes_in, hash_out, 1, 2, *constants, config);
+    poseidon.hash_many(hashes_in, hash_out, 1, 2, 1, config);
     level_hash = hash_out[0];
   }
   return proof_hash[0] == level_hash;
@@ -109,16 +110,15 @@ int main(int argc, char* argv[])
     d = d + scalar_t::one();
   }
   std::cout << "Hashing blocks into tree leaves..." << std::endl;
-  PoseidonConstants<scalar_t> constants;
-  bn254_init_optimized_poseidon_constants_cuda(data_arity, ctx, &constants);
-  PoseidonConfig config = default_poseidon_config(data_arity + 1);
-  bn254_poseidon_hash_cuda(data, &tree[tree_index(leaf_level, 0)], tree_width, 4, constants, config);
+
+  Poseidon<scalar_t> poseidon(data_arity, ctx);
+  SpongeConfig config = default_sponge_config(ctx); 
+  poseidon.hash_many(data, &tree[tree_index(leaf_level, 0)], tree_width, data_arity, 1, config);
 
   std::cout << "3. Building Merkle tree" << std::endl;
-  PoseidonConstants<scalar_t> tree_constants;
-  bn254_init_optimized_poseidon_constants_cuda(tree_arity, ctx, &tree_constants);
-  PoseidonConfig tree_config = default_poseidon_config(tree_arity + 1);
-  build_tree(tree_height, tree, &tree_constants, tree_config);
+  Poseidon<scalar_t> tree_poseidon(tree_arity, ctx);
+  SpongeConfig tree_config = default_sponge_config(ctx);
+  build_tree(tree_height, tree, tree_poseidon, tree_config);
 
   std::cout << "4. Generate membership proof" << std::endl;
   uint32_t position = tree_width - 1;
@@ -133,13 +133,13 @@ int main(int argc, char* argv[])
   std::cout << "5. Validate the hash membership" << std::endl;
   uint32_t validated;
   const scalar_t hash = tree[tree_index(leaf_level, query_position)];
-  validated = validate_proof(hash, tree_height, proof_lr, proof_hash, &tree_constants, tree_config);
+  validated = validate_proof(hash, tree_height, proof_lr, proof_hash, tree_poseidon, tree_config);
   std::cout << "Validated: " << validated << std::endl;
 
   std::cout << "6. Tamper the hash" << std::endl;
   const scalar_t tampered_hash = hash + scalar_t::one();
-  validated = validate_proof(tampered_hash, tree_height, proof_lr, proof_hash, &tree_constants, tree_config);
-
+  validated = validate_proof(tampered_hash, tree_height, proof_lr, proof_hash, tree_poseidon, tree_config);
+  
   std::cout << "7. Invalidate tamper hash membership" << std::endl;
   std::cout << "Validated: " << validated << std::endl;
   return 0;

examples/rust/poseidon/src/main.rs

@@ -2,7 +2,8 @@ use icicle_bls12_381::curve::ScalarField as F;
 
 use icicle_cuda_runtime::device_context::DeviceContext;
 
-use icicle_core::poseidon::{load_optimized_poseidon_constants, poseidon_hash_many, PoseidonConfig};
+use icicle_core::hash::{SpongeHash, SpongeConfig};
+use icicle_core::poseidon::Poseidon;
 use icicle_core::traits::FieldImpl;
 use icicle_cuda_runtime::memory::HostSlice;
 
@@ -24,14 +25,14 @@ fn main() {
     let test_size = 1 << size;
 
     println!("Running Icicle Examples: Rust Poseidon Hash");
-    let arity = 2u32;
+    let arity = 2;
     println!(
         "---------------------- Loading optimized Poseidon constants for arity={} ------------------------",
         arity
     );
     let ctx = DeviceContext::default();
-    let constants = load_optimized_poseidon_constants::<F>(arity, &ctx).unwrap();
-    let config = PoseidonConfig::default();
+    let poseidon = Poseidon::load(arity, &ctx).unwrap();
+    let config = SpongeConfig::default();
 
     println!(
         "---------------------- Input size 2^{}={} ------------------------",
@@ -45,12 +46,12 @@ fn main() {
     println!("Executing BLS12-381 Poseidon Hash on device...");
     #[cfg(feature = "profile")]
     let start = Instant::now();
-    poseidon_hash_many::<F>(
+    poseidon.hash_many(
         input_slice,
         output_slice,
-        test_size as u32,
-        arity as u32,
-        &constants,
+        test_size,
+        arity,
+        1,
         &config,
     )
     .unwrap();