From db51239c610fa54bebaf110da3d34408732009da Mon Sep 17 00:00:00 2001
From: Yaman Umuroglu <maltanar@gmail.com>
Date: Wed, 8 Apr 2020 21:56:46 +0100
Subject: [PATCH] [Test] use get_nodes_by_op_type in end2end tests

---
 tests/end2end/test_end2end_cnv_w1a1.py | 50 +++++++++-----------------
 tests/end2end/test_end2end_tfc_w1a1.py | 29 ++++++++-------
 tests/end2end/test_end2end_tfc_w1a2.py | 29 ++++++++-------
 tests/end2end/test_end2end_tfc_w2a2.py | 29 ++++++++-------
 4 files changed, 59 insertions(+), 78 deletions(-)

diff --git a/tests/end2end/test_end2end_cnv_w1a1.py b/tests/end2end/test_end2end_cnv_w1a1.py
index 668ff267e..30440bdaa 100644
--- a/tests/end2end/test_end2end_cnv_w1a1.py
+++ b/tests/end2end/test_end2end_cnv_w1a1.py
@@ -124,7 +124,8 @@ def test_end2end_cnv_w1a1_create_dataflow_partition():
     model = ModelWrapper(build_dir + "/end2end_cnv_w1a1_hls_layers.onnx")
     parent_model = model.transform(CreateDataflowPartition())
     parent_model.save(build_dir + "/end2end_cnv_w1a1_dataflow_parent.onnx")
-    sdp_node = getCustomOp(parent_model.graph.node[2])
+    sdp_node = parent_model.get_nodes_by_op_type("StreamingDataflowPartition")[0]
+    sdp_node = getCustomOp(sdp_node)
     dataflow_model_filename = sdp_node.get_nodeattr("model")
     dataflow_model = ModelWrapper(dataflow_model_filename)
     dataflow_model.save(build_dir + "/end2end_cnv_w1a1_dataflow_model.onnx")
@@ -132,49 +133,32 @@ def test_end2end_cnv_w1a1_create_dataflow_partition():
 
 def test_end2end_cnv_w1a1_fold_and_tlastmarker():
     model = ModelWrapper(build_dir + "/end2end_cnv_w1a1_dataflow_model.onnx")
-    fc0 = model.graph.node[1]
-    fc1 = model.graph.node[3]
-    fc2 = model.graph.node[6]
-    fc3 = model.graph.node[8]
-    fc4 = model.graph.node[11]
-    fc5 = model.graph.node[13]
-    fc6 = model.graph.node[14]
-    fc7 = model.graph.node[15]
-    fc8 = model.graph.node[16]
-    fc0w = getCustomOp(fc0)
-    fc1w = getCustomOp(fc1)
-    fc2w = getCustomOp(fc2)
-    fc3w = getCustomOp(fc3)
-    fc4w = getCustomOp(fc4)
-    fc5w = getCustomOp(fc5)
-    fc6w = getCustomOp(fc6)
-    fc7w = getCustomOp(fc7)
-    fc8w = getCustomOp(fc8)
-
+    fc_layers = model.get_nodes_by_op_type("StreamingFCLayer_Batch")
+    fc0w = getCustomOp(fc_layers[0])
+    fc1w = getCustomOp(fc_layers[1])
+    fc2w = getCustomOp(fc_layers[2])
+    fc3w = getCustomOp(fc_layers[3])
+    fc4w = getCustomOp(fc_layers[4])
+    fc5w = getCustomOp(fc_layers[5])
+    fc6w = getCustomOp(fc_layers[6])
+    fc7w = getCustomOp(fc_layers[7])
+    fc8w = getCustomOp(fc_layers[8])
     fc0w.set_nodeattr("SIMD", 27)
     fc0w.set_nodeattr("PE", 8)
-
     fc1w.set_nodeattr("SIMD", 32)
     fc1w.set_nodeattr("PE", 8)
-
     fc2w.set_nodeattr("SIMD", 32)
     fc2w.set_nodeattr("PE", 16)
-
     fc3w.set_nodeattr("SIMD", 32)
     fc3w.set_nodeattr("PE", 16)
