dmlc--dgl
Quan (Andy) Gan
929d863447
* test * more stuff * add test * fixes * optimize algo * replace unordered_map with arrays * lint * lint x2 * oops * disable gpu csrmm tests * remove gpu invocation * optimize with openmp * remove python functions * add back with docstrings * lint * lint * update python interface * functionize * functionize * lint * lint
210 行
7.1 KiB
C++
210 行
7.1 KiB
C++
#include <gtest/gtest.h>
|
|
#include <dgl/array.h>
|
|
#include <dgl/kernel.h>
|
|
#include "../../src/array/cpu/array_utils.h" // PairHash
|
|
#include "./common.h"
|
|
|
|
using namespace dgl;
|
|
using namespace dgl::runtime;
|
|
|
|
namespace {
|
|
|
|
// Unit tests:
|
|
// CSRMM(A, B) == A_mm_B
|
|
// CSRSum({A, C}) == A_plus_C
|
|
// CSRMask(A, C) = A_mask_C
|
|
|
|
template <typename IdType, typename DType>
|
|
std::unordered_map<std::pair<IdType, IdType>, DType, aten::PairHash> COOToMap(
|
|
aten::COOMatrix coo, NDArray weights) {
|
|
std::unordered_map<std::pair<IdType, IdType>, DType, aten::PairHash> map;
|
|
|
|
for (int64_t i = 0; i < coo.row->shape[0]; ++i) {
|
|
IdType irow = aten::IndexSelect<IdType>(coo.row, i);
|
|
IdType icol = aten::IndexSelect<IdType>(coo.col, i);
|
|
IdType ieid = aten::COOHasData(coo) ? aten::IndexSelect<IdType>(coo.data, i) : i;
|
|
DType idata = aten::IndexSelect<DType>(weights, ieid);
|
|
map.insert({{irow, icol}, idata});
|
|
}
|
|
return map;
|
|
}
|
|
|
|
template <typename IdType, typename DType>
|
|
bool CSRIsClose(
|
|
aten::CSRMatrix A,
|
|
aten::CSRMatrix B,
|
|
NDArray A_weights,
|
|
NDArray B_weights,
|
|
DType rtol,
|
|
DType atol) {
|
|
auto Amap = COOToMap<IdType, DType>(CSRToCOO(A, false), A_weights);
|
|
auto Bmap = COOToMap<IdType, DType>(CSRToCOO(B, false), B_weights);
|
|
|
|
if (Amap.size() != Bmap.size())
|
|
return false;
|
|
|
|
for (auto itA : Amap) {
|
|
auto itB = Bmap.find(itA.first);
|
|
if (itB == Bmap.end())
|
|
return false;
|
|
if (fabs(itA.second - itB->second) >= rtol * fabs(itA.second) + atol)
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
template <typename IdType, typename DType>
|
|
std::pair<aten::CSRMatrix, NDArray> CSR_A(DLContext ctx = CTX) {
|
|
// matrix([[0. , 0. , 1. , 0.7, 0. ],
|
|
// [0. , 0. , 0.5, 0.+, 0. ],
|
|
// [0.4, 0.7, 0. , 0.2, 0. ],
|
|
// [0. , 0. , 0. , 0. , 0.2]])
|
|
// (0.+ indicates that the entry exists but the value is 0.)
|
|
auto csr = aten::CSRMatrix(
|
|
4, 5,
|
|
NDArray::FromVector(std::vector<IdType>({0, 2, 4, 7, 8}), ctx),
|
|
NDArray::FromVector(std::vector<IdType>({2, 3, 2, 3, 0, 1, 3, 4}), ctx));
|
|
auto weights = NDArray::FromVector(
|
|
std::vector<DType>({1.0, 0.7, 0.5, 0.0, 0.4, 0.7, 0.2, 0.2}), ctx);
|
|
return {csr, weights};
|
|
}
|
|
|
|
template <typename IdType, typename DType>
|
|
std::pair<aten::CSRMatrix, NDArray> CSR_B(DLContext ctx = CTX) {
|
|
// matrix([[0. , 0.9, 0. , 0.6, 0. , 0.3],
|
|
// [0. , 0. , 0. , 0. , 0. , 0.4],
|
|
// [0.+, 0. , 0. , 0. , 0. , 0.9],
|
|
// [0.8, 0.2, 0.3, 0.2, 0. , 0. ],
|
|
// [0.2, 0.4, 0. , 0. , 0. , 0. ]])
|
|
// (0.+ indicates that the entry exists but the value is 0.)
|
|
auto csr = aten::CSRMatrix(
|
|
5, 6,
|
|
NDArray::FromVector(std::vector<IdType>({0, 3, 4, 6, 10, 12}), ctx),
|
|
NDArray::FromVector(std::vector<IdType>({1, 3, 5, 5, 0, 5, 0, 1, 2, 3, 0, 1}), ctx));
|
|
auto weights = NDArray::FromVector(
|
|
std::vector<DType>({0.9, 0.6, 0.3, 0.4, 0.0, 0.9, 0.8, 0.2, 0.3, 0.2, 0.2, 0.4}), ctx);
|
|
return {csr, weights};
|
|
}
|
|
|
|
template <typename IdType, typename DType>
|
|
std::pair<aten::CSRMatrix, NDArray> CSR_C(DLContext ctx = CTX) {
|
|
// matrix([[0. , 0. , 0. , 0.2, 0. ],
|
|
// [0. , 0. , 0. , 0.5, 0.4],
|
|
// [0. , 0.2, 0. , 0.9, 0.2],
|
|
// [0. , 1. , 0. , 0.7, 0. ]])
|
|
auto csr = aten::CSRMatrix(
|
|
4, 5,
|
|
NDArray::FromVector(std::vector<IdType>({0, 1, 3, 6, 8}), ctx),
|
|
NDArray::FromVector(std::vector<IdType>({3, 3, 4, 1, 3, 4, 1, 3}), ctx));
|
|
auto weights = NDArray::FromVector(
|
|
std::vector<DType>({0.2, 0.5, 0.4, 0.2, 0.9, 0.2, 1. , 0.7}), ctx);
|
|
return {csr, weights};
|
|
}
|
|
|
|
template <typename IdType, typename DType>
|
|
std::pair<aten::CSRMatrix, NDArray> CSR_A_mm_B(DLContext ctx = CTX) {
|
|
// matrix([[0.56, 0.14, 0.21, 0.14, 0. , 0.9 ],
|
|
// [0.+ , 0.+ , 0.+ , 0.+ , 0. , 0.45],
|
|
// [0.16, 0.4 , 0.06, 0.28, 0. , 0.4 ],
|
|
// [0.04, 0.08, 0. , 0. , 0. , 0. ]])
|
|
// (0.+ indicates that the entry exists but the value is 0.)
|
|
auto csr = aten::CSRMatrix(
|
|
4, 6,
|
|
NDArray::FromVector(std::vector<IdType>({0, 5, 10, 15, 17}), ctx),
|
|
NDArray::FromVector(std::vector<IdType>(
|
|
{0, 1, 2, 3, 5, 0, 1, 2, 3, 5, 0, 1, 2, 3, 5, 0, 1}), ctx));
|
|
auto weights = NDArray::FromVector(
|
|
std::vector<DType>({
|
|
0.56, 0.14, 0.21, 0.14, 0.9 , 0. , 0. , 0. , 0. , 0.45, 0.16, 0.4 , 0.06, 0.28, 0.4 ,
|
|
0.04, 0.08}), ctx);
|
|
return {csr, weights};
|
|
}
|
|
|
|
template <typename IdType, typename DType>
|
|
std::pair<aten::CSRMatrix, NDArray> CSR_A_plus_C(DLContext ctx = CTX) {
|
|
auto csr = aten::CSRMatrix(
|
|
4, 5,
|
|
NDArray::FromVector(std::vector<IdType>({0, 2, 5, 9, 12}), ctx),
|
|
NDArray::FromVector(std::vector<IdType>({2, 3, 2, 3, 4, 0, 1, 3, 4, 1, 3, 4}), ctx));
|
|
auto weights = NDArray::FromVector(
|
|
std::vector<DType>({1. , 0.9, 0.5, 0.5, 0.4, 0.4, 0.9, 1.1, 0.2, 1. , 0.7, 0.2}), ctx);
|
|
return {csr, weights};
|
|
}
|
|
|
|
template <typename DType>
|
|
NDArray CSR_A_mask_C(DLContext ctx = CTX) {
|
|
return NDArray::FromVector(std::vector<DType>({0.7, 0.0, 0.0, 0.7, 0.2, 0.0, 0.0, 0.0}), ctx);
|
|
}
|
|
|
|
template <typename IdType, typename DType>
|
|
void _TestCsrmm(DLContext ctx = CTX) {
|
|
auto A = CSR_A<IdType, DType>(ctx);
|
|
auto B = CSR_B<IdType, DType>(ctx);
|
|
auto A_mm_B = aten::CSRMM(A.first, A.second, B.first, B.second);
|
|
auto A_mm_B2 = CSR_A_mm_B<IdType, DType>(ctx);
|
|
bool result = CSRIsClose<IdType, DType>(A_mm_B.first, A_mm_B2.first, A_mm_B.second, A_mm_B2.second, 1e-4, 1e-4);
|
|
ASSERT_TRUE(result);
|
|
}
|
|
|
|
template <typename IdType, typename DType>
|
|
void _TestCsrsum(DLContext ctx = CTX) {
|
|
auto A = CSR_A<IdType, DType>(ctx);
|
|
auto C = CSR_C<IdType, DType>(ctx);
|
|
auto A_plus_C = aten::CSRSum({A.first, C.first}, {A.second, C.second});
|
|
auto A_plus_C2 = CSR_A_plus_C<IdType, DType>(ctx);
|
|
bool result = CSRIsClose<IdType, DType>(
|
|
A_plus_C.first, A_plus_C2.first, A_plus_C.second, A_plus_C2.second, 1e-4, 1e-4);
|
|
ASSERT_TRUE(result);
|
|
}
|
|
|
|
template <typename IdType, typename DType>
|
|
void _TestCsrmask(DLContext ctx = CTX) {
|
|
auto A = CSR_A<IdType, DType>(ctx);
|
|
auto C = CSR_C<IdType, DType>(ctx);
|
|
auto A_mask_C = aten::CSRMask(A.first, A.second, C.first);
|
|
auto A_mask_C2 = CSR_A_mask_C<DType>(ctx);
|
|
ASSERT_TRUE(ArrayEQ<DType>(A_mask_C, A_mask_C2));
|
|
}
|
|
|
|
TEST(CsrmmTest, TestCsrmm) {
|
|
_TestCsrmm<int32_t, float>(CPU);
|
|
_TestCsrmm<int32_t, double>(CPU);
|
|
_TestCsrmm<int64_t, float>(CPU);
|
|
_TestCsrmm<int64_t, double>(CPU);
|
|
#ifdef DGL_USE_CUDA
|
|
_TestCsrmm<int32_t, float>(GPU);
|
|
_TestCsrmm<int32_t, double>(GPU);
|
|
_TestCsrmm<int64_t, float>(GPU);
|
|
_TestCsrmm<int64_t, double>(GPU);
|
|
#endif
|
|
}
|
|
|
|
TEST(CsrmmTest, TestCsrsum) {
|
|
_TestCsrsum<int32_t, float>(CPU);
|
|
_TestCsrsum<int32_t, double>(CPU);
|
|
_TestCsrsum<int64_t, float>(CPU);
|
|
_TestCsrsum<int64_t, double>(CPU);
|
|
#ifdef DGL_USE_CUDA
|
|
_TestCsrsum<int32_t, float>(GPU);
|
|
_TestCsrsum<int32_t, double>(GPU);
|
|
_TestCsrsum<int64_t, float>(GPU);
|
|
_TestCsrsum<int64_t, double>(GPU);
|
|
#endif
|
|
}
|
|
|
|
TEST(CsrmmTest, TestCsrmask) {
|
|
_TestCsrmask<int32_t, float>(CPU);
|
|
_TestCsrmask<int32_t, double>(CPU);
|
|
_TestCsrmask<int64_t, float>(CPU);
|
|
_TestCsrmask<int64_t, double>(CPU);
|
|
#ifdef DGL_USE_CUDA
|
|
_TestCsrmask<int32_t, float>(GPU);
|
|
_TestCsrmask<int32_t, double>(GPU);
|
|
_TestCsrmask<int64_t, float>(GPU);
|
|
_TestCsrmask<int64_t, double>(GPU);
|
|
#endif
|
|
}
|
|
|
|
}; // namespace
|