feat: Add single operator decomposition support

graph_net/tools/single_operator_decompose.sh

@@ -0,0 +1,48 @@
+#!/bin/bash
+set -x
+
+GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(os.path.dirname(os.path.dirname(graph_net.__file__)))")
+DECOMPOSE_WORKSPACE=/tmp/workspace_single_operator_decompose
+
+mkdir -p "$DECOMPOSE_WORKSPACE"
+
+model_list="$GRAPH_NET_ROOT/graph_net/config/small10_torch_samples_list.txt"
+
+python3 -m graph_net.apply_sample_pass \
+    --model-path-list $model_list \
+    --sample-pass-file-path $GRAPH_NET_ROOT/graph_net/torch/sample_pass/op_names_extractor.py \
+    --sample-pass-class-name OpNamesExtractor \
+    --sample-pass-config=$(base64 -w 0 <<EOF
+{
+    "resume": true,
+    "model_path_prefix": "$GRAPH_NET_ROOT",
+    "output_dir": "$DECOMPOSE_WORKSPACE"
+}
+EOF
+)
+
+SINGLE_OP_RANGES_WORKSPACE=$DECOMPOSE_WORKSPACE/workspace_single_operator_ranges
+
+python3 -m graph_net.torch.single_sequence_split_points \
+    --model-list "$model_list" \
+    --op-names-path-prefix "$DECOMPOSE_WORKSPACE" \
+    --output-dir "$SINGLE_OP_RANGES_WORKSPACE" \
+    --subgraph-ranges-file-name "subgraph_ranges.json" \
+    --subgraph-ranges-json "$DECOMPOSE_WORKSPACE/subgraph_ranges.json" \
+    --output-json "$DECOMPOSE_WORKSPACE/split_results.json"
+
+python3 -m graph_net.apply_sample_pass \
+    --model-path-list $model_list \
+    --sample-pass-file-path $GRAPH_NET_ROOT/graph_net/torch/sample_pass/subgraph_generator.py \
+    --sample-pass-class-name SubgraphGenerator \
+    --sample-pass-config=$(base64 -w 0 <<EOF
+{
+    "resume": false,
+    "model_path_prefix": "$GRAPH_NET_ROOT",
+    "output_dir": "$DECOMPOSE_WORKSPACE",
+    "subgraph_ranges_json_root": "$SINGLE_OP_RANGES_WORKSPACE",
+    "group_head_and_tail": false,
+    "chain_style": false
+}
+EOF
+)

graph_net/torch/single_sequence_split_points.py

@@ -0,0 +1,114 @@
+import argparse
+import json
+from pathlib import Path
+
+
+def create_single_operator_ranges(
+    op_names_path_prefix,
+    model_list,
+    output_dir,
+    subgraph_ranges_file_name="subgraph_ranges.json",
+    subgraph_ranges_json=None,
+    output_json=None,
+):
+    with open(model_list, "r") as f:
+        rel_model_paths = [
+            line.strip() for line in f if line.strip() and not line.startswith("#")
+        ]
+
+    split_positions_json = {}
+    subgraph_ranges_json_data = {}
+    for rel_model_path in rel_model_paths:
+        txt_path = Path(op_names_path_prefix) / rel_model_path / "op_names.txt"
+        if not txt_path.exists():
+            print(f"File not found: {txt_path}")
+            continue
+
+        with open(txt_path, "r") as f:
+            seq = [line.strip() for line in f if line.strip()]
+
+        if not seq:
+            print(f"Empty sequence in: {txt_path}")
+            continue
+
+        num_ops = len(seq)
+        model_name = Path(rel_model_path).name
+        subgraph_ranges = [[i, i + 1] for i in range(num_ops)]
+        model_output_dir = Path(output_dir) / rel_model_path
+        model_output_dir.mkdir(parents=True, exist_ok=True)
+        output_file = model_output_dir / subgraph_ranges_file_name
+        with open(output_file, "w") as f:
+            json.dump({"subgraph_ranges": subgraph_ranges}, f, indent=4)
+
+        split_positions = sorted(
+            set(pos for start, end in subgraph_ranges for pos in (start, end))
+        )
+        split_positions_json[str(rel_model_path)] = {
+            "model_name": model_name,
+            "split_positions": split_positions,
+            "total_length": num_ops,
+        }
+        subgraph_ranges_json_data[str(rel_model_path)] = {
+            "model_name": model_name,
+            "subgraph_ranges": subgraph_ranges,
+            "total_length": num_ops,
+        }
+        print(
+            f"Created single operator ranges for {model_name}: {len(subgraph_ranges)} subgraphs"
+        )
+
+    if output_json:
+        with open(output_json, "w") as f:
+            json.dump(split_positions_json, f, indent=4)
+
+    if subgraph_ranges_json:
+        with open(subgraph_ranges_json, "w") as f:
+            json.dump(subgraph_ranges_json_data, f, indent=4)
+
+    return split_positions_json, subgraph_ranges_json_data
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Create single operator subgraph ranges"
+    )
+    parser.add_argument(
+        "--model-list", type=str, required=True, help="Path to model list file"
+    )
+    parser.add_argument(
+        "--op-names-path-prefix",
+        type=str,
+        default="./",
+        help="Prefix for op_names.txt files",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        required=True,
+        help="Output directory for subgraph ranges",
+    )
+    parser.add_argument(
+        "--subgraph-ranges-file-name",
+        type=str,
+        default="subgraph_ranges.json",
+        help="Name of subgraph ranges file",
+    )
+    parser.add_argument(
+        "--subgraph-ranges-json", type=str, help="Path to save combined subgraph ranges"
+    )
+    parser.add_argument("--output-json", type=str, help="Path to save analysis results")
+
+    args = parser.parse_args()
+
+    create_single_operator_ranges(
+        args.op_names_path_prefix,
+        args.model_list,
+        args.output_dir,
+        args.subgraph_ranges_file_name,
+        args.subgraph_ranges_json,
+        args.output_json,
+    )
+
+
+if __name__ == "__main__":
+    main()