From d132d85b0d6d0e5820978461ea3dfdf4f8c275cc Mon Sep 17 00:00:00 2001
From: Bo Wun Cheng <bwcheng@stanford.edu>
Date: Wed, 13 Dec 2023 19:26:14 -0800
Subject: [PATCH] added mapping and routing support for fp_max, fp_add and
 faddiexp instruction for the alu

---
 sam/onyx/hw_nodes/compute_node.py      | 16 +++++++++++++---
 sam/onyx/hw_nodes/read_scanner_node.py | 10 +++++++++-
 sam/onyx/hw_nodes/reduce_node.py       |  2 +-
 sam/onyx/parse_dot.py                  |  4 ++--
 4 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/sam/onyx/hw_nodes/compute_node.py b/sam/onyx/hw_nodes/compute_node.py
index b3d375b1..eac07348 100644
--- a/sam/onyx/hw_nodes/compute_node.py
+++ b/sam/onyx/hw_nodes/compute_node.py
@@ -119,10 +119,16 @@ def connect(self, other, edge, kwargs=None):
             other_conn = other.get_num_inputs()
             pe = self.get_name()
             # TODO: remove hack eventually
-            if 'Max' in other.op:
+            if 'Max 0' in other.op:
                 other_conn = 1
-            else:
-                other_conn = other.get_num_inputs()
+            elif 'Faddiexp' in other.op:
+                comment = edge.get_attributes()["comment"].strip('"')
+                if 'fp' in comment:
+                    other_conn = 0
+                elif 'exp' in comment:
+                    other_conn = 1
+                else: 
+                    assert 0 & "edge connected to faddiexp has to have comment specified to either 'exp' or 'fp'"
             new_conns = {
                 f'pe_to_pe_{other_conn}': [
                     ([(pe, "res"), (other_pe, f"data{other_conn}")], 17),
@@ -186,6 +192,10 @@ def configure(self, attributes):
             op_code = 8
         elif c_op == 'faddiexp':
             op_code = 9
+        elif c_op == 'fp_max':
+            op_code = 10
+        elif c_op == 'fp_add':
+            op_code = 11
 
         rb_const = None
         if "rb_const" in attributes:
diff --git a/sam/onyx/hw_nodes/read_scanner_node.py b/sam/onyx/hw_nodes/read_scanner_node.py
index 14c2bd18..72afd45e 100644
--- a/sam/onyx/hw_nodes/read_scanner_node.py
+++ b/sam/onyx/hw_nodes/read_scanner_node.py
@@ -204,7 +204,15 @@ def connect(self, other, edge, kwargs=None):
             # Can use dynamic information to assign inputs to compute nodes
             # since add/mul are commutative
             compute_conn = other.get_num_inputs()
-
+            # TODO: get rid of this hack
+            if 'Faddiexp' in other.op:
+                comment = edge.get_attributes()["comment"].strip('"')
+                if 'fp' in comment:
+                    compute_conn = 0
+                elif 'exp' in comment:
+                    compute_conn = 1
+                else: 
+                    assert 0 & "edge connected to faddiexp has to have comment specified to either 'exp' or 'fp'"
             new_conns = {
                 f'rd_scan_to_compute_{compute_conn}': [
                     ([(rd_scan, "coord_out"), (compute, f"data{compute_conn}")], 17),
diff --git a/sam/onyx/hw_nodes/reduce_node.py b/sam/onyx/hw_nodes/reduce_node.py
index 2c904309..b19cf5e2 100644
--- a/sam/onyx/hw_nodes/reduce_node.py
+++ b/sam/onyx/hw_nodes/reduce_node.py
@@ -68,7 +68,7 @@ def connect(self, other, edge, kwargs=None):
             raise NotImplementedError(f'Cannot connect ReduceNode to {other_type}')
         elif other_type == ComputeNode:
             pe = other.get_name()
-            if 'Max' in other.op:
+            if 'Max 0' in other.op:
                 other_conn = 1
             else:
                 other_conn = other.get_num_inputs()
diff --git a/sam/onyx/parse_dot.py b/sam/onyx/parse_dot.py
index 0da63382..fa18a557 100644
--- a/sam/onyx/parse_dot.py
+++ b/sam/onyx/parse_dot.py
@@ -99,11 +99,11 @@ def map_nodes(self):
                     hw_nt = f"HWNodeType.RepSigGen"
                 elif n_type == "repeat":
                     hw_nt = f"HWNodeType.Repeat"
-                elif n_type == "mul" or n_type == "add" or n_type == "max":
+                elif n_type == "mul" or n_type == "add" or n_type == "max" or n_type == "and":
                     hw_nt = f"HWNodeType.Compute"
                 elif n_type == "fgetfint" or n_type == "fgetffrac" or n_type == "faddiexp":
                     hw_nt = f"HWNodeType.Compute"
-                elif n_type == "fp_mul":
+                elif n_type == "fp_mul" or n_type == "fp_max" or n_type == "fp_add":
                     hw_nt = f"HWNodeType.Compute"
                 elif n_type == "reduce":
                     hw_nt = f"HWNodeType.Reduce"