已合并
【fix】fix empty_like output contiguous tensor in copy h2d/d2h #32904
chengpeng25创建于 4月1日
【fix】fix empty_like output contiguous tensor in copy h2d/d2h #32904
已合并
共 3 个文件变更+142-4
| @@ -0,0 +1,138 @@ | |||
| 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 | +from torch_npu.testing.common_utils import create_common_tensor | ||
| 8 | + | ||
| 9 | + | ||
| 10 | +class TestCopyKernelMemoryFormat(TestCase): | ||
| 11 | + def test_h2d_copy_contiguous_tensor(self): | ||
| 12 | + dtype_list = [np.float16, np.float32, np.int32, np.int64] | ||
| 13 | + shape_list = [[10, 20], [32, 64, 128], [2, 3, 4, 5]] | ||
| 14 | + shape_format = [ | ||
| 15 | + [dtype, 2, shape] for dtype in dtype_list for shape in shape_list | ||
| 16 | + ] | ||
| 17 | + | ||
| 18 | + for item in shape_format: | ||
| 19 | + cpu_input, npu_input = create_common_tensor(item, -100, 100) | ||
| 20 | + npu_input_copy = cpu_input.npu() | ||
| 21 | + self.assertRtolEqual(npu_input_copy.cpu().numpy(), cpu_input.numpy()) | ||
| 22 | + | ||
| 23 | + def test_h2d_copy_non_contiguous_tensor(self): | ||
| 24 | + dtype_list = [np.float16, np.float32] | ||
| 25 | + shape_list = [[32, 64], [16, 32, 64]] | ||
| 26 | + shape_format = [ | ||
| 27 | + [dtype, 2, shape] for dtype in dtype_list for shape in shape_list | ||
| 28 | + ] | ||
| 29 | + | ||
| 30 | + for item in shape_format: | ||
| 31 | + cpu_input, npu_input = create_common_tensor(item, -100, 100) | ||
| 32 | + cpu_transposed = cpu_input.transpose(-1, -2) | ||
| 33 | + npu_transposed = cpu_transposed.npu() | ||
| 34 | + npu_contiguous = npu_transposed.contiguous() | ||
| 35 | + self.assertRtolEqual(npu_contiguous.cpu().numpy(), cpu_transposed.contiguous().numpy()) | ||
| 36 | + | ||
| 37 | + def test_d2h_copy_contiguous_tensor(self): | ||
| 38 | + dtype_list = [np.float16, np.float32, np.int32, np.int64] | ||
| 39 | + shape_list = [[10, 20], [32, 64, 128], [2, 3, 4, 5]] | ||
| 40 | + shape_format = [ | ||
| 41 | + [dtype, 2, shape] for dtype in dtype_list for shape in shape_list | ||
| 42 | + ] | ||
| 43 | + | ||
| 44 | + for item in shape_format: | ||
| 45 | + cpu_input, npu_input = create_common_tensor(item, -100, 100) | ||
| 46 | + cpu_output = npu_input.cpu() | ||
| 47 | + self.assertRtolEqual(cpu_output.numpy(), cpu_input.numpy()) | ||
| 48 | + | ||
| 49 | + def test_d2h_copy_non_contiguous_tensor(self): | ||
| 50 | + dtype_list = [np.float16, np.float32] | ||
| 51 | + shape_list = [[32, 64], [16, 32, 64]] | ||
| 52 | + shape_format = [ | ||
| 53 | + [dtype, 2, shape] for dtype in dtype_list for shape in shape_list | ||
| 54 | + ] | ||
| 55 | + | ||
| 56 | + for item in shape_format: | ||
| 57 | + cpu_input, npu_input = create_common_tensor(item, -100, 100) | ||
| 58 | + npu_transposed = npu_input.transpose(-1, -2) | ||
| 59 | + cpu_output = npu_transposed.cpu() | ||
| 60 | + self.assertRtolEqual(cpu_output.numpy(), cpu_input.transpose(-1, -2).contiguous().numpy()) | ||
| 61 | + | ||
| 62 | + def test_h2d_copy_different_dtype(self): | ||
| 63 | + src_dtype_list = [np.float32, np.float16] | ||
| 64 | + dst_dtype_list = [torch.float16, torch.float32] | ||
| 65 | + shape = [32, 64] | ||
| 66 | + | ||
| 67 | + for src_dtype, dst_dtype in zip(src_dtype_list, dst_dtype_list): | ||
| 68 | + cpu_input = torch.randn(shape, dtype=torch.float32) * 100 | ||
| 69 | + cpu_input = cpu_input.to(torch.from_numpy(np.array([])).dtype if src_dtype == np.float32 else torch.float16) | ||
| 70 | + | ||
| 71 | + npu_input = cpu_input.npu() | ||
| 72 | + npu_output = npu_input.to(dst_dtype) | ||
| 73 | + | ||
| 74 | + cpu_output = cpu_input.to(dst_dtype) | ||
| 75 | + self.assertRtolEqual(npu_output.cpu().numpy(), cpu_output.numpy()) | ||
| 76 | + | ||
| 77 | + def test_d2h_copy_different_dtype(self): | ||
| 78 | + dtype_pairs = [ | ||
| 79 | + (np.float16, torch.float32), | ||
| 80 | + (np.float32, torch.float16), | ||
| 81 | + ] | ||
| 82 | + shape = [32, 64] | ||
| 83 | + | ||
| 84 | + for src_dtype, dst_dtype in dtype_pairs: | ||
| 85 | + cpu_input, npu_input = create_common_tensor([src_dtype, 0, shape], -100, 100) | ||
| 86 | + cpu_output = npu_input.cpu().to(dst_dtype) | ||
| 87 | + | ||
| 88 | + expected = cpu_input.to(dst_dtype) | ||
| 89 | + self.assertRtolEqual(cpu_output.numpy(), expected.numpy()) | ||
| 90 | + | ||
| 91 | + def test_h2d_copy_slice_tensor(self): | ||
| 92 | + shape = [64, 128] | ||
| 93 | + cpu_input = torch.randn(shape) | ||
| 94 | + | ||
| 95 | + cpu_slice = cpu_input[10:30, 20:60] | ||
| 96 | + npu_slice = cpu_slice.npu() | ||
| 97 | + | ||
| 98 | + npu_contiguous = npu_slice.contiguous() | ||
| 99 | + self.assertRtolEqual(npu_contiguous.cpu().numpy(), cpu_slice.contiguous().numpy()) | ||
| 100 | + | ||
| 101 | + def test_d2h_copy_slice_tensor(self): | ||
| 102 | + shape = [64, 128] | ||
| 103 | + cpu_input = torch.randn(shape) | ||
| 104 | + npu_input = cpu_input.npu() | ||
| 105 | + | ||
| 106 | + npu_slice = npu_input[10:30, 20:60] | ||
| 107 | + cpu_slice = npu_slice.cpu() | ||
| 108 | + self.assertRtolEqual(cpu_slice.numpy(), cpu_input[10:30, 20:60].contiguous().numpy()) | ||
| 109 | + | ||
| 110 | + def test_d2h_copy_broadcast_tensor(self): | ||
| 111 | + shape = [1, 64, 1] | ||
| 112 | + cpu_input = torch.randn(shape) | ||
| 113 | + npu_input = cpu_input.npu() | ||
| 114 | + | ||
| 115 | + npu_broadcast = npu_input.expand(4, 64, 128) | ||
| 116 | + cpu_output = npu_broadcast.cpu() | ||
| 117 | + self.assertRtolEqual(cpu_output.numpy(), cpu_input.expand(4, 64, 128).contiguous().numpy()) | ||
| 118 | + | ||
| 119 | + def test_h2d_copy_permute_tensor(self): | ||
| 120 | + shape = [32, 64, 128] | ||
| 121 | + cpu_input = torch.randn(shape) | ||
| 122 | + cpu_permuted = cpu_input.permute(2, 0, 1) | ||
| 123 | + npu_permuted = cpu_permuted.npu() | ||
| 124 | + npu_contiguous = npu_permuted.contiguous() | ||
| 125 | + self.assertRtolEqual(npu_contiguous.cpu().numpy(), cpu_permuted.contiguous().numpy()) | ||
| 126 | + | ||
| 127 | + def test_d2h_copy_permute_tensor(self): | ||
| 128 | + shape = [32, 64, 128] | ||
| 129 | + cpu_input = torch.randn(shape) | ||
| 130 | + npu_input = cpu_input.npu() | ||
| 131 | + | ||
| 132 | + npu_permuted = npu_input.permute(2, 0, 1) | ||
| 133 | + cpu_output = npu_permuted.cpu() | ||
| 134 | + self.assertRtolEqual(cpu_output.numpy(), cpu_input.permute(2, 0, 1).contiguous().numpy()) | ||
| 135 | + | ||
| 136 | + | ||
| 137 | +if __name__ == "__main__": | ||
| 138 | + run_tests() | ||
| @@ -187,7 +187,7 @@ void copy_h2d_baseformat( | |||
| 187 | return; | 187 | return; |
| 188 | } | 188 | } |
| 189 | 189 | ||
| 190 | - at::Tensor dst_contig = dst_is_contiguous ? dst : at::empty_like(dst); | 190 | + at::Tensor dst_contig = dst_is_contiguous ? dst : at::empty_like(dst, LEGACY_CONTIGUOUS_MEMORY_FORMAT); |
| 191 | at::Tensor src_contig; | 191 | at::Tensor src_contig; |
| 192 | if (!same_type) { | 192 | if (!same_type) { |
| 193 | src_contig = src.to(dst.dtype()).expand_as(dst).contiguous(); | 193 | src_contig = src.to(dst.dtype()).expand_as(dst).contiguous(); |
| @@ -216,7 +216,7 @@ void copy_d2h_baseformat(at::Tensor& dst, const at::Tensor& src, bool non_blocki | |||
| 216 | return; | 216 | return; |
| 217 | } | 217 | } |
| 218 | at::Tensor dst_contig = | 218 | at::Tensor dst_contig = |
| 219 | - (dst_is_contiguous && same_type) ? dst : at::empty_like(dst, src.dtype()); | 219 | + (dst_is_contiguous && same_type) ? dst : at::empty_like(dst, src.dtype(), LEGACY_CONTIGUOUS_MEMORY_FORMAT); |
| 220 | at::Tensor src_contig = src.expand_as(dst).contiguous(); | 220 | at::Tensor src_contig = src.expand_as(dst).contiguous(); |
| 221 | // perform a same-dtype copy on contiguous tensors | 221 | // perform a same-dtype copy on contiguous tensors |
| 222 | TORCH_INTERNAL_ASSERT(dst_contig.sizes().equals(src_contig.sizes())); | 222 | TORCH_INTERNAL_ASSERT(dst_contig.sizes().equals(src_contig.sizes())); |
| @@ -125,7 +125,7 @@ void copy_h2d_baseformat_opapi(at::Tensor& dst, const at::Tensor& src, bool non_ | |||
| 125 | return; | 125 | return; |
| 126 | } | 126 | } |
| 127 | 127 | ||
| 128 | - at::Tensor dst_contig = dst_is_contiguous ? dst : at::empty_like(dst); | 128 | + at::Tensor dst_contig = dst_is_contiguous ? dst : at::empty_like(dst, LEGACY_CONTIGUOUS_MEMORY_FORMAT); |
| 129 | at::Tensor src_contig; | 129 | at::Tensor src_contig; |
| 130 | if (!same_type) { | 130 | if (!same_type) { |
| 131 | src_contig = src.to(dst.dtype()).expand_as(dst).contiguous(); | 131 | src_contig = src.to(dst.dtype()).expand_as(dst).contiguous(); |
| @@ -154,7 +154,7 @@ void copy_d2h_baseformat_opapi(at::Tensor& dst, const at::Tensor& src, bool non_ | |||
| 154 | copy_d2h_baseformat_dtype_contigous_opapi(dst, src, non_blocking); | 154 | copy_d2h_baseformat_dtype_contigous_opapi(dst, src, non_blocking); |
| 155 | return; | 155 | return; |
| 156 | } | 156 | } |
| 157 | - at::Tensor dst_contig = (dst_is_contiguous && same_type) ? dst : at::empty_like(dst, src.dtype()); | 157 | + at::Tensor dst_contig = (dst_is_contiguous && same_type) ? dst : at::empty_like(dst, src.dtype(), LEGACY_CONTIGUOUS_MEMORY_FORMAT); |
| 158 | at::Tensor src_contig = src.expand_as(dst).contiguous(); | 158 | at::Tensor src_contig = src.expand_as(dst).contiguous(); |
| 159 | // perform a same-dtype copy on contiguous tensors | 159 | // perform a same-dtype copy on contiguous tensors |
| 160 | TORCH_INTERNAL_ASSERT(dst_contig.sizes().equals(src_contig.sizes()), OPS_ERROR(ErrCode::VALUE)); | 160 | TORCH_INTERNAL_ASSERT(dst_contig.sizes().equals(src_contig.sizes()), OPS_ERROR(ErrCode::VALUE)); |