[SOT] fix sot export bugs (#60707)

feifei-111 · web-flow · commit f31985185dbb · 2024-01-23T14:55:29.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/function_graph.py b/python/paddle/jit/sot/opcode_translator/executor/function_graph.py
@@ -108,16 +108,23 @@ def func(x):
     return map_variables(func, inputs)
 
 
-def get_symbol_meta_map(inputs):
-    output = {}
+def record_symbols(SIR, *args, **kwargs):
+    symbol_meta_map = {}
+    params = set()
+    non_params = set()
 
-    def func(x):
-        if isinstance(x, TensorVariable):
-            output[x.get_symbol()] = x.meta
-        return x
+    def fn(value):
+        if isinstance(value, TensorVariable):
+            symbol_meta_map[value.get_symbol()] = value.meta
+            if isinstance(value, ParameterVariable):
+                params.add(value.get_symbol())
+            else:
+                non_params.add(value.get_symbol())
+        return value
 
-    map_variables(func, inputs)
-    return output
+    map_variables(fn, [args, kwargs])
+    SIR.set_symbol_meta_map(symbol_meta_map)
+    SIR.set_parameter_info(params, non_params)
 
 
 def get_params_and_non_param_symbol(*args, **kwargs):
@@ -573,10 +580,7 @@ def symbolic_call(self, infer_meta_fn, compute_fn, func, *args, **kwargs):
             convert_to_symbol(kwargs),
         )
 
-        self.sir_ctx.TOS.set_symbol_meta_map(get_symbol_meta_map(args))
-        self.sir_ctx.TOS.set_symbol_meta_map(get_symbol_meta_map(kwargs))
-        params, non_params = get_params_and_non_param_symbol(*args, **kwargs)
-        self.sir_ctx.TOS.set_parameter_info(params, non_params)
+        record_symbols(self.sir_ctx.TOS, *args, **kwargs)
 
         log(3, f"         inputs : {inputs_symbols}", "\n")
 
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py b/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py
@@ -769,7 +769,7 @@ def call_function(self, /, *args, **kwargs):
         assert self.check_no_weight_and_buffers(
             new_layer
         ), "You have created a layer in to_static function which may have Potential bugs. please create it in __init__/main function."
-        return PaddleLayerVariable(
+        return VariableFactory.from_value(
             new_layer, self.graph, CreateLayerTracker(self, args, kwargs)
         )
 
diff --git a/python/paddle/jit/sot/symbolic/compile_cache.py b/python/paddle/jit/sot/symbolic/compile_cache.py
@@ -65,6 +65,7 @@ def __init__(self, compiled_fn, SIR, is_training: bool):
         self.concrete_program = None
         self.SIR = SIR  # for debug
         self.is_training = is_training
+        self.exported = False
 
     def amp_cast_inputs(self, args, kwargs):
         """Prepare inputs for amp, cast float16 into float32 if needed."""
@@ -145,8 +146,9 @@ def __call__(self, *args, **kwargs):
                 4,
                 lambda: print("[CompileCache] run sir forward success."),
             )
-            if ENV_SOT_EXPORT.get() != "":
+            if ENV_SOT_EXPORT.get() != "" and not self.exported:
                 export(self.SIR, ENV_SOT_EXPORT.get())
+                self.exported = True
 
             return outputs
 
diff --git a/python/paddle/jit/sot/symbolic/export.py b/python/paddle/jit/sot/symbolic/export.py
@@ -15,6 +15,7 @@
 import os
 from itertools import chain
 
+import paddle
 from paddle.utils import flatten
 
 from ..utils import ConstTypes, ExportError, NameGenerator
@@ -43,14 +44,74 @@ def __str__(self):
         return "\n".join(self.get_lines())
 
 
