tensor_util.h 12.0 KB
Newer Older
D
dzhwinter 已提交
1 2
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

L
Luo Tao 已提交
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
D
dzhwinter 已提交
6

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
D
dzhwinter 已提交
8

L
Luo Tao 已提交
9 10 11 12 13
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. */
D
dzhwinter 已提交
14 15

#pragma once
Y
Yang Yu 已提交
16 17
#include "paddle/framework/data_type.h"
#include "paddle/framework/eigen.h"
18
#include "paddle/framework/framework.pb.h"
D
dzhwinter 已提交
19
#include "paddle/framework/tensor.h"
Y
Yang Yu 已提交
20
#include "paddle/platform/device_context.h"
D
dzhwinter 已提交
21 22 23 24 25 26 27 28 29 30 31

namespace paddle {
namespace framework {

/**
 * @brief   Copy the content of external tensor to a new place.
 *
 * @param[in] src        The external tensor.
 * @param[in] dst_place  The dst place.
 * @param[in] ctx        The device context contains device resources.
 *
32
 * @note    Copy supports CPU <-> GPU, GPU <-> GPU.
D
dzhwinter 已提交
33
 */
34 35
inline void Copy(const Tensor& src, const platform::Place& dst_place,
                 const platform::DeviceContext& ctx, Tensor* dst) {
36 37
  VLOG(3) << "Copy " << src.dims() << " from " << src.place() << " to "
          << dst_place;
D
dzhwinter 已提交
38 39 40
  src.check_memory_size();

  dst->Resize(src.dims());
D
dzhwinter 已提交
41
  dst->set_layout(src.layout());
D
dzhwinter 已提交
42 43 44 45 46 47 48 49 50 51 52 53 54 55
  auto src_place = src.place();
  auto src_ptr = src.data<void>();

  auto dst_ptr = dst->mutable_data(dst_place, src.type());

  auto size = src.numel() * SizeOfType(src.type());

  if (platform::is_cpu_place(src_place) && platform::is_cpu_place(dst_place)) {
    memory::Copy(boost::get<platform::CPUPlace>(dst_place), dst_ptr,
                 boost::get<platform::CPUPlace>(src_place), src_ptr, size);
  }
#ifdef PADDLE_WITH_CUDA
  else if (platform::is_gpu_place(src_place) &&  // NOLINT
           platform::is_cpu_place(dst_place)) {
D
dzhwinter 已提交
56
    auto src_gpu_place = boost::get<platform::CUDAPlace>(src_place);
D
dzhwinter 已提交
57 58 59
    auto dst_cpu_place = boost::get<platform::CPUPlace>(dst_place);
    auto ctx_place = ctx.GetPlace();
    PADDLE_ENFORCE(platform::is_gpu_place(ctx_place));
D
dzhwinter 已提交
60
    auto ctx_gpu_place = boost::get<platform::CUDAPlace>(ctx_place);
D
dzhwinter 已提交
61 62 63 64 65 66 67
    PADDLE_ENFORCE_EQ(src_gpu_place, ctx_gpu_place);
    memory::Copy(
        dst_cpu_place, dst_ptr, src_gpu_place, src_ptr, size,
        reinterpret_cast<const platform::CUDADeviceContext&>(ctx).stream());
  } else if (platform::is_cpu_place(src_place) &&
             platform::is_gpu_place(dst_place)) {
    auto src_cpu_place = boost::get<platform::CPUPlace>(src_place);
D
dzhwinter 已提交
68
    auto dst_gpu_place = boost::get<platform::CUDAPlace>(dst_place);
D
dzhwinter 已提交
69 70
    auto ctx_place = ctx.GetPlace();
    PADDLE_ENFORCE(platform::is_gpu_place(ctx_place));
D
dzhwinter 已提交
71
    auto ctx_gpu_place = boost::get<platform::CUDAPlace>(ctx_place);
D
dzhwinter 已提交
72 73 74 75 76 77
    PADDLE_ENFORCE_EQ(dst_gpu_place, ctx_gpu_place);
    memory::Copy(
        dst_gpu_place, dst_ptr, src_cpu_place, src_ptr, size,
        reinterpret_cast<const platform::CUDADeviceContext&>(ctx).stream());
  } else if (platform::is_gpu_place(src_place) &&
             platform::is_gpu_place(dst_place)) {
D
dzhwinter 已提交
78 79
    auto src_gpu_place = boost::get<platform::CUDAPlace>(src_place);
    auto dst_gpu_place = boost::get<platform::CUDAPlace>(dst_place);
D
dzhwinter 已提交
80 81
    auto ctx_place = ctx.GetPlace();
    PADDLE_ENFORCE(platform::is_gpu_place(ctx_place));
D
dzhwinter 已提交
82
    auto ctx_gpu_place = boost::get<platform::CUDAPlace>(ctx_place);
D
dzhwinter 已提交
83 84 85 86 87 88 89 90
    PADDLE_ENFORCE_EQ(src_gpu_place, ctx_gpu_place);
    memory::Copy(
        dst_gpu_place, dst_ptr, src_gpu_place, src_ptr, size,
        reinterpret_cast<const platform::CUDADeviceContext&>(ctx).stream());
  }
#endif
}

D
dzhwinter 已提交
91
/**
92 93 94 95 96 97 98 99
 * @brief Wrapper on
 *     Copy(const Tensor& src, const platform::Place& dst_place,
 *              const platform::DeviceContext& ctx, Tensor* dst);
 *
 * @param[in] src        The external tensor.
 * @param[in] dst_place  The dst place.
 *
 * @note    Copy supports CPU <-> GPU, GPU <-> GPU.
D
dzhwinter 已提交
100
 */
101 102
inline void Copy(const Tensor& src, const platform::Place& dst_place,
                 Tensor* dst) {
103 104 105 106 107 108 109 110
  platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
  const platform::DeviceContext* dev_ctx;
  if (platform::is_gpu_place(src.place())) {
    dev_ctx = pool.Get(src.place());
  } else {
    dev_ctx = pool.Get(dst_place);
  }
  Copy(src, dst_place, *dev_ctx, dst);
D
dzhwinter 已提交
111 112
}

D
dzhwinter 已提交
113 114 115 116 117 118
/**
 * @brief   Copy the content of an external vector to a tensor.
 *
 * @param[in] src        The external tensor.
 * @param[in] ctx        The device context contains device resources.
 *
X
xuwei06 已提交
119 120
 * * @note    CopyFromVector will resize dst to an 1D tensor with the same
 *            size as src.
D
dzhwinter 已提交
121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138
 */
template <typename T>
inline void CopyFromVector(const std::vector<T>& src,
                           const platform::DeviceContext& ctx, Tensor* dst) {
  auto dst_place = ctx.GetPlace();
  auto src_ptr = static_cast<const void*>(src.data());
  platform::CPUPlace src_place;
  dst->Resize({static_cast<int64_t>(src.size())});
  auto dst_ptr = static_cast<void*>(dst->mutable_data<T>(dst_place));
  auto size = src.size() * sizeof(T);

  if (platform::is_cpu_place(dst_place)) {
    memory::Copy(boost::get<platform::CPUPlace>(dst_place), dst_ptr, src_place,
                 src_ptr, size);
  }
#ifdef PADDLE_WITH_CUDA
  else if (platform::is_gpu_place(dst_place)) {  // NOLINT
    memory::Copy(
D
dzhwinter 已提交
139
        boost::get<platform::CUDAPlace>(dst_place), dst_ptr, src_place, src_ptr,
D
dzhwinter 已提交
140 141 142 143 144 145
        size,
        reinterpret_cast<const platform::CUDADeviceContext&>(ctx).stream());
  }
#endif
}

D
dzhwinter 已提交
146 147 148 149 150 151 152 153 154 155 156 157 158 159 160
/**
 * @brief CopyFromVector CPU vector -> CPU Tensor
 */
template <typename T>
inline void CopyFromVector(const std::vector<T>& src, Tensor* dst) {
  platform::CPUPlace dst_place = platform::CPUPlace();
  auto src_ptr = static_cast<const void*>(src.data());
  platform::CPUPlace src_place;
  dst->Resize({static_cast<int64_t>(src.size())});
  auto dst_ptr = static_cast<void*>(dst->mutable_data<T>(dst_place));
  auto size = src.size() * sizeof(T);

  memory::Copy(dst_place, dst_ptr, src_place, src_ptr, size);
}

D
dzhwinter 已提交
161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180
/**
 * @brief   Copy the content of a tensor to a vector
 *
 * @param[in] src        The external tensor.
 * @param[in] ctx        The device context contains device resources.
 *
 * * @note    CopyFromVector assumes that the tensor has been resized
 *            before invoking.
 */
template <typename T>
inline void CopyToVector(const Tensor& src, const platform::DeviceContext& ctx,
                         std::vector<T>* dst) {
  auto src_ptr = static_cast<const void*>(src.data<T>());
  auto size = src.numel() * sizeof(T);

  platform::CPUPlace dst_place;
  dst->resize(src.numel());
  auto dst_ptr = static_cast<void*>(dst->data());

  if (platform::is_cpu_place(src.place())) {
181 182
    memory::Copy(dst_place, dst_ptr,
                 boost::get<platform::CPUPlace>(src.place()), src_ptr, size);
D
dzhwinter 已提交
183 184 185 186
  }
#ifdef PADDLE_WITH_CUDA
  else if (platform::is_gpu_place(src.place())) {  // NOLINT
    memory::Copy(
D
dzhwinter 已提交
187
        dst_place, dst_ptr, boost::get<platform::CUDAPlace>(src.place()),
188
        src_ptr, size,
D
dzhwinter 已提交
189 190 191 192 193
        reinterpret_cast<const platform::CUDADeviceContext&>(ctx).stream());
  }
#endif
}

D
dzhwinter 已提交
194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211
/**
 * @brief CopyToVector CPUTensor <-> CPU Vector
 */
template <typename T>
inline void CopyToVector(const Tensor& src, std::vector<T>* dst) {
  auto src_ptr = static_cast<const void*>(src.data<T>());
  auto size = src.numel() * sizeof(T);

  platform::CPUPlace dst_place;
  dst->resize(src.numel());
  auto dst_ptr = static_cast<void*>(dst->data());

  PADDLE_ENFORCE(platform::is_cpu_place(src.place()));

  memory::Copy(dst_place, dst_ptr, boost::get<platform::CPUPlace>(src.place()),
               src_ptr, size);
}

Y
Yang Yu 已提交
212
// Returns true if a tensor contains NAN, i.e., Not A Number.
Y
Yang Yu 已提交
213
bool HasNAN(const framework::Tensor& tensor);
Y
Yang Yu 已提交
214 215

// Returns true if a tensor contains Inf, i.e., Infinity.
Y
Yang Yu 已提交
216
bool HasInf(const framework::Tensor& tensor);
Y
Yang Yu 已提交
217

218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272
inline void SerializeToStream(std::ostream& os, const Tensor& tensor,
                              const platform::DeviceContext& dev_ctx) {
  // TODO(typhoonzero): serialize to ostream
  {  // 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::TensorDesc desc;
    desc.set_data_type(framework::ToDataType(tensor.type()));
    auto dims = framework::vectorize(tensor.dims());
    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
    uint64_t size = tensor.memory_size();
    auto* data_ptr = tensor.data<void>();
    PADDLE_ENFORCE(size < std::numeric_limits<std::streamsize>::max(),
                   "Index overflow when writing tensor");
    if (platform::is_gpu_place(tensor.place())) {
#ifdef PADDLE_WITH_CUDA
      constexpr size_t kBufSize = 1024 * 1024 * 64;  // 64MB
      std::unique_ptr<char[]> buf(new char[kBufSize]);
      auto& gpu_dev_ctx =
          static_cast<const platform::CUDADeviceContext&>(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));
        memory::Copy(cpu, buf.get(),
                     boost::get<platform::CUDAPlace>(tensor.place()),
                     reinterpret_cast<const void*>(data), size_to_write,
                     gpu_dev_ctx.stream());
        gpu_dev_ctx.Wait();
        os.write(buf.get(), size_to_write);
        data += size_to_write;
        size -= size_to_write;
      }
#else
      PADDLE_THROW("Unexpected branch");
#endif
    } else {
      os.write(static_cast<const char*>(data_ptr),
               static_cast<std::streamsize>(size));
    }
  }
}

