wehub-resource-sync
138 lines
3.5 KiB
C++
138 lines
3.5 KiB
C++
/**
|
|
* Copyright (c) 2020 by Contributors
|
|
* @file dgl/aten/spmat.h
|
|
* @brief Sparse matrix definitions
|
|
*/
|
|
#ifndef DGL_ATEN_SPMAT_H_
|
|
#define DGL_ATEN_SPMAT_H_
|
|
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "../runtime/object.h"
|
|
#include "./types.h"
|
|
|
|
namespace dgl {
|
|
|
|
/**
|
|
* @brief Sparse format.
|
|
*/
|
|
enum class SparseFormat {
|
|
kCOO = 1,
|
|
kCSR = 2,
|
|
kCSC = 3,
|
|
};
|
|
|
|
/**
|
|
* @brief Sparse format codes
|
|
*/
|
|
const dgl_format_code_t ALL_CODE = 0x7;
|
|
const dgl_format_code_t ANY_CODE = 0x0;
|
|
const dgl_format_code_t COO_CODE = 0x1;
|
|
const dgl_format_code_t CSR_CODE = 0x2;
|
|
const dgl_format_code_t CSC_CODE = 0x4;
|
|
|
|
// Parse sparse format from string.
|
|
inline SparseFormat ParseSparseFormat(const std::string& name) {
|
|
if (name == "coo")
|
|
return SparseFormat::kCOO;
|
|
else if (name == "csr")
|
|
return SparseFormat::kCSR;
|
|
else if (name == "csc")
|
|
return SparseFormat::kCSC;
|
|
else
|
|
LOG(FATAL) << "Sparse format not recognized";
|
|
return SparseFormat::kCOO;
|
|
}
|
|
|
|
// Create string from sparse format.
|
|
inline std::string ToStringSparseFormat(SparseFormat sparse_format) {
|
|
if (sparse_format == SparseFormat::kCOO)
|
|
return std::string("coo");
|
|
else if (sparse_format == SparseFormat::kCSR)
|
|
return std::string("csr");
|
|
else
|
|
return std::string("csc");
|
|
}
|
|
|
|
inline std::vector<SparseFormat> CodeToSparseFormats(dgl_format_code_t code) {
|
|
std::vector<SparseFormat> ret;
|
|
if (code & COO_CODE) ret.push_back(SparseFormat::kCOO);
|
|
if (code & CSR_CODE) ret.push_back(SparseFormat::kCSR);
|
|
if (code & CSC_CODE) ret.push_back(SparseFormat::kCSC);
|
|
return ret;
|
|
}
|
|
|
|
inline dgl_format_code_t SparseFormatsToCode(
|
|
const std::vector<SparseFormat>& formats) {
|
|
dgl_format_code_t ret = 0;
|
|
for (auto format : formats) {
|
|
switch (format) {
|
|
case SparseFormat::kCOO:
|
|
ret |= COO_CODE;
|
|
break;
|
|
case SparseFormat::kCSR:
|
|
ret |= CSR_CODE;
|
|
break;
|
|
case SparseFormat::kCSC:
|
|
ret |= CSC_CODE;
|
|
break;
|
|
default:
|
|
LOG(FATAL) << "Only support COO/CSR/CSC formats.";
|
|
}
|
|
}
|
|
return ret;
|
|
}
|
|
|
|
inline std::string CodeToStr(dgl_format_code_t code) {
|
|
std::string ret = "";
|
|
if (code & COO_CODE) ret += "coo ";
|
|
if (code & CSR_CODE) ret += "csr ";
|
|
if (code & CSC_CODE) ret += "csc ";
|
|
return ret;
|
|
}
|
|
|
|
inline SparseFormat DecodeFormat(dgl_format_code_t code) {
|
|
if (code & COO_CODE) return SparseFormat::kCOO;
|
|
if (code & CSC_CODE) return SparseFormat::kCSC;
|
|
return SparseFormat::kCSR;
|
|
}
|
|
|
|
// Sparse matrix object that is exposed to python API.
|
|
struct SparseMatrix : public runtime::Object {
|
|
// Sparse format.
|
|
int32_t format = 0;
|
|
|
|
// Shape of this matrix.
|
|
int64_t num_rows = 0, num_cols = 0;
|
|
|
|
// Index arrays. For CSR, it is {indptr, indices, data}. For COO, it is {row,
|
|
// col, data}.
|
|
std::vector<IdArray> indices;
|
|
|
|
// Boolean flags.
|
|
// TODO(minjie): We might revisit this later to provide a more general
|
|
// solution. Currently, we only consider aten::COOMatrix and aten::CSRMatrix.
|
|
std::vector<bool> flags;
|
|
|
|
SparseMatrix() {}
|
|
|
|
SparseMatrix(
|
|
int32_t fmt, int64_t nrows, int64_t ncols,
|
|
const std::vector<IdArray>& idx, const std::vector<bool>& flg)
|
|
: format(fmt),
|
|
num_rows(nrows),
|
|
num_cols(ncols),
|
|
indices(idx),
|
|
flags(flg) {}
|
|
|
|
static constexpr const char* _type_key = "aten.SparseMatrix";
|
|
DGL_DECLARE_OBJECT_TYPE_INFO(SparseMatrix, runtime::Object);
|
|
};
|
|
// Define SparseMatrixRef
|
|
DGL_DEFINE_OBJECT_REF(SparseMatrixRef, SparseMatrix);
|
|
|
|
} // namespace dgl
|
|
|
|
#endif // DGL_ATEN_SPMAT_H_
|