Update the process of typical_sequence_decomposer

JewelRoam · JewelRoam · commit 3d24089f667a · 2025-12-08T19:24:01.000+08:00
diff --git a/graph_net/test/typical_sequence_decomposer_test.sh b/graph_net/test/typical_sequence_decomposer_test.sh
@@ -1,60 +1,57 @@
 #!/bin/bash
 
 GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(os.path.dirname(os.path.dirname(graph_net.__file__)))")
-DECOMPOSE_PATH=$GRAPH_NET_ROOT/decompose_workspace
+DECOMPOSE_PATH=/tmp/decompose_workspace
 
 mkdir -p "$DECOMPOSE_PATH"
 
-temp_model_list=$(mktemp)
-cat "$GRAPH_NET_ROOT/graph_net/config/torch_samples_list.txt" > "$temp_model_list"
+model_list="$GRAPH_NET_ROOT/graph_net/config/small100_torch_samples_list.txt"
 
 python3 -m graph_net.torch.typical_sequence_split_points \
-    --model-list "$temp_model_list" \
+    --model-list "$model_list" \
     --device "cuda" \
     --window-size 10 \
+    --fold-policy default \
+    --fold-times 10 \
     --output-json "$DECOMPOSE_PATH/split_results.json"
 
-while IFS= read -r MODEL_PATH_IN_SAMPLES; do
-    if [[ -n "$MODEL_PATH_IN_SAMPLES" ]]; then
-        MODEL_FULL_PATH="$GRAPH_NET_ROOT/$MODEL_PATH_IN_SAMPLES"
-        MODEL_NAME=$(basename "$MODEL_PATH_IN_SAMPLES")
-
-        echo "== Decomposing $MODEL_PATH_IN_SAMPLES. =="
-
-        decomposer_config_json_str=$(cat <<EOF
+decompose_config_json_str=$(cat <<EOF
 {
-    "split_results_path": "$DECOMPOSE_PATH/split_results.json",
-    "workspace_path": "$DECOMPOSE_PATH",
-    "chain_style": true,
-    "target_model_name": "$MODEL_NAME"
+    "handler_path": "$GRAPH_NET_ROOT/graph_net/torch/naive_graph_decomposer.py",
+    "handler_class_name": "RangeDecomposerExtractor",
+    "handler_config": {
+        "model_path_prefix": "$GRAPH_NET_ROOT",
+        "output_dir": "$DECOMPOSE_PATH",
+        "split_results_path": "$DECOMPOSE_PATH/split_results.json",
+        "group_head_and_tail": true,
+        "chain_style": true
+    }
 }
 EOF
-        )
-        DECOMPOSER_CONFIG=$(echo $decomposer_config_json_str | base64 -w 0)
-
-        python3 -m graph_net.torch.test_compiler \
-            --model-path "$MODEL_FULL_PATH" \
-            --compiler range_decomposer \
-            --device cuda \
-            --config="$DECOMPOSER_CONFIG"
-
-        cp -r "$MODEL_FULL_PATH" "$DECOMPOSE_PATH/"
-
-        echo "== Validating $MODEL_PATH_IN_SAMPLES. =="
+)
+DECOMPOSE_CONFIG=$(echo $decompose_config_json_str | base64 -w 0)
 
-        python3 -m graph_net.torch.test_compiler \
-            --model-path "$DECOMPOSE_PATH/$MODEL_NAME" \
-            --compiler range_decomposer_validator \
-            --device cuda > "$DECOMPOSE_PATH/${MODEL_NAME}_validation.log" 2>&1
+python3 -m graph_net.model_path_handler \
+    --model-path-list $model_list \
+    --handler-config=$DECOMPOSE_CONFIG \
+    --use-subprocess
 
-        echo "== Finished processing $MODEL_PATH_IN_SAMPLES. =="
-    fi
-done < $temp_model_list
-
-rm -f "$temp_model_list"
+test_compiler_config_json_str=$(cat <<EOF
+{
+    "decomposed_root": "$DECOMPOSE_PATH"
+}
+EOF
+)
+TEST_COMPILER_CONFIG=$(echo $test_compiler_config_json_str | base64 -w 0)
 
