Merge branch 'develop' into support_meta_restore

Author: Xreki

graph_net/optional.py

@@ -0,0 +1,25 @@
+from typing import TypeVar, Generic, Union
+
+T = TypeVar("T")
+
+
+class Optional(Generic[T]):
+    def __init__(self, value: Union[T, None]):
+        self._value = value
+
+    def reset(self, that):
+        assert isinstance(that, Optional)
+        self._value = that._value
+
+    def is_some(self) -> bool:
+        return self._value is not None
+
+    def unwrap(self) -> T:
+        """Returns the value or raises an error if None."""
+        if self._value is None:
+            raise ValueError("Tried to unwrap a None value!")
+        return self._value
+
+    def unwrap_or(self, default: T) -> T:
+        """Returns the value or a default if None."""
+        return self._value if self._value is not None else default

graph_net/sample_pass/agent_unittest_generator.py

@@ -224,7 +224,7 @@ def __init__(
         self.output_dir = Path(output_dir)
         self.device = self._choose_device(device)
         self.generate_main = generate_main
-        self.try_run = try_run and generate_main
+        self.try_run = try_run
         self.data_input_predicator = self._make_data_input_predicator(
             data_input_predicator_filepath, data_input_predicator_class_name
         )
@@ -244,20 +244,26 @@ def generate(self):
             input_tensor_metas,
             weight_tensor_metas,
         ) = self._get_input_and_weight_tensor_metas(input_arg_names, weight_arg_names)
-        graph_module_desc = GraphModuleDescriptor(
-            device=self.device,
-            generate_main=self.generate_main,
-            model_name=model_name,
-            input_arg_names=input_arg_names,
-            input_tensor_metas=input_tensor_metas,
-            weight_arg_names=weight_arg_names,
-            weight_tensor_metas=weight_tensor_metas,
-            forward_body=self._get_forward_body(
-                graph_module, input_arg_names, weight_arg_names
-            ),
-        )
-        unittest = self._render_template(graph_module_desc)
-        if self._try_to_run_unittest(unittest):
+
+        def _generate_unittest(generate_main):
+            graph_module_desc = GraphModuleDescriptor(
+                device=self.device,
+                generate_main=generate_main,
+                model_name=model_name,
+                input_arg_names=input_arg_names,
+                input_tensor_metas=input_tensor_metas,
+                weight_arg_names=weight_arg_names,
+                weight_tensor_metas=weight_tensor_metas,
+                forward_body=self._get_forward_body(
+                    graph_module, input_arg_names, weight_arg_names
+                ),
+            )
+            return self._render_template(graph_module_desc)
+
+        # Generate unittest with main for try-run.
+        unittest_for_try_run = _generate_unittest(generate_main=self.try_run)
+        if self._try_to_run_unittest(unittest_for_try_run):
+            unittest = _generate_unittest(generate_main=self.generate_main)
             self._write_to_file(unittest, self.output_dir)
 
     def _choose_device(self, device) -> str:

graph_net/sample_pass/fusible_subgraph_ranges_generator.py

