Implement NaiveDecomposer for paddle.

Author: Xreki

graph_net/paddle/extractor.py

@@ -11,6 +11,34 @@
 from graph_net.paddle import utils
 
 
+# used as configuration of python -m graph_net.paddle.run_model
+class RunModelDecorator:
+    def __init__(self, config):
+        self.config = self.make_config(**config)
+
+    def __call__(self, model):
+        return extract(**self.config)(model)
+
+    def make_config(
+        self,
+        name=None,
+        dynamic=False,
+        input_spec=None,
+        custom_extractor_path: str = None,
+        custom_extractor_config: dict = None,
+    ):
+        assert name is not None
+        return {
+            "name": name,
+            "dynamic": dynamic,
+            "input_spec": input_spec,
+            "extractor_config": {
+                "custom_extractor_path": custom_extractor_path,
+                "custom_extractor_config": custom_extractor_config,
+            },
+        }
+
+
 class GraphExtractor:
     def __init__(
         self,
@@ -26,45 +54,39 @@ def __init__(
         self.input_spec = input_spec
         assert not self.dynamic, "dynamic=True is not supported now!"
 
-        self.subgraph_counter = 0
-        self.dump_path = os.environ.get("GRAPH_NET_PIR_DUMP_WORKSPACE", "/tmp")
-        self.workspace_path = (
+        self.num_subgraphs = 0
+        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)
+
+        workspace_path = (
             workspace_path
             if workspace_path is not None
             else os.environ.get("GRAPH_NET_EXTRACT_WORKSPACE")
         )
+        self.workspace_path = os.path.abspath(workspace_path)
         if not self.workspace_path:
             raise EnvironmentError(
                 "Environment variable 'GRAPH_NET_EXTRACT_WORKSPACE' is not set."
             )
 
     def prepare_to_extract(self, model_dump_path):
         os.makedirs(model_dump_path, exist_ok=True)
-        old_flags = paddle.get_flags(
-            [
-                "FLAGS_logging_trunc_pir_py_code",
-                "FLAGS_logging_pir_py_code_int_tensor_element_limit",
-                "FLAGS_logging_pir_py_code_dir",
-            ]
-        )
+        new_flags = {
+            "FLAGS_logging_trunc_pir_py_code": 1,
+            "FLAGS_logging_pir_py_code_int_tensor_element_limit": 64,
+            "FLAGS_logging_pir_py_code_dir": model_dump_path,
+        }
+        old_flags = paddle.get_flags(list(new_flags.keys()))
+
         print(f"Set pir dumping path to {model_dump_path}")
-        paddle.set_flags(
-            {
-                "FLAGS_logging_trunc_pir_py_code": 1,
-                "FLAGS_logging_pir_py_code_int_tensor_element_limit": 64,
-                "FLAGS_logging_pir_py_code_dir": model_dump_path,
-            }
-        )
+        paddle.set_flags(new_flags)
         return old_flags
 
-    def write_to_file(self, filepath, content):
-        print(f"Write to {filepath}")
-        with open(filepath, "w") as f:
-            f.write(content)
-
-    def __call__(self, **input_dict):
-        # 1. Get model dump path
-        model_dump_path = os.path.join(self.dump_path, self.name)
+    def run_model_with_dump_enabled(self, model_dump_path, **input_dict):
+        # Get model dump path
         old_flags = self.prepare_to_extract(model_dump_path)
 
         if self.input_spec is None:
@@ -74,13 +96,19 @@ def __call__(self, **input_dict):
                 if isinstance(value, paddle.Tensor)
             ]
 
-        # 2. Run the model to dump pir programs
+        # Run the static model
         static_model = paddle.jit.to_static(
             self.model, input_spec=self.input_spec, full_graph=True
         )
         static_model(**input_dict)
 
-        # 3. Convert pir programs to graphnet samples
+        # Restore the environment
+        paddle.set_flags(old_flags)
+        return static_model
+
+    def translate_pir_program_to_sample_codes(
+        self, model_dump_path, split_positions=None
+    ):
         ir_programs_path = os.path.join(model_dump_path, "exec_programs.py")
         example_inputs_path = os.path.join(
             model_dump_path, "programs_example_input_tensor_meta.py"
@@ -92,43 +120,108 @@ def __call__(self, **input_dict):
             example_inputs_path
         ), f"{example_inputs_path} is not a regular file."
 
+        # Arguments for graph decomposer
+        op_example_inputs_path = (
+            os.path.join(model_dump_path, "op_example_input_tensor_meta.py")
+            if split_positions
+            else None
+        )
         graphnet_samples = generate_samples(
             model_name=self.name,
             ir_programs=ir_programs_path,
             example_inputs=example_inputs_path,
+            op_example_inputs=op_example_inputs_path,
+            split_positions=split_positions,
             eval_mode=True,
         )
 
-        # 4. Save to model_path
+        self.subgraph_idx2samples = {}
+        for sample in graphnet_samples:
+            if sample.subgraph_idx not in self.subgraph_idx2samples.keys():
+                self.subgraph_idx2samples[sample.subgraph_idx] = []
+            self.subgraph_idx2samples[sample.subgraph_idx].append(sample)
+
+        self.num_subgraphs = len(self.subgraph_idx2samples)
+        self.num_samples_of_all_subgraphs = len(graphnet_samples)
+        assert self.num_subgraphs > 0
+        return self.subgraph_idx2samples
+
+    def write_sample_to_file(self, subgraph_path, sample):
+        def write_to_file(filepath, content):
+            print(f"Write to {filepath}")
+            with open(filepath, "w") as f:
+                f.write(content)
+
+        if not os.path.exists(subgraph_path):
+            os.makedirs(subgraph_path, exist_ok=True)
+        write_to_file(f"{subgraph_path}/model.py", sample.model)
+        write_to_file(f"{subgraph_path}/weight_meta.py", sample.weight_meta)
+        write_to_file(f"{subgraph_path}/input_meta.py", sample.input_meta)
+        with open(os.path.join(subgraph_path, "graph_net.json"), "w") as f:
+            json.dump(sample.metadata, f, indent=4)
+
+    def __call__(self, **input_dict):
+        # 1. Run the model to dump pir programs
+        model_dump_path = os.path.join(self.dump_path, self.name)
+        static_model = self.run_model_with_dump_enabled(model_dump_path, **input_dict)
+
+        # 2. Convert pir programs to graphnet samples
+        self.translate_pir_program_to_sample_codes(
+            model_dump_path, split_positions=None
+        )
+
+        # 3. Save to model_path
         model_path = os.path.join(self.workspace_path, self.name)
-        self.subgraph_counter = len(graphnet_samples)
-        for i, sample in enumerate(graphnet_samples):
-            subgraph_path = (
-                model_path
-                if self.subgraph_counter == 1
-                else os.path.join(model_path, f"subgraph_{i}")
-            )
-            if not os.path.exists(subgraph_path):
-                os.makedirs(subgraph_path, exist_ok=True)
-            self.write_to_file(f"{subgraph_path}/model.py", sample.model)
-            self.write_to_file(f"{subgraph_path}/weight_meta.py", sample.weight_meta)
-            self.write_to_file(f"{subgraph_path}/input_meta.py", sample.input_meta)
-            with open(os.path.join(subgraph_path, "graph_net.json"), "w") as f:
-                json.dump(sample.metadata, f, indent=4)
+        for subgraph_idx, samples in self.subgraph_idx2samples.items():
+            assert len(samples) == 1
+            if self.num_samples_of_all_subgraphs == 1:
+                subgraph_path = model_path
+            else:
+                subgraph_path = os.path.join(model_path, f"subgraph_{subgraph_idx}")
+            self.write_sample_to_file(subgraph_path, samples[0])
 
         print(
             f"Graph and tensors for '{self.name}' extracted successfully to: {model_path}"
         )
-
-        # 5. Restore the environment
-        paddle.set_flags(old_flags)
         return static_model
 
 
-def extract(name, dynamic=False, input_spec=None):
+def extract(name, dynamic=False, input_spec=None, extractor_config: dict = None):
+    """
+    Extract computation graphs from PaddlePaddle nn.Layer.
+    The extracted computation graphs will be saved into directory of env var $GRAPH_NET_EXTRACT_WORKSPACE.
+
+    Args:
+        name (str): The name of the model, used as the directory name for saving.
+        dynamic (bool): Enable dynamic shape support in paddle.jit.to_static.
+        input_spec (list[InputSpec] | tuple[InputSpec]): InputSpec for input tensors, which includes tensor's name, shape and dtype.
+            When dynamic is False, input_spec can be inferred automatically.
+
+    Returns:
+        wrapper or decorator
+    """
+
+    extractor_config = make_extractor_config(extractor_config)
+
+    def get_graph_extractor_maker():
+        custom_extractor_path = extractor_config["custom_extractor_path"]
+        custom_extractor_config = extractor_config["custom_extractor_config"]
+        if custom_extractor_path is None:
+            return GraphExtractor
+        import importlib.util as imp
+
+        print(f"Import graph_extractor from {custom_extractor_path}")
+        # import custom_extractor_path as graph_extractor
+        spec = imp.spec_from_file_location("graph_extractor", custom_extractor_path)
+        graph_extractor = imp.module_from_spec(spec)
+        spec.loader.exec_module(graph_extractor)
+        cls = graph_extractor.GraphExtractor
+        return lambda *args, **kwargs: cls(custom_extractor_config, *args, **kwargs)
+
     def wrapper(model: paddle.nn.Layer):
         assert isinstance(model, paddle.nn.Layer), f"{type(model)=}"
-        return GraphExtractor(model, name, dynamic, input_spec)
+        extractor = get_graph_extractor_maker()(model, name, dynamic, input_spec)
+        return extractor
 
     def decorator(module_class):
         def constructor(*args, **kwargs):
@@ -147,3 +240,18 @@ def decorator_or_wrapper(obj):
             )
 
     return decorator_or_wrapper
