Skip to content

Commit bfc34ac

Browse files
sneaxiysneaxiy
authored
Merge pull request PaddlePaddle#14536 from sneaxiy/dlpack_integration
Add dlpack support
2 parents 00b9e9a + 488610a commit bfc34ac

File tree

6 files changed

+320
-0
lines changed

6 files changed

+320
-0
lines changed

‎CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,7 @@ include(external/pybind11) # download pybind11
205205
include(external/cares)
206206
include(external/cub)
207207
include(external/xxhash) # download xxhash
208+
include(external/dlpack)
208209
include(external/snappy) # download snappy
209210
include(external/snappystream) # download snappystream
210211

‎cmake/external/dlpack.cmake

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
include(ExternalProject)
2+
3+
set(DLPACK_SOURCE_DIR ${THIRD_PARTY_PATH}/dlpack)
4+
set(DLPACK_INCLUDE_DIR ${DLPACK_SOURCE_DIR}/src/extern_dlpack/include)
5+
6+
include_directories(${DLPACK_INCLUDE_DIR})
7+
8+
ExternalProject_Add(
9+
extern_dlpack
10+
${EXTERNAL_PROJECT_LOG_ARGS}
11+
GIT_REPOSITORY "https://github.com/dmlc/dlpack.git"
12+
GIT_TAG "v0.2"
13+
PREFIX ${DLPACK_SOURCE_DIR}
14+
UPDATE_COMMAND ""
15+
CONFIGURE_COMMAND ""
16+
BUILD_COMMAND ""
17+
INSTALL_COMMAND ""
18+
TEST_COMMAND ""
19+
)
20+
21+
if(${CMAKE_VERSION} VERSION_LESS "3.3.0")
22+
set(dummyfile ${CMAKE_CURRENT_BINARY_DIR}/dlpack_dummy.c)
23+
file(WRITE ${dummyfile} "const char *dummy = \"${dummyfile}\";")
24+
add_library(dlpack STATIC ${dummyfile})
25+
else()
26+
add_library(dlpack INTERFACE)
27+
endif()
28+
29+
add_dependencies(dlpack extern_dlpack)
30+
31+
LIST(APPEND externl_project_dependencies dlpack)

‎paddle/fluid/framework/CMakeLists.txt

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -192,3 +192,6 @@ cc_test(tuple_test SRCS tuple_test.cc )
192192
if (NOT WIN32)
193193
cc_test(rw_lock_test SRCS rw_lock_test.cc)
194194
endif (NOT WIN32)
195+
196+
cc_library(dlpack_tensor SRCS dlpack_tensor.cc DEPS tensor dlpack)
197+
cc_test(dlpack_tensor_test SRCS dlpack_tensor_test.cc DEPS dlpack_tensor glog)

