已合并
support_list_for_register_meta_npu #39246
cuiduo创建于 6月25日
support_list_for_register_meta_npu #39246
已合并
共 4 个文件变更+60-9
| @@ -0,0 +1,33 @@ | |||
| 1 | +import torch | ||
| 2 | +from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests | ||
| 3 | +from testutils import TestUtils | ||
| 4 | +import torch_npu | ||
| 5 | +import torch_npu._inductor | ||
| 6 | +from torch._inductor.utils import run_and_get_code | ||
| 7 | + | ||
| 8 | + | ||
| 9 | +class TestNanquantileDefault(TestUtils): | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + def test_nanquantile(self, shape, dtype, dynamic): | ||
| 15 | + a = torch.randn(shape, dtype=torch.float32, device="npu") | ||
 | |||
| 16 | + a[0, 0] = float('nan') | ||
| 17 | + q = torch.tensor([0.25, 0.5, 0.75], dtype=torch.float32, device="npu") | ||
| 18 | + | ||
| 19 | + def fn(a, q): | ||
| 20 | + return torch.ops.aten.nanquantile.default(a, q, dim=0, keepdim=False) | ||
| 21 | + | ||
| 22 | + r1 = fn(a, q) | ||
| 23 | + func = torch.compile(fn, backend="inductor", dynamic=dynamic) | ||
| 24 | + r, codes = run_and_get_code(func, a, q) | ||
| 25 | + self.assertEqual(r, r1, atol=1e-3, rtol=1e-3) | ||
| 26 | + self.assertTrue('nanquantile' in codes[0]) | ||
| 27 | + | ||
| 28 | + | ||
| 29 | +instantiate_parametrized_tests(TestNanquantileDefault) | ||
| 30 | + | ||
| 31 | + | ||
| 32 | +if __name__ == "__main__": | ||
| 33 | + run_tests() | ||
| @@ -160,6 +160,10 @@ def _load_triton_backend(): | |||
| 160 | 160 | ||
| 161 | _patch_flex_attention_singleton_sort() | 161 | _patch_flex_attention_singleton_sort() |
| 162 | 162 | ||
| 163 | + from ._npu_meta_registration import npu_patch_meta | ||
| 164 | + | ||
| 165 | + npu_patch_meta() | ||
| 166 | + | ||
| 163 | def _inductor_register_backend_for_device(): | 167 | def _inductor_register_backend_for_device(): |
| 164 | from .codegen.cpp_wrapper import CppWrapperNpu | 168 | from .codegen.cpp_wrapper import CppWrapperNpu |
| 165 | from .codegen.npu_combined_scheduling import NPUCombinedScheduling | 169 | from .codegen.npu_combined_scheduling import NPUCombinedScheduling |
| @@ -17,6 +17,7 @@ from torch._dynamo.symbolic_convert import break_graph_if_unsupported, Instructi | |||
| 17 | from torch._dynamo.exc import Unsupported | 17 | from torch._dynamo.exc import Unsupported |
| 18 | from torch._dynamo.variables.lists import TupleVariable | 18 | from torch._dynamo.variables.lists import TupleVariable |
| 19 | from torch._dynamo.variables.nn_module import NNModuleVariable | 19 | from torch._dynamo.variables.nn_module import NNModuleVariable |
| 20 | +from torch._meta_registrations import device_hint | ||
| 20 | import torch_npu | 21 | import torch_npu |
| 21 | 22 | ||
| 22 | aten = torch.ops.aten | 23 | aten = torch.ops.aten |
| @@ -94,7 +95,9 @@ def patch_torch_decomp_decompositions(): | |||
| 94 | 95 | ||
| 95 | def register_meta_npu(op, avoid_fallback_flag=False, inductor_decomp=False): | 96 | def register_meta_npu(op, avoid_fallback_flag=False, inductor_decomp=False): |
| 96 | def meta_decorator(fn: Callable): | 97 | def meta_decorator(fn: Callable): |
| 97 | - _add_op_to_meta_table(op, fn, avoid_fallback_flag, inductor_decomp) | 98 | + ops = op if isinstance(op, list) else [op] |
| 99 | + for single_op in ops: | ||
| 100 | + _add_op_to_meta_table(single_op, fn, avoid_fallback_flag, inductor_decomp) | ||
| 98 | return fn | 101 | return fn |
| 99 | 102 | ||
| 100 | return meta_decorator | 103 | return meta_decorator |
| @@ -120,8 +123,26 @@ def npu_patch_meta(): | |||
| 120 | patch_torch_decomp_decompositions() | 123 | patch_torch_decomp_decompositions() |
| 121 | patch_torch_inductor_decompositions() | 124 | patch_torch_inductor_decompositions() |
| 122 | 125 | ||
| 126 | + | ||
| 127 | + [ | ||
| 128 | + aten.sort.default, | ||
| 129 | + aten.sort.stable, | ||
| 130 | + aten.sort.values, | ||
| 131 | + aten.sort.values_stable, | ||
| 132 | + ] | ||
| 133 | + , avoid_fallback_flag=True | ||
| 134 | +) | ||
| 135 | +def meta_sort(self, stable=None, dim=-1, descending=False, values=None, indices=None): | ||
| 136 | + if device_hint(self) == "npu": | ||
| 137 | + v = torch.empty(self.shape, dtype=self.dtype, device=self.device) | ||
| 138 | + i = torch.empty(self.shape, dtype=torch.int64, device=self.device) | ||
| 123 | 139 | ||
| 124 | - | 140 | + return v, i |
| 141 | + else: | ||
| 142 | + from torch._meta_registrations import meta_sort | ||
| 143 | + meta_sort(self, stable=stable, dim=dim, descending=descending, values=values, indices=indices) | ||
| 144 | + | ||
| 145 | + | ||
| 125 | 146 | ||
| 126 | def meta_index_put_patch(self, indices, values, accumulate=False): | 147 | def meta_index_put_patch(self, indices, values, accumulate=False): |
| 127 | return self.new_empty(self.shape) | 148 | return self.new_empty(self.shape) |
| @@ -13,10 +13,3 @@ def apply_npu_format_patch(): | |||
| 13 | from torch_npu.npu._format import _apply_npu_format_patch | 13 | from torch_npu.npu._format import _apply_npu_format_patch |
| 14 | 14 | ||
| 15 | _apply_npu_format_patch() | 15 | _apply_npu_format_patch() |
| 16 | - | ||
| 17 | - | ||
| 18 | - | ||
| 19 | -def apply_npu_meta_patch(): | ||
| 20 | - from torch_npu.utils._npu_meta_registration import npu_patch_meta | ||
| 21 | - | ||
| 22 | - npu_patch_meta() | ||
🔵 Low Priority
第 14 行
@parametrize('dtype', ['float32'])将dtype='float32'传入测试方法，但第 16 行a = torch.randn(shape, dtype=torch.float32, device="npu")直接硬编码了dtype=torch.float32，完全忽略了dtype形参。第 18 行的q同样硬编码为torch.float32。后果：
@parametrize('dtype', ...)是无效装饰——无论传入什么 dtype 字符串，测试始终只测 float32。如果后续向列表中增加其他 dtype，测试不会真正覆盖它们，造成虚假的测试覆盖率。建议：要么移除无用的
@parametrize('dtype', ...)装饰器并将dtype从函数签名中删除，要么使用dtype参数动态构造 dtype（如getattr(torch, dtype)），使参数化真正生效。