Comments

Author: daniil-lyakhov

nncf/experimental/torch/fx/constant_folding.py

@@ -248,30 +248,29 @@ def constant_fold(
     :param constraint_fn: Constraint function which takes a node and returs the constraint:
         should the node be constant folded or not.
     """
-    with torch.no_grad():
-        with torch.utils._python_dispatch._disable_current_modes():
-            cf = ConstantFolder(gm)
-            cf.run()
-
-            device = get_model_device(gm)
-            for node, constant in cf.node_replacements.items():
-                if constraint_fn is not None and not constraint_fn(node):
-                    continue
-                constant = constant.to(device)
-                _replace_node_with_constant(gm, node, constant)
-
-            erased_params = []
-            for node in gm.graph.find_nodes(op="get_attr"):
-                if len(node.users) == 0:
-                    if hasattr(gm, node.target):
-                        delattr(gm, node.target)
-                    erased_params.append(node)
-
-            for node in erased_params:
-                gm.graph.erase_node(node)
-
-            # Custom _is_impure function allows to eliminate all layers with zero
-            # users including inplace ops like relu_ besides output and placeholders.
-            gm.graph.eliminate_dead_code(_is_impure)
-            gm.graph.lint()
-            gm.recompile()
+    with torch.no_grad(), torch.utils._python_dispatch._disable_current_modes():
+        cf = ConstantFolder(gm)
+        cf.run()
+
+        device = get_model_device(gm)
+        for node, constant in cf.node_replacements.items():
+            if constraint_fn is not None and not constraint_fn(node):
+                continue
+            constant = constant.to(device)
+            _replace_node_with_constant(gm, node, constant)
+
+        erased_params = []
+        for node in gm.graph.find_nodes(op="get_attr"):
+            if len(node.users) == 0:
+                if hasattr(gm, node.target):
+                    delattr(gm, node.target)
+                erased_params.append(node)
+
+        for node in erased_params:
+            gm.graph.erase_node(node)
+
+        # Custom _is_impure function allows to eliminate all layers with zero
+        # users including inplace ops like relu_ besides output and placeholders.
+        gm.graph.eliminate_dead_code(_is_impure)
+        gm.graph.lint()
+        gm.recompile()

tests/torch/data/reference_graphs/fx/transformed/folded_scalar_clone_model.dot

@@ -1,20 +1,15 @@
 strict digraph  {
 "0 linear_weight" [id=0, type=get_attr];
 "1 linear_bias" [id=1, type=get_attr];
-"2 lifted_tensor_0" [id=2, type=get_attr];
-"3 x" [id=3, type=input];
-"4 lift_fresh_copy" [id=4, type=lift_fresh_copy];
-"5 detach_" [id=5, type=detach_];
-"6 _frozen_param0" [id=6, type=get_attr];
-"7 linear" [id=7, type=linear];
-"8 add" [id=8, type=add];
-"9 output" [id=9, type=output];
-"0 linear_weight" -> "7 linear"  [label="(3, 3)", style=solid];
-"1 linear_bias" -> "7 linear"  [label="(3,)", style=solid];
-"2 lifted_tensor_0" -> "4 lift_fresh_copy"  [label="()", style=solid];
-"3 x" -> "7 linear"  [label="(1, 3, 3, 3)", style=solid];
-"4 lift_fresh_copy" -> "5 detach_"  [label="()", style=solid];
-"6 _frozen_param0" -> "8 add"  [label="()", style=solid];
-"7 linear" -> "8 add"  [label="(1, 3, 3, 3)", style=solid];
-"8 add" -> "9 output"  [label="(1, 3, 3, 3)", style=solid];
+"2 x" [id=2, type=input];
+"3 _frozen_param0" [id=3, type=get_attr];
+"4 linear" [id=4, type=linear];
+"5 add" [id=5, type=add];
+"6 output" [id=6, type=output];
+"0 linear_weight" -> "4 linear"  [label="(3, 3)", style=solid];
+"1 linear_bias" -> "4 linear"  [label="(3,)", style=solid];
+"2 x" -> "4 linear"  [label="(1, 3, 3, 3)", style=solid];
+"3 _frozen_param0" -> "5 add"  [label="()", style=solid];
+"4 linear" -> "5 add"  [label="(1, 3, 3, 3)", style=solid];
+"5 add" -> "6 output"  [label="(1, 3, 3, 3)", style=solid];
 }