From b1a4ab5fb0d8cdfa7cfd6d9d07c8f8e6004afd11 Mon Sep 17 00:00:00 2001
From: Amin Abdulrahman <amin@abdulrahman.de>
Date: Fri, 11 Oct 2024 16:01:21 +0200
Subject: [PATCH 1/8] More general structure for loop parsing

---
 slothy/core/slothy.py                  |  10 +--
 slothy/targets/aarch64/aarch64_neon.py | 118 ++++++++++++++-----------
 slothy/targets/arm_v81m/arch_v81m.py   | 108 +++++++++++++++-------
 3 files changed, 146 insertions(+), 90 deletions(-)

diff --git a/slothy/core/slothy.py b/slothy/core/slothy.py
index 3fb67a41..bcff4108 100644
--- a/slothy/core/slothy.py
+++ b/slothy/core/slothy.py
@@ -367,12 +367,11 @@ def fusion_loop(self, loop_lbl):
         """Run fusion callbacks on loop body"""
         logger = self.logger.getChild(f"ssa_loop_{loop_lbl}")
 
-        pre , body, post, _, other_data = \
+        pre , body, post, _, other_data, loop = \
             self.arch.Loop.extract(self.source, loop_lbl)
-        (loop_cnt, _, _) = other_data
+        loop_cnt = other_data['cnt']
         indentation = AsmHelper.find_indentation(body)
 
-        loop = self.arch.Loop(lbl_start=loop_lbl)
         body_ssa = SourceLine.read_multiline(loop.start(loop_cnt)) + \
             SourceLine.apply_indentation(self._fusion_core(pre, body, logger), indentation) + \
             SourceLine.read_multiline(loop.end(other_data))
@@ -398,9 +397,9 @@ def optimize_loop(self, loop_lbl, postamble_label=None):
 
         logger = self.logger.getChild(loop_lbl)
 
-        early, body, late, _, other_data = \
+        early, body, late, _, other_data, loop = \
             self.arch.Loop.extract(self.source, loop_lbl)
-        (loop_cnt, _, _) = other_data
+        loop_cnt = other_data['cnt']
 
         # Check if the body has a dominant indentation
         indentation = AsmHelper.find_indentation(body)
@@ -464,7 +463,6 @@ def loop_lbl_iter(i):
             for i in range(1, num_exceptional):
                 optimized_code += indented(self.arch.Branch.if_equal(loop_cnt, i, loop_lbl_iter(i)))
 
-        loop = self.arch.Loop(lbl_start=loop_lbl)
         optimized_code += indented(preamble_code)
 
         if self.config.sw_pipelining.unknown_iteration_count:
diff --git a/slothy/targets/aarch64/aarch64_neon.py b/slothy/targets/aarch64/aarch64_neon.py
index b2e8f36f..fad73cc0 100644
--- a/slothy/targets/aarch64/aarch64_neon.py
+++ b/slothy/targets/aarch64/aarch64_neon.py
@@ -43,6 +43,7 @@ class which generates instruction parsers and writers from instruction templates
 import math
 from enum import Enum
 from functools import cache
+from abc import ABC, abstractmethod
 
 from sympy import simplify
 
@@ -169,70 +170,46 @@ def unconditional(lbl):
         """Emit unconditional branch"""
         yield f"b {lbl}"
 
-class Loop:
-    """Helper functions for parsing and writing simple loops in AArch64
-
-    TODO: Generalize; current implementation too specific about shape of loop"""
-
+class Loop(ABC):
     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
 
+    @abstractmethod
     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}:"
+        # TODO: Use different type of fixup for cmp vs. subs loops
+        pass
 
+    @abstractmethod
     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}sub {reg0}, {reg1}, {imm}"
