more dimension generalization pass

Author: lixinqi

graph_net/config/empty_cstr_torch_samples_list.txt

@@ -25,7 +25,7 @@ def __init__(self, config=None):
         self.model_runnable_predicator = self._make_model_runnable_predicator(
             self.config
         )
-        self.num_successful_handled_models = 0
+        self.num_handled_models = 0
 
     def _make_data_input_predicator(self, config):
         module = load_module(config["data_input_predicator_filepath"])
@@ -51,7 +51,7 @@ def _make_config(
         model_path_prefix="",
         resume=False,
         last_model_log_file=None,
-        limits_successfully_handled_models=None,
+        limits_handled_models=None,
     ):
         if data_input_predicator_config is None:
             data_input_predicator_config = {}
@@ -72,7 +72,7 @@ def _make_config(
             "dimension_generalizer_class_name": dimension_generalizer_class_name,
             "dimension_generalizer_config": dimension_generalizer_config,
             "last_model_log_file": last_model_log_file,
-            "limits_successfully_handled_models": limits_successfully_handled_models,
+            "limits_handled_models": limits_handled_models,
         }
 
     def __call__(self, model_path):
@@ -125,16 +125,15 @@ def is_dyn_dim_cstr_feasible(dyn_dim_cstr):
         )
         self._save_dyn_dim_cstr(dyn_dim_cstr, model_path)
         self._save_dim_gen_pass_names(dim_gen_pass_names, model_path)
-        if len(dyn_dim_cstr.symbols) > 0:
-            self.num_successful_handled_models += 1
-            limits = self.config["limits_successfully_handled_models"]
-            if limits is not None:
-                if self.num_successful_handled_models > limits:
-                    print(
-                        "`num_successful_handled_models` exceeds config `limits_successfully_handled_models`",
-                        file=sys.stderr,
-                    )
-                    sys.exit(0)
+        self.num_handled_models += 1
+        limits = self.config["limits_handled_models"]
+        if limits is not None:
+            if self.num_handled_models >= limits:
+                print(
+                    "`num_handled_models` exceeds config `limits_handled_models`",
+                    file=sys.stderr,
+                )
+                sys.exit(0)
 
     def get_dimension_generalizer(self):
         if hasattr(self, "_dim_generalizer"):
@@ -159,6 +158,7 @@ def get_model(self, model_path):
     def _try_dimension_generalization(self, dim_axes_pairs, model_path, inputs):
         logging.warning(f"enter _try_dimension_generalization")
         if self.config["dimension_generalizer_filepath"] is None:
+            self._save_model_to_log_file(model_path)
             yield model_path, ()
             return
         model = self.get_model(model_path)
@@ -168,6 +168,7 @@ def _try_dimension_generalization(self, dim_axes_pairs, model_path, inputs):
         need_rewrite = dim_gen_pass.need_rewrite(inputs)
         logging.warning(f"after need_rewrite")
         if not need_rewrite:
+            self._save_model_to_log_file(model_path)
             yield model_path, ()
             return
 
@@ -177,11 +178,14 @@ def _try_dimension_generalization(self, dim_axes_pairs, model_path, inputs):
         with tempfile.TemporaryDirectory() as tmp_dir:
             shutil.copytree(Path(model_path), Path(tmp_dir), dirs_exist_ok=True)
             dim_gen_pass.save_graph_module(graph_module, tmp_dir)
-            if self.config["last_model_log_file"] is not None:
-                log_file = Path(self.config["last_model_log_file"])
-                shutil.copy(Path(tmp_dir) / "model.py", log_file)
+            self._save_model_to_log_file(tmp_dir)
             yield tmp_dir, dim_gen_pass.get_pass_names()
 
+    def _save_model_to_log_file(self, model_path):
+        if self.config["last_model_log_file"] is not None:
+            log_file = Path(self.config["last_model_log_file"])
+            shutil.copy(Path(model_path) / "model.py", log_file)
+
     def _save_dim_gen_pass_names(self, dim_gen_pass_names, model_path):
         from graph_net.graph_net_json_file_util import kDimensionGeneralizationPasses
 
@@ -324,7 +328,7 @@ def append_dim_gen_pass_names(dim_gen_pass_names):
         )
 
     for i, picked_dim in enumerate(unqiue_dims):
-        logging.warning(f"{i=} {picked_dim=}")
+        logging.warning(f"{i=} {picked_dim=} {dim2axes[picked_dim]=}")
         cur_dyn_dim_cstr = copy.deepcopy(dyn_dim_cstr)
 
         def filter_fn(input_name, input_idx, axis, dim):

graph_net/constraint_util.py

@@ -19,11 +19,27 @@ config_json_str=$(cat <<EOF
         "model_runnable_predicator_class_name": "$model_runnable_predicator",
         "dimension_generalizer_filepath": "$GRAPH_NET_ROOT/torch/static_to_dynamic.py",
         "dimension_generalizer_class_name": "StaticToDynamic",
+        "dimension_generalizer_config": {
+            "pass_names": [
+                "batch_call_method_view_pass",
+                "tuple_arg_call_method_view_pass",
+                "naive_call_method_reshape_pass",
+                "naive_call_method_expand_pass",
+                "non_batch_call_method_expand_pass",
+                "non_batch_call_function_arange_pass",
+                "non_batch_call_function_getitem_slice_pass",
+                "non_batch_call_function_full_pass",
+                "non_batch_call_function_full_plus_one_pass",
+                "non_batch_call_function_zeros_pass",
+                "non_batch_call_function_arange_plus_one_pass"
+            ]
+        },
+        "limits_handled_models": 1,
         "last_model_log_file": "/tmp/a.py"
     }
 }
 EOF
 )
 CONFIG=$(echo $config_json_str | base64 -w 0)
 
