"""
测试目的:验证 torch.Tensor.copy_ 接口功能正确性
API 名称:torch.Tensor.copy_
API 签名:copy_(src, non_blocking=False) -> Tensor
覆盖维度表:
| 覆盖维度 | 说明 | 覆盖情况 |
|------------------|--------------------------------------------------------------|------------------------------------------------|
| 空/非空 | size-0 张量互拷 | 已覆盖 |
| 枚举选项 | non_blocking 为 False / True | 已覆盖 |
| 参数类型 | src 为 Tensor(含标量张量)、与 self dtype 可不同 | 已覆盖 |
| 传参与不传参 | non_blocking 省略与显式传入 | 已覆盖 |
| 等价类/边界值 | 同形、可广播、非连续目标、跨 CPU/NPU | 已覆盖 |
| 精度/数值正确性 | mixed-dtype host-device 路径下,同步/异步 copy 结果一致 | 已覆盖 |
| 正常传参场景 | NPU 上 copy 后 self 的 shape/dtype 不变;返回 self | 已覆盖 |
| 异常传参场景 | 不可广播的 shape | 已覆盖 |
未覆盖项及原因:
- 无
注意:本测试除了验证功能正确性(调用不报错、tensor 结构属性符合预期),
也对 mixed-dtype host-device 路径补充了同步/异步 copy 结果一致性校验。
"""
import torch
import torch_npu
from torch_npu.testing.common_utils import SupportedDevices
try:
from torch_npu.testing.testcase import TestCase, run_tests
except ImportError:
import sys
import unittest
from unittest import TestCase
def run_tests():
unittest.main(argv=sys.argv)
class TestTensorCopy_(TestCase):
"""Functional tests for torch.Tensor.copy_ on NPU."""
def setUp(self):
super().setUp()
self.device_name = torch._C._get_privateuse1_backend_name()
self.assertEqual(
self.device_name,
"npu",
f"Expected device 'npu', got '{self.device_name}'",
)
self.device = torch.device(self.device_name)
self.dtype_cast_pairs = [
(torch.int32, torch.float32),
(torch.int64, torch.float32),
(torch.float16, torch.float32),
(torch.float32, torch.float16),
]
self.aclnn_cast_fallback_dtypes = [
torch.float8_e5m2,
torch.float8_e4m3fn,
]
self.precision_compare_cases = [
(torch.int8, torch.float16, [-127, -31, -1, 0, 1, 7, 42, 127]),
(torch.int16, torch.float32, [-32768, -1025, -1, 0, 1, 255, 1024, 32767]),
(torch.int64, torch.float32, [-(2 ** 20), -12345, -1, 0, 1, 12345, 4096, 2 ** 20]),
(torch.float16, torch.float32, [-2048.0, -7.5, -0.125, 0.0, 0.125, 1.5, 33.25, 2048.0]),
(torch.bfloat16, torch.float32, [-1.0e4, -7.5, -0.125, 0.0, 0.125, 1.5, 33.25, 1.0e4]),
(torch.float32, torch.float16, [-65504.0, -255.5, -0.33325195, 0.0, 0.33325195, 17.625, 255.5, 65504.0]),
(torch.float32, torch.bfloat16, [-1.0e8, -255.5, -0.33325195, 0.0, 0.33325195, 17.625, 255.5, 1.0e8]),
]
def _make_host_source(self, dtype, pin_memory=False):
src = torch.tensor(
[[-7.5, -3.25, -1.0, 0.0], [1.5, 2.25, 5.0, 9.75]],
dtype=torch.float32,
).to(dtype)
return src.pin_memory() if pin_memory else src
def _make_device_source(self, dtype):
return self._make_host_source(dtype).to(self.device)
def _assert_dtype_cast_copy_keeps_async(self, dst, src):
gate_stream = torch_npu.npu.Stream(device=self.device)
copy_stream = torch_npu.npu.Stream(device=self.device)
gate_event = torch_npu.npu.Event()
done_event = torch_npu.npu.Event()
torch_npu.npu.synchronize()
gate_a = torch.ones((4096, 4096), device=self.device, dtype=torch.float32)
gate_b = torch.ones((4096, 4096), device=self.device, dtype=torch.float32)
with torch_npu.npu.stream(gate_stream):
gate_c = gate_a @ gate_b
gate_c = gate_c @ gate_b
gate_event.record()
with torch_npu.npu.stream(copy_stream):
copy_stream.wait_event(gate_event)
ret = dst.copy_(src, non_blocking=True)
done_event.record()
self.assertIs(ret, dst)
self.assertFalse(done_event.query())
done_event.synchronize()
def _assert_copy_matches_cast(self, dst, src):
expected = src.cpu().to(dtype=dst.dtype)
actual = dst.cpu() if dst.device.type == self.device_name else dst
self.assertEqual(actual, expected)
def _to_cpu_if_needed(self, tensor):
return tensor.cpu() if tensor.device.type == self.device_name else tensor
def _assert_non_blocking_matches_blocking(self, async_dst, sync_dst, src, async_base=None, sync_base=None):
sync_ret = sync_dst.copy_(src, non_blocking=False)
async_ret = async_dst.copy_(src, non_blocking=True)
self.assertIs(sync_ret, sync_dst)
self.assertIs(async_ret, async_dst)
torch_npu.npu.synchronize()
self.assertEqual(self._to_cpu_if_needed(async_dst), self._to_cpu_if_needed(sync_dst))
if async_base is not None and sync_base is not None:
self.assertEqual(self._to_cpu_if_needed(async_base), self._to_cpu_if_needed(sync_base))
def test_copy_npu_same_device_same_shape(self):
dst = torch.empty(3, 4, device=self.device, dtype=torch.float32)
src = torch.randn(3, 4, device=self.device, dtype=torch.float32)
before_shape = dst.shape
before_dtype = dst.dtype
ret = dst.copy_(src)
self.assertIs(ret, dst)
self.assertEqual(dst.shape, before_shape)
self.assertEqual(dst.dtype, before_dtype)
self.assertEqual(dst.device.type, self.device_name)
def test_copy_npu_broadcast_src(self):
dst = torch.empty(4, 3, device=self.device)
src = torch.randn(1, 3, device=self.device)
dst.copy_(src)
self.assertEqual(dst.shape, torch.Size([4, 3]))
def test_copy_npu_from_cpu_src(self):
dst = torch.empty(2, 5, device=self.device)
src = torch.randn(2, 5)
dst.copy_(src)
self.assertEqual(dst.device.type, self.device_name)
self.assertEqual(dst.shape, torch.Size([2, 5]))
def test_copy_npu_non_blocking_false(self):
dst = torch.empty(2, 2, device=self.device)
src = torch.ones(2, 2, device=self.device)
ret = dst.copy_(src, non_blocking=False)
self.assertIs(ret, dst)
def test_copy_npu_non_blocking_true(self):
dst = torch.empty(2, 2, device=self.device)
src = torch.ones(2, 2, device=self.device)
ret = dst.copy_(src, non_blocking=True)
self.assertIs(ret, dst)
self.assertEqual(dst.shape, torch.Size([2, 2]))
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_npu_from_pinned_cpu_src_dtype_cast_non_blocking(self):
for src_dtype, dst_dtype in self.dtype_cast_pairs:
dst = torch.empty(2, 4, dtype=dst_dtype, device=self.device)
src = self._make_host_source(src_dtype, pin_memory=True)
self._assert_dtype_cast_copy_keeps_async(dst, src)
self._assert_copy_matches_cast(dst, src)
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_pinned_cpu_from_npu_src_dtype_cast_non_blocking(self):
for src_dtype, dst_dtype in self.dtype_cast_pairs:
dst = torch.empty(2, 4, dtype=dst_dtype, pin_memory=True)
src = self._make_device_source(src_dtype)
self._assert_dtype_cast_copy_keeps_async(dst, src)
self.assertTrue(dst.is_pinned())
self._assert_copy_matches_cast(dst, src)
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_pinned_cpu_from_non_contiguous_npu_src_dtype_cast_non_blocking(self):
src = torch.arange(8, dtype=torch.int32, device=self.device).reshape(4, 2).t()
dst = torch.empty(2, 4, dtype=torch.float32, pin_memory=True)
self.assertFalse(src.is_contiguous())
self._assert_dtype_cast_copy_keeps_async(dst, src)
self.assertTrue(dst.is_pinned())
self._assert_copy_matches_cast(dst, src)
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_pinned_cpu_non_contiguous_dst_dtype_cast_preserves_strided_layout(self):
base = torch.full((3, 6), -99.0, dtype=torch.float32).pin_memory()
dst = base[:, 1::2]
src = torch.arange(9, dtype=torch.int32, device=self.device).reshape(3, 3)
expected_base = torch.full((3, 6), -99.0, dtype=torch.float32)
expected_base[:, 1::2] = src.cpu().to(dtype=dst.dtype)
self.assertFalse(dst.is_contiguous())
self.assertTrue(dst.is_pinned())
self.assertNotEqual(dst.storage_offset(), 0)
ret = dst.copy_(src, non_blocking=True)
torch_npu.npu.synchronize()
self.assertIs(ret, dst)
self.assertEqual(dst, src.cpu().to(dtype=dst.dtype))
self.assertEqual(base, expected_base)
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_npu_non_contiguous_dst_dtype_cast_non_blocking(self):
base = torch.empty(4, 2, dtype=torch.float32, device=self.device)
dst = base.t()
src = self._make_host_source(torch.int32, pin_memory=True)
self.assertFalse(dst.is_contiguous())
self._assert_dtype_cast_copy_keeps_async(dst, src)
self._assert_copy_matches_cast(dst, src)
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_npu_non_contiguous_dst_dtype_cast_preserves_strided_layout(self):
base = torch.full((3, 6), -99.0, dtype=torch.float32, device=self.device)
dst = base[:, 1::2]
src = torch.arange(9, dtype=torch.int32).reshape(3, 3).pin_memory()
expected_base = torch.full((3, 6), -99.0, dtype=torch.float32)
expected_base[:, 1::2] = src.to(dtype=dst.dtype)
self.assertFalse(dst.is_contiguous())
self.assertNotEqual(dst.storage_offset(), 0)
ret = dst.copy_(src, non_blocking=True)
torch_npu.npu.synchronize()
self.assertIs(ret, dst)
self.assertEqual(dst.cpu(), src.to(dtype=dst.dtype))
self.assertEqual(base.cpu(), expected_base)
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_npu_from_pinned_cpu_broadcast_src_dtype_cast_non_blocking(self):
dst = torch.empty(3, 4, dtype=torch.float32, device=self.device)
src = torch.arange(4, dtype=torch.int32).reshape(1, 4).pin_memory()
expected = src.to(dtype=dst.dtype).expand(3, 4)
ret = dst.copy_(src, non_blocking=True)
torch_npu.npu.synchronize()
self.assertIs(ret, dst)
self.assertEqual(dst.cpu(), expected)
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_pinned_cpu_from_npu_broadcast_src_dtype_cast_non_blocking(self):
dst = torch.empty(3, 4, dtype=torch.float32, pin_memory=True)
src = torch.arange(4, dtype=torch.int32, device=self.device).reshape(1, 4)
expected = src.cpu().to(dtype=dst.dtype).expand_as(dst)
ret = dst.copy_(src, non_blocking=True)
torch_npu.npu.synchronize()
self.assertIs(ret, dst)
self.assertEqual(dst, expected)
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_npu_from_pinned_cpu_aclnn_cast_unsupported_src_dtype_fallback(self):
for src_dtype in self.aclnn_cast_fallback_dtypes:
dst = torch.empty(2, 4, dtype=torch.float32, device=self.device)
src = torch.arange(8, dtype=torch.float32).reshape(2, 4).to(src_dtype).pin_memory()
ret = dst.copy_(src, non_blocking=True)
torch_npu.npu.synchronize()
self.assertIs(ret, dst)
self.assertEqual(dst.cpu(), src.to(dtype=dst.dtype))
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_pinned_cpu_from_npu_aclnn_cast_unsupported_dst_dtype_fallback(self):
for dst_dtype in self.aclnn_cast_fallback_dtypes:
dst = torch.empty(2, 4, dtype=dst_dtype).pin_memory()
src = torch.arange(8, dtype=torch.float32, device=self.device).reshape(2, 4)
ret = dst.copy_(src, non_blocking=True)
torch_npu.npu.synchronize()
self.assertIs(ret, dst)
self.assertEqual(dst.dtype, dst_dtype)
self.assertEqual(dst.to(dtype=src.dtype), src.cpu())
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_npu_from_pinned_cpu_slice_dtype_cast_non_blocking(self):
src_base = torch.arange(9, dtype=torch.int32).pin_memory()
src = src_base[1:].reshape(2, 4)
dst = torch.empty(2, 4, dtype=torch.float32, device=self.device)
self.assertTrue(src.is_pinned())
self.assertNotEqual(src.data_ptr(), src.untyped_storage().data_ptr())
self._assert_dtype_cast_copy_keeps_async(dst, src)
self._assert_copy_matches_cast(dst, src)
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_npu_from_pinned_cpu_dtype_cast_non_blocking_matches_blocking(self):
for src_dtype, dst_dtype, values in self.precision_compare_cases:
with self.subTest(direction="h2d", src_dtype=src_dtype, dst_dtype=dst_dtype):
src = torch.tensor(values, dtype=src_dtype).reshape(2, 4).pin_memory()
async_dst = torch.empty(2, 4, dtype=dst_dtype, device=self.device)
sync_dst = torch.empty(2, 4, dtype=dst_dtype, device=self.device)
self._assert_non_blocking_matches_blocking(async_dst, sync_dst, src)
self._assert_copy_matches_cast(async_dst, src)
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_pinned_cpu_from_npu_dtype_cast_non_blocking_matches_blocking(self):
for src_dtype, dst_dtype, values in self.precision_compare_cases:
with self.subTest(direction="d2h", src_dtype=src_dtype, dst_dtype=dst_dtype):
src = torch.tensor(values, dtype=src_dtype).reshape(2, 4).to(self.device)
async_dst = torch.empty(2, 4, dtype=dst_dtype, pin_memory=True)
sync_dst = torch.empty(2, 4, dtype=dst_dtype, pin_memory=True)
self._assert_non_blocking_matches_blocking(async_dst, sync_dst, src)
self._assert_copy_matches_cast(async_dst, src)
self.assertTrue(async_dst.is_pinned())
self.assertTrue(sync_dst.is_pinned())
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_mixed_dtype_non_blocking_matches_blocking_for_layout_variants(self):
layout_cases = [
{
"name": "h2d_non_contiguous_dst",
"src_dtype": torch.float32,
"dst_dtype": torch.float16,
"make_src": lambda: torch.tensor(
[-63.5, -7.25, -0.5, 0.0, 0.5, 3.25, 17.75, 63.5, -19.5],
dtype=torch.float32,
).reshape(3, 3).pin_memory(),
"make_async_dst": lambda: torch.full((3, 6), -99.0, dtype=torch.float16, device=self.device),
"make_sync_dst": lambda: torch.full((3, 6), -99.0, dtype=torch.float16, device=self.device),
"select_dst_view": lambda base: base[:, 1::2],
"expected": lambda src, dst: src.to(dtype=dst.dtype),
},
{
"name": "d2h_non_contiguous_dst",
"src_dtype": torch.int32,
"dst_dtype": torch.float32,
"make_src": lambda: torch.arange(9, dtype=torch.int32, device=self.device).reshape(3, 3) - 4,
"make_async_dst": lambda: torch.full((3, 6), -99.0, dtype=torch.float32).pin_memory(),
"make_sync_dst": lambda: torch.full((3, 6), -99.0, dtype=torch.float32).pin_memory(),
"select_dst_view": lambda base: base[:, 1::2],
"expected": lambda src, dst: src.cpu().to(dtype=dst.dtype),
},
{
"name": "h2d_broadcast_src",
"src_dtype": torch.int16,
"dst_dtype": torch.float32,
"make_src": lambda: torch.tensor([-32768, -17, 9, 32767], dtype=torch.int16).reshape(1, 4).pin_memory(),
"make_async_dst": lambda: torch.empty(3, 4, dtype=torch.float32, device=self.device),
"make_sync_dst": lambda: torch.empty(3, 4, dtype=torch.float32, device=self.device),
"select_dst_view": lambda base: base,
"expected": lambda src, dst: src.to(dtype=dst.dtype).expand_as(dst),
},
{
"name": "d2h_non_contiguous_src",
"src_dtype": torch.float16,
"dst_dtype": torch.float32,
"make_src": lambda: torch.tensor(
[-7.5, -1.25, 0.0, 1.25, 3.5, 7.75, 15.5, 31.0],
dtype=torch.float16,
device=self.device,
).reshape(4, 2).t(),
"make_async_dst": lambda: torch.empty(2, 4, dtype=torch.float32, pin_memory=True),
"make_sync_dst": lambda: torch.empty(2, 4, dtype=torch.float32, pin_memory=True),
"select_dst_view": lambda base: base,
"expected": lambda src, dst: src.cpu().to(dtype=dst.dtype),
},
]
for case in layout_cases:
with self.subTest(case=case["name"], src_dtype=case["src_dtype"], dst_dtype=case["dst_dtype"]):
src = case["make_src"]()
async_base = case["make_async_dst"]()
sync_base = case["make_sync_dst"]()
async_dst = case["select_dst_view"](async_base)
sync_dst = case["select_dst_view"](sync_base)
self._assert_non_blocking_matches_blocking(
async_dst,
sync_dst,
src,
async_base=async_base,
sync_base=sync_base,
)
self.assertEqual(
self._to_cpu_if_needed(async_dst),
case["expected"](src, async_dst),
)
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_dtype_cast_non_blocking_temporary_lifetime(self):
h2d_dst = torch.empty(2, 4, dtype=torch.float32, device=self.device)
d2h_dst = torch.empty(2, 4, dtype=torch.float32, pin_memory=True)
expected = None
for i in range(32):
host_src = self._make_host_source(torch.int32) + i
h2d_src = host_src.pin_memory()
h2d_dst.copy_(h2d_src, non_blocking=True)
d2h_src = host_src.to(self.device)
d2h_dst.copy_(d2h_src, non_blocking=True)
expected = host_src.to(dtype=d2h_dst.dtype)
torch_npu.npu.synchronize()
self.assertEqual(h2d_dst.cpu(), h2d_src.to(dtype=h2d_dst.dtype))
self.assertEqual(d2h_dst, expected)
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_npu_from_cpu_src_dtype_cast_blocking(self):
for src_dtype, dst_dtype in self.dtype_cast_pairs:
dst = torch.empty(2, 4, dtype=dst_dtype, device=self.device)
src = self._make_host_source(src_dtype)
ret = dst.copy_(src, non_blocking=False)
self.assertIs(ret, dst)
self._assert_copy_matches_cast(dst, src)
@SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
def test_copy_cpu_from_npu_src_dtype_cast_blocking(self):
for src_dtype, dst_dtype in self.dtype_cast_pairs:
dst = torch.empty(2, 4, dtype=dst_dtype)
src = self._make_device_source(src_dtype)
ret = dst.copy_(src, non_blocking=False)
self.assertIs(ret, dst)
self._assert_copy_matches_cast(dst, src)
def test_copy_npu_src_int_dtype_cast(self):
dst = torch.empty(2, 2, dtype=torch.float32, device=self.device)
src = torch.ones(2, 2, dtype=torch.int32, device=self.device)
dst.copy_(src)
self.assertEqual(dst.dtype, torch.float32)
def test_copy_npu_non_contiguous_dst(self):
base = torch.empty(6, 4, device=self.device)
dst = base.t()
self.assertFalse(dst.is_contiguous())
src = torch.randn(4, 6, device=self.device)
dst.copy_(src)
self.assertEqual(dst.shape, torch.Size([4, 6]))
def test_copy_npu_empty_tensor(self):
dst = torch.empty(0, 3, device=self.device)
src = torch.empty(0, 3, device=self.device)
dst.copy_(src)
self.assertEqual(dst.shape, torch.Size([0, 3]))
def test_copy_npu_float16(self):
dst = torch.empty(2, 3, dtype=torch.float16, device=self.device)
src = torch.randn(2, 3, dtype=torch.float16, device=self.device)
dst.copy_(src)
self.assertEqual(dst.dtype, torch.float16)
def test_copy_npu_bfloat16(self):
dst = torch.empty(2, 3, dtype=torch.bfloat16, device=self.device)
src = torch.randn(2, 3, dtype=torch.bfloat16, device=self.device)
dst.copy_(src)
self.assertEqual(dst.dtype, torch.bfloat16)
def test_copy_npu_incompatible_shape_raises(self):
dst = torch.empty(3, 4, device=self.device)
src = torch.randn(2, 3, device=self.device)
with self.assertRaises(RuntimeError):
out = dst.copy_(src)
out.cpu()
def test_copy_cpu_baseline(self):
dst = torch.empty(3, 4)
src = torch.randn(3, 4)
ret = dst.copy_(src)
self.assertIs(ret, dst)
self.assertEqual(dst.shape, torch.Size([3, 4]))
def test_copy_cpu_baseline_broadcast(self):
dst = torch.empty(2, 4)
src = torch.randn(1, 4)
dst.copy_(src)
self.assertEqual(dst.shape, torch.Size([2, 4]))
if __name__ == "__main__":
run_tests()