sampler_core.cpp

#include <iostream>
#include <string>
#include <cstdlib>
#include <random>
#include <omp.h>
#include <math.h>
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>

namespace py = pybind11;

typedef int NodeIDType;
typedef int EdgeIDType;
typedef float TimeStampType;

class TemporalGraphBlock
{
    public:
        std::vector<NodeIDType> row;
        std::vector<NodeIDType> col;
        std::vector<EdgeIDType> eid;
        std::vector<TimeStampType> ts;
        std::vector<TimeStampType> dts;
        std::vector<NodeIDType> nodes;
        NodeIDType dim_in, dim_out;
        double ptr_time = 0;
        double search_time = 0;
        double sample_time = 0;
        double tot_time = 0;
        double coo_time = 0;

        TemporalGraphBlock(){}

        TemporalGraphBlock(std::vector<NodeIDType> &_row, std::vector<NodeIDType> &_col,
                           std::vector<EdgeIDType> &_eid, std::vector<TimeStampType> &_ts,
                           std::vector<TimeStampType> &_dts, std::vector<NodeIDType> &_nodes, 
                           NodeIDType _dim_in, NodeIDType _dim_out) :
                           row(_row), col(_col), eid(_eid), ts(_ts), dts(_dts),
                           nodes(_nodes), dim_in(_dim_in), dim_out(_dim_out) {}
};

class ParallelSampler
{
    public:
        std::vector<EdgeIDType> indptr;
        std::vector<EdgeIDType> indices;
        std::vector<EdgeIDType> eid;
        std::vector<TimeStampType> ts;
        NodeIDType num_nodes;
        EdgeIDType num_edges;
        int num_thread_per_worker;
        int num_workers;
        int num_threads;
        int num_layers;
        std::vector<int> num_neighbors;
        bool recent;
        bool prop_time;
        int num_history;
        TimeStampType window_duration;
        std::vector<std::vector<std::vector<EdgeIDType>::size_type>> ts_ptr;
        omp_lock_t *ts_ptr_lock;
        std::vector<TemporalGraphBlock> ret;

        ParallelSampler(std::vector<EdgeIDType> &_indptr, std::vector<EdgeIDType> &_indices,
                        std::vector<EdgeIDType> &_eid, std::vector<TimeStampType> &_ts,
                        int _num_thread_per_worker, int _num_workers, int _num_layers,
                        std::vector<int> &_num_neighbors, bool _recent, bool _prop_time,
                        int _num_history, TimeStampType _window_duration) :
                        indptr(_indptr), indices(_indices), eid(_eid), ts(_ts), prop_time(_prop_time),
                        num_thread_per_worker(_num_thread_per_worker), num_workers(_num_workers),
                        num_layers(_num_layers), num_neighbors(_num_neighbors), recent(_recent),
                        num_history(_num_history), window_duration(_window_duration)
        {
            omp_set_num_threads(num_thread_per_worker * num_workers);
            num_threads = num_thread_per_worker * num_workers;
            num_nodes = indptr.size() - 1;
            num_edges = indices.size();
            ts_ptr_lock = (omp_lock_t *)malloc(num_nodes * sizeof(omp_lock_t));
            for (int i = 0; i < num_nodes; i++)
                omp_init_lock(&ts_ptr_lock[i]);
            ts_ptr.resize(num_history + 1);
            for (auto it = ts_ptr.begin(); it != ts_ptr.end(); it++)
            {
                it->resize(indptr.size() - 1);
#pragma omp parallel for
                for (auto itt = indptr.begin(); itt < indptr.end() - 1; itt++)
                    (*it)[itt - indptr.begin()] = *itt;
            }
        }

        void reset()
        {
            for (auto it = ts_ptr.begin(); it != ts_ptr.end(); it++)
            {
                it->resize(indptr.size() - 1);
#pragma omp parallel for
                for (auto itt = indptr.begin(); itt < indptr.end() - 1; itt++)
                    (*it)[itt - indptr.begin()] = *itt;
            }
        }

        void update_ts_ptr(int slc, std::vector<NodeIDType> &root_nodes, 
                           std::vector<TimeStampType> &root_ts, float offset)
        {
#pragma omp parallel for schedule(static, int(ceil(static_cast<float>(root_nodes.size()) / num_threads)))
            for (std::vector<NodeIDType>::size_type i = 0; i < root_nodes.size(); i++)
            {
                NodeIDType n = root_nodes[i];
                omp_set_lock(&(ts_ptr_lock[n]));
                for (std::vector<EdgeIDType>::size_type j = ts_ptr[slc][n]; j < indptr[n + 1]; j++)
                {
                    // std::cout << "comparing " << ts[j] << " with " << root_ts[i] << std::endl;
                    if (ts[j] > (root_ts[i] + offset - 1e-7f))
                    {
                        if (j != ts_ptr[slc][n])
                            ts_ptr[slc][n] = j - 1;
                        break;
                    }
                    if (j == indptr[n + 1] - 1)
                    {
                        ts_ptr[slc][n] = j;
                    }
                }
                omp_unset_lock(&(ts_ptr_lock[n]));
            }
        }

