Generalize sequence dimension (#386)

* 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

* save dimension generalization pass names into graph_net.json

* Generalize sequence dimension

Author: lixinqi

graph_net/constraint_util.py

@@ -11,6 +11,7 @@
 import tempfile
 import shutil
 from pathlib import Path
+import json
 
 
 class UpdateInputTensorConstraints:
@@ -109,12 +110,13 @@ def is_dyn_dim_cstr_feasible(dyn_dim_cstr):
             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, dim_gen_pass_names = 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)
+        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"]
@@ -129,7 +131,7 @@ def is_dyn_dim_cstr_feasible(dyn_dim_cstr):
     @contextmanager
     def _try_dimension_generalization(self, dim_axes_pairs, model_path, tensor_metas):
         if self.config["dimension_generalizer_filepath"] is None:
-            yield model_path
+            yield model_path, ()
             return
         py_module = load_module(os.path.join(model_path, "model.py"))
         GraphModule = getattr(py_module, "GraphModule")
@@ -142,7 +144,7 @@ def _try_dimension_generalization(self, dim_axes_pairs, model_path, tensor_metas
         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
+            yield model_path, ()
             return
         from dataclasses import asdict
 
@@ -156,8 +158,15 @@ def _try_dimension_generalization(self, dim_axes_pairs, model_path, tensor_metas
             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)
+            yield tmp_dir, dim_gen_pass.get_pass_names()
+
+    def _save_dim_gen_pass_names(self, dim_gen_pass_names, model_path):
+        from graph_net.graph_net_json_file_util import kDimensionGeneralizationPasses
+
+        graph_net_json_file_path = Path(f"{model_path}/graph_net.json")
+        graph_net_json = json.loads(graph_net_json_file_path.read_text())
+        graph_net_json[kDimensionGeneralizationPasses] = list(dim_gen_pass_names)
+        graph_net_json_file_path.write_text(json.dumps(graph_net_json))
 
     def _save_dyn_dim_cstr(self, dyn_dim_cstr, model_path):
         cstr_code = dyn_dim_cstr.serialize_to_py_str()
@@ -219,9 +228,12 @@ 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: Callable[[DynamicDimConstraints], bool],
+        dim_gen_pass_names: tuple[str],
     ):
         self.is_dyn_dim_cstr_feasible = is_dyn_dim_cstr_feasible
+        self.dim_gen_pass_names = dim_gen_pass_names
 
     def __call__(self, dyn_dim_cstr: DynamicDimConstraints) -> bool:
         return self.is_dyn_dim_cstr_feasible(dyn_dim_cstr)