-        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|bnz|bne) <cnt>, loop_lbl
-           ```
+        pass
+    
+    def _extract(self, source, lbl):
+        """Locate a loop with start label `lbl` in `source`.```
 
         """
         assert isinstance(source, list)
+        
+        additional_data = None
 
         pre  = []
         body = []
         post = []
-        loop_lbl_regexp_txt = r"^\s*(?P<label>\w+)\s*:(?P<remainder>.*)$"
+        loop_lbl_regexp_txt = self.lbl_regex
         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|bnz|bne)\s+(?P<reg0>\w+),\s*{lbl}")
+        # end_regex shall contain group cnt as the counter variable
+        loop_end_regexp_txt = self.end_regex
         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
+        loop_end_ctr = 0
         while True:
             if not keep:
                 l = next(lines, None)
@@ -251,28 +228,63 @@ def extract(source, lbl):
                     pre.append(l)
                 continue
             if state == 1:
-                p = loop_end_regexp[0].match(l_str)
+                p = loop_end_regexp[loop_end_ctr].match(l_str)
                 if p is not None:
-                    reg0 = p.group("reg0")
-                    reg1 = p.group("reg1")
-                    imm = p.group("imm")
-                    state = 2
+                    if additional_data is None:
+                        additional_data = p.groupdict()
+                    loop_end_ctr += 1
+                    if loop_end_ctr == len(loop_end_regexp):
+                        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:
+        if state < 2:
             raise FatalParsingException(f"Couldn't identify loop {lbl}")
-        return pre, body, post, lbl, (reg0, reg1, imm)
+        return pre, body, post, lbl, additional_data
+
+    @staticmethod 
+    def extract(source, lbl):
+        for loop_type in Loop.__subclasses__():
+            try:
+                l = loop_type(lbl)
+                return l._extract(source, lbl) + (l,)
+            except FatalParsingException:
+                logging.debug("Parsing loop type '%s'failed", loop_type)
+                pass
+                
+        raise FatalParsingException(f"Couldn't identify loop {lbl}")
+
+class SubsLoop(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*sub[s]?\s+(?P<cnt>\w+),\s*(?P<reg1>\w+),\s*(?P<imm>#1)",
+                               rf"^\s*(cbnz|bnz|bne)\s+(?P<reg0>\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
+        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"
+
+        yield f"{indent}sub {other['cnt']}, {other['reg1']}, {other['imm']}"
+        yield f"{indent}cbnz {other['cnt']}, {lbl_start}"
 
 class FatalParsingException(Exception):
     """A fatal error happened during instruction parsing"""
diff --git a/slothy/targets/arm_v81m/arch_v81m.py b/slothy/targets/arm_v81m/arch_v81m.py
index 204f3f62..8349ee80 100644
--- a/slothy/targets/arm_v81m/arch_v81m.py
+++ b/slothy/targets/arm_v81m/arch_v81m.py
@@ -36,6 +36,7 @@
 import inspect
 import re
 import math
+from abc import ABC, abstractmethod
 
 from sympy import simplify
 from enum import Enum
@@ -109,46 +110,44 @@ def default_reserved():
         """Return the list of registers that should be reserved by default"""
         return set(["r14"])
 
-class Loop:
-
-    def __init__(self, lbl_start="1", lbl_end="2"):
+class Loop(ABC):
+    def __init__(self, lbl_start="1", lbl_end="2", loop_init="lr"):
         self.lbl_start = lbl_start
         self.lbl_end   = lbl_end
-        pass
+        self.loop_init = loop_init
 
-    def start(self,reg,indentation=0, fixup=0, unroll=1, jump_if_empty=None):
-        assert reg == "lr"
-        indent = ' ' * indentation
-        if unroll > 1:
-            if not unroll in [1,2,4,8,16,32]:
-                raise Exception("unsupported unrolling")
-            yield f"{indent}lsr lr, lr, #{int(math.log2(unroll))}"
-        if fixup != 0:
-            yield f"{indent}sub lr, lr, #{fixup}"
-        yield f".p2align 2"
-        yield f"{self.lbl_start}:"
-
-    def end(self,unused,indentation=0):
-        indent = ' ' * indentation
-        lbl_start = self.lbl_start
-        if lbl_start.isdigit():
-            lbl_start += "b"
-        yield f"{indent}le lr, {lbl_start}"
+    @abstractmethod
+    def start(self, loop_cnt, indentation=0, fixup=0, unroll=1, jump_if_empty=None):
+        """Emit starting instruction(s) and jump label for loop"""
+        # TODO: Use different type of fixup for cmp vs. subs loops
+        pass
 
-    def extract(source, lbl):
+    @abstractmethod
+    def end(self, other, indentation=0):
+        """Emit compare-and-branch at the end of the loop"""
+        pass
+    
+    def _extract(self, source, lbl):
+        """Locate a loop with start label `lbl` in `source`."""
         assert isinstance(source, list)
+        
+        additional_data = None
 
         pre  = []
         body = []
         post = []
-        loop_lbl_regexp_txt = r"^\s*(?P<label>\w+)\s*:(?P<remainder>.*)$"
+        loop_lbl_regexp_txt = self.lbl_regex
         loop_lbl_regexp = re.compile(loop_lbl_regexp_txt)
-        loop_end_regexp_txt = rf"^\s*le\s+(lr|r14)\s*,\s*{lbl}"
-        loop_end_regexp = re.compile(loop_end_regexp_txt)
+
+        # TODO: Allow other forms of looping
+        # end_regex shall contain group cnt as the counter variable
+        loop_end_regexp_txt = self.end_regex
+        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
+        loop_end_ctr = 0
         while True:
             if not keep:
                 l = next(lines, None)
@@ -160,16 +159,20 @@ def extract(source, lbl):
             if state == 0:
                 p = loop_lbl_regexp.match(l_str)
                 if p is not None and p.group("label") == lbl:
-                    l.set_text(p.group("remainder"))
+                    l = l.copy().set_text(p.group("remainder"))
                     keep = True
                     state = 1
                 else:
                     pre.append(l)
                 continue
             if state == 1:
-                p = loop_end_regexp.match(l_str)
+                p = loop_end_regexp[loop_end_ctr].match(l_str)
                 if p is not None:
-                    state = 2
+                    if additional_data is None:
+                        additional_data = p.groupdict()
+                    loop_end_ctr += 1
+                    if loop_end_ctr == len(loop_end_regexp):
+                        state = 2
                     continue
                 body.append(l)
                 continue
@@ -177,8 +180,51 @@ def extract(source, lbl):
                 post.append(l)
                 continue
         if state < 2:
-            raise Exception(f"Couldn't identify loop {lbl}")
-        return pre, body, post, lbl, ("lr", "lr", 0)
+            raise FatalParsingException(f"Couldn't identify loop {lbl}")
+        return pre, body, post, lbl, additional_data
+
+    @staticmethod 
+    def extract(source, lbl):
+        for loop_type in Loop.__subclasses__():
+            try:
+                l = loop_type(lbl)
+                return l._extract(source, lbl) + (l,)
+            except FatalParsingException:
+                logging.debug("Parsing loop type '%s'failed", loop_type)
+                pass
+                
+        raise FatalParsingException(f"Couldn't identify loop {lbl}")
+
+class LeLoop(Loop):
+
+    def __init__(self, lbl_start="1", lbl_end="2"):
+        super().__init__(lbl_start=lbl_start, lbl_end=lbl_end)
+        self.lbl_regex = r"^\s*(?P<label>\w+)\s*:(?P<remainder>.*)$"
+        self.end_regex = (rf"^\s*le\s+((?P<cnt>\w+)|r14)\s*,\s*{lbl_start}",)
+
+    def start(self,reg,indentation=0, fixup=0, unroll=1, jump_if_empty=None):
+        assert reg == "lr"
+        indent = ' ' * indentation
+        if unroll > 1:
+            if not unroll in [1,2,4,8,16,32]:
+                raise Exception("unsupported unrolling")
+            yield f"{indent}lsr lr, lr, #{int(math.log2(unroll))}"
+        if fixup != 0:
+            yield f"{indent}sub lr, lr, #{fixup}"
+        yield f".p2align 2"
+        yield f"{self.lbl_start}:"
+
+    def end(self,unused,indentation=0):
+        indent = ' ' * indentation
+        lbl_start = self.lbl_start
+        if lbl_start.isdigit():
+            lbl_start += "b"
+        yield f"{indent}le lr, {lbl_start}"
+
+
+class FatalParsingException(Exception):
+    """A fatal error happened during instruction parsing"""
+
 
 class Instruction:
 

From 466e3f9a261d886755d03364a451aea253fa7391 Mon Sep 17 00:00:00 2001
From: Amin Abdulrahman <amin@abdulrahman.de>
Date: Tue, 15 Oct 2024 13:21:04 +0200
Subject: [PATCH 2/8] Move Loop class * Add file for common code between models

---
 slothy/helper.py                       |  90 ++++++++++++++++++++++
 slothy/targets/aarch64/aarch64_neon.py | 102 ++-----------------------
 slothy/targets/arm_v81m/arch_v81m.py   |  89 +--------------------
 slothy/targets/common.py               |   8 ++
 4 files changed, 107 insertions(+), 182 deletions(-)
 create mode 100644 slothy/targets/common.py

diff --git a/slothy/helper.py b/slothy/helper.py
index 5ef6c50c..1eb357ad 100644
--- a/slothy/helper.py
+++ b/slothy/helper.py
@@ -28,6 +28,9 @@
 import re
 import subprocess
 import logging
+from abc import ABC, abstractmethod
+
+from slothy.targets.common import *
 
 class SourceLine:
     """Representation of a single line of source code"""
@@ -1089,3 +1092,90 @@ def forward_to_file(self, log_label, filename, lvl=logging.DEBUG):
         h.setLevel(lvl)
         l.addHandler(h)
         self.forward(l)
+
+
+class Loop(ABC):
+    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
+
+    @abstractmethod
+    def start(self, loop_cnt, indentation=0, fixup=0, unroll=1, jump_if_empty=None):
+        """Emit starting instruction(s) and jump label for loop"""
+        pass
+
+    @abstractmethod
+    def end(self, other, indentation=0):
+        """Emit compare-and-branch at the end of the loop"""
+        pass
+    
+    def _extract(self, source, lbl):
+        """Locate a loop with start label `lbl` in `source`.```
+
+        """
+        assert isinstance(source, list)
+        
+        additional_data = {}
+
+        pre  = []
+        body = []
+        post = []
+        loop_lbl_regexp_txt = self.lbl_regex
+        loop_lbl_regexp = re.compile(loop_lbl_regexp_txt)
+
+        # TODO: Allow other forms of looping
+        # end_regex shall contain group cnt as the counter variable
+        loop_end_regexp_txt = self.end_regex
+        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
+        loop_end_ctr = 0
+        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[loop_end_ctr].match(l_str)
+                if p is not None:
+                    # collect all named groups
+                    additional_data = additional_data | p.groupdict()
+                    loop_end_ctr += 1
+                    if loop_end_ctr == len(loop_end_regexp):
+                        state = 2
+                    continue
+                body.append(l)
+                continue
+            if state == 2:
+                post.append(l)
+                continue
+        if state < 2:
+            raise FatalParsingException(f"Couldn't identify loop {lbl}")
+        return pre, body, post, lbl, additional_data
+
+    @staticmethod 
+    def extract(source, lbl):
+        for loop_type in Loop.__subclasses__():
+            try:
+                l = loop_type(lbl)
+                return l._extract(source, lbl) + (l,)
+            except FatalParsingException:
+                logging.debug("Parsing loop type '%s'failed", loop_type)
+                pass
+                
+        raise FatalParsingException(f"Couldn't identify loop {lbl}")
\ No newline at end of file
diff --git a/slothy/targets/aarch64/aarch64_neon.py b/slothy/targets/aarch64/aarch64_neon.py
index fad73cc0..7df0106d 100644
--- a/slothy/targets/aarch64/aarch64_neon.py
+++ b/slothy/targets/aarch64/aarch64_neon.py
@@ -43,10 +43,12 @@ class which generates instruction parsers and writers from instruction templates
 import math
 from enum import Enum
 from functools import cache
-from abc import ABC, abstractmethod
 
 from sympy import simplify
 
+from slothy.targets.common import *
+from slothy.helper import Loop
+
 arch_name = "Arm_AArch64"
 llvm_mca_arch = "aarch64"
 
@@ -170,99 +172,15 @@ def unconditional(lbl):
         """Emit unconditional branch"""
         yield f"b {lbl}"
 
-class Loop(ABC):
-    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
-
-    @abstractmethod
-    def start(self, loop_cnt, indentation=0, fixup=0, unroll=1, jump_if_empty=None):
-        """Emit starting instruction(s) and jump label for loop"""
-        # TODO: Use different type of fixup for cmp vs. subs loops
-        pass
-
-    @abstractmethod
-    def end(self, other, indentation=0):
-        """Emit compare-and-branch at the end of the loop"""
-        pass
-    
-    def _extract(self, source, lbl):
-        """Locate a loop with start label `lbl` in `source`.```
-
-        """
-        assert isinstance(source, list)
-        
-        additional_data = None
-
-        pre  = []
-        body = []
-        post = []
-        loop_lbl_regexp_txt = self.lbl_regex
-        loop_lbl_regexp = re.compile(loop_lbl_regexp_txt)
-
-        # TODO: Allow other forms of looping
-        # end_regex shall contain group cnt as the counter variable
-        loop_end_regexp_txt = self.end_regex
-        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
-        loop_end_ctr = 0
-        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[loop_end_ctr].match(l_str)
-                if p is not None:
-                    if additional_data is None:
-                        additional_data = p.groupdict()
-                    loop_end_ctr += 1
-                    if loop_end_ctr == len(loop_end_regexp):
-                        state = 2
-                    continue
-                body.append(l)
-                continue
-            if state == 2:
-                post.append(l)
-                continue
-        if state < 2:
-            raise FatalParsingException(f"Couldn't identify loop {lbl}")
-        return pre, body, post, lbl, additional_data
-
-    @staticmethod 
-    def extract(source, lbl):
-        for loop_type in Loop.__subclasses__():
-            try:
-                l = loop_type(lbl)
-                return l._extract(source, lbl) + (l,)
-            except FatalParsingException:
-                logging.debug("Parsing loop type '%s'failed", loop_type)
-                pass
-                
-        raise FatalParsingException(f"Couldn't identify loop {lbl}")
 
 class SubsLoop(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)
+        # The group naming in the regex should be consistent; give same group
+        # names to the same registers
         self.lbl_regex = r"^\s*(?P<label>\w+)\s*:(?P<remainder>.*)$"
         self.end_regex = (r"^\s*sub[s]?\s+(?P<cnt>\w+),\s*(?P<reg1>\w+),\s*(?P<imm>#1)",
-                               rf"^\s*(cbnz|bnz|bne)\s+(?P<reg0>\w+),\s*{lbl}")
+                               rf"^\s*(cbnz|bnz|bne)\s+(?P<cnt>\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"""
@@ -286,14 +204,6 @@ def end(self, other, indentation=0):
         yield f"{indent}sub {other['cnt']}, {other['reg1']}, {other['imm']}"
         yield f"{indent}cbnz {other['cnt']}, {lbl_start}"
 
