From 4d6a35ccf5a081d53ec21a76ef418117f0f379b2 Mon Sep 17 00:00:00 2001
From: Mamy Ratsimbazafy <mamy_github@numforge.co>
Date: Fri, 17 Jan 2025 17:41:57 +0100
Subject: [PATCH] llvm: prepare generalizing inline ASM codegen to LLVM IR

---
 constantine/mac/mac_poly1305.nim              | 10 +++----
 constantine/math/arithmetic/finite_fields.nim |  4 +--
 constantine/math/arithmetic/limbs.nim         |  4 +--
 .../math/arithmetic/limbs_crandall.nim        |  4 +--
 constantine/math/arithmetic/limbs_extmul.nim  |  6 ++--
 .../math/arithmetic/limbs_montgomery.nim      | 12 ++++----
 .../assembly/fp2_asm_x86_adx_bmi2.nim         |  2 +-
 constantine/math/extension_fields/towers.nim  |  2 +-
 constantine/math_asm/README.md                | 28 +++++++++++++++++++
 .../limbs_asm_bigint_arm64.nim                |  4 ---
 .../limbs_asm_bigint_x86.nim                  |  0
 .../limbs_asm_crandall_x86.nim                |  0
 .../limbs_asm_crandall_x86_adx_bmi2.nim       |  0
 .../limbs_asm_modular_arm64.nim               |  0
 .../limbs_asm_modular_dbl_prec_x86.nim        |  0
 .../limbs_asm_modular_x86.nim                 |  0
 .../limbs_asm_mul_arm64.nim                   |  0
 .../limbs_asm_mul_mont_arm64.nim              |  4 ---
 .../limbs_asm_mul_mont_x86.nim                |  0
 .../limbs_asm_mul_mont_x86_adx_bmi2.nim       |  0
 .../limbs_asm_mul_x86.nim                     |  0
 .../limbs_asm_mul_x86_adx_bmi2.nim            |  0
 .../limbs_asm_redc_mont_arm64.nim             |  0
 .../limbs_asm_redc_mont_x86.nim               |  0
 .../limbs_asm_redc_mont_x86_adx_bmi2.nim      |  0
 .../platforms/isa_x86/macro_assembler_x86.nim |  2 +-
 .../isa_x86/macro_assembler_x86_att.nim       |  2 +-
 .../isa_x86/macro_assembler_x86_intel.nim     |  2 +-
 28 files changed, 53 insertions(+), 33 deletions(-)
 create mode 100644 constantine/math_asm/README.md
 rename constantine/{math/arithmetic/assembly => math_asm}/limbs_asm_bigint_arm64.nim (98%)
 rename constantine/{math/arithmetic/assembly => math_asm}/limbs_asm_bigint_x86.nim (100%)
 rename constantine/{math/arithmetic/assembly => math_asm}/limbs_asm_crandall_x86.nim (100%)
 rename constantine/{math/arithmetic/assembly => math_asm}/limbs_asm_crandall_x86_adx_bmi2.nim (100%)
 rename constantine/{math/arithmetic/assembly => math_asm}/limbs_asm_modular_arm64.nim (100%)
 rename constantine/{math/arithmetic/assembly => math_asm}/limbs_asm_modular_dbl_prec_x86.nim (100%)
 rename constantine/{math/arithmetic/assembly => math_asm}/limbs_asm_modular_x86.nim (100%)
 rename constantine/{math/arithmetic/assembly => math_asm}/limbs_asm_mul_arm64.nim (100%)
 rename constantine/{math/arithmetic/assembly => math_asm}/limbs_asm_mul_mont_arm64.nim (99%)
 rename constantine/{math/arithmetic/assembly => math_asm}/limbs_asm_mul_mont_x86.nim (100%)
 rename constantine/{math/arithmetic/assembly => math_asm}/limbs_asm_mul_mont_x86_adx_bmi2.nim (100%)
 rename constantine/{math/arithmetic/assembly => math_asm}/limbs_asm_mul_x86.nim (100%)
 rename constantine/{math/arithmetic/assembly => math_asm}/limbs_asm_mul_x86_adx_bmi2.nim (100%)
 rename constantine/{math/arithmetic/assembly => math_asm}/limbs_asm_redc_mont_arm64.nim (100%)
 rename constantine/{math/arithmetic/assembly => math_asm}/limbs_asm_redc_mont_x86.nim (100%)
 rename constantine/{math/arithmetic/assembly => math_asm}/limbs_asm_redc_mont_x86_adx_bmi2.nim (100%)