@@ -242,16 +254,17 @@ def __init__(
     def __call__(
         self, dim_axes_pairs
     ) -> AbstractContextManager[DynDimCstrFeasibilityPredicator]:
-        with self.get_tmp_model_path_ctx_mgr(dim_axes_pairs) as tmp_model_apth:
+        ctx_mgr = self.get_tmp_model_path_ctx_mgr
+        with ctx_mgr(dim_axes_pairs) as (tmp_model_apth, dg_pass_names):
             predicator = self.get_predicator_is_dyn_dim_cstr_feasible(tmp_model_apth)
-            yield DynDimCstrFeasibilityPredicator(predicator)
+            yield DynDimCstrFeasibilityPredicator(predicator, dg_pass_names)
 
 
 def symbolize_data_input_dims(
     dyn_dim_cstr: DynamicDimConstraints,
     is_data_input: Callable[[str], bool],
     dyn_dim_cstr_feasibility_ctx_mgr: DynDimCstrFeasibilityContextManager,
-) -> DynamicDimConstraints | None:
+) -> (DynamicDimConstraints | None, tuple[str]):
     """
     is_data_input: Callable[["input_var_name:str"], bool]
     Symbolizes data input dimensions as much as possible.
@@ -273,6 +286,21 @@ def dumpy_filter_fn(input_name, input_idx, axis, dim):
 
     # Collect input dimensions into `unqiue_dims`
     assert dyn_dim_cstr.symbolize(dumpy_filter_fn) is None
+    total_dim_gen_pass_names = ()
+
+    def append_dim_gen_pass_names(dim_gen_pass_names):
+        nonlocal total_dim_gen_pass_names
+        total_dim_gen_pass_names = tuple(
+            [
+                *total_dim_gen_pass_names,
+                *(
+                    pass_name
+                    for pass_name in dim_gen_pass_names
+                    if pass_name not in total_dim_gen_pass_names
+                ),
+            ]
+        )
+
     for i, picked_dim in enumerate(unqiue_dims):
         cur_dyn_dim_cstr = copy.deepcopy(dyn_dim_cstr)
 
@@ -292,12 +320,12 @@ def filter_fn(input_name, input_idx, axis, dim):
         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:
+        ctx_mgr = dyn_dim_cstr_feasibility_ctx_mgr
+        with ctx_mgr(dim_axes_pairs) as dyn_dim_cstr_feasibility:
             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):
+            if not dyn_dim_cstr_feasibility(tmp_dyn_dim_cstr):
                 continue
             dyn_dim_cstr = cur_dyn_dim_cstr
-    return dyn_dim_cstr
+            append_dim_gen_pass_names(dyn_dim_cstr_feasibility.dim_gen_pass_names)
+    return dyn_dim_cstr, total_dim_gen_pass_names

graph_net/graph_net_json_file_util.py

@@ -0,0 +1 @@
+kDimensionGeneralizationPasses = "dimension_generalization_passes"

graph_net/torch/dim_gen_passes/naive_call_method_expand_pass.py

@@ -1,14 +1,6 @@
 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):

graph_net/torch/dim_gen_passes/naive_call_method_reshape_pass.py

@@ -1,14 +1,6 @@
 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):

graph_net/torch/dim_gen_passes/naive_call_method_view_pass.py

@@ -1,14 +1,6 @@
 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):

graph_net/torch/dim_gen_passes/non_batch_call_function_arange_pass.py

@@ -0,0 +1,88 @@
+import torch
+import torch.fx as fx
+from graph_net.torch.dim_gen_passes import DimensionGeneralizationPass
+from collections import namedtuple
+
+
+class ConcretePass(DimensionGeneralizationPass):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def need_rewrite(self, traced_module: fx.GraphModule) -> bool:
+        # non batch
+        if 0 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:
+        return (
+            node.op == "call_function"
+            and node.target == torch.arange
+            and len(node.args) >= 2
+            and node.args[0] == 0
+            and node.args[1] == self.dim
+        )
+
+    def rewrite(self, traced_module: fx.GraphModule) -> fx.GraphModule:
+        """
+        Fx Pass: Replaces hardcoded constants in 'torch.arange' ops that match an input tensor dimension
+        with a dynamic 'size()' call. The primary goal is to dynamicize the batch size.
+        """
+        # 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 = {}
+
+        NodeAxis = namedtuple("NodeAxis", ["node", "shape_axis"])
+        last_node_axis = None
+
+        def try_reset_last_node_axis(node, new_node):
+            nonlocal last_node_axis
+            node_meta = node.meta.get("tensor_meta")
+            if node_meta is None:
+                return
+            if not hasattr(node_meta, "shape"):
+                return
+            for axis, dim in enumerate(node_meta.shape):
+                if not (dim == self.dim and axis > 0):
+                    continue
+                last_node_axis = NodeAxis(node=new_node, shape_axis=axis)
+                return
+
+        def get_new_node_arg(i, arg):
+            if i != 1:
+                return val_map[arg] if arg in val_map else arg
+            # i == 1
+            assert arg == self.dim
+
+            # Use the size() method to retrieve the dynamic dimension
+            size_node = new_graph.call_method(
+                "size", args=(last_node_axis.node, last_node_axis.shape_axis)
+            )
+
+            return size_node
+
+        def create_new_node(node):
+            if not (self._node_need_rewrite(node) and last_node_axis is not None):
+                # Copy other nodes to the new graph
+                new_node = new_graph.node_copy(node, lambda x: val_map[x])
+                try_reset_last_node_axis(node=node, new_node=new_node)
+                return new_node
+
+            new_arange_args = tuple(
+                get_new_node_arg(i, arg) for i, arg in enumerate(node.args)
+            )
+
+            # --- Rebuild the torch.arange node ---
+            new_node = new_graph.call_function(torch.arange, args=new_arange_args)
+
+            return new_node
+
+        for node in traced_module.graph.nodes:
+            val_map[node] = create_new_node(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/non_batch_call_function_getitem_slice_pass.py

@@ -0,0 +1,113 @@
+import torch
+import torch.fx as fx
+from graph_net.torch.dim_gen_passes import DimensionGeneralizationPass
+from collections import namedtuple
+import operator
+import copy
+
+
+class ConcretePass(DimensionGeneralizationPass):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def need_rewrite(self, traced_module: fx.GraphModule) -> bool:
+        # non batch
+        if 0 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_function"):
+            return False
+        if not (node.target == operator.getitem):
+            return False
+        if not isinstance(node.args[1], tuple):
+            return False
+        if not any(self._slice_need_rewrite(slice_obj) for slice_obj in node.args[1]):
+            return False
+        return True
+
+    def _slice_need_rewrite(self, slice_obj) -> bool:
+        if not isinstance(slice_obj, slice):
+            return False
+        return (
+            slice_obj.stop == self.dim
+            and (slice_obj.start is None or slice_obj.start == 0)
+            and (slice_obj.step is None or slice_obj.step == 1)
+        )
+
+    def rewrite(self, traced_module: fx.GraphModule) -> fx.GraphModule:
+        """
+        Fx Pass: Replaces hardcoded constants in 'operator.getitem' ops that match an input tensor dimension
+        with a dynamic 'size()' call. The primary goal is to dynamicize the batch size.
+        """
+        # 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 = {}
+
+        NodeAxis = namedtuple("NodeAxis", ["node", "shape_axis"])
+        last_node_axis = None
+
+        def try_reset_last_node_axis(node, new_node):
+            nonlocal last_node_axis
+            node_meta = node.meta.get("tensor_meta")
+            if node_meta is None:
+                return
+            if not hasattr(node_meta, "shape"):
+                return
+            for axis, dim in enumerate(node_meta.shape):
+                if not (dim == self.dim and axis > 0):
+                    continue
+                last_node_axis = NodeAxis(node=new_node, shape_axis=axis)
+                return
+
+        def val_map_contains(key):
+            if isinstance(key, slice):
+                return False
+            return key in val_map
+
+        def get_new_getitem_tuple_elem(elem):
+            if not (
+                isinstance(elem, slice)
+                and self._slice_need_rewrite(elem)
+                and last_node_axis is not None
+            ):
+                return val_map[elem] if val_map_contains(elem) else elem
+
+            slice_obj = copy.deepcopy(elem)
+            assert slice_obj.stop == self.dim
+
+            # Use the size() method to retrieve the dynamic dimension
+            size_node = new_graph.call_method(
+                "size", args=(last_node_axis.node, last_node_axis.shape_axis)
+            )
+            return slice(slice_obj.start, size_node, slice_obj.step)
+
+        def get_new_getitem_arg(arg):
+            if not isinstance(arg, tuple):
+                return val_map[arg] if arg in val_map else arg
+            return tuple(get_new_getitem_tuple_elem(elem) for elem in arg)
+
+        def create_new_node(node):
+            if not self._node_need_rewrite(node):
+                # Copy other nodes to the new graph
+                new_node = new_graph.node_copy(node, lambda x: val_map[x])
+                try_reset_last_node_axis(node=node, new_node=new_node)
+                return new_node
+
+            new_gettiem_args = tuple(get_new_getitem_arg(arg) for arg in node.args)
+
+            # --- Rebuild the operator.getitem node ---
+            new_node = new_graph.call_function(operator.getitem, args=new_gettiem_args)
+
+            return new_node
+
+        for node in traced_module.graph.nodes:
+            val_map[node] = create_new_node(node)
+
+        # Replace the old graph with the new graph and return
+        traced_module.graph = new_graph
+        traced_module.recompile()
+        return traced_module