add sample_pass/fusible_subgraph_ranges_generator.py

Author: lixinqi

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/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
+)

graph_net/test/workspace_cumsum_num_kernels/samples/timm/resnet18/cumsum_num_kernels.json

@@ -0,0 +1,145 @@
+{
+    "start_offset_in_original_graph": 0,
+    "num_subgraph_kernels": [
+        4,
+        4,
+        4,
+        5,
+        11,
+        11,
+        11,
+        16,
+        16,
+        15,
+        15,
+        21,
+        21,
+        21,
+        26,
+        26,
+        25,
+        25,
+        31,
+        31,
+        31,
+        36,
+        36,
+        37,
+        37,
+        36,
+        37,
+        42,
+        42,
+        42,
+        47,
+        47,
+        46,
+        46,
+        51,
+        51,
+        51,
+        56,
+        56,
+        57,
+        57,
+        56,
+        57,
+        62,
+        62,
+        62,
+        67,
+        67,
+        66,
+        66,
+        71,
+        71,
+        71,
+        76,
+        76,
+        79,
+        79,
+        78,
+        79,
+        84,
+        84,
+        84,
+        89,
+        89,
+        88,
+        88,
+        88,
+        88,
+        89
+    ],
+    "num_subgraph_ops": [
+        1,
+        2,
+        3,
+        4,
+        5,
+        6,
+        7,
+        8,
+        9,
+        10,
+        11,
+        12,
+        13,
+        14,
+        15,
+        16,
+        17,
+        18,
+        19,
+        20,
+        21,
+        22,
+        23,
+        24,
+        25,
+        26,
+        27,
+        28,
+        29,
+        30,
+        31,
+        32,
+        33,
+        34,
+        35,
+        36,
+        37,
+        38,
+        39,
+        40,
+        41,
+        42,
+        43,
+        44,
+        45,
+        46,
+        47,
+        48,
+        49,
+        50,
+        51,
+        52,
+        53,
+        54,
+        55,
+        56,
+        57,
+        58,
+        59,
+        60,
+        61,
+        62,
+        63,
+        64,
+        65,
+        66,
+        67,
+        68,
+        69
+    ]
+}

graph_net/torch/sample_passes/cumsum_num_kernels_generator.py

@@ -24,7 +24,7 @@ def declare_config(
         model_path_prefix: str,
         output_dir: str,
         resume: bool = False,
-        start_offset_in_origin_graph: int = 0,
+        start_offset_in_original_graph: int = 0,
         limits_handled_models: int = None,
         output_json_file_name: str = "cumsum_num_kernels.json",
     ):
@@ -39,8 +39,8 @@ def sample_handled(self, rel_model_path: str) -> bool:
 
     def resume(self, rel_model_path: str):
         model_path = Path(self.config["model_path_prefix"]) / rel_model_path
-        start_offset_in_origin_graph = self.config["start_offset_in_origin_graph"]
-        analyzer = CumsumNumKernelsAnalyzer(model_path, start_offset_in_origin_graph)
+        start_offset_in_original_graph = self.config["start_offset_in_original_graph"]
+        analyzer = CumsumNumKernelsAnalyzer(model_path, start_offset_in_original_graph)
         cumsum_num_kernels = analyzer.analyze()
         cumsum_num_kernels_json = json.dumps(cumsum_num_kernels, indent=4)
         output_dir_path = Path(self.config["output_dir"]) / rel_model_path
@@ -50,15 +50,17 @@ def resume(self, rel_model_path: str):
 
 
 class CumsumNumKernelsAnalyzer:
-    def __init__(self, model_path: Path, start_offset_in_origin_graph: int):
+    def __init__(self, model_path: Path, start_offset_in_original_graph: int):
         self.model_path = model_path
-        self.start_offset_in_origin_graph = start_offset_in_origin_graph
+        self.start_offset_in_original_graph = start_offset_in_original_graph
 
     def analyze(self):
         triples = list(self._get_cumsum_num_kernels())
         data = {
-            "num_kernels": [num_kernels for start, end, num_kernels in triples],
-            "start_offset_in_origin_graph": self.start_offset_in_origin_graph,
+            "start_offset_in_original_graph": self.start_offset_in_original_graph,
+            "num_subgraph_kernels": [
+                num_kernels for start, end, num_kernels in triples
+            ],
             "num_subgraph_ops": [end for start, end, num_kernels in triples],
         }
         return data