Changes related to moving to PyTorch core

Makefile

@@ -33,16 +33,11 @@ overhead: develop
 format:
 	isort $(PY_FILES)
 	black $(PY_FILES)
-	! which $(CLANG_FORMAT) >/dev/null 2>&1 || $(CLANG_FORMAT) -i $(C_FILES)
 
 lint:
 	black --check --diff $(PY_FILES)
 	isort --check --diff $(PY_FILES)
 	flake8 $(PY_FILES)
-	mypy
-	! which $(CLANG_TIDY) >/dev/null 2>&1 || $(CLANG_TIDY) $(C_FILES) -- \
-		-I`python -c 'from distutils.sysconfig import get_python_inc as X; print(X())'` \
-		`python -c 'from torch.utils.cpp_extension import include_paths; print(" ".join(map("-I{}".format, include_paths())))'`
 
 lint-deps:
 	grep -E '(black|flake8|isort|click|torch|mypy)' requirements.txt | xargs $(PIP) install

benchmarks/common.py

@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 import argparse
 import collections
 import copy
@@ -891,6 +891,33 @@ def scale(self, loss):
         return loss
 
 
+def maybe_fresh_cache(fn):
+    def inner(self, *args, **kwargs):
+        cache_minder = NullContext()
+        if self.args.cold_start_latency:
+            cache_entries = {}
+            cache_minder = fresh_triton_cache(cache_entries)
+
+        try:
+            with cache_minder:
+                return fn(self, *args, **kwargs)
+        finally:
+            dump_cache = False
+            if dump_cache and self.args.cold_start_latency:
+                output_csv(
+                    output_filename[:-4] + "_triton_cache.csv",
+                    ["dev", "name", "batch_size", "triton_cache"],
+                    [
+                        current_device,
+                        current_name,
+                        current_batch_size,
+                        cache_entries,
+                    ],
+                )
+
+    return inner
+
+
 class BenchmarkRunner:
     def __init__(self):
         self.model_iter_fn = None
@@ -1285,32 +1312,6 @@ def compare_branches(
                 "--diff_main called on main branch, what are you diffing?"
             )
 
