[Feature Enhancement] Optimize decomposition and evaluation implementation to support test_devices. (#392)

* [Feature Enhancement] Implement iterative subgraph decomposition and evaluation pipeline.

* Update paddle decomposer to support group_head_and_tail.

* Support paddle.

* Use system temp directory if dump path is not set.

* Support separate testing of reference and target device.

* Run test_target_device successfully.

* Update to use the original model_path to decompose.

* Remove some changes.

* Fix model_name and refine some codes.

* Minor change to format codes.

* Update commit of Athena.

---------

Co-authored-by: TelGome <[email protected]>

Author: Xreki

graph_net/paddle/extractor.py

@@ -1,11 +1,10 @@
 import os
 import json
-import importlib.util
+import tempfile
 
 import paddle
-from athena.graphnet_samples import GraphnetSample, RunGeneration
+from athena.graphnet_samples import RunGeneration
 from graph_net import imp_util
-from graph_net.paddle import utils
 
 
 def load_class_from_file(file_path: str, class_name: str):
@@ -102,8 +101,8 @@ def __init__(
         self.num_samples_of_all_subgraphs = 0
         self.subgraph_idx2samples = None
 
-        dump_path = os.environ.get("GRAPH_NET_PIR_DUMP_WORKSPACE", "/tmp")
-        self.dump_path = os.path.abspath(dump_path)
+        dump_path = os.environ.get("GRAPH_NET_PIR_DUMP_WORKSPACE", None)
+        self.dump_path = os.path.abspath(dump_path) if dump_path else tempfile.mkdtemp()
 
         workspace_path = (
             workspace_path
@@ -167,7 +166,7 @@ def run_model_with_dump_enabled(self, model_dump_path, **input_dict):
             backend=None,
         )
         static_model.eval()
-        program = static_model.forward.concrete_program.main_program
+        # program = static_model.forward.concrete_program.main_program
         # print(program)
         static_model(**data_dict)
 
@@ -176,7 +175,10 @@ def run_model_with_dump_enabled(self, model_dump_path, **input_dict):
         return static_model
 
     def translate_pir_program_to_sample_codes(
-        self, model_dump_path, split_positions=None
+        self,
+        model_dump_path,
+        split_positions=None,
+        group_head_and_tail=True,
     ):
         ir_programs_path = os.path.join(model_dump_path, "exec_programs.py")
         example_inputs_path = os.path.join(
@@ -201,7 +203,9 @@ def translate_pir_program_to_sample_codes(
             example_inputs=example_inputs_path,
             op_example_inputs=op_example_inputs_path,
             split_positions=split_positions,
+            group_head_and_tail=group_head_and_tail,
             eval_mode=True,
+            tmp_dir=model_dump_path,
         )
 
         self.subgraph_idx2samples = {}

graph_net/paddle/naive_graph_decomposer.py

@@ -56,6 +56,9 @@ def __init__(self, parent_graph_extractor):
             workspace_path=self.parent_graph_extractor.config["output_dir"],
         )
         self.split_positions = self.parent_graph_extractor.config["split_positions"]
+        self.group_head_and_tail = self.parent_graph_extractor.config[
+            "group_head_and_tail"
+        ]
         self.post_process = self.make_post_process(self.parent_graph_extractor.config)
 
     def do_extract(self, **input_dict):
@@ -69,7 +72,9 @@ def do_extract(self, **input_dict):
 
         # 2. Convert pir programs to graphnet samples
         self.builtin_extractor.translate_pir_program_to_sample_codes(
-            model_dump_path, split_positions=self.split_positions
+            model_dump_path,
+            split_positions=self.split_positions,
+            group_head_and_tail=self.group_head_and_tail,
         )
 
         # 3. Save to model_path

graph_net/path_utils.py

@@ -7,7 +7,9 @@ def get_graphnet_root():
 
 
 def is_single_model_dir(model_dir):
-    return os.path.isfile(f"{model_dir}/graph_net.json")
+    return os.path.isfile(f"{model_dir}/graph_net.json") and os.path.isfile(
+        f"{model_dir}/model.py"
+    )
 
 
 def get_recursively_model_path(root_dir):