support --subgraph-ranges-json for typical_sequence_split_points.py

Author: lixinqi

graph_net/tools/typical_sequence_decompose.sh

@@ -24,13 +24,15 @@ EOF
 )
 
 python3 -m graph_net.torch.typical_sequence_split_points \
-    --enable-resume \
     --model-list "$model_list" \
     --op-names-path-prefix "$DECOMPOSE_WORKSPACE" \
     --device "cuda" \
     --window-size 10 \
     --fold-policy default \
     --fold-times 10 \
+    --min-seq-ops 4 \
+    --max-seq-ops 16 \
+    --subgraph-ranges-json "$DECOMPOSE_WORKSPACE/subgraph_ranges.json" \
     --output-json "$DECOMPOSE_WORKSPACE/split_results.json"
 
 python3 -m graph_net.model_path_handler \
@@ -40,10 +42,11 @@ python3 -m graph_net.model_path_handler \
     "handler_path": "$GRAPH_NET_ROOT/graph_net/torch/graph_decomposer.py",
     "handler_class_name": "RangeDecomposerExtractor",
     "handler_config": {
-        "resume": true,
+        "resume": false,
         "model_path_prefix": "$GRAPH_NET_ROOT",
         "output_dir": "$DECOMPOSE_WORKSPACE",
         "split_results_path": "$DECOMPOSE_WORKSPACE/split_results.json",
+        "subgraph_ranges_path": "$DECOMPOSE_WORKSPACE/subgraph_ranges.json",
         "group_head_and_tail": true,
         "chain_style": false
     }

graph_net/torch/decompose_util.py

@@ -8,6 +8,7 @@
 def convert_to_submodules_graph(
     gm: torch.fx.GraphModule,
     split_positions: list[int],
+    subgraph_ranges: list[(int, int)] = None,
     submodule_hook=None,
     submodule_name_prefix="extracted_submodule",
     chain_style=False,
@@ -26,21 +27,42 @@ def convert_to_submodules_graph(
             "output",
         }
     ]