        inline void add_neighbor(std::vector<NodeIDType> *_row, std::vector<NodeIDType> *_col,
                                 std::vector<EdgeIDType> *_eid, std::vector<TimeStampType> *_ts,
                                 std::vector<TimeStampType> *_dts, std::vector<NodeIDType> *_nodes, 
                                 EdgeIDType &k, TimeStampType &src_ts, int &row_id)
        {
            _row->push_back(row_id);
            _col->push_back(_nodes->size());
            _eid->push_back(eid[k]);
            if (prop_time)
                _ts->push_back(src_ts);
            else
                _ts->push_back(ts[k]);
            _dts->push_back(src_ts - ts[k]);
            _nodes->push_back(indices[k]);
            // _row.push_back(0);
            // _col.push_back(0);
            // _eid.push_back(0);
            // if (prop_time)
            //     _ts.push_back(src_ts);
            // else
            //     _ts.push_back(10000);
            // _nodes.push_back(100);
        }

        inline void combine_coo(TemporalGraphBlock &_ret, std::vector<NodeIDType> **_row, 
                                std::vector<NodeIDType> **_col, 
                                std::vector<EdgeIDType> **_eid, 
                                std::vector<TimeStampType> **_ts, 
                                std::vector<TimeStampType> **_dts,
                                std::vector<NodeIDType> **_nodes,
                                std::vector<int> &_out_nodes)
        {
            std::vector<EdgeIDType> cum_row, cum_col;
            cum_row.push_back(0);
            cum_col.push_back(0);
            for (int tid = 0; tid < num_threads; tid++)
            {
                // std::cout<<tid<<" here "<<_out_nodes[tid]<<std::endl;  
                cum_row.push_back(cum_row.back() + _out_nodes[tid]);
                cum_col.push_back(cum_col.back() + _col[tid]->size());
            }
            int num_root_nodes = _ret.nodes.size();
            _ret.row.resize(cum_col.back());
            _ret.col.resize(cum_col.back());
            _ret.eid.resize(cum_col.back());
            _ret.ts.resize(cum_col.back() + num_root_nodes);
            _ret.dts.resize(cum_col.back() + num_root_nodes);
            _ret.nodes.resize(cum_col.back() + num_root_nodes);
#pragma omp parallel for schedule(static, 1)
            for (int tid = 0; tid < num_threads; tid++)
            {
                std::transform(_row[tid]->begin(), _row[tid]->end(), _row[tid]->begin(),
                               [&](auto &v){ return v + cum_row[tid]; });
                std::transform(_col[tid]->begin(), _col[tid]->end(), _col[tid]->begin(),
                               [&](auto &v){ return v + cum_col[tid] + num_root_nodes; });
                std::copy(_row[tid]->begin(), _row[tid]->end(), _ret.row.begin() + cum_col[tid]);
                std::copy(_col[tid]->begin(), _col[tid]->end(), _ret.col.begin() + cum_col[tid]);
                std::copy(_eid[tid]->begin(), _eid[tid]->end(), _ret.eid.begin() + cum_col[tid]);
                std::copy(_ts[tid]->begin(), _ts[tid]->end(), _ret.ts.begin() + cum_col[tid] + num_root_nodes);
                std::copy(_dts[tid]->begin(), _dts[tid]->end(), _ret.dts.begin() + cum_col[tid] + num_root_nodes);
                std::copy(_nodes[tid]->begin(), _nodes[tid]->end(), _ret.nodes.begin() + cum_col[tid] + num_root_nodes);
                delete _row[tid];
                delete _col[tid];
                delete _eid[tid];
                delete _ts[tid];
                delete _dts[tid];
                delete _nodes[tid];
            }
            _ret.dim_in = _ret.nodes.size();
            _ret.dim_out = cum_row.back();
        }

        void sample_layer(std::vector<NodeIDType> &_root_nodes, std::vector<TimeStampType> &_root_ts,
                          int neighs, bool use_ptr, bool from_root)
        {
            double t_s = omp_get_wtime();
            std::vector<NodeIDType> *root_nodes;
            std::vector<TimeStampType> *root_ts;
            if (from_root)
            {
                root_nodes = &_root_nodes;
                root_ts = &_root_ts;
            }
            double t_ptr_s = omp_get_wtime();
            if (use_ptr)
                update_ts_ptr(num_history, *root_nodes, *root_ts, 0);
            ret[0].ptr_time += omp_get_wtime() - t_ptr_s;
            for (int i = 0; i < num_history; i++)
            {
                if (!from_root)
                {
                    root_nodes = &(ret[ret.size() - 1 - i - num_history].nodes);
                    root_ts = &(ret[ret.size() - 1 - i - num_history].ts);
                }
                TimeStampType offset = -i * window_duration;
                t_ptr_s = omp_get_wtime();
                if ((use_ptr) && (std::abs(window_duration) > 1e-7f))
                    update_ts_ptr(num_history - 1 - i, *root_nodes, *root_ts, offset - window_duration);
                ret[0].ptr_time += omp_get_wtime() - t_ptr_s;
                std::vector<NodeIDType> *_row[num_threads];
                std::vector<NodeIDType> *_col[num_threads];
                std::vector<EdgeIDType> *_eid[num_threads];
                std::vector<TimeStampType> *_ts[num_threads];
                std::vector<TimeStampType> *_dts[num_threads];
                std::vector<NodeIDType> *_nodes[num_threads];
                std::vector<int> _out_node(num_threads, 0);
                int reserve_capacity = int(ceil((*root_nodes).size() / num_threads)) * neighs;
#pragma omp parallel
                {
                    int tid = omp_get_thread_num();
                    unsigned int loc_seed = tid;
                    _row[tid] = new std::vector<NodeIDType>;
                    _col[tid] = new std::vector<NodeIDType>;
                    _eid[tid] = new std::vector<EdgeIDType>;
                    _ts[tid] = new std::vector<TimeStampType>;
                    _dts[tid] = new std::vector<TimeStampType>;
                    _nodes[tid] = new std::vector<NodeIDType>;
                    _row[tid]->reserve(reserve_capacity);
                    _col[tid]->reserve(reserve_capacity);
                    _eid[tid]->reserve(reserve_capacity);
                    _ts[tid]->reserve(reserve_capacity);
                    _dts[tid]->reserve(reserve_capacity);
                    _nodes[tid]->reserve(reserve_capacity);
// #pragma omp critical
//                     std::cout<<tid<<" sampling: "<<root_nodes->size()<<" "<<int(ceil((*root_nodes).size() / num_threads))<<std::endl;
#pragma omp for schedule(static, int(ceil(static_cast<float>((*root_nodes).size()) / num_threads)))
                    for (std::vector<NodeIDType>::size_type j = 0; j < (*root_nodes).size(); j++)
                    {
                        NodeIDType n = (*root_nodes)[j];
                        // if (tid == 16)
                        //     std::cout << _out_node[tid] << " " <<j << " " << n << std::endl;
                        TimeStampType nts = (*root_ts)[j];
                        EdgeIDType s_search, e_search;
                        if (use_ptr)
                        {
                            s_search = ts_ptr[num_history - 1 - i][n];
                            e_search = ts_ptr[num_history - i][n];
                        }
                        else
                        {
                            // search for start and end pointer
                            double t_search_s = omp_get_wtime();
                            if (num_history == 1)
                            {
                                // TGAT style
                                s_search = indptr[n];
                                auto e_it = std::upper_bound(ts.begin() + indptr[n], 
                                                             ts.begin() + indptr[n + 1], nts);
                                e_search = std::max(int(e_it - ts.begin()) - 1, s_search);
                            }
                            else
                            {
                                // DySAT style
                                auto s_it = std::upper_bound(ts.begin() + indptr[n],
                                                             ts.begin() + indptr[n + 1],
                                                             nts + offset - window_duration);
                                s_search = std::max(int(s_it - ts.begin()) - 1, indptr[n]);
                                auto e_it = std::upper_bound(ts.begin() + indptr[n],
                                                             ts.begin() + indptr[n + 1], nts + offset);
                                e_search = std::max(int(e_it - ts.begin()) - 1, s_search);
                            }
                            if (tid == 0)
                                ret[0].search_time += omp_get_wtime() - t_search_s;
                        }
                        // std::cout << n << " " << s_search << " " << e_search << std::endl;
                        double t_sample_s = omp_get_wtime();
                        if ((recent) || (e_search - s_search < neighs))
                        {                            
                            // no sampling, pick recent neighbors
                            for (EdgeIDType k = e_search; k > std::max(s_search, e_search - neighs); k--)
                            {
                                if (ts[k] < nts + offset - 1e-7f)
                                {
                                    add_neighbor(_row[tid], _col[tid], _eid[tid], _ts[tid], 
                                                 _dts[tid], _nodes[tid], k, nts, _out_node[tid]);
                                }
                            }
                        }
                        else
                        {
                            // random sampling within ptr
                            for (int _i = 0; _i < neighs; _i++)
                            {
                                EdgeIDType picked = s_search + rand_r(&loc_seed) % (e_search - s_search + 1);
                                if (ts[picked] < nts + offset - 1e-7f)
                                {
                                    add_neighbor(_row[tid], _col[tid], _eid[tid], _ts[tid], 
                                                 _dts[tid], _nodes[tid], picked, nts, _out_node[tid]);
                                }
                            }
                        }
                        _out_node[tid] += 1;
                        if (tid == 0)
                            ret[0].sample_time += omp_get_wtime() - t_sample_s;
                    }
                }
                double t_coo_s = omp_get_wtime();
                ret[ret.size() - 1 - i].ts.insert(ret[ret.size() - 1 - i].ts.end(), 
                                                  root_ts->begin(), root_ts->end());
                ret[ret.size() - 1 - i].nodes.insert(ret[ret.size() - 1 - i].nodes.end(), 
                                                     root_nodes->begin(), root_nodes->end());
                ret[ret.size() - 1 - i].dts.resize(root_nodes->size());
                combine_coo(ret[ret.size() - 1 - i], _row, _col, _eid, _ts, _dts, _nodes, _out_node);
                ret[0].coo_time += omp_get_wtime() - t_coo_s;
            }
            ret[0].tot_time += omp_get_wtime() - t_s;
        }