@@ -0,0 +1,132 @@
+from graph_net.sample_pass.sample_pass import SamplePass
+from graph_net.sample_pass.resumable_sample_pass_mixin import ResumableSamplePassMixin
+from pathlib import Path
+import json
+from itertools import groupby
+
+
+class FusibleSubgraphRangesGenerator(SamplePass, ResumableSamplePassMixin):
+    def __init__(self, config):
+        super().__init__(config)
+
+    def declare_config(
+        self,
+        model_path_prefix: str,
+        output_dir: str,
+        input_json_file_name: str,
+        resume: bool = False,
+        limits_handled_models: int = None,
+        output_json_file_name: str = "fusible_subgraph_ranges.json",
+    ):
+        pass
+
+    def __call__(self, rel_model_path: str):
+        self.resumable_handle_sample(rel_model_path)
+
+    def sample_handled(self, rel_model_path: str) -> bool:
+        file_name = self.config["output_json_file_name"]
+        return self.naive_sample_handled(rel_model_path, search_file_name=file_name)
+
+    def resume(self, rel_model_path: str):
+        analyzer = self._make_analyzer(rel_model_path)
+        output_obj = analyzer.analyze()
+        self._save_output(rel_model_path, output_obj)
+
+    def _save_output(self, rel_model_path, output_obj):
+        output_json = json.dumps(output_obj, indent=4)
+        output_dir_path = Path(self.config["output_dir"]) / rel_model_path
+        output_dir_path.mkdir(parents=True, exist_ok=True)
+        output_file_path = output_dir_path / self.config["output_json_file_name"]
+        output_file_path.write_text(output_json)
+
+    def _make_analyzer(self, rel_model_path: str):
+        model_path = (
+            Path(self.config["model_path_prefix"])
+            / rel_model_path
+            / self.config["input_json_file_name"]
+        )
+        json_ctx = self._make_json_ctx(model_path)
+        return FusibleSubgraphRangesAnalyzer(
+            num_subgraph_kernels_list=self._get_num_subgraph_kernels_list(json_ctx),
+            num_subgraph_ops_list=self._get_num_subgraph_ops_list(json_ctx),
+            start_offset_in_original_graph=self._get_start_offset_in_original_graph(
+                json_ctx
+            ),
+        )
+
+    def _get_start_offset_in_original_graph(self, json_ctx):
+        return json_ctx["start_offset_in_original_graph"]
+
+    def _get_num_subgraph_kernels_list(self, json_ctx):
+        return json_ctx["num_subgraph_kernels"]
+
+    def _get_num_subgraph_ops_list(self, json_ctx):
+        return json_ctx["num_subgraph_ops"]
+
+    def _make_json_ctx(self, model_path: Path):
+        obj = json.loads(model_path.read_text())
+        assert len(obj["num_subgraph_kernels"]) == len(obj["num_subgraph_ops"])
+        return obj
+
+
+class FusibleSubgraphRangesAnalyzer:
+    def __init__(
+        self,
+        num_subgraph_kernels_list: list[int],
+        num_subgraph_ops_list: list[int],
+        start_offset_in_original_graph: int,
+    ):
+        assert len(num_subgraph_kernels_list) == len(num_subgraph_ops_list)
+        self.num_subgraph_kernels_list = num_subgraph_kernels_list
+        self.num_subgraph_ops_list = num_subgraph_ops_list
+        self.start_offset_in_original_graph = start_offset_in_original_graph
+
+    def analyze(self):
+        num_kernels_and_num_ops_list: list[
+            (int, list[int])
+        ] = self._make_num_kernels_and_num_ops_list()
+        num_kernels_and_num_ops_list = sorted(
+            num_kernels_and_num_ops_list, key=lambda pair: pair[0]
+        )
+        num_ops_lists = [
+            sorted(num_ops_list)
+            for _, num_ops_list in num_kernels_and_num_ops_list
+            if len(set(num_ops_list)) > 1
+        ]
+        fusible_subgraph_ranges = [
+            (start, end)
+            for num_ops_list in num_ops_lists
+            for start in [num_ops_list[0] - 1]
+            for end in [num_ops_list[-1]]
+        ]
+        # sorted by `start`
+        fusible_subgraph_ranges = sorted(
+            fusible_subgraph_ranges, key=lambda pair: pair[0]
+        )
+        # remove shadowed
+        fusible_subgraph_ranges = [
+            fusible_subgraph_ranges[i]
+            for i in range(len(fusible_subgraph_ranges))
+            if i == 0
+            or (fusible_subgraph_ranges[i][0] >= fusible_subgraph_ranges[i - 1][1])
+        ]
+        return fusible_subgraph_ranges
+
+    def _make_num_kernels_and_num_ops_list(self):
+        num_kernels_and_num_ops = zip(
+            self.num_subgraph_kernels_list,
+            self.num_subgraph_ops_list,
+        )
+
+        def get_num_kernels(pair):
+            return pair[0]
+
+        num_kernels_and_num_ops = sorted(num_kernels_and_num_ops, key=get_num_kernels)
+        grouped_num_kernels_and_num_ops = groupby(
+            num_kernels_and_num_ops, key=get_num_kernels
+        )
+        num_kernels_and_num_ops_list = [
+            (num_kernels, [num_ops for _, num_ops in group])
+            for num_kernels, group in grouped_num_kernels_and_num_ops
+        ]
+        return num_kernels_and_num_ops_list

graph_net/sample_pass/resumable_sample_pass_mixin.py

@@ -18,11 +18,15 @@ def declare_config(
     ):
         pass
 
+    @abc.abstractmethod
     def sample_handled(self, rel_model_path: str) -> bool:
+        raise NotImplementedError()
+
+    def naive_sample_handled(self, rel_model_path: str, search_file_name: str) -> bool:
         dst_model_path = Path(self.config["output_dir"]) / rel_model_path
         if not dst_model_path.exists():
             return False
-        num_model_py_files = len(list(dst_model_path.rglob("model.py")))
+        num_model_py_files = len(list(dst_model_path.rglob(search_file_name)))
         assert num_model_py_files <= 1
         return num_model_py_files == 1
 

graph_net/test/cumsum_num_kernels_test.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(os.path.dirname(os.path.dirname(graph_net.__file__)))")
+
+python3 -m graph_net.model_path_handler \
+    --model-path-list "$GRAPH_NET_ROOT/graph_net/test/dev_model_list/cumsum_num_kernels_sample_list.txt" \
+    --handler-config $(base64 -w 0 <<EOF
+{
+    "handler_path": "$GRAPH_NET_ROOT/graph_net/torch/sample_passes/cumsum_num_kernels_generator.py",
+    "handler_class_name": "CumSumNumKernelsGenerator",
+    "handler_config": {
+        "output_json_file_name": "cumsum_num_kernels.json",
+        "resume": false,
+        "model_path_prefix": "$GRAPH_NET_ROOT",
+        "output_dir": "/tmp/cumsum_num_kernels_workspace"
+    }
+}
+EOF
+)

graph_net/test/dev_model_list/cumsum_num_kernels_sample_list.txt

@@ -0,0 +1 @@
+samples/timm/resnet18

graph_net/test/fusible_subgraph_ranges_test.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(os.path.dirname(os.path.dirname(graph_net.__file__)))")
+
+python3 -m graph_net.model_path_handler \
+    --model-path-list "$GRAPH_NET_ROOT/graph_net/test/dev_model_list/cumsum_num_kernels_sample_list.txt" \
+    --handler-config $(base64 -w 0 <<EOF
+{
+    "handler_path": "$GRAPH_NET_ROOT/graph_net/sample_pass/fusible_subgraph_ranges_generator.py",
+    "handler_class_name": "FusibleSubgraphRangesGenerator",
+    "handler_config": {
+        "resume": false,
+        "model_path_prefix": "$GRAPH_NET_ROOT/graph_net/test/workspace_cumsum_num_kernels",
+        "input_json_file_name": "cumsum_num_kernels.json",
+        "output_json_file_name": "fusible_subgraph_ranges.json",
+        "output_dir": "/tmp/workspace_fusible_subgraph_ranges"
+    }
+}
+EOF
+)