From fac21e232706daf3944d2d0b0c190972d9e8e3d8 Mon Sep 17 00:00:00 2001
From: Amin Abdulrahman <amin@abdulrahman.de>
Date: Tue, 15 Oct 2024 15:15:20 +0200
Subject: [PATCH] Add Loop parser + examples for armv7m

---
 example.py                         |  35 ++++++++
 examples/naive/armv7m/loop_cmp.s   |  10 +++
 examples/naive/armv7m/loop_subs.s  |   4 +
 slothy/targets/arm_v7m/arch_v7m.py | 129 ++++++++++-------------------
 4 files changed, 91 insertions(+), 87 deletions(-)
 create mode 100644 examples/naive/armv7m/loop_cmp.s
 create mode 100644 examples/naive/armv7m/loop_subs.s

diff --git a/example.py b/example.py
index d7dbd347..1240e9e3 100644
--- a/example.py
+++ b/example.py
@@ -293,6 +293,39 @@ def __init__(self, var="", arch=AArch64_Neon, target=Target_CortexA55):
     def core(self,slothy):
         slothy.config.variable_size=True
         slothy.optimize_loop("start")
+        
+class Armv7mLoopSubs(Example):
+    def __init__(self, var="", arch=Arch_Armv7M, target=Target_CortexM7):
+        name = "loop_subs"
+        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)
+
+    def core(self,slothy):
+        slothy.config.variable_size=True
+        slothy.optimize_loop("start")
+
+class Armv7mLoopCmp(Example):
+    def __init__(self, var="", arch=Arch_Armv7M, target=Target_CortexM7):
+        name = "loop_cmp"
+        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)
+
+    def core(self,slothy):
+        slothy.config.variable_size=True
+        slothy.config.outputs = ["r6"]
+        slothy.optimize_loop("start")
 
 class CRT(Example):
     def __init__(self):
@@ -2106,6 +2139,8 @@ def main():
                 # Loop examples
                  AArch64LoopSubs(),
                  LoopLe(),
+                 Armv7mLoopSubs(),
+                 Armv7mLoopCmp(),
 
                  CRT(),
 
diff --git a/examples/naive/armv7m/loop_cmp.s b/examples/naive/armv7m/loop_cmp.s
new file mode 100644
index 00000000..be8d09a7
--- /dev/null
+++ b/examples/naive/armv7m/loop_cmp.s
@@ -0,0 +1,10 @@
+/* For example, r5 represents an address where we will stop iterating and r6 is
+the actual pointer which is incremented inside the loop. */
+
+mov.w r6, #0
+add.w r5, r6, #64
+
+start:
+    add r6, r6, #4
+    cmp.w r6, r5
+    bne.w start
\ No newline at end of file
diff --git a/examples/naive/armv7m/loop_subs.s b/examples/naive/armv7m/loop_subs.s
new file mode 100644
index 00000000..5dd05c9c
--- /dev/null
+++ b/examples/naive/armv7m/loop_subs.s
@@ -0,0 +1,4 @@
+movw r5, #16
+start:
+    subs.w r5, #1
+    bne.w start
\ No newline at end of file
diff --git a/slothy/targets/arm_v7m/arch_v7m.py b/slothy/targets/arm_v7m/arch_v7m.py
index 1dfc4ae6..1abc41b1 100644
--- a/slothy/targets/arm_v7m/arch_v7m.py
+++ b/slothy/targets/arm_v7m/arch_v7m.py
@@ -5,7 +5,7 @@
 from enum import Enum
 from functools import cache
 
-from slothy.helper import SourceLine
+from slothy.helper import SourceLine, Loop
 from sympy import simplify
 
 llvm_mca_arch = "arm"  # TODO
@@ -126,16 +126,13 @@ def unconditional(lbl):
         yield f"b {lbl}"
 
 
-class Loop:
-    """Helper functions for parsing and writing simple loops in armv7m
-
-    TODO: Generalize; current implementation too specific about shape of loop"""
-
-    def __init__(self, lbl_start="1", lbl_end="2", loop_init="lr"):
-        self.lbl_start = lbl_start
-        self.lbl_end   = lbl_end
-        self.loop_init = loop_init
-
+class CmpLoop(Loop):
+    def __init__(self, lbl="lbl", lbl_start="1", lbl_end="2", loop_init="lr") -> None:
+        super().__init__(lbl_start=lbl_start, lbl_end=lbl_end, loop_init=loop_init)
+        self.lbl_regex = r"^\s*(?P<label>\w+)\s*:(?P<remainder>.*)$"
+        self.end_regex = (r"^\s*cmp(?:\.w)?\s+(?P<cnt>\w+),\s*(?P<reg1>\w+)",
+                               rf"^\s*(cbnz|cbz|bne)(?:\.w)?\s+{lbl}")
+    
     def start(self, loop_cnt, indentation=0, fixup=0, unroll=1, jump_if_empty=None):
         """Emit starting instruction(s) and jump label for loop"""
         indent = ' ' * indentation
@@ -147,89 +144,47 @@ def start(self, loop_cnt, indentation=0, fixup=0, unroll=1, jump_if_empty=None):
         if jump_if_empty is not None:
             yield f"cbz {loop_cnt}, {jump_if_empty}"
         yield f"{self.lbl_start}:"