+class NameGener:
+    def __init__(self, SIR):
+        self.SIR = SIR
+        self.name_map = {}
+        self.param_name_generator = NameGenerator("self.parameter_")
+        self.non_param_name_generator = NameGenerator("var_")
+
+    def __call__(self, var):
+        return self.get_str(var)
+
+    def get_str(self, var):
+        if isinstance(var, list):
+            return self.get_list_str(var)
+        elif isinstance(var, tuple):
+            return self.get_tuple_str(var)
+        elif isinstance(var, dict):
+            return self.get_dict_str(var)
+        elif isinstance(var, set):
+            return self.get_set_str(var)
+        else:
+            return self.get_obj_str(var)
+
+    def get_list_str(self, list_):
+        return "[{}]".format(", ".join(self.get_str(var) for var in list_))
+
+    def get_tuple_str(self, tuple_):
+        return "({},)".format(", ".join(self.get_str(var) for var in tuple_))
+
+    def get_dict_str(self, dict_):
+        return "{{{},}}".format(
+            ", ".join(
+                f"{self.get_str(k)}: {self.get_str(v)}"
+                for k, v in dict_.items()
+            )
+        )
+
+    def get_set_str(self, set_):
+        return "{{{},}}".format(", ".join(self.get_str(var) for var in set_))
+
+    def get_obj_str(self, var):
+        if isinstance(var, Symbol):
+            if var not in self.name_map:
+                self.register_symbol(var)
+            return self.name_map[var]
+
+        elif isinstance(var, str):
+            return f"'{var}'"
+        else:
+            return str(var)
+
+    def register_symbol(self, symbol):
+        if symbol in self.SIR.param_symbol:
+            name = self.param_name_generator.next()
+        else:
+            name = self.non_param_name_generator.next()
+        self.name_map[symbol] = name
+
+
 class PyFileGen:
     def __init__(self, SIR):
         self.SIR = SIR
         self.roots = []
 
-        self.layer_name_map = {}
-        self.layer_name_generator = NameGenerator("_")
-        self.SIR_name = SIR.name.replace("_", "")
+        self.name_gener = NameGener(self.SIR)
+
+        self.SIR_sig = "||".join(
+            f"{stmt.type}:{stmt.name}" for stmt in SIR.statements
+        )
 
     def new_root(self, *args):
         stmt = PyStatement(*args)
@@ -93,21 +154,22 @@ def check_exportable(self):
 
     def create_header(self):
         self.new_root(
+            f"# {self.SIR_sig}",
             "import paddle",
             "import unittest",
             "import numpy as np",
         )
 
     def create_layer(self):
-        layer_class = self.new_root(f"class {self.SIR_name}(paddle.nn.Layer):")
+        layer_class = self.new_root("class LayerCase(paddle.nn.Layer):")
 
         init_fn = layer_class.add_sub("def __init__(self):")
         init_fn.add_sub("super().__init__()")
 
         for param in self.SIR.param_symbol:
