tensor_util.cc 49.1 KB
Newer Older
Y
Yang Yu 已提交
1 2
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.

3 4 5
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
Y
Yang Yu 已提交
6

7 8 9 10 11 12 13
    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */
Y
Yang Yu 已提交
14

15 16
#include "paddle/fluid/framework/tensor_util.h"

C
chengduoZH 已提交
17 18
#include <algorithm>
#include <limits>
C
chengduo 已提交
19
#include <memory>
20
#include <string>
C
chengduo 已提交
21
#include <utility>
C
chengduoZH 已提交
22
#include <vector>
23

24
#include "paddle/fluid/framework/convert_utils.h"
Y
yuyang18 已提交
25
#include "paddle/fluid/framework/data_type.h"
26
#include "paddle/fluid/platform/complex.h"
27
#include "paddle/fluid/platform/profiler/event_tracing.h"
28
#include "paddle/phi/core/dense_tensor.h"
29

30
#ifdef PADDLE_WITH_MKLDNN
31
#include "dnnl_debug.h"  // NOLINT
32
#endif
Y
Yang Yu 已提交
33 34 35

namespace paddle {
namespace framework {
Y
Yi Wang 已提交
36

37
template <typename TENSOR>
38 39 40 41
void TensorCopyImpl(const TENSOR& src,
                    const platform::Place& dst_place,
                    const platform::DeviceContext& ctx,
                    TENSOR* dst) {
42 43
  if (&src == dst) {
    auto src_copy = src;
44
    TensorCopyImpl(src_copy, dst_place, ctx, dst);
45 46 47
    return;
  }

M
minqiyang 已提交
48 49
  VLOG(3) << "TensorCopy " << src.dims() << " from " << src.place() << " to "
          << dst_place;
Y
Yi Wang 已提交
50 51 52 53
  src.check_memory_size();
  dst->Resize(src.dims());
  dst->set_layout(src.layout());
  auto src_place = src.place();
54
  auto src_ptr = src.data();
55
#ifdef PADDLE_WITH_MKLDNN
56
  dst->set_mem_desc(src.mem_desc());
57 58 59
  // oneDNN tensors due to padding may be of bigger size
  // than numel()*size(type())
  auto dst_ptr =
60
      src.layout() == DataLayout::ONEDNN
61 62
          ? dst->mutable_data(dst_place, src.dtype(), src.memory_size())
          : dst->mutable_data(dst_place, src.dtype());
63
#else
64
  auto dst_ptr = dst->mutable_data(dst_place, src.dtype());
65
#endif
66
  dst->set_layout(src.layout());
67 68 69 70 71
  if (src_ptr == dst_ptr && src_place == dst_place) {
    VLOG(3) << "Skip copy the same data async from " << src_place << " to "
            << dst_place;
    return;
  }
72
  VLOG(4) << "src:" << src_ptr << ", dst:" << dst_ptr;
73

74
#ifdef PADDLE_WITH_MKLDNN
75
  auto size = src.layout() == DataLayout::ONEDNN
76
                  ? src.memory_size()
77
                  : src.numel() * framework::DataTypeSize(src.dtype());
78
#else
79
  auto size = src.numel() * framework::DataTypeSize(src.dtype());
80
#endif
Y
Yi Wang 已提交
81 82

  if (platform::is_cpu_place(src_place) && platform::is_cpu_place(dst_place)) {
83
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
Y
Yi Wang 已提交
84
  }
85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107
#ifdef PADDLE_WITH_CUSTOM_DEVICE
  else if (platform::is_custom_place(src_place) &&  // NOLINT
           platform::is_cpu_place(dst_place)) {
    auto stream =
        reinterpret_cast<const platform::CustomDeviceContext&>(ctx).stream();
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size, stream);
  } else if (platform::is_cpu_place(src_place) &&  // NOLINT
             platform::is_custom_place(dst_place)) {
    auto stream =
        reinterpret_cast<const platform::CustomDeviceContext&>(ctx).stream();
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size, stream);
  } else if (platform::is_custom_place(src_place) &&  // NOLINT
             platform::is_custom_place(dst_place)) {
    if (src_ptr == dst_ptr) {
      VLOG(3) << "Skip copy the same data async from " << src_place << " to "
              << dst_place;
      return;
    }
    auto stream =
        reinterpret_cast<const platform::CustomDeviceContext&>(ctx).stream();
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size, stream);
  }
#endif
108 109 110
#ifdef PADDLE_WITH_XPU
  else if (platform::is_xpu_place(src_place) &&  // NOLINT
           platform::is_cpu_place(dst_place)) {
111
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
112 113
  } else if (platform::is_cpu_place(src_place) &&
             platform::is_xpu_place(dst_place)) {
114
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
115 116 117 118 119 120 121
  } else if (platform::is_xpu_place(src_place) &&
             platform::is_xpu_place(dst_place)) {
    if (src_ptr == dst_ptr) {
      VLOG(3) << "Skip copy the same data async from " << src_place << " to "
              << dst_place;
      return;
    }
122
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
123 124 125 126 127
  } else {
    PADDLE_THROW(platform::errors::Unimplemented(
        "Copy from %s to %s is not supported.", src_place, dst_place));
  }
#endif
128 129 130 131 132 133
#ifdef PADDLE_WITH_ASCEND_CL
  // TODO(zhiqiu): handle different condition like CUDA code below
  else if (platform::is_npu_place(src_place) &&  // NOLINT
           platform::is_cpu_place(dst_place)) {
    auto stream =
        reinterpret_cast<const platform::NPUDeviceContext&>(ctx).stream();
134
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size, stream);
135 136 137
  }
  else if (platform::is_cpu_place(src_place) &&  // NOLINT
           platform::is_npu_place(dst_place)) {
138 139
    //  1. cpu tensor -> npu pinned tensor
    platform::NPUPinnedPlace npu_pinned_place;
140
    phi::DenseTensor npu_pinned_tensor;
141 142
    npu_pinned_tensor.Resize(src.dims());
    auto npu_pinned_ptr =
143
        npu_pinned_tensor.mutable_data(npu_pinned_place, src.dtype());
144
    memory::Copy(npu_pinned_place, npu_pinned_ptr, src_place, src_ptr, size);
145 146 147

    //  2. async copy npu pinned tensor -> npu tensor
    memory::Copy(
148 149 150 151 152
        dst_place,
        dst_ptr,
        npu_pinned_place,
        npu_pinned_ptr,
        size,
153 154 155 156 157 158 159 160
        reinterpret_cast<const platform::NPUDeviceContext&>(ctx).stream());

    //  3. record event
    auto npu_pinned_allocator =
        static_cast<paddle::memory::allocation::NPUPinnedAllocator*>(
            paddle::memory::allocation::AllocatorFacade::Instance()
                .GetAllocator(npu_pinned_place)
                .get());
161
    phi::Allocation* allocation = npu_pinned_tensor.Holder().get();
162 163 164
    npu_pinned_allocator->RecordEvent(
        allocation,
        reinterpret_cast<const platform::NPUDeviceContext&>(ctx).stream());
165 166 167 168 169 170 171 172 173 174
  }
  else if (platform::is_npu_place(src_place) &&  // NOLINT
           platform::is_npu_place(dst_place)) {
    if (src_ptr == dst_ptr) {
      VLOG(3) << "Skip copy the same data async from " << src_place << " to "
              << dst_place;
      return;
    }
    auto stream =
        reinterpret_cast<const platform::NPUDeviceContext&>(ctx).stream();
175
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size, stream);
176
  }
