flop count: do not ignore scripted function

Reviewed By: bxiong1202

Differential Revision: D32233496

fbshipit-source-id: d7cdbb2e1107bd5a12a952ea44f2eb56355de2df

Author: ppwwyyxx

fvcore/nn/jit_analysis.py

@@ -575,14 +575,8 @@ def _analyze(self) -> "Statistics":
                 ancestors = self._get_all_ancestors(scope_names[-1])
                 all_seen.update(ancestors)
             if kind not in self._op_handles:
-                # Ignore all prim:: operators. However, prim::PythonOp can be
-                # a user-implemented `torch.autograd.Function` so we shouldn't
-                # ignore it.
-                if kind in self._ignored_ops or (
-                    kind.startswith("prim::") and not kind.startswith("prim::PythonOp")
-                ):
+                if self._should_ignore_node(node):
                     continue
-
                 for name in ancestors:
                     unsupported_ops[name][kind] += 1
             else:
@@ -640,3 +634,16 @@ def _has_forward(self, mod_name: str) -> bool:
             if module_type.forward is mod.forward:  # pyre-ignore[16]
                 return False
         return True
+
+    def _should_ignore_node(self, node) -> bool:
+        kind = node.kind()
+        if kind in self._ignored_ops:
+            return True
+        # Ignore all prim:: operators, with two exceptions:
+        # * prim::PythonOp can be a user-implemented `torch.autograd.Function`
+        # * prim::CallFunction an be a call to scripted module/function.
+        if kind.startswith("prim::PythonOp") or kind.startswith("prim::CallFunction"):
+            return False
+        if kind.startswith("prim::"):
+            return True
+        return False

tests/test_flop_count.py

@@ -829,6 +829,21 @@ def forward(self, x):
         )
         self.assertEqual(flop.total(), 42)
 
+    def test_scripted_function(self):
+        # Scripted function is not yet supported. It should produce a warning
+
+        def func(x):
+            return x @ x
+
+        class Mod(nn.Module):
+            def forward(self, x):
+                f = torch.jit.script(func)
+                return f(x * x)
+
+        flop = FlopCountAnalysis(Mod(), (torch.rand(5, 5),))
+        _ = flop.total()
+        self.assertIn("prim::CallFunction", flop.unsupported_ops())
+
 
 class TestFlopCountHandles(unittest.TestCase):
     def _count_function(self, func, inputs, name) -> Tuple[Any, Any]: