-
Notifications
You must be signed in to change notification settings - Fork 6
Expand file tree
/
Copy pathkernel_pybind.cpp
More file actions
128 lines (101 loc) · 3.57 KB
/
Copy pathkernel_pybind.cpp
File metadata and controls
128 lines (101 loc) · 3.57 KB
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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
#include <pybind11/pybind11.h>
#include<pybind11/numpy.h>
#include <iostream>
#include "dlpack.h"
#include "kernel.h"
#include "csr.h"
using std::cout;
using std::endl;
namespace py = pybind11;
array1d_t<float> capsule_to_array1d(py::capsule& capsule)
{
//DLManagedTensor * dlMTensor = (DLManagedTensor *)PyCapsule_GetPointer(&input, "dltesnor");
//The worst static_cast, this is becacuse of 'operator T*()' in pybind11 code
//Pybind11 should have simply provided GetPointer() API.
DLManagedTensor * dlMTensor = static_cast<DLManagedTensor*>(capsule);
assert(dlMTensor);
DLTensor* tensor = &dlMTensor->dl_tensor;
int shape0 = tensor->shape[0];
float* data_ptr = (float*)tensor->data;
array1d_t<float> array(data_ptr, shape0);
return array;
}
array2d_t<float> capsule_to_array2d(py::capsule& capsule)
{
DLManagedTensor * dlMTensor = static_cast<DLManagedTensor*>(capsule);
assert(dlMTensor);
DLTensor* tensor = &dlMTensor->dl_tensor;
int shape0 = tensor->shape[0];
int shape1 = tensor->shape[1];
float* data_ptr = (float*)tensor->data;
array2d_t<float> array(data_ptr, shape0, shape1);
return array;
}
array3d_t<float> capsule_to_array3d(py::capsule& capsule)
{
DLManagedTensor * dlMTensor = static_cast<DLManagedTensor*>(capsule);
assert(dlMTensor);
DLTensor* tensor = &dlMTensor->dl_tensor;
int shape0 = tensor->shape[0];
int shape1 = tensor->shape[1];
int shape2 = tensor->shape[2];
float* data_ptr = (float*)tensor->data;
array3d_t<float> array(data_ptr, shape0, shape1, shape2);
return array;
}
array4d_t<float> capsule_to_array4d(py::capsule& capsule)
{
DLManagedTensor * dlMTensor = static_cast<DLManagedTensor*>(capsule);
assert(dlMTensor);
DLTensor* tensor = &dlMTensor->dl_tensor;
int shape0 = tensor->shape[0];
int shape1 = tensor->shape[1];
int shape2 = tensor->shape[2];
int shape3 = tensor->shape[3];
float* data_ptr = (float*)tensor->data;
array4d_t<float> array(data_ptr, shape0, shape1, shape2, shape3);
return array;
}
#include "generated_pybind.h"
PYBIND11_MODULE(graphpy, m) {
py::enum_<op_t>(m, "op_t", py::arithmetic(), "Operations Enums")
.value("eSUM", eSUM, "sum")
.value("eMAX", eMAX, "max")
.value("eMIN", eMIN, "min")
.value("eSUB", eSUB, "sub")
.value("eMUL", eMUL, "mul")
.value("eDIV", eDIV, "div")
.export_values();
py::class_<graph_t>(m, "graph_t")
.def(py::init<>())
.def("save_graph", &graph_t::save_graph)
.def("get_vcount", &graph_t::get_vcount)
.def("get_edge_count", &graph_t::get_ecount)
;
m.def("init_graph",
[](py::array offset_csr, py::array nebrs_csr, py::array offset_csc, py::array nebrs_csc, int flag, int num_vcount) {
graph_t* graph = new graph_t;
//cout<< offset_csr.shape(0) - 1<< "num_vcount"<< endl;
graph->init(offset_csr.shape(0) - 1, nebrs_csr.itemsize(),
offset_csr.request().ptr, nebrs_csr.request().ptr,
offset_csc.request().ptr, nebrs_csc.request().ptr, flag, num_vcount);
//THD_COUNT = thd_count;
return graph;
}
);
m.def("load_graph",
[](const string& odir) {
graph_t* graph = new graph_t;
graph->load_graph(odir);
return graph;
}
);
m.def("load_graph_noeid",
[](const string& odir) {
graph_t* graph = new graph_t;
graph->load_graph_noeid(odir);
return graph;
}
);
export_kernel(m);
}