W
WangXi 已提交
177 178
  else if (platform::is_npu_pinned_place(src_place) &&  // NOLINT
           platform::is_npu_place(dst_place)) {         /* npu_pinned->npu */
179 180
    auto src_npu_pinned_place = src_place;
    auto dst_npu_place = dst_place;
W
WangXi 已提交
181
    auto ctx_place = ctx.GetPlace();
182 183 184 185 186 187 188
    PADDLE_ENFORCE_EQ(
        platform::is_npu_place(ctx_place),
        true,
        platform::errors::PreconditionNotMet(
            "Device context place mismatch. When copying phi::DenseTensor "
            "data from NPU Pinned memory to NPU memory, current "
            "device context place should be NPU."));
189
    auto ctx_npu_place = ctx_place;
190 191
    PADDLE_ENFORCE_EQ(dst_npu_place,
                      ctx_npu_place,
W
WangXi 已提交
192 193 194 195
                      platform::errors::PreconditionNotMet(
                          "The target NPU device and current device context do "
                          "not match. The target NPU device number is %d, but "
                          "device context NPU number is %d.",
196 197
                          dst_npu_place.device,
                          ctx_npu_place.device));
W
WangXi 已提交
198 199
    auto stream =
        reinterpret_cast<const platform::NPUDeviceContext&>(ctx).stream();
200 201
    memory::Copy(
        dst_npu_place, dst_ptr, src_npu_pinned_place, src_ptr, size, stream);
W
WangXi 已提交
202 203 204
  }
  else if (platform::is_npu_place(src_place) &&        // NOLINT
           platform::is_npu_pinned_place(dst_place)) { /* npu->npu_pinned */
205 206
    auto src_npu_place = src_place;
    auto dst_npu_pinned_place = dst_place;
W
WangXi 已提交
207
    auto ctx_place = ctx.GetPlace();
208 209 210 211 212 213 214
    PADDLE_ENFORCE_EQ(
        platform::is_npu_place(ctx_place),
        true,
        platform::errors::PreconditionNotMet(
            "Device context place mismatch. When copying phi::DenseTensor "
            "data from NPU memory to NPU Pinned memory, current "
            "device context place should be NPU."));
215
    auto ctx_npu_place = ctx_place;
216 217
    PADDLE_ENFORCE_EQ(src_place,
                      ctx_npu_place,
W
WangXi 已提交
218 219 220 221
                      platform::errors::PreconditionNotMet(
                          "The source NPU device and current device context do "
                          "not match. The source NPU device number is %d, but "
                          "device context NPU number is %d.",
222 223
                          src_npu_place.device,
                          ctx_npu_place.device));
W
WangXi 已提交
224 225
    auto stream =
        reinterpret_cast<const platform::NPUDeviceContext&>(ctx).stream();
226 227
    memory::Copy(
        dst_npu_pinned_place, dst_ptr, src_npu_place, src_ptr, size, stream);
W
WangXi 已提交
228
  }
229 230 231 232 233
  else {  // NOLINT
    PADDLE_THROW(platform::errors::Unimplemented(
        "Copy from %s to %s is not supported.", src_place, dst_place));
  }
#endif
234
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
235 236
  else if (platform::is_cuda_pinned_place(src_place) &&  // NOLINT
           platform::is_cuda_pinned_place(dst_place)) {
237
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
238
  }
239
  else if (platform::is_cuda_pinned_place(src_place) &&  // NOLINT
Y
Yi Wang 已提交
240
           platform::is_cpu_place(dst_place)) {
241
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
242 243 244
  }
  else if (platform::is_cpu_place(src_place) &&  // NOLINT
           platform::is_cuda_pinned_place(dst_place)) {
245
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
246 247 248
  }
  else if (platform::is_gpu_place(src_place) &&  // NOLINT
           platform::is_cpu_place(dst_place)) {
249 250
    auto src_gpu_place = src_place;
    auto dst_cpu_place = dst_place;
Y
Yi Wang 已提交
251
    auto ctx_place = ctx.GetPlace();
252
    PADDLE_ENFORCE_EQ(
253 254
        platform::is_gpu_place(ctx_place),
        true,
255 256 257
        platform::errors::PreconditionNotMet(
            "Context place error, excepted GPUPlace, but actually %s.",
            ctx_place));
258
    auto ctx_gpu_place = ctx_place;
259 260
    PADDLE_ENFORCE_EQ(src_gpu_place,
                      ctx_gpu_place,
261 262 263
                      platform::errors::Unavailable(
                          "Source place and context place do not match, source "
                          "place is %s, context place is %s.",
264 265
                          src_gpu_place,
                          ctx_gpu_place));
L
Leo Chen 已提交
266
    auto stream = reinterpret_cast<const phi::GPUContext&>(ctx).stream();
267
    memory::Copy(dst_cpu_place, dst_ptr, src_gpu_place, src_ptr, size, stream);
268 269 270
  }
  else if (platform::is_cpu_place(src_place) &&  // NOLINT
           platform::is_gpu_place(dst_place)) {
271 272
    auto src_cpu_place = src_place;
    auto dst_gpu_place = dst_place;
Y
Yi Wang 已提交
273
    auto ctx_place = ctx.GetPlace();
274
    PADDLE_ENFORCE_EQ(
275 276
        platform::is_gpu_place(ctx_place),
        true,
277 278 279
        platform::errors::PreconditionNotMet(
            "Context place error, excepted GPUPlace, but actually %s.",
            ctx_place));
280
    auto ctx_gpu_place = ctx_place;
281 282
    PADDLE_ENFORCE_EQ(dst_gpu_place,
                      ctx_gpu_place,
283 284 285
                      platform::errors::Unavailable(
                          "Destination place and context place do not match, "
                          "destination place is %s, context place is %s.",
286 287
                          dst_gpu_place,
                          ctx_gpu_place));
L
Leo Chen 已提交
288
    auto stream = reinterpret_cast<const phi::GPUContext&>(ctx).stream();
289
    memory::Copy(dst_gpu_place, dst_ptr, src_cpu_place, src_ptr, size, stream);
290 291 292
  }
  else if (platform::is_gpu_place(src_place) &&  // NOLINT
           platform::is_cuda_pinned_place(dst_place)) {
293 294
    auto src_gpu_place = src_place;
    auto dst_cuda_pinned_place = dst_place;
295
    auto ctx_place = ctx.GetPlace();
296 297 298 299 300 301 302
    PADDLE_ENFORCE_EQ(
        platform::is_gpu_place(ctx_place),
        true,
        platform::errors::PreconditionNotMet(
            "Device context place mismatch. When copying phi::DenseTensor "
            "data from GPU memory to CUDA Pinned memory, current "
            "device context place should be GPU."));
303
    auto ctx_gpu_place = ctx_place;
304 305
    PADDLE_ENFORCE_EQ(src_gpu_place,
                      ctx_gpu_place,
306 307 308 309
                      platform::errors::PreconditionNotMet(
                          "The source GPU device and current device context do "
                          "not match. The source GPU device number is %d, but "
                          "device context GPU number is %d.",
310 311
                          src_gpu_place.device,
                          ctx_gpu_place.device));
L
Leo Chen 已提交
312
    auto stream = reinterpret_cast<const phi::GPUContext&>(ctx).stream();
313 314
    memory::Copy(
        dst_cuda_pinned_place, dst_ptr, src_gpu_place, src_ptr, size, stream);
315 316 317
  }
  else if (platform::is_cuda_pinned_place(src_place) &&  // NOLINT
           platform::is_gpu_place(dst_place)) {
318 319
    auto src_cuda_pinned_place = src_place;
    auto dst_gpu_place = dst_place;
320
    auto ctx_place = ctx.GetPlace();
321 322 323 324 325 326 327
    PADDLE_ENFORCE_EQ(
        platform::is_gpu_place(ctx_place),
        true,
        platform::errors::PreconditionNotMet(
            "Device context place mismatch. When copying phi::DenseTensor "
            "data from CUDA Pinned memory to GPU memory, current "
            "device context place should be GPU."));
328
    auto ctx_gpu_place = ctx_place;
329 330
    PADDLE_ENFORCE_EQ(dst_gpu_place,
                      ctx_gpu_place,
331 332 333 334
                      platform::errors::PreconditionNotMet(
                          "The target GPU device and current device context do "
                          "not match. The target GPU device number is %d, but "
                          "device context GPU number is %d.",
335 336
                          dst_gpu_place.device,
                          ctx_gpu_place.device));
L
Leo Chen 已提交
337
    auto stream = reinterpret_cast<const phi::GPUContext&>(ctx).stream();
338 339
    memory::Copy(
        dst_gpu_place, dst_ptr, src_cuda_pinned_place, src_ptr, size, stream);
340 341 342
  }
  else if (platform::is_gpu_place(src_place) &&  // NOLINT
           platform::is_gpu_place(dst_place)) {
343 344
    auto src_gpu_place = src_place;
    auto dst_gpu_place = dst_place;
Y
Yi Wang 已提交
345
    auto ctx_place = ctx.GetPlace();
346
    PADDLE_ENFORCE_EQ(
347 348
        platform::is_gpu_place(ctx_place),
        true,
349 350 351
        platform::errors::PreconditionNotMet(
            "Context place error, excepted GPUPlace, but actually %s.",
            ctx_place));
L
Leo Chen 已提交
352
    auto stream = reinterpret_cast<const phi::GPUContext&>(ctx).stream();
C
chengduo 已提交
353
    if (platform::is_same_place(src_place, dst_place)) {
354 355
      memory::Copy(
          dst_gpu_place, dst_ptr, src_gpu_place, src_ptr, size, stream);
C
chengduo 已提交
356 357
    } else {
      if (platform::is_same_place(ctx_place, src_place)) {
358 359
        memory::Copy(
            dst_gpu_place, dst_ptr, src_gpu_place, src_ptr, size, stream);
C
chengduo 已提交
360
        platform::DeviceContextPool::Instance().Get(src.place())->Wait();
C
chengduo 已提交
361
      } else if (platform::is_same_place(ctx_place, dst_place)) {
C
chengduo 已提交
362
        platform::DeviceContextPool::Instance().Get(src.place())->Wait();
363 364
        memory::Copy(
            dst_gpu_place, dst_ptr, src_gpu_place, src_ptr, size, stream);
C
chengduo 已提交
365
      } else {
366 367
        PADDLE_THROW(platform::errors::Unavailable(
            "Context place dose not match the source and destination place."));
C
chengduo 已提交
368 369
      }
    }
370 371
  }
  else {  // NOLINT
372 373
    PADDLE_THROW(platform::errors::Unimplemented(
        "Copying from %s to %s is not supported.", src_place, dst_place));
Y
Yi Wang 已提交
374 375
  }
