generate initial input_tensor_constraints.py

Author: lixinqi

graph_net/constraint_util.py

@@ -1,41 +1,174 @@
 from graph_net.dynamic_dim_constraints import DynamicDimConstraints
+from graph_net.imp_util import load_module
+from graph_net.tensor_meta import TensorMeta
 from typing import Callable
 import copy
+import sys
+import os
+
+
+class UpdateInputTensorConstraints:
+    def __init__(self, config=None):
+        if config is None:
+            config = {}
+        self.config = self._make_config(**config)
+        self.data_input_predicator = self._make_data_input_predicator(self.config)
+        self.model_runnable_predicator = self._make_model_runnable_predicator(
+            self.config
+        )
+
+    def _make_data_input_predicator(self, config):
+        module = load_module(config["data_input_predicator_filepath"])
+        cls = getattr(module, config["data_input_predicator_class_name"])
+        return cls(config["data_input_predicator_config"])
+
+    def _make_model_runnable_predicator(self, config):
+        module = load_module(config["model_runnable_predicator_filepath"])
+        cls = getattr(module, config["model_runnable_predicator_class_name"])
+        return cls(config["model_runnable_predicator_config"])
+
+    def _make_config(
+        self,
+        data_input_predicator_filepath,
+        model_runnable_predicator_filepath,
+        data_input_predicator_class_name="DataInputPredicator",
+        data_input_predicator_config=None,
+        model_runnable_predicator_class_name="ModelRunner",
+        model_runnable_predicator_config=None,
+    ):
+        if data_input_predicator_config is None:
+            data_input_predicator_config = {}
+        if model_runnable_predicator_config is None:
+            model_runnable_predicator_config = {}
+        return {
+            "data_input_predicator_filepath": data_input_predicator_filepath,
+            "data_input_predicator_class_name": data_input_predicator_class_name,
+            "data_input_predicator_config": data_input_predicator_config,
+            "model_runnable_predicator_filepath": model_runnable_predicator_filepath,
+            "model_runnable_predicator_class_name": model_runnable_predicator_class_name,
+            "model_runnable_predicator_config": model_runnable_predicator_config,
+        }
+
+    def __call__(self, model_path):
+        tensor_metas = self._get_tensor_metas(model_path)
+        dyn_dim_cstr = make_dyn_dim_cstr_from_tensor_metas(tensor_metas)
+
+        def data_input_predicator(input_var_name):
+            return self.data_input_predicator(model_path, input_var_name)
+
+        def is_dyn_dim_cstr_feasible(dyn_dim_cstr):
+            return self._is_dyn_dim_cstr_feasible(
+                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)
+
+    def _save_dyn_dim_cstr(self, dyn_dim_cstr, model_path):
+        cstr_code = dyn_dim_cstr.serialize_to_py_str()
+        with open(os.path.join(model_path, "input_tensor_constraints.py"), "w") as fp:
+            fp.write(cstr_code)
+
+    def _get_tensor_metas(self, model_path):
+        make = TensorMeta.unserialize_from_py_file
+        return [
+            *make(os.path.join(model_path, "input_meta.py")),
+            *make(os.path.join(model_path, "weight_meta.py")),
+        ]
+
+    def _is_dyn_dim_cstr_feasible(
+        self, model_path, tensor_metas, dyn_dim_cstr: DynamicDimConstraints
+    ):
+        tensor_metas = copy.deepcopy(tensor_metas)
+        update_tensor_metas_by_dyn_dim_cstr(tensor_metas, dyn_dim_cstr)
+        weight_meta_code = "\n".join(
+            tensor_meta.serialize_to_py_str() for tensor_meta in tensor_metas
+        )
+        import tempfile
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            for filename in ["graph_net.json", "model.py"]:
+                with open(os.path.join(tmpdir, filename), "w") as f:
+                    f.write(open(os.path.join(model_path, filename)).read())
+            with open(os.path.join(tmpdir, "input_meta.py"), "w") as f:
+                f.write("")
+            with open(os.path.join(tmpdir, "weight_meta.py"), "w") as f:
+                f.write(weight_meta_code)
+            return self.model_runnable_predicator(tmpdir)
+
+
+def update_tensor_metas_by_dyn_dim_cstr(
+    tensor_metas: list[TensorMeta], dyn_dim_cstr: DynamicDimConstraints
+):
+    input_shapes = dyn_dim_cstr.get_reified_input_shapes()
+    input_max_values = dyn_dim_cstr.get_reified_input_max_values()
+    assert len(tensor_metas) == len(input_shapes)
+    assert len(tensor_metas) == len(input_max_values)
+    for i, tensor_meta in enumerate(tensor_metas):
+        tensor_meta.shape = input_shapes[i]
+        tensor_meta.max_val = input_max_values[i]
+
+
+def make_dyn_dim_cstr_from_tensor_metas(tensor_metas: list[TensorMeta]):
+    named_shapes = [
+        (shape, name)
+        for tensor_meta in tensor_metas
+        for name in [tensor_meta.name]
+        for shape in [tensor_meta.shape]
+    ]
+    named_max_values = [
+        (max_val, name)
+        for tensor_meta in tensor_metas
+        for name in [tensor_meta.name]
+        for max_val in [tensor_meta.max_val]
+    ]
+    return DynamicDimConstraints.make_by_named_inputs(
+        named_shapes=named_shapes,
+        named_max_values=named_max_values,
+    )
 
 
 def symbolize_data_input_dims(
     dyn_dim_cstr: DynamicDimConstraints,
-    is_data_input: Callable[["input_var_name:str"], bool],
-    is_input_shape_valid: Callable[[DynamicDimConstraints], bool],
+    is_data_input: Callable[[str], bool],
+    is_dyn_dim_cstr_feasible: Callable[[DynamicDimConstraints], bool],
 ) -> DynamicDimConstraints | None:
     """
+    is_data_input: Callable[["input_var_name:str"], bool]
     Symbolizes data input dimensions as much as possible.
     Returns new DynamicDimConstraints if success.
     Returns None if no symbolicable dim .
     """
     unqiue_dims = set()
 
     def dumpy_filter_fn(input_name, input_idx, axis, dim):