Y
Yancey 已提交
273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289
struct DeserializedDataFunctor {
  DeserializedDataFunctor(void** buf, Tensor* tensor,
                          const platform::Place& place)
      : buf_(buf), tensor_(tensor), place_(place) {}

  template <typename T>
  void operator()() {
    *buf_ = tensor_->mutable_data<T>(place_);
  }

  void** buf_;
  Tensor* tensor_;
  platform::Place place_;
};

inline void DeserializeFromStream(std::istream& is, Tensor* tensor,
                                  const platform::DeviceContext& dev_ctx) {
290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308
  uint32_t version;
  is.read(reinterpret_cast<char*>(&version), sizeof(version));
  PADDLE_ENFORCE_EQ(version, 0U, "Only version 0 is supported");
  proto::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(desc.ParseFromArray(buf.get(), size),
                   "Cannot parse tensor desc");
  }
  {  // 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));
    tensor->Resize(framework::make_ddim(dims));
    void* buf;
Y
Yancey 已提交
309 310 311 312 313 314 315 316 317
    auto ctx = platform::CPUDeviceContext();
    if (platform::is_gpu_place(dev_ctx.GetPlace())) {
#ifdef PADDLE_WITH_CUDA
      Tensor cpu_tensor;
      cpu_tensor.Resize(framework::make_ddim(dims));
      framework::VisitDataType(
          desc.data_type(),
          DeserializedDataFunctor(&buf, &cpu_tensor, ctx.GetPlace()));
      is.read(static_cast<char*>(buf), cpu_tensor.memory_size());
Y
Yancey1989 已提交
318 319
      auto dst_place = dev_ctx.GetPlace();
      framework::Copy(cpu_tensor, dst_place, dev_ctx, tensor);
Y
Yancey 已提交
320 321 322 323 324 325 326 327
#else
      PADDLE_THROW("Unexpected branch");
#endif
    } else {
      framework::VisitDataType(
          desc.data_type(),
          DeserializedDataFunctor(&buf, tensor, ctx.GetPlace()));
      is.read(static_cast<char*>(buf), tensor->memory_size());
328 329 330 331
    }
  }
}

D
dzhwinter 已提交
332 333
}  // namespace framework
}  // namespace paddle