已合并
【community issue】torch inductor patch add #36598
kkjocker创建于 5月25日
【community issue】torch inductor patch add #36598
已合并
共 3 个文件变更+67-8
| @@ -1,5 +1,5 @@ | |||
| 1 | diff --git a/test/nn/test_pooling.py b/test/nn/test_pooling.py | 1 | diff --git a/test/nn/test_pooling.py b/test/nn/test_pooling.py |
| 2 | -index 82f6ca2..a71a15a 100644 | 2 | +index 82f6ca2..c53e0a0 100644 |
| 3 | --- a/test/nn/test_pooling.py | 3 | --- a/test/nn/test_pooling.py |
| 4 | +++ b/test/nn/test_pooling.py | 4 | +++ b/test/nn/test_pooling.py |
| 5 | from functools import partial, reduce | 5 | from functools import partial, reduce |
| @@ -11,6 +11,24 @@ index 82f6ca2..a71a15a 100644 | |||
| 11 | import torch.nn as nn | 11 | import torch.nn as nn |
| 12 | import torch.nn.functional as F | 12 | import torch.nn.functional as F |
| 13 | from torch import inf, nan | 13 | from torch import inf, nan |
| 14 | + class TestPoolingNN(NNTestCase): | ||
| 15 | + ) | ||
| 16 | + | ||
| 17 | + @unittest.skipIf(not TEST_CUDA, "CUDA unavailable") | ||
| 18 | +- @largeTensorTest("12GB", device="cuda") | ||
| 19 | ++ @largeTensorTest("12GB", device="npu") | ||
| 20 | + def test_adaptive_pooling_avg_nhwc_launch_config_backward(self): | ||
| 21 | + input = torch.randint( | ||
| 22 | + 1, 10, (1, 32, 2**17 + 1, 32), dtype=torch.float32, device="cuda" | ||
| 23 | + class TestPoolingNN(NNTestCase): | ||
| 24 | + self.assertEqual(input.grad, ref_input.grad) | ||
| 25 | + | ||
| 26 | + @unittest.skipIf(not TEST_CUDA, "CUDA unavailable") | ||
| 27 | +- @largeTensorTest("12GB", device="cuda") | ||
| 28 | ++ @largeTensorTest("12GB", device="npu") | ||
| 29 | + def test_adaptive_pooling_avg_nhwc_launch_config_forward(self): | ||
| 30 | + input = torch.randint( | ||
| 31 | + 1, 10, (1, 32, 16, 16), dtype=torch.float32, device="cuda" | ||
| 14 | class TestPoolingNNDeviceType(NNTestCase): | 32 | class TestPoolingNNDeviceType(NNTestCase): |
| 15 | self.assertEqual(inp.grad, torch.zeros_like(inp)) | 33 | self.assertEqual(inp.grad, torch.zeros_like(inp)) |
| 16 | self.assertEqual(unpool_out, torch.zeros_like(unpool_out)) | 34 | self.assertEqual(unpool_out, torch.zeros_like(unpool_out)) |
| @@ -64,4 +82,4 @@ index 82f6ca2..a71a15a 100644 | |||
| 64 | + # @slowTest | 82 | + # @slowTest |
| 65 | def test_adaptive_pool_odd_size(self, device): | 83 | def test_adaptive_pool_odd_size(self, device): |
| 66 | # See https://github.com/pytorch/pytorch/issues/81409 | 84 | # See https://github.com/pytorch/pytorch/issues/81409 |
| 67 | - Ih, Iw, Oh, Ow = 5873, 3693, 3527, 2219 | 85 | + Ih, Iw, Oh, Ow = 5873, 3693, 3527, 2219 |
| @@ -0,0 +1,25 @@ | |||
| 1 | +diff --git a/torch/_inductor/utils.py b/torch/_inductor/utils.py | ||
| 2 | +index 8e92150..6c8a831 100644 | ||
| 3 | +--- a/torch/_inductor/utils.py | ||
| 4 | ++++ b/torch/_inductor/utils.py | ||
| 5 | + if TYPE_CHECKING: | ||
| 6 | + from .scheduler import BaseSchedulerNode, SchedulerBuffer | ||
| 7 | + | ||
| 8 | + | ||
| 9 | +-GPU_TYPES = ["cuda", "mps", "xpu"] | ||
| 10 | ++GPU_TYPES = ["cuda", "mps", "xpu", "npu"] | ||
| 11 | + T = TypeVar("T") | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + T = TypeVar("T") | ||
| 15 | + # when get_gpu_type is imported from dynamo | ||
| 16 | + @functools.lru_cache(None) | ||
| 17 | + def get_gpu_type() -> str: | ||
| 18 | +- avail_gpus = [x for x in GPU_TYPES if getattr(torch, x).is_available()] | ||
| 19 | +- assert len(avail_gpus) <= 1 | ||
| 20 | +- gpu_type = "cuda" if len(avail_gpus) == 0 else avail_gpus.pop() | ||
| 21 | +- return gpu_type | ||
| 22 | ++ return "npu" | ||
| 23 | + | ||
| 24 | + | ||
| 25 | + from torch._dynamo.device_interface import get_interface_for_device | ||
| @@ -1,5 +1,5 @@ | |||
| 1 | diff --git a/torch/testing/_internal/common_device_type.py b/torch/testing/_internal/common_device_type.py | 1 | diff --git a/torch/testing/_internal/common_device_type.py b/torch/testing/_internal/common_device_type.py |
| 2 | -index 834e0ed..c05a88e 100644 | 2 | +index 834e0ed..1970b4d 100644 |
| 3 | --- a/torch/testing/_internal/common_device_type.py | 3 | --- a/torch/testing/_internal/common_device_type.py |
| 4 | +++ b/torch/testing/_internal/common_device_type.py | 4 | +++ b/torch/testing/_internal/common_device_type.py |
| 5 | def filter_desired_device_types(device_type_test_bases, except_for=None, only_fo | 5 | def filter_desired_device_types(device_type_test_bases, except_for=None, only_fo |
| @@ -73,24 +73,31 @@ index 834e0ed..c05a88e 100644 | |||
| 73 | return list(device_type_test_bases) | 73 | return list(device_type_test_bases) |
| 74 | 74 | ||
| 75 | 75 | ||
| 76 | -@@ -1301,14 +1326,14 @@ class skipPRIVATEUSE1If(skipIf): | 76 | +@@ -1301,17 +1326,17 @@ class skipPRIVATEUSE1If(skipIf): |
| 77 | 77 | ||
| 78 | 78 | ||
| 79 | def _has_sufficient_memory(device, size): | 79 | def _has_sufficient_memory(device, size): |
| 80 | - if torch.device(device).type == "cuda": | 80 | - if torch.device(device).type == "cuda": |
| 81 | +- if not torch.cuda.is_available(): | ||
| 81 | + if torch.device(device).type == "npu": | 82 | + if torch.device(device).type == "npu": |
| 82 | - if not torch.cuda.is_available(): | 83 | ++ if not torch.npu.is_available(): |
| 83 | return False | 84 | return False |
| 84 | gc.collect() | 85 | gc.collect() |
| 85 | - torch.cuda.empty_cache() | 86 | +- torch.cuda.empty_cache() |
| 87 | ++ torch.npu.empty_cache() | ||
| 86 | # torch.cuda.mem_get_info, aka cudaMemGetInfo, returns a tuple of (free memory, total memory) of a GPU | 88 | # torch.cuda.mem_get_info, aka cudaMemGetInfo, returns a tuple of (free memory, total memory) of a GPU |
| 87 | - if device == "cuda": | 89 | - if device == "cuda": |
| 88 | - device = "cuda:0" | 90 | - device = "cuda:0" |
| 89 | + if device == "npu": | 91 | + if device == "npu": |
| 90 | + device = "npu:0" | 92 | + device = "npu:0" |
| 91 | return ( | 93 | return ( |
| 92 | - torch.cuda.memory.mem_get_info(device)[0] | 94 | +- torch.cuda.memory.mem_get_info(device)[0] |
| 93 | - * torch.cuda.memory.get_per_process_memory_fraction(device) | 95 | +- * torch.cuda.memory.get_per_process_memory_fraction(device) |
| 96 | ++ torch.npu.mem_get_info(device)[0] | ||
| 97 | ++ * torch.npu.get_per_process_memory_fraction(device) | ||
| 98 | + ) >= size | ||
| 99 | + | ||
| 100 | + if device == "xla": | ||
| 94 | def largeTensorTest(size, device=None, inductor=TEST_WITH_TORCHINDUCTOR): | 101 | def largeTensorTest(size, device=None, inductor=TEST_WITH_TORCHINDUCTOR): |
| 95 | if inductor and torch._inductor.config.cpp_wrapper and _device != "cpu": | 102 | if inductor and torch._inductor.config.cpp_wrapper and _device != "cpu": |
| 96 | size_bytes *= 2 | 103 | size_bytes *= 2 |
| @@ -122,3 +129,12 @@ index 834e0ed..c05a88e 100644 | |||
| 122 | def onlyMPS(fn): | 129 | def onlyMPS(fn): |
| 123 | return onlyOn("mps")(fn) | 130 | return onlyOn("mps")(fn) |
| 124 | 131 | ||
| 132 | + def skipPRIVATEUSE1(fn): | ||
| 133 | + # TODO: the "all" in the name isn't true anymore for quite some time as we have also have for example XLA and MPS now. | ||
| 134 | + # This should probably enumerate all available device type test base classes. | ||
| 135 | + def get_all_device_types() -> list[str]: | ||
| 136 | +- return ["cpu"] if not torch.cuda.is_available() else ["cpu", "cuda"] | ||
| 137 | ++ return ["cpu"] if not torch.cuda.is_available() else ["cpu", "cuda", "npu"] | ||
| 138 | + | ||
| 139 | + | ||
| 140 | + flex_attention_supported_platform = unittest.skipUnless( | ||