From bb37e0013b9c3a2bbe833aa4f56368f66531a50c Mon Sep 17 00:00:00 2001
From: root <owhsu@stanford.edu>
Date: Mon, 28 Oct 2024 16:10:44 -0700
Subject: [PATCH] Clean up SAM repo

---
 sam/onyx/asplos.py     | 4 ++--
 sam/onyx/fiber_tree.py | 2 +-
 sam/onyx/parse_dot.py  | 9 +++++----
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/sam/onyx/asplos.py b/sam/onyx/asplos.py
index 35d7bd0c..151c9c35 100644
--- a/sam/onyx/asplos.py
+++ b/sam/onyx/asplos.py
@@ -31,13 +31,13 @@
     parser = argparse.ArgumentParser(description='ASPLOS argparser')
     parser.add_argument('--sam_graph',
                         type=str,
-                        default="/home/max/Documents/SPARSE/sam/compiler/sam-outputs/dot/")
+                        default="./sam/compiler/sam-outputs/dot/")
     parser.add_argument('--output_png',
                         type=str,
                         default="output.png")
     parser.add_argument('--output_graph',
                         type=str,
-                        default="/home/max/Documents/SPARSE/sam/mek.gv")
+                        default="output.gv")
     args = parser.parse_args()
 
     sam_graph = args.sam_graph
diff --git a/sam/onyx/fiber_tree.py b/sam/onyx/fiber_tree.py
index 66525311..1b10ffb6 100644
--- a/sam/onyx/fiber_tree.py
+++ b/sam/onyx/fiber_tree.py
@@ -91,7 +91,7 @@ def __str__(self):
 if __name__ == "__main__":
     random.seed(10)
     numpy.random.seed(10)
-    mg = MatrixGenerator(name='B', shape=[10, 10], dump_dir='/home/max/Documents/SPARSE/sam/OUTPUTS_DUMP', sparsity=0.8)
+    mg = MatrixGenerator(name='B', shape=[10, 10], dump_dir='./OUTPUTS_DUMP', sparsity=0.8)
     array = mg.get_matrix()
     print(array)
     ft = FiberTree(tensor=array)
diff --git a/sam/onyx/parse_dot.py b/sam/onyx/parse_dot.py
index b56a8c38..21730f09 100644
--- a/sam/onyx/parse_dot.py
+++ b/sam/onyx/parse_dot.py
@@ -42,7 +42,6 @@ def __init__(self, filename=None, local_mems=True, use_fork=False,
         self.stage2_count = {}
 
         self.annotate_IO_nodes()
-        self.graph.write_png('mek.png')
         # exit()
         # print(self.graph)
 
@@ -1451,14 +1450,16 @@ def parse_graph(graph):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='SAM DOT Parser')
     parser.add_argument('--sam_graph',
+                        type=str)
+    parser.add_argument('--collat_dir',
                         type=str,
-                        default="/home/max/Documents/SPARSE/sam/compiler/sam-outputs/dot/mat_identity.gv")
+                        default="temp/")
     parser.add_argument('--output_png',
                         type=str,
                         default="output.png")
     parser.add_argument('--output_graph',
                         type=str,
-                        default="/home/max/Documents/SPARSE/sam/mek.gv")
+                        default="output.gv")
     parser.add_argument('--unroll',
                         type=int,
                         default=1)
@@ -1469,7 +1470,7 @@ def parse_graph(graph):
     output_graph = args.output_graph
     unroll = args.unroll
     sdg = SAMDotGraph(filename=sam_graph, use_fork=True,
-                      unroll=unroll)
+                      unroll=unroll, collat_dir=args.collat_dir)
     graph = sdg.get_graph()
     print(graph)
     # parse_graph(graph)