Namespace rej_uniform_asm

- There are some define in this files, namesapce them by MLK

Signed-off-by: willieyz <willie.zhao@chelpis.com>

Author: willieyz

dev/x86_64/src/rej_uniform_asm.S

@@ -24,38 +24,38 @@
     !defined(MLK_CONFIG_MULTILEVEL_NO_SHARED)
 /* simpasm: header-end */
 
-#define in %rsi
-#define out %rdi
-#define len %rdx
-#define tab %rcx
-
-#define cnt %rax
-#define ecnt %eax
-#define pos %r8
-
-#define good %r11
-#define pext_mask %r9
-#define table_idx %r10
-
-#define bound %xmm0
-#define temp0 %xmm1
-#define temp1 %xmm3
-#define vals %xmm2
-#define shuffle_out_mask %xmm3
-#define shuffle_in_mask %xmm4
-#define and_mask %xmm5
+#define MLK_IN %rsi
+#define MLK_OUT %rdi
+#define MLK_LEN %rdx
+#define MLK_TAB %rcx
+
+#define MLK_CNT %rax
+#define MLK_ECNT %eax
+#define MLK_POS %r8
+
+#define MLK_GOOD %r11
+#define MLK_PEXT_MASK %r9
+#define MLK_TABLE_IDX %r10
+
+#define MLK_BOUND %xmm0
+#define MLK_TEMP0 %xmm1
+#define MLK_TEMP1 %xmm3
+#define MLK_VALS %xmm2
+#define MLK_SHUFFLE_OUT_MASK %xmm3
+#define MLK_SHUFFLE_IN_MASK %xmm4
+#define MLK_AND_MASK %xmm5
 
 //  High level overview of the algorithm:
 //  For every 96 bits (12 bytes) of the input:
 //    1. Split 96 bits into eight 12-bit integers where each integer
-//       occupies a corresponding 16-bit element of `vals` xmm register,
-//    2. Compute an 8-bit value `good` such that
-//         good[i] = vals[i] < MLKEM_Q ? 1 : 0, for i in [0, 7],
-//    3. Shuffle the elements in `vals` such that all good elements
+//       occupies a corresponding 16-bit element of `MLK_VALS` xmm register,
+//    2. Compute an 8-bit value `MLK_GOOD` such that
+//         MLK_GOOD[i] = MLK_VALS[i] < MLKEM_Q ? 1 : 0, for i in [0, 7],
+//    3. Shuffle the elements in `MLK_VALS` such that all MLK_GOOD elements
 //       are ordered consecutivelly, and store them.
 //
 //  Notes:
-//    - We exit early if we find the required number of good values,
+//    - We exit early if we find the required number of MLK_GOOD values,
 //    - We use the stack as a temporary storage and copy to the actual
 //      output buffer only in the end. This is because the algorithm
 //      can overwrite up to 14 bytes (we use 16B for alignment), 
@@ -74,77 +74,82 @@ MLK_ASM_FN_SYMBOL(rej_uniform_asm)
   testq len, len
   jz rej_uniform_asm_end
 
+  // Return if input length is 0
+  xorl ecnt, ecnt
+  testq len, len
+  jz rej_uniform_asm_end
+
   // Broadcast MLKEM_Q (3329) to all 16-bit elements of bound reg.
   movq $0x0D010D010D010D01, %rax
-  movq %rax, bound
-  pinsrq $1, %rax, bound
+  movq %rax, MLK_BOUND
+  pinsrq $1, %rax, MLK_BOUND
 
-  // Broadcast 12-bit mask 0xFFF to all 16-bit elements of bound reg.
+  // Broadcast 12-bit mask 0xFFF to all 16-bit elements of MLK_BOUND reg.
   movq $0x0FFF0FFF0FFF0FFF, %rax
-  movq %rax, and_mask
-  pinsrq $1, %rax, and_mask
+  movq %rax, MLK_AND_MASK
+  pinsrq $1, %rax, MLK_AND_MASK
 
   // Load shuffle mask:
   //   0, 1, 1, 2, 3, 4, 4, 5, 6, 7, 7, 8, 9, 10, 10, 11.
   movq $0x0504040302010100, %rax
-  movq %rax, shuffle_in_mask
+  movq %rax, MLK_SHUFFLE_IN_MASK
   movq $0x0B0A0A0908070706, %rax
-  pinsrq $1, %rax, shuffle_in_mask
+  pinsrq $1, %rax, MLK_SHUFFLE_IN_MASK
 
-  movq $0, cnt // cnt counts the number of good values we've found.
-  movq $0, pos // pos is the current position in the input buffer.
-  movq $0x5555, pext_mask // 0x5555 mask to extract every second bit.
+  movq $0, MLK_CNT // MLK_CNT counts the number of MLK_GOOD values we've found.
+  movq $0, MLK_POS // MLK_POS is the current position in the input buffer.
+  movq $0x5555, MLK_PEXT_MASK // 0x5555 mask to extract every second bit.
 
 rej_uniform_asm_loop_start:
   // 1. Split 96 bits into eight 12-bit integers where each integer.
