Enable ASAN in CUDA tests (#147812)

Fixes #ISSUE_NUMBER

Pull Request resolved: https://github.com/pytorch/pytorch/pull/147812
Approved by: https://github.com/janeyx99

test/inductor/test_binary_folding.py

@@ -24,7 +24,6 @@ from inductor.test_torchinductor import (  # @manual=fbcode//caffe2/test/inducto
     check_model_gpu,
     copy_tests,
 )
-from torch.testing._internal.common_utils import TEST_WITH_ASAN
 from torch.testing._internal.inductor_utils import skipCUDAIf
 
 
@@ -331,7 +330,7 @@ if HAS_CPU and not torch.backends.mps.is_available():
 
     copy_tests(BinaryFoldingTemplate, FreezingCpuTests, "cpu")
 
-if HAS_GPU and not TEST_WITH_ASAN:
+if HAS_GPU:
 
     class FreezingGpuTests(TestCase):
         common = check_model_gpu

test/inductor/test_cudagraph_trees.py

@@ -36,7 +36,6 @@ from torch.testing._internal.common_utils import (
     parametrize,
     skipIfRocm,
     TEST_CUDA_GRAPH,
-    TEST_WITH_ASAN,
 )
 from torch.utils._mode_utils import no_dispatch
 from torch.utils._python_dispatch import TorchDispatchMode
@@ -122,7 +121,7 @@ class TestCase(InductorTestCase):
         torch._dynamo.reset()
 
 
-if HAS_CUDA and not TEST_WITH_ASAN:
+if HAS_CUDA:
 
     def get_all_cudagraph_segments():
         segments = torch.cuda.memory_snapshot()

test/inductor/test_cudagraph_trees_expandable_segments.py

@@ -7,14 +7,14 @@ import sys
 
 import torch
 from torch.testing._internal.common_cuda import IS_JETSON, IS_WINDOWS
-from torch.testing._internal.common_utils import run_tests, TEST_WITH_ASAN
+from torch.testing._internal.common_utils import run_tests
 from torch.testing._internal.inductor_utils import HAS_CUDA
 
 
 pytorch_test_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
 sys.path.append(pytorch_test_dir)
 
-if HAS_CUDA and not TEST_WITH_ASAN:
+if HAS_CUDA:
     try:
         from .test_cudagraph_trees import CudaGraphTreeTests
     except ImportError:
@@ -32,13 +32,7 @@ from tools.stats.import_test_stats import get_disabled_tests  # @manual
 sys.path.remove(str(REPO_ROOT))
 
 if __name__ == "__main__":
-    if (
-        torch.cuda.is_available()
-        and not IS_JETSON
-        and not IS_WINDOWS
-        and HAS_CUDA
-        and not TEST_WITH_ASAN
-    ):
+    if torch.cuda.is_available() and not IS_JETSON and not IS_WINDOWS and HAS_CUDA:
         get_disabled_tests(".")
 
         torch.cuda.memory._set_allocator_settings("expandable_segments:True")

test/inductor/test_gpu_cpp_wrapper.py

@@ -11,7 +11,7 @@ from torch._inductor.utils import is_gpu
 from torch.testing._internal.common_device_type import (
     get_desired_device_type_test_bases,
 )
-from torch.testing._internal.common_utils import slowTest, TEST_WITH_ASAN
+from torch.testing._internal.common_utils import slowTest
 from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_GPU
 
 
@@ -40,10 +40,8 @@ except unittest.SkipTest:
 
 
 _desired_test_bases = get_desired_device_type_test_bases(allow_xpu=True)
-RUN_GPU = (
-    HAS_GPU
-    and any(is_gpu(getattr(x, "device_type", "")) for x in _desired_test_bases)
-    and not TEST_WITH_ASAN
+RUN_GPU = HAS_GPU and any(
+    is_gpu(getattr(x, "device_type", "")) for x in _desired_test_bases
 )
 
 

test/inductor/test_inductor_freezing.py

@@ -28,7 +28,7 @@ from inductor.test_torchinductor import (  # @manual=fbcode//caffe2/test/inducto
     check_model_gpu,
     copy_tests,
 )
-from torch.testing._internal.common_utils import TEST_WITH_ASAN, TEST_WITH_ROCM
+from torch.testing._internal.common_utils import TEST_WITH_ROCM
 
 
 importlib.import_module("functorch")
@@ -968,7 +968,7 @@ if HAS_CPU and not torch.backends.mps.is_available():
 
     copy_tests(OptimizeForInferenceTemplate, FreezingCpuTests, "cpu")
 
-if HAS_GPU and not TEST_WITH_ASAN:
+if HAS_GPU:
 
     class FreezingGpuTests(TestCase):
         common = check_model_gpu

test/inductor/test_torchinductor.py

@@ -12827,7 +12827,7 @@ if HAS_CPU:
 
     copy_tests(CommonTemplate, CpuTests, "cpu")
 
-if HAS_GPU and not TEST_WITH_ASAN:
+if HAS_GPU:
 
     class SweepInputsGPUTest(SweepInputs2, TestCase):
         gen = InputGen(10, GPU_TYPE)