-python3 -m graph_net.model_path_handler --model-path-list $GRAPH_NET_ROOT/config/torch_samples_list.txt --handler-config=$CONFIG
+python3 -m graph_net.model_path_handler --model-path-list $GRAPH_NET_ROOT/config/empty_cstr_torch_samples_list.txt --handler-config=$CONFIG

graph_net/tools/batch_init_input_tensor_constraints.sh

@@ -8,7 +8,9 @@ def __init__(self, config):
         self.config = config
 
     def __call__(self, model_path, input_var_name: str) -> bool:
-        return not ("_self_" in input_var_name)
+        return not (
+            "_self_" in input_var_name or "_instance_modules_" in input_var_name
+        )
 
 
 class ModelRunnablePredicator:

graph_net/torch/constraint_util.py

@@ -0,0 +1,104 @@
+import torch
+import torch.fx as fx
+from graph_net.torch.dim_gen_passes import DimensionGeneralizationPass
+import os
+
+
+class ConcretePass(DimensionGeneralizationPass):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def get_pass_name(cls) -> bool:
+        return os.path.basename(__file__)[:-3]
+
+    def need_rewrite(self, traced_module: fx.GraphModule) -> bool:
+        if 0 not in self.axes:
+            return False
+        return any(self._node_need_rewrite(node) for node in traced_module.graph.nodes)
+
+    def _node_need_rewrite(self, node) -> bool:
+        if not (node.op == "call_method"):
+            return False
+        if not (node.target == "view"):
+            return False
+        if not (len(node.args) >= 2):
+            return False
+        if not (isinstance(node.args[1], int)):
+            return False
+        if not (self.dim == node.args[1]):
+            return False
+        return True
+
+    def rewrite(self, traced_module: fx.GraphModule) -> fx.GraphModule:
+        """
+        Fx Pass: Replaces hardcoded constants in 'view' ops that match an input tensor dimension
+        with a dynamic 'size()' call. The primary goal is to dynamicize the batch size (axis 0).
+        """
+        # Create a new graph to hold the rewritten nodes
+        new_graph = fx.Graph()
+
+        # Create a map to link nodes from the old graph to nodes in the new graph
+        val_map = {}
+
+        def get_new_tuple_args(input_tensor_node, view_args):
+            # --- Dependency on ShapeProp Results ---
+            # input_shape is the static shape (e.g., batch_size, C, H, W)
+            input_meta = input_tensor_node.meta.get("tensor_meta")
+            if input_meta is None:
+                raise RuntimeError(
+                    f"Node {input_tensor_node.name} lacks tensor_meta. Did ShapeProp run?"
+                )
+
+            input_shape = input_meta.shape
+
+            # Find the new list of view arguments
+            new_view_args = []
+            for axis_idx, target_dim in enumerate(view_args):
+                if not isinstance(target_dim, int) or target_dim < 1:
+                    new_view_args.append(
+                        val_map[target_dim] if target_dim in val_map else target_dim
+                    )
+                    continue
+
+                if axis_idx == 0 and target_dim == input_shape[axis_idx]:
+                    new_input_node = val_map[input_tensor_node]
+                    size_node = new_graph.call_method(
+                        "size", args=(new_input_node, axis_idx)
+                    )
+                    best_match = size_node
+                else:
+                    best_match = target_dim
+                new_view_args.append(best_match)
+            return tuple(new_view_args)
+
+        for node in traced_module.graph.nodes:
+            if self._node_need_rewrite(node):
+                # Get the input tensor node
+                input_tensor_node = node.args[0]
+                # Get the target shape arguments for view (e.g., 1, -1, 6, 64)
+                view_args = node.args[1:]
+                print(f"{view_args=}")
+                new_view_args = get_new_tuple_args(input_tensor_node, view_args)
+
+                # --- Rebuild the view node ---
+                # 1. Map the input tensor node to the new graph node
+                new_input_node = val_map[input_tensor_node]
+
+                # 2. Insert the new view node into the new graph
+                # with new_graph.inserting_after(new_input_node):
+                new_node = new_graph.call_method(
+                    "view", args=(new_input_node, *new_view_args)
+                )
+
+                # 3. Map the old node to the new node
+                val_map[node] = new_node
+
+            else:
+                # Copy other nodes to the new graph
+                new_node = new_graph.node_copy(node, lambda x: val_map[x])
+                val_map[node] = new_node
+
+        # Replace the old graph with the new graph and return
+        traced_module.graph = new_graph
+        traced_module.recompile()
+        return traced_module

graph_net/torch/dim_gen_passes/batch_call_method_view_pass.py

@@ -14,10 +14,17 @@ def get_pass_name(cls) -> bool:
     def need_rewrite(self, traced_module: fx.GraphModule) -> bool:
         if 0 not in self.axes:
             return False
-        for node in traced_module.graph.nodes:
-            if node.op == "call_method" and node.target == "view":
-                return True
-        return False
+        return any(self._node_need_rewrite(node) for node in traced_module.graph.nodes)
+
+    def _node_need_rewrite(self, node) -> bool:
+        if not (node.op == "call_method"):
+            return False
+        if not (node.target == "view"):
+            return False
+        print(f"{self.dim=} {node.args[1:]=}")
+        if not (self.dim in node.args[1:]):
+            return False
+        return True
 
     def rewrite(self, traced_module: fx.GraphModule) -> fx.GraphModule:
         """