DimensionGeneralizationPass (#384)

* support checking model redundancy

* revert change of vision_model_test

* reformat python code.

* reformat bert_model_test.py and utils.py

* minor fix

* fix failed check by comparing directories after os.path.realpath()

* fix bugs in check_validate.sh

* set dynamic=False in single_device_runner.py

* reset graph hash

* add robustness code for generating input tensor constraints

* Introduce input_tensor_constraints.py using shape propagation logic.

* support dimension generalization for torch.Tensor.view and torch.Tensor.reshape

* 1) support dimension generalization for torch.Tensor.expand(); 2) fix bugs in generalization for torch.Tensor.view and torch.Tensor.reshape

* dimension_generalization_passes

* Refactored DimensionGeneralizationPass.__init__ to accept argument dim_axes_pairs, enabling targeted configuration for specific use cases

Author: lixinqi

graph_net/config/small_torch_samples_list.txt

@@ -1,3 +1,3 @@
 # samples/timm/resnetaa50d.d_in12k
-# samples/transformers-auto-model/opus-mt-en-gmw
-samples/transformers-auto-model/Michielo_mt5-small_nl-en_translation
+samples/transformers-auto-model/opus-mt-en-gmw
+# samples/transformers-auto-model/Michielo_mt5-small_nl-en_translation

graph_net/constraint_util.py

@@ -1,4 +1,5 @@
 from graph_net.dynamic_dim_constraints import DynamicDimConstraints
+from contextlib import AbstractContextManager
 from graph_net.imp_util import load_module
 from graph_net.tensor_meta import TensorMeta
 from typing import Callable
@@ -21,6 +22,7 @@ def __init__(self, config=None):
         self.model_runnable_predicator = self._make_model_runnable_predicator(
             self.config
         )
+        self.num_successful_handled_models = 0
 
     def _make_data_input_predicator(self, config):
         module = load_module(config["data_input_predicator_filepath"])
@@ -45,6 +47,8 @@ def _make_config(
         dimension_generalizer_config=None,
         model_path_prefix="",
         resume=False,
+        last_model_log_file=None,
+        limits_successfully_handled_models=None,
     ):
         if data_input_predicator_config is None:
             data_input_predicator_config = {}
@@ -64,6 +68,8 @@ def _make_config(
             "dimension_generalizer_filepath": dimension_generalizer_filepath,
             "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,
         }
 
     def __call__(self, model_path):
@@ -86,24 +92,42 @@ def __call__(self, model_path):
         def data_input_predicator(input_var_name):
             return self.data_input_predicator(model_path, input_var_name)
 
-        with self._try_dimension_generalization(
-            model_path, tensor_metas
-        ) as tmp_model_path:
+        def get_tmp_model_path_ctx_mgr(dim_axes_pairs):
+            return self._try_dimension_generalization(
+                dim_axes_pairs, model_path, tensor_metas
+            )
 
+        def get_predicator_is_dyn_dim_cstr_feasible(tmp_model_path):
             def is_dyn_dim_cstr_feasible(dyn_dim_cstr):
                 return self._is_dyn_dim_cstr_feasible(
                     tmp_model_path, tensor_metas, dyn_dim_cstr
                 )
 
-            dyn_dim_cstr = symbolize_data_input_dims(
-                dyn_dim_cstr,
-                is_data_input=data_input_predicator,
-                is_dyn_dim_cstr_feasible=is_dyn_dim_cstr_feasible,
-            )
-            self._save_dyn_dim_cstr(dyn_dim_cstr, model_path)
+            return is_dyn_dim_cstr_feasible
+
+        dyn_dim_cstr_feasibility_ctx_mgr = DynDimCstrFeasibilityContextManager(
+            get_tmp_model_path_ctx_mgr=get_tmp_model_path_ctx_mgr,
+            get_predicator_is_dyn_dim_cstr_feasible=get_predicator_is_dyn_dim_cstr_feasible,
+        )
+        dyn_dim_cstr = symbolize_data_input_dims(
+            dyn_dim_cstr,
+            is_data_input=data_input_predicator,
+            dyn_dim_cstr_feasibility_ctx_mgr=dyn_dim_cstr_feasibility_ctx_mgr,
+        )
+        self._save_dyn_dim_cstr(dyn_dim_cstr, 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)
 
     @contextmanager
-    def _try_dimension_generalization(self, model_path, tensor_metas):
+    def _try_dimension_generalization(self, dim_axes_pairs, model_path, tensor_metas):
         if self.config["dimension_generalizer_filepath"] is None:
             yield model_path
             return
@@ -115,20 +139,23 @@ def _try_dimension_generalization(self, model_path, tensor_metas):
             load_module(self.config["dimension_generalizer_filepath"]),
             self.config["dimension_generalizer_class_name"],
         )
-        pass_obj = decorator_cls(self.config["dimension_generalizer_config"])(model)
-        if not pass_obj.need_rewrite():
+        dim_generalizer = decorator_cls(self.config["dimension_generalizer_config"])
+        dim_gen_pass = dim_generalizer(model, dim_axes_pairs)
+        if not dim_gen_pass.need_rewrite():
             yield model_path
             return
         from dataclasses import asdict
 
         tensor_meta_attrs_list = [asdict(tensor_meta) for tensor_meta in tensor_metas]
