diff --git a/graph_net/tools/generate_subgraph_dataset.sh b/graph_net/tools/generate_subgraph_dataset.sh index b729d5453..78e2b872c 100755 --- a/graph_net/tools/generate_subgraph_dataset.sh +++ b/graph_net/tools/generate_subgraph_dataset.sh @@ -32,19 +32,22 @@ DIM_GENERALIZED_FUSIBLE_SUBGRAPH_DIR=$DECOMPOSE_WORKSPACE/11_dimension_generaliz RENAMED_DIM_GENERALIZED_FUSIBLE_SUBGRAPH_DIR=$DECOMPOSE_WORKSPACE/12_renamed_dimension_generalized_fusible_subgraphs DEDUP_DIM_GENERALIZED_FUSIBLE_SUBGRAPH_DIR=$DECOMPOSE_WORKSPACE/13_deduplicated_dimension_generalized_fusible_subgraphs DTYPE_GENERALIZED_FUSIBLE_SUBGRAPH_DIR=$DECOMPOSE_WORKSPACE/14_dtype_generalized_fusible_subgraphs -# FUSIBLE_SUBGRAPH_UNITTEST_DIR=$DECOMPOSE_WORKSPACE/15_fusible_subgraphs_unittests +BACKWARD_GRAPH_OUTPUT_DIR=$DECOMPOSE_WORKSPACE/15_backward_graph_extracted +FUSIBLE_SUBGRAPH_UNITTEST_DIR=$DECOMPOSE_WORKSPACE/16_fusible_subgraphs_unittests # typical_subgraphs DIM_GENERALIZED_TYPICAL_SUBGRAPH_DIR=$DECOMPOSE_WORKSPACE/2-08_dimension_generalized_typical_subgraphs RENAMED_DIM_GENERALIZED_TYPICAL_SUBGRAPH_DIR=$DECOMPOSE_WORKSPACE/2-09_renamed_dimension_generalized_typical_subgraphs DEDUP_DIM_GENERALIZED_TYPICAL_SUBGRAPH_DIR=$DECOMPOSE_WORKSPACE/2-10_deduplicated_dimension_generalized_typical_subgraphs DTYPE_GENERALIZED_TYPICAL_SUBGRAPH_DIR=$DECOMPOSE_WORKSPACE/2-11_dtype_generalized_typical_subgraphs -# TYPICAL_SUBGRAPH_UNITTEST_DIR=$DECOMPOSE_WORKSPACE/2-12_typical_kernelbench_unittests +BACKWARD_GRAPH_TYPICAL_OUTPUT_DIR=$DECOMPOSE_WORKSPACE/2-13_backward_graph_extracted +TYPICAL_SUBGRAPH_UNITTEST_DIR=$DECOMPOSE_WORKSPACE/2-14_typical_kernelbench_unittests mkdir -p $DECOMPOSE_WORKSPACE mkdir -p $OUTPUT_DIR -model_list="$GRAPH_NET_ROOT/graph_net/config/small100_torch_samples_list.txt" +model_list="$GRAPH_NET_ROOT/graph_net/config/torch_samples_list.txt" +DB_PATH=$DECOMPOSE_WORKSPACE/small100_torch_samples.db device_rewrited_sample_list=${DECOMPOSE_WORKSPACE}/device_rewrited_sample_list.txt range_decomposed_subgraph_list=${DECOMPOSE_WORKSPACE}/range_decomposed_subgraph_sample_list.txt @@ -405,8 +408,28 @@ EOF ) } +function backward_graph_extractor() { + echo ">>> [13] Extract backward graphs for samples under ${DTYPE_GENERALIZED_FUSIBLE_SUBGRAPH_DIR}." + echo ">>>" + python3 -m graph_net.model_path_handler \ + --model-path-list ${dtype_generalized_subgraphs_list} \ + --handler-config=$(base64 -w 0 <>> [13] Generate unittests for subgraph samples under ${DTYPE_GENERALIZED_FUSIBLE_SUBGRAPH_DIR}. " + echo ">>> [14] Generate unittests for subgraph samples under ${BACKWARD_GRAPH_OUTPUT_DIR}. " echo ">>>" python3 -m graph_net.model_path_handler \ --model-path-list ${dtype_generalized_subgraphs_list} \ @@ -416,7 +439,7 @@ function generate_unittests() { "handler_class_name": "AgentUnittestGeneratorPass", "handler_config": { "framework": "torch", - "model_path_prefix": "${DTYPE_GENERALIZED_FUSIBLE_SUBGRAPH_DIR}", + "model_path_prefix": "${BACKWARD_GRAPH_OUTPUT_DIR}", "output_dir": "${FUSIBLE_SUBGRAPH_UNITTEST_DIR}", "device": "cuda", "generate_main": false, @@ -520,8 +543,28 @@ EOF ) } +function backward_graph_extractor_typical() { + echo ">>> [2-13] Extract backward graphs for typical subgraph samples under ${DTYPE_GENERALIZED_TYPICAL_SUBGRAPH_DIR}." + echo ">>>" + python3 -m graph_net.model_path_handler \ + --model-path-list ${dtype_generalized_typical_subgraph_list} \ + --handler-config=$(base64 -w 0 <>> [2-13] Generate unittests for subgraph samples under ${DTYPE_GENERALIZED_TYPICAL_SUBGRAPH_DIR}. " + echo ">>> [2-14] Generate unittests for subgraph samples under ${BACKWARD_GRAPH_TYPICAL_OUTPUT_DIR}. " echo ">>>" python3 -m graph_net.model_path_handler \ --model-path-list ${dtype_generalized_typical_subgraph_list} \ @@ -531,7 +574,7 @@ function generate_unittest_for_typical_subgraphs() { "handler_class_name": "AgentUnittestGeneratorPass", "handler_config": { "framework": "torch", - "model_path_prefix": "${DTYPE_GENERALIZED_TYPICAL_SUBGRAPH_DIR}", + "model_path_prefix": "${BACKWARD_GRAPH_TYPICAL_OUTPUT_DIR}", "output_dir": "${TYPICAL_SUBGRAPH_UNITTEST_DIR}", "device": "cuda", "generate_main": false, @@ -586,6 +629,9 @@ function generate_fusible_subgraphs() { dtype_generalizer 2>&1 | tee ${DECOMPOSE_WORKSPACE}/log_dtype_generalizer_fusible_subgraphs_${suffix}.txt generate_generalized_subgraph_list ${DTYPE_GENERALIZED_FUSIBLE_SUBGRAPH_DIR} ${dtype_generalized_subgraphs_list} + # extract backward graphs (train_forward, train_backward, eval_forward) + backward_graph_extractor 2>&1 | tee ${DECOMPOSE_WORKSPACE}/log_backward_graph_extractor_${suffix}.txt + # generate kernelbench format unittest # generate_unittests 2>&1 | tee ${DECOMPOSE_WORKSPACE}/log_unittests_${suffix}.txt } @@ -606,6 +652,9 @@ function generate_typical_subgraphs() { dtype_generalizer_for_typical_subgraphs 2>&1 | tee ${DECOMPOSE_WORKSPACE}/log_dtype_generalizer_typical_subgraphs_${suffix}.txt generate_generalized_subgraph_list ${DTYPE_GENERALIZED_TYPICAL_SUBGRAPH_DIR} ${dtype_generalized_typical_subgraph_list} + # extract backward graphs (train_forward, train_backward, eval_forward) + backward_graph_extractor_typical 2>&1 | tee ${DECOMPOSE_WORKSPACE}/log_backward_graph_extractor_typical_${suffix}.txt + # generate kernelbench format unittest # generate_unittest_for_typical_subgraphs 2>&1 | tee ${DECOMPOSE_WORKSPACE}/log_unittests_typical_subgraphs_${suffix}.txt } @@ -643,6 +692,7 @@ function main() { sample_type="typical_graph" generate_typical_subgraphs + insert_graph_sample ${DEDUP_TYPICAL_SUBGRAPH_DIR} "github_torch_samples" "typical_graph" ${deduplicated_typical_subgraph_list} cp -rf $DTYPE_GENERALIZED_TYPICAL_SUBGRAPH_DIR $OUTPUT_DIR/$sample_type cp -rf $dtype_generalized_typical_subgraph_list $OUTPUT_DIR/$sample_type/sample_list.txt @@ -729,6 +779,25 @@ function summary() { done echo "" + num_backward_graph_subgraphs=`find ${BACKWARD_GRAPH_OUTPUT_DIR} -name "model.py" | wc -l` + echo "- [Fusible - 6] backward graph extraction: successed=${num_backward_graph_subgraphs}" + for mode in eval_forward train_forward train_backward + do + num_backward_graph_subgraphs_mode=`find ${BACKWARD_GRAPH_OUTPUT_DIR}/${mode} -name "model.py" | wc -l` + echo " ${mode}, successed=${num_backward_graph_subgraphs_mode}" + done + echo "" + + num_successed_unittests=`find ${FUSIBLE_SUBGRAPH_UNITTEST_DIR} -name "*_test.py" | wc -l` + unittest_successed_percent=$((num_successed_unittests * 100 / num_backward_graph_subgraphs)) + echo "- [Fusible - 7] generate unittest: successed=${num_successed_unittests}, percent=${unittest_successed_percent}%" + for dtype in float32 float16 bfloat16 + do + num_successed_unittests=`find ${FUSIBLE_SUBGRAPH_UNITTEST_DIR}/${dtype} -name "*_test.py" | wc -l` + echo " ${dtype}, successed=${num_successed_unittests}" + done + echo "" + # typical subgraphs num_dim_generalized_typical_subgraphs=`find ${DIM_GENERALIZED_TYPICAL_SUBGRAPH_DIR} -name "model.py" | wc -l` echo "- [Typical - 1] subgraph dimension generalization: successed=${num_dim_generalized_typical_subgraphs}" @@ -754,8 +823,17 @@ function summary() { done echo "" + num_backward_graph_typical_subgraphs=`find ${BACKWARD_GRAPH_TYPICAL_OUTPUT_DIR} -name "model.py" | wc -l` + echo "- [Typical - 4] backward graph extraction: successed=${num_backward_graph_typical_subgraphs}" + for mode in eval_forward train_forward train_backward + do + num_backward_graph_typical_mode=`find ${BACKWARD_GRAPH_TYPICAL_OUTPUT_DIR}/${mode} -name "model.py" | wc -l` + echo " ${mode}, successed=${num_backward_graph_typical_mode}" + done + echo "" + num_dtype_generalized_typical_subgraphs=`find ${DTYPE_GENERALIZED_TYPICAL_SUBGRAPH_DIR} -name "model.py" | wc -l` - echo "- [Typical - 4] dtype generalization: successed=${num_dtype_generalized_typical_subgraphs}" + echo "- [Typical - 5] dtype generalization: successed=${num_dtype_generalized_typical_subgraphs}" for dtype in float32 float16 bfloat16 do num_dtype_generalized_typical_subgraphs_index=`find ${DTYPE_GENERALIZED_TYPICAL_SUBGRAPH_DIR}/${dtype} -name "model.py" | wc -l` diff --git a/graph_net/torch/sample_pass/backward_graph_extractor.py b/graph_net/torch/sample_pass/backward_graph_extractor.py index f27b1da9b..e50d3b7e9 100644 --- a/graph_net/torch/sample_pass/backward_graph_extractor.py +++ b/graph_net/torch/sample_pass/backward_graph_extractor.py @@ -16,11 +16,12 @@ class BackwardGraphExtractor: - def __init__(self, model_name, model_path, output_dir, device): + def __init__(self, model_name, model_path, output_dir, device, rel_model_path): self.model_name = model_name self.model_path = model_path self.output_dir = output_dir self.device = device + self.rel_model_path = rel_model_path def __call__(self): module, forward_inputs = get_torch_module_and_inputs( @@ -28,24 +29,27 @@ def __call__(self): ) module.train() - original_name_dir = os.path.join(self.output_dir, self.model_name) - if not os.path.exists(original_name_dir): - shutil.copytree(self.model_path, original_name_dir) + eval_forward_dir = os.path.join( + self.output_dir, "eval_forward", self.rel_model_path + ) + if not os.path.exists(eval_forward_dir): + shutil.copytree(self.model_path, eval_forward_dir) forward_inputs = self.set_requires_grad_for_forward_inputs( self.model_path, module, forward_inputs ) gm_holder, backward_inputs = self.capture_graph(module, forward_inputs) - self.get_extractor("forward")(gm_holder["forward_gm"], forward_inputs) - self.get_extractor("backward")(gm_holder["backward_gm"], backward_inputs) + self.get_extractor("train_forward")(gm_holder["forward_gm"], forward_inputs) + self.get_extractor("train_backward")(gm_holder["backward_gm"], backward_inputs) def get_extractor(self, suffix): + workspace_path = os.path.join(self.output_dir, suffix, self.rel_model_path) return BuiltinGraphExtractor( name=f"{self.model_name}_{suffix}", dynamic=False, mut_graph_codes=[], placeholder_auto_rename=False, - workspace_path=self.output_dir, + workspace_path=workspace_path, ) def capture_graph(self, module, forward_inputs): @@ -180,13 +184,11 @@ def resume(self, rel_model_path: str): model_path_prefix = Path(self.config["model_path_prefix"]) model_name = f"{os.path.basename(rel_model_path)}" model_path = model_path_prefix / rel_model_path - output_dir = ( - Path(self.config["output_dir"]) - / os.path.dirname(rel_model_path) - / model_name - ) + output_dir = Path(self.config["output_dir"]) device = self._choose_device(self.config["device"]) - extractor = BackwardGraphExtractor(model_name, model_path, output_dir, device) + extractor = BackwardGraphExtractor( + model_name, model_path, output_dir, device, rel_model_path + ) extractor() def _choose_device(self, device) -> str: