Refactor torchdynamo imports to be relative (#1494)

Author: jansel

.flake8

@@ -1,3 +1,4 @@
 [flake8]
+exclude = minifier_launcher.py,repro.py
 ignore = E203,W503,C101,C4,EXE,B,Y
 max-line-length = 120

Makefile

@@ -1,8 +1,10 @@
 .PHONY: default develop test torchbench format lint setup clean
 
 PY_FILES := $(wildcard *.py) $(wildcard torchdynamo/*.py) $(wildcard torchdynamo/*/*.py) \
-            $(wildcard test/*.py) $(wildcard torchinductor/*.py) $(wildcard torchinductor/*/*.py) \
-            $(wildcard benchmarks/*.py) $(wildcard benchmarks/*/*.py) $(wildcard .circleci/*.py)
+            $(wildcard torchinductor/*.py) $(wildcard torchinductor/*/*.py)  \
+            $(wildcard benchmarks/*.py) $(wildcard benchmarks/*/*.py)  \
+            $(wildcard test/*.py) $(wildcard test/*/*.py)  \
+            $(wildcard .circleci/*.py)
 C_FILES := $(wildcard torchdynamo/*.c torchdynamo/*.cpp)
 CLANG_TIDY ?= clang-tidy-10
 CLANG_FORMAT ?= clang-format-10
@@ -74,7 +76,6 @@ clone-deps:
 		&& (test -e pytorch || git clone --recursive https://github.com/pytorch/pytorch pytorch) \
 		&& (test -e torchvision || git clone --recursive https://github.com/pytorch/vision torchvision) \
 		&& (test -e torchtext || git clone --recursive https://github.com/pytorch/text torchtext) \
-		&& (test -e torchaudio || git clone --recursive https://github.com/pytorch/audio torchaudio) \
 		&& (test -e detectron2 || git clone --recursive https://github.com/facebookresearch/detectron2) \
 		&& (test -e torchbenchmark || git clone --recursive https://github.com/pytorch/benchmark torchbenchmark) \
 		&& (test -e triton || git clone --recursive https://github.com/openai/triton.git) \
@@ -84,7 +85,6 @@ pull-deps:
 	(cd ../pytorch        && git pull && git submodule update --init --recursive)
 	(cd ../torchvision    && git pull && git submodule update --init --recursive)
 	(cd ../torchtext      && git pull && git submodule update --init --recursive)
-	(cd ../torchaudio     && git pull && git submodule update --init --recursive)
 	(cd ../detectron2     && git pull && git submodule update --init --recursive)
 	(cd ../torchbenchmark && git pull && git submodule update --init --recursive)
 	(cd ../triton         && git checkout master && git pull && git checkout $(TRITON_VERSION) && git submodule update --init --recursive)
@@ -102,7 +102,6 @@ build-deps: clone-deps
 	(cd ../pytorch     && python setup.py clean && python setup.py develop)
 	(cd ../torchvision && python setup.py clean && python setup.py develop)
 	(cd ../torchtext   && python setup.py clean && python setup.py develop)
-	(cd ../torchaudio  && python setup.py clean && python setup.py develop)
 	(cd ../detectron2  && python setup.py clean && python setup.py develop)
 	(cd ../torchbenchmark && python install.py --continue_on_fail)
 	(cd ../triton/python && python setup.py clean && python setup.py develop)

test/mock_modules/mock_module2.py

@@ -14,6 +14,6 @@ def method2(self, x):
 
 
 def method1(x, y):
-    z = torch.ones(1, 1)
+    z = torch.ones(1, 1)  # noqa
     x.append(y)
     return x

test/mock_modules/mock_module3.py

@@ -2,6 +2,6 @@
 
 
 def method1(x, y):
-    z = torch.ones(1, 1)
+    z = torch.ones(1, 1)  # noqa
     x.append(y)
     return x

torchdynamo/codegen.py

@@ -7,8 +7,6 @@
 
 import torch.nn
 
-import torchdynamo
-
 from .bytecode_transformation import Instruction
 from .bytecode_transformation import create_instruction
 from .exc import unimplemented
@@ -42,7 +40,7 @@ class PyCodegen(object):
 
     def __init__(
         self,
-        tx: "torchdynamo.symbolic_convert.InstructionTranslator" = None,
+        tx=None,
         root: torch.nn.Module = None,
         graph_output_var: str = None,
         tempvars=None,

torchdynamo/config.py

@@ -7,8 +7,6 @@
 
 import torch
 
-import torchdynamo.utils
-
 try:
     import torch._prims
     import torch._refs
@@ -141,21 +139,28 @@
 # to allow DDP comm/compute overlap
 optimize_ddp = False
 
-
 # If True, raises exception if TorchDynamo is called with a context manager
 raise_on_ctx_manager_usage = True
 
 # If True, raise when aot autograd is unsafe to use
 raise_on_unsafe_aot_autograd = False
 
+# How to import torchdynamo, either torchdynamo or torch.dynamo
+dynamo_import = __name__.replace(".config", "")
+
+# How to import torchinductor, either torchinductor or torch.inductor
+inductor_import = dynamo_import.replace("dynamo", "inductor")
+
 
 class _AccessLimitingConfig(ModuleType):
     def __setattr__(self, name, value):
         if name not in _allowed_config_names:
             raise AttributeError(f"{__name__}.{name} does not exist")
         # automatically set logger level whenever config.log_level is modified
         if name == "log_level":
-            torchdynamo.utils.set_loggers_level(value)
+            from .utils import set_loggers_level
+
+            set_loggers_level(value)
         return object.__setattr__(self, name, value)
 
 

torchdynamo/convert_frame.py

@@ -199,7 +199,7 @@ def replay_record_msg():
             and record_filename is not None
         ):
             return f"\nLast frame execution written to {record_filename}. To run only this frame while debugging, run\
- torchdynamo.replay('{record_filename}').\n"
+ {config.dynamo_import}.replay('{record_filename}').\n"
         else:
             return ""
 
@@ -238,7 +238,9 @@ def replay_record_msg():
 
         msg += replay_record_msg()
 
-        msg += "\nSet torchdynamo.config.verbose=True for more information\n"
+        msg += (
+            f"\nSet {config.dynamo_import}.config.verbose=True for more information\n"
+        )
     msg += "=" * 10
     return msg
 

torchdynamo/debug_utils.py

@@ -12,10 +12,9 @@
 import torch
 import torch.fx as fx
 
-import torchdynamo
-from torchdynamo import config
-from torchdynamo.optimizations.backends import register_backend
-from torchdynamo.utils import clone_inputs
+from . import config
+from .optimizations.backends import register_backend
+from .utils import clone_inputs
 
 log = logging.getLogger(__name__)
 
@@ -132,11 +131,11 @@ def _cuda_system_info_comment():
 
 def generate_compiler_repro_string(gm, args):
     model_str = textwrap.dedent(
-        """
+        f"""
         import torch
         from torch import tensor, device
         import torch.fx as fx
-        from torchdynamo.testing import rand_strided
+        from {config.dynamo_import}.testing import rand_strided
         from math import inf
         from torch.fx.experimental.proxy_tensor import make_fx
 
@@ -180,7 +179,7 @@ def dump_compiler_graph_state(gm, args, compiler_name):
     print(f"Writing checkpoint with {len(gm.graph.nodes)} nodes to {file_name}")
     with open(file_name, "w") as fd:
         save_graph_repro(fd, gm, args, compiler_name)
-    repro_path = os.path.join(torchdynamo.config.base_dir, "repro.py")
+    repro_path = os.path.join(config.base_dir, "repro.py")
     shutil.copyfile(file_name, repro_path)
 
 
@@ -210,7 +209,7 @@ def isolate_fails(fx_g, args, compiler_name: str, env=None):
         fd.write(
             textwrap.dedent(
                 f"""
-                from torchdynamo.debug_utils import {fail_fn}
+                from {__name__} import {fail_fn}
                 """
             )
         )
@@ -290,7 +289,7 @@ def helper_for_dump_minify(contents):
         log.exception(e)
         raise NotImplementedError("Could not write to {minified_repro_path}")
 
-    local_path = os.path.join(torchdynamo.config.base_dir, "minifier_launcher.py")
+    local_path = os.path.join(config.base_dir, "minifier_launcher.py")
     try:
         shutil.copyfile(minified_repro_path, local_path)
         log.warning(
@@ -308,7 +307,7 @@ def dump_to_minify(gm, args, compiler_name: str):
 {generate_compiler_repro_string(gm, args)}
 
 from functools import partial
-from torchdynamo.debug_utils import (
+from {__name__} import (
     isolate_fails,
     dump_compiler_graph_state,
 )
@@ -385,9 +384,9 @@ def run_fwd_maybe_bwd(gm, args):
     """
     Runs a forward and possibly backward iteration for a given mod and args.
     """
-    from torchdynamo.testing import collect_results
-    from torchdynamo.testing import reduce_to_scalar_loss
-    from torchdynamo.testing import requires_bwd_pass
+    from .testing import collect_results
+    from .testing import reduce_to_scalar_loss
+    from .testing import requires_bwd_pass
 
     gm = copy.deepcopy(gm)
     args = clone_inputs(args)
@@ -406,7 +405,7 @@ def same_two_models(gm, opt_gm, example_inputs):
     """
     Check two models have same accuracy.
     """
-    from torchdynamo.utils import same
+    from .utils import same
 
     ref = run_fwd_maybe_bwd(gm, example_inputs)
 
@@ -447,21 +446,21 @@ def generate_dynamo_fx_repro_string(model_str, args, compiler_name):
 
     return textwrap.dedent(
         f"""
+from math import inf
 import torch
-import torchdynamo
 from torch import tensor, device
 import torch.fx as fx
-from torchdynamo.testing import rand_strided
-from math import inf
-from torchdynamo.debug_utils import run_fwd_maybe_bwd
+import {config.dynamo_import}
+from {config.dynamo_import}.testing import rand_strided
+from {config.dynamo_import}.debug_utils import run_fwd_maybe_bwd
 
 args = {[(tuple(a.shape), tuple(a.stride()), a.dtype, a.device.type, a.requires_grad) for a in args]}
 args = [rand_strided(sh, st, dt, dev).requires_grad_(rg) for (sh, st, dt, dev, rg) in args]
 
 {model_str}
 
 mod = Repro().cuda()
-opt_mod = torchdynamo.optimize("{compiler_name}")(mod)
+opt_mod = {config.dynamo_import}.optimize("{compiler_name}")(mod)
 
 with torch.cuda.amp.autocast(enabled={torch.is_autocast_enabled()}):
     ref = run_fwd_maybe_bwd(mod, args)
@@ -487,7 +486,7 @@ def dump_backend_repro_as_file(gm, args, compiler_name):
     log.warning(f"Copying {file_name} to {latest_repro} for convenience")
     shutil.copyfile(file_name, latest_repro)
 
-    local_path = os.path.join(torchdynamo.config.base_dir, "repro.py")
+    local_path = os.path.join(config.base_dir, "repro.py")
     try:
         shutil.copyfile(file_name, local_path)
         log.warning(
@@ -542,11 +541,11 @@ def dump_backend_repro_as_file(gm, args, compiler_name):
 #             )
 #         )
 
-#     local_dir = os.path.join(torchdynamo.config.base_dir, "repro")
+#     local_dir = os.path.join(config.base_dir, "repro")
 #     if os.path.exists(local_dir):
 #         shutil.rmtree(local_dir)
 #     shutil.copytree(tmp_dir, local_dir)
-#     local_tar_file = os.path.join(torchdynamo.config.base_dir, "repro.tar.gz")
+#     local_tar_file = os.path.join(config.base_dir, "repro.tar.gz")
 #     print(f"Writing checkpoint with {len(gm.graph.nodes)} locally to {local_tar_file}")
 #     with tarfile.open(local_tar_file, "w:gz") as tar:
 #         tar.add(local_dir, arcname=os.path.basename(local_dir))
@@ -595,18 +594,18 @@ def dump_to_minify_after_dynamo(gm, args, compiler_name):
 
     contents = textwrap.dedent(
         f"""
+import os
+from math import inf
 import torch
-import torchdynamo
 from torch import tensor, device
 import torch.fx as fx
-from torchdynamo.testing import rand_strided
-from math import inf
-from torchdynamo.debug_utils import run_fwd_maybe_bwd
-from torchdynamo.optimizations.backends import BACKENDS
 import functools
-import os
+import {config.dynamo_import}
+from {config.dynamo_import}.debug_utils import run_fwd_maybe_bwd
+from {config.dynamo_import}.optimizations.backends import BACKENDS
+from {config.dynamo_import}.testing import rand_strided
 
-torchdynamo.config.repro_dir = \"{minifier_dir()}\"
+{config.dynamo_import}.config.repro_dir = \"{minifier_dir()}\"
 
 args = {[(tuple(a.shape), tuple(a.stride()), a.dtype, a.device.type, a.requires_grad) for a in args]}
 args = [rand_strided(sh, st, dt, dev).requires_grad_(rg) for (sh, st, dt, dev, rg) in args]
@@ -620,7 +619,7 @@ def dump_to_minify_after_dynamo(gm, args, compiler_name):
     compiler_fn,
     compiler_name="{compiler_name}",
 )
-opt_mod = torchdynamo.optimize(dynamo_minifier_backend)(mod)
+opt_mod = {config.dynamo_import}.optimize(dynamo_minifier_backend)(mod)
 
 opt_mod(*args)
         """
@@ -678,7 +677,7 @@ def debug_wrapper(gm, example_inputs, **kwargs):
 def dynamo_minifier_backend(gm, example_inputs, compiler_name):
     from functorch.compile import minifier
 
-    from torchdynamo.optimizations.backends import BACKENDS
+    from .optimizations.backends import BACKENDS
 
     if compiler_name == "inductor":
         from torchinductor.compile_fx import compile_fx

torchdynamo/eval_frame.py

@@ -15,19 +15,17 @@
 from torch.fx.experimental.proxy_tensor import make_fx
 from torch.nn.parallel.distributed import DistributedDataParallel
 
-import torchdynamo
-from torchdynamo.optimizations.distributed import DDPOptimizer
-from torchdynamo.utils import checkpoint_params
-from torchdynamo.utils import clone_inputs
-from torchdynamo.utils import compile_times
-from torchdynamo.utils import same
-
 from . import config
 from . import convert_frame
 from . import skipfiles
 from . import utils
 from .exc import ResetRequired
 from .mutation_guard import install_generation_tagging_init
+from .optimizations.distributed import DDPOptimizer
+from .utils import checkpoint_params
+from .utils import clone_inputs
+from .utils import compile_times
+from .utils import same
 
 log = logging.getLogger(__name__)
 
@@ -65,7 +63,7 @@ def remove_from_cache(f):
     elif hasattr(getattr(f, "forward", None), "__code__"):
         reset_code(f.forward.__code__)
     else:
-        from torchdynamo import reset
+        from . import reset
 
         reset()
         log.warning("could not determine __code__ for %s", f)
@@ -301,7 +299,7 @@ def __call__(self, gm: torch.fx.GraphModule, example_inputs):
 
 
 def get_compiler_fn(compiler_fn):
-    from torchdynamo.debug_utils import wrap_backend_debug
+    from .debug_utils import wrap_backend_debug
 
     """Expand backend strings to functions"""
     compiler_str = compiler_fn if isinstance(compiler_fn, str) else None
@@ -368,7 +366,9 @@ def toy_example(a, b):
 @patch("torchdynamo.symbolic_convert.explain", True)
 def explain(f, *args, **kwargs):
     # TODO(voz): Do we want a decorator for this?
-    torchdynamo.reset()
+    from . import reset
+
+    reset()
 
     out_guards = []
     graphs = []
@@ -428,7 +428,7 @@ def guard_export_print(guards):
     explanation += compile_times()
 
     # TODO(voz): Do we want a decorator for this?
-    torchdynamo.reset()
+    reset()
     return explanation, out_guards, graphs, ops_per_graph, break_reasons
 
 

torchdynamo/exc.py

@@ -2,7 +2,7 @@
 import os
 import textwrap
 
-from torchdynamo.utils import counters
+from .utils import counters
 
 
 class TorchDynamoException(RuntimeError):

torchdynamo/guards.py

@@ -20,9 +20,8 @@
 import numpy as np
 import torch
 
-from torchdynamo import convert_frame
-
 from . import config
+from . import convert_frame
 from . import mutation_guard
 from ._guards import TensorGuards
 from ._guards import check_obj_id