-cat $DECOMPOSE_PATH/*_validation.log >> $DECOMPOSE_PATH/combined.log
+python3 -m graph_net.torch.test_compiler \
+    --allow-list $model_list \
+    --compiler range_decomposer_validator \
+    --device cuda \
+    --config $TEST_COMPILER_CONFIG \
+    --model-path-prefix $GRAPH_NET_ROOT \
+    > "$DECOMPOSE_PATH/validation.log" 2>&1
 
 python3 -m graph_net.plot_ESt \
-    --benchmark-path "$DECOMPOSE_PATH/combined.log" \
+    --benchmark-path "$DECOMPOSE_PATH/validation.log" \
     --output-dir "$DECOMPOSE_PATH"
diff --git a/graph_net/torch/naive_graph_decomposer.py b/graph_net/torch/naive_graph_decomposer.py
@@ -1,12 +1,19 @@
 import os
 import torch
+import json
 from graph_net.torch.decompose_util import convert_to_submodules_graph
 from graph_net.torch.extractor import GraphExtractor as BuiltinGraphExtractor
 import graph_net.imp_util as imp_util
 from graph_net.torch.fx_graph_module_util import get_torch_module_and_inputs
 from graph_net.torch.fx_graph_parse_util import parse_sole_graph_module
 
 
+def load_json(file_path):
+    with open(file_path, "r", encoding="utf-8") as file:
+        data_dict = json.load(file)
+    return data_dict
+
+
 class GraphExtractor:
     """
     Used by graph_net.torch.run_model
@@ -151,6 +158,83 @@ def fn(submodule, seq_no):
         return fn
 
 
+class RangeDecomposerExtractor:
+    """
+    Used by graph_net.model_path_handler
+    """
+
+    def __init__(self, config: dict = None):
+        if config is None:
+            config = {}
+        self.config = self._make_config(**config)
+
+    def _make_config(
+        self,
+        split_results_path=None,
+        group_head_and_tail=False,
+        chain_style=False,
+        output_dir="./tmp/naive_decomposer_dir",
+        filter_path=None,
+        filter_config=None,
+        post_extract_process_path=None,
+        post_extract_process_class_name=None,
+        post_extract_process_config=None,
+        model_path_prefix="",
+        **kwargs,
+    ):
+        if os.path.isfile(split_results_path) and split_results_path.endswith(".json"):
+            pass
+        else:
+            raise ValueError(
+                f"split_results_path should be a valid JSON file path, but got {split_results_path=}"
+            )
+        if post_extract_process_config is None:
+            post_extract_process_config = {}
+        return {
+            "split_results_path": split_results_path,
+            "group_head_and_tail": group_head_and_tail,
+            "chain_style": chain_style,
+            "output_dir": output_dir,
+            "filter_path": filter_path,
+            "filter_config": filter_config if filter_config is not None else {},
+            "post_extract_process_path": post_extract_process_path,
+            "post_extract_process_class_name": post_extract_process_class_name,
+            "post_extract_process_config": post_extract_process_config,
+            "model_path_prefix": model_path_prefix,
+        }
+
+    def __call__(self, rel_model_path):
+        model_path = os.path.join(self.config["model_path_prefix"], rel_model_path)
+        split_results = load_json(self.config["split_results_path"])
+        split_positions = split_results[os.path.basename(rel_model_path)][
+            "split_points"
+        ]
+        config = {
+            "split_positions": split_positions,
+            "group_head_and_tail": self.config.get("group_head_and_tail", False),
+            "chain_style": self.config.get("chain_style", False),
+        }
+        module, inputs = get_torch_module_and_inputs(model_path)
+        gm = parse_sole_graph_module(module, inputs)
+        rewrited_gm: torch.fx.GraphModule = convert_to_submodules_graph(
+            gm,
+            submodule_hook=self.get_naive_decomposer_extractor(model_path),
+            **config,
+        )
+        rewrited_gm(*inputs)
+
+    def get_naive_decomposer_extractor(self, model_path):
+        def fn(submodule, seq_no):
+            return NaiveDecomposerExtractorModule(
+                config=self.config,
+                parent_graph_name=os.path.basename(model_path),
+                submodule=submodule,
+                seq_no=seq_no,
+            )
+
+        return fn
+
+
 class NaiveDecomposerExtractorModule(torch.nn.Module):
     def __init__(
         self,
diff --git a/graph_net/torch/typical_sequence_split_points.py b/graph_net/torch/typical_sequence_split_points.py
@@ -82,8 +82,12 @@ def get_input_dict(self, model_path: str, device: str) -> Dict[str, torch.Tensor
 
 
 class SplitAnalyzer:
-    def __init__(self, window_size: int = 10):
+    def __init__(
+        self, window_size: int = 10, fold_policy: str = "default", fold_times: int = 0
+    ):
         self.window_size = window_size
+        self.fold_policy = fold_policy
+        self.fold_times = fold_times
 
     def _resolve_token_to_ops(
         self, tid, num_primitives, token_id2primitive_id, symbol_map
@@ -169,7 +173,9 @@ def analyze(self, model_paths_file: str, device: str) -> Dict[str, Dict]:
             return {}
 
         rp_parser = RpExprParser(
-            window_size=self.window_size, fold_policy="default", fold_times=0
+            window_size=self.window_size,
+            fold_policy=self.fold_policy,
+            fold_times=self.fold_times,
         )
         rp_expr, token_id2primitive_id = rp_parser(inputs_seqs)
         rp_expr.try_unwrap_body_of_sole_symbol_token()
@@ -253,7 +259,11 @@ def _print_analysis(self, name, path, splits, total_len, full_ops):
 
 
 def main(args):
-    analyzer = SplitAnalyzer(window_size=args.window_size)
+    analyzer = SplitAnalyzer(
+        window_size=args.window_size,
+        fold_policy=args.fold_policy,
+        fold_times=args.fold_times,
+    )
     results = analyzer.analyze(args.model_list, args.device)
     if args.output_json:
         with open(args.output_json, "w") as f:
@@ -279,6 +289,18 @@ def main(args):
     parser.add_argument(
         "--window-size", type=int, default=10, help="Window size for RP Parser."
     )
+    parser.add_argument(
+        "--fold-policy",
+        type=str,
+        default="default",
+        help="Policy for split analysis, one of 'default' or 'longest'",
+    )
+    parser.add_argument(
+        "--fold-times",
+        type=int,
+        default=0,
+        help="How many times to fold tokens. If 0, then no folding is done.",
+    )
     parser.add_argument(
         "--output-json",
         type=str,
