Add inverse NTTs for Kyber & Dilithium

example.py

@@ -601,6 +601,28 @@ def core(self, slothy):
         slothy.optimize_loop("layer123_start")
         slothy.optimize_loop("layer4567_start")
 
+class intt_kyber_123_4567(Example):
+    def __init__(self, var="", arch=AArch64_Neon, target=Target_CortexA55, timeout=None):
+        name = "intt_kyber_123_4567"
+        infile = name
+
+        if var != "":
+            name += f"_{var}"
+            infile += f"_{var}"
+        name += f"_{target_label_dict[target]}"
+
+        super().__init__(infile, name, rename=True, arch=arch, target=target, timeout=timeout)
+
+    def core(self, slothy):
+        slothy.config.sw_pipelining.enabled = True
+        slothy.config.inputs_are_outputs = True
+        slothy.config.sw_pipelining.minimize_overlapping = False
+        slothy.config.variable_size = True
+        slothy.config.reserved_regs = [f"x{i}" for i in range(0, 7)] + ["x30", "sp"]
+        slothy.config.constraints.stalls_first_attempt = 64
+        slothy.optimize_loop("layer4567_start")
+        slothy.optimize_loop("layer123_start")
+
 
 class ntt_kyber_123(Example):
     def __init__(self, var="", arch=AArch64_Neon, target=Target_CortexA55):
@@ -1030,6 +1052,39 @@ def core(self, slothy):
         slothy.optimize_loop("layer45678_start")
 
 
+class intt_dilithium_123_45678(Example):
+    def __init__(self, var="", arch=AArch64_Neon, target=Target_CortexA55, timeout=None):
+        name = f"intt_dilithium_123_45678"
+        infile = name
+
+        if var != "":
+            name += f"_{var}"
+            infile += f"_{var}"
+        name += f"_{target_label_dict[target]}"
+
+        super().__init__(infile, name, rename=True, arch=arch, target=target, timeout=timeout)
+
+    def core(self, slothy):
+        slothy.config.sw_pipelining.enabled = True
+        slothy.config.sw_pipelining.minimize_overlapping = False
+        slothy.config.inputs_are_outputs = True
+
+        slothy.config.reserved_regs = [
+            f"x{i}" for i in range(0, 7)] + ["v8", "x30", "sp"]
+        slothy.config.reserved_regs += self.target_reserved
+        slothy.config.constraints.stalls_first_attempt = 40
+        slothy.optimize_loop("layer45678_start")
+
+        slothy.config.reserved_regs = [
+            f"x{i}" for i in range(0, 7)] + ["v8", "x30", "sp"]
+        slothy.config.reserved_regs += self.target_reserved
+        slothy.config.inputs_are_outputs = True
+        slothy.config.constraints.stalls_first_attempt = 110
+        slothy.optimize_loop("layer123_start")
+
+
+
+
 class ntt_dilithium_123(Example):
     def __init__(self, var="", arch=AArch64_Neon, target=Target_CortexA55):
         name = "ntt_dilithium_123"
@@ -1112,7 +1167,33 @@ def core(self, slothy):
 
         if self.timeout is not None:
             slothy.config.timeout = self.timeout * 12
+        if self.timeout is not None:
+            slothy.config.timeout = self.timeout * 12
+
+        slothy.config.reserved_regs = [
+            f"x{i}" for i in range(0, 6)] + ["x30", "sp"]
+        slothy.config.inputs_are_outputs = True
+        slothy.config.reserved_regs += self.target_reserved
+        slothy.config.sw_pipelining.enabled = True
+        slothy.config.sw_pipelining.minimize_overlapping = False
+        slothy.config.sw_pipelining.halving_heuristic = False
+        slothy.config.split_heuristic = False
+        slothy.optimize_loop("layer5678_start")
+
+
+class intt_dilithium_1234_5678(Example):
+    def __init__(self, var="", arch=AArch64_Neon, target=Target_CortexA72, timeout=None):
+        name = f"intt_dilithium_1234_5678"
+        infile = name
+
+        if var != "":
+            name += f"_{var}"
+            infile += f"_{var}"
+        name += f"_{target_label_dict[target]}"
+
+        super().__init__(infile, name, rename=True, arch=arch, target=target, timeout=timeout)
 
+    def core(self, slothy):
         slothy.config.reserved_regs = [
             f"x{i}" for i in range(0, 6)] + ["x30", "sp"]
         slothy.config.inputs_are_outputs = True
@@ -1123,6 +1204,25 @@ def core(self, slothy):
         slothy.config.split_heuristic = False
         slothy.optimize_loop("layer5678_start")
 