-class FatalParsingException(Exception):
-    """A fatal error happened during instruction parsing"""
-
-class UnknownInstruction(Exception):
-    """The parent instruction class for the given object could not be found"""
-
-class UnknownRegister(Exception):
-    """The register could not be found"""
 
 class Instruction:
 
diff --git a/slothy/targets/arm_v81m/arch_v81m.py b/slothy/targets/arm_v81m/arch_v81m.py
index 8349ee80..a1e71775 100644
--- a/slothy/targets/arm_v81m/arch_v81m.py
+++ b/slothy/targets/arm_v81m/arch_v81m.py
@@ -36,11 +36,13 @@
 import inspect
 import re
 import math
-from abc import ABC, abstractmethod
 
 from sympy import simplify
 from enum import Enum
 
+from slothy.targets.common import *
+from slothy.helper import Loop
+
 arch_name = "Arm_v81M"
 llvm_mca_arch = "arm"
 
@@ -110,91 +112,6 @@ def default_reserved():
         """Return the list of registers that should be reserved by default"""
         return set(["r14"])
 
-class Loop(ABC):
-    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
-
-    @abstractmethod
-    def start(self, loop_cnt, indentation=0, fixup=0, unroll=1, jump_if_empty=None):
-        """Emit starting instruction(s) and jump label for loop"""
-        # TODO: Use different type of fixup for cmp vs. subs loops
-        pass
-
-    @abstractmethod
-    def end(self, other, indentation=0):
-        """Emit compare-and-branch at the end of the loop"""
-        pass
-    
-    def _extract(self, source, lbl):
-        """Locate a loop with start label `lbl` in `source`."""
-        assert isinstance(source, list)
-        
-        additional_data = None
-
-        pre  = []
-        body = []
-        post = []
-        loop_lbl_regexp_txt = self.lbl_regex
-        loop_lbl_regexp = re.compile(loop_lbl_regexp_txt)
-
-        # TODO: Allow other forms of looping
-        # end_regex shall contain group cnt as the counter variable
-        loop_end_regexp_txt = self.end_regex
-        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
-        loop_end_ctr = 0
-        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[loop_end_ctr].match(l_str)
-                if p is not None:
-                    if additional_data is None:
-                        additional_data = p.groupdict()
-                    loop_end_ctr += 1
-                    if loop_end_ctr == len(loop_end_regexp):
-                        state = 2
-                    continue
-                body.append(l)
-                continue
-            if state == 2:
-                post.append(l)
-                continue
-        if state < 2:
-            raise FatalParsingException(f"Couldn't identify loop {lbl}")
-        return pre, body, post, lbl, additional_data
-
-    @staticmethod 
-    def extract(source, lbl):
-        for loop_type in Loop.__subclasses__():
-            try:
-                l = loop_type(lbl)
-                return l._extract(source, lbl) + (l,)
-            except FatalParsingException:
-                logging.debug("Parsing loop type '%s'failed", loop_type)
-                pass
-                
-        raise FatalParsingException(f"Couldn't identify loop {lbl}")
-
 class LeLoop(Loop):
 
     def __init__(self, lbl_start="1", lbl_end="2"):
