已合并
move_has_triton_to_dynamo #43316
cuiduo创建于 7月30日
move_has_triton_to_dynamo #43316
已合并
共 3 个文件变更+44-39
| @@ -20,13 +20,11 @@ class TestHasTriton(TestCase): | |||
| 20 | # no triton library found, skip test_has_triton | 20 | # no triton library found, skip test_has_triton |
| 21 | return | 21 | return |
| 22 | 22 | ||
| 23 | - self.assertFalse(has_triton()) | ||
| 24 | self.assertFalse(has_triton_tma()) | 23 | self.assertFalse(has_triton_tma()) |
| 25 | self.assertFalse(has_triton_tma_device()) | 24 | self.assertFalse(has_triton_tma_device()) |
| 26 | 25 | ||
| 27 | from torch_npu.contrib import transfer_to_npu | 26 | from torch_npu.contrib import transfer_to_npu |
| 28 | 27 | ||
| 29 | - self.assertFalse(has_triton()) | ||
| 30 | self.assertFalse(has_triton_tma()) | 28 | self.assertFalse(has_triton_tma()) |
| 31 | self.assertFalse(has_triton_tma_device()) | 29 | self.assertFalse(has_triton_tma_device()) |
| 32 | 30 | ||
| @@ -34,4 +32,4 @@ class TestHasTriton(TestCase): | |||
| 34 | 32 | ||
| 35 | 33 | ||
| 36 | if __name__ == "__main__": | 34 | if __name__ == "__main__": |
| 37 | - run_tests() | 35 | + run_tests() |
| @@ -21,43 +21,8 @@ def patch_is_gpu(): | |||
| 21 | 21 | ||
| 22 | 22 | ||
| 23 | def patch_has_triton(): | 23 | def patch_has_triton(): |
| 24 | - from torch.utils._triton import has_triton_package | 24 | + from torch_npu.utils._dynamo import has_triton |
| 25 | 25 | ||
| 26 | - | ||
| 27 | - def has_triton() -> bool: | ||
| 28 | - if not has_triton_package(): | ||
| 29 | - return False | ||
| 30 | - | ||
| 31 | - from torch._dynamo.device_interface import get_interface_for_device | ||
| 32 | - | ||
| 33 | - def cuda_extra_check(device_interface): | ||
| 34 | - return True | ||
| 35 | - | ||
| 36 | - def cpu_extra_check(device_interface): | ||
| 37 | - import triton.backends | ||
| 38 | - | ||
| 39 | - return "cpu" in triton.backends.backends | ||
| 40 | - | ||
| 41 | - def _return_true(device_interface): | ||
| 42 | - return True | ||
| 43 | - | ||
| 44 | - triton_supported_devices = { | ||
| 45 | - "cuda": cuda_extra_check, | ||
| 46 | - "xpu": _return_true, | ||
| 47 | - "cpu": cpu_extra_check, | ||
| 48 | - "npu": _return_true, | ||
| 49 | - } | ||
| 50 | - | ||
| 51 | - def is_device_compatible_with_triton(): | ||
| 52 | - for device, extra_check in triton_supported_devices.items(): | ||
| 53 | - device_interface = get_interface_for_device(device) | ||
| 54 | - if device_interface.is_available() and extra_check(device_interface): | ||
| 55 | - return True | ||
| 56 | - return False | ||
| 57 | - | ||
| 58 | - return is_device_compatible_with_triton() | ||
| 59 | - | ||
| 60 | - torch.utils._triton.has_triton = has_triton | ||
| 61 | torch._inductor.scheduler.has_triton = has_triton | 26 | torch._inductor.scheduler.has_triton = has_triton |
| 62 | torch._inductor.compile_fx.has_triton = has_triton | 27 | torch._inductor.compile_fx.has_triton = has_triton |
| 63 | 28 | ||
| @@ -488,7 +488,49 @@ def add_dynamo_methods_init(): | |||
| 488 | patch_npu_stream_context() | 488 | patch_npu_stream_context() |
| 489 | 489 | ||
| 490 | 490 | ||
| 491 | + | ||
| 492 | +def has_triton() -> bool: | ||
| 493 | + from torch.utils._triton import has_triton_package | ||
| 494 | + | ||
| 495 | + if not has_triton_package(): | ||
| 496 | + return False | ||
| 497 | + | ||
| 498 | + from torch._dynamo.device_interface import get_interface_for_device | ||
| 499 | + | ||
| 500 | + def cuda_extra_check(device_interface): | ||
| 501 | + return True | ||
| 502 | + | ||
| 503 | + def cpu_extra_check(device_interface): | ||
| 504 | + import triton.backends | ||
| 505 | + | ||
| 506 | + return "cpu" in triton.backends.backends | ||
| 507 | + | ||
| 508 | + def _return_true(device_interface): | ||
| 509 | + return True | ||
| 510 | + | ||
| 511 | + triton_supported_devices = { | ||
| 512 | + "cuda": cuda_extra_check, | ||
| 513 | + "xpu": _return_true, | ||
| 514 | + "cpu": cpu_extra_check, | ||
| 515 | + "npu": _return_true, | ||
| 516 | + } | ||
| 517 | + | ||
| 518 | + def is_device_compatible_with_triton(): | ||
| 519 | + for device, extra_check in triton_supported_devices.items(): | ||
| 520 | + device_interface = get_interface_for_device(device) | ||
| 521 | + if device_interface.is_available() and extra_check(device_interface): | ||
| 522 | + return True | ||
| 523 | + return False | ||
| 524 | + | ||
| 525 | + return is_device_compatible_with_triton() | ||
| 526 | + | ||
| 527 | + | ||
| 528 | +def patch_has_triton(): | ||
| 529 | + torch.utils._triton.has_triton = has_triton | ||
| 530 | + | ||
| 531 | + | ||
| 491 | def add_dynamo_methods(): | 532 | def add_dynamo_methods(): |
| 492 | patch_dynamo_optimize() | 533 | patch_dynamo_optimize() |
| 493 | patch_builtin_variable() | 534 | patch_builtin_variable() |
| 494 | patch_inductor_wrapper() | 535 | patch_inductor_wrapper() |
| 536 | + patch_has_triton() | ||