From 59953c68fae6b680fa0b332fc56cb80905de0390 Mon Sep 17 00:00:00 2001
From: Christoph Berganski <christoph.berganski@gmail.com>
Date: Tue, 28 Jan 2025 15:33:15 +0100
Subject: [PATCH 1/3] [Streamline] Extend set of invariants for
 MoveScalarLinearPastInvariants

---
 src/finn/transformation/streamline/reorder.py | 31 ++++++++++++++-----
 1 file changed, 24 insertions(+), 7 deletions(-)

diff --git a/src/finn/transformation/streamline/reorder.py b/src/finn/transformation/streamline/reorder.py
index 9a7e9d0723..d31bfc9c60 100644
--- a/src/finn/transformation/streamline/reorder.py
+++ b/src/finn/transformation/streamline/reorder.py
@@ -606,6 +606,19 @@ class MoveScalarLinearPastInvariants(Transformation):
     GlobalAveragePool
     """
 
+    # Op-types of currently supported invariants
+    SUPPORTED_INVARIANTS = {
+        "GlobalAveragePool",
+        "Identity"
+        "Reshape",
+        "Transpose",
+        "Flatten",
+        "Expand"
+        "Slice",
+        "Squeeze",
+        "Unsqueeze",
+    }
+
     def apply(self, model):
         graph = model.graph
         node_ind = 0
@@ -618,13 +631,7 @@ def apply(self, model):
                 # Extract mode and scales and input shape
                 mode = get_by_name(n.attribute, "mode").s.decode("ascii")
                 is_nearest_neighbor_resample = mode == "nearest"
-            if (
-                n.op_type == "GlobalAveragePool"
-                or n.op_type == "Reshape"
-                or n.op_type == "Transpose"
-                or n.op_type == "Flatten"
-                or is_nearest_neighbor_resample
-            ):
+            if n.op_type in self.SUPPORTED_INVARIANTS or is_nearest_neighbor_resample:
                 in0 = n.input[0]
                 if in0 is None:
                     continue
@@ -634,6 +641,16 @@ def apply(self, model):
                     continue
 
                 if prod0.op_type in ["Mul", "Add", "Div"]:
+                    # Cannot handle fork-nodes, try MoveLinearPastFork first
+                    if model.is_fork_node(prod0):
+                        warnings.warn(
+                            f"{self.__class__.__name__}:"
+                            f" Skipping near match: {prod0.name} is a fork-node,"
+                            f" try MoveLinearPastFork first"
+                        )
+                        # Skip transforming this node as moving this would lead
+                        # to messed up or detached graph
+                        continue
                     # check if second input of producer is an initializer
                     init0 = model.get_initializer(prod0.input[1])
                     # if either initializer is None, skip

From dd0a6e14c8dc89a9a786fae50d257d78b2d7a831 Mon Sep 17 00:00:00 2001
From: Christoph Berganski <christoph.berganski@gmail.com>
Date: Tue, 28 Jan 2025 15:37:48 +0100
Subject: [PATCH 2/3] [Streamline] Enable "Sub" for
 MoveScalarLinearPastInvariants

---
 src/finn/transformation/streamline/reorder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/finn/transformation/streamline/reorder.py b/src/finn/transformation/streamline/reorder.py
index d31bfc9c60..09645d097e 100644
--- a/src/finn/transformation/streamline/reorder.py
+++ b/src/finn/transformation/streamline/reorder.py
@@ -640,7 +640,7 @@ def apply(self, model):
                 if prod0 is None:
                     continue
 
-                if prod0.op_type in ["Mul", "Add", "Div"]:
+                if prod0.op_type in ["Mul", "Div", "Add", "Sub"]:
                     # Cannot handle fork-nodes, try MoveLinearPastFork first
                     if model.is_fork_node(prod0):
                         warnings.warn(

From 1b5b465bf4136416bbea082412de6d8668b0422b Mon Sep 17 00:00:00 2001
From: Christoph Berganski <christoph.berganski@gmail.com>
Date: Tue, 28 Jan 2025 16:49:07 +0100
Subject: [PATCH 3/3] Add missing commas

---
 src/finn/transformation/streamline/reorder.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/finn/transformation/streamline/reorder.py b/src/finn/transformation/streamline/reorder.py
index 09645d097e..b67758b0d6 100644
--- a/src/finn/transformation/streamline/reorder.py
+++ b/src/finn/transformation/streamline/reorder.py
@@ -609,11 +609,11 @@ class MoveScalarLinearPastInvariants(Transformation):
     # Op-types of currently supported invariants
     SUPPORTED_INVARIANTS = {
         "GlobalAveragePool",
-        "Identity"
+        "Identity",
         "Reshape",
         "Transpose",
         "Flatten",
-        "Expand"
+        "Expand",
         "Slice",
         "Squeeze",
         "Unsqueeze",