-
     fc4w.set_nodeattr("SIMD", 32)
     fc4w.set_nodeattr("PE", 32)
-
     fc5w.set_nodeattr("SIMD", 64)
     fc5w.set_nodeattr("PE", 16)
-
     fc6w.set_nodeattr("SIMD", 32)
     fc6w.set_nodeattr("PE", 16)
-
     fc7w.set_nodeattr("SIMD", 64)
     fc7w.set_nodeattr("PE", 8)
-
     fc8w.set_nodeattr("SIMD", 16)
     fc8w.set_nodeattr("PE", 10)
 
@@ -213,10 +197,9 @@ def test_end2end_cnv_w1a1_verify_dataflow_part():
     res_npysim = ret_npysim[out_name]
     # node-by-node rtlsim
     model = model.transform(SetExecMode("rtlsim"))
-    for node in model.graph.node:
-        if node.op_type == "StreamingFCLayer_Batch":
-            inst = getCustomOp(node)
-            inst.set_nodeattr("rtlsim_trace", "default")
+    fc_layers = model.get_nodes_by_op_type("StreamingFCLayer_Batch")
+    for fcl in fc_layers:
+        getCustomOp(fcl).set_nodeattr("rtlsim_trace", "default")
     model.save(build_dir + "/end2end_cnv_w1a1_ipgen_nodebynode_rtlsim.onnx")
     ret_rtlsim_nodebynode = execute_onnx(model, inp_dict, True)
     res_rtlsim_nodebynode = ret_rtlsim_nodebynode[out_name]
@@ -332,7 +315,8 @@ def test_end2end_cnv_w1a1_run_on_pynq():
         if ip == "":
             pytest.skip("PYNQ board IP address not specified")
         # produce results with npysim
-        sdp_node = getCustomOp(parent_model.graph.node[2])
+        sdp_node = parent_model.get_nodes_by_op_type("StreamingDataflowPartition")[0]
+        sdp_node = getCustomOp(sdp_node)
         sdp_node.set_nodeattr("model", build_dir + "/end2end_cnv_w1a1_pynq_deploy.onnx")
         ret = execute_onnx(parent_model, {iname: x}, True)
         y = ret[oname]
diff --git a/tests/end2end/test_end2end_tfc_w1a1.py b/tests/end2end/test_end2end_tfc_w1a1.py
index 9cd338caa..d61ea8d27 100644
--- a/tests/end2end/test_end2end_tfc_w1a1.py
+++ b/tests/end2end/test_end2end_tfc_w1a1.py
@@ -117,7 +117,8 @@ def test_end2end_tfc_w1a1_create_dataflow_partition():
     model = ModelWrapper(build_dir + "/end2end_tfc_w1a1_hls_layers.onnx")
     parent_model = model.transform(CreateDataflowPartition())
     parent_model.save(build_dir + "/end2end_tfc_w1a1_dataflow_parent.onnx")
-    sdp_node = getCustomOp(parent_model.graph.node[2])
+    sdp_node = parent_model.get_nodes_by_op_type("StreamingDataflowPartition")[0]
+    sdp_node = getCustomOp(sdp_node)
     dataflow_model_filename = sdp_node.get_nodeattr("model")
     dataflow_model = ModelWrapper(dataflow_model_filename)
     dataflow_model.save(build_dir + "/end2end_tfc_w1a1_dataflow_model.onnx")
@@ -125,14 +126,11 @@ def test_end2end_tfc_w1a1_create_dataflow_partition():
 
 def test_end2end_tfc_w1a1_fold_and_tlastmarker():
     model = ModelWrapper(build_dir + "/end2end_tfc_w1a1_dataflow_model.onnx")
-    fc0 = model.graph.node[0]
-    fc1 = model.graph.node[1]
-    fc2 = model.graph.node[2]
-    fc3 = model.graph.node[3]
-    fc0w = getCustomOp(fc0)
-    fc1w = getCustomOp(fc1)
-    fc2w = getCustomOp(fc2)
-    fc3w = getCustomOp(fc3)
+    fc_layers = model.get_nodes_by_op_type("StreamingFCLayer_Batch")
+    fc0w = getCustomOp(fc_layers[0])
+    fc1w = getCustomOp(fc_layers[1])
+    fc2w = getCustomOp(fc_layers[2])
+    fc3w = getCustomOp(fc_layers[3])
     fc0w.set_nodeattr("inFIFODepth", 50)
     fc0w.set_nodeattr("SIMD", 16)
     fc0w.set_nodeattr("PE", 16)
