Revert most churn changes and add simpler solution

Replaces C-fallback with cleaner assembly-based misalignment handling.

Signed-off-by: Brendan Moran <brendan.moran@arm.com>

Author: bremoran

dev/fips202/armv81m/mve.h

@@ -17,8 +17,7 @@
 
 #if !defined(__ASSEMBLER__)
 #include "../api.h"
-
-#define MLK_KECCAKF1600_X4_BITINTERLEAVED_ODD_WORDS 50u
+#include "src/fips202_native_armv81m.h"
 
 /*
  * Native x4 permutation
@@ -34,238 +33,33 @@ static MLK_INLINE int mlk_keccak_f1600_x4_native(uint64_t *state)
   return mlk_keccak_f1600_x4_native_impl(state);
 }
 
-/*
- * Slow path for buffers that the MVE word gather/scatter main loops cannot
- * access directly. Returning MLK_NATIVE_FUNC_FALLBACK is not safe here because
- * the Armv8.1-M x4 state is bit-interleaved, while the generic C x4 fallback
- * uses four contiguous scalar Keccak states.
- */
-static MLK_INLINE uint32_t mlk_keccakf1600x4_bitinterleave_even_u8(uint8_t x)
-{
-  uint32_t r;
-
-  r = x & 0x01u;
-  r |= (x & 0x04u) >> 1;
-  r |= (x & 0x10u) >> 2;
-  r |= (x & 0x40u) >> 3;
-  return r;
-}
-
-static MLK_INLINE uint32_t mlk_keccakf1600x4_bitinterleave_odd_u8(uint8_t x)
-{
-  uint32_t r;
-
-  r = (x & 0x02u) >> 1;
-  r |= (x & 0x08u) >> 2;
-  r |= (x & 0x20u) >> 3;
-  r |= (x & 0x80u) >> 4;
-  return r;
-}
-
-static MLK_INLINE uint8_t mlk_keccakf1600x4_bitdeinterleave_u8(uint32_t even,
-                                                               uint32_t odd)
-{
-  uint8_t r;
-
-  r = (uint8_t)(even & 0x01u);
-  r = (uint8_t)(r | ((odd & 0x01u) << 1));
-  r = (uint8_t)(r | ((even & 0x02u) << 1));
-  r = (uint8_t)(r | ((odd & 0x02u) << 2));
-  r = (uint8_t)(r | ((even & 0x04u) << 2));
-  r = (uint8_t)(r | ((odd & 0x04u) << 3));
-  r = (uint8_t)(r | ((even & 0x08u) << 3));
-  r = (uint8_t)(r | ((odd & 0x08u) << 4));
-  return r;
-}
-
-static MLK_INLINE void mlk_keccakf1600x4_xor_bitinterleaved_byte(
-    uint64_t *state, uint8_t data, unsigned byte_offset, unsigned instance)
-{
-  unsigned lane;
-  unsigned byte_in_lane;
-  unsigned state_index;
-  unsigned shift;
-  uint64_t even;
-  uint64_t odd;
-
-  lane = byte_offset >> 3;
-  byte_in_lane = byte_offset & 7u;
-  state_index = (2u * lane) + (instance >> 1);
-  shift = 32u * (instance & 1u);
-
-  even = (uint64_t)mlk_keccakf1600x4_bitinterleave_even_u8(data)
-         << (4u * byte_in_lane);
-  odd = (uint64_t)mlk_keccakf1600x4_bitinterleave_odd_u8(data)
-        << (4u * byte_in_lane);
-
-  state[state_index] ^= even << shift;
-  state[MLK_KECCAKF1600_X4_BITINTERLEAVED_ODD_WORDS + state_index] ^= odd
-                                                                      << shift;
-}
-
-static MLK_INLINE uint8_t mlk_keccakf1600x4_extract_bitinterleaved_byte(
-    uint64_t *state, unsigned byte_offset, unsigned instance)
-{
-  unsigned lane;
-  unsigned byte_in_lane;
-  unsigned state_index;
-  unsigned shift;
-  uint32_t even;
-  uint32_t odd;
-
-  lane = byte_offset >> 3;
-  byte_in_lane = byte_offset & 7u;
-  state_index = (2u * lane) + (instance >> 1);
-  shift = 32u * (instance & 1u);
-
-  even =
-      (uint32_t)((state[state_index] >> shift) >> (4u * byte_in_lane)) & 0x0Fu;
-  odd = (uint32_t)((state[MLK_KECCAKF1600_X4_BITINTERLEAVED_ODD_WORDS +
-                          state_index] >>
-                    shift) >>
-                   (4u * byte_in_lane)) &
-        0x0Fu;
-
-  return mlk_keccakf1600x4_bitdeinterleave_u8(even, odd);
-}
-
-static MLK_INLINE void mlk_keccakf1600_xor_bytes_x4_bitinterleaved_c(
-    uint64_t *state, const uint8_t *data0, const uint8_t *data1,
-    const uint8_t *data2, const uint8_t *data3, unsigned offset,
-    unsigned length)
-{
-  unsigned i;
-
-  for (i = 0; i < length; i++)
-  {
-    unsigned byte_offset;
-
-    byte_offset = offset + i;
-    mlk_keccakf1600x4_xor_bitinterleaved_byte(state, data0[i], byte_offset, 0);
-    mlk_keccakf1600x4_xor_bitinterleaved_byte(state, data1[i], byte_offset, 1);
-    mlk_keccakf1600x4_xor_bitinterleaved_byte(state, data2[i], byte_offset, 2);
-    mlk_keccakf1600x4_xor_bitinterleaved_byte(state, data3[i], byte_offset, 3);
-  }
-}
-
-static MLK_INLINE void mlk_keccakf1600_extract_bytes_x4_bitinterleaved_c(
-    uint64_t *state, uint8_t *data0, uint8_t *data1, uint8_t *data2,
-    uint8_t *data3, unsigned offset, unsigned length)
-{
-  unsigned i;
-
-  for (i = 0; i < length; i++)
-  {
-    unsigned byte_offset;
-
-    byte_offset = offset + i;
-    data0[i] =
-        mlk_keccakf1600x4_extract_bitinterleaved_byte(state, byte_offset, 0);
-    data1[i] =
-        mlk_keccakf1600x4_extract_bitinterleaved_byte(state, byte_offset, 1);
-    data2[i] =
-        mlk_keccakf1600x4_extract_bitinterleaved_byte(state, byte_offset, 2);
-    data3[i] =
-        mlk_keccakf1600x4_extract_bitinterleaved_byte(state, byte_offset, 3);
-  }
-}
-
 /*
  * Native x4 XOR bytes (with on-the-fly bit interleaving)
  */