‎paddle/fluid/framework/dlpack_tensor.cc

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#include "paddle/fluid/framework/dlpack_tensor.h"
16+
17+
namespace paddle {
18+
namespace framework {
19+
20+
namespace internal {
21+
template <typename T>
22+
static ::DLDataType GetDLDataTypeCode() {
23+
::DLDataType dtype;
24+
if (std::is_same<T, platform::float16>::value ||
25+
std::is_floating_point<T>::value) {
26+
dtype.code = kDLFloat;
27+
} else if (std::is_unsigned<T>::value) {
28+
dtype.code = kDLUInt;
29+
} else if (std::is_integral<T>::value) {
30+
dtype.code = kDLInt;
31+
} else {
32+
PADDLE_THROW("Unsupported data type %s", typeid(T).name());
33+
}
34+
dtype.bits = 8 * sizeof(T);
35+
dtype.lanes = 1;
36+
return dtype;
37+
}
38+
39+
static DLDataType GetDLDataTypeFromTypeIndex(const std::type_index &type) {
40+
#define REG_DL_DATA_TYPE(type) \
41+
{ std::type_index(typeid(type)), GetDLDataTypeCode<type>() }
42+
static const std::unordered_map<std::type_index, ::DLDataType>
43+
type_to_dtype_map({
44+
REG_DL_DATA_TYPE(platform::float16), // NOLINT
45+
REG_DL_DATA_TYPE(float), // NOLINT
46+
REG_DL_DATA_TYPE(double), // NOLINT
47+
REG_DL_DATA_TYPE(int), // NOLINT
48+
REG_DL_DATA_TYPE(int64_t), // NOLINT
49+
REG_DL_DATA_TYPE(bool), // NOLINT
50+
REG_DL_DATA_TYPE(size_t), // NOLINT
51+
REG_DL_DATA_TYPE(int16_t), // NOLINT
52+
REG_DL_DATA_TYPE(uint8_t), // NOLINT
53+
REG_DL_DATA_TYPE(int8_t) // NOLINT
54+
});
55+
static auto type_to_dtype_map_end_it = type_to_dtype_map.end();
56+
auto it = type_to_dtype_map.find(type);
57+
PADDLE_ENFORCE(it != type_to_dtype_map_end_it, "Unsupported data type %s",
58+
type.name());
59+
return it->second;
60+
#undef REG_DL_DATA_TYPE
61+
}
62+
63+
struct DLContextVisitor : public boost::static_visitor<::DLContext> {
64+
inline ::DLContext operator()(const platform::CPUPlace &place) const {
65+
DLContext ctx;
66+
ctx.device_type = kDLCPU;
67+
ctx.device_id = 0;
68+
return ctx;
69+
}
70+
71+
inline ::DLContext operator()(const platform::CUDAPlace &place) const {
72+
#ifdef PADDLE_WITH_CUDA
73+
DLContext ctx;
74+
ctx.device_type = kDLGPU;
75+
ctx.device_id = place.device;
76+
return ctx;
77+
#else
78+
PADDLE_THROW("platform::CUDAPlace is not supported in CPU only version");
79+
#endif
80+
}
81+
82+
inline ::DLContext operator()(const platform::CUDAPinnedPlace &place) const {
83+
#ifdef PADDLE_WITH_CUDA
84+
DLContext ctx;
85+
ctx.device_type = kDLCPUPinned;
86+
ctx.device_id = 0;
87+
return ctx;
88+
#else
89+
PADDLE_THROW(
90+
"platform::CUDAPinnedPlace is not supported in CPU only version");
91+
#endif
92+
}
93+
};
94+
} // namespace internal
95+
96+
DLPackTensor::DLPackTensor(const Tensor &tensor, LaneType lanes) {
97+
// init data, data buffer
98+
t_.data = const_cast<void *>(tensor.data<void>());
99+
100+
// init ctx, DLContext type with device_type and device_id
101+
auto place = tensor.place();
102+
t_.ctx = boost::apply_visitor(internal::DLContextVisitor(), place);
103+
104+
// init dtype
105+
t_.dtype = internal::GetDLDataTypeFromTypeIndex(tensor.type());
106+
t_.dtype.lanes = lanes;
107+
108+
// init ndim, tensor rank
109+
auto &dims = tensor.dims();
110+
using DimType = decltype(t_.ndim); // int
111+
t_.ndim = static_cast<DimType>(dims.size());
112+
113+
// init shape, tensor dims
114+
t_.shape = shape_;
115+
for (DimType i = 0; i < t_.ndim; ++i) {
116+
t_.shape[i] = dims[i];
117+
}
118+
119+
// init strides, nullptr means the tensor is compact
120+
t_.strides = nullptr;
121+
122+
// init byte_offset
123+
t_.byte_offset = 0;
124+
}
125+
126+
} // namespace framework
127+
} // namespace paddle

‎paddle/fluid/framework/dlpack_tensor.h

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#pragma once
16+
17+
#include <dlpack/dlpack.h>
18+
#include "paddle/fluid/framework/tensor.h"
19+
20+
namespace paddle {
21+
namespace framework {
22+
23+
class DLPackTensor {
24+
public:
25+
using LaneType = decltype(::DLTensor::dtype.lanes); // uint16_t
26+
using ShapeType =
27+
std::remove_reference<decltype(::DLTensor::shape[0])>::type; // int64_t
28+
29+
// lanes is only used in CPU to enable vectorization
30+
explicit DLPackTensor(const Tensor& tensor, LaneType lanes = 1);
31+
32+
inline operator const ::DLTensor&() const { return t_; }
33+
34+
inline operator ::DLTensor&() { return t_; }
35+
36+
private:
37+
::DLTensor t_;
38+
39+
// The shape in DLTensor is defined as int64_t*
40+
// Add this member to make TVMTensor init without heap allocation
41+
ShapeType shape_[9];
42+
};
43+
44+
} // namespace framework
45+
} // namespace paddle

