| 445 | } |
| 446 | |
| 447 | void TensorToStream(std::ostream& os, |
| 448 | const DenseTensor& tensor, |
| 449 | const phi::DeviceContext& dev_ctx) { |
| 450 | const auto ensure_contiguous = [](const DenseTensor& tensor) { |
| 451 | if (tensor.meta().is_contiguous()) { |
| 452 | return tensor; |
| 453 | } |
| 454 | return paddle::experimental::Trans2Contiguous(tensor); |
| 455 | }; |
| 456 | const DenseTensor& contiguous_tensor = ensure_contiguous(tensor); |
| 457 | { // the 1st field, uint32_t version |
| 458 | constexpr uint32_t version = 0; |
| 459 | os.write(reinterpret_cast<const char*>(&version), sizeof(version)); |
| 460 | } |
| 461 | { // the 2nd field, tensor description |
| 462 | // int32_t size |
| 463 | // void* protobuf message |
| 464 | proto::VarType::TensorDesc desc; |
| 465 | desc.set_data_type( |
| 466 | framework::TransToProtoVarType(contiguous_tensor.dtype())); |
| 467 | auto dims = common::vectorize(contiguous_tensor.dims()); |
| 468 | auto* pb_dims = desc.mutable_dims(); |
| 469 | pb_dims->Resize(static_cast<int>(dims.size()), 0); |
| 470 | std::copy(dims.begin(), dims.end(), pb_dims->begin()); |
| 471 | int32_t size = desc.ByteSize(); |
| 472 | os.write(reinterpret_cast<const char*>(&size), sizeof(size)); |
| 473 | auto out = desc.SerializeAsString(); |
| 474 | os.write(out.data(), size); |
| 475 | } |
| 476 | { // the 3rd field, tensor data |
| 477 | uint64_t size = |
| 478 | contiguous_tensor.numel() * phi::SizeOf(contiguous_tensor.dtype()); |
| 479 | |
| 480 | auto* data_ptr = contiguous_tensor.data(); |
| 481 | PADDLE_ENFORCE_LT(size, |
| 482 | (std::numeric_limits<std::streamsize>::max)(), |
| 483 | common::errors::ResourceExhausted( |
| 484 | "tensor size %d overflow when writing tensor", size)); |
| 485 | if (phi::is_gpu_place(contiguous_tensor.place())) { |
| 486 | #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) |
| 487 | constexpr size_t kBufSize = 1024 * 1024 * 64; // 64MB |
| 488 | std::unique_ptr<char[]> buf(new char[kBufSize]); |
| 489 | auto& gpu_dev_ctx = static_cast<const phi::GPUContext&>(dev_ctx); |
| 490 | CPUPlace cpu; |
| 491 | uintptr_t data = reinterpret_cast<uintptr_t>(data_ptr); |
| 492 | while (size != 0) { |
| 493 | size_t size_to_write = std::min(kBufSize, static_cast<size_t>(size)); |
| 494 | memory::Copy(cpu, |
| 495 | buf.get(), |
| 496 | contiguous_tensor.place(), |
| 497 | reinterpret_cast<const void*>(data), // NOLINT |
| 498 | size_to_write, |
| 499 | gpu_dev_ctx.stream()); |
| 500 | gpu_dev_ctx.Wait(); |
| 501 | os.write(buf.get(), size_to_write); |
| 502 | data += size_to_write; |
| 503 | size -= size_to_write; |
| 504 | } |