Last active
April 2, 2019 18:13
-
-
Save cwpearson/e4fe1ecd60451d49003d9e47e1043ddd to your computer and use it in GitHub Desktop.
Example COO/CSR implementation. Read the comments at the bottom.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#pragma once | |
#include <set> | |
#include "coo.hpp" | |
#ifdef __CUDACC__ | |
#define PANGOLIN_HOST __host__ | |
#define DEVICE __device__ | |
#else | |
#define PANGOLIN_HOST | |
#define DEVICE | |
#endif | |
namespace pangolin { | |
template <typename Index> COO<Index>::COO() {} | |
template <typename Index> PANGOLIN_HOST DEVICE uint64_t COO<Index>::num_rows() const { | |
if (rowPtr_.size() == 0) { | |
return 0; | |
} else { | |
return rowPtr_.size() - 1; | |
} | |
} | |
template <typename Index> uint64_t COO<Index>::num_nodes() const { | |
std::set<Index> nodes; | |
// add all dsts | |
for (Index ci = 0; ci < colInd_.size(); ++ci) { | |
nodes.insert(colInd_[ci]); | |
} | |
// add non-zero sources | |
for (Index i = 0; i < rowPtr_.size() - 1; ++i) { | |
Index row_start = rowPtr_[i]; | |
Index row_end = rowPtr_[i + 1]; | |
if (row_start != row_end) { | |
nodes.insert(i); | |
} | |
} | |
return nodes.size(); | |
} | |
/*! Build a COO from a sequence of edges | |
*/ | |
template <typename Index> | |
template <typename EdgeIter> | |
COO<Index> COO<Index>::from_edges(EdgeIter begin, EdgeIter end, std::function<bool(EdgeTy<Index>)> f) { | |
COO<Index> coo; | |
if (begin == end) { | |
LOG(warn, "constructing from empty edge sequence"); | |
return coo; | |
} | |
for (auto ei = begin; ei != end; ++ei) { | |
EdgeTy<Index> edge = *ei; | |
const Index src = edge.first; | |
const Index dst = edge.second; | |
// edge has a new src and should be in a new row | |
// even if the edge is filtered out, we need to add empty rows | |
while (coo.rowPtr_.size() != size_t(src + 1)) { | |
// expecting inputs to be sorted by src, so it should be at least | |
// as big as the current largest row we have recored | |
assert(src >= coo.rowPtr_.size() && "are edges not ordered by source?"); | |
coo.rowPtr_.push_back(coo.colInd_.size()); | |
} | |
if (f(edge)) { | |
coo.rowInd_.push_back(src); | |
coo.colInd_.push_back(dst); | |
} else { | |
continue; | |
} | |
} | |
// add the final length of the non-zeros to the offset array | |
coo.rowPtr_.push_back(coo.colInd_.size()); | |
assert(coo.rowInd_.size() == coo.colInd_.size()); | |
return coo; | |
} | |
} // namespace pangolin | |
#undef PANGOLIN_HOST | |
#undef DEVICE |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#pragma once | |
#include <functional> | |
#include "pangolin/dense/vector.hu" | |
#include "pangolin/edge_list.hpp" | |
#include "pangolin/types.hpp" | |
#ifdef __CUDACC__ | |
#define HOST __host__ | |
#define DEVICE __device__ | |
#else | |
#define HOST | |
#define DEVICE | |
#endif | |
namespace pangolin { | |
template <typename Index> class COO; | |
/*! \brief a read-only view of a COO suitable for passing to a GPU kernel by | |
value. | |
Any modifications to the underlying COO may invalidate this view. | |
*/ | |
template <typename Index> class COOView { | |
friend class COO<Index>; | |
private: | |
uint64_t nnz_; //!< number of non-zeros | |
uint64_t num_rows_; //!< length of rowOffset - 1 | |
public: | |
typedef Index index_type; | |
const Index *rowPtr_; //!< offset in col_ that each row starts at | |
const Index *rowInd_; //!< non-zero row indices | |
const Index *colInd_; //!< non-zero column indices | |
HOST DEVICE uint64_t nnz() const { return nnz_; } | |
HOST DEVICE uint64_t num_rows() const { return num_rows_; } | |
const Index *row_ptr() const { return rowPtr_; } //!< row offset array | |
const Index *col_ind() const { return colInd_; } //!< column index array | |
const Index *row_ind() const { return rowInd_; } //<! row index array | |
HOST DEVICE const Index *device_row_ptr() const { return rowPtr_; } //!< row offset array | |
HOST DEVICE const Index *device_col_ind() const { return colInd_; } //!< column index array | |
HOST DEVICE const Index *device_row_ind() const { return rowInd_; } //<! row index array | |
}; | |
/*! \brief A COO matrix backed by CUDA Unified Memory, with a CSR rowPtr | |
Copying to a GPU kernel by value will cause the underling memory to be copied as | |
well. For read-only GPU access, use the view() method to get a lightweight | |
reference to the data. | |
*/ | |
template <typename Index> class COO { | |
private: | |
Index maxCol_; | |
public: | |
typedef Index index_type; | |
COO(); //!< empty CSR | |
Vector<Index> rowPtr_; //!< offset in col_ that each row starts at | |
Vector<Index> colInd_; //!< non-zero column indices | |
Vector<Index> rowInd_; //!< non-zero row indices | |
HOST DEVICE uint64_t nnz() const { | |
assert(colInd_.size() == rowInd_.size()); | |
return colInd_.size(); | |
} //!< number of non-zeros | |
uint64_t num_nodes() const; //!< number of unique row/col indices | |
HOST DEVICE uint64_t num_rows() const; //!< number of matrix rows | |
/*! Build a COO from a sequence of edges | |
Only include edges where f is true (default = all edges) | |
*/ | |
template <typename EdgeIter> | |
static COO<Index> from_edges(EdgeIter begin, EdgeIter end, | |
std::function<bool(EdgeTy<Index>)> f = [](EdgeTy<Index> e) { return true; }); | |
COOView<Index> view() const; //!< create a COOView for this COO | |
const Index *row_ptr() const { return rowPtr_.data(); } //!< row offset array | |
const Index *col_ind() const { return colInd_.data(); } //!< column index array | |
const Index *row_ind() const { return rowInd_.data(); } //<! row index array | |
HOST DEVICE const Index *device_row_ptr() const { return rowPtr_.data(); } //!< row offset array | |
HOST DEVICE const Index *device_col_ind() const { return colInd_.data(); } //!< column index array | |
HOST DEVICE const Index *device_row_ind() const { return rowInd_.data(); } //<! row index array | |
}; | |
} // namespace pangolin | |
#undef HOST | |
#undef DEVICE | |
#include "coo-impl.hpp" |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Probably will not compile as-is, but should give a general idea:
pangolin
namespace.