‎paddle/fluid/framework/dlpack_tensor_test.cc

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#include "paddle/fluid/framework/dlpack_tensor.h"
16+
#include <glog/logging.h>
17+
#include <gtest/gtest.h>
18+
#include <vector>
19+
20+
namespace paddle {
21+
namespace framework {
22+
23+
namespace { // NOLINT
24+
template <typename T>
25+
constexpr uint8_t GetDLDataTypeCode() {
26+
return std::is_same<platform::float16, T>::value ||
27+
std::is_floating_point<T>::value
28+
? static_cast<uint8_t>(kDLFloat)
29+
: (std::is_unsigned<T>::value
30+
? static_cast<uint8_t>(kDLUInt)
31+
: (std::is_integral<T>::value ? static_cast<uint8_t>(kDLInt)
32+
: static_cast<uint8_t>(-1)));
33+
}
34+
} // NOLINT
35+
36+
template <typename T>
37+
void TestMain(const platform::Place &place, uint16_t lanes) {
38+
DDim dims{4, 5, 6, 7};
39+
Tensor tensor;
40+
tensor.Resize(dims);
41+
void *p = tensor.mutable_data<T>(place);
42+
43+
DLPackTensor dlpack_tensor(tensor, lanes);
44+
::DLTensor &dl_tensor = dlpack_tensor;
45+
46+
CHECK_EQ(p, dl_tensor.data);
47+
if (platform::is_cpu_place(place)) {
48+
CHECK_EQ(kDLCPU, dl_tensor.ctx.device_type);
49+
CHECK_EQ(0, dl_tensor.ctx.device_id);
50+
} else if (platform::is_gpu_place(place)) {
51+
CHECK_EQ(kDLGPU, dl_tensor.ctx.device_type);
52+
CHECK_EQ(boost::get<platform::CUDAPlace>(place).device,
53+
dl_tensor.ctx.device_id);
54+
} else if (platform::is_cuda_pinned_place(place)) {
55+
CHECK_EQ(kDLCPUPinned, dl_tensor.ctx.device_type);
56+
CHECK_EQ(0, dl_tensor.ctx.device_id);
57+
} else {
58+
CHECK_EQ(false, true);
59+
}
60+
61+
CHECK_EQ(dims.size(), dl_tensor.ndim);
62+
for (auto i = 0; i < dims.size(); ++i) {
63+
CHECK_EQ(dims[i], dl_tensor.shape[i]);
64+
}
65+
66+
CHECK_EQ(dl_tensor.strides == nullptr, true);
67+
CHECK_EQ(static_cast<uint64_t>(0), dl_tensor.byte_offset);
68+
69+
CHECK_EQ(lanes, dl_tensor.dtype.lanes);
70+
CHECK_EQ(sizeof(T) * 8, dl_tensor.dtype.bits);
71+
72+
CHECK_EQ(GetDLDataTypeCode<T>(), dl_tensor.dtype.code);
73+
}
74+
75+
template <typename T>
76+
void TestMainLoop() {
77+
#ifdef PADDLE_WITH_CUDA
78+
std::vector<platform::Place> places{platform::CPUPlace(),
79+
platform::CUDAPlace(0),
80+
platform::CUDAPinnedPlace()};
81+
if (platform::GetCUDADeviceCount() > 1) {
82+
places.emplace_back(platform::CUDAPlace(1));
83+
}
84+
#else
85+
std::vector<platform::Place> places{platform::CPUPlace()};
86+
#endif
87+
std::vector<uint16_t> lanes{1, 2};
88+
for (auto &p : places) {
89+
for (auto &l : lanes) {
90+
TestMain<T>(p, l);
91+
}
92+
}
93+
}
94+
95+
#define PADDLE_DLPACK_TEST(type) \
96+
TEST(dlpack, test_##type) { TestMainLoop<type>(); }
97+
98+
using float16 = platform::float16;
99+
PADDLE_DLPACK_TEST(float16);
100+
PADDLE_DLPACK_TEST(float);
101+
PADDLE_DLPACK_TEST(double);
102+
PADDLE_DLPACK_TEST(int);
103+
PADDLE_DLPACK_TEST(int64_t);
104+
PADDLE_DLPACK_TEST(bool);
105+
PADDLE_DLPACK_TEST(size_t);
106+
PADDLE_DLPACK_TEST(int16_t);
107+
PADDLE_DLPACK_TEST(uint8_t);
108+
PADDLE_DLPACK_TEST(int8_t);
109+
110+
#undef PADDLE_DLPACK_TEST
111+
112+
} // namespace framework
113+
} // namespace paddle

0 commit comments

Comments
 (0)