已合并
_construct_NPU_Tensor_From_Storage_And_Metadata support npu_format #27908
wang-pierre-jiacheng创建于 2025年12月12日
_construct_NPU_Tensor_From_Storage_And_Metadata support npu_format #27908
已合并
共 2 个文件变更+36-4
| @@ -365,6 +365,30 @@ class TestNPUGraphNode: | |||
| 365 | outputs = basic_npu_graph_node.reconstruct_outputs() | 365 | outputs = basic_npu_graph_node.reconstruct_outputs() |
| 366 | assert len(outputs) == 1 | 366 | assert len(outputs) == 1 |
| 367 | 367 | ||
| 368 | + def test_reconstruct_outputs_with_format(self, basic_npu_graph_node): | ||
| 369 | + # Setup mock metadata and storage info | ||
| 370 | + basic_npu_graph_node.outputs_metadata = [ | ||
| 371 | + { | ||
| 372 | + "nbytes": 4, | ||
| 373 | + "data_ptr": 1234, | ||
| 374 | + "size": (1,), | ||
| 375 | + "stride": (1,), | ||
| 376 | + "dtype": torch.float32, | ||
| 377 | + "device": "npu", | ||
| 378 | + "npu_format": 29, | ||
| 379 | + "storage_offset": 0, | ||
| 380 | + } | ||
| 381 | + ] | ||
| 382 | + basic_npu_graph_node.output_weakrefs = [MagicMock()] | ||
| 383 | + basic_npu_graph_node.output_storage_alias = [UnaliasedStorage] | ||
| 384 | + basic_npu_graph_node.cached_tensor_outputs = [MagicMock()] | ||
| 385 | + | ||
| 386 | + with patch( | ||
| 387 | + "torch_npu._C._construct_NPU_Tensor_From_Storage_And_Metadata" | ||
| 388 | + ) as mock_construct: | ||
| 389 | + outputs = basic_npu_graph_node.reconstruct_outputs() | ||
| 390 | + assert len(outputs) == 1 | ||
| 391 | + | ||
| 368 | def test_aliased_output_reconstruction(self, basic_npu_graph_node): | 392 | def test_aliased_output_reconstruction(self, basic_npu_graph_node): |
| 369 | basic_npu_graph_node.outputs_metadata = [ | 393 | basic_npu_graph_node.outputs_metadata = [ |
| 370 | { | 394 | { |
| @@ -496,10 +496,18 @@ void RegisterNpuPluggableAllocator(PyObject* module) | |||
| 496 | at::Tensor tensor = at::detail::make_tensor_base<c10::TensorImpl>( | 496 | at::Tensor tensor = at::detail::make_tensor_base<c10::TensorImpl>( |
| 497 | std::move(s), npu_dks, meta); | 497 | std::move(s), npu_dks, meta); |
| 498 | 498 | ||
| 499 | - at_npu::native::StorageDescHelper::SetDesc( | 499 | + if (metadata.contains("npu_format")) { |
| 500 | - tensor, | 500 | + at_npu::native::StorageDescHelper::SetDesc( |
| 501 | - metadata["size"].cast<std::vector<int64_t>>(), | 501 | + tensor, |
| 502 | - metadata["stride"].cast<std::vector<int64_t>>()); | 502 | + metadata["size"].cast<std::vector<int64_t>>(), |
| 503 | + metadata["stride"].cast<std::vector<int64_t>>(), | ||
| 504 | + static_cast<aclFormat>(metadata["npu_format"].cast<int64_t>())); | ||
| 505 | + } else { | ||
| 506 | + at_npu::native::StorageDescHelper::SetDesc( | ||
| 507 | + tensor, | ||
| 508 | + metadata["size"].cast<std::vector<int64_t>>(), | ||
| 509 | + metadata["stride"].cast<std::vector<int64_t>>()); | ||
| 510 | + } | ||
| 503 | tensor.unsafeGetTensorImpl()->set_sizes_and_strides( | 511 | tensor.unsafeGetTensorImpl()->set_sizes_and_strides( |
| 504 | metadata["size"].cast<std::vector<int64_t>>(), | 512 | metadata["size"].cast<std::vector<int64_t>>(), |
| 505 | metadata["stride"].cast<std::vector<int64_t>>()); | 513 | metadata["stride"].cast<std::vector<int64_t>>()); |