-
+    
     def end(self, other, indentation=0):
         """Emit compare-and-branch at the end of the loop"""
-        (reg0, reg1, imm) = other
         indent = ' ' * indentation
         lbl_start = self.lbl_start
         if lbl_start.isdigit():
             lbl_start += "b"
+        
+        yield f'{indent}cmp {other["cnt"]}, {other["reg1"]}'
+        yield f'{indent}bne {lbl_start}'
+        
+class SubsLoop(Loop):
+    def __init__(self, lbl_start="1", lbl_end="2", loop_init="lr") -> None:
+        super().__init__(lbl_start=lbl_start, lbl_end=lbl_end, loop_init=loop_init)
+        self.lbl_regex = r"^\s*(?P<label>\w+)\s*:(?P<remainder>.*)$"
+        self.end_regex = (r"^\s*sub[s]?(?:\.w)?\s+(?P<cnt>\w+),(?:\s*(?P<reg1>\w+),)?\s*(?P<imm>#1)",
+                               rf"^\s*(cbnz|cbz|bne)(?:\.w)?\s+{lbl_start}")
 
-        yield f"{indent}subs {reg0}, {reg1}, {imm}"  # `subs` sets flags
-        yield f"{indent}cbnz {reg0}, {lbl_start}"
-
-    @staticmethod
-    def extract(source, lbl):
-        """Locate a loop with start label `lbl` in `source`.
-
-        We currently only support the following loop forms:
-
-           ```
-           loop_lbl:
-               {code}
-               sub[s] <cnt>, <cnt>, #1
-               (cbnz|cbz|bne) <cnt>, loop_lbl
-           ```
-
-        """
-        assert isinstance(source, list)
-
-        pre  = []
-        body = []
-        post = []
-        loop_lbl_regexp_txt = r"^\s*(?P<label>\w+)\s*:(?P<remainder>.*)$"
-        loop_lbl_regexp = re.compile(loop_lbl_regexp_txt)
-
-        # TODO: Allow other forms of looping
-
-        loop_end_regexp_txt = (r"^\s*sub[s]?\s+(?P<reg0>\w+),\s*(?P<reg1>\w+),\s*(?P<imm>#1)",
-                               rf"^\s*(cbnz|cbz|bne)\s+(?P<reg0>\w+),\s*{lbl}")
-        loop_end_regexp = [re.compile(txt) for txt in loop_end_regexp_txt]
-        lines = iter(source)
-        l = None
-        keep = False
-        state = 0 # 0: haven't found loop yet, 1: extracting loop, 2: after loop
-        while True:
-            if not keep:
-                l = next(lines, None)
-            keep = False
-            if l is None:
-                break
-            l_str = l.text
-            assert isinstance(l, str) is False
-            if state == 0:
-                p = loop_lbl_regexp.match(l_str)
-                if p is not None and p.group("label") == lbl:
-                    l = l.copy().set_text(p.group("remainder"))
-                    keep = True
-                    state = 1
-                else:
-                    pre.append(l)
-                continue
-            if state == 1:
-                p = loop_end_regexp[0].match(l_str)
-                if p is not None:
-                    reg0 = p.group("reg0")
-                    reg1 = p.group("reg1")
-                    imm = p.group("imm")
-                    state = 2
-                    continue
-                body.append(l)
-                continue
-            if state == 2:
-                p = loop_end_regexp[1].match(l_str)
-                if p is not None:
-                    state = 3
-                    continue
-                body.append(l)
-                continue
-            if state == 3:
-                post.append(l)
-                continue
-        if state < 3:
-            raise FatalParsingException(f"Couldn't identify loop {lbl}")
-        return pre, body, post, lbl, (reg0, reg1, imm)
+    def start(self, loop_cnt, indentation=0, fixup=0, unroll=1, jump_if_empty=None):
+        """Emit starting instruction(s) and jump label for loop"""
+        indent = ' ' * indentation
+        if unroll > 1:
+            assert unroll in [1,2,4,8,16,32]
+            yield f"{indent}lsr {loop_cnt}, {loop_cnt}, #{int(math.log2(unroll))}"
+        if fixup != 0:
+            yield f"{indent}sub {loop_cnt}, {loop_cnt}, #{fixup}"
+        if jump_if_empty is not None:
+            yield f"cbz {loop_cnt}, {jump_if_empty}"
+        yield f"{self.lbl_start}:"
+    
+    def end(self, other, indentation=0):
+        """Emit compare-and-branch at the end of the loop"""
+        indent = ' ' * indentation
+        lbl_start = self.lbl_start
+        if lbl_start.isdigit():
+            lbl_start += "b"
+        if other["reg1"] is None:
+            yield f'{indent}subs {other["cnt"]}, #1'
+        else:
+            yield f'{indent}subs {other["cnt"]}, {other["reg1"]}, {other["imm"]}'  # `subs` sets flags
+        yield f'{indent}bne {lbl_start}'
     
 class FatalParsingException(Exception):
     """A fatal error happened during instruction parsing"""