#endif
376 377 378
#ifdef PADDLE_WITH_MLU
  else if (platform::is_mlu_place(src_place) &&  // NOLINT
           platform::is_cpu_place(dst_place)) {
379 380
    auto src_mlu_place = src_place;
    auto dst_cpu_place = dst_place;
381 382 383 384 385 386
    auto stream =
        reinterpret_cast<const platform::MLUDeviceContext&>(ctx).stream();
    memory::Copy(dst_cpu_place, dst_ptr, src_mlu_place, src_ptr, size, stream);
  }
  else if (platform::is_cpu_place(src_place) &&  // NOLINT
           platform::is_mlu_place(dst_place)) {
387 388
    auto src_cpu_place = src_place;
    auto dst_mlu_place = dst_place;
389 390 391 392 393 394
    auto stream =
        reinterpret_cast<const platform::MLUDeviceContext&>(ctx).stream();
    memory::Copy(dst_mlu_place, dst_ptr, src_cpu_place, src_ptr, size, stream);
  }
  else if (platform::is_mlu_place(src_place) &&  // NOLINT
           platform::is_mlu_place(dst_place)) {
395 396
    auto src_mlu_place = src_place;
    auto dst_mlu_place = dst_place;
397 398 399 400 401 402 403 404 405
    auto stream =
        reinterpret_cast<const platform::MLUDeviceContext&>(ctx).stream();
    memory::Copy(dst_mlu_place, dst_ptr, src_mlu_place, src_ptr, size, stream);
  }
  else {  // NOLINT
    PADDLE_THROW(platform::errors::Unimplemented(
        "Copying from %s to %s is not supported.", src_place, dst_place));
  }
#endif
A
Allen Guo 已提交
406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428
#ifdef PADDLE_WITH_IPU
  else if (platform::is_ipu_place(src_place) &&  // NOLINT
           platform::is_cpu_place(dst_place)) {
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
  }
  else if (platform::is_cpu_place(src_place) &&  // NOLINT
           platform::is_ipu_place(dst_place)) {
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
  }
  else if (platform::is_ipu_place(src_place) &&  // NOLINT
           platform::is_ipu_place(dst_place)) {
    if (src_ptr == dst_ptr) {
      VLOG(3) << "Skip copy the same data sync from " << src_place << " to "
              << dst_place;
      return;
    }
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
  }
  else {  // NOLINT
    PADDLE_THROW(platform::errors::Unimplemented(
        "Copying from %s to %s is not supported.", src_place, dst_place));
  }
#endif
Y
Yi Wang 已提交
429 430
}

431
template <typename TENSOR>
432 433
void TensorCopyImpl(const TENSOR& src,
                    const platform::Place& dst_place,
434
                    TENSOR* dst) {
Y
Yi Wang 已提交
435 436
  platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
  const platform::DeviceContext* dev_ctx;
437
  if (platform::is_gpu_place(dst_place) || platform::is_npu_place(dst_place) ||
438 439
      platform::is_mlu_place(dst_place) ||
      platform::is_custom_place(dst_place)) {
Y
Yi Wang 已提交
440
    dev_ctx = pool.Get(dst_place);
C
chengduo 已提交
441 442
  } else {
    dev_ctx = pool.Get(src.place());
Y
Yi Wang 已提交
443
  }
444 445 446
  TensorCopyImpl(src, dst_place, *dev_ctx, dst);
}

447
void TensorCopy(const phi::DenseTensor& src,
448
                const platform::Place& dst_place,
449 450
                phi::DenseTensor* dst) {
  TensorCopyImpl<phi::DenseTensor>(src, dst_place, dst);
451
}
452
void TensorCopy(const phi::DenseTensor& src,
453 454
                const platform::Place& dst_place,
                const platform::DeviceContext& ctx,
455 456
                phi::DenseTensor* dst) {
  TensorCopyImpl<phi::DenseTensor>(src, dst_place, ctx, dst);
457
}
Y
Yi Wang 已提交
458

459
void TensorCopySync(const phi::DenseTensor& src,
460
                    const platform::Place& dst_place,
461
                    phi::DenseTensor* dst) {
462 463 464 465 466 467
  if (&src == dst) {
    auto src_copy = src;
    TensorCopySync(src_copy, dst_place, dst);
    return;
  }

M
minqiyang 已提交
468 469
  VLOG(3) << "TensorCopySync " << src.dims() << " from " << src.place()
          << " to " << dst_place;
F
fengjiayi 已提交
470 471 472
  src.check_memory_size();
  dst->Resize(src.dims());
  dst->set_layout(src.layout());
J
Jacek Czaja 已提交
473
#ifdef PADDLE_WITH_MKLDNN
474
  if (src.layout() == DataLayout::ONEDNN) {
475 476
    dst->set_mem_desc(src.mem_desc());
  }
J
Jacek Czaja 已提交
477
#endif
F
fengjiayi 已提交
478
  auto src_place = src.place();
479
  auto src_ptr = src.data();
480
  auto dst_ptr = dst->mutable_data(dst_place, src.dtype());
481
  VLOG(4) << "src:" << src_ptr << ", dst:" << dst_ptr;
482 483 484 485 486 487 488

  if (src_ptr == dst_ptr && src_place == dst_place) {
    VLOG(3) << "Skip copy the same data from " << src_place << " to "
            << dst_place;
    return;
  }

489
  auto size = src.numel() * framework::DataTypeSize(src.dtype());
F
fengjiayi 已提交
490
  if (platform::is_cpu_place(src_place) && platform::is_cpu_place(dst_place)) {
491
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
F
fengjiayi 已提交
492
  }
493 494 495 496
#ifdef PADDLE_WITH_CUSTOM_DEVICE
  else if (platform::is_custom_place(src_place) &&  // NOLINT
           platform::is_cpu_place(dst_place)) {     /* custom_device -> cpu*/
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size, nullptr);
A
Allen Guo 已提交
497
  }                                                // NOLINT
