已合并
[feat]TensorTo support preserve_format consistent with GPU #35479
culechan创建于 5月13日
[feat]TensorTo support preserve_format consistent with GPU #35479
已合并
共 4 个文件变更+1817-913
| @@ -1,815 +1,1089 @@ | |||
| 1 | -from itertools import product | 1 | +import collections |
| 2 | -import collections | 2 | +import gc |
| 3 | -import gc | 3 | +from itertools import product |
| 4 | -import numpy as np | 4 | + |
| 5 | - | 5 | +import numpy as np |
| 6 | -import torch | 6 | +import torch |
| 7 | -from torch.autograd import Variable | 7 | +from torch.autograd import Variable |
| 8 | -import torch_npu | 8 | + |
| 9 | - | 9 | +import torch_npu |
| 10 | -from torch_npu.testing.testcase import TestCase, run_tests | 10 | +from torch_npu.testing.common_utils import freeze_rng_state |
| 11 | -from torch_npu.testing.common_utils import freeze_rng_state | 11 | +from torch_npu.testing.testcase import run_tests, TestCase |
| 12 | - | 12 | + |
| 13 | - | 13 | + |
| 14 | -class TestNpu(TestCase): | 14 | +class TestNpu(TestCase): |
| 15 | - | 15 | + FIFTY_MIL_CYCLES = 50000000 |
| 16 | - FIFTY_MIL_CYCLES = 50000000 | 16 | + |
| 17 | - | 17 | + def _check_memory_stat_consistency(self): |
| 18 | - def _check_memory_stat_consistency(self): | 18 | + snapshot = torch_npu.npu.memory_snapshot() |
| 19 | - snapshot = torch_npu.npu.memory_snapshot() | 19 | + |
| 20 | - | 20 | + expected_each_device = collections.defaultdict( |
| 21 | - expected_each_device = collections.defaultdict(lambda: collections.defaultdict(int)) | 21 | + lambda: collections.defaultdict(int) |
| 22 | - | 22 | + ) |
| 23 | - for segment in snapshot: | 23 | + |
| 24 | - expected = expected_each_device[segment["device"]] | 24 | + for segment in snapshot: |
| 25 | - pool_str = segment["segment_type"] + "_pool" | 25 | + expected = expected_each_device[segment["device"]] |
| 26 | - | 26 | + pool_str = segment["segment_type"] + "_pool" |
| 27 | - expected["segment.all.current"] += 1 | 27 | + |
| 28 | - expected["segment." + pool_str + ".current"] += 1 | 28 | + expected["segment.all.current"] += 1 |
| 29 | - | 29 | + expected["segment." + pool_str + ".current"] += 1 |
| 30 | - expected["allocated_bytes.all.current"] += segment["allocated_size"] | 30 | + |
| 31 | - expected["allocated_bytes." + pool_str + ".current"] += segment["allocated_size"] | 31 | + expected["allocated_bytes.all.current"] += segment["allocated_size"] |
| 32 | - | 32 | + expected["allocated_bytes." + pool_str + ".current"] += segment[ |
| 33 | - expected["reserved_bytes.all.current"] += segment["total_size"] | 33 | + "allocated_size" |
| 34 | - expected["reserved_bytes." + pool_str + ".current"] += segment["total_size"] | 34 | + ] |
| 35 | - | 35 | + |
| 36 | - expected["active_bytes.all.current"] += segment["active_size"] | 36 | + expected["reserved_bytes.all.current"] += segment["total_size"] |
| 37 | - expected["active_bytes." + pool_str + ".current"] += segment["active_size"] | 37 | + expected["reserved_bytes." + pool_str + ".current"] += segment["total_size"] |
| 38 | - | 38 | + |
| 39 | - is_split = len(segment["blocks"]) > 1 | 39 | + expected["active_bytes.all.current"] += segment["active_size"] |
| 40 | - for block in segment["blocks"]: | 40 | + expected["active_bytes." + pool_str + ".current"] += segment["active_size"] |
| 41 | - if block["state"] == "active_allocated": | 41 | + |
| 42 | - expected["allocation.all.current"] += 1 | 42 | + is_split = len(segment["blocks"]) > 1 |
| 43 | - expected["allocation." + pool_str + ".current"] += 1 | 43 | + for block in segment["blocks"]: |
| 44 | - | 44 | + if block["state"] == "active_allocated": |
| 45 | - if block["state"].startswith("active_"): | 45 | + expected["allocation.all.current"] += 1 |
| 46 | - expected["active.all.current"] += 1 | 46 | + expected["allocation." + pool_str + ".current"] += 1 |
| 47 | - expected["active." + pool_str + ".current"] += 1 | 47 | + |
| 48 | - | 48 | + if block["state"].startswith("active_"): |
| 49 | - if block["state"] == "inactive" and is_split: | 49 | + expected["active.all.current"] += 1 |
| 50 | - expected["inactive_split.all.current"] += 1 | 50 | + expected["active." + pool_str + ".current"] += 1 |
| 51 | - expected["inactive_split." + pool_str + ".current"] += 1 | 51 | + |
| 52 | - expected["inactive_split_bytes.all.current"] += block["size"] | 52 | + if block["state"] == "inactive" and is_split: |
| 53 | - expected["inactive_split_bytes." + pool_str + ".current"] += block["size"] | 53 | + expected["inactive_split.all.current"] += 1 |
| 54 | - | 54 | + expected["inactive_split." + pool_str + ".current"] += 1 |
| 55 | - for device, expected in expected_each_device.items(): | 55 | + expected["inactive_split_bytes.all.current"] += block["size"] |
| 56 | - stats = torch_npu.npu.memory_stats(device) | 56 | + expected["inactive_split_bytes." + pool_str + ".current"] += block[ |
| 57 | - for k, v in expected.items(): | 57 | + "size" |
| 58 | - self.assertEqual(v, stats[k]) | 58 | + ] |
| 59 | - | 59 | + |
| 60 | - def test_memory_summary_format(self, device=None): | 60 | + for device, expected in expected_each_device.items(): |
| 61 | - summary = torch_npu.npu.memory_summary() | 61 | + stats = torch_npu.npu.memory_stats(device) |
| 62 | - | 62 | + for k, v in expected.items(): |
| 63 | - device = torch_npu.npu._get_device_index(device, optional=True) | 63 | + self.assertEqual(v, stats[k]) |
| 64 | - stats = torch_npu.npu.memory_stats(device=device) | 64 | + |
| 65 | - fmt_dict = {"_": "", "device": device} | 65 | + def test_memory_summary_format(self, device=None): |
| 66 | - for k, v in stats.items(): | 66 | + summary = torch_npu.npu.memory_summary() |
| 67 | - fmt_dict[k.replace(".", "-")] = v | 67 | + |
| 68 | - | 68 | + device = torch_npu.npu._get_device_index(device, optional=True) |
| 69 | - expected_head = [] | 69 | + stats = torch_npu.npu.memory_stats(device=device) |
| 70 | - expected_head.append("=" * 75) | 70 | + fmt_dict = {"_": "", "device": device} |
| 71 | - expected_head.append(" {_:16} PyTorch NPU memory summary, device ID {device:<18d} ") | 71 | + for k, v in stats.items(): |
| 72 | - expected_head.append("-" * 75) | 72 | + fmt_dict[k.replace(".", "-")] = v |
| 73 | - expected_head.append(" {_:9} NPU OOMs: {num_ooms:<13d} | {_:6} npuMalloc retries: {num_alloc_retries:<9d} ") | 73 | + |
| 74 | - expected_head.append("=" * 75) | 74 | + expected_head = [] |
| 75 | - expected_head.append(" Metric | Cur Usage | Peak Usage | Tot Alloc | Tot Freed ") | 75 | + expected_head.append("=" * 75) |
| 76 | - | 76 | + expected_head.append( |
| 77 | - expected_head_str = "|" + "|\n|".join(expected_head).format(**fmt_dict) + "|\n" | 77 | + " {_:16} PyTorch NPU memory summary, device ID {device:<18d} " |
| 78 | - assert_len = len(expected_head_str) | 78 | + ) |
| 79 | - | 79 | + expected_head.append("-" * 75) |
| 80 | - self.assertEqual(expected_head_str, summary[:assert_len]) | 80 | + expected_head.append( |
| 81 | - | 81 | + " {_:9} NPU OOMs: {num_ooms:<13d} | {_:6} npuMalloc retries: {num_alloc_retries:<9d} " |
| 82 | - @staticmethod | 82 | + ) |
| 83 | - def _test_memory_stats_generator(self, device=None, N=35): | 83 | + expected_head.append("=" * 75) |
| 84 | - if device is None: | 84 | + expected_head.append( |
| 85 | - device = torch_npu.npu.current_device() | 85 | + " Metric | Cur Usage | Peak Usage | Tot Alloc | Tot Freed " |
| 86 | - | 86 | + ) |
| 87 | - m0 = torch_npu.npu.memory_allocated(device) | 87 | + |
| 88 | - last_m_arr = [torch_npu.npu.memory_allocated(device)] | 88 | + expected_head_str = "|" + "|\n|".join(expected_head).format(**fmt_dict) + "|\n" |
| 89 | - max_m_arr = [torch_npu.npu.max_memory_allocated(device)] | 89 | + assert_len = len(expected_head_str) |
| 90 | - last_r_arr = [torch_npu.npu.memory_reserved(device)] | 90 | + |
| 91 | - max_r_arr = [torch_npu.npu.max_memory_reserved(device)] | 91 | + self.assertEqual(expected_head_str, summary[:assert_len]) |
| 92 | - | 92 | + |
| 93 | - def alloc(*size): | 93 | + @staticmethod |
| 94 | - with torch_npu.npu.device(device): | 94 | + def _test_memory_stats_generator(self, device=None, N=35): |
| 95 | - # NOTE: do **not** use methods that can have additional | 95 | + if device is None: |
| 96 | - # memory overhead, e.g., inplace random sampling methods. | 96 | + device = torch_npu.npu.current_device() |
| 97 | - # they can leave some memory occupied even after being | 97 | + |
| 98 | - # deallocated, e.g., initialized RNG state, causing some | 98 | + m0 = torch_npu.npu.memory_allocated(device) |
| 99 | - # memory checks below to fail. | 99 | + last_m_arr = [torch_npu.npu.memory_allocated(device)] |
| 100 | - return torch.npu.FloatTensor(*size) | 100 | + max_m_arr = [torch_npu.npu.max_memory_allocated(device)] |
| 101 | - | 101 | + last_r_arr = [torch_npu.npu.memory_reserved(device)] |
| 102 | - def assert_change(comp=1, empty_cache=False, reset_peak=False): | 102 | + max_r_arr = [torch_npu.npu.max_memory_reserved(device)] |
| 103 | - # comp > 0: increased | 103 | + |
| 104 | - # comp = 0: equal | 104 | + def alloc(*size): |
| 105 | - # comp < 0: decreased | 105 | + with torch_npu.npu.device(device): |
| 106 | - new_m = torch_npu.npu.memory_allocated(device) | 106 | + # NOTE: do **not** use methods that can have additional |
| 107 | - new_max_m = torch_npu.npu.max_memory_allocated(device) | 107 | + # memory overhead, e.g., inplace random sampling methods. |
| 108 | - if comp > 0: | 108 | + # they can leave some memory occupied even after being |
| 109 | - self.assertGreater(new_m, last_m_arr[0]) | 109 | + # deallocated, e.g., initialized RNG state, causing some |
| 110 | - elif comp < 0: | 110 | + # memory checks below to fail. |
| 111 | - self.assertLess(new_m, last_m_arr[0]) | 111 | + return torch.npu.FloatTensor(*size) |
| 112 | - else: | 112 | + |
| 113 | - self.assertEqual(new_m, last_m_arr[0]) | 113 | + def assert_change(comp=1, empty_cache=False, reset_peak=False): |
| 114 | - self.assertLessEqual(new_m, new_max_m) | 114 | + # comp > 0: increased |
| 115 | - self.assertGreaterEqual(new_max_m, max_m_arr[0]) | 115 | + # comp = 0: equal |
| 116 | - last_m_arr[0] = new_m | 116 | + # comp < 0: decreased |
| 117 | - max_m_arr[0] = new_max_m | 117 | + new_m = torch_npu.npu.memory_allocated(device) |
| 118 | - | 118 | + new_max_m = torch_npu.npu.max_memory_allocated(device) |
| 119 | - new_r = torch_npu.npu.memory_reserved(device) | 119 | + if comp > 0: |
| 120 | - new_max_r = torch_npu.npu.max_memory_reserved(device) | 120 | + self.assertGreater(new_m, last_m_arr[0]) |
| 121 | - # emptying cache may happen (due to allocation or empty_cache), so | 121 | + elif comp < 0: |
| 122 | - # we can't assert new_c >= last_c | 122 | + self.assertLess(new_m, last_m_arr[0]) |
| 123 | - self.assertLessEqual(new_r, new_max_r) | 123 | + else: |
| 124 | - self.assertGreaterEqual(new_max_r, max_r_arr[0]) | 124 | + self.assertEqual(new_m, last_m_arr[0]) |
| 125 | - last_r_arr[0] = new_r | 125 | + self.assertLessEqual(new_m, new_max_m) |
| 126 | - max_r_arr[0] = new_max_r | 126 | + self.assertGreaterEqual(new_max_m, max_m_arr[0]) |
| 127 | - | 127 | + last_m_arr[0] = new_m |
| 128 | - if empty_cache: | 128 | + max_m_arr[0] = new_max_m |
| 129 | - torch_npu.npu.empty_cache() | 129 | + |
| 130 | - new_r = torch_npu.npu.memory_reserved(device) | 130 | + new_r = torch_npu.npu.memory_reserved(device) |
| 131 | - new_max_r = torch_npu.npu.max_memory_reserved(device) | 131 | + new_max_r = torch_npu.npu.max_memory_reserved(device) |
| 132 | - self.assertLessEqual(new_r, last_r_arr[0]) | 132 | + # emptying cache may happen (due to allocation or empty_cache), so |
| 133 | - self.assertLessEqual(new_r, new_max_r) | 133 | + # we can't assert new_c >= last_c |
| 134 | - self.assertEqual(new_max_r, max_r_arr[0]) | 134 | + self.assertLessEqual(new_r, new_max_r) |
| 135 | - last_r_arr[0] = new_r | 135 | + self.assertGreaterEqual(new_max_r, max_r_arr[0]) |
| 136 | - | 136 | + last_r_arr[0] = new_r |
| 137 | - if reset_peak: | 137 | + max_r_arr[0] = new_max_r |
| 138 | - torch_npu.npu.reset_peak_memory_stats(device) | 138 | + |
| 139 | - self.assertEqual(torch_npu.npu.memory_allocated(device), last_m_arr[0]) | 139 | + if empty_cache: |
| 140 | - self.assertEqual(torch_npu.npu.max_memory_allocated(device), last_m_arr[0]) | 140 | + torch_npu.npu.empty_cache() |
| 141 | - max_m_arr[0] = last_m_arr[0] | 141 | + new_r = torch_npu.npu.memory_reserved(device) |
| 142 | - self.assertEqual(torch_npu.npu.memory_reserved(device), last_r_arr[0]) | 142 | + new_max_r = torch_npu.npu.max_memory_reserved(device) |
| 143 | - self.assertEqual(torch_npu.npu.max_memory_reserved(device), last_r_arr[0]) | 143 | + self.assertLessEqual(new_r, last_r_arr[0]) |
| 144 | - max_r_arr[0] = last_r_arr[0] | 144 | + self.assertLessEqual(new_r, new_max_r) |
| 145 | - | 145 | + self.assertEqual(new_max_r, max_r_arr[0]) |
| 146 | - assert_change(0) | 146 | + last_r_arr[0] = new_r |
| 147 | - assert_change(0, reset_peak=True) | 147 | + |
| 148 | - assert_change(0, empty_cache=True) | 148 | + if reset_peak: |
| 149 | - assert_change(0, reset_peak=True) | 149 | + torch_npu.npu.reset_peak_memory_stats(device) |
| 150 | - assert_change(0) | 150 | + self.assertEqual(torch_npu.npu.memory_allocated(device), last_m_arr[0]) |
| 151 | - yield | 151 | + self.assertEqual( |
| 152 | - | 152 | + torch_npu.npu.max_memory_allocated(device), last_m_arr[0] |
| 153 | - def assert_change_by_tensor(): | 153 | + ) |
| 154 | - tensors1 = [alloc(1), alloc(10, 20), alloc(200, 300, 2000)] | 154 | + max_m_arr[0] = last_m_arr[0] |
| 155 | - m1 = torch_npu.npu.memory_allocated(device) | 155 | + self.assertEqual(torch_npu.npu.memory_reserved(device), last_r_arr[0]) |
| 156 | - assert_change(1) | 156 | + self.assertEqual( |
| 157 | - yield | 157 | + torch_npu.npu.max_memory_reserved(device), last_r_arr[0] |
| 158 | - | 158 | + ) |
| 159 | - tensors2 = [] | 159 | + max_r_arr[0] = last_r_arr[0] |
| 160 | - | 160 | + |
| 161 | - for i in range(1, int(N / 2) + 1): | 161 | + assert_change(0) |
| 162 | - # small ones | 162 | + assert_change(0, reset_peak=True) |
| 163 | - tensors2.append(alloc(i, i * 4)) | 163 | + assert_change(0, empty_cache=True) |
| 164 | - assert_change(1) | 164 | + assert_change(0, reset_peak=True) |
| 165 | - yield | 165 | + assert_change(0) |
| 166 | - | 166 | + yield |
| 167 | - for i in range(5, int(N / 2) + 5): | 167 | + |
| 168 | - # large ones | 168 | + def assert_change_by_tensor(): |
| 169 | - tensors2.append(alloc(i, i * 7, i * 9, i * 11)) | 169 | + tensors1 = [alloc(1), alloc(10, 20), alloc(200, 300, 2000)] |
| 170 | - assert_change(1, reset_peak=(i % 2 == 0)) | 170 | + m1 = torch_npu.npu.memory_allocated(device) |
| 171 | - yield | 171 | + assert_change(1) |
| 172 | - | 172 | + yield |
| 173 | - tensors2.append(alloc(0, 0, 0)) | 173 | + |
| 174 | - assert_change(0) | 174 | + tensors2 = [] |
| 175 | - yield | 175 | + |
| 176 | - | 176 | + for i in range(1, int(N / 2) + 1): |
| 177 | - permute = [] | 177 | + # small ones |
| 178 | - for i in torch.randperm(len(tensors2)): | 178 | + tensors2.append(alloc(i, i * 4)) |
| 179 | - permute.append(tensors2[i]) | 179 | + assert_change(1) |
| 180 | - assert_change(0) | 180 | + yield |
| 181 | - yield | 181 | + |
| 182 | - | 182 | + for i in range(5, int(N / 2) + 5): |
| 183 | - del tensors2 | 183 | + # large ones |
| 184 | - assert_change(0) | 184 | + tensors2.append(alloc(i, i * 7, i * 9, i * 11)) |
| 185 | - yield | 185 | + assert_change(1, reset_peak=(i % 2 == 0)) |
| 186 | - tensors2 = permute | 186 | + yield |
| 187 | - assert_change(0) | 187 | + |
| 188 | - yield | 188 | + tensors2.append(alloc(0, 0, 0)) |
| 189 | - del permute | 189 | + assert_change(0) |
| 190 | - assert_change(0, reset_peak=True) | 190 | + yield |
| 191 | - yield | 191 | + |
| 192 | - | 192 | + permute = [] |
| 193 | - for i in range(int(N / 2)): | 193 | + for i in torch.randperm(len(tensors2)): |
| 194 | - x = tensors2[i].numel() | 194 | + permute.append(tensors2[i]) |
| 195 | - del tensors2[i] | 195 | + assert_change(0) |
| 196 | - assert_change(-x) # in case that tensors2[i] is empty | 196 | + yield |
| 197 | - yield | 197 | + |
| 198 | - | 198 | + del tensors2 |
| 199 | - for i in range(2, int(2 * N / 3) + 2): | 199 | + assert_change(0) |
| 200 | - tensors2.append(alloc(i, i * 3, i * 8)) | 200 | + yield |
| 201 | - assert_change(1) | 201 | + tensors2 = permute |
| 202 | - yield | 202 | + assert_change(0) |
| 203 | - | 203 | + yield |
| 204 | - del tensors2 | 204 | + del permute |
| 205 | - assert_change(-1, reset_peak=True) | 205 | + assert_change(0, reset_peak=True) |
| 206 | - assert_change(0) | 206 | + yield |
| 207 | - self.assertEqual(torch_npu.npu.memory_allocated(device), m1) | 207 | + |
| 208 | - yield True | 208 | + for i in range(int(N / 2)): |
| 209 | - | 209 | + x = tensors2[i].numel() |
| 210 | - del tensors1 | 210 | + del tensors2[i] |
| 211 | - assert_change(-1, reset_peak=True) | 211 | + assert_change(-x) # in case that tensors2[i] is empty |
| 212 | - self.assertEqual(torch_npu.npu.memory_allocated(device), m0) | 212 | + yield |
| 213 | - | 213 | + |
| 214 | - assert_change_by_tensor() | 214 | + for i in range(2, int(2 * N / 3) + 2): |
| 215 | - | 215 | + tensors2.append(alloc(i, i * 3, i * 8)) |
| 216 | - # test empty_cache and reset_peak | 216 | + assert_change(1) |
| 217 | - assert_change(0, empty_cache=True) | 217 | + yield |
| 218 | - assert_change(0, reset_peak=True) | 218 | + |
| 219 | - | 219 | + del tensors2 |
| 220 | - def test_memory_stats(self): | 220 | + assert_change(-1, reset_peak=True) |
| 221 | - gc.collect() | 221 | + assert_change(0) |
| 222 | - torch_npu.npu.empty_cache() | 222 | + self.assertEqual(torch_npu.npu.memory_allocated(device), m1) |
| 223 | - for _ in self._test_memory_stats_generator(self): | 223 | + yield True |
| 224 | - self._check_memory_stat_consistency() | 224 | + |
| 225 | - | 225 | + del tensors1 |
| 226 | - def test_memory_allocation(self): | 226 | + assert_change(-1, reset_peak=True) |
| 227 | - gc.collect() | 227 | + self.assertEqual(torch_npu.npu.memory_allocated(device), m0) |
| 228 | - torch_npu.npu.empty_cache() | 228 | + |
| 229 | - mem = None | 229 | + assert_change_by_tensor() |
| 230 | - size = 1 | 230 | + |
| 231 | - prev = 0 | 231 | + # test empty_cache and reset_peak |
| 232 | - try: | 232 | + assert_change(0, empty_cache=True) |
| 233 | - prev = torch_npu.npu.memory_allocated() | 233 | + assert_change(0, reset_peak=True) |
| 234 | - mem = torch_npu.npu.caching_allocator_alloc(size) | 234 | + |
| 235 | - self.assertGreater(torch_npu.npu.memory_allocated(), prev) | 235 | + def test_memory_stats(self): |
| 236 | - finally: | 236 | + gc.collect() |
| 237 | - if mem is not None: | 237 | + torch_npu.npu.empty_cache() |
| 238 | - torch_npu.npu.caching_allocator_delete(mem) | 238 | + for _ in self._test_memory_stats_generator(self): |
| 239 | - self.assertEqual(torch_npu.npu.memory_allocated(), prev) | 239 | + self._check_memory_stat_consistency() |
| 240 | - | 240 | + |
| 241 | - def test_out_of_memory(self): | 241 | + def test_memory_allocation(self): |
| 242 | - tensor = torch.zeros(1024, device='npu') | 242 | + gc.collect() |
| 243 | - | 243 | + torch_npu.npu.empty_cache() |
| 244 | - with self.assertRaisesRegex(RuntimeError, "Tried to allocate more than 1EB memory"): | 244 | + mem = None |
| 245 | - torch.empty(1024 * 1024 * 1024 * 8000000000, dtype=torch.int8, device='npu') | 245 | + size = 1 |
| 246 | - | 246 | + prev = 0 |
| 247 | - # ensure out of memory error doesn't disturb subsequent kernel | 247 | + try: |
| 248 | - tensor.fill_(1) | 248 | + prev = torch_npu.npu.memory_allocated() |
| 249 | - self.assertTrue((tensor == 1).all()) | 249 | + mem = torch_npu.npu.caching_allocator_alloc(size) |
| 250 | - | 250 | + self.assertGreater(torch_npu.npu.memory_allocated(), prev) |
| 251 | - def test_set_per_process_memory_fraction(self): | 251 | + finally: |
| 252 | - # test invalid fraction value. | 252 | + if mem is not None: |
| 253 | - with self.assertRaisesRegex(TypeError, "Invalid type"): | 253 | + torch_npu.npu.caching_allocator_delete(mem) |
| 254 | - torch_npu.npu.set_per_process_memory_fraction(int(1)) | 254 | + self.assertEqual(torch_npu.npu.memory_allocated(), prev) |
| 255 | - with self.assertRaisesRegex(ValueError, "Invalid fraction value"): | 255 | + |
| 256 | - torch_npu.npu.set_per_process_memory_fraction(-0.1) | 256 | + def test_out_of_memory(self): |
| 257 | - with self.assertRaisesRegex(ValueError, "Invalid fraction value"): | 257 | + tensor = torch.zeros(1024, device="npu") |
| 258 | - torch_npu.npu.set_per_process_memory_fraction(2.0) | 258 | + |
| 259 | - | 259 | + with self.assertRaisesRegex( |
| 260 | - tensor = torch.zeros(1024, device='npu') | 260 | + RuntimeError, "Tried to allocate more than 1EB memory" |
| 261 | - torch_npu.npu.empty_cache() | 261 | + ): |
| 262 | - total_memory = torch_npu.npu.get_device_properties(0).total_memory | 262 | + torch.empty(1024 * 1024 * 1024 * 8000000000, dtype=torch.int8, device="npu") |
| 263 | - torch_npu.npu.set_per_process_memory_fraction(0.5, 0) | 263 | + |
| 264 | - | 264 | + # ensure out of memory error doesn't disturb subsequent kernel |
| 265 | - # test 0.499 allocation is ok. | 265 | + tensor.fill_(1) |
| 266 | - application = int(total_memory * 0.499) - torch_npu.npu.max_memory_reserved() | 266 | + self.assertTrue((tensor == 1).all()) |
| 267 | - tmp_tensor = torch.empty(application, dtype=torch.int8, device='npu') | 267 | + |
| 268 | - del tmp_tensor | 268 | + def test_set_per_process_memory_fraction(self): |
| 269 | - torch_npu.npu.empty_cache() | 269 | + # test invalid fraction value. |
| 270 | - | 270 | + with self.assertRaisesRegex(TypeError, "Invalid type"): |
| 271 | - application = int(total_memory * 0.5) | 271 | + torch_npu.npu.set_per_process_memory_fraction(1) |
| 272 | - # it will get OOM when try to allocate more than half memory. | 272 | + with self.assertRaisesRegex(ValueError, "Invalid fraction value"): |
| 273 | - with self.assertRaisesRegex(RuntimeError, "out of memory"): | 273 | + torch_npu.npu.set_per_process_memory_fraction(-0.1) |
| 274 | - torch.empty(application, dtype=torch.int8, device='npu') | 274 | + with self.assertRaisesRegex(ValueError, "Invalid fraction value"): |
| 275 | - | 275 | + torch_npu.npu.set_per_process_memory_fraction(2.0) |
| 276 | - # ensure out of memory error doesn't disturb subsequent kernel | 276 | + |
| 277 | - tensor.fill_(1) | 277 | + tensor = torch.zeros(1024, device="npu") |
| 278 | - self.assertTrue((tensor == 1).all()) | 278 | + torch_npu.npu.empty_cache() |
| 279 | - | 279 | + total_memory = torch_npu.npu.get_device_properties(0).total_memory |
| 280 | - def _test_copy_sync_current_stream(self, x, y): | 280 | + torch_npu.npu.set_per_process_memory_fraction(0.5, 0) |
| 281 | - x_plus_one = x + 1 | 281 | + |
| 282 | - s0 = torch_npu.npu.Stream(device=x.device) | 282 | + # test 0.499 allocation is ok. |
| 283 | - s1 = torch_npu.npu.Stream(device=y.device) | 283 | + application = int(total_memory * 0.499) - torch_npu.npu.max_memory_reserved() |
| 284 | - s2 = torch_npu.npu.Stream(device=x.device) | 284 | + tmp_tensor = torch.empty(application, dtype=torch.int8, device="npu") |
| 285 | - s3 = torch_npu.npu.Stream(device=y.device) | 285 | + del tmp_tensor |
| 286 | - | 286 | + torch_npu.npu.empty_cache() |
| 287 | - # same dst stream different src streams | 287 | + |
| 288 | - with torch_npu.npu.stream(s0): | 288 | + application = int(total_memory * 0.5) |
| 289 | - with torch_npu.npu.stream(s1): | 289 | + # it will get OOM when try to allocate more than half memory. |
| 290 | - y.copy_(x_plus_one) | 290 | + with self.assertRaisesRegex(RuntimeError, "out of memory"): |
| 291 | - | 291 | + torch.empty(application, dtype=torch.int8, device="npu") |
| 292 | - with torch_npu.npu.stream(s2), torch_npu.npu.stream(s1): | 292 | + |
| 293 | - y.copy_(x) | 293 | + # ensure out of memory error doesn't disturb subsequent kernel |
| 294 | - | 294 | + tensor.fill_(1) |
| 295 | - s1.synchronize() | 295 | + self.assertTrue((tensor == 1).all()) |
| 296 | - # The copy() is synchronized on the current streams of both src and dst. | 296 | + |
| 297 | - # In the above test, the _sleep() op on s0 will not block the copy() on | 297 | + def _test_copy_sync_current_stream(self, x, y): |
| 298 | - # s2, but both copies are synchronized on s1 in the dst device. Hence, | 298 | + x_plus_one = x + 1 |
| 299 | - # x is copied to y after x_plus_one is copied to y. If x and y are on | 299 | + s0 = torch_npu.npu.Stream(device=x.device) |
| 300 | - # the same device, both copy() ops are synchronized on s1. | 300 | + s1 = torch_npu.npu.Stream(device=y.device) |
| 301 | - self.assertEqual(y, x) | 301 | + s2 = torch_npu.npu.Stream(device=x.device) |
| 302 | - | 302 | + s3 = torch_npu.npu.Stream(device=y.device) |
| 303 | - # same src stream different dst streams | 303 | + |
| 304 | - with torch_npu.npu.stream(s1): | 304 | + # same dst stream different src streams |
| 305 | - with torch_npu.npu.stream(s0): | 305 | + with torch_npu.npu.stream(s0): |
| 306 | - y.copy_(x_plus_one) | 306 | + with torch_npu.npu.stream(s1): |
| 307 | - | 307 | + y.copy_(x_plus_one) |
| 308 | - with torch_npu.npu.stream(s3), torch_npu.npu.stream(s0): | 308 | + |
| 309 | - y.copy_(x) | 309 | + with torch_npu.npu.stream(s2), torch_npu.npu.stream(s1): |
| 310 | - | 310 | + y.copy_(x) |
| 311 | - s0.synchronize() | 311 | + |
| 312 | - # Similarly, both copy() ops are synchronized on s0. | 312 | + s1.synchronize() |
| 313 | - self.assertEqual(y, x) | 313 | + # The copy() is synchronized on the current streams of both src and dst. |
| 314 | - | 314 | + # In the above test, the _sleep() op on s0 will not block the copy() on |
| 315 | - def test_copy_non_blocking(self): | 315 | + # s2, but both copies are synchronized on s1 in the dst device. Hence, |
| 316 | - def _test_copy_non_blocking(a, b): | 316 | + # x is copied to y after x_plus_one is copied to y. If x and y are on |
| 317 | - event = torch_npu.npu.Event() | 317 | + # the same device, both copy() ops are synchronized on s1. |
| 318 | - a.copy_(b, non_blocking=True) | 318 | + self.assertEqual(y, x) |
| 319 | - event.record() | 319 | + |
| 320 | - event.synchronize() | 320 | + # same src stream different dst streams |
| 321 | - self.assertEqual(a, b) | 321 | + with torch_npu.npu.stream(s1): |
| 322 | - | 322 | + with torch_npu.npu.stream(s0): |
| 323 | - # 10MB copies | 323 | + y.copy_(x_plus_one) |
| 324 | - x = torch.ones(10000000, dtype=torch.uint8).npu() | 324 | + |
| 325 | - y = torch.zeros(10000000, dtype=torch.uint8).pin_memory() | 325 | + with torch_npu.npu.stream(s3), torch_npu.npu.stream(s0): |
| 326 | - _test_copy_non_blocking(x, y) | 326 | + y.copy_(x) |
| 327 | - | 327 | + |
| 328 | - x = torch.zeros(10000000, dtype=torch.uint8).pin_memory() | 328 | + s0.synchronize() |
| 329 | - y = torch.ones(10000000, dtype=torch.uint8).npu() | 329 | + # Similarly, both copy() ops are synchronized on s0. |
| 330 | - _test_copy_non_blocking(x, y) | 330 | + self.assertEqual(y, x) |
| 331 | - | 331 | + |
| 332 | - def test_to_non_blocking(self): | 332 | + def test_copy_non_blocking(self): |
| 333 | - stream = torch_npu.npu.current_stream() | 333 | + def _test_copy_non_blocking(a, b): |
| 334 | - | 334 | + event = torch_npu.npu.Event() |
| 335 | - def _test_to_non_blocking(a, non_blocking, dst): | 335 | + a.copy_(b, non_blocking=True) |
| 336 | - torch_npu.npu.synchronize() | 336 | + event.record() |
| 337 | - # Pushes an 0.1 second spin to stream so if the copy is non blocking, | 337 | + event.synchronize() |
| 338 | - # stream will almost surely be active when we query(). | 338 | + self.assertEqual(a, b) |
| 339 | - b = a.to(device=dst, non_blocking=non_blocking) | 339 | + |
| 340 | - stream.synchronize() | 340 | + # 10MB copies |
| 341 | - self.assertEqual(a, b) | 341 | + x = torch.ones(10000000, dtype=torch.uint8).npu() |
| 342 | - self.assertTrue(b.is_pinned() == (non_blocking and dst == "cpu")) | 342 | + y = torch.zeros(10000000, dtype=torch.uint8).pin_memory() |
| 343 | - | 343 | + _test_copy_non_blocking(x, y) |
| 344 | - for dst, try_non_blocking in product(("npu", "cpu"), (True, False)): | 344 | + |
| 345 | - # Creates source on the opposite device from destination. | 345 | + x = torch.zeros(10000000, dtype=torch.uint8).pin_memory() |
| 346 | - src = torch.randn(1000, 1000, 2, 100, | 346 | + y = torch.ones(10000000, dtype=torch.uint8).npu() |
| 347 | - device="npu" if dst == "cpu" else "cpu", | 347 | + _test_copy_non_blocking(x, y) |
| 348 | - pin_memory=True if dst == "npu" else False) | 348 | + |
| 349 | - _test_to_non_blocking(src, try_non_blocking, dst) | 349 | + def test_to_non_blocking(self): |
| 350 | - | 350 | + stream = torch_npu.npu.current_stream() |
| 351 | - def test_to_cpu_blocking_by_default(self): | 351 | + |
| 352 | - src = torch.randn(1000000, device="npu") | 352 | + def _test_to_non_blocking(a, non_blocking, dst): |
| 353 | - torch_npu.npu.synchronize() | 353 | + torch_npu.npu.synchronize() |
| 354 | - dst = src.to(device="cpu") | 354 | + # Pushes an 0.1 second spin to stream so if the copy is non blocking, |
| 355 | - self.assertEqual(src, dst) | 355 | + # stream will almost surely be active when we query(). |
| 356 | - self.assertFalse(dst.is_pinned()) | 356 | + b = a.to(device=dst, non_blocking=non_blocking) |
| 357 | - | 357 | + stream.synchronize() |
| 358 | - def test_torch_manual_seed_seeds_npu_devices(self): | 358 | + self.assertEqual(a, b) |
| 359 | - with freeze_rng_state(): | 359 | + self.assertTrue(b.is_pinned() == (non_blocking and dst == "cpu")) |
| 360 | - x = torch.zeros(4, 4).float() # Not support device RNG (.npu()). | 360 | + |
| 361 | - torch.manual_seed(2) | 361 | + for dst, try_non_blocking in product(("npu", "cpu"), (True, False)): |
| 362 | - self.assertEqual(torch_npu.npu.initial_seed(), 2) | 362 | + # Creates source on the opposite device from destination. |
| 363 | - x.uniform_() | 363 | + src = torch.randn( |
| 364 | - torch.manual_seed(2) | 364 | + 1000, |
| 365 | - y = x.clone().uniform_() | 365 | + 1000, |
| 366 | - self.assertEqual(x, y) | 366 | + 2, |
| 367 | - self.assertEqual(torch_npu.npu.initial_seed(), 2) | 367 | + 100, |
| 368 | - | 368 | + device="npu" if dst == "cpu" else "cpu", |
| 369 | - def test_manual_seed(self): | 369 | + pin_memory=(dst == "npu"), |
| 370 | - with freeze_rng_state(): | 370 | + ) |
| 371 | - x = torch.zeros(4, 4).float() # Not support device RNG (.npu()). | 371 | + _test_to_non_blocking(src, try_non_blocking, dst) |
| 372 | - torch_npu.npu.manual_seed(2) | 372 | + |
| 373 | - torch.manual_seed(2) | 373 | + def test_to_cpu_blocking_by_default(self): |
| 374 | - self.assertEqual(torch_npu.npu.initial_seed(), 2) | 374 | + src = torch.randn(1000000, device="npu") |
| 375 | - x.uniform_() | 375 | + torch_npu.npu.synchronize() |
| 376 | - a = torch.bernoulli(torch.full_like(x, 0.5)) | 376 | + dst = src.to(device="cpu") |
| 377 | - torch.manual_seed(2) | 377 | + self.assertEqual(src, dst) |
| 378 | - y = x.clone().uniform_() | 378 | + self.assertFalse(dst.is_pinned()) |
| 379 | - b = torch.bernoulli(torch.full_like(x, 0.5)) | 379 | + |
| 380 | - self.assertEqual(x, y) | 380 | + def test_torch_manual_seed_seeds_npu_devices(self): |
| 381 | - self.assertEqual(a, b) | 381 | + with freeze_rng_state(): |
| 382 | - self.assertEqual(torch_npu.npu.initial_seed(), 2) | 382 | + x = torch.zeros(4, 4).float() # Not support device RNG (.npu()). |
| 383 | - | 383 | + torch.manual_seed(2) |
| 384 | - def test_get_set_rng_state(self): | 384 | + self.assertEqual(torch_npu.npu.initial_seed(), 2) |
| 385 | - with freeze_rng_state(): | 385 | + x.uniform_() |
| 386 | - torch.manual_seed(3) | 386 | + torch.manual_seed(2) |
| 387 | - cpu_state = torch.get_rng_state() | 387 | + y = x.clone().uniform_() |
| 388 | - npu_state = torch_npu.npu.get_rng_state() | 388 | + self.assertEqual(x, y) |
| 389 | - self.assertEqual(int(cpu_state[0]), 3) | 389 | + self.assertEqual(torch_npu.npu.initial_seed(), 2) |
| 390 | - self.assertEqual(cpu_state[0], npu_state[0]) | 390 | + |
| 391 | - torch_npu.npu.manual_seed(2) | 391 | + def test_manual_seed(self): |
| 392 | - cpu_state_new = torch.get_rng_state() | 392 | + with freeze_rng_state(): |
| 393 | - npu_state = torch_npu.npu.get_rng_state() | 393 | + x = torch.zeros(4, 4).float() # Not support device RNG (.npu()). |
| 394 | - self.assertEqual(cpu_state, cpu_state_new) | 394 | + torch_npu.npu.manual_seed(2) |
| 395 | - self.assertEqual(int(npu_state[0]), 2) | 395 | + torch.manual_seed(2) |
| 396 | - | 396 | + self.assertEqual(torch_npu.npu.initial_seed(), 2) |
| 397 | - def test_get_set_rng_state_input_device(self): | 397 | + x.uniform_() |
| 398 | - npu_state = torch_npu.npu.get_rng_state() | 398 | + a = torch.bernoulli(torch.full_like(x, 0.5)) |
| 399 | - torch_npu.npu.set_rng_state(npu_state) | 399 | + torch.manual_seed(2) |
| 400 | - devices = ["npu", 0, torch.device("npu:0")] | 400 | + y = x.clone().uniform_() |
| 401 | - for device in devices: | 401 | + b = torch.bernoulli(torch.full_like(x, 0.5)) |
| 402 | - npu_state = torch_npu.npu.get_rng_state(device) | 402 | + self.assertEqual(x, y) |
| 403 | - torch_npu.npu.set_rng_state(npu_state, device) | 403 | + self.assertEqual(a, b) |
| 404 | - | 404 | + self.assertEqual(torch_npu.npu.initial_seed(), 2) |
| 405 | - def test_get_device_index(self): | 405 | + |
| 406 | - from torch_npu.npu import _get_device_index | 406 | + def test_get_set_rng_state(self): |
| 407 | - with self.assertRaisesRegex(RuntimeError, "Invalid device string"): | 407 | + with freeze_rng_state(): |
| 408 | - _get_device_index('npu0', optional=True) | 408 | + torch.manual_seed(3) |
| 409 | - | 409 | + cpu_state = torch.get_rng_state() |
| 410 | - with self.assertRaisesRegex(ValueError, "Expected a npu device"): | 410 | + npu_state = torch_npu.npu.get_rng_state() |
| 411 | - cpu_device = torch.device('cpu') | 411 | + self.assertEqual(int(cpu_state[0]), 3) |
| 412 | - _get_device_index(cpu_device, optional=True) | 412 | + self.assertEqual(cpu_state[0], npu_state[0]) |
| 413 | - | 413 | + torch_npu.npu.manual_seed(2) |
| 414 | - def test_npu_synchronize(self): | 414 | + cpu_state_new = torch.get_rng_state() |
| 415 | - torch_npu.npu.synchronize() | 415 | + npu_state = torch_npu.npu.get_rng_state() |
| 416 | - torch_npu.npu.synchronize('npu') | 416 | + self.assertEqual(cpu_state, cpu_state_new) |
| 417 | - torch_npu.npu.synchronize('npu:0') | 417 | + self.assertEqual(int(npu_state[0]), 2) |
| 418 | - torch_npu.npu.synchronize(0) | 418 | + |
| 419 | - torch_npu.npu.synchronize(torch.device('npu:0')) | 419 | + def test_get_set_rng_state_input_device(self): |
| 420 | - | 420 | + npu_state = torch_npu.npu.get_rng_state() |
| 421 | - with self.assertRaisesRegex(ValueError, "Expected a npu device, but"): | 421 | + torch_npu.npu.set_rng_state(npu_state) |
| 422 | - torch_npu.npu.synchronize(torch.device("cpu")) | 422 | + devices = ["npu", 0, torch.device("npu:0")] |
| 423 | - | 423 | + for device in devices: |
| 424 | - with self.assertRaisesRegex(ValueError, "Expected a npu device, but"): | 424 | + npu_state = torch_npu.npu.get_rng_state(device) |
| 425 | - torch_npu.npu.synchronize("cpu") | 425 | + torch_npu.npu.set_rng_state(npu_state, device) |
| 426 | - | 426 | + |
| 427 | - def test_streams(self): | 427 | + def test_get_device_index(self): |
| 428 | - default_stream = torch_npu.npu.current_stream() | 428 | + from torch_npu.npu import _get_device_index |
| 429 | - user_stream = torch_npu.npu.Stream() | 429 | + |
| 430 | - self.assertEqual(torch_npu.npu.current_stream(), default_stream) | 430 | + with self.assertRaisesRegex(RuntimeError, "Invalid device string"): |
| 431 | - self.assertNotEqual(default_stream, user_stream) | 431 | + _get_device_index("npu0", optional=True) |
| 432 | - self.assertNotEqual(user_stream.npu_stream, 0) | 432 | + |
| 433 | - with torch_npu.npu.stream(user_stream): | 433 | + with self.assertRaisesRegex(ValueError, "Expected a npu device"): |
| 434 | - self.assertEqual(torch_npu.npu.current_stream(), user_stream) | 434 | + cpu_device = torch.device("cpu") |
| 435 | - | 435 | + _get_device_index(cpu_device, optional=True) |
| 436 | - def test_stream_event_repr(self): | 436 | + |
| 437 | - s = torch_npu.npu.current_stream() | 437 | + def test_npu_synchronize(self): |
| 438 | - self.assertTrue("torch_npu.npu.Stream" in s.__repr__()) | 438 | + torch_npu.npu.synchronize() |
| 439 | - e = torch_npu.npu.Event() | 439 | + torch_npu.npu.synchronize("npu") |
| 440 | - self.assertTrue("torch_npu.npu.Event" in e.__repr__()) | 440 | + torch_npu.npu.synchronize("npu:0") |
| 441 | - s.record_event(e) | 441 | + torch_npu.npu.synchronize(0) |
| 442 | - self.assertTrue("torch_npu.npu.Event" in e.__repr__()) | 442 | + torch_npu.npu.synchronize(torch.device("npu:0")) |
| 443 | - | 443 | + |
| 444 | - def test_events(self): | 444 | + with self.assertRaisesRegex(ValueError, "Expected a npu device, but"): |
| 445 | - stream = torch_npu.npu.current_stream() | 445 | + torch_npu.npu.synchronize(torch.device("cpu")) |
| 446 | - event = torch_npu.npu.Event(enable_timing=True) | 446 | + |
| 447 | - self.assertTrue(event.query()) | 447 | + with self.assertRaisesRegex(ValueError, "Expected a npu device, but"): |
| 448 | - start_event = torch_npu.npu.Event(enable_timing=True) | 448 | + torch_npu.npu.synchronize("cpu") |
| 449 | - stream.record_event(start_event) | 449 | + |
| 450 | - stream.record_event(event) | 450 | + def test_streams(self): |
| 451 | - event.synchronize() | 451 | + default_stream = torch_npu.npu.current_stream() |
| 452 | - self.assertTrue(event.query()) | 452 | + user_stream = torch_npu.npu.Stream() |
| 453 | - self.assertGreater(start_event.elapsed_time(event), 0) | 453 | + self.assertEqual(torch_npu.npu.current_stream(), default_stream) |
| 454 | - | 454 | + self.assertNotEqual(default_stream, user_stream) |
| 455 | - def test_record_stream(self): | 455 | + self.assertNotEqual(user_stream.npu_stream, 0) |
| 456 | - t = torch.FloatTensor([1, 2, 3, 4]).pin_memory() | 456 | + with torch_npu.npu.stream(user_stream): |
| 457 | - result = torch_npu.npu.FloatTensor(t.size()) | 457 | + self.assertEqual(torch_npu.npu.current_stream(), user_stream) |
| 458 | - stream = torch_npu.npu.Stream() # stream to record tensor copy | 458 | + |
| 459 | - alrm_stream = torch_npu.npu.Stream() # alarm stream as npu not support stream._sleep | 459 | + def test_stream_event_repr(self): |
| 460 | - event = torch_npu.npu.Event() # alarm event | 460 | + s = torch_npu.npu.current_stream() |
| 461 | - ptr = [None] | 461 | + self.assertTrue("torch_npu.npu.Stream" in s.__repr__()) |
| 462 | - | 462 | + e = torch_npu.npu.Event() |
| 463 | - # Performs the CPU->NPU copy in a background stream | 463 | + self.assertTrue("torch_npu.npu.Event" in e.__repr__()) |
| 464 | - with torch_npu.npu.stream(stream): | 464 | + s.record_event(e) |
| 465 | - tmp = t.npu(non_blocking=True) | 465 | + self.assertTrue("torch_npu.npu.Event" in e.__repr__()) |
| 466 | - ptr[0] = tmp.data_ptr() | 466 | + |
| 467 | - torch_npu.npu.current_stream().wait_stream(stream) # wait for copy to complete | 467 | + def test_events(self): |
| 468 | - torch_npu.npu.current_stream().wait_event(event) # wait for alarm event to be recorded for mocking of cuda delay | 468 | + stream = torch_npu.npu.current_stream() |
| 469 | - tmp.record_stream(torch_npu.npu.current_stream()) | 469 | + event = torch_npu.npu.Event(enable_timing=True) |
| 470 | - result.copy_(tmp) | 470 | + self.assertTrue(event.query()) |
| 471 | - with torch_npu.npu.stream(stream): | 471 | + start_event = torch_npu.npu.Event(enable_timing=True) |
| 472 | - tmp2 = torch_npu.npu.FloatTensor(t.size()) | 472 | + stream.record_event(start_event) |
| 473 | - tmp2.zero_() | 473 | + stream.record_event(event) |
| 474 | - # ptr of tmp will not be re-used util alarm event is recorded | 474 | + event.synchronize() |
| 475 | - self.assertNotEqual(tmp2.data_ptr(), ptr[0], message='allocation re-used to soon') | 475 | + self.assertTrue(event.query()) |
| 476 | - alrm_stream.record_event(event) | 476 | + self.assertGreater(start_event.elapsed_time(event), 0) |
| 477 | - | 477 | + |
| 478 | - self.assertEqual(result.tolist(), [1, 2, 3, 4]) | 478 | + def test_record_stream(self): |
| 479 | - | 479 | + t = torch.FloatTensor([1, 2, 3, 4]).pin_memory() |
| 480 | - def test_erase_stream(self): | 480 | + result = torch_npu.npu.FloatTensor(t.size()) |
| 481 | - stream1 = torch_npu.npu.Stream() | 481 | + stream = torch_npu.npu.Stream() # stream to record tensor copy |
| 482 | - stream2 = torch_npu.npu.Stream() | 482 | + alrm_stream = ( |
| 483 | - | 483 | + torch_npu.npu.Stream() |
| 484 | - with torch_npu.npu.stream(stream2): | 484 | + ) # alarm stream as npu not support stream._sleep |
| 485 | - matrix1 = torch.ones(1000, 1000, device='npu') | 485 | + event = torch_npu.npu.Event() # alarm event |
| 486 | - matrix2 = torch.ones(1000, 1000, device='npu') | 486 | + ptr = [None] |
| 487 | - tensor1 = torch.matmul(matrix1, matrix2) | 487 | + |
| 488 | - data_ptr1 = tensor1.data_ptr() | 488 | + # Performs the CPU->NPU copy in a background stream |
| 489 | - | 489 | + with torch_npu.npu.stream(stream): |
| 490 | - tensor1.record_stream(stream1) | 490 | + tmp = t.npu(non_blocking=True) |
| 491 | - torch_npu.erase_stream(tensor1, stream1) | 491 | + ptr[0] = tmp.data_ptr() |
| 492 | - del tensor1 | 492 | + torch_npu.npu.current_stream().wait_stream(stream) # wait for copy to complete |
| 493 | - | 493 | + torch_npu.npu.current_stream().wait_event( |
| 494 | - tensor2 = torch.ones(1000, 1000, device='npu') | 494 | + event |
| 495 | - self.assertEqual(tensor2.data_ptr(), data_ptr1) | 495 | + ) # wait for alarm event to be recorded for mocking of cuda delay |
| 496 | - | 496 | + tmp.record_stream(torch_npu.npu.current_stream()) |
| 497 | - @staticmethod | 497 | + result.copy_(tmp) |
| 498 | - def _stream_synchronize(self, spin_time_cycles): | 498 | + with torch_npu.npu.stream(stream): |
| 499 | - s = torch_npu.npu.current_stream() | 499 | + tmp2 = torch_npu.npu.FloatTensor(t.size()) |
| 500 | - e_tik = torch_npu.npu.Event(enable_timing=True) | 500 | + tmp2.zero_() |
| 501 | - e_tok = torch_npu.npu.Event(enable_timing=True) | 501 | + # ptr of tmp will not be reused util alarm event is recorded |
| 502 | - | 502 | + self.assertNotEqual( |
| 503 | - e_tik.record(s) | 503 | + tmp2.data_ptr(), ptr[0], message="allocation reused to soon" |
| 504 | - e_tok.record(s) | 504 | + ) |
| 505 | - s.synchronize() | 505 | + alrm_stream.record_event(event) |
| 506 | - | 506 | + |
| 507 | - self.assertTrue(s.query()) | 507 | + self.assertEqual(result.tolist(), [1, 2, 3, 4]) |
| 508 | - | 508 | + |
| 509 | - # not necessary to check e_tik and e_tok, as elapsed_time would throw | 509 | + def test_erase_stream(self): |
| 510 | - # exception if otherwise. | 510 | + stream1 = torch_npu.npu.Stream() |
| 511 | - return e_tik.elapsed_time(e_tok) | 511 | + stream2 = torch_npu.npu.Stream() |
| 512 | - | 512 | + |
| 513 | - @staticmethod | 513 | + with torch_npu.npu.stream(stream2): |
| 514 | - def _event_synchronize(self, spin_time_cycles): | 514 | + matrix1 = torch.ones(1000, 1000, device="npu") |
| 515 | - s = torch_npu.npu.current_stream() | 515 | + matrix2 = torch.ones(1000, 1000, device="npu") |
| 516 | - e_tik = torch_npu.npu.Event(enable_timing=True) | 516 | + tensor1 = torch.matmul(matrix1, matrix2) |
| 517 | - e_tok = torch_npu.npu.Event(enable_timing=True) | 517 | + data_ptr1 = tensor1.data_ptr() |
| 518 | - | 518 | + |
| 519 | - e_tik.record(s) | 519 | + tensor1.record_stream(stream1) |
| 520 | - s.record_event(e_tok) | 520 | + torch_npu.erase_stream(tensor1, stream1) |
| 521 | - e_tok.synchronize() | 521 | + del tensor1 |
| 522 | - | 522 | + |
| 523 | - self.assertTrue(s.query()) | 523 | + tensor2 = torch.ones(1000, 1000, device="npu") |
| 524 | - | 524 | + self.assertEqual(tensor2.data_ptr(), data_ptr1) |
| 525 | - # not necessary to check e_tik and e_tok, as elapsed_time would throw | 525 | + |
| 526 | - # exception if otherwise. | 526 | + @staticmethod |
| 527 | - return e_tik.elapsed_time(e_tok) | 527 | + def _stream_synchronize(self, spin_time_cycles): |
| 528 | - | 528 | + s = torch_npu.npu.current_stream() |
| 529 | - @staticmethod | 529 | + e_tik = torch_npu.npu.Event(enable_timing=True) |
| 530 | - def _event_wait(self, spin_time_cycles): | 530 | + e_tok = torch_npu.npu.Event(enable_timing=True) |
| 531 | - s0 = torch_npu.npu.current_stream() | 531 | + |
| 532 | - s1 = torch_npu.npu.Stream() | 532 | + e_tik.record(s) |
| 533 | - e_tik = torch_npu.npu.Event(blocking=True, enable_timing=True) | 533 | + e_tok.record(s) |
| 534 | - e_tok = torch_npu.npu.Event(blocking=True, enable_timing=True) | 534 | + s.synchronize() |
| 535 | - | 535 | + |
| 536 | - e_tik.record(s0) | 536 | + self.assertTrue(s.query()) |
| 537 | - e_sync = torch_npu.npu.Event(blocking=True) | 537 | + |
| 538 | - e_sync.record() | 538 | + # not necessary to check e_tik and e_tok, as elapsed_time would throw |
| 539 | - e_sync.wait(s1) | 539 | + # exception if otherwise. |
| 540 | - s1.synchronize() | 540 | + return e_tik.elapsed_time(e_tok) |
| 541 | - e_tok.record() | 541 | + |
| 542 | - e_tok.synchronize() | 542 | + @staticmethod |
| 543 | - | 543 | + def _event_synchronize(self, spin_time_cycles): |
| 544 | - self.assertTrue(s0.query()) | 544 | + s = torch_npu.npu.current_stream() |
| 545 | - self.assertTrue(s1.query()) | 545 | + e_tik = torch_npu.npu.Event(enable_timing=True) |
| 546 | - self.assertTrue(e_sync.query()) | 546 | + e_tok = torch_npu.npu.Event(enable_timing=True) |
| 547 | - | 547 | + |
| 548 | - # not necessary to check e_tik and e_tok, as elapsed_time would throw | 548 | + e_tik.record(s) |
| 549 | - # exception if otherwise. | 549 | + s.record_event(e_tok) |
| 550 | - return e_tik.elapsed_time(e_tok) | 550 | + e_tok.synchronize() |
| 551 | - | 551 | + |
| 552 | - @staticmethod | 552 | + self.assertTrue(s.query()) |
| 553 | - def _test_stream_event_nogil(self, sync_func, p2c, c2p): | 553 | + |
| 554 | - with torch_npu.npu.device('npu:1'): | 554 | + # not necessary to check e_tik and e_tok, as elapsed_time would throw |
| 555 | - c2p.put(0) | 555 | + # exception if otherwise. |
| 556 | - p2c.get() | 556 | + return e_tik.elapsed_time(e_tok) |
| 557 | - c2p.put(sync_func(self, TestNpu.FIFTY_MIL_CYCLES)) | 557 | + |
| 558 | - | 558 | + @staticmethod |
| 559 | - def test_noncontiguous_pinned_memory(self): | 559 | + def _event_wait(self, spin_time_cycles): |
| 560 | - # See issue #3266 | 560 | + s0 = torch_npu.npu.current_stream() |
| 561 | - x = torch.arange(0, 10).view((2, 5)) | 561 | + s1 = torch_npu.npu.Stream() |
| 562 | - self.assertEqual(x.t(), x.t().pin_memory()) | 562 | + e_tik = torch_npu.npu.Event(blocking=True, enable_timing=True) |
| 563 | - | 563 | + e_tok = torch_npu.npu.Event(blocking=True, enable_timing=True) |
| 564 | - def test_caching_pinned_memory(self): | 564 | + |
| 565 | - | 565 | + e_tik.record(s0) |
| 566 | - # check that allocations are re-used after deletion | 566 | + e_sync = torch_npu.npu.Event(blocking=True) |
| 567 | - t = torch.FloatTensor([1]).pin_memory() | 567 | + e_sync.record() |
| 568 | - ptr = t.data_ptr() | 568 | + e_sync.wait(s1) |
| 569 | - del t | 569 | + s1.synchronize() |
| 570 | - t = torch.FloatTensor([1]).pin_memory() | 570 | + e_tok.record() |
| 571 | - self.assertEqual(t.data_ptr(), ptr) | 571 | + e_tok.synchronize() |
| 572 | - | 572 | + |
| 573 | - # check that the allocation is not re-used if it's in-use by a copy | 573 | + self.assertTrue(s0.query()) |
| 574 | - npu_tensor = torch.npu.FloatTensor([0]) | 574 | + self.assertTrue(s1.query()) |
| 575 | - npu_tensor.copy_(t, non_blocking=True) | 575 | + self.assertTrue(e_sync.query()) |
| 576 | - del t | 576 | + |
| 577 | - t = torch.FloatTensor([1]).pin_memory() | 577 | + # not necessary to check e_tik and e_tok, as elapsed_time would throw |
| 578 | - self.assertEqual(list(npu_tensor), [1]) | 578 | + # exception if otherwise. |
| 579 | - | 579 | + return e_tik.elapsed_time(e_tok) |
| 580 | - def test_function_torch_empty_and_to(self): | 580 | + |
| 581 | - x = torch.empty((2, 3), dtype=torch.float16, device='npu') | 581 | + @staticmethod |
| 582 | - x_int32 = x.to(torch.int32) | 582 | + def _test_stream_event_nogil(self, sync_func, p2c, c2p): |
| 583 | - res = x_int32 + 1 | 583 | + with torch_npu.npu.device("npu:1"): |
| 584 | - | 584 | + c2p.put(0) |
| 585 | - def test_function_npu(self): | 585 | + p2c.get() |
| 586 | - x = torch.empty((2, 3), dtype=torch.float16, device='cpu') | 586 | + c2p.put(sync_func(self, TestNpu.FIFTY_MIL_CYCLES)) |
| 587 | - x_npu = x.npu() | 587 | + |
| 588 | - res = x_npu + 1 | 588 | + def test_noncontiguous_pinned_memory(self): |
| 589 | - | 589 | + # See issue #3266 |
| 590 | - def test_function_torch_empty_with_format(self): | 590 | + x = torch.arange(0, 10).view((2, 5)) |
| 591 | - x = torch_npu.empty_with_format((2, 3), dtype=torch.float32, device='npu') | 591 | + self.assertEqual(x.t(), x.t().pin_memory()) |
| 592 | - res = x + 1 | 592 | + |
| 593 | - | 593 | + def test_caching_pinned_memory(self): |
| 594 | - def test_function_torch_empty_like(self): | 594 | + # check that allocations are reused after deletion |
| 595 | - x = torch.empty((2, 3), dtype=torch.float32, device='npu') | 595 | + t = torch.FloatTensor([1]).pin_memory() |
| 596 | - x_like = torch.empty_like(x) | 596 | + ptr = t.data_ptr() |
| 597 | - res = x_like + 1 | 597 | + del t |
| 598 | - | 598 | + t = torch.FloatTensor([1]).pin_memory() |
| 599 | - def test_function_torch_empty_like_with_stride(self): | 599 | + self.assertEqual(t.data_ptr(), ptr) |
| 600 | - # if a is contiguous, stride of b should be same as a | 600 | + |
| 601 | - a = torch.empty([16, 32], device='npu') | 601 | + # check that the allocation is not reused if it's in-use by a copy |
| 602 | - self.assertTrue(a.is_contiguous()) | 602 | + npu_tensor = torch.npu.FloatTensor([0]) |
| 603 | - self.assertEqual(a.stride(), (32, 1)) | 603 | + npu_tensor.copy_(t, non_blocking=True) |
| 604 | - | 604 | + del t |
| 605 | - b = torch.empty_like(a) | 605 | + t = torch.FloatTensor([1]).pin_memory() # noqa: F841 |
| 606 | - self.assertTrue(b.is_contiguous()) | 606 | + self.assertEqual(list(npu_tensor), [1]) |
| 607 | - self.assertEqual(b.stride(), (32, 1)) | 607 | + |
| 608 | - | 608 | + def test_function_torch_empty_and_to(self): |
| 609 | - b = torch.empty_like(a, memory_format=torch.preserve_format) | 609 | + x = torch.empty((2, 3), dtype=torch.float16, device="npu") |
| 610 | - self.assertTrue(b.is_contiguous()) | 610 | + x_int32 = x.to(torch.int32) |
| 611 | - self.assertEqual(b.stride(), (32, 1)) | 611 | + res = x_int32 + 1 # noqa: F841 |
| 612 | - | 612 | + |
| 613 | - b = torch.empty_like(a, memory_format=torch.contiguous_format) | 613 | + def test_function_npu(self): |
| 614 | - self.assertTrue(b.is_contiguous()) | 614 | + x = torch.empty((2, 3), dtype=torch.float16, device="cpu") |
| 615 | - self.assertEqual(b.stride(), (32, 1)) | 615 | + x_npu = x.npu() |
| 616 | - | 616 | + res = x_npu + 1 # noqa: F841 |
| 617 | - # if a is Not contiguous, stride of b should be same as a when memory_format=torch.preserve_format(default) | 617 | + |
| 618 | - a = torch.empty([16, 32], device='npu').T | 618 | + def test_function_torch_empty_with_format(self): |
| 619 | - self.assertFalse(a.is_contiguous()) | 619 | + x = torch_npu.empty_with_format((2, 3), dtype=torch.float32, device="npu") |
| 620 | - self.assertEqual(a.stride(), (1, 32)) | 620 | + res = x + 1 # noqa: F841 |
| 621 | - | 621 | + |
| 622 | - b = torch.empty_like(a) | 622 | + def test_function_torch_empty_like(self): |
| 623 | - self.assertFalse(b.is_contiguous()) | 623 | + x = torch.empty((2, 3), dtype=torch.float32, device="npu") |
| 624 | - self.assertEqual(b.stride(), (1, 32)) | 624 | + x_like = torch.empty_like(x) |
| 625 | - | 625 | + res = x_like + 1 # noqa: F841 |
| 626 | - b = torch.empty_like(a, memory_format=torch.preserve_format) | 626 | + |
| 627 | - self.assertFalse(b.is_contiguous()) | 627 | + def test_function_torch_empty_like_with_stride(self): |
| 628 | - self.assertEqual(b.stride(), (1, 32)) | 628 | + # if a is contiguous, stride of b should be same as a |
| 629 | - | 629 | + a = torch.empty([16, 32], device="npu") |
| 630 | - b = torch.empty_like(a, memory_format=torch.contiguous_format) | 630 | + self.assertTrue(a.is_contiguous()) |
| 631 | - self.assertTrue(b.is_contiguous()) | 631 | + self.assertEqual(a.stride(), (32, 1)) |
| 632 | - self.assertEqual(b.stride(), (16, 1)) | 632 | + |
| 633 | - | 633 | + b = torch.empty_like(a) |
| 634 | - def test_function_torch_empty_like_in_fake_tensor_mode(self): | 634 | + self.assertTrue(b.is_contiguous()) |
| 635 | - with torch._subclasses.fake_tensor.FakeTensorMode(): | 635 | + self.assertEqual(b.stride(), (32, 1)) |
| 636 | - x = torch.rand(3, 3).npu() | 636 | + |
| 637 | - with torch.utils._mode_utils.no_dispatch(): | 637 | + b = torch.empty_like(a, memory_format=torch.preserve_format) |
| 638 | - x_like = torch.empty_like(x) | 638 | + self.assertTrue(b.is_contiguous()) |
| 639 | - self.assertEqual(torch_npu.get_npu_format(x_like), 2) | 639 | + self.assertEqual(b.stride(), (32, 1)) |
| 640 | - | 640 | + |
| 641 | - def test_function_torch_empty_strided(self): | 641 | + b = torch.empty_like(a, memory_format=torch.contiguous_format) |
| 642 | - x = torch.empty_strided((2, 3), (1, 2), dtype=torch.int8, device='npu') | 642 | + self.assertTrue(b.is_contiguous()) |
| 643 | - | 643 | + self.assertEqual(b.stride(), (32, 1)) |
| 644 | - def test_function_tensor_new_empty(self): | 644 | + |
| 645 | - x = torch.ones(()).npu() | 645 | + # if a is Not contiguous, stride of b should be same as a when memory_format=torch.preserve_format(default) |
| 646 | - x_new_empty = x.new_empty((2, 3), dtype=torch.float16, device='npu') | 646 | + a = torch.empty([16, 32], device="npu").T |
| 647 | - res = x_new_empty + 1 | 647 | + self.assertFalse(a.is_contiguous()) |
| 648 | - x_new_empty = x.new_empty(size=(2, 3), dtype=torch.float16, device='npu') | 648 | + self.assertEqual(a.stride(), (1, 32)) |
| 649 | - res = x_new_empty + 1 | 649 | + |
| 650 | - | 650 | + b = torch.empty_like(a) |
| 651 | - def test_function_tensor_new_empty_strided(self): | 651 | + self.assertFalse(b.is_contiguous()) |
| 652 | - x = torch.ones(()).npu() | 652 | + self.assertEqual(b.stride(), (1, 32)) |
| 653 | - x_new = x.new_empty_strided([2, 3], [3, 1], dtype=torch.float32, device='npu') | 653 | + |
| 654 | - res = x_new + 1 | 654 | + b = torch.empty_like(a, memory_format=torch.preserve_format) |
| 655 | - | 655 | + self.assertFalse(b.is_contiguous()) |
| 656 | - def test_function_tensor_data_npu(self): | 656 | + self.assertEqual(b.stride(), (1, 32)) |
| 657 | - x = torch.ones(()) | 657 | + |
| 658 | - x.data = x.data.npu() | 658 | + b = torch.empty_like(a, memory_format=torch.contiguous_format) |
| 659 | - | 659 | + self.assertTrue(b.is_contiguous()) |
| 660 | - def test_function_tensor_new_full(self): | 660 | + self.assertEqual(b.stride(), (16, 1)) |
| 661 | - x_cpu = torch.tensor((), dtype=torch.float32) | 661 | + |
| 662 | - cpu_out = x_cpu.new_full((2, 3), 3.1) | 662 | + def test_function_torch_empty_like_in_fake_tensor_mode(self): |
| 663 | - | 663 | + with torch._subclasses.fake_tensor.FakeTensorMode(): |
| 664 | - x = torch.tensor((), dtype=torch.float32).npu() | 664 | + x = torch.rand(3, 3).npu() |
| 665 | - npu_output1 = x.new_full((2, 3), 3.1, device=None, requires_grad=False) | 665 | + with torch.utils._mode_utils.no_dispatch(): |
| 666 | - npu_output2 = x.new_full((2, 3), 3.1, device='cpu', requires_grad=False) | 666 | + x_like = torch.empty_like(x) |
| 667 | - npu_output3 = x.new_full((2, 3), 3.1, device='npu', requires_grad=False) | 667 | + self.assertEqual(torch_npu.get_npu_format(x_like), 2) |
| 668 | - self.assertRtolEqual(cpu_out.numpy(), npu_output1.cpu().numpy()) | 668 | + |
| 669 | - self.assertRtolEqual(cpu_out.numpy(), npu_output2.cpu().numpy()) | 669 | + @classmethod |
| 670 | - self.assertRtolEqual(cpu_out.numpy(), npu_output3.cpu().numpy()) | 670 | + def setUpClass(cls): |
| 671 | - | 671 | + cls.npu_available = torch.npu.is_available() |
| 672 | - def test_function_tensor_new_ones(self): | 672 | + |
| 673 | - x_cpu = torch.tensor((), dtype=torch.float32) | 673 | + def _compare_with_cpu(self, cpu_tensor, npu_tensor, memory_format=None): |
| 674 | - cpu_out = x_cpu.new_ones((2, 3)) | 674 | + cpu_result = torch.empty_like(cpu_tensor, memory_format=memory_format) |
| 675 | - | 675 | + npu_result = torch.empty_like(npu_tensor, memory_format=memory_format) |
| 676 | - x = torch.tensor((), dtype=torch.float32).npu() | 676 | + self.assertEqual( |
| 677 | - npu_output1 = x.new_ones((2, 3), device=None, requires_grad=False) | 677 | + cpu_result.shape, |
| 678 | - npu_output2 = x.new_ones((2, 3), device='cpu', requires_grad=False) | 678 | + npu_result.shape, |
| 679 | - npu_output3 = x.new_ones((2, 3), device='npu', requires_grad=False) | 679 | + f"Shape mismatch: CPU {cpu_result.shape} vs NPU {npu_result.shape}", |
| 680 | - npu_output4 = x.new_ones(size=(2, 3), device='npu', requires_grad=False) | 680 | + ) |
| 681 | - self.assertRtolEqual(cpu_out.numpy(), npu_output1.cpu().numpy()) | 681 | + self.assertEqual( |
| 682 | - self.assertRtolEqual(cpu_out.numpy(), npu_output2.cpu().numpy()) | 682 | + cpu_result.dtype, |
| 683 | - self.assertRtolEqual(cpu_out.numpy(), npu_output3.cpu().numpy()) | 683 | + npu_result.dtype, |
| 684 | - self.assertRtolEqual(cpu_out.numpy(), npu_output4.cpu().numpy()) | 684 | + f"Dtype mismatch: CPU {cpu_result.dtype} vs NPU {npu_result.dtype}", |
| 685 | - | 685 | + ) |
| 686 | - def test_function_tensor_new_tensor(self): | 686 | + self.assertEqual( |
| 687 | - x_cpu = torch.tensor((), dtype=torch.float32) | 687 | + cpu_result.stride(), |
| 688 | - x = torch.tensor((), dtype=torch.float32).npu() | 688 | + npu_result.stride(), |
| 689 | - | 689 | + f"Stride mismatch: CPU {cpu_result.stride()} vs NPU {npu_result.stride()}", |
| 690 | - list_input = [[1, 2, 3], [4, 5, 6]] | 690 | + ) |
| 691 | - cpu_out = x_cpu.new_tensor(list_input) | 691 | + self.assertEqual( |
| 692 | - npu_output1 = x.new_tensor(list_input, device=None, requires_grad=False) | 692 | + cpu_result.is_contiguous(), |
| 693 | - npu_output2 = x.new_tensor(list_input, device='cpu', requires_grad=False) | 693 | + npu_result.is_contiguous(), |
| 694 | - npu_output3 = x.new_tensor(list_input, device='npu', requires_grad=False) | 694 | + f"Contiguity mismatch: CPU {cpu_result.is_contiguous()} vs NPU {npu_result.is_contiguous()}", |
| 695 | - self.assertRtolEqual(cpu_out.numpy(), npu_output1.cpu().numpy()) | 695 | + ) |
| 696 | - self.assertRtolEqual(cpu_out.numpy(), npu_output2.cpu().numpy()) | 696 | + |
| 697 | - print(cpu_out.numpy().dtype, npu_output3.cpu().numpy().dtype) | 697 | + # ---------------------------------------------------------------- |
| 698 | - self.assertRtolEqual(cpu_out.numpy(), npu_output3.cpu().numpy()) | 698 | + # Scenario 1: Transposed tensor (non-contiguous, strides reversed) |
| 699 | - | 699 | + # ---------------------------------------------------------------- |
| 700 | - np_input = np.array(list_input) | 700 | + def test_transposed_2d(self): |
| 701 | - cpu_out = x_cpu.new_tensor(np_input) | 701 | + """Transposed 2D tensor: strides should be preserved on both CPU and NPU.""" |
| 702 | - npu_output1 = x.new_tensor(np_input, device=None, requires_grad=False) | 702 | + cpu_t = torch.randn(4, 6).t() |
| 703 | - npu_output2 = x.new_tensor(np_input, device='cpu', requires_grad=False) | 703 | + npu_t = cpu_t.npu() |
| 704 | - npu_output3 = x.new_tensor(np_input, device='npu', requires_grad=False) | 704 | + self.assertFalse(cpu_t.is_contiguous()) |
| 705 | - self.assertRtolEqual(cpu_out.numpy(), npu_output1.cpu().numpy()) | 705 | + self.assertFalse(npu_t.is_contiguous()) |
| 706 | - self.assertRtolEqual(cpu_out.numpy(), npu_output2.cpu().numpy()) | 706 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) |
| 707 | - self.assertRtolEqual(cpu_out.numpy(), npu_output3.cpu().numpy()) | 707 | + |
| 708 | - | 708 | + def test_transposed_3d(self): |
| 709 | - tensor_input = torch.tensor(list_input) | 709 | + """Transposed 3D tensor: strides should be preserved.""" |
| 710 | - cpu_out = x_cpu.new_tensor(tensor_input) | 710 | + cpu_t = torch.randn(2, 3, 4).transpose(0, 2) |
| 711 | - npu_output1 = x.new_tensor(tensor_input, device=None, requires_grad=False) | 711 | + npu_t = cpu_t.npu() |
| 712 | - npu_output2 = x.new_tensor(tensor_input, device='cpu', requires_grad=False) | 712 | + self.assertFalse(cpu_t.is_contiguous()) |
| 713 | - npu_output3 = x.new_tensor(tensor_input, device='npu', requires_grad=False) | 713 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) |
| 714 | - self.assertRtolEqual(cpu_out.numpy(), npu_output1.cpu().numpy()) | 714 | + |
| 715 | - self.assertRtolEqual(cpu_out.numpy(), npu_output2.cpu().numpy()) | 715 | + # ---------------------------------------------------------------- |
| 716 | - self.assertRtolEqual(cpu_out.numpy(), npu_output3.cpu().numpy()) | 716 | + # Scenario 2: Sliced tensor (non-contiguous via slicing) |
| 717 | - | 717 | + # ---------------------------------------------------------------- |
| 718 | - def test_function_tensor_new_zeros(self): | 718 | + def test_sliced_dim0(self): |
| 719 | - x_cpu = torch.tensor((), dtype=torch.float32) | 719 | + """Slice along dim 0: every other row, non-contiguous.""" |
| 720 | - cpu_out = x_cpu.new_zeros((2, 3)) | 720 | + cpu_t = torch.randn(8, 5)[::2] |
| 721 | - | 721 | + npu_t = cpu_t.npu() |
| 722 | - x = torch.tensor((), dtype=torch.float32).npu() | 722 | + self.assertFalse(cpu_t.is_contiguous()) |
| 723 | - npu_output1 = x.new_zeros((2, 3), device=None, requires_grad=False) | 723 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) |
| 724 | - npu_output2 = x.new_zeros((2, 3), device='cpu', requires_grad=False) | 724 | + |
| 725 | - npu_output3 = x.new_zeros((2, 3), device='npu', requires_grad=False) | 725 | + def test_sliced_dim1(self): |
| 726 | - npu_output4 = x.new_zeros(size=(2, 3), device='npu', requires_grad=False) | 726 | + """Slice along dim 1: every other column, non-contiguous.""" |
| 727 | - self.assertRtolEqual(cpu_out.numpy(), npu_output1.cpu().numpy()) | 727 | + cpu_t = torch.randn(4, 8)[:, ::2] |
| 728 | - self.assertRtolEqual(cpu_out.numpy(), npu_output2.cpu().numpy()) | 728 | + npu_t = cpu_t.npu() |
| 729 | - self.assertRtolEqual(cpu_out.numpy(), npu_output3.cpu().numpy()) | 729 | + self.assertFalse(cpu_t.is_contiguous()) |
| 730 | - self.assertRtolEqual(cpu_out.numpy(), npu_output4.cpu().numpy()) | 730 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) |
| 731 | - | 731 | + |
| 732 | - def test_type_conversions_npu(self): | 732 | + # ---------------------------------------------------------------- |
| 733 | - x = torch.randn(5, 5) | 733 | + # Scenario 3: Narrowed tensor |
| 734 | - self.assertIsInstance(x.float(), torch.FloatTensor) | 734 | + # ---------------------------------------------------------------- |
| 735 | - self.assertIsInstance(x.double().npu(), torch.npu.DoubleTensor) | 735 | + def test_narrowed(self): |
| 736 | - self.assertIsInstance(x.npu().float(), torch.npu.FloatTensor) | 736 | + """Narrowed tensor: non-contiguous via narrow.""" |
| 737 | - self.assertIsInstance(x.npu().float().cpu(), torch.FloatTensor) | 737 | + cpu_t = torch.randn(6, 8).narrow(1, 1, 5) |
| 738 | - self.assertIsInstance(x.npu().float().cpu().int(), torch.IntTensor) | 738 | + npu_t = cpu_t.npu() |
| 739 | - | 739 | + self.assertFalse(cpu_t.is_contiguous()) |
| 740 | - def _test_type_conversion_backward(self, t): | 740 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) |
| 741 | - fvar = Variable(t(torch.randn(5, 5).float()), requires_grad=True) | 741 | + |
| 742 | - fvar.double().sum().backward() | 742 | + # ---------------------------------------------------------------- |
| 743 | - self.assertEqual(fvar.grad, torch.ones_like(fvar)) | 743 | + # Scenario 4: Expanded tensor (stride 0 dimension) |
| 744 | - self.assertEqual(type(fvar.grad), type(fvar)) | 744 | + # ---------------------------------------------------------------- |
| 745 | - dvar = Variable(t(torch.randn(5, 5).double()), requires_grad=True) | 745 | + def test_expanded(self): |
| 746 | - dvar.float().sum().backward() | 746 | + """Expanded tensor: has stride=0 dimension, non-contiguous.""" |
| 747 | - self.assertEqual(dvar.grad, torch.ones_like(dvar)) | 747 | + cpu_t = torch.randn(1, 5).expand(4, 5) |
| 748 | - self.assertEqual(type(dvar.grad), type(dvar)) | 748 | + npu_t = cpu_t.npu() |
| 749 | - | 749 | + self.assertFalse(cpu_t.is_contiguous()) |
| 750 | - def test_type_conversions(self): | 750 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) |
| 751 | - x = torch.randn(5, 5) | 751 | + |
| 752 | - self.assertIsInstance(x.float(), torch.FloatTensor) | 752 | + def test_expanded_3d(self): |
| 753 | - self.assertIsInstance(x.int(), torch.IntTensor) | 753 | + """Expanded 3D tensor with stride=0.""" |
| 754 | - if torch.npu.is_available(): | 754 | + cpu_t = torch.randn(2, 1, 4).expand(2, 3, 4) |
| 755 | - self.assertIsInstance(x.float().npu(), torch.npu.FloatTensor) | 755 | + npu_t = cpu_t.npu() |
| 756 | - self.assertIsInstance(x.double().npu(), torch.npu.DoubleTensor) | 756 | + self.assertFalse(cpu_t.is_contiguous()) |
| 757 | - self.assertIsInstance(x.int().npu(), torch.npu.IntTensor) | 757 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) |
| 758 | - self.assertIsInstance(x.int().npu().cpu(), torch.IntTensor) | 758 | + |
| 759 | - | 759 | + # ---------------------------------------------------------------- |
| 760 | - tensor_types = [torch.DoubleTensor, torch.FloatTensor, torch.IntTensor, torch.ByteTensor] | 760 | + # Scenario 5: Permuted tensor |
| 761 | - for t, y_var in product(tensor_types, (True, False)): | 761 | + # ---------------------------------------------------------------- |
| 762 | - y = torch.randint(5, (5, 5), dtype=t.dtype) | 762 | + def test_permuted(self): |
| 763 | - y = Variable(y) if y_var else y | 763 | + """Permuted tensor: dimension order changed.""" |
| 764 | - self.assertIsInstance(x.type(t), t) | 764 | + cpu_t = torch.randn(2, 3, 4, 5).permute(3, 1, 0, 2) |
| 765 | - self.assertIsInstance(x.type_as(y), t) | 765 | + npu_t = cpu_t.npu() |
| 766 | - | 766 | + self.assertFalse(cpu_t.is_contiguous()) |
| 767 | - t_dtype = t().dtype | 767 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) |
| 768 | - self.assertIsInstance(x.type(t_dtype), t) | 768 | + |
| 769 | - self.assertIs(t_dtype, x.type(t_dtype).dtype) | 769 | + # ---------------------------------------------------------------- |
| 770 | - self.assertEqual(y.data_ptr(), y.type(t).data_ptr()) | 770 | + # Scenario 6: Select on a non-leading dimension |
| 771 | - | 771 | + # ---------------------------------------------------------------- |
| 772 | - self._test_type_conversion_backward(lambda x: x) | 772 | + def test_select_dim1(self): |
| 773 | - if torch.npu.is_available(): | 773 | + """Select along dim 1: produces non-contiguous tensor.""" |
| 774 | - self._test_type_conversion_backward(lambda x: x.npu()) | 774 | + cpu_t = torch.randn(3, 5, 4).select(1, 2) |
| 775 | - | 775 | + npu_t = cpu_t.npu() |
| 776 | - def test_get_allocator_backend(self): | 776 | + self.assertFalse(cpu_t.is_contiguous()) |
| 777 | - npu_allocator_name = torch.npu.get_allocator_backend() | 777 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) |
| 778 | - self.assertEqual(npu_allocator_name, "native") | 778 | + |
| 779 | - | 779 | + # ---------------------------------------------------------------- |
| 780 | - def test_contiguous(self): | 780 | + # Scenario 7: Diagonal / as_strided |
| 781 | - def run_once(): | 781 | + # ---------------------------------------------------------------- |
| 782 | - x = torch.randn(4, 3, 8, 8).npu() | 782 | + def test_as_strided_non_contiguous(self): |
| 783 | - x = x.permute(0, 2, 1, 3) | 783 | + """Manually constructed non-contiguous strides via as_strided.""" |
| 784 | - x = x.contiguous() | 784 | + cpu_t = torch.randn(12).as_strided((3, 3), (4, 1)) |
| 785 | - return x | 785 | + npu_t = cpu_t.npu() |
| 786 | - | 786 | + self.assertFalse(cpu_t.is_contiguous()) |
| 787 | - with torch._subclasses.fake_tensor.FakeTensorMode(): | 787 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) |
| 788 | - y = run_once() | 788 | + |
| 789 | - self.assertTrue(y.is_contiguous()) | 789 | + # ---------------------------------------------------------------- |
| 790 | - | 790 | + # Scenario 8: memory_format=None (default Preserve) |
| 791 | - x = torch.randn(1, 16, 5, 5).npu() | 791 | + # Non-contiguous input but no preserve_format requested |
| 792 | - self.assertTrue(x.is_contiguous()) | 792 | + # ---------------------------------------------------------------- |
| 793 | - stride = list(x.stride()) | 793 | + def test_transposed_memory_format_none(self): |
| 794 | - stride[0] = 20 | 794 | + """Transposed tensor with memory_format=None (default Contiguous). |
| 795 | - # change the stride in dimension 0. the tensor is still contiguous because size[0] is 1 | 795 | + Both CPU and NPU should produce contiguous result.""" |
| 796 | - x.set_(x.storage(), 0, x.size(), stride) | 796 | + cpu_t = torch.randn(4, 6).t() |
| 797 | - self.assertTrue(x.is_contiguous()) | 797 | + npu_t = cpu_t.npu() |
| 798 | - | 798 | + self.assertFalse(cpu_t.is_contiguous()) |
| 799 | - x.contiguous(memory_format=torch.contiguous_format) | 799 | + |
| 800 | - x.contiguous(memory_format=torch.preserve_format) | 800 | + cpu_result = torch.empty_like(cpu_t) |
| 801 | - | 801 | + npu_result = torch.empty_like(npu_t) |
| 802 | - with self.assertRaisesRegex(RuntimeError, "ERR01007 OPS feature not supported"): | 802 | + |
| 803 | - x.contiguous(memory_format=torch.channels_last) | 803 | + self.assertFalse( |
| 804 | - | 804 | + cpu_result.is_contiguous(), |
| 805 | - with self.assertRaisesRegex(RuntimeError, "ERR01007 OPS feature not supported"): | 805 | + "CPU result should be non-ontiguous with default memory_format", |
| 806 | - x.contiguous(memory_format=torch.channels_last_3d) | 806 | + ) |
| 807 | - | 807 | + self.assertFalse( |
| 808 | - def test_uuid(self): | 808 | + npu_result.is_contiguous(), |
| 809 | - uuid = torch.npu.get_device_properties(0).uuid | 809 | + "NPU result should be non-contiguous with default memory_format", |
| 810 | - self.assertEqual(len(str(uuid)), 36) | 810 | + ) |
| 811 | - self.assertEqual(len(uuid.bytes), 16) | 811 | + self.assertEqual(cpu_result.shape, npu_result.shape) |
| 812 | - | 812 | + self.assertEqual(cpu_result.stride(), npu_result.stride()) |
| 813 | - | 813 | + |
| 814 | -if __name__ == '__main__': | 814 | + # ---------------------------------------------------------------- |
| 815 | - run_tests() | 815 | + # Scenario 9: memory_format=Contiguous on non-contiguous input |
| 816 | + # ---------------------------------------------------------------- | ||
| 817 | + def test_transposed_memory_format_contiguous(self): | ||
| 818 | + """Transposed tensor with memory_format=Contiguous. | ||
| 819 | + Both should produce contiguous result.""" | ||
| 820 | + cpu_t = torch.randn(4, 6).t() | ||
| 821 | + npu_t = cpu_t.npu() | ||
| 822 | + | ||
| 823 | + cpu_result = torch.empty_like(cpu_t, memory_format=torch.contiguous_format) | ||
| 824 | + npu_result = torch.empty_like(npu_t, memory_format=torch.contiguous_format) | ||
| 825 | + | ||
| 826 | + self.assertTrue(cpu_result.is_contiguous()) | ||
| 827 | + self.assertTrue(npu_result.is_contiguous()) | ||
| 828 | + self.assertEqual(cpu_result.shape, npu_result.shape) | ||
| 829 | + | ||
| 830 | + # ---------------------------------------------------------------- | ||
| 831 | + # Scenario 10: Contiguous input with preserve_format (baseline) | ||
| 832 | + # ---------------------------------------------------------------- | ||
| 833 | + def test_contiguous_preserve(self): | ||
| 834 | + """Contiguous tensor with preserve_format: should remain contiguous.""" | ||
| 835 | + cpu_t = torch.randn(4, 6) | ||
| 836 | + npu_t = cpu_t.npu() | ||
| 837 | + self.assertTrue(cpu_t.is_contiguous()) | ||
| 838 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 839 | + | ||
| 840 | + # ---------------------------------------------------------------- | ||
| 841 | + # Scenario 11: 1D non-contiguous (via as_strided) | ||
| 842 | + # ---------------------------------------------------------------- | ||
| 843 | + def test_1d_non_contiguous(self): | ||
| 844 | + """1D non-contiguous tensor via as_strided with stride > 1.""" | ||
| 845 | + cpu_t = torch.randn(10).as_strided((4,), (2,)) | ||
| 846 | + npu_t = cpu_t.npu() | ||
| 847 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 848 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 849 | + | ||
| 850 | + # ---------------------------------------------------------------- | ||
| 851 | + # Scenario 12: Double transposed (back to contiguous but different strides) | ||
| 852 | + # ---------------------------------------------------------------- | ||
| 853 | + def test_double_transpose(self): | ||
| 854 | + """Double transpose: logically contiguous but stride order differs from default.""" | ||
| 855 | + cpu_t = torch.randn(3, 4).t().t() | ||
| 856 | + npu_t = cpu_t.npu() | ||
| 857 | + self.assertTrue(cpu_t.is_contiguous()) | ||
| 858 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 859 | + | ||
| 860 | + # ---------------------------------------------------------------- | ||
| 861 | + # Scenario 13: Non-contiguous with different dtypes | ||
| 862 | + # ---------------------------------------------------------------- | ||
| 863 | + def test_transposed_float16(self): | ||
| 864 | + """Non-contiguous float16 tensor.""" | ||
| 865 | + cpu_t = torch.randn(4, 6, dtype=torch.float16).t() | ||
| 866 | + npu_t = cpu_t.npu() | ||
| 867 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 868 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 869 | + | ||
| 870 | + def test_transposed_int64(self): | ||
| 871 | + """Non-contiguous int64 tensor.""" | ||
| 872 | + cpu_t = torch.randint(0, 10, (4, 6), dtype=torch.int64).t() | ||
| 873 | + npu_t = cpu_t.npu() | ||
| 874 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 875 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 876 | + | ||
| 877 | + # ---------------------------------------------------------------- | ||
| 878 | + # Scenario 14: Non-contiguous with explicit dtype override | ||
| 879 | + # ---------------------------------------------------------------- | ||
| 880 | + def test_transposed_dtype_override(self): | ||
| 881 | + """Non-contiguous tensor with dtype override in empty_like.""" | ||
| 882 | + cpu_t = torch.randn(4, 6).t() | ||
| 883 | + npu_t = cpu_t.npu() | ||
| 884 | + | ||
| 885 | + cpu_result = torch.empty_like( | ||
| 886 | + cpu_t, dtype=torch.float32, memory_format=torch.preserve_format | ||
| 887 | + ) | ||
| 888 | + npu_result = torch.empty_like( | ||
| 889 | + npu_t, dtype=torch.float32, memory_format=torch.preserve_format | ||
| 890 | + ) | ||
| 891 | + | ||
| 892 | + self.assertEqual(cpu_result.shape, npu_result.shape) | ||
| 893 | + self.assertEqual( | ||
| 894 | + cpu_result.stride(), | ||
| 895 | + npu_result.stride(), | ||
| 896 | + f"Stride mismatch with dtype override: CPU {cpu_result.stride()} vs NPU {npu_result.stride()}", | ||
| 897 | + ) | ||
| 898 | + self.assertEqual(cpu_result.dtype, npu_result.dtype) | ||
| 899 | + | ||
| 900 | + # ---------------------------------------------------------------- | ||
| 901 | + # Scenario 15: 5D non-contiguous tensor | ||
| 902 | + # ---------------------------------------------------------------- | ||
| 903 | + def test_5d_transposed(self): | ||
| 904 | + """5D non-contiguous tensor via transpose.""" | ||
| 905 | + cpu_t = torch.randn(2, 3, 4, 5, 6).transpose(1, 3) | ||
| 906 | + npu_t = cpu_t.npu() | ||
| 907 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 908 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 909 | + | ||
| 910 | + def test_function_torch_empty_strided(self): | ||
| 911 | + x = torch.empty_strided((2, 3), (1, 2), dtype=torch.int8, device="npu") # noqa: F841 | ||
| 912 | + | ||
| 913 | + def test_function_tensor_new_empty(self): | ||
| 914 | + x = torch.ones(()).npu() | ||
| 915 | + x_new_empty = x.new_empty((2, 3), dtype=torch.float16, device="npu") | ||
| 916 | + res = x_new_empty + 1 | ||
| 917 | + x_new_empty = x.new_empty(size=(2, 3), dtype=torch.float16, device="npu") | ||
| 918 | + res = x_new_empty + 1 # noqa: F841 | ||
| 919 | + | ||
| 920 | + def test_function_tensor_new_empty_strided(self): | ||
| 921 | + x = torch.ones(()).npu() | ||
| 922 | + x_new = x.new_empty_strided([2, 3], [3, 1], dtype=torch.float32, device="npu") | ||
| 923 | + res = x_new + 1 # noqa: F841 | ||
| 924 | + | ||
| 925 | + def test_function_tensor_data_npu(self): | ||
| 926 | + x = torch.ones(()) | ||
| 927 | + x.data = x.data.npu() | ||
| 928 | + | ||
| 929 | + def test_function_tensor_new_full(self): | ||
| 930 | + x_cpu = torch.tensor((), dtype=torch.float32) | ||
| 931 | + cpu_out = x_cpu.new_full((2, 3), 3.1) | ||
| 932 | + | ||
| 933 | + x = torch.tensor((), dtype=torch.float32).npu() | ||
| 934 | + npu_output1 = x.new_full((2, 3), 3.1, device=None, requires_grad=False) | ||
| 935 | + npu_output2 = x.new_full((2, 3), 3.1, device="cpu", requires_grad=False) | ||
| 936 | + npu_output3 = x.new_full((2, 3), 3.1, device="npu", requires_grad=False) | ||
| 937 | + self.assertRtolEqual(cpu_out.numpy(), npu_output1.cpu().numpy()) | ||
| 938 | + self.assertRtolEqual(cpu_out.numpy(), npu_output2.cpu().numpy()) | ||
| 939 | + self.assertRtolEqual(cpu_out.numpy(), npu_output3.cpu().numpy()) | ||
| 940 | + | ||
| 941 | + def test_function_tensor_new_ones(self): | ||
| 942 | + x_cpu = torch.tensor((), dtype=torch.float32) | ||
| 943 | + cpu_out = x_cpu.new_ones((2, 3)) | ||
| 944 | + | ||
| 945 | + x = torch.tensor((), dtype=torch.float32).npu() | ||
| 946 | + npu_output1 = x.new_ones((2, 3), device=None, requires_grad=False) | ||
| 947 | + npu_output2 = x.new_ones((2, 3), device="cpu", requires_grad=False) | ||
| 948 | + npu_output3 = x.new_ones((2, 3), device="npu", requires_grad=False) | ||
| 949 | + npu_output4 = x.new_ones(size=(2, 3), device="npu", requires_grad=False) | ||
| 950 | + self.assertRtolEqual(cpu_out.numpy(), npu_output1.cpu().numpy()) | ||
| 951 | + self.assertRtolEqual(cpu_out.numpy(), npu_output2.cpu().numpy()) | ||
| 952 | + self.assertRtolEqual(cpu_out.numpy(), npu_output3.cpu().numpy()) | ||
| 953 | + self.assertRtolEqual(cpu_out.numpy(), npu_output4.cpu().numpy()) | ||
| 954 | + | ||
| 955 | + def test_function_tensor_new_tensor(self): | ||
| 956 | + x_cpu = torch.tensor((), dtype=torch.float32) | ||
| 957 | + x = torch.tensor((), dtype=torch.float32).npu() | ||
| 958 | + | ||
| 959 | + list_input = [[1, 2, 3], [4, 5, 6]] | ||
| 960 | + cpu_out = x_cpu.new_tensor(list_input) | ||
| 961 | + npu_output1 = x.new_tensor(list_input, device=None, requires_grad=False) | ||
| 962 | + npu_output2 = x.new_tensor(list_input, device="cpu", requires_grad=False) | ||
| 963 | + npu_output3 = x.new_tensor(list_input, device="npu", requires_grad=False) | ||
| 964 | + self.assertRtolEqual(cpu_out.numpy(), npu_output1.cpu().numpy()) | ||
| 965 | + self.assertRtolEqual(cpu_out.numpy(), npu_output2.cpu().numpy()) | ||
| 966 | + print(cpu_out.numpy().dtype, npu_output3.cpu().numpy().dtype) | ||
| 967 | + self.assertRtolEqual(cpu_out.numpy(), npu_output3.cpu().numpy()) | ||
| 968 | + | ||
| 969 | + np_input = np.array(list_input) | ||
| 970 | + cpu_out = x_cpu.new_tensor(np_input) | ||
| 971 | + npu_output1 = x.new_tensor(np_input, device=None, requires_grad=False) | ||
| 972 | + npu_output2 = x.new_tensor(np_input, device="cpu", requires_grad=False) | ||
| 973 | + npu_output3 = x.new_tensor(np_input, device="npu", requires_grad=False) | ||
| 974 | + self.assertRtolEqual(cpu_out.numpy(), npu_output1.cpu().numpy()) | ||
| 975 | + self.assertRtolEqual(cpu_out.numpy(), npu_output2.cpu().numpy()) | ||
| 976 | + self.assertRtolEqual(cpu_out.numpy(), npu_output3.cpu().numpy()) | ||
| 977 | + | ||
| 978 | + tensor_input = torch.tensor(list_input) | ||
| 979 | + cpu_out = x_cpu.new_tensor(tensor_input) | ||
| 980 | + npu_output1 = x.new_tensor(tensor_input, device=None, requires_grad=False) | ||
| 981 | + npu_output2 = x.new_tensor(tensor_input, device="cpu", requires_grad=False) | ||
| 982 | + npu_output3 = x.new_tensor(tensor_input, device="npu", requires_grad=False) | ||
| 983 | + self.assertRtolEqual(cpu_out.numpy(), npu_output1.cpu().numpy()) | ||
| 984 | + self.assertRtolEqual(cpu_out.numpy(), npu_output2.cpu().numpy()) | ||
| 985 | + self.assertRtolEqual(cpu_out.numpy(), npu_output3.cpu().numpy()) | ||
| 986 | + | ||
| 987 | + def test_function_tensor_new_zeros(self): | ||
| 988 | + x_cpu = torch.tensor((), dtype=torch.float32) | ||
| 989 | + cpu_out = x_cpu.new_zeros((2, 3)) | ||
| 990 | + | ||
| 991 | + x = torch.tensor((), dtype=torch.float32).npu() | ||
| 992 | + npu_output1 = x.new_zeros((2, 3), device=None, requires_grad=False) | ||
| 993 | + npu_output2 = x.new_zeros((2, 3), device="cpu", requires_grad=False) | ||
| 994 | + npu_output3 = x.new_zeros((2, 3), device="npu", requires_grad=False) | ||
| 995 | + npu_output4 = x.new_zeros(size=(2, 3), device="npu", requires_grad=False) | ||
| 996 | + self.assertRtolEqual(cpu_out.numpy(), npu_output1.cpu().numpy()) | ||
| 997 | + self.assertRtolEqual(cpu_out.numpy(), npu_output2.cpu().numpy()) | ||
| 998 | + self.assertRtolEqual(cpu_out.numpy(), npu_output3.cpu().numpy()) | ||
| 999 | + self.assertRtolEqual(cpu_out.numpy(), npu_output4.cpu().numpy()) | ||
| 1000 | + | ||
| 1001 | + def test_type_conversions_npu(self): | ||
| 1002 | + x = torch.randn(5, 5) | ||
| 1003 | + self.assertIsInstance(x.float(), torch.FloatTensor) | ||
| 1004 | + self.assertIsInstance(x.double().npu(), torch.npu.DoubleTensor) | ||
| 1005 | + self.assertIsInstance(x.npu().float(), torch.npu.FloatTensor) | ||
| 1006 | + self.assertIsInstance(x.npu().float().cpu(), torch.FloatTensor) | ||
| 1007 | + self.assertIsInstance(x.npu().float().cpu().int(), torch.IntTensor) | ||
| 1008 | + | ||
| 1009 | + def _test_type_conversion_backward(self, t): | ||
| 1010 | + fvar = Variable(t(torch.randn(5, 5).float()), requires_grad=True) | ||
| 1011 | + fvar.double().sum().backward() | ||
| 1012 | + self.assertEqual(fvar.grad, torch.ones_like(fvar)) | ||
| 1013 | + self.assertEqual(type(fvar.grad), type(fvar)) | ||
| 1014 | + dvar = Variable(t(torch.randn(5, 5).double()), requires_grad=True) | ||
| 1015 | + dvar.float().sum().backward() | ||
| 1016 | + self.assertEqual(dvar.grad, torch.ones_like(dvar)) | ||
| 1017 | + self.assertEqual(type(dvar.grad), type(dvar)) | ||
| 1018 | + | ||
| 1019 | + def test_type_conversions(self): | ||
| 1020 | + x = torch.randn(5, 5) | ||
| 1021 | + self.assertIsInstance(x.float(), torch.FloatTensor) | ||
| 1022 | + self.assertIsInstance(x.int(), torch.IntTensor) | ||
| 1023 | + if torch.npu.is_available(): | ||
| 1024 | + self.assertIsInstance(x.float().npu(), torch.npu.FloatTensor) | ||
| 1025 | + self.assertIsInstance(x.double().npu(), torch.npu.DoubleTensor) | ||
| 1026 | + self.assertIsInstance(x.int().npu(), torch.npu.IntTensor) | ||
| 1027 | + self.assertIsInstance(x.int().npu().cpu(), torch.IntTensor) | ||
| 1028 | + | ||
| 1029 | + tensor_types = [ | ||
| 1030 | + torch.DoubleTensor, | ||
| 1031 | + torch.FloatTensor, | ||
| 1032 | + torch.IntTensor, | ||
| 1033 | + torch.ByteTensor, | ||
| 1034 | + ] | ||
| 1035 | + for t, y_var in product(tensor_types, (True, False)): | ||
| 1036 | + y = torch.randint(5, (5, 5), dtype=t.dtype) | ||
| 1037 | + y = Variable(y) if y_var else y | ||
| 1038 | + self.assertIsInstance(x.type(t), t) | ||
| 1039 | + self.assertIsInstance(x.type_as(y), t) | ||
| 1040 | + | ||
| 1041 | + t_dtype = t().dtype | ||
| 1042 | + self.assertIsInstance(x.type(t_dtype), t) | ||
| 1043 | + self.assertIs(t_dtype, x.type(t_dtype).dtype) | ||
| 1044 | + self.assertEqual(y.data_ptr(), y.type(t).data_ptr()) | ||
| 1045 | + | ||
| 1046 | + self._test_type_conversion_backward(lambda x: x) | ||
| 1047 | + if torch.npu.is_available(): | ||
| 1048 | + self._test_type_conversion_backward(lambda x: x.npu()) | ||
| 1049 | + | ||
| 1050 | + def test_get_allocator_backend(self): | ||
| 1051 | + npu_allocator_name = torch.npu.get_allocator_backend() | ||
| 1052 | + self.assertEqual(npu_allocator_name, "native") | ||
| 1053 | + | ||
| 1054 | + def test_contiguous(self): | ||
| 1055 | + def run_once(): | ||
| 1056 | + x = torch.randn(4, 3, 8, 8).npu() | ||
| 1057 | + x = x.permute(0, 2, 1, 3) | ||
| 1058 | + x = x.contiguous() | ||
| 1059 | + return x | ||
| 1060 | + | ||
| 1061 | + with torch._subclasses.fake_tensor.FakeTensorMode(): | ||
| 1062 | + y = run_once() | ||
| 1063 | + self.assertTrue(y.is_contiguous()) | ||
| 1064 | + | ||
| 1065 | + x = torch.randn(1, 16, 5, 5).npu() | ||
| 1066 | + self.assertTrue(x.is_contiguous()) | ||
| 1067 | + stride = list(x.stride()) | ||
| 1068 | + stride[0] = 20 | ||
| 1069 | + # change the stride in dimension 0. the tensor is still contiguous because size[0] is 1 | ||
| 1070 | + x.set_(x.storage(), 0, x.size(), stride) | ||
| 1071 | + self.assertTrue(x.is_contiguous()) | ||
| 1072 | + | ||
| 1073 | + x.contiguous(memory_format=torch.contiguous_format) | ||
| 1074 | + x.contiguous(memory_format=torch.preserve_format) | ||
| 1075 | + | ||
| 1076 | + with self.assertRaisesRegex(RuntimeError, "ERR01007 OPS feature not supported"): | ||
| 1077 | + x.contiguous(memory_format=torch.channels_last) | ||
| 1078 | + | ||
| 1079 | + with self.assertRaisesRegex(RuntimeError, "ERR01007 OPS feature not supported"): | ||
| 1080 | + x.contiguous(memory_format=torch.channels_last_3d) | ||
| 1081 | + | ||
| 1082 | + def test_uuid(self): | ||
| 1083 | + uuid = torch.npu.get_device_properties(0).uuid | ||
| 1084 | + self.assertEqual(len(str(uuid)), 36) | ||
| 1085 | + self.assertEqual(len(uuid.bytes), 16) | ||
| 1086 | + | ||
| 1087 | + | ||
| 1088 | +if __name__ == "__main__": | ||
| 1089 | + run_tests() | ||
| @@ -4,14 +4,14 @@ import doctest | |||
| 4 | import functools | 4 | import functools |
| 5 | import importlib | 5 | import importlib |
| 6 | import inspect | 6 | import inspect |
| 7 | -import itertools | ||
| 8 | import math | 7 | import math |
| 9 | import os | 8 | import os |
| 10 | import re | 9 | import re |
| 11 | import subprocess | 10 | import subprocess |
| 12 | import sys | 11 | import sys |
| 13 | import unittest.mock | 12 | import unittest.mock |
| 14 | -from typing import Any, Callable, Iterator, List, Tuple | 13 | +from typing import Any |
| 14 | +from collections.abc import Callable, Iterator | ||
| 15 | import operator | 15 | import operator |
| 16 | import torch | 16 | import torch |
| 17 | import torch_npu | 17 | import torch_npu |
| @@ -22,7 +22,7 @@ from torch.testing._internal.common_utils import \ | |||
| 22 | parametrize, subtest, instantiate_parametrized_tests, dtype_name, TEST_WITH_ROCM, decorateIf) | 22 | parametrize, subtest, instantiate_parametrized_tests, dtype_name, TEST_WITH_ROCM, decorateIf) |
| 23 | from torch.testing._internal.common_device_type import \ | 23 | from torch.testing._internal.common_device_type import \ |
| 24 | (PYTORCH_TESTING_DEVICE_EXCEPT_FOR_KEY, PYTORCH_TESTING_DEVICE_ONLY_FOR_KEY, dtypes, | 24 | (PYTORCH_TESTING_DEVICE_EXCEPT_FOR_KEY, PYTORCH_TESTING_DEVICE_ONLY_FOR_KEY, dtypes, |
| 25 | - get_device_type_test_bases, instantiate_device_type_tests, onlyCPU, onlyCUDA, onlyNativeDeviceTypes, | 25 | + get_device_type_test_bases, instantiate_device_type_tests, onlyCUDA, onlyNativeDeviceTypes, |
| 26 | deviceCountAtLeast, ops, expectedFailureMeta, OpDTypes) | 26 | deviceCountAtLeast, ops, expectedFailureMeta, OpDTypes) |
| 27 | from torch.testing._internal.common_methods_invocations import op_db | 27 | from torch.testing._internal.common_methods_invocations import op_db |
| 28 | from torch.testing._internal import opinfo | 28 | from torch.testing._internal import opinfo |
| @@ -30,8 +30,9 @@ from torch.testing._internal.common_dtype import all_types_and_complex_and, floa | |||
| 30 | from torch.testing._internal.common_modules import modules, module_db, ModuleInfo | 30 | from torch.testing._internal.common_modules import modules, module_db, ModuleInfo |
| 31 | from torch.testing._internal.opinfo.core import SampleInput, DecorateInfo, OpInfo | 31 | from torch.testing._internal.opinfo.core import SampleInput, DecorateInfo, OpInfo |
| 32 | 32 | ||
| 33 | +DEVICE_NAME = torch_npu.npu.get_device_name(0)[:10] | ||
| 33 | 34 | ||
| 34 | -# For testing TestCase methods and torch.testing functions | 35 | +# For testing TestCase methods and torch.testing functions |
| 35 | class TestTesting(TestCase): | 36 | class TestTesting(TestCase): |
| 36 | # Ensure that assertEqual handles numpy arrays properly | 37 | # Ensure that assertEqual handles numpy arrays properly |
| 37 | 38 | ||
| @@ -454,7 +455,7 @@ if __name__ == '__main__': | |||
| 454 | # Test without setting env var should run everything. | 455 | # Test without setting env var should run everything. |
| 455 | env = dict(os.environ) | 456 | env = dict(os.environ) |
| 456 | for k in ['CI', PYTORCH_TESTING_DEVICE_ONLY_FOR_KEY, PYTORCH_TESTING_DEVICE_EXCEPT_FOR_KEY]: | 457 | for k in ['CI', PYTORCH_TESTING_DEVICE_ONLY_FOR_KEY, PYTORCH_TESTING_DEVICE_EXCEPT_FOR_KEY]: |
| 457 | - if k in env.keys(): | 458 | + if k in env: |
| 458 | del env[k] | 459 | del env[k] |
| 459 | _, stderr = TestCase.run_process_no_exception(test_filter_file_template, env=env) | 460 | _, stderr = TestCase.run_process_no_exception(test_filter_file_template, env=env) |
| 460 | self.assertIn(f'Ran {test_bases_count} test', stderr.decode('ascii')) | 461 | self.assertIn(f'Ran {test_bases_count} test', stderr.decode('ascii')) |
| @@ -476,7 +477,7 @@ if __name__ == '__main__': | |||
| 476 | self.assertNotIn('OK', stderr.decode('ascii')) | 477 | self.assertNotIn('OK', stderr.decode('ascii')) |
| 477 | 478 | ||
| 478 | 479 | ||
| 479 | -def make_assert_close_inputs(actual: Any, expected: Any) -> List[Tuple[Any, Any]]: | 480 | +def make_assert_close_inputs(actual: Any, expected: Any) -> list[tuple[Any, Any]]: |
| 480 | """Makes inputs for :func:`torch.testing.assert_close` functions based on two examples. | 481 | """Makes inputs for :func:`torch.testing.assert_close` functions based on two examples. |
| 481 | 482 | ||
| 482 | Args: | 483 | Args: |
| @@ -794,7 +795,6 @@ class TestAssertClose(TestCase): | |||
| 794 | the test should mock a component to raise this instead of the regular behavior. We avoid using a builtin | 795 | the test should mock a component to raise this instead of the regular behavior. We avoid using a builtin |
| 795 | exception here to avoid triggering possible handling of them. | 796 | exception here to avoid triggering possible handling of them. |
| 796 | """ | 797 | """ |
| 797 | - pass | ||
| 798 | 798 | ||
| 799 | 799 | ||
| 800 | def test_unexpected_error_originate(self, _): | 800 | def test_unexpected_error_originate(self, _): |
| @@ -1366,6 +1366,304 @@ class TestAssertCloseQuantized(TestCase): | |||
| 1366 | for fn in assert_close_with_inputs(actual, expected): | 1366 | for fn in assert_close_with_inputs(actual, expected): |
| 1367 | fn() | 1367 | fn() |
| 1368 | 1368 | ||
| 1369 | + | ||
| 1370 | + def setUpClass(cls): | ||
| 1371 | + cls.npu_available = torch.npu.is_available() | ||
| 1372 | + | ||
| 1373 | + | ||
| 1374 | + "Ops clone not support non-contiguous output on A1 device.") | ||
| 1375 | + def _compare_with_cpu(self, cpu_tensor, npu_tensor, memory_format=None): | ||
| 1376 | + cpu_result = cpu_tensor.clone(memory_format=memory_format) | ||
| 1377 | + npu_result = npu_tensor.clone(memory_format=memory_format) | ||
| 1378 | + | ||
| 1379 | + self.assertEqual(cpu_result.shape, npu_result.shape, | ||
| 1380 | + f"Shape mismatch: CPU {cpu_result.shape} vs NPU {npu_result.shape}") | ||
| 1381 | + self.assertEqual(cpu_result.dtype, npu_result.dtype, | ||
| 1382 | + f"Dtype mismatch: CPU {cpu_result.dtype} vs NPU {npu_result.dtype}") | ||
| 1383 | + self.assertEqual(cpu_result.stride(), npu_result.stride(), | ||
| 1384 | + f"Stride mismatch: CPU {cpu_result.stride()} vs NPU {npu_result.stride()}") | ||
| 1385 | + self.assertEqual(cpu_result.is_contiguous(), npu_result.is_contiguous(), | ||
| 1386 | + f"Contiguity mismatch: CPU contiguous={cpu_result.is_contiguous()}" | ||
| 1387 | + f" vs NPU contiguous={npu_result.is_contiguous()}") | ||
| 1388 | + | ||
| 1389 | + cpu_data = cpu_result.cpu() | ||
| 1390 | + npu_data = npu_result.cpu() | ||
| 1391 | + self.assertTrue(torch.equal(cpu_data, npu_data), | ||
| 1392 | + "Data mismatch: clone result differs between CPU and NPU") | ||
| 1393 | + | ||
| 1394 | + # ---------------------------------------------------------------- | ||
| 1395 | + # Scenario 1: Transposed 2D tensor | ||
| 1396 | + # is_non_overlapping_and_dense=True → empty_strided path | ||
| 1397 | + # ---------------------------------------------------------------- | ||
| 1398 | + def test_transposed_2d(self): | ||
| 1399 | + """Transposed 2D: non_overlapping_and_dense=True, strides preserved.""" | ||
| 1400 | + cpu_t = torch.randn(4, 6).t() | ||
| 1401 | + npu_t = cpu_t.npu() | ||
| 1402 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 1403 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1404 | + | ||
| 1405 | + # ---------------------------------------------------------------- | ||
| 1406 | + # Scenario 2: Transposed 3D tensor | ||
| 1407 | + # is_non_overlapping_and_dense=True → empty_strided path | ||
| 1408 | + # ---------------------------------------------------------------- | ||
| 1409 | + def test_transposed_3d(self): | ||
| 1410 | + """Transposed 3D: non_overlapping_and_dense=True, strides preserved.""" | ||
| 1411 | + cpu_t = torch.randn(2, 3, 4).transpose(0, 2) | ||
| 1412 | + npu_t = cpu_t.npu() | ||
| 1413 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 1414 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1415 | + | ||
| 1416 | + # ---------------------------------------------------------------- | ||
| 1417 | + # Scenario 3: Sliced tensor (every other row) | ||
| 1418 | + # is_non_overlapping_and_dense=True → empty_strided path | ||
| 1419 | + # ---------------------------------------------------------------- | ||
| 1420 | + def test_sliced_rows(self): | ||
| 1421 | + """Slice rows [::2]: non_overlapping_and_dense=True, strides preserved.""" | ||
| 1422 | + cpu_t = torch.randn(8, 5)[::2] | ||
| 1423 | + npu_t = cpu_t.npu() | ||
| 1424 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 1425 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1426 | + | ||
| 1427 | + # ---------------------------------------------------------------- | ||
| 1428 | + # Scenario 4: Sliced tensor (every other column) | ||
| 1429 | + # is_non_overlapping_and_dense=True → empty_strided path | ||
| 1430 | + # ---------------------------------------------------------------- | ||
| 1431 | + def test_sliced_cols(self): | ||
| 1432 | + """Slice columns [:, ::2]: non_overlapping_and_dense=True, strides preserved.""" | ||
| 1433 | + cpu_t = torch.randn(4, 8)[:, ::2] | ||
| 1434 | + npu_t = cpu_t.npu() | ||
| 1435 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 1436 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1437 | + | ||
| 1438 | + # ---------------------------------------------------------------- | ||
| 1439 | + # Scenario 5: Narrowed tensor | ||
| 1440 | + # is_non_overlapping_and_dense=True → empty_strided path | ||
| 1441 | + # ---------------------------------------------------------------- | ||
| 1442 | + def test_narrowed(self): | ||
| 1443 | + """Narrowed tensor: non_overlapping_and_dense=True, strides preserved.""" | ||
| 1444 | + cpu_t = torch.randn(6, 8).narrow(1, 1, 5) | ||
| 1445 | + npu_t = cpu_t.npu() | ||
| 1446 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 1447 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1448 | + | ||
| 1449 | + # ---------------------------------------------------------------- | ||
| 1450 | + # Scenario 6: Expanded tensor (stride=0) | ||
| 1451 | + # is_non_overlapping_and_dense=False → apply_tensor_without_format path | ||
| 1452 | + # ---------------------------------------------------------------- | ||
| 1453 | + def test_expanded_2d(self): | ||
| 1454 | + """Expanded 2D (stride=0): non_overlapping_and_dense=False, falls to else path.""" | ||
| 1455 | + cpu_t = torch.randn(1, 5).expand(4, 5) | ||
| 1456 | + npu_t = cpu_t.npu() | ||
| 1457 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 1458 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1459 | + | ||
| 1460 | + def test_expanded_3d(self): | ||
| 1461 | + """Expanded 3D (stride=0): non_overlapping_and_dense=False.""" | ||
| 1462 | + cpu_t = torch.randn(2, 1, 4).expand(2, 3, 4) | ||
| 1463 | + npu_t = cpu_t.npu() | ||
| 1464 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 1465 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1466 | + | ||
| 1467 | + # ---------------------------------------------------------------- | ||
| 1468 | + # Scenario 7: Permuted tensor | ||
| 1469 | + # is_non_overlapping_and_dense=True → empty_strided path | ||
| 1470 | + # ---------------------------------------------------------------- | ||
| 1471 | + def test_permuted(self): | ||
| 1472 | + """Permuted 4D tensor: non_overlapping_and_dense=True, strides preserved.""" | ||
| 1473 | + cpu_t = torch.randn(2, 3, 4, 5).permute(3, 1, 0, 2) | ||
| 1474 | + npu_t = cpu_t.npu() | ||
| 1475 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 1476 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1477 | + | ||
| 1478 | + # ---------------------------------------------------------------- | ||
| 1479 | + # Scenario 8: Select on non-leading dimension | ||
| 1480 | + # is_non_overlapping_and_dense=True → empty_strided path | ||
| 1481 | + # ---------------------------------------------------------------- | ||
| 1482 | + def test_select_dim1(self): | ||
| 1483 | + """Select dim=1: non_overlapping_and_dense=True, strides preserved.""" | ||
| 1484 | + cpu_t = torch.randn(3, 5, 4).select(1, 2) | ||
| 1485 | + npu_t = cpu_t.npu() | ||
| 1486 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 1487 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1488 | + | ||
| 1489 | + # ---------------------------------------------------------------- | ||
| 1490 | + # Scenario 9: as_strided non-contiguous | ||
| 1491 | + # May or may not be non_overlapping_and_dense depending on strides | ||
| 1492 | + # ---------------------------------------------------------------- | ||
| 1493 | + def test_as_strided_dense(self): | ||
| 1494 | + """as_strided with dense strides: non_overlapping_and_dense=True.""" | ||
| 1495 | + cpu_t = torch.randn(24).as_strided((2, 3, 4), (12, 4, 1)) | ||
| 1496 | + npu_t = cpu_t.npu() | ||
| 1497 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1498 | + | ||
| 1499 | + def test_as_strided_non_dense(self): | ||
| 1500 | + """as_strided with overlapping strides: non_overlapping_and_dense=False.""" | ||
| 1501 | + cpu_t = torch.randn(10).as_strided((3, 3), (3, 1)) | ||
| 1502 | + npu_t = cpu_t.npu() | ||
| 1503 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1504 | + | ||
| 1505 | + # ---------------------------------------------------------------- | ||
| 1506 | + # Scenario 10: 1D non-contiguous via as_strided | ||
| 1507 | + # is_non_overlapping_and_dense=True → empty_strided path | ||
| 1508 | + # ---------------------------------------------------------------- | ||
| 1509 | + def test_1d_non_contiguous(self): | ||
| 1510 | + """1D non-contiguous via as_strided: non_overlapping_and_dense=True.""" | ||
| 1511 | + cpu_t = torch.randn(10).as_strided((4,), (2,)) | ||
| 1512 | + npu_t = cpu_t.npu() | ||
| 1513 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 1514 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1515 | + | ||
| 1516 | + # ---------------------------------------------------------------- | ||
| 1517 | + # Scenario 11: Contiguous input with preserve_format (baseline) | ||
| 1518 | + # is_non_overlapping_and_dense=True → empty_strided path | ||
| 1519 | + # ---------------------------------------------------------------- | ||
| 1520 | + def test_contiguous_preserve(self): | ||
| 1521 | + """Contiguous tensor with preserve_format: should remain identical.""" | ||
| 1522 | + cpu_t = torch.randn(4, 6) | ||
| 1523 | + npu_t = cpu_t.npu() | ||
| 1524 | + self.assertTrue(cpu_t.is_contiguous()) | ||
| 1525 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1526 | + | ||
| 1527 | + # ---------------------------------------------------------------- | ||
| 1528 | + # Scenario 12: memory_format=Contiguous on non-contiguous input | ||
| 1529 | + # → apply_tensor_without_format path (contiguous result) | ||
| 1530 | + # ---------------------------------------------------------------- | ||
| 1531 | + | ||
| 1532 | + "Ops clone not support non-contiguous output on A1 device.") | ||
| 1533 | + def test_transposed_contiguous_format(self): | ||
| 1534 | + """Transposed tensor with contiguous_format: result should be contiguous.""" | ||
| 1535 | + cpu_t = torch.randn(4, 6).t() | ||
| 1536 | + npu_t = cpu_t.npu() | ||
| 1537 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 1538 | + | ||
| 1539 | + cpu_result = cpu_t.clone(memory_format=torch.contiguous_format) | ||
| 1540 | + npu_result = npu_t.clone(memory_format=torch.contiguous_format) | ||
| 1541 | + | ||
| 1542 | + self.assertTrue(cpu_result.is_contiguous()) | ||
| 1543 | + self.assertTrue(npu_result.is_contiguous()) | ||
| 1544 | + self.assertEqual(cpu_result.shape, npu_result.shape) | ||
| 1545 | + self.assertTrue(torch.equal(cpu_result.cpu(), npu_result.cpu())) | ||
| 1546 | + | ||
| 1547 | + # ---------------------------------------------------------------- | ||
| 1548 | + # Scenario 13: memory_format=None (default Preserve) on non-contiguous | ||
| 1549 | + # Same as preserve_format | ||
| 1550 | + # ---------------------------------------------------------------- | ||
| 1551 | + | ||
| 1552 | + "Ops clone not support non-contiguous output on A1 device.") | ||
| 1553 | + def test_transposed_default_format(self): | ||
| 1554 | + """Transposed tensor with default memory_format (Preserve).""" | ||
| 1555 | + cpu_t = torch.randn(4, 6).t() | ||
| 1556 | + npu_t = cpu_t.npu() | ||
| 1557 | + | ||
| 1558 | + cpu_result = cpu_t.clone() | ||
| 1559 | + npu_result = npu_t.clone() | ||
| 1560 | + | ||
| 1561 | + self.assertEqual(cpu_result.stride(), npu_result.stride(), | ||
| 1562 | + f"Stride mismatch: CPU {cpu_result.stride()} vs NPU {npu_result.stride()}") | ||
| 1563 | + self.assertEqual(cpu_result.is_contiguous(), npu_result.is_contiguous()) | ||
| 1564 | + self.assertTrue(torch.equal(cpu_result.cpu(), npu_result.cpu())) | ||
| 1565 | + | ||
| 1566 | + # ---------------------------------------------------------------- | ||
| 1567 | + # Scenario 14: Non-contiguous with different dtypes | ||
| 1568 | + # ---------------------------------------------------------------- | ||
| 1569 | + def test_transposed_float16(self): | ||
| 1570 | + """Non-contiguous float16 tensor clone.""" | ||
| 1571 | + cpu_t = torch.randn(4, 6, dtype=torch.float16).t() | ||
| 1572 | + npu_t = cpu_t.npu() | ||
| 1573 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1574 | + | ||
| 1575 | + def test_transposed_bfloat16(self): | ||
| 1576 | + """Non-contiguous bfloat16 tensor clone.""" | ||
| 1577 | + cpu_t = torch.randn(4, 6, dtype=torch.bfloat16).t() | ||
| 1578 | + npu_t = cpu_t.npu() | ||
| 1579 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1580 | + | ||
| 1581 | + # ---------------------------------------------------------------- | ||
| 1582 | + # Scenario 15: 5D non-contiguous tensor | ||
| 1583 | + # is_non_overlapping_and_dense=True → empty_strided path | ||
| 1584 | + # ---------------------------------------------------------------- | ||
| 1585 | + def test_5d_transposed(self): | ||
| 1586 | + """5D non-contiguous tensor via transpose.""" | ||
| 1587 | + cpu_t = torch.randn(2, 3, 4, 5, 6).transpose(1, 3) | ||
| 1588 | + npu_t = cpu_t.npu() | ||
| 1589 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 1590 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1591 | + | ||
| 1592 | + # ---------------------------------------------------------------- | ||
| 1593 | + # Scenario 16: Double transposed (back to contiguous) | ||
| 1594 | + # is_non_overlapping_and_dense=True → empty_strided path | ||
| 1595 | + # ---------------------------------------------------------------- | ||
| 1596 | + def test_double_transpose(self): | ||
| 1597 | + """Double transpose: logically contiguous, stride order differs from default.""" | ||
| 1598 | + cpu_t = torch.randn(3, 4).t().t() | ||
| 1599 | + npu_t = cpu_t.npu() | ||
| 1600 | + self.assertTrue(cpu_t.is_contiguous()) | ||
| 1601 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1602 | + | ||
| 1603 | + # ---------------------------------------------------------------- | ||
| 1604 | + # Scenario 17: Verify clone is a true deep copy (data independence) | ||
| 1605 | + # ---------------------------------------------------------------- | ||
| 1606 | + | ||
| 1607 | + "Ops clone not support non-contiguous output on A1 device.") | ||
| 1608 | + def test_clone_independence(self): | ||
| 1609 | + """Cloned tensor should not share storage with original.""" | ||
| 1610 | + cpu_t = torch.randn(4, 6).t() | ||
| 1611 | + npu_t = cpu_t.npu() | ||
| 1612 | + | ||
| 1613 | + cpu_clone = cpu_t.clone(memory_format=torch.preserve_format) | ||
| 1614 | + npu_clone = npu_t.clone(memory_format=torch.preserve_format) | ||
| 1615 | + | ||
| 1616 | + cpu_clone.fill_(0) | ||
| 1617 | + npu_clone.fill_(0) | ||
| 1618 | + | ||
| 1619 | + self.assertFalse(torch.equal(cpu_t, cpu_clone), | ||
| 1620 | + "CPU: clone should be independent from original") | ||
| 1621 | + self.assertFalse(torch.equal(npu_t.cpu(), npu_clone.cpu()), | ||
| 1622 | + "NPU: clone should be independent from original") | ||
| 1623 | + | ||
| 1624 | + # ---------------------------------------------------------------- | ||
| 1625 | + # Scenario 18: Expanded then clone (stride=0 → else path) | ||
| 1626 | + # Verify data correctness even when strides are not preserved | ||
| 1627 | + # ---------------------------------------------------------------- | ||
| 1628 | + | ||
| 1629 | + "Ops clone not support non-contiguous output on A1 device.") | ||
| 1630 | + def test_expanded_data_correctness(self): | ||
| 1631 | + """Expanded tensor clone: data must be correct even if strides differ.""" | ||
| 1632 | + cpu_t = torch.randn(1, 5).expand(3, 5) | ||
| 1633 | + npu_t = cpu_t.npu() | ||
| 1634 | + | ||
| 1635 | + cpu_result = cpu_t.clone(memory_format=torch.preserve_format) | ||
| 1636 | + npu_result = npu_t.clone(memory_format=torch.preserve_format) | ||
| 1637 | + | ||
| 1638 | + self.assertTrue(torch.equal(cpu_result.cpu(), npu_result.cpu()), | ||
| 1639 | + "Data mismatch for expanded tensor clone") | ||
| 1640 | + | ||
| 1641 | + # ---------------------------------------------------------------- | ||
| 1642 | + # Scenario 19: Non-contiguous + memory_format=Preserve with | ||
| 1643 | + # is_non_overlapping_and_dense=True | ||
| 1644 | + # Verify strides are exactly preserved (not just "equivalent") | ||
| 1645 | + # ---------------------------------------------------------------- | ||
| 1646 | + | ||
| 1647 | + "Ops clone not support non-contiguous output on A1 device.") | ||
| 1648 | + def test_strides_exactly_preserved(self): | ||
| 1649 | + """Verify clone preserves exact stride values, not just contiguity.""" | ||
| 1650 | + cpu_t = torch.randn(3, 4, 5).transpose(0, 1) | ||
| 1651 | + npu_t = cpu_t.npu() | ||
| 1652 | + expected_stride = cpu_t.stride() | ||
| 1653 | + | ||
| 1654 | + cpu_result = cpu_t.clone(memory_format=torch.preserve_format) | ||
| 1655 | + npu_result = npu_t.clone(memory_format=torch.preserve_format) | ||
| 1656 | + | ||
| 1657 | + self.assertEqual(cpu_result.stride(), expected_stride, | ||
| 1658 | + f"CPU stride not preserved: got {cpu_result.stride()}, expected {expected_stride}") | ||
| 1659 | + self.assertEqual(npu_result.stride(), expected_stride, | ||
| 1660 | + f"NPU stride not preserved: got {npu_result.stride()}, expected {expected_stride}") | ||
| 1661 | + | ||
| 1662 | + def test_slice_dense(self): | ||
| 1663 | + """slice""" | ||
| 1664 | + cpu_t = torch.randn(24).t()[::2] | ||
| 1665 | + npu_t = cpu_t.npu() | ||
| 1666 | + self._compare_with_cpu(cpu_t, npu_t, memory_format=torch.preserve_format) | ||
| 1369 | 1667 | ||
| 1370 | class TestMakeTensor(TestCase): | 1668 | class TestMakeTensor(TestCase): |
| 1371 | supported_dtypes = dtypes( | 1669 | supported_dtypes = dtypes( |
| @@ -1958,8 +2256,10 @@ class TestTestParametrizationDeviceType(TestCase): | |||
| 1958 | for op in op_db: | 2256 | for op in op_db: |
| 1959 | for dtype in op.supported_dtypes(torch.device(device).type): | 2257 | for dtype in op.supported_dtypes(torch.device(device).type): |
| 1960 | for flag_part in ('flag_disabled', 'flag_enabled'): | 2258 | for flag_part in ('flag_disabled', 'flag_enabled'): |
| 1961 | - expected_name = '{}.test_op_parametrized_{}_{}_{}_{}'.format( | 2259 | + expected_name = ( |
| 1962 | - device_cls.__name__, op.formatted_name, flag_part, device, dtype_name(dtype)) | 2260 | + f'{device_cls.__name__}.test_op_parametrized_' |
| 2261 | + f'{op.formatted_name}_{flag_part}_{device}_{dtype_name(dtype)}' | ||
| 2262 | + ) | ||
| 1963 | expected_test_names.append(expected_name) | 2263 | expected_test_names.append(expected_name) |
| 1964 | 2264 | ||
| 1965 | test_names = _get_test_names_for_test_class(device_cls) | 2265 | test_names = _get_test_names_for_test_class(device_cls) |
| @@ -2302,15 +2602,15 @@ class TestOpInfos(TestCase): | |||
| 2302 | 2602 | ||
| 2303 | # Construction with natural syntax | 2603 | # Construction with natural syntax |
| 2304 | s = SampleInput(a, b, c, d=d, e=e) | 2604 | s = SampleInput(a, b, c, d=d, e=e) |
| 2305 | - assert s.input is a | 2605 | + self.assertIs(s.input, a) |
| 2306 | - assert s.args == (b, c) | 2606 | + self.assertEqual(s.args, (b, c)) |
| 2307 | - assert s.kwargs == dict(d=d, e=e) | 2607 | + self.assertEqual(s.kwargs, dict(d=d, e=e)) |
| 2308 | 2608 | ||
| 2309 | # Construction with explicit args and kwargs | 2609 | # Construction with explicit args and kwargs |
| 2310 | s = SampleInput(a, args=(b,), kwargs=dict(c=c, d=d, e=e)) | 2610 | s = SampleInput(a, args=(b,), kwargs=dict(c=c, d=d, e=e)) |
| 2311 | - assert s.input is a | 2611 | + self.assertIs(s.input, a) |
| 2312 | - assert s.args == (b,) | 2612 | + self.assertEqual(s.args, (b,)) |
| 2313 | - assert s.kwargs == dict(c=c, d=d, e=e) | 2613 | + self.assertEqual(s.kwargs, dict(c=c, d=d, e=e)) |
| 2314 | 2614 | ||
| 2315 | # Construction with a mixed form will error | 2615 | # Construction with a mixed form will error |
| 2316 | with self.assertRaises(AssertionError): | 2616 | with self.assertRaises(AssertionError): |
| @@ -2338,8 +2638,8 @@ class TestOpInfos(TestCase): | |||
| 2338 | # But when only input is given, metadata is allowed for backward | 2638 | # But when only input is given, metadata is allowed for backward |
| 2339 | # compatibility | 2639 | # compatibility |
| 2340 | s = SampleInput(a, broadcasts_input=True) | 2640 | s = SampleInput(a, broadcasts_input=True) |
| 2341 | - assert s.input is a | 2641 | + self.assertIs(s.input, a) |
| 2342 | - assert s.broadcasts_input | 2642 | + self.assertTrue(s.broadcasts_input) |
| 2343 | 2643 | ||
| 2344 | def test_sample_input_metadata(self) -> None: | 2644 | def test_sample_input_metadata(self) -> None: |
| 2345 | a, b = (object() for _ in range(2)) | 2645 | a, b = (object() for _ in range(2)) |
| @@ -0,0 +1,321 @@ | |||
| 1 | +import os | ||
| 2 | +import unittest | ||
| 3 | +import torch | ||
| 4 | +import torch_npu | ||
| 5 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 6 | +from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU | ||
| 7 | + | ||
| 8 | +os.environ["PYTORCH_NPU_ALLOC_CONF"] = 'expandable_segments:True' | ||
| 9 | + | ||
| 10 | +class TestTensorToPreserveFormat(TestCase): | ||
| 11 | + def _make_non_contiguous_tensors(self): | ||
| 12 | + """ | ||
| 13 | + Generate various types of non-contiguous tensors, returns list: | ||
| 14 | + (name, cpu_tensor, is_non_overlapping_and_dense) | ||
| 15 | + """ | ||
| 16 | + tensors = [] | ||
| 17 | + # 1. transpose 2D | ||
| 18 | + t = torch.randn(4, 6).t() | ||
| 19 | + tensors.append(("transpose_2d", t, True)) | ||
| 20 | + # 2. permute 4D | ||
| 21 | + t = torch.randn(2, 3, 4, 5).permute(3, 1, 0, 2) | ||
| 22 | + tensors.append(("permute_4d", t, True)) | ||
| 23 | + # 3. slice dim0 (step > 1) | ||
| 24 | + t = torch.randn(8, 5)[::2] | ||
| 25 | + tensors.append(("slice_dim0", t, True)) | ||
| 26 | + # 4. slice dim1 (step > 1) | ||
| 27 | + t = torch.randn(4, 8)[:, ::2] | ||
| 28 | + tensors.append(("slice_dim1", t, True)) | ||
| 29 | + # 5. narrow | ||
| 30 | + t = torch.randn(6, 8).narrow(1, 1, 5) | ||
| 31 | + tensors.append(("narrow", t, True)) | ||
| 32 | + # 6. select (reduces dim) | ||
| 33 | + t = torch.randn(3, 5, 4).select(1, 2) | ||
| 34 | + tensors.append(("select_dim1", t, True)) | ||
| 35 | + # 7. 1D non-contiguous (as_strided stride > 1) | ||
| 36 | + t = torch.randn(10).as_strided((4,), (2,)) | ||
| 37 | + tensors.append(("1d_as_strided", t, True)) | ||
| 38 | + # 8. expand 2D (stride=0) | ||
| 39 | + t = torch.randn(1, 5).expand(4, 5) | ||
| 40 | + tensors.append(("expand_2d", t, False)) | ||
| 41 | + # 9. expand 3D (stride=0) | ||
| 42 | + t = torch.randn(2, 1, 4).expand(2, 3, 4) | ||
| 43 | + tensors.append(("expand_3d", t, False)) | ||
| 44 | + # 10. as_strided with overlap (elements overlap, non-overlapping=False) | ||
| 45 | + t = torch.randn(12).as_strided((3, 3), (4, 1)) | ||
| 46 | + tensors.append(("as_strided_overlap", t, False)) | ||
| 47 | + | ||
| 48 | + return tensors | ||
| 49 | + | ||
| 50 | + def _verify_preserve_format_result(self, src, result, scenario_name): | ||
| 51 | + """Verify basic correctness of preserve_format result""" | ||
| 52 | + self.assertEqual(list(src.shape), list(result.shape)) | ||
| 53 | + | ||
| 54 | + def _verify_stride_preserved(self, src, result, scenario_name, is_nod): | ||
| 55 | + """ | ||
| 56 | + Verify stride behavior: | ||
| 57 | + - is_non_overlapping_and_dense=True -> stride should be preserved | ||
| 58 | + - is_non_overlapping_and_dense=False -> stride may differ (fallback to contiguous) | ||
| 59 | + """ | ||
| 60 | + if is_nod: | ||
| 61 | + self.assertEqual(src.stride(), result.stride()) | ||
| 62 | + else: | ||
| 63 | + pass | ||
| 64 | + | ||
| 65 | + # ================================================================ | ||
| 66 | + # H2D scenario: CPU -> NPU | ||
| 67 | + # ================================================================ | ||
| 68 | + | ||
| 69 | + def test_h2d_transpose_2d(self): | ||
| 70 | + """H2D: transpose 2D non-contiguous -> preserve_format preserves stride""" | ||
| 71 | + cpu_t = torch.randn(4, 6).t() | ||
| 72 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 73 | + | ||
| 74 | + | ||
| 75 | + npu_t = cpu_t.to("npu", memory_format=torch.preserve_format) | ||
| 76 | + self._verify_preserve_format_result(cpu_t, npu_t, "H2D-transpose_2d") | ||
| 77 | + self.assertEqual(str(npu_t.device), "npu:0") | ||
| 78 | + self.assertEqual(cpu_t.stride(), npu_t.stride(), | ||
| 79 | + "H2D-transpose_2d: stride not preserved") | ||
| 80 | + | ||
| 81 | + def test_h2d_permute_4d(self): | ||
| 82 | + """H2D: permute 4D non-contiguous -> preserve_format preserves stride""" | ||
| 83 | + cpu_t = torch.randn(2, 3, 4, 5).permute(3, 1, 0, 2) | ||
| 84 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 85 | + npu_t = cpu_t.to("npu", memory_format=torch.preserve_format) | ||
| 86 | + self._verify_preserve_format_result(cpu_t, npu_t, "H2D-permute_4d") | ||
| 87 | + self.assertEqual(cpu_t.stride(), npu_t.stride(), | ||
| 88 | + "H2D-permute_4d: stride not preserved") | ||
| 89 | + | ||
| 90 | + def _get_dense_strides(self, src): | ||
| 91 | + """Get dense storage strides""" | ||
| 92 | + dummy = torch.empty_like(src) | ||
| 93 | + return dummy.stride() | ||
| 94 | + | ||
| 95 | + def test_h2d_slice_dim0(self): | ||
| 96 | + """H2D: slice dim0 non-contiguous -> preserve_format preserves stride""" | ||
| 97 | + cpu_t = torch.randn(8, 5)[::2] | ||
| 98 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 99 | + | ||
| 100 | + | ||
| 101 | + npu_t = cpu_t.to("npu", memory_format=torch.preserve_format) | ||
| 102 | + expec_strides = self._get_dense_strides(cpu_t) | ||
| 103 | + self.assertEqual(expec_strides, npu_t.stride(), | ||
| 104 | + "H2D-slice_dim0: stride not preserved") | ||
| 105 | + | ||
| 106 | + def test_h2d_slice_dim1(self): | ||
| 107 | + """H2D: slice dim1 non-contiguous -> preserve_format preserves stride""" | ||
| 108 | + cpu_t = torch.randn(4, 8)[:, ::2] | ||
| 109 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 110 | + | ||
| 111 | + | ||
| 112 | + npu_t = cpu_t.to("npu", memory_format=torch.preserve_format) | ||
| 113 | + self._verify_preserve_format_result(cpu_t, npu_t, "H2D-slice_dim1") | ||
| 114 | + expec_strides = self._get_dense_strides(cpu_t) | ||
| 115 | + self.assertEqual(expec_strides, npu_t.stride(), | ||
| 116 | + "H2D-slice_dim1: stride not preserved") | ||
| 117 | + | ||
| 118 | + def test_h2d_narrow(self): | ||
| 119 | + """H2D: narrow non-contiguous -> preserve_format preserves stride""" | ||
| 120 | + cpu_t = torch.randn(6, 8).narrow(1, 1, 5) | ||
| 121 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 122 | + | ||
| 123 | + | ||
| 124 | + npu_t = cpu_t.to("npu", memory_format=torch.preserve_format) | ||
| 125 | + self._verify_preserve_format_result(cpu_t, npu_t, "H2D-narrow") | ||
| 126 | + expec_strides = self._get_dense_strides(cpu_t) | ||
| 127 | + self.assertEqual(expec_strides, npu_t.stride(), | ||
| 128 | + "H2D-narrow: stride not preserved") | ||
| 129 | + | ||
| 130 | + def test_h2d_select(self): | ||
| 131 | + """H2D: select (reduces dim) non-contiguous -> preserve_format preserves stride""" | ||
| 132 | + cpu_t = torch.randn(3, 5, 4).select(1, 2) | ||
| 133 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 134 | + | ||
| 135 | + npu_t = cpu_t.to("npu", memory_format=torch.preserve_format) | ||
| 136 | + self._verify_preserve_format_result(cpu_t, npu_t, "H2D-select") | ||
| 137 | + expec_strides = self._get_dense_strides(cpu_t) | ||
| 138 | + self.assertEqual(expec_strides, npu_t.stride(), | ||
| 139 | + "H2D-select: stride not preserved") | ||
| 140 | + | ||
| 141 | + def test_h2d_1d_as_strided(self): | ||
| 142 | + """H2D: 1D as_strided non-contiguous -> preserve_format preserves stride""" | ||
| 143 | + cpu_t = torch.randn(10).as_strided((4,), (2,)) | ||
| 144 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 145 | + | ||
| 146 | + npu_t = cpu_t.to("npu", memory_format=torch.preserve_format) | ||
| 147 | + self._verify_preserve_format_result(cpu_t, npu_t, "H2D-1d_as_strided") | ||
| 148 | + expec_strides = self._get_dense_strides(cpu_t) | ||
| 149 | + self.assertEqual(expec_strides, npu_t.stride(), | ||
| 150 | + "H2D-1d_as_strided: stride not preserved") | ||
| 151 | + | ||
| 152 | + def test_h2d_expand_2d(self): | ||
| 153 | + """H2D: expand 2D (stride=0) -> preserve_format falls back to suggest_memory_format""" | ||
| 154 | + cpu_t = torch.randn(1, 5).expand(4, 5) | ||
| 155 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 156 | + | ||
| 157 | + npu_t = cpu_t.to("npu", memory_format=torch.preserve_format) | ||
| 158 | + self._verify_preserve_format_result(cpu_t, npu_t, "H2D-expand_2d") | ||
| 159 | + # expand stride is not preserved, falls back to contiguous | ||
| 160 | + self.assertTrue(npu_t.is_contiguous(), | ||
| 161 | + "H2D-expand_2d: non-overlapping-and-dense should fall back to contiguous") | ||
| 162 | + | ||
| 163 | + def test_h2d_expand_3d(self): | ||
| 164 | + """H2D: expand 3D (stride=0) -> preserve_format falls back to suggest_memory_format""" | ||
| 165 | + cpu_t = torch.randn(2, 1, 4).expand(2, 3, 4) | ||
| 166 | + self.assertFalse(cpu_t.is_contiguous()) | ||
| 167 | + | ||
| 168 | + | ||
| 169 | + npu_t = cpu_t.to("npu", memory_format=torch.preserve_format) | ||
| 170 | + self._verify_preserve_format_result(cpu_t, npu_t, "H2D-expand_3d") | ||
| 171 | + self.assertTrue(npu_t.is_contiguous(), | ||
| 172 | + "H2D-expand_3d: non-overlapping-and-dense should fall back to contiguous") | ||
| 173 | + | ||
| 174 | + def test_h2d_as_strided_overlap(self): | ||
| 175 | + """H2D: as_strided with overlap -> preserve_format falls back to suggest_memory_format""" | ||
| 176 | + cpu_t = torch.randn(12).as_strided((3, 3), (4, 1)) | ||
| 177 | + | ||
| 178 | + | ||
| 179 | + npu_t = cpu_t.to("npu", memory_format=torch.preserve_format) | ||
| 180 | + self._verify_preserve_format_result(cpu_t, npu_t, "H2D-as_strided_overlap") | ||
| 181 | + # overlapping tensor falls back, result should be contiguous | ||
| 182 | + self.assertTrue(npu_t.is_contiguous(), | ||
| 183 | + "H2D-as_strided_overlap: non-overlapping-and-dense should fall back to contiguous") | ||
| 184 | + | ||
| 185 | + # ================================================================ | ||
| 186 | + # D2H scenario: NPU -> CPU | ||
| 187 | + # ================================================================ | ||
| 188 | + | ||
| 189 | + def test_d2h_transpose_2d(self): | ||
| 190 | + """D2H: transpose 2D non-contiguous -> preserve_format preserves stride""" | ||
| 191 | + npu_t = torch.randn(4, 6, device="npu").t() | ||
| 192 | + self.assertFalse(npu_t.is_contiguous()) | ||
| 193 | + | ||
| 194 | + cpu_t = npu_t.to("cpu", memory_format=torch.preserve_format) | ||
| 195 | + self._verify_preserve_format_result(npu_t, cpu_t, "D2H-transpose_2d") | ||
| 196 | + self.assertEqual(npu_t.stride(), cpu_t.stride(), | ||
| 197 | + "D2H-transpose_2d: stride not preserved") | ||
| 198 | + | ||
| 199 | + def test_d2h_permute_4d(self): | ||
| 200 | + """D2H: permute 4D non-contiguous -> preserve_format preserves stride""" | ||
| 201 | + npu_t = torch.randn(2, 3, 4, 5, device="npu").permute(3, 1, 0, 2) | ||
| 202 | + self.assertFalse(npu_t.is_contiguous()) | ||
| 203 | + | ||
| 204 | + cpu_t = npu_t.to("cpu", memory_format=torch.preserve_format) | ||
| 205 | + self._verify_preserve_format_result(npu_t, cpu_t, "D2H-permute_4d") | ||
| 206 | + self.assertEqual(npu_t.stride(), cpu_t.stride(), | ||
| 207 | + "D2H-permute_4d: stride not preserved") | ||
| 208 | + | ||
| 209 | + def test_d2h_slice_dim0(self): | ||
| 210 | + """D2H: slice dim0 non-contiguous -> preserve_format preserves stride""" | ||
| 211 | + npu_t = torch.randn(8, 5, device="npu")[::2] | ||
| 212 | + self.assertFalse(npu_t.is_contiguous()) | ||
| 213 | + | ||
| 214 | + cpu_t = npu_t.to("cpu", memory_format=torch.preserve_format) | ||
| 215 | + self._verify_preserve_format_result(npu_t, cpu_t, "D2H-slice_dim0") | ||
| 216 | + expec_strides = self._get_dense_strides(npu_t) | ||
| 217 | + self.assertEqual(expec_strides, cpu_t.stride(), | ||
| 218 | + "D2H-slice_dim0: stride not preserved") | ||
| 219 | + | ||
| 220 | + def test_d2h_slice_dim1(self): | ||
| 221 | + """D2H: slice dim1 non-contiguous -> preserve_format preserves stride""" | ||
| 222 | + npu_t = torch.randn(4, 8, device="npu")[:, ::2] | ||
| 223 | + self.assertFalse(npu_t.is_contiguous()) | ||
| 224 | + | ||
| 225 | + cpu_t = npu_t.to("cpu", memory_format=torch.preserve_format) | ||
| 226 | + self._verify_preserve_format_result(npu_t, cpu_t, "D2H-slice_dim1") | ||
| 227 | + expec_strides = self._get_dense_strides(npu_t) | ||
| 228 | + self.assertEqual(expec_strides, cpu_t.stride(), | ||
| 229 | + "D2H-slice_dim1: stride not preserved") | ||
| 230 | + | ||
| 231 | + def test_d2h_narrow(self): | ||
| 232 | + """D2H: narrow non-contiguous -> preserve_format preserves stride""" | ||
| 233 | + npu_t = torch.randn(6, 8, device="npu").narrow(1, 1, 5) | ||
| 234 | + self.assertFalse(npu_t.is_contiguous()) | ||
| 235 | + | ||
| 236 | + cpu_t = npu_t.to("cpu", memory_format=torch.preserve_format) | ||
| 237 | + self._verify_preserve_format_result(npu_t, cpu_t, "D2H-narrow") | ||
| 238 | + expec_strides = self._get_dense_strides(npu_t) | ||
| 239 | + self.assertEqual(expec_strides, cpu_t.stride(), | ||
| 240 | + "D2H-narrow: stride not preserved") | ||
| 241 | + | ||
| 242 | + def test_d2h_select(self): | ||
| 243 | + """D2H: select (reduces dim) non-contiguous -> preserve_format preserves stride""" | ||
| 244 | + npu_t = torch.randn(3, 5, 4, device="npu").select(1, 2) | ||
| 245 | + self.assertFalse(npu_t.is_contiguous()) | ||
| 246 | + | ||
| 247 | + cpu_t = npu_t.to("cpu", memory_format=torch.preserve_format) | ||
| 248 | + self._verify_preserve_format_result(npu_t, cpu_t, "D2H-select") | ||
| 249 | + expec_strides = self._get_dense_strides(npu_t) | ||
| 250 | + self.assertEqual(expec_strides, cpu_t.stride(), | ||
| 251 | + "D2H-select: stride not preserved") | ||
| 252 | + | ||
| 253 | + def test_d2h_1d_as_strided(self): | ||
| 254 | + """D2H: 1D as_strided non-contiguous -> preserve_format preserves stride""" | ||
| 255 | + npu_t = torch.randn(10, device="npu").as_strided((4,), (2,)) | ||
| 256 | + self.assertFalse(npu_t.is_contiguous()) | ||
| 257 | + | ||
| 258 | + cpu_t = npu_t.to("cpu", memory_format=torch.preserve_format) | ||
| 259 | + self._verify_preserve_format_result(npu_t, cpu_t, "D2H-1d_as_strided") | ||
| 260 | + expec_strides = self._get_dense_strides(npu_t) | ||
| 261 | + self.assertEqual(expec_strides, cpu_t.stride(), | ||
| 262 | + "D2H-1d_as_strided: stride not preserved") | ||
| 263 | + | ||
| 264 | + def test_d2h_expand_2d(self): | ||
| 265 | + """D2H: expand 2D (stride=0) -> preserve_format falls back to suggest_memory_format""" | ||
| 266 | + npu_t = torch.randn(1, 5, device="npu").expand(4, 5) | ||
| 267 | + self.assertFalse(npu_t.is_contiguous()) | ||
| 268 | + | ||
| 269 | + cpu_t = npu_t.to("cpu", memory_format=torch.preserve_format) | ||
| 270 | + self._verify_preserve_format_result(npu_t, cpu_t, "D2H-expand_2d") | ||
| 271 | + self.assertTrue(cpu_t.is_contiguous(), | ||
| 272 | + "D2H-expand_2d: non-overlapping-and-dense should fall back to contiguous") | ||
| 273 | + | ||
| 274 | + def test_d2h_expand_3d(self): | ||
| 275 | + """D2H: expand 3D (stride=0) -> preserve_format falls back to suggest_memory_format""" | ||
| 276 | + npu_t = torch.randn(2, 1, 4, device="npu").expand(2, 3, 4) | ||
| 277 | + self.assertFalse(npu_t.is_contiguous()) | ||
| 278 | + cpu_t = npu_t.to("cpu", memory_format=torch.preserve_format) | ||
| 279 | + self._verify_preserve_format_result(npu_t, cpu_t, "D2H-expand_3d") | ||
| 280 | + self.assertTrue(cpu_t.is_contiguous(), | ||
| 281 | + "D2H-expand_3d: non-overlapping-and-dense should fall back to contiguous") | ||
| 282 | + | ||
| 283 | + def test_d2h_as_strided_overlap(self): | ||
| 284 | + """D2H: as_strided with overlap -> preserve_format falls back to suggest_memory_format""" | ||
| 285 | + npu_t = torch.randn(12, device="npu").as_strided((3, 3), (4, 1)) | ||
| 286 | + | ||
| 287 | + cpu_t = npu_t.to("cpu", memory_format=torch.preserve_format) | ||
| 288 | + self._verify_preserve_format_result(npu_t, cpu_t, "D2H-as_strided_overlap") | ||
| 289 | + self.assertTrue(cpu_t.is_contiguous(), | ||
| 290 | + "D2H-as_strided_overlap: non-overlapping-and-dense should fall back to contiguous") | ||
| 291 | + | ||
| 292 | + # ================================================================ | ||
| 293 | + # D2D scenario: NPU:0 -> NPU:1 (different device_index) | ||
| 294 | + # ================================================================ | ||
| 295 | + | ||
| 296 | + | ||
| 297 | + def test_d2d_transpose_2d(self): | ||
| 298 | + """D2D: transpose 2D non-contiguous -> preserve_format preserves stride""" | ||
| 299 | + npu0_t = torch.randn(4, 6, device="npu:0").t() | ||
| 300 | + self.assertFalse(npu0_t.is_contiguous()) | ||
| 301 | + | ||
| 302 | + npu1_t = npu0_t.to("npu:1", memory_format=torch.preserve_format) | ||
| 303 | + self._verify_preserve_format_result(npu0_t, npu1_t, "D2D-transpose_2d") | ||
| 304 | + self.assertEqual(str(npu1_t.device), "npu:1") | ||
| 305 | + self.assertEqual(npu0_t.stride(), npu1_t.stride(), | ||
| 306 | + "D2D-transpose_2d: stride not preserved") | ||
| 307 | + | ||
| 308 | + | ||
| 309 | + def test_d2d_permute_4d(self): | ||
| 310 | + """D2D: permute 4D non-contiguous -> preserve_format preserves stride""" | ||
| 311 | + npu0_t = torch.randn(2, 3, 4, 5, device="npu:0").permute(3, 1, 0, 2) | ||
| 312 | + self.assertFalse(npu0_t.is_contiguous()) | ||
| 313 | + | ||
| 314 | + npu1_t = npu0_t.to("npu:1", memory_format=torch.preserve_format) | ||
| 315 | + self._verify_preserve_format_result(npu0_t, npu1_t, "D2D-permute_4d") | ||
| 316 | + self.assertEqual(str(npu1_t.device), "npu:1") | ||
| 317 | + self.assertEqual(npu0_t.stride(), npu1_t.stride(), | ||
| 318 | + "D2D-permute_4d: stride not preserved") | ||
| 319 | + | ||
| 320 | +if __name__ == "__main__": | ||
| 321 | + run_tests() | ||
| @@ -1,11 +1,11 @@ | |||
| 1 | 1 | ||
| 2 | 2 | ||
| 3 | - | ||
| 4 | 3 | ||
| 4 | + | ||
| 5 | 5 | ||
| 6 | -#include "torch_npu/csrc/framework/utils/OpAdapter.h" | 6 | +#include <torch_npu/csrc/aten/CustomFunctions.h> |
| 7 | -#include "torch_npu/csrc/aten/CustomFunctions.h" | 7 | +#include <torch_npu/csrc/aten/NPUNativeFunctions.h> |
| 8 | -#include "torch_npu/csrc/aten/NPUNativeFunctions.h" | 8 | +#include <torch_npu/csrc/framework/utils/OpAdapter.h> |
| 9 | 9 | ||
| 10 | namespace at_npu { | 10 | namespace at_npu { |
| 11 | namespace native { | 11 | namespace native { |
| @@ -14,14 +14,13 @@ namespace native { | |||
| 14 | // representing the current device) and return the corresponding c10::Device | 14 | // representing the current device) and return the corresponding c10::Device |
| 15 | // according to the actual device at the time of this function call. No-op | 15 | // according to the actual device at the time of this function call. No-op |
| 16 | // if the device_index is set. | 16 | // if the device_index is set. |
| 17 | -static inline c10::Device ensure_has_index(c10::Device device) | 17 | +static inline c10::Device ensure_has_index(c10::Device device) { |
| 18 | -{ | 18 | + if (device.is_cpu() || device.has_index()) { |
| 19 | - if (device.is_cpu() || device.has_index()) { | 19 | + return device; |
| 20 | - return device; | 20 | + } |
| 21 | - } | 21 | + const c10::impl::DeviceGuardImplInterface* impl = |
| 22 | - const c10::impl::DeviceGuardImplInterface* impl = | 22 | + c10::impl::getDeviceGuardImpl(device.type()); |
| 23 | - c10::impl::getDeviceGuardImpl(device.type()); | 23 | + return impl->getDevice(); |
| 24 | - return impl->getDevice(); | ||
| 25 | } | 24 | } |
| 26 | 25 | ||
| 27 | at::Tensor NPUNativeFunctions::_to_copy( | 26 | at::Tensor NPUNativeFunctions::_to_copy( |
| @@ -31,83 +30,93 @@ at::Tensor NPUNativeFunctions::_to_copy( | |||
| 31 | c10::optional<c10::Device> device, | 30 | c10::optional<c10::Device> device, |
| 32 | c10::optional<bool> pin_memory, | 31 | c10::optional<bool> pin_memory, |
| 33 | bool non_blocking, | 32 | bool non_blocking, |
| 34 | - c10::optional<c10::MemoryFormat> optional_memory_format) | 33 | + c10::optional<c10::MemoryFormat> optional_memory_format) { |
| 35 | -{ | 34 | + if (dtype.has_value() && !layout.has_value() && !device.has_value()) { |
| 36 | - if (dtype.has_value() && !layout.has_value() && !device.has_value()) { | 35 | + // _to_copy is used by to(..., copy=True). Same dtype tensors must |
| 37 | - // _to_copy is used by to(..., copy=True). Same dtype tensors must | 36 | + // still fall through to the existing copy path instead of returning self. |
| 38 | - // still fall through to the existing copy path instead of returning self. | 37 | + if (self.dtype() != dtype) { |
| 39 | - if (self.dtype() != dtype) { | 38 | + if (dtype == at::ScalarType::Double) { |
| 40 | - if (dtype == at::ScalarType::Double) { | 39 | + TORCH_NPU_WARN_ONCE( |
| 41 | - TORCH_NPU_WARN_ONCE( | 40 | + "Device do not support double dtype now, " |
| 42 | - "Device do not support double dtype now, " | 41 | + "dtype cast replace with float."); |
| 43 | - "dtype cast replace with float."); | 42 | + } |
| 44 | - } | 43 | + dtype = (dtype == at::ScalarType::Double) ? at::ScalarType::Float : dtype; |
| 45 | - dtype = (dtype == at::ScalarType::Double) ? at::ScalarType::Float : dtype; | ||
| 46 | - } | ||
| 47 | } | 44 | } |
| 45 | + } | ||
| 48 | 46 | ||
| 49 | - c10::TensorOptions options_ = c10::TensorOptions() | 47 | + c10::TensorOptions options_ = |
| 50 | - .dtype(dtype) | 48 | + c10::TensorOptions().dtype(dtype).layout(layout).device(device); |
| 51 | - .layout(layout) | ||
| 52 | - .device(device); | ||
| 53 | 49 | ||
| 54 | - auto options = self.options().merge_in(options_); | 50 | + auto options = self.options().merge_in(options_); |
| 55 | 51 | ||
| 56 | - if (layout.has_value()) { | 52 | + if (layout.has_value()) { |
| 57 | - TORCH_CHECK( | ||
| 58 | - self.layout() == layout.value(), | ||
| 59 | - "to(options) doesn't support converting to a different layout, " | ||
| 60 | - "but got self.layout being ", | ||
| 61 | - self.layout(), | ||
| 62 | - " and options.layout set as ", | ||
| 63 | - layout.value(), OPS_ERROR(ErrCode::NOT_SUPPORT)); | ||
| 64 | - } | ||
| 65 | - | ||
| 66 | - if (device.has_value()) { | ||
| 67 | - options = options.device(ensure_has_index(device.value())); | ||
| 68 | - } | ||
| 69 | - | ||
| 70 | - if (optional_memory_format.has_value()) { | ||
| 71 | - TORCH_CHECK( | ||
| 72 | - optional_memory_format.value() == c10::MemoryFormat::Preserve || | ||
| 73 | - optional_memory_format.value() == c10::MemoryFormat::Contiguous, | ||
| 74 | - "Only contiguous_format or preserve_format is supported.", OPS_ERROR(ErrCode::NOT_SUPPORT)); | ||
| 75 | - options = options.memory_format(optional_memory_format.value()); | ||
| 76 | - } else { | ||
| 77 | - if (torch_npu::utils::is_npu(self)) { | ||
| 78 | - options = options.memory_format(c10::MemoryFormat::Contiguous); | ||
| 79 | - } else { | ||
| 80 | - // keep the same as cpu default memory format: Preserve | ||
| 81 | - options = options.memory_format(c10::MemoryFormat::Preserve); | ||
| 82 | - } | ||
| 83 | - } | ||
| 84 | TORCH_CHECK( | 53 | TORCH_CHECK( |
| 85 | - options.requires_grad_opt() == c10::nullopt, | 54 | + self.layout() == layout.value(), |
| 86 | - "to(options) expects unset requires_grad flag, but got " | 55 | + "to(options) doesn't support converting to a different layout, " |
| 87 | - "options.requires_grad set as ", | 56 | + "but got self.layout being ", |
| 88 | - options.requires_grad(), OPS_ERROR(ErrCode::PARAM)); | 57 | + self.layout(), |
| 58 | + " and options.layout set as ", | ||
| 59 | + layout.value(), | ||
| 60 | + OPS_ERROR(ErrCode::NOT_SUPPORT)); | ||
| 61 | + } | ||
| 89 | 62 | ||
| 90 | - bool pin_out = non_blocking && torch_npu::utils::is_npu(self) && options.device().is_cpu() && | 63 | + if (device.has_value()) { |
| 91 | - (options.layout() == c10::kStrided); | 64 | + options = options.device(ensure_has_index(device.value())); |
| 65 | + } | ||
| 92 | 66 | ||
| 93 | - c10::MemoryFormat memory_format = options.memory_format_opt().value_or(c10::MemoryFormat::Contiguous); | 67 | + if (optional_memory_format.has_value()) { |
| 94 | - if (memory_format == c10::MemoryFormat::Preserve) { | 68 | + TORCH_CHECK( |
| 95 | - if (self.is_non_overlapping_and_dense()) { | 69 | + optional_memory_format.value() == c10::MemoryFormat::Preserve || |
| 96 | - // Copy all strides | 70 | + optional_memory_format.value() == c10::MemoryFormat::Contiguous, |
| 97 | - auto r = at::empty_strided( | 71 | + "Only contiguous_format or preserve_format is supported.", |
| 98 | - self.sizes(), self.strides(), options.memory_format(c10::nullopt).pinned_memory(pin_out)); | 72 | + OPS_ERROR(ErrCode::NOT_SUPPORT)); |
| 99 | - r.copy_(self, non_blocking); | 73 | + options = options.memory_format(optional_memory_format.value()); |
| 100 | - return r; | 74 | + } else { |
| 101 | - } else { | 75 | + options = options.memory_format(c10::MemoryFormat::Preserve); |
| 102 | - memory_format = self.suggest_memory_format(); | 76 | + } |
| 103 | - } | 77 | + TORCH_CHECK( |
| 78 | + options.requires_grad_opt() == c10::nullopt, | ||
| 79 | + "to(options) expects unset requires_grad flag, but got " | ||
| 80 | + "options.requires_grad set as ", | ||
| 81 | + options.requires_grad(), | ||
| 82 | + OPS_ERROR(ErrCode::PARAM)); | ||
| 83 | + | ||
| 84 | + bool pin_out = non_blocking && torch_npu::utils::is_npu(self) && | ||
| 85 | + options.device().is_cpu() && (options.layout() == c10::kStrided); | ||
| 86 | + | ||
| 87 | + c10::MemoryFormat memory_format = | ||
| 88 | + options.memory_format_opt().value_or(c10::MemoryFormat::Contiguous); | ||
| 89 | + if (memory_format == c10::MemoryFormat::Preserve) { | ||
| 90 | + if (options.device().supports_as_strided()) { | ||
| 91 | + if (self.is_non_overlapping_and_dense()) { | ||
| 92 | + // Copy all strides | ||
| 93 | + auto r = at::empty_strided( | ||
| 94 | + self.sizes(), | ||
| 95 | + self.strides(), | ||
| 96 | + options.memory_format(c10::nullopt).pinned_memory(pin_out)); | ||
| 97 | + r.copy_(self, non_blocking); | ||
| 98 | + return r; | ||
| 99 | + } else if (!self.is_quantized() && self.layout() == c10::kStrided) { | ||
| 100 | + auto strides = at::infer_dense_strides(self.sizes(), self.strides()); | ||
| 101 | + auto r = at::empty_strided( | ||
| 102 | + self.sizes(), strides, options.pinned_memory(pin_out)); | ||
| 103 | + r.copy_(self, non_blocking); | ||
| 104 | + return r; | ||
| 105 | + } else { | ||
| 106 | + memory_format = self.suggest_memory_format(); | ||
| 107 | + } | ||
| 108 | + } else { | ||
| 109 | + memory_format = self.suggest_memory_format(); | ||
| 104 | } | 110 | } |
| 105 | - | 111 | + } |
| 106 | - // See Note [Explicit nullopt c10::MemoryFormat argument] | 112 | + |
| 107 | - auto r = at::empty( | 113 | + // See Note [Explicit nullopt c10::MemoryFormat argument] |
| 108 | - self.sizes(), options.memory_format(memory_format).pinned_memory(pin_out), c10::nullopt); | 114 | + auto r = at::empty( |
| 109 | - r.copy_(self, non_blocking); | 115 | + self.sizes(), |
| 110 | - return r; | 116 | + options.memory_format(memory_format).pinned_memory(pin_out), |
| 117 | + c10::nullopt); | ||
| 118 | + r.copy_(self, non_blocking); | ||
| 119 | + return r; | ||
| 111 | } | 120 | } |
| 112 | 121 | ||
| 113 | } // namespace native | 122 | } // namespace native |