1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
|
/**
* Copyright 2019-2024, XGBoost Contributors
*/
#include <gtest/gtest.h>
#include <thrust/device_vector.h>
#include <thrust/execution_policy.h> // for device
#include <thrust/sequence.h> // for sequence
#include <xgboost/data.h>
#include <xgboost/json.h>
namespace xgboost {
template <typename T>
Json GenerateDenseColumn(std::string const& typestr, size_t kRows,
thrust::device_vector<T>* out_d_data) {
auto& d_data = *out_d_data;
d_data.resize(kRows);
Json column { Object() };
std::vector<Json> j_shape {Json(Integer(static_cast<Integer::Int>(kRows)))};
column["shape"] = Array(j_shape);
column["strides"] = Array(std::vector<Json>{Json(Integer(static_cast<Integer::Int>(sizeof(T))))});
column["stream"] = nullptr;
d_data.resize(kRows);
thrust::sequence(thrust::device, d_data.begin(), d_data.end(), 0.0f, 2.0f);
auto p_d_data = d_data.data().get();
std::vector<Json> j_data {
Json(Integer(reinterpret_cast<Integer::Int>(p_d_data))),
Json(Boolean(false))};
column["data"] = j_data;
column["version"] = 3;
column["typestr"] = String(typestr);
return column;
}
template <typename T>
Json GenerateSparseColumn(std::string const& typestr, size_t kRows,
thrust::device_vector<T>* out_d_data) {
auto& d_data = *out_d_data;
Json column { Object() };
std::vector<Json> j_shape {Json(Integer(static_cast<Integer::Int>(kRows)))};
column["shape"] = Array(j_shape);
column["strides"] = Array(std::vector<Json>{Json(Integer(static_cast<Integer::Int>(sizeof(T))))});
column["stream"] = nullptr;
d_data.resize(kRows);
for (size_t i = 0; i < d_data.size(); ++i) {
d_data[i] = i * 2.0;
}
auto p_d_data = d_data.data().get();
std::vector<Json> j_data {
Json(Integer(reinterpret_cast<Integer::Int>(p_d_data))),
Json(Boolean(false))};
column["data"] = j_data;
column["version"] = 3;
column["typestr"] = String(typestr);
return column;
}
template <typename T>
Json Generate2dArrayInterface(int rows, int cols, std::string typestr,
thrust::device_vector<T> *p_data) {
auto& data = *p_data;
thrust::sequence(data.begin(), data.end());
Json array_interface{Object()};
std::vector<Json> shape = {Json(static_cast<Integer::Int>(rows)),
Json(static_cast<Integer::Int>(cols))};
array_interface["shape"] = Array(shape);
std::vector<Json> j_data{
Json(Integer(reinterpret_cast<Integer::Int>(data.data().get()))),
Json(Boolean(false))};
array_interface["data"] = j_data;
array_interface["version"] = 3;
array_interface["typestr"] = String(typestr);
array_interface["stream"] = nullptr;
return array_interface;
}
} // namespace xgboost
|
