[3/N] Merge skipfiles.check rules (#111451)

yanboliang · pytorchmergebot · commit bf01a7b023c1 · 2023-10-22T04:35:15.000Z
This major change in this PR is to consolidate the skipfiles.check rules, the major thing done is merging the original ```FILE_INLINELIST``` with ```SUBMOD_INLINELIST``` into new ```MOD_INLINELIST``` and a legacy ```LEGACY_MOD_INLINELIST```. Let's use the following example to illustrate what is the expected behavior for this force inline list: https://github.com/pytorch/pytorch/blob/fa995626a8e181e3666b27fdb4edbe6116b22ee3/torch/_dynamo/skipfiles.py#L344-L369 The handling logic is: * If f2 is inlined, we will check both ```MOD_INLINELIST``` and ```LEGACY_MOD_INLINELIST``` to consultant force inline rules for f3. * If f2 is skipped, we will check ```LEGACY_MOD_INLINELIST``` only for inline rules for f3. The reason behind this design is: if f2 is skipped, if we always trace all recursively called functions, we will go to the very low level functions (e.g, ```super().__init__```) which caused graph breaks. We treated this as a signal that all functions that f2 recursively called should be skipped as well if f2 is skipped. This is also a feature that many PyTorch developers requested, they just want to skip all recursive functions if they mark the upper level functions as skipped. For PyTorch developers, we should only use ```MOD_INLINELIST``` going forward. I think most of the modules in the ```LEGACY_MOD_INLINELIST``` are legacy things to workaround when we didn't have a good skip/inline API. Pull Request resolved: #111451 Approved by: https://github.com/ezyang
diff --git a/test/dynamo/test_allow_inline_skip.py b/test/dynamo/test_allow_inline_skip.py
@@ -6,9 +6,9 @@
 import torch
 import torch._dynamo.test_case
 from torch._dynamo.skipfiles import (
-    FILE_INLINELIST,
     FUNC_INLINELIST,
-    SUBMODULE_INLINELIST,
+    LEGACY_MOD_INLINELIST,
+    MOD_INLINELIST,
 )
 from torch._dynamo.utils import istype
 
@@ -36,7 +36,7 @@ class AllowInlineSkipTests(torch._dynamo.test_case.TestCase):
     # this unit test is to make sure the functions/modules can be correctly imported
     # or loaded in case there is typo in the strings.
     def test_skipfiles_inlinelist_correctness(self):
-        for m in FILE_INLINELIST.union(SUBMODULE_INLINELIST):
+        for m in LEGACY_MOD_INLINELIST.union(MOD_INLINELIST):
             self.assertTrue(isinstance(importlib.import_module(m), types.ModuleType))
         for f in FUNC_INLINELIST:
             module_name, fn_name = f.rsplit(".", 1)
@@ -54,11 +54,9 @@ def fn(x):
         func_inlinelist.add("torch._dynamo.utils.istype")
 
         self.assertTrue(
-            "torch._dynamo.utils" not in torch._dynamo.skipfiles.FILE_INLINELIST
-        )
-        self.assertTrue(
-            "torch._dynamo" not in torch._dynamo.skipfiles.SUBMODULE_INLINELIST
+            "torch._dynamo" not in torch._dynamo.skipfiles.LEGACY_MOD_INLINELIST
         )
+        self.assertTrue("torch._dynamo" not in torch._dynamo.skipfiles.MOD_INLINELIST)
 
         with unittest.mock.patch(
             "torch._dynamo.skipfiles.get_func_inlinelist",
diff --git a/torch/_dynamo/skipfiles.py b/torch/_dynamo/skipfiles.py
@@ -64,10 +64,10 @@
     * BUILTIN_SKIPLIST contains builtin python modules, such as abc, collections, etc.
     * THIRDPARTY_SKIPLIST contains common third party libraries, such as numpy, pandas, etc.
 * Functions in these two SKIPLISTs are always skipped, except when they are explicitly
-    put into the three INLINELIST: FUNC_INLINELIST, FILE_INLINELIST and SUBMODULE_INLINELIST.
+    put into the two INLINELIST: FUNC_INLINELIST and MOD_INLINELIST.
 * PyTorch(torch) is in the BUILTIN_SKIPLIST by default, but there are many cases
     where we want inline the functions under torch namespace. We should add them
-    into one of the three *_INLINELIST to make dynamo inline those functions.
+    into one of the two *_INLINELIST to make dynamo inline those functions.
 * If you call functions under skipped modules/files, Dynamo will wrap these functions
     as SkipFilesVariable. There are a few functions(e.g, collections.OrderedDict) that
     we have special handling at SkipFilesVariable.call_function.
@@ -76,17 +76,14 @@
 
 To figure out what the behavior is, check the following list in order:
 * FUNC_INLINELIST (Inline if YES)
-* FILE_INLINELIST (Inline if YES)
-* SUBMODULE_INLINELIST (Inline if YES)
+* MOD_INLINELIST (Inline if YES)
 * BUILTIN_SKIPLIST & THIRDPARTY_SKIPLIST (Skip if YES)
 * Inline by default
 
 In general, if you want to force inline a function or module, please consider adding
-the function's file or python module to FILE_INLINELIST first.
-Use the FUNC_INLINELIST only when there are other functions under the same file that
-you don't want to inline.
-In the future, we will consolidate FILE_INLINELIST and SUBMODULE_INLINELIST into one list
-as we use the same logic (filename.startswith) to determine if a file or module is skipped.
+the function's python module to MOD_INLINELIST first.
+Use the FUNC_INLINELIST only when there are other functions under the same module that
+you don't want to inline them.
 """
 
 
@@ -114,7 +111,7 @@
     tempfile,
     threading,
     tokenize,
-    torch,  # torch/* is skipped by default unless specified in FILE_INLINELIST or SUBMODULE_INLINELIST
+    torch,  # torch/* is skipped by default unless specified in FUNC_INLINELIST or MOD_INLINELIST
     traceback,
     types,
     typing,
@@ -165,36 +162,25 @@ def _module_dir(m: types.ModuleType):
 }
 
 
-# Force inline functions in these files or directories, even they are in *_SKIPLIST.
-# We are using python module name instead of file or directory object to avoid circular dependency.
-# Please keep this sorted alphabetically.
-# TODO: Merge FILE_INLINELIST into SUBMODULE_INLINELIST.
-FILE_INLINELIST = {
-    "torch._dynamo._trace_wrapped_higher_order_op",
-    "torch._dynamo.comptime",
+# These are legacy workarounds, don't add new modules to this list.
+# Please use the MOD_INLINELIST instead to force inline functions under particular modules.
+LEGACY_MOD_INLINELIST = {
     "torch._dynamo.external_utils",
-    "torch._dynamo.polyfill",
     "torch._export.db.examples",
     "torch._export.wrappers",
     "torch._functorch.apis",
     "torch._functorch.deprecated",
     "torch._higher_order_ops.cond",
-    "torch._inductor.test_operators",
     "torch.ao.quantization.pt2e.eval_utils",
     "torch.ao.quantization.pt2e.qat_utils",
     "torch.ao.quantization.pt2e.representation.rewrite",
     "torch.ao.quantization.pt2e.utils",
     "torch.ao.quantization.quantizer.xnnpack_quantizer",
-    "torch.nn.modules.container",
-    "torch.optim._functional",
-    "torch.random",
-    "torch.utils._content_store",
-    "torch.utils._foreach_utils",
+    "torch.optim",
 }
 
-
 if torch.distributed.is_available():
-    FILE_INLINELIST |= {
+    LEGACY_MOD_INLINELIST |= {
         "torch.distributed._tensor.api",
         "torch.distributed._tensor.device_mesh",
         "torch.distributed.algorithms._checkpoint.checkpoint_wrapper",
@@ -203,31 +189,35 @@ def _module_dir(m: types.ModuleType):
         "torch.distributed.tensor.parallel.style",
     }
 
-# Include optimizer code for tracing
-FILE_INLINELIST |= {
-    str(obj.__module__) for obj in torch.optim.__dict__.values() if inspect.isclass(obj)
-}
 
-# TODO: consolidate SUBMODULE_INLINELIST and FILE_INLINELIST into one list
-# Force inline functions under these modules, even the modules is in *_SKIPLIST.
-SUBMODULE_INLINELIST = {
+# Force inline functions under these modules, even they are in *_SKIPLIST.
+# We are using python module name instead of file or directory object to avoid circular dependency.
+# Please keep this sorted alphabetically.
+MOD_INLINELIST = {
     "torch._refs",
     "torch._prims",
     "torch._decomp",
+    "torch._dynamo._trace_wrapped_higher_order_op",
+    "torch._dynamo.comptime",
+    "torch._dynamo.polyfill",
+    "torch._inductor.test_operators",
     "torch.ao.nn",
     "torch.distributions",
     "torch.fx._pytree",
     "torch.nn",
+    "torch.random",
     "torch.sparse",
     "torch.testing",
+    "torch.utils._content_store",
     "torch.utils._contextlib",
+    "torch.utils._foreach_utils",
     "torch.utils._pytree",
 }
 
 
 if torch.distributed.is_available():
-    SUBMODULE_INLINELIST.add("torch.distributed")
-    SUBMODULE_INLINELIST.add("torch.distributed._functional_collectives")
+    MOD_INLINELIST.add("torch.distributed")
+    MOD_INLINELIST.add("torch.distributed._functional_collectives")
 
 
 # TODO: support adding bound method into this list
@@ -243,17 +233,17 @@ def get_func_inlinelist():
 
 
 @functools.lru_cache(None)
-def get_file_inlinelist():
+def get_legacy_mod_inlinelist():
     inlinelist = set()
-    for f in FILE_INLINELIST:
-        inlinelist.add(_module_dir(torch) + f[len("torch.") :].replace(".", "/"))
+    for m in LEGACY_MOD_INLINELIST:
+        inlinelist.add(_module_dir(torch) + m[len("torch.") :].replace(".", "/"))
     return inlinelist
 
 
 @functools.lru_cache(None)
-def get_submodule_inlinelist():
+def get_mod_inlinelist():
     inlinelist = set()
-    for m in SUBMODULE_INLINELIST:
+    for m in MOD_INLINELIST:
         inlinelist.add(_module_dir(torch) + m[len("torch.") :].replace(".", "/"))
     return inlinelist
 
@@ -303,21 +293,19 @@ class SkipResult:
     reason: Optional[str]
 
 
-# TODO(ybliang): This is a temp function, we should consolidate this with check_file.
-def _check_file_inner(filename, allow_torch=False):
+def check_file(filename, allow_torch=False):
     """Should skip this file?"""
     if filename is None:
         return SkipResult(True, "filename is None")
-    if any(filename.startswith(d) for d in get_file_inlinelist()):
+    if any(filename.startswith(d) for d in get_legacy_mod_inlinelist()):
         return SkipResult(
             False,
-            "inlined according skipfiles.FILE_INLINELIST",
+            "inlined according skipfiles.LEGACY_MOD_INLINELIST",
         )
-    # TODO(ybliang): the is_torch check should be consolidate with is_torch_inline_allowed
-    if allow_torch and is_torch(filename):
+    if allow_torch and is_torch_inline_allowed(filename):
         return SkipResult(
             False,
-            "inlined according skipfiles.is_torch",
+            "inlined according skipfiles.MOD_INLINELIST",
         )
     if is_fbcode and bool(FBCODE_SKIP_DIRS_RE.match(filename)):
         return SkipResult(
@@ -330,17 +318,6 @@ def _check_file_inner(filename, allow_torch=False):
         return SkipResult(False, "inlined by default")
 
 
-def check_file(filename, allow_torch=False, extra_check=False):
-    result = _check_file_inner(filename, allow_torch)
-    if extra_check and result.skipped and is_torch_inline_allowed(filename):
-        return SkipResult(
-            False,
-            "inlined according skipfiles.is_torch_inline_allowed returning True",
-        )
-    else:
-        return result
-
-
 """
 This is the main entry point to determine whether an object (function) should be inlined or skipped.
 Let's illustrate the logic with an example:
@@ -366,10 +343,13 @@ def f3(x, y):
     * If f2 is skipped by Dynamo, when evaluating the frame of f3, Dynamo need the inline/skip check again
       and the call site is in catch_errors_wrapper.catch_errors of eval_frame.py.
 * For global variables and function arguments, Dynamo needs to decide if they are wrapped as SkipFilesVariable in builder.py.
+
+allow_torch is used to indicate whether we are checking the MOD_INLINELIST (torch modules), we only do this check when
+f2 is not skipped.
 """
 
 
-def check_verbose(obj, allow_torch=False, extra_check=False):
+def check_verbose(obj, allow_torch=False):
     if isinstance(
         obj, (UserFunctionVariable, UserMethodVariable, NestedUserFunctionVariable)
     ):
@@ -387,11 +367,11 @@ def check_verbose(obj, allow_torch=False, extra_check=False):
             False,
             "inlined according skipfiles.FUNC_INLINELIST",
         )
-    return check_file(filename, allow_torch, extra_check)
+    return check_file(filename, allow_torch)
 
 
-def check(obj, allow_torch=False, extra_check=False):
-    return check_verbose(obj, allow_torch, extra_check).skipped
+def check(obj, allow_torch=False):
+    return check_verbose(obj, allow_torch).skipped
 
 
 # skip common third party libs
@@ -402,7 +382,7 @@ def check(obj, allow_torch=False, extra_check=False):
 
 
 def is_torch_inline_allowed(filename):
-    return any(filename.startswith(d) for d in get_submodule_inlinelist())
+    return any(filename.startswith(d) for d in get_mod_inlinelist())
 
 
 @functools.lru_cache(None)
diff --git a/torch/_dynamo/symbolic_convert.py b/torch/_dynamo/symbolic_convert.py
@@ -2257,7 +2257,7 @@ def check_inlineable(func):
         except NotImplementedError:
             pass  # closures
 
-        result = skipfiles.check_verbose(func, extra_check=True)
+        result = skipfiles.check_verbose(func, allow_torch=True)
         if result.skipped:
             from torch._dynamo.variables.misc import (
                 produce_trampoline_autograd_apply,
diff --git a/torch/_dynamo/variables/builder.py b/torch/_dynamo/variables/builder.py
@@ -727,6 +727,7 @@ def index_source(key):
             istype(value, (type, types.FunctionType))
             and skipfiles.check(value, allow_torch=True)
             and not inspect.getattr_static(value, "_torchdynamo_inline", False)
+            and not inspect.getattr_static(value, "__script_if_tracing_wrapper", False)
         ):
             return SkipFilesVariable(
                 value,