diff --git a/constantine/mac/mac_poly1305.nim b/constantine/mac/mac_poly1305.nim
index 4839b6c0b..90ca42d8b 100644
--- a/constantine/mac/mac_poly1305.nim
+++ b/constantine/mac/mac_poly1305.nim
@@ -7,13 +7,13 @@
 # at your option. This file may not be copied, modified, or distributed except according to those terms.
 
 import
-  ../platforms/[abstractions, views],
-  ../math/arithmetic/bigints,
-  ../math/arithmetic/[limbs, limbs_extmul],
-  ../math/io/io_bigints
+  constantine/platforms/[abstractions, views],
+  constantine/math/arithmetic/bigints,
+  constantine/math/arithmetic/[limbs, limbs_extmul],
+  constantine/math/io/io_bigints
 
 when UseASM_X86_64:
-  import ../math/arithmetic/assembly/limbs_asm_modular_x86
+  import constantine/math_asm/limbs_asm_modular_x86
 
 # No exceptions allowed
 {.push raises: [].}
diff --git a/constantine/math/arithmetic/finite_fields.nim b/constantine/math/arithmetic/finite_fields.nim
index 206c0389e..902a7bbe3 100644
--- a/constantine/math/arithmetic/finite_fields.nim
+++ b/constantine/math/arithmetic/finite_fields.nim
@@ -35,9 +35,9 @@ import
   ./bigints_crandall
 
 when UseASM_X86_64:
-  import ./assembly/limbs_asm_modular_x86
+  import constantine/math_asm/limbs_asm_modular_x86
 when UseASM_ARM64:
-  import ./assembly/limbs_asm_modular_arm64
+  import constantine/math_asm/limbs_asm_modular_arm64
 
 when nimvm:
   from constantine/named/deriv/precompute import montyResidue_precompute
diff --git a/constantine/math/arithmetic/limbs.nim b/constantine/math/arithmetic/limbs.nim
index 47c48c99e..c8ae33071 100644
--- a/constantine/math/arithmetic/limbs.nim
+++ b/constantine/math/arithmetic/limbs.nim
@@ -9,9 +9,9 @@
 import constantine/platforms/abstractions
 
 when UseASM_X86_32:
-  import ./assembly/limbs_asm_bigint_x86
+  import constantine/math_asm/limbs_asm_bigint_x86
 when UseASM_ARM64:
-  import ./assembly/limbs_asm_bigint_arm64
+  import constantine/math_asm/limbs_asm_bigint_arm64
 
 # ############################################################
 #
diff --git a/constantine/math/arithmetic/limbs_crandall.nim b/constantine/math/arithmetic/limbs_crandall.nim
index 83915d6c6..00238e780 100644
--- a/constantine/math/arithmetic/limbs_crandall.nim
+++ b/constantine/math/arithmetic/limbs_crandall.nim
@@ -12,8 +12,8 @@ import
 
 when UseASM_X86_32:
   import
-    ./assembly/limbs_asm_crandall_x86,
-    ./assembly/limbs_asm_crandall_x86_adx_bmi2
+    constantine/math_asm/limbs_asm_crandall_x86,
+    constantine/math_asm/limbs_asm_crandall_x86_adx_bmi2
 
 # No exceptions allowed
 {.push raises: [], checks: off.}