-    def maybe_fresh_cache(fn):
-        def inner(self, *args, **kwargs):
-            cache_minder = NullContext()
-            if self.args.cold_start_latency:
-                cache_entries = {}
-                cache_minder = fresh_triton_cache(cache_entries)
-
-            try:
-                with cache_minder:
-                    return fn(self, *args, **kwargs)
-            finally:
-                dump_cache = False
-                if dump_cache and self.args.cold_start_latency:
-                    output_csv(
-                        output_filename[:-4] + "_triton_cache.csv",
-                        ["dev", "name", "batch_size", "triton_cache"],
-                        [
-                            current_device,
-                            current_name,
-                            current_batch_size,
-                            cache_entries,
-                        ],
-                    )
-
-        return inner
-
     @maybe_fresh_cache
     def run_one_model(
         self,

benchmarks/huggingface.py

@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 import importlib
 import logging
 import os

benchmarks/microbenchmarks/bench_conv_fusion.py

@@ -1,3 +1,4 @@
+# flake8: noqa
 import model
 import torch
 import triton

benchmarks/microbenchmarks/bench_mm_fusion.py

@@ -1,3 +1,4 @@
+# flake8: noqa
 import torch
 import triton
 from prettytable import PrettyTable

benchmarks/microbenchmarks/benchmark_helper.py

@@ -1,7 +1,8 @@
 from torch.utils.benchmark import Timer
 
 
-def time_with_torch_timer(fn, args, kwargs={}, iters=100):
+def time_with_torch_timer(fn, args, kwargs=None, iters=100):
+    kwargs = kwargs or {}
     env = {"args": args, "kwargs": kwargs, "fn": fn}
     fn_call = "fn(*args, **kwargs)"
 

benchmarks/microbenchmarks/microbench.py

@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 import argparse
 import inspect
 import sys

benchmarks/microbenchmarks/operator_inp_utils.py

@@ -85,7 +85,7 @@ def serialize_sparse_tensor(e):
 
 
 def deserialize_sparse_tensor(size, dtype, layout, is_coalesced, nnz=None):
-    assert False, "NYI"
+    raise NotImplementedError()
 
 
 def deserialize_tensor(size, dtype, stride=None):

benchmarks/microbenchmarks/operatorbench.py

@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 import click
 import numpy as np
 import torch

benchmarks/runner.py

@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 
 """
 A wrapper over the benchmark infrastructure to generate commonly used commands,

benchmarks/timm_models.py

@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 import importlib
 import logging
 import os
@@ -30,9 +30,8 @@ def pip_install(package):
     from timm.models import create_model
 
 TIMM_MODELS = dict()
-filename = "timm_models_list.txt"
-if os.path.exists("benchmarks"):
-    filename = "benchmarks/" + filename
+filename = os.path.join(os.path.dirname(__file__), "timm_models_list.txt")
+
 with open(filename, "r") as fh:
     lines = fh.readlines()
     lines = [line.rstrip() for line in lines]

benchmarks/torchbench.py

@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 import gc
 import importlib
 import logging
@@ -240,6 +240,8 @@ def load_model(
         if batch_size is None and is_training and model_name in USE_SMALL_BATCH_SIZE:
             batch_size = USE_SMALL_BATCH_SIZE[model_name]
 
+        # workaround "RuntimeError: not allowed to set torch.backends.cudnn flags"
+        torch.backends.__allow_nonbracketed_mutation_flag = True
         if is_training:
             benchmark = benchmark_cls(
                 test="train", device=device, jit=False, batch_size=batch_size

copy_to_core.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+set -ex
+
+rsync -ra ~/torchdynamo/torchdynamo/ ~/pytorch/torch/_dynamo
+rsync -ra ~/torchdynamo/torchinductor/ ~/pytorch/torch/_inductor
+rsync -ra ~/torchdynamo/test/{dynamo,inductor} ~/pytorch/test/
+rsync -ra ~/torchdynamo/benchmarks/ ~/pytorch/benchmarks/dynamo
+
+for DIR in ~/pytorch/test/{dynamo,inductor} ~/pytorch/benchmarks/dynamo
+do
+  find $DIR -name '*.py' | xargs -n1 -- sed -i 's/torchdynamo/torch._dynamo/g'
+  find $DIR -name '*.py' | xargs -n1 -- sed -i 's/torchinductor/torch._inductor/g'
+  find $DIR -name '*.py' | xargs -n1 -- sed -i 's/_torch[.]_inductor/_torchinductor/g'
+  find $DIR -name '*.py' | xargs -n1 -- sed -i 's@pytorch/torch[.]_dynamo@pytorch/torchdynamo@g'
+done
+
+(cd ~/pytorch && (lintrunner -a || lintrunner -a))

setup.py

@@ -1,10 +1,8 @@
 #!/usr/bin/env python
 import sys
 
-from setuptools import Extension
 from setuptools import find_packages
 from setuptools import setup
-from torch.utils.cpp_extension import CppExtension
 
 long_description = """
 TorchDynamo is a Python-level JIT compiler designed to make unmodified
@@ -31,7 +29,7 @@
     ]
 
 install_requires = [
-    "torch>=1.12.0",
+    "torch>=1.13.0",
     "numpy",
     "tabulate",
     "pyyaml",
@@ -43,13 +41,13 @@
 
 setup(
     name="torchdynamo",
-    version="1.13.0.dev0",
+    version="1.14.0.dev0",
     url="https://github.com/pytorch/torchdynamo",
     description="A Python-level JIT compiler designed to make unmodified PyTorch programs faster.",
     long_description=long_description,
     long_description_content_type="text/markdown",
     author="Jason Ansel",
-    author_email="jansel@fb.com",
+    author_email="jansel@meta.com",
     license="BSD-3",
     keywords="pytorch machine learning compilers",
     python_requires=">=3.7, <3.11",
@@ -59,16 +57,4 @@
         "torchinductor.codegen": ["*.h", "*.j2"],
     },
     zip_safe=False,
-    ext_modules=[
-        Extension(
-            "torchdynamo._eval_frame",
-            ["torchdynamo/_eval_frame.c"],
-            extra_compile_args=["-Wall"],
-        ),
-        CppExtension(
-            name="torchdynamo._guards",
-            sources=["torchdynamo/_guards.cpp"],
-            extra_compile_args=["-std=c++14"],
-        ),
-    ],
 )