-        unqiue_dims.add(dim)
+        if is_data_input(input_name):
+            print("data_input", input_name, input_idx, axis, dim)
+            unqiue_dims.add(dim)
         # No symbolization because of returning True
         return False
 
     # Collect input dimensions into `unqiue_dims`
     assert dyn_dim_cstr.symbolize(dumpy_filter_fn) is None
     for picked_dim in unqiue_dims:
-        tmp_dyn_dim_cstr = copy.deepcopy(dyn_dim_cstr)
+        cur_dyn_dim_cstr = copy.deepcopy(dyn_dim_cstr)
 
         def filter_fn(input_name, input_idx, axis, dim):
             return is_data_input(input_name) and dim == picked_dim
 
-        symbol = tmp_dyn_dim_cstr.symbolize(filter_fn)
+        symbol = cur_dyn_dim_cstr.symbolize(filter_fn)
         if symbol is None:
             continue
         sym2example_value = {symbol: picked_dim + 1}
-        if not tmp_dyn_dim_cstr.check_delta_symbol2example_value(sym2example_value):
+        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_input_shape_valid(tmp_dyn_dim_cstr):
+        if not is_dyn_dim_cstr_feasible(tmp_dyn_dim_cstr):
             continue
-        dyn_dim_cstr = tmp_dyn_dim_cstr
+        dyn_dim_cstr = cur_dyn_dim_cstr
     return dyn_dim_cstr

graph_net/dynamic_dim_constraints.py

@@ -9,7 +9,7 @@
 
 @dataclass
 class DynamicDimConstraints:
-    kSymbolVarNamePrefix = "s"
+    kSymbolVarNamePrefix = "S"
 
     symbols: list[sympy.Symbol]
     kSymbols = "dynamic_dim_constraint_symbols"
@@ -22,27 +22,38 @@ class DynamicDimConstraints:
     kRelations = "dynamic_dim_constraint_relations"
 
     # len(input_shapes) equals number of Model.forward arguments
-    input_shapes: list[("var-name", tuple[sympy.Expr | int])]
+    input_shapes: list[(tuple[sympy.Expr | int], "var-name")]
     kInputShapes = "dynamic_dim_constraint_input_shapes"
 
     # len(input_shapes) equals number of Model.forward arguments
-    input_max_values: list[("var-name", tuple[sympy.Expr | int | None])]
+    input_max_values: list[(tuple[sympy.Expr | int | None], "var-name")]
     kInputMaxValues = "dynamic_dim_constraint_input_max_values"
 
