fix output node order bug in torch/decompose_util.py

Author: lixinqi

graph_net/bash_templates/model_path_handler_sh.txt

@@ -1,7 +1,10 @@
 #!/bin/bash
 
 GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(os.path.dirname(os.path.dirname(graph_net.__file__)))")
-model_path_handler_config_json_str=$(cat <<EOF
+
+python3 -m graph_net.model_path_handler \
+    --model-path-list "customize_your_model_path_list" \
+    --handler-config $(base64 -w 0 <<EOF
 {
     "handler_path": "$GRAPH_NET_ROOT/graph_net/customize_your_sample_pass.py",
     "handler_class_name": "customize_your_class_name",
@@ -13,14 +16,3 @@ model_path_handler_config_json_str=$(cat <<EOF
 }
 EOF
 )
-
-model_path_handler_model_path_list="customize_your_model_path_list"
-MODEL_PATH_HANDLER_CONFIG=$(echo $model_path_handler_config_json_str | base64 -w 0)
-
-python3 -m graph_net.model_path_handler \
-    --model-path-list $model_path_handler_model_path_list \
-    --handler-config $MODEL_PATH_HANDLER_CONFIG \
-
-unset model_path_handler_model_path_list
-unset MODEL_PATH_HANDLER_CONFIG
-

graph_net/test/typical_sequence_decomposer_test.sh

@@ -1,4 +1,5 @@
 #!/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_PATH=/tmp/decompose_workspace
@@ -9,7 +10,9 @@ mkdir -p "$DECOMPOSE_PATH"
 # model_list="$GRAPH_NET_ROOT/graph_net/config/small100_torch_samples_list.txt"
 model_list="$GRAPH_NET_ROOT/graph_net/test/dev_model_list/validation_error_model_list.txt"
 
-op_names_extractor_config_json_str=$(cat <<EOF
+python3 -m graph_net.model_path_handler \
+    --model-path-list $model_list \
+    --handler-config=$(base64 -w 0 <<EOF
 {
     "handler_path": "$GRAPH_NET_ROOT/graph_net/torch/typical_sequence_split_points.py",
     "handler_class_name": "OpNamesExtractor",
@@ -21,11 +24,6 @@ op_names_extractor_config_json_str=$(cat <<EOF
 }
 EOF
 )
-OP_NAMES_EXTRACTOR_CONFIG=$(echo $op_names_extractor_config_json_str | base64 -w 0)
-
-python3 -m graph_net.model_path_handler \
-    --model-path-list $model_list \
-    --handler-config=$OP_NAMES_EXTRACTOR_CONFIG \
 
 python3 -m graph_net.torch.typical_sequence_split_points \
     --enable-resume \
@@ -37,7 +35,9 @@ python3 -m graph_net.torch.typical_sequence_split_points \
     --fold-times 10 \
     --output-json "$DECOMPOSE_PATH/split_results.json"
 
-decompose_config_json_str=$(cat <<EOF
+python3 -m graph_net.model_path_handler \
+    --model-path-list $model_list \
+    --handler-config=$(base64 -w 0 <<EOF
 {
     "handler_path": "$GRAPH_NET_ROOT/graph_net/torch/graph_decomposer.py",
     "handler_class_name": "RangeDecomposerExtractor",
@@ -52,27 +52,46 @@ decompose_config_json_str=$(cat <<EOF
 }
 EOF
 )
-DECOMPOSE_CONFIG=$(echo $decompose_config_json_str | base64 -w 0)
 
-python3 -m graph_net.model_path_handler \
-    --model-path-list $model_list \
-    --handler-config=$DECOMPOSE_CONFIG \
+device_rewrite_sample_list=$DECOMPOSE_PATH/device_rewrite_sample_list.txt
+cat $model_list \
+    | grep -v '# ' \
+    | xargs -I {} find $DECOMPOSE_PATH/{} -name "model.py" \
+    | xargs dirname \
+    | xargs realpath --relative-to=$DECOMPOSE_PATH \
+    | tee $device_rewrite_sample_list
 
-test_compiler_config_json_str=$(cat <<EOF
+DEVICE_REWRITE_WORKSPACE=$DECOMPOSE_PATH/device_rewrite
+
+python3 -m graph_net.model_path_handler \
+    --model-path-list $device_rewrite_sample_list \
+    --handler-config $(base64 -w 0 <<EOF
 {
-    "model_path_prefix": "$GRAPH_NET_ROOT",
-    "decomposed_root": "$DECOMPOSE_PATH"
+    "handler_path": "$GRAPH_NET_ROOT/graph_net/torch/sample_passes/device_rewrite_sample_pass.py",
+    "handler_class_name": "DeviceRewriteSamplePass",
+    "handler_config": {
+        "device": "cuda",
+        "resume": false,
+        "model_path_prefix": "$DECOMPOSE_PATH",
+        "output_dir": "$DEVICE_REWRITE_WORKSPACE"
+    }
 }
 EOF
 )
-TEST_COMPILER_CONFIG=$(echo $test_compiler_config_json_str | base64 -w 0)
+
 
 python3 -m graph_net.torch.test_compiler \
+    --model-path-prefix $GRAPH_NET_ROOT \
     --allow-list $model_list \
     --compiler range_decomposer_validator \
     --device cuda \
-    --config $TEST_COMPILER_CONFIG \
-    --model-path-prefix $GRAPH_NET_ROOT \
+    --config $(base64 -w 0 <<EOF
+{
+    "model_path_prefix": "$GRAPH_NET_ROOT",
+    "decomposed_root": "$DEVICE_REWRITE_WORKSPACE"
+}
+EOF
+) \
     2>&1 | tee "$DECOMPOSE_PATH/validation.log"
 
 python3 -m graph_net.plot_ESt \

graph_net/torch/decompose_util.py

@@ -237,6 +237,7 @@ def _get_minimal_submodule_inputs_and_outputs(
         defaultdict(int),
     )
     node_list = list(gm.graph.nodes)
+    assert end_node_idx <= len(node_list)
 
     def get_args_node(arg):
         if isinstance(arg, torch.fx.Node):
@@ -274,17 +275,44 @@ def get_args_node_and_self_node(node):
         for related_node in get_args_node_and_self_node(node):
             count_ctx.node2after_output[related_node] += 1
 
-    input_nodes = [
-        node
-        for node in node_list
-        if count_ctx.node2before_input[node] > 0
-        if count_ctx.node2body[node] > 0
-    ]
-    output_nodes = [
-        node
-        for node in node_list
-        if not (count_ctx.node2before_input[node] > 0)
-        if count_ctx.node2body[node] > 0
-        if count_ctx.node2after_output[node] > 0
-    ]
+    if _is_node_idx_out_of_range(gm, start_node_idx):
+        input_nodes = list(_get_return_nodes(gm))
+    else:
+        input_nodes = [
+            node
+            for node in node_list
+            if count_ctx.node2before_input[node] > 0
+            if count_ctx.node2body[node] > 0
+        ]
+    if _is_node_idx_out_of_range(gm, end_node_idx):
+        output_nodes = list(_get_return_nodes(gm))
+    else:
+        output_nodes = [
+            node
+            for node in node_list
+            if not (count_ctx.node2before_input[node] > 0)
+            if count_ctx.node2body[node] > 0
+            if count_ctx.node2after_output[node] > 0
+        ]
     return input_nodes, output_nodes
+
+
+def _get_return_nodes(gm):
+    for node in gm.graph.nodes:
+        if node.op != "output":
+            continue
+        for arg in node.args:
+            if isinstance(arg, (tuple, list)):
+                yield from arg
+            else:
+                yield arg
+
+
+def _is_node_idx_out_of_range(gm, node_idx: int):
+    node_list = list(gm.graph.nodes)
+    num_nodes = len(node_list)
+    if node_idx < 0:
+        return True
+    if node_idx >= num_nodes:
+        return True
+    return node_list[node_idx].op in {"output", "placeholder"}