test/dynamo/mock_modules/mock_module2.py

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

test/dynamo/mock_modules/mock_module3.py

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

test/dynamo/test_aot_autograd.py

@@ -1,9 +1,10 @@
-#!/usr/bin/env pytest
+# Owner(s): ["module: dynamo"]
 import functools
 
 import torch
 
 import torchdynamo
+import torchdynamo.test_case
 from torchdynamo.optimizations.training import is_aot_autograd_safe_to_run
 from torchdynamo.testing import rand_strided
 
@@ -13,14 +14,14 @@ def compiler_safe_fn(gm, example_inputs, is_safe):
     return gm.forward
 
 
-class AotAutogradFallbackTests(torchdynamo.testing.TestCase):
+class AotAutogradFallbackTests(torchdynamo.test_case.TestCase):
     def test_LSTM(self):
         # https://github.com/pytorch/torchdynamo/issues/1147
         class Repro(torch.nn.Module):
             def __init__(self):
                 super().__init__()
                 self.self_mod_model_lstm_lstm = torch.nn.LSTM(
-                    2048, 2048, num_layers=2, bidirectional=True
+                    64, 64, num_layers=2, bidirectional=True
                 )
 
             def forward(self, permute: torch.Tensor):
@@ -32,7 +33,7 @@ def forward(self, permute: torch.Tensor):
         compiler_fn = functools.partial(compiler_safe_fn, is_safe=is_safe)
         aot_mod = torchdynamo.optimize(compiler_fn)(mod)
 
-        args = [((92, 4, 2048), (1, 188416, 92), torch.float32, "cpu", False)]
+        args = [((92, 4, 64), (1, 5888, 92), torch.float32, "cpu", False)]
         args = [
             rand_strided(sh, st, dt, dev).requires_grad_(rg)
             for (sh, st, dt, dev, rg) in args
@@ -130,3 +131,9 @@ def fn(x, y):
         aot_fn = torchdynamo.optimize(compiler_fn)(fn)
         aot_fn(x, y)
         self.assertTrue(is_safe[0])
+
+
+if __name__ == "__main__":
+    from torchdynamo.test_case import run_tests
+
+    run_tests()

test/dynamo/test_aot_cudagraphs.py

@@ -1,4 +1,3 @@
-#!/usr/bin/env pytest
 # Owner(s): ["module: cuda graphs"]
 
 import functools
@@ -8,6 +7,7 @@
 import torch
 
 import torchdynamo
+import torchdynamo.test_case
 import torchdynamo.testing
 from torchdynamo.testing import same
 
@@ -53,7 +53,7 @@ def patch_all(ok=True):
 
 
 @unittest.skipIf(not torch.cuda.is_available(), "these tests require cuda")
-class TestAotCudagraphs(torchdynamo.testing.TestCase):
+class TestAotCudagraphs(torchdynamo.test_case.TestCase):
     @patch_all()
     def test_basic(self):
         def model(x, y):
@@ -202,4 +202,6 @@ def fn(x):
 
 
 if __name__ == "__main__":
-    unittest.main()
+    from torchdynamo.test_case import run_tests
+
+    run_tests()