tensor_py.h 5.6 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
2

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
6

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
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. */
14 15

#pragma once
Q
qijun 已提交
16
#include <string>
Y
Yi Wang 已提交
17 18 19
#include "paddle/fluid/framework/lod_tensor.h"
#include "paddle/fluid/memory/memcpy.h"
#include "paddle/fluid/platform/device_context.h"
Q
qijun 已提交
20 21
#include "pybind11/numpy.h"
#include "pybind11/pybind11.h"
22 23 24 25 26

namespace py = pybind11;

namespace paddle {

27
namespace pybind {
28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45

namespace details {

template <bool less, size_t I, typename... ARGS>
struct CastToPyBufferImpl;

template <size_t I, typename... ARGS>
struct CastToPyBufferImpl<false, I, ARGS...> {
  py::buffer_info operator()(framework::Tensor &tensor) {
    PADDLE_THROW("This type of tensor cannot be expose to Python");
    return py::buffer_info();
  }
};

template <size_t I, typename... ARGS>
struct CastToPyBufferImpl<true, I, ARGS...> {
  using CUR_TYPE = typename std::tuple_element<I, std::tuple<ARGS...>>::type;
  py::buffer_info operator()(framework::Tensor &tensor) {
Y
Yu Yang 已提交
46
    if (std::type_index(typeid(CUR_TYPE)) == tensor.type()) {
47 48 49 50 51 52 53 54 55 56 57 58
      auto dim_vec = framework::vectorize(tensor.dims());
      std::vector<size_t> dims_outside;
      std::vector<size_t> strides;
      dims_outside.resize(dim_vec.size());
      strides.resize(dim_vec.size());

      size_t prod = 1;
      for (size_t i = dim_vec.size(); i != 0; --i) {
        dims_outside[i - 1] = (size_t)dim_vec[i - 1];
        strides[i - 1] = sizeof(CUR_TYPE) * prod;
        prod *= dims_outside[i - 1];
      }
Q
qijun 已提交
59
      framework::Tensor dst_tensor;
Y
Yu Yang 已提交
60
      if (paddle::platform::is_gpu_place(tensor.place())) {
61 62 63 64
#ifdef PADDLE_WITH_CUDA
        auto *src_ptr = static_cast<const void *>(tensor.data<CUR_TYPE>());
        auto *dst_ptr = static_cast<void *>(dst_tensor.mutable_data<CUR_TYPE>(
            tensor.dims(), platform::CPUPlace()));
D
dzhwinter 已提交
65

Y
Yang Yu 已提交
66
        platform::DeviceContextPool &pool =
Y
Yang Yu 已提交
67
            platform::DeviceContextPool::Instance();
D
dzhwinter 已提交
68
        auto dev_ctx = static_cast<const platform::CUDADeviceContext *>(
Y
Yang Yu 已提交
69
            pool.Get(tensor.place()));
D
dzhwinter 已提交
70 71 72 73

        paddle::platform::GpuMemcpyAsync(
            dst_ptr, src_ptr, sizeof(CUR_TYPE) * tensor.numel(),
            cudaMemcpyDeviceToHost, dev_ctx->stream());
74
#else
D
dzhwinter 已提交
75
        PADDLE_THROW("'CUDAPlace' is not supported in CPU only device.");
76
#endif
Y
Yu Yang 已提交
77
      } else if (paddle::platform::is_cpu_place(tensor.place())) {
Q
qijun 已提交
78 79
        dst_tensor = tensor;
      }
Y
Yang Yu 已提交
80 81 82 83
      return py::buffer_info(dst_tensor.data<CUR_TYPE>(), sizeof(CUR_TYPE),
                             py::format_descriptor<CUR_TYPE>::format(),
                             (size_t)framework::arity(dst_tensor.dims()),
                             dims_outside, strides);
84 85 86 87 88 89 90 91
    } else {
      constexpr bool less = I + 1 < std::tuple_size<std::tuple<ARGS...>>::value;
      return CastToPyBufferImpl<less, I + 1, ARGS...>()(tensor);
    }
  }
};
}  // namespace details
inline py::buffer_info CastToPyBuffer(framework::Tensor &tensor) {
92
  auto buffer_info =
Y
Yu Yang 已提交
93
      details::CastToPyBufferImpl<true, 0, float, int, double, int64_t, bool>()(
94
          tensor);
95 96 97
  return buffer_info;
}

98 99
template <typename T>
T TensorGetElement(framework::Tensor &self, size_t offset) {
100 101 102 103 104 105 106
  if (platform::is_cpu_place(self.place())) {
    return self.data<T>()[offset];
  } else {
    std::shared_ptr<framework::Tensor> dst(new framework::Tensor);
    framework::Copy(self, platform::CPUPlace(), dst.get());
    return dst->data<T>()[offset];
  }
107 108
}

109
// TODO(dzhwinter) : fix the redundent Tensor allocate and free
110 111
template <typename T>
void TensorSetElement(framework::Tensor &self, size_t offset, T elem) {
112 113 114 115 116 117 118 119 120
  if (platform::is_gpu_place(self.place())) {
    std::shared_ptr<framework::Tensor> dst(new framework::Tensor);
    framework::Copy(self, platform::CPUPlace(), dst.get());
    dst->data<T>()[offset] = elem;
    framework::Copy(*dst.get(), self.place(), &self);

  } else if (platform::is_cpu_place(self.place())) {
    self.data<T>()[offset] = elem;
  }
121 122
}

123
template <typename T>
Q
qijun 已提交
124
void PyCPUTensorSetFromArray(
125
    framework::Tensor &self,
Q
qijun 已提交
126 127
    py::array_t<T, py::array::c_style | py::array::forcecast> array,
    paddle::platform::CPUPlace &place) {
Q
qijun 已提交
128
  std::vector<int64_t> dims;
129 130 131 132 133
  dims.reserve(array.ndim());
  for (size_t i = 0; i < array.ndim(); ++i) {
    dims.push_back((int)array.shape()[i]);
  }

F
fengjiayi 已提交
134
  self.Resize(framework::make_ddim(dims));
Q
qijun 已提交
135
  auto *dst = self.mutable_data<T>(place);
136 137 138
  std::memcpy(dst, array.data(), sizeof(T) * array.size());
}

139
#ifdef PADDLE_WITH_CUDA
Q
qijun 已提交
140 141 142 143
template <typename T>
void PyCUDATensorSetFromArray(
    framework::Tensor &self,
    py::array_t<T, py::array::c_style | py::array::forcecast> array,
D
dzhwinter 已提交
144
    paddle::platform::CUDAPlace &place) {
Q
qijun 已提交
145
  std::vector<int64_t> dims;
Q
qijun 已提交
146 147 148
  dims.reserve(array.ndim());
  for (size_t i = 0; i < array.ndim(); ++i) {
    dims.push_back((int)array.shape()[i]);
Q
qijun 已提交
149
  }
Q
qijun 已提交
150 151 152

  self.Resize(framework::make_ddim(dims));
  auto *dst = self.mutable_data<T>(place);
D
dzhwinter 已提交
153

Y
Yang Yu 已提交
154
  platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
D
dzhwinter 已提交
155
  auto dev_ctx =
Y
Yang Yu 已提交
156
      static_cast<const platform::CUDADeviceContext *>(pool.Get(place));
D
dzhwinter 已提交
157 158
  paddle::platform::GpuMemcpyAsync(dst, array.data(), sizeof(T) * array.size(),
                                   cudaMemcpyHostToDevice, dev_ctx->stream());
159
}
Q
qijun 已提交
160
#endif
161 162 163

}  // namespace pybind
}  // namespace paddle