| 577 | }; |
| 578 | |
| 579 | void TensorFromStream(std::istream& is, |
| 580 | DenseTensor* tensor, |
| 581 | const phi::DeviceContext& dev_ctx, |
| 582 | const size_t& seek, |
| 583 | const std::vector<int64_t>& shape) { |
| 584 | uint32_t version = 0; |
| 585 | is.read(reinterpret_cast<char*>(&version), sizeof(version)); |
| 586 | |
| 587 | PADDLE_ENFORCE_EQ( |
| 588 | version, |
| 589 | 0U, |
| 590 | common::errors::InvalidArgument( |
| 591 | "tensor version %u is not supported, Only version 0 is supported", |
| 592 | version)); |
| 593 | |
| 594 | proto::VarType::TensorDesc desc; |
| 595 | { // int32_t size |
| 596 | // proto buffer |
| 597 | int32_t size = 0; |
| 598 | is.read(reinterpret_cast<char*>(&size), sizeof(size)); |
| 599 | std::unique_ptr<char[]> buf(new char[size]); // NOLINT |
| 600 | is.read(reinterpret_cast<char*>(buf.get()), size); |
| 601 | PADDLE_ENFORCE_EQ( |
| 602 | desc.ParseFromArray(buf.get(), size), |
| 603 | true, |
| 604 | common::errors::InvalidArgument("Cannot parse tensor desc")); |
| 605 | } |
| 606 | { // read tensor |
| 607 | tensor->Resize(common::make_ddim(shape)); |
| 608 | size_t seekg = seek * framework::SizeOfType(desc.data_type()); |
| 609 | is.seekg(seekg, is.cur); // NOLINT |
| 610 | |
| 611 | void* buf = nullptr; |
| 612 | phi::CPUContext ctx; |
| 613 | size_t size = tensor->numel() * framework::SizeOfType(desc.data_type()); |
| 614 | if (phi::is_gpu_place(dev_ctx.GetPlace()) || |
| 615 | phi::is_xpu_place(dev_ctx.GetPlace()) || |
| 616 | phi::is_custom_place(dev_ctx.GetPlace())) { |
| 617 | #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \ |
| 618 | defined(PADDLE_WITH_XPU) || defined(PADDLE_WITH_CUSTOM_DEVICE) |
| 619 | DenseTensor cpu_tensor; |
| 620 | cpu_tensor.Resize(common::make_ddim(shape)); |
| 621 | framework::VisitDataType( |
| 622 | desc.data_type(), |
| 623 | DeserializedDataFunctor(&buf, &cpu_tensor, ctx.GetPlace())); |
| 624 | is.read(static_cast<char*>(buf), size); // NOLINT |
| 625 | auto dst_place = dev_ctx.GetPlace(); |
| 626 | framework::TensorCopy(cpu_tensor, dst_place, dev_ctx, tensor); |
| 627 | if (phi::is_custom_place(dev_ctx.GetPlace())) { |
| 628 | dev_ctx.Wait(); |
| 629 | } |
| 630 | #else |
| 631 | if (phi::is_gpu_place(dev_ctx.GetPlace())) { |
| 632 | PADDLE_THROW(common::errors::Unimplemented( |
| 633 | "CUDAPlace is not supported when not compiled with CUDA")); |
| 634 | } else if (phi::is_xpu_place(dev_ctx.GetPlace())) { |
| 635 | PADDLE_THROW(common::errors::Unimplemented( |
| 636 | "XPUPlace is not supported when not compiled with XPU")); |