已合并
add eager dvm testcase and remove _TORCH_NPU_ENABLE_DVM #40316
hbhu_bin创建于 7月7日
add eager dvm testcase and remove _TORCH_NPU_ENABLE_DVM #40316
已合并
共 4 个文件变更+83-7
| @@ -16,9 +16,6 @@ from sysconfig import get_paths | |||
| 16 | # Disable autoloading before running 'import torch' to avoid circular dependencies | 16 | # Disable autoloading before running 'import torch' to avoid circular dependencies |
| 17 | os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" | 17 | os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" |
| 18 | 18 | ||
| 19 | -# Opt op-plugin into compiling the DVM lazy-fusion sources on this branch | ||
| 20 | -os.environ.setdefault("_TORCH_NPU_ENABLE_DVM", "1") | ||
| 21 | - | ||
| 22 | from torchnpugen.utils import PathManager | 19 | from torchnpugen.utils import PathManager |
| 23 | 20 | ||
| 24 | 21 | ||
| @@ -28,9 +28,6 @@ from wheel.bdist_wheel import bdist_wheel | |||
| 28 | # Disable autoloading before running 'import torch' to avoid circular dependencies | 28 | # Disable autoloading before running 'import torch' to avoid circular dependencies |
| 29 | os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" | 29 | os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" |
| 30 | 30 | ||
| 31 | -# Opt op-plugin into compiling the DVM lazy-fusion sources on this branch | ||
| 32 | -os.environ.setdefault("_TORCH_NPU_ENABLE_DVM", "1") | ||
| 33 | - | ||
| 34 | from torchnpugen.utils import PathManager | 31 | from torchnpugen.utils import PathManager |
| 35 | 32 | ||
| 36 | BASE_DIR = os.path.dirname(os.path.realpath(__file__)) | 33 | BASE_DIR = os.path.dirname(os.path.realpath(__file__)) |
| @@ -692,6 +692,45 @@ result = out.cpu() | |||
| 692 | """) | 692 | """) |
| 693 | self._assert_match(r_off, r_on) | 693 | self._assert_match(r_off, r_on) |
| 694 | 694 | ||
| 695 | + def test_viewstore_inplace_strided_3d_view(self): | ||
| 696 | + """Inplace writes to a non-contiguous innermost-contiguous view use viewstore.""" | ||
| 697 | + r_off, r_on = compare_dvm_on_off(""" | ||
| 698 | +base = torch.randn(2, 5, 8, dtype=torch.float32).npu() | ||
| 699 | +other = torch.randn(2, 3, 5, dtype=torch.float32).npu() | ||
| 700 | +view = base[:, 1:4, 2:7] | ||
| 701 | + | ||
| 702 | +view.add_(other, alpha=0.25) | ||
| 703 | +view.tanh_() | ||
| 704 | + | ||
| 705 | +result = { | ||
| 706 | + "base": base.cpu(), | ||
| 707 | + "view": view.cpu(), | ||
| 708 | +} | ||
| 709 | +""") | ||
| 710 | + self._assert_match(r_off, r_on, atol=1e-6) | ||
| 711 | + | ||
| 712 | + def test_viewstore_inplace_sibling_views(self): | ||
| 713 | + """Multiple non-contiguous sibling views can be updated through viewstore.""" | ||
| 714 | + r_off, r_on = compare_dvm_on_off(""" | ||
| 715 | +base = torch.randn(4, 10, dtype=torch.float32).npu() | ||
| 716 | +left = base[:, 1:5] | ||
| 717 | +right = base[:, 5:9] | ||
| 718 | +left_src = torch.randn(4, 4, dtype=torch.float32).npu() | ||
| 719 | +right_src = torch.randn(4, 4, dtype=torch.float32).npu() | ||
| 720 | + | ||
| 721 | +left.mul_(1.5) | ||
| 722 | +left.add_(left_src) | ||
| 723 | +right.sub_(right_src, alpha=0.5) | ||
| 724 | +right.relu_() | ||
| 725 | + | ||
| 726 | +result = { | ||
| 727 | + "base": base.cpu(), | ||
| 728 | + "left": left.cpu(), | ||
| 729 | + "right": right.cpu(), | ||
| 730 | +} | ||
| 731 | +""") | ||
| 732 | + self._assert_match(r_off, r_on, atol=1e-6) | ||
| 733 | + | ||
| 695 | class TestDvmFusionPatterns(_DvmTestBase): | 734 | class TestDvmFusionPatterns(_DvmTestBase): |
| 696 | """Complex fusion patterns: chains, RMSNorm, SwiGLU, attention, etc.""" | 735 | """Complex fusion patterns: chains, RMSNorm, SwiGLU, attention, etc.""" |
| 697 | 736 | ||
| @@ -921,5 +960,48 @@ result = out.cpu() | |||
| 921 | msg=f" shape={shape} with_weight={with_weight} with_bias={with_bias}" | 960 | msg=f" shape={shape} with_weight={with_weight} with_bias={with_bias}" |
| 922 | ) | 961 | ) |
| 923 | 962 | ||
| 963 | + def test_native_batch_norm_backward_train_weight_only(self): | ||
| 964 | + r_off, r_on = compare_dvm_on_off(""" | ||
| 965 | +x = torch.randn(4, 3, 5, 5, dtype=torch.float32).npu() | ||
| 966 | +grad = torch.randn_like(x) | ||
| 967 | +weight = torch.randn(3, dtype=torch.float32).npu() | ||
| 968 | +save_mean = torch.randn(3, dtype=torch.float32).npu() | ||
| 969 | +save_var = (torch.rand(3, dtype=torch.float32) + 0.2).npu() | ||
| 970 | + | ||
| 971 | +out = torch.ops.aten.native_batch_norm_backward.default( | ||
| 972 | + grad, x, weight, None, None, save_mean, save_var, True, 1e-5, | ||
| 973 | + [False, True, False], | ||
| 974 | +) | ||
| 975 | +result = tuple(None if t is None else t.cpu() for t in out) | ||
| 976 | +""") | ||
| 977 | + self._assert_match(r_off, r_on, atol=1e-4) | ||
| 978 | + | ||
| 979 | + def test_native_batch_norm_backward_train_masks(self): | ||
| 980 | + cases = [ | ||
| 981 | + ((4, 3, 5, 5), [True, False, False]), | ||
| 982 | + ((4, 3, 5, 5), [False, False, True]), | ||
| 983 | + ((2, 4, 6), [True, True, True]), | ||
| 984 | + ] | ||
| 985 | + for shape, mask in cases: | ||
| 986 | + with self.subTest(shape=shape, mask=mask): | ||
| 987 | + r_off, r_on = compare_dvm_on_off(f""" | ||
| 988 | +shape = {shape} | ||
| 989 | +mask = {mask} | ||
| 990 | +c = shape[1] | ||
| 991 | + | ||
| 992 | +x = torch.randn(shape, dtype=torch.float32).npu() | ||
| 993 | +grad = torch.randn_like(x) | ||
| 994 | +weight = torch.randn(c, dtype=torch.float32).npu() | ||
| 995 | +save_mean = torch.randn(c, dtype=torch.float32).npu() | ||
| 996 | +save_var = (torch.rand(c, dtype=torch.float32) + 0.2).npu() | ||
| 997 | + | ||
| 998 | +out = torch.ops.aten.native_batch_norm_backward.default( | ||
| 999 | + grad, x, weight, None, None, save_mean, save_var, True, 1e-5, | ||
| 1000 | + mask, | ||
| 1001 | +) | ||
| 1002 | +result = tuple(None if t is None else t.cpu() for t in out) | ||
| 1003 | +""") | ||
| 1004 | + self._assert_match(r_off, r_on, atol=1e-4, msg=f" shape={shape} mask={mask}") | ||
| 1005 | + | ||
| 924 | if __name__ == "__main__": | 1006 | if __name__ == "__main__": |
| 925 | unittest.main() | 1007 | unittest.main() |