-    split_positions = (
-        [0, *split_positions, len(submodules_body_nodes)]
-        if group_head_and_tail
-        else split_positions
+
+    def get_range_idx2range_by_split_positions():
+        nonlocal split_positions
+        split_positions = (
+            [0, *split_positions, len(submodules_body_nodes)]
+            if group_head_and_tail
+            else split_positions
+        )
+        split_positions = [
+            max(0, min(pos, len(submodules_body_nodes))) for pos in split_positions
+        ]
+        return [
+            (start, end)
+            for i in range(len(split_positions) - 1)
+            for start in [split_positions[i]]
+            for end in [split_positions[i + 1]]
+            if end > start
+        ]
+
+    def get_range_idx2range_by_subgraph_ranges():
+        assert submodules_body_nodes is not None
+        num_nodes = len(submodules_body_nodes)
+        for i in range(len(subgraph_ranges)):
+            start, end = subgraph_ranges[i]
+            assert start >= 0
+            assert start < end
+            assert end <= num_nodes
+            # check disjoint
+            assert i == 0 or start >= subgraph_ranges[i - 1][1], f"{i=}"
+        return subgraph_ranges
+
+    range_idx2range = (
+        get_range_idx2range_by_split_positions()
+        if (chain_style or submodules_body_nodes is None)
+        else get_range_idx2range_by_subgraph_ranges()
     )
-    split_positions = [
-        max(0, min(pos, len(submodules_body_nodes))) for pos in split_positions
-    ]
-    range_idx2range = [
-        (start, end)
-        for i in range(len(split_positions) - 1)
-        for start in [split_positions[i]]
-        for end in [split_positions[i + 1]]
-        if end > start
-    ]
     range_idx2submodule_body_nodes = [
         submodules_body_nodes[start:end] for start, end in range_idx2range
     ]

graph_net/torch/graph_decomposer.py

@@ -170,6 +170,7 @@ def _make_config(
         self,
         resume: bool = False,
         split_results_path=None,
+        subgraph_ranges_path=None,
         group_head_and_tail=False,
         chain_style=False,
         output_dir="./tmp/naive_decomposer_dir",
@@ -180,13 +181,18 @@ def _make_config(
     ):
         if os.path.isfile(split_results_path) and split_results_path.endswith(".json"):
             pass
+        elif os.path.isfile(subgraph_ranges_path) and subgraph_ranges_path.endswith(
+            ".json"
+        ):
+            pass
         else:
             raise ValueError(
                 f"split_results_path should be a valid JSON file path, but got {split_results_path=}"
             )
         return {
             "resume": resume,
             "split_results_path": split_results_path,
+            "subgraph_ranges_path": subgraph_ranges_path,
             "group_head_and_tail": group_head_and_tail,
             "chain_style": chain_style,
             "output_dir": output_dir,
@@ -195,8 +201,19 @@ def _make_config(
             "model_path_prefix": model_path_prefix,
         }
 
-    def _is_model_handled(self, rel_model_path, split_positions):
-        num_subgraphs = len(split_positions) + 1
+    def _is_model_handled(self, rel_model_path, split_positions, subgraph_ranges):
+        if self.config["chain_style"]:
+            return self._has_enough_subgraphs(
+                rel_model_path,
+                num_subgraphs=len(split_positions) + 1,
+            )
+        else:
+            return self._has_enough_subgraphs(
+                rel_model_path,
+                num_subgraphs=len(subgraph_ranges),
+            )
+
+    def _has_enough_subgraphs(self, rel_model_path, num_subgraphs):
         decomposed_model_path = Path(self.config["output_dir"]) / rel_model_path
         num_decomposed = len(list(decomposed_model_path.rglob("model.py")))
         if num_decomposed > 0 and num_subgraphs != num_decomposed:
@@ -209,21 +226,26 @@ def __call__(self, rel_model_path):
             self.config["output_dir"]
         )
         model_path = os.path.join(self.config["model_path_prefix"], rel_model_path)
-        split_results = load_json(self.config["split_results_path"])
+        split_positions_json = load_json(self.config["split_results_path"])
+        subgraph_ranges_json = load_json(self.config["subgraph_ranges_path"])
         if (
-            split_results[rel_model_path]["split_positions"] is None
-            or len(split_results[rel_model_path]["split_positions"]) == 0
+            split_positions_json[rel_model_path]["split_positions"] is None
+            or len(split_positions_json[rel_model_path]["split_positions"]) == 0
+            or subgraph_ranges_json[rel_model_path]["subgraph_ranges"] is None
+            or len(subgraph_ranges_json[rel_model_path]["subgraph_ranges"]) == 0
         ):
             sys.stderr.write(f"Error: {rel_model_path} has no split positions.\n")
             return
-        split_positions = split_results[rel_model_path]["split_positions"]
+        split_positions = split_positions_json[rel_model_path]["split_positions"]
+        subgraph_ranges = subgraph_ranges_json[rel_model_path]["subgraph_ranges"]
         if self.config["resume"] and self._is_model_handled(
-            rel_model_path, split_positions
+            rel_model_path, split_positions, subgraph_ranges
         ):
             return
         torch.cuda.empty_cache()
         config = {
             "split_positions": split_positions,
+            "subgraph_ranges": subgraph_ranges,
             "group_head_and_tail": self.config.get("group_head_and_tail", False),
             "chain_style": self.config.get("chain_style", False),
         }

graph_net/torch/typical_sequence_split_points.py

@@ -201,7 +201,8 @@ def get_debug_sprintf():
         symbol_map = dict(zip(rp_expr.symbol_token_ids, rp_expr.symbol_token_tensors))
         token2len = self._calculate_token_lengths(rp_expr, num_primitives, symbol_map)
 
-        results = {}
+        split_positions_json = {}
+        subgraph_ranges_json = {}
 
         for i, (model_name, original_path) in enumerate(valid_models):
             if i >= len(rp_expr.body_rp_expr):
@@ -222,12 +223,14 @@ def get_debug_sprintf():
 
             total_len = sum(token2len.get(t, 1) for t in seq_tokens)
 
+            subgraph_ranges = list(
+                tree.FilterSubTreeRangeBySize(self.min_seq_ops, self.max_seq_ops)
+            )
+
             sorted_splits = sorted(
                 set(
                     split_pos
-                    for start, end in tree.FilterSubTreeRangeBySize(
-                        self.min_seq_ops, self.max_seq_ops
-                    )
+                    for start, end in subgraph_ranges
                     for split_pos in (start, end)
                     if end - start > 1
                 )
@@ -237,13 +240,31 @@ def get_debug_sprintf():
                 model_name, str(original_path), sorted_splits, total_len, full_model_ops
             )
 
-            results[str(original_path)] = {
+            split_positions_json[str(original_path)] = {
                 "model_name": model_name,
                 "split_positions": sorted_splits,
                 "total_length": total_len,
             }
 
-        return results
+            sorted_subgraph_ranges = sorted(
+                set((start, end) for start, end in subgraph_ranges if end - start > 1)
+            )
+
+            # make sorted_subgraph_ranges is a disjoint set
+            sorted_subgraph_ranges = [
+                sorted_subgraph_ranges[i]
+                for i in range(len(sorted_subgraph_ranges))
+                if i == 0
+                or sorted_subgraph_ranges[i][0] >= sorted_subgraph_ranges[i - 1][1]
+            ]
+
+            subgraph_ranges_json[str(original_path)] = {
+                "model_name": model_name,
+                "subgraph_ranges": sorted_subgraph_ranges,
+                "total_length": total_len,
+            }
+
+        return split_positions_json, subgraph_ranges_json
 
     def _print_analysis(self, name, path, splits, total_len, full_ops):
         print("=" * 60)
@@ -292,10 +313,16 @@ def main(args):
         min_seq_ops=args.min_seq_ops,
         max_seq_ops=args.max_seq_ops,
     )
-    results = analyzer.analyze(args.op_names_path_prefix, args.model_list, args.device)
+    split_positions_json, subgraph_ranges_json = analyzer.analyze(
+        args.op_names_path_prefix, args.model_list, args.device
+    )
     if args.output_json:
         with open(args.output_json, "w") as f:
-            json.dump(results, f, indent=4)
+            json.dump(split_positions_json, f, indent=4)
+    print(f"{args.subgraph_ranges_json=}")
+    if args.subgraph_ranges_json:
+        with open(args.subgraph_ranges_json, "w") as f:
+            json.dump(subgraph_ranges_json, f, indent=4)
 
 
 if __name__ == "__main__":
@@ -335,6 +362,12 @@ def main(args):
         default=0,
         help="How many times to fold tokens. If 0, then no folding is done.",
     )
+    parser.add_argument(
+        "--subgraph-ranges-json",
+        type=str,
+        default="subgraph_ranges.json",
+        help="Path to save the subgraph ranges in JSON format.",
+    )
     parser.add_argument(
         "--output-json",
         type=str,