-#define mlk_keccak_f1600_x4_state_xor_bytes \
-  MLK_NAMESPACE(keccak_f1600_x4_state_xor_bytes_asm)
-void mlk_keccak_f1600_x4_state_xor_bytes(void *state, const uint8_t *data0,
-                                         const uint8_t *data1,
-                                         const uint8_t *data2,
-                                         const uint8_t *data3, unsigned offset,
-                                         unsigned length);
-
 MLK_MUST_CHECK_RETURN_VALUE
 static MLK_INLINE int mlk_keccakf1600_xor_bytes_x4_native(
     uint64_t *state, const uint8_t *data0, const uint8_t *data1,
     const uint8_t *data2, const uint8_t *data3, unsigned offset,
     unsigned length)
 {
-  unsigned offset_mod8;
-  unsigned byte_prefix;
-  unsigned main_length;
-
-  offset_mod8 = offset & 7u;
-  byte_prefix = 0;
-  if (offset_mod8 != 0u)
-  {
-    byte_prefix = 8u - offset_mod8;
-    if (byte_prefix > length)
-    {
-      byte_prefix = length;
-    }
-  }
-
-  main_length = length - byte_prefix;
-  if (main_length >= 8u &&
-      ((((uintptr_t)data0 + byte_prefix) | ((uintptr_t)data1 + byte_prefix) |
-        ((uintptr_t)data2 + byte_prefix) | ((uintptr_t)data3 + byte_prefix)) &
-       3u) != 0u)
-  {
-    mlk_keccakf1600_xor_bytes_x4_bitinterleaved_c(state, data0, data1, data2,
-                                                  data3, offset, length);
-    return MLK_NATIVE_FUNC_SUCCESS;
-  }
-
-  mlk_keccak_f1600_x4_state_xor_bytes(state, data0, data1, data2, data3, offset,
-                                      length);
+  mlk_keccak_f1600_x4_state_xor_bytes_asm(state, data0, data1, data2, data3,
+                                          offset, length);
   return MLK_NATIVE_FUNC_SUCCESS;
 }
 
 /*
  * Native x4 extract bytes (with on-the-fly bit de-interleaving)
  */