-            meta = self.SIR.symbol_meta_map[param.name]
+            meta = self.SIR.symbol_meta_map[param]
             init_fn.add_sub(
-                f"self.{param.name} = self.create_parameter(",
+                f"{self.name_gener(param)} = self.create_parameter(",
                 f"   shape={meta.shape},",
                 f"   dtype={meta.dtype},",
                 ")",
@@ -116,20 +178,16 @@ def create_layer(self):
         for stmt in self.SIR.statements:
             if stmt.type == "layer":
                 layer = stmt.layer()
-                if id(layer) not in self.layer_name_map:
-                    layer_name = (
-                        layer.__class__.__name__
-                        + self.layer_name_generator.next()
-                    )
-                    self.layer_name_map[id(layer)] = layer_name
-                    init_fn.add_sub(self.init_sub_layer(layer, layer_name))
+                init_fn.add_sub(self.init_sub_layer(layer))
 
         forward_definition = ["def forward(", "    self,"]
 
         for inp in self.SIR.inputs:
             if inp in self.SIR.non_param_symbol:
-                meta = self.SIR.symbol_meta_map[inp.name]
-                forward_definition.append(f"    {inp.name},    # {str(meta)}")
+                meta = self.SIR.symbol_meta_map[inp]
+                forward_definition.append(
+                    f"    {self.name_gener(inp)},    # {str(meta)}"
+                )
         forward_definition.append("):")
 
         forward_fn = layer_class.add_sub(*forward_definition)
@@ -139,14 +197,12 @@ def create_layer(self):
 
         forward_fn.add_sub(
             "return {}".format(
-                ", ".join(self.true_name(out) for out in self.SIR.outputs)
+                ", ".join(self.name_gener(out) for out in self.SIR.outputs)
             )
         )
 
     def create_test(self):
-        test_class = self.new_root(
-            f"class Test{self.SIR_name}(unittest.TestCase):"
-        )
+        test_class = self.new_root("class TestLayer(unittest.TestCase):")
 
         setup = test_class.add_sub("def setUp(self):")
         test_inputs = [
@@ -155,16 +211,29 @@ def create_test(self):
         for inp in self.SIR.inputs:
             if inp in self.SIR.non_param_symbol:
                 meta = self.SIR.symbol_meta_map[inp.name]
-                test_inputs.append(
-                    f"    paddle.rand(shape={meta.shape}, dtype={meta.dtype}),"
-                )
+                shape_str = "[1]" if len(meta.shape) == 0 else str(meta.shape)
+                if meta.dtype in (
+                    paddle.int8,
+                    paddle.int16,
+                    paddle.int32,
+                    paddle.int64,
+                ):
+                    test_inputs.append(
+                        f"    paddle.randint(low=0, high=10, shape={shape_str}, dtype={meta.dtype}),"
+                    )
+                else:
+                    test_inputs.append(
+                        f"    paddle.rand(shape={shape_str}, dtype={meta.dtype}),"
+                    )
         test_inputs.append(")")
         setup.add_sub(*test_inputs)
+        setup.add_sub("self.net = LayerCase()")
 
         train = test_class.add_sub(
-            "def train(self, net, to_static, with_cinn=False):"
+            "def train(self, net, to_static, with_prim=False, with_cinn=False):"
         )
         train.add_sub(
+            "paddle.set_flags({'FLAGS_prim_all': with_prim})",
             "if to_static:",
             "    if with_cinn:",
             "        build_strategy = paddle.static.BuildStrategy()",
@@ -176,24 +245,14 @@ def create_test(self):
             "return outs",
         )
 
-        test_ast_static = test_class.add_sub("def test_ast_static(self):")
-        test_ast_static.add_sub(
-            "net = SIR0()",
-            "dy_out = self.train(net, to_static=False)",
-            "st_out = self.train(net, to_static=True, with_cinn=False)",
-            "for dy, st in zip(paddle.utils.flatten(dy_out), paddle.utils.flatten(st_out)):",
-            "    np.testing.assert_allclose(dy.numpy(), st.numpy(), atol=1e-8)",
-        )
-
         test_ast_cinn_static = test_class.add_sub(
-            "def test_ast_cinn_static(self):"
+            "def test_ast_prim_cinn(self):"
         )
         test_ast_cinn_static.add_sub(
-            "net = SIR0()",
-            "dy_out = self.train(net, to_static=False)",
-            "st_out = self.train(net, to_static=True, with_cinn=True)",
-            "for dy, st in zip(paddle.utils.flatten(dy_out), paddle.utils.flatten(st_out)):",
-            "    np.testing.assert_allclose(dy.numpy(), st.numpy(), atol=1e-8)",
+            "st_out = self.train(self.net, to_static=True)",
+            "cinn_out = self.train(self.net, to_static=True, with_prim=True, with_cinn=True)",
+            "for st, cinn in zip(paddle.utils.flatten(st_out), paddle.utils.flatten(cinn_out)):",
+            "    np.testing.assert_allclose(st.numpy(), cinn.numpy(), atol=1e-8)",
         )
 
     def create_tail(self):
@@ -202,15 +261,6 @@ def create_tail(self):
             "    unittest.main()",
         )
 
-    def true_name(self, var):
-        if isinstance(var, Symbol):
-            if var in self.SIR.param_symbol:
-                return "self." + var.name
-            else:
-                return var.name
-        else:
-            return str(var)
-
     def init_sub_layer(self, layer, layer_name):
         # TODO @wuzhanfei need more effecient way to create a sub layer
         # now, we just close call_Layer behavior
@@ -219,8 +269,8 @@ def init_sub_layer(self, layer, layer_name):
     def create_input_string(self, args, kwargs):
         return ", ".join(
             chain(
-                (self.true_name(arg) for arg in args),
-                (f"{k}={self.true_name(v)}" for k, v in kwargs.items()),
+                (self.name_gener(arg) for arg in args),
+                (f"{k}={self.name_gener(v)}" for k, v in kwargs.items()),
             )
         )
 
@@ -234,7 +284,7 @@ def search(outputs, path, result):
             elif isinstance(outputs, dict):
                 search_dict(outputs, path, result)
             elif isinstance(outputs, Symbol):
-                result.append(self.true_name(outputs) + " = " + "".join(path))
+                result.append(self.name_gener(outputs) + " = " + "".join(path))
 
         def search_sequnce(outputs, path, result):
             for idx, out in enumerate(outputs):
@@ -260,7 +310,7 @@ def create_api_stmt(self, stmt):
         api = stmt.api
         api_str = api.__module__ + "." + api.__name__
         if isinstance(stmt.outputs, Symbol):
-            return [f"{stmt.outputs.name} = {api_str}({input_str})"]
+            return [f"{self.name_gener(stmt.outputs)} = {api_str}({input_str})"]
         else:
             compute_code = f"out = {api_str}({input_str})"
             unpack_codes = self.create_unpack_output_string(stmt.outputs)
@@ -269,36 +319,28 @@ def create_api_stmt(self, stmt):
     def create_method_stmt(self, stmt):
         args, kwargs = stmt.inputs
         input_str = self.create_input_string(args[1:], kwargs)
-        method_str = args[0].name + "." + stmt.method
+        method_str = self.name_gener(args[0]) + "." + stmt.method
         if isinstance(stmt.outputs, Symbol):
-            return [f"{stmt.outputs.name} = {method_str}({input_str})"]
+            return [
+                f"{self.name_gener(stmt.outputs)} = {method_str}({input_str})"
+            ]
         else:
             compute_code = f"out = {method_str}({input_str})"
             unpack_codes = self.create_unpack_output_string(stmt.outputs)
             return [compute_code] + unpack_codes
 
-    def create_layer_stmt(self, stmt):
-        args, kwargs = stmt.inputs
-        input_str = self.create_input_string(args, kwargs)
-        layer_str = "self." + self.layer_name_map[id(stmt.layer())]
-        if isinstance(stmt.outputs, Symbol):
-            return [f"{stmt.outputs.name} = {layer_str}({input_str})"]
-        else:
-            compute_code = f"out = {layer_str}({input_str})"
-            unpack_codes = self.create_unpack_output_string(stmt.outputs)
-            return [compute_code] + unpack_codes
-
 
 def export(SIR, path):
     try:
         pygen = PyFileGen(SIR)
         string = pygen.gen_py_codes()
     except ExportError as e:
-        print("[SOT] Export SIR Failed:", e)
+        print(f"[SOT] Export {SIR.name} Failed:", e)
         return
 
     if not os.path.exists(path):
         os.makedirs(path)
 
     with open(os.path.join(path, f"{SIR.name}.py"), "w") as f:
         f.write(string)
+        print(f"[SOT] Export {SIR.name} Sucess with size {len(SIR.statements)}")
diff --git a/python/paddle/jit/sot/symbolic/statement_ir.py b/python/paddle/jit/sot/symbolic/statement_ir.py
@@ -135,7 +135,7 @@ def __init__(
         stacks: list[str],
     ):
         super().__init__(
-            "api", "paddle." + api.__name__, inputs, outputs, stacks
+            "api", api.__module__ + "." + api.__name__, inputs, outputs, stacks
         )
         self.api = api
 
@@ -228,6 +228,8 @@ def __deepcopy__(self, memo=None):
         new_sir.outputs = list(self.outputs)
         new_sir.statements = list(self.statements)
         new_sir.symbol_meta_map = dict(self.symbol_meta_map.items())
+        new_sir.param_symbol = set(self.param_symbol)
+        new_sir.non_param_symbol = set(self.non_param_symbol)
         return new_sir
 
     def set_parameter_info(self, params, non_params):
diff --git a/test/sot/test_sot_export.py b/test/sot/test_sot_export.py

Original file line number	Diff line number	Diff line change
`@@ -769,7 +769,7 @@ def call_function(self, /, args, *kwargs):`
`769`	`769`	`assert self.check_no_weight_and_buffers(`
`770`	`770`	`new_layer`
`771`	`771`	`), "You have created a layer in to_static function which may have Potential bugs. please create it in __init__/main function."`
`772`		`- return PaddleLayerVariable(`
	`772`	`+ return VariableFactory.from_value(`
`773`	`773`	`new_layer, self.graph, CreateLayerTracker(self, args, kwargs)`
`774`	`774`	`)`
`775`	`775`