diff --git a/slothy/targets/common.py b/slothy/targets/common.py
new file mode 100644
index 00000000..a8ca4b40
--- /dev/null
+++ b/slothy/targets/common.py
@@ -0,0 +1,8 @@
+class FatalParsingException(Exception):
+    """A fatal error happened during instruction parsing"""
+
+class UnknownInstruction(Exception):
+    """The parent instruction class for the given object could not be found"""
+
+class UnknownRegister(Exception):
+    """The register could not be found"""
\ No newline at end of file

From a5a4ce6c96035f94ebb85c4a7632cac240c49b63 Mon Sep 17 00:00:00 2001
From: Amin Abdulrahman <amin@abdulrahman.de>
Date: Tue, 15 Oct 2024 13:44:46 +0200
Subject: [PATCH 3/8] Improve documentation for loop abstraction

---
 slothy/helper.py                       |  9 ++++++---
 slothy/targets/aarch64/aarch64_neon.py | 14 +++++++++++++-
 slothy/targets/arm_v81m/arch_v81m.py   | 13 ++++++++++++-
 3 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/slothy/helper.py b/slothy/helper.py
index 1eb357ad..c3e5b640 100644
--- a/slothy/helper.py
+++ b/slothy/helper.py
@@ -1111,11 +1111,11 @@ def end(self, other, indentation=0):
         pass
     
     def _extract(self, source, lbl):