498 499 500
  else if (platform::is_cpu_place(src_place) &&    // NOLINT
           platform::is_custom_place(dst_place)) { /* cpu -> custom_device*/
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size, nullptr);
A
Allen Guo 已提交
501
  }                                                 // NOLINT
502 503 504 505 506 507 508 509 510 511 512
  else if (platform::is_custom_place(src_place) &&  // NOLINT
           platform::is_custom_place(
               dst_place)) { /* custom_device -> custom_device*/
    if (src_ptr == dst_ptr) {
      VLOG(3) << "Skip copy the same data sync from " << src_place << " to "
              << dst_place;
      return;
    }
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size, nullptr);
  }
#endif
513 514 515
#ifdef PADDLE_WITH_XPU
  else if (platform::is_xpu_place(src_place) &&  // NOLINT
           platform::is_cpu_place(dst_place)) {
516
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
A
Allen Guo 已提交
517
  }                                              // NOLINT
J
jianghaicheng 已提交
518 519
  else if (platform::is_cpu_place(src_place) &&  // NOLINT
           platform::is_xpu_place(dst_place)) {
520
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
A
Allen Guo 已提交
521
  }                                              // NOLINT
J
jianghaicheng 已提交
522 523
  else if (platform::is_xpu_place(src_place) &&  // NOLINT
           platform::is_xpu_place(dst_place)) {
524 525 526 527 528
    if (src_ptr == dst_ptr) {
      VLOG(3) << "Skip copy the same data async from " << src_place << " to "
              << dst_place;
      return;
    }
529 530 531
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
    platform::XPUPlace xpu_dst_place = dst_place;
    platform::XPUPlace xpu_src_place = src_place;
532 533 534 535
    if (xpu_dst_place.device == xpu_src_place.device) {
      auto xpu_ctx = platform::DeviceContextPool::Instance().Get(xpu_dst_place);
      xpu_ctx->Wait();
    }
A
Allen Guo 已提交
536
  }       // NOLINT
J
jianghaicheng 已提交
537
  else {  // NOLINT
538 539 540 541
    PADDLE_THROW(platform::errors::Unimplemented(
        "Copy from %s to %s is not supported.", src_place, dst_place));
  }
#endif
542 543 544
#ifdef PADDLE_WITH_ASCEND_CL
  else if (platform::is_npu_place(src_place) &&  // NOLINT
           platform::is_cpu_place(dst_place)) {  /* npu -> cpu*/
545
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size, nullptr);
546 547 548
  }
  else if (platform::is_cpu_place(src_place) &&  // NOLINT
           platform::is_npu_place(dst_place)) {  /* cpu -> npu*/
549
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size, nullptr);
550 551 552 553 554 555 556 557
  }
  else if (platform::is_npu_place(src_place) &&  // NOLINT
           platform::is_npu_place(dst_place)) {  /* npu -> npu*/
    if (src_ptr == dst_ptr) {
      VLOG(3) << "Skip copy the same data sync from " << src_place << " to "
              << dst_place;
      return;
    }
558
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size, nullptr);
559 560 561 562 563 564
  }
  else {  // NOLINT
    PADDLE_THROW(platform::errors::Unimplemented(
        "Copy from %s to %s is not supported.", src_place, dst_place));
  }
#endif
565
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
566 567
  else if (platform::is_cuda_pinned_place(src_place) &&  // NOLINT
           platform::is_cuda_pinned_place(dst_place)) {
568
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
569
  }
570
  else if (platform::is_cuda_pinned_place(src_place) &&  // NOLINT
F
fengjiayi 已提交
571
           platform::is_cpu_place(dst_place)) {
572
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
573 574 575
  }
  else if (platform::is_cpu_place(src_place) &&  // NOLINT
           platform::is_cuda_pinned_place(dst_place)) {
576
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
577 578 579
  }
  else if (platform::is_gpu_place(src_place) &&  // NOLINT
           platform::is_cuda_pinned_place(dst_place)) {
580
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size, nullptr);
581 582 583
  }
  else if (platform::is_gpu_place(src_place) &&  // NOLINT
           platform::is_cpu_place(dst_place)) {
584 585
    auto src_gpu_place = src_place;
    auto dst_cpu_place = dst_place;
F
fengjiayi 已提交
586
    memory::Copy(dst_cpu_place, dst_ptr, src_gpu_place, src_ptr, size, nullptr);
587 588 589
  }
  else if (platform::is_cpu_place(src_place) &&  // NOLINT
           platform::is_gpu_place(dst_place)) {
590 591
    auto src_cpu_place = src_place;
    auto dst_gpu_place = dst_place;
F
fengjiayi 已提交
592
    memory::Copy(dst_gpu_place, dst_ptr, src_cpu_place, src_ptr, size, nullptr);
593 594 595
  }
  else if (platform::is_gpu_place(src_place) &&  // NOLINT
           platform::is_gpu_place(dst_place)) {
596 597
    auto src_gpu_place = src_place;
    auto dst_gpu_place = dst_place;
F
fengjiayi 已提交
598
    memory::Copy(dst_gpu_place, dst_ptr, src_gpu_place, src_ptr, size, nullptr);
599 600 601
  }
  else if (platform::is_cuda_pinned_place(src_place) &&  // NOLINT
           platform::is_gpu_place(dst_place)) {
602 603
    auto src_pinned_place = src_place;
    auto dst_gpu_place = dst_place;
604 605
    memory::Copy(
        dst_gpu_place, dst_ptr, src_pinned_place, src_ptr, size, nullptr);
606 607
  }
  else {  // NOLINT
608 609
    PADDLE_THROW(platform::errors::Unimplemented(
        "Copy from %s to %s is not supported.", src_place, dst_place));
F
fengjiayi 已提交
610 611
  }
#endif
612 613 614
#ifdef PADDLE_WITH_MLU
  else if (platform::is_mlu_place(src_place) &&  // NOLINT
           platform::is_cpu_place(dst_place)) {
615
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size, nullptr);
616 617 618
  }
  else if (platform::is_cpu_place(src_place) &&  // NOLINT
           platform::is_mlu_place(dst_place)) {
619
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size, nullptr);
620 621 622 623 624 625 626 627
  }
  else if (platform::is_mlu_place(src_place) &&  // NOLINT
           platform::is_mlu_place(dst_place)) {
    if (src_ptr == dst_ptr) {
      VLOG(3) << "Skip copy the same data async from " << src_place << " to "
              << dst_place;
      return;
    }
628
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size, nullptr);
629 630 631 632 633 634
  }
  else {  // NOLINT
    PADDLE_THROW(platform::errors::Unimplemented(
        "Copy from %s to %s is not supported.", src_place, dst_place));
  }
#endif
A
Allen Guo 已提交
635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657
#ifdef PADDLE_WITH_IPU
  else if (platform::is_ipu_place(src_place) &&  // NOLINT
           platform::is_cpu_place(dst_place)) {
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
  }
  else if (platform::is_cpu_place(src_place) &&  // NOLINT
           platform::is_ipu_place(dst_place)) {
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
  }
  else if (platform::is_ipu_place(src_place) &&  // NOLINT
           platform::is_ipu_place(dst_place)) {
    if (src_ptr == dst_ptr) {
      VLOG(3) << "Skip copy the same data sync from " << src_place << " to "
              << dst_place;
      return;
    }
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
  }
  else {  // NOLINT
    PADDLE_THROW(platform::errors::Unimplemented(
        "Copy from %s to %s is not supported.", src_place, dst_place));
  }
