/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve. 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 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. */ #pragma once #include #include "paddle/framework/tensor.h" #include "paddle/memory/memcpy.h" #include "paddle/platform/device_context.h" #include "pybind11/numpy.h" #include "pybind11/pybind11.h" namespace py = pybind11; namespace paddle { namespace pybind { namespace details { template struct CastToPyBufferImpl; template struct CastToPyBufferImpl { py::buffer_info operator()(framework::Tensor &tensor) { PADDLE_THROW("This type of tensor cannot be expose to Python"); return py::buffer_info(); } }; template struct CastToPyBufferImpl { using CUR_TYPE = typename std::tuple_element>::type; py::buffer_info operator()(framework::Tensor &tensor) { if (std::type_index(typeid(CUR_TYPE)) == tensor.type()) { auto dim_vec = framework::vectorize(tensor.dims()); std::vector dims_outside; std::vector 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]; } framework::Tensor dst_tensor; if (paddle::platform::is_gpu_place(tensor.place())) { #ifdef PADDLE_WITH_CUDA auto *src_ptr = static_cast(tensor.data()); auto *dst_ptr = static_cast(dst_tensor.mutable_data( tensor.dims(), platform::CPUPlace())); platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance; auto dev_ctx = static_cast( pool.Get(tensor.place())); paddle::platform::GpuMemcpyAsync( dst_ptr, src_ptr, sizeof(CUR_TYPE) * tensor.numel(), cudaMemcpyDeviceToHost, dev_ctx->stream()); #else PADDLE_THROW("'CUDAPlace' is not supported in CPU only device."); #endif } else if (paddle::platform::is_cpu_place(tensor.place())) { dst_tensor = tensor; } return py::buffer_info( dst_tensor.mutable_data(dst_tensor.place()), sizeof(CUR_TYPE), py::format_descriptor::format(), (size_t)framework::arity(dst_tensor.dims()), dims_outside, strides); } else { constexpr bool less = I + 1 < std::tuple_size>::value; return CastToPyBufferImpl()(tensor); } } }; } // namespace details inline py::buffer_info CastToPyBuffer(framework::Tensor &tensor) { auto buffer_info = details::CastToPyBufferImpl()( tensor); return buffer_info; } template T TensorGetElement(framework::Tensor &self, size_t offset) { PADDLE_ENFORCE(platform::is_cpu_place(self.place())); return self.data()[offset]; } template void TensorSetElement(framework::Tensor &self, size_t offset, T elem) { PADDLE_ENFORCE(platform::is_cpu_place(self.place())); self.data()[offset] = elem; } template void PyCPUTensorSetFromArray( framework::Tensor &self, py::array_t array, paddle::platform::CPUPlace &place) { std::vector dims; dims.reserve(array.ndim()); for (size_t i = 0; i < array.ndim(); ++i) { dims.push_back((int)array.shape()[i]); } self.Resize(framework::make_ddim(dims)); auto *dst = self.mutable_data(place); std::memcpy(dst, array.data(), sizeof(T) * array.size()); } #ifdef PADDLE_WITH_CUDA template void PyCUDATensorSetFromArray( framework::Tensor &self, py::array_t array, paddle::platform::CUDAPlace &place) { std::vector dims; dims.reserve(array.ndim()); for (size_t i = 0; i < array.ndim(); ++i) { dims.push_back((int)array.shape()[i]); } self.Resize(framework::make_ddim(dims)); auto *dst = self.mutable_data(place); platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance(); auto dev_ctx = static_cast(pool.Get(place)); paddle::platform::GpuMemcpyAsync(dst, array.data(), sizeof(T) * array.size(), cudaMemcpyHostToDevice, dev_ctx->stream()); } #endif } // namespace pybind } // namespace paddle