-        """Locate a loop with start label `lbl` in `source`.```
-
-        """
+        """Locate a loop with start label `lbl` in `source`.```"""
         assert isinstance(source, list)
         
+        # additional_data will be assigned according to the capture groups from
+        # loop_end_regexp. 
         additional_data = {}
 
         pre  = []
@@ -1173,6 +1173,9 @@ def extract(source, lbl):
         for loop_type in Loop.__subclasses__():
             try:
                 l = loop_type(lbl)
+                # concatenate the extracted loop with an instance of the
+                # identified loop_type, (l,) creates a tuple with one element to
+                # merge with the tuple retuned by _extract
                 return l._extract(source, lbl) + (l,)
             except FatalParsingException:
                 logging.debug("Parsing loop type '%s'failed", loop_type)
diff --git a/slothy/targets/aarch64/aarch64_neon.py b/slothy/targets/aarch64/aarch64_neon.py
index 7df0106d..0c553a46 100644
--- a/slothy/targets/aarch64/aarch64_neon.py
+++ b/slothy/targets/aarch64/aarch64_neon.py
@@ -174,6 +174,18 @@ def unconditional(lbl):
 
 
 class SubsLoop(Loop):
+    """
+    Loop ending in a flag setting subtraction and a branch.
+    
+    Example:
+    ```
+           loop_lbl:
+               {code}
+               sub[s] <cnt>, <cnt>, #1
+               (cbnz|bnz|bne) <cnt>, loop_lbl
+    ```
+    where cnt is the loop counter in lr.
+    """
     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)
         # The group naming in the regex should be consistent; give same group
@@ -201,7 +213,7 @@ def end(self, other, indentation=0):
         if lbl_start.isdigit():
             lbl_start += "b"
 
-        yield f"{indent}sub {other['cnt']}, {other['reg1']}, {other['imm']}"
+        yield f"{indent}sub {other['cnt']}, {other['cnt']}, {other['imm']}"
         yield f"{indent}cbnz {other['cnt']}, {lbl_start}"
 
 
diff --git a/slothy/targets/arm_v81m/arch_v81m.py b/slothy/targets/arm_v81m/arch_v81m.py
index a1e71775..1e8f72f1 100644
--- a/slothy/targets/arm_v81m/arch_v81m.py
+++ b/slothy/targets/arm_v81m/arch_v81m.py
@@ -113,7 +113,18 @@ def default_reserved():
         return set(["r14"])
 
 class LeLoop(Loop):
-
+    """
+    Loop ending in a le instruction.
+    
+    Example:
+    ```
+           loop_lbl:
+               {code}
+               le <cnt>, loop_lbl
+    ```
+    
+    where cnt is the loop counter in lr.
+    """
     def __init__(self, lbl_start="1", lbl_end="2"):
         super().__init__(lbl_start=lbl_start, lbl_end=lbl_end)
         self.lbl_regex = r"^\s*(?P<label>\w+)\s*:(?P<remainder>.*)$"

From f276e8c5e1ce3f04f6376d15a19c3638e1f813a3 Mon Sep 17 00:00:00 2001
From: Amin Abdulrahman <amin@abdulrahman.de>
Date: Tue, 15 Oct 2024 13:58:18 +0200
Subject: [PATCH 4/8] Add examples for loop types

---
 example.py                                 | 35 ++++++++++++++++++++++
 examples/naive/aarch64/aarch64_loop_subs.s |  9 ++++++
 examples/naive/loop_le.s                   |  4 +++
 3 files changed, 48 insertions(+)
 create mode 100644 examples/naive/aarch64/aarch64_loop_subs.s
 create mode 100644 examples/naive/loop_le.s

diff --git a/example.py b/example.py
index aa3ea54b..89a295d6 100644
--- a/example.py
+++ b/example.py
@@ -196,6 +196,37 @@ def core(self, slothy):
         slothy.config.inputs_are_outputs = True
         slothy.optimize_loop("start")
 