diff --git a/constantine/math/arithmetic/limbs_extmul.nim b/constantine/math/arithmetic/limbs_extmul.nim
index 6633965f7..4e05fa1d5 100644
--- a/constantine/math/arithmetic/limbs_extmul.nim
+++ b/constantine/math/arithmetic/limbs_extmul.nim
@@ -11,10 +11,10 @@ import
   ./limbs
 
 when UseASM_X86_64:
-  import ./assembly/limbs_asm_mul_x86
-  import ./assembly/limbs_asm_mul_x86_adx_bmi2
+  import constantine/math_asm/limbs_asm_mul_x86
+  import constantine/math_asm/limbs_asm_mul_x86_adx_bmi2
 when UseASM_ARM64:
-  import ./assembly/limbs_asm_mul_arm64
+  import constantine/math_asm/limbs_asm_mul_arm64
 
 # ############################################################
 #
diff --git a/constantine/math/arithmetic/limbs_montgomery.nim b/constantine/math/arithmetic/limbs_montgomery.nim
index d459a88c6..6620aa43e 100644
--- a/constantine/math/arithmetic/limbs_montgomery.nim
+++ b/constantine/math/arithmetic/limbs_montgomery.nim
@@ -12,16 +12,16 @@ import
   ./limbs, ./limbs_extmul
 
 when UseASM_X86_32:
-  import ./assembly/limbs_asm_redc_mont_x86
+  import constantine/math_asm/limbs_asm_redc_mont_x86
 when UseASM_X86_64:
   import
-    ./assembly/limbs_asm_mul_mont_x86,
-    ./assembly/limbs_asm_mul_mont_x86_adx_bmi2,
-    ./assembly/limbs_asm_redc_mont_x86_adx_bmi2
+    constantine/math_asm/limbs_asm_mul_mont_x86,
+    constantine/math_asm/limbs_asm_mul_mont_x86_adx_bmi2,
+    constantine/math_asm/limbs_asm_redc_mont_x86_adx_bmi2
 when UseASM_ARM64:
   import
-    ./assembly/limbs_asm_mul_mont_arm64,
-    ./assembly/limbs_asm_redc_mont_arm64
+    constantine/math_asm/limbs_asm_mul_mont_arm64,
+    constantine/math_asm/limbs_asm_redc_mont_arm64
 
 # ############################################################
 #
diff --git a/constantine/math/extension_fields/assembly/fp2_asm_x86_adx_bmi2.nim b/constantine/math/extension_fields/assembly/fp2_asm_x86_adx_bmi2.nim
index ac25e67e4..d9135f831 100644
--- a/constantine/math/extension_fields/assembly/fp2_asm_x86_adx_bmi2.nim
+++ b/constantine/math/extension_fields/assembly/fp2_asm_x86_adx_bmi2.nim
@@ -11,7 +11,7 @@ import
   constantine/platforms/abstractions,
   constantine/named/algebras,
   constantine/math/arithmetic,