#endif
F
fengjiayi 已提交
658 659
}

660
void TensorToStream(std::ostream& os,
661
                    const phi::DenseTensor& tensor,
Y
Yi Wang 已提交
662 663 664 665 666 667 668 669 670
                    const platform::DeviceContext& dev_ctx) {
  {  // the 1st field, uint32_t version
    constexpr uint32_t version = 0;
    os.write(reinterpret_cast<const char*>(&version), sizeof(version));
  }
  {  // the 2nd field, tensor description
     // int32_t  size
     // void*    protobuf message
    proto::VarType::TensorDesc desc;
671
    desc.set_data_type(framework::TransToProtoVarType(tensor.dtype()));
672
    auto dims = phi::vectorize(tensor.dims());
Y
Yi Wang 已提交
673 674 675 676 677 678 679 680 681
    auto* pb_dims = desc.mutable_dims();
    pb_dims->Resize(static_cast<int>(dims.size()), 0);
    std::copy(dims.begin(), dims.end(), pb_dims->begin());
    int32_t size = desc.ByteSize();
    os.write(reinterpret_cast<const char*>(&size), sizeof(size));
    auto out = desc.SerializeAsString();
    os.write(out.data(), size);
  }
  {  // the 3rd field, tensor data
682
    uint64_t size = tensor.numel() * framework::DataTypeSize(tensor.dtype());
Y
yuyang18 已提交
683

684
    auto* data_ptr = tensor.data();
685 686
    PADDLE_ENFORCE_LT(size,
                      (std::numeric_limits<std::streamsize>::max)(),
T
tangwei12 已提交
687 688
                      platform::errors::ResourceExhausted(
                          "tensor size %d overflow when writing tensor", size));
Y
Yi Wang 已提交
689
    if (platform::is_gpu_place(tensor.place())) {
690
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
Y
Yi Wang 已提交
691 692
      constexpr size_t kBufSize = 1024 * 1024 * 64;  // 64MB
      std::unique_ptr<char[]> buf(new char[kBufSize]);
L
Leo Chen 已提交
693
      auto& gpu_dev_ctx = static_cast<const phi::GPUContext&>(dev_ctx);
Y
Yi Wang 已提交
694 695 696 697
      platform::CPUPlace cpu;
      uintptr_t data = reinterpret_cast<uintptr_t>(data_ptr);
      while (size != 0) {
        size_t size_to_write = std::min(kBufSize, static_cast<size_t>(size));
698 699 700 701 702
        memory::Copy(cpu,
                     buf.get(),
                     tensor.place(),
                     reinterpret_cast<const void*>(data),
                     size_to_write,
Y
Yi Wang 已提交
703 704 705 706 707 708 709
                     gpu_dev_ctx.stream());
        gpu_dev_ctx.Wait();
        os.write(buf.get(), size_to_write);
        data += size_to_write;
        size -= size_to_write;
      }
#else
T
tangwei12 已提交
710 711
      PADDLE_THROW(platform::errors::Unimplemented(
          "CUDAPlace is not supported when not compiled with CUDA"));
712 713 714 715 716 717 718 719 720 721 722
#endif
    } else if (platform::is_xpu_place(tensor.place())) {
#ifdef PADDLE_WITH_XPU
      constexpr size_t kBufSize = 1024 * 1024 * 64;  // 64MB
      std::unique_ptr<char[]> buf(new char[kBufSize]);
      auto& xpu_dev_ctx =
          static_cast<const platform::XPUDeviceContext&>(dev_ctx);
      platform::CPUPlace cpu;
      uintptr_t data = reinterpret_cast<uintptr_t>(data_ptr);
      while (size != 0) {
        size_t size_to_write = std::min(kBufSize, static_cast<size_t>(size));
723 724 725 726 727
        memory::Copy(cpu,
                     buf.get(),
                     tensor.place(),
                     reinterpret_cast<const void*>(data),
                     size_to_write);
728 729 730 731 732 733 734 735
        xpu_dev_ctx.Wait();
        os.write(buf.get(), size_to_write);
        data += size_to_write;
        size -= size_to_write;
      }
#else
      PADDLE_THROW(platform::errors::Unimplemented(
          "XPUPlace is not supported when not compiled with XPU"));
736 737 738 739 740 741 742 743 744 745 746
#endif
    } else if (platform::is_mlu_place(tensor.place())) {
#ifdef PADDLE_WITH_MLU
      constexpr size_t kBufSize = 1024 * 1024 * 64;  // 64MB
      std::unique_ptr<char[]> buf(new char[kBufSize]);
      auto& mlu_dev_ctx =
          static_cast<const platform::MLUDeviceContext&>(dev_ctx);
      platform::CPUPlace cpu;
      uintptr_t data = reinterpret_cast<uintptr_t>(data_ptr);
      while (size != 0) {
        size_t size_to_write = std::min(kBufSize, static_cast<size_t>(size));
747 748 749 750 751
        memory::Copy(cpu,
                     buf.get(),
                     tensor.place(),
                     reinterpret_cast<const void*>(data),
                     size_to_write,
752 753 754 755 756 757 758 759 760
                     mlu_dev_ctx.stream());
        mlu_dev_ctx.Wait();
        os.write(buf.get(), size_to_write);
        data += size_to_write;
        size -= size_to_write;
      }
#else
      PADDLE_THROW(platform::errors::Unimplemented(
          "MLUPlace is not supported when not compiled with MLU"));
761 762 763 764 765 766 767 768 769 770 771
#endif
    } else if (platform::is_npu_place(tensor.place())) {
#ifdef PADDLE_WITH_ASCEND_CL
      constexpr size_t kBufSize = 1024 * 1024 * 64;  // 64MB
      std::unique_ptr<char[]> buf(new char[kBufSize]);
      auto& npu_dev_ctx =
          static_cast<const platform::NPUDeviceContext&>(dev_ctx);
      platform::CPUPlace cpu;
      uintptr_t data = reinterpret_cast<uintptr_t>(data_ptr);
      while (size != 0) {
        size_t size_to_write = std::min(kBufSize, static_cast<size_t>(size));
772 773 774 775 776
        memory::Copy(cpu,
                     buf.get(),
                     tensor.place(),
                     reinterpret_cast<const void*>(data),
                     size_to_write,
777 778 779 780 781 782 783 784 785
                     npu_dev_ctx.stream());
        npu_dev_ctx.Wait();
        os.write(buf.get(), size_to_write);
        data += size_to_write;
        size -= size_to_write;
      }
#else
      PADDLE_THROW(platform::errors::Unimplemented(
          "NPUPlace is not supported when not compiled with NPU"));
786 787 788 789 790 791 792 793 794 795 796
#endif
    } else if (platform::is_custom_place(tensor.place())) {
#ifdef PADDLE_WITH_CUSTOM_DEVICE
      constexpr size_t kBufSize = 1024 * 1024 * 64;  // 64MB
      std::unique_ptr<char[]> buf(new char[kBufSize]);
      auto& custom_device_context =
          static_cast<const platform::CustomDeviceContext&>(dev_ctx);
      platform::CPUPlace cpu;
      uintptr_t data = reinterpret_cast<uintptr_t>(data_ptr);
      while (size != 0) {
        size_t size_to_write = std::min(kBufSize, static_cast<size_t>(size));
797 798 799 800 801
        memory::Copy(cpu,
                     buf.get(),
                     tensor.place(),
                     reinterpret_cast<const void*>(data),
                     size_to_write,
802 803 804 805 806 807 808 809 810 811
                     custom_device_context.stream());
        custom_device_context.Wait();
        os.write(buf.get(), size_to_write);
        data += size_to_write;
        size -= size_to_write;
      }
#else
      PADDLE_THROW(platform::errors::Unimplemented(
          "CustomPlace is not supported when not compiled with "
          "CustomDevice"));
Y
Yi Wang 已提交
812 813 814 815 816 817 818 819 820
#endif
    } else {
      os.write(static_cast<const char*>(data_ptr),
               static_cast<std::streamsize>(size));
    }
  }
}