+class LoopLe(Example):
+    def __init__(self, var="", arch=Arch_Armv81M, target=Target_CortexM55r1):
+        name = "loop_le"
+        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 AArch64LoopSubs(Example):
+    def __init__(self, var="", arch=AArch64_Neon, target=Target_CortexA55):
+        name = "aarch64_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 CRT(Example):
     def __init__(self):
@@ -1394,6 +1425,10 @@ def main():
                  AArch64Example2(),
                  AArch64Example2(target=Target_CortexA72),
 
+                # Loop examples
+                 AArch64LoopSubs(),
+                 LoopLe(),
+
                  CRT(),
 
                  ntt_n256_l6_s32("bar"),
diff --git a/examples/naive/aarch64/aarch64_loop_subs.s b/examples/naive/aarch64/aarch64_loop_subs.s
new file mode 100644
index 00000000..949aa6ac
--- /dev/null
+++ b/examples/naive/aarch64/aarch64_loop_subs.s
@@ -0,0 +1,9 @@
+count .req x2
+
+mov count, #16
+start:
+
+    nop
+
+    subs count, count, #1
+    cbnz count, start
diff --git a/examples/naive/loop_le.s b/examples/naive/loop_le.s
new file mode 100644
index 00000000..db8d62d2
--- /dev/null
+++ b/examples/naive/loop_le.s
@@ -0,0 +1,4 @@
+mov lr, #16
+start:
+    nop
+    le lr, start
\ No newline at end of file

From 2a900db1e77777f04e7861057bf736bfa94aca6e Mon Sep 17 00:00:00 2001
From: Amin Abdulrahman <amin@abdulrahman.de>
Date: Tue, 15 Oct 2024 15:31:41 +0200
Subject: [PATCH 5/8] Add option to give factor for loops with non-one
 in/decrements

---
 slothy/core/slothy.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/slothy/core/slothy.py b/slothy/core/slothy.py
index bcff4108..47022750 100644
--- a/slothy/core/slothy.py
+++ b/slothy/core/slothy.py
@@ -392,7 +392,7 @@ def fusion_loop(self, loop_lbl):
 
         assert SourceLine.is_source(self.source)
 
-    def optimize_loop(self, loop_lbl, postamble_label=None):
+    def optimize_loop(self, loop_lbl, postamble_label=None, cnt_iter_factor=1):
         """Optimize the loop starting at a given label"""
 
         logger = self.logger.getChild(loop_lbl)
@@ -475,7 +475,7 @@ def loop_lbl_iter(i):
         optimized_code += SourceLine.read_multiline(loop.start(
             loop_cnt,
             indentation=self.config.indentation,
-            fixup=num_exceptional,
+            fixup=num_exceptional * cnt_iter_factor,
             unroll=self.config.sw_pipelining.unroll,
             jump_if_empty=jump_if_empty))
         optimized_code += indented(kernel_code)

From 43058969ea9f4a877a9e5836487f8e585d4030df Mon Sep 17 00:00:00 2001
From: Amin Abdulrahman <amin@abdulrahman.de>
Date: Wed, 16 Oct 2024 15:04:44 +0200
Subject: [PATCH 6/8] Switch strategy for loop counter offset modification

---
 slothy/core/slothy.py                  | 8 +++++---
 slothy/targets/aarch64/aarch64_neon.py | 2 +-
 slothy/targets/arm_v81m/arch_v81m.py   | 2 +-
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/slothy/core/slothy.py b/slothy/core/slothy.py
index 47022750..3d643e14 100644
--- a/slothy/core/slothy.py
+++ b/slothy/core/slothy.py
@@ -392,7 +392,7 @@ def fusion_loop(self, loop_lbl):
 
         assert SourceLine.is_source(self.source)
 
-    def optimize_loop(self, loop_lbl, postamble_label=None, cnt_iter_factor=1):
+    def optimize_loop(self, loop_lbl, postamble_label=None):
         """Optimize the loop starting at a given label"""
 
         logger = self.logger.getChild(loop_lbl)
@@ -475,9 +475,11 @@ def loop_lbl_iter(i):
         optimized_code += SourceLine.read_multiline(loop.start(
             loop_cnt,
             indentation=self.config.indentation,
-            fixup=num_exceptional * cnt_iter_factor,
+            fixup=num_exceptional,
             unroll=self.config.sw_pipelining.unroll,
-            jump_if_empty=jump_if_empty))
+            jump_if_empty=jump_if_empty,
+            preamble_code=preamble_code,
+            postamble_code=postamble_code))
         optimized_code += indented(kernel_code)
         optimized_code += SourceLine.read_multiline(loop.end(other_data,
             indentation=self.config.indentation))
diff --git a/slothy/targets/aarch64/aarch64_neon.py b/slothy/targets/aarch64/aarch64_neon.py
index 0c553a46..97ebe689 100644
--- a/slothy/targets/aarch64/aarch64_neon.py
+++ b/slothy/targets/aarch64/aarch64_neon.py
@@ -194,7 +194,7 @@ def __init__(self, lbl="lbl", lbl_start="1", lbl_end="2", loop_init="lr") -> Non
         self.end_regex = (r"^\s*sub[s]?\s+(?P<cnt>\w+),\s*(?P<reg1>\w+),\s*(?P<imm>#1)",
                                rf"^\s*(cbnz|bnz|bne)\s+(?P<cnt>\w+),\s*{lbl}")
     