        void sample(std::vector<NodeIDType> &root_nodes, std::vector<TimeStampType> &root_ts)
        {
            // a weird bug, dgl library seems to modify the total number of threads
            omp_set_num_threads(num_threads);
            ret.resize(0);
            bool first_layer = true;
            bool use_ptr = false;
            for (int i = 0; i < num_layers; i++)
            {
                ret.resize(ret.size() + num_history);
                if ((first_layer) || ((prop_time) && num_history == 1) || (recent))
                {
                    first_layer = false;
                    use_ptr = true;
                }
                else
                    use_ptr = false;
                if (i==0)
                    sample_layer(root_nodes, root_ts, num_neighbors[i], use_ptr, true);
                else
                    sample_layer(root_nodes, root_ts, num_neighbors[i], use_ptr, false);
            }
        }
};

template<typename T>
inline py::array vec2npy(const std::vector<T> &vec)
{
    // need to let python garbage collector handle C++ vector memory 
    // see https://github.com/pybind/pybind11/issues/1042
    auto v = new std::vector<T>(vec);
    auto capsule = py::capsule(v, [](void *v)
                               { delete reinterpret_cast<std::vector<T> *>(v); });
    return py::array(v->size(), v->data(), capsule);
    // return py::array(vec.size(), vec.data());
}

PYBIND11_MODULE(sampler_core, m)
{
    py::class_<TemporalGraphBlock>(m, "TemporalGraphBlock")
        .def(py::init<std::vector<NodeIDType> &, std::vector<NodeIDType> &,
                      std::vector<EdgeIDType> &, std::vector<TimeStampType> &,
                      std::vector<TimeStampType> &, std::vector<NodeIDType> &,
                      NodeIDType, NodeIDType>())
        .def("row", [](const TemporalGraphBlock &tgb) { return vec2npy(tgb.row); })
        .def("col", [](const TemporalGraphBlock &tgb) { return vec2npy(tgb.col); })
        .def("eid", [](const TemporalGraphBlock &tgb) { return vec2npy(tgb.eid); })
        .def("ts", [](const TemporalGraphBlock &tgb) { return vec2npy(tgb.ts); })
        .def("dts", [](const TemporalGraphBlock &tgb) { return vec2npy(tgb.dts); })
        .def("nodes", [](const TemporalGraphBlock &tgb) { return vec2npy(tgb.nodes); })
        .def("dim_in", [](const TemporalGraphBlock &tgb) { return tgb.dim_in; })
        .def("dim_out", [](const TemporalGraphBlock &tgb) { return tgb.dim_out; })
        .def("tot_time", [](const TemporalGraphBlock &tgb) { return tgb.tot_time; })
        .def("ptr_time", [](const TemporalGraphBlock &tgb) { return tgb.ptr_time; })
        .def("search_time", [](const TemporalGraphBlock &tgb) { return tgb.search_time; })
        .def("sample_time", [](const TemporalGraphBlock &tgb) { return tgb.sample_time; })
        .def("coo_time", [](const TemporalGraphBlock &tgb) { return tgb.coo_time; });
    py::class_<ParallelSampler>(m, "ParallelSampler")
        .def(py::init<std::vector<EdgeIDType> &, std::vector<EdgeIDType> &,
                      std::vector<EdgeIDType> &, std::vector<TimeStampType> &,
                      int, int, int, std::vector<int> &, bool, bool,
                      int, TimeStampType>())
        .def("sample", &ParallelSampler::sample)
        .def("reset", &ParallelSampler::reset)
        .def("get_ret", [](const ParallelSampler &ps) { return ps.ret; });
}