struct DeserializedDataFunctor {
821
  DeserializedDataFunctor(void** buf,
822
                          phi::DenseTensor* tensor,
Y
Yi Wang 已提交
823 824 825 826
                          const platform::Place& place)
      : buf_(buf), tensor_(tensor), place_(place) {}

  template <typename T>
D
dzhwinter 已提交
827
  void apply() {
Y
Yi Wang 已提交
828 829 830 831
    *buf_ = tensor_->mutable_data<T>(place_);
  }

  void** buf_;
832
  phi::DenseTensor* tensor_;
Y
Yi Wang 已提交
833 834 835
  platform::Place place_;
};

836
void TensorFromStream(std::istream& is,
837
                      phi::DenseTensor* tensor,
T
tangwei12 已提交
838
                      const platform::DeviceContext& dev_ctx,
839 840
                      const size_t& seek,
                      const std::vector<int64_t>& shape) {
T
tangwei12 已提交
841 842 843 844
  uint32_t version;
  is.read(reinterpret_cast<char*>(&version), sizeof(version));

  PADDLE_ENFORCE_EQ(
845 846
      version,
      0U,
T
tangwei12 已提交
847 848 849 850 851 852 853 854 855 856 857 858
      platform::errors::InvalidArgument(
          "tensor version %u is not supported, Only version 0 is supported",
          version));

  proto::VarType::TensorDesc desc;
  {  // int32_t size
    // proto buffer
    int32_t size;
    is.read(reinterpret_cast<char*>(&size), sizeof(size));
    std::unique_ptr<char[]> buf(new char[size]);
    is.read(reinterpret_cast<char*>(buf.get()), size);
    PADDLE_ENFORCE_EQ(
859 860
        desc.ParseFromArray(buf.get(), size),
        true,
T
tangwei12 已提交
861 862 863
        platform::errors::InvalidArgument("Cannot parse tensor desc"));
  }
  {  // read tensor
864
    tensor->Resize(phi::make_ddim(shape));
T
tangwei12 已提交
865 866 867 868
    size_t seekg = seek * framework::SizeOfType(desc.data_type());
    is.seekg(seekg, is.cur);

    void* buf;
L
Leo Chen 已提交
869
    phi::CPUContext ctx;
T
tangwei12 已提交
870
    size_t size = tensor->numel() * framework::SizeOfType(desc.data_type());
871
    if (platform::is_gpu_place(dev_ctx.GetPlace()) ||
872
        platform::is_xpu_place(dev_ctx.GetPlace()) ||
873
        platform::is_mlu_place(dev_ctx.GetPlace()) ||
874 875
        platform::is_npu_place(dev_ctx.GetPlace()) ||
        platform::is_custom_place(dev_ctx.GetPlace())) {
876
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
877
    defined(PADDLE_WITH_XPU) || defined(PADDLE_WITH_MLU) ||  \
878
    defined(PADDLE_WITH_ASCEND_CL) || defined(PADDLE_WITH_CUSTOM_DEVICE)
879
      phi::DenseTensor cpu_tensor;
880
      cpu_tensor.Resize(phi::make_ddim(shape));
T
tangwei12 已提交
881 882 883 884 885 886
      framework::VisitDataType(
          desc.data_type(),
          DeserializedDataFunctor(&buf, &cpu_tensor, ctx.GetPlace()));
      is.read(static_cast<char*>(buf), size);
      auto dst_place = dev_ctx.GetPlace();
      framework::TensorCopy(cpu_tensor, dst_place, dev_ctx, tensor);
887 888
      if (platform::is_npu_place(dev_ctx.GetPlace()) ||
          platform::is_custom_place(dev_ctx.GetPlace())) {
889 890
        dev_ctx.Wait();
      }
T
tangwei12 已提交
891
#else
892 893 894
      if (platform::is_gpu_place(dev_ctx.GetPlace())) {
        PADDLE_THROW(platform::errors::Unimplemented(
            "CUDAPlace is not supported when not compiled with CUDA"));
895
      } else if (platform::is_xpu_place(dev_ctx.GetPlace())) {
896 897
        PADDLE_THROW(platform::errors::Unimplemented(
            "XPUPlace is not supported when not compiled with XPU"));
898 899 900
      } else if (platform::is_mlu_place(dev_ctx.GetPlace())) {
        PADDLE_THROW(platform::errors::Unimplemented(
            "MLUPlace is not supported when not compiled with MLU"));
901 902 903
      } else {
        PADDLE_THROW(platform::errors::Unimplemented(
            "NPUPlace is not supported when not compiled with NPU"));
904
      }
T
tangwei12 已提交
905 906 907 908 909 910 911 912 913 914
#endif
    } else {
      framework::VisitDataType(
          desc.data_type(),
          DeserializedDataFunctor(&buf, tensor, ctx.GetPlace()));
      is.read(static_cast<char*>(buf), size);
    }
  }
}