+        slothy.config = Config(self.arch, self.target)
+
+        if self.timeout is not None:
+            slothy.config.timeout = self.timeout // 12
+
+        slothy.config.sw_pipelining.enabled = True
+        slothy.config.sw_pipelining.minimize_overlapping = False
+        slothy.config.reserved_regs = [
+            f"x{i}" for i in range(0, 6)] + ["x30", "sp"]
+        slothy.config.reserved_regs += self.target_reserved
+        slothy.config.inputs_are_outputs = True
+        slothy.config.sw_pipelining.halving_heuristic = True
+        slothy.config.split_heuristic = True
+        slothy.config.split_heuristic_factor = 2
+        slothy.config.split_heuristic_repeat = 4
+        slothy.config.split_heuristic_stepsize = 0.1
+        slothy.config.constraints.stalls_first_attempt = 14
+        slothy.optimize_loop("layer1234_start")
+
 
 class ntt_dilithium_1234(Example):
     def __init__(self, var="", arch=AArch64_Neon, target=Target_CortexA72):
@@ -1296,13 +1396,17 @@ def main():
                  ntt_kyber_123_4567(var="scalar_load_store"),
                  ntt_kyber_123_4567(var="manual_st4"),
                  ntt_kyber_1234_567(),
+                 intt_kyber_123_4567(),
+                 intt_kyber_123_4567(var="manual_ld4"),
                  # Cortex-A72
                  ntt_kyber_123_4567(target=Target_CortexA72),
                  ntt_kyber_123_4567(var="scalar_load", target=Target_CortexA72),
                  ntt_kyber_123_4567(var="scalar_store", target=Target_CortexA72),
                  ntt_kyber_123_4567(var="scalar_load_store", target=Target_CortexA72),
                  ntt_kyber_123_4567(var="manual_st4", target=Target_CortexA72),
                  ntt_kyber_1234_567(target=Target_CortexA72),
+                 intt_kyber_123_4567(target=Target_CortexA72),
+                 intt_kyber_123_4567(var="manual_ld4", target=Target_CortexA72),
                 #  # Apple M1 Firestorm
                  ntt_kyber_123_4567(target=Target_AppleM1_firestorm, timeout=3600),
                  ntt_kyber_123_4567(var="scalar_load", target=Target_AppleM1_firestorm, timeout=3600),
@@ -1311,6 +1415,8 @@ def main():
                  ntt_kyber_123_4567(var="manual_st4", target=Target_AppleM1_firestorm, timeout=3600),
                  ntt_kyber_1234_567(target=Target_AppleM1_firestorm, timeout=300),
                  ntt_kyber_1234_567(var="manual_st4", target=Target_AppleM1_firestorm, timeout=300),
+                 intt_kyber_123_4567(target=Target_AppleM1_firestorm, timeout=3600),
+                 intt_kyber_123_4567(var="manual_ld4", target=Target_AppleM1_firestorm, timeout=3600),
                  # Apple M1 Icestorm
                  ntt_kyber_123_4567(target=Target_AppleM1_icestorm, timeout=3600),
                  ntt_kyber_123_4567(var="scalar_load", target=Target_AppleM1_icestorm, timeout=3600),
@@ -1319,6 +1425,8 @@ def main():
                  ntt_kyber_123_4567(var="manual_st4", target=Target_AppleM1_icestorm, timeout=3600),
                  ntt_kyber_1234_567(target=Target_AppleM1_icestorm, timeout=300),
                  ntt_kyber_1234_567(var="manual_st4", target=Target_AppleM1_icestorm, timeout=300),
+                 intt_kyber_123_4567(target=Target_AppleM1_icestorm, timeout=3600),
+                 intt_kyber_123_4567(var="manual_ld4", target=Target_AppleM1_icestorm, timeout=3600),
                  # Kyber InvNTT
                  # Cortex-M55
                  intt_kyber_1_23_45_67(),
@@ -1340,24 +1448,40 @@ def main():
                  ntt_dilithium_123_45678(var="manual_st4"),
                  ntt_dilithium_1234_5678(),
                  ntt_dilithium_1234_5678(var="manual_st4"),
+                 intt_dilithium_123_45678(),
+                 intt_dilithium_123_45678(var="manual_ld4"),
+                 intt_dilithium_1234_5678(),
+                 intt_dilithium_1234_5678(var="manual_ld4"),
                  # Cortex-A72
                  ntt_dilithium_123_45678(target=Target_CortexA72),
                  ntt_dilithium_123_45678(var="w_scalar", target=Target_CortexA72),
                  ntt_dilithium_123_45678(var="manual_st4", target=Target_CortexA72),
                  ntt_dilithium_1234_5678(target=Target_CortexA72),
                  ntt_dilithium_1234_5678(var="manual_st4", target=Target_CortexA72),
+                 intt_dilithium_123_45678(target=Target_CortexA72),
+                 intt_dilithium_123_45678(var="manual_ld4", target=Target_CortexA72),
+                 intt_dilithium_1234_5678(target=Target_CortexA72),
+                 intt_dilithium_1234_5678(var="manual_ld4", target=Target_CortexA72),
                  # Apple M1 Firestorm