-  //    We explain the algorithm by considering the lowest 64 bits of vals.
-  movdqu  (in, pos), vals
-  //    vals: [ 63..48 | 47..32 | 31..16 | 15..0 ]
-  pshufb  shuffle_in_mask, vals
-  //    vals: [ 47..32 | 39..24 | 23..8 | 15..0 ]
-  movdqa  vals, temp1
+  //    We explain the algorithm by considering the lowest 64 bits of MLK_VALS.
+  movdqu  (MLK_IN, MLK_POS), MLK_VALS
+  //    MLK_VALS: [ 63..48 | 47..32 | 31..16 | 15..0 ]
+  pshufb  MLK_SHUFFLE_IN_MASK, MLK_VALS
+  //    MLK_VALS: [ 47..32 | 39..24 | 23..8 | 15..0 ]
+  movdqa  MLK_VALS, MLK_TEMP1
   //    temp: [ 47..32 | 39..24 | 23..8 | 15..0 ]
-  psrlw   $4, temp1
+  psrlw   $4, MLK_TEMP1
   //    temp: [ 47..36 | 39..28 | 23..12 | 15..4 ]
-  pblendw $0xAA, temp1, vals
-  //    vals: [ 47..36 | 39..24 | 23..12 | 15..0]
-  pand    and_mask, vals
-  //    vals: [ 47..36 | 35..24 | 23..12 | 12..0]
-
-  // 2. Compute an 8-bit value `good` such that
-  //      good[i] = vals[i] < MLKEM_Q ? 1 : 0, for i in [0, 7],
-  movdqa   bound, temp0
-  pcmpgtw  vals, temp0
-  pmovmskb temp0, good
-  pext     pext_mask, good, good
-
-  // 3. Shuffle the elements in `vals` such that all good elements
+  pblendw $0xAA, MLK_TEMP1, MLK_VALS
+  //    MLK_VALS: [ 47..36 | 39..24 | 23..12 | 15..0]
+  pand    MLK_AND_MASK, MLK_VALS
+  //    MLK_VALS: [ 47..36 | 35..24 | 23..12 | 12..0]
+
+  // 2. Compute an 8-bit value `MLK_GOOD` such that
+  //      MLK_GOOD[i] = MLK_VALS[i] < MLKEM_Q ? 1 : 0, for i in [0, 7],
+  movdqa   MLK_BOUND, MLK_TEMP0
+  pcmpgtw  MLK_VALS, MLK_TEMP0
+  pmovmskb MLK_TEMP0, MLK_GOOD
+  pext     MLK_PEXT_MASK, MLK_GOOD, MLK_GOOD
+
+  // 3. Shuffle the elements in `MLK_VALS` such that all MLK_GOOD elements
   //    are ordered consecutivelly, and store them.
-  movq   good, table_idx
-  shl    $4, table_idx
-  movdqu (tab, table_idx), shuffle_out_mask
-  pshufb shuffle_out_mask, vals
-  movdqu vals, (%rsp, cnt, 2)
+  movq   MLK_GOOD, MLK_TABLE_IDX
+  shl    $4, MLK_TABLE_IDX
+  movdqu (MLK_TAB, MLK_TABLE_IDX), MLK_SHUFFLE_OUT_MASK
+  pshufb MLK_SHUFFLE_OUT_MASK, MLK_VALS
+  movdqu MLK_VALS, (%rsp, MLK_CNT, 2)
 
   // Update the counter and check if we are done.
-  popcnt good, good
-  addq good, cnt
+  popcnt MLK_GOOD, MLK_GOOD
+  addq MLK_GOOD, MLK_CNT
 
-  cmpq $256, cnt
+  cmpq $256, MLK_CNT
   jnb rej_uniform_asm_final_copy
 
-  addq $12, pos
-  cmpq pos, len
+  addq $12, MLK_POS
+  cmpq MLK_POS, MLK_LEN
   ja rej_uniform_asm_loop_start
 
 rej_uniform_asm_final_copy:
-  // Copy up to 256 values to the output: min(cnt, 256).
+  // Copy up to 256 values to the output: min(MLK_CNT, 256).
   mov   $256, %rcx
-  cmp   $256, cnt
-  cmova %rcx, cnt
+  cmp   $256, MLK_CNT
+  cmova %rcx, MLK_CNT
 
   movq %rsp, %rsi
-  movq cnt, %rcx
+  movq MLK_CNT, %rcx
   shlq $1, %rcx
   rep  movsb
 
@@ -155,23 +160,23 @@ rej_uniform_asm_end:
 
 /* To facilitate single-compilation-unit (SCU) builds, undefine all macros.
  * Don't modify by hand -- this is auto-generated by scripts/autogen. */
-#undef in
-#undef out
-#undef len
-#undef tab
-#undef cnt
-#undef ecnt
-#undef pos
-#undef good
-#undef pext_mask
-#undef table_idx
-#undef bound
-#undef temp0
-#undef temp1
-#undef vals
-#undef shuffle_out_mask
-#undef shuffle_in_mask
-#undef and_mask
+#undef MLK_IN
+#undef MLK_OUT
+#undef MLK_LEN
+#undef MLK_TAB
+#undef MLK_CNT
+#undef MLK_ECNT
+#undef MLK_POS
+#undef MLK_GOOD
+#undef MLK_PEXT_MASK
+#undef MLK_TABLE_IDX
+#undef MLK_BOUND
+#undef MLK_TEMP0
+#undef MLK_TEMP1
+#undef MLK_VALS
+#undef MLK_SHUFFLE_OUT_MASK
+#undef MLK_SHUFFLE_IN_MASK
+#undef MLK_AND_MASK
 #undef MLK_STACK_SIZE
 
 /* simpasm: footer-start */