[inductor] Enable FX graph caching on another round of inductor tests (pytorch#121994)

Summary: Enabling caching for these tests was blocked by pytorch#121686

Pull Request resolved: pytorch#121994
Approved by: https://github.com/eellison

Author: masnesral

test/inductor/test_compiled_optimizers.py

@@ -14,6 +14,8 @@
 import torch._inductor.cudagraph_trees
 from torch._inductor import config
 
+from torch._inductor.test_case import TestCase
+
 from torch.optim import (
     Adadelta,
     Adagrad,
@@ -39,8 +41,6 @@
     optim_db,
     optims,
 )
-
-from torch.testing._internal.common_utils import TestCase
 from torch.testing._internal.inductor_utils import HAS_CPU, HAS_CUDA, has_triton
 from torch.testing._internal.triton_utils import requires_cuda
 
@@ -598,7 +598,7 @@ def loop(opt, c):
 instantiate_device_type_tests(CompiledOptimizerParityTests, globals())
 
 if __name__ == "__main__":
-    from torch._dynamo.test_case import run_tests
+    from torch._inductor.test_case import run_tests
 
     if HAS_CPU or HAS_CUDA:
         run_tests(needs="filelock")

test/inductor/test_cutlass_backend.py

@@ -5,9 +5,9 @@
 from typing import Callable, List
 
 import torch
-from torch._dynamo.test_case import run_tests, TestCase
 from torch._dynamo.utils import counters
 from torch._inductor import config
+from torch._inductor.test_case import run_tests, TestCase
 from torch.testing._internal.common_cuda import SM75OrLater, SM80OrLater, SM90OrLater
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,

test/inductor/test_group_batch_fusion.py

@@ -5,8 +5,8 @@
 
 import torch
 import torch._inductor
-from torch._dynamo.test_case import run_tests, TestCase
 from torch._dynamo.utils import counters, optimus_scuba_log
+from torch._inductor.test_case import run_tests, TestCase
 from torch.testing._internal.inductor_utils import HAS_CUDA
 
 try:

test/inductor/test_torchinductor_dynamic_shapes.py

@@ -11,6 +11,7 @@
 import torch._custom_ops as custom_ops
 import torch.library
 from torch._dynamo.testing import make_test_cls_with_patches
+from torch._inductor.test_case import TestCase
 from torch.testing._internal.common_device_type import (
     instantiate_device_type_tests,
     onlyCPU,
@@ -24,7 +25,6 @@
     TEST_CUDA_MEM_LEAK_CHECK,
     TEST_WITH_ASAN,
     TEST_WITH_ROCM,
-    TestCaseBase as TestCase,
 )
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_CPU, HAS_GPU
 
@@ -613,7 +613,7 @@ def func(x, fn, a):
 instantiate_device_type_tests(TestInductorDynamic, globals())
 
 if __name__ == "__main__":
-    from torch._dynamo.test_case import run_tests
+    from torch._inductor.test_case import run_tests
 
     # Slow on ASAN after https://github.com/pytorch/pytorch/pull/94068
     if (HAS_CPU or HAS_GPU) and not TEST_WITH_ASAN:

test/inductor/test_torchinductor_opinfo.py

@@ -13,7 +13,7 @@
 import torch
 
 from torch._dispatch.python import enable_python_dispatcher
-from torch._dynamo.test_case import run_tests
+from torch._inductor.test_case import run_tests, TestCase
 from torch._subclasses.fake_tensor import (
     DataDependentOutputException,
     DynamicOutputShapeException,
@@ -40,7 +40,6 @@
     TEST_MKL,
     TEST_WITH_ASAN,
     TEST_WITH_ROCM,
-    TestCase,
 )
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_CPU, HAS_CUDA
 from torch.utils._python_dispatch import TorchDispatchMode

test/inductor/test_unbacked_symints.py

@@ -6,14 +6,15 @@
 
 from torch._dynamo import config as dynamo_config
 from torch._inductor import config as inductor_config
+from torch._inductor.test_case import TestCase as InductorTestCase
 from torch._inductor.utils import is_big_gpu
 from torch.testing import make_tensor
 from torch.testing._internal.common_device_type import instantiate_device_type_tests
-from torch.testing._internal.common_utils import IS_LINUX, TestCase as TorchTestCase
+from torch.testing._internal.common_utils import IS_LINUX
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_CUDA, skipCUDAIf
 
 
-class TestUnbackedSymints(TorchTestCase):
+class TestUnbackedSymints(InductorTestCase):
     @skipCUDAIf(not HAS_CUDA, "requires cuda")
     @dynamo_config.patch({"capture_dynamic_output_shape_ops": True})
     def test_expand(self, device):
@@ -147,7 +148,7 @@ def fn(x):
 )
 
 if __name__ == "__main__":
-    from torch._dynamo.test_case import run_tests
+    from torch._inductor.test_case import run_tests
 
     if IS_LINUX and HAS_CUDA and is_big_gpu(0):
         run_tests()