-  constantine/math/arithmetic/assembly/[
+  constantine/math_asm/[
     limbs_asm_mul_x86_adx_bmi2,
     limbs_asm_mul_mont_x86_adx_bmi2,
     limbs_asm_redc_mont_x86_adx_bmi2
diff --git a/constantine/math/extension_fields/towers.nim b/constantine/math/extension_fields/towers.nim
index 7113094fe..3b8aa22fa 100644
--- a/constantine/math/extension_fields/towers.nim
+++ b/constantine/math/extension_fields/towers.nim
@@ -15,7 +15,7 @@ export Fp
 
 when UseASM_X86_64:
   import
-    ./assembly/fp2_asm_x86_adx_bmi2
+    constantine/math_asm/fp2_asm_x86_adx_bmi2
 
 # Note: to avoid burdening the Nim compiler, we rely on generic extension
 # to complain if the base field procedures don't exist
diff --git a/constantine/math_asm/README.md b/constantine/math_asm/README.md
new file mode 100644
index 000000000..79d39b3d0
--- /dev/null
+++ b/constantine/math_asm/README.md
@@ -0,0 +1,28 @@
+# Assembly code generator for mathematical primitives
+
+This folder holds code generators for inline assembly in Nim and LLVM IR.
+
+Inline assembly is necessary for security, ensure constant-time from a high-level language, and performance as certain instructions cannot be emitted by a compiler (ADOX/ADCX) despite offering a large performance advantage (up to 70% for ADOX/ADCX).
+
+Even when using LLVM IR and in the case where all instructions can be emitted (ARM64),
+and the number of compute instructions between inline assembly and LLVM IR is the same
+stack usage might be significantly worse due to bad register allocation and regular stack spill.
+
+For example on ARM64, with LLVM IR that mirrors inline assembly we get the following
+breakdown on 6 limbs (CodeGenLevelDefault):
+-  inline ASM     vs pure LLVM IR
+-  64 bytes stack vs      368
+-   4 stp         vs       23
+-  10 ldp         vs       35
+-   6 ldr         vs       61
+-   6 str         vs       43
+-   6 mov         vs       24
+-  78 mul         vs       78
+-  72 umulh       vs       72
+-  17 adds        vs       17
+- 103 adcs        vs      103
+-  23 adc         vs       12   -> the ADC have become cset to save the carry/borrow flag in register
+-   6 cmn         vs        6
+-   0 cset        vs       11
+
+And generating single instruction in LLVM inline assembly doesn't solve register spilling to the stack.
\ No newline at end of file
diff --git a/constantine/math/arithmetic/assembly/limbs_asm_bigint_arm64.nim b/constantine/math_asm/limbs_asm_bigint_arm64.nim
similarity index 98%
rename from constantine/math/arithmetic/assembly/limbs_asm_bigint_arm64.nim
rename to constantine/math_asm/limbs_asm_bigint_arm64.nim
index 8f63943c7..eedc85a5c 100644
--- a/constantine/math/arithmetic/assembly/limbs_asm_bigint_arm64.nim
+++ b/constantine/math_asm/limbs_asm_bigint_arm64.nim
@@ -78,10 +78,6 @@ macro ccopy_gen[N: static int](a_PIR: var Limbs[N], b_PIR: Limbs[N], ctl: Secret
   # Codegen
   result.add ctx.generate()
 
-  debugEcho "======Transfo====="
-  debugEcho getImplTransformed(result).repr()
-  debugEcho "======"
-
 func ccopy_asm*(a: var Limbs, b: Limbs, ctl: SecretBool) =
   ## Constant-time conditional copy
   ## If ctl is true: b is copied into a
diff --git a/constantine/math/arithmetic/assembly/limbs_asm_bigint_x86.nim b/constantine/math_asm/limbs_asm_bigint_x86.nim
similarity index 100%
rename from constantine/math/arithmetic/assembly/limbs_asm_bigint_x86.nim
rename to constantine/math_asm/limbs_asm_bigint_x86.nim
diff --git a/constantine/math/arithmetic/assembly/limbs_asm_crandall_x86.nim b/constantine/math_asm/limbs_asm_crandall_x86.nim
similarity index 100%
rename from constantine/math/arithmetic/assembly/limbs_asm_crandall_x86.nim
rename to constantine/math_asm/limbs_asm_crandall_x86.nim
diff --git a/constantine/math/arithmetic/assembly/limbs_asm_crandall_x86_adx_bmi2.nim b/constantine/math_asm/limbs_asm_crandall_x86_adx_bmi2.nim
similarity index 100%
rename from constantine/math/arithmetic/assembly/limbs_asm_crandall_x86_adx_bmi2.nim
rename to constantine/math_asm/limbs_asm_crandall_x86_adx_bmi2.nim
diff --git a/constantine/math/arithmetic/assembly/limbs_asm_modular_arm64.nim b/constantine/math_asm/limbs_asm_modular_arm64.nim
similarity index 100%
rename from constantine/math/arithmetic/assembly/limbs_asm_modular_arm64.nim
rename to constantine/math_asm/limbs_asm_modular_arm64.nim
diff --git a/constantine/math/arithmetic/assembly/limbs_asm_modular_dbl_prec_x86.nim b/constantine/math_asm/limbs_asm_modular_dbl_prec_x86.nim
similarity index 100%
rename from constantine/math/arithmetic/assembly/limbs_asm_modular_dbl_prec_x86.nim
rename to constantine/math_asm/limbs_asm_modular_dbl_prec_x86.nim
diff --git a/constantine/math/arithmetic/assembly/limbs_asm_modular_x86.nim b/constantine/math_asm/limbs_asm_modular_x86.nim
similarity index 100%
rename from constantine/math/arithmetic/assembly/limbs_asm_modular_x86.nim
rename to constantine/math_asm/limbs_asm_modular_x86.nim
diff --git a/constantine/math/arithmetic/assembly/limbs_asm_mul_arm64.nim b/constantine/math_asm/limbs_asm_mul_arm64.nim
similarity index 100%
rename from constantine/math/arithmetic/assembly/limbs_asm_mul_arm64.nim
rename to constantine/math_asm/limbs_asm_mul_arm64.nim
diff --git a/constantine/math/arithmetic/assembly/limbs_asm_mul_mont_arm64.nim b/constantine/math_asm/limbs_asm_mul_mont_arm64.nim
similarity index 99%
rename from constantine/math/arithmetic/assembly/limbs_asm_mul_mont_arm64.nim
rename to constantine/math_asm/limbs_asm_mul_mont_arm64.nim
index 7731b08cb..72dd2196a 100644
--- a/constantine/math/arithmetic/assembly/limbs_asm_mul_mont_arm64.nim
+++ b/constantine/math_asm/limbs_asm_mul_mont_arm64.nim
@@ -236,9 +236,6 @@ macro mulMont_CIOS_sparebit_gen[N: static int](
       ctx.str t[i], r[i]
 
   result.add ctx.generate()
-  debugEcho "======Transfo====="
-  debugEcho getImplTransformed(result).repr()
-  debugEcho "======"
 
 func mulMont_CIOS_sparebit_asm*(r: var Limbs, a, b, M: Limbs, m0ninv: BaseType, lazyReduce: static bool = false) =
   ## Constant-time Montgomery multiplication
@@ -412,7 +409,6 @@ macro sumprodMont_CIOS_spare2bits_gen[N, K: static int](
       ctx.mulhiadd_cio(t[j], m, M[j], t[j])
     ctx.mulhiadd_ci(t[N-1], m, M[N-1], t[N-1])
 
-
   if lazyReduce:
     for i in 0 ..< N:
       ctx.str t[i], r[i]
diff --git a/constantine/math/arithmetic/assembly/limbs_asm_mul_mont_x86.nim b/constantine/math_asm/limbs_asm_mul_mont_x86.nim
similarity index 100%
rename from constantine/math/arithmetic/assembly/limbs_asm_mul_mont_x86.nim
rename to constantine/math_asm/limbs_asm_mul_mont_x86.nim
diff --git a/constantine/math/arithmetic/assembly/limbs_asm_mul_mont_x86_adx_bmi2.nim b/constantine/math_asm/limbs_asm_mul_mont_x86_adx_bmi2.nim
similarity index 100%
rename from constantine/math/arithmetic/assembly/limbs_asm_mul_mont_x86_adx_bmi2.nim
rename to constantine/math_asm/limbs_asm_mul_mont_x86_adx_bmi2.nim
diff --git a/constantine/math/arithmetic/assembly/limbs_asm_mul_x86.nim b/constantine/math_asm/limbs_asm_mul_x86.nim
similarity index 100%
rename from constantine/math/arithmetic/assembly/limbs_asm_mul_x86.nim
rename to constantine/math_asm/limbs_asm_mul_x86.nim
diff --git a/constantine/math/arithmetic/assembly/limbs_asm_mul_x86_adx_bmi2.nim b/constantine/math_asm/limbs_asm_mul_x86_adx_bmi2.nim
similarity index 100%
rename from constantine/math/arithmetic/assembly/limbs_asm_mul_x86_adx_bmi2.nim
rename to constantine/math_asm/limbs_asm_mul_x86_adx_bmi2.nim
diff --git a/constantine/math/arithmetic/assembly/limbs_asm_redc_mont_arm64.nim b/constantine/math_asm/limbs_asm_redc_mont_arm64.nim
similarity index 100%
rename from constantine/math/arithmetic/assembly/limbs_asm_redc_mont_arm64.nim
rename to constantine/math_asm/limbs_asm_redc_mont_arm64.nim
diff --git a/constantine/math/arithmetic/assembly/limbs_asm_redc_mont_x86.nim b/constantine/math_asm/limbs_asm_redc_mont_x86.nim
similarity index 100%
rename from constantine/math/arithmetic/assembly/limbs_asm_redc_mont_x86.nim
rename to constantine/math_asm/limbs_asm_redc_mont_x86.nim
diff --git a/constantine/math/arithmetic/assembly/limbs_asm_redc_mont_x86_adx_bmi2.nim b/constantine/math_asm/limbs_asm_redc_mont_x86_adx_bmi2.nim
similarity index 100%
rename from constantine/math/arithmetic/assembly/limbs_asm_redc_mont_x86_adx_bmi2.nim
rename to constantine/math_asm/limbs_asm_redc_mont_x86_adx_bmi2.nim
diff --git a/constantine/platforms/isa_x86/macro_assembler_x86.nim b/constantine/platforms/isa_x86/macro_assembler_x86.nim
index 57fa6dcd3..7a617fd91 100644
--- a/constantine/platforms/isa_x86/macro_assembler_x86.nim
+++ b/constantine/platforms/isa_x86/macro_assembler_x86.nim
@@ -6,7 +6,7 @@
 #   * Apache v2 license (license terms in the root directory or at http://www.apache.org/licenses/LICENSE-2.0).
 # at your option. This file may not be copied, modified, or distributed except according to those terms.
 
-import ../config
+import constantine/platforms/config
 
 when UseAsmSyntaxIntel:
   # We need Intel syntax.
diff --git a/constantine/platforms/isa_x86/macro_assembler_x86_att.nim b/constantine/platforms/isa_x86/macro_assembler_x86_att.nim
index e95f3eafa..c4df585bb 100644
--- a/constantine/platforms/isa_x86/macro_assembler_x86_att.nim
+++ b/constantine/platforms/isa_x86/macro_assembler_x86_att.nim
@@ -8,7 +8,7 @@
 
 import
   std/[macros, strutils, sets, hashes, algorithm, sequtils, enumutils],
-  ../[config, bithacks]
+  constantine/platforms/[config, bithacks]
 
 # A compile-time inline assembler
 
diff --git a/constantine/platforms/isa_x86/macro_assembler_x86_intel.nim b/constantine/platforms/isa_x86/macro_assembler_x86_intel.nim
index e9bb3c443..a87129e47 100644
--- a/constantine/platforms/isa_x86/macro_assembler_x86_intel.nim
+++ b/constantine/platforms/isa_x86/macro_assembler_x86_intel.nim
@@ -8,7 +8,7 @@
 
 import
   std/[macros, strutils, sets, hashes, algorithm, sequtils, enumutils],
-  ../[config, bithacks]
+  constantine/platforms/[config, bithacks]
 
 # A compile-time inline assembler