已合并
fix: route internal-format copy_ to native path and guard NPUStorageDesc access #44450
wuyouqi1创建于 11 天前
fix: route internal-format copy_ to native path and guard NPUStorageDesc access #44450
已合并
共 3 个文件变更+126-6
| @@ -35,6 +35,17 @@ NZ_ROUNDTRIP_DTYPES = [ | |||
| 35 | ] | 35 | ] |
| 36 | D2H_TEST_DTYPES = [torch.float16, torch.bfloat16, torch.int8, torch.int32] | 36 | D2H_TEST_DTYPES = [torch.float16, torch.bfloat16, torch.int8, torch.int32] |
| 37 | 37 | ||
| 38 | +# Ascend950 materializes a FRACTAL_NZ cast as an NZ_C0 variant (50-54) since C0 | ||
| 39 | +# variants are output-only; accept the NZ family when checking the cast result. | ||
| 40 | +NZ_FORMAT_FAMILY = ( | ||
| 41 | + FORMAT_INFO["FRACTAL_NZ"], | ||
| 42 | + torch_npu.Format.FRACTAL_NZ_C0_16, | ||
| 43 | + torch_npu.Format.FRACTAL_NZ_C0_32, | ||
| 44 | + torch_npu.Format.FRACTAL_NZ_C0_2, | ||
| 45 | + torch_npu.Format.FRACTAL_NZ_C0_4, | ||
| 46 | + torch_npu.Format.FRACTAL_NZ_C0_8, | ||
| 47 | +) | ||
| 48 | + | ||
| 38 | 49 | ||
| 39 | def save_tensor(tensor, path, acl_format): | 50 | def save_tensor(tensor, path, acl_format): |
| 40 | x = torch_npu.npu_format_cast(tensor.npu(), acl_format) | 51 | x = torch_npu.npu_format_cast(tensor.npu(), acl_format) |
| @@ -129,11 +140,8 @@ class TestSerializationFormatAscend950(TestCase): | |||
| 129 | def test_nz_d2h_and_repr(self): | 140 | def test_nz_d2h_and_repr(self): |
| 130 | """D2H, repr, str, print on private-format tensor must not crash.""" | 141 | """D2H, repr, str, print on private-format tensor must not crash.""" |
| 131 | for dt in D2H_TEST_DTYPES: | 142 | for dt in D2H_TEST_DTYPES: |
| 132 | - try: | 143 | + x = torch.randn(64, 64, dtype=torch.float32).to(dt).npu() |
| 133 | - x = torch.randn(64, 64, dtype=torch.float32).to(dt).npu() | 144 | + x = torch_npu.npu_format_cast(x, torch_npu.Format.FRACTAL_NZ) |
| 134 | - x = torch_npu.npu_format_cast(x, torch_npu.Format.FRACTAL_NZ) | ||
| 135 | - except Exception: | ||
| 136 | - continue | ||
| 137 | 145 | ||
| 138 | c = x.cpu() | 146 | c = x.cpu() |
| 139 | self.assertEqual(c.device.type, "cpu") | 147 | self.assertEqual(c.device.type, "cpu") |
| @@ -142,5 +150,97 @@ class TestSerializationFormatAscend950(TestCase): | |||
| 142 | print(x) | 150 | print(x) |
| 143 | 151 | ||
| 144 | 152 | ||
| 153 | +def cast_copy_tensor(tensor, acl_format): | ||
| 154 | + # Without allow_internal_format the cast silently downgrades internal formats to ND. | ||
| 155 | + if acl_format != FORMAT_INFO["ND"]: | ||
| 156 | + torch_npu.npu.config.allow_internal_format = True | ||
| 157 | + return torch_npu.npu_format_cast(tensor, acl_format) | ||
| 158 | + | ||
| 159 | + | ||
| 160 | + | ||
| 161 | + "Ascend950 copy behavior differs; see TestCopyFormatAscend950") | ||
| 162 | +class TestCopyFormat(TestCase): | ||
| 163 | + """A2/A3: copy_ across all FORMAT_INFO formats in H2D/D2H/D2D directions.""" | ||
| 164 | + | ||
| 165 | + def test_copy_formats_h2d_d2h_d2d(self): | ||
| 166 | + for fmt_name, fmt in FORMAT_INFO.items(): | ||
| 167 | + src_cpu = torch.randn(2, 3, 7, 7) | ||
| 168 | + | ||
| 169 | + # h2d: NPU dst (fmt) <- CPU src | ||
| 170 | + dst_h2d = cast_copy_tensor(torch.zeros(2, 3, 7, 7).npu(), fmt) | ||
| 171 | + dst_h2d.copy_(src_cpu) | ||
| 172 | + self.assertTrue(torch.equal(dst_h2d.cpu(), src_cpu), | ||
| 173 | + f"h2d dst={fmt_name}") | ||
| 174 | + | ||
| 175 | + # d2h: CPU dst <- NPU src (fmt) | ||
| 176 | + src_d2h = cast_copy_tensor(src_cpu.npu(), fmt) | ||
| 177 | + dst_d2h = torch.zeros(2, 3, 7, 7) | ||
| 178 | + dst_d2h.copy_(src_d2h) | ||
| 179 | + self.assertTrue(torch.equal(dst_d2h, src_cpu), | ||
| 180 | + f"d2h src={fmt_name}") | ||
| 181 | + | ||
| 182 | + # d2d: dst (fmt) <- src (other format) | ||
| 183 | + for other_name, other in FORMAT_INFO.items(): | ||
| 184 | + dst_d2d = cast_copy_tensor(torch.zeros(2, 3, 7, 7).npu(), fmt) | ||
| 185 | + fmt_before = torch_npu.get_npu_format(dst_d2d) | ||
| 186 | + src_d2d = cast_copy_tensor(src_cpu.npu(), other) | ||
| 187 | + dst_d2d.copy_(src_d2d) | ||
| 188 | + self.assertTrue(torch.equal(dst_d2d.cpu(), src_cpu), | ||
| 189 | + f"d2d dst={fmt_name} src={other_name}") | ||
| 190 | + self.assertEqual(torch_npu.get_npu_format(dst_d2d), fmt_before, | ||
| 191 | + f"d2d dst={fmt_name} format changed") | ||
| 192 | + | ||
| 193 | + | ||
| 194 | + | ||
| 195 | +class TestCopyFormatAscend950(TestCase): | ||
| 196 | + """Ascend950: internal-format copy_ is only supported device-to-host.""" | ||
| 197 | + | ||
| 198 | + COPY_TEST_DTYPE = torch.float16 | ||
| 199 | + | ||
| 200 | + def test_copy_base_format_h2d_d2h_d2d(self): | ||
| 201 | + src_cpu = torch.randn(8, 8, dtype=self.COPY_TEST_DTYPE) | ||
| 202 | + dst_h2d = torch.zeros(8, 8, dtype=self.COPY_TEST_DTYPE).npu() | ||
| 203 | + dst_h2d.copy_(src_cpu) | ||
| 204 | + self.assertTrue(torch.equal(dst_h2d.cpu(), src_cpu)) | ||
| 205 | + | ||
| 206 | + src_d2h = src_cpu.npu() | ||
| 207 | + dst_d2h = torch.zeros(8, 8, dtype=self.COPY_TEST_DTYPE) | ||
| 208 | + dst_d2h.copy_(src_d2h) | ||
| 209 | + self.assertTrue(torch.equal(dst_d2h, src_cpu)) | ||
| 210 | + | ||
| 211 | + dst_d2d = torch.zeros(8, 8, dtype=self.COPY_TEST_DTYPE).npu() | ||
| 212 | + dst_d2d.copy_(src_d2h) | ||
| 213 | + self.assertTrue(torch.equal(dst_d2d.cpu(), src_cpu)) | ||
| 214 | + | ||
| 215 | + def test_copy_nz_d2h(self): | ||
| 216 | + for dt in D2H_TEST_DTYPES: | ||
| 217 | + src_cpu = torch.randn(8, 8).to(dt) | ||
| 218 | + src_nz = cast_copy_tensor(src_cpu.npu(), FORMAT_INFO["FRACTAL_NZ"]) | ||
| 219 | + self.assertIn(torch_npu.get_npu_format(src_nz), NZ_FORMAT_FAMILY, | ||
| 220 | + f"expected FRACTAL_NZ family, got {torch_npu.get_npu_format(src_nz)}") | ||
| 221 | + | ||
| 222 | + dst_cpu = torch.zeros(8, 8, dtype=dt) | ||
| 223 | + dst_cpu.copy_(src_nz) | ||
| 224 | + self.assertTrue(torch.equal(dst_cpu, src_cpu), f"d2h {dt}") | ||
| 225 | + | ||
| 226 | + def test_copy_nz_h2d_not_supported(self): | ||
| 227 | + dst_nz = cast_copy_tensor(torch.zeros(8, 8, dtype=self.COPY_TEST_DTYPE).npu(), | ||
| 228 | + FORMAT_INFO["FRACTAL_NZ"]) | ||
| 229 | + with self.assertRaisesRegex(RuntimeError, "not supported on Ascend950"): | ||
| 230 | + dst_nz.copy_(torch.randn(8, 8, dtype=self.COPY_TEST_DTYPE)) | ||
| 231 | + | ||
| 232 | + def test_copy_nz_d2d_not_supported(self): | ||
| 233 | + src_nz = cast_copy_tensor(torch.randn(8, 8, dtype=self.COPY_TEST_DTYPE).npu(), | ||
| 234 | + FORMAT_INFO["FRACTAL_NZ"]) | ||
| 235 | + dst_nz = cast_copy_tensor(torch.zeros(8, 8, dtype=self.COPY_TEST_DTYPE).npu(), | ||
| 236 | + FORMAT_INFO["FRACTAL_NZ"]) | ||
| 237 | + with self.assertRaisesRegex(RuntimeError, "not supported on Ascend950"): | ||
| 238 | + dst_nz.copy_(src_nz) | ||
| 239 | + | ||
| 240 | + dst_nd = torch.zeros(8, 8, dtype=self.COPY_TEST_DTYPE).npu() | ||
| 241 | + with self.assertRaisesRegex(RuntimeError, "not supported on Ascend950"): | ||
| 242 | + dst_nd.copy_(src_nz) | ||
| 243 | + | ||
| 244 | + | ||
| 145 | if __name__ == "__main__": | 245 | if __name__ == "__main__": |
| 146 | run_tests() | 246 | run_tests() |
| @@ -239,6 +239,19 @@ at::Tensor& NPUNativeOpApiFunctions::copy_(at::Tensor& self, const at::Tensor& s | |||
| 239 | 239 | ||
| 240 | auto maybe_outnames = at::namedinference::compute_broadcast_outnames(self, src); | 240 | auto maybe_outnames = at::namedinference::compute_broadcast_outnames(self, src); |
| 241 | 241 | ||
| 242 | + // aclnnInplaceCopy corrupts internal-format storage: on A2/A3 fall back to | ||
| 243 | + // the native copy_, on Ascend950 only d2h is supported (cast to base first). | ||
| 244 | + const bool self_is_base = FormatHelper::IsOpInputBaseFormat(self); | ||
| 245 | + const bool src_is_base = FormatHelper::IsOpInputBaseFormat(src); | ||
| 246 | + if (!self_is_base || !src_is_base) { | ||
| 247 | + if (!c10_npu::IsAclnnOnly()) { | ||
| 248 | + return NPUNativeFunctions::copy_(self, src, non_blocking); | ||
| 249 | + } | ||
| 250 | + TORCH_CHECK(!torch_npu::utils::is_npu(self), | ||
| 251 | + "The copy_ operator with internal format tensors is not supported on Ascend950, " | ||
| 252 | + "only device-to-host copies are supported", OPS_ERROR(ErrCode::NOT_SUPPORT)); | ||
| 253 | + } | ||
| 254 | + | ||
| 242 | if (torch_npu::utils::is_npu(self)) { | 255 | if (torch_npu::utils::is_npu(self)) { |
| 243 | if (torch_npu::utils::is_npu(src)) { | 256 | if (torch_npu::utils::is_npu(src)) { |
| 244 | copy_d2d_baseformat_opapi(self, src, non_blocking); | 257 | copy_d2d_baseformat_opapi(self, src, non_blocking); |
| @@ -1,3 +1,5 @@ | |||
| 1 | + | ||
| 2 | + | ||
| 1 | 3 | ||
| 2 | 4 | ||
| 3 | 5 | ||
| @@ -19,7 +21,12 @@ NPUStorageImpl *NPUBridge::GetNpuStorageImpl(const at::Tensor &tensor) | |||
| 19 | 21 | ||
| 20 | NPUStorageDesc &NPUBridge::GetNpuStorageImplDesc(const at::Tensor &tensor) | 22 | NPUStorageDesc &NPUBridge::GetNpuStorageImplDesc(const at::Tensor &tensor) |
| 21 | { | 23 | { |
| 22 | - return static_cast<NPUStorageImpl *>(tensor.storage().unsafeGetStorageImpl())->npu_desc_; | 24 | + // from_blob tensors (legacy serialization _write_file) carry a plain |
| 25 | + // c10::StorageImpl; reading npu_desc_ on them is out-of-bounds. | ||
| 26 | + auto *storage_impl = tensor.storage().unsafeGetStorageImpl(); | ||
| 27 | + TORCH_CHECK(typeid(*storage_impl) == typeid(NPUStorageImpl), | ||
| 28 | + "The npu storage desc is unavailable: the tensor's storage is not an NPUStorageImpl."); | ||
| 29 | + return static_cast<NPUStorageImpl *>(storage_impl)->npu_desc_; | ||
| 23 | } | 30 | } |
| 24 | 31 | ||
| 25 | 32 | ||