已合并
Align NPU mixed-dtype async host-device copy behavior with CUDA #39577
zzhongmin创建于 6月30日
Align NPU mixed-dtype async host-device copy behavior with CUDA #39577
已合并
共 3 个文件变更+427-12
| @@ -7,10 +7,9 @@ import torch | |||
| 7 | from torch.autograd import Variable | 7 | from torch.autograd import Variable |
| 8 | 8 | ||
| 9 | import torch_npu | 9 | import torch_npu |
| 10 | -from torch_npu.testing.common_utils import freeze_rng_state | 10 | +from torch_npu.testing.common_utils import SupportedDevices, freeze_rng_state |
| 11 | from torch_npu.testing.testcase import run_tests, TestCase | 11 | from torch_npu.testing.testcase import run_tests, TestCase |
| 12 | 12 | ||
| 13 | - | ||
| 14 | class TestNpu(TestCase): | 13 | class TestNpu(TestCase): |
| 15 | FIFTY_MIL_CYCLES = 50000000 | 14 | FIFTY_MIL_CYCLES = 50000000 |
| 16 | 15 | ||
| @@ -370,6 +369,24 @@ class TestNpu(TestCase): | |||
| 370 | ) | 369 | ) |
| 371 | _test_to_non_blocking(src, try_non_blocking, dst) | 370 | _test_to_non_blocking(src, try_non_blocking, dst) |
| 372 | 371 | ||
| 372 | + | ||
| 373 | + def test_to_non_blocking_different_dtype(self): | ||
| 374 | + stream = torch_npu.npu.current_stream() | ||
| 375 | + | ||
| 376 | + def _test_to_non_blocking_different_dtype(src, non_blocking, dst, dtype): | ||
| 377 | + torch_npu.npu.synchronize() | ||
| 378 | + out = src.to(device=dst, dtype=dtype, non_blocking=non_blocking) | ||
| 379 | + stream.synchronize() | ||
| 380 | + self.assertEqual(src.to(dtype=dtype), out) | ||
| 381 | + self.assertTrue(out.is_pinned() == (non_blocking and dst == "cpu")) | ||
| 382 | + | ||
| 383 | + src_cpu = torch.arange(1024, dtype=torch.int32).reshape(128, 8).pin_memory() | ||
| 384 | + src_npu = torch.arange(1024, dtype=torch.int32, device="npu").reshape(128, 8) | ||
| 385 | + | ||
| 386 | + for non_blocking in (True, False): | ||
| 387 | + _test_to_non_blocking_different_dtype(src_cpu, non_blocking, "npu", torch.float32) | ||
| 388 | + _test_to_non_blocking_different_dtype(src_npu, non_blocking, "cpu", torch.float32) | ||
| 389 | + | ||
| 373 | def test_to_cpu_blocking_by_default(self): | 390 | def test_to_cpu_blocking_by_default(self): |
| 374 | src = torch.randn(1000000, device="npu") | 391 | src = torch.randn(1000000, device="npu") |
| 375 | torch_npu.npu.synchronize() | 392 | torch_npu.npu.synchronize() |
| @@ -12,17 +12,19 @@ API 签名：copy_(src, non_blocking=False) -> Tensor | |||
| 12 | | 参数类型 | src 为 Tensor（含标量张量）、与 self dtype 可不同 | 已覆盖 | | 12 | | 参数类型 | src 为 Tensor（含标量张量）、与 self dtype 可不同 | 已覆盖 | |
| 13 | | 传参与不传参 | non_blocking 省略与显式传入 | 已覆盖 | | 13 | | 传参与不传参 | non_blocking 省略与显式传入 | 已覆盖 | |
| 14 | | 等价类/边界值 | 同形、可广播、非连续目标、跨 CPU/NPU | 已覆盖 | | 14 | | 等价类/边界值 | 同形、可广播、非连续目标、跨 CPU/NPU | 已覆盖 | |
| 15 | +| 精度/数值正确性 | mixed-dtype host-device 路径下，同步/异步 copy 结果一致 | 已覆盖 | | ||
| 15 | | 正常传参场景 | NPU 上 copy 后 self 的 shape/dtype 不变；返回 self | 已覆盖 | | 16 | | 正常传参场景 | NPU 上 copy 后 self 的 shape/dtype 不变；返回 self | 已覆盖 | |
| 16 | | 异常传参场景 | 不可广播的 shape | 已覆盖 | | 17 | | 异常传参场景 | 不可广播的 shape | 已覆盖 | |
| 17 | 18 | ||
| 18 | 未覆盖项及原因： | 19 | 未覆盖项及原因： |
| 19 | - 无 | 20 | - 无 |
| 20 | 21 | ||
| 21 | -注意：本测试仅验证功能正确性（调用不报错、tensor 结构属性符合预期）， | 22 | +注意：本测试除了验证功能正确性（调用不报错、tensor 结构属性符合预期）， |
| 22 | - 不做精度和数值正确性校验。 | 23 | + 也对 mixed-dtype host-device 路径补充了同步/异步 copy 结果一致性校验。 |
| 23 | """ | 24 | """ |
| 24 | import torch | 25 | import torch |
| 25 | import torch_npu # noqa: F401 | 26 | import torch_npu # noqa: F401 |
| 27 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 26 | 28 | ||
| 27 | try: | 29 | try: |
| 28 | from torch_npu.testing.testcase import TestCase, run_tests | 30 | from torch_npu.testing.testcase import TestCase, run_tests |
| @@ -34,7 +36,6 @@ except ImportError: | |||
| 34 | def run_tests(): | 36 | def run_tests(): |
| 35 | unittest.main(argv=sys.argv) | 37 | unittest.main(argv=sys.argv) |
| 36 | 38 | ||
| 37 | - | ||
| 38 | class TestTensorCopy_(TestCase): | 39 | class TestTensorCopy_(TestCase): |
| 39 | """Functional tests for torch.Tensor.copy_ on NPU.""" | 40 | """Functional tests for torch.Tensor.copy_ on NPU.""" |
| 40 | 41 | ||
| @@ -47,6 +48,82 @@ class TestTensorCopy_(TestCase): | |||
| 47 | f"Expected device 'npu', got '{self.device_name}'", | 48 | f"Expected device 'npu', got '{self.device_name}'", |
| 48 | ) | 49 | ) |
| 49 | self.device = torch.device(self.device_name) | 50 | self.device = torch.device(self.device_name) |
| 51 | + self.dtype_cast_pairs = [ | ||
| 52 | + (torch.int32, torch.float32), | ||
| 53 | + (torch.int64, torch.float32), | ||
| 54 | + (torch.float16, torch.float32), | ||
| 55 | + (torch.float32, torch.float16), | ||
| 56 | + ] | ||
| 57 | + self.aclnn_cast_fallback_dtypes = [ | ||
| 58 | + torch.float8_e5m2, | ||
| 59 | + torch.float8_e4m3fn, | ||
| 60 | + ] | ||
| 61 | + self.precision_compare_cases = [ | ||
| 62 | + (torch.int8, torch.float16, [-127, -31, -1, 0, 1, 7, 42, 127]), | ||
| 63 | + (torch.int16, torch.float32, [-32768, -1025, -1, 0, 1, 255, 1024, 32767]), | ||
| 64 | + (torch.int64, torch.float32, [-(2 ** 20), -12345, -1, 0, 1, 12345, 4096, 2 ** 20]), | ||
| 65 | + (torch.float16, torch.float32, [-2048.0, -7.5, -0.125, 0.0, 0.125, 1.5, 33.25, 2048.0]), | ||
| 66 | + (torch.bfloat16, torch.float32, [-1.0e4, -7.5, -0.125, 0.0, 0.125, 1.5, 33.25, 1.0e4]), | ||
| 67 | + (torch.float32, torch.float16, [-65504.0, -255.5, -0.33325195, 0.0, 0.33325195, 17.625, 255.5, 65504.0]), | ||
| 68 | + (torch.float32, torch.bfloat16, [-1.0e8, -255.5, -0.33325195, 0.0, 0.33325195, 17.625, 255.5, 1.0e8]), | ||
| 69 | + ] | ||
| 70 | + | ||
| 71 | + def _make_host_source(self, dtype, pin_memory=False): | ||
| 72 | + src = torch.tensor( | ||
| 73 | + [[-7.5, -3.25, -1.0, 0.0], [1.5, 2.25, 5.0, 9.75]], | ||
| 74 | + dtype=torch.float32, | ||
| 75 | + ).to(dtype) | ||
| 76 | + return src.pin_memory() if pin_memory else src | ||
| 77 | + | ||
| 78 | + def _make_device_source(self, dtype): | ||
| 79 | + return self._make_host_source(dtype).to(self.device) | ||
| 80 | + | ||
| 81 | + def _assert_dtype_cast_copy_keeps_async(self, dst, src): | ||
| 82 | + gate_stream = torch_npu.npu.Stream(device=self.device) | ||
| 83 | + copy_stream = torch_npu.npu.Stream(device=self.device) | ||
| 84 | + gate_event = torch_npu.npu.Event() | ||
| 85 | + done_event = torch_npu.npu.Event() | ||
| 86 | + | ||
| 87 | + torch_npu.npu.synchronize() | ||
| 88 | + | ||
| 89 | + # Keep copy_stream pending behind work on gate_stream. A synchronous | ||
| 90 | + # fallback in copy_ would wait for the gate before returning. | ||
| 91 | + gate_a = torch.ones((4096, 4096), device=self.device, dtype=torch.float32) | ||
| 92 | + gate_b = torch.ones((4096, 4096), device=self.device, dtype=torch.float32) | ||
| 93 | + with torch_npu.npu.stream(gate_stream): | ||
| 94 | + gate_c = gate_a @ gate_b | ||
| 95 | + gate_c = gate_c @ gate_b | ||
| 96 | + gate_event.record() | ||
| 97 | + | ||
| 98 | + with torch_npu.npu.stream(copy_stream): | ||
| 99 | + copy_stream.wait_event(gate_event) | ||
| 100 | + ret = dst.copy_(src, non_blocking=True) | ||
| 101 | + done_event.record() | ||
| 102 | + | ||
| 103 | + self.assertIs(ret, dst) | ||
| 104 | + self.assertFalse(done_event.query()) | ||
| 105 | + done_event.synchronize() | ||
| 106 | + | ||
| 107 | + def _assert_copy_matches_cast(self, dst, src): | ||
| 108 | + expected = src.cpu().to(dtype=dst.dtype) | ||
| 109 | + actual = dst.cpu() if dst.device.type == self.device_name else dst | ||
| 110 | + self.assertEqual(actual, expected) | ||
| 111 | + | ||
| 112 | + def _to_cpu_if_needed(self, tensor): | ||
| 113 | + return tensor.cpu() if tensor.device.type == self.device_name else tensor | ||
| 114 | + | ||
| 115 | + def _assert_non_blocking_matches_blocking(self, async_dst, sync_dst, src, async_base=None, sync_base=None): | ||
| 116 | + sync_ret = sync_dst.copy_(src, non_blocking=False) | ||
| 117 | + async_ret = async_dst.copy_(src, non_blocking=True) | ||
| 118 | + | ||
| 119 | + self.assertIs(sync_ret, sync_dst) | ||
| 120 | + self.assertIs(async_ret, async_dst) | ||
| 121 | + | ||
| 122 | + torch_npu.npu.synchronize() | ||
| 123 | + self.assertEqual(self._to_cpu_if_needed(async_dst), self._to_cpu_if_needed(sync_dst)) | ||
| 124 | + | ||
| 125 | + if async_base is not None and sync_base is not None: | ||
| 126 | + self.assertEqual(self._to_cpu_if_needed(async_base), self._to_cpu_if_needed(sync_base)) | ||
| 50 | 127 | ||
| 51 | def test_copy_npu_same_device_same_shape(self): | 128 | def test_copy_npu_same_device_same_shape(self): |
| 52 | dst = torch.empty(3, 4, device=self.device, dtype=torch.float32) | 129 | dst = torch.empty(3, 4, device=self.device, dtype=torch.float32) |
| @@ -87,6 +164,268 @@ class TestTensorCopy_(TestCase): | |||
| 87 | self.assertIs(ret, dst) | 164 | self.assertIs(ret, dst) |
| 88 | self.assertEqual(dst.shape, torch.Size([2, 2])) | 165 | self.assertEqual(dst.shape, torch.Size([2, 2])) |
| 89 | 166 | ||
| 167 | + | ||
| 168 | + def test_copy_npu_from_pinned_cpu_src_dtype_cast_non_blocking(self): | ||
| 169 | + for src_dtype, dst_dtype in self.dtype_cast_pairs: | ||
| 170 | + dst = torch.empty(2, 4, dtype=dst_dtype, device=self.device) | ||
| 171 | + src = self._make_host_source(src_dtype, pin_memory=True) | ||
| 172 | + self._assert_dtype_cast_copy_keeps_async(dst, src) | ||
| 173 | + self._assert_copy_matches_cast(dst, src) | ||
| 174 | + | ||
| 175 | + | ||
| 176 | + def test_copy_pinned_cpu_from_npu_src_dtype_cast_non_blocking(self): | ||
| 177 | + for src_dtype, dst_dtype in self.dtype_cast_pairs: | ||
| 178 | + dst = torch.empty(2, 4, dtype=dst_dtype, pin_memory=True) | ||
| 179 | + src = self._make_device_source(src_dtype) | ||
| 180 | + self._assert_dtype_cast_copy_keeps_async(dst, src) | ||
| 181 | + self.assertTrue(dst.is_pinned()) | ||
| 182 | + self._assert_copy_matches_cast(dst, src) | ||
| 183 | + | ||
| 184 | + | ||
| 185 | + def test_copy_pinned_cpu_from_non_contiguous_npu_src_dtype_cast_non_blocking(self): | ||
| 186 | + src = torch.arange(8, dtype=torch.int32, device=self.device).reshape(4, 2).t() | ||
| 187 | + dst = torch.empty(2, 4, dtype=torch.float32, pin_memory=True) | ||
| 188 | + self.assertFalse(src.is_contiguous()) | ||
| 189 | + self._assert_dtype_cast_copy_keeps_async(dst, src) | ||
| 190 | + self.assertTrue(dst.is_pinned()) | ||
| 191 | + self._assert_copy_matches_cast(dst, src) | ||
| 192 | + | ||
| 193 | + | ||
| 194 | + def test_copy_pinned_cpu_non_contiguous_dst_dtype_cast_preserves_strided_layout(self): | ||
| 195 | + base = torch.full((3, 6), -99.0, dtype=torch.float32).pin_memory() | ||
| 196 | + dst = base[:, 1::2] | ||
| 197 | + src = torch.arange(9, dtype=torch.int32, device=self.device).reshape(3, 3) | ||
| 198 | + expected_base = torch.full((3, 6), -99.0, dtype=torch.float32) | ||
| 199 | + expected_base[:, 1::2] = src.cpu().to(dtype=dst.dtype) | ||
| 200 | + | ||
| 201 | + self.assertFalse(dst.is_contiguous()) | ||
| 202 | + self.assertTrue(dst.is_pinned()) | ||
| 203 | + self.assertNotEqual(dst.storage_offset(), 0) | ||
| 204 | + ret = dst.copy_(src, non_blocking=True) | ||
| 205 | + torch_npu.npu.synchronize() | ||
| 206 | + | ||
| 207 | + self.assertIs(ret, dst) | ||
| 208 | + self.assertEqual(dst, src.cpu().to(dtype=dst.dtype)) | ||
| 209 | + self.assertEqual(base, expected_base) | ||
| 210 | + | ||
| 211 | + | ||
| 212 | + def test_copy_npu_non_contiguous_dst_dtype_cast_non_blocking(self): | ||
| 213 | + base = torch.empty(4, 2, dtype=torch.float32, device=self.device) | ||
| 214 | + dst = base.t() | ||
| 215 | + src = self._make_host_source(torch.int32, pin_memory=True) | ||
| 216 | + self.assertFalse(dst.is_contiguous()) | ||
| 217 | + self._assert_dtype_cast_copy_keeps_async(dst, src) | ||
| 218 | + self._assert_copy_matches_cast(dst, src) | ||
| 219 | + | ||
| 220 | + | ||
| 221 | + def test_copy_npu_non_contiguous_dst_dtype_cast_preserves_strided_layout(self): | ||
| 222 | + base = torch.full((3, 6), -99.0, dtype=torch.float32, device=self.device) | ||
| 223 | + dst = base[:, 1::2] | ||
| 224 | + src = torch.arange(9, dtype=torch.int32).reshape(3, 3).pin_memory() | ||
| 225 | + expected_base = torch.full((3, 6), -99.0, dtype=torch.float32) | ||
| 226 | + expected_base[:, 1::2] = src.to(dtype=dst.dtype) | ||
| 227 | + | ||
| 228 | + self.assertFalse(dst.is_contiguous()) | ||
| 229 | + self.assertNotEqual(dst.storage_offset(), 0) | ||
| 230 | + ret = dst.copy_(src, non_blocking=True) | ||
| 231 | + torch_npu.npu.synchronize() | ||
| 232 | + | ||
| 233 | + self.assertIs(ret, dst) | ||
| 234 | + self.assertEqual(dst.cpu(), src.to(dtype=dst.dtype)) | ||
| 235 | + self.assertEqual(base.cpu(), expected_base) | ||
| 236 | + | ||
| 237 | + | ||
| 238 | + def test_copy_npu_from_pinned_cpu_broadcast_src_dtype_cast_non_blocking(self): | ||
| 239 | + dst = torch.empty(3, 4, dtype=torch.float32, device=self.device) | ||
| 240 | + src = torch.arange(4, dtype=torch.int32).reshape(1, 4).pin_memory() | ||
| 241 | + expected = src.to(dtype=dst.dtype).expand(3, 4) | ||
| 242 | + | ||
| 243 | + ret = dst.copy_(src, non_blocking=True) | ||
| 244 | + torch_npu.npu.synchronize() | ||
| 245 | + | ||
| 246 | + self.assertIs(ret, dst) | ||
| 247 | + self.assertEqual(dst.cpu(), expected) | ||
| 248 | + | ||
| 249 | + | ||
| 250 | + def test_copy_pinned_cpu_from_npu_broadcast_src_dtype_cast_non_blocking(self): | ||
| 251 | + dst = torch.empty(3, 4, dtype=torch.float32, pin_memory=True) | ||
| 252 | + src = torch.arange(4, dtype=torch.int32, device=self.device).reshape(1, 4) | ||
| 253 | + expected = src.cpu().to(dtype=dst.dtype).expand_as(dst) | ||
| 254 | + | ||
| 255 | + ret = dst.copy_(src, non_blocking=True) | ||
| 256 | + torch_npu.npu.synchronize() | ||
| 257 | + | ||
| 258 | + self.assertIs(ret, dst) | ||
| 259 | + self.assertEqual(dst, expected) | ||
| 260 | + | ||
| 261 | + | ||
| 262 | + def test_copy_npu_from_pinned_cpu_aclnn_cast_unsupported_src_dtype_fallback(self): | ||
| 263 | + for src_dtype in self.aclnn_cast_fallback_dtypes: | ||
| 264 | + dst = torch.empty(2, 4, dtype=torch.float32, device=self.device) | ||
| 265 | + src = torch.arange(8, dtype=torch.float32).reshape(2, 4).to(src_dtype).pin_memory() | ||
| 266 | + | ||
| 267 | + ret = dst.copy_(src, non_blocking=True) | ||
| 268 | + torch_npu.npu.synchronize() | ||
| 269 | + | ||
| 270 | + self.assertIs(ret, dst) | ||
| 271 | + self.assertEqual(dst.cpu(), src.to(dtype=dst.dtype)) | ||
| 272 | + | ||
| 273 | + | ||
| 274 | + def test_copy_pinned_cpu_from_npu_aclnn_cast_unsupported_dst_dtype_fallback(self): | ||
| 275 | + for dst_dtype in self.aclnn_cast_fallback_dtypes: | ||
| 276 | + dst = torch.empty(2, 4, dtype=dst_dtype).pin_memory() | ||
| 277 | + src = torch.arange(8, dtype=torch.float32, device=self.device).reshape(2, 4) | ||
| 278 | + | ||
| 279 | + ret = dst.copy_(src, non_blocking=True) | ||
| 280 | + torch_npu.npu.synchronize() | ||
| 281 | + | ||
| 282 | + self.assertIs(ret, dst) | ||
| 283 | + self.assertEqual(dst.dtype, dst_dtype) | ||
| 284 | + self.assertEqual(dst.to(dtype=src.dtype), src.cpu()) | ||
| 285 | + | ||
| 286 | + | ||
| 287 | + def test_copy_npu_from_pinned_cpu_slice_dtype_cast_non_blocking(self): | ||
| 288 | + src_base = torch.arange(9, dtype=torch.int32).pin_memory() | ||
| 289 | + src = src_base[1:].reshape(2, 4) | ||
| 290 | + dst = torch.empty(2, 4, dtype=torch.float32, device=self.device) | ||
| 291 | + self.assertTrue(src.is_pinned()) | ||
| 292 | + self.assertNotEqual(src.data_ptr(), src.untyped_storage().data_ptr()) | ||
| 293 | + self._assert_dtype_cast_copy_keeps_async(dst, src) | ||
| 294 | + self._assert_copy_matches_cast(dst, src) | ||
| 295 | + | ||
| 296 | + | ||
| 297 | + def test_copy_npu_from_pinned_cpu_dtype_cast_non_blocking_matches_blocking(self): | ||
| 298 | + for src_dtype, dst_dtype, values in self.precision_compare_cases: | ||
| 299 | + with self.subTest(direction="h2d", src_dtype=src_dtype, dst_dtype=dst_dtype): | ||
| 300 | + src = torch.tensor(values, dtype=src_dtype).reshape(2, 4).pin_memory() | ||
| 301 | + async_dst = torch.empty(2, 4, dtype=dst_dtype, device=self.device) | ||
| 302 | + sync_dst = torch.empty(2, 4, dtype=dst_dtype, device=self.device) | ||
| 303 | + | ||
| 304 | + self._assert_non_blocking_matches_blocking(async_dst, sync_dst, src) | ||
| 305 | + self._assert_copy_matches_cast(async_dst, src) | ||
| 306 | + | ||
| 307 | + | ||
| 308 | + def test_copy_pinned_cpu_from_npu_dtype_cast_non_blocking_matches_blocking(self): | ||
| 309 | + for src_dtype, dst_dtype, values in self.precision_compare_cases: | ||
| 310 | + with self.subTest(direction="d2h", src_dtype=src_dtype, dst_dtype=dst_dtype): | ||
| 311 | + src = torch.tensor(values, dtype=src_dtype).reshape(2, 4).to(self.device) | ||
| 312 | + async_dst = torch.empty(2, 4, dtype=dst_dtype, pin_memory=True) | ||
| 313 | + sync_dst = torch.empty(2, 4, dtype=dst_dtype, pin_memory=True) | ||
| 314 | + | ||
| 315 | + self._assert_non_blocking_matches_blocking(async_dst, sync_dst, src) | ||
| 316 | + self._assert_copy_matches_cast(async_dst, src) | ||
| 317 | + self.assertTrue(async_dst.is_pinned()) | ||
| 318 | + self.assertTrue(sync_dst.is_pinned()) | ||
| 319 | + | ||
| 320 | + | ||
| 321 | + def test_copy_mixed_dtype_non_blocking_matches_blocking_for_layout_variants(self): | ||
| 322 | + layout_cases = [ | ||
| 323 | + { | ||
| 324 | + "name": "h2d_non_contiguous_dst", | ||
| 325 | + "src_dtype": torch.float32, | ||
| 326 | + "dst_dtype": torch.float16, | ||
| 327 | + "make_src": lambda: torch.tensor( | ||
| 328 | + [-63.5, -7.25, -0.5, 0.0, 0.5, 3.25, 17.75, 63.5, -19.5], | ||
| 329 | + dtype=torch.float32, | ||
| 330 | + ).reshape(3, 3).pin_memory(), | ||
| 331 | + "make_async_dst": lambda: torch.full((3, 6), -99.0, dtype=torch.float16, device=self.device), | ||
| 332 | + "make_sync_dst": lambda: torch.full((3, 6), -99.0, dtype=torch.float16, device=self.device), | ||
| 333 | + "select_dst_view": lambda base: base[:, 1::2], | ||
| 334 | + "expected": lambda src, dst: src.to(dtype=dst.dtype), | ||
| 335 | + }, | ||
| 336 | + { | ||
| 337 | + "name": "d2h_non_contiguous_dst", | ||
| 338 | + "src_dtype": torch.int32, | ||
| 339 | + "dst_dtype": torch.float32, | ||
| 340 | + "make_src": lambda: torch.arange(9, dtype=torch.int32, device=self.device).reshape(3, 3) - 4, | ||
| 341 | + "make_async_dst": lambda: torch.full((3, 6), -99.0, dtype=torch.float32).pin_memory(), | ||
| 342 | + "make_sync_dst": lambda: torch.full((3, 6), -99.0, dtype=torch.float32).pin_memory(), | ||
| 343 | + "select_dst_view": lambda base: base[:, 1::2], | ||
| 344 | + "expected": lambda src, dst: src.cpu().to(dtype=dst.dtype), | ||
| 345 | + }, | ||
| 346 | + { | ||
| 347 | + "name": "h2d_broadcast_src", | ||
| 348 | + "src_dtype": torch.int16, | ||
| 349 | + "dst_dtype": torch.float32, | ||
| 350 | + "make_src": lambda: torch.tensor([-32768, -17, 9, 32767], dtype=torch.int16).reshape(1, 4).pin_memory(), | ||
| 351 | + "make_async_dst": lambda: torch.empty(3, 4, dtype=torch.float32, device=self.device), | ||
| 352 | + "make_sync_dst": lambda: torch.empty(3, 4, dtype=torch.float32, device=self.device), | ||
| 353 | + "select_dst_view": lambda base: base, | ||
| 354 | + "expected": lambda src, dst: src.to(dtype=dst.dtype).expand_as(dst), | ||
| 355 | + }, | ||
| 356 | + { | ||
| 357 | + "name": "d2h_non_contiguous_src", | ||
| 358 | + "src_dtype": torch.float16, | ||
| 359 | + "dst_dtype": torch.float32, | ||
| 360 | + "make_src": lambda: torch.tensor( | ||
| 361 | + [-7.5, -1.25, 0.0, 1.25, 3.5, 7.75, 15.5, 31.0], | ||
| 362 | + dtype=torch.float16, | ||
| 363 | + device=self.device, | ||
| 364 | + ).reshape(4, 2).t(), | ||
| 365 | + "make_async_dst": lambda: torch.empty(2, 4, dtype=torch.float32, pin_memory=True), | ||
| 366 | + "make_sync_dst": lambda: torch.empty(2, 4, dtype=torch.float32, pin_memory=True), | ||
| 367 | + "select_dst_view": lambda base: base, | ||
| 368 | + "expected": lambda src, dst: src.cpu().to(dtype=dst.dtype), | ||
| 369 | + }, | ||
| 370 | + ] | ||
| 371 | + | ||
| 372 | + for case in layout_cases: | ||
| 373 | + with self.subTest(case=case["name"], src_dtype=case["src_dtype"], dst_dtype=case["dst_dtype"]): | ||
| 374 | + src = case["make_src"]() | ||
| 375 | + async_base = case["make_async_dst"]() | ||
| 376 | + sync_base = case["make_sync_dst"]() | ||
| 377 | + async_dst = case["select_dst_view"](async_base) | ||
| 378 | + sync_dst = case["select_dst_view"](sync_base) | ||
| 379 | + | ||
| 380 | + self._assert_non_blocking_matches_blocking( | ||
| 381 | + async_dst, | ||
| 382 | + sync_dst, | ||
| 383 | + src, | ||
| 384 | + async_base=async_base, | ||
| 385 | + sync_base=sync_base, | ||
| 386 | + ) | ||
| 387 | + | ||
| 388 | + self.assertEqual( | ||
| 389 | + self._to_cpu_if_needed(async_dst), | ||
| 390 | + case["expected"](src, async_dst), | ||
| 391 | + ) | ||
| 392 | + | ||
| 393 | + | ||
| 394 | + def test_copy_dtype_cast_non_blocking_temporary_lifetime(self): | ||
| 395 | + h2d_dst = torch.empty(2, 4, dtype=torch.float32, device=self.device) | ||
| 396 | + d2h_dst = torch.empty(2, 4, dtype=torch.float32, pin_memory=True) | ||
| 397 | + expected = None | ||
| 398 | + | ||
| 399 | + for i in range(32): | ||
| 400 | + host_src = self._make_host_source(torch.int32) + i | ||
| 401 | + h2d_src = host_src.pin_memory() | ||
| 402 | + h2d_dst.copy_(h2d_src, non_blocking=True) | ||
| 403 | + d2h_src = host_src.to(self.device) | ||
| 404 | + d2h_dst.copy_(d2h_src, non_blocking=True) | ||
| 405 | + expected = host_src.to(dtype=d2h_dst.dtype) | ||
| 406 | + | ||
| 407 | + torch_npu.npu.synchronize() | ||
| 408 | + self.assertEqual(h2d_dst.cpu(), h2d_src.to(dtype=h2d_dst.dtype)) | ||
| 409 | + self.assertEqual(d2h_dst, expected) | ||
| 410 | + | ||
| 411 | + | ||
| 412 | + def test_copy_npu_from_cpu_src_dtype_cast_blocking(self): | ||
| 413 | + for src_dtype, dst_dtype in self.dtype_cast_pairs: | ||
| 414 | + dst = torch.empty(2, 4, dtype=dst_dtype, device=self.device) | ||
| 415 | + src = self._make_host_source(src_dtype) | ||
| 416 | + ret = dst.copy_(src, non_blocking=False) | ||
| 417 | + self.assertIs(ret, dst) | ||
| 418 | + self._assert_copy_matches_cast(dst, src) | ||
| 419 | + | ||
| 420 | + | ||
| 421 | + def test_copy_cpu_from_npu_src_dtype_cast_blocking(self): | ||
| 422 | + for src_dtype, dst_dtype in self.dtype_cast_pairs: | ||
| 423 | + dst = torch.empty(2, 4, dtype=dst_dtype) | ||
| 424 | + src = self._make_device_source(src_dtype) | ||
| 425 | + ret = dst.copy_(src, non_blocking=False) | ||
| 426 | + self.assertIs(ret, dst) | ||
| 427 | + self._assert_copy_matches_cast(dst, src) | ||
| 428 | + | ||
| 90 | def test_copy_npu_src_int_dtype_cast(self): | 429 | def test_copy_npu_src_int_dtype_cast(self): |
| 91 | dst = torch.empty(2, 2, dtype=torch.float32, device=self.device) | 430 | dst = torch.empty(2, 2, dtype=torch.float32, device=self.device) |
| 92 | src = torch.ones(2, 2, dtype=torch.int32, device=self.device) | 431 | src = torch.ones(2, 2, dtype=torch.int32, device=self.device) |
| @@ -16,10 +16,14 @@ | |||
| 16 | 16 | ||
| 17 | 17 | ||
| 18 | 18 | ||
| 19 | + | ||
| 19 | 20 | ||
| 21 | + | ||
| 20 | 22 | ||
| 23 | + | ||
| 21 | 24 | ||
| 22 | 25 | ||
| 26 | + | ||
| 23 | 27 | ||
| 24 | 28 | ||
| 25 | 29 | ||
| @@ -31,6 +35,36 @@ | |||
| 31 | namespace at_npu { | 35 | namespace at_npu { |
| 32 | namespace native { | 36 | namespace native { |
| 33 | 37 | ||
| 38 | +namespace { | ||
| 39 | + | ||
| 40 | +bool is_aclnn_cast_unsupported_dtype(const at::Tensor& tensor) | ||
| 41 | +{ | ||
| 42 | + // On A2-and-later products, aclnnCast rejects these dtype families. | ||
| 43 | + aclDataType dtype = c10_npu::GetAclDataType(static_cast<int64_t>(tensor.scalar_type())); | ||
| 44 | + return dtype == aclDataType::ACL_COMPLEX32 || | ||
| 45 | + dtype == aclDataType::ACL_HIFLOAT8 || | ||
| 46 | + dtype == aclDataType::ACL_FLOAT8_E5M2 || | ||
| 47 | + dtype == aclDataType::ACL_FLOAT8_E4M3FN || | ||
| 48 | + dtype == aclDataType::ACL_FLOAT4_E2M1 || | ||
| 49 | + dtype == aclDataType::ACL_FLOAT4_E1M2 || | ||
| 50 | + dtype == aclDataType::ACL_INT4; | ||
| 51 | +} | ||
| 52 | + | ||
| 53 | +bool should_fallback_to_cpu_cast(const at::Tensor& dst, const at::Tensor& src) | ||
| 54 | +{ | ||
| 55 | + const auto soc = c10_npu::GetSocVersion(); | ||
| 56 | + const bool is_a2_or_later = | ||
| 57 | + ((soc >= c10_npu::SocVersion::Ascend910B1 && soc < c10_npu::SocVersion::Ascend310B1) || | ||
| 58 | + (soc >= c10_npu::SocVersion::Ascend910_9391)); | ||
| 59 | + if (!is_a2_or_later) { | ||
| 60 | + return false; | ||
| 61 | + } | ||
| 62 | + return is_aclnn_cast_unsupported_dtype(src) || | ||
| 63 | + is_aclnn_cast_unsupported_dtype(dst); | ||
| 64 | +} | ||
| 65 | + | ||
| 66 | +} // namespace | ||
| 67 | + | ||
| 34 | // the format of dst and src is base format now | 68 | // the format of dst and src is base format now |
| 35 | // the dtype of dst and src is same | 69 | // the dtype of dst and src is same |
| 36 | // and src and dst are contiguous | 70 | // and src and dst are contiguous |
| @@ -111,6 +145,14 @@ void copy_d2h_baseformat_dtype_contigous_opapi(at::Tensor& dst, const at::Tensor | |||
| 111 | copy_between_host_and_device_opapi(dst, src, kind, non_blocking); | 145 | copy_between_host_and_device_opapi(dst, src, kind, non_blocking); |
| 112 | } | 146 | } |
| 113 | 147 | ||
| 148 | +void cast_dtype_out_baseformat_opapi(at::Tensor& dst, const at::Tensor& src) | ||
| 149 | +{ | ||
| 150 | + TORCH_INTERNAL_ASSERT(dst.sizes().equals(src.sizes()), OPS_ERROR(ErrCode::VALUE)); | ||
| 151 | + TORCH_INTERNAL_ASSERT(dst.device() == src.device(), OPS_ERROR(ErrCode::VALUE)); | ||
| 152 | + auto dst_scalar_type = dst.scalar_type(); | ||
| 153 | + EXEC_NPU_CMD(aclnnCast, src, dst_scalar_type, dst); | ||
| 154 | +} | ||
| 155 | + | ||
| 114 | // the format of dst and src is baseformat now | 156 | // the format of dst and src is baseformat now |
| 115 | void copy_h2d_baseformat_opapi(at::Tensor& dst, const at::Tensor& src, bool non_blocking, | 157 | void copy_h2d_baseformat_opapi(at::Tensor& dst, const at::Tensor& src, bool non_blocking, |
| 116 | bool dst_must_be_contiguous = false) | 158 | bool dst_must_be_contiguous = false) |
| @@ -124,12 +166,16 @@ void copy_h2d_baseformat_opapi(at::Tensor& dst, const at::Tensor& src, bool non_ | |||
| 124 | return; | 166 | return; |
| 125 | } | 167 | } |
| 126 | 168 | ||
| 127 | - at::Tensor dst_contig = dst_is_contiguous ? dst : at::empty_like(dst, LEGACY_CONTIGUOUS_MEMORY_FORMAT); | 169 | + at::Tensor dst_contig; |
| 128 | at::Tensor src_contig; | 170 | at::Tensor src_contig; |
| 129 | - if (!same_type) { | 171 | + if (!same_type && non_blocking && !should_fallback_to_cpu_cast(dst, src)) { |
| 130 | - src_contig = src.to(dst.dtype()).expand_as(dst).contiguous(); | 172 | + // keep the H2D leg same-dtype, then cast on device. |
| 131 | - } else { | 173 | + dst_contig = at::empty_like(dst, src.dtype(), LEGACY_CONTIGUOUS_MEMORY_FORMAT); |
| 132 | src_contig = src.expand_as(dst).contiguous(); | 174 | src_contig = src.expand_as(dst).contiguous(); |
| 175 | + } else { | ||
| 176 | + dst_contig = dst_is_contiguous ? dst : at::empty_like(dst, LEGACY_CONTIGUOUS_MEMORY_FORMAT); | ||
| 177 | + src_contig = !same_type ? src.to(dst.dtype()).expand_as(dst).contiguous() | ||
| 178 | + : src.expand_as(dst).contiguous(); | ||
| 133 | } | 179 | } |
| 134 | // perform a same-dtype copy on contiguous tensors | 180 | // perform a same-dtype copy on contiguous tensors |
| 135 | TORCH_INTERNAL_ASSERT(dst_contig.sizes().equals(src_contig.sizes()), OPS_ERROR(ErrCode::VALUE)); | 181 | TORCH_INTERNAL_ASSERT(dst_contig.sizes().equals(src_contig.sizes()), OPS_ERROR(ErrCode::VALUE)); |
| @@ -138,7 +184,11 @@ void copy_h2d_baseformat_opapi(at::Tensor& dst, const at::Tensor& src, bool non_ | |||
| 138 | // if necessary, copy back into dst | 184 | // if necessary, copy back into dst |
| 139 | if (!dst_contig.is_same(dst)) { | 185 | if (!dst_contig.is_same(dst)) { |
| 140 | TORCH_INTERNAL_ASSERT(dst_contig.device() == dst.device(), OPS_ERROR(ErrCode::VALUE)); | 186 | TORCH_INTERNAL_ASSERT(dst_contig.device() == dst.device(), OPS_ERROR(ErrCode::VALUE)); |
| 141 | - copy_d2d_baseformat_opapi(dst, dst_contig, non_blocking); | 187 | + if (dst_contig.scalar_type() == dst.scalar_type()) { |
| 188 | + copy_d2d_baseformat_opapi(dst, dst_contig, non_blocking); | ||
H | |||
| 189 | + } else { | ||
| 190 | + cast_dtype_out_baseformat_opapi(dst, dst_contig); | ||
| 191 | + } | ||
| 142 | } | 192 | } |
| 143 | } | 193 | } |
| 144 | 194 | ||
| @@ -153,8 +203,17 @@ void copy_d2h_baseformat_opapi(at::Tensor& dst, const at::Tensor& src, bool non_ | |||
| 153 | copy_d2h_baseformat_dtype_contigous_opapi(dst, src, non_blocking); | 203 | copy_d2h_baseformat_dtype_contigous_opapi(dst, src, non_blocking); |
| 154 | return; | 204 | return; |
| 155 | } | 205 | } |
| 156 | - at::Tensor dst_contig = (dst_is_contiguous && same_type) ? dst : at::empty_like(dst, src.dtype(), LEGACY_CONTIGUOUS_MEMORY_FORMAT); | 206 | + at::Tensor dst_contig; |
| 157 | - at::Tensor src_contig = src.expand_as(dst).contiguous(); | 207 | + at::Tensor src_contig; |
| 208 | + if (!same_type && non_blocking && !should_fallback_to_cpu_cast(dst, src)) { | ||
| 209 | + // cast on device before the D2H leg. | ||
| 210 | + dst_contig = dst_is_contiguous ? dst : at::empty_like(dst, LEGACY_CONTIGUOUS_MEMORY_FORMAT); | ||
| 211 | + at::Tensor src_cast_input = NpuUtils::check_match(&src) ? src : NpuUtils::format_contiguous(src); | ||
| 212 | + src_contig = custom_ops::_npu_dtype_cast(src_cast_input, dst.scalar_type()).expand_as(dst).contiguous(); | ||
| 213 | + } else { | ||
| 214 | + dst_contig = (dst_is_contiguous && same_type) ? dst : at::empty_like(dst, src.dtype(), LEGACY_CONTIGUOUS_MEMORY_FORMAT); | ||
| 215 | + src_contig = src.expand_as(dst).contiguous(); | ||
| 216 | + } | ||
| 158 | // perform a same-dtype copy on contiguous tensors | 217 | // perform a same-dtype copy on contiguous tensors |
| 159 | TORCH_INTERNAL_ASSERT(dst_contig.sizes().equals(src_contig.sizes()), OPS_ERROR(ErrCode::VALUE)); | 218 | TORCH_INTERNAL_ASSERT(dst_contig.sizes().equals(src_contig.sizes()), OPS_ERROR(ErrCode::VALUE)); |
| 160 | TORCH_INTERNAL_ASSERT(dst_contig.scalar_type() == src_contig.scalar_type(), OPS_ERROR(ErrCode::VALUE)); | 219 | TORCH_INTERNAL_ASSERT(dst_contig.scalar_type() == src_contig.scalar_type(), OPS_ERROR(ErrCode::VALUE)); |
这个算子是不支持跨dtype吗