915
void TensorFromStream(std::istream& is,
916
                      phi::DenseTensor* tensor,
Y
Yi Wang 已提交
917 918 919
                      const platform::DeviceContext& dev_ctx) {
  uint32_t version;
  is.read(reinterpret_cast<char*>(&version), sizeof(version));
T
tangwei12 已提交
920
  PADDLE_ENFORCE_EQ(
921 922
      version,
      0U,
T
tangwei12 已提交
923 924 925
      platform::errors::InvalidArgument(
          "tensor version %u is not supported, Only version 0 is supported",
          version));
Y
Yi Wang 已提交
926 927 928
  proto::VarType::TensorDesc desc;
  {  // int32_t size
     // proto buffer
Z
zlsh80826 已提交
929
    int32_t size = -1;
Y
Yi Wang 已提交
930
    is.read(reinterpret_cast<char*>(&size), sizeof(size));
931
    PADDLE_ENFORCE_EQ(
932 933
        is.good(),
        true,
934
        platform::errors::Unavailable("Cannot read tensor desc size"));
935 936 937 938
    PADDLE_ENFORCE_GE(size,
                      0,
                      platform::errors::InvalidArgument(
                          "phi::DenseTensor desc size should >= 0"));
Y
Yi Wang 已提交
939 940
    std::unique_ptr<char[]> buf(new char[size]);
    is.read(reinterpret_cast<char*>(buf.get()), size);
T
tangwei12 已提交
941
    PADDLE_ENFORCE_EQ(
942 943
        desc.ParseFromArray(buf.get(), size),
        true,
T
tangwei12 已提交
944
        platform::errors::InvalidArgument("Cannot parse tensor desc"));
Y
Yi Wang 已提交
945 946 947 948 949
  }
  {  // read tensor
    std::vector<int64_t> dims;
    dims.reserve(static_cast<size_t>(desc.dims().size()));
    std::copy(desc.dims().begin(), desc.dims().end(), std::back_inserter(dims));
950
    tensor->Resize(phi::make_ddim(dims));
Y
Yi Wang 已提交
951
    void* buf;
L
Leo Chen 已提交
952
    phi::CPUContext ctx;
Y
Yu Yang 已提交
953
    size_t size = tensor->numel() * framework::SizeOfType(desc.data_type());
954
    if (platform::is_gpu_place(dev_ctx.GetPlace()) ||
955
        platform::is_xpu_place(dev_ctx.GetPlace()) ||
956
        platform::is_mlu_place(dev_ctx.GetPlace()) ||
957 958
        platform::is_npu_place(dev_ctx.GetPlace()) ||
        platform::is_custom_place(dev_ctx.GetPlace())) {
959
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
960
    defined(PADDLE_WITH_XPU) || defined(PADDLE_WITH_MLU) ||  \
961
    defined(PADDLE_WITH_ASCEND_CL) || defined(PADDLE_WITH_CUSTOM_DEVICE)
962
      phi::DenseTensor cpu_tensor;
963
      cpu_tensor.Resize(phi::make_ddim(dims));
Y
Yi Wang 已提交
964 965 966
      framework::VisitDataType(
          desc.data_type(),
          DeserializedDataFunctor(&buf, &cpu_tensor, ctx.GetPlace()));
Y
yuyang18 已提交
967
      is.read(static_cast<char*>(buf), size);
Y
Yi Wang 已提交
968 969
      auto dst_place = dev_ctx.GetPlace();
      framework::TensorCopy(cpu_tensor, dst_place, dev_ctx, tensor);
970 971
      if (platform::is_npu_place(dev_ctx.GetPlace()) ||
          platform::is_custom_place(dev_ctx.GetPlace())) {
972 973
        dev_ctx.Wait();
      }
Y
Yi Wang 已提交
974
#else
975 976 977
      if (platform::is_gpu_place(dev_ctx.GetPlace())) {
        PADDLE_THROW(platform::errors::Unimplemented(
            "CUDAPlace is not supported when not compiled with CUDA"));
978
      } else if (platform::is_xpu_place(dev_ctx.GetPlace())) {
979 980
        PADDLE_THROW(platform::errors::Unimplemented(
            "XPUPlace is not supported when not compiled with XPU"));
981 982 983
      } else if (platform::is_mlu_place(dev_ctx.GetPlace())) {
        PADDLE_THROW(platform::errors::Unimplemented(
            "MLUPlace is not supported when not compiled with MLU"));
984
      } else if (platform::is_npu_place(dev_ctx.GetPlace())) {
985 986
        PADDLE_THROW(platform::errors::Unimplemented(
            "NPUPlace is not supported when not compiled with NPU"));
987 988 989
      } else {
        PADDLE_THROW(platform::errors::Unimplemented(
            "CutomPlace is not supported when not compiled with CustomDevice"));
990
      }
Y
Yi Wang 已提交
991 992 993 994 995
#endif
    } else {
      framework::VisitDataType(
          desc.data_type(),
          DeserializedDataFunctor(&buf, tensor, ctx.GetPlace()));
Y
yuyang18 已提交
996
      is.read(static_cast<char*>(buf), size);
Y
Yi Wang 已提交
997 998 999 1000
    }
  }
}

6
633WHU 已提交
1001
// get tensor data point by DLDataType
1002
void* GetDstPtrByDLDataType(DLDataType type,
1003
                            phi::DenseTensor* dst,
6
633WHU 已提交
1004 1005
                            const platform::Place& dst_place) {
  // vector types not currently supported
1006 1007
  PADDLE_ENFORCE_LE(type.lanes,
                    1,
1008 1009
                    platform::errors::Unimplemented(
                        "Vector type is not supported currently."));
6
633WHU 已提交
1010 1011 1012 1013 1014 1015 1016

  switch (type.bits) {
    case 8:
      if (type.code == kDLInt)
        return static_cast<void*>(dst->mutable_data<int8_t>(dst_place));
      if (type.code == kDLUInt)
        return static_cast<void*>(dst->mutable_data<uint8_t>(dst_place));
1017 1018
      PADDLE_THROW(platform::errors::Unimplemented(
          "DLDataType code <%d> is illegal when DLDataType.bits is <%d>.",
1019 1020
          type.code,
          type.bits));
6
633WHU 已提交
1021 1022 1023 1024 1025 1026
    case 16:
      if (type.code == kDLInt)
        return static_cast<void*>(dst->mutable_data<int16_t>(dst_place));
      if (type.code == kDLFloat)
        return static_cast<void*>(
            dst->mutable_data<paddle::platform::float16>(dst_place));
S
Siming Dai 已提交
1027 1028 1029
      if (type.code == kDLBfloat)
        return static_cast<void*>(
            dst->mutable_data<paddle::platform::bfloat16>(dst_place));
1030 1031
      PADDLE_THROW(platform::errors::Unimplemented(
          "DLDataType code <%d> is illegal when DLDataType.bits is <%d>.",
1032 1033
          type.code,
          type.bits));
6
633WHU 已提交
1034 1035 1036 1037 1038
    case 32:
      if (type.code == kDLInt)
        return static_cast<void*>(dst->mutable_data<int32_t>(dst_place));
      if (type.code == kDLFloat)
        return static_cast<void*>(dst->mutable_data<float>(dst_place));
1039 1040
      PADDLE_THROW(platform::errors::Unimplemented(
          "DLDataType code <%d> is illegal when DLDataType.bits is <%d>.",
1041 1042
          type.code,
          type.bits));
6
633WHU 已提交
1043 1044 1045 1046 1047
    case 64:
      if (type.code == kDLInt)
        return static_cast<void*>(dst->mutable_data<int64_t>(dst_place));
      if (type.code == kDLFloat)
        return static_cast<void*>(dst->mutable_data<double>(dst_place));
S
Siming Dai 已提交
1048 1049 1050 1051 1052
      if (type.code == kDLComplex)
        return static_cast<void*>(
            dst->mutable_data<paddle::platform::complex<float>>(dst_place));
      PADDLE_THROW(platform::errors::Unimplemented(
          "DLDataType code <%d> is illegal when DLDataType.bits is <%d>.",
1053 1054
          type.code,
          type.bits));
S
Siming Dai 已提交
1055 1056 1057 1058
    case 128:
      if (type.code == kDLComplex)
        return static_cast<void*>(
            dst->mutable_data<paddle::platform::complex<double>>(dst_place));
1059 1060
      PADDLE_THROW(platform::errors::Unimplemented(
          "DLDataType code <%d> is illegal when DLDataType.bits is <%d>.",
1061 1062
          type.code,
          type.bits));
6
633WHU 已提交
1063
    default:
1064 1065
      PADDLE_THROW(platform::errors::Unimplemented(
          "Unsupported DLDataType.bits %d.", type.bits));
6
633WHU 已提交
1066 1067 1068
  }
}

1069
void TensorFromDLPack(const ::DLTensor& dl_tensor, phi::DenseTensor* dst) {
6
633WHU 已提交
1070 1071 1072 1073
  platform::CPUPlace dst_place = platform::CPUPlace();
  platform::CPUPlace src_place = platform::CPUPlace();

  std::vector<int64_t> vec;
1074 1075
  std::copy(dl_tensor.shape,
            dl_tensor.shape + dl_tensor.ndim,
6
633WHU 已提交
1076 1077
            std::back_inserter(vec));

1078
  framework::DDim vddim = phi::make_ddim(vec);
6
633WHU 已提交
1079 1080 1081 1082 1083 1084

  dst->Resize(vddim);
  ::DLDataType type = dl_tensor.dtype;
  void* dst_ptr = GetDstPtrByDLDataType(type, dst, dst_place);

  auto src_ptr = static_cast<const void*>(dl_tensor.data);
1085
  auto size = phi::product(vddim) * type.bits / 8;
6
633WHU 已提交
1086

S
Siming Dai 已提交
1087
  if (dl_tensor.device.device_type == kDLCPU) {
6
633WHU 已提交
1088 1089
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
  }
1090
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
S
Siming Dai 已提交
1091
  if (dl_tensor.device.device_type == kDLGPU) {
6
633WHU 已提交
1092
    platform::CUDAPlace dst_place =
S
Siming Dai 已提交
1093
        platform::CUDAPlace(dl_tensor.device.device_id);
6
633WHU 已提交
1094
    platform::CUDAPlace src_place =
S
Siming Dai 已提交
1095
        platform::CUDAPlace(dl_tensor.device.device_id);
6
633WHU 已提交
1096 1097
    dst_ptr = GetDstPtrByDLDataType(type, dst, dst_place);
    auto* ctx = platform::DeviceContextPool::Instance().GetByPlace(dst_place);
L
Leo Chen 已提交
1098 1099 1100 1101 1102 1103
    memory::Copy(dst_place,
                 dst_ptr,
                 src_place,
                 src_ptr,
                 size,
                 reinterpret_cast<const phi::GPUContext&>(*ctx).stream());
6
633WHU 已提交
1104 1105
  }
#endif
1106 1107 1108
#ifdef PADDLE_WITH_XPU
  PADDLE_THROW(platform::errors::Unimplemented("XPUPlace is not supported"));
#endif
6
633WHU 已提交
1109 1110
}

