已合并
【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
已合并
chengpeng25创建于 4月1日
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 tensors221 // 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 tensors159 // 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));