@@ -181,10 +179,9 @@ def test_end2end_tfc_w1a1_verify_dataflow_part():
     res_npysim = ret_npysim[out_name]
     # node-by-node rtlsim
     model = model.transform(SetExecMode("rtlsim"))
-    getCustomOp(model.graph.node[0]).set_nodeattr("rtlsim_trace", "default")
-    getCustomOp(model.graph.node[1]).set_nodeattr("rtlsim_trace", "default")
-    getCustomOp(model.graph.node[2]).set_nodeattr("rtlsim_trace", "default")
-    getCustomOp(model.graph.node[3]).set_nodeattr("rtlsim_trace", "default")
+    fc_layers = model.get_nodes_by_op_type("StreamingFCLayer_Batch")
+    for fcl in fc_layers:
+        getCustomOp(fcl).set_nodeattr("rtlsim_trace", "default")
     model.save(build_dir + "/end2end_tfc_w1a1_ipstitch_nodebynode_rtlsim.onnx")
     ret_rtlsim_nodebynode = execute_onnx(model, inp_dict, True)
     res_rtlsim_nodebynode = ret_rtlsim_nodebynode[out_name]
@@ -215,7 +212,8 @@ def test_end2end_tfc_w1a1_verify_all():
     iname = parent_model.graph.input[0].name
     oname = parent_model.graph.output[0].name
     # produce results with npysim
-    sdp_node = getCustomOp(parent_model.graph.node[2])
+    sdp_node = parent_model.get_nodes_by_op_type("StreamingDataflowPartition")[0]
+    sdp_node = getCustomOp(sdp_node)
     sdp_node.set_nodeattr("model", build_dir + "/end2end_tfc_w1a1_ipstitch_npysim.onnx")
     ret_npysim = execute_onnx(parent_model, {iname: x}, True)
     y_npysim = ret_npysim[oname]
@@ -292,7 +290,8 @@ def test_end2end_tfc_w1a1_run_on_pynq():
         if ip == "":
             pytest.skip("PYNQ board IP address not specified")
         # produce results with npysim
-        sdp_node = getCustomOp(parent_model.graph.node[2])
+        sdp_node = parent_model.get_nodes_by_op_type("StreamingDataflowPartition")[0]
+        sdp_node = getCustomOp(sdp_node)
         sdp_node.set_nodeattr("model", build_dir + "/end2end_tfc_w1a1_pynq_deploy.onnx")
         ret = execute_onnx(parent_model, {iname: x}, True)
         y = ret[oname]
diff --git a/tests/end2end/test_end2end_tfc_w1a2.py b/tests/end2end/test_end2end_tfc_w1a2.py
index e3eead845..2ac344af5 100644
--- a/tests/end2end/test_end2end_tfc_w1a2.py
+++ b/tests/end2end/test_end2end_tfc_w1a2.py
@@ -109,7 +109,8 @@ def test_end2end_tfc_w1a2_create_dataflow_partition():
     model = ModelWrapper(build_dir + "/end2end_tfc_w1a2_hls_layers.onnx")
     parent_model = model.transform(CreateDataflowPartition())
     parent_model.save(build_dir + "/end2end_tfc_w1a2_dataflow_parent.onnx")
-    sdp_node = getCustomOp(parent_model.graph.node[2])
+    sdp_node = parent_model.get_nodes_by_op_type("StreamingDataflowPartition")[0]
+    sdp_node = getCustomOp(sdp_node)
     dataflow_model_filename = sdp_node.get_nodeattr("model")
     dataflow_model = ModelWrapper(dataflow_model_filename)
     dataflow_model.save(build_dir + "/end2end_tfc_w1a2_dataflow_model.onnx")
@@ -117,14 +118,11 @@ def test_end2end_tfc_w1a2_create_dataflow_partition():
 
 def test_end2end_tfc_w1a2_fold_and_tlastmarker():
     model = ModelWrapper(build_dir + "/end2end_tfc_w1a2_dataflow_model.onnx")