-#define mlk_keccak_f1600_x4_state_extract_bytes \
-  MLK_NAMESPACE(keccak_f1600_x4_state_extract_bytes_asm)
-void mlk_keccak_f1600_x4_state_extract_bytes(void *state, uint8_t *data0,
-                                             uint8_t *data1, uint8_t *data2,
-                                             uint8_t *data3, unsigned offset,
-                                             unsigned length);
-
 MLK_MUST_CHECK_RETURN_VALUE
 static MLK_INLINE int mlk_keccakf1600_extract_bytes_x4_native(
     uint64_t *state, uint8_t *data0, uint8_t *data1, uint8_t *data2,
     uint8_t *data3, unsigned offset, unsigned length)
 {
-  unsigned offset_mod8;
-  unsigned byte_prefix;
-  unsigned main_length;
-
-  offset_mod8 = offset & 7u;
-  byte_prefix = 0;
-  if (offset_mod8 != 0u)
-  {
-    byte_prefix = 8u - offset_mod8;
-    if (byte_prefix > length)
-    {
-      byte_prefix = length;
-    }
-  }
-
-  main_length = length - byte_prefix;
-  if (main_length >= 8u &&
-      ((((uintptr_t)data0 + byte_prefix) | ((uintptr_t)data1 + byte_prefix) |
-        ((uintptr_t)data2 + byte_prefix) | ((uintptr_t)data3 + byte_prefix)) &
-       3u) != 0u)
-  {
-    mlk_keccakf1600_extract_bytes_x4_bitinterleaved_c(
-        state, data0, data1, data2, data3, offset, length);
-    return MLK_NATIVE_FUNC_SUCCESS;
-  }
-
-  mlk_keccak_f1600_x4_state_extract_bytes(state, data0, data1, data2, data3,
-                                          offset, length);
+  mlk_keccak_f1600_x4_state_extract_bytes_asm(state, data0, data1, data2, data3,
+                                              offset, length);
   return MLK_NATIVE_FUNC_SUCCESS;
 }
 
 #endif /* !__ASSEMBLER__ */
 
-#undef MLK_KECCAKF1600_X4_BITINTERLEAVED_ODD_WORDS
-
 #endif /* !MLK_DEV_FIPS202_ARMV81M_MVE_H */

dev/fips202/armv81m/src/state_xor_bytes_x4_mve.S

@@ -51,8 +51,9 @@
 // ---------------------------------------------------------------------------
 //   Prologue -- if offset is not 8-byte aligned, absorb
 //               min(length, 8-(offset%8)) bytes via predicated byte loads.
-//   Main     -- process full 8-byte groups via word-level gather loads,
-//               bit-interleave, then VEOR into even/odd state halves.
+//   Main     -- process full 8-byte groups via vldrw.u32 gather loads when
+//               32-bit aligned, otherwise byte loads; bit-interleave, then
+//               VEOR into even/odd state halves.
 //   Tail     -- absorb remaining <8 bytes via predicated byte loads.
 
 #include "../../../../common.h"
@@ -303,13 +304,13 @@ MLK_ASM_FN_SYMBOL(keccak_f1600_x4_state_xor_bytes_asm)
     lsl mask, mask, off
     vmsr p0, mask
 
-    // Predicated byte loads of nB bytes per instance; post-increment of 4
-    // positions qP for the main loop's [qP, #4]! loads.
+    // Predicated byte loads of nB bytes per instance; post-increment to the
+    // next lane start.
     vpstttt
-    vldrbt.u8 qd0, [dp0], #4
-    vldrbt.u8 qd1, [dp1], #4
-    vldrbt.u8 qd2, [dp2], #4
-    vldrbt.u8 qd3, [dp3], #4
+    vldrbt.u8 qd0, [dp0], #8
+    vldrbt.u8 qd1, [dp1], #8
+    vldrbt.u8 qd2, [dp2], #8
+    vldrbt.u8 qd3, [dp3], #8
 
     // Transpose from per-instance layout to per-lane layout
     transpose_streams_to_lanes
@@ -320,20 +321,22 @@ MLK_ASM_FN_SYMBOL(keccak_f1600_x4_state_xor_bytes_asm)
     // XOR into state (post-increments rSE/rSO by 16)
     xor_lane_and_store_postinc q0, q1, qS0, qS1, rSE, rSO
 
-    // Build qP = {d0,d1,d2,d3} as u32 lanes for gather loads
-    vmov    qP[2], qP[0],  dp0, dp2
-    vmov    qP[3], qP[1],  dp1, dp3
-    b       keccak_f1600_x4_state_xor_bytes_asm_main_body
-
 keccak_f1600_x4_state_xor_bytes_asm_pre_main:
+    // dpN points at the next lane start; vldrw requires 32-bit alignment.
+    orr     tmp, dp0, dp1
+    orr     tmp, tmp, dp2
+    orr     tmp, tmp, dp3
+    tst     tmp, #3
+    bne     keccak_f1600_x4_state_xor_bytes_asm_main_body_byte
+
     vmov    qP[2], qP[0],  dp0, dp2
     vmov    qP[3], qP[1],  dp1, dp3
     mov     tmp, #4
     vsub.u32    qP, qP, tmp        // pre-bias so first [qP,#4]! lands at original ptr
 
 keccak_f1600_x4_state_xor_bytes_asm_main_body:
     // -----------------------------------------------------------------------