-                 ntt_dilithium_123_45678(target=Target_AppleM1_firestorm, timeout=3600),
+                ntt_dilithium_123_45678(target=Target_AppleM1_firestorm, timeout=3600),
                  ntt_dilithium_123_45678(var="w_scalar", target=Target_AppleM1_firestorm, timeout=3600),
                  ntt_dilithium_123_45678(var="manual_st4", target=Target_AppleM1_firestorm, timeout=3600),
                  ntt_dilithium_1234_5678(target=Target_AppleM1_firestorm, timeout=300),
                  ntt_dilithium_1234_5678(var="manual_st4", target=Target_AppleM1_firestorm, timeout=300),
+                 intt_dilithium_123_45678(target=Target_AppleM1_firestorm, timeout=3600),
+                 intt_dilithium_123_45678(var="manual_ld4", target=Target_AppleM1_firestorm, timeout=3600),
+                 intt_dilithium_1234_5678(target=Target_AppleM1_firestorm, timeout=3600),
+                 intt_dilithium_1234_5678(var="manual_ld4", target=Target_AppleM1_firestorm, timeout=3600),
                  # Apple M1 Icestorm
                  ntt_dilithium_123_45678(target=Target_AppleM1_icestorm, timeout=3600),
                  ntt_dilithium_123_45678(var="w_scalar", target=Target_AppleM1_icestorm, timeout=3600),
                  ntt_dilithium_123_45678(var="manual_st4", target=Target_AppleM1_icestorm, timeout=3600),
                  ntt_dilithium_1234_5678(target=Target_AppleM1_icestorm, timeout=300),
                  ntt_dilithium_1234_5678(var="manual_st4", target=Target_AppleM1_icestorm, timeout=300),
+                 intt_dilithium_123_45678(target=Target_AppleM1_icestorm, timeout=3600),
+                 intt_dilithium_123_45678(var="manual_ld4", target=Target_AppleM1_icestorm, timeout=3600),
+                 intt_dilithium_1234_5678(target=Target_AppleM1_icestorm, timeout=3600),
+                 intt_dilithium_1234_5678(var="manual_ld4", target=Target_AppleM1_icestorm, timeout=3600),
                  # Dilithium invNTT
                  # Cortex-M55
                  intt_dilithium_12_34_56_78(),

examples/misc/gen_roots.py

@@ -405,6 +405,13 @@ def _main():
     ntt_kyber_l123.export("../naive/ntt_kyber_123_45_67_twiddles.s")
     ntt_kyber_l123.export("../opt/ntt_kyber_123_45_67_twiddles.s")
 
+    # For intt_kyber_123_4567.s
+    intt_kyber_l123 = NttRootGen(size=256,modulus=3329,root=17,layers=7,iters=[(0,3),(3,2),(5,2)], 
+                                pad=[0,3], print_label=True, widen_single_twiddles_to_words=False,
+                                inverse=True)
+    intt_kyber_l123.export("../naive/aarch64/intt_kyber_123_45_67_twiddles.s")
+    intt_kyber_l123.export("../opt/aarch64/intt_kyber_123_45_67_twiddles.s")
+
     ntt_kyber = NttRootGen(size=256,modulus=3329,root=17,layers=7)
     ntt_kyber.export("../naive/ntt_kyber_1_23_45_67_twiddles.s")
     ntt_kyber.export("../opt/ntt_kyber_1_23_45_67_twiddles.s")
@@ -428,6 +435,11 @@ def _main():
     ntt_dilithium_l123.export("../naive/ntt_dilithium_123_456_78_twiddles.s")
     ntt_dilithium_l123.export("../opt/ntt_dilithium_123_456_78_twiddles.s")
 
+    intt_dilithium_l123 = NttRootGen(size=256,inverse=True,bitsize=32,modulus=8380417,root=1753,layers=8,
+                                    print_label=True, pad=[0,3], iters=[(0,3),(3,3),(6,2)])
+    intt_dilithium_l123.export("../naive/aarch64/intt_dilithium_123_456_78_twiddles.s")
+    intt_dilithium_l123.export("../opt/aarch64/intt_dilithium_123_456_78_twiddles.s")
+
     ntt_dilithium_l123 = NttRootGen(size=256,bitsize=32,modulus=8380417,root=1753,layers=8,
                                     print_label=True, pad=[0,3], iters=[(0,3),(3,3),(6,2)])
     ntt_dilithium_l123.export("../naive/aarch64/ntt_dilithium_123_456_78_twiddles.s")

examples/naive/aarch64/intt_dilithium_1234_5678.s

@@ -334,10 +334,14 @@ _intt_dilithium_1234_5678:
 
         .p2align 2
 layer5678_start:
-        ldr_vo data0, inp, (16*0)
-        ldr_vo data1, inp, (16*1)
-        ldr_vo data2, inp, (16*2)
-        ldr_vo data3, inp, (16*3)
+        // manual_ld4
+        // ldr_vo data0, inp, (16*0)
+        // ldr_vo data1, inp, (16*1)
+        // ldr_vo data2, inp, (16*2)
+        // ldr_vo data3, inp, (16*3)
+        // transpose4 data
+
+        ld4 {data0.4S, data1.4S, data2.4S, data3.4S}, [inp]
 
         load_next_roots_78 root0, root0_tw, root1, root1_tw, root2, root2_tw, r_ptr0