-    fc0 = model.graph.node[0]
-    fc1 = model.graph.node[1]
-    fc2 = model.graph.node[2]
-    fc3 = model.graph.node[3]
-    fc0w = getCustomOp(fc0)
-    fc1w = getCustomOp(fc1)
-    fc2w = getCustomOp(fc2)
-    fc3w = getCustomOp(fc3)
+    fc_layers = model.get_nodes_by_op_type("StreamingFCLayer_Batch")
+    fc0w = getCustomOp(fc_layers[0])
+    fc1w = getCustomOp(fc_layers[1])
+    fc2w = getCustomOp(fc_layers[2])
+    fc3w = getCustomOp(fc_layers[3])
     fc0w.set_nodeattr("inFIFODepth", 50)
     fc0w.set_nodeattr("SIMD", 8)
     fc0w.set_nodeattr("PE", 16)
@@ -172,10 +170,9 @@ def test_end2end_tfc_w1a2_verify_dataflow_part():
     res_npysim = ret_npysim[out_name]
     # node-by-node rtlsim
     model = model.transform(SetExecMode("rtlsim"))
-    getCustomOp(model.graph.node[0]).set_nodeattr("rtlsim_trace", "default")
-    getCustomOp(model.graph.node[1]).set_nodeattr("rtlsim_trace", "default")
-    getCustomOp(model.graph.node[2]).set_nodeattr("rtlsim_trace", "default")
-    getCustomOp(model.graph.node[3]).set_nodeattr("rtlsim_trace", "default")
+    fc_layers = model.get_nodes_by_op_type("StreamingFCLayer_Batch")
+    for fcl in fc_layers:
+        getCustomOp(fcl).set_nodeattr("rtlsim_trace", "default")
     model.save(build_dir + "/end2end_tfc_w1a2_ipstitch_nodebynode_rtlsim.onnx")
     ret_rtlsim_nodebynode = execute_onnx(model, inp_dict, True)
     res_rtlsim_nodebynode = ret_rtlsim_nodebynode[out_name]
@@ -206,7 +203,8 @@ def test_end2end_tfc_w1a2_verify_all():
     iname = parent_model.graph.input[0].name
     oname = parent_model.graph.output[0].name
     # produce results with npysim
-    sdp_node = getCustomOp(parent_model.graph.node[2])
+    sdp_node = parent_model.get_nodes_by_op_type("StreamingDataflowPartition")[0]
+    sdp_node = getCustomOp(sdp_node)
     sdp_node.set_nodeattr("model", build_dir + "/end2end_tfc_w1a2_ipstitch_npysim.onnx")
     ret_npysim = execute_onnx(parent_model, {iname: x}, True)
     y_npysim = ret_npysim[oname]
@@ -283,7 +281,8 @@ def test_end2end_tfc_w1a2_run_on_pynq():
         if ip == "":
             pytest.skip("PYNQ board IP address not specified")
         # produce results with npysim
-        sdp_node = getCustomOp(parent_model.graph.node[2])
+        sdp_node = parent_model.get_nodes_by_op_type("StreamingDataflowPartition")[0]
+        sdp_node = getCustomOp(sdp_node)
         sdp_node.set_nodeattr("model", build_dir + "/end2end_tfc_w1a2_pynq_deploy.onnx")
         ret = execute_onnx(parent_model, {iname: x}, True)
         y = ret[oname]
diff --git a/tests/end2end/test_end2end_tfc_w2a2.py b/tests/end2end/test_end2end_tfc_w2a2.py
index 84133bb6c..acc01b568 100644
--- a/tests/end2end/test_end2end_tfc_w2a2.py
+++ b/tests/end2end/test_end2end_tfc_w2a2.py
@@ -109,7 +109,8 @@ def test_end2end_tfc_w2a2_create_dataflow_partition():
     model = ModelWrapper(build_dir + "/end2end_tfc_w2a2_hls_layers.onnx")
     parent_model = model.transform(CreateDataflowPartition())
     parent_model.save(build_dir + "/end2end_tfc_w2a2_dataflow_parent.onnx")