-    def start(self, loop_cnt, indentation=0, fixup=0, unroll=1, jump_if_empty=None):
+    def start(self, loop_cnt, indentation=0, fixup=0, unroll=1, jump_if_empty=None, preamble_code=None, postamble_code=None):
         """Emit starting instruction(s) and jump label for loop"""
         indent = ' ' * indentation
         if unroll > 1:
diff --git a/slothy/targets/arm_v81m/arch_v81m.py b/slothy/targets/arm_v81m/arch_v81m.py
index 1e8f72f1..60cfdd14 100644
--- a/slothy/targets/arm_v81m/arch_v81m.py
+++ b/slothy/targets/arm_v81m/arch_v81m.py
@@ -130,7 +130,7 @@ def __init__(self, lbl_start="1", lbl_end="2"):
         self.lbl_regex = r"^\s*(?P<label>\w+)\s*:(?P<remainder>.*)$"
         self.end_regex = (rf"^\s*le\s+((?P<cnt>\w+)|r14)\s*,\s*{lbl_start}",)
 
-    def start(self,reg,indentation=0, fixup=0, unroll=1, jump_if_empty=None):
+    def start(self,reg,indentation=0, fixup=0, unroll=1, jump_if_empty=None, preamble_code=None, postamble_code=None):
         assert reg == "lr"
         indent = ' ' * indentation
         if unroll > 1:

From 6461088ab53f2bf447bba0a1e30c8cffd9d6df89 Mon Sep 17 00:00:00 2001
From: Amin Abdulrahman <amin@abdulrahman.de>
Date: Wed, 16 Oct 2024 15:10:10 +0200
Subject: [PATCH 7/8] Make additional loop data part of class

---
 slothy/helper.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/slothy/helper.py b/slothy/helper.py
index c3e5b640..996d6db6 100644
--- a/slothy/helper.py
+++ b/slothy/helper.py
@@ -1099,6 +1099,7 @@ 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
+        self.additional_data = {}
 
     @abstractmethod
     def start(self, loop_cnt, indentation=0, fixup=0, unroll=1, jump_if_empty=None):
@@ -1116,7 +1117,6 @@ def _extract(self, source, lbl):
         
         # additional_data will be assigned according to the capture groups from
         # loop_end_regexp. 
-        additional_data = {}
 
         pre  = []
         body = []
@@ -1154,7 +1154,7 @@ def _extract(self, source, lbl):
                 p = loop_end_regexp[loop_end_ctr].match(l_str)
                 if p is not None:
                     # collect all named groups
-                    additional_data = additional_data | p.groupdict()
+                    self.additional_data = self.additional_data | p.groupdict()
                     loop_end_ctr += 1
                     if loop_end_ctr == len(loop_end_regexp):
                         state = 2
@@ -1166,7 +1166,7 @@ def _extract(self, source, lbl):
                 continue
         if state < 2:
             raise FatalParsingException(f"Couldn't identify loop {lbl}")
-        return pre, body, post, lbl, additional_data
+        return pre, body, post, lbl, self.additional_data
 
     @staticmethod 
     def extract(source, lbl):

From 899328bfc691402cfe05afddec8d705ad067faec Mon Sep 17 00:00:00 2001
From: Amin Abdulrahman <amin@abdulrahman.de>
Date: Thu, 21 Nov 2024 11:44:13 +0100
Subject: [PATCH 8/8] Loop parsing improvments

---
 slothy/core/slothy.py                  |  8 ++++++--
 slothy/helper.py                       | 27 +++++++++++++++++++-------
 slothy/targets/aarch64/aarch64_neon.py | 13 ++++++++-----
 slothy/targets/arm_v81m/arch_v81m.py   |  2 +-
 4 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/slothy/core/slothy.py b/slothy/core/slothy.py
index 3d643e14..b0b776ee 100644
--- a/slothy/core/slothy.py
+++ b/slothy/core/slothy.py
@@ -393,7 +393,9 @@ def fusion_loop(self, loop_lbl):
         assert SourceLine.is_source(self.source)
 
     def optimize_loop(self, loop_lbl, postamble_label=None):
-        """Optimize the loop starting at a given label"""
+        """Optimize the loop starting at a given label
+            The postamble_label marks the end of the loop kernel.
+        """
 
         logger = self.logger.getChild(loop_lbl)
 
@@ -479,7 +481,9 @@ def loop_lbl_iter(i):
             unroll=self.config.sw_pipelining.unroll,
             jump_if_empty=jump_if_empty,
             preamble_code=preamble_code,
-            postamble_code=postamble_code))
+            body_code=kernel_code,
+            postamble_code=postamble_code,
+            register_aliases=c.register_aliases))
         optimized_code += indented(kernel_code)
         optimized_code += SourceLine.read_multiline(loop.end(other_data,
             indentation=self.config.indentation))
diff --git a/slothy/helper.py b/slothy/helper.py
index 996d6db6..14bac9b8 100644
--- a/slothy/helper.py
+++ b/slothy/helper.py
@@ -1110,21 +1110,21 @@ def start(self, loop_cnt, indentation=0, fixup=0, unroll=1, jump_if_empty=None):
     def end(self, other, indentation=0):
         """Emit compare-and-branch at the end of the loop"""
         pass
-    
+
     def _extract(self, source, lbl):
         """Locate a loop with start label `lbl` in `source`.```"""
         assert isinstance(source, list)
-        
-        # additional_data will be assigned according to the capture groups from
-        # loop_end_regexp. 
 