-        graph_module = pass_obj.rewrite_with_tensor_meta_attrs_list(
-            tensor_meta_attrs_list
+        graph_module = dim_gen_pass.rewrite_with_tensor_meta_attrs_list(
+            tensor_meta_attrs_list=tensor_meta_attrs_list,
         )
         with tempfile.TemporaryDirectory() as tmp_dir:
             shutil.copytree(Path(model_path), Path(tmp_dir), dirs_exist_ok=True)
-            pass_obj.save_graph_module(graph_module, tmp_dir)
-            shutil.copy(Path(tmp_dir) / "model.py", Path("/tmp/a.py"))
+            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)
             yield tmp_dir
             # shutil.copytree(Path(tmp_dir), Path(model_path), dirs_exist_ok=True)
 
@@ -190,10 +217,40 @@ def make_dyn_dim_cstr_from_tensor_metas(tensor_metas: list[TensorMeta]):
     )
 
 
+class DynDimCstrFeasibilityPredicator:
+    def __init__(
+        self, is_dyn_dim_cstr_feasible: Callable[[DynamicDimConstraints], bool]
+    ):
+        self.is_dyn_dim_cstr_feasible = is_dyn_dim_cstr_feasible
+
+    def __call__(self, dyn_dim_cstr: DynamicDimConstraints) -> bool:
+        return self.is_dyn_dim_cstr_feasible(dyn_dim_cstr)
+
+
+class DynDimCstrFeasibilityContextManager:
+    def __init__(
+        self,
+        get_tmp_model_path_ctx_mgr,
+        get_predicator_is_dyn_dim_cstr_feasible,
+    ):
+        self.get_tmp_model_path_ctx_mgr = get_tmp_model_path_ctx_mgr
+        self.get_predicator_is_dyn_dim_cstr_feasible = (
+            get_predicator_is_dyn_dim_cstr_feasible
+        )
+
+    @contextmanager
+    def __call__(
+        self, dim_axes_pairs
+    ) -> AbstractContextManager[DynDimCstrFeasibilityPredicator]:
+        with self.get_tmp_model_path_ctx_mgr(dim_axes_pairs) as tmp_model_apth:
+            predicator = self.get_predicator_is_dyn_dim_cstr_feasible(tmp_model_apth)
+            yield DynDimCstrFeasibilityPredicator(predicator)
+
+
 def symbolize_data_input_dims(
     dyn_dim_cstr: DynamicDimConstraints,
     is_data_input: Callable[[str], bool],
-    is_dyn_dim_cstr_feasible: Callable[[DynamicDimConstraints], bool],
+    dyn_dim_cstr_feasibility_ctx_mgr: DynDimCstrFeasibilityContextManager,
 ) -> DynamicDimConstraints | None:
     """
     is_data_input: Callable[["input_var_name:str"], bool]
@@ -202,18 +259,21 @@ def symbolize_data_input_dims(
     Returns None if no symbolicable dim .
     """
     unqiue_dims = []
+    dim2axes = {}
 
     def dumpy_filter_fn(input_name, input_idx, axis, dim):
         if is_data_input(input_name):
             print("data_input", input_name, input_idx, axis, dim)
             if dim not in unqiue_dims:
                 unqiue_dims.append(dim)
-        # No symbolization because of returning True
+                dim2axes[dim] = []
+            dim2axes[dim].append(axis)
+        # No symbolization by returning False
         return False
 
     # Collect input dimensions into `unqiue_dims`
     assert dyn_dim_cstr.symbolize(dumpy_filter_fn) is None
-    for picked_dim in unqiue_dims:
+    for i, picked_dim in enumerate(unqiue_dims):
         cur_dyn_dim_cstr = copy.deepcopy(dyn_dim_cstr)
 
         def filter_fn(input_name, input_idx, axis, dim):
@@ -229,9 +289,15 @@ def filter_fn(input_name, input_idx, axis, dim):
         sym2example_value = {symbol: picked_dim + 1}
         if not cur_dyn_dim_cstr.check_delta_symbol2example_value(sym2example_value):
             continue
-        tmp_dyn_dim_cstr = copy.deepcopy(cur_dyn_dim_cstr)
-        tmp_dyn_dim_cstr.update_symbol2example_value(sym2example_value)
-        if not is_dyn_dim_cstr_feasible(tmp_dyn_dim_cstr):
-            continue
-        dyn_dim_cstr = cur_dyn_dim_cstr
+        dim_axes_pairs = tuple(
+            (dim, axes) for dim in unqiue_dims[: i + 1] for axes in [dim2axes[dim]]
+        )
+        with dyn_dim_cstr_feasibility_ctx_mgr(
+            dim_axes_pairs
+        ) as is_dyn_dim_cstr_feasible:
+            tmp_dyn_dim_cstr = copy.deepcopy(cur_dyn_dim_cstr)
+            tmp_dyn_dim_cstr.update_symbol2example_value(sym2example_value)
+            if not is_dyn_dim_cstr_feasible(tmp_dyn_dim_cstr):
+                continue
+            dyn_dim_cstr = cur_dyn_dim_cstr
     return dyn_dim_cstr