-    sdp_node = getCustomOp(parent_model.graph.node[2])
+    sdp_node = parent_model.get_nodes_by_op_type("StreamingDataflowPartition")[0]
+    sdp_node = getCustomOp(sdp_node)
     dataflow_model_filename = sdp_node.get_nodeattr("model")
     dataflow_model = ModelWrapper(dataflow_model_filename)
     dataflow_model.save(build_dir + "/end2end_tfc_w2a2_dataflow_model.onnx")
@@ -117,14 +118,11 @@ def test_end2end_tfc_w2a2_create_dataflow_partition():
 
 def test_end2end_tfc_w2a2_fold_and_tlastmarker():
     model = ModelWrapper(build_dir + "/end2end_tfc_w2a2_dataflow_model.onnx")
-    fc0 = model.graph.node[0]
-    fc1 = model.graph.node[1]
-    fc2 = model.graph.node[2]
-    fc3 = model.graph.node[3]
-    fc0w = getCustomOp(fc0)
-    fc1w = getCustomOp(fc1)
-    fc2w = getCustomOp(fc2)
-    fc3w = getCustomOp(fc3)
+    fc_layers = model.get_nodes_by_op_type("StreamingFCLayer_Batch")
+    fc0w = getCustomOp(fc_layers[0])
+    fc1w = getCustomOp(fc_layers[1])
+    fc2w = getCustomOp(fc_layers[2])
+    fc3w = getCustomOp(fc_layers[3])
     fc0w.set_nodeattr("inFIFODepth", 50)
     fc0w.set_nodeattr("SIMD", 8)
     fc0w.set_nodeattr("PE", 16)
@@ -172,10 +170,9 @@ def test_end2end_tfc_w2a2_verify_dataflow_part():
     res_npysim = ret_npysim[out_name]
     # node-by-node rtlsim
     model = model.transform(SetExecMode("rtlsim"))
-    getCustomOp(model.graph.node[0]).set_nodeattr("rtlsim_trace", "default")
-    getCustomOp(model.graph.node[1]).set_nodeattr("rtlsim_trace", "default")
-    getCustomOp(model.graph.node[2]).set_nodeattr("rtlsim_trace", "default")
-    getCustomOp(model.graph.node[3]).set_nodeattr("rtlsim_trace", "default")
+    fc_layers = model.get_nodes_by_op_type("StreamingFCLayer_Batch")
+    for fcl in fc_layers:
+        getCustomOp(fcl).set_nodeattr("rtlsim_trace", "default")
     model.save(build_dir + "/end2end_tfc_w2a2_ipstitch_nodebynode_rtlsim.onnx")
     ret_rtlsim_nodebynode = execute_onnx(model, inp_dict, True)
     res_rtlsim_nodebynode = ret_rtlsim_nodebynode[out_name]
@@ -206,7 +203,8 @@ def test_end2end_tfc_w2a2_verify_all():
     iname = parent_model.graph.input[0].name
     oname = parent_model.graph.output[0].name
     # produce results with npysim
-    sdp_node = getCustomOp(parent_model.graph.node[2])
+    sdp_node = parent_model.get_nodes_by_op_type("StreamingDataflowPartition")[0]
+    sdp_node = getCustomOp(sdp_node)
     sdp_node.set_nodeattr("model", build_dir + "/end2end_tfc_w2a2_ipstitch_npysim.onnx")
     ret_npysim = execute_onnx(parent_model, {iname: x}, True)
     y_npysim = ret_npysim[oname]
@@ -283,7 +281,8 @@ def test_end2end_tfc_w2a2_run_on_pynq():
         if ip == "":
             pytest.skip("PYNQ board IP address not specified")
         # produce results with npysim
-        sdp_node = getCustomOp(parent_model.graph.node[2])
+        sdp_node = parent_model.get_nodes_by_op_type("StreamingDataflowPartition")[0]
+        sdp_node = getCustomOp(sdp_node)
         sdp_node.set_nodeattr("model", build_dir + "/end2end_tfc_w2a2_pynq_deploy.onnx")
         ret = execute_onnx(parent_model, {iname: x}, True)
         y = ret[oname]
-- 
GitLab