已合并
fix: handle zero tensor src to avoid null pointer in aclrtMemcpy #33383
wuyouqi1创建于 4月9日
fix: handle zero tensor src to avoid null pointer in aclrtMemcpy #33383
已合并
从已删除 :v2.9.0-issue146合入到Ascend/pytorchv2.9.0
共 3 个文件变更+74-0
| @@ -0,0 +1,68 @@ | |||
| 1 | +import unittest | ||
| 2 | +import torch | ||
| 3 | +import numpy as np | ||
| 4 | + | ||
| 5 | +import torch_npu | ||
| 6 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 7 | + | ||
| 8 | + | ||
| 9 | +class TestCopyZeroTensor(TestCase): | ||
| 10 | + """Test copying zero tensor to NPU tensor. | ||
| 11 | + | ||
| 12 | + This test verifies the fix for issue 146 where copying a CPU zero tensor | ||
| 13 | + (created by _efficientzerotensor) to NPU would fail with null pointer error | ||
| 14 | + in aclrtMemcpy. The fix checks _is_zerotensor() and calls zero_() instead. | ||
| 15 | + """ | ||
| 16 | + | ||
| 17 | + def test_copy_zero_tensor_to_npu(self): | ||
| 18 | + """Test copying CPU zero tensor to NPU tensor.""" | ||
| 19 | + zero_tensor = torch._efficientzerotensor((3, 4), dtype=torch.float32) | ||
| 20 | + | ||
| 21 | + self.assertTrue(zero_tensor._is_zerotensor()) | ||
| 22 | + npu_tensor = torch.randn(3, 4).npu() | ||
| 23 | + npu_tensor.copy_(zero_tensor) | ||
| 24 | + | ||
| 25 | + expected = torch.zeros(3, 4) | ||
| 26 | + self.assertRtolEqual(npu_tensor.cpu().numpy(), expected.numpy()) | ||
| 27 | + | ||
| 28 | + def test_copy_zero_tensor_various_shapes(self): | ||
| 29 | + """Test copying zero tensor with various shapes.""" | ||
| 30 | + shapes = [(2, 3), (5,), (2, 3, 4), (10, 5)] | ||
| 31 | + | ||
| 32 | + for shape in shapes: | ||
| 33 | + zero_tensor = torch._efficientzerotensor(shape, dtype=torch.float32) | ||
| 34 | + npu_tensor = torch.ones(shape, dtype=torch.float32).npu() | ||
| 35 | + npu_tensor.copy_(zero_tensor) | ||
| 36 | + expected = torch.zeros(shape, dtype=torch.float32) | ||
| 37 | + self.assertRtolEqual(npu_tensor.cpu().numpy(), expected.numpy()) | ||
| 38 | + | ||
| 39 | + def test_copy_zero_tensor_various_dtypes(self): | ||
| 40 | + """Test copying zero tensor with various dtypes.""" | ||
| 41 | + dtypes = [torch.float32, torch.float16, torch.int32] | ||
| 42 | + shape = (3, 4) | ||
| 43 | + | ||
| 44 | + for dtype in dtypes: | ||
| 45 | + zero_tensor = torch._efficientzerotensor(shape, dtype=dtype) | ||
| 46 | + npu_tensor = torch.ones(shape, dtype=dtype).npu() | ||
| 47 | + npu_tensor.copy_(zero_tensor) | ||
| 48 | + expected = torch.zeros(shape, dtype=dtype) | ||
| 49 | + self.assertRtolEqual(npu_tensor.cpu().numpy(), expected.numpy()) | ||
| 50 | + | ||
| 51 | + def test_copy_zero_tensor_from_functorch(self): | ||
| 52 | + """Test the original failing case from functorch test. | ||
| 53 | + | ||
| 54 | + This reproduces the scenario from test_linearize_composition_grad_npu_float32 | ||
| 55 | + where _efficientzerotensor falls back to CPU and then gets copied to NPU. | ||
| 56 | + """ | ||
| 57 | + shape = (2, 3) | ||
| 58 | + dtype = torch.float32 | ||
| 59 | + zero_tensor = torch._efficientzerotensor(shape, dtype=dtype) | ||
| 60 | + npu_tensor = torch.randn(shape, dtype=dtype).npu() | ||
| 61 | + | ||
| 62 | + npu_tensor.copy_(zero_tensor) | ||
| 63 | + expected = torch.zeros(shape, dtype=dtype) | ||
| 64 | + self.assertRtolEqual(npu_tensor.cpu().numpy(), expected.numpy()) | ||
| 65 | + | ||
| 66 | + | ||
| 67 | +if __name__ == "__main__": | ||
| 68 | + run_tests() | ||
| @@ -436,6 +436,9 @@ at::Tensor& NPUNativeFunctions::copy_(at::Tensor& self, const at::Tensor& src, b | |||
| 436 | if (self.numel() == 0) { | 436 | if (self.numel() == 0) { |
| 437 | return self; | 437 | return self; |
| 438 | } | 438 | } |
| 439 | + if (src._is_zerotensor()) { | ||
| 440 | + return self.zero_(); | ||
| 441 | + } | ||
| 439 | // save tensor dim name | 442 | // save tensor dim name |
| 440 | c10::optional<at::DimnameList> names = src.opt_names(); | 443 | c10::optional<at::DimnameList> names = src.opt_names(); |
| 441 | if (names.has_value()) { | 444 | if (names.has_value()) { |
| @@ -173,6 +173,9 @@ at::Tensor& NPUNativeOpApiFunctions::copy_(at::Tensor& self, const at::Tensor& s | |||
| 173 | if (self.numel() == 0) { | 173 | if (self.numel() == 0) { |
| 174 | return self; | 174 | return self; |
| 175 | } | 175 | } |
| 176 | + if (src._is_zerotensor()) { | ||
| 177 | + return self.zero_(); | ||
| 178 | + } | ||
| 176 | 179 | ||
| 177 | auto maybe_outnames = at::namedinference::compute_broadcast_outnames(self, src); | 180 | auto maybe_outnames = at::namedinference::compute_broadcast_outnames(self, src); |
| 178 | 181 | ||