+    @classmethod
+    def make_by_named_inputs(cls, named_shapes, named_max_values):
+        return cls(
+            symbols=[],
+            symbol2example_value={},
+            relations=[],
+            input_shapes=named_shapes,
+            input_max_values=named_max_values,
+        )
+
     def symbolize(
         self,
-        filter_fn=Callable[
-            ["input_name:str", "input_idx:int", "axis:int", "dim:int"], bool
-        ],
+        filter_fn: Callable[[str, int, int, int], bool],
     ) -> sympy.Symbol | None:
         """
+        filter_fn: Callable[
+            ["input_name:str", "input_idx:int", "axis:int", "dim:int"], bool
+        ]
         Returns created symbol.
         """
         InputDim = namedtuple("InputDim", ["input_idx", "axis", "dim"])
         input_dims = [
             InputDim(input_idx, axis, dim)
             for input_idx, namedshape in enumerate(self.input_shapes)
-            for input_name, shape in [namedshape]
+            for shape, input_name in [namedshape]
             for axis, dim in enumerate(shape)
             if isinstance(dim, int)
             if filter_fn(input_name, input_idx, axis, dim)
@@ -54,41 +65,34 @@ def symbolize(
         dim = list(unique_dims)[0]
         new_sym = self._new_symbol(example_value=dim)
         for input_dix, axis, _ in input_dims:
-            self.input_shapes[input_dix][1][axis] = new_sym
+            self.input_shapes[input_dix][0][axis] = new_sym
         return new_sym
 
-    def _new_symbol(self, example_value):
-        max_existed_seq_no = max(
-            [
-                -1,
-                *(
-                    seq_no
-                    for symbol in self.symbols
-                    for seq_no in [int(symbol.name[1:])]
-                ),
-            ]
+    def update_symbol2example_value(self, symbol2example_value: dict):
+        self.symbol2example_value = self._merge_symbol2example_value(
+            symbol2example_value
         )
-        seq_no = max_existed_seq_no + 1
-        symbol = sympy.Symbol(f"{self.kSymbolVarNamePrefix}{seq_no}")
-        self.symbol2example_value[symbol] = example_value
-        self.symbols.append(symbol)
-        return symbol
-
-    def update_symbol2example_value(self, sym2example_value: dict):
-        self.symbol2example_value = self.merge_symbol2example_value(sym2example_value)
         return self
 
-    def merge_symbol2example_value(self, sym2example_value: dict):
-        return {
-            k: v
-            for k, v in [*self.symbol2example_value.items(), *sym2example_value.items()]
-        }
+    def get_reified_input_shapes(self):
+        return [
+            [self._try_reify(dim) for dim in shape] for shape, name in self.input_shapes
+        ]
+
+    def get_reified_input_max_values(self):
+        return [self._try_reify(max_value) for max_value, name in self.input_max_values]
 
-    def check_delta_symbol2example_value(self, sym2example_value: dict):
-        if len(sym2example_value) == 0:
+    def _try_reify(self, dim):
+        if isinstance(dim, sympy.Expr):
+            dim = int(dim.subs(self.symbol2example_value))
+        assert isinstance(dim, (int, type(None))), f"{type(dim)=} {dim=}"
+        return dim
+
+    def check_delta_symbol2example_value(self, symbol2example_value: dict):
+        if len(symbol2example_value) == 0:
             return True
 
-        sym2example_value = self.merge_symbol2example_value(sym2example_value)
+        symbol2example_value = self._merge_symbol2example_value(symbol2example_value)
 
         sym_exprs = [
             *self._get_sym_exprs_from_input_shapes(),
@@ -98,22 +102,7 @@ def check_delta_symbol2example_value(self, sym2example_value: dict):
         relations = [*self.relations, *(sym_expr > 0 for sym_expr in sym_exprs)]
 
         return all(
-            relation.subs(sym2example_value) == sympy.true for relation in relations
-        )
-
-    def _get_sym_exprs_from_input_shapes(self):
-        yield from (
-            sym_dim
-            for name, shape in self.input_shapes
-            for sym_dim in shape
-            if isinstance(sym_dim, sympy.Expr)
-        )
-
-    def _get_sym_exprs_from_input_max_values(self):
-        yield from (
-            sym_value
-            for name, sym_value in self.input_max_values
-            if isinstance(sym_value, sympy.Expr)
+            relation.subs(symbol2example_value) == sympy.true for relation in relations
         )
 
     def serialize_to_py_str(self):
@@ -188,6 +177,47 @@ def load_module(cls, path, name="unamed"):
         spec.loader.exec_module(module)
         return module
 
+    def _new_symbol(self, example_value):
+        max_existed_seq_no = max(
+            [
+                -1,
+                *(
+                    seq_no
+                    for symbol in self.symbols
+                    for seq_no in [int(symbol.name[len(self.kSymbolVarNamePrefix) :])]
+                ),
+            ]
+        )
+        seq_no = max_existed_seq_no + 1
+        symbol = sympy.Symbol(f"{self.kSymbolVarNamePrefix}{seq_no}")
+        self.symbol2example_value[symbol] = example_value
+        self.symbols.append(symbol)
+        return symbol
+
+    def _merge_symbol2example_value(self, symbol2example_value: dict):
+        return {
+            k: v
+            for k, v in [
+                *self.symbol2example_value.items(),
+                *symbol2example_value.items(),
+            ]
+        }
+
+    def _get_sym_exprs_from_input_shapes(self):
+        yield from (
+            sym_dim
+            for shape, name in self.input_shapes
+            for sym_dim in shape
+            if isinstance(sym_dim, sympy.Expr)
+        )
+
+    def _get_sym_exprs_from_input_max_values(self):
+        yield from (
+            sym_value
+            for sym_value, name in self.input_max_values
+            if isinstance(sym_value, sympy.Expr)
+        )
+
 
 if __name__ == "__main__":
     cstr_code = """