S
Siming Dai 已提交
1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146 1147 1148 1149 1150 1151 1152
void TensorFromDLPack(const DLManagedTensor* src, phi::DenseTensor* dst) {
  std::vector<int64_t> vec;
  std::copy(src->dl_tensor.shape,
            src->dl_tensor.shape + src->dl_tensor.ndim,
            std::back_inserter(vec));

  framework::DDim vddim = phi::make_ddim(vec);
  dst->Resize(vddim);
  ::DLDataType type = src->dl_tensor.dtype;

  auto src_ptr = static_cast<const void*>(src->dl_tensor.data);
  auto size = phi::product(vddim) * type.bits / 8;

  if (src->dl_tensor.device.device_type == kDLCPU) {
    platform::CPUPlace dst_place = platform::CPUPlace();
    platform::CPUPlace src_place = platform::CPUPlace();
    void* dst_ptr = GetDstPtrByDLDataType(type, dst, dst_place);
    memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
  }
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
  if (src->dl_tensor.device.device_type == kDLGPU) {
    platform::CUDAPlace dst_place =
        platform::CUDAPlace(src->dl_tensor.device.device_id);
    platform::CUDAPlace src_place =
        platform::CUDAPlace(src->dl_tensor.device.device_id);
    void* dst_ptr = GetDstPtrByDLDataType(type, dst, dst_place);
    auto* ctx = platform::DeviceContextPool::Instance().GetByPlace(dst_place);
    // Fix copy by share allocation.
    memory::Copy(dst_place,
                 dst_ptr,
                 src_place,
                 src_ptr,
                 size,
                 reinterpret_cast<const phi::GPUContext&>(*ctx).stream());
  }
#endif
  src->deleter(const_cast<DLManagedTensor*>(src));
#ifdef PADDLE_WITH_XPU
  PADDLE_THROW(platform::errors::Unimplemented("XPUPlace is not supported"));
#endif
}

1153
template <typename T>
1154
std::string format_tensor(const phi::DenseTensor& tensor) {
1155 1156 1157 1158
  // TODO(zhiqiu): use the print option to format tensor.
  return "NOT IMPLEMENTED";
}

1159
template <typename T>
1160
std::ostream& print_tensor(std::ostream& os, const phi::DenseTensor& tensor) {
1161 1162 1163
  auto inspect = tensor.data<T>();
  auto element_num = tensor.numel();

1164
  os << "  - data: [";
1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178
  // Note: int8_t && uint8_t is typedf of char, ostream unable to print properly
  if (typeid(int8_t) == typeid(T) || typeid(uint8_t) == typeid(T)) {
    if (element_num > 0) {
      os << signed(inspect[0]);
      for (int j = 1; j < element_num; ++j) {
        os << " " << signed(inspect[j]);
      }
    }
  } else {
    if (element_num > 0) {
      os << inspect[0];
      for (int j = 1; j < element_num; ++j) {
        os << " " << inspect[j];
      }
1179 1180 1181 1182 1183 1184
    }
  }
  os << "]";
  return os;
}

1185
template <>
1186
std::ostream& print_tensor<paddle::platform::complex<float>>(
1187
    std::ostream& os, const phi::DenseTensor& tensor) {
1188
  auto inspect = tensor.data<paddle::platform::complex<float>>();
1189 1190 1191 1192
  auto element_num = tensor.numel();

  os << "  - data: [";
  if (element_num > 0) {
1193
    os << signed(inspect[0].real) << "+" << signed(inspect[0].imag) << "j";
1194
    for (int j = 1; j < element_num; ++j) {
1195 1196
      os << " " << signed(inspect[j].real) << "+" << signed(inspect[j].imag)
         << "j";
1197 1198 1199 1200 1201 1202 1203
    }
  }
  os << "]";
  return os;
}

template <>
1204
std::ostream& print_tensor<paddle::platform::complex<double>>(
1205
    std::ostream& os, const phi::DenseTensor& tensor) {
1206
  auto inspect = tensor.data<paddle::platform::complex<double>>();
1207 1208 1209 1210
  auto element_num = tensor.numel();

  os << "  - data: [";
  if (element_num > 0) {
1211
    os << signed(inspect[0].real) << "+" << signed(inspect[0].imag) << "j";
1212
    for (int j = 1; j < element_num; ++j) {
1213 1214
      os << " " << signed(inspect[j].real) << "+" << signed(inspect[j].imag)
         << "j";
1215 1216 1217 1218 1219 1220
    }
  }
  os << "]";
  return os;
}

1221
std::ostream& operator<<(std::ostream& os, const LoD& lod) {
1222 1223
  // NOTE(xiongkun):
  // https://stackoverflow.com/questions/5195512/namespaces-and-operator-resolution
1224
  // if we don't redefine, the operator << of phi / framework LoD is not found.
1225
  paddle::string::operator<<(os, lod);
1226 1227 1228
  return os;
}

1229 1230 1231
}  // namespace framework
}  // namespace paddle

1232
namespace phi {
1233

1234 1235 1236 1237 1238
std::ostream& operator<<(std::ostream& os, const LoD& lod) {
  paddle::string::operator<<(os, lod);
  return os;
}

1239
std::ostream& operator<<(std::ostream& os, const phi::DenseTensor& t) {
1240 1241 1242 1243
  if (t.lod().size() > 0) {
    os << "  - lod: " << t.lod() << "\n";
  }

1244 1245
  os << "  - place: " << t.place() << "\n";
  os << "  - shape: [" << t.dims() << "]\n";
1246
  os << "  - layout: " << phi::DataLayoutToString(t.layout()) << "\n";
1247

1248
  DenseTensor tensor;
1249
  tensor.Resize(t.dims());
1250
  if (paddle::platform::is_cpu_place(t.place())) {
1251 1252
    tensor.ShareDataWith(t);
  } else {
1253 1254 1255 1256
    paddle::platform::CPUPlace place;
    paddle::framework::TensorCopy(t, place, &tensor);
    paddle::platform::DeviceContextPool& pool =
        paddle::platform::DeviceContextPool::Instance();
1257 1258 1259 1260
    auto& dev_ctx = *pool.Get(t.place());
    dev_ctx.Wait();
  }

1261 1262 1263 1264 1265 1266 1267 1268
#define PrintTensorCallback(cpp_type, proto_type)                 \
  do {                                                            \
    if (paddle::framework::TransToProtoVarType(tensor.dtype()) == \
        proto_type) {                                             \
      os << "  - dtype: " << proto_type << "\n";                  \
      paddle::framework::print_tensor<cpp_type>(os, tensor);      \
      return os;                                                  \
    }                                                             \
1269 1270 1271 1272 1273 1274
  } while (0)

  _ForEachDataType_(PrintTensorCallback);
  VLOG(1) << "PrintVar: unrecognized data type:" << t.type();
  return os;
}
1275
}  // namespace phi