+
+
+def make_extractor_config(extractor_config):
+    kwargs = extractor_config if extractor_config is not None else {}
+    return make_extractor_config_impl(**kwargs)
+
+
+def make_extractor_config_impl(
+    custom_extractor_path: str = None, custom_extractor_config: dict = None
+):
+    config = custom_extractor_config if custom_extractor_config is not None else {}
+    return {
+        "custom_extractor_path": custom_extractor_path,
+        "custom_extractor_config": config,
+    }

graph_net/paddle/naive_graph_decomposer.py

@@ -0,0 +1,120 @@
+import os
+from graph_net.paddle.extractor import GraphExtractor as BuiltinGraphExtractor
+
+
+class GraphExtractor:
+    def __init__(
+        self,
+        config: dict,
+        model,
+        name,
+        dynamic,
+        input_spec=None,
+    ):
+        self.model = model
+        self.name = name
+        self.dynamic = dynamic
+        self.input_spec = input_spec
+        self.config = self.make_config(**config)
+
+    def make_config(
+        self,
+        split_positions=(),
+        group_head_and_tail=False,
+        chain_style=False,
+        output_dir="./tmp/naive_decomposer_dir",
+    ):
+        for pos in split_positions:
+            assert isinstance(
+                pos, int
+            ), f"split_positions should be list of int, {split_positions=}"
+        return {
+            "split_positions": split_positions,
+            "group_head_and_tail": group_head_and_tail,
+            "chain_style": chain_style,
+            "output_dir": output_dir,
+        }
+
+    def __call__(self, **input_dict):
+        extracted_model = self.get_naive_decomposer_extractor()(**input_dict)
+        return extracted_model
+
+    def get_naive_decomposer_extractor(self):
+        return NaiveDecomposerExtractor(self)
+
+
+class NaiveDecomposerExtractor:
+    def __init__(self, parent_graph_extractor):
+        super().__init__()
+        self.parent_graph_extractor = parent_graph_extractor
+        self.extracted = False
+        self.builtin_extractor = BuiltinGraphExtractor(
+            model=parent_graph_extractor.model,
+            name=parent_graph_extractor.name,
+            dynamic=parent_graph_extractor.dynamic,
+            input_spec=parent_graph_extractor.input_spec,
+            workspace_path=self.parent_graph_extractor.config["output_dir"],
+        )
+        self.split_positions = self.parent_graph_extractor.config["split_positions"]
+        self.post_process = self.make_post_process(self.parent_graph_extractor.config)
+
+    def do_extract(self, **input_dict):
+        # 1. Run the model to dump pir programs
+        model_dump_path = os.path.join(
+            self.builtin_extractor.dump_path, self.builtin_extractor.name
+        )
+        static_model = self.builtin_extractor.run_model_with_dump_enabled(
+            model_dump_path, **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
+        )
+
+        # 3. Save to model_path
+        self.subgraph_path_list = []
+        model_path = os.path.join(
+            self.builtin_extractor.workspace_path, self.builtin_extractor.name
+        )
+        for (
+            subgraph_idx,
+            samples,
+        ) in self.builtin_extractor.subgraph_idx2samples.items():
+            for seq_idx in range(len(samples)):
+                if (
+                    self.builtin_extractor.num_samples_of_all_subgraphs == 1
+                    and len(samples) == 1
+                ):
+                    subgraph_path = model_path
+                elif len(samples) == 1:
+                    subgraph_path = os.path.join(model_path, f"subgraph_{subgraph_idx}")
+                else:
+                    subgraph_path = os.path.join(
+                        model_path, f"subgraph_{subgraph_idx}_{seq_idx}"
+                    )
+                self.subgraph_path_list.append(subgraph_path)
+                self.builtin_extractor.write_sample_to_file(
+                    subgraph_path, samples[seq_idx]
+                )
+        print(
+            f"Graph and tensors for '{self.builtin_extractor.name}' extracted successfully to: {model_path}"
+        )
+        return static_model
+
+    def __call__(self, **input_dict):
+        extracted_model = None
+        if not self.extracted:
+            extracted_model = self.do_extract(**input_dict)
+            self.extracted = True
+        # if self.extracted:
+        #    for subgraph_path in self.subgraph_path_list:
+        #        self.post_process(subgraph_path)
+        return extracted_model
+
+    def make_post_process(self, config):
+        return None
+        # if config["post_process_path"] is None:
+        #    return None
+        # module = imp_util.load_module(config["post_process_path"])
+        # return module.PostExtractProcess(config["post_process_config"])