Implement polyvecl_uniform_gamma1 using 4-way Keccak

This commit adds poly_uniform_gamma1_4x which implements uniform
sampling in [-(MLDSA_GAMMA1 - 1), MLDSA_GAMMA1] using 4-way batched
Keccak. It then implements polyvecl_uniform_gamma1 using
the new batched function.

L = 4 is implemented using one call to poly_uniform_gamma1_4x.
L = 5 is implemented using one call to poly_uniform_gamma1_4x and one call to
poly_uniform_gamma1
L = 7 is implemented using two call to poly_uniform_gamma1_4x (with one
polynomial being wasted).

Signed-off-by: Matthias J. Kannwischer <[email protected]>

Author: mkannwischer

mldsa/poly.c

@@ -561,6 +561,44 @@ void poly_uniform_gamma1(poly *a, const uint8_t seed[MLDSA_CRHBYTES],
   polyz_unpack(a, buf);
 }
 
+void poly_uniform_gamma1_4x(poly *r0, poly *r1, poly *r2, poly *r3,
+                            const uint8_t seed[MLDSA_CRHBYTES], uint16_t nonce0,
+                            uint16_t nonce1, uint16_t nonce2, uint16_t nonce3)
+{
+  /* Temporary buffers for XOF output before rejection sampling */
+  MLD_ALIGN uint8_t
+      buf[4][MLD_ALIGN_UP(POLY_UNIFORM_GAMMA1_NBLOCKS * STREAM256_BLOCKBYTES)];
+
+  MLD_ALIGN uint8_t extseed[4][MLD_ALIGN_UP(MLDSA_CRHBYTES + 2)];
+
+  /* Tracks the number of coefficients we have already sampled */
+  mld_xof256_x4_ctx state;
+
+  memcpy(extseed[0], seed, MLDSA_CRHBYTES);
+  memcpy(extseed[1], seed, MLDSA_CRHBYTES);
+  memcpy(extseed[2], seed, MLDSA_CRHBYTES);
+  memcpy(extseed[3], seed, MLDSA_CRHBYTES);
+  extseed[0][MLDSA_CRHBYTES] = nonce0 & 0xFF;
+  extseed[1][MLDSA_CRHBYTES] = nonce1 & 0xFF;
+  extseed[2][MLDSA_CRHBYTES] = nonce2 & 0xFF;
+  extseed[3][MLDSA_CRHBYTES] = nonce3 & 0xFF;
+  extseed[0][MLDSA_CRHBYTES + 1] = nonce0 >> 8;
+  extseed[1][MLDSA_CRHBYTES + 1] = nonce1 >> 8;
+  extseed[2][MLDSA_CRHBYTES + 1] = nonce2 >> 8;
+  extseed[3][MLDSA_CRHBYTES + 1] = nonce3 >> 8;
+
+  mld_xof256_x4_init(&state);
+  mld_xof256_x4_absorb(&state, extseed, MLDSA_CRHBYTES + 2);
+  mld_xof256_x4_squeezeblocks(buf, POLY_UNIFORM_GAMMA1_NBLOCKS, &state);
+
+  polyz_unpack(r0, buf[0]);
+  polyz_unpack(r1, buf[1]);
+  polyz_unpack(r2, buf[2]);
+  polyz_unpack(r3, buf[3]);
+  mld_xof256_x4_release(&state);
+}
+
+
 void poly_challenge(poly *c, const uint8_t seed[MLDSA_CTILDEBYTES])
 {
   unsigned int i, j, pos;

mldsa/poly.h

@@ -346,7 +346,7 @@ void poly_uniform_eta_4x(poly *r0, poly *r1, poly *r2, poly *r3,
 
 #define poly_uniform_gamma1 MLD_NAMESPACE(poly_uniform_gamma1)
 /*************************************************
- * Name:        poly_uniform_gamma1m1
+ * Name:        poly_uniform_gamma1
  *
  * Description: Sample polynomial with uniformly random coefficients
  *              in [-(MLDSA_GAMMA1 - 1), MLDSA_GAMMA1] by unpacking output
@@ -366,6 +366,24 @@ __contract__(
   ensures(array_bound(a->coeffs, 0, MLDSA_N, -(MLDSA_GAMMA1 - 1), MLDSA_GAMMA1 + 1))
 );
 
+
+#define poly_uniform_gamma1 MLD_NAMESPACE(poly_uniform_gamma1)
+/*************************************************
+ * Name:        poly_uniform_gamma1_4x
+ *
+ * Description: Sample polynomial with uniformly random coefficients
+ *              in [-(MLDSA_GAMMA1 - 1), MLDSA_GAMMA1] by unpacking output
+ *              stream of SHAKE256(seed|nonce)
+ *
+ * Arguments:   - poly *a: pointer to output polynomial
+ *              - const uint8_t seed[]: byte array with seed of length
+ *                MLDSA_CRHBYTES
+ *              - uint16_t nonce: 16-bit nonce
+ **************************************************/
+void poly_uniform_gamma1_4x(poly *r0, poly *r1, poly *r2, poly *r3,
+                            const uint8_t seed[MLDSA_CRHBYTES], uint16_t nonce0,
+                            uint16_t nonce1, uint16_t nonce2, uint16_t nonce3);
+
 #define poly_challenge MLD_NAMESPACE(poly_challenge)
 /*************************************************
  * Name:        poly_challenge

mldsa/polyvec.c

@@ -85,12 +85,21 @@ void polyvec_matrix_pointwise_montgomery(polyveck *t,
 void polyvecl_uniform_gamma1(polyvecl *v, const uint8_t seed[MLDSA_CRHBYTES],
                              uint16_t nonce)
 {
-  unsigned int i;
-
-  for (i = 0; i < MLDSA_L; ++i)
-  {
-    poly_uniform_gamma1(&v->vec[i], seed, MLDSA_L * nonce + i);
-  }
+  nonce = MLDSA_L * nonce;
+#if MLDSA_L == 4
+  poly_uniform_gamma1_4x(&v->vec[0], &v->vec[1], &v->vec[2], &v->vec[3], seed,
+                         nonce, nonce + 1, nonce + 2, nonce + 3);
+#elif MLDSA_L == 5
+  poly_uniform_gamma1_4x(&v->vec[0], &v->vec[1], &v->vec[2], &v->vec[3], seed,
+                         nonce, nonce + 1, nonce + 2, nonce + 3);
+  poly_uniform_gamma1(&v->vec[4], seed, nonce + 4);
+#elif MLDSA_L == 7
+  poly_uniform_gamma1_4x(&v->vec[0], &v->vec[1], &v->vec[2],
+                         &v->vec[3 /* irrelevant */], seed, nonce, nonce + 1,
+                         nonce + 2, 0xFF /* irrelevant */);
+  poly_uniform_gamma1_4x(&v->vec[3], &v->vec[4], &v->vec[5], &v->vec[6], seed,
+                         nonce + 3, nonce + 4, nonce + 5, nonce + 6);
+#endif /* MLDSA_L == 7 */
 }
 
 void polyvecl_reduce(polyvecl *v)