-    // MAIN: process full 8-byte lanes
+    // MAIN: process full 8-byte lanes with aligned gather-word loads
     // -----------------------------------------------------------------------
     lsr     lr, length, #3
     wls     lr, lr, keccak_f1600_x4_state_xor_bytes_asm_main_loop_end
@@ -351,9 +354,6 @@ keccak_f1600_x4_state_xor_bytes_asm_main_loop_start:
     le      lr, keccak_f1600_x4_state_xor_bytes_asm_main_loop_start
 keccak_f1600_x4_state_xor_bytes_asm_main_loop_end:
 
-    // -----------------------------------------------------------------------
-    // TAIL: absorb <8 remaining bytes at lane offset 0
-    // -----------------------------------------------------------------------
     ands    length, length, #7
     beq     keccak_f1600_x4_state_xor_bytes_asm_exit
 
@@ -363,12 +363,25 @@ keccak_f1600_x4_state_xor_bytes_asm_main_loop_end:
     vmov    dp0, dp2, qP[2], qP[0]
     vmov    dp1, dp3, qP[3], qP[1]
 
-    vctp.8  length
+keccak_f1600_x4_state_xor_bytes_asm_main_body_byte:
+    // -----------------------------------------------------------------------
+    // MAIN/TAIL: process remaining lanes using predicated byte loads
+    // -----------------------------------------------------------------------
+    cmp     length, #0
+    beq     keccak_f1600_x4_state_xor_bytes_asm_exit
+keccak_f1600_x4_state_xor_bytes_asm_main_loop_byte_start:
+    // Predicate with min(length, 8): a plain MVE byte vector load would
+    // over-read the Keccak lane, and the last iteration may be a partial tail.
+    mov     nB, #8
+    cmp     length, #8
+    it      lo
+    movlo   nB, length
+    vctp.8  nB
     vpstttt
-    vldrbt.u8   qd0, [dp0]
-    vldrbt.u8   qd1, [dp1]
-    vldrbt.u8   qd2, [dp2]
-    vldrbt.u8   qd3, [dp3]
+    vldrbt.u8 qd0, [dp0], #8
+    vldrbt.u8 qd1, [dp1], #8
+    vldrbt.u8 qd2, [dp2], #8
+    vldrbt.u8 qd3, [dp3], #8
 
     // Transpose from per-instance layout to per-lane layout
     transpose_streams_to_lanes
@@ -377,7 +390,10 @@ keccak_f1600_x4_state_xor_bytes_asm_main_loop_end:
     to_bit_interleaving_4x q0, q1, q2, q3, q4
 
     // XOR into state (post-increments rSE/rSO by 16)
-    xor_lane_and_store_postinc qd0, qd1, qS0, qS1, rSE, rSO
+    xor_lane_and_store_postinc q0, q1, qS0, qS1, rSE, rSO
+
+    subs    length, length, nB
+    bne     keccak_f1600_x4_state_xor_bytes_asm_main_loop_byte_start
 
 keccak_f1600_x4_state_xor_bytes_asm_exit:
     vpop    {d8-d15}

mlkem/mlkem_native.c

@@ -504,7 +504,6 @@
 #undef MLK_FIPS202_ARMV81M_NEED_X4
 #undef MLK_FIPS202_NATIVE_ARMV81M
 #undef MLK_FIPS202_NATIVE_ARMV81M_MVE_H
-#undef MLK_KECCAKF1600_X4_BITINTERLEAVED_ODD_WORDS
 #undef MLK_USE_FIPS202_X4_EXTRACT_BYTES_NATIVE
 #undef MLK_USE_FIPS202_X4_NATIVE
 #undef MLK_USE_FIPS202_X4_XOR_BYTES_NATIVE

mlkem/mlkem_native_asm.S

@@ -528,7 +528,6 @@
 #undef MLK_FIPS202_ARMV81M_NEED_X4
 #undef MLK_FIPS202_NATIVE_ARMV81M
 #undef MLK_FIPS202_NATIVE_ARMV81M_MVE_H
-#undef MLK_KECCAKF1600_X4_BITINTERLEAVED_ODD_WORDS
 #undef MLK_USE_FIPS202_X4_EXTRACT_BYTES_NATIVE
 #undef MLK_USE_FIPS202_X4_NATIVE
 #undef MLK_USE_FIPS202_X4_XOR_BYTES_NATIVE