xref: /aosp_15_r20/external/pytorch/torch/testing/_internal/inductor_utils.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1# mypy: ignore-errors
2
3import logging
4import torch
5import re
6import unittest
7import functools
8import os
9from subprocess import CalledProcessError
10import sys
11import torch._inductor.async_compile  # noqa: F401 required to warm up AsyncCompile pools
12from torch._inductor.codecache import CppCodeCache
13from torch._inductor.utils import get_gpu_shared_memory, is_big_gpu
14from torch._inductor.utils import GPU_TYPES, get_gpu_type
15from torch.utils._triton import has_triton
16from torch.testing._internal.common_utils import (
17    LazyVal,
18    IS_FBCODE,
19)
20from torch.testing._internal.common_utils import (
21    TestCase,
22    IS_CI,
23    IS_WINDOWS,
24)
25
26log: logging.Logger = logging.getLogger(__name__)
27
28def test_cpu():
29    try:
30        CppCodeCache.load("")
31        return not IS_FBCODE
32    except (
33        CalledProcessError,
34        OSError,
35        torch._inductor.exc.InvalidCxxCompiler,
36        torch._inductor.exc.CppCompileError,
37    ):
38        return False
39
40HAS_CPU = LazyVal(test_cpu)
41
42HAS_CUDA = torch.cuda.is_available() and has_triton()
43
44HAS_XPU = torch.xpu.is_available() and has_triton()
45
46HAS_GPU = HAS_CUDA or HAS_XPU
47
48GPU_TYPE = get_gpu_type()
49
50HAS_MULTIGPU = any(
51    getattr(torch, gpu).is_available() and getattr(torch, gpu).device_count() >= 2
52    for gpu in GPU_TYPES
53)
54
55def _check_has_dynamic_shape(
56    self: TestCase,
57    code,
58):
59    for_loop_found = False
60    has_dynamic = False
61    lines = code.split("\n")
62    for line in lines:
63        if "for(" in line:
64            for_loop_found = True
65            if re.search(r";.*ks.*;", line) is not None:
66                has_dynamic = True
67                break
68    self.assertTrue(
69        has_dynamic, msg=f"Failed to find dynamic for loop variable\n{code}"
70    )
71    self.assertTrue(for_loop_found, f"Failed to find for loop\n{code}")
72
73
74def skipDeviceIf(cond, msg, *, device):
75    if cond:
76        def decorate_fn(fn):
77            def inner(self, *args, **kwargs):
78                if not hasattr(self, "device"):
79                    warn_msg = "Expect the test class to have attribute device but not found. "
80                    if hasattr(self, "device_type"):
81                        warn_msg += "Consider using the skip device decorators in common_device_type.py"
82                    log.warning(warn_msg)
83                if self.device == device:
84                    raise unittest.SkipTest(msg)
85                return fn(self, *args, **kwargs)
86            return inner
87    else:
88        def decorate_fn(fn):
89            return fn
90
91    return decorate_fn
92
93def skip_windows_ci(name: str, file: str) -> None:
94    if IS_WINDOWS and IS_CI:
95        module = os.path.basename(file).strip(".py")
96        sys.stderr.write(
97            f"Windows CI does not have necessary dependencies for {module} tests yet\n"
98        )
99        if name == "__main__":
100            sys.exit(0)
101        raise unittest.SkipTest("requires sympy/functorch/filelock")
102
103requires_gpu = functools.partial(unittest.skipIf, not HAS_GPU, "requires gpu")
104
105skipCUDAIf = functools.partial(skipDeviceIf, device="cuda")
106skipXPUIf = functools.partial(skipDeviceIf, device="xpu")
107skipCPUIf = functools.partial(skipDeviceIf, device="cpu")
108
109IS_A100 = LazyVal(
110    lambda: HAS_CUDA
111    and get_gpu_shared_memory() == 166912
112)
113
114IS_H100 = LazyVal(
115    lambda: HAS_CUDA
116    and get_gpu_shared_memory() == 232448
117)
118
119IS_BIG_GPU = LazyVal(lambda: HAS_CUDA and is_big_gpu(0))
120