graph_net/test/batch_init_input_tensor_constraints_test.sh

@@ -18,7 +18,8 @@ config_json_str=$(cat <<EOF
         "model_runnable_predicator_filepath": "$GRAPH_NET_ROOT/torch/constraint_util.py",
         "model_runnable_predicator_class_name": "ModelRunnablePredicator",
         "dimension_generalizer_filepath": "$GRAPH_NET_ROOT/torch/static_to_dynamic.py",
-        "dimension_generalizer_class_name": "StaticToDynamic"
+        "dimension_generalizer_class_name": "StaticToDynamic",
+        "last_model_log_file": "/tmp/a.py"
     }
 }
 EOF

graph_net/test/shape_prop_batch_init_input_tensor_constraints_test.sh

@@ -18,7 +18,8 @@ config_json_str=$(cat <<EOF
         "model_runnable_predicator_filepath": "$GRAPH_NET_ROOT/torch/constraint_util.py",
         "model_runnable_predicator_class_name": "ShapePropagatablePredicator",
         "dimension_generalizer_filepath": "$GRAPH_NET_ROOT/torch/static_to_dynamic.py",
-        "dimension_generalizer_class_name": "StaticToDynamic"
+        "dimension_generalizer_class_name": "StaticToDynamic",
+        "last_model_log_file": "/tmp/a.py"
     }
 }
 EOF

graph_net/torch/dim_gen_passes/__init__.py

@@ -0,0 +1 @@
+from graph_net.torch.dim_gen_passes.pass_base import DimensionGeneralizationPass

graph_net/torch/dim_gen_passes/naive_call_method_expand_pass.py

@@ -0,0 +1,122 @@
+import torch
+import torch.fx as fx
+from graph_net.torch.dim_gen_passes import DimensionGeneralizationPass
+from torch.fx.passes.shape_prop import ShapeProp
+from graph_net.torch.utils import apply_templates
+from pathlib import Path
+import inspect
+from typing import Any
+from contextlib import contextmanager
+from torch.export import export
+from graph_net.torch.fx_graph_parse_util import parse_sole_graph_module
+
+
+class ConcretePass(DimensionGeneralizationPass):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    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 == "expand":
+                return True
+        return False
+
+    def rewrite(self, traced_module: fx.GraphModule) -> fx.GraphModule:
+        """
+        Fx Pass: Replaces hardcoded constants in 'expand' 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 = {}
+
+        for node in traced_module.graph.nodes:
+            if node.op == "call_method" and node.target == "expand":
+                # Get the input tensor node
+                input_tensor_node = node.args[0]
+                # Get the target shape arguments for expand (e.g., 1, 4, 6, 64)
+                expand_args = node.args[1:]
+
+                # --- 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 expand arguments
+                new_expand_args = []
+
+                # Iterate over the target dimensions of expand (dim0, dim1, ...)
+                for i, target_dim in enumerate(expand_args):
+                    # 1. Handle dynamic dimensions (e.g., -1 or non-integer values)
+                    if not isinstance(target_dim, int) or target_dim < 1:
+                        new_expand_args.append(
+                            val_map[target_dim] if target_dim in val_map else target_dim
+                        )
+                        continue
+
+                    # 2. Handle hardcoded constants (e.g., 1, 6, 64)
+
+                    # --- Core Logic: Find the matching dynamic axis ---
+
+                    # Default: Keep the hardcoded constant if no matching dynamic axis is found
+                    best_match = target_dim
+                    matched_axis = -1
+
+                    axis_idx = i
+                    input_dim_size = input_shape[i]
+                    if target_dim == input_dim_size:
+                        if axis_idx == 0:
+                            matched_axis = axis_idx
+                        elif axis_idx > 0 and input_dim_size > 1:
+                            matched_axis = axis_idx
+                        else:
+                            # Do nothing.
+                            pass
+
+                    if matched_axis != -1:
+                        # Found a matching dynamic axis (matched_axis), replace it with a size() call
+
+                        # 1. Create a call to size(axis) in the new graph
+                        # NOTE: input_tensor_node must first be mapped to a new graph node via val_map
+                        new_input_node = val_map[input_tensor_node]
+
+                        # Use the size() method to retrieve the dynamic dimension
+                        size_node = new_graph.call_method(
+                            "size", args=(new_input_node, matched_axis)
+                        )
+
+                        best_match = size_node
+
+                    new_expand_args.append(best_match)
+
+                # --- Rebuild the expand node ---
+                # 1. Map the input tensor node to the new graph node
+                new_input_node = val_map[input_tensor_node]
+
+                # 2. Insert the new expand node into the new graph
+                # with new_graph.inserting_after(new_input_node):
+                new_node = new_graph.call_method(
+                    "expand", args=(new_input_node, *new_expand_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