+        # additional_data will be assigned according to the capture groups from
+        # loop_end_regexp.
         pre  = []
         body = []
         post = []
+        # candidate lines for the end of the loop
+        loop_end_candidates = []
         loop_lbl_regexp_txt = self.lbl_regex
         loop_lbl_regexp = re.compile(loop_lbl_regexp_txt)
 
-        # TODO: Allow other forms of looping
         # end_regex shall contain group cnt as the counter variable
         loop_end_regexp_txt = self.end_regex
         loop_end_regexp = [re.compile(txt) for txt in loop_end_regexp_txt]
@@ -1153,22 +1153,35 @@ def _extract(self, source, lbl):
             if state == 1:
                 p = loop_end_regexp[loop_end_ctr].match(l_str)
                 if p is not None:
+                    # Case: We may have encountered part of the loop end
                     # collect all named groups
                     self.additional_data = self.additional_data | p.groupdict()
                     loop_end_ctr += 1
+                    loop_end_candidates.append(l)
                     if loop_end_ctr == len(loop_end_regexp):
                         state = 2
                     continue
+                elif loop_end_ctr > 0 and l_str != "":
+                    # Case: The sequence of loop end candidates was interrupted
+                    #       i.e., we found a false-positive or this is not a proper loop
+                    
+                    # The loop end candidates are not part of the loop, meaning
+                    # they belonged to the body
+                    body += loop_end_candidates
+                    self.additional_data = {}
+                    loop_end_ctr = 0
+                    loop_end_candidates = []
                 body.append(l)
                 continue
             if state == 2:
+                loop_end_candidates = []
                 post.append(l)
                 continue
         if state < 2:
             raise FatalParsingException(f"Couldn't identify loop {lbl}")
         return pre, body, post, lbl, self.additional_data
 
-    @staticmethod 
+    @staticmethod
     def extract(source, lbl):
         for loop_type in Loop.__subclasses__():
             try:
@@ -1180,5 +1193,5 @@ def extract(source, lbl):
             except FatalParsingException:
                 logging.debug("Parsing loop type '%s'failed", loop_type)
                 pass
-                
+
         raise FatalParsingException(f"Couldn't identify loop {lbl}")
\ No newline at end of file
diff --git a/slothy/targets/aarch64/aarch64_neon.py b/slothy/targets/aarch64/aarch64_neon.py
index 97ebe689..68fdf3c0 100644
--- a/slothy/targets/aarch64/aarch64_neon.py
+++ b/slothy/targets/aarch64/aarch64_neon.py
@@ -176,12 +176,12 @@ def unconditional(lbl):
 class SubsLoop(Loop):
     """
     Loop ending in a flag setting subtraction and a branch.
-    
+
     Example:
     ```
            loop_lbl:
                {code}
-               sub[s] <cnt>, <cnt>, #1
+               sub[s] <cnt>, <cnt>, #<imm>
                (cbnz|bnz|bne) <cnt>, loop_lbl
     ```
     where cnt is the loop counter in lr.
@@ -191,16 +191,19 @@ def __init__(self, lbl="lbl", lbl_start="1", lbl_end="2", loop_init="lr") -> Non
         # The group naming in the regex should be consistent; give same group
         # names to the same registers
         self.lbl_regex = r"^\s*(?P<label>\w+)\s*:(?P<remainder>.*)$"
-        self.end_regex = (r"^\s*sub[s]?\s+(?P<cnt>\w+),\s*(?P<reg1>\w+),\s*(?P<imm>#1)",
+        self.end_regex = (r"^\s*sub[s]?\s+(?P<cnt>\w+),\s*(?P<reg1>\w+),\s*#(?P<imm>\d+)",
                                rf"^\s*(cbnz|bnz|bne)\s+(?P<cnt>\w+),\s*{lbl}")
-    
-    def start(self, loop_cnt, indentation=0, fixup=0, unroll=1, jump_if_empty=None, preamble_code=None, postamble_code=None):
+
+    def start(self, loop_cnt, indentation=0, fixup=0, unroll=1, jump_if_empty=None, preamble_code=None, body_code=None, postamble_code=None, register_aliases=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:
+            # In case the immediate is >1, we need to scale the fixup. This
+            # allows for loops that do not use an increment of 1
+            fixup *= self.additional_data['imm']
             yield f"{indent}sub {loop_cnt}, {loop_cnt}, #{fixup}"
         if jump_if_empty is not None:
             yield f"cbz {loop_cnt}, {jump_if_empty}"
diff --git a/slothy/targets/arm_v81m/arch_v81m.py b/slothy/targets/arm_v81m/arch_v81m.py
index 60cfdd14..f14e47c8 100644
--- a/slothy/targets/arm_v81m/arch_v81m.py
+++ b/slothy/targets/arm_v81m/arch_v81m.py
@@ -130,7 +130,7 @@ def __init__(self, lbl_start="1", lbl_end="2"):
         self.lbl_regex = r"^\s*(?P<label>\w+)\s*:(?P<remainder>.*)$"
         self.end_regex = (rf"^\s*le\s+((?P<cnt>\w+)|r14)\s*,\s*{lbl_start}",)
 
-    def start(self,reg,indentation=0, fixup=0, unroll=1, jump_if_empty=None, preamble_code=None, postamble_code=None):
+    def start(self,reg,indentation=0, fixup=0, unroll=1, jump_if_empty=None, preamble_code=None, body_code=None, postamble_code=None, register_aliases=None):
         assert reg == "lr"
         indent = ' ' * indentation
         if unroll > 1: