项目文件夹

文件
Minjie Wang 870da747ea [CUDA][Kernel] More CUDA kernels; Standardize the behavior for sorted COO/CSR (#1704)
* add cub; array cumsum

* CSRSliceRows

* fix warning

* operator << for ndarray; CSRSliceRows

* add CSRIsSorted

* add csr_sort

* inplace coosort and outplace csrsort

* WIP: coo is sorted

* mv cuda_utils

* add AllTrue utility

* csr sort

* coo sort

* coo2csr for sorted coo arrays

* CSRToCOO from sorted

* pass tests for the new kernel changes

* cannot use inplace sort

* lint

* try fix msvc error

* Fix g.copy_to and g.asnumbits; ToBlock no longer uses CSC

* stash

* revert some hack

* revert some changes

* address comments

* fix

* fix to_block unittest

* add todo note
2020-06-28 18:37:28 +08:00

63 行
1.6 KiB
C++

#ifndef TEST_COMMON_H_
#define TEST_COMMON_H_
#include <dgl/runtime/ndarray.h>
static constexpr DLContext CTX = DLContext{kDLCPU, 0};
static constexpr DLContext CPU = DLContext{kDLCPU, 0};
#ifdef DGL_USE_CUDA
static constexpr DLContext GPU = DLContext{kDLGPU, 0};
#endif
template <typename T>
inline T* Ptr(dgl::runtime::NDArray nd) {
return static_cast<T*>(nd->data);
}
inline int64_t* PI64(dgl::runtime::NDArray nd) {
return static_cast<int64_t*>(nd->data);
}
inline int32_t* PI32(dgl::runtime::NDArray nd) {
return static_cast<int32_t*>(nd->data);
}
inline int64_t Len(dgl::runtime::NDArray nd) {
return nd->shape[0];
}
template <typename T>
inline bool ArrayEQ(dgl::runtime::NDArray a1, dgl::runtime::NDArray a2) {
if (a1->ndim != a2->ndim) return false;
if (a1->dtype != a2->dtype) return false;
if (a1->ctx != a2->ctx) return false;
if (a1.NumElements() != a2.NumElements()) return false;
if (a1.NumElements() == 0) return true;
int64_t num = 1;
for (int i = 0; i < a1->ndim; ++i) {
if (a1->shape[i] != a2->shape[i])
return false;
num *= a1->shape[i];
}
if (a1->ctx != a2->ctx) return false;
a1 = a1.CopyTo(CPU);
a2 = a2.CopyTo(CPU);
for (int64_t i = 0; i < num; ++i)
if (static_cast<T*>(a1->data)[i] != static_cast<T*>(a2->data)[i])
return false;
return true;
}
template <typename T>
inline bool IsInArray(dgl::runtime::NDArray a, T x) {
if (!a.defined() || a->shape[0] == 0)
return false;
for (int64_t i = 0; i < a->shape[0]; ++i) {
if (x == static_cast<T*>(a->data)[i])
return true;
}
return false;
}
#endif // TEST_COMMON_H_