Poseidon2 - add sponge function and padding support (#743)

Co-authored-by: LeonHibnik <[email protected]>

Author: danny-shterman

docs/docs/icicle/primitives/hash.md

@@ -58,23 +58,27 @@ The optional `domain_tag` pointer parameter enables domain separation, allowing
 ### Poseidon2
 
 [Poseidon2](https://eprint.iacr.org/2023/323.pdf) is a cryptographic hash function designed specifically for field elements.
-It is an improved version of the original [Poseidon](https://eprint.iacr.org/2019/458) hash, offering better performance on modern hardware. Poseidon2 is optimized for use with elliptic curve cryptography and finite fields, making it ideal for decentralized systems like blockchain. Its main advantage is balancing strong security with efficient computation, which is crucial for applications that require fast, reliable hashing.
+It is an improved version of the original [Poseidon](https://eprint.iacr.org/2019/458) hash, offering better performance on modern hardware. 
+Poseidon2 is optimized for use with elliptic curve cryptography and finite fields, making it ideal for decentralized systems like blockchain. 
+Its main advantage is balancing strong security with efficient computation, which is crucial for applications that require fast, reliable hashing.
 
 The optional `domain_tag` pointer parameter enables domain separation, allowing isolation of hash outputs across different contexts or applications.
 
-:::info
+The supported values of state size ***t*** as defined in [eprint 2023/323](https://eprint.iacr.org/2023/323.pdf) are 2, 3, 4, 8, 12, 16, 20 and 24. 
+Note that ***t*** sizes 8, 12, 16, 20 and 24 are supported only for small fields (babybear and m31).
 
-The supported values of state size ***t*** as defined in [eprint 2023/323](https://eprint.iacr.org/2023/323.pdf) are 2, 3, 4, 8, 12, 16, 20 and 24. Note that ***t*** sizes 8, 12, 16, 20 and 24 are supported only for small fields (babybear and m31).
+The S box power alpha, number of full rounds and partial rounds, rounds constants, MDS matrix, and partial matrix for each field and ***t*** can be 
+found in this [folder](https://github.com/ingonyama-zk/icicle/tree/9b1506cda9eab30fc6a8d0a338e2cfab877402f7/icicle/include/icicle/hash/poseidon2_constants/constants).
 
-:::
+There are two modes for using the Poseidon2 hash - sponge function and non-sponge (merkle tree) function. The key difference between these 
+modes is their execution pattern. The sponge function is inherently serial (each hash must wait for the previous hash to complete before 
+starting its own process), while the non-sponge function (which consists of multiple independent hashes that don't share inputs) runs in 
+parallel using GPU threads, with the number of threads equal to config.batch.
+Another difference between two modes is that currently padding is supported for the sponge function and is not supported for the non-sponge.
+For the sponge function the config.batch should be equal one.
 
-:::info
-
-The S box power alpha, number of full rounds and partial rounds, rounds constants, MDS matrix, and partial matrix for each field and ***t*** can be found in this [folder](https://github.com/ingonyama-zk/icicle/tree/9b1506cda9eab30fc6a8d0a338e2cfab877402f7/icicle/include/icicle/hash/poseidon2_constants/constants).
-
-:::
-
-In the current version the padding is not supported and should be performed by the user.
+The hash function automatically chooses between these modes based on the input size. It runs in sponge mode if the input size (including the 
+domain_tag if present) is greater than the single hash width (in this case, config.batch should be set to one). Otherwise, it uses the non-sponge mode.
 
 ## Using Hash API
 
@@ -177,10 +181,6 @@ eIcicleErr err = keccak256.hash(input.data(), input.size() / config.batch, confi
 
 Currently the poseidon sponge mode (sponge function description could be found in Sec 2.1 of [eprint 2019/458](https://eprint.iacr.org/2019/458.pdf)) isn't implemented.
 
-### 5. Poseidon2 sponge function
-
-Currently the poseidon2 is implemented in compression mode, the sponge mode discussed in [eprint 2023/323](https://eprint.iacr.org/2023/323.pdf) is not implemented.
-
 ### Supported Bindings
 
 - [Rust](../rust-bindings/hash)

docs/versioned_docs/version-2.8.0/icicle/primitives/msm.md

@@ -79,7 +79,7 @@ The MSM supports batch mode - running multiple MSMs in parallel. It's always bet
   struct MSMConfig {
     device_context::DeviceContext ctx; /**< Details related to the device such as its id and stream id. */
     int points_size;         /**< Number of points in the MSM. If a batch of MSMs needs to be computed, this should be
-                              *   a number of different points. So, if each MSM re-uses the same set of points, this
+                              *   a number of different points. So, if each MSM reuses the same set of points, this
                               *   variable is set equal to the MSM size. And if every MSM uses a distinct set of
                               *   points, it should be set to the product of MSM size and [batch_size](@ref
                               *   batch_size). Default value: 0 (meaning it's equal to the MSM size). */

examples/c++/polynomial-multiplication/example.cpp

@@ -23,7 +23,7 @@ void incremental_values(scalar_t* res, uint32_t count)
   }
 }
 
-// calcaulting polynomial multiplication A*B via NTT,pointwise-multiplication and INTT
+// calculating polynomial multiplication A*B via NTT,pointwise-multiplication and INTT
 // (1) allocate A,B on HOST. Randomize first half, zero second half
 // (2) allocate A,B,Res on device
 // (3) calc NTT for A and for B from host to device
@@ -99,4 +99,4 @@ int main(int argc, char** argv)
   ntt_release_domain<scalar_t>();
 
   return 0;
-}
+}

icicle/backend/cpu/src/hash/cpu_poseidon2.cpp

@@ -196,33 +196,176 @@ namespace icicle {
 
     // For merkle tree size should be equal to the arity of a single hasher multiplier by sizeof(S).
     // For sponge function it could be any number.
+    // Size parameter here is in bytes.
     eIcicleError hash(const std::byte* input, uint64_t size, const HashConfig& config, std::byte* output) const override
     {
-      unsigned int arity = m_use_domain_tag ? m_t - 1 : m_t;
+      const unsigned arity = m_use_domain_tag ? m_t - 1 : m_t;
+      bool is_sponge = false;
+      int input_size_in_scalars = size / sizeof(S);
+      if ((config.batch == 1) && (input_size_in_scalars != (m_use_domain_tag ? m_t - 1 : m_t))) { // Check if sponge
+                                                                                                  // function.
+        is_sponge = true;
+        if (config.batch != 1) {
+          ICICLE_LOG_ERROR << "The only supported value of config.batch for sponge functions is 1.\n";
+          return eIcicleError::INVALID_ARGUMENT;
+        }
+      } // sponge function
+      else { // Non-sponge function.
+        if ((m_use_domain_tag ? input_size_in_scalars : input_size_in_scalars - 1) % (m_t - 1) != 0) {
+          ICICLE_LOG_ERROR << "Padding isn't supported for non-sponge function hash. The following should be true: "
+                              "((m_use_domain_tag ? size : size-1) % (m_t-1) != 0).\n";
+          return eIcicleError::INVALID_ARGUMENT;
+        }
+      } // Non-sponge function.
 
-      // Currently sponge and padding functionalities are not supported.
-      if (size != arity * sizeof(S)) {
-        ICICLE_LOG_ERROR
-          << "Sponge function still isn't supported. The following should be true: (size == T) but it is not.\n";
-        return eIcicleError::INVALID_ARGUMENT;
+      const unsigned int T = m_t;
+      bool is_unsupported_T_for_this_field = poseidon2_constants[T].nof_upper_full_rounds == 0;
+      if (is_unsupported_T_for_this_field) {
+        ICICLE_LOG_ERROR << "Unsupported poseidon width (t = " << T << ") for this field! Planned for next version";
+        return eIcicleError::API_NOT_IMPLEMENTED;
       }
-      // Call hash_single config.batch times.
-      for (int batch_hash_idx = 0; batch_hash_idx < config.batch; batch_hash_idx++) {
-        eIcicleError err = hash_single(input, output);
+
+      int alpha = poseidon2_constants[T].alpha;
+      int nof_upper_full_rounds = poseidon2_constants[T].nof_upper_full_rounds;
+      int nof_partial_rounds = poseidon2_constants[T].nof_partial_rounds;
+      int nof_bottom_full_rounds = poseidon2_constants[T].nof_bottom_full_rounds;
+      S* rounds_constants = poseidon2_constants[T].rounds_constants;
+      S* mds_matrix = poseidon2_constants[T].mds_matrix;
+      S* partial_matrix_diagonal_m1 = poseidon2_constants[T].partial_matrix_diagonal_m1;
+
+      // Allocate temporary memory for intermediate calcs and in order not to change the input.
+      // int sponge_nof_hashers = m_use_domain_tag ? (input_size_in_scalars / arity) :  ((input_size_in_scalars - 1) /
+      // (arity - 1)); int tmp_fields_nof_scalars = is_sponge ? (T * sponge_nof_hashers) : (T * config.batch); S*
+      // tmp_fields = new S[tmp_fields_nof_scalars];
+      S* tmp_fields;
+      S* tmp_fields_init_ptr; // This pointer to keep initial tmp_fields value to perform a easy rollback when needed.
+      int sponge_nof_hashers;
+      const S* in_fields = (S*)(input);
+      int padding_size = 0;
+      S* padding;
+      if (is_sponge) {
+        if (input_size_in_scalars < T) { // Single hasher in the chain.
+          sponge_nof_hashers = 1;
+          padding_size = T - (input_size_in_scalars + (m_use_domain_tag == true));
+        } else if (input_size_in_scalars >= T) { // More than a single hasher in the chain.
+          sponge_nof_hashers = (input_size_in_scalars - !(m_use_domain_tag == true) + (T - 2)) / (T - 1);
+          bool is_padding_needed = (input_size_in_scalars - !(m_use_domain_tag == true)) % (T - 1);
+          if (is_padding_needed) {
+            padding_size = (T - 1) - ((input_size_in_scalars - !(m_use_domain_tag == true)) % (T - 1));
+          }
+        }
+        if (padding_size > 0) { // Fill padding array with 1,0,0,...
+          padding = new S[padding_size];
+          padding[0] = S::from(1);
+          for (int i = 1; i < padding_size; i++) {
+            padding[i] = S::from(0);
+          }
+        }
+        tmp_fields = new S[T * sponge_nof_hashers];
+        tmp_fields_init_ptr = tmp_fields;
+        // Take care of hasher 0. It's done separately of the rest of the hashers because of the domain tag.
+        if (m_use_domain_tag) {
+          // Domain tag exists only for the first hasher. For the rest of the hashers this
+          // input is undefined at this stage and its value will be set later.
+          // tmp_fields = {{dt, in0}, {undef, in1}, {undef, in2}, etc.}
+          memcpy(tmp_fields, &m_domain_tag, sizeof(S));
+        } else {
+          // tmp_fields = {{in0 (T inputs)}, {undef, in1 (T-1 inputs)}, {under, in2 (T-1 inputs)}, etc.}
+          memcpy(tmp_fields, &in_fields[0], sizeof(S));
+          in_fields += 1;
+        }
+        tmp_fields += 1;
+        // Take care of rest of the hashers (T-1 scalar to each hasher).
+        for (int hasher_idx = 0; hasher_idx < sponge_nof_hashers; hasher_idx++) {
+          if (hasher_idx == sponge_nof_hashers - 1 && padding_size > 0) {
+            // Last hasher in the chain. Take care of padding if needed.
+            memcpy(tmp_fields, in_fields, (T - padding_size - 1) * sizeof(S));
+            memcpy(tmp_fields + T - padding_size - 1, padding, padding_size * sizeof(S));
+          } else { // Not a last hasher in the chain. There is no padding.
+            memcpy(tmp_fields, in_fields, (T - 1) * sizeof(S));
+          }
+          in_fields += (T - 1);
+          tmp_fields += T;
+        }
+        tmp_fields = tmp_fields_init_ptr; // Rollback to initial value.
+      } // if (is_sponge) {
+      else { // Not a sponge function. The is no padding.
+        // Input of each hash should have domain tag at its input.
+        // tmp_fields = {{dt, in0 (T-1 inputs)}, {dt, in1 (T-1 inputs)}, {dt, in2 (T-1 inputs)}, etc.}
+        tmp_fields = new S[T * config.batch];
+        tmp_fields_init_ptr = tmp_fields; // Keep tmp_fields pointer for delete.
+        if (m_use_domain_tag) {
+          for (int batch_idx = 0; batch_idx < config.batch; batch_idx++) {
+            memcpy(tmp_fields, &m_domain_tag, sizeof(S));
+            memcpy(tmp_fields + 1, in_fields, (T - 1) * sizeof(S));
+            in_fields += (T - 1);
+            tmp_fields += T;
+          }
+          tmp_fields = tmp_fields_init_ptr; // Rollback to initial value.
+        } else {
+          // tmp_fields = {{in0 (T inputs)}, {in1 (T inputs)}, {in2 (T inputs)}, etc.}
+          memcpy(tmp_fields, in_fields, T * config.batch * sizeof(S));
+        }
+      }
+
+      // Hashes processing.
+      if (is_sponge) {
+        // Call hash_single for hasher[0]
+        eIcicleError err = hash_single(
+          tmp_fields /* input */, tmp_fields /* output */, alpha, nof_upper_full_rounds, nof_partial_rounds,
+          nof_bottom_full_rounds, rounds_constants, mds_matrix, partial_matrix_diagonal_m1);
+        S* tmp_fields_tmp_ptr = tmp_fields; // Save current pointer in order to access prev output.
         if (err != eIcicleError::SUCCESS) return err;
-        input += arity * sizeof(S);
-        output += sizeof(S);
+        if (sponge_nof_hashers != 1) {
+          tmp_fields[T] = tmp_fields[0]; // Current first output is an input to the next hasher.
+        }
+        tmp_fields += T;
+        // Process rest of the hashers.
+        for (int hasher_idx = 1; hasher_idx < sponge_nof_hashers; hasher_idx++) {
+          // The first output of the prev hasher is the first input of the current hasher.
+          // The T-1 new inputs of the current hasher should be added to the T-1 outputs of the
+          // prev hasher (starting from index 1).
+          for (int i = 1; i < T; i++) {
+            tmp_fields[i] = tmp_fields_tmp_ptr[i] + tmp_fields[i];
+          }
+          eIcicleError err = hash_single(
+            tmp_fields /* input */, tmp_fields /* output */, alpha, nof_upper_full_rounds, nof_partial_rounds,
+            nof_bottom_full_rounds, rounds_constants, mds_matrix, partial_matrix_diagonal_m1);
+          tmp_fields_tmp_ptr = tmp_fields; // Save current pointer in order to access prev output.
+          if (err != eIcicleError::SUCCESS) return err;
+          if (hasher_idx != sponge_nof_hashers - 1) // Not to do in the last loop to prevent mem leak.
+            tmp_fields[T] = tmp_fields[0];          // Fill first scalar of the input to the next hasher.
+          tmp_fields += T; // Now tmp_fields points to input of the next hasher before the addition.
+        } // for (int hasher_idx = 1; hasher_idx < sponge_nof_hashers; hasher_idx++) {
+        tmp_fields -= T; // Rollback to the last hasher output.
+        memcpy(output, (std::byte*)(&tmp_fields[1]), sizeof(S));
+        tmp_fields = tmp_fields_init_ptr; // Rollback to initial value.
+      } else {                            // Not a sponge function.
+        for (int batch_hash_idx = 0; batch_hash_idx < config.batch; batch_hash_idx++) {
+          eIcicleError err = hash_single(
+            tmp_fields /* input */, tmp_fields /* output */, alpha, nof_upper_full_rounds, nof_partial_rounds,
+            nof_bottom_full_rounds, rounds_constants, mds_matrix, partial_matrix_diagonal_m1);
+          if (err != eIcicleError::SUCCESS) return err;
+          memcpy(output, (std::byte*)(&tmp_fields[1]), sizeof(S));
+          tmp_fields += T;
+          output += sizeof(S);
+        }
+        tmp_fields = tmp_fields_init_ptr; // Rollback to initial value.
       }
 
+      delete[] tmp_fields;
+      if (padding_size != 0) delete[] padding;
+      tmp_fields = nullptr;
+
       return eIcicleError::SUCCESS;
-    }
+    } // eIcicleError hash(const std::byte* input, uint64_t size, const HashConfig& config, std::byte* output) const
+      // override
 
   private:
     // // DEBUG start. Do not remove!!!
-    // void print_state(std::string str, S* state_to_print)  const {
+    // void print_state(std::string str, const S* state_to_print, int count)  const {
     //   std::cout << str << std::endl;
-    //   unsigned int T = m_t;
-    //   for (int state_idx = 0; state_idx < T; state_idx++) { // Columns of matrix.
+    //   for (int state_idx = 0; state_idx < count; state_idx++) { // Columns of matrix.
     //     std::cout << std::hex << state_to_print[state_idx] << std::endl;
     //   }
     // }
@@ -236,40 +379,21 @@ namespace icicle {
     // // DEBUG end
 
     // This function performs a single hash according to parameters in the poseidon2_constants[] struct.
-    eIcicleError hash_single(const std::byte* input, std::byte* output) const
+    // eIcicleError hash_single(const std::byte* input, std::byte* output) const
+    eIcicleError hash_single(
+      S* tmp_fields,
+      S* hasher_output,
+      int alpha,
+      int nof_upper_full_rounds,
+      int nof_partial_rounds,
+      int nof_bottom_full_rounds,
+      S* rounds_constants,
+      S* mds_matrix,
+      S* partial_matrix_diagonal_m1) const
     {
       const unsigned int T = m_t;
-      bool is_unsupported_T_for_this_field = poseidon2_constants[T].nof_upper_full_rounds == 0;
-      if (is_unsupported_T_for_this_field) {
-        ICICLE_LOG_ERROR << "Unsupported poseidon width (t=" << T << ") for this field! Planned for next version";
-        return eIcicleError::API_NOT_IMPLEMENTED;
-      }
 
-      unsigned int alpha = poseidon2_constants[T].alpha;
-      unsigned int nof_upper_full_rounds = poseidon2_constants[T].nof_upper_full_rounds;
-      unsigned int nof_partial_rounds = poseidon2_constants[T].nof_partial_rounds;
-      unsigned int nof_bottom_full_rounds = poseidon2_constants[T].nof_bottom_full_rounds;
-      // S* rounds_constants = poseidon2_constants[T].rounds_constants;
-      S* rounds_constants = poseidon2_constants[T].rounds_constants;
-      S* mds_matrix = poseidon2_constants[T].mds_matrix;
-      // S* partial_matrix_diagonal = poseidon2_constants[T].partial_matrix_diagonal;
-      S* partial_matrix_diagonal_m1 = poseidon2_constants[T].partial_matrix_diagonal_m1;
-      // Allocate temporary memory for intermediate calcs.
-      S* tmp_fields = new S[T];
-      // Casting from bytes to scalar.
-      const S* in_fields = (S*)(input);
-      // Copy input scalar to the output (as a temp storage) to be used in the rounds.
-      // *tmp_fields are used as a temp storage during the calculations in this function.
-      if (m_use_domain_tag) {
-        // in that case we hash [domain_tag, t-1 field elements]
-        memcpy(tmp_fields, &m_domain_tag, sizeof(S));
-        memcpy(tmp_fields + 1, in_fields, (T - 1) * sizeof(S));
-      } else {
-        // in that case we hash [t field elements]
-        memcpy(tmp_fields, in_fields, T * sizeof(S));
-      }
-
-      // Pre-rounds full maatrix multiplication.
+      // Pre-rounds full matrix multiplication.
       full_matrix_mul_by_vector(tmp_fields, mds_matrix, tmp_fields);
 
       // Upper full rounds.
@@ -288,10 +412,8 @@ namespace icicle {
       // Bottom full rounds.
       full_rounds(nof_bottom_full_rounds, tmp_fields, rounds_constants);
 
-      memcpy(output, (std::byte*)(&tmp_fields[1]), sizeof(S));
-
-      delete[] tmp_fields;
-      tmp_fields = nullptr;
+      memcpy(hasher_output, (std::byte*)(tmp_fields), T * sizeof(S));
+      // memcpy(output, (std::byte*)(&tmp_fields[1]), sizeof(S));
 
       return eIcicleError::SUCCESS;
     } // eIcicleError hash_single(const std::byte* input, std::byte* output) const