chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,150 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file api/api_container.cc
|
||||
* @brief Runtime container APIs. (reference: tvm/src/api/api_lang.cc)
|
||||
*/
|
||||
#include <dgl/packed_func_ext.h>
|
||||
#include <dgl/runtime/container.h>
|
||||
#include <dgl/runtime/ndarray.h>
|
||||
#include <dgl/runtime/registry.h>
|
||||
|
||||
namespace dgl {
|
||||
namespace runtime {
|
||||
|
||||
DGL_REGISTER_GLOBAL("_List").set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
auto ret_obj = std::make_shared<runtime::ListObject>();
|
||||
for (int i = 0; i < args.size(); ++i) {
|
||||
ret_obj->data.push_back(args[i].obj_sptr());
|
||||
}
|
||||
*rv = ret_obj;
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("_ListGetItem").set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
auto& sptr = args[0].obj_sptr();
|
||||
CHECK(sptr->is_type<ListObject>());
|
||||
auto* o = static_cast<const ListObject*>(sptr.get());
|
||||
int64_t i = args[1];
|
||||
CHECK_LT(i, o->data.size()) << "list out of bound";
|
||||
*rv = o->data[i];
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("_ListSize").set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
auto& sptr = args[0].obj_sptr();
|
||||
CHECK(sptr->is_type<ListObject>());
|
||||
auto* o = static_cast<const ListObject*>(sptr.get());
|
||||
*rv = static_cast<int64_t>(o->data.size());
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("_Map").set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
CHECK_EQ(args.size() % 2, 0);
|
||||
if (args.size() != 0 && args[0].type_code() == kStr) {
|
||||
// StrMap
|
||||
StrMapObject::ContainerType data;
|
||||
for (int i = 0; i < args.size(); i += 2) {
|
||||
CHECK(args[i].type_code() == kStr) << "The key of the map must be string";
|
||||
CHECK(args[i + 1].type_code() == kObjectHandle)
|
||||
<< "The value of the map must be an object type";
|
||||
data.emplace(std::make_pair(
|
||||
args[i].operator std::string(), args[i + 1].obj_sptr()));
|
||||
}
|
||||
auto obj = std::make_shared<StrMapObject>();
|
||||
obj->data = std::move(data);
|
||||
*rv = obj;
|
||||
} else {
|
||||
// object container
|
||||
MapObject::ContainerType data;
|
||||
for (int i = 0; i < args.size(); i += 2) {
|
||||
CHECK(args[i].type_code() == kObjectHandle)
|
||||
<< "The key of the map must be an object type";
|
||||
CHECK(args[i + 1].type_code() == kObjectHandle)
|
||||
<< "The value of the map must be an object type";
|
||||
data.emplace(std::make_pair(args[i].obj_sptr(), args[i + 1].obj_sptr()));
|
||||
}
|
||||
auto obj = std::make_shared<MapObject>();
|
||||
obj->data = std::move(data);
|
||||
*rv = obj;
|
||||
}
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("_EmptyStrMap").set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
StrMapObject::ContainerType data;
|
||||
auto obj = std::make_shared<StrMapObject>();
|
||||
obj->data = std::move(data);
|
||||
*rv = obj;
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("_MapSize").set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
auto& sptr = args[0].obj_sptr();
|
||||
if (sptr->is_type<MapObject>()) {
|
||||
auto* o = static_cast<const MapObject*>(sptr.get());
|
||||
*rv = static_cast<int64_t>(o->data.size());
|
||||
} else {
|
||||
CHECK(sptr->is_type<StrMapObject>());
|
||||
auto* o = static_cast<const StrMapObject*>(sptr.get());
|
||||
*rv = static_cast<int64_t>(o->data.size());
|
||||
}
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("_MapGetItem").set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
auto& sptr = args[0].obj_sptr();
|
||||
if (sptr->is_type<MapObject>()) {
|
||||
auto* o = static_cast<const MapObject*>(sptr.get());
|
||||
auto it = o->data.find(args[1].obj_sptr());
|
||||
CHECK(it != o->data.end()) << "cannot find the key in the map";
|
||||
*rv = (*it).second;
|
||||
} else {
|
||||
CHECK(sptr->is_type<StrMapObject>());
|
||||
auto* o = static_cast<const StrMapObject*>(sptr.get());
|
||||
auto it = o->data.find(args[1].operator std::string());
|
||||
CHECK(it != o->data.end()) << "cannot find the key in the map";
|
||||
*rv = (*it).second;
|
||||
}
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("_MapItems").set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
auto& sptr = args[0].obj_sptr();
|
||||
if (sptr->is_type<MapObject>()) {
|
||||
auto* o = static_cast<const MapObject*>(sptr.get());
|
||||
auto rkvs = std::make_shared<ListObject>();
|
||||
for (const auto& kv : o->data) {
|
||||
rkvs->data.push_back(kv.first);
|
||||
rkvs->data.push_back(kv.second);
|
||||
}
|
||||
*rv = rkvs;
|
||||
} else {
|
||||
CHECK(sptr->is_type<StrMapObject>());
|
||||
auto* o = static_cast<const StrMapObject*>(sptr.get());
|
||||
auto rkvs = std::make_shared<ListObject>();
|
||||
for (const auto& kv : o->data) {
|
||||
rkvs->data.push_back(MakeValue(kv.first));
|
||||
rkvs->data.push_back(kv.second);
|
||||
}
|
||||
*rv = rkvs;
|
||||
}
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("_MapCount").set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
auto& sptr = args[0].obj_sptr();
|
||||
if (sptr->is_type<MapObject>()) {
|
||||
auto* o = static_cast<const MapObject*>(sptr.get());
|
||||
*rv = static_cast<int64_t>(o->data.count(args[1].obj_sptr()));
|
||||
} else {
|
||||
CHECK(sptr->is_type<StrMapObject>());
|
||||
auto* o = static_cast<const StrMapObject*>(sptr.get());
|
||||
*rv = static_cast<int64_t>(o->data.count(args[1].operator std::string()));
|
||||
}
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("_Value").set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
*rv = MakeValue(args[0]);
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("_ValueGet").set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
auto& sptr = args[0].obj_sptr();
|
||||
CHECK(sptr->is_type<ValueObject>());
|
||||
auto* o = static_cast<const ValueObject*>(sptr.get());
|
||||
*rv = o->data;
|
||||
});
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,50 @@
|
||||
/**
|
||||
* Copyright (c) 2022 by Contributors
|
||||
* @file api/api_test.cc
|
||||
* @brief C APIs for testing FFI
|
||||
*/
|
||||
#include <dgl/packed_func_ext.h>
|
||||
#include <dgl/runtime/container.h>
|
||||
#include <dgl/runtime/ndarray.h>
|
||||
#include <dgl/runtime/registry.h>
|
||||
|
||||
#include <thread>
|
||||
|
||||
namespace dgl {
|
||||
namespace runtime {
|
||||
|
||||
// Register an internal API for testing python callback.
|
||||
// It receives two arguments:
|
||||
// - The python callback function.
|
||||
// - The argument to pass to the python callback
|
||||
// It returns what python callback returns
|
||||
DGL_REGISTER_GLOBAL("_TestPythonCallback")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
LOG(INFO) << "Inside C API";
|
||||
PackedFunc fn = args[0];
|
||||
DGLArgs cb_args(args.values + 1, args.type_codes + 1, 1);
|
||||
fn.CallPacked(cb_args, rv);
|
||||
});
|
||||
|
||||
// Register an internal API for testing python callback.
|
||||
// It receives two arguments:
|
||||
// - The python callback function.
|
||||
// - The argument to pass to the python callback
|
||||
// It returns what python callback returns
|
||||
//
|
||||
// The API runs the python callback in a separate thread to test
|
||||
// python GIL is properly released.
|
||||
DGL_REGISTER_GLOBAL("_TestPythonCallbackThread")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
LOG(INFO) << "Inside C API";
|
||||
PackedFunc fn = args[0];
|
||||
auto thr = std::make_shared<std::thread>([fn, args, rv]() {
|
||||
LOG(INFO) << "Callback thread " << std::this_thread::get_id();
|
||||
DGLArgs cb_args(args.values + 1, args.type_codes + 1, 1);
|
||||
fn.CallPacked(cb_args, rv);
|
||||
});
|
||||
thr->join();
|
||||
});
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,109 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file array/arith.h
|
||||
* @brief Arithmetic functors
|
||||
*/
|
||||
#ifndef DGL_ARRAY_ARITH_H_
|
||||
#define DGL_ARRAY_ARITH_H_
|
||||
|
||||
#ifdef __CUDACC__
|
||||
#define DGLDEVICE __device__
|
||||
#define DGLINLINE __forceinline__
|
||||
#else
|
||||
#define DGLDEVICE
|
||||
#define DGLINLINE inline
|
||||
#endif // __CUDACC__
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace arith {
|
||||
|
||||
struct Add {
|
||||
template <typename T>
|
||||
static DGLINLINE DGLDEVICE T Call(const T& t1, const T& t2) {
|
||||
return t1 + t2;
|
||||
}
|
||||
};
|
||||
|
||||
struct Sub {
|
||||
template <typename T>
|
||||
static DGLINLINE DGLDEVICE T Call(const T& t1, const T& t2) {
|
||||
return t1 - t2;
|
||||
}
|
||||
};
|
||||
|
||||
struct Mul {
|
||||
template <typename T>
|
||||
static DGLINLINE DGLDEVICE T Call(const T& t1, const T& t2) {
|
||||
return t1 * t2;
|
||||
}
|
||||
};
|
||||
|
||||
struct Div {
|
||||
template <typename T>
|
||||
static DGLINLINE DGLDEVICE T Call(const T& t1, const T& t2) {
|
||||
return t1 / t2;
|
||||
}
|
||||
};
|
||||
|
||||
struct Mod {
|
||||
template <typename T>
|
||||
static DGLINLINE DGLDEVICE T Call(const T& t1, const T& t2) {
|
||||
return t1 % t2;
|
||||
}
|
||||
};
|
||||
|
||||
struct GT {
|
||||
template <typename T>
|
||||
static DGLINLINE DGLDEVICE bool Call(const T& t1, const T& t2) {
|
||||
return t1 > t2;
|
||||
}
|
||||
};
|
||||
|
||||
struct LT {
|
||||
template <typename T>
|
||||
static DGLINLINE DGLDEVICE bool Call(const T& t1, const T& t2) {
|
||||
return t1 < t2;
|
||||
}
|
||||
};
|
||||
|
||||
struct GE {
|
||||
template <typename T>
|
||||
static DGLINLINE DGLDEVICE bool Call(const T& t1, const T& t2) {
|
||||
return t1 >= t2;
|
||||
}
|
||||
};
|
||||
|
||||
struct LE {
|
||||
template <typename T>
|
||||
static DGLINLINE DGLDEVICE bool Call(const T& t1, const T& t2) {
|
||||
return t1 <= t2;
|
||||
}
|
||||
};
|
||||
|
||||
struct EQ {
|
||||
template <typename T>
|
||||
static DGLINLINE DGLDEVICE bool Call(const T& t1, const T& t2) {
|
||||
return t1 == t2;
|
||||
}
|
||||
};
|
||||
|
||||
struct NE {
|
||||
template <typename T>
|
||||
static DGLINLINE DGLDEVICE bool Call(const T& t1, const T& t2) {
|
||||
return t1 != t2;
|
||||
}
|
||||
};
|
||||
|
||||
struct Neg {
|
||||
template <typename T>
|
||||
static DGLINLINE DGLDEVICE T Call(const T& t1) {
|
||||
return -t1;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace arith
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_ARITH_H_
|
||||
+1309
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,211 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file array/array_aritch.cc
|
||||
* @brief DGL array arithmetic operations
|
||||
*/
|
||||
#include <dgl/packed_func_ext.h>
|
||||
#include <dgl/runtime/container.h>
|
||||
#include <dgl/runtime/ndarray.h>
|
||||
|
||||
#include "../c_api_common.h"
|
||||
#include "./arith.h"
|
||||
#include "./array_op.h"
|
||||
|
||||
using namespace dgl::runtime;
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
|
||||
// Generate operators with both operations being NDArrays.
|
||||
#define BINARY_ELEMENT_OP(name, op) \
|
||||
IdArray name(IdArray lhs, IdArray rhs) { \
|
||||
IdArray ret; \
|
||||
CHECK_SAME_DTYPE(lhs, rhs); \
|
||||
CHECK_SAME_CONTEXT(lhs, rhs); \
|
||||
ATEN_XPU_SWITCH_CUDA(lhs->ctx.device_type, XPU, #name, { \
|
||||
ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, { \
|
||||
ret = impl::BinaryElewise<XPU, IdType, arith::op>(lhs, rhs); \
|
||||
}); \
|
||||
}); \
|
||||
return ret; \
|
||||
}
|
||||
|
||||
// Generate operators with only lhs being NDArray.
|
||||
#define BINARY_ELEMENT_OP_L(name, op) \
|
||||
IdArray name(IdArray lhs, int64_t rhs) { \
|
||||
IdArray ret; \
|
||||
ATEN_XPU_SWITCH_CUDA(lhs->ctx.device_type, XPU, #name, { \
|
||||
ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, { \
|
||||
ret = impl::BinaryElewise<XPU, IdType, arith::op>(lhs, rhs); \
|
||||
}); \
|
||||
}); \
|
||||
return ret; \
|
||||
}
|
||||
|
||||
// Generate operators with only lhs being NDArray.
|
||||
#define BINARY_ELEMENT_OP_R(name, op) \
|
||||
IdArray name(int64_t lhs, IdArray rhs) { \
|
||||
IdArray ret; \
|
||||
ATEN_XPU_SWITCH_CUDA(rhs->ctx.device_type, XPU, #name, { \
|
||||
ATEN_ID_TYPE_SWITCH(rhs->dtype, IdType, { \
|
||||
ret = impl::BinaryElewise<XPU, IdType, arith::op>(lhs, rhs); \
|
||||
}); \
|
||||
}); \
|
||||
return ret; \
|
||||
}
|
||||
|
||||
// Generate operators with only lhs being NDArray.
|
||||
#define UNARY_ELEMENT_OP(name, op) \
|
||||
IdArray name(IdArray lhs) { \
|
||||
IdArray ret; \
|
||||
ATEN_XPU_SWITCH_CUDA(lhs->ctx.device_type, XPU, #name, { \
|
||||
ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, { \
|
||||
ret = impl::UnaryElewise<XPU, IdType, arith::op>(lhs); \
|
||||
}); \
|
||||
}); \
|
||||
return ret; \
|
||||
}
|
||||
|
||||
BINARY_ELEMENT_OP(Add, Add)
|
||||
BINARY_ELEMENT_OP(Sub, Sub)
|
||||
BINARY_ELEMENT_OP(Mul, Mul)
|
||||
BINARY_ELEMENT_OP(Div, Div)
|
||||
BINARY_ELEMENT_OP(Mod, Mod)
|
||||
BINARY_ELEMENT_OP(GT, GT)
|
||||
BINARY_ELEMENT_OP(LT, LT)
|
||||
BINARY_ELEMENT_OP(GE, GE)
|
||||
BINARY_ELEMENT_OP(LE, LE)
|
||||
BINARY_ELEMENT_OP(EQ, EQ)
|
||||
BINARY_ELEMENT_OP(NE, NE)
|
||||
|
||||
BINARY_ELEMENT_OP_L(Add, Add)
|
||||
BINARY_ELEMENT_OP_L(Sub, Sub)
|
||||
BINARY_ELEMENT_OP_L(Mul, Mul)
|
||||
BINARY_ELEMENT_OP_L(Div, Div)
|
||||
BINARY_ELEMENT_OP_L(Mod, Mod)
|
||||
BINARY_ELEMENT_OP_L(GT, GT)
|
||||
BINARY_ELEMENT_OP_L(LT, LT)
|
||||
BINARY_ELEMENT_OP_L(GE, GE)
|
||||
BINARY_ELEMENT_OP_L(LE, LE)
|
||||
BINARY_ELEMENT_OP_L(EQ, EQ)
|
||||
BINARY_ELEMENT_OP_L(NE, NE)
|
||||
|
||||
BINARY_ELEMENT_OP_R(Add, Add)
|
||||
BINARY_ELEMENT_OP_R(Sub, Sub)
|
||||
BINARY_ELEMENT_OP_R(Mul, Mul)
|
||||
BINARY_ELEMENT_OP_R(Div, Div)
|
||||
BINARY_ELEMENT_OP_R(Mod, Mod)
|
||||
BINARY_ELEMENT_OP_R(GT, GT)
|
||||
BINARY_ELEMENT_OP_R(LT, LT)
|
||||
BINARY_ELEMENT_OP_R(GE, GE)
|
||||
BINARY_ELEMENT_OP_R(LE, LE)
|
||||
BINARY_ELEMENT_OP_R(EQ, EQ)
|
||||
BINARY_ELEMENT_OP_R(NE, NE)
|
||||
|
||||
UNARY_ELEMENT_OP(Neg, Neg)
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
///////////////// Operator overloading for NDArray /////////////////
|
||||
NDArray operator+(const NDArray& lhs, const NDArray& rhs) {
|
||||
return dgl::aten::Add(lhs, rhs);
|
||||
}
|
||||
NDArray operator-(const NDArray& lhs, const NDArray& rhs) {
|
||||
return dgl::aten::Sub(lhs, rhs);
|
||||
}
|
||||
NDArray operator*(const NDArray& lhs, const NDArray& rhs) {
|
||||
return dgl::aten::Mul(lhs, rhs);
|
||||
}
|
||||
NDArray operator/(const NDArray& lhs, const NDArray& rhs) {
|
||||
return dgl::aten::Div(lhs, rhs);
|
||||
}
|
||||
NDArray operator%(const NDArray& lhs, const NDArray& rhs) {
|
||||
return dgl::aten::Mod(lhs, rhs);
|
||||
}
|
||||
NDArray operator+(const NDArray& lhs, int64_t rhs) {
|
||||
return dgl::aten::Add(lhs, rhs);
|
||||
}
|
||||
NDArray operator-(const NDArray& lhs, int64_t rhs) {
|
||||
return dgl::aten::Sub(lhs, rhs);
|
||||
}
|
||||
NDArray operator*(const NDArray& lhs, int64_t rhs) {
|
||||
return dgl::aten::Mul(lhs, rhs);
|
||||
}
|
||||
NDArray operator/(const NDArray& lhs, int64_t rhs) {
|
||||
return dgl::aten::Div(lhs, rhs);
|
||||
}
|
||||
NDArray operator%(const NDArray& lhs, int64_t rhs) {
|
||||
return dgl::aten::Mod(lhs, rhs);
|
||||
}
|
||||
NDArray operator+(int64_t lhs, const NDArray& rhs) {
|
||||
return dgl::aten::Add(lhs, rhs);
|
||||
}
|
||||
NDArray operator-(int64_t lhs, const NDArray& rhs) {
|
||||
return dgl::aten::Sub(lhs, rhs);
|
||||
}
|
||||
NDArray operator*(int64_t lhs, const NDArray& rhs) {
|
||||
return dgl::aten::Mul(lhs, rhs);
|
||||
}
|
||||
NDArray operator/(int64_t lhs, const NDArray& rhs) {
|
||||
return dgl::aten::Div(lhs, rhs);
|
||||
}
|
||||
NDArray operator%(int64_t lhs, const NDArray& rhs) {
|
||||
return dgl::aten::Mod(lhs, rhs);
|
||||
}
|
||||
NDArray operator-(const NDArray& array) { return dgl::aten::Neg(array); }
|
||||
|
||||
NDArray operator>(const NDArray& lhs, const NDArray& rhs) {
|
||||
return dgl::aten::GT(lhs, rhs);
|
||||
}
|
||||
NDArray operator<(const NDArray& lhs, const NDArray& rhs) {
|
||||
return dgl::aten::LT(lhs, rhs);
|
||||
}
|
||||
NDArray operator>=(const NDArray& lhs, const NDArray& rhs) {
|
||||
return dgl::aten::GE(lhs, rhs);
|
||||
}
|
||||
NDArray operator<=(const NDArray& lhs, const NDArray& rhs) {
|
||||
return dgl::aten::LE(lhs, rhs);
|
||||
}
|
||||
NDArray operator==(const NDArray& lhs, const NDArray& rhs) {
|
||||
return dgl::aten::EQ(lhs, rhs);
|
||||
}
|
||||
NDArray operator!=(const NDArray& lhs, const NDArray& rhs) {
|
||||
return dgl::aten::NE(lhs, rhs);
|
||||
}
|
||||
NDArray operator>(const NDArray& lhs, int64_t rhs) {
|
||||
return dgl::aten::GT(lhs, rhs);
|
||||
}
|
||||
NDArray operator<(const NDArray& lhs, int64_t rhs) {
|
||||
return dgl::aten::LT(lhs, rhs);
|
||||
}
|
||||
NDArray operator>=(const NDArray& lhs, int64_t rhs) {
|
||||
return dgl::aten::GE(lhs, rhs);
|
||||
}
|
||||
NDArray operator<=(const NDArray& lhs, int64_t rhs) {
|
||||
return dgl::aten::LE(lhs, rhs);
|
||||
}
|
||||
NDArray operator==(const NDArray& lhs, int64_t rhs) {
|
||||
return dgl::aten::EQ(lhs, rhs);
|
||||
}
|
||||
NDArray operator!=(const NDArray& lhs, int64_t rhs) {
|
||||
return dgl::aten::NE(lhs, rhs);
|
||||
}
|
||||
NDArray operator>(int64_t lhs, const NDArray& rhs) {
|
||||
return dgl::aten::GT(lhs, rhs);
|
||||
}
|
||||
NDArray operator<(int64_t lhs, const NDArray& rhs) {
|
||||
return dgl::aten::LT(lhs, rhs);
|
||||
}
|
||||
NDArray operator>=(int64_t lhs, const NDArray& rhs) {
|
||||
return dgl::aten::GE(lhs, rhs);
|
||||
}
|
||||
NDArray operator<=(int64_t lhs, const NDArray& rhs) {
|
||||
return dgl::aten::LE(lhs, rhs);
|
||||
}
|
||||
NDArray operator==(int64_t lhs, const NDArray& rhs) {
|
||||
return dgl::aten::EQ(lhs, rhs);
|
||||
}
|
||||
NDArray operator!=(int64_t lhs, const NDArray& rhs) {
|
||||
return dgl::aten::NE(lhs, rhs);
|
||||
}
|
||||
@@ -0,0 +1,354 @@
|
||||
/**
|
||||
* Copyright (c) 2019-2022 by Contributors
|
||||
* @file array/array_op.h
|
||||
* @brief Array operator templates
|
||||
*/
|
||||
#ifndef DGL_ARRAY_ARRAY_OP_H_
|
||||
#define DGL_ARRAY_ARRAY_OP_H_
|
||||
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/graph_traversal.h>
|
||||
|
||||
#include <tuple>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray Full(IdType val, int64_t length, DGLContext ctx);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray Range(IdType low, IdType high, DGLContext ctx);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray AsNumBits(IdArray arr, uint8_t bits);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename Op>
|
||||
IdArray BinaryElewise(IdArray lhs, IdArray rhs);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename Op>
|
||||
IdArray BinaryElewise(IdArray lhs, IdType rhs);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename Op>
|
||||
IdArray BinaryElewise(IdType lhs, IdArray rhs);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename Op>
|
||||
IdArray UnaryElewise(IdArray array);
|
||||
|
||||
template <DGLDeviceType XPU, typename DType, typename IdType>
|
||||
NDArray IndexSelect(NDArray array, IdArray index);
|
||||
|
||||
template <DGLDeviceType XPU, typename DType>
|
||||
DType IndexSelect(NDArray array, int64_t index);
|
||||
|
||||
template <DGLDeviceType XPU, typename DType>
|
||||
IdArray NonZero(BoolArray bool_arr);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray NonZero(NDArray array);
|
||||
|
||||
template <DGLDeviceType XPU, typename DType>
|
||||
std::pair<IdArray, IdArray> Sort(IdArray array, int num_bits);
|
||||
|
||||
template <DGLDeviceType XPU, typename DType, typename IdType>
|
||||
NDArray Scatter(NDArray array, IdArray indices);
|
||||
|
||||
template <DGLDeviceType XPU, typename DType, typename IdType>
|
||||
void Scatter_(IdArray index, NDArray value, NDArray out);
|
||||
|
||||
template <DGLDeviceType XPU, typename DType, typename IdType>
|
||||
NDArray Repeat(NDArray array, IdArray repeats);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray Relabel_(const std::vector<IdArray>& arrays);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
NDArray Concat(const std::vector<IdArray>& arrays);
|
||||
|
||||
template <DGLDeviceType XPU, typename DType>
|
||||
std::tuple<NDArray, IdArray, IdArray> Pack(NDArray array, DType pad_value);
|
||||
|
||||
template <DGLDeviceType XPU, typename DType, typename IdType>
|
||||
std::pair<NDArray, IdArray> ConcatSlices(NDArray array, IdArray lengths);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray CumSum(IdArray array, bool prepend_zero);
|
||||
|
||||
// sparse arrays
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
bool CSRIsNonZero(CSRMatrix csr, int64_t row, int64_t col);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
runtime::NDArray CSRIsNonZero(
|
||||
CSRMatrix csr, runtime::NDArray row, runtime::NDArray col);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
bool CSRHasDuplicate(CSRMatrix csr);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
int64_t CSRGetRowNNZ(CSRMatrix csr, int64_t row);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
runtime::NDArray CSRGetRowNNZ(CSRMatrix csr, runtime::NDArray row);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
runtime::NDArray CSRGetRowColumnIndices(CSRMatrix csr, int64_t row);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
runtime::NDArray CSRGetRowData(CSRMatrix csr, int64_t row);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
bool CSRIsSorted(CSRMatrix csr);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename DType>
|
||||
runtime::NDArray CSRGetData(
|
||||
CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols,
|
||||
bool return_eids, runtime::NDArray weights, DType filler);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename DType>
|
||||
runtime::NDArray CSRGetData(
|
||||
CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols,
|
||||
runtime::NDArray weights, DType filler) {
|
||||
return CSRGetData<XPU, IdType, DType>(
|
||||
csr, rows, cols, false, weights, filler);
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
NDArray CSRGetData(CSRMatrix csr, NDArray rows, NDArray cols) {
|
||||
return CSRGetData<XPU, IdType, IdType>(
|
||||
csr, rows, cols, true, NullArray(rows->dtype), -1);
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::vector<runtime::NDArray> CSRGetDataAndIndices(
|
||||
CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix CSRTranspose(CSRMatrix csr);
|
||||
|
||||
// Convert CSR to COO
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix CSRToCOO(CSRMatrix csr);
|
||||
|
||||
// Convert CSR to COO using data array as order
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix CSRToCOODataAsOrder(CSRMatrix csr);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix CSRSliceRows(CSRMatrix csr, int64_t start, int64_t end);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix CSRSliceRows(CSRMatrix csr, runtime::NDArray rows);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix CSRSliceMatrix(
|
||||
CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
void CSRSort_(CSRMatrix* csr);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename TagType>
|
||||
std::pair<CSRMatrix, NDArray> CSRSortByTag(
|
||||
const CSRMatrix& csr, IdArray tag_array, int64_t num_tags);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix CSRReorder(
|
||||
CSRMatrix csr, runtime::NDArray new_row_ids, runtime::NDArray new_col_ids);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix COOReorder(
|
||||
COOMatrix coo, runtime::NDArray new_row_ids, runtime::NDArray new_col_ids);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix CSRRemove(CSRMatrix csr, IdArray entries);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename FloatType>
|
||||
std::pair<COOMatrix, FloatArray> CSRLaborSampling(
|
||||
CSRMatrix mat, IdArray rows, int64_t num_samples, FloatArray prob,
|
||||
int importance_sampling, IdArray random_seed, float seed2_contribution,
|
||||
IdArray NIDs);
|
||||
|
||||
// FloatType is the type of probability data.
|
||||
template <DGLDeviceType XPU, typename IdType, typename DType>
|
||||
COOMatrix CSRRowWiseSampling(
|
||||
CSRMatrix mat, IdArray rows, int64_t num_samples, NDArray prob_or_mask,
|
||||
bool replace);
|
||||
|
||||
// FloatType is the type of probability data.
|
||||
template <
|
||||
DGLDeviceType XPU, typename IdxType, typename DType, bool map_seed_nodes>
|
||||
std::pair<CSRMatrix, IdArray> CSRRowWiseSamplingFused(
|
||||
CSRMatrix mat, IdArray rows, IdArray seed_mapping,
|
||||
std::vector<IdxType>* new_seed_nodes, int64_t num_samples,
|
||||
NDArray prob_or_mask, bool replace);
|
||||
|
||||
// FloatType is the type of probability data.
|
||||
template <DGLDeviceType XPU, typename IdType, typename DType>
|
||||
COOMatrix CSRRowWisePerEtypeSampling(
|
||||
CSRMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
|
||||
const std::vector<int64_t>& num_samples,
|
||||
const std::vector<NDArray>& prob_or_mask, bool replace,
|
||||
bool rowwise_etype_sorted);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix CSRRowWiseSamplingUniform(
|
||||
CSRMatrix mat, IdArray rows, int64_t num_samples, bool replace);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, bool map_seed_nodes>
|
||||
std::pair<CSRMatrix, IdArray> CSRRowWiseSamplingUniformFused(
|
||||
CSRMatrix mat, IdArray rows, IdArray seed_mapping,
|
||||
std::vector<IdType>* new_seed_nodes, int64_t num_samples, bool replace);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix CSRRowWisePerEtypeSamplingUniform(
|
||||
CSRMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
|
||||
const std::vector<int64_t>& num_samples, bool replace,
|
||||
bool rowwise_etype_sorted);
|
||||
|
||||
// FloatType is the type of weight data.
|
||||
template <DGLDeviceType XPU, typename IdType, typename DType>
|
||||
COOMatrix CSRRowWiseTopk(
|
||||
CSRMatrix mat, IdArray rows, int64_t k, NDArray weight, bool ascending);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename FloatType>
|
||||
COOMatrix CSRRowWiseSamplingBiased(
|
||||
CSRMatrix mat, IdArray rows, int64_t num_samples, NDArray tag_offset,
|
||||
FloatArray bias, bool replace);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::pair<IdArray, IdArray> CSRGlobalUniformNegativeSampling(
|
||||
const CSRMatrix& csr, int64_t num_samples, int num_trials,
|
||||
bool exclude_self_loops, bool replace, double redundancy);
|
||||
|
||||
// Union CSRMatrixes
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix UnionCsr(const std::vector<CSRMatrix>& csrs);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::tuple<CSRMatrix, IdArray, IdArray> CSRToSimple(CSRMatrix csr);
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
bool COOIsNonZero(COOMatrix coo, int64_t row, int64_t col);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
runtime::NDArray COOIsNonZero(
|
||||
COOMatrix coo, runtime::NDArray row, runtime::NDArray col);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
bool COOHasDuplicate(COOMatrix coo);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
int64_t COOGetRowNNZ(COOMatrix coo, int64_t row);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
runtime::NDArray COOGetRowNNZ(COOMatrix coo, runtime::NDArray row);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::pair<runtime::NDArray, runtime::NDArray> COOGetRowDataAndIndices(
|
||||
COOMatrix coo, int64_t row);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::vector<runtime::NDArray> COOGetDataAndIndices(
|
||||
COOMatrix coo, runtime::NDArray rows, runtime::NDArray cols);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
runtime::NDArray COOGetData(
|
||||
COOMatrix mat, runtime::NDArray rows, runtime::NDArray cols);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix COOTranspose(COOMatrix coo);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix COOToCSR(COOMatrix coo);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix COOSliceRows(COOMatrix coo, int64_t start, int64_t end);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix COOSliceRows(COOMatrix coo, runtime::NDArray rows);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix COOSliceMatrix(
|
||||
COOMatrix coo, runtime::NDArray rows, runtime::NDArray cols);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::pair<COOMatrix, IdArray> COOCoalesce(COOMatrix coo);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix DisjointUnionCoo(const std::vector<COOMatrix>& coos);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
void COOSort_(COOMatrix* mat, bool sort_column);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::pair<bool, bool> COOIsSorted(COOMatrix coo);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix COORemove(COOMatrix coo, IdArray entries);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename FloatType>
|
||||
std::pair<COOMatrix, FloatArray> COOLaborSampling(
|
||||
COOMatrix mat, IdArray rows, int64_t num_samples, FloatArray prob,
|
||||
int importance_sampling, IdArray random_seed, float seed2_contribution,
|
||||
IdArray NIDs);
|
||||
|
||||
// FloatType is the type of probability data.
|
||||
template <DGLDeviceType XPU, typename IdType, typename DType>
|
||||
COOMatrix COORowWiseSampling(
|
||||
COOMatrix mat, IdArray rows, int64_t num_samples, NDArray prob_or_mask,
|
||||
bool replace);
|
||||
|
||||
// FloatType is the type of probability data.
|
||||
template <DGLDeviceType XPU, typename IdType, typename DType>
|
||||
COOMatrix COORowWisePerEtypeSampling(
|
||||
COOMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
|
||||
const std::vector<int64_t>& num_samples,
|
||||
const std::vector<NDArray>& prob_or_mask, bool replace);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix COORowWiseSamplingUniform(
|
||||
COOMatrix mat, IdArray rows, int64_t num_samples, bool replace);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix COORowWisePerEtypeSamplingUniform(
|
||||
COOMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
|
||||
const std::vector<int64_t>& num_samples, bool replace);
|
||||
|
||||
// FloatType is the type of weight data.
|
||||
template <DGLDeviceType XPU, typename IdType, typename FloatType>
|
||||
COOMatrix COORowWiseTopk(
|
||||
COOMatrix mat, IdArray rows, int64_t k, FloatArray weight, bool ascending);
|
||||
|
||||
///////////////////////// Graph Traverse routines //////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
Frontiers BFSNodesFrontiers(const CSRMatrix& csr, IdArray source);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
Frontiers BFSEdgesFrontiers(const CSRMatrix& csr, IdArray source);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
Frontiers TopologicalNodesFrontiers(const CSRMatrix& csr);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
Frontiers DGLDFSEdges(const CSRMatrix& csr, IdArray source);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
Frontiers DGLDFSLabeledEdges(
|
||||
const CSRMatrix& csr, IdArray source, const bool has_reverse_edge,
|
||||
const bool has_nontree_edge, const bool return_labels);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix COOLineGraph(const COOMatrix& coo, bool backtracking);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_ARRAY_OP_H_
|
||||
@@ -0,0 +1,61 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file array/check.h
|
||||
* @brief DGL check utilities
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CHECK_H_
|
||||
#define DGL_ARRAY_CHECK_H_
|
||||
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/ndarray.h>
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
|
||||
// Check whether the given arguments have the same context.
|
||||
inline void CheckCtx(
|
||||
const DGLContext& ctx, const std::vector<NDArray>& arrays,
|
||||
const std::vector<std::string>& names) {
|
||||
for (size_t i = 0; i < arrays.size(); ++i) {
|
||||
if (IsNullArray(arrays[i])) continue;
|
||||
CHECK_EQ(ctx, arrays[i]->ctx)
|
||||
<< "Expected device context " << ctx << ". But got " << arrays[i]->ctx
|
||||
<< " for " << names[i] << ".";
|
||||
}
|
||||
}
|
||||
|
||||
// Check whether input tensors are contiguous.
|
||||
inline void CheckContiguous(
|
||||
const std::vector<NDArray>& arrays, const std::vector<std::string>& names) {
|
||||
for (size_t i = 0; i < arrays.size(); ++i) {
|
||||
if (IsNullArray(arrays[i])) continue;
|
||||
CHECK(arrays[i].IsContiguous())
|
||||
<< "Expect " << names[i] << " to be a contiguous tensor";
|
||||
}
|
||||
}
|
||||
|
||||
// Check whether input tensors have valid shape.
|
||||
inline void CheckShape(
|
||||
const std::vector<uint64_t>& gdim, const std::vector<int>& uev_idx,
|
||||
const std::vector<NDArray>& arrays, const std::vector<std::string>& names) {
|
||||
for (size_t i = 0; i < arrays.size(); ++i) {
|
||||
if (IsNullArray(arrays[i])) continue;
|
||||
CHECK_GE(arrays[i]->ndim, 2)
|
||||
<< "Expect " << names[i] << " to have ndim >= 2, "
|
||||
<< "Note that for scalar feature we expand its "
|
||||
<< "dimension with an additional dimension of "
|
||||
<< "length one.";
|
||||
CHECK_EQ(gdim[uev_idx[i]], arrays[i]->shape[0])
|
||||
<< "Expect " << names[i] << " to have size " << gdim[uev_idx[i]]
|
||||
<< " on the first dimension, "
|
||||
<< "but got " << arrays[i]->shape[0];
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CHECK_H_
|
||||
@@ -0,0 +1,41 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/array_cumsum.cc
|
||||
* @brief Array cumsum CPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray CumSum(IdArray array, bool prepend_zero) {
|
||||
const int64_t len = array.NumElements();
|
||||
if (len == 0)
|
||||
return !prepend_zero ? array
|
||||
: aten::Full(0, 1, array->dtype.bits, array->ctx);
|
||||
if (prepend_zero) {
|
||||
IdArray ret = aten::NewIdArray(len + 1, array->ctx, array->dtype.bits);
|
||||
const IdType* in_d = array.Ptr<IdType>();
|
||||
IdType* out_d = ret.Ptr<IdType>();
|
||||
out_d[0] = 0;
|
||||
for (int64_t i = 0; i < len; ++i) out_d[i + 1] = out_d[i] + in_d[i];
|
||||
return ret;
|
||||
} else {
|
||||
IdArray ret = aten::NewIdArray(len, array->ctx, array->dtype.bits);
|
||||
const IdType* in_d = array.Ptr<IdType>();
|
||||
IdType* out_d = ret.Ptr<IdType>();
|
||||
out_d[0] = in_d[0];
|
||||
for (int64_t i = 1; i < len; ++i) out_d[i] = out_d[i - 1] + in_d[i];
|
||||
return ret;
|
||||
}
|
||||
}
|
||||
|
||||
template IdArray CumSum<kDGLCPU, int32_t>(IdArray, bool);
|
||||
template IdArray CumSum<kDGLCPU, int64_t>(IdArray, bool);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,54 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file array/cpu/array_index_select.cc
|
||||
* @brief Array index select CPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename DType, typename IdType>
|
||||
NDArray IndexSelect(NDArray array, IdArray index) {
|
||||
CHECK_EQ(array->shape[0], array.NumElements())
|
||||
<< "Only support tensor"
|
||||
<< " whose first dimension equals number of elements, e.g. (5,), (5, 1)";
|
||||
|
||||
const DType* array_data = static_cast<DType*>(array->data);
|
||||
const IdType* idx_data = static_cast<IdType*>(index->data);
|
||||
const int64_t arr_len = array->shape[0];
|
||||
const int64_t len = index->shape[0];
|
||||
NDArray ret = NDArray::Empty({len}, array->dtype, array->ctx);
|
||||
DType* ret_data = static_cast<DType*>(ret->data);
|
||||
for (int64_t i = 0; i < len; ++i) {
|
||||
CHECK_LT(idx_data[i], arr_len) << "Index out of range.";
|
||||
ret_data[i] = array_data[idx_data[i]];
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
template NDArray IndexSelect<kDGLCPU, int32_t, int32_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelect<kDGLCPU, int32_t, int64_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelect<kDGLCPU, int64_t, int32_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelect<kDGLCPU, int64_t, int64_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelect<kDGLCPU, float, int32_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelect<kDGLCPU, float, int64_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelect<kDGLCPU, double, int32_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelect<kDGLCPU, double, int64_t>(NDArray, IdArray);
|
||||
|
||||
template <DGLDeviceType XPU, typename DType>
|
||||
DType IndexSelect(NDArray array, int64_t index) {
|
||||
const DType* data = static_cast<DType*>(array->data);
|
||||
return data[index];
|
||||
}
|
||||
|
||||
template int32_t IndexSelect<kDGLCPU, int32_t>(NDArray array, int64_t index);
|
||||
template int64_t IndexSelect<kDGLCPU, int64_t>(NDArray array, int64_t index);
|
||||
template float IndexSelect<kDGLCPU, float>(NDArray array, int64_t index);
|
||||
template double IndexSelect<kDGLCPU, double>(NDArray array, int64_t index);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,27 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/array_nonzero.cc
|
||||
* @brief Array nonzero CPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray NonZero(IdArray array) {
|
||||
std::vector<int64_t> ret;
|
||||
const IdType* data = array.Ptr<IdType>();
|
||||
for (int64_t i = 0; i < array->shape[0]; ++i)
|
||||
if (data[i] != 0) ret.push_back(i);
|
||||
return NDArray::FromVector(ret, array->ctx);
|
||||
}
|
||||
|
||||
template IdArray NonZero<kDGLCPU, int32_t>(IdArray);
|
||||
template IdArray NonZero<kDGLCPU, int64_t>(IdArray);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,309 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file array/cpu/array_op_impl.cc
|
||||
* @brief Array operator CPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/ndarray.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "../arith.h"
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
using runtime::parallel_for;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
///////////////////////////// AsNumBits /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray AsNumBits(IdArray arr, uint8_t bits) {
|
||||
CHECK(bits == 32 || bits == 64) << "invalid number of integer bits";
|
||||
if (sizeof(IdType) * 8 == bits) {
|
||||
return arr;
|
||||
}
|
||||
const int64_t len = arr->shape[0];
|
||||
IdArray ret = NewIdArray(len, arr->ctx, bits);
|
||||
const IdType* arr_data = static_cast<IdType*>(arr->data);
|
||||
if (bits == 32) {
|
||||
int32_t* ret_data = static_cast<int32_t*>(ret->data);
|
||||
for (int64_t i = 0; i < len; ++i) {
|
||||
ret_data[i] = arr_data[i];
|
||||
}
|
||||
} else {
|
||||
int64_t* ret_data = static_cast<int64_t*>(ret->data);
|
||||
for (int64_t i = 0; i < len; ++i) {
|
||||
ret_data[i] = arr_data[i];
|
||||
}
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
template IdArray AsNumBits<kDGLCPU, int32_t>(IdArray arr, uint8_t bits);
|
||||
template IdArray AsNumBits<kDGLCPU, int64_t>(IdArray arr, uint8_t bits);
|
||||
|
||||
///////////////////////////// BinaryElewise /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename Op>
|
||||
IdArray BinaryElewise(IdArray lhs, IdArray rhs) {
|
||||
IdArray ret = NewIdArray(lhs->shape[0], lhs->ctx, lhs->dtype.bits);
|
||||
const IdType* lhs_data = static_cast<IdType*>(lhs->data);
|
||||
const IdType* rhs_data = static_cast<IdType*>(rhs->data);
|
||||
IdType* ret_data = static_cast<IdType*>(ret->data);
|
||||
// TODO(BarclayII): this usually incurs lots of overhead in thread spawning,
|
||||
// scheduling, etc., especially since the workload is very light. Need to
|
||||
// replace with parallel_for.
|
||||
for (int64_t i = 0; i < lhs->shape[0]; i++) {
|
||||
ret_data[i] = Op::Call(lhs_data[i], rhs_data[i]);
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::Add>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::Sub>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::Mul>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::Div>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::Mod>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::GT>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::LT>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::GE>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::LE>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::EQ>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::NE>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::Add>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::Sub>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::Mul>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::Div>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::Mod>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::GT>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::LT>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::GE>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::LE>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::EQ>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::NE>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename Op>
|
||||
IdArray BinaryElewise(IdArray lhs, IdType rhs) {
|
||||
IdArray ret = NewIdArray(lhs->shape[0], lhs->ctx, lhs->dtype.bits);
|
||||
const IdType* lhs_data = static_cast<IdType*>(lhs->data);
|
||||
IdType* ret_data = static_cast<IdType*>(ret->data);
|
||||
// TODO(BarclayII): this usually incurs lots of overhead in thread spawning,
|
||||
// scheduling, etc., especially since the workload is very light. Need to
|
||||
// replace with parallel_for.
|
||||
for (int64_t i = 0; i < lhs->shape[0]; i++) {
|
||||
ret_data[i] = Op::Call(lhs_data[i], rhs);
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::Add>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::Sub>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::Mul>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::Div>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::Mod>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::GT>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::LT>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::GE>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::LE>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::EQ>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::NE>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::Add>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::Sub>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::Mul>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::Div>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::Mod>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::GT>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::LT>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::GE>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::LE>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::EQ>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::NE>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename Op>
|
||||
IdArray BinaryElewise(IdType lhs, IdArray rhs) {
|
||||
IdArray ret = NewIdArray(rhs->shape[0], rhs->ctx, rhs->dtype.bits);
|
||||
const IdType* rhs_data = static_cast<IdType*>(rhs->data);
|
||||
IdType* ret_data = static_cast<IdType*>(ret->data);
|
||||
// TODO(BarclayII): this usually incurs lots of overhead in thread spawning,
|
||||
// scheduling, etc., especially since the workload is very light. Need to
|
||||
// replace with parallel_for.
|
||||
for (int64_t i = 0; i < rhs->shape[0]; i++) {
|
||||
ret_data[i] = Op::Call(lhs, rhs_data[i]);
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::Add>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::Sub>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::Mul>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::Div>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::Mod>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::GT>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::LT>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::GE>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::LE>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::EQ>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int32_t, arith::NE>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::Add>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::Sub>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::Mul>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::Div>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::Mod>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::GT>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::LT>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::GE>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::LE>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::EQ>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCPU, int64_t, arith::NE>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename Op>
|
||||
IdArray UnaryElewise(IdArray lhs) {
|
||||
IdArray ret = NewIdArray(lhs->shape[0], lhs->ctx, lhs->dtype.bits);
|
||||
const IdType* lhs_data = static_cast<IdType*>(lhs->data);
|
||||
IdType* ret_data = static_cast<IdType*>(ret->data);
|
||||
// TODO(BarclayII): this usually incurs lots of overhead in thread spawning,
|
||||
// scheduling, etc., especially since the workload is very light. Need to
|
||||
// replace with parallel_for.
|
||||
for (int64_t i = 0; i < lhs->shape[0]; i++) {
|
||||
ret_data[i] = Op::Call(lhs_data[i]);
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
template IdArray UnaryElewise<kDGLCPU, int32_t, arith::Neg>(IdArray lhs);
|
||||
template IdArray UnaryElewise<kDGLCPU, int64_t, arith::Neg>(IdArray lhs);
|
||||
|
||||
///////////////////////////// Full /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename DType>
|
||||
NDArray Full(DType val, int64_t length, DGLContext ctx) {
|
||||
NDArray ret = NDArray::Empty({length}, DGLDataTypeTraits<DType>::dtype, ctx);
|
||||
DType* ret_data = static_cast<DType*>(ret->data);
|
||||
std::fill(ret_data, ret_data + length, val);
|
||||
return ret;
|
||||
}
|
||||
|
||||
template NDArray Full<kDGLCPU, int32_t>(
|
||||
int32_t val, int64_t length, DGLContext ctx);
|
||||
template NDArray Full<kDGLCPU, int64_t>(
|
||||
int64_t val, int64_t length, DGLContext ctx);
|
||||
template NDArray Full<kDGLCPU, float>(
|
||||
float val, int64_t length, DGLContext ctx);
|
||||
template NDArray Full<kDGLCPU, double>(
|
||||
double val, int64_t length, DGLContext ctx);
|
||||
|
||||
///////////////////////////// Range /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray Range(IdType low, IdType high, DGLContext ctx) {
|
||||
CHECK(high >= low) << "high must be bigger than low";
|
||||
IdArray ret = NewIdArray(high - low, ctx, sizeof(IdType) * 8);
|
||||
IdType* ret_data = static_cast<IdType*>(ret->data);
|
||||
std::iota(ret_data, ret_data + high - low, low);
|
||||
return ret;
|
||||
}
|
||||
|
||||
template IdArray Range<kDGLCPU, int32_t>(int32_t, int32_t, DGLContext);
|
||||
template IdArray Range<kDGLCPU, int64_t>(int64_t, int64_t, DGLContext);
|
||||
|
||||
///////////////////////////// Relabel_ /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray Relabel_(const std::vector<IdArray>& arrays) {
|
||||
// build map & relabel
|
||||
IdType newid = 0;
|
||||
std::unordered_map<IdType, IdType> oldv2newv;
|
||||
for (IdArray arr : arrays) {
|
||||
for (int64_t i = 0; i < arr->shape[0]; ++i) {
|
||||
const IdType id = static_cast<IdType*>(arr->data)[i];
|
||||
if (!oldv2newv.count(id)) {
|
||||
oldv2newv[id] = newid++;
|
||||
}
|
||||
static_cast<IdType*>(arr->data)[i] = oldv2newv[id];
|
||||
}
|
||||
}
|
||||
// map array
|
||||
IdArray maparr =
|
||||
NewIdArray(newid, DGLContext{kDGLCPU, 0}, sizeof(IdType) * 8);
|
||||
IdType* maparr_data = static_cast<IdType*>(maparr->data);
|
||||
for (const auto& kv : oldv2newv) {
|
||||
maparr_data[kv.second] = kv.first;
|
||||
}
|
||||
return maparr;
|
||||
}
|
||||
|
||||
template IdArray Relabel_<kDGLCPU, int32_t>(const std::vector<IdArray>& arrays);
|
||||
template IdArray Relabel_<kDGLCPU, int64_t>(const std::vector<IdArray>& arrays);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,97 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file array/cpu/array_index_select.cc
|
||||
* @brief Array index select CPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
|
||||
#include <tuple>
|
||||
#include <utility>
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
using runtime::parallel_for;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename DType, typename IdType>
|
||||
std::pair<NDArray, IdArray> ConcatSlices(NDArray array, IdArray lengths) {
|
||||
const int64_t rows = lengths->shape[0];
|
||||
const int64_t cols = (array->ndim == 1 ? array->shape[0] : array->shape[1]);
|
||||
const int64_t stride = (array->ndim == 1 ? 0 : cols);
|
||||
const DType *array_data = static_cast<DType *>(array->data);
|
||||
const IdType *length_data = static_cast<IdType *>(lengths->data);
|
||||
|
||||
IdArray offsets = NewIdArray(rows, array->ctx, sizeof(IdType) * 8);
|
||||
IdType *offsets_data = static_cast<IdType *>(offsets->data);
|
||||
for (int64_t i = 0; i < rows; ++i)
|
||||
offsets_data[i] = (i == 0 ? 0 : length_data[i - 1] + offsets_data[i - 1]);
|
||||
const int64_t total_length = offsets_data[rows - 1] + length_data[rows - 1];
|
||||
|
||||
NDArray concat = NDArray::Empty({total_length}, array->dtype, array->ctx);
|
||||
DType *concat_data = static_cast<DType *>(concat->data);
|
||||
|
||||
parallel_for(0, rows, [=](size_t b, size_t e) {
|
||||
for (auto i = b; i < e; ++i) {
|
||||
for (int64_t j = 0; j < length_data[i]; ++j)
|
||||
concat_data[offsets_data[i] + j] = array_data[i * stride + j];
|
||||
}
|
||||
});
|
||||
|
||||
return std::make_pair(concat, offsets);
|
||||
}
|
||||
|
||||
template std::pair<NDArray, IdArray> ConcatSlices<kDGLCPU, int32_t, int32_t>(
|
||||
NDArray, IdArray);
|
||||
template std::pair<NDArray, IdArray> ConcatSlices<kDGLCPU, int64_t, int32_t>(
|
||||
NDArray, IdArray);
|
||||
template std::pair<NDArray, IdArray> ConcatSlices<kDGLCPU, float, int32_t>(
|
||||
NDArray, IdArray);
|
||||
template std::pair<NDArray, IdArray> ConcatSlices<kDGLCPU, double, int32_t>(
|
||||
NDArray, IdArray);
|
||||
template std::pair<NDArray, IdArray> ConcatSlices<kDGLCPU, int32_t, int64_t>(
|
||||
NDArray, IdArray);
|
||||
template std::pair<NDArray, IdArray> ConcatSlices<kDGLCPU, int64_t, int64_t>(
|
||||
NDArray, IdArray);
|
||||
template std::pair<NDArray, IdArray> ConcatSlices<kDGLCPU, float, int64_t>(
|
||||
NDArray, IdArray);
|
||||
template std::pair<NDArray, IdArray> ConcatSlices<kDGLCPU, double, int64_t>(
|
||||
NDArray, IdArray);
|
||||
|
||||
template <DGLDeviceType XPU, typename DType>
|
||||
std::tuple<NDArray, IdArray, IdArray> Pack(NDArray array, DType pad_value) {
|
||||
CHECK_NDIM(array, 2, "array");
|
||||
const DType *array_data = static_cast<DType *>(array->data);
|
||||
const int64_t rows = array->shape[0];
|
||||
const int64_t cols = array->shape[1];
|
||||
|
||||
IdArray length = NewIdArray(rows, array->ctx);
|
||||
int64_t *length_data = static_cast<int64_t *>(length->data);
|
||||
parallel_for(0, rows, [=](size_t b, size_t e) {
|
||||
for (auto i = b; i < e; ++i) {
|
||||
int64_t j;
|
||||
for (j = 0; j < cols; ++j) {
|
||||
const DType val = array_data[i * cols + j];
|
||||
if (val == pad_value) break;
|
||||
}
|
||||
length_data[i] = j;
|
||||
}
|
||||
});
|
||||
|
||||
auto ret = ConcatSlices<XPU, DType, int64_t>(array, length);
|
||||
return std::make_tuple(ret.first, length, ret.second);
|
||||
}
|
||||
|
||||
template std::tuple<NDArray, IdArray, IdArray> Pack<kDGLCPU, int32_t>(
|
||||
NDArray, int32_t);
|
||||
template std::tuple<NDArray, IdArray, IdArray> Pack<kDGLCPU, int64_t>(
|
||||
NDArray, int64_t);
|
||||
template std::tuple<NDArray, IdArray, IdArray> Pack<kDGLCPU, float>(
|
||||
NDArray, float);
|
||||
template std::tuple<NDArray, IdArray, IdArray> Pack<kDGLCPU, double>(
|
||||
NDArray, double);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,51 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/array_repeat.cc
|
||||
* @brief Array repeat CPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename DType, typename IdType>
|
||||
NDArray Repeat(NDArray array, IdArray repeats) {
|
||||
CHECK(array->shape[0] == repeats->shape[0])
|
||||
<< "shape of array and repeats mismatch";
|
||||
|
||||
const int64_t len = array->shape[0];
|
||||
const DType *array_data = static_cast<DType *>(array->data);
|
||||
const IdType *repeats_data = static_cast<IdType *>(repeats->data);
|
||||
|
||||
IdType num_elements = 0;
|
||||
for (int64_t i = 0; i < len; ++i) num_elements += repeats_data[i];
|
||||
|
||||
NDArray result = NDArray::Empty({num_elements}, array->dtype, array->ctx);
|
||||
DType *result_data = static_cast<DType *>(result->data);
|
||||
IdType curr = 0;
|
||||
for (int64_t i = 0; i < len; ++i) {
|
||||
std::fill(
|
||||
result_data + curr, result_data + curr + repeats_data[i],
|
||||
array_data[i]);
|
||||
curr += repeats_data[i];
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
template NDArray Repeat<kDGLCPU, int32_t, int32_t>(NDArray, IdArray);
|
||||
template NDArray Repeat<kDGLCPU, int64_t, int32_t>(NDArray, IdArray);
|
||||
template NDArray Repeat<kDGLCPU, float, int32_t>(NDArray, IdArray);
|
||||
template NDArray Repeat<kDGLCPU, double, int32_t>(NDArray, IdArray);
|
||||
template NDArray Repeat<kDGLCPU, int32_t, int64_t>(NDArray, IdArray);
|
||||
template NDArray Repeat<kDGLCPU, int64_t, int64_t>(NDArray, IdArray);
|
||||
template NDArray Repeat<kDGLCPU, float, int64_t>(NDArray, IdArray);
|
||||
template NDArray Repeat<kDGLCPU, double, int64_t>(NDArray, IdArray);
|
||||
|
||||
}; // namespace impl
|
||||
}; // namespace aten
|
||||
}; // namespace dgl
|
||||
@@ -0,0 +1,62 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file array/cpu/array_scatter.cc
|
||||
* @brief Array scatter CPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename DType, typename IdType>
|
||||
NDArray Scatter(NDArray array, IdArray indices) {
|
||||
NDArray result =
|
||||
NDArray::Empty({indices->shape[0]}, array->dtype, array->ctx);
|
||||
|
||||
const DType *array_data = static_cast<DType *>(array->data);
|
||||
const IdType *indices_data = static_cast<IdType *>(indices->data);
|
||||
DType *result_data = static_cast<DType *>(result->data);
|
||||
|
||||
for (int64_t i = 0; i < indices->shape[0]; ++i)
|
||||
result_data[indices_data[i]] = array_data[i];
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
template NDArray Scatter<kDGLCPU, int32_t, int32_t>(NDArray, IdArray);
|
||||
template NDArray Scatter<kDGLCPU, int64_t, int32_t>(NDArray, IdArray);
|
||||
template NDArray Scatter<kDGLCPU, float, int32_t>(NDArray, IdArray);
|
||||
template NDArray Scatter<kDGLCPU, double, int32_t>(NDArray, IdArray);
|
||||
template NDArray Scatter<kDGLCPU, int32_t, int64_t>(NDArray, IdArray);
|
||||
template NDArray Scatter<kDGLCPU, int64_t, int64_t>(NDArray, IdArray);
|
||||
template NDArray Scatter<kDGLCPU, float, int64_t>(NDArray, IdArray);
|
||||
template NDArray Scatter<kDGLCPU, double, int64_t>(NDArray, IdArray);
|
||||
|
||||
template <DGLDeviceType XPU, typename DType, typename IdType>
|
||||
void Scatter_(IdArray index, NDArray value, NDArray out) {
|
||||
const int64_t len = index->shape[0];
|
||||
const IdType *idx = index.Ptr<IdType>();
|
||||
const DType *val = value.Ptr<DType>();
|
||||
DType *outd = out.Ptr<DType>();
|
||||
runtime::parallel_for(0, len, [&](size_t b, size_t e) {
|
||||
for (auto i = b; i < e; ++i) {
|
||||
outd[idx[i]] = val[i];
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template void Scatter_<kDGLCPU, int32_t, int32_t>(IdArray, NDArray, NDArray);
|
||||
template void Scatter_<kDGLCPU, int64_t, int32_t>(IdArray, NDArray, NDArray);
|
||||
template void Scatter_<kDGLCPU, float, int32_t>(IdArray, NDArray, NDArray);
|
||||
template void Scatter_<kDGLCPU, double, int32_t>(IdArray, NDArray, NDArray);
|
||||
template void Scatter_<kDGLCPU, int32_t, int64_t>(IdArray, NDArray, NDArray);
|
||||
template void Scatter_<kDGLCPU, int64_t, int64_t>(IdArray, NDArray, NDArray);
|
||||
template void Scatter_<kDGLCPU, float, int64_t>(IdArray, NDArray, NDArray);
|
||||
template void Scatter_<kDGLCPU, double, int64_t>(IdArray, NDArray, NDArray);
|
||||
|
||||
}; // namespace impl
|
||||
}; // namespace aten
|
||||
}; // namespace dgl
|
||||
@@ -0,0 +1,167 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/array_sort.cc
|
||||
* @brief Array sort CPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
#ifdef PARALLEL_ALGORITHMS
|
||||
#include <parallel/algorithm>
|
||||
#endif
|
||||
#include <algorithm>
|
||||
#include <iterator>
|
||||
|
||||
namespace {
|
||||
|
||||
template <typename V1, typename V2>
|
||||
struct PairRef {
|
||||
PairRef() = delete;
|
||||
PairRef(const PairRef& other) = default;
|
||||
PairRef(PairRef&& other) = default;
|
||||
PairRef(V1* const r, V2* const c) : row(r), col(c) {}
|
||||
|
||||
PairRef& operator=(const PairRef& other) {
|
||||
*row = *other.row;
|
||||
*col = *other.col;
|
||||
return *this;
|
||||
}
|
||||
PairRef& operator=(const std::pair<V1, V2>& val) {
|
||||
*row = std::get<0>(val);
|
||||
*col = std::get<1>(val);
|
||||
return *this;
|
||||
}
|
||||
|
||||
operator std::pair<V1, V2>() const { return std::make_pair(*row, *col); }
|
||||
|
||||
void Swap(const PairRef& other) const {
|
||||
std::swap(*row, *other.row);
|
||||
std::swap(*col, *other.col);
|
||||
}
|
||||
|
||||
V1* row;
|
||||
V2* col;
|
||||
};
|
||||
|
||||
using std::swap;
|
||||
template <typename V1, typename V2>
|
||||
void swap(const PairRef<V1, V2>& r1, const PairRef<V1, V2>& r2) {
|
||||
r1.Swap(r2);
|
||||
}
|
||||
|
||||
template <typename V1, typename V2>
|
||||
struct PairIterator
|
||||
: public std::iterator<
|
||||
std::random_access_iterator_tag, std::pair<V1, V2>, std::ptrdiff_t,
|
||||
std::pair<V1*, V2*>, PairRef<V1, V2>> {
|
||||
PairIterator() = default;
|
||||
PairIterator(const PairIterator& other) = default;
|
||||
PairIterator(PairIterator&& other) = default;
|
||||
PairIterator(V1* r, V2* c) : row(r), col(c) {}
|
||||
|
||||
PairIterator& operator=(const PairIterator& other) = default;
|
||||
PairIterator& operator=(PairIterator&& other) = default;
|
||||
~PairIterator() = default;
|
||||
|
||||
bool operator==(const PairIterator& other) const { return row == other.row; }
|
||||
|
||||
bool operator!=(const PairIterator& other) const { return row != other.row; }
|
||||
|
||||
bool operator<(const PairIterator& other) const { return row < other.row; }
|
||||
|
||||
bool operator>(const PairIterator& other) const { return row > other.row; }
|
||||
|
||||
bool operator<=(const PairIterator& other) const { return row <= other.row; }
|
||||
|
||||
bool operator>=(const PairIterator& other) const { return row >= other.row; }
|
||||
|
||||
PairIterator& operator+=(const std::ptrdiff_t& movement) {
|
||||
row += movement;
|
||||
col += movement;
|
||||
return *this;
|
||||
}
|
||||
|
||||
PairIterator& operator-=(const std::ptrdiff_t& movement) {
|
||||
row -= movement;
|
||||
col -= movement;
|
||||
return *this;
|
||||
}
|
||||
|
||||
PairIterator& operator++() { return operator+=(1); }
|
||||
|
||||
PairIterator& operator--() { return operator-=(1); }
|
||||
|
||||
PairIterator operator++(int) {
|
||||
PairIterator ret(*this);
|
||||
operator++();
|
||||
return ret;
|
||||
}
|
||||
|
||||
PairIterator operator--(int) {
|
||||
PairIterator ret(*this);
|
||||
operator--();
|
||||
return ret;
|
||||
}
|
||||
|
||||
PairIterator operator+(const std::ptrdiff_t& movement) const {
|
||||
PairIterator ret(*this);
|
||||
ret += movement;
|
||||
return ret;
|
||||
}
|
||||
|
||||
PairIterator operator-(const std::ptrdiff_t& movement) const {
|
||||
PairIterator ret(*this);
|
||||
ret -= movement;
|
||||
return ret;
|
||||
}
|
||||
|
||||
std::ptrdiff_t operator-(const PairIterator& other) const {
|
||||
return row - other.row;
|
||||
}
|
||||
|
||||
PairRef<V1, V2> operator*() const { return PairRef<V1, V2>(row, col); }
|
||||
PairRef<V1, V2> operator*() { return PairRef<V1, V2>(row, col); }
|
||||
|
||||
// required for random access iterators in VS2019
|
||||
PairRef<V1, V2> operator[](size_t offset) const {
|
||||
return PairRef<V1, V2>(row + offset, col + offset);
|
||||
}
|
||||
|
||||
V1* row;
|
||||
V2* col;
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::pair<IdArray, IdArray> Sort(IdArray array, int /* num_bits */) {
|
||||
const int64_t nitem = array->shape[0];
|
||||
IdArray val = array.Clone();
|
||||
IdArray idx = aten::Range(0, nitem, 64, array->ctx);
|
||||
IdType* val_data = val.Ptr<IdType>();
|
||||
int64_t* idx_data = idx.Ptr<int64_t>();
|
||||
typedef std::pair<IdType, int64_t> Pair;
|
||||
#ifdef PARALLEL_ALGORITHMS
|
||||
__gnu_parallel::sort(
|
||||
#else
|
||||
std::sort(
|
||||
#endif
|
||||
PairIterator<IdType, int64_t>(val_data, idx_data),
|
||||
PairIterator<IdType, int64_t>(val_data, idx_data) + nitem,
|
||||
[](const Pair& a, const Pair& b) {
|
||||
return std::get<0>(a) < std::get<0>(b);
|
||||
});
|
||||
return std::make_pair(val, idx);
|
||||
}
|
||||
|
||||
template std::pair<IdArray, IdArray> Sort<kDGLCPU, int32_t>(
|
||||
IdArray, int num_bits);
|
||||
template std::pair<IdArray, IdArray> Sort<kDGLCPU, int64_t>(
|
||||
IdArray, int num_bits);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,124 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file dgl/array_utils.h
|
||||
* @brief Utility classes and functions for DGL arrays.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CPU_ARRAY_UTILS_H_
|
||||
#define DGL_ARRAY_CPU_ARRAY_UTILS_H_
|
||||
|
||||
#include <dgl/aten/types.h>
|
||||
#include <tsl/robin_map.h>
|
||||
|
||||
#include <unordered_map>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "../../c_api_common.h"
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
|
||||
/**
|
||||
* @brief A hashmap that maps each ids in the given array to new ids starting
|
||||
* from zero.
|
||||
*
|
||||
* Useful for relabeling integers and finding unique integers.
|
||||
*
|
||||
* Usually faster than std::unordered_map in existence checking.
|
||||
*/
|
||||
template <typename IdType>
|
||||
class IdHashMap {
|
||||
public:
|
||||
// default ctor
|
||||
IdHashMap() : filter_(kFilterSize, false) {}
|
||||
|
||||
// Construct the hashmap using the given id array.
|
||||
// The id array could contain duplicates.
|
||||
// If the id array has no duplicates, the array will be relabeled to
|
||||
// consecutive integers starting from 0.
|
||||
explicit IdHashMap(IdArray ids) : filter_(kFilterSize, false) {
|
||||
oldv2newv_.reserve(ids->shape[0]);
|
||||
Update(ids);
|
||||
}
|
||||
|
||||
// copy ctor
|
||||
IdHashMap(const IdHashMap& other) = default;
|
||||
|
||||
void Reserve(const int64_t size) { oldv2newv_.reserve(size); }
|
||||
|
||||
// Update the hashmap with given id array.
|
||||
// The id array could contain duplicates.
|
||||
void Update(IdArray ids) {
|
||||
const IdType* ids_data = static_cast<IdType*>(ids->data);
|
||||
const int64_t len = ids->shape[0];
|
||||
for (int64_t i = 0; i < len; ++i) {
|
||||
const IdType id = ids_data[i];
|
||||
// Insertion will not happen if the key already exists.
|
||||
oldv2newv_.insert({id, oldv2newv_.size()});
|
||||
filter_[id & kFilterMask] = true;
|
||||
}
|
||||
}
|
||||
|
||||
// Return true if the given id is contained in this hashmap.
|
||||
bool Contains(IdType id) const {
|
||||
return filter_[id & kFilterMask] && oldv2newv_.count(id);
|
||||
}
|
||||
|
||||
// Return the new id of the given id. If the given id is not contained
|
||||
// in the hash map, returns the default_val instead.
|
||||
IdType Map(IdType id, IdType default_val) const {
|
||||
if (filter_[id & kFilterMask]) {
|
||||
auto it = oldv2newv_.find(id);
|
||||
return (it == oldv2newv_.end()) ? default_val : it->second;
|
||||
} else {
|
||||
return default_val;
|
||||
}
|
||||
}
|
||||
|
||||
// Return the new id of each id in the given array.
|
||||
IdArray Map(IdArray ids, IdType default_val) const {
|
||||
const IdType* ids_data = static_cast<IdType*>(ids->data);
|
||||
const int64_t len = ids->shape[0];
|
||||
IdArray values = NewIdArray(len, ids->ctx, ids->dtype.bits);
|
||||
IdType* values_data = static_cast<IdType*>(values->data);
|
||||
for (int64_t i = 0; i < len; ++i)
|
||||
values_data[i] = Map(ids_data[i], default_val);
|
||||
return values;
|
||||
}
|
||||
|
||||
// Return all the old ids collected so far, ordered by new id.
|
||||
IdArray Values() const {
|
||||
IdArray values = NewIdArray(
|
||||
oldv2newv_.size(), DGLContext{kDGLCPU, 0}, sizeof(IdType) * 8);
|
||||
IdType* values_data = static_cast<IdType*>(values->data);
|
||||
for (auto pair : oldv2newv_) values_data[pair.second] = pair.first;
|
||||
return values;
|
||||
}
|
||||
|
||||
inline size_t Size() const { return oldv2newv_.size(); }
|
||||
|
||||
private:
|
||||
static constexpr int32_t kFilterMask = 0xFFFFFF;
|
||||
static constexpr int32_t kFilterSize = kFilterMask + 1;
|
||||
// This bitmap is used as a bloom filter to remove some lookups.
|
||||
// Hashtable is very slow. Using bloom filter can significantly speed up
|
||||
// lookups.
|
||||
std::vector<bool> filter_;
|
||||
// The hashmap from old vid to new vid
|
||||
tsl::robin_map<IdType, IdType> oldv2newv_;
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Hash type for building maps/sets with pairs as keys.
|
||||
*/
|
||||
struct PairHash {
|
||||
template <class T1, class T2>
|
||||
std::size_t operator()(const std::pair<T1, T2>& pair) const {
|
||||
return std::hash<T1>()(pair.first) ^ std::hash<T2>()(pair.second);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CPU_ARRAY_UTILS_H_
|
||||
@@ -0,0 +1,249 @@
|
||||
/**
|
||||
* Copyright (c) 2023 by Contributors
|
||||
* @file array/cpu/concurrent_id_hash_map.cc
|
||||
* @brief Class about id hash map
|
||||
*/
|
||||
|
||||
#include "concurrent_id_hash_map.h"
|
||||
|
||||
#ifdef _MSC_VER
|
||||
#include <intrin.h>
|
||||
#endif // _MSC_VER
|
||||
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/device_api.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
|
||||
#include <cmath>
|
||||
#include <numeric>
|
||||
|
||||
using namespace dgl::runtime;
|
||||
|
||||
namespace {
|
||||
static constexpr int64_t kEmptyKey = -1;
|
||||
static constexpr int kGrainSize = 256;
|
||||
|
||||
// The formula is established from experience which is used
|
||||
// to get the hashmap size from the input array size.
|
||||
inline size_t GetMapSize(size_t num) {
|
||||
size_t capacity = 1;
|
||||
return capacity << static_cast<size_t>(1 + std::log2(num * 3));
|
||||
}
|
||||
} // namespace
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
|
||||
template <typename IdType>
|
||||
IdType ConcurrentIdHashMap<IdType>::CompareAndSwap(
|
||||
IdType* ptr, IdType old_val, IdType new_val) {
|
||||
#ifdef _MSC_VER
|
||||
if (sizeof(IdType) == 4) {
|
||||
return _InterlockedCompareExchange(
|
||||
reinterpret_cast<LONG*>(ptr), new_val, old_val);
|
||||
} else if (sizeof(IdType) == 8) {
|
||||
return _InterlockedCompareExchange64(
|
||||
reinterpret_cast<LONGLONG*>(ptr), new_val, old_val);
|
||||
} else {
|
||||
LOG(FATAL) << "ID can only be int32 or int64";
|
||||
}
|
||||
#elif __GNUC__ // _MSC_VER
|
||||
return __sync_val_compare_and_swap(ptr, old_val, new_val);
|
||||
#else // _MSC_VER
|
||||
#error "CompareAndSwap is not supported on this platform."
|
||||
#endif // _MSC_VER
|
||||
}
|
||||
|
||||
template <typename IdType>
|
||||
ConcurrentIdHashMap<IdType>::ConcurrentIdHashMap() : mask_(0) {
|
||||
// Used to deallocate the memory in hash_map_ with device api
|
||||
// when the pointer is freed.
|
||||
auto deleter = [](Mapping* mappings) {
|
||||
if (mappings != nullptr) {
|
||||
DGLContext ctx = DGLContext{kDGLCPU, 0};
|
||||
auto device = DeviceAPI::Get(ctx);
|
||||
device->FreeWorkspace(ctx, mappings);
|
||||
}
|
||||
};
|
||||
hash_map_ = {nullptr, deleter};
|
||||
}
|
||||
|
||||
template <typename IdType>
|
||||
IdArray ConcurrentIdHashMap<IdType>::Init(
|
||||
const IdArray& ids, size_t num_seeds) {
|
||||
CHECK_EQ(ids.defined(), true);
|
||||
const IdType* ids_data = ids.Ptr<IdType>();
|
||||
const size_t num_ids = static_cast<size_t>(ids->shape[0]);
|
||||
// Make sure `ids` is not 0 dim.
|
||||
CHECK_GE(num_seeds, 0);
|
||||
CHECK_GE(num_ids, num_seeds);
|
||||
size_t capacity = GetMapSize(num_ids);
|
||||
mask_ = static_cast<IdType>(capacity - 1);
|
||||
|
||||
auto ctx = DGLContext{kDGLCPU, 0};
|
||||
auto device = DeviceAPI::Get(ctx);
|
||||
hash_map_.reset(static_cast<Mapping*>(
|
||||
device->AllocWorkspace(ctx, sizeof(Mapping) * capacity)));
|
||||
memset(hash_map_.get(), -1, sizeof(Mapping) * capacity);
|
||||
|
||||
// This code block is to fill the ids into hash_map_.
|
||||
IdArray unique_ids = NewIdArray(num_ids, ctx, sizeof(IdType) * 8);
|
||||
IdType* unique_ids_data = unique_ids.Ptr<IdType>();
|
||||
// Fill in the first `num_seeds` ids.
|
||||
parallel_for(0, num_seeds, kGrainSize, [&](int64_t s, int64_t e) {
|
||||
for (int64_t i = s; i < e; i++) {
|
||||
InsertAndSet(ids_data[i], static_cast<IdType>(i));
|
||||
}
|
||||
});
|
||||
// Place the first `num_seeds` ids.
|
||||
device->CopyDataFromTo(
|
||||
ids_data, 0, unique_ids_data, 0, sizeof(IdType) * num_seeds, ctx, ctx,
|
||||
ids->dtype);
|
||||
|
||||
// An auxiliary array indicates whether the corresponding elements
|
||||
// are inserted into hash map or not. Use `int16_t` instead of `bool` as
|
||||
// vector<bool> is unsafe when updating different elements from different
|
||||
// threads. See https://en.cppreference.com/w/cpp/container#Thread_safety.
|
||||
std::vector<int16_t> valid(num_ids);
|
||||
auto thread_num = compute_num_threads(0, num_ids, kGrainSize);
|
||||
std::vector<size_t> block_offset(thread_num + 1, 0);
|
||||
// Insert all elements in this loop.
|
||||
parallel_for(num_seeds, num_ids, kGrainSize, [&](int64_t s, int64_t e) {
|
||||
size_t count = 0;
|
||||
for (int64_t i = s; i < e; i++) {
|
||||
valid[i] = Insert(ids_data[i]);
|
||||
count += valid[i];
|
||||
}
|
||||
block_offset[omp_get_thread_num() + 1] = count;
|
||||
});
|
||||
|
||||
// Get ExclusiveSum of each block.
|
||||
std::partial_sum(
|
||||
block_offset.begin() + 1, block_offset.end(), block_offset.begin() + 1);
|
||||
unique_ids->shape[0] = num_seeds + block_offset.back();
|
||||
|
||||
// Get unique array from ids and set value for hash map.
|
||||
parallel_for(num_seeds, num_ids, kGrainSize, [&](int64_t s, int64_t e) {
|
||||
auto tid = omp_get_thread_num();
|
||||
auto pos = block_offset[tid] + num_seeds;
|
||||
for (int64_t i = s; i < e; i++) {
|
||||
if (valid[i]) {
|
||||
unique_ids_data[pos] = ids_data[i];
|
||||
Set(ids_data[i], pos);
|
||||
pos = pos + 1;
|
||||
}
|
||||
}
|
||||
});
|
||||
return unique_ids;
|
||||
}
|
||||
|
||||
template <typename IdType>
|
||||
IdArray ConcurrentIdHashMap<IdType>::MapIds(const IdArray& ids) const {
|
||||
CHECK_EQ(ids.defined(), true);
|
||||
const IdType* ids_data = ids.Ptr<IdType>();
|
||||
const size_t num_ids = static_cast<size_t>(ids->shape[0]);
|
||||
CHECK_GT(num_ids, 0);
|
||||
|
||||
DGLContext ctx = DGLContext{kDGLCPU, 0};
|
||||
IdArray new_ids = NewIdArray(num_ids, ctx, sizeof(IdType) * 8);
|
||||
IdType* values_data = new_ids.Ptr<IdType>();
|
||||
|
||||
parallel_for(0, num_ids, kGrainSize, [&](int64_t s, int64_t e) {
|
||||
for (int64_t i = s; i < e; i++) {
|
||||
values_data[i] = MapId(ids_data[i]);
|
||||
}
|
||||
});
|
||||
return new_ids;
|
||||
}
|
||||
|
||||
template <typename IdType>
|
||||
inline void ConcurrentIdHashMap<IdType>::Next(
|
||||
IdType* pos, IdType* delta) const {
|
||||
// Use Quadric probing.
|
||||
*pos = (*pos + (*delta) * (*delta)) & mask_;
|
||||
*delta = *delta + 1;
|
||||
}
|
||||
|
||||
template <typename IdType>
|
||||
inline IdType ConcurrentIdHashMap<IdType>::MapId(IdType id) const {
|
||||
IdType pos = (id & mask_), delta = 1;
|
||||
IdType empty_key = static_cast<IdType>(kEmptyKey);
|
||||
while (hash_map_[pos].key != empty_key && hash_map_[pos].key != id) {
|
||||
Next(&pos, &delta);
|
||||
}
|
||||
return hash_map_[pos].value;
|
||||
}
|
||||
|
||||
template <typename IdType>
|
||||
bool ConcurrentIdHashMap<IdType>::Insert(IdType id) {
|
||||
IdType pos = (id & mask_), delta = 1;
|
||||
InsertState state = AttemptInsertAt(pos, id);
|
||||
while (state == InsertState::OCCUPIED) {
|
||||
Next(&pos, &delta);
|
||||
state = AttemptInsertAt(pos, id);
|
||||
}
|
||||
|
||||
return state == InsertState::INSERTED;
|
||||
}
|
||||
|
||||
template <typename IdType>
|
||||
inline void ConcurrentIdHashMap<IdType>::Set(IdType key, IdType value) {
|
||||
IdType pos = (key & mask_), delta = 1;
|
||||
while (hash_map_[pos].key != key) {
|
||||
Next(&pos, &delta);
|
||||
}
|
||||
|
||||
hash_map_[pos].value = value;
|
||||
}
|
||||
|
||||
template <typename IdType>
|
||||
inline void ConcurrentIdHashMap<IdType>::InsertAndSet(IdType id, IdType value) {
|
||||
IdType pos = (id & mask_), delta = 1;
|
||||
while (AttemptInsertAt(pos, id) == InsertState::OCCUPIED) {
|
||||
Next(&pos, &delta);
|
||||
}
|
||||
|
||||
hash_map_[pos].value = value;
|
||||
}
|
||||
|
||||
template <typename IdType>
|
||||
inline typename ConcurrentIdHashMap<IdType>::InsertState
|
||||
ConcurrentIdHashMap<IdType>::AttemptInsertAt(int64_t pos, IdType key) {
|
||||
IdType empty_key = static_cast<IdType>(kEmptyKey);
|
||||
IdType old_val = CompareAndSwap(&(hash_map_[pos].key), empty_key, key);
|
||||
if (old_val == empty_key) {
|
||||
return InsertState::INSERTED;
|
||||
} else if (old_val == key) {
|
||||
return InsertState::EXISTED;
|
||||
} else {
|
||||
return InsertState::OCCUPIED;
|
||||
}
|
||||
}
|
||||
|
||||
template class ConcurrentIdHashMap<int32_t>;
|
||||
template class ConcurrentIdHashMap<int64_t>;
|
||||
|
||||
template <typename IdType>
|
||||
bool BoolCompareAndSwap(IdType* ptr) {
|
||||
#ifdef _MSC_VER
|
||||
if (sizeof(IdType) == 4) {
|
||||
return _InterlockedCompareExchange(reinterpret_cast<LONG*>(ptr), 0, -1) ==
|
||||
-1;
|
||||
} else if (sizeof(IdType) == 8) {
|
||||
return _InterlockedCompareExchange64(
|
||||
reinterpret_cast<LONGLONG*>(ptr), 0, -1) == -1;
|
||||
} else {
|
||||
LOG(FATAL) << "ID can only be int32 or int64";
|
||||
}
|
||||
#elif __GNUC__ // _MSC_VER
|
||||
return __sync_bool_compare_and_swap(ptr, -1, 0);
|
||||
#else // _MSC_VER
|
||||
#error "CompareAndSwap is not supported on this platform."
|
||||
#endif // _MSC_VER
|
||||
}
|
||||
|
||||
template bool BoolCompareAndSwap<int32_t>(int32_t*);
|
||||
template bool BoolCompareAndSwap<int64_t>(int64_t*);
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,204 @@
|
||||
/**
|
||||
* Copyright (c) 2023 by Contributors
|
||||
* @file array/cpu/concurrent_id_hash_map.h
|
||||
* @brief Class about concurrent id hash map
|
||||
*/
|
||||
|
||||
#ifndef DGL_ARRAY_CPU_CONCURRENT_ID_HASH_MAP_H_
|
||||
#define DGL_ARRAY_CPU_CONCURRENT_ID_HASH_MAP_H_
|
||||
|
||||
#include <dgl/aten/types.h>
|
||||
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
|
||||
/**
|
||||
* @brief A CPU targeted hashmap for mapping duplicate and non-consecutive ids
|
||||
* in the provided array to unique and consecutive ones. It utilizes
|
||||
* multi-threading to accelerate the insert and search speed. Currently it is
|
||||
* only designed to be used in `ToBlockCpu` for optimizing, so it only support
|
||||
* key insertions once with Init function, and it does not support key deletion.
|
||||
*
|
||||
* The hash map should be prepared in two phases before using. With the first
|
||||
* being creating the hashmap, and then initialize it with an id array which is
|
||||
* divided into 2 parts: [`seed ids`, `sampled ids`]. `Seed ids` refer to
|
||||
* a set ids chosen as the input for sampling process and `sampled ids` are the
|
||||
* ids new sampled from the process (note the the `seed ids` might also be
|
||||
* sampled in the process and included in the `sampled ids`). In result `seed
|
||||
* ids` are mapped to [0, num_seed_ids) and `sampled ids` to [num_seed_ids,
|
||||
* num_unique_ids). Notice that mapping order is stable for `seed ids` while not
|
||||
* for the `sampled ids`.
|
||||
*
|
||||
* For example, for an array `A` having 4 seed ids with following entries:
|
||||
* [99, 98, 100, 97, 97, 101, 101, 102, 101]
|
||||
* Create the hashmap `H` with:
|
||||
* `H = ConcurrentIdHashMap()` (1)
|
||||
* And Init it with:
|
||||
* `U = H.Init(A)` (2) (U is an id array used to store the unqiue
|
||||
* ids in A).
|
||||
* Then `U` should be (U is not exclusive as the overall mapping is not stable):
|
||||
* [99, 98, 100, 97, 102, 101]
|
||||
* And the hashmap should generate following mappings:
|
||||
* * [
|
||||
* {key: 99, value: 0},
|
||||
* {key: 98, value: 1},
|
||||
* {key: 100, value: 2},
|
||||
* {key: 97, value: 3},
|
||||
* {key: 102, value: 4},
|
||||
* {key: 101, value: 5}
|
||||
* ]
|
||||
* Search the hashmap with array `I`=[98, 99, 102]:
|
||||
* R = H.Map(I) (3)
|
||||
* R should be:
|
||||
* [1, 0, 4]
|
||||
**/
|
||||
template <typename IdType>
|
||||
class ConcurrentIdHashMap {
|
||||
private:
|
||||
/**
|
||||
* @brief The result state of an attempt to insert.
|
||||
*/
|
||||
enum class InsertState {
|
||||
OCCUPIED, // Indicates that the space where an insertion is being
|
||||
// attempted is already occupied by another element.
|
||||
EXISTED, // Indicates that the element being inserted already exists in the
|
||||
// map, and thus no insertion is performed.
|
||||
INSERTED // Indicates that the insertion was successful and a new element
|
||||
// was added to the map.
|
||||
};
|
||||
|
||||
public:
|
||||
/**
|
||||
* @brief An entry in the hashtable.
|
||||
*/
|
||||
struct Mapping {
|
||||
/**
|
||||
* @brief The ID of the item inserted.
|
||||
*/
|
||||
IdType key;
|
||||
/**
|
||||
* @brief The value of the item inserted.
|
||||
*/
|
||||
IdType value;
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Cross platform CAS operation.
|
||||
* It is an atomic operation that compares the contents of a memory
|
||||
* location with a given value and, only if they are the same, modifies
|
||||
* the contents of that memory location to a new given value.
|
||||
*
|
||||
* @param ptr The pointer to the object to test and modify .
|
||||
* @param old_val The value expected to be found in `ptr`.
|
||||
* @param new_val The value to store in `ptr` if it is as expected.
|
||||
*
|
||||
* @return Old value pointed by the `ptr`.
|
||||
*/
|
||||
static IdType CompareAndSwap(IdType* ptr, IdType old_val, IdType new_val);
|
||||
|
||||
ConcurrentIdHashMap();
|
||||
|
||||
ConcurrentIdHashMap(const ConcurrentIdHashMap& other) = delete;
|
||||
ConcurrentIdHashMap& operator=(const ConcurrentIdHashMap& other) = delete;
|
||||
|
||||
/**
|
||||
* @brief Initialize the hashmap with an array of ids. The first `num_seeds`
|
||||
* ids are unique and must be mapped to a contiguous array starting
|
||||
* from 0. The left can be duplicated and the mapping result is not stable.
|
||||
*
|
||||
* @param ids The array of the ids to be inserted.
|
||||
* @param num_seeds The number of seed ids.
|
||||
*
|
||||
* @return Unique ids from the input `ids`.
|
||||
*/
|
||||
IdArray Init(const IdArray& ids, size_t num_seeds);
|
||||
|
||||
/**
|
||||
* @brief Find mappings of given keys.
|
||||
*
|
||||
* @param ids The keys to map for.
|
||||
*
|
||||
* @return Mapping results corresponding to `ids`.
|
||||
*/
|
||||
IdArray MapIds(const IdArray& ids) const;
|
||||
|
||||
private:
|
||||
/**
|
||||
* @brief Get the next position and delta for probing.
|
||||
*
|
||||
* @param[in,out] pos Calculate the next position with quadric probing.
|
||||
* @param[in,out] delta Calculate the next delta by adding 1.
|
||||
*/
|
||||
inline void Next(IdType* pos, IdType* delta) const;
|
||||
|
||||
/**
|
||||
* @brief Find the mapping of a given key.
|
||||
*
|
||||
* @param id The key to map for.
|
||||
*
|
||||
* @return Mapping result corresponding to `id`.
|
||||
*/
|
||||
inline IdType MapId(const IdType id) const;
|
||||
|
||||
/**
|
||||
* @brief Insert an id into the hash map.
|
||||
*
|
||||
* @param id The id to be inserted.
|
||||
*
|
||||
* @return Whether the `id` is inserted or not.
|
||||
*/
|
||||
inline bool Insert(IdType id);
|
||||
|
||||
/**
|
||||
* @brief Set the value for the key in the hash map.
|
||||
*
|
||||
* @param key The key to set for.
|
||||
* @param value The value to be set for the `key`.
|
||||
*
|
||||
* @warning Key must exist.
|
||||
*/
|
||||
inline void Set(IdType key, IdType value);
|
||||
|
||||
/**
|
||||
* @brief Insert a key into the hash map.
|
||||
*
|
||||
* @param id The key to be inserted.
|
||||
* @param value The value to be set for the `key`.
|
||||
*
|
||||
*/
|
||||
inline void InsertAndSet(IdType key, IdType value);
|
||||
|
||||
/**
|
||||
* @brief Attempt to insert the key into the hash map at the given position.
|
||||
*
|
||||
* @param pos The position in the hash map to be inserted at.
|
||||
* @param key The key to be inserted.
|
||||
*
|
||||
* @return The state of the insertion.
|
||||
*/
|
||||
inline InsertState AttemptInsertAt(int64_t pos, IdType key);
|
||||
|
||||
private:
|
||||
/**
|
||||
* @brief Hash maps which is used to store all elements.
|
||||
*/
|
||||
std::unique_ptr<Mapping[], std::function<void(Mapping*)>> hash_map_;
|
||||
|
||||
/**
|
||||
* @brief Mask which is assisted to get the position in the table
|
||||
* for a key by performing `&` operation with it.
|
||||
*/
|
||||
IdType mask_;
|
||||
};
|
||||
|
||||
template <typename IdType>
|
||||
bool BoolCompareAndSwap(IdType* ptr);
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CPU_CONCURRENT_ID_HASH_MAP_H_
|
||||
@@ -0,0 +1,56 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file array/cpu/coo_coalesce.cc
|
||||
* @brief COO coalescing
|
||||
*/
|
||||
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include <vector>
|
||||
|
||||
namespace dgl {
|
||||
|
||||
namespace aten {
|
||||
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::pair<COOMatrix, IdArray> COOCoalesce(COOMatrix coo) {
|
||||
const int64_t nnz = coo.row->shape[0];
|
||||
const IdType* coo_row_data = static_cast<IdType*>(coo.row->data);
|
||||
const IdType* coo_col_data = static_cast<IdType*>(coo.col->data);
|
||||
|
||||
if (!coo.row_sorted || !coo.col_sorted) coo = COOSort(coo, true);
|
||||
|
||||
std::vector<IdType> new_row, new_col, count;
|
||||
IdType prev_row = -1, prev_col = -1;
|
||||
for (int64_t i = 0; i < nnz; ++i) {
|
||||
const IdType curr_row = coo_row_data[i];
|
||||
const IdType curr_col = coo_col_data[i];
|
||||
if (curr_row == prev_row && curr_col == prev_col) {
|
||||
++count[count.size() - 1];
|
||||
} else {
|
||||
new_row.push_back(curr_row);
|
||||
new_col.push_back(curr_col);
|
||||
count.push_back(1);
|
||||
prev_row = curr_row;
|
||||
prev_col = curr_col;
|
||||
}
|
||||
}
|
||||
|
||||
COOMatrix coo_result = COOMatrix{
|
||||
coo.num_rows,
|
||||
coo.num_cols,
|
||||
NDArray::FromVector(new_row),
|
||||
NDArray::FromVector(new_col),
|
||||
NullArray(),
|
||||
true};
|
||||
return std::make_pair(coo_result, NDArray::FromVector(count));
|
||||
}
|
||||
|
||||
template std::pair<COOMatrix, IdArray> COOCoalesce<kDGLCPU, int32_t>(COOMatrix);
|
||||
template std::pair<COOMatrix, IdArray> COOCoalesce<kDGLCPU, int64_t>(COOMatrix);
|
||||
|
||||
}; // namespace impl
|
||||
}; // namespace aten
|
||||
}; // namespace dgl
|
||||
@@ -0,0 +1,59 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/coo_line_graph.cc
|
||||
* @brief COO LineGraph
|
||||
*/
|
||||
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <iterator>
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix COOLineGraph(const COOMatrix& coo, bool backtracking) {
|
||||
const int64_t nnz = coo.row->shape[0];
|
||||
IdType* coo_row = coo.row.Ptr<IdType>();
|
||||
IdType* coo_col = coo.col.Ptr<IdType>();
|
||||
IdArray data = COOHasData(coo)
|
||||
? coo.data
|
||||
: Range(0, nnz, coo.row->dtype.bits, coo.row->ctx);
|
||||
IdType* data_data = data.Ptr<IdType>();
|
||||
std::vector<IdType> new_row;
|
||||
std::vector<IdType> new_col;
|
||||
|
||||
for (int64_t i = 0; i < nnz; ++i) {
|
||||
IdType u = coo_row[i];
|
||||
IdType v = coo_col[i];
|
||||
for (int64_t j = 0; j < nnz; ++j) {
|
||||
// no self-loop
|
||||
if (i == j) continue;
|
||||
|
||||
// succ_u == v
|
||||
// if not backtracking succ_u != u
|
||||
if (v == coo_row[j] && (backtracking || u != coo_col[j])) {
|
||||
new_row.push_back(data_data[i]);
|
||||
new_col.push_back(data_data[j]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
COOMatrix res = COOMatrix(
|
||||
nnz, nnz, NDArray::FromVector(new_row), NDArray::FromVector(new_col),
|
||||
NullArray(), false, false);
|
||||
return res;
|
||||
}
|
||||
|
||||
template COOMatrix COOLineGraph<kDGLCPU, int32_t>(
|
||||
const COOMatrix& coo, bool backtracking);
|
||||
template COOMatrix COOLineGraph<kDGLCPU, int64_t>(
|
||||
const COOMatrix& coo, bool backtracking);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,100 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/coo_remove.cc
|
||||
* @brief COO matrix remove entries CPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "array_utils.h"
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
namespace {
|
||||
|
||||
/** @brief COORemove implementation for COOMatrix with default consecutive edge
|
||||
* IDs */
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
void COORemoveConsecutive(
|
||||
COOMatrix coo, IdArray entries, std::vector<IdType> *new_rows,
|
||||
std::vector<IdType> *new_cols, std::vector<IdType> *new_eids) {
|
||||
const int64_t nnz = coo.row->shape[0];
|
||||
const int64_t n_entries = entries->shape[0];
|
||||
const IdType *row_data = static_cast<IdType *>(coo.row->data);
|
||||
const IdType *col_data = static_cast<IdType *>(coo.col->data);
|
||||
const IdType *entry_data = static_cast<IdType *>(entries->data);
|
||||
|
||||
std::vector<IdType> entry_data_sorted(entry_data, entry_data + n_entries);
|
||||
std::sort(entry_data_sorted.begin(), entry_data_sorted.end());
|
||||
|
||||
int64_t j = 0;
|
||||
for (int64_t i = 0; i < nnz; ++i) {
|
||||
if (j < n_entries && entry_data_sorted[j] == i) {
|
||||
// Move on to the next different entry
|
||||
while (j < n_entries && entry_data_sorted[j] == i) ++j;
|
||||
continue;
|
||||
}
|
||||
new_rows->push_back(row_data[i]);
|
||||
new_cols->push_back(col_data[i]);
|
||||
new_eids->push_back(i);
|
||||
}
|
||||
}
|
||||
|
||||
/** @brief COORemove implementation for COOMatrix with shuffled edge IDs */
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
void COORemoveShuffled(
|
||||
COOMatrix coo, IdArray entries, std::vector<IdType> *new_rows,
|
||||
std::vector<IdType> *new_cols, std::vector<IdType> *new_eids) {
|
||||
const int64_t nnz = coo.row->shape[0];
|
||||
const IdType *row_data = static_cast<IdType *>(coo.row->data);
|
||||
const IdType *col_data = static_cast<IdType *>(coo.col->data);
|
||||
const IdType *eid_data = static_cast<IdType *>(coo.data->data);
|
||||
|
||||
IdHashMap<IdType> eid_map(entries);
|
||||
|
||||
for (int64_t i = 0; i < nnz; ++i) {
|
||||
const IdType eid = eid_data[i];
|
||||
if (eid_map.Contains(eid)) continue;
|
||||
new_rows->push_back(row_data[i]);
|
||||
new_cols->push_back(col_data[i]);
|
||||
new_eids->push_back(eid);
|
||||
}
|
||||
}
|
||||
|
||||
}; // namespace
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix COORemove(COOMatrix coo, IdArray entries) {
|
||||
const int64_t nnz = coo.row->shape[0];
|
||||
const int64_t n_entries = entries->shape[0];
|
||||
if (n_entries == 0) return coo;
|
||||
|
||||
std::vector<IdType> new_rows, new_cols, new_eids;
|
||||
new_rows.reserve(nnz - n_entries);
|
||||
new_cols.reserve(nnz - n_entries);
|
||||
new_eids.reserve(nnz - n_entries);
|
||||
|
||||
if (COOHasData(coo))
|
||||
COORemoveShuffled<XPU, IdType>(
|
||||
coo, entries, &new_rows, &new_cols, &new_eids);
|
||||
else
|
||||
// Removing from COO ordered by eid has more efficient implementation.
|
||||
COORemoveConsecutive<XPU, IdType>(
|
||||
coo, entries, &new_rows, &new_cols, &new_eids);
|
||||
|
||||
return COOMatrix(
|
||||
coo.num_rows, coo.num_cols, IdArray::FromVector(new_rows),
|
||||
IdArray::FromVector(new_cols), IdArray::FromVector(new_eids));
|
||||
}
|
||||
|
||||
template COOMatrix COORemove<kDGLCPU, int32_t>(COOMatrix coo, IdArray entries);
|
||||
template COOMatrix COORemove<kDGLCPU, int64_t>(COOMatrix coo, IdArray entries);
|
||||
|
||||
}; // namespace impl
|
||||
}; // namespace aten
|
||||
}; // namespace dgl
|
||||
@@ -0,0 +1,219 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/coo_sort.cc
|
||||
* @brief COO sorting
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
#ifdef PARALLEL_ALGORITHMS
|
||||
#include <parallel/algorithm>
|
||||
#endif
|
||||
#include <algorithm>
|
||||
#include <iterator>
|
||||
#include <numeric>
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
|
||||
template <typename IdType>
|
||||
struct TupleRef {
|
||||
TupleRef() = delete;
|
||||
TupleRef(const TupleRef& other) = default;
|
||||
TupleRef(TupleRef&& other) = default;
|
||||
TupleRef(IdType* const r, IdType* const c, IdType* const d)
|
||||
: row(r), col(c), data(d) {}
|
||||
|
||||
TupleRef& operator=(const TupleRef& other) {
|
||||
*row = *other.row;
|
||||
*col = *other.col;
|
||||
*data = *other.data;
|
||||
return *this;
|
||||
}
|
||||
TupleRef& operator=(const std::tuple<IdType, IdType, IdType>& val) {
|
||||
*row = std::get<0>(val);
|
||||
*col = std::get<1>(val);
|
||||
*data = std::get<2>(val);
|
||||
return *this;
|
||||
}
|
||||
|
||||
operator std::tuple<IdType, IdType, IdType>() const {
|
||||
return std::make_tuple(*row, *col, *data);
|
||||
}
|
||||
|
||||
void Swap(const TupleRef& other) const {
|
||||
std::swap(*row, *other.row);
|
||||
std::swap(*col, *other.col);
|
||||
std::swap(*data, *other.data);
|
||||
}
|
||||
|
||||
IdType *row, *col, *data;
|
||||
};
|
||||
|
||||
using std::swap;
|
||||
template <typename IdType>
|
||||
void swap(const TupleRef<IdType>& r1, const TupleRef<IdType>& r2) {
|
||||
r1.Swap(r2);
|
||||
}
|
||||
|
||||
template <typename IdType>
|
||||
struct CooIterator
|
||||
: public std::iterator<
|
||||
std::random_access_iterator_tag, std::tuple<IdType, IdType, IdType>,
|
||||
std::ptrdiff_t, std::tuple<IdType*, IdType*, IdType*>,
|
||||
TupleRef<IdType>> {
|
||||
CooIterator() = default;
|
||||
CooIterator(const CooIterator& other) = default;
|
||||
CooIterator(CooIterator&& other) = default;
|
||||
CooIterator(IdType* r, IdType* c, IdType* d) : row(r), col(c), data(d) {}
|
||||
|
||||
CooIterator& operator=(const CooIterator& other) = default;
|
||||
CooIterator& operator=(CooIterator&& other) = default;
|
||||
~CooIterator() = default;
|
||||
|
||||
bool operator==(const CooIterator& other) const { return row == other.row; }
|
||||
|
||||
bool operator!=(const CooIterator& other) const { return row != other.row; }
|
||||
|
||||
bool operator<(const CooIterator& other) const { return row < other.row; }
|
||||
|
||||
bool operator>(const CooIterator& other) const { return row > other.row; }
|
||||
|
||||
bool operator<=(const CooIterator& other) const { return row <= other.row; }
|
||||
|
||||
bool operator>=(const CooIterator& other) const { return row >= other.row; }
|
||||
|
||||
CooIterator& operator+=(const std::ptrdiff_t& movement) {
|
||||
row += movement;
|
||||
col += movement;
|
||||
data += movement;
|
||||
return *this;
|
||||
}
|
||||
|
||||
CooIterator& operator-=(const std::ptrdiff_t& movement) {
|
||||
row -= movement;
|
||||
col -= movement;
|
||||
data -= movement;
|
||||
return *this;
|
||||
}
|
||||
|
||||
CooIterator& operator++() { return operator+=(1); }
|
||||
|
||||
CooIterator& operator--() { return operator-=(1); }
|
||||
|
||||
CooIterator operator++(int) {
|
||||
CooIterator ret(*this);
|
||||
operator++();
|
||||
return ret;
|
||||
}
|
||||
|
||||
CooIterator operator--(int) {
|
||||
CooIterator ret(*this);
|
||||
operator--();
|
||||
return ret;
|
||||
}
|
||||
|
||||
CooIterator operator+(const std::ptrdiff_t& movement) const {
|
||||
CooIterator ret(*this);
|
||||
ret += movement;
|
||||
return ret;
|
||||
}
|
||||
|
||||
CooIterator operator-(const std::ptrdiff_t& movement) const {
|
||||
CooIterator ret(*this);
|
||||
ret -= movement;
|
||||
return ret;
|
||||
}
|
||||
|
||||
std::ptrdiff_t operator-(const CooIterator& other) const {
|
||||
return row - other.row;
|
||||
}
|
||||
|
||||
TupleRef<IdType> operator*() const {
|
||||
return TupleRef<IdType>(row, col, data);
|
||||
}
|
||||
TupleRef<IdType> operator*() { return TupleRef<IdType>(row, col, data); }
|
||||
|
||||
// required for random access iterators in VS2019
|
||||
TupleRef<IdType> operator[](size_t offset) const {
|
||||
return TupleRef<IdType>(row + offset, col + offset, data + offset);
|
||||
}
|
||||
|
||||
IdType *row, *col, *data;
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
///////////////////////////// COOSort_ /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
void COOSort_(COOMatrix* coo, bool sort_column) {
|
||||
const int64_t nnz = coo->row->shape[0];
|
||||
IdType* coo_row = coo->row.Ptr<IdType>();
|
||||
IdType* coo_col = coo->col.Ptr<IdType>();
|
||||
if (!COOHasData(*coo))
|
||||
coo->data = aten::Range(0, nnz, coo->row->dtype.bits, coo->row->ctx);
|
||||
IdType* coo_data = coo->data.Ptr<IdType>();
|
||||
|
||||
typedef std::tuple<IdType, IdType, IdType> Tuple;
|
||||
|
||||
// Arg sort
|
||||
if (sort_column) {
|
||||
#ifdef PARALLEL_ALGORITHMS
|
||||
__gnu_parallel::sort(
|
||||
#else
|
||||
std::sort(
|
||||
#endif
|
||||
CooIterator<IdType>(coo_row, coo_col, coo_data),
|
||||
CooIterator<IdType>(coo_row, coo_col, coo_data) + nnz,
|
||||
[](const Tuple& a, const Tuple& b) {
|
||||
return (std::get<0>(a) != std::get<0>(b))
|
||||
? (std::get<0>(a) < std::get<0>(b))
|
||||
: (std::get<1>(a) < std::get<1>(b));
|
||||
});
|
||||
} else {
|
||||
#ifdef PARALLEL_ALGORITHMS
|
||||
__gnu_parallel::sort(
|
||||
#else
|
||||
std::sort(
|
||||
#endif
|
||||
CooIterator<IdType>(coo_row, coo_col, coo_data),
|
||||
CooIterator<IdType>(coo_row, coo_col, coo_data) + nnz,
|
||||
[](const Tuple& a, const Tuple& b) {
|
||||
return std::get<0>(a) < std::get<0>(b);
|
||||
});
|
||||
}
|
||||
|
||||
coo->row_sorted = true;
|
||||
coo->col_sorted = sort_column;
|
||||
}
|
||||
|
||||
template void COOSort_<kDGLCPU, int32_t>(COOMatrix*, bool);
|
||||
template void COOSort_<kDGLCPU, int64_t>(COOMatrix*, bool);
|
||||
|
||||
///////////////////////////// COOIsSorted /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::pair<bool, bool> COOIsSorted(COOMatrix coo) {
|
||||
const int64_t nnz = coo.row->shape[0];
|
||||
IdType* row = coo.row.Ptr<IdType>();
|
||||
IdType* col = coo.col.Ptr<IdType>();
|
||||
bool row_sorted = true;
|
||||
bool col_sorted = true;
|
||||
for (int64_t i = 1; row_sorted && i < nnz; ++i) {
|
||||
row_sorted = (row[i - 1] <= row[i]);
|
||||
col_sorted = col_sorted && (row[i - 1] < row[i] || col[i - 1] <= col[i]);
|
||||
}
|
||||
if (!row_sorted) col_sorted = false;
|
||||
return {row_sorted, col_sorted};
|
||||
}
|
||||
|
||||
template std::pair<bool, bool> COOIsSorted<kDGLCPU, int32_t>(COOMatrix coo);
|
||||
template std::pair<bool, bool> COOIsSorted<kDGLCPU, int64_t>(COOMatrix coo);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,144 @@
|
||||
/**
|
||||
* Copyright (c) 2021 by Contributors
|
||||
* @file array/cpu/csr_get_data.cc
|
||||
* @brief Retrieve entries of a CSR matrix
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
|
||||
#include <numeric>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#include "array_utils.h"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using runtime::NDArray;
|
||||
using runtime::parallel_for;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
void CollectDataFromSorted(
|
||||
const IdType* indices_data, const IdType* data, const IdType start,
|
||||
const IdType end, const IdType col, std::vector<IdType>* ret_vec) {
|
||||
const IdType* start_ptr = indices_data + start;
|
||||
const IdType* end_ptr = indices_data + end;
|
||||
auto it = std::lower_bound(start_ptr, end_ptr, col);
|
||||
// This might be a multi-graph. We need to collect all of the matched
|
||||
// columns.
|
||||
for (; it != end_ptr; it++) {
|
||||
// If the col exist
|
||||
if (*it == col) {
|
||||
IdType idx = it - indices_data;
|
||||
ret_vec->push_back(data ? data[idx] : idx);
|
||||
} else {
|
||||
// If we find a column that is different, we can stop searching now.
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename DType>
|
||||
NDArray CSRGetData(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, DType filler) {
|
||||
const int64_t rowlen = rows->shape[0];
|
||||
const int64_t collen = cols->shape[0];
|
||||
|
||||
CHECK((rowlen == collen) || (rowlen == 1) || (collen == 1))
|
||||
<< "Invalid row and col id array.";
|
||||
|
||||
const int64_t row_stride = (rowlen == 1 && collen != 1) ? 0 : 1;
|
||||
const int64_t col_stride = (collen == 1 && rowlen != 1) ? 0 : 1;
|
||||
const IdType* row_data = static_cast<IdType*>(rows->data);
|
||||
const IdType* col_data = static_cast<IdType*>(cols->data);
|
||||
|
||||
const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
|
||||
const IdType* indices_data = static_cast<IdType*>(csr.indices->data);
|
||||
const IdType* data =
|
||||
CSRHasData(csr) ? static_cast<IdType*>(csr.data->data) : nullptr;
|
||||
|
||||
const int64_t retlen = std::max(rowlen, collen);
|
||||
const DType* weight_data = return_eids ? nullptr : weights.Ptr<DType>();
|
||||
if (return_eids)
|
||||
BUG_IF_FAIL(DGLDataTypeTraits<DType>::dtype == rows->dtype)
|
||||
<< "DType does not match row's dtype.";
|
||||
|
||||
NDArray ret = Full(filler, retlen, rows->ctx);
|
||||
DType* ret_data = ret.Ptr<DType>();
|
||||
|
||||
// NOTE: In most cases, the input csr is already sorted. If not, we might need
|
||||
// to
|
||||
// consider sorting it especially when the number of (row, col) pairs is
|
||||
// large. Need more benchmarks to justify the choice.
|
||||
|
||||
if (csr.sorted) {
|
||||
// use binary search on each row
|
||||
parallel_for(0, retlen, [&](size_t b, size_t e) {
|
||||
for (auto p = b; p < e; ++p) {
|
||||
const IdType row_id = row_data[p * row_stride],
|
||||
col_id = col_data[p * col_stride];
|
||||
CHECK(row_id >= 0 && row_id < csr.num_rows)
|
||||
<< "Invalid row index: " << row_id;
|
||||
CHECK(col_id >= 0 && col_id < csr.num_cols)
|
||||
<< "Invalid col index: " << col_id;
|
||||
const IdType* start_ptr = indices_data + indptr_data[row_id];
|
||||
const IdType* end_ptr = indices_data + indptr_data[row_id + 1];
|
||||
auto it = std::lower_bound(start_ptr, end_ptr, col_id);
|
||||
if (it != end_ptr && *it == col_id) {
|
||||
const IdType idx = it - indices_data;
|
||||
IdType eid = data ? data[idx] : idx;
|
||||
ret_data[p] = return_eids ? eid : weight_data[eid];
|
||||
}
|
||||
}
|
||||
});
|
||||
} else {
|
||||
// linear search on each row
|
||||
parallel_for(0, retlen, [&](size_t b, size_t e) {
|
||||
for (auto p = b; p < e; ++p) {
|
||||
const IdType row_id = row_data[p * row_stride],
|
||||
col_id = col_data[p * col_stride];
|
||||
CHECK(row_id >= 0 && row_id < csr.num_rows)
|
||||
<< "Invalid row index: " << row_id;
|
||||
CHECK(col_id >= 0 && col_id < csr.num_cols)
|
||||
<< "Invalid col index: " << col_id;
|
||||
for (IdType idx = indptr_data[row_id]; idx < indptr_data[row_id + 1];
|
||||
++idx) {
|
||||
if (indices_data[idx] == col_id) {
|
||||
IdType eid = data ? data[idx] : idx;
|
||||
ret_data[p] = return_eids ? eid : weight_data[eid];
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
template NDArray CSRGetData<kDGLCPU, int32_t, float>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, float filler);
|
||||
template NDArray CSRGetData<kDGLCPU, int64_t, float>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, float filler);
|
||||
template NDArray CSRGetData<kDGLCPU, int32_t, double>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, double filler);
|
||||
template NDArray CSRGetData<kDGLCPU, int64_t, double>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, double filler);
|
||||
|
||||
// For CSRGetData<XPU, IdType>(CSRMatrix, NDArray, NDArray)
|
||||
template NDArray CSRGetData<kDGLCPU, int32_t, int32_t>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, int32_t filler);
|
||||
template NDArray CSRGetData<kDGLCPU, int64_t, int64_t>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, int64_t filler);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,138 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file array/cpu/csr_mm.cc
|
||||
* @brief CSR Matrix Multiplication
|
||||
*/
|
||||
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
#include <tsl/robin_map.h>
|
||||
#include <tsl/robin_set.h>
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "array_utils.h"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using dgl::runtime::NDArray;
|
||||
using dgl::runtime::parallel_for;
|
||||
|
||||
namespace aten {
|
||||
|
||||
namespace {
|
||||
|
||||
// TODO(BarclayII): avoid using map for sorted CSRs
|
||||
template <typename IdType>
|
||||
void CountNNZPerRow(
|
||||
const IdType* A_indptr, const IdType* A_indices, const IdType* B_indptr,
|
||||
const IdType* B_indices, IdType* C_indptr_data, int64_t M) {
|
||||
parallel_for(0, M, [=](size_t b, size_t e) {
|
||||
for (auto i = b; i < e; ++i) {
|
||||
tsl::robin_set<IdType> set;
|
||||
for (IdType u = A_indptr[i]; u < A_indptr[i + 1]; ++u) {
|
||||
IdType w = A_indices[u];
|
||||
for (IdType v = B_indptr[w]; v < B_indptr[w + 1]; ++v)
|
||||
set.insert(B_indices[v]);
|
||||
}
|
||||
C_indptr_data[i] = set.size();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename IdType>
|
||||
int64_t ComputeIndptrInPlace(IdType* C_indptr_data, int64_t M) {
|
||||
int64_t nnz = 0;
|
||||
IdType len = 0;
|
||||
for (IdType i = 0; i < M; ++i) {
|
||||
len = C_indptr_data[i];
|
||||
C_indptr_data[i] = nnz;
|
||||
nnz += len;
|
||||
}
|
||||
C_indptr_data[M] = nnz;
|
||||
return nnz;
|
||||
}
|
||||
|
||||
template <typename IdType, typename DType>
|
||||
void ComputeIndicesAndData(
|
||||
const IdType* A_indptr, const IdType* A_indices, const IdType* A_eids,
|
||||
const DType* A_data, const IdType* B_indptr, const IdType* B_indices,
|
||||
const IdType* B_eids, const DType* B_data, const IdType* C_indptr_data,
|
||||
IdType* C_indices_data, DType* C_weights_data, int64_t M) {
|
||||
parallel_for(0, M, [=](size_t b, size_t e) {
|
||||
for (auto i = b; i < e; ++i) {
|
||||
tsl::robin_map<IdType, DType> map;
|
||||
for (IdType u = A_indptr[i]; u < A_indptr[i + 1]; ++u) {
|
||||
IdType w = A_indices[u];
|
||||
DType vA = A_data[A_eids ? A_eids[u] : u];
|
||||
for (IdType v = B_indptr[w]; v < B_indptr[w + 1]; ++v) {
|
||||
IdType t = B_indices[v];
|
||||
DType vB = B_data[B_eids ? B_eids[v] : v];
|
||||
map[t] += vA * vB;
|
||||
}
|
||||
}
|
||||
|
||||
IdType v = C_indptr_data[i];
|
||||
for (auto it : map) {
|
||||
C_indices_data[v] = it.first;
|
||||
C_weights_data[v] = it.second;
|
||||
++v;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
}; // namespace
|
||||
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
std::pair<CSRMatrix, NDArray> CSRMM(
|
||||
const CSRMatrix& A, NDArray A_weights, const CSRMatrix& B,
|
||||
NDArray B_weights) {
|
||||
CHECK_EQ(A.num_cols, B.num_rows)
|
||||
<< "A's number of columns must equal to B's number of rows";
|
||||
const bool A_has_eid = !IsNullArray(A.data);
|
||||
const bool B_has_eid = !IsNullArray(B.data);
|
||||
const IdType* A_indptr = A.indptr.Ptr<IdType>();
|
||||
const IdType* A_indices = A.indices.Ptr<IdType>();
|
||||
const IdType* A_eids = A_has_eid ? A.data.Ptr<IdType>() : nullptr;
|
||||
const IdType* B_indptr = B.indptr.Ptr<IdType>();
|
||||
const IdType* B_indices = B.indices.Ptr<IdType>();
|
||||
const IdType* B_eids = B_has_eid ? B.data.Ptr<IdType>() : nullptr;
|
||||
const DType* A_data = A_weights.Ptr<DType>();
|
||||
const DType* B_data = B_weights.Ptr<DType>();
|
||||
const int64_t M = A.num_rows;
|
||||
const int64_t P = B.num_cols;
|
||||
|
||||
IdArray C_indptr = IdArray::Empty({M + 1}, A.indptr->dtype, A.indptr->ctx);
|
||||
IdType* C_indptr_data = C_indptr.Ptr<IdType>();
|
||||
|
||||
CountNNZPerRow<IdType>(
|
||||
A_indptr, A_indices, B_indptr, B_indices, C_indptr_data, M);
|
||||
int64_t nnz = ComputeIndptrInPlace<IdType>(C_indptr_data, M);
|
||||
// Allocate indices and weights array
|
||||
IdArray C_indices = IdArray::Empty({nnz}, A.indices->dtype, A.indices->ctx);
|
||||
NDArray C_weights = NDArray::Empty({nnz}, A_weights->dtype, A_weights->ctx);
|
||||
IdType* C_indices_data = C_indices.Ptr<IdType>();
|
||||
DType* C_weights_data = C_weights.Ptr<DType>();
|
||||
|
||||
ComputeIndicesAndData<IdType, DType>(
|
||||
A_indptr, A_indices, A_eids, A_data, B_indptr, B_indices, B_eids, B_data,
|
||||
C_indptr_data, C_indices_data, C_weights_data, M);
|
||||
|
||||
return {
|
||||
CSRMatrix(
|
||||
M, P, C_indptr, C_indices, NullArray(C_indptr->dtype, C_indptr->ctx)),
|
||||
C_weights};
|
||||
}
|
||||
|
||||
template std::pair<CSRMatrix, NDArray> CSRMM<kDGLCPU, int32_t, float>(
|
||||
const CSRMatrix&, NDArray, const CSRMatrix&, NDArray);
|
||||
template std::pair<CSRMatrix, NDArray> CSRMM<kDGLCPU, int64_t, float>(
|
||||
const CSRMatrix&, NDArray, const CSRMatrix&, NDArray);
|
||||
template std::pair<CSRMatrix, NDArray> CSRMM<kDGLCPU, int32_t, double>(
|
||||
const CSRMatrix&, NDArray, const CSRMatrix&, NDArray);
|
||||
template std::pair<CSRMatrix, NDArray> CSRMM<kDGLCPU, int64_t, double>(
|
||||
const CSRMatrix&, NDArray, const CSRMatrix&, NDArray);
|
||||
|
||||
}; // namespace aten
|
||||
}; // namespace dgl
|
||||
@@ -0,0 +1,104 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/coo_remove.cc
|
||||
* @brief CSR matrix remove entries CPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "array_utils.h"
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
namespace {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
void CSRRemoveConsecutive(
|
||||
CSRMatrix csr, IdArray entries, std::vector<IdType> *new_indptr,
|
||||
std::vector<IdType> *new_indices, std::vector<IdType> *new_eids) {
|
||||
CHECK_SAME_DTYPE(csr.indices, entries);
|
||||
const int64_t n_entries = entries->shape[0];
|
||||
const IdType *indptr_data = static_cast<IdType *>(csr.indptr->data);
|
||||
const IdType *indices_data = static_cast<IdType *>(csr.indices->data);
|
||||
const IdType *entry_data = static_cast<IdType *>(entries->data);
|
||||
|
||||
std::vector<IdType> entry_data_sorted(entry_data, entry_data + n_entries);
|
||||
std::sort(entry_data_sorted.begin(), entry_data_sorted.end());
|
||||
|
||||
int64_t k = 0;
|
||||
new_indptr->push_back(0);
|
||||
for (int64_t i = 0; i < csr.num_rows; ++i) {
|
||||
for (IdType j = indptr_data[i]; j < indptr_data[i + 1]; ++j) {
|
||||
if (k < n_entries && entry_data_sorted[k] == j) {
|
||||
// Move on to the next different entry
|
||||
while (k < n_entries && entry_data_sorted[k] == j) ++k;
|
||||
continue;
|
||||
}
|
||||
new_indices->push_back(indices_data[j]);
|
||||
new_eids->push_back(k);
|
||||
}
|
||||
new_indptr->push_back(new_indices->size());
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
void CSRRemoveShuffled(
|
||||
CSRMatrix csr, IdArray entries, std::vector<IdType> *new_indptr,
|
||||
std::vector<IdType> *new_indices, std::vector<IdType> *new_eids) {
|
||||
CHECK_SAME_DTYPE(csr.indices, entries);
|
||||
const IdType *indptr_data = static_cast<IdType *>(csr.indptr->data);
|
||||
const IdType *indices_data = static_cast<IdType *>(csr.indices->data);
|
||||
const IdType *eid_data = static_cast<IdType *>(csr.data->data);
|
||||
|
||||
IdHashMap<IdType> eid_map(entries);
|
||||
|
||||
new_indptr->push_back(0);
|
||||
for (int64_t i = 0; i < csr.num_rows; ++i) {
|
||||
for (IdType j = indptr_data[i]; j < indptr_data[i + 1]; ++j) {
|
||||
const IdType eid = eid_data ? eid_data[j] : j;
|
||||
if (eid_map.Contains(eid)) continue;
|
||||
new_indices->push_back(indices_data[j]);
|
||||
new_eids->push_back(eid);
|
||||
}
|
||||
new_indptr->push_back(new_indices->size());
|
||||
}
|
||||
}
|
||||
|
||||
}; // namespace
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix CSRRemove(CSRMatrix csr, IdArray entries) {
|
||||
CHECK_SAME_DTYPE(csr.indices, entries);
|
||||
const int64_t nnz = csr.indices->shape[0];
|
||||
const int64_t n_entries = entries->shape[0];
|
||||
if (n_entries == 0) return csr;
|
||||
|
||||
std::vector<IdType> new_indptr, new_indices, new_eids;
|
||||
new_indptr.reserve(nnz - n_entries);
|
||||
new_indices.reserve(nnz - n_entries);
|
||||
new_eids.reserve(nnz - n_entries);
|
||||
|
||||
if (CSRHasData(csr))
|
||||
CSRRemoveShuffled<XPU, IdType>(
|
||||
csr, entries, &new_indptr, &new_indices, &new_eids);
|
||||
else
|
||||
// Removing from CSR ordered by eid has more efficient implementation
|
||||
CSRRemoveConsecutive<XPU, IdType>(
|
||||
csr, entries, &new_indptr, &new_indices, &new_eids);
|
||||
|
||||
return CSRMatrix(
|
||||
csr.num_rows, csr.num_cols, IdArray::FromVector(new_indptr),
|
||||
IdArray::FromVector(new_indices), IdArray::FromVector(new_eids));
|
||||
}
|
||||
|
||||
template CSRMatrix CSRRemove<kDGLCPU, int32_t>(CSRMatrix csr, IdArray entries);
|
||||
template CSRMatrix CSRRemove<kDGLCPU, int64_t>(CSRMatrix csr, IdArray entries);
|
||||
|
||||
}; // namespace impl
|
||||
}; // namespace aten
|
||||
}; // namespace dgl
|
||||
@@ -0,0 +1,159 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/csr_sort.cc
|
||||
* @brief CSR sorting
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
///////////////////////////// CSRIsSorted /////////////////////////////
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
bool CSRIsSorted(CSRMatrix csr) {
|
||||
const IdType *indptr = csr.indptr.Ptr<IdType>();
|
||||
const IdType *indices = csr.indices.Ptr<IdType>();
|
||||
return runtime::parallel_reduce(
|
||||
0, csr.num_rows, 1, 1,
|
||||
[indptr, indices](size_t b, size_t e, bool ident) {
|
||||
for (size_t row = b; row < e; ++row) {
|
||||
for (IdType i = indptr[row] + 1; i < indptr[row + 1]; ++i) {
|
||||
if (indices[i - 1] > indices[i]) return false;
|
||||
}
|
||||
}
|
||||
return ident;
|
||||
},
|
||||
[](bool a, bool b) { return a && b; });
|
||||
}
|
||||
|
||||
template bool CSRIsSorted<kDGLCPU, int64_t>(CSRMatrix csr);
|
||||
template bool CSRIsSorted<kDGLCPU, int32_t>(CSRMatrix csr);
|
||||
|
||||
///////////////////////////// CSRSort /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
void CSRSort_(CSRMatrix *csr) {
|
||||
typedef std::pair<IdType, IdType> ShufflePair;
|
||||
const int64_t num_rows = csr->num_rows;
|
||||
const int64_t nnz = csr->indices->shape[0];
|
||||
const IdType *indptr_data = static_cast<IdType *>(csr->indptr->data);
|
||||
IdType *indices_data = static_cast<IdType *>(csr->indices->data);
|
||||
|
||||
if (CSRIsSorted(*csr)) {
|
||||
csr->sorted = true;
|
||||
return;
|
||||
}
|
||||
|
||||
if (!CSRHasData(*csr)) {
|
||||
csr->data = aten::Range(0, nnz, csr->indptr->dtype.bits, csr->indptr->ctx);
|
||||
}
|
||||
IdType *eid_data = static_cast<IdType *>(csr->data->data);
|
||||
|
||||
runtime::parallel_for(0, num_rows, [=](size_t b, size_t e) {
|
||||
for (auto row = b; row < e; ++row) {
|
||||
const int64_t num_cols = indptr_data[row + 1] - indptr_data[row];
|
||||
std::vector<ShufflePair> reorder_vec(num_cols);
|
||||
IdType *col = indices_data + indptr_data[row];
|
||||
IdType *eid = eid_data + indptr_data[row];
|
||||
|
||||
for (int64_t i = 0; i < num_cols; i++) {
|
||||
reorder_vec[i].first = col[i];
|
||||
reorder_vec[i].second = eid[i];
|
||||
}
|
||||
std::sort(
|
||||
reorder_vec.begin(), reorder_vec.end(),
|
||||
[](const ShufflePair &e1, const ShufflePair &e2) {
|
||||
return e1.first < e2.first;
|
||||
});
|
||||
for (int64_t i = 0; i < num_cols; i++) {
|
||||
col[i] = reorder_vec[i].first;
|
||||
eid[i] = reorder_vec[i].second;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
csr->sorted = true;
|
||||
}
|
||||
|
||||
template void CSRSort_<kDGLCPU, int64_t>(CSRMatrix *csr);
|
||||
template void CSRSort_<kDGLCPU, int32_t>(CSRMatrix *csr);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename TagType>
|
||||
std::pair<CSRMatrix, NDArray> CSRSortByTag(
|
||||
const CSRMatrix &csr, const IdArray tag_array, int64_t num_tags) {
|
||||
const auto indptr_data = static_cast<const IdType *>(csr.indptr->data);
|
||||
const auto indices_data = static_cast<const IdType *>(csr.indices->data);
|
||||
const auto eid_data = aten::CSRHasData(csr)
|
||||
? static_cast<const IdType *>(csr.data->data)
|
||||
: nullptr;
|
||||
const auto tag_data = static_cast<const TagType *>(tag_array->data);
|
||||
const int64_t num_rows = csr.num_rows;
|
||||
|
||||
NDArray tag_pos = NDArray::Empty(
|
||||
{csr.num_rows, num_tags + 1}, csr.indptr->dtype, csr.indptr->ctx);
|
||||
auto tag_pos_data = static_cast<IdType *>(tag_pos->data);
|
||||
std::fill(tag_pos_data, tag_pos_data + csr.num_rows * (num_tags + 1), 0);
|
||||
|
||||
aten::CSRMatrix output(
|
||||
csr.num_rows, csr.num_cols, csr.indptr.Clone(), csr.indices.Clone(),
|
||||
NDArray::Empty(
|
||||
{csr.indices->shape[0]}, csr.indices->dtype, csr.indices->ctx),
|
||||
csr.sorted);
|
||||
|
||||
auto out_indices_data = static_cast<IdType *>(output.indices->data);
|
||||
auto out_eid_data = static_cast<IdType *>(output.data->data);
|
||||
|
||||
runtime::parallel_for(0, num_rows, [&](size_t b, size_t e) {
|
||||
for (auto src = b; src < e; ++src) {
|
||||
const IdType start = indptr_data[src];
|
||||
const IdType end = indptr_data[src + 1];
|
||||
|
||||
auto tag_pos_row = tag_pos_data + src * (num_tags + 1);
|
||||
std::vector<IdType> pointer(num_tags, 0);
|
||||
|
||||
for (IdType ptr = start; ptr < end; ++ptr) {
|
||||
const IdType eid = eid_data ? eid_data[ptr] : ptr;
|
||||
const TagType tag = tag_data[eid];
|
||||
CHECK_LT(tag, num_tags);
|
||||
++tag_pos_row[tag + 1];
|
||||
} // count
|
||||
|
||||
for (TagType tag = 1; tag <= num_tags; ++tag) {
|
||||
tag_pos_row[tag] += tag_pos_row[tag - 1];
|
||||
} // cumulate
|
||||
|
||||
for (IdType ptr = start; ptr < end; ++ptr) {
|
||||
const IdType dst = indices_data[ptr];
|
||||
const IdType eid = eid_data ? eid_data[ptr] : ptr;
|
||||
const TagType tag = tag_data[eid];
|
||||
const IdType offset = tag_pos_row[tag] + pointer[tag];
|
||||
CHECK_LT(offset, tag_pos_row[tag + 1]);
|
||||
++pointer[tag];
|
||||
|
||||
out_indices_data[start + offset] = dst;
|
||||
out_eid_data[start + offset] = eid;
|
||||
}
|
||||
}
|
||||
});
|
||||
output.sorted = false;
|
||||
return std::make_pair(output, tag_pos);
|
||||
}
|
||||
|
||||
template std::pair<CSRMatrix, NDArray> CSRSortByTag<kDGLCPU, int64_t, int64_t>(
|
||||
const CSRMatrix &csr, const IdArray tag, int64_t num_tags);
|
||||
template std::pair<CSRMatrix, NDArray> CSRSortByTag<kDGLCPU, int64_t, int32_t>(
|
||||
const CSRMatrix &csr, const IdArray tag, int64_t num_tags);
|
||||
template std::pair<CSRMatrix, NDArray> CSRSortByTag<kDGLCPU, int32_t, int64_t>(
|
||||
const CSRMatrix &csr, const IdArray tag, int64_t num_tags);
|
||||
template std::pair<CSRMatrix, NDArray> CSRSortByTag<kDGLCPU, int32_t, int32_t>(
|
||||
const CSRMatrix &csr, const IdArray tag, int64_t num_tags);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,146 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file array/cpu/csr_sum.cc
|
||||
* @brief CSR Summation
|
||||
*/
|
||||
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
#include <tsl/robin_map.h>
|
||||
#include <tsl/robin_set.h>
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "array_utils.h"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using dgl::runtime::NDArray;
|
||||
|
||||
namespace aten {
|
||||
|
||||
namespace {
|
||||
|
||||
// TODO(BarclayII): avoid using map for sorted CSRs
|
||||
template <typename IdType>
|
||||
void CountNNZPerRow(
|
||||
const std::vector<const IdType*>& A_indptr,
|
||||
const std::vector<const IdType*>& A_indices, IdType* C_indptr_data,
|
||||
int64_t M) {
|
||||
int64_t n = A_indptr.size();
|
||||
|
||||
runtime::parallel_for(0, M, [=](size_t b, size_t e) {
|
||||
for (size_t i = b; i < e; ++i) {
|
||||
tsl::robin_set<IdType> set;
|
||||
for (int64_t k = 0; k < n; ++k) {
|
||||
for (IdType u = A_indptr[k][i]; u < A_indptr[k][i + 1]; ++u)
|
||||
set.insert(A_indices[k][u]);
|
||||
}
|
||||
C_indptr_data[i] = set.size();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
template <typename IdType>
|
||||
int64_t ComputeIndptrInPlace(IdType* C_indptr_data, int64_t M) {
|
||||
int64_t nnz = 0;
|
||||
IdType len = 0;
|
||||
for (IdType i = 0; i < M; ++i) {
|
||||
len = C_indptr_data[i];
|
||||
C_indptr_data[i] = nnz;
|
||||
nnz += len;
|
||||
}
|
||||
C_indptr_data[M] = nnz;
|
||||
return nnz;
|
||||
}
|
||||
|
||||
template <typename IdType, typename DType>
|
||||
void ComputeIndicesAndData(
|
||||
const std::vector<const IdType*>& A_indptr,
|
||||
const std::vector<const IdType*>& A_indices,
|
||||
const std::vector<const IdType*>& A_eids,
|
||||
const std::vector<const DType*>& A_data, const IdType* C_indptr_data,
|
||||
IdType* C_indices_data, DType* C_weights_data, int64_t M) {
|
||||
int64_t n = A_indptr.size();
|
||||
runtime::parallel_for(0, M, [=](size_t b, size_t e) {
|
||||
for (auto i = b; i < e; ++i) {
|
||||
tsl::robin_map<IdType, DType> map;
|
||||
for (int64_t k = 0; k < n; ++k) {
|
||||
for (IdType u = A_indptr[k][i]; u < A_indptr[k][i + 1]; ++u) {
|
||||
IdType kA = A_indices[k][u];
|
||||
DType vA = A_data[k][A_eids[k] ? A_eids[k][u] : u];
|
||||
map[kA] += vA;
|
||||
}
|
||||
}
|
||||
IdType j = C_indptr_data[i];
|
||||
for (auto it : map) {
|
||||
C_indices_data[j] = it.first;
|
||||
C_weights_data[j] = it.second;
|
||||
++j;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
}; // namespace
|
||||
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
std::pair<CSRMatrix, NDArray> CSRSum(
|
||||
const std::vector<CSRMatrix>& A, const std::vector<NDArray>& A_weights) {
|
||||
CHECK(A.size() > 0) << "List of matrices can't be empty.";
|
||||
CHECK_EQ(A.size(), A_weights.size())
|
||||
<< "List of matrices and weights must have same length";
|
||||
const int64_t M = A[0].num_rows;
|
||||
const int64_t N = A[0].num_cols;
|
||||
const int64_t n = A.size();
|
||||
|
||||
std::vector<bool> A_has_eid(n);
|
||||
std::vector<const IdType*> A_indptr(n);
|
||||
std::vector<const IdType*> A_indices(n);
|
||||
std::vector<const IdType*> A_eids(n);
|
||||
std::vector<const DType*> A_data(n);
|
||||
|
||||
for (int64_t i = 0; i < n; ++i) {
|
||||
const CSRMatrix& csr = A[i];
|
||||
const NDArray& data = A_weights[i];
|
||||
A_has_eid[i] = !IsNullArray(csr.data);
|
||||
A_indptr[i] = csr.indptr.Ptr<IdType>();
|
||||
A_indices[i] = csr.indices.Ptr<IdType>();
|
||||
A_eids[i] = A_has_eid[i] ? csr.data.Ptr<IdType>() : nullptr;
|
||||
A_data[i] = data.Ptr<DType>();
|
||||
}
|
||||
|
||||
IdArray C_indptr =
|
||||
IdArray::Empty({M + 1}, A[0].indptr->dtype, A[0].indptr->ctx);
|
||||
IdType* C_indptr_data = C_indptr.Ptr<IdType>();
|
||||
|
||||
CountNNZPerRow<IdType>(A_indptr, A_indices, C_indptr_data, M);
|
||||
IdType nnz = ComputeIndptrInPlace<IdType>(C_indptr_data, M);
|
||||
// Allocate indices and weights array
|
||||
IdArray C_indices =
|
||||
IdArray::Empty({nnz}, A[0].indices->dtype, A[0].indices->ctx);
|
||||
NDArray C_weights =
|
||||
NDArray::Empty({nnz}, A_weights[0]->dtype, A_weights[0]->ctx);
|
||||
IdType* C_indices_data = C_indices.Ptr<IdType>();
|
||||
DType* C_weights_data = C_weights.Ptr<DType>();
|
||||
ComputeIndicesAndData<IdType, DType>(
|
||||
A_indptr, A_indices, A_eids, A_data, C_indptr_data, C_indices_data,
|
||||
C_weights_data, M);
|
||||
|
||||
return {
|
||||
CSRMatrix(
|
||||
M, N, C_indptr, C_indices, NullArray(C_indptr->dtype, C_indptr->ctx)),
|
||||
C_weights};
|
||||
}
|
||||
|
||||
template std::pair<CSRMatrix, NDArray> CSRSum<kDGLCPU, int32_t, float>(
|
||||
const std::vector<CSRMatrix>&, const std::vector<NDArray>&);
|
||||
template std::pair<CSRMatrix, NDArray> CSRSum<kDGLCPU, int64_t, float>(
|
||||
const std::vector<CSRMatrix>&, const std::vector<NDArray>&);
|
||||
template std::pair<CSRMatrix, NDArray> CSRSum<kDGLCPU, int32_t, double>(
|
||||
const std::vector<CSRMatrix>&, const std::vector<NDArray>&);
|
||||
template std::pair<CSRMatrix, NDArray> CSRSum<kDGLCPU, int64_t, double>(
|
||||
const std::vector<CSRMatrix>&, const std::vector<NDArray>&);
|
||||
|
||||
}; // namespace aten
|
||||
}; // namespace dgl
|
||||
@@ -0,0 +1,73 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/csr_to_simple.cc
|
||||
* @brief CSR sorting
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::tuple<CSRMatrix, IdArray, IdArray> CSRToSimple(CSRMatrix csr) {
|
||||
if (!csr.sorted) csr = CSRSort(csr);
|
||||
|
||||
const IdType *indptr_data = static_cast<IdType *>(csr.indptr->data);
|
||||
const IdType *indices_data = static_cast<IdType *>(csr.indices->data);
|
||||
|
||||
std::vector<IdType> indptr;
|
||||
std::vector<IdType> indices;
|
||||
std::vector<IdType> count;
|
||||
indptr.resize(csr.indptr->shape[0]);
|
||||
indptr[0] = 0;
|
||||
|
||||
for (int64_t i = 1; i < csr.indptr->shape[0]; ++i) {
|
||||
if (indptr_data[i - 1] == indptr_data[i]) {
|
||||
indptr[i] = indptr[i - 1];
|
||||
continue;
|
||||
}
|
||||
|
||||
int64_t cnt = 1;
|
||||
int64_t dup_cnt = 1;
|
||||
indices.push_back(indices_data[indptr_data[i - 1]]);
|
||||
for (int64_t j = indptr_data[i - 1] + 1; j < indptr_data[i]; ++j) {
|
||||
if (indices_data[j - 1] == indices_data[j]) {
|
||||
++dup_cnt;
|
||||
continue;
|
||||
}
|
||||
count.push_back(dup_cnt);
|
||||
dup_cnt = 1;
|
||||
indices.push_back(indices_data[j]);
|
||||
++cnt;
|
||||
}
|
||||
count.push_back(dup_cnt);
|
||||
indptr[i] = indptr[i - 1] + cnt;
|
||||
}
|
||||
|
||||
CSRMatrix res_csr = CSRMatrix(
|
||||
csr.num_rows, csr.num_cols, IdArray::FromVector(indptr),
|
||||
IdArray::FromVector(indices), NullArray(), true);
|
||||
|
||||
const IdArray &edge_count = IdArray::FromVector(count);
|
||||
const IdArray new_eids =
|
||||
Range(0, res_csr.indices->shape[0], sizeof(IdType) * 8, csr.indptr->ctx);
|
||||
const IdArray eids_remapped =
|
||||
CSRHasData(csr) ? Scatter(Repeat(new_eids, edge_count), csr.data)
|
||||
: Repeat(new_eids, edge_count);
|
||||
|
||||
return std::make_tuple(res_csr, edge_count, eids_remapped);
|
||||
}
|
||||
|
||||
template std::tuple<CSRMatrix, IdArray, IdArray> CSRToSimple<kDGLCPU, int32_t>(
|
||||
CSRMatrix);
|
||||
template std::tuple<CSRMatrix, IdArray, IdArray> CSRToSimple<kDGLCPU, int64_t>(
|
||||
CSRMatrix);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,114 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/coo_sort.cc
|
||||
* @brief COO sorting
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <iterator>
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix UnionCsr(const std::vector<CSRMatrix> &csrs) {
|
||||
std::vector<IdType> res_indptr;
|
||||
std::vector<IdType> res_indices;
|
||||
std::vector<IdType> res_data;
|
||||
|
||||
// some preprocess
|
||||
// we assume the number of csrs is not large in common cases
|
||||
std::vector<IdArray> data;
|
||||
std::vector<IdType *> data_data;
|
||||
std::vector<IdType *> indptr_data;
|
||||
std::vector<IdType *> indices_data;
|
||||
int64_t num_edges = 0;
|
||||
bool sorted = true;
|
||||
for (size_t i = 0; i < csrs.size(); ++i) {
|
||||
// eids of csrs[0] remains unchanged
|
||||
// eids of csrs[1] will be increased by number of edges of csrs[0], etc.
|
||||
data.push_back(
|
||||
CSRHasData(csrs[i])
|
||||
? csrs[i].data + num_edges
|
||||
: Range(
|
||||
num_edges, num_edges + csrs[i].indices->shape[0],
|
||||
csrs[i].indptr->dtype.bits, csrs[i].indptr->ctx));
|
||||
data_data.push_back(data[i].Ptr<IdType>());
|
||||
indptr_data.push_back(csrs[i].indptr.Ptr<IdType>());
|
||||
indices_data.push_back(csrs[i].indices.Ptr<IdType>());
|
||||
num_edges += csrs[i].indices->shape[0];
|
||||
sorted &= csrs[i].sorted;
|
||||
}
|
||||
|
||||
res_indptr.resize(csrs[0].num_rows + 1);
|
||||
res_indices.resize(num_edges);
|
||||
res_data.resize(num_edges);
|
||||
res_indptr[0] = 0;
|
||||
|
||||
if (sorted) { // all csrs are sorted
|
||||
#pragma omp for
|
||||
for (int64_t i = 1; i <= csrs[0].num_rows; ++i) {
|
||||
std::vector<int64_t> indices_off;
|
||||
res_indptr[i] = indptr_data[0][i];
|
||||
|
||||
indices_off.push_back(indptr_data[0][i - 1]);
|
||||
for (size_t j = 1; j < csrs.size(); ++j) {
|
||||
res_indptr[i] += indptr_data[j][i];
|
||||
indices_off.push_back(indptr_data[j][i - 1]);
|
||||
}
|
||||
|
||||
IdType off = res_indptr[i - 1];
|
||||
while (off < res_indptr[i]) {
|
||||
IdType min = csrs[0].num_cols + 1;
|
||||
int64_t min_idx = -1;
|
||||
for (size_t j = 0; j < csrs.size(); ++j) {
|
||||
if (indices_off[j] < indptr_data[j][i]) {
|
||||
if (min <= indices_data[j][indices_off[j]]) {
|
||||
continue;
|
||||
} else {
|
||||
min = indices_data[j][indices_off[j]];
|
||||
min_idx = j;
|
||||
}
|
||||
} // for check out of bound
|
||||
} // for
|
||||
res_indices[off] = min;
|
||||
res_data[off] = data_data[min_idx][indices_off[min_idx]];
|
||||
indices_off[min_idx] += 1;
|
||||
++off;
|
||||
} // while
|
||||
} // omp for
|
||||
} else { // some csrs are not sorted
|
||||
#pragma omp for
|
||||
for (int64_t i = 1; i <= csrs[0].num_rows; ++i) {
|
||||
IdType off = res_indptr[i - 1];
|
||||
res_indptr[i] = 0;
|
||||
|
||||
for (size_t j = 0; j < csrs.size(); ++j) {
|
||||
std::memcpy(
|
||||
&res_indices[off], &indices_data[j][indptr_data[j][i - 1]],
|
||||
sizeof(IdType) * (indptr_data[j][i] - indptr_data[j][i - 1]));
|
||||
std::memcpy(
|
||||
&res_data[off], &data_data[j][indptr_data[j][i - 1]],
|
||||
sizeof(IdType) * (indptr_data[j][i] - indptr_data[j][i - 1]));
|
||||
off += indptr_data[j][i] - indptr_data[j][i - 1];
|
||||
}
|
||||
res_indptr[i] = off;
|
||||
} // omp for
|
||||
}
|
||||
|
||||
return CSRMatrix(
|
||||
csrs[0].num_rows, csrs[0].num_cols, IdArray::FromVector(res_indptr),
|
||||
IdArray::FromVector(res_indices), IdArray::FromVector(res_data), sorted);
|
||||
}
|
||||
|
||||
template CSRMatrix UnionCsr<kDGLCPU, int64_t>(const std::vector<CSRMatrix> &);
|
||||
template CSRMatrix UnionCsr<kDGLCPU, int32_t>(const std::vector<CSRMatrix> &);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,130 @@
|
||||
/**
|
||||
* Copyright (c) 2022, NVIDIA CORPORATION.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
* @file array/cpu/disjoint_union.cc
|
||||
* @brief Disjoint union CPU implementation.
|
||||
*/
|
||||
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
|
||||
#include <tuple>
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::tuple<IdArray, IdArray, IdArray> _ComputePrefixSums(
|
||||
const std::vector<COOMatrix>& coos) {
|
||||
IdArray prefix_src_arr =
|
||||
NewIdArray(coos.size(), coos[0].row->ctx, coos[0].row->dtype.bits);
|
||||
IdArray prefix_dst_arr =
|
||||
NewIdArray(coos.size(), coos[0].row->ctx, coos[0].row->dtype.bits);
|
||||
IdArray prefix_elm_arr =
|
||||
NewIdArray(coos.size(), coos[0].row->ctx, coos[0].row->dtype.bits);
|
||||
|
||||
auto prefix_src = prefix_src_arr.Ptr<IdType>();
|
||||
auto prefix_dst = prefix_dst_arr.Ptr<IdType>();
|
||||
auto prefix_elm = prefix_elm_arr.Ptr<IdType>();
|
||||
|
||||
dgl::runtime::parallel_for(0, coos.size(), [&](IdType b, IdType e) {
|
||||
for (IdType i = b; i < e; ++i) {
|
||||
prefix_src[i] = coos[i].num_rows;
|
||||
prefix_dst[i] = coos[i].num_cols;
|
||||
prefix_elm[i] = coos[i].row->shape[0];
|
||||
}
|
||||
});
|
||||
|
||||
return std::make_tuple(
|
||||
CumSum(prefix_src_arr, true), CumSum(prefix_dst_arr, true),
|
||||
CumSum(prefix_elm_arr, true));
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix DisjointUnionCoo(const std::vector<COOMatrix>& coos) {
|
||||
bool has_data = false;
|
||||
bool row_sorted = true;
|
||||
bool col_sorted = true;
|
||||
// check if data index array
|
||||
for (size_t i = 0; i < coos.size(); ++i) {
|
||||
CHECK_SAME_DTYPE(coos[0].row, coos[i].row);
|
||||
CHECK_SAME_CONTEXT(coos[0].row, coos[i].row);
|
||||
has_data |= COOHasData(coos[i]);
|
||||
}
|
||||
|
||||
auto prefixes = _ComputePrefixSums<XPU, IdType>(coos);
|
||||
auto prefix_src = static_cast<IdArray>(std::get<0>(prefixes)).Ptr<IdType>();
|
||||
auto prefix_dst = static_cast<IdArray>(std::get<1>(prefixes)).Ptr<IdType>();
|
||||
auto prefix_elm = static_cast<IdArray>(std::get<2>(prefixes)).Ptr<IdType>();
|
||||
|
||||
IdArray result_src = NewIdArray(
|
||||
prefix_elm[coos.size()], coos[0].row->ctx, coos[0].row->dtype.bits);
|
||||
IdArray result_dst = NewIdArray(
|
||||
prefix_elm[coos.size()], coos[0].col->ctx, coos[0].col->dtype.bits);
|
||||
IdArray result_dat = NullArray();
|
||||
if (has_data) {
|
||||
result_dat = NewIdArray(
|
||||
prefix_elm[coos.size()], coos[0].row->ctx, coos[0].row->dtype.bits);
|
||||
}
|
||||
|
||||
auto res_src_data = result_src.Ptr<IdType>();
|
||||
auto res_dst_data = result_dst.Ptr<IdType>();
|
||||
auto res_dat_data = result_dat.Ptr<IdType>();
|
||||
|
||||
// 32 is a number obtained from experience. If a user set the grain size
|
||||
// explicitly via env, use that value instead.
|
||||
size_t grain_size = dgl::runtime::DefaultGrainSizeT(32)();
|
||||
dgl::runtime::parallel_for(
|
||||
0, coos.size(), grain_size, [&](IdType b, IdType e) {
|
||||
for (IdType i = b; i < e; ++i) {
|
||||
const aten::COOMatrix& coo = coos[i];
|
||||
if (!coo.row_sorted) row_sorted = false;
|
||||
if (!coo.col_sorted) col_sorted = false;
|
||||
|
||||
auto edges_src = coo.row.Ptr<IdType>();
|
||||
auto edges_dst = coo.col.Ptr<IdType>();
|
||||
auto edges_dat = coo.data.Ptr<IdType>();
|
||||
|
||||
for (IdType j = 0; j < coo.row->shape[0]; j++) {
|
||||
res_src_data[prefix_elm[i] + j] = edges_src[j] + prefix_src[i];
|
||||
}
|
||||
|
||||
for (IdType j = 0; j < coo.row->shape[0]; j++) {
|
||||
res_dst_data[prefix_elm[i] + j] = edges_dst[j] + prefix_dst[i];
|
||||
}
|
||||
|
||||
if (has_data) {
|
||||
for (IdType j = 0; j < coo.row->shape[0]; j++) {
|
||||
const auto d = (!COOHasData(coo)) ? j : edges_dat[j];
|
||||
res_dat_data[prefix_elm[i] + j] = d + prefix_elm[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
return COOMatrix(
|
||||
prefix_src[coos.size()], prefix_dst[coos.size()], result_src, result_dst,
|
||||
result_dat, row_sorted, col_sorted);
|
||||
}
|
||||
|
||||
template COOMatrix DisjointUnionCoo<kDGLCPU, int32_t>(
|
||||
const std::vector<COOMatrix>& coos);
|
||||
template COOMatrix DisjointUnionCoo<kDGLCPU, int64_t>(
|
||||
const std::vector<COOMatrix>& coos);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,114 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file kernel/cpu/gaher_mm.cc
|
||||
* @brief GatherMM C APIs and definitions.
|
||||
*/
|
||||
#include "./gather_mm.h"
|
||||
|
||||
#include <dgl/array.h>
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
|
||||
/** @brief Generalized SegmentMM. */
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SegmentMM(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans) {
|
||||
LOG(FATAL) << "Unsupported CPU kernel for SegmentMM.";
|
||||
}
|
||||
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SegmentMMBackwardB(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen) {
|
||||
LOG(FATAL) << "Unsupported CPU kernel for SegmentMMBackwardB.";
|
||||
}
|
||||
|
||||
/** @brief Generalized GatherMM. */
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void GatherMM(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b) {
|
||||
LOG(FATAL) << "Unsupported CPU kernel for GatherMM.";
|
||||
}
|
||||
|
||||
/** @brief Generalized GatherMM_scatter. */
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void GatherMMScatter(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c) {
|
||||
LOG(FATAL) << "Unsupported CPU kernel for GatherMM.";
|
||||
}
|
||||
|
||||
template void GatherMM<kDGLCPU, int32_t, BFloat16>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b);
|
||||
template void GatherMM<kDGLCPU, int64_t, BFloat16>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b);
|
||||
template void GatherMM<kDGLCPU, int32_t, float>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b);
|
||||
template void GatherMM<kDGLCPU, int64_t, float>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b);
|
||||
template void GatherMM<kDGLCPU, int32_t, double>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b);
|
||||
template void GatherMM<kDGLCPU, int64_t, double>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b);
|
||||
|
||||
template void GatherMMScatter<kDGLCPU, int32_t, BFloat16>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c);
|
||||
template void GatherMMScatter<kDGLCPU, int64_t, BFloat16>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c);
|
||||
template void GatherMMScatter<kDGLCPU, int32_t, float>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c);
|
||||
template void GatherMMScatter<kDGLCPU, int64_t, float>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c);
|
||||
template void GatherMMScatter<kDGLCPU, int32_t, double>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c);
|
||||
template void GatherMMScatter<kDGLCPU, int64_t, double>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c);
|
||||
|
||||
template void SegmentMM<kDGLCPU, int32_t, BFloat16>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans);
|
||||
template void SegmentMM<kDGLCPU, int64_t, BFloat16>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans);
|
||||
template void SegmentMM<kDGLCPU, int32_t, float>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans);
|
||||
template void SegmentMM<kDGLCPU, int64_t, float>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans);
|
||||
template void SegmentMM<kDGLCPU, int32_t, double>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans);
|
||||
template void SegmentMM<kDGLCPU, int64_t, double>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans);
|
||||
|
||||
template void SegmentMMBackwardB<kDGLCPU, int32_t, BFloat16>(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen);
|
||||
template void SegmentMMBackwardB<kDGLCPU, int64_t, BFloat16>(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen);
|
||||
template void SegmentMMBackwardB<kDGLCPU, int32_t, float>(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen);
|
||||
template void SegmentMMBackwardB<kDGLCPU, int64_t, float>(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen);
|
||||
template void SegmentMMBackwardB<kDGLCPU, int32_t, double>(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen);
|
||||
template void SegmentMMBackwardB<kDGLCPU, int64_t, double>(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen);
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,115 @@
|
||||
/**
|
||||
* Copyright (c) 2022 by Contributors
|
||||
* @file array/cpu/gather_mm.h
|
||||
* @brief GATHER_MM CPU kernel function header.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CPU_GATHER_MM_H_
|
||||
#define DGL_ARRAY_CPU_GATHER_MM_H_
|
||||
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/bcast.h>
|
||||
|
||||
#include <utility>
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace cpu {
|
||||
|
||||
template <typename DType>
|
||||
void transpose(const DType *in, DType *out, const int N, const int M) {
|
||||
#pragma omp parallel for
|
||||
for (int n = 0; n < N * M; n++) {
|
||||
int i = n / N;
|
||||
int j = n % N;
|
||||
out[n] = in[M * j + i];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename DType>
|
||||
void matmul(
|
||||
const DType *A, const DType *B, DType *C, const int M, const int N,
|
||||
const int K) {
|
||||
#pragma omp parallel
|
||||
{
|
||||
int i, j, k;
|
||||
#pragma omp for
|
||||
for (i = 0; i < M; i++) {
|
||||
for (j = 0; j < N; j++) {
|
||||
DType local_accum = 0;
|
||||
for (k = 0; k < K; k++) {
|
||||
local_accum += A[i * K + k] * B[k * N + j];
|
||||
}
|
||||
C[i * N + j] = local_accum;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CPU kernel of Gather_mm. The input matrix A is expected to be
|
||||
* sorted according to relation type.
|
||||
* @param A The input dense matrix of dimension m x k
|
||||
* @param B The input dense matrix of dimension k x n
|
||||
* @param C The output dense matrix od dimension m x n
|
||||
* @param A_dim1_per_rel The number of rows in each relation in A
|
||||
* @param B_dim1_per_rel The number of rows in each relation in B
|
||||
* @param a_trans Matrix A to be transposed
|
||||
* @param b_trans Matrix B to be transposed
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void gatherMM_SortedEtype(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray A_dim1_per_rel,
|
||||
const NDArray B_dim1_per_rel, bool a_trans, bool b_trans) {
|
||||
assert(A_dim1_per_rel.NumElements() == B_dim1_per_rel.NumElements());
|
||||
int64_t num_rel = A_dim1_per_rel.NumElements();
|
||||
const DType *A_data = A.Ptr<DType>();
|
||||
const DType *B_data = B.Ptr<DType>();
|
||||
const IdType *A_rel_data = A_dim1_per_rel.Ptr<IdType>();
|
||||
const IdType *B_rel_data = B_dim1_per_rel.Ptr<IdType>();
|
||||
DType *C_data = C.Ptr<DType>();
|
||||
|
||||
int64_t A_offset = 0, B_offset = 0, C_offset = 0;
|
||||
int64_t m, n, k, h_col, w_row;
|
||||
for (int etype = 0; etype < num_rel; ++etype) {
|
||||
assert(
|
||||
(a_trans) ? A_rel_data[etype]
|
||||
: A->shape[1] == (b_trans) ? B->shape[1]
|
||||
: B_rel_data[etype]);
|
||||
m = A_rel_data[etype]; // rows of A
|
||||
n = B->shape[1]; // cols of B
|
||||
k = B_rel_data[etype]; // rows of B == cols of A
|
||||
|
||||
NDArray A_trans, B_trans;
|
||||
if (a_trans) {
|
||||
A_trans = NDArray::Empty({m * k}, A->dtype, A->ctx);
|
||||
transpose<DType>(
|
||||
A_data + A_offset, static_cast<DType *>(A_trans->data), m, k);
|
||||
}
|
||||
if (b_trans) {
|
||||
B_trans = NDArray::Empty({k * n}, B->dtype, B->ctx);
|
||||
transpose<DType>(
|
||||
B_data + B_offset, static_cast<DType *>(B_trans->data), k, n);
|
||||
}
|
||||
if (a_trans || b_trans) {
|
||||
int64_t tmp = k;
|
||||
if (a_trans) std::swap(m, k);
|
||||
if (b_trans) {
|
||||
k = tmp;
|
||||
std::swap(n, k);
|
||||
}
|
||||
}
|
||||
matmul<DType>(
|
||||
(a_trans) ? static_cast<DType *>(A_trans->data) : A_data + A_offset,
|
||||
(b_trans) ? static_cast<DType *>(B_trans->data) : B_data + B_offset,
|
||||
C_data + C_offset, m, n, k);
|
||||
A_offset += m * k;
|
||||
B_offset += k * n;
|
||||
C_offset += m * n;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace cpu
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CPU_GATHER_MM_H_
|
||||
@@ -0,0 +1,320 @@
|
||||
/**
|
||||
* Copyright (c) 2022, NVIDIA Corporation
|
||||
* Copyright (c) 2022, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
|
||||
* All rights reserved.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
* @file array/cpu/labor_pick.h
|
||||
* @brief Template implementation for layerwise pick operators.
|
||||
*/
|
||||
|
||||
#ifndef DGL_ARRAY_CPU_LABOR_PICK_H_
|
||||
#define DGL_ARRAY_CPU_LABOR_PICK_H_
|
||||
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/random.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
#include <dmlc/omp.h>
|
||||
#include <tsl/robin_map.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <numeric>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "../../random/continuous_seed.h"
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
using dgl::random::continuous_seed;
|
||||
|
||||
template <typename K, typename V>
|
||||
using map_t = tsl::robin_map<K, V>;
|
||||
template <typename iterator>
|
||||
auto& mutable_value_ref(iterator it) {
|
||||
return it.value();
|
||||
}
|
||||
|
||||
constexpr double eps = 0.0001;
|
||||
|
||||
template <typename IdxType, typename FloatType>
|
||||
auto compute_importance_sampling_probabilities(
|
||||
DGLContext ctx, DGLDataType dtype, const IdxType max_degree,
|
||||
const IdxType num_rows, const int importance_sampling, const bool weighted,
|
||||
const IdxType* rows_data, const IdxType* indptr, const FloatType* A,
|
||||
const IdxType* indices, const IdxType num_picks, const FloatType* ds,
|
||||
FloatType* cs) {
|
||||
constexpr FloatType ONE = 1;
|
||||
// ps stands for \pi in arXiv:2210.13339
|
||||
FloatArray ps_array = NDArray::Empty({max_degree + 1}, dtype, ctx);
|
||||
FloatType* ps = ps_array.Ptr<FloatType>();
|
||||
|
||||
double prev_ex_nodes = max_degree * num_rows;
|
||||
|
||||
map_t<IdxType, FloatType> hop_map, hop_map2;
|
||||
for (int iters = 0; iters < importance_sampling || importance_sampling < 0;
|
||||
iters++) {
|
||||
// NOTE(mfbalin) When the graph is unweighted, the first c values in
|
||||
// the first iteration can be computed in O(1) as k / d where k is fanout
|
||||
// and d is the degree.
|
||||
|
||||
// If the graph is weighted, the first c values are computed in the inner
|
||||
// for loop instead. Therefore the importance_sampling argument should be
|
||||
// increased by one in the caller.
|
||||
|
||||
// The later iterations will have correct c values so the if block will be
|
||||
// executed.
|
||||
|
||||
if (!weighted || iters) {
|
||||
hop_map2.clear();
|
||||
for (int64_t i = 0; i < num_rows; ++i) {
|
||||
const FloatType c = cs[i];
|
||||
const IdxType rid = rows_data[i];
|
||||
for (auto j = indptr[rid]; j < indptr[rid + 1]; j++) {
|
||||
const auto ct = c * (weighted && iters == 1 ? A[j] : 1);
|
||||
auto itb = hop_map2.emplace(indices[j], ct);
|
||||
if (!itb.second) {
|
||||
mutable_value_ref(itb.first) = std::max(ct, itb.first->second);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (hop_map.empty())
|
||||
hop_map = std::move(hop_map2);
|
||||
else
|
||||
// Update the pi array according to Eq 18.
|
||||
for (auto it : hop_map2) hop_map[it.first] *= it.second;
|
||||
}
|
||||
|
||||
// Compute c_s according to Equation (15), (17) is slower because sorting is
|
||||
// required.
|
||||
for (int64_t i = 0; i < num_rows; ++i) {
|
||||
const IdxType rid = rows_data[i];
|
||||
const auto d = indptr[rid + 1] - indptr[rid];
|
||||
if (d == 0) continue;
|
||||
|
||||
const auto k = std::min(num_picks, d);
|
||||
|
||||
if (hop_map.empty()) { // weighted first iter, pi = A
|
||||
for (auto j = indptr[rid]; j < indptr[rid + 1]; j++)
|
||||
ps[j - indptr[rid]] = A[j];
|
||||
} else {
|
||||
for (auto j = indptr[rid]; j < indptr[rid + 1]; j++)
|
||||
ps[j - indptr[rid]] = hop_map[indices[j]];
|
||||
}
|
||||
|
||||
// stands for RHS of Equation (22) in arXiv:2210.13339 after moving the
|
||||
// other terms without c_s to RHS.
|
||||
double var_target = ds[i] * ds[i] / k;
|
||||
if (weighted) {
|
||||
var_target -= ds[i] * ds[i] / d;
|
||||
for (auto j = indptr[rid]; j < indptr[rid + 1]; j++)
|
||||
var_target += A[j] * A[j];
|
||||
}
|
||||
FloatType c = cs[i];
|
||||
// stands for left handside of Equation (22) in arXiv:2210.13339 after
|
||||
// moving the other terms without c_s to RHS.
|
||||
double var_1;
|
||||
// Compute c_s in Equation (22) via fixed-point iteration.
|
||||
do {
|
||||
var_1 = 0;
|
||||
if (weighted) {
|
||||
for (auto j = indptr[rid]; j < indptr[rid + 1]; j++)
|
||||
// The check for zero is necessary for numerical stability
|
||||
var_1 += A[j] > 0
|
||||
? A[j] * A[j] / std::min(ONE, c * ps[j - indptr[rid]])
|
||||
: 0;
|
||||
} else {
|
||||
for (auto j = indptr[rid]; j < indptr[rid + 1]; j++)
|
||||
var_1 += ONE / std::min(ONE, c * ps[j - indptr[rid]]);
|
||||
}
|
||||
|
||||
c *= var_1 / var_target;
|
||||
} while (std::min(var_1, var_target) / std::max(var_1, var_target) <
|
||||
1 - eps);
|
||||
|
||||
cs[i] = c;
|
||||
}
|
||||
|
||||
// Check convergence
|
||||
if (!weighted || iters) {
|
||||
double cur_ex_nodes = 0;
|
||||
for (auto it : hop_map) cur_ex_nodes += std::min((FloatType)1, it.second);
|
||||
if (cur_ex_nodes / prev_ex_nodes >= 1 - eps) break;
|
||||
prev_ex_nodes = cur_ex_nodes;
|
||||
}
|
||||
}
|
||||
|
||||
return hop_map;
|
||||
}
|
||||
|
||||
// Template for picking non-zero values row-wise.
|
||||
template <typename IdxType, typename FloatType>
|
||||
std::pair<COOMatrix, FloatArray> CSRLaborPick(
|
||||
CSRMatrix mat, IdArray rows, int64_t num_picks, FloatArray prob,
|
||||
int importance_sampling, IdArray random_seed_arr, float seed2_contribution,
|
||||
IdArray NIDs) {
|
||||
using namespace aten;
|
||||
const IdxType* indptr = mat.indptr.Ptr<IdxType>();
|
||||
const IdxType* indices = mat.indices.Ptr<IdxType>();
|
||||
const IdxType* data = CSRHasData(mat) ? mat.data.Ptr<IdxType>() : nullptr;
|
||||
const IdxType* rows_data = rows.Ptr<IdxType>();
|
||||
const IdxType* nids = IsNullArray(NIDs) ? nullptr : NIDs.Ptr<IdxType>();
|
||||
const auto num_rows = rows->shape[0];
|
||||
const auto& ctx = mat.indptr->ctx;
|
||||
|
||||
const bool weighted = !IsNullArray(prob);
|
||||
// O(1) c computation not possible, so one more iteration is needed.
|
||||
if (importance_sampling >= 0) importance_sampling += weighted;
|
||||
// A stands for the same notation in arXiv:2210.13339, i.e. the edge weights.
|
||||
auto A_arr = prob;
|
||||
FloatType* A = A_arr.Ptr<FloatType>();
|
||||
constexpr FloatType ONE = 1;
|
||||
|
||||
constexpr auto dtype = DGLDataTypeTraits<FloatType>::dtype;
|
||||
|
||||
// cs stands for c_s in arXiv:2210.13339
|
||||
FloatArray cs_array = NDArray::Empty({num_rows}, dtype, ctx);
|
||||
FloatType* cs = cs_array.Ptr<FloatType>();
|
||||
// ds stands for A_{*s} in arXiv:2210.13339
|
||||
FloatArray ds_array = NDArray::Empty({num_rows}, dtype, ctx);
|
||||
FloatType* ds = ds_array.Ptr<FloatType>();
|
||||
|
||||
IdxType max_degree = 1;
|
||||
IdxType hop_size = 0;
|
||||
for (int64_t i = 0; i < num_rows; ++i) {
|
||||
const IdxType rid = rows_data[i];
|
||||
const auto act_degree = indptr[rid + 1] - indptr[rid];
|
||||
max_degree = std::max(act_degree, max_degree);
|
||||
double d = weighted
|
||||
? std::accumulate(A + indptr[rid], A + indptr[rid + 1], 0.0)
|
||||
: act_degree;
|
||||
// O(1) c computation, samples more than needed for weighted case, mentioned
|
||||
// in the sentence between (10) and (11) in arXiv:2210.13339
|
||||
cs[i] = num_picks / d;
|
||||
ds[i] = d;
|
||||
hop_size += act_degree;
|
||||
}
|
||||
|
||||
map_t<IdxType, FloatType> hop_map;
|
||||
|
||||
if (importance_sampling)
|
||||
hop_map = compute_importance_sampling_probabilities<IdxType, FloatType>(
|
||||
ctx, dtype, max_degree, num_rows, importance_sampling, weighted,
|
||||
rows_data, indptr, A, indices, (IdxType)num_picks, ds, cs);
|
||||
|
||||
constexpr auto vidtype = DGLDataTypeTraits<IdxType>::dtype;
|
||||
|
||||
IdArray picked_row = NDArray::Empty({hop_size}, vidtype, ctx);
|
||||
IdArray picked_col = NDArray::Empty({hop_size}, vidtype, ctx);
|
||||
IdArray picked_idx = NDArray::Empty({hop_size}, vidtype, ctx);
|
||||
FloatArray picked_imp = importance_sampling
|
||||
? NDArray::Empty({hop_size}, dtype, ctx)
|
||||
: NullArray();
|
||||
IdxType* picked_rdata = picked_row.Ptr<IdxType>();
|
||||
IdxType* picked_cdata = picked_col.Ptr<IdxType>();
|
||||
IdxType* picked_idata = picked_idx.Ptr<IdxType>();
|
||||
FloatType* picked_imp_data = picked_imp.Ptr<FloatType>();
|
||||
|
||||
const continuous_seed random_seed =
|
||||
IsNullArray(random_seed_arr)
|
||||
? continuous_seed(RandomEngine::ThreadLocal()->RandInt(1000000000))
|
||||
: continuous_seed(random_seed_arr, seed2_contribution);
|
||||
|
||||
// compute number of edges first and do sampling
|
||||
IdxType num_edges = 0;
|
||||
for (int64_t i = 0; i < num_rows; i++) {
|
||||
const IdxType rid = rows_data[i];
|
||||
const auto c = cs[i];
|
||||
|
||||
FloatType norm_inv_p = 0;
|
||||
const auto off = num_edges;
|
||||
for (auto j = indptr[rid]; j < indptr[rid + 1]; j++) {
|
||||
const auto v = indices[j];
|
||||
const uint64_t t = nids ? nids[v] : v; // t in the paper
|
||||
// rolled random number r_t is a function of the random_seed and t
|
||||
const auto rnd = random_seed.uniform(t);
|
||||
const auto w = (weighted ? A[j] : 1);
|
||||
// if hop_map is initialized, get ps from there, otherwise get it from the
|
||||
// alternative.
|
||||
const auto ps = std::min(
|
||||
ONE, importance_sampling - weighted ? c * hop_map[v] : c * w);
|
||||
if (rnd <= ps) {
|
||||
picked_rdata[num_edges] = rid;
|
||||
picked_cdata[num_edges] = v;
|
||||
picked_idata[num_edges] = data ? data[j] : j;
|
||||
if (importance_sampling) {
|
||||
const auto edge_weight = w / ps;
|
||||
norm_inv_p += edge_weight;
|
||||
picked_imp_data[num_edges] = edge_weight;
|
||||
}
|
||||
num_edges++;
|
||||
}
|
||||
}
|
||||
|
||||
if (importance_sampling) {
|
||||
const auto norm_factor = (num_edges - off) / norm_inv_p;
|
||||
for (auto i = off; i < num_edges; i++)
|
||||
// so that fn.mean can be used
|
||||
picked_imp_data[i] *= norm_factor;
|
||||
}
|
||||
}
|
||||
|
||||
picked_row = picked_row.CreateView({num_edges}, picked_row->dtype);
|
||||
picked_col = picked_col.CreateView({num_edges}, picked_col->dtype);
|
||||
picked_idx = picked_idx.CreateView({num_edges}, picked_idx->dtype);
|
||||
if (importance_sampling)
|
||||
picked_imp = picked_imp.CreateView({num_edges}, picked_imp->dtype);
|
||||
|
||||
return std::make_pair(
|
||||
COOMatrix(mat.num_rows, mat.num_cols, picked_row, picked_col, picked_idx),
|
||||
picked_imp);
|
||||
}
|
||||
|
||||
// Template for picking non-zero values row-wise. The implementation first
|
||||
// slices out the corresponding rows and then converts it to CSR format. It then
|
||||
// performs row-wise pick on the CSR matrix and rectifies the returned results.
|
||||
template <typename IdxType, typename FloatType>
|
||||
std::pair<COOMatrix, FloatArray> COOLaborPick(
|
||||
COOMatrix mat, IdArray rows, int64_t num_picks, FloatArray prob,
|
||||
int importance_sampling, IdArray random_seed, float seed2_contribution,
|
||||
IdArray NIDs) {
|
||||
using namespace aten;
|
||||
const auto& csr = COOToCSR(COOSliceRows(mat, rows));
|
||||
const IdArray new_rows =
|
||||
Range(0, rows->shape[0], rows->dtype.bits, rows->ctx);
|
||||
const auto&& picked_importances = CSRLaborPick<IdxType, FloatType>(
|
||||
csr, new_rows, num_picks, prob, importance_sampling, random_seed,
|
||||
seed2_contribution, NIDs);
|
||||
const auto& picked = picked_importances.first;
|
||||
const auto& importances = picked_importances.second;
|
||||
return std::make_pair(
|
||||
COOMatrix(
|
||||
mat.num_rows, mat.num_cols,
|
||||
IndexSelect(
|
||||
rows, picked.row), // map the row index to the correct one
|
||||
picked.col, picked.data),
|
||||
importances);
|
||||
}
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CPU_LABOR_PICK_H_
|
||||
@@ -0,0 +1,79 @@
|
||||
/*!
|
||||
* Copyright (c) 2022, NVIDIA Corporation
|
||||
* Copyright (c) 2022, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
|
||||
* All rights reserved.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
* \file array/cuda/labor_sampling.cc
|
||||
* \brief labor sampling
|
||||
*/
|
||||
#include "./labor_pick.h"
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
/////////////////////////////// CSR ///////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdxType, typename FloatType>
|
||||
std::pair<COOMatrix, FloatArray> CSRLaborSampling(
|
||||
CSRMatrix mat, IdArray rows, int64_t num_samples, FloatArray prob,
|
||||
int importance_sampling, IdArray random_seed, float seed2_contribution,
|
||||
IdArray NIDs) {
|
||||
return CSRLaborPick<IdxType, FloatType>(
|
||||
mat, rows, num_samples, prob, importance_sampling, random_seed,
|
||||
seed2_contribution, NIDs);
|
||||
}
|
||||
|
||||
template std::pair<COOMatrix, FloatArray>
|
||||
CSRLaborSampling<kDGLCPU, int32_t, float>(
|
||||
CSRMatrix, IdArray, int64_t, FloatArray, int, IdArray, float, IdArray);
|
||||
template std::pair<COOMatrix, FloatArray>
|
||||
CSRLaborSampling<kDGLCPU, int64_t, float>(
|
||||
CSRMatrix, IdArray, int64_t, FloatArray, int, IdArray, float, IdArray);
|
||||
template std::pair<COOMatrix, FloatArray>
|
||||
CSRLaborSampling<kDGLCPU, int32_t, double>(
|
||||
CSRMatrix, IdArray, int64_t, FloatArray, int, IdArray, float, IdArray);
|
||||
template std::pair<COOMatrix, FloatArray>
|
||||
CSRLaborSampling<kDGLCPU, int64_t, double>(
|
||||
CSRMatrix, IdArray, int64_t, FloatArray, int, IdArray, float, IdArray);
|
||||
|
||||
/////////////////////////////// COO ///////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdxType, typename FloatType>
|
||||
std::pair<COOMatrix, FloatArray> COOLaborSampling(
|
||||
COOMatrix mat, IdArray rows, int64_t num_samples, FloatArray prob,
|
||||
int importance_sampling, IdArray random_seed, float seed2_contribution,
|
||||
IdArray NIDs) {
|
||||
return COOLaborPick<IdxType, FloatType>(
|
||||
mat, rows, num_samples, prob, importance_sampling, random_seed,
|
||||
seed2_contribution, NIDs);
|
||||
}
|
||||
|
||||
template std::pair<COOMatrix, FloatArray>
|
||||
COOLaborSampling<kDGLCPU, int32_t, float>(
|
||||
COOMatrix, IdArray, int64_t, FloatArray, int, IdArray, float, IdArray);
|
||||
template std::pair<COOMatrix, FloatArray>
|
||||
COOLaborSampling<kDGLCPU, int64_t, float>(
|
||||
COOMatrix, IdArray, int64_t, FloatArray, int, IdArray, float, IdArray);
|
||||
template std::pair<COOMatrix, FloatArray>
|
||||
COOLaborSampling<kDGLCPU, int32_t, double>(
|
||||
COOMatrix, IdArray, int64_t, FloatArray, int, IdArray, float, IdArray);
|
||||
template std::pair<COOMatrix, FloatArray>
|
||||
COOLaborSampling<kDGLCPU, int64_t, double>(
|
||||
COOMatrix, IdArray, int64_t, FloatArray, int, IdArray, float, IdArray);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,76 @@
|
||||
/**
|
||||
* Copyright (c) 2021 by Contributors
|
||||
* @file array/cpu/negative_sampling.cc
|
||||
* @brief Uniform negative sampling on CSR.
|
||||
*/
|
||||
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/array_iterator.h>
|
||||
#include <dgl/random.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <utility>
|
||||
|
||||
using namespace dgl::runtime;
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::pair<IdArray, IdArray> CSRGlobalUniformNegativeSampling(
|
||||
const CSRMatrix& csr, int64_t num_samples, int num_trials,
|
||||
bool exclude_self_loops, bool replace, double redundancy) {
|
||||
const int64_t num_row = csr.num_rows;
|
||||
const int64_t num_col = csr.num_cols;
|
||||
const int64_t num_actual_samples =
|
||||
static_cast<int64_t>(num_samples * (1 + redundancy));
|
||||
IdArray row = Full<IdType>(-1, num_actual_samples, csr.indptr->ctx);
|
||||
IdArray col = Full<IdType>(-1, num_actual_samples, csr.indptr->ctx);
|
||||
IdType* row_data = row.Ptr<IdType>();
|
||||
IdType* col_data = col.Ptr<IdType>();
|
||||
|
||||
parallel_for(0, num_actual_samples, 1, [&](int64_t b, int64_t e) {
|
||||
for (int64_t i = b; i < e; ++i) {
|
||||
for (int trial = 0; trial < num_trials; ++trial) {
|
||||
IdType u = RandomEngine::ThreadLocal()->RandInt(num_row);
|
||||
IdType v = RandomEngine::ThreadLocal()->RandInt(num_col);
|
||||
if (!(exclude_self_loops && (u == v)) && !CSRIsNonZero(csr, u, v)) {
|
||||
row_data[i] = u;
|
||||
col_data[i] = v;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
PairIterator<IdType> begin(row_data, col_data);
|
||||
PairIterator<IdType> end = std::remove_if(
|
||||
begin, begin + num_actual_samples,
|
||||
[](const std::pair<IdType, IdType>& val) { return val.first == -1; });
|
||||
if (!replace) {
|
||||
std::sort(
|
||||
begin, end,
|
||||
[](const std::pair<IdType, IdType>& a,
|
||||
const std::pair<IdType, IdType>& b) {
|
||||
return a.first < b.first ||
|
||||
(a.first == b.first && a.second < b.second);
|
||||
});
|
||||
end = std::unique(begin, end);
|
||||
}
|
||||
int64_t num_sampled =
|
||||
std::min(static_cast<int64_t>(end - begin), num_samples);
|
||||
return {
|
||||
row.CreateView({num_sampled}, row->dtype),
|
||||
col.CreateView({num_sampled}, col->dtype)};
|
||||
}
|
||||
|
||||
template std::pair<IdArray, IdArray> CSRGlobalUniformNegativeSampling<
|
||||
kDGLCPU, int32_t>(const CSRMatrix&, int64_t, int, bool, bool, double);
|
||||
template std::pair<IdArray, IdArray> CSRGlobalUniformNegativeSampling<
|
||||
kDGLCPU, int64_t>(const CSRMatrix&, int64_t, int, bool, bool, double);
|
||||
|
||||
}; // namespace impl
|
||||
}; // namespace aten
|
||||
}; // namespace dgl
|
||||
@@ -0,0 +1,583 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/rowwise_pick.h
|
||||
* @brief Template implementation for rowwise pick operators.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CPU_ROWWISE_PICK_H_
|
||||
#define DGL_ARRAY_CPU_ROWWISE_PICK_H_
|
||||
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
#include <dmlc/omp.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
// User-defined function for picking elements from one row.
|
||||
//
|
||||
// The column indices of the given row are stored in
|
||||
// [col + off, col + off + len)
|
||||
//
|
||||
// Similarly, the data indices are stored in
|
||||
// [data + off, data + off + len)
|
||||
// Data index pointer could be NULL, which means data[i] == i
|
||||
//
|
||||
// *ATTENTION*: This function will be invoked concurrently. Please make sure
|
||||
// it is thread-safe.
|
||||
//
|
||||
// @param rowid The row to pick from.
|
||||
// @param off Starting offset of this row.
|
||||
// @param len NNZ of the row.
|
||||
// @param num_picks Number of picks on the row.
|
||||
// @param col Pointer of the column indices.
|
||||
// @param data Pointer of the data indices.
|
||||
// @param out_idx Picked indices in [off, off + len).
|
||||
template <typename IdxType>
|
||||
using PickFn = std::function<void(
|
||||
IdxType rowid, IdxType off, IdxType len, IdxType num_picks,
|
||||
const IdxType* col, const IdxType* data, IdxType* out_idx)>;
|
||||
|
||||
// User-defined function for determining the number of elements to pick from one
|
||||
// row.
|
||||
//
|
||||
// The column indices of the given row are stored in
|
||||
// [col + off, col + off + len)
|
||||
//
|
||||
// Similarly, the data indices are stored in
|
||||
// [data + off, data + off + len)
|
||||
// Data index pointer could be NULL, which means data[i] == i
|
||||
//
|
||||
// *ATTENTION*: This function will be invoked concurrently. Please make sure
|
||||
// it is thread-safe.
|
||||
//
|
||||
// @param rowid The row to pick from.
|
||||
// @param off Starting offset of this row.
|
||||
// @param len NNZ of the row.
|
||||
// @param col Pointer of the column indices.
|
||||
// @param data Pointer of the data indices.
|
||||
template <typename IdxType>
|
||||
using NumPicksFn = std::function<IdxType(
|
||||
IdxType rowid, IdxType off, IdxType len, const IdxType* col,
|
||||
const IdxType* data)>;
|
||||
|
||||
// User-defined function for picking elements from a range within a row.
|
||||
//
|
||||
// The column indices of each element is in
|
||||
// off + et_idx[et_offset+i]), where i is in [et_offset, et_offset+et_len)
|
||||
//
|
||||
// Similarly, the data indices are stored in
|
||||
// data[off+et_idx[et_offset+i])]
|
||||
// Data index pointer could be NULL, which means data[i] ==
|
||||
// off+et_idx[et_offset+i])
|
||||
//
|
||||
// *ATTENTION*: This function will be invoked concurrently. Please make sure
|
||||
// it is thread-safe.
|
||||
//
|
||||
// @param off Starting offset of this row.
|
||||
// @param et_offset Starting offset of this range.
|
||||
// @param cur_et The edge type.
|
||||
// @param et_len Length of the range.
|
||||
// @param et_idx A map from local idx to column id.
|
||||
// @param et_eid Edge-type-specific id array.
|
||||
// @param eid Pointer of the homogenized edge id array.
|
||||
// @param out_idx Picked indices in [et_offset, et_offset + et_len).
|
||||
template <typename IdxType>
|
||||
using EtypeRangePickFn = std::function<void(
|
||||
IdxType off, IdxType et_offset, IdxType cur_et, IdxType et_len,
|
||||
const std::vector<IdxType>& et_idx, const std::vector<IdxType>& et_eid,
|
||||
const IdxType* eid, IdxType* out_idx)>;
|
||||
|
||||
template <typename IdxType, bool map_seed_nodes>
|
||||
std::pair<CSRMatrix, IdArray> CSRRowWisePickFused(
|
||||
CSRMatrix mat, IdArray rows, IdArray seed_mapping,
|
||||
std::vector<IdxType>* new_seed_nodes, int64_t num_picks, bool replace,
|
||||
PickFn<IdxType> pick_fn, NumPicksFn<IdxType> num_picks_fn) {
|
||||
using namespace aten;
|
||||
|
||||
const IdxType* indptr = static_cast<IdxType*>(mat.indptr->data);
|
||||
const IdxType* indices = static_cast<IdxType*>(mat.indices->data);
|
||||
const IdxType* data =
|
||||
CSRHasData(mat) ? static_cast<IdxType*>(mat.data->data) : nullptr;
|
||||
const IdxType* rows_data = static_cast<IdxType*>(rows->data);
|
||||
const int64_t num_rows = rows->shape[0];
|
||||
const auto& ctx = mat.indptr->ctx;
|
||||
const auto& idtype = mat.indptr->dtype;
|
||||
IdxType* seed_mapping_data = nullptr;
|
||||
if (map_seed_nodes) seed_mapping_data = seed_mapping.Ptr<IdxType>();
|
||||
|
||||
const int num_threads = runtime::compute_num_threads(0, num_rows, 1);
|
||||
std::vector<int64_t> global_prefix(num_threads + 1, 0);
|
||||
|
||||
IdArray picked_col, picked_idx, picked_coo_rows;
|
||||
|
||||
IdArray block_csr_indptr = IdArray::Empty({num_rows + 1}, idtype, ctx);
|
||||
IdxType* block_csr_indptr_data = block_csr_indptr.Ptr<IdxType>();
|
||||
|
||||
#pragma omp parallel num_threads(num_threads)
|
||||
{
|
||||
const int thread_id = omp_get_thread_num();
|
||||
|
||||
const int64_t start_i =
|
||||
thread_id * (num_rows / num_threads) +
|
||||
std::min(static_cast<int64_t>(thread_id), num_rows % num_threads);
|
||||
const int64_t end_i =
|
||||
(thread_id + 1) * (num_rows / num_threads) +
|
||||
std::min(static_cast<int64_t>(thread_id + 1), num_rows % num_threads);
|
||||
assert(thread_id + 1 < num_threads || end_i == num_rows);
|
||||
|
||||
const int64_t num_local = end_i - start_i;
|
||||
|
||||
std::unique_ptr<int64_t[]> local_prefix(new int64_t[num_local + 1]);
|
||||
local_prefix[0] = 0;
|
||||
for (int64_t i = start_i; i < end_i; ++i) {
|
||||
// build prefix-sum
|
||||
const int64_t local_i = i - start_i;
|
||||
const IdxType rid = rows_data[i];
|
||||
if (map_seed_nodes) seed_mapping_data[rid] = i;
|
||||
|
||||
IdxType len = num_picks_fn(
|
||||
rid, indptr[rid], indptr[rid + 1] - indptr[rid], indices, data);
|
||||
local_prefix[local_i + 1] = local_prefix[local_i] + len;
|
||||
}
|
||||
global_prefix[thread_id + 1] = local_prefix[num_local];
|
||||
|
||||
#pragma omp barrier
|
||||
#pragma omp master
|
||||
{
|
||||
for (int t = 0; t < num_threads; ++t) {
|
||||
global_prefix[t + 1] += global_prefix[t];
|
||||
}
|
||||
picked_col = IdArray::Empty({global_prefix[num_threads]}, idtype, ctx);
|
||||
picked_idx = IdArray::Empty({global_prefix[num_threads]}, idtype, ctx);
|
||||
picked_coo_rows =
|
||||
IdArray::Empty({global_prefix[num_threads]}, idtype, ctx);
|
||||
}
|
||||
|
||||
#pragma omp barrier
|
||||
IdxType* picked_cdata = picked_col.Ptr<IdxType>();
|
||||
IdxType* picked_idata = picked_idx.Ptr<IdxType>();
|
||||
IdxType* picked_rows = picked_coo_rows.Ptr<IdxType>();
|
||||
|
||||
const IdxType thread_offset = global_prefix[thread_id];
|
||||
|
||||
for (int64_t i = start_i; i < end_i; ++i) {
|
||||
const IdxType rid = rows_data[i];
|
||||
const int64_t local_i = i - start_i;
|
||||
block_csr_indptr_data[i] = local_prefix[local_i] + thread_offset;
|
||||
|
||||
const IdxType off = indptr[rid];
|
||||
const IdxType len = indptr[rid + 1] - off;
|
||||
if (len == 0) continue;
|
||||
|
||||
const int64_t row_offset = local_prefix[local_i] + thread_offset;
|
||||
const int64_t num_picks =
|
||||
local_prefix[local_i + 1] + thread_offset - row_offset;
|
||||
|
||||
pick_fn(
|
||||
rid, off, len, num_picks, indices, data, picked_idata + row_offset);
|
||||
for (int64_t j = 0; j < num_picks; ++j) {
|
||||
const IdxType picked = picked_idata[row_offset + j];
|
||||
picked_cdata[row_offset + j] = indices[picked];
|
||||
picked_idata[row_offset + j] = data ? data[picked] : picked;
|
||||
picked_rows[row_offset + j] = i;
|
||||
}
|
||||
}
|
||||
}
|
||||
block_csr_indptr_data[num_rows] = global_prefix.back();
|
||||
|
||||
const IdxType num_cols = picked_col->shape[0];
|
||||
if (map_seed_nodes) {
|
||||
(*new_seed_nodes).resize(num_rows);
|
||||
memcpy((*new_seed_nodes).data(), rows_data, sizeof(IdxType) * num_rows);
|
||||
}
|
||||
|
||||
return std::make_pair(
|
||||
CSRMatrix(num_rows, num_cols, block_csr_indptr, picked_col, picked_idx),
|
||||
picked_coo_rows);
|
||||
}
|
||||
|
||||
// Template for picking non-zero values row-wise. The implementation utilizes
|
||||
// OpenMP parallelization on rows because each row performs computation
|
||||
// independently.
|
||||
template <typename IdxType>
|
||||
COOMatrix CSRRowWisePick(
|
||||
CSRMatrix mat, IdArray rows, int64_t num_picks, bool replace,
|
||||
PickFn<IdxType> pick_fn, NumPicksFn<IdxType> num_picks_fn) {
|
||||
using namespace aten;
|
||||
const IdxType* indptr = static_cast<IdxType*>(mat.indptr->data);
|
||||
const IdxType* indices = static_cast<IdxType*>(mat.indices->data);
|
||||
const IdxType* data =
|
||||
CSRHasData(mat) ? static_cast<IdxType*>(mat.data->data) : nullptr;
|
||||
const IdxType* rows_data = static_cast<IdxType*>(rows->data);
|
||||
const int64_t num_rows = rows->shape[0];
|
||||
const auto& ctx = mat.indptr->ctx;
|
||||
const auto& idtype = mat.indptr->dtype;
|
||||
|
||||
// To leverage OMP parallelization, we create two arrays to store
|
||||
// picked src and dst indices. Each array is of length num_rows * num_picks.
|
||||
// For rows whose nnz < num_picks, the indices are padded with -1.
|
||||
//
|
||||
// We check whether all the given rows
|
||||
// have at least num_picks number of nnz when replace is false.
|
||||
//
|
||||
// If the check holds, remove -1 elements by remove_if operation, which simply
|
||||
// moves valid elements to the head of arrays and create a view of the
|
||||
// original array. The implementation consumes a little extra memory than the
|
||||
// actual requirement.
|
||||
//
|
||||
// Otherwise, directly use the row and col arrays to construct the result COO
|
||||
// matrix.
|
||||
//
|
||||
// [02/29/2020 update]: OMP is disabled for now since batch-wise parallelism
|
||||
// is more
|
||||
// significant. (minjie)
|
||||
|
||||
// Do not use omp_get_max_threads() since that doesn't work for compiling
|
||||
// without OpenMP.
|
||||
const int num_threads = runtime::compute_num_threads(0, num_rows, 1);
|
||||
std::vector<int64_t> global_prefix(num_threads + 1, 0);
|
||||
|
||||
// TODO(BarclayII) Using OMP parallel directly instead of using
|
||||
// runtime::parallel_for does not handle exceptions well (directly aborts when
|
||||
// an exception pops up). It runs faster though because there is less
|
||||
// scheduling. Need to handle exceptions better.
|
||||
IdArray picked_row, picked_col, picked_idx;
|
||||
#pragma omp parallel num_threads(num_threads)
|
||||
{
|
||||
const int thread_id = omp_get_thread_num();
|
||||
|
||||
const int64_t start_i =
|
||||
thread_id * (num_rows / num_threads) +
|
||||
std::min(static_cast<int64_t>(thread_id), num_rows % num_threads);
|
||||
const int64_t end_i =
|
||||
(thread_id + 1) * (num_rows / num_threads) +
|
||||
std::min(static_cast<int64_t>(thread_id + 1), num_rows % num_threads);
|
||||
assert(thread_id + 1 < num_threads || end_i == num_rows);
|
||||
|
||||
const int64_t num_local = end_i - start_i;
|
||||
|
||||
// make sure we don't have to pay initialization cost
|
||||
std::unique_ptr<int64_t[]> local_prefix(new int64_t[num_local + 1]);
|
||||
local_prefix[0] = 0;
|
||||
for (int64_t i = start_i; i < end_i; ++i) {
|
||||
// build prefix-sum
|
||||
const int64_t local_i = i - start_i;
|
||||
const IdxType rid = rows_data[i];
|
||||
IdxType len = num_picks_fn(
|
||||
rid, indptr[rid], indptr[rid + 1] - indptr[rid], indices, data);
|
||||
local_prefix[local_i + 1] = local_prefix[local_i] + len;
|
||||
}
|
||||
global_prefix[thread_id + 1] = local_prefix[num_local];
|
||||
|
||||
#pragma omp barrier
|
||||
#pragma omp master
|
||||
{
|
||||
for (int t = 0; t < num_threads; ++t) {
|
||||
global_prefix[t + 1] += global_prefix[t];
|
||||
}
|
||||
picked_row = IdArray::Empty({global_prefix[num_threads]}, idtype, ctx);
|
||||
picked_col = IdArray::Empty({global_prefix[num_threads]}, idtype, ctx);
|
||||
picked_idx = IdArray::Empty({global_prefix[num_threads]}, idtype, ctx);
|
||||
}
|
||||
|
||||
#pragma omp barrier
|
||||
IdxType* picked_rdata = picked_row.Ptr<IdxType>();
|
||||
IdxType* picked_cdata = picked_col.Ptr<IdxType>();
|
||||
IdxType* picked_idata = picked_idx.Ptr<IdxType>();
|
||||
|
||||
const IdxType thread_offset = global_prefix[thread_id];
|
||||
|
||||
for (int64_t i = start_i; i < end_i; ++i) {
|
||||
const IdxType rid = rows_data[i];
|
||||
|
||||
const IdxType off = indptr[rid];
|
||||
const IdxType len = indptr[rid + 1] - off;
|
||||
if (len == 0) continue;
|
||||
|
||||
const int64_t local_i = i - start_i;
|
||||
const int64_t row_offset = thread_offset + local_prefix[local_i];
|
||||
const int64_t num_picks =
|
||||
thread_offset + local_prefix[local_i + 1] - row_offset;
|
||||
|
||||
pick_fn(
|
||||
rid, off, len, num_picks, indices, data, picked_idata + row_offset);
|
||||
for (int64_t j = 0; j < num_picks; ++j) {
|
||||
const IdxType picked = picked_idata[row_offset + j];
|
||||
picked_rdata[row_offset + j] = rid;
|
||||
picked_cdata[row_offset + j] = indices[picked];
|
||||
picked_idata[row_offset + j] = data ? data[picked] : picked;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const int64_t new_len = global_prefix.back();
|
||||
|
||||
return COOMatrix(
|
||||
mat.num_rows, mat.num_cols,
|
||||
picked_row.CreateView({new_len}, picked_row->dtype),
|
||||
picked_col.CreateView({new_len}, picked_row->dtype),
|
||||
picked_idx.CreateView({new_len}, picked_row->dtype));
|
||||
}
|
||||
|
||||
// Template for picking non-zero values row-wise. The implementation utilizes
|
||||
// OpenMP parallelization on rows because each row performs computation
|
||||
// independently.
|
||||
template <typename IdxType, typename DType>
|
||||
COOMatrix CSRRowWisePerEtypePick(
|
||||
CSRMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
|
||||
const std::vector<int64_t>& num_picks, bool replace,
|
||||
bool rowwise_etype_sorted, EtypeRangePickFn<IdxType> pick_fn,
|
||||
const std::vector<NDArray>& prob_or_mask) {
|
||||
using namespace aten;
|
||||
const IdxType* indptr = mat.indptr.Ptr<IdxType>();
|
||||
const IdxType* indices = mat.indices.Ptr<IdxType>();
|
||||
const IdxType* eid = CSRHasData(mat) ? mat.data.Ptr<IdxType>() : nullptr;
|
||||
const IdxType* rows_data = rows.Ptr<IdxType>();
|
||||
const int64_t num_rows = rows->shape[0];
|
||||
const auto& ctx = mat.indptr->ctx;
|
||||
const int64_t num_etypes = num_picks.size();
|
||||
const bool has_probs = (prob_or_mask.size() > 0);
|
||||
std::vector<IdArray> picked_rows(rows->shape[0]);
|
||||
std::vector<IdArray> picked_cols(rows->shape[0]);
|
||||
std::vector<IdArray> picked_idxs(rows->shape[0]);
|
||||
|
||||
// Check if the number of picks have the same value.
|
||||
// If so, we can potentially speed up if we have a node with total number of
|
||||
// neighbors less than the given number of picks with replace=False.
|
||||
bool same_num_pick = true;
|
||||
int64_t num_pick_value = num_picks[0];
|
||||
for (int64_t num_pick : num_picks) {
|
||||
if (num_pick_value != num_pick) {
|
||||
same_num_pick = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
runtime::parallel_for(0, num_rows, [&](size_t b, size_t e) {
|
||||
for (size_t i = b; i < e; ++i) {
|
||||
const IdxType rid = rows_data[i];
|
||||
CHECK_LT(rid, mat.num_rows);
|
||||
const IdxType off = indptr[rid];
|
||||
const IdxType len = indptr[rid + 1] - off;
|
||||
|
||||
// do something here
|
||||
if (len == 0) {
|
||||
picked_rows[i] = NewIdArray(0, ctx, sizeof(IdxType) * 8);
|
||||
picked_cols[i] = NewIdArray(0, ctx, sizeof(IdxType) * 8);
|
||||
picked_idxs[i] = NewIdArray(0, ctx, sizeof(IdxType) * 8);
|
||||
continue;
|
||||
}
|
||||
|
||||
// fast path
|
||||
if (same_num_pick && len <= num_pick_value && !replace) {
|
||||
IdArray rows = Full(rid, len, sizeof(IdxType) * 8, ctx);
|
||||
IdArray cols = Full(-1, len, sizeof(IdxType) * 8, ctx);
|
||||
IdArray idx = Full(-1, len, sizeof(IdxType) * 8, ctx);
|
||||
IdxType* cdata = cols.Ptr<IdxType>();
|
||||
IdxType* idata = idx.Ptr<IdxType>();
|
||||
|
||||
int64_t k = 0;
|
||||
for (int64_t j = 0; j < len; ++j) {
|
||||
const IdxType homogenized_eid = eid ? eid[off + j] : off + j;
|
||||
auto it = std::upper_bound(
|
||||
eid2etype_offset.begin(), eid2etype_offset.end(),
|
||||
homogenized_eid);
|
||||
const IdxType heterogenized_etype = it - eid2etype_offset.begin() - 1;
|
||||
const IdxType heterogenized_eid =
|
||||
homogenized_eid - eid2etype_offset[heterogenized_etype];
|
||||
|
||||
if (!has_probs || IsNullArray(prob_or_mask[heterogenized_etype])) {
|
||||
// No probability array, select all
|
||||
cdata[k] = indices[off + j];
|
||||
idata[k] = homogenized_eid;
|
||||
++k;
|
||||
} else {
|
||||
// Select the entries with non-zero probability
|
||||
const NDArray& p = prob_or_mask[heterogenized_etype];
|
||||
const DType* pdata = p.Ptr<DType>();
|
||||
if (pdata[heterogenized_eid] > 0) {
|
||||
cdata[k] = indices[off + j];
|
||||
idata[k] = homogenized_eid;
|
||||
++k;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
picked_rows[i] = rows.CreateView({k}, rows->dtype);
|
||||
picked_cols[i] = cols.CreateView({k}, cols->dtype);
|
||||
picked_idxs[i] = idx.CreateView({k}, idx->dtype);
|
||||
} else {
|
||||
// need to do per edge type sample
|
||||
std::vector<IdxType> rows;
|
||||
std::vector<IdxType> cols;
|
||||
std::vector<IdxType> idx;
|
||||
|
||||
std::vector<IdxType> et(len);
|
||||
std::vector<IdxType> et_idx(len);
|
||||
std::vector<IdxType> et_eid(len);
|
||||
std::iota(et_idx.begin(), et_idx.end(), 0);
|
||||
for (int64_t j = 0; j < len; ++j) {
|
||||
const IdxType homogenized_eid = eid ? eid[off + j] : off + j;
|
||||
auto it = std::upper_bound(
|
||||
eid2etype_offset.begin(), eid2etype_offset.end(),
|
||||
homogenized_eid);
|
||||
const IdxType heterogenized_etype = it - eid2etype_offset.begin() - 1;
|
||||
const IdxType heterogenized_eid =
|
||||
homogenized_eid - eid2etype_offset[heterogenized_etype];
|
||||
et[j] = heterogenized_etype;
|
||||
et_eid[j] = heterogenized_eid;
|
||||
}
|
||||
if (!rowwise_etype_sorted) // the edge type is sorted, not need to sort
|
||||
// it
|
||||
std::sort(
|
||||
et_idx.begin(), et_idx.end(),
|
||||
[&et](IdxType i1, IdxType i2) { return et[i1] < et[i2]; });
|
||||
CHECK_LT(et[et_idx[len - 1]], num_etypes)
|
||||
<< "etype values exceed the number of fanouts";
|
||||
|
||||
IdxType cur_et = et[et_idx[0]];
|
||||
int64_t et_offset = 0;
|
||||
int64_t et_len = 1;
|
||||
for (int64_t j = 0; j < len; ++j) {
|
||||
CHECK((j + 1 == len) || (et[et_idx[j]] <= et[et_idx[j + 1]]))
|
||||
<< "Edge type is not sorted. Please sort in advance or specify "
|
||||
"'rowwise_etype_sorted' as false.";
|
||||
if ((j + 1 == len) || cur_et != et[et_idx[j + 1]]) {
|
||||
// 1 end of the current etype
|
||||
// 2 end of the row
|
||||
// random pick for current etype
|
||||
if ((num_picks[cur_et] == -1) ||
|
||||
(et_len <= num_picks[cur_et] && !replace)) {
|
||||
// fast path, select all
|
||||
for (int64_t k = 0; k < et_len; ++k) {
|
||||
const IdxType eid_offset = off + et_idx[et_offset + k];
|
||||
const IdxType homogenized_eid =
|
||||
eid ? eid[eid_offset] : eid_offset;
|
||||
auto it = std::upper_bound(
|
||||
eid2etype_offset.begin(), eid2etype_offset.end(),
|
||||
homogenized_eid);
|
||||
const IdxType heterogenized_etype =
|
||||
it - eid2etype_offset.begin() - 1;
|
||||
const IdxType heterogenized_eid =
|
||||
homogenized_eid - eid2etype_offset[heterogenized_etype];
|
||||
|
||||
if (!has_probs ||
|
||||
IsNullArray(prob_or_mask[heterogenized_etype])) {
|
||||
// No probability, select all
|
||||
rows.push_back(rid);
|
||||
cols.push_back(indices[eid_offset]);
|
||||
idx.push_back(homogenized_eid);
|
||||
} else {
|
||||
// Select the entries with non-zero probability
|
||||
const NDArray& p = prob_or_mask[heterogenized_etype];
|
||||
const DType* pdata = p.Ptr<DType>();
|
||||
if (pdata[heterogenized_eid] > 0) {
|
||||
rows.push_back(rid);
|
||||
cols.push_back(indices[eid_offset]);
|
||||
idx.push_back(homogenized_eid);
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
IdArray picked_idx =
|
||||
Full(-1, num_picks[cur_et], sizeof(IdxType) * 8, ctx);
|
||||
IdxType* picked_idata = picked_idx.Ptr<IdxType>();
|
||||
|
||||
// need call random pick
|
||||
pick_fn(
|
||||
off, et_offset, cur_et, et_len, et_idx, et_eid, eid,
|
||||
picked_idata);
|
||||
for (int64_t k = 0; k < num_picks[cur_et]; ++k) {
|
||||
const IdxType picked = picked_idata[k];
|
||||
if (picked == -1) continue;
|
||||
rows.push_back(rid);
|
||||
cols.push_back(indices[off + et_idx[et_offset + picked]]);
|
||||
if (eid) {
|
||||
idx.push_back(eid[off + et_idx[et_offset + picked]]);
|
||||
} else {
|
||||
idx.push_back(off + et_idx[et_offset + picked]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (j + 1 == len) break;
|
||||
// next etype
|
||||
cur_et = et[et_idx[j + 1]];
|
||||
et_offset = j + 1;
|
||||
et_len = 1;
|
||||
} else {
|
||||
et_len++;
|
||||
}
|
||||
}
|
||||
|
||||
picked_rows[i] = VecToIdArray(rows, sizeof(IdxType) * 8, ctx);
|
||||
picked_cols[i] = VecToIdArray(cols, sizeof(IdxType) * 8, ctx);
|
||||
picked_idxs[i] = VecToIdArray(idx, sizeof(IdxType) * 8, ctx);
|
||||
} // end processing one row
|
||||
|
||||
CHECK_EQ(picked_rows[i]->shape[0], picked_cols[i]->shape[0]);
|
||||
CHECK_EQ(picked_rows[i]->shape[0], picked_idxs[i]->shape[0]);
|
||||
} // end processing all rows
|
||||
});
|
||||
|
||||
IdArray picked_row = Concat(picked_rows);
|
||||
IdArray picked_col = Concat(picked_cols);
|
||||
IdArray picked_idx = Concat(picked_idxs);
|
||||
return COOMatrix(
|
||||
mat.num_rows, mat.num_cols, picked_row, picked_col, picked_idx);
|
||||
}
|
||||
|
||||
// Template for picking non-zero values row-wise. The implementation first
|
||||
// slices out the corresponding rows and then converts it to CSR format. It then
|
||||
// performs row-wise pick on the CSR matrix and rectifies the returned results.
|
||||
template <typename IdxType>
|
||||
COOMatrix COORowWisePick(
|
||||
COOMatrix mat, IdArray rows, int64_t num_picks, bool replace,
|
||||
PickFn<IdxType> pick_fn, NumPicksFn<IdxType> num_picks_fn) {
|
||||
using namespace aten;
|
||||
const auto& csr = COOToCSR(COOSliceRows(mat, rows));
|
||||
const IdArray new_rows =
|
||||
Range(0, rows->shape[0], rows->dtype.bits, rows->ctx);
|
||||
const auto& picked = CSRRowWisePick<IdxType>(
|
||||
csr, new_rows, num_picks, replace, pick_fn, num_picks_fn);
|
||||
return COOMatrix(
|
||||
mat.num_rows, mat.num_cols,
|
||||
IndexSelect(rows, picked.row), // map the row index to the correct one
|
||||
picked.col, picked.data);
|
||||
}
|
||||
|
||||
// Template for picking non-zero values row-wise. The implementation first
|
||||
// slices out the corresponding rows and then converts it to CSR format. It then
|
||||
// performs row-wise pick on the CSR matrix and rectifies the returned results.
|
||||
template <typename IdxType, typename DType>
|
||||
COOMatrix COORowWisePerEtypePick(
|
||||
COOMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
|
||||
const std::vector<int64_t>& num_picks, bool replace,
|
||||
EtypeRangePickFn<IdxType> pick_fn,
|
||||
const std::vector<NDArray>& prob_or_mask) {
|
||||
using namespace aten;
|
||||
const auto& csr = COOToCSR(COOSliceRows(mat, rows));
|
||||
const IdArray new_rows =
|
||||
Range(0, rows->shape[0], rows->dtype.bits, rows->ctx);
|
||||
const auto& picked = CSRRowWisePerEtypePick<IdxType, DType>(
|
||||
csr, new_rows, eid2etype_offset, num_picks, replace, false, pick_fn,
|
||||
prob_or_mask);
|
||||
return COOMatrix(
|
||||
mat.num_rows, mat.num_cols,
|
||||
IndexSelect(rows, picked.row), // map the row index to the correct one
|
||||
picked.col, picked.data);
|
||||
}
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CPU_ROWWISE_PICK_H_
|
||||
@@ -0,0 +1,531 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/rowwise_sampling.cc
|
||||
* @brief rowwise sampling
|
||||
*/
|
||||
#include <dgl/random.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "./rowwise_pick.h"
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
namespace {
|
||||
// Equivalent to numpy expression: array[idx[off:off + len]]
|
||||
template <typename IdxType, typename FloatType>
|
||||
inline FloatArray DoubleSlice(
|
||||
FloatArray array, const IdxType* idx_data, IdxType off, IdxType len) {
|
||||
const FloatType* array_data = static_cast<FloatType*>(array->data);
|
||||
FloatArray ret = FloatArray::Empty({len}, array->dtype, array->ctx);
|
||||
FloatType* ret_data = static_cast<FloatType*>(ret->data);
|
||||
for (int64_t j = 0; j < len; ++j) {
|
||||
if (idx_data)
|
||||
ret_data[j] = array_data[idx_data[off + j]];
|
||||
else
|
||||
ret_data[j] = array_data[off + j];
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
template <typename IdxType, typename DType>
|
||||
inline NumPicksFn<IdxType> GetSamplingNumPicksFn(
|
||||
int64_t num_samples, NDArray prob_or_mask, bool replace) {
|
||||
NumPicksFn<IdxType> num_picks_fn = [prob_or_mask, num_samples, replace](
|
||||
IdxType rowid, IdxType off,
|
||||
IdxType len, const IdxType* col,
|
||||
const IdxType* data) {
|
||||
const int64_t max_num_picks = (num_samples == -1) ? len : num_samples;
|
||||
const DType* prob_or_mask_data = prob_or_mask.Ptr<DType>();
|
||||
IdxType nnz = 0;
|
||||
for (IdxType i = off; i < off + len; ++i) {
|
||||
const IdxType eid = data ? data[i] : i;
|
||||
if (prob_or_mask_data[eid] > 0) {
|
||||
++nnz;
|
||||
}
|
||||
}
|
||||
|
||||
if (replace) {
|
||||
return static_cast<IdxType>(nnz == 0 ? 0 : max_num_picks);
|
||||
} else {
|
||||
return std::min(static_cast<IdxType>(max_num_picks), nnz);
|
||||
}
|
||||
};
|
||||
return num_picks_fn;
|
||||
}
|
||||
|
||||
template <typename IdxType, typename DType>
|
||||
inline PickFn<IdxType> GetSamplingPickFn(
|
||||
int64_t num_samples, NDArray prob_or_mask, bool replace) {
|
||||
PickFn<IdxType> pick_fn = [prob_or_mask, num_samples, replace](
|
||||
IdxType rowid, IdxType off, IdxType len,
|
||||
IdxType num_picks, const IdxType* col,
|
||||
const IdxType* data, IdxType* out_idx) {
|
||||
NDArray prob_or_mask_selected =
|
||||
DoubleSlice<IdxType, DType>(prob_or_mask, data, off, len);
|
||||
RandomEngine::ThreadLocal()->Choice<IdxType, DType>(
|
||||
num_picks, prob_or_mask_selected, out_idx, replace);
|
||||
for (int64_t j = 0; j < num_picks; ++j) {
|
||||
out_idx[j] += off;
|
||||
}
|
||||
};
|
||||
return pick_fn;
|
||||
}
|
||||
|
||||
template <typename IdxType, typename FloatType>
|
||||
inline EtypeRangePickFn<IdxType> GetSamplingRangePickFn(
|
||||
const std::vector<int64_t>& num_samples,
|
||||
const std::vector<FloatArray>& prob, bool replace) {
|
||||
EtypeRangePickFn<IdxType> pick_fn =
|
||||
[prob, num_samples, replace](
|
||||
IdxType off, IdxType et_offset, IdxType cur_et, IdxType et_len,
|
||||
const std::vector<IdxType>& et_idx,
|
||||
const std::vector<IdxType>& et_eid, const IdxType* eid,
|
||||
IdxType* out_idx) {
|
||||
const FloatArray& p = prob[cur_et];
|
||||
const FloatType* p_data = IsNullArray(p) ? nullptr : p.Ptr<FloatType>();
|
||||
FloatArray probs = FloatArray::Empty({et_len}, p->dtype, p->ctx);
|
||||
FloatType* probs_data = probs.Ptr<FloatType>();
|
||||
for (int64_t j = 0; j < et_len; ++j) {
|
||||
const IdxType cur_eid = et_eid[et_idx[et_offset + j]];
|
||||
probs_data[j] = p_data ? p_data[cur_eid] : static_cast<FloatType>(1.);
|
||||
}
|
||||
|
||||
RandomEngine::ThreadLocal()->Choice<IdxType, FloatType>(
|
||||
num_samples[cur_et], probs, out_idx, replace);
|
||||
};
|
||||
return pick_fn;
|
||||
}
|
||||
|
||||
template <typename IdxType>
|
||||
inline NumPicksFn<IdxType> GetSamplingUniformNumPicksFn(
|
||||
int64_t num_samples, bool replace) {
|
||||
NumPicksFn<IdxType> num_picks_fn = [num_samples, replace](
|
||||
IdxType rowid, IdxType off,
|
||||
IdxType len, const IdxType* col,
|
||||
const IdxType* data) {
|
||||
const int64_t max_num_picks = (num_samples == -1) ? len : num_samples;
|
||||
if (replace) {
|
||||
return static_cast<IdxType>(len == 0 ? 0 : max_num_picks);
|
||||
} else {
|
||||
return std::min(static_cast<IdxType>(max_num_picks), len);
|
||||
}
|
||||
};
|
||||
return num_picks_fn;
|
||||
}
|
||||
|
||||
template <typename IdxType>
|
||||
inline PickFn<IdxType> GetSamplingUniformPickFn(
|
||||
int64_t num_samples, bool replace) {
|
||||
PickFn<IdxType> pick_fn = [num_samples, replace](
|
||||
IdxType rowid, IdxType off, IdxType len,
|
||||
IdxType num_picks, const IdxType* col,
|
||||
const IdxType* data, IdxType* out_idx) {
|
||||
RandomEngine::ThreadLocal()->UniformChoice<IdxType>(
|
||||
num_picks, len, out_idx, replace);
|
||||
for (int64_t j = 0; j < num_picks; ++j) {
|
||||
out_idx[j] += off;
|
||||
}
|
||||
};
|
||||
return pick_fn;
|
||||
}
|
||||
|
||||
template <typename IdxType>
|
||||
inline EtypeRangePickFn<IdxType> GetSamplingUniformRangePickFn(
|
||||
const std::vector<int64_t>& num_samples, bool replace) {
|
||||
EtypeRangePickFn<IdxType> pick_fn =
|
||||
[num_samples, replace](
|
||||
IdxType off, IdxType et_offset, IdxType cur_et, IdxType et_len,
|
||||
const std::vector<IdxType>& et_idx,
|
||||
const std::vector<IdxType>& et_eid, const IdxType* data,
|
||||
IdxType* out_idx) {
|
||||
RandomEngine::ThreadLocal()->UniformChoice<IdxType>(
|
||||
num_samples[cur_et], et_len, out_idx, replace);
|
||||
};
|
||||
return pick_fn;
|
||||
}
|
||||
|
||||
template <typename IdxType, typename FloatType>
|
||||
inline NumPicksFn<IdxType> GetSamplingBiasedNumPicksFn(
|
||||
int64_t num_samples, IdArray split, FloatArray bias, bool replace) {
|
||||
NumPicksFn<IdxType> num_picks_fn = [num_samples, split, bias, replace](
|
||||
IdxType rowid, IdxType off,
|
||||
IdxType len, const IdxType* col,
|
||||
const IdxType* data) {
|
||||
const int64_t max_num_picks = (num_samples == -1) ? len : num_samples;
|
||||
const int64_t num_tags = split->shape[1] - 1;
|
||||
const IdxType* tag_offset = split.Ptr<IdxType>() + rowid * split->shape[1];
|
||||
const FloatType* bias_data = bias.Ptr<FloatType>();
|
||||
IdxType nnz = 0;
|
||||
for (int64_t j = 0; j < num_tags; ++j) {
|
||||
if (bias_data[j] > 0) {
|
||||
nnz += tag_offset[j + 1] - tag_offset[j];
|
||||
}
|
||||
}
|
||||
|
||||
if (replace) {
|
||||
return static_cast<IdxType>(nnz == 0 ? 0 : max_num_picks);
|
||||
} else {
|
||||
return std::min(static_cast<IdxType>(max_num_picks), nnz);
|
||||
}
|
||||
};
|
||||
return num_picks_fn;
|
||||
}
|
||||
|
||||
template <typename IdxType, typename FloatType>
|
||||
inline PickFn<IdxType> GetSamplingBiasedPickFn(
|
||||
int64_t num_samples, IdArray split, FloatArray bias, bool replace) {
|
||||
PickFn<IdxType> pick_fn = [num_samples, split, bias, replace](
|
||||
IdxType rowid, IdxType off, IdxType len,
|
||||
IdxType num_picks, const IdxType* col,
|
||||
const IdxType* data, IdxType* out_idx) {
|
||||
const IdxType* tag_offset = split.Ptr<IdxType>() + rowid * split->shape[1];
|
||||
RandomEngine::ThreadLocal()->BiasedChoice<IdxType, FloatType>(
|
||||
num_picks, tag_offset, bias, out_idx, replace);
|
||||
for (int64_t j = 0; j < num_picks; ++j) {
|
||||
out_idx[j] += off;
|
||||
}
|
||||
};
|
||||
return pick_fn;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
/////////////////////////////// CSR ///////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdxType, typename DType>
|
||||
COOMatrix CSRRowWiseSampling(
|
||||
CSRMatrix mat, IdArray rows, int64_t num_samples, NDArray prob_or_mask,
|
||||
bool replace) {
|
||||
// If num_samples is -1, select all neighbors without replacement.
|
||||
replace = (replace && num_samples != -1);
|
||||
CHECK(prob_or_mask.defined());
|
||||
auto num_picks_fn =
|
||||
GetSamplingNumPicksFn<IdxType, DType>(num_samples, prob_or_mask, replace);
|
||||
auto pick_fn =
|
||||
GetSamplingPickFn<IdxType, DType>(num_samples, prob_or_mask, replace);
|
||||
return CSRRowWisePick(mat, rows, num_samples, replace, pick_fn, num_picks_fn);
|
||||
}
|
||||
|
||||
template COOMatrix CSRRowWiseSampling<kDGLCPU, int32_t, float>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix CSRRowWiseSampling<kDGLCPU, int64_t, float>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix CSRRowWiseSampling<kDGLCPU, int32_t, double>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix CSRRowWiseSampling<kDGLCPU, int64_t, double>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix CSRRowWiseSampling<kDGLCPU, int32_t, int8_t>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix CSRRowWiseSampling<kDGLCPU, int64_t, int8_t>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix CSRRowWiseSampling<kDGLCPU, int32_t, uint8_t>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix CSRRowWiseSampling<kDGLCPU, int64_t, uint8_t>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, bool);
|
||||
|
||||
template <
|
||||
DGLDeviceType XPU, typename IdxType, typename DType, bool map_seed_nodes>
|
||||
std::pair<CSRMatrix, IdArray> CSRRowWiseSamplingFused(
|
||||
CSRMatrix mat, IdArray rows, IdArray seed_mapping,
|
||||
std::vector<IdxType>* new_seed_nodes, int64_t num_samples,
|
||||
NDArray prob_or_mask, bool replace) {
|
||||
// If num_samples is -1, select all neighbors without replacement.
|
||||
replace = (replace && num_samples != -1);
|
||||
CHECK(prob_or_mask.defined());
|
||||
auto num_picks_fn =
|
||||
GetSamplingNumPicksFn<IdxType, DType>(num_samples, prob_or_mask, replace);
|
||||
auto pick_fn =
|
||||
GetSamplingPickFn<IdxType, DType>(num_samples, prob_or_mask, replace);
|
||||
return CSRRowWisePickFused<IdxType, map_seed_nodes>(
|
||||
mat, rows, seed_mapping, new_seed_nodes, num_samples, replace, pick_fn,
|
||||
num_picks_fn);
|
||||
}
|
||||
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingFused<kDGLCPU, int32_t, float, true>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int32_t>*, int64_t, NDArray, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingFused<kDGLCPU, int64_t, float, true>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int64_t>*, int64_t, NDArray, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingFused<kDGLCPU, int32_t, double, true>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int32_t>*, int64_t, NDArray, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingFused<kDGLCPU, int64_t, double, true>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int64_t>*, int64_t, NDArray, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingFused<kDGLCPU, int32_t, int8_t, true>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int32_t>*, int64_t, NDArray, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingFused<kDGLCPU, int64_t, int8_t, true>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int64_t>*, int64_t, NDArray, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingFused<kDGLCPU, int32_t, uint8_t, true>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int32_t>*, int64_t, NDArray, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingFused<kDGLCPU, int64_t, uint8_t, true>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int64_t>*, int64_t, NDArray, bool);
|
||||
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingFused<kDGLCPU, int32_t, float, false>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int32_t>*, int64_t, NDArray, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingFused<kDGLCPU, int64_t, float, false>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int64_t>*, int64_t, NDArray, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingFused<kDGLCPU, int32_t, double, false>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int32_t>*, int64_t, NDArray, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingFused<kDGLCPU, int64_t, double, false>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int64_t>*, int64_t, NDArray, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingFused<kDGLCPU, int32_t, int8_t, false>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int32_t>*, int64_t, NDArray, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingFused<kDGLCPU, int64_t, int8_t, false>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int64_t>*, int64_t, NDArray, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingFused<kDGLCPU, int32_t, uint8_t, false>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int32_t>*, int64_t, NDArray, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingFused<kDGLCPU, int64_t, uint8_t, false>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int64_t>*, int64_t, NDArray, bool);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdxType, typename DType>
|
||||
COOMatrix CSRRowWisePerEtypeSampling(
|
||||
CSRMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
|
||||
const std::vector<int64_t>& num_samples,
|
||||
const std::vector<NDArray>& prob_or_mask, bool replace,
|
||||
bool rowwise_etype_sorted) {
|
||||
CHECK(prob_or_mask.size() == num_samples.size())
|
||||
<< "the number of probability tensors does not match the number of edge "
|
||||
"types.";
|
||||
for (auto& p : prob_or_mask) CHECK(p.defined());
|
||||
auto pick_fn = GetSamplingRangePickFn<IdxType, DType>(
|
||||
num_samples, prob_or_mask, replace);
|
||||
return CSRRowWisePerEtypePick<IdxType, DType>(
|
||||
mat, rows, eid2etype_offset, num_samples, replace, rowwise_etype_sorted,
|
||||
pick_fn, prob_or_mask);
|
||||
}
|
||||
|
||||
template COOMatrix CSRRowWisePerEtypeSampling<kDGLCPU, int32_t, float>(
|
||||
CSRMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, const std::vector<NDArray>&, bool, bool);
|
||||
template COOMatrix CSRRowWisePerEtypeSampling<kDGLCPU, int64_t, float>(
|
||||
CSRMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, const std::vector<NDArray>&, bool, bool);
|
||||
template COOMatrix CSRRowWisePerEtypeSampling<kDGLCPU, int32_t, double>(
|
||||
CSRMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, const std::vector<NDArray>&, bool, bool);
|
||||
template COOMatrix CSRRowWisePerEtypeSampling<kDGLCPU, int64_t, double>(
|
||||
CSRMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, const std::vector<NDArray>&, bool, bool);
|
||||
template COOMatrix CSRRowWisePerEtypeSampling<kDGLCPU, int32_t, int8_t>(
|
||||
CSRMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, const std::vector<NDArray>&, bool, bool);
|
||||
template COOMatrix CSRRowWisePerEtypeSampling<kDGLCPU, int64_t, int8_t>(
|
||||
CSRMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, const std::vector<NDArray>&, bool, bool);
|
||||
template COOMatrix CSRRowWisePerEtypeSampling<kDGLCPU, int32_t, uint8_t>(
|
||||
CSRMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, const std::vector<NDArray>&, bool, bool);
|
||||
template COOMatrix CSRRowWisePerEtypeSampling<kDGLCPU, int64_t, uint8_t>(
|
||||
CSRMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, const std::vector<NDArray>&, bool, bool);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdxType>
|
||||
COOMatrix CSRRowWiseSamplingUniform(
|
||||
CSRMatrix mat, IdArray rows, int64_t num_samples, bool replace) {
|
||||
// If num_samples is -1, select all neighbors without replacement.
|
||||
replace = (replace && num_samples != -1);
|
||||
auto num_picks_fn =
|
||||
GetSamplingUniformNumPicksFn<IdxType>(num_samples, replace);
|
||||
auto pick_fn = GetSamplingUniformPickFn<IdxType>(num_samples, replace);
|
||||
return CSRRowWisePick(mat, rows, num_samples, replace, pick_fn, num_picks_fn);
|
||||
}
|
||||
|
||||
template COOMatrix CSRRowWiseSamplingUniform<kDGLCPU, int32_t>(
|
||||
CSRMatrix, IdArray, int64_t, bool);
|
||||
template COOMatrix CSRRowWiseSamplingUniform<kDGLCPU, int64_t>(
|
||||
CSRMatrix, IdArray, int64_t, bool);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdxType, bool map_seed_nodes>
|
||||
std::pair<CSRMatrix, IdArray> CSRRowWiseSamplingUniformFused(
|
||||
CSRMatrix mat, IdArray rows, IdArray seed_mapping,
|
||||
std::vector<IdxType>* new_seed_nodes, int64_t num_samples, bool replace) {
|
||||
// If num_samples is -1, select all neighbors without replacement.
|
||||
replace = (replace && num_samples != -1);
|
||||
auto num_picks_fn =
|
||||
GetSamplingUniformNumPicksFn<IdxType>(num_samples, replace);
|
||||
auto pick_fn = GetSamplingUniformPickFn<IdxType>(num_samples, replace);
|
||||
return CSRRowWisePickFused<IdxType, map_seed_nodes>(
|
||||
mat, rows, seed_mapping, new_seed_nodes, num_samples, replace, pick_fn,
|
||||
num_picks_fn);
|
||||
}
|
||||
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingUniformFused<kDGLCPU, int32_t, true>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int32_t>*, int64_t, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingUniformFused<kDGLCPU, int64_t, true>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int64_t>*, int64_t, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingUniformFused<kDGLCPU, int32_t, false>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int32_t>*, int64_t, bool);
|
||||
template std::pair<CSRMatrix, IdArray>
|
||||
CSRRowWiseSamplingUniformFused<kDGLCPU, int64_t, false>(
|
||||
CSRMatrix, IdArray, IdArray, std::vector<int64_t>*, int64_t, bool);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdxType>
|
||||
COOMatrix CSRRowWisePerEtypeSamplingUniform(
|
||||
CSRMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
|
||||
const std::vector<int64_t>& num_samples, bool replace,
|
||||
bool rowwise_etype_sorted) {
|
||||
auto pick_fn = GetSamplingUniformRangePickFn<IdxType>(num_samples, replace);
|
||||
return CSRRowWisePerEtypePick<IdxType, float>(
|
||||
mat, rows, eid2etype_offset, num_samples, replace, rowwise_etype_sorted,
|
||||
pick_fn, {});
|
||||
}
|
||||
|
||||
template COOMatrix CSRRowWisePerEtypeSamplingUniform<kDGLCPU, int32_t>(
|
||||
CSRMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, bool, bool);
|
||||
template COOMatrix CSRRowWisePerEtypeSamplingUniform<kDGLCPU, int64_t>(
|
||||
CSRMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, bool, bool);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdxType, typename FloatType>
|
||||
COOMatrix CSRRowWiseSamplingBiased(
|
||||
CSRMatrix mat, IdArray rows, int64_t num_samples, NDArray tag_offset,
|
||||
FloatArray bias, bool replace) {
|
||||
// If num_samples is -1, select all neighbors without replacement.
|
||||
replace = (replace && num_samples != -1);
|
||||
auto num_picks_fn = GetSamplingBiasedNumPicksFn<IdxType, FloatType>(
|
||||
num_samples, tag_offset, bias, replace);
|
||||
auto pick_fn = GetSamplingBiasedPickFn<IdxType, FloatType>(
|
||||
num_samples, tag_offset, bias, replace);
|
||||
return CSRRowWisePick(mat, rows, num_samples, replace, pick_fn, num_picks_fn);
|
||||
}
|
||||
|
||||
template COOMatrix CSRRowWiseSamplingBiased<kDGLCPU, int32_t, float>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, FloatArray, bool);
|
||||
|
||||
template COOMatrix CSRRowWiseSamplingBiased<kDGLCPU, int64_t, float>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, FloatArray, bool);
|
||||
|
||||
template COOMatrix CSRRowWiseSamplingBiased<kDGLCPU, int32_t, double>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, FloatArray, bool);
|
||||
|
||||
template COOMatrix CSRRowWiseSamplingBiased<kDGLCPU, int64_t, double>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, FloatArray, bool);
|
||||
|
||||
/////////////////////////////// COO ///////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdxType, typename DType>
|
||||
COOMatrix COORowWiseSampling(
|
||||
COOMatrix mat, IdArray rows, int64_t num_samples, NDArray prob_or_mask,
|
||||
bool replace) {
|
||||
// If num_samples is -1, select all neighbors without replacement.
|
||||
replace = (replace && num_samples != -1);
|
||||
CHECK(prob_or_mask.defined());
|
||||
auto num_picks_fn =
|
||||
GetSamplingNumPicksFn<IdxType, DType>(num_samples, prob_or_mask, replace);
|
||||
auto pick_fn =
|
||||
GetSamplingPickFn<IdxType, DType>(num_samples, prob_or_mask, replace);
|
||||
return COORowWisePick(mat, rows, num_samples, replace, pick_fn, num_picks_fn);
|
||||
}
|
||||
|
||||
template COOMatrix COORowWiseSampling<kDGLCPU, int32_t, float>(
|
||||
COOMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix COORowWiseSampling<kDGLCPU, int64_t, float>(
|
||||
COOMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix COORowWiseSampling<kDGLCPU, int32_t, double>(
|
||||
COOMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix COORowWiseSampling<kDGLCPU, int64_t, double>(
|
||||
COOMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix COORowWiseSampling<kDGLCPU, int32_t, int8_t>(
|
||||
COOMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix COORowWiseSampling<kDGLCPU, int64_t, int8_t>(
|
||||
COOMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix COORowWiseSampling<kDGLCPU, int32_t, uint8_t>(
|
||||
COOMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix COORowWiseSampling<kDGLCPU, int64_t, uint8_t>(
|
||||
COOMatrix, IdArray, int64_t, NDArray, bool);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdxType, typename DType>
|
||||
COOMatrix COORowWisePerEtypeSampling(
|
||||
COOMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
|
||||
const std::vector<int64_t>& num_samples,
|
||||
const std::vector<NDArray>& prob_or_mask, bool replace) {
|
||||
CHECK(prob_or_mask.size() == num_samples.size())
|
||||
<< "the number of probability tensors do not match the number of edge "
|
||||
"types.";
|
||||
for (auto& p : prob_or_mask) CHECK(p.defined());
|
||||
auto pick_fn = GetSamplingRangePickFn<IdxType, DType>(
|
||||
num_samples, prob_or_mask, replace);
|
||||
return COORowWisePerEtypePick<IdxType, DType>(
|
||||
mat, rows, eid2etype_offset, num_samples, replace, pick_fn, prob_or_mask);
|
||||
}
|
||||
|
||||
template COOMatrix COORowWisePerEtypeSampling<kDGLCPU, int32_t, float>(
|
||||
COOMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, const std::vector<NDArray>&, bool);
|
||||
template COOMatrix COORowWisePerEtypeSampling<kDGLCPU, int64_t, float>(
|
||||
COOMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, const std::vector<NDArray>&, bool);
|
||||
template COOMatrix COORowWisePerEtypeSampling<kDGLCPU, int32_t, double>(
|
||||
COOMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, const std::vector<NDArray>&, bool);
|
||||
template COOMatrix COORowWisePerEtypeSampling<kDGLCPU, int64_t, double>(
|
||||
COOMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, const std::vector<NDArray>&, bool);
|
||||
template COOMatrix COORowWisePerEtypeSampling<kDGLCPU, int32_t, int8_t>(
|
||||
COOMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, const std::vector<NDArray>&, bool);
|
||||
template COOMatrix COORowWisePerEtypeSampling<kDGLCPU, int64_t, int8_t>(
|
||||
COOMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, const std::vector<NDArray>&, bool);
|
||||
template COOMatrix COORowWisePerEtypeSampling<kDGLCPU, int32_t, uint8_t>(
|
||||
COOMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, const std::vector<NDArray>&, bool);
|
||||
template COOMatrix COORowWisePerEtypeSampling<kDGLCPU, int64_t, uint8_t>(
|
||||
COOMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, const std::vector<NDArray>&, bool);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdxType>
|
||||
COOMatrix COORowWiseSamplingUniform(
|
||||
COOMatrix mat, IdArray rows, int64_t num_samples, bool replace) {
|
||||
// If num_samples is -1, select all neighbors without replacement.
|
||||
replace = (replace && num_samples != -1);
|
||||
auto num_picks_fn =
|
||||
GetSamplingUniformNumPicksFn<IdxType>(num_samples, replace);
|
||||
auto pick_fn = GetSamplingUniformPickFn<IdxType>(num_samples, replace);
|
||||
return COORowWisePick(mat, rows, num_samples, replace, pick_fn, num_picks_fn);
|
||||
}
|
||||
|
||||
template COOMatrix COORowWiseSamplingUniform<kDGLCPU, int32_t>(
|
||||
COOMatrix, IdArray, int64_t, bool);
|
||||
template COOMatrix COORowWiseSamplingUniform<kDGLCPU, int64_t>(
|
||||
COOMatrix, IdArray, int64_t, bool);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdxType>
|
||||
COOMatrix COORowWisePerEtypeSamplingUniform(
|
||||
COOMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
|
||||
const std::vector<int64_t>& num_samples, bool replace) {
|
||||
auto pick_fn = GetSamplingUniformRangePickFn<IdxType>(num_samples, replace);
|
||||
return COORowWisePerEtypePick<IdxType, float>(
|
||||
mat, rows, eid2etype_offset, num_samples, replace, pick_fn, {});
|
||||
}
|
||||
|
||||
template COOMatrix COORowWisePerEtypeSamplingUniform<kDGLCPU, int32_t>(
|
||||
COOMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, bool);
|
||||
template COOMatrix COORowWisePerEtypeSamplingUniform<kDGLCPU, int64_t>(
|
||||
COOMatrix, IdArray, const std::vector<int64_t>&,
|
||||
const std::vector<int64_t>&, bool);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,122 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/rowwise_topk.cc
|
||||
* @brief rowwise topk
|
||||
*/
|
||||
#include <algorithm>
|
||||
#include <numeric>
|
||||
|
||||
#include "./rowwise_pick.h"
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
namespace {
|
||||
|
||||
template <typename IdxType>
|
||||
inline NumPicksFn<IdxType> GetTopkNumPicksFn(int64_t k) {
|
||||
NumPicksFn<IdxType> num_picks_fn = [k](IdxType rowid, IdxType off,
|
||||
IdxType len, const IdxType* col,
|
||||
const IdxType* data) {
|
||||
const int64_t max_num_picks = (k == -1) ? len : k;
|
||||
return std::min(static_cast<IdxType>(max_num_picks), len);
|
||||
};
|
||||
return num_picks_fn;
|
||||
}
|
||||
|
||||
template <typename IdxType, typename DType>
|
||||
inline PickFn<IdxType> GetTopkPickFn(NDArray weight, bool ascending) {
|
||||
const DType* wdata = static_cast<DType*>(weight->data);
|
||||
PickFn<IdxType> pick_fn = [ascending, wdata](
|
||||
IdxType rowid, IdxType off, IdxType len,
|
||||
IdxType num_picks, const IdxType* col,
|
||||
const IdxType* data, IdxType* out_idx) {
|
||||
std::function<bool(IdxType, IdxType)> compare_fn;
|
||||
if (ascending) {
|
||||
if (data) {
|
||||
compare_fn = [wdata, data](IdxType i, IdxType j) {
|
||||
return wdata[data[i]] < wdata[data[j]];
|
||||
};
|
||||
} else {
|
||||
compare_fn = [wdata](IdxType i, IdxType j) {
|
||||
return wdata[i] < wdata[j];
|
||||
};
|
||||
}
|
||||
} else {
|
||||
if (data) {
|
||||
compare_fn = [wdata, data](IdxType i, IdxType j) {
|
||||
return wdata[data[i]] > wdata[data[j]];
|
||||
};
|
||||
} else {
|
||||
compare_fn = [wdata](IdxType i, IdxType j) {
|
||||
return wdata[i] > wdata[j];
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<IdxType> idx(len);
|
||||
std::iota(idx.begin(), idx.end(), off);
|
||||
std::sort(idx.begin(), idx.end(), compare_fn);
|
||||
for (int64_t j = 0; j < num_picks; ++j) {
|
||||
out_idx[j] = idx[j];
|
||||
}
|
||||
};
|
||||
|
||||
return pick_fn;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
template <DGLDeviceType XPU, typename IdxType, typename DType>
|
||||
COOMatrix CSRRowWiseTopk(
|
||||
CSRMatrix mat, IdArray rows, int64_t k, NDArray weight, bool ascending) {
|
||||
auto num_picks_fn = GetTopkNumPicksFn<IdxType>(k);
|
||||
auto pick_fn = GetTopkPickFn<IdxType, DType>(weight, ascending);
|
||||
return CSRRowWisePick(mat, rows, k, false, pick_fn, num_picks_fn);
|
||||
}
|
||||
|
||||
template COOMatrix CSRRowWiseTopk<kDGLCPU, int32_t, int32_t>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix CSRRowWiseTopk<kDGLCPU, int64_t, int32_t>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix CSRRowWiseTopk<kDGLCPU, int32_t, int64_t>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix CSRRowWiseTopk<kDGLCPU, int64_t, int64_t>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix CSRRowWiseTopk<kDGLCPU, int32_t, float>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix CSRRowWiseTopk<kDGLCPU, int64_t, float>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix CSRRowWiseTopk<kDGLCPU, int32_t, double>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix CSRRowWiseTopk<kDGLCPU, int64_t, double>(
|
||||
CSRMatrix, IdArray, int64_t, NDArray, bool);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdxType, typename DType>
|
||||
COOMatrix COORowWiseTopk(
|
||||
COOMatrix mat, IdArray rows, int64_t k, NDArray weight, bool ascending) {
|
||||
auto num_picks_fn = GetTopkNumPicksFn<IdxType>(k);
|
||||
auto pick_fn = GetTopkPickFn<IdxType, DType>(weight, ascending);
|
||||
return COORowWisePick(mat, rows, k, false, pick_fn, num_picks_fn);
|
||||
}
|
||||
|
||||
template COOMatrix COORowWiseTopk<kDGLCPU, int32_t, int32_t>(
|
||||
COOMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix COORowWiseTopk<kDGLCPU, int64_t, int32_t>(
|
||||
COOMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix COORowWiseTopk<kDGLCPU, int32_t, int64_t>(
|
||||
COOMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix COORowWiseTopk<kDGLCPU, int64_t, int64_t>(
|
||||
COOMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix COORowWiseTopk<kDGLCPU, int32_t, float>(
|
||||
COOMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix COORowWiseTopk<kDGLCPU, int64_t, float>(
|
||||
COOMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix COORowWiseTopk<kDGLCPU, int32_t, double>(
|
||||
COOMatrix, IdArray, int64_t, NDArray, bool);
|
||||
template COOMatrix COORowWiseTopk<kDGLCPU, int64_t, double>(
|
||||
COOMatrix, IdArray, int64_t, NDArray, bool);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,230 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file aten/cpu/sddmm.cc
|
||||
* @brief SDDMM C APIs and definitions.
|
||||
*/
|
||||
#include "./sddmm.h"
|
||||
|
||||
#include <dgl/array.h>
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
|
||||
#define SWITCH_RHS(rhs_target, RhsTarget, ...) \
|
||||
do { \
|
||||
if ((rhs_target) == 0) { \
|
||||
constexpr int RhsTarget = 0; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((rhs_target) == 1) { \
|
||||
constexpr int RhsTarget = 1; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((rhs_target) == 2) { \
|
||||
constexpr int RhsTarget = 2; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else { \
|
||||
LOG(INFO) << "Invalid rhs target: " << (rhs_target); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
#define SWITCH_TARGET(lhs_target, rhs_target, LhsTarget, RhsTarget, ...) \
|
||||
do { \
|
||||
if ((lhs_target) == 0) { \
|
||||
constexpr int LhsTarget = 0; \
|
||||
SWITCH_RHS(rhs_target, RhsTarget, __VA_ARGS__); \
|
||||
} else if ((lhs_target) == 1) { \
|
||||
constexpr int LhsTarget = 1; \
|
||||
SWITCH_RHS(rhs_target, RhsTarget, __VA_ARGS__); \
|
||||
} else if ((lhs_target) == 2) { \
|
||||
constexpr int LhsTarget = 2; \
|
||||
SWITCH_RHS(rhs_target, RhsTarget, __VA_ARGS__); \
|
||||
} else { \
|
||||
LOG(INFO) << "Invalid lhs target: " << (lhs_target); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
/** @brief Generalized SDDMM on Csr format. */
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SDDMMCsr(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target) {
|
||||
SWITCH_OP(op, Op, {
|
||||
SWITCH_TARGET(lhs_target, rhs_target, LhsTarget, RhsTarget, {
|
||||
cpu::SDDMMCsr<IdType, DType, Op, LhsTarget, RhsTarget>(
|
||||
bcast, csr, lhs, rhs, out);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** @brief Generalized SDDMM on Csr format with Heterograph support. */
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SDDMMCsrHetero(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& vec_lhs,
|
||||
const std::vector<NDArray>& vec_rhs, std::vector<NDArray> vec_out,
|
||||
int lhs_target, int rhs_target, const std::vector<dgl_type_t>& lhs_nid,
|
||||
const std::vector<dgl_type_t>& rhs_nid) {
|
||||
SWITCH_OP(op, Op, {
|
||||
SWITCH_TARGET(lhs_target, rhs_target, LhsTarget, RhsTarget, {
|
||||
/* Call SDDMM for each relation type */
|
||||
for (dgl_type_t etype = 0; etype < lhs_nid.size(); ++etype) {
|
||||
CSRMatrix csr = vec_csr[etype];
|
||||
NDArray lhs = vec_lhs[lhs_nid[etype]];
|
||||
NDArray rhs = vec_rhs[rhs_nid[etype]];
|
||||
NDArray out = vec_out[etype];
|
||||
cpu::SDDMMCsr<IdType, DType, Op, LhsTarget, RhsTarget>(
|
||||
bcast, csr, lhs, rhs, out);
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
template void SDDMMCsr<kDGLCPU, int32_t, BFloat16>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCsr<kDGLCPU, int64_t, BFloat16>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCsr<kDGLCPU, int32_t, float>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCsr<kDGLCPU, int64_t, float>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCsr<kDGLCPU, int32_t, double>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCsr<kDGLCPU, int64_t, double>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
|
||||
template void SDDMMCsrHetero<kDGLCPU, int32_t, BFloat16>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCsrHetero<kDGLCPU, int64_t, BFloat16>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCsrHetero<kDGLCPU, int32_t, float>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCsrHetero<kDGLCPU, int64_t, float>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCsrHetero<kDGLCPU, int32_t, double>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCsrHetero<kDGLCPU, int64_t, double>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
|
||||
/** @brief Generalized SDDMM on Coo format. */
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SDDMMCoo(
|
||||
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target) {
|
||||
SWITCH_OP(op, Op, {
|
||||
SWITCH_TARGET(lhs_target, rhs_target, LhsTarget, RhsTarget, {
|
||||
cpu::SDDMMCoo<IdType, DType, Op, LhsTarget, RhsTarget>(
|
||||
bcast, coo, lhs, rhs, out);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** @brief Generalized SDDMM on Coo format with Heterograph support. */
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SDDMMCooHetero(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& vec_lhs,
|
||||
const std::vector<NDArray>& vec_rhs, std::vector<NDArray> vec_out,
|
||||
int lhs_target, int rhs_target, const std::vector<dgl_type_t>& lhs_nid,
|
||||
const std::vector<dgl_type_t>& rhs_nid) {
|
||||
SWITCH_OP(op, Op, {
|
||||
SWITCH_TARGET(lhs_target, rhs_target, LhsTarget, RhsTarget, {
|
||||
/* Call SDDMM for each relation type */
|
||||
for (dgl_type_t etype = 0; etype < lhs_nid.size(); ++etype) {
|
||||
COOMatrix coo = vec_coo[etype];
|
||||
NDArray lhs = vec_lhs[lhs_nid[etype]];
|
||||
NDArray rhs = vec_rhs[rhs_nid[etype]];
|
||||
NDArray out = vec_out[etype];
|
||||
cpu::SDDMMCoo<IdType, DType, Op, LhsTarget, RhsTarget>(
|
||||
bcast, coo, lhs, rhs, out);
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
template void SDDMMCoo<kDGLCPU, int32_t, BFloat16>(
|
||||
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCoo<kDGLCPU, int64_t, BFloat16>(
|
||||
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCoo<kDGLCPU, int32_t, float>(
|
||||
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCoo<kDGLCPU, int64_t, float>(
|
||||
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCoo<kDGLCPU, int32_t, double>(
|
||||
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCoo<kDGLCPU, int64_t, double>(
|
||||
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
|
||||
template void SDDMMCooHetero<kDGLCPU, int32_t, BFloat16>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCooHetero<kDGLCPU, int64_t, BFloat16>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCooHetero<kDGLCPU, int32_t, float>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCooHetero<kDGLCPU, int64_t, float>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCooHetero<kDGLCPU, int32_t, double>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCooHetero<kDGLCPU, int64_t, double>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,229 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/sddmm.h
|
||||
* @brief SDDMM CPU kernel function header.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CPU_SDDMM_H_
|
||||
#define DGL_ARRAY_CPU_SDDMM_H_
|
||||
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/bcast.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
|
||||
#include "../selector.h"
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace cpu {
|
||||
|
||||
/**
|
||||
* @brief CPU kernel of g-SDDMM on Csr format.
|
||||
* @param bcast Broadcast information.
|
||||
* @param csr The Csr matrix.
|
||||
* @param lhs The left hand side operand feature.
|
||||
* @param rhs The right hand size operand feature.
|
||||
* @param out The result feature on edges.
|
||||
* @note it uses node parallel strategy, different threads are responsible
|
||||
* for the computation of different nodes.
|
||||
*/
|
||||
template <
|
||||
typename IdType, typename DType, typename Op, int LhsTarget = 0,
|
||||
int RhsTarget = 2>
|
||||
void SDDMMCsr(
|
||||
const BcastOff& bcast, const CSRMatrix& csr, NDArray lhs, NDArray rhs,
|
||||
NDArray out) {
|
||||
const bool has_idx = !IsNullArray(csr.data);
|
||||
const IdType* indptr = csr.indptr.Ptr<IdType>();
|
||||
const IdType* indices = csr.indices.Ptr<IdType>();
|
||||
const IdType* edges = csr.data.Ptr<IdType>();
|
||||
const DType* X = lhs.Ptr<DType>();
|
||||
const DType* Y = rhs.Ptr<DType>();
|
||||
const int64_t dim = bcast.out_len, lhs_dim = bcast.lhs_len,
|
||||
rhs_dim = bcast.rhs_len, reduce_size = bcast.reduce_size;
|
||||
DType* O = out.Ptr<DType>();
|
||||
runtime::parallel_for(0, csr.num_rows, [=](IdType b, IdType e) {
|
||||
for (auto rid = b; rid < e; ++rid) {
|
||||
const IdType row_start = indptr[rid], row_end = indptr[rid + 1];
|
||||
for (IdType j = row_start; j < row_end; ++j) {
|
||||
const IdType cid = indices[j];
|
||||
const IdType eid = has_idx ? edges[j] : j;
|
||||
DType* out_off = O + eid * dim;
|
||||
for (int64_t k = 0; k < dim; ++k) {
|
||||
const int64_t lhs_add = bcast.use_bcast ? bcast.lhs_offset[k] : k;
|
||||
const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
|
||||
const DType* lhs_off =
|
||||
Op::use_lhs
|
||||
? X + Selector<LhsTarget>::Call(rid, eid, cid) * lhs_dim +
|
||||
lhs_add * reduce_size
|
||||
: nullptr;
|
||||
const DType* rhs_off =
|
||||
Op::use_rhs
|
||||
? Y + Selector<RhsTarget>::Call(rid, eid, cid) * rhs_dim +
|
||||
rhs_add * reduce_size
|
||||
: nullptr;
|
||||
out_off[k] = Op::Call(lhs_off, rhs_off, reduce_size);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CPU kernel of g-SDDMM on Coo format.
|
||||
* @param bcast Broadcast information.
|
||||
* @param coo The COO matrix.
|
||||
* @param lhs The left hand side operand feature.
|
||||
* @param rhs The right hand size operand feature.
|
||||
* @param out The result feature on edges.
|
||||
* @note it uses edge parallel strategy, different threads are responsible
|
||||
* for the computation of different edges.
|
||||
*/
|
||||
template <
|
||||
typename IdType, typename DType, typename Op, int LhsTarget = 0,
|
||||
int RhsTarget = 2>
|
||||
void SDDMMCoo(
|
||||
const BcastOff& bcast, const COOMatrix& coo, NDArray lhs, NDArray rhs,
|
||||
NDArray out) {
|
||||
const bool has_idx = !IsNullArray(coo.data);
|
||||
const IdType* row = coo.row.Ptr<IdType>();
|
||||
const IdType* col = coo.col.Ptr<IdType>();
|
||||
const IdType* edges = coo.data.Ptr<IdType>();
|
||||
const DType* X = lhs.Ptr<DType>();
|
||||
const DType* Y = rhs.Ptr<DType>();
|
||||
const int64_t dim = bcast.out_len, lhs_dim = bcast.lhs_len,
|
||||
rhs_dim = bcast.rhs_len, reduce_size = bcast.reduce_size;
|
||||
DType* O = out.Ptr<DType>();
|
||||
#pragma omp parallel for
|
||||
for (int64_t i = 0; i < coo.row->shape[0]; ++i) {
|
||||
const IdType rid = row[i];
|
||||
const IdType cid = col[i];
|
||||
const IdType eid = has_idx ? edges[i] : i;
|
||||
DType* out_off = O + eid * dim;
|
||||
for (int64_t k = 0; k < dim; ++k) {
|
||||
const int64_t lhs_add = bcast.use_bcast ? bcast.lhs_offset[k] : k;
|
||||
const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
|
||||
const DType* lhs_off =
|
||||
Op::use_lhs ? X + Selector<LhsTarget>::Call(rid, eid, cid) * lhs_dim +
|
||||
lhs_add * reduce_size
|
||||
: nullptr;
|
||||
const DType* rhs_off =
|
||||
Op::use_rhs ? Y + Selector<RhsTarget>::Call(rid, eid, cid) * rhs_dim +
|
||||
rhs_add * reduce_size
|
||||
: nullptr;
|
||||
out_off[k] = Op::Call(lhs_off, rhs_off, bcast.reduce_size);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
namespace op {
|
||||
|
||||
////////////////////////// binary operators on CPU /////////////////////////////
|
||||
template <typename DType>
|
||||
struct Add {
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = true;
|
||||
inline static DType Call(
|
||||
const DType* lhs_off, const DType* rhs_off, int64_t len = 1) {
|
||||
return *lhs_off + *rhs_off;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename DType>
|
||||
struct Sub {
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = true;
|
||||
inline static DType Call(
|
||||
const DType* lhs_off, const DType* rhs_off, int64_t len = 1) {
|
||||
return *lhs_off - *rhs_off;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename DType>
|
||||
struct Mul {
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = true;
|
||||
inline static DType Call(
|
||||
const DType* lhs_off, const DType* rhs_off, int64_t len = 1) {
|
||||
return *lhs_off * *rhs_off;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename DType>
|
||||
struct Div {
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = true;
|
||||
inline static DType Call(
|
||||
const DType* lhs_off, const DType* rhs_off, int64_t len = 1) {
|
||||
return *lhs_off / *rhs_off;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename DType>
|
||||
struct CopyLhs {
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = false;
|
||||
inline static DType Call(
|
||||
const DType* lhs_off, const DType*, int64_t len = 1) {
|
||||
return *lhs_off;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename DType>
|
||||
struct CopyRhs {
|
||||
static constexpr bool use_lhs = false;
|
||||
static constexpr bool use_rhs = true;
|
||||
inline static DType Call(
|
||||
const DType*, const DType* rhs_off, int64_t len = 1) {
|
||||
return *rhs_off;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename DType>
|
||||
struct Dot {
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = true;
|
||||
inline static DType Call(
|
||||
const DType* lhs_off, const DType* rhs_off, int64_t len = 1) {
|
||||
DType rst = 0;
|
||||
for (int64_t l = 0; l < len; ++l) {
|
||||
rst += lhs_off[l] * rhs_off[l];
|
||||
}
|
||||
return rst;
|
||||
}
|
||||
};
|
||||
|
||||
#define SWITCH_OP(op, Op, ...) \
|
||||
do { \
|
||||
if ((op) == "add") { \
|
||||
typedef dgl::aten::cpu::op::Add<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "sub") { \
|
||||
typedef dgl::aten::cpu::op::Sub<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "mul") { \
|
||||
typedef dgl::aten::cpu::op::Mul<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "div") { \
|
||||
typedef dgl::aten::cpu::op::Div<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "copy_lhs") { \
|
||||
typedef dgl::aten::cpu::op::CopyLhs<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "copy_rhs") { \
|
||||
typedef dgl::aten::cpu::op::CopyRhs<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "dot") { \
|
||||
typedef dgl::aten::cpu::op::Dot<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else { \
|
||||
LOG(FATAL) << "Unsupported SDDMM binary operator: " << op; \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
} // namespace op
|
||||
|
||||
} // namespace cpu
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CPU_SDDMM_H_
|
||||
@@ -0,0 +1,142 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file kernel/cpu/segment_reduce.cc
|
||||
* @brief Segment reduce C APIs and definitions.
|
||||
*/
|
||||
#include "./segment_reduce.h"
|
||||
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include <string>
|
||||
|
||||
#include "./spmm_binary_ops.h"
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
|
||||
/** @brief Segment Reduce operator. */
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SegmentReduce(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg) {
|
||||
if (op == "sum") {
|
||||
cpu::SegmentSum<IdType, DType>(feat, offsets, out);
|
||||
} else if (op == "max" || op == "min") {
|
||||
if (op == "max") {
|
||||
cpu::SegmentCmp<IdType, DType, cpu::op::Max<DType>>(
|
||||
feat, offsets, out, arg);
|
||||
} else {
|
||||
cpu::SegmentCmp<IdType, DType, cpu::op::Min<DType>>(
|
||||
feat, offsets, out, arg);
|
||||
}
|
||||
} else {
|
||||
LOG(FATAL) << "Unsupported reduce function " << op;
|
||||
}
|
||||
}
|
||||
|
||||
/** @brief Scatter Add.*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void ScatterAdd(NDArray feat, NDArray idx, NDArray out) {
|
||||
cpu::ScatterAdd<IdType, DType>(feat, idx, out);
|
||||
}
|
||||
|
||||
/** @brief Update gradients for reduce operator max/min on heterogeneous
|
||||
* graph.*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void UpdateGradMinMax_hetero(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out) {
|
||||
cpu::UpdateGradMinMax_hetero<IdType, DType>(g, op, feat, idx, idx_etype, out);
|
||||
}
|
||||
|
||||
/** @brief Backward function of segment cmp.*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void BackwardSegmentCmp(NDArray feat, NDArray arg, NDArray out) {
|
||||
cpu::BackwardSegmentCmp<IdType, DType>(feat, arg, out);
|
||||
}
|
||||
|
||||
template void SegmentReduce<kDGLCPU, int32_t, BFloat16>(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg);
|
||||
template void SegmentReduce<kDGLCPU, int64_t, BFloat16>(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg);
|
||||
template void SegmentReduce<kDGLCPU, int32_t, float>(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg);
|
||||
template void SegmentReduce<kDGLCPU, int64_t, float>(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg);
|
||||
template void SegmentReduce<kDGLCPU, int32_t, double>(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg);
|
||||
template void SegmentReduce<kDGLCPU, int64_t, double>(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg);
|
||||
|
||||
template <>
|
||||
void ScatterAdd<kDGLCPU, int32_t, BFloat16>(
|
||||
NDArray feat, NDArray idx, NDArray out) {
|
||||
LOG(FATAL) << "Unsupported CPU kernel for ScatterAdd for BF16.";
|
||||
}
|
||||
template <>
|
||||
void ScatterAdd<kDGLCPU, int64_t, BFloat16>(
|
||||
NDArray feat, NDArray idx, NDArray out) {
|
||||
LOG(FATAL) << "Unsupported CPU kernel for ScatterAdd for BF16.";
|
||||
}
|
||||
template void ScatterAdd<kDGLCPU, int32_t, float>(
|
||||
NDArray feat, NDArray idx, NDArray out);
|
||||
template void ScatterAdd<kDGLCPU, int64_t, float>(
|
||||
NDArray feat, NDArray idx, NDArray out);
|
||||
template void ScatterAdd<kDGLCPU, int32_t, double>(
|
||||
NDArray feat, NDArray idx, NDArray out);
|
||||
template void ScatterAdd<kDGLCPU, int64_t, double>(
|
||||
NDArray feat, NDArray arg, NDArray out);
|
||||
|
||||
template <>
|
||||
void UpdateGradMinMax_hetero<kDGLCPU, int32_t, BFloat16>(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out) {
|
||||
LOG(FATAL) << "Unsupported CPU kernel for UpdateGradMinMax_hetero for BF16.";
|
||||
}
|
||||
template <>
|
||||
void UpdateGradMinMax_hetero<kDGLCPU, int64_t, BFloat16>(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out) {
|
||||
LOG(FATAL) << "Unsupported CPU kernel for UpdateGradMinMax_hetero for BF16.";
|
||||
}
|
||||
template void UpdateGradMinMax_hetero<kDGLCPU, int32_t, float>(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out);
|
||||
template void UpdateGradMinMax_hetero<kDGLCPU, int64_t, float>(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out);
|
||||
template void UpdateGradMinMax_hetero<kDGLCPU, int32_t, double>(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out);
|
||||
template void UpdateGradMinMax_hetero<kDGLCPU, int64_t, double>(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out);
|
||||
|
||||
template void BackwardSegmentCmp<kDGLCPU, int32_t, BFloat16>(
|
||||
NDArray feat, NDArray arg, NDArray out);
|
||||
template void BackwardSegmentCmp<kDGLCPU, int64_t, BFloat16>(
|
||||
NDArray feat, NDArray arg, NDArray out);
|
||||
template void BackwardSegmentCmp<kDGLCPU, int32_t, float>(
|
||||
NDArray feat, NDArray arg, NDArray out);
|
||||
template void BackwardSegmentCmp<kDGLCPU, int64_t, float>(
|
||||
NDArray feat, NDArray arg, NDArray out);
|
||||
template void BackwardSegmentCmp<kDGLCPU, int32_t, double>(
|
||||
NDArray feat, NDArray arg, NDArray out);
|
||||
template void BackwardSegmentCmp<kDGLCPU, int64_t, double>(
|
||||
NDArray feat, NDArray arg, NDArray out);
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,194 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/spmm.h
|
||||
* @brief Segment reduce kernel function header.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CPU_SEGMENT_REDUCE_H_
|
||||
#define DGL_ARRAY_CPU_SEGMENT_REDUCE_H_
|
||||
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/base_heterograph.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace cpu {
|
||||
|
||||
/**
|
||||
* @brief CPU kernel of segment sum.
|
||||
* @param feat The input tensor.
|
||||
* @param offsets The offset tensor storing the ranges of segments.
|
||||
* @param out The output tensor.
|
||||
*/
|
||||
template <typename IdType, typename DType>
|
||||
void SegmentSum(NDArray feat, NDArray offsets, NDArray out) {
|
||||
if (std::is_same<DType, BFloat16>::value)
|
||||
LOG(FATAL) << "Unsupported CPU kernel for SegmentSum for BF16.";
|
||||
int n = out->shape[0];
|
||||
int dim = 1;
|
||||
for (int i = 1; i < out->ndim; ++i) dim *= out->shape[i];
|
||||
const DType* feat_data = feat.Ptr<DType>();
|
||||
const IdType* offsets_data = offsets.Ptr<IdType>();
|
||||
DType* out_data = out.Ptr<DType>();
|
||||
runtime::parallel_for(0, n, [=](int b, int e) {
|
||||
for (auto i = b; i < e; ++i) {
|
||||
for (IdType j = offsets_data[i]; j < offsets_data[i + 1]; ++j) {
|
||||
for (int k = 0; k < dim; ++k) {
|
||||
out_data[i * dim + k] += feat_data[j * dim + k];
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CPU kernel of segment min/max.
|
||||
* @param feat The input tensor.
|
||||
* @param offsets The offset tensor storing the ranges of segments.
|
||||
* @param out The output tensor.
|
||||
* @param arg An auxiliary tensor storing the argmin/max information
|
||||
* used in backward phase.
|
||||
*/
|
||||
template <typename IdType, typename DType, typename Cmp>
|
||||
void SegmentCmp(NDArray feat, NDArray offsets, NDArray out, NDArray arg) {
|
||||
int n = out->shape[0];
|
||||
int dim = 1;
|
||||
for (int i = 1; i < out->ndim; ++i) dim *= out->shape[i];
|
||||
const DType* feat_data = feat.Ptr<DType>();
|
||||
const IdType* offsets_data = offsets.Ptr<IdType>();
|
||||
DType* out_data = out.Ptr<DType>();
|
||||
IdType* arg_data = arg.Ptr<IdType>();
|
||||
std::fill(out_data, out_data + out.NumElements(), Cmp::zero);
|
||||
std::fill(arg_data, arg_data + arg.NumElements(), -1);
|
||||
runtime::parallel_for(0, n, [=](int b, int e) {
|
||||
for (auto i = b; i < e; ++i) {
|
||||
for (IdType j = offsets_data[i]; j < offsets_data[i + 1]; ++j) {
|
||||
for (int k = 0; k < dim; ++k) {
|
||||
const DType val = feat_data[j * dim + k];
|
||||
if (Cmp::Call(out_data[i * dim + k], val)) {
|
||||
out_data[i * dim + k] = val;
|
||||
arg_data[i * dim + k] = j;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CPU kernel of Scatter Add (on first dimension) operator.
|
||||
* @note math equation: out[idx[i], *] += feat[i, *]
|
||||
* @param feat The input tensor.
|
||||
* @param idx The indices tensor.
|
||||
* @param out The output tensor.
|
||||
*/
|
||||
template <typename IdType, typename DType>
|
||||
void ScatterAdd(NDArray feat, NDArray idx, NDArray out) {
|
||||
int n = feat->shape[0];
|
||||
int dim = 1;
|
||||
for (int i = 1; i < out->ndim; ++i) dim *= out->shape[i];
|
||||
const DType* feat_data = feat.Ptr<DType>();
|
||||
const IdType* idx_data = idx.Ptr<IdType>();
|
||||
DType* out_data = out.Ptr<DType>();
|
||||
#pragma omp parallel for
|
||||
for (int i = 0; i < n; ++i) {
|
||||
const int write_row = idx_data[i];
|
||||
for (int k = 0; k < dim; ++k) {
|
||||
#pragma omp atomic
|
||||
out_data[write_row * dim + k] += feat_data[i * dim + k];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CPU kernel to update gradients for reduce op max/min
|
||||
* @param graph The input heterogeneous graph.
|
||||
* @param op The binary operator, could be `copy_u`, `copy_e'.
|
||||
* @param list_feat List of the input tensors.
|
||||
* @param list_idx List of the indices tensors.
|
||||
* @param list_idx_etype List of the node- or edge-type tensors.
|
||||
* @param list_out List of the output tensors.
|
||||
*/
|
||||
template <typename IdType, typename DType>
|
||||
void UpdateGradMinMax_hetero(
|
||||
HeteroGraphPtr graph, const std::string& op,
|
||||
const std::vector<NDArray>& list_feat, const std::vector<NDArray>& list_idx,
|
||||
const std::vector<NDArray>& list_idx_types,
|
||||
std::vector<NDArray>* list_out) {
|
||||
if (op == "copy_lhs" || op == "copy_rhs") {
|
||||
std::vector<std::vector<dgl_id_t>> src_dst_ntypes(
|
||||
graph->NumVertexTypes(), std::vector<dgl_id_t>());
|
||||
|
||||
for (dgl_type_t etype = 0; etype < graph->NumEdgeTypes(); ++etype) {
|
||||
auto pair = graph->meta_graph()->FindEdge(etype);
|
||||
const dgl_id_t dst_ntype = pair.first; // graph is reversed
|
||||
const dgl_id_t src_ntype = pair.second;
|
||||
auto same_src_dst_ntype = std::find(
|
||||
std::begin(src_dst_ntypes[dst_ntype]),
|
||||
std::end(src_dst_ntypes[dst_ntype]), src_ntype);
|
||||
// if op is "copy_lhs", relation type with same src and dst node type will
|
||||
// be updated once
|
||||
if (op == "copy_lhs" &&
|
||||
same_src_dst_ntype != std::end(src_dst_ntypes[dst_ntype]))
|
||||
continue;
|
||||
src_dst_ntypes[dst_ntype].push_back(src_ntype);
|
||||
const DType* feat_data = list_feat[dst_ntype].Ptr<DType>();
|
||||
const IdType* idx_data = list_idx[dst_ntype].Ptr<IdType>();
|
||||
const IdType* idx_type_data = list_idx_types[dst_ntype].Ptr<IdType>();
|
||||
int type = (op == "copy_lhs") ? src_ntype : etype;
|
||||
DType* out_data = (*list_out)[type].Ptr<DType>();
|
||||
int dim = 1;
|
||||
for (int i = 1; i < (*list_out)[type]->ndim; ++i)
|
||||
dim *= (*list_out)[type]->shape[i];
|
||||
int n = list_feat[dst_ntype]->shape[0];
|
||||
#pragma omp parallel for
|
||||
for (int i = 0; i < n; ++i) {
|
||||
for (int k = 0; k < dim; ++k) {
|
||||
if (type == idx_type_data[i * dim + k]) {
|
||||
const int write_row = idx_data[i * dim + k];
|
||||
#pragma omp atomic
|
||||
out_data[write_row * dim + k] +=
|
||||
feat_data[i * dim + k]; // feat = dZ
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
LOG(FATAL) << "Unsupported binary operator: " << op;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CPU kernel of backward phase of segment min/max.
|
||||
* @note math equation: out[arg[i, k], k] = feat[i, k]
|
||||
* @param feat The input tensor.
|
||||
* @param arg The argmin/argmax tensor.
|
||||
* @param out The output tensor.
|
||||
*/
|
||||
template <typename IdType, typename DType>
|
||||
void BackwardSegmentCmp(NDArray feat, NDArray arg, NDArray out) {
|
||||
int n = feat->shape[0];
|
||||
int dim = 1;
|
||||
for (int i = 1; i < out->ndim; ++i) dim *= out->shape[i];
|
||||
const DType* feat_data = feat.Ptr<DType>();
|
||||
const IdType* arg_data = arg.Ptr<IdType>();
|
||||
DType* out_data = out.Ptr<DType>();
|
||||
runtime::parallel_for(0, n, [=](int b, int e) {
|
||||
for (auto i = b; i < e; ++i) {
|
||||
for (int k = 0; k < dim; ++k) {
|
||||
int write_row = arg_data[i * dim + k];
|
||||
if (write_row >= 0)
|
||||
out_data[write_row * dim + k] = feat_data[i * dim + k];
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
} // namespace cpu
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CPU_SEGMENT_REDUCE_H_
|
||||
@@ -0,0 +1,925 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file array/cpu/spmat_op_impl.cc
|
||||
* @brief CPU implementation of COO sparse matrix operators
|
||||
*/
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
#include <dmlc/omp.h>
|
||||
|
||||
#include <numeric>
|
||||
#include <tuple>
|
||||
#include <unordered_map>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#include "array_utils.h"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using runtime::NDArray;
|
||||
using runtime::parallel_for;
|
||||
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
/**
|
||||
* TODO(BarclayII):
|
||||
* For row-major sorted COOs, we have faster implementation with binary search,
|
||||
* sorted search, etc. Later we should benchmark how much we can gain with
|
||||
* sorted COOs on hypersparse graphs.
|
||||
*/
|
||||
|
||||
///////////////////////////// COOIsNonZero /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
bool COOIsNonZero(COOMatrix coo, int64_t row, int64_t col) {
|
||||
CHECK(row >= 0 && row < coo.num_rows) << "Invalid row index: " << row;
|
||||
CHECK(col >= 0 && col < coo.num_cols) << "Invalid col index: " << col;
|
||||
const IdType *coo_row_data = static_cast<IdType *>(coo.row->data);
|
||||
const IdType *coo_col_data = static_cast<IdType *>(coo.col->data);
|
||||
for (int64_t i = 0; i < coo.row->shape[0]; ++i) {
|
||||
if (coo_row_data[i] == row && coo_col_data[i] == col) return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
template bool COOIsNonZero<kDGLCPU, int32_t>(COOMatrix, int64_t, int64_t);
|
||||
template bool COOIsNonZero<kDGLCPU, int64_t>(COOMatrix, int64_t, int64_t);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
NDArray COOIsNonZero(COOMatrix coo, NDArray row, NDArray col) {
|
||||
const auto rowlen = row->shape[0];
|
||||
const auto collen = col->shape[0];
|
||||
const auto rstlen = std::max(rowlen, collen);
|
||||
NDArray rst = NDArray::Empty({rstlen}, row->dtype, row->ctx);
|
||||
IdType *rst_data = static_cast<IdType *>(rst->data);
|
||||
const IdType *row_data = static_cast<IdType *>(row->data);
|
||||
const IdType *col_data = static_cast<IdType *>(col->data);
|
||||
const int64_t row_stride = (rowlen == 1 && collen != 1) ? 0 : 1;
|
||||
const int64_t col_stride = (collen == 1 && rowlen != 1) ? 0 : 1;
|
||||
const int64_t kmax = std::max(rowlen, collen);
|
||||
parallel_for(0, kmax, [=](size_t b, size_t e) {
|
||||
for (auto k = b; k < e; ++k) {
|
||||
int64_t i = row_stride * k;
|
||||
int64_t j = col_stride * k;
|
||||
rst_data[k] =
|
||||
COOIsNonZero<XPU, IdType>(coo, row_data[i], col_data[j]) ? 1 : 0;
|
||||
}
|
||||
});
|
||||
return rst;
|
||||
}
|
||||
|
||||
template NDArray COOIsNonZero<kDGLCPU, int32_t>(COOMatrix, NDArray, NDArray);
|
||||
template NDArray COOIsNonZero<kDGLCPU, int64_t>(COOMatrix, NDArray, NDArray);
|
||||
|
||||
///////////////////////////// COOHasDuplicate /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
bool COOHasDuplicate(COOMatrix coo) {
|
||||
std::unordered_set<std::pair<IdType, IdType>, PairHash> hashmap;
|
||||
const IdType *src_data = static_cast<IdType *>(coo.row->data);
|
||||
const IdType *dst_data = static_cast<IdType *>(coo.col->data);
|
||||
const auto nnz = coo.row->shape[0];
|
||||
for (IdType eid = 0; eid < nnz; ++eid) {
|
||||
const auto &p = std::make_pair(src_data[eid], dst_data[eid]);
|
||||
if (hashmap.count(p)) {
|
||||
return true;
|
||||
} else {
|
||||
hashmap.insert(p);
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
template bool COOHasDuplicate<kDGLCPU, int32_t>(COOMatrix coo);
|
||||
template bool COOHasDuplicate<kDGLCPU, int64_t>(COOMatrix coo);
|
||||
|
||||
///////////////////////////// COOGetRowNNZ /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
int64_t COOGetRowNNZ(COOMatrix coo, int64_t row) {
|
||||
CHECK(row >= 0 && row < coo.num_rows) << "Invalid row index: " << row;
|
||||
const IdType *coo_row_data = static_cast<IdType *>(coo.row->data);
|
||||
int64_t result = 0;
|
||||
for (int64_t i = 0; i < coo.row->shape[0]; ++i) {
|
||||
if (coo_row_data[i] == row) ++result;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
template int64_t COOGetRowNNZ<kDGLCPU, int32_t>(COOMatrix, int64_t);
|
||||
template int64_t COOGetRowNNZ<kDGLCPU, int64_t>(COOMatrix, int64_t);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
NDArray COOGetRowNNZ(COOMatrix coo, NDArray rows) {
|
||||
CHECK_SAME_DTYPE(coo.col, rows);
|
||||
const auto len = rows->shape[0];
|
||||
const IdType *vid_data = static_cast<IdType *>(rows->data);
|
||||
NDArray rst = NDArray::Empty({len}, rows->dtype, rows->ctx);
|
||||
IdType *rst_data = static_cast<IdType *>(rst->data);
|
||||
#pragma omp parallel for
|
||||
for (int64_t i = 0; i < len; ++i) {
|
||||
rst_data[i] = COOGetRowNNZ<XPU, IdType>(coo, vid_data[i]);
|
||||
}
|
||||
return rst;
|
||||
}
|
||||
|
||||
template NDArray COOGetRowNNZ<kDGLCPU, int32_t>(COOMatrix, NDArray);
|
||||
template NDArray COOGetRowNNZ<kDGLCPU, int64_t>(COOMatrix, NDArray);
|
||||
|
||||
////////////////////////// COOGetRowDataAndIndices /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::pair<NDArray, NDArray> COOGetRowDataAndIndices(
|
||||
COOMatrix coo, int64_t row) {
|
||||
CHECK(row >= 0 && row < coo.num_rows) << "Invalid row index: " << row;
|
||||
|
||||
const IdType *coo_row_data = static_cast<IdType *>(coo.row->data);
|
||||
const IdType *coo_col_data = static_cast<IdType *>(coo.col->data);
|
||||
const IdType *coo_data =
|
||||
COOHasData(coo) ? static_cast<IdType *>(coo.data->data) : nullptr;
|
||||
|
||||
std::vector<IdType> indices;
|
||||
std::vector<IdType> data;
|
||||
|
||||
for (int64_t i = 0; i < coo.row->shape[0]; ++i) {
|
||||
if (coo_row_data[i] == row) {
|
||||
indices.push_back(coo_col_data[i]);
|
||||
data.push_back(coo_data ? coo_data[i] : i);
|
||||
}
|
||||
}
|
||||
|
||||
return std::make_pair(
|
||||
NDArray::FromVector(data), NDArray::FromVector(indices));
|
||||
}
|
||||
|
||||
template std::pair<NDArray, NDArray> COOGetRowDataAndIndices<kDGLCPU, int32_t>(
|
||||
COOMatrix, int64_t);
|
||||
template std::pair<NDArray, NDArray> COOGetRowDataAndIndices<kDGLCPU, int64_t>(
|
||||
COOMatrix, int64_t);
|
||||
|
||||
///////////////////////////// COOGetData /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray COOGetData(COOMatrix coo, IdArray rows, IdArray cols) {
|
||||
const int64_t rowlen = rows->shape[0];
|
||||
const int64_t collen = cols->shape[0];
|
||||
CHECK((rowlen == collen) || (rowlen == 1) || (collen == 1))
|
||||
<< "Invalid row and col Id array:" << rows << " " << cols;
|
||||
const int64_t row_stride = (rowlen == 1 && collen != 1) ? 0 : 1;
|
||||
const int64_t col_stride = (collen == 1 && rowlen != 1) ? 0 : 1;
|
||||
const IdType *row_data = rows.Ptr<IdType>();
|
||||
const IdType *col_data = cols.Ptr<IdType>();
|
||||
|
||||
const IdType *coo_row = coo.row.Ptr<IdType>();
|
||||
const IdType *coo_col = coo.col.Ptr<IdType>();
|
||||
const IdType *data = COOHasData(coo) ? coo.data.Ptr<IdType>() : nullptr;
|
||||
const int64_t nnz = coo.row->shape[0];
|
||||
|
||||
const int64_t retlen = std::max(rowlen, collen);
|
||||
IdArray ret = Full(-1, retlen, rows->dtype.bits, rows->ctx);
|
||||
IdType *ret_data = ret.Ptr<IdType>();
|
||||
|
||||
// TODO(minjie): We might need to consider sorting the COO beforehand
|
||||
// especially when the number of (row, col) pairs is large. Need more
|
||||
// benchmarks to justify the choice.
|
||||
|
||||
if (coo.row_sorted) {
|
||||
parallel_for(0, retlen, [&](size_t b, size_t e) {
|
||||
for (auto p = b; p < e; ++p) {
|
||||
const IdType row_id = row_data[p * row_stride],
|
||||
col_id = col_data[p * col_stride];
|
||||
auto it = std::lower_bound(coo_row, coo_row + nnz, row_id);
|
||||
for (; it < coo_row + nnz && *it == row_id; ++it) {
|
||||
const auto idx = it - coo_row;
|
||||
if (coo_col[idx] == col_id) {
|
||||
ret_data[p] = data ? data[idx] : idx;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
} else {
|
||||
#pragma omp parallel for
|
||||
for (int64_t p = 0; p < retlen; ++p) {
|
||||
const IdType row_id = row_data[p * row_stride],
|
||||
col_id = col_data[p * col_stride];
|
||||
for (int64_t idx = 0; idx < nnz; ++idx) {
|
||||
if (coo_row[idx] == row_id && coo_col[idx] == col_id) {
|
||||
ret_data[p] = data ? data[idx] : idx;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return ret;
|
||||
}
|
||||
|
||||
template IdArray COOGetData<kDGLCPU, int32_t>(COOMatrix, IdArray, IdArray);
|
||||
template IdArray COOGetData<kDGLCPU, int64_t>(COOMatrix, IdArray, IdArray);
|
||||
|
||||
///////////////////////////// COOGetDataAndIndices /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::vector<NDArray> COOGetDataAndIndices(
|
||||
COOMatrix coo, NDArray rows, NDArray cols) {
|
||||
CHECK_SAME_DTYPE(coo.col, rows);
|
||||
CHECK_SAME_DTYPE(coo.col, cols);
|
||||
const int64_t rowlen = rows->shape[0];
|
||||
const int64_t collen = cols->shape[0];
|
||||
const int64_t len = std::max(rowlen, collen);
|
||||
|
||||
CHECK((rowlen == collen) || (rowlen == 1) || (collen == 1))
|
||||
<< "Invalid row and col id array.";
|
||||
|
||||
const int64_t row_stride = (rowlen == 1 && collen != 1) ? 0 : 1;
|
||||
const int64_t col_stride = (collen == 1 && rowlen != 1) ? 0 : 1;
|
||||
const IdType *row_data = static_cast<IdType *>(rows->data);
|
||||
const IdType *col_data = static_cast<IdType *>(cols->data);
|
||||
|
||||
const IdType *coo_row_data = static_cast<IdType *>(coo.row->data);
|
||||
const IdType *coo_col_data = static_cast<IdType *>(coo.col->data);
|
||||
const IdType *data =
|
||||
COOHasData(coo) ? static_cast<IdType *>(coo.data->data) : nullptr;
|
||||
|
||||
std::vector<IdType> ret_rows, ret_cols;
|
||||
std::vector<IdType> ret_data;
|
||||
ret_rows.reserve(len);
|
||||
ret_cols.reserve(len);
|
||||
ret_data.reserve(len);
|
||||
|
||||
// NOTE(BarclayII): With a small number of lookups, linear scan is faster.
|
||||
// The threshold 200 comes from benchmarking both algorithms on a P3.8x
|
||||
// instance. I also tried sorting plus binary search. The speed gain is only
|
||||
// significant for medium-sized graphs and lookups, so I didn't include it.
|
||||
if (len >= 200) {
|
||||
// TODO(BarclayII) Ideally we would want to cache this object. However I'm
|
||||
// not sure what is the best way to do so since this object is valid for CPU
|
||||
// only.
|
||||
std::unordered_multimap<std::pair<IdType, IdType>, IdType, PairHash>
|
||||
pair_map;
|
||||
pair_map.reserve(coo.row->shape[0]);
|
||||
for (int64_t k = 0; k < coo.row->shape[0]; ++k)
|
||||
pair_map.emplace(
|
||||
std::make_pair(coo_row_data[k], coo_col_data[k]), data ? data[k] : k);
|
||||
|
||||
for (int64_t i = 0, j = 0; i < rowlen && j < collen;
|
||||
i += row_stride, j += col_stride) {
|
||||
const IdType row_id = row_data[i], col_id = col_data[j];
|
||||
CHECK(row_id >= 0 && row_id < coo.num_rows)
|
||||
<< "Invalid row index: " << row_id;
|
||||
CHECK(col_id >= 0 && col_id < coo.num_cols)
|
||||
<< "Invalid col index: " << col_id;
|
||||
auto range = pair_map.equal_range({row_id, col_id});
|
||||
for (auto it = range.first; it != range.second; ++it) {
|
||||
ret_rows.push_back(row_id);
|
||||
ret_cols.push_back(col_id);
|
||||
ret_data.push_back(it->second);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for (int64_t i = 0, j = 0; i < rowlen && j < collen;
|
||||
i += row_stride, j += col_stride) {
|
||||
const IdType row_id = row_data[i], col_id = col_data[j];
|
||||
CHECK(row_id >= 0 && row_id < coo.num_rows)
|
||||
<< "Invalid row index: " << row_id;
|
||||
CHECK(col_id >= 0 && col_id < coo.num_cols)
|
||||
<< "Invalid col index: " << col_id;
|
||||
for (int64_t k = 0; k < coo.row->shape[0]; ++k) {
|
||||
if (coo_row_data[k] == row_id && coo_col_data[k] == col_id) {
|
||||
ret_rows.push_back(row_id);
|
||||
ret_cols.push_back(col_id);
|
||||
ret_data.push_back(data ? data[k] : k);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
NDArray::FromVector(ret_rows), NDArray::FromVector(ret_cols),
|
||||
NDArray::FromVector(ret_data)};
|
||||
}
|
||||
|
||||
template std::vector<NDArray> COOGetDataAndIndices<kDGLCPU, int32_t>(
|
||||
COOMatrix coo, NDArray rows, NDArray cols);
|
||||
template std::vector<NDArray> COOGetDataAndIndices<kDGLCPU, int64_t>(
|
||||
COOMatrix coo, NDArray rows, NDArray cols);
|
||||
|
||||
///////////////////////////// COOTranspose /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix COOTranspose(COOMatrix coo) {
|
||||
return COOMatrix{coo.num_cols, coo.num_rows, coo.col, coo.row, coo.data};
|
||||
}
|
||||
|
||||
template COOMatrix COOTranspose<kDGLCPU, int32_t>(COOMatrix coo);
|
||||
template COOMatrix COOTranspose<kDGLCPU, int64_t>(COOMatrix coo);
|
||||
|
||||
///////////////////////////// COOToCSR /////////////////////////////
|
||||
namespace {
|
||||
|
||||
template <class IdType>
|
||||
CSRMatrix SortedCOOToCSR(const COOMatrix &coo) {
|
||||
const int64_t N = coo.num_rows;
|
||||
const int64_t NNZ = coo.row->shape[0];
|
||||
const IdType *const row_data = static_cast<IdType *>(coo.row->data);
|
||||
const IdType *const data =
|
||||
COOHasData(coo) ? static_cast<IdType *>(coo.data->data) : nullptr;
|
||||
|
||||
NDArray ret_indptr = NDArray::Empty({N + 1}, coo.row->dtype, coo.row->ctx);
|
||||
NDArray ret_indices = coo.col;
|
||||
NDArray ret_data = data == nullptr
|
||||
? NDArray::Empty({NNZ}, coo.row->dtype, coo.row->ctx)
|
||||
: coo.data;
|
||||
|
||||
// compute indptr
|
||||
IdType *const Bp = static_cast<IdType *>(ret_indptr->data);
|
||||
Bp[0] = 0;
|
||||
|
||||
IdType *const fill_data =
|
||||
data ? nullptr : static_cast<IdType *>(ret_data->data);
|
||||
|
||||
if (NNZ > 0) {
|
||||
auto num_threads = omp_get_max_threads();
|
||||
parallel_for(0, num_threads, [&](int b, int e) {
|
||||
for (auto thread_id = b; thread_id < e; ++thread_id) {
|
||||
// We partition the set the of non-zeros among the threads
|
||||
const int64_t nz_chunk = (NNZ + num_threads - 1) / num_threads;
|
||||
const int64_t nz_start = thread_id * nz_chunk;
|
||||
const int64_t nz_end = std::min(NNZ, nz_start + nz_chunk);
|
||||
|
||||
// Each thread searchs the row array for a change, and marks it's
|
||||
// location in Bp. Threads, other than the first, start at the last
|
||||
// index covered by the previous, in order to detect changes in the row
|
||||
// array between thread partitions. This means that each thread after
|
||||
// the first, searches the range [nz_start-1, nz_end). That is,
|
||||
// if we had 10 non-zeros, and 4 threads, the indexes searched by each
|
||||
// thread would be:
|
||||
// 0: [0, 1, 2]
|
||||
// 1: [2, 3, 4, 5]
|
||||
// 2: [5, 6, 7, 8]
|
||||
// 3: [8, 9]
|
||||
//
|
||||
// That way, if the row array were [0, 0, 1, 2, 2, 2, 4, 5, 5, 6], each
|
||||
// change in row would be captured by one thread:
|
||||
//
|
||||
// 0: [0, 0, 1] - row 0
|
||||
// 1: [1, 2, 2, 2] - row 1
|
||||
// 2: [2, 4, 5, 5] - rows 2, 3, and 4
|
||||
// 3: [5, 6] - rows 5 and 6
|
||||
//
|
||||
int64_t row = 0;
|
||||
if (nz_start < nz_end) {
|
||||
row = nz_start == 0 ? 0 : row_data[nz_start - 1];
|
||||
for (int64_t i = nz_start; i < nz_end; ++i) {
|
||||
while (row != row_data[i]) {
|
||||
++row;
|
||||
Bp[row] = i;
|
||||
}
|
||||
}
|
||||
|
||||
// We will not detect the row change for the last row, nor any empty
|
||||
// rows at the end of the matrix, so the last active thread needs
|
||||
// mark all remaining rows in Bp with NNZ.
|
||||
if (nz_end == NNZ) {
|
||||
while (row < N) {
|
||||
++row;
|
||||
Bp[row] = NNZ;
|
||||
}
|
||||
}
|
||||
|
||||
if (fill_data) {
|
||||
// TODO(minjie): Many of our current implementation assumes that CSR
|
||||
// must have
|
||||
// a data array. This is a temporary workaround. Remove this
|
||||
// after:
|
||||
// - The old immutable graph implementation is deprecated.
|
||||
// - The old binary reduce kernel is deprecated.
|
||||
std::iota(fill_data + nz_start, fill_data + nz_end, nz_start);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
} else {
|
||||
std::fill(Bp, Bp + N + 1, 0);
|
||||
}
|
||||
|
||||
return CSRMatrix(
|
||||
coo.num_rows, coo.num_cols, ret_indptr, ret_indices, ret_data,
|
||||
coo.col_sorted);
|
||||
}
|
||||
|
||||
template <class IdType>
|
||||
CSRMatrix UnSortedSparseCOOToCSR(const COOMatrix &coo) {
|
||||
// Unsigned version of the original integer index data type.
|
||||
// It avoids overflow in (N + num_threads) and (n_start + n_chunk) below.
|
||||
typedef typename std::make_unsigned<IdType>::type UIdType;
|
||||
|
||||
const UIdType N = coo.num_rows;
|
||||
const int64_t NNZ = coo.row->shape[0];
|
||||
const IdType *const row_data = static_cast<IdType *>(coo.row->data);
|
||||
const IdType *const col_data = static_cast<IdType *>(coo.col->data);
|
||||
const IdType *const data =
|
||||
COOHasData(coo) ? static_cast<IdType *>(coo.data->data) : nullptr;
|
||||
|
||||
NDArray ret_indptr = NDArray::Empty(
|
||||
{static_cast<int64_t>(N) + 1}, coo.row->dtype, coo.row->ctx);
|
||||
NDArray ret_indices = NDArray::Empty({NNZ}, coo.row->dtype, coo.row->ctx);
|
||||
NDArray ret_data = NDArray::Empty({NNZ}, coo.row->dtype, coo.row->ctx);
|
||||
IdType *const Bp = static_cast<IdType *>(ret_indptr->data);
|
||||
Bp[N] = 0;
|
||||
IdType *const Bi = static_cast<IdType *>(ret_indices->data);
|
||||
IdType *const Bx = static_cast<IdType *>(ret_data->data);
|
||||
|
||||
// store sorted data and original index.
|
||||
NDArray sorted_data = NDArray::Empty({NNZ}, coo.row->dtype, coo.row->ctx);
|
||||
NDArray sorted_data_pos = NDArray::Empty({NNZ}, coo.row->dtype, coo.row->ctx);
|
||||
IdType *const Sx = static_cast<IdType *>(sorted_data->data);
|
||||
IdType *const Si = static_cast<IdType *>(sorted_data_pos->data);
|
||||
|
||||
// Lower number of threads if cost of parallelization is grater than gain
|
||||
// from making calculation parallel.
|
||||
const int64_t min_chunk_size = 1000;
|
||||
const int64_t num_threads_for_batch = 2 + (NNZ + N) / min_chunk_size;
|
||||
const int num_threads_required = std::min(
|
||||
static_cast<int64_t>(omp_get_max_threads()), num_threads_for_batch);
|
||||
|
||||
// record row_idx in each thread.
|
||||
std::vector<std::vector<int64_t>> p_sum(
|
||||
num_threads_required, std::vector<int64_t>(num_threads_required));
|
||||
|
||||
#pragma omp parallel num_threads(num_threads_required)
|
||||
{
|
||||
const int num_threads = omp_get_num_threads();
|
||||
const int thread_id = omp_get_thread_num();
|
||||
CHECK_LT(thread_id, num_threads);
|
||||
|
||||
const int64_t nz_chunk = (NNZ + num_threads - 1) / num_threads;
|
||||
const int64_t nz_start = thread_id * nz_chunk;
|
||||
const int64_t nz_end = std::min(NNZ, nz_start + nz_chunk);
|
||||
|
||||
const UIdType n_chunk = (N + num_threads - 1) / num_threads;
|
||||
const UIdType n_start = thread_id * n_chunk;
|
||||
const UIdType n_end = std::min(N, n_start + n_chunk);
|
||||
|
||||
for (auto i = n_start; i < n_end; ++i) {
|
||||
Bp[i] = 0;
|
||||
}
|
||||
|
||||
// iterate on NNZ data and count row_idx.
|
||||
for (auto i = nz_start; i < nz_end; ++i) {
|
||||
const IdType row_idx = row_data[i];
|
||||
const IdType row_thread_id = row_idx / n_chunk;
|
||||
++p_sum[thread_id][row_thread_id];
|
||||
}
|
||||
|
||||
#pragma omp barrier
|
||||
#pragma omp master
|
||||
// accumulate row_idx.
|
||||
{
|
||||
int64_t cum = 0;
|
||||
for (int j = 0; j < num_threads; ++j) {
|
||||
for (int i = 0; i < num_threads; ++i) {
|
||||
auto tmp = p_sum[i][j];
|
||||
p_sum[i][j] = cum;
|
||||
cum += tmp;
|
||||
}
|
||||
}
|
||||
CHECK_EQ(cum, NNZ);
|
||||
}
|
||||
#pragma omp barrier
|
||||
const int64_t i_start = p_sum[0][thread_id];
|
||||
const int64_t i_end =
|
||||
thread_id + 1 == num_threads ? NNZ : p_sum[0][thread_id + 1];
|
||||
#pragma omp barrier
|
||||
|
||||
// sort data by row_idx and place into Sx/Si.
|
||||
auto &data_pos = p_sum[thread_id];
|
||||
for (auto i = nz_start; i < nz_end; ++i) {
|
||||
const IdType row_idx = row_data[i];
|
||||
const IdType row_thread_id = row_idx / n_chunk;
|
||||
const int64_t pos = data_pos[row_thread_id]++;
|
||||
Sx[pos] = data == nullptr ? i : data[i];
|
||||
Si[pos] = i;
|
||||
}
|
||||
|
||||
#pragma omp barrier
|
||||
|
||||
// Now we're able to do coo2csr on sorted data in each thread in parallel.
|
||||
// compute data number on each row_idx.
|
||||
for (auto i = i_start; i < i_end; ++i) {
|
||||
const UIdType row_idx = row_data[Si[i]];
|
||||
++Bp[row_idx + 1];
|
||||
}
|
||||
|
||||
// accumulate on each row
|
||||
IdType cumsum = i_start;
|
||||
for (auto i = n_start + 1; i <= n_end; ++i) {
|
||||
const auto tmp = Bp[i];
|
||||
Bp[i] = cumsum;
|
||||
cumsum += tmp;
|
||||
}
|
||||
|
||||
// update Bi/Bp/Bx
|
||||
for (auto i = i_start; i < i_end; ++i) {
|
||||
const UIdType row_idx = row_data[Si[i]];
|
||||
const int64_t dest = (Bp[row_idx + 1]++);
|
||||
Bi[dest] = col_data[Si[i]];
|
||||
Bx[dest] = Sx[i];
|
||||
}
|
||||
}
|
||||
return CSRMatrix(
|
||||
coo.num_rows, coo.num_cols, ret_indptr, ret_indices, ret_data,
|
||||
coo.col_sorted);
|
||||
}
|
||||
|
||||
template <class IdType>
|
||||
CSRMatrix UnSortedDenseCOOToCSR(const COOMatrix &coo) {
|
||||
// Unsigned version of the original integer index data type.
|
||||
// It avoids overflow in (N + num_threads) and (n_start + n_chunk) below.
|
||||
typedef typename std::make_unsigned<IdType>::type UIdType;
|
||||
|
||||
const UIdType N = coo.num_rows;
|
||||
const int64_t NNZ = coo.row->shape[0];
|
||||
const IdType *const row_data = static_cast<IdType *>(coo.row->data);
|
||||
const IdType *const col_data = static_cast<IdType *>(coo.col->data);
|
||||
const IdType *const data =
|
||||
COOHasData(coo) ? static_cast<IdType *>(coo.data->data) : nullptr;
|
||||
|
||||
NDArray ret_indptr = NDArray::Empty(
|
||||
{static_cast<int64_t>(N) + 1}, coo.row->dtype, coo.row->ctx);
|
||||
NDArray ret_indices = NDArray::Empty({NNZ}, coo.row->dtype, coo.row->ctx);
|
||||
NDArray ret_data = NDArray::Empty({NNZ}, coo.row->dtype, coo.row->ctx);
|
||||
IdType *const Bp = static_cast<IdType *>(ret_indptr->data);
|
||||
Bp[0] = 0;
|
||||
IdType *const Bi = static_cast<IdType *>(ret_indices->data);
|
||||
IdType *const Bx = static_cast<IdType *>(ret_data->data);
|
||||
|
||||
// the offset within each row, that each thread will write to
|
||||
std::vector<std::vector<IdType>> local_ptrs;
|
||||
std::vector<int64_t> thread_prefixsum;
|
||||
|
||||
#pragma omp parallel
|
||||
{
|
||||
const int num_threads = omp_get_num_threads();
|
||||
const int thread_id = omp_get_thread_num();
|
||||
CHECK_LT(thread_id, num_threads);
|
||||
|
||||
const int64_t nz_chunk = (NNZ + num_threads - 1) / num_threads;
|
||||
const int64_t nz_start = thread_id * nz_chunk;
|
||||
const int64_t nz_end = std::min(NNZ, nz_start + nz_chunk);
|
||||
|
||||
const UIdType n_chunk = (N + num_threads - 1) / num_threads;
|
||||
const UIdType n_start = thread_id * n_chunk;
|
||||
const UIdType n_end = std::min(N, n_start + n_chunk);
|
||||
|
||||
#pragma omp master
|
||||
{
|
||||
local_ptrs.resize(num_threads);
|
||||
thread_prefixsum.resize(num_threads + 1);
|
||||
}
|
||||
|
||||
#pragma omp barrier
|
||||
local_ptrs[thread_id].resize(N, 0);
|
||||
|
||||
for (int64_t i = nz_start; i < nz_end; ++i) {
|
||||
++local_ptrs[thread_id][row_data[i]];
|
||||
}
|
||||
|
||||
#pragma omp barrier
|
||||
// compute prefixsum in parallel
|
||||
int64_t sum = 0;
|
||||
for (UIdType i = n_start; i < n_end; ++i) {
|
||||
IdType tmp = 0;
|
||||
for (int j = 0; j < num_threads; ++j) {
|
||||
auto previous = local_ptrs[j][i];
|
||||
local_ptrs[j][i] = tmp;
|
||||
tmp += previous;
|
||||
}
|
||||
sum += tmp;
|
||||
Bp[i + 1] = sum;
|
||||
}
|
||||
thread_prefixsum[thread_id + 1] = sum;
|
||||
|
||||
#pragma omp barrier
|
||||
#pragma omp master
|
||||
{
|
||||
for (int i = 0; i < num_threads; ++i) {
|
||||
thread_prefixsum[i + 1] += thread_prefixsum[i];
|
||||
}
|
||||
CHECK_EQ(thread_prefixsum[num_threads], NNZ);
|
||||
}
|
||||
#pragma omp barrier
|
||||
|
||||
sum = thread_prefixsum[thread_id];
|
||||
for (UIdType i = n_start; i < n_end; ++i) {
|
||||
Bp[i + 1] += sum;
|
||||
}
|
||||
|
||||
#pragma omp barrier
|
||||
for (int64_t i = nz_start; i < nz_end; ++i) {
|
||||
const IdType r = row_data[i];
|
||||
const int64_t index = Bp[r] + local_ptrs[thread_id][r]++;
|
||||
Bi[index] = col_data[i];
|
||||
Bx[index] = data ? data[i] : i;
|
||||
}
|
||||
}
|
||||
CHECK_EQ(Bp[N], NNZ);
|
||||
|
||||
return CSRMatrix(
|
||||
coo.num_rows, coo.num_cols, ret_indptr, ret_indices, ret_data,
|
||||
coo.col_sorted);
|
||||
}
|
||||
|
||||
// complexity: time O(NNZ), space O(1)
|
||||
template <typename IdType>
|
||||
CSRMatrix UnSortedSmallCOOToCSR(COOMatrix coo) {
|
||||
const int64_t N = coo.num_rows;
|
||||
const int64_t NNZ = coo.row->shape[0];
|
||||
const IdType *row_data = static_cast<IdType *>(coo.row->data);
|
||||
const IdType *col_data = static_cast<IdType *>(coo.col->data);
|
||||
const IdType *data =
|
||||
COOHasData(coo) ? static_cast<IdType *>(coo.data->data) : nullptr;
|
||||
NDArray ret_indptr = NDArray::Empty({N + 1}, coo.row->dtype, coo.row->ctx);
|
||||
NDArray ret_indices = NDArray::Empty({NNZ}, coo.row->dtype, coo.row->ctx);
|
||||
NDArray ret_data = NDArray::Empty({NNZ}, coo.row->dtype, coo.row->ctx);
|
||||
IdType *Bp = static_cast<IdType *>(ret_indptr->data);
|
||||
IdType *Bi = static_cast<IdType *>(ret_indices->data);
|
||||
IdType *Bx = static_cast<IdType *>(ret_data->data);
|
||||
|
||||
// Count elements in each row
|
||||
std::fill(Bp, Bp + N, 0);
|
||||
for (int64_t i = 0; i < NNZ; ++i) {
|
||||
Bp[row_data[i]]++;
|
||||
}
|
||||
|
||||
// Convert to indexes
|
||||
for (IdType i = 0, cumsum = 0; i < N; ++i) {
|
||||
const IdType temp = Bp[i];
|
||||
Bp[i] = cumsum;
|
||||
cumsum += temp;
|
||||
}
|
||||
|
||||
for (int64_t i = 0; i < NNZ; ++i) {
|
||||
const IdType r = row_data[i];
|
||||
Bi[Bp[r]] = col_data[i];
|
||||
Bx[Bp[r]] = data ? data[i] : i;
|
||||
Bp[r]++;
|
||||
}
|
||||
|
||||
// Restore the indptr
|
||||
for (int64_t i = N; i > 0; --i) {
|
||||
Bp[i] = Bp[i - 1];
|
||||
}
|
||||
Bp[0] = 0;
|
||||
|
||||
return CSRMatrix(
|
||||
coo.num_rows, coo.num_cols, ret_indptr, ret_indices, ret_data,
|
||||
coo.col_sorted);
|
||||
}
|
||||
|
||||
enum class COOToCSRAlg {
|
||||
sorted = 0,
|
||||
unsortedSmall,
|
||||
unsortedSparse,
|
||||
unsortedDense
|
||||
};
|
||||
|
||||
/**
|
||||
* Chose COO to CSR format conversion algorithm for given COO matrix according
|
||||
* to heuristic based on measured performance.
|
||||
*
|
||||
* Implementation and complexity details. N: num_nodes, NNZ: num_edges, P:
|
||||
* num_threads.
|
||||
* 1. If row is sorted in COO, SortedCOOToCSR<> is applied. Time: O(NNZ/P),
|
||||
* space: O(1).
|
||||
* 2 If row is NOT sorted in COO and graph is small (small number of NNZ),
|
||||
* UnSortedSmallCOOToCSR<> is applied. Time: O(NNZ), space O(N).
|
||||
* 3 If row is NOT sorted in COO and graph is sparse (low average degree),
|
||||
* UnSortedSparseCOOToCSR<> is applied. Time: O(NNZ/P + N/P + P^2),
|
||||
* space O(NNZ + P^2).
|
||||
* 4. If row is NOT sorted in COO and graph is dense (medium/high average
|
||||
* degree), UnSortedDenseCOOToCSR<> is applied. Time: O(NNZ/P + N/P),
|
||||
* space O(NNZ + N*P).
|
||||
*
|
||||
* Note:
|
||||
* If you change this function, change also _TestCOOToCSRAlgs in
|
||||
* tests/cpp/test_spmat_coo.cc
|
||||
*/
|
||||
template <typename IdType>
|
||||
inline COOToCSRAlg WhichCOOToCSR(const COOMatrix &coo) {
|
||||
if (coo.row_sorted) {
|
||||
return COOToCSRAlg::sorted;
|
||||
} else {
|
||||
#ifdef _WIN32
|
||||
// On Windows omp_get_max_threads() gives larger value than later OMP can
|
||||
// spawn.
|
||||
int64_t num_threads;
|
||||
#pragma omp parallel
|
||||
#pragma master
|
||||
{ num_threads = omp_get_num_threads(); }
|
||||
#else
|
||||
const int64_t num_threads = omp_get_max_threads();
|
||||
#endif
|
||||
const int64_t N = coo.num_rows;
|
||||
const int64_t NNZ = coo.row->shape[0];
|
||||
// Parameters below are heuristically chosen according to measured
|
||||
// performance.
|
||||
const int64_t type_scale = sizeof(IdType) >> 1;
|
||||
const int64_t small = 50 * num_threads * type_scale * type_scale;
|
||||
if (NNZ < small || num_threads == 1) {
|
||||
// For relatively small number of non zero elements cost of spread
|
||||
// algorithm between threads is bigger than improvements from using
|
||||
// many cores
|
||||
return COOToCSRAlg::unsortedSmall;
|
||||
} else if (type_scale * NNZ < num_threads * N) {
|
||||
// For relatively small number of non zero elements in matrix, sparse
|
||||
// parallel version of algorithm is more efficient than dense.
|
||||
return COOToCSRAlg::unsortedSparse;
|
||||
}
|
||||
return COOToCSRAlg::unsortedDense;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix COOToCSR(COOMatrix coo) {
|
||||
CHECK_NO_OVERFLOW(coo.row->dtype, coo.row->shape[0]);
|
||||
switch (WhichCOOToCSR<IdType>(coo)) {
|
||||
case COOToCSRAlg::sorted:
|
||||
return SortedCOOToCSR<IdType>(coo);
|
||||
case COOToCSRAlg::unsortedSmall:
|
||||
default:
|
||||
return UnSortedSmallCOOToCSR<IdType>(coo);
|
||||
case COOToCSRAlg::unsortedSparse:
|
||||
return UnSortedSparseCOOToCSR<IdType>(coo);
|
||||
case COOToCSRAlg::unsortedDense:
|
||||
return UnSortedDenseCOOToCSR<IdType>(coo);
|
||||
}
|
||||
}
|
||||
|
||||
template CSRMatrix COOToCSR<kDGLCPU, int32_t>(COOMatrix coo);
|
||||
template CSRMatrix COOToCSR<kDGLCPU, int64_t>(COOMatrix coo);
|
||||
|
||||
///////////////////////////// COOSliceRows /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix COOSliceRows(COOMatrix coo, int64_t start, int64_t end) {
|
||||
// TODO(minjie): use binary search when coo.row_sorted is true
|
||||
CHECK(start >= 0 && start < coo.num_rows) << "Invalid start row " << start;
|
||||
CHECK(end > 0 && end <= coo.num_rows) << "Invalid end row " << end;
|
||||
|
||||
const IdType *coo_row_data = static_cast<IdType *>(coo.row->data);
|
||||
const IdType *coo_col_data = static_cast<IdType *>(coo.col->data);
|
||||
const IdType *coo_data =
|
||||
COOHasData(coo) ? static_cast<IdType *>(coo.data->data) : nullptr;
|
||||
|
||||
std::vector<IdType> ret_row, ret_col;
|
||||
std::vector<IdType> ret_data;
|
||||
|
||||
for (int64_t i = 0; i < coo.row->shape[0]; ++i) {
|
||||
const IdType row_id = coo_row_data[i];
|
||||
const IdType col_id = coo_col_data[i];
|
||||
if (row_id < end && row_id >= start) {
|
||||
ret_row.push_back(row_id - start);
|
||||
ret_col.push_back(col_id);
|
||||
ret_data.push_back(coo_data ? coo_data[i] : i);
|
||||
}
|
||||
}
|
||||
return COOMatrix(
|
||||
end - start, coo.num_cols, NDArray::FromVector(ret_row),
|
||||
NDArray::FromVector(ret_col), NDArray::FromVector(ret_data),
|
||||
coo.row_sorted, coo.col_sorted);
|
||||
}
|
||||
|
||||
template COOMatrix COOSliceRows<kDGLCPU, int32_t>(COOMatrix, int64_t, int64_t);
|
||||
template COOMatrix COOSliceRows<kDGLCPU, int64_t>(COOMatrix, int64_t, int64_t);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix COOSliceRows(COOMatrix coo, NDArray rows) {
|
||||
const IdType *coo_row_data = static_cast<IdType *>(coo.row->data);
|
||||
const IdType *coo_col_data = static_cast<IdType *>(coo.col->data);
|
||||
const IdType *coo_data =
|
||||
COOHasData(coo) ? static_cast<IdType *>(coo.data->data) : nullptr;
|
||||
|
||||
std::vector<IdType> ret_row, ret_col;
|
||||
std::vector<IdType> ret_data;
|
||||
|
||||
IdHashMap<IdType> hashmap(rows);
|
||||
|
||||
for (int64_t i = 0; i < coo.row->shape[0]; ++i) {
|
||||
const IdType row_id = coo_row_data[i];
|
||||
const IdType col_id = coo_col_data[i];
|
||||
const IdType mapped_row_id = hashmap.Map(row_id, -1);
|
||||
if (mapped_row_id != -1) {
|
||||
ret_row.push_back(mapped_row_id);
|
||||
ret_col.push_back(col_id);
|
||||
ret_data.push_back(coo_data ? coo_data[i] : i);
|
||||
}
|
||||
}
|
||||
|
||||
return COOMatrix{
|
||||
rows->shape[0],
|
||||
coo.num_cols,
|
||||
NDArray::FromVector(ret_row),
|
||||
NDArray::FromVector(ret_col),
|
||||
NDArray::FromVector(ret_data),
|
||||
coo.row_sorted,
|
||||
coo.col_sorted};
|
||||
}
|
||||
|
||||
template COOMatrix COOSliceRows<kDGLCPU, int32_t>(COOMatrix, NDArray);
|
||||
template COOMatrix COOSliceRows<kDGLCPU, int64_t>(COOMatrix, NDArray);
|
||||
|
||||
///////////////////////////// COOSliceMatrix /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix COOSliceMatrix(
|
||||
COOMatrix coo, runtime::NDArray rows, runtime::NDArray cols) {
|
||||
const IdType *coo_row_data = static_cast<IdType *>(coo.row->data);
|
||||
const IdType *coo_col_data = static_cast<IdType *>(coo.col->data);
|
||||
const IdType *coo_data =
|
||||
COOHasData(coo) ? static_cast<IdType *>(coo.data->data) : nullptr;
|
||||
|
||||
IdHashMap<IdType> row_map(rows), col_map(cols);
|
||||
|
||||
std::vector<IdType> ret_row, ret_col;
|
||||
std::vector<IdType> ret_data;
|
||||
|
||||
for (int64_t i = 0; i < coo.row->shape[0]; ++i) {
|
||||
const IdType row_id = coo_row_data[i];
|
||||
const IdType col_id = coo_col_data[i];
|
||||
const IdType mapped_row_id = row_map.Map(row_id, -1);
|
||||
if (mapped_row_id != -1) {
|
||||
const IdType mapped_col_id = col_map.Map(col_id, -1);
|
||||
if (mapped_col_id != -1) {
|
||||
ret_row.push_back(mapped_row_id);
|
||||
ret_col.push_back(mapped_col_id);
|
||||
ret_data.push_back(coo_data ? coo_data[i] : i);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return COOMatrix(
|
||||
rows->shape[0], cols->shape[0], NDArray::FromVector(ret_row),
|
||||
NDArray::FromVector(ret_col), NDArray::FromVector(ret_data),
|
||||
coo.row_sorted, coo.col_sorted);
|
||||
}
|
||||
|
||||
template COOMatrix COOSliceMatrix<kDGLCPU, int32_t>(
|
||||
COOMatrix coo, runtime::NDArray rows, runtime::NDArray cols);
|
||||
template COOMatrix COOSliceMatrix<kDGLCPU, int64_t>(
|
||||
COOMatrix coo, runtime::NDArray rows, runtime::NDArray cols);
|
||||
|
||||
///////////////////////////// COOReorder /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix COOReorder(
|
||||
COOMatrix coo, runtime::NDArray new_row_id_arr,
|
||||
runtime::NDArray new_col_id_arr) {
|
||||
CHECK_SAME_DTYPE(coo.row, new_row_id_arr);
|
||||
CHECK_SAME_DTYPE(coo.col, new_col_id_arr);
|
||||
|
||||
// Input COO
|
||||
const IdType *in_rows = static_cast<IdType *>(coo.row->data);
|
||||
const IdType *in_cols = static_cast<IdType *>(coo.col->data);
|
||||
int64_t num_rows = coo.num_rows;
|
||||
int64_t num_cols = coo.num_cols;
|
||||
int64_t nnz = coo.row->shape[0];
|
||||
CHECK_EQ(num_rows, new_row_id_arr->shape[0])
|
||||
<< "The new row Id array needs to be the same as the number of rows of "
|
||||
"COO";
|
||||
CHECK_EQ(num_cols, new_col_id_arr->shape[0])
|
||||
<< "The new col Id array needs to be the same as the number of cols of "
|
||||
"COO";
|
||||
|
||||
// New row/col Ids.
|
||||
const IdType *new_row_ids = static_cast<IdType *>(new_row_id_arr->data);
|
||||
const IdType *new_col_ids = static_cast<IdType *>(new_col_id_arr->data);
|
||||
|
||||
// Output COO
|
||||
NDArray out_row_arr = NDArray::Empty({nnz}, coo.row->dtype, coo.row->ctx);
|
||||
NDArray out_col_arr = NDArray::Empty({nnz}, coo.col->dtype, coo.col->ctx);
|
||||
NDArray out_data_arr = COOHasData(coo) ? coo.data : NullArray();
|
||||
IdType *out_row = static_cast<IdType *>(out_row_arr->data);
|
||||
IdType *out_col = static_cast<IdType *>(out_col_arr->data);
|
||||
|
||||
parallel_for(0, nnz, [=](size_t b, size_t e) {
|
||||
for (auto i = b; i < e; ++i) {
|
||||
out_row[i] = new_row_ids[in_rows[i]];
|
||||
out_col[i] = new_col_ids[in_cols[i]];
|
||||
}
|
||||
});
|
||||
return COOMatrix(num_rows, num_cols, out_row_arr, out_col_arr, out_data_arr);
|
||||
}
|
||||
|
||||
template COOMatrix COOReorder<kDGLCPU, int64_t>(
|
||||
COOMatrix csr, runtime::NDArray new_row_ids, runtime::NDArray new_col_ids);
|
||||
template COOMatrix COOReorder<kDGLCPU, int32_t>(
|
||||
COOMatrix csr, runtime::NDArray new_row_ids, runtime::NDArray new_col_ids);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,634 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file array/cpu/spmat_op_impl_csr.cc
|
||||
* @brief CSR matrix operator CPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
|
||||
#include <atomic>
|
||||
#include <numeric>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#include "array_utils.h"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using runtime::NDArray;
|
||||
using runtime::parallel_for;
|
||||
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
///////////////////////////// CSRIsNonZero /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
bool CSRIsNonZero(CSRMatrix csr, int64_t row, int64_t col) {
|
||||
const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
|
||||
const IdType* indices_data = static_cast<IdType*>(csr.indices->data);
|
||||
if (csr.sorted) {
|
||||
const IdType* start = indices_data + indptr_data[row];
|
||||
const IdType* end = indices_data + indptr_data[row + 1];
|
||||
return std::binary_search(start, end, col);
|
||||
} else {
|
||||
for (IdType i = indptr_data[row]; i < indptr_data[row + 1]; ++i) {
|
||||
if (indices_data[i] == col) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
template bool CSRIsNonZero<kDGLCPU, int32_t>(CSRMatrix, int64_t, int64_t);
|
||||
template bool CSRIsNonZero<kDGLCPU, int64_t>(CSRMatrix, int64_t, int64_t);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
NDArray CSRIsNonZero(CSRMatrix csr, NDArray row, NDArray col) {
|
||||
const auto rowlen = row->shape[0];
|
||||
const auto collen = col->shape[0];
|
||||
const auto rstlen = std::max(rowlen, collen);
|
||||
NDArray rst = NDArray::Empty({rstlen}, row->dtype, row->ctx);
|
||||
IdType* rst_data = static_cast<IdType*>(rst->data);
|
||||
const IdType* row_data = static_cast<IdType*>(row->data);
|
||||
const IdType* col_data = static_cast<IdType*>(col->data);
|
||||
const int64_t row_stride = (rowlen == 1 && collen != 1) ? 0 : 1;
|
||||
const int64_t col_stride = (collen == 1 && rowlen != 1) ? 0 : 1;
|
||||
runtime::parallel_for(
|
||||
0, std::max(rowlen, collen), 1, [=](int64_t b, int64_t e) {
|
||||
int64_t i = (row_stride == 0) ? 0 : b;
|
||||
int64_t j = (col_stride == 0) ? 0 : b;
|
||||
for (int64_t k = b; i < e && j < e;
|
||||
i += row_stride, j += col_stride, ++k)
|
||||
rst_data[k] =
|
||||
CSRIsNonZero<XPU, IdType>(csr, row_data[i], col_data[j]) ? 1 : 0;
|
||||
});
|
||||
return rst;
|
||||
}
|
||||
|
||||
template NDArray CSRIsNonZero<kDGLCPU, int32_t>(CSRMatrix, NDArray, NDArray);
|
||||
template NDArray CSRIsNonZero<kDGLCPU, int64_t>(CSRMatrix, NDArray, NDArray);
|
||||
|
||||
///////////////////////////// CSRHasDuplicate /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
bool CSRHasDuplicate(CSRMatrix csr) {
|
||||
const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
|
||||
const IdType* indices_data = static_cast<IdType*>(csr.indices->data);
|
||||
for (IdType src = 0; src < csr.num_rows; ++src) {
|
||||
std::unordered_set<IdType> hashmap;
|
||||
for (IdType eid = indptr_data[src]; eid < indptr_data[src + 1]; ++eid) {
|
||||
const IdType dst = indices_data[eid];
|
||||
if (hashmap.count(dst)) {
|
||||
return true;
|
||||
} else {
|
||||
hashmap.insert(dst);
|
||||
}
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
template bool CSRHasDuplicate<kDGLCPU, int32_t>(CSRMatrix csr);
|
||||
template bool CSRHasDuplicate<kDGLCPU, int64_t>(CSRMatrix csr);
|
||||
|
||||
///////////////////////////// CSRGetRowNNZ /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
int64_t CSRGetRowNNZ(CSRMatrix csr, int64_t row) {
|
||||
const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
|
||||
return indptr_data[row + 1] - indptr_data[row];
|
||||
}
|
||||
|
||||
template int64_t CSRGetRowNNZ<kDGLCPU, int32_t>(CSRMatrix, int64_t);
|
||||
template int64_t CSRGetRowNNZ<kDGLCPU, int64_t>(CSRMatrix, int64_t);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
NDArray CSRGetRowNNZ(CSRMatrix csr, NDArray rows) {
|
||||
CHECK_SAME_DTYPE(csr.indices, rows);
|
||||
const auto len = rows->shape[0];
|
||||
const IdType* vid_data = static_cast<IdType*>(rows->data);
|
||||
const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
|
||||
NDArray rst = NDArray::Empty({len}, rows->dtype, rows->ctx);
|
||||
IdType* rst_data = static_cast<IdType*>(rst->data);
|
||||
for (int64_t i = 0; i < len; ++i) {
|
||||
const auto vid = vid_data[i];
|
||||
rst_data[i] = indptr_data[vid + 1] - indptr_data[vid];
|
||||
}
|
||||
return rst;
|
||||
}
|
||||
|
||||
template NDArray CSRGetRowNNZ<kDGLCPU, int32_t>(CSRMatrix, NDArray);
|
||||
template NDArray CSRGetRowNNZ<kDGLCPU, int64_t>(CSRMatrix, NDArray);
|
||||
|
||||
/////////////////////////// CSRGetRowColumnIndices /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
NDArray CSRGetRowColumnIndices(CSRMatrix csr, int64_t row) {
|
||||
const int64_t len = impl::CSRGetRowNNZ<XPU, IdType>(csr, row);
|
||||
const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
|
||||
const int64_t offset = indptr_data[row] * sizeof(IdType);
|
||||
return csr.indices.CreateView({len}, csr.indices->dtype, offset);
|
||||
}
|
||||
|
||||
template NDArray CSRGetRowColumnIndices<kDGLCPU, int32_t>(CSRMatrix, int64_t);
|
||||
template NDArray CSRGetRowColumnIndices<kDGLCPU, int64_t>(CSRMatrix, int64_t);
|
||||
|
||||
///////////////////////////// CSRGetRowData /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
NDArray CSRGetRowData(CSRMatrix csr, int64_t row) {
|
||||
const int64_t len = impl::CSRGetRowNNZ<XPU, IdType>(csr, row);
|
||||
const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
|
||||
const int64_t offset = indptr_data[row] * sizeof(IdType);
|
||||
if (CSRHasData(csr))
|
||||
return csr.data.CreateView({len}, csr.data->dtype, offset);
|
||||
else
|
||||
return aten::Range(
|
||||
offset, offset + len, csr.indptr->dtype.bits, csr.indptr->ctx);
|
||||
}
|
||||
|
||||
template NDArray CSRGetRowData<kDGLCPU, int32_t>(CSRMatrix, int64_t);
|
||||
template NDArray CSRGetRowData<kDGLCPU, int64_t>(CSRMatrix, int64_t);
|
||||
|
||||
///////////////////////////// CSRGetData /////////////////////////////
|
||||
///////////////////////////// CSRGetDataAndIndices /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
void CollectDataIndicesFromSorted(
|
||||
const IdType* indices_data, const IdType* data, const IdType start,
|
||||
const IdType end, const IdType col, std::vector<IdType>* col_vec,
|
||||
std::vector<IdType>* ret_vec) {
|
||||
const IdType* start_ptr = indices_data + start;
|
||||
const IdType* end_ptr = indices_data + end;
|
||||
auto it = std::lower_bound(start_ptr, end_ptr, col);
|
||||
// This might be a multi-graph. We need to collect all of the matched
|
||||
// columns.
|
||||
for (; it != end_ptr; it++) {
|
||||
// If the col exist
|
||||
if (*it == col) {
|
||||
IdType idx = it - indices_data;
|
||||
col_vec->push_back(indices_data[idx]);
|
||||
ret_vec->push_back(data[idx]);
|
||||
} else {
|
||||
// If we find a column that is different, we can stop searching now.
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::vector<NDArray> CSRGetDataAndIndices(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols) {
|
||||
// TODO(minjie): more efficient implementation for matrix without duplicate
|
||||
// entries
|
||||
const int64_t rowlen = rows->shape[0];
|
||||
const int64_t collen = cols->shape[0];
|
||||
|
||||
CHECK((rowlen == collen) || (rowlen == 1) || (collen == 1))
|
||||
<< "Invalid row and col id array.";
|
||||
|
||||
const int64_t row_stride = (rowlen == 1 && collen != 1) ? 0 : 1;
|
||||
const int64_t col_stride = (collen == 1 && rowlen != 1) ? 0 : 1;
|
||||
const IdType* row_data = static_cast<IdType*>(rows->data);
|
||||
const IdType* col_data = static_cast<IdType*>(cols->data);
|
||||
|
||||
const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
|
||||
const IdType* indices_data = static_cast<IdType*>(csr.indices->data);
|
||||
const IdType* data =
|
||||
CSRHasData(csr) ? static_cast<IdType*>(csr.data->data) : nullptr;
|
||||
|
||||
std::vector<IdType> ret_rows, ret_cols;
|
||||
std::vector<IdType> ret_data;
|
||||
|
||||
for (int64_t i = 0, j = 0; i < rowlen && j < collen;
|
||||
i += row_stride, j += col_stride) {
|
||||
const IdType row_id = row_data[i], col_id = col_data[j];
|
||||
CHECK(row_id >= 0 && row_id < csr.num_rows)
|
||||
<< "Invalid row index: " << row_id;
|
||||
CHECK(col_id >= 0 && col_id < csr.num_cols)
|
||||
<< "Invalid col index: " << col_id;
|
||||
if (csr.sorted) {
|
||||
// Here we collect col indices and data.
|
||||
CollectDataIndicesFromSorted<XPU, IdType>(
|
||||
indices_data, data, indptr_data[row_id], indptr_data[row_id + 1],
|
||||
col_id, &ret_cols, &ret_data);
|
||||
// We need to add row Ids.
|
||||
while (ret_rows.size() < ret_data.size()) {
|
||||
ret_rows.push_back(row_id);
|
||||
}
|
||||
} else {
|
||||
for (IdType i = indptr_data[row_id]; i < indptr_data[row_id + 1]; ++i) {
|
||||
if (indices_data[i] == col_id) {
|
||||
ret_rows.push_back(row_id);
|
||||
ret_cols.push_back(col_id);
|
||||
ret_data.push_back(data ? data[i] : i);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
NDArray::FromVector(ret_rows, csr.indptr->ctx),
|
||||
NDArray::FromVector(ret_cols, csr.indptr->ctx),
|
||||
NDArray::FromVector(ret_data, csr.data->ctx)};
|
||||
}
|
||||
|
||||
template std::vector<NDArray> CSRGetDataAndIndices<kDGLCPU, int32_t>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols);
|
||||
template std::vector<NDArray> CSRGetDataAndIndices<kDGLCPU, int64_t>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols);
|
||||
|
||||
///////////////////////////// CSRTranspose /////////////////////////////
|
||||
|
||||
// for a matrix of shape (N, M) and NNZ
|
||||
// complexity: time O(NNZ + max(N, M)), space O(1)
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix CSRTranspose(CSRMatrix csr) {
|
||||
const int64_t N = csr.num_rows;
|
||||
const int64_t M = csr.num_cols;
|
||||
const int64_t nnz = csr.indices->shape[0];
|
||||
const IdType* Ap = static_cast<IdType*>(csr.indptr->data);
|
||||
const IdType* Aj = static_cast<IdType*>(csr.indices->data);
|
||||
const IdType* Ax =
|
||||
CSRHasData(csr) ? static_cast<IdType*>(csr.data->data) : nullptr;
|
||||
NDArray ret_indptr =
|
||||
NDArray::Empty({M + 1}, csr.indptr->dtype, csr.indptr->ctx);
|
||||
NDArray ret_indices =
|
||||
NDArray::Empty({nnz}, csr.indices->dtype, csr.indices->ctx);
|
||||
NDArray ret_data = NDArray::Empty({nnz}, csr.indptr->dtype, csr.indptr->ctx);
|
||||
IdType* Bp = static_cast<IdType*>(ret_indptr->data);
|
||||
IdType* Bi = static_cast<IdType*>(ret_indices->data);
|
||||
IdType* Bx = static_cast<IdType*>(ret_data->data);
|
||||
|
||||
std::fill(Bp, Bp + M, 0);
|
||||
|
||||
for (int64_t j = 0; j < nnz; ++j) {
|
||||
Bp[Aj[j]]++;
|
||||
}
|
||||
|
||||
// cumsum
|
||||
for (int64_t i = 0, cumsum = 0; i < M; ++i) {
|
||||
const IdType temp = Bp[i];
|
||||
Bp[i] = cumsum;
|
||||
cumsum += temp;
|
||||
}
|
||||
Bp[M] = nnz;
|
||||
|
||||
for (int64_t i = 0; i < N; ++i) {
|
||||
for (IdType j = Ap[i]; j < Ap[i + 1]; ++j) {
|
||||
const IdType dst = Aj[j];
|
||||
Bi[Bp[dst]] = i;
|
||||
Bx[Bp[dst]] = Ax ? Ax[j] : j;
|
||||
Bp[dst]++;
|
||||
}
|
||||
}
|
||||
|
||||
// correct the indptr
|
||||
for (int64_t i = 0, last = 0; i <= M; ++i) {
|
||||
IdType temp = Bp[i];
|
||||
Bp[i] = last;
|
||||
last = temp;
|
||||
}
|
||||
|
||||
return CSRMatrix{
|
||||
csr.num_cols, csr.num_rows, ret_indptr, ret_indices, ret_data};
|
||||
}
|
||||
|
||||
template CSRMatrix CSRTranspose<kDGLCPU, int32_t>(CSRMatrix csr);
|
||||
template CSRMatrix CSRTranspose<kDGLCPU, int64_t>(CSRMatrix csr);
|
||||
|
||||
///////////////////////////// CSRToCOO /////////////////////////////
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix CSRToCOO(CSRMatrix csr) {
|
||||
const int64_t nnz = csr.indices->shape[0];
|
||||
const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
|
||||
NDArray ret_row = NDArray::Empty({nnz}, csr.indices->dtype, csr.indices->ctx);
|
||||
IdType* ret_row_data = static_cast<IdType*>(ret_row->data);
|
||||
parallel_for(0, csr.indptr->shape[0] - 1, 10000, [=](int64_t b, int64_t e) {
|
||||
for (auto i = b; i < e; ++i) {
|
||||
std::fill(
|
||||
ret_row_data + indptr_data[i], ret_row_data + indptr_data[i + 1], i);
|
||||
}
|
||||
});
|
||||
return COOMatrix(
|
||||
csr.num_rows, csr.num_cols, ret_row, csr.indices, csr.data, true,
|
||||
csr.sorted);
|
||||
}
|
||||
|
||||
template COOMatrix CSRToCOO<kDGLCPU, int32_t>(CSRMatrix csr);
|
||||
template COOMatrix CSRToCOO<kDGLCPU, int64_t>(CSRMatrix csr);
|
||||
|
||||
// complexity: time O(NNZ), space O(1)
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix CSRToCOODataAsOrder(CSRMatrix csr) {
|
||||
const int64_t N = csr.num_rows;
|
||||
const int64_t M = csr.num_cols;
|
||||
const int64_t nnz = csr.indices->shape[0];
|
||||
const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
|
||||
const IdType* indices_data = static_cast<IdType*>(csr.indices->data);
|
||||
// data array should have the same type as the indices arrays
|
||||
const IdType* data =
|
||||
CSRHasData(csr) ? static_cast<IdType*>(csr.data->data) : nullptr;
|
||||
NDArray ret_row = NDArray::Empty({nnz}, csr.indices->dtype, csr.indices->ctx);
|
||||
NDArray ret_col = NDArray::Empty({nnz}, csr.indices->dtype, csr.indices->ctx);
|
||||
IdType* ret_row_data = static_cast<IdType*>(ret_row->data);
|
||||
IdType* ret_col_data = static_cast<IdType*>(ret_col->data);
|
||||
// scatter using the indices in the data array
|
||||
parallel_for(0, N, 10000, [=](int64_t b, int64_t e) {
|
||||
for (auto row = b; row < e; ++row) {
|
||||
for (IdType j = indptr_data[row]; j < indptr_data[row + 1]; ++j) {
|
||||
const IdType col = indices_data[j];
|
||||
ret_row_data[data ? data[j] : j] = row;
|
||||
ret_col_data[data ? data[j] : j] = col;
|
||||
}
|
||||
}
|
||||
});
|
||||
return COOMatrix(N, M, ret_row, ret_col);
|
||||
}
|
||||
|
||||
template COOMatrix CSRToCOODataAsOrder<kDGLCPU, int32_t>(CSRMatrix csr);
|
||||
template COOMatrix CSRToCOODataAsOrder<kDGLCPU, int64_t>(CSRMatrix csr);
|
||||
|
||||
///////////////////////////// CSRSliceRows /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix CSRSliceRows(CSRMatrix csr, int64_t start, int64_t end) {
|
||||
const IdType* indptr = static_cast<IdType*>(csr.indptr->data);
|
||||
const int64_t num_rows = end - start;
|
||||
const int64_t nnz = indptr[end] - indptr[start];
|
||||
IdArray ret_indptr =
|
||||
IdArray::Empty({num_rows + 1}, csr.indptr->dtype, csr.indices->ctx);
|
||||
IdType* r_indptr = static_cast<IdType*>(ret_indptr->data);
|
||||
for (int64_t i = start; i < end + 1; ++i) {
|
||||
r_indptr[i - start] = indptr[i] - indptr[start];
|
||||
}
|
||||
// indices and data can be view arrays
|
||||
IdArray ret_indices = csr.indices.CreateView(
|
||||
{nnz}, csr.indices->dtype, indptr[start] * sizeof(IdType));
|
||||
IdArray ret_data;
|
||||
if (CSRHasData(csr))
|
||||
ret_data = csr.data.CreateView(
|
||||
{nnz}, csr.data->dtype, indptr[start] * sizeof(IdType));
|
||||
else
|
||||
ret_data = aten::Range(
|
||||
indptr[start], indptr[end], csr.indptr->dtype.bits, csr.indptr->ctx);
|
||||
return CSRMatrix(
|
||||
num_rows, csr.num_cols, ret_indptr, ret_indices, ret_data, csr.sorted);
|
||||
}
|
||||
|
||||
template CSRMatrix CSRSliceRows<kDGLCPU, int32_t>(CSRMatrix, int64_t, int64_t);
|
||||
template CSRMatrix CSRSliceRows<kDGLCPU, int64_t>(CSRMatrix, int64_t, int64_t);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix CSRSliceRows(CSRMatrix csr, NDArray rows) {
|
||||
CHECK_SAME_DTYPE(csr.indices, rows);
|
||||
const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
|
||||
const IdType* indices_data = static_cast<IdType*>(csr.indices->data);
|
||||
const IdType* data =
|
||||
CSRHasData(csr) ? static_cast<IdType*>(csr.data->data) : nullptr;
|
||||
const auto len = rows->shape[0];
|
||||
const IdType* rows_data = static_cast<IdType*>(rows->data);
|
||||
int64_t nnz = 0;
|
||||
|
||||
CSRMatrix ret;
|
||||
ret.num_rows = len;
|
||||
ret.num_cols = csr.num_cols;
|
||||
ret.indptr = NDArray::Empty({len + 1}, csr.indptr->dtype, csr.indices->ctx);
|
||||
|
||||
IdType* ret_indptr_data = static_cast<IdType*>(ret.indptr->data);
|
||||
ret_indptr_data[0] = 0;
|
||||
|
||||
std::vector<IdType> sums;
|
||||
|
||||
std::atomic_flag err_flag = ATOMIC_FLAG_INIT;
|
||||
bool err = false;
|
||||
std::stringstream err_msg_stream;
|
||||
|
||||
// Perform two-round parallel prefix sum using OpenMP
|
||||
#pragma omp parallel
|
||||
{
|
||||
int64_t tid = omp_get_thread_num();
|
||||
int64_t num_threads = omp_get_num_threads();
|
||||
|
||||
#pragma omp single
|
||||
{
|
||||
sums.resize(num_threads + 1);
|
||||
sums[0] = 0;
|
||||
}
|
||||
|
||||
int64_t sum = 0;
|
||||
|
||||
// First round of parallel prefix sum. All threads perform local prefix sums.
|
||||
#pragma omp for schedule(static) nowait
|
||||
for (int64_t i = 0; i < len; ++i) {
|
||||
int64_t rid = rows_data[i];
|
||||
if (rid >= csr.num_rows) {
|
||||
if (!err_flag.test_and_set()) {
|
||||
err_msg_stream << "expect row ID " << rid
|
||||
<< " to be less than number of rows " << csr.num_rows;
|
||||
err = true;
|
||||
}
|
||||
} else {
|
||||
sum += indptr_data[rid + 1] - indptr_data[rid];
|
||||
ret_indptr_data[i + 1] = sum;
|
||||
}
|
||||
}
|
||||
sums[tid + 1] = sum;
|
||||
#pragma omp barrier
|
||||
|
||||
#pragma omp single
|
||||
{
|
||||
for (int64_t i = 1; i < num_threads; ++i) sums[i] += sums[i - 1];
|
||||
}
|
||||
|
||||
int64_t offset = sums[tid];
|
||||
|
||||
// Second round of parallel prefix sum. Update the local prefix sums.
|
||||
#pragma omp for schedule(static)
|
||||
for (int64_t i = 0; i < len; ++i) ret_indptr_data[i + 1] += offset;
|
||||
}
|
||||
if (err) {
|
||||
LOG(FATAL) << err_msg_stream.str();
|
||||
return ret;
|
||||
}
|
||||
|
||||
// After the prefix sum, the last element of ret_indptr_data holds the
|
||||
// sum of all elements
|
||||
nnz = ret_indptr_data[len];
|
||||
|
||||
ret.indices = NDArray::Empty({nnz}, csr.indices->dtype, csr.indices->ctx);
|
||||
ret.data = NDArray::Empty({nnz}, csr.indptr->dtype, csr.indptr->ctx);
|
||||
ret.sorted = csr.sorted;
|
||||
|
||||
IdType* ret_indices_data = static_cast<IdType*>(ret.indices->data);
|
||||
IdType* ret_data = static_cast<IdType*>(ret.data->data);
|
||||
|
||||
parallel_for(0, len, [=](int64_t b, int64_t e) {
|
||||
for (auto i = b; i < e; ++i) {
|
||||
const IdType rid = rows_data[i];
|
||||
// note: zero is allowed
|
||||
std::copy(
|
||||
indices_data + indptr_data[rid], indices_data + indptr_data[rid + 1],
|
||||
ret_indices_data + ret_indptr_data[i]);
|
||||
if (data)
|
||||
std::copy(
|
||||
data + indptr_data[rid], data + indptr_data[rid + 1],
|
||||
ret_data + ret_indptr_data[i]);
|
||||
else
|
||||
std::iota(
|
||||
ret_data + ret_indptr_data[i], ret_data + ret_indptr_data[i + 1],
|
||||
indptr_data[rid]);
|
||||
}
|
||||
});
|
||||
return ret;
|
||||
}
|
||||
|
||||
template CSRMatrix CSRSliceRows<kDGLCPU, int32_t>(CSRMatrix, NDArray);
|
||||
template CSRMatrix CSRSliceRows<kDGLCPU, int64_t>(CSRMatrix, NDArray);
|
||||
|
||||
///////////////////////////// CSRSliceMatrix /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix CSRSliceMatrix(
|
||||
CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols) {
|
||||
IdHashMap<IdType> hashmap(cols);
|
||||
const int64_t new_nrows = rows->shape[0];
|
||||
const int64_t new_ncols = cols->shape[0];
|
||||
const IdType* rows_data = static_cast<IdType*>(rows->data);
|
||||
const bool has_data = CSRHasData(csr);
|
||||
|
||||
const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
|
||||
const IdType* indices_data = static_cast<IdType*>(csr.indices->data);
|
||||
const IdType* data =
|
||||
has_data ? static_cast<IdType*>(csr.data->data) : nullptr;
|
||||
|
||||
std::vector<IdType> sub_indptr, sub_indices;
|
||||
std::vector<IdType> sub_data;
|
||||
sub_indptr.resize(new_nrows + 1, 0);
|
||||
const IdType kInvalidId = new_ncols + 1;
|
||||
for (int64_t i = 0; i < new_nrows; ++i) {
|
||||
// NOTE: newi == i
|
||||
const IdType oldi = rows_data[i];
|
||||
CHECK(oldi >= 0 && oldi < csr.num_rows) << "Invalid row index: " << oldi;
|
||||
for (IdType p = indptr_data[oldi]; p < indptr_data[oldi + 1]; ++p) {
|
||||
const IdType oldj = indices_data[p];
|
||||
const IdType newj = hashmap.Map(oldj, kInvalidId);
|
||||
if (newj != kInvalidId) {
|
||||
++sub_indptr[i];
|
||||
sub_indices.push_back(newj);
|
||||
sub_data.push_back(has_data ? data[p] : p);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// cumsum sub_indptr
|
||||
for (int64_t i = 0, cumsum = 0; i < new_nrows; ++i) {
|
||||
const IdType temp = sub_indptr[i];
|
||||
sub_indptr[i] = cumsum;
|
||||
cumsum += temp;
|
||||
}
|
||||
sub_indptr[new_nrows] = sub_indices.size();
|
||||
|
||||
const int64_t nnz = sub_data.size();
|
||||
NDArray sub_data_arr =
|
||||
NDArray::Empty({nnz}, csr.indptr->dtype, csr.indptr->ctx);
|
||||
IdType* ptr = static_cast<IdType*>(sub_data_arr->data);
|
||||
std::copy(sub_data.begin(), sub_data.end(), ptr);
|
||||
return CSRMatrix{
|
||||
new_nrows, new_ncols, NDArray::FromVector(sub_indptr, csr.indptr->ctx),
|
||||
NDArray::FromVector(sub_indices, csr.indptr->ctx), sub_data_arr};
|
||||
}
|
||||
|
||||
template CSRMatrix CSRSliceMatrix<kDGLCPU, int32_t>(
|
||||
CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols);
|
||||
template CSRMatrix CSRSliceMatrix<kDGLCPU, int64_t>(
|
||||
CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols);
|
||||
|
||||
///////////////////////////// CSRReorder /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix CSRReorder(
|
||||
CSRMatrix csr, runtime::NDArray new_row_id_arr,
|
||||
runtime::NDArray new_col_id_arr) {
|
||||
CHECK_SAME_DTYPE(csr.indices, new_row_id_arr);
|
||||
CHECK_SAME_DTYPE(csr.indices, new_col_id_arr);
|
||||
|
||||
// Input CSR
|
||||
const IdType* in_indptr = static_cast<IdType*>(csr.indptr->data);
|
||||
const IdType* in_indices = static_cast<IdType*>(csr.indices->data);
|
||||
const IdType* in_data = static_cast<IdType*>(csr.data->data);
|
||||
int64_t num_rows = csr.num_rows;
|
||||
int64_t num_cols = csr.num_cols;
|
||||
int64_t nnz = csr.indices->shape[0];
|
||||
CHECK_EQ(nnz, in_indptr[num_rows]);
|
||||
CHECK_EQ(num_rows, new_row_id_arr->shape[0])
|
||||
<< "The new row Id array needs to be the same as the number of rows of "
|
||||
"CSR";
|
||||
CHECK_EQ(num_cols, new_col_id_arr->shape[0])
|
||||
<< "The new col Id array needs to be the same as the number of cols of "
|
||||
"CSR";
|
||||
|
||||
// New row/col Ids.
|
||||
const IdType* new_row_ids = static_cast<IdType*>(new_row_id_arr->data);
|
||||
const IdType* new_col_ids = static_cast<IdType*>(new_col_id_arr->data);
|
||||
|
||||
// Output CSR
|
||||
NDArray out_indptr_arr =
|
||||
NDArray::Empty({num_rows + 1}, csr.indptr->dtype, csr.indptr->ctx);
|
||||
NDArray out_indices_arr =
|
||||
NDArray::Empty({nnz}, csr.indices->dtype, csr.indices->ctx);
|
||||
NDArray out_data_arr = NDArray::Empty({nnz}, csr.data->dtype, csr.data->ctx);
|
||||
IdType* out_indptr = static_cast<IdType*>(out_indptr_arr->data);
|
||||
IdType* out_indices = static_cast<IdType*>(out_indices_arr->data);
|
||||
IdType* out_data = static_cast<IdType*>(out_data_arr->data);
|
||||
|
||||
// Compute the length of rows for the new matrix.
|
||||
std::vector<IdType> new_row_lens(num_rows, -1);
|
||||
parallel_for(0, num_rows, [=, &new_row_lens](size_t b, size_t e) {
|
||||
for (auto i = b; i < e; ++i) {
|
||||
int64_t new_row_id = new_row_ids[i];
|
||||
new_row_lens[new_row_id] = in_indptr[i + 1] - in_indptr[i];
|
||||
}
|
||||
});
|
||||
// Compute the starting location of each row in the new matrix.
|
||||
out_indptr[0] = 0;
|
||||
// This is sequential. It should be pretty fast.
|
||||
for (int64_t i = 0; i < num_rows; i++) {
|
||||
CHECK_GE(new_row_lens[i], 0);
|
||||
out_indptr[i + 1] = out_indptr[i] + new_row_lens[i];
|
||||
}
|
||||
CHECK_EQ(out_indptr[num_rows], nnz);
|
||||
// Copy indieces and data with the new order.
|
||||
// Here I iterate rows in the order of the old matrix.
|
||||
parallel_for(0, num_rows, [=](size_t b, size_t e) {
|
||||
for (auto i = b; i < e; ++i) {
|
||||
const IdType* in_row = in_indices + in_indptr[i];
|
||||
const IdType* in_row_data = in_data + in_indptr[i];
|
||||
|
||||
int64_t new_row_id = new_row_ids[i];
|
||||
IdType* out_row = out_indices + out_indptr[new_row_id];
|
||||
IdType* out_row_data = out_data + out_indptr[new_row_id];
|
||||
|
||||
int64_t row_len = new_row_lens[new_row_id];
|
||||
// Here I iterate col indices in a row in the order of the old matrix.
|
||||
for (int64_t j = 0; j < row_len; j++) {
|
||||
out_row[j] = new_col_ids[in_row[j]];
|
||||
out_row_data[j] = in_row_data[j];
|
||||
}
|
||||
// TODO(zhengda) maybe we should sort the column indices.
|
||||
}
|
||||
});
|
||||
return CSRMatrix(
|
||||
num_rows, num_cols, out_indptr_arr, out_indices_arr, out_data_arr);
|
||||
}
|
||||
|
||||
template CSRMatrix CSRReorder<kDGLCPU, int64_t>(
|
||||
CSRMatrix csr, runtime::NDArray new_row_ids, runtime::NDArray new_col_ids);
|
||||
template CSRMatrix CSRReorder<kDGLCPU, int32_t>(
|
||||
CSRMatrix csr, runtime::NDArray new_row_ids, runtime::NDArray new_col_ids);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,304 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file kernel/cpu/spmm.cc
|
||||
* @brief SPMM C APIs and definitions.
|
||||
*/
|
||||
#include "./spmm.h"
|
||||
|
||||
#include <dgl/array.h>
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
|
||||
/** @brief Generalized SpMM on Csr format. */
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SpMMCsr(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux) {
|
||||
const int64_t dim = bcast.out_len;
|
||||
if (reduce == "sum") {
|
||||
SWITCH_OP(op, Op, {
|
||||
cpu::SpMMSumCsr<IdType, DType, Op>(bcast, csr, ufeat, efeat, out);
|
||||
});
|
||||
} else if (reduce == "max" || reduce == "min") {
|
||||
SWITCH_OP(op, Op, {
|
||||
DType* out_off = out.Ptr<DType>();
|
||||
if (reduce == "max") {
|
||||
std::fill(
|
||||
out_off, out_off + csr.num_rows * dim, cpu::op::Max<DType>::zero);
|
||||
cpu::SpMMCmpCsr<IdType, DType, Op, cpu::op::Max<DType>>(
|
||||
bcast, csr, ufeat, efeat, out, out_aux[0], out_aux[1]);
|
||||
} else {
|
||||
std::fill(
|
||||
out_off, out_off + csr.num_rows * dim, cpu::op::Min<DType>::zero);
|
||||
cpu::SpMMCmpCsr<IdType, DType, Op, cpu::op::Min<DType>>(
|
||||
bcast, csr, ufeat, efeat, out, out_aux[0], out_aux[1]);
|
||||
}
|
||||
});
|
||||
} else {
|
||||
LOG(FATAL) << "Unsupported SpMM reducer: " << reduce;
|
||||
}
|
||||
}
|
||||
|
||||
/** @brief Generalized SpMM on Csr format. */
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SpMMCsrHetero(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr,
|
||||
const std::vector<NDArray>& vec_ufeat,
|
||||
const std::vector<NDArray>& vec_efeat, std::vector<NDArray>* vec_out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_node_tids,
|
||||
const std::vector<dgl_type_t>& out_node_tids) {
|
||||
const int64_t dim = bcast.out_len;
|
||||
if (reduce == "sum") {
|
||||
SWITCH_OP(op, Op, {
|
||||
/* Call SpMM for each relation type */
|
||||
for (dgl_type_t etype = 0; etype < ufeat_node_tids.size(); ++etype) {
|
||||
const dgl_type_t src_id = ufeat_node_tids[etype];
|
||||
const dgl_type_t dst_id = out_node_tids[etype];
|
||||
CSRMatrix csr = vec_csr[etype];
|
||||
NDArray ufeat =
|
||||
(vec_ufeat.size() == 0) ? NullArray() : vec_ufeat[src_id];
|
||||
NDArray efeat =
|
||||
(vec_efeat.size() == 0) ? NullArray() : vec_efeat[etype];
|
||||
NDArray out = (*vec_out)[dst_id];
|
||||
cpu::SpMMSumCsr<IdType, DType, Op>(bcast, csr, ufeat, efeat, out);
|
||||
}
|
||||
});
|
||||
} else if (reduce == "max" || reduce == "min") {
|
||||
SWITCH_OP(op, Op, {
|
||||
std::vector<bool> updated((*vec_out).size(), false);
|
||||
// TODO(Israt): use vector updated to fill(out...) too
|
||||
for (dgl_type_t etype = 0; etype < ufeat_node_tids.size(); ++etype) {
|
||||
DType* out_off = (*vec_out)[out_node_tids[etype]].Ptr<DType>();
|
||||
if (reduce == "max")
|
||||
std::fill(
|
||||
out_off, out_off + vec_csr[etype].num_rows * dim,
|
||||
cpu::op::Max<DType>::zero);
|
||||
else
|
||||
std::fill(
|
||||
out_off, out_off + vec_csr[etype].num_rows * dim,
|
||||
cpu::op::Min<DType>::zero);
|
||||
const dgl_type_t dst_id = out_node_tids[etype];
|
||||
if (!updated[dst_id]) {
|
||||
updated[dst_id] = true;
|
||||
if (Op::use_lhs) {
|
||||
IdType* argu_ntype = (*out_aux)[2][dst_id].Ptr<IdType>();
|
||||
std::fill(
|
||||
argu_ntype, argu_ntype + vec_csr[etype].num_rows * dim, -1);
|
||||
}
|
||||
if (Op::use_rhs) {
|
||||
IdType* arge_etype = (*out_aux)[3][dst_id].Ptr<IdType>();
|
||||
std::fill(
|
||||
arge_etype, arge_etype + vec_csr[etype].num_rows * dim, -1);
|
||||
}
|
||||
}
|
||||
}
|
||||
/* Call SpMM for each relation type */
|
||||
for (dgl_type_t etype = 0; etype < ufeat_node_tids.size(); ++etype) {
|
||||
const dgl_type_t src_id = ufeat_node_tids[etype];
|
||||
const dgl_type_t dst_id = out_node_tids[etype];
|
||||
CSRMatrix csr = vec_csr[etype];
|
||||
NDArray ufeat =
|
||||
(vec_ufeat.size() == 0) ? NullArray() : vec_ufeat[src_id];
|
||||
NDArray efeat =
|
||||
(vec_efeat.size() == 0) ? NullArray() : vec_efeat[etype];
|
||||
NDArray out = (*vec_out)[dst_id];
|
||||
if (reduce == "max") {
|
||||
cpu::SpMMCmpCsrHetero<IdType, DType, Op, cpu::op::Max<DType>>(
|
||||
bcast, csr, ufeat, efeat, out, (*out_aux)[0][dst_id],
|
||||
(*out_aux)[1][dst_id], (*out_aux)[2][dst_id],
|
||||
(*out_aux)[3][dst_id], src_id, etype);
|
||||
} else {
|
||||
cpu::SpMMCmpCsrHetero<IdType, DType, Op, cpu::op::Min<DType>>(
|
||||
bcast, csr, ufeat, efeat, out, (*out_aux)[0][dst_id],
|
||||
(*out_aux)[1][dst_id], (*out_aux)[2][dst_id],
|
||||
(*out_aux)[3][dst_id], src_id, etype);
|
||||
}
|
||||
}
|
||||
});
|
||||
} else {
|
||||
LOG(FATAL) << "Unsupported SpMM reducer: " << reduce;
|
||||
}
|
||||
}
|
||||
|
||||
template void SpMMCsr<kDGLCPU, int32_t, BFloat16>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCsr<kDGLCPU, int64_t, BFloat16>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCsr<kDGLCPU, int32_t, float>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCsr<kDGLCPU, int64_t, float>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCsr<kDGLCPU, int32_t, double>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCsr<kDGLCPU, int64_t, double>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
|
||||
template void SpMMCsrHetero<kDGLCPU, int32_t, BFloat16>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
|
||||
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_node_tids,
|
||||
const std::vector<dgl_type_t>& out_node_tids);
|
||||
template void SpMMCsrHetero<kDGLCPU, int64_t, BFloat16>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
|
||||
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_node_tids,
|
||||
const std::vector<dgl_type_t>& out_node_tids);
|
||||
template void SpMMCsrHetero<kDGLCPU, int32_t, float>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
|
||||
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_node_tids,
|
||||
const std::vector<dgl_type_t>& out_node_tids);
|
||||
template void SpMMCsrHetero<kDGLCPU, int64_t, float>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
|
||||
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_node_tids,
|
||||
const std::vector<dgl_type_t>& out_node_tids);
|
||||
template void SpMMCsrHetero<kDGLCPU, int32_t, double>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
|
||||
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_node_tids,
|
||||
const std::vector<dgl_type_t>& out_node_tids);
|
||||
template void SpMMCsrHetero<kDGLCPU, int64_t, double>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
|
||||
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_node_tids,
|
||||
const std::vector<dgl_type_t>& out_node_tids);
|
||||
|
||||
/** @brief Edge_softmax_csr forward op on Csr format. */
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void Edge_softmax_csr_forward(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray ufeat, NDArray efeat, NDArray out) {
|
||||
SWITCH_OP(op, Op, {
|
||||
cpu::Edge_softmax_csr_forward<IdType, DType, Op>(
|
||||
bcast, csr, ufeat, efeat, out);
|
||||
});
|
||||
}
|
||||
|
||||
/** @brief Edge_softmax_csr backward op on Csr format. */
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void Edge_softmax_csr_backward(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray out, NDArray sds, NDArray back_out) {
|
||||
SWITCH_OP(op, Op, {
|
||||
cpu::Edge_softmax_csr_backward<IdType, DType, Op>(
|
||||
bcast, csr, out, sds, back_out);
|
||||
});
|
||||
}
|
||||
template void Edge_softmax_csr_forward<kDGLCPU, int32_t, BFloat16>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray ufeat, NDArray efeat, NDArray out);
|
||||
template void Edge_softmax_csr_forward<kDGLCPU, int64_t, BFloat16>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray ufeat, NDArray efeat, NDArray out);
|
||||
template void Edge_softmax_csr_forward<kDGLCPU, int32_t, float>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray ufeat, NDArray efeat, NDArray out);
|
||||
template void Edge_softmax_csr_forward<kDGLCPU, int64_t, float>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray ufeat, NDArray efeat, NDArray out);
|
||||
template void Edge_softmax_csr_forward<kDGLCPU, int32_t, double>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray ufeat, NDArray efeat, NDArray out);
|
||||
template void Edge_softmax_csr_forward<kDGLCPU, int64_t, double>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray ufeat, NDArray efeat, NDArray out);
|
||||
|
||||
template void Edge_softmax_csr_backward<kDGLCPU, int32_t, BFloat16>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray ufeat, NDArray efeat, NDArray out);
|
||||
template void Edge_softmax_csr_backward<kDGLCPU, int64_t, BFloat16>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray ufeat, NDArray efeat, NDArray out);
|
||||
template void Edge_softmax_csr_backward<kDGLCPU, int32_t, float>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray ufeat, NDArray efeat, NDArray out);
|
||||
template void Edge_softmax_csr_backward<kDGLCPU, int64_t, float>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray ufeat, NDArray efeat, NDArray out);
|
||||
template void Edge_softmax_csr_backward<kDGLCPU, int32_t, double>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray ufeat, NDArray efeat, NDArray out);
|
||||
template void Edge_softmax_csr_backward<kDGLCPU, int64_t, double>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray ufeat, NDArray efeat, NDArray out);
|
||||
|
||||
/** @brief Generalized SpMM on Coo format. */
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SpMMCoo(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux) {
|
||||
if (reduce == "sum") {
|
||||
SWITCH_OP(op, Op, {
|
||||
cpu::SpMMSumCoo<IdType, DType, Op>(bcast, coo, ufeat, efeat, out);
|
||||
});
|
||||
} else if (reduce == "max" || reduce == "min") {
|
||||
SWITCH_OP(op, Op, {
|
||||
if (reduce == "max")
|
||||
cpu::SpMMCmpCoo<IdType, DType, Op, cpu::op::Max<DType>>(
|
||||
bcast, coo, ufeat, efeat, out, out_aux[0], out_aux[1]);
|
||||
else
|
||||
cpu::SpMMCmpCoo<IdType, DType, Op, cpu::op::Min<DType>>(
|
||||
bcast, coo, ufeat, efeat, out, out_aux[0], out_aux[1]);
|
||||
});
|
||||
} else {
|
||||
LOG(FATAL) << "Unsupported SpMM reducer: " << reduce;
|
||||
}
|
||||
}
|
||||
|
||||
template void SpMMCoo<kDGLCPU, int32_t, BFloat16>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCoo<kDGLCPU, int64_t, BFloat16>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCoo<kDGLCPU, int32_t, float>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCoo<kDGLCPU, int64_t, float>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCoo<kDGLCPU, int32_t, double>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCoo<kDGLCPU, int64_t, double>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,576 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/spmm.h
|
||||
* @brief SPMM CPU kernel function header.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CPU_SPMM_H_
|
||||
#define DGL_ARRAY_CPU_SPMM_H_
|
||||
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/bcast.h>
|
||||
#include <dgl/runtime/config.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
#include <math.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <limits>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "spmm_binary_ops.h"
|
||||
#if !defined(_WIN32)
|
||||
#ifdef USE_LIBXSMM
|
||||
#include "spmm_blocking_libxsmm.h"
|
||||
#endif // USE_LIBXSMM
|
||||
#endif // _WIN32
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace cpu {
|
||||
|
||||
template <typename DType>
|
||||
using AccType = typename std::conditional<
|
||||
std::is_same<DType, BFloat16>::value, float, DType>::type;
|
||||
|
||||
/**
|
||||
* @brief Naive CPU kernel of SpMM on Csr format.
|
||||
* @param cpu_spec JIT'ed kernel
|
||||
* @param bcast Broadcast information.
|
||||
* @param csr The Csr matrix.
|
||||
* @param X The feature on source nodes.
|
||||
* @param W The feature on edges.
|
||||
* @param O The result feature on destination nodes.
|
||||
* @note it uses node parallel strategy, different threads are responsible
|
||||
* for the computation of different nodes.
|
||||
*/
|
||||
template <typename IdType, typename DType, typename Op>
|
||||
typename std::enable_if<!std::is_same<DType, BFloat16>::value, void>::type
|
||||
SpMMSumCsrNaive(
|
||||
const BcastOff& bcast, const CSRMatrix& csr, const DType* X, const DType* W,
|
||||
DType* O) {
|
||||
const bool has_idx = !IsNullArray(csr.data);
|
||||
const IdType* indptr = csr.indptr.Ptr<IdType>();
|
||||
const IdType* indices = csr.indices.Ptr<IdType>();
|
||||
const IdType* edges = csr.data.Ptr<IdType>();
|
||||
int64_t dim = bcast.out_len, lhs_dim = bcast.lhs_len, rhs_dim = bcast.rhs_len;
|
||||
runtime::parallel_for(0, csr.num_rows, [&](size_t b, size_t e) {
|
||||
for (auto rid = b; rid < e; ++rid) {
|
||||
const IdType row_start = indptr[rid], row_end = indptr[rid + 1];
|
||||
DType* out_off = O + rid * dim;
|
||||
for (IdType j = row_start; j < row_end; ++j) {
|
||||
const IdType cid = indices[j];
|
||||
const IdType eid = has_idx ? edges[j] : j;
|
||||
for (int64_t k = 0; k < dim; ++k) {
|
||||
const int64_t lhs_add = bcast.use_bcast ? bcast.lhs_offset[k] : k;
|
||||
const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
|
||||
const DType* lhs_off =
|
||||
Op::use_lhs ? X + cid * lhs_dim + lhs_add : nullptr;
|
||||
const DType* rhs_off =
|
||||
Op::use_rhs ? W + eid * rhs_dim + rhs_add : nullptr;
|
||||
out_off[k] += Op::Call(lhs_off, rhs_off);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// Naive implementation with additional accumulator, which prevents accuracy
|
||||
// degradation in less precise data types, like bfloat16.
|
||||
template <typename IdType, typename DType, typename Op>
|
||||
typename std::enable_if<std::is_same<DType, BFloat16>::value, void>::type
|
||||
SpMMSumCsrNaive(
|
||||
const BcastOff& bcast, const CSRMatrix& csr, const DType* X, const DType* W,
|
||||
DType* O) {
|
||||
const bool has_idx = !IsNullArray(csr.data);
|
||||
const IdType* indptr = csr.indptr.Ptr<IdType>();
|
||||
const IdType* indices = csr.indices.Ptr<IdType>();
|
||||
const IdType* edges = csr.data.Ptr<IdType>();
|
||||
int64_t dim = bcast.out_len, lhs_dim = bcast.lhs_len, rhs_dim = bcast.rhs_len;
|
||||
runtime::parallel_for(0, csr.num_rows, [&](size_t b, size_t e) {
|
||||
for (auto rid = b; rid < e; ++rid) {
|
||||
const IdType row_start = indptr[rid], row_end = indptr[rid + 1];
|
||||
DType* out_off = O + rid * dim;
|
||||
for (int64_t k = 0; k < dim; ++k) {
|
||||
AccType<DType> acc = 0.;
|
||||
for (IdType j = row_start; j < row_end; ++j) {
|
||||
const IdType cid = indices[j];
|
||||
const IdType eid = has_idx ? edges[j] : j;
|
||||
const int64_t lhs_add = bcast.use_bcast ? bcast.lhs_offset[k] : k;
|
||||
const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
|
||||
const DType* lhs_off =
|
||||
Op::use_lhs ? X + cid * lhs_dim + lhs_add : nullptr;
|
||||
const DType* rhs_off =
|
||||
Op::use_rhs ? W + eid * rhs_dim + rhs_add : nullptr;
|
||||
acc += Op::Call(lhs_off, rhs_off);
|
||||
}
|
||||
out_off[k] += acc;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CPU kernel of SpMM on Csr format.
|
||||
* @param bcast Broadcast information.
|
||||
* @param csr The Csr matrix.
|
||||
* @param ufeat The feature on source nodes.
|
||||
* @param efeat The feature on edges.
|
||||
* @param out The result feature on destination nodes.
|
||||
* @note it uses node parallel strategy, different threads are responsible
|
||||
* for the computation of different nodes.
|
||||
*/
|
||||
template <typename IdType, typename DType, typename Op>
|
||||
void SpMMSumCsr(
|
||||
const BcastOff& bcast, const CSRMatrix& csr, NDArray ufeat, NDArray efeat,
|
||||
NDArray out) {
|
||||
const bool has_idx = !IsNullArray(csr.data);
|
||||
const IdType* indptr = csr.indptr.Ptr<IdType>();
|
||||
const IdType* indices = csr.indices.Ptr<IdType>();
|
||||
const IdType* edges = csr.data.Ptr<IdType>();
|
||||
const DType* X = ufeat.Ptr<DType>();
|
||||
const DType* W = efeat.Ptr<DType>();
|
||||
DType* O = out.Ptr<DType>();
|
||||
CHECK_NOTNULL(indptr);
|
||||
CHECK_NOTNULL(O);
|
||||
if (Op::use_lhs) {
|
||||
CHECK_NOTNULL(indices);
|
||||
CHECK_NOTNULL(X);
|
||||
}
|
||||
if (Op::use_rhs) {
|
||||
if (has_idx) CHECK_NOTNULL(edges);
|
||||
CHECK_NOTNULL(W);
|
||||
}
|
||||
#if !defined(_WIN32)
|
||||
#ifdef USE_LIBXSMM
|
||||
int cpu_id = libxsmm_cpuid_x86();
|
||||
const bool no_libxsmm =
|
||||
bcast.use_bcast || std::is_same<DType, double>::value ||
|
||||
(std::is_same<DType, BFloat16>::value && cpu_id < LIBXSMM_X86_AVX512) ||
|
||||
!dgl::runtime::Config::Global()->IsLibxsmmAvailable();
|
||||
if (!no_libxsmm) {
|
||||
SpMMSumCsrLibxsmm<IdType, DType, Op>(bcast, csr, ufeat, efeat, out);
|
||||
} else {
|
||||
#endif // USE_LIBXSMM
|
||||
#endif // _WIN32
|
||||
SpMMSumCsrNaive<IdType, DType, Op>(bcast, csr, X, W, O);
|
||||
#if !defined(_WIN32)
|
||||
#ifdef USE_LIBXSMM
|
||||
}
|
||||
#endif // USE_LIBXSMM
|
||||
#endif // _WIN32
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CPU kernel of SpMM on Coo format.
|
||||
* @param bcast Broadcast information.
|
||||
* @param coo The Coo matrix.
|
||||
* @param ufeat The feature on source nodes.
|
||||
* @param efeat The feature on edges.
|
||||
* @param out The result feature on destination nodes.
|
||||
* @note it uses node parallel strategy, different threads are responsible
|
||||
* for the computation of different nodes. To avoid possible data hazard,
|
||||
* we use atomic operators in the reduction phase.
|
||||
*/
|
||||
template <typename IdType, typename DType, typename Op>
|
||||
typename std::enable_if<!std::is_same<DType, BFloat16>::value, void>::type
|
||||
SpMMSumCoo(
|
||||
const BcastOff& bcast, const COOMatrix& coo, NDArray ufeat, NDArray efeat,
|
||||
NDArray out) {
|
||||
const bool has_idx = !IsNullArray(coo.data);
|
||||
const IdType* row = coo.row.Ptr<IdType>();
|
||||
const IdType* col = coo.col.Ptr<IdType>();
|
||||
const IdType* edges = coo.data.Ptr<IdType>();
|
||||
const DType* X = ufeat.Ptr<DType>();
|
||||
const DType* W = efeat.Ptr<DType>();
|
||||
int64_t dim = bcast.out_len, lhs_dim = bcast.lhs_len, rhs_dim = bcast.rhs_len;
|
||||
DType* O = out.Ptr<DType>();
|
||||
const int64_t nnz = coo.row->shape[0];
|
||||
// fill zero elements
|
||||
memset(O, 0, out.GetSize());
|
||||
// spmm
|
||||
#pragma omp parallel for
|
||||
for (IdType i = 0; i < nnz; ++i) {
|
||||
const IdType rid = row[i];
|
||||
const IdType cid = col[i];
|
||||
const IdType eid = has_idx ? edges[i] : i;
|
||||
DType* out_off = O + cid * dim;
|
||||
for (int64_t k = 0; k < dim; ++k) {
|
||||
const int64_t lhs_add = bcast.use_bcast ? bcast.lhs_offset[k] : k;
|
||||
const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
|
||||
const DType* lhs_off =
|
||||
Op::use_lhs ? X + rid * lhs_dim + lhs_add : nullptr;
|
||||
const DType* rhs_off =
|
||||
Op::use_rhs ? W + eid * rhs_dim + rhs_add : nullptr;
|
||||
const DType val = Op::Call(lhs_off, rhs_off);
|
||||
if (val != 0) {
|
||||
#pragma omp atomic
|
||||
out_off[k] += val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename IdType, typename DType, typename Op>
|
||||
typename std::enable_if<std::is_same<DType, BFloat16>::value, void>::type
|
||||
SpMMSumCoo(
|
||||
const BcastOff& bcast, const COOMatrix& coo, NDArray ufeat, NDArray efeat,
|
||||
NDArray out) {
|
||||
LOG(FATAL) << "Unsupported CPU kernel for SpMMSumCoo for BF16.";
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CPU kernel of SpMM-Min/Max on Csr format.
|
||||
* @param bcast Broadcast information.
|
||||
* @param csr The Csr matrix.
|
||||
* @param ufeat The feature on source nodes.
|
||||
* @param efeat The feature on edges.
|
||||
* @param out The result feature on destination nodes.
|
||||
* @param argu Arg-Min/Max on source nodes, which refers the source node indices
|
||||
* correspond to the minimum/maximum values of reduction result on
|
||||
* destination nodes. It's useful in computing gradients of Min/Max
|
||||
* reducer.
|
||||
* @param arge Arg-Min/Max on edges. which refers the source node indices
|
||||
correspond to the minimum/maximum values of reduction result on
|
||||
* destination nodes. It's useful in computing gradients of Min/Max
|
||||
* reducer.
|
||||
* @note It uses node parallel strategy, different threads are responsible for
|
||||
* the computation of different nodes.
|
||||
* @note The result will contain infinity for zero-degree nodes.
|
||||
*/
|
||||
template <typename IdType, typename DType, typename Op, typename Cmp>
|
||||
void SpMMCmpCsr(
|
||||
const BcastOff& bcast, const CSRMatrix& csr, NDArray ufeat, NDArray efeat,
|
||||
NDArray out, NDArray argu, NDArray arge) {
|
||||
const bool has_idx = !IsNullArray(csr.data);
|
||||
const IdType* indptr = static_cast<IdType*>(csr.indptr->data);
|
||||
const IdType* indices = static_cast<IdType*>(csr.indices->data);
|
||||
const IdType* edges =
|
||||
has_idx ? static_cast<IdType*>(csr.data->data) : nullptr;
|
||||
const DType* X = Op::use_lhs ? static_cast<DType*>(ufeat->data) : nullptr;
|
||||
const DType* W = Op::use_rhs ? static_cast<DType*>(efeat->data) : nullptr;
|
||||
const int64_t dim = bcast.out_len, lhs_dim = bcast.lhs_len,
|
||||
rhs_dim = bcast.rhs_len;
|
||||
DType* O = static_cast<DType*>(out->data);
|
||||
IdType* argX = Op::use_lhs ? static_cast<IdType*>(argu->data) : nullptr;
|
||||
IdType* argW = Op::use_rhs ? static_cast<IdType*>(arge->data) : nullptr;
|
||||
CHECK_NOTNULL(indptr);
|
||||
CHECK_NOTNULL(O);
|
||||
if (Op::use_lhs) {
|
||||
CHECK_NOTNULL(indices);
|
||||
CHECK_NOTNULL(X);
|
||||
CHECK_NOTNULL(argX);
|
||||
}
|
||||
if (Op::use_rhs) {
|
||||
if (has_idx) CHECK_NOTNULL(edges);
|
||||
CHECK_NOTNULL(W);
|
||||
CHECK_NOTNULL(argW);
|
||||
}
|
||||
#if !defined(_WIN32)
|
||||
#ifdef USE_LIBXSMM
|
||||
int cpu_id = libxsmm_cpuid_x86();
|
||||
const bool no_libxsmm = bcast.use_bcast ||
|
||||
std::is_same<DType, double>::value ||
|
||||
cpu_id < LIBXSMM_X86_AVX512 ||
|
||||
!dgl::runtime::Config::Global()->IsLibxsmmAvailable();
|
||||
if (!no_libxsmm) {
|
||||
SpMMCmpCsrLibxsmm<IdType, DType, Op, Cmp>(
|
||||
bcast, csr, ufeat, efeat, out, argu, arge);
|
||||
} else {
|
||||
#endif // USE_LIBXSMM
|
||||
#endif // _WIN32
|
||||
|
||||
runtime::parallel_for(0, csr.num_rows, [&](size_t b, size_t e) {
|
||||
for (auto rid = b; rid < e; ++rid) {
|
||||
const IdType row_start = indptr[rid], row_end = indptr[rid + 1];
|
||||
DType* out_off = O + rid * dim;
|
||||
IdType* argx_off = argX + rid * dim;
|
||||
IdType* argw_off = argW + rid * dim;
|
||||
for (IdType j = row_start; j < row_end; ++j) {
|
||||
const IdType cid = indices[j];
|
||||
const IdType eid = has_idx ? edges[j] : j;
|
||||
for (int64_t k = 0; k < dim; ++k) {
|
||||
const int64_t lhs_add = bcast.use_bcast ? bcast.lhs_offset[k] : k;
|
||||
const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
|
||||
const DType* lhs_off =
|
||||
Op::use_lhs ? X + cid * lhs_dim + lhs_add : nullptr;
|
||||
const DType* rhs_off =
|
||||
Op::use_rhs ? W + eid * rhs_dim + rhs_add : nullptr;
|
||||
const DType val = Op::Call(lhs_off, rhs_off);
|
||||
if (Cmp::Call(out_off[k], val)) {
|
||||
out_off[k] = val;
|
||||
if (Op::use_lhs) argx_off[k] = cid;
|
||||
if (Op::use_rhs) argw_off[k] = eid;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
#if !defined(_WIN32)
|
||||
#ifdef USE_LIBXSMM
|
||||
}
|
||||
#endif // USE_LIBXSMM
|
||||
#endif // _WIN32
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CPU kernel of SpMM-Min/Max on Csr format.
|
||||
* @param bcast Broadcast information.
|
||||
* @param csr The Csr matrix.
|
||||
* @param ufeat The feature on source nodes.
|
||||
* @param efeat The feature on edges.
|
||||
* @param out The result feature on destination nodes.
|
||||
* @param argu Arg-Min/Max on source nodes, which refers the source node indices
|
||||
* correspond to the minimum/maximum values of reduction result on
|
||||
* destination nodes. It's useful in computing gradients of Min/Max
|
||||
* reducer.
|
||||
* @param arge Arg-Min/Max on edges. which refers the source node indices
|
||||
* correspond to the minimum/maximum values of reduction result on
|
||||
* destination nodes. It's useful in computing gradients of Min/Max
|
||||
* reducer.
|
||||
* @param argu_ntype Node type of the arg-Min/Max on source nodes, which refers
|
||||
* the source node types correspond to the minimum/maximum values of
|
||||
* reduction result on destination nodes. It's useful in computing
|
||||
* gradients of Min/Max reducer.
|
||||
* @param arge_etype Edge-type of the arg-Min/Max on edges. which refers the
|
||||
* source node indices correspond to the minimum/maximum values of
|
||||
* reduction result on destination nodes. It's useful in computing
|
||||
* gradients of Min/Max reducer.
|
||||
* @param src_type Node type of the source nodes of an etype
|
||||
* @param etype Edge type
|
||||
*/
|
||||
template <typename IdType, typename DType, typename Op, typename Cmp>
|
||||
void SpMMCmpCsrHetero(
|
||||
const BcastOff& bcast, const CSRMatrix& csr, NDArray ufeat, NDArray efeat,
|
||||
NDArray out, NDArray argu, NDArray arge, NDArray argu_ntype,
|
||||
NDArray arge_etype, const int ntype, const int etype) {
|
||||
const bool has_idx = !IsNullArray(csr.data);
|
||||
const IdType* indptr = static_cast<IdType*>(csr.indptr->data);
|
||||
const IdType* indices = static_cast<IdType*>(csr.indices->data);
|
||||
const IdType* edges =
|
||||
has_idx ? static_cast<IdType*>(csr.data->data) : nullptr;
|
||||
const DType* X = Op::use_lhs ? static_cast<DType*>(ufeat->data) : nullptr;
|
||||
const DType* W = Op::use_rhs ? static_cast<DType*>(efeat->data) : nullptr;
|
||||
const int64_t dim = bcast.out_len, lhs_dim = bcast.lhs_len,
|
||||
rhs_dim = bcast.rhs_len;
|
||||
DType* O = static_cast<DType*>(out->data);
|
||||
IdType* argX = Op::use_lhs ? static_cast<IdType*>(argu->data) : nullptr;
|
||||
IdType* argW = Op::use_rhs ? static_cast<IdType*>(arge->data) : nullptr;
|
||||
IdType* argX_ntype =
|
||||
Op::use_lhs ? static_cast<IdType*>(argu_ntype->data) : nullptr;
|
||||
IdType* argW_etype =
|
||||
Op::use_rhs ? static_cast<IdType*>(arge_etype->data) : nullptr;
|
||||
CHECK_NOTNULL(indptr);
|
||||
CHECK_NOTNULL(O);
|
||||
if (Op::use_lhs) {
|
||||
CHECK_NOTNULL(indices);
|
||||
CHECK_NOTNULL(X);
|
||||
CHECK_NOTNULL(argX);
|
||||
}
|
||||
if (Op::use_rhs) {
|
||||
if (has_idx) CHECK_NOTNULL(edges);
|
||||
CHECK_NOTNULL(W);
|
||||
CHECK_NOTNULL(argW);
|
||||
}
|
||||
// TODO(Israt): Use LIBXSMM. Homogeneous graph uses LIBXMM when enabled.
|
||||
runtime::parallel_for(0, csr.num_rows, [&](size_t b, size_t e) {
|
||||
for (auto rid = b; rid < e; ++rid) {
|
||||
const IdType row_start = indptr[rid], row_end = indptr[rid + 1];
|
||||
DType* out_off = O + rid * dim;
|
||||
IdType* argx_off = argX + rid * dim;
|
||||
IdType* argw_off = argW + rid * dim;
|
||||
IdType* argx_ntype = argX_ntype + rid * dim;
|
||||
IdType* argw_etype = argW_etype + rid * dim;
|
||||
for (IdType j = row_start; j < row_end; ++j) {
|
||||
const IdType cid = indices[j];
|
||||
const IdType eid = has_idx ? edges[j] : j;
|
||||
for (int64_t k = 0; k < dim; ++k) {
|
||||
const int64_t lhs_add = bcast.use_bcast ? bcast.lhs_offset[k] : k;
|
||||
const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
|
||||
const DType* lhs_off =
|
||||
Op::use_lhs ? X + cid * lhs_dim + lhs_add : nullptr;
|
||||
const DType* rhs_off =
|
||||
Op::use_rhs ? W + eid * rhs_dim + rhs_add : nullptr;
|
||||
const DType val = Op::Call(lhs_off, rhs_off);
|
||||
if (Cmp::Call(out_off[k], val)) {
|
||||
out_off[k] = val;
|
||||
if (Op::use_lhs) {
|
||||
argx_off[k] = cid;
|
||||
argx_ntype[k] = ntype;
|
||||
}
|
||||
if (Op::use_rhs) {
|
||||
argw_off[k] = eid;
|
||||
argw_etype[k] = etype;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CPU kernel of SpMM-Min/Max on Coo format.
|
||||
* @param bcast Broadcast information.
|
||||
* @param coo The Coo matrix.
|
||||
* @param ufeat The feature on source nodes.
|
||||
* @param efeat The feature on edges.
|
||||
* @param out The result feature on destination nodes.
|
||||
* @param argu Arg-Min/Max on source nodes, which refers the source node indices
|
||||
* correspond to the minimum/maximum values of reduction result on
|
||||
* destination nodes. It's useful in computing gradients of Min/Max
|
||||
* reducer.
|
||||
* @param arge Arg-Min/Max on edges. which refers the source node indices
|
||||
* correspond to the minimum/maximum values of reduction result on
|
||||
* destination nodes. It's useful in computing gradients of Min/Max
|
||||
* reducer.
|
||||
* @note it uses node parallel strategy, different threads are responsible for
|
||||
* the computation of different nodes. To avoid possible data hazard, we
|
||||
* use atomic operators in the reduction phase.
|
||||
* @note The result will contain infinity for zero-degree nodes.
|
||||
*/
|
||||
template <typename IdType, typename DType, typename Op, typename Cmp>
|
||||
void SpMMCmpCoo(
|
||||
const BcastOff& bcast, const COOMatrix& coo, NDArray ufeat, NDArray efeat,
|
||||
NDArray out, NDArray argu, NDArray arge) {
|
||||
const bool has_idx = !IsNullArray(coo.data);
|
||||
const IdType* row = static_cast<IdType*>(coo.row->data);
|
||||
const IdType* col = static_cast<IdType*>(coo.col->data);
|
||||
const IdType* edges =
|
||||
has_idx ? static_cast<IdType*>(coo.data->data) : nullptr;
|
||||
const DType* X = Op::use_lhs ? static_cast<DType*>(ufeat->data) : nullptr;
|
||||
const DType* W = Op::use_rhs ? static_cast<DType*>(efeat->data) : nullptr;
|
||||
const int64_t dim = bcast.out_len, lhs_dim = bcast.lhs_len,
|
||||
rhs_dim = bcast.rhs_len;
|
||||
DType* O = static_cast<DType*>(out->data);
|
||||
IdType* argX = Op::use_lhs ? static_cast<IdType*>(argu->data) : nullptr;
|
||||
IdType* argW = Op::use_rhs ? static_cast<IdType*>(arge->data) : nullptr;
|
||||
const int64_t nnz = coo.row->shape[0];
|
||||
// fill zero elements
|
||||
std::fill(O, O + out.NumElements(), Cmp::zero);
|
||||
// spmm
|
||||
#pragma omp parallel for
|
||||
for (IdType i = 0; i < nnz; ++i) {
|
||||
const IdType rid = row[i];
|
||||
const IdType cid = col[i];
|
||||
const IdType eid = has_idx ? edges[i] : i;
|
||||
DType* out_off = O + cid * dim;
|
||||
IdType* argx_off = Op::use_lhs ? argX + cid * dim : nullptr;
|
||||
IdType* argw_off = Op::use_rhs ? argW + cid * dim : nullptr;
|
||||
for (int64_t k = 0; k < dim; ++k) {
|
||||
const int64_t lhs_add = bcast.use_bcast ? bcast.lhs_offset[k] : k;
|
||||
const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
|
||||
const DType* lhs_off =
|
||||
Op::use_lhs ? X + rid * lhs_dim + lhs_add : nullptr;
|
||||
const DType* rhs_off =
|
||||
Op::use_rhs ? W + eid * rhs_dim + rhs_add : nullptr;
|
||||
const DType val = Op::Call(lhs_off, rhs_off);
|
||||
#pragma omp critical
|
||||
if (Cmp::Call(out_off[k], val)) {
|
||||
out_off[k] = val;
|
||||
if (Op::use_lhs) argx_off[k] = rid;
|
||||
if (Op::use_rhs) argw_off[k] = eid;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CPU kernel of Edge_softmax_csr_forward on Csr format.
|
||||
* @param bcast Broadcast information.
|
||||
* @param csr The Csr matrix.
|
||||
* @param ufeat The feature on source nodes.
|
||||
* @param efeat The feature on edges.
|
||||
* @param out The result of edge_softmax_forward.
|
||||
*/
|
||||
template <typename IdType, typename DType, typename Op>
|
||||
void Edge_softmax_csr_forward(
|
||||
const BcastOff& bcast, const CSRMatrix& csr, NDArray ufeat, NDArray efeat,
|
||||
NDArray out) {
|
||||
const bool has_idx = !IsNullArray(csr.data);
|
||||
const IdType* indptr = static_cast<IdType*>(csr.indptr->data);
|
||||
const IdType* edges =
|
||||
has_idx ? static_cast<IdType*>(csr.data->data) : nullptr;
|
||||
const DType* W = Op::use_rhs ? static_cast<DType*>(efeat->data) : nullptr;
|
||||
const int64_t dim = bcast.out_len, rhs_dim = bcast.rhs_len;
|
||||
runtime::parallel_for(0, csr.num_rows, [&](size_t b, size_t e) {
|
||||
for (auto rid = b; rid < e; ++rid) {
|
||||
const IdType row_start = indptr[rid], row_end = indptr[rid + 1];
|
||||
std::vector<AccType<DType>> data_e(row_end - row_start, 0);
|
||||
std::vector<IdType> num(row_end - row_start, 0);
|
||||
for (int64_t k = 0; k < dim; ++k) {
|
||||
DType max_v = -std::numeric_limits<DType>::infinity();
|
||||
for (IdType j = row_start; j < row_end; ++j) {
|
||||
const IdType eid = has_idx ? edges[j] : j;
|
||||
const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
|
||||
const DType* rhs_off =
|
||||
Op::use_rhs ? W + eid * rhs_dim + rhs_add : nullptr;
|
||||
data_e[j - row_start] = *rhs_off;
|
||||
num[j - row_start] = eid * rhs_dim + rhs_add;
|
||||
max_v = std::max<DType>(max_v, (*rhs_off));
|
||||
}
|
||||
DType exp_sum = 0;
|
||||
for (auto& element : data_e) {
|
||||
element -= max_v;
|
||||
element = std::exp(element);
|
||||
exp_sum += element;
|
||||
}
|
||||
for (int i = 0; i < row_end - row_start; i++) {
|
||||
out.Ptr<DType>()[num[i]] = data_e[i] / exp_sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CPU kernel of Edge_softmax_csr_backward on Csr format.
|
||||
* @param bcast Broadcast information.
|
||||
* @param csr The Csr matrix.
|
||||
* @param out The result of forward.
|
||||
* @param sds The result of gradiet * out.
|
||||
* @param back_out The result of edge_softmax_backward.
|
||||
*/
|
||||
template <typename IdType, typename DType, typename Op>
|
||||
void Edge_softmax_csr_backward(
|
||||
const BcastOff& bcast, const CSRMatrix& csr, NDArray out, NDArray sds,
|
||||
NDArray back_out) {
|
||||
typedef typename std::conditional<
|
||||
std::is_same<DType, BFloat16>::value, float, DType>::type AccType;
|
||||
const bool has_idx = !IsNullArray(csr.data);
|
||||
const IdType* indptr = static_cast<IdType*>(csr.indptr->data);
|
||||
const IdType* edges =
|
||||
has_idx ? static_cast<IdType*>(csr.data->data) : nullptr;
|
||||
const DType* W_out = Op::use_rhs ? static_cast<DType*>(out->data) : nullptr;
|
||||
const DType* W_sds = Op::use_rhs ? static_cast<DType*>(sds->data) : nullptr;
|
||||
const int64_t dim = bcast.out_len, rhs_dim = bcast.rhs_len;
|
||||
runtime::parallel_for(0, csr.num_rows, [&](size_t b, size_t e) {
|
||||
for (auto rid = b; rid < e; ++rid) {
|
||||
const IdType row_start = indptr[rid], row_end = indptr[rid + 1];
|
||||
for (int64_t k = 0; k < dim; ++k) {
|
||||
AccType sum_sds = 0;
|
||||
for (IdType j = row_start; j < row_end; ++j) {
|
||||
const IdType eid = has_idx ? edges[j] : j;
|
||||
const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
|
||||
const DType* rhs_off_sds =
|
||||
Op::use_rhs ? W_sds + eid * rhs_dim + rhs_add : nullptr;
|
||||
sum_sds += (*rhs_off_sds);
|
||||
}
|
||||
for (IdType j = row_start; j < row_end; ++j) {
|
||||
const IdType eid = has_idx ? edges[j] : j;
|
||||
const int64_t rhs_add = bcast.use_bcast ? bcast.rhs_offset[k] : k;
|
||||
const DType* rhs_off_out =
|
||||
Op::use_rhs ? W_out + eid * rhs_dim + rhs_add : nullptr;
|
||||
const DType* rhs_off_sds =
|
||||
Op::use_rhs ? W_sds + eid * rhs_dim + rhs_add : nullptr;
|
||||
back_out.Ptr<DType>()[eid * rhs_dim + rhs_add] =
|
||||
(*rhs_off_sds) - sum_sds * (*rhs_off_out);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
} // namespace cpu
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CPU_SPMM_H_
|
||||
@@ -0,0 +1,172 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/spmm_binary_ops.h
|
||||
* @brief SPMM CPU Binary ops.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CPU_SPMM_BINARY_OPS_H_
|
||||
#define DGL_ARRAY_CPU_SPMM_BINARY_OPS_H_
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/bcast.h>
|
||||
|
||||
#include <limits>
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace cpu {
|
||||
namespace op {
|
||||
|
||||
//////////////////////////////// binary operators on CPU
|
||||
///////////////////////////////////
|
||||
template <typename DType>
|
||||
struct Add {
|
||||
typedef DType type;
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = true;
|
||||
inline static DType Call(const DType* lhs_off, const DType* rhs_off) {
|
||||
return *lhs_off + *rhs_off;
|
||||
}
|
||||
};
|
||||
template <typename DType>
|
||||
constexpr bool Add<DType>::use_lhs;
|
||||
template <typename DType>
|
||||
constexpr bool Add<DType>::use_rhs;
|
||||
|
||||
template <typename DType>
|
||||
struct Sub {
|
||||
typedef DType type;
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = true;
|
||||
inline static DType Call(const DType* lhs_off, const DType* rhs_off) {
|
||||
return *lhs_off - *rhs_off;
|
||||
}
|
||||
};
|
||||
template <typename DType>
|
||||
constexpr bool Sub<DType>::use_lhs;
|
||||
template <typename DType>
|
||||
constexpr bool Sub<DType>::use_rhs;
|
||||
|
||||
template <typename DType>
|
||||
struct Mul {
|
||||
typedef DType type;
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = true;
|
||||
inline static DType Call(const DType* lhs_off, const DType* rhs_off) {
|
||||
return *lhs_off * *rhs_off;
|
||||
}
|
||||
};
|
||||
template <typename DType>
|
||||
constexpr bool Mul<DType>::use_lhs;
|
||||
template <typename DType>
|
||||
constexpr bool Mul<DType>::use_rhs;
|
||||
|
||||
template <typename DType>
|
||||
struct Div {
|
||||
typedef DType type;
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = true;
|
||||
inline static DType Call(const DType* lhs_off, const DType* rhs_off) {
|
||||
return *lhs_off / *rhs_off;
|
||||
}
|
||||
};
|
||||
template <typename DType>
|
||||
constexpr bool Div<DType>::use_lhs;
|
||||
template <typename DType>
|
||||
constexpr bool Div<DType>::use_rhs;
|
||||
|
||||
template <typename DType>
|
||||
struct CopyLhs {
|
||||
typedef DType type;
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = false;
|
||||
inline static DType Call(const DType* lhs_off, const DType*) {
|
||||
return *lhs_off;
|
||||
}
|
||||
};
|
||||
template <typename DType>
|
||||
constexpr bool CopyLhs<DType>::use_lhs;
|
||||
template <typename DType>
|
||||
constexpr bool CopyLhs<DType>::use_rhs;
|
||||
|
||||
template <typename DType>
|
||||
struct CopyRhs {
|
||||
typedef DType type;
|
||||
static constexpr bool use_lhs = false;
|
||||
static constexpr bool use_rhs = true;
|
||||
inline static DType Call(const DType*, const DType* rhs_off) {
|
||||
return *rhs_off;
|
||||
}
|
||||
};
|
||||
template <typename DType>
|
||||
constexpr bool CopyRhs<DType>::use_lhs;
|
||||
template <typename DType>
|
||||
constexpr bool CopyRhs<DType>::use_rhs;
|
||||
|
||||
//////////////////////////////// Reduce operators on CPU
|
||||
///////////////////////////////////
|
||||
template <typename DType>
|
||||
constexpr DType MinDType() {
|
||||
if (std::is_same<DType, BFloat16>::value)
|
||||
return BFloat16::Min();
|
||||
else
|
||||
return -std::numeric_limits<DType>::infinity();
|
||||
}
|
||||
|
||||
template <typename DType>
|
||||
struct Max {
|
||||
typedef DType type;
|
||||
static constexpr DType zero = MinDType<DType>();
|
||||
// return true if accum should be replaced
|
||||
inline static DType Call(DType accum, DType val) { return accum < val; }
|
||||
};
|
||||
template <typename DType>
|
||||
constexpr DType Max<DType>::zero;
|
||||
|
||||
template <typename DType>
|
||||
constexpr DType MaxDType() {
|
||||
if (std::is_same<DType, BFloat16>::value)
|
||||
return BFloat16::Max();
|
||||
else
|
||||
return std::numeric_limits<DType>::infinity();
|
||||
}
|
||||
|
||||
template <typename DType>
|
||||
struct Min {
|
||||
typedef DType type;
|
||||
static constexpr DType zero = MaxDType<DType>();
|
||||
// return true if accum should be replaced
|
||||
inline static DType Call(DType accum, DType val) { return accum > val; }
|
||||
};
|
||||
template <typename DType>
|
||||
constexpr DType Min<DType>::zero;
|
||||
|
||||
#define SWITCH_OP(op, Op, ...) \
|
||||
do { \
|
||||
if ((op) == "add") { \
|
||||
typedef dgl::aten::cpu::op::Add<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "sub") { \
|
||||
typedef dgl::aten::cpu::op::Sub<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "mul") { \
|
||||
typedef dgl::aten::cpu::op::Mul<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "div") { \
|
||||
typedef dgl::aten::cpu::op::Div<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "copy_lhs") { \
|
||||
typedef dgl::aten::cpu::op::CopyLhs<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "copy_rhs") { \
|
||||
typedef dgl::aten::cpu::op::CopyRhs<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else { \
|
||||
LOG(FATAL) << "Unsupported SpMM binary operator: " << op; \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
} // namespace op
|
||||
|
||||
} // namespace cpu
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CPU_SPMM_BINARY_OPS_H_
|
||||
@@ -0,0 +1,603 @@
|
||||
/**
|
||||
* Copyright (c) 2021 Intel Corporation
|
||||
* @file array/cpu/spmm.h
|
||||
* @brief SPMM CPU kernel function header.
|
||||
* @author Sanchit Misra <sanchit.misra@intel.com>,
|
||||
* Ramanarayan Mohanty <ramanarayan.mohanty@intel.com>,
|
||||
* Vasimuddin Md <vasimuddin.md@intel.com>,
|
||||
* Sasikanth Avancha <sasikanth.avancha@intel.com>
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CPU_SPMM_BLOCKING_LIBXSMM_H_
|
||||
#define DGL_ARRAY_CPU_SPMM_BLOCKING_LIBXSMM_H_
|
||||
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/bcast.h>
|
||||
#include <dmlc/logging.h>
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
#if !defined(_WIN32)
|
||||
#ifdef USE_LIBXSMM
|
||||
#include <libxsmm_source.h>
|
||||
#include <unistd.h>
|
||||
#ifdef DEBUG
|
||||
#include <x86intrin.h>
|
||||
#endif // DEBUG
|
||||
#include <dmlc/omp.h>
|
||||
|
||||
#define NUM_BLOCKS_PER_THREAD 20
|
||||
#define BLOCKING_HEURISTIC_PARAM 500
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace cpu {
|
||||
|
||||
template <typename IdType, typename DType>
|
||||
struct CSRMatrixInternal {
|
||||
IdType num_rows;
|
||||
IdType num_cols;
|
||||
IdType *indptr;
|
||||
IdType *indices;
|
||||
DType *data;
|
||||
};
|
||||
|
||||
int32_t GetLLCSize() {
|
||||
#ifdef _SC_LEVEL3_CACHE_SIZE
|
||||
int32_t cache_size = sysconf(_SC_LEVEL3_CACHE_SIZE);
|
||||
if (cache_size < 0) cache_size = DGL_CPU_LLC_SIZE;
|
||||
#else
|
||||
int32_t cache_size = DGL_CPU_LLC_SIZE;
|
||||
#endif
|
||||
return cache_size;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Tile the CSR matrix to roughly make sure that the column tiles and
|
||||
* corresponding neighbor features fit into LLC and the row tiles
|
||||
* are assigned to OMP threads.
|
||||
* @param csr The Csr matrix.
|
||||
* @param block_csr_array The array containing csr matrices of all blocks.
|
||||
* @param num_M_blocks Number of blocks to create along the rows of adjacency
|
||||
* matrix.
|
||||
* @param num_K_blocks Number of blocks to create along the columns of adjacency
|
||||
* matrix.
|
||||
* @param M_block_size block size along the rows of adjacency matrix.
|
||||
* @param K_block_size block size along the columns of adjacency matrix.
|
||||
* @param use_lhs Whether to use lhs.
|
||||
* @param use_rhs Whether to use rhs.
|
||||
*/
|
||||
template <typename IdType>
|
||||
inline void SpMMCreateBlocks(
|
||||
const CSRMatrix &csr, CSRMatrixInternal<IdType, IdType> *block_csr_array,
|
||||
IdType num_M_blocks, IdType num_K_blocks, IdType M_block_size,
|
||||
IdType K_block_size, bool use_lhs, bool use_rhs) {
|
||||
const IdType M = csr.num_rows;
|
||||
const IdType K = csr.num_cols;
|
||||
IdType *indptr = csr.indptr.Ptr<IdType>();
|
||||
IdType *indices = csr.indices.Ptr<IdType>();
|
||||
IdType *edges = csr.data.Ptr<IdType>();
|
||||
CHECK_NOTNULL(indptr);
|
||||
if (use_lhs) CHECK_NOTNULL(indices);
|
||||
if (use_rhs) CHECK_NOTNULL(edges);
|
||||
|
||||
if (num_K_blocks > 1) {
|
||||
IdType *indptr_block_buf = reinterpret_cast<IdType *>(aligned_alloc(
|
||||
64, (M_block_size + 1) * num_M_blocks * num_K_blocks * sizeof(IdType)));
|
||||
IdType *indices_block_buf = nullptr;
|
||||
if (use_lhs) {
|
||||
indices_block_buf = reinterpret_cast<IdType *>(
|
||||
aligned_alloc(64, indptr[M] * sizeof(IdType)));
|
||||
}
|
||||
IdType *edges_block_buf = nullptr;
|
||||
if (use_rhs) {
|
||||
edges_block_buf = reinterpret_cast<IdType *>(
|
||||
aligned_alloc(64, indptr[M] * sizeof(IdType)));
|
||||
}
|
||||
|
||||
#pragma omp parallel
|
||||
{
|
||||
IdType *my_cur_col_id = reinterpret_cast<IdType *>(
|
||||
aligned_alloc(64, 2 * M_block_size * sizeof(IdType)));
|
||||
|
||||
#pragma omp for
|
||||
for (IdType m = 0; m < num_M_blocks; m++) {
|
||||
const IdType M_start = m * M_block_size;
|
||||
const IdType M_end = std::min((m + 1) * M_block_size, M);
|
||||
const IdType nnz = indptr[M_end] - indptr[M_start];
|
||||
|
||||
IdType cur_indices_id = 0;
|
||||
IdType *my_indices_block_buf, *my_edges_block_buf;
|
||||
if (use_lhs) my_indices_block_buf = indices_block_buf + indptr[M_start];
|
||||
if (use_rhs) my_edges_block_buf = edges_block_buf + indptr[M_start];
|
||||
|
||||
for (IdType i = M_start; i < M_end; i++) {
|
||||
my_cur_col_id[(i - M_start) * 2] = indptr[i];
|
||||
my_cur_col_id[(i - M_start) * 2 + 1] = indptr[i + 1];
|
||||
}
|
||||
for (IdType k = 0; k < num_K_blocks; k++) {
|
||||
const IdType K_start = k * K_block_size;
|
||||
const IdType K_end = std::min((k + 1) * K_block_size, K);
|
||||
CSRMatrixInternal<IdType, IdType> cur_csr;
|
||||
cur_csr.num_rows = M_end - M_start;
|
||||
cur_csr.num_cols = K_end - K_start;
|
||||
// Create csr_ij
|
||||
IdType *cur_csr_indptr =
|
||||
indptr_block_buf + (m * num_K_blocks + k) * (M_block_size + 1);
|
||||
IdType *cur_csr_indices = nullptr, *cur_csr_edges = nullptr;
|
||||
if (use_lhs) cur_csr_indices = my_indices_block_buf + cur_indices_id;
|
||||
if (use_rhs) cur_csr_edges = my_edges_block_buf + cur_indices_id;
|
||||
IdType cur_nnz = 0;
|
||||
for (IdType i = M_start; i < M_end; i++) {
|
||||
const IdType row_start = my_cur_col_id[(i - M_start) * 2];
|
||||
const IdType row_end = my_cur_col_id[(i - M_start) * 2 + 1];
|
||||
cur_csr_indptr[i - M_start] = cur_nnz;
|
||||
IdType eid;
|
||||
for (eid = row_start; eid < row_end; eid++) {
|
||||
const IdType src = indices[eid];
|
||||
const IdType edge = edges[eid];
|
||||
if (src >= K_end) {
|
||||
break;
|
||||
}
|
||||
CHECK_LT(cur_indices_id + cur_nnz, nnz);
|
||||
if (use_lhs) cur_csr_indices[cur_nnz] = src;
|
||||
if (use_rhs) cur_csr_edges[cur_nnz] = edge;
|
||||
cur_nnz++;
|
||||
}
|
||||
my_cur_col_id[(i - M_start) * 2] = eid;
|
||||
}
|
||||
cur_csr_indptr[cur_csr.num_rows] = cur_nnz;
|
||||
cur_indices_id += cur_nnz;
|
||||
cur_csr.indptr = cur_csr_indptr;
|
||||
if (use_lhs) cur_csr.indices = cur_csr_indices;
|
||||
if (use_rhs) cur_csr.data = cur_csr_edges;
|
||||
block_csr_array[m * num_K_blocks + k] = cur_csr;
|
||||
}
|
||||
CHECK_EQ(nnz, cur_indices_id);
|
||||
}
|
||||
free(my_cur_col_id);
|
||||
}
|
||||
} else {
|
||||
for (IdType m = 0; m < num_M_blocks; m++) {
|
||||
const IdType M_start = m * M_block_size;
|
||||
const IdType M_end = std::min((m + 1) * M_block_size, M);
|
||||
|
||||
CSRMatrixInternal<IdType, IdType> cur_csr;
|
||||
cur_csr.num_rows = M_end - M_start;
|
||||
cur_csr.num_cols = K;
|
||||
cur_csr.indptr = indptr + M_start;
|
||||
cur_csr.indices = indices;
|
||||
cur_csr.data = edges;
|
||||
|
||||
block_csr_array[m] = cur_csr;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Create libxsmm kernel.
|
||||
* @param has_idx For the edge features, are there indices available.
|
||||
* @param N Feature size.
|
||||
* @param redop_flag Flag specifying the reduction operation.
|
||||
* @param is_cmp Is the reduction operation a compare operation.
|
||||
* @note libxsmm_dispatch_meltw_opreduce_vecs_idx creates a JIT'ed kernel.
|
||||
* Given a node u, the kernel performs an elementwise "Op" on the
|
||||
* features of the neighbors and/or the edges incident on u.
|
||||
* Subsequently, it performs an elementwise "Redop" on all such
|
||||
* features created and stores into the feature of node u.
|
||||
* It uses a SIMD and a cache efficient design and also provides
|
||||
* support to enable software prefetching if needed. For IdType,
|
||||
* it supports INT32 and INT64. For DType, it supports BF16 and FP32.
|
||||
* It supports all the "Ops" and "Redops" supported by DGL. Once a
|
||||
* kernel is generated by libxsmm_dispatch_meltw_opreduce_vecs_idx,
|
||||
* it is cached for the entire duration of the execution of a program
|
||||
* so that subsequently if the kernel is needed again, it just returns
|
||||
* the cached copy.
|
||||
*/
|
||||
template <typename IdType, typename DType, typename Op>
|
||||
inline libxsmm_meltwfunction_opreduce_vecs_idx SpMMCreateLibxsmmKernel(
|
||||
bool has_idx, IdType N, libxsmm_meltw_opreduce_vecs_flags redop_flag,
|
||||
bool is_cmp) {
|
||||
int _ld = N;
|
||||
libxsmm_meltw_opreduce_vecs_flags opredop_flags;
|
||||
// First, set the Op in the opredop_flags
|
||||
if (std::is_same<Op, op::Add<DType>>::value) {
|
||||
opredop_flags = LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_OP_ADD;
|
||||
} else if (std::is_same<Op, op::Sub<DType>>::value) {
|
||||
opredop_flags = LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_OP_SUB;
|
||||
} else if (std::is_same<Op, op::Mul<DType>>::value) {
|
||||
opredop_flags = LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_OP_MUL;
|
||||
} else if (std::is_same<Op, op::Div<DType>>::value) {
|
||||
opredop_flags = LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_OP_DIV;
|
||||
} else if (std::is_same<Op, op::CopyLhs<DType>>::value) {
|
||||
opredop_flags = LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_OP_COPY;
|
||||
} else if (std::is_same<Op, op::CopyRhs<DType>>::value) {
|
||||
opredop_flags = LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_OP_COPY;
|
||||
}
|
||||
// Second, set which of lhs or rhs is considered first and second operand.
|
||||
// This is needed since libxsmm assumes that the copy operation always copies
|
||||
// the first operand. So, if we need to copy rhs, we need to set that as the
|
||||
// first operand. For rhs, we also set whether to use implicit indices or
|
||||
// provided indices.
|
||||
// TODO(Steve): fix this long line in a separate PR.
|
||||
if (std::is_same<Op, op::CopyLhs<DType>>::value) {
|
||||
opredop_flags =
|
||||
(libxsmm_meltw_opreduce_vecs_flags)(opredop_flags | LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_OPORDER_VECIDX_VECIN); // NOLINT
|
||||
} else if (std::is_same<Op, op::CopyRhs<DType>>::value) {
|
||||
opredop_flags =
|
||||
(libxsmm_meltw_opreduce_vecs_flags)(opredop_flags | LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_OPORDER_VECIN_VECIDX); // NOLINT
|
||||
if (!has_idx) {
|
||||
opredop_flags =
|
||||
(libxsmm_meltw_opreduce_vecs_flags)(opredop_flags | LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_IMPLICIT_INDEXED_VECIDX); // NOLINT
|
||||
}
|
||||
} else {
|
||||
opredop_flags =
|
||||
(libxsmm_meltw_opreduce_vecs_flags)(opredop_flags | LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_OPORDER_VECIDX_VECIN); // NOLINT
|
||||
if (has_idx) {
|
||||
opredop_flags =
|
||||
(libxsmm_meltw_opreduce_vecs_flags)(opredop_flags | LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_INDEXED_VEC); // NOLINT
|
||||
} else {
|
||||
opredop_flags =
|
||||
(libxsmm_meltw_opreduce_vecs_flags)(opredop_flags | LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_IMPLICIT_INDEXED_VEC); // NOLINT
|
||||
}
|
||||
}
|
||||
// Third, we set the Redop in the opredop_flags
|
||||
opredop_flags =
|
||||
(libxsmm_meltw_opreduce_vecs_flags)(opredop_flags | redop_flag);
|
||||
// Fourth, in case of Cmp Redop, set whether to record argmax/argmin for
|
||||
// lhs/rhs
|
||||
if (is_cmp) {
|
||||
if (Op::use_lhs) {
|
||||
opredop_flags =
|
||||
(libxsmm_meltw_opreduce_vecs_flags)(opredop_flags | LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_RECORD_ARGOP_OFF_VEC_0); // NOLINT
|
||||
}
|
||||
if (Op::use_rhs) {
|
||||
opredop_flags =
|
||||
(libxsmm_meltw_opreduce_vecs_flags)(opredop_flags | LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_RECORD_ARGOP_OFF_VEC_1); // NOLINT
|
||||
}
|
||||
}
|
||||
libxsmm_meltwfunction_opreduce_vecs_idx kernel = nullptr;
|
||||
if (std::is_same<DType, float>::value) {
|
||||
kernel = libxsmm_dispatch_meltw_opreduce_vecs_idx(
|
||||
N, &_ld, &_ld, LIBXSMM_DATATYPE_F32, LIBXSMM_DATATYPE_F32,
|
||||
(sizeof(IdType) == 8) ? LIBXSMM_DATATYPE_I64 : LIBXSMM_DATATYPE_I32,
|
||||
opredop_flags, 0);
|
||||
} else { // assume bf16
|
||||
kernel = libxsmm_dispatch_meltw_opreduce_vecs_idx(
|
||||
N, &_ld, &_ld, LIBXSMM_DATATYPE_BF16, LIBXSMM_DATATYPE_BF16,
|
||||
(sizeof(IdType) == 8) ? LIBXSMM_DATATYPE_I64 : LIBXSMM_DATATYPE_I32,
|
||||
opredop_flags, 0);
|
||||
}
|
||||
|
||||
if (kernel == nullptr) {
|
||||
LOG(FATAL) << "Failed to generate libxsmm kernel for the SpMM operation."
|
||||
"To disable libxsmm, use dgl.use_libxsmm(false).";
|
||||
}
|
||||
return kernel;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Use libxsmm to perform SpMM-Sum on all blocks.
|
||||
* @param block_csr_array The array containing csr matrices of all blocks.
|
||||
* @param B The feature on source nodes.
|
||||
* @param E The feature on edges.
|
||||
* @param C The result feature on destination nodes.
|
||||
* @param has_idx For the edge features, are there indices available.
|
||||
* @param N Feature size.
|
||||
* @param num_M_blocks Number of blocks to create along the rows of adjacency
|
||||
* matrix.
|
||||
* @param num_K_blocks Number of blocks to create along the columns of adjacency
|
||||
* matrix.
|
||||
* @param M_block_size block size along the rows of adjacency matrix.
|
||||
* @param kernel The libxsmm kernel.
|
||||
*/
|
||||
template <typename IdType, typename DType>
|
||||
inline void SpMMBlockwiseOpSum(
|
||||
CSRMatrixInternal<IdType, IdType> *block_csr_array, const DType *B,
|
||||
const DType *E, DType *C, bool has_idx, IdType N, IdType num_M_blocks,
|
||||
IdType num_K_blocks, IdType M_block_size,
|
||||
libxsmm_meltwfunction_opreduce_vecs_idx kernel) {
|
||||
const DType *in_matrix1 = B;
|
||||
const DType *in_matrix2 = E;
|
||||
DType *output = C;
|
||||
#pragma omp parallel
|
||||
{
|
||||
for (IdType k = 0; k < num_K_blocks; k++) {
|
||||
#pragma omp for schedule(dynamic)
|
||||
for (IdType m = 0; m < num_M_blocks; m++) {
|
||||
CSRMatrixInternal<IdType, IdType> cur_csr =
|
||||
block_csr_array[m * num_K_blocks + k];
|
||||
|
||||
const IdType M_start = m * M_block_size;
|
||||
for (IdType i = 0; i < cur_csr.num_rows; i++) {
|
||||
const IdType row_start = cur_csr.indptr[i];
|
||||
const IdType row_end = cur_csr.indptr[i + 1];
|
||||
const IdType dst = i + M_start;
|
||||
|
||||
libxsmm_meltw_opreduce_vecs_idx_param params;
|
||||
params.n = row_end - row_start;
|
||||
params.indices = &cur_csr.indices[row_start];
|
||||
params.in_matrix = in_matrix1;
|
||||
params.out_vec = &output[dst * N];
|
||||
params.scale_vals = nullptr;
|
||||
if (has_idx) {
|
||||
params.in_matrix2 = in_matrix2;
|
||||
params.indices2 = &cur_csr.data[row_start];
|
||||
} else {
|
||||
params.in_matrix2 = &in_matrix2[row_start * N];
|
||||
}
|
||||
kernel(¶ms);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Use libxsmm to perform SpMM-Max/Min on all blocks.
|
||||
* @param block_csr_array The array containing csr matrices of all blocks.
|
||||
* @param B The feature on source nodes.
|
||||
* @param E The feature on edges.
|
||||
* @param C The result feature on destination nodes.
|
||||
* @param argB Arg-Min/Max on source nodes.
|
||||
* @param argE Arg-Min/Max on edges.
|
||||
* @param has_idx For the edge features, are there indices available.
|
||||
* @param N Feature size.
|
||||
* @param num_M_blocks Number of blocks to create along the rows of adjacency
|
||||
* matrix.
|
||||
* @param num_K_blocks Number of blocks to create along the columns of adjacency
|
||||
* matrix.
|
||||
* @param M_block_size block size along the rows of adjacency matrix.
|
||||
* @param kernel The libxsmm kernel.
|
||||
*/
|
||||
template <typename IdType, typename DType, typename Op, typename Cmp>
|
||||
inline void SpMMBlockwiseOpCmp(
|
||||
CSRMatrixInternal<IdType, IdType> *block_csr_array, const DType *B,
|
||||
const DType *E, DType *C, IdType *argB, IdType *argE, bool has_idx,
|
||||
IdType N, IdType num_M_blocks, IdType num_K_blocks, IdType M_block_size,
|
||||
libxsmm_meltwfunction_opreduce_vecs_idx kernel) {
|
||||
const DType *in_matrix1 = B;
|
||||
const DType *in_matrix2 = E;
|
||||
DType *output = C;
|
||||
IdType *out_matrix1 = argB;
|
||||
IdType *out_matrix2 = argE;
|
||||
|
||||
#pragma omp parallel
|
||||
{
|
||||
for (IdType k = 0; k < num_K_blocks; k++) {
|
||||
#pragma omp for schedule(dynamic)
|
||||
for (IdType m = 0; m < num_M_blocks; m++) {
|
||||
CSRMatrixInternal<IdType, IdType> cur_csr =
|
||||
block_csr_array[m * num_K_blocks + k];
|
||||
|
||||
const IdType M_start = m * M_block_size;
|
||||
for (IdType i = 0; i < cur_csr.num_rows; i++) {
|
||||
const IdType row_start = cur_csr.indptr[i];
|
||||
const IdType row_end = cur_csr.indptr[i + 1];
|
||||
const IdType dst = i + M_start;
|
||||
|
||||
libxsmm_meltw_opreduce_vecs_idx_param params;
|
||||
params.n = row_end - row_start;
|
||||
params.indices = &cur_csr.indices[row_start];
|
||||
params.in_matrix = in_matrix1;
|
||||
params.out_vec = &output[dst * N];
|
||||
params.argop_off_vec_0 = &out_matrix1[dst * N];
|
||||
params.argop_off_vec_1 = &out_matrix2[dst * N];
|
||||
params.scale_vals = nullptr;
|
||||
if (has_idx) {
|
||||
params.in_matrix2 = in_matrix2;
|
||||
params.indices2 = &cur_csr.data[row_start];
|
||||
} else {
|
||||
params.in_matrix2 = &in_matrix2[row_start * N];
|
||||
}
|
||||
kernel(¶ms);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Free the tiled CSR matrix data.
|
||||
* @param block_csr_array The array containing csr matrices of all blocks.
|
||||
* @param num_M_blocks Number of blocks to create along the rows of adjacency
|
||||
* matrix.
|
||||
* @param num_K_blocks Number of blocks to create along the columns of adjacency
|
||||
* matrix.
|
||||
* @param use_lhs Whether to use lhs.
|
||||
* @param use_rhs Whether to use rhs.
|
||||
*/
|
||||
template <typename IdType>
|
||||
inline void SpMMFreeBlocks(
|
||||
CSRMatrixInternal<IdType, IdType> *block_csr_array, IdType num_M_blocks,
|
||||
IdType num_K_blocks, bool use_lhs, bool use_rhs) {
|
||||
if (num_K_blocks > 1) {
|
||||
free(block_csr_array[0].indptr);
|
||||
if (use_lhs) free(block_csr_array[0].indices);
|
||||
if (use_rhs) free(block_csr_array[0].data);
|
||||
}
|
||||
free(block_csr_array);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Optimized CPU kernel of SpMM-Sum/Max/Min on Csr format.
|
||||
* @param bcast Broadcast information.
|
||||
* @param csr The Csr matrix.
|
||||
* @param ufeat The feature on source nodes.
|
||||
* @param efeat The feature on edges.
|
||||
* @param out The result feature on destination nodes.
|
||||
* @param argu Arg-Min/Max on source nodes.
|
||||
* @param arge Arg-Min/Max on edges.
|
||||
* @note it uses libxsmm, blocking and dynamic thread scheduling.
|
||||
*/
|
||||
template <typename IdType, typename DType, typename Op, typename Redop>
|
||||
void SpMMRedopCsrOpt(
|
||||
const BcastOff &bcast, const CSRMatrix &csr, NDArray ufeat, NDArray efeat,
|
||||
NDArray out, NDArray argu, NDArray arge) {
|
||||
int32_t llc_size = GetLLCSize();
|
||||
|
||||
#ifdef DEBUG
|
||||
uint64_t startTick, endTick;
|
||||
startTick = __rdtsc();
|
||||
#endif // DEBUG
|
||||
|
||||
const bool has_idx = !IsNullArray(csr.data);
|
||||
|
||||
DType *C = out.Ptr<DType>();
|
||||
const DType *B = ufeat.Ptr<DType>();
|
||||
const DType *E = efeat.Ptr<DType>();
|
||||
IdType *argB, *argE;
|
||||
if (std::is_same<Redop, op::Max<DType>>::value ||
|
||||
std::is_same<Redop, op::Min<DType>>::value) {
|
||||
argB = argu.Ptr<IdType>();
|
||||
argE = arge.Ptr<IdType>();
|
||||
}
|
||||
|
||||
const int nthreads = omp_get_max_threads();
|
||||
const IdType M = csr.num_rows;
|
||||
const IdType N = bcast.out_len;
|
||||
const IdType K = csr.num_cols;
|
||||
const IdType *indptr = csr.indptr.Ptr<IdType>();
|
||||
CHECK_NOTNULL(indptr);
|
||||
const IdType total_nnz = indptr[M];
|
||||
if (M <= 0 || K <= 0 || N <= 0 || total_nnz <= 0) return;
|
||||
|
||||
const double avg_degree = total_nnz * 1.0 / M;
|
||||
const double nnz_prob = avg_degree / K;
|
||||
|
||||
IdType K_block_size = std::min(
|
||||
(int64_t)K,
|
||||
(int64_t)(llc_size / (N * sizeof(DType) * nnz_prob * BLOCKING_HEURISTIC_PARAM))); // NOLINT
|
||||
IdType M_block_size = M / (nthreads * NUM_BLOCKS_PER_THREAD);
|
||||
if (M_block_size == 0) M_block_size = 1;
|
||||
if (K_block_size == 0) K_block_size = 1;
|
||||
|
||||
IdType num_M_blocks = (M + M_block_size - 1) / M_block_size;
|
||||
IdType num_K_blocks = (K + K_block_size - 1) / K_block_size;
|
||||
|
||||
CSRMatrixInternal<IdType, IdType> *block_csr_array =
|
||||
(CSRMatrixInternal<IdType, IdType> *)aligned_alloc(
|
||||
64, sizeof(CSRMatrixInternal<IdType, IdType>) * num_M_blocks *
|
||||
num_K_blocks);
|
||||
|
||||
#ifdef DEBUG
|
||||
endTick = __rdtsc();
|
||||
if (std::is_same<Redop, op::Max<DType>>::value) {
|
||||
LOG(INFO) << "Redop = Max";
|
||||
} else if (std::is_same<Redop, op::Min<DType>>::value) {
|
||||
LOG(INFO) << "Redop = Min";
|
||||
} else if (std::is_same<Redop, op::Add<DType>>::value) {
|
||||
LOG(INFO) << "Redop = Add";
|
||||
}
|
||||
LOG(INFO) << "nthreads = " << nthreads << ", llc_size = " << llc_size;
|
||||
LOG(INFO) << "M = " << M << ", K = " << K << ", N = " << N;
|
||||
LOG(INFO) << "use_lhs = " << Op::use_lhs << ", use_rhs = " << Op::use_rhs;
|
||||
LOG(INFO) << "total_nnz = " << total_nnz << ", avg_degree = " << avg_degree;
|
||||
LOG(INFO) << "has_idx = " << has_idx;
|
||||
LOG(INFO) << "nnz_prob = " << nnz_prob;
|
||||
LOG(INFO) << "K_block_size = " << K_block_size
|
||||
<< ", M_block_size = " << M_block_size;
|
||||
LOG(INFO) << "num_K_blocks = " << num_K_blocks
|
||||
<< ", num_M_blocks = " << num_M_blocks;
|
||||
LOG(INFO) << "stage0 ticks = " << (endTick - startTick);
|
||||
startTick = __rdtsc();
|
||||
#endif // DEBUG
|
||||
|
||||
SpMMCreateBlocks(
|
||||
csr, block_csr_array, num_M_blocks, num_K_blocks, M_block_size,
|
||||
K_block_size, Op::use_lhs, Op::use_rhs);
|
||||
|
||||
#ifdef DEBUG
|
||||
endTick = __rdtsc();
|
||||
LOG(INFO) << "stage1 ticks = " << (endTick - startTick);
|
||||
startTick = __rdtsc();
|
||||
#endif // DEBUG
|
||||
|
||||
libxsmm_meltwfunction_opreduce_vecs_idx kernel = nullptr;
|
||||
if (std::is_same<Redop, op::Max<DType>>::value) {
|
||||
kernel = SpMMCreateLibxsmmKernel<IdType, DType, Op>(
|
||||
has_idx, N, LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_REDOP_MAX, true);
|
||||
} else if (std::is_same<Redop, op::Min<DType>>::value) {
|
||||
kernel = SpMMCreateLibxsmmKernel<IdType, DType, Op>(
|
||||
has_idx, N, LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_REDOP_MIN, true);
|
||||
} else if (std::is_same<Redop, op::Add<DType>>::value) {
|
||||
kernel = SpMMCreateLibxsmmKernel<IdType, DType, Op>(
|
||||
has_idx, N, LIBXSMM_MELTW_FLAG_OPREDUCE_VECS_REDOP_SUM, false);
|
||||
}
|
||||
|
||||
#ifdef DEBUG
|
||||
endTick = __rdtsc();
|
||||
LOG(INFO) << "stage2 ticks = " << (endTick - startTick);
|
||||
startTick = __rdtsc();
|
||||
#endif // DEBUG
|
||||
|
||||
if (std::is_same<Redop, op::Max<DType>>::value ||
|
||||
std::is_same<Redop, op::Min<DType>>::value) {
|
||||
SpMMBlockwiseOpCmp<IdType, DType, Op, Redop>(
|
||||
block_csr_array, B, E, C, argB, argE, has_idx, N, num_M_blocks,
|
||||
num_K_blocks, M_block_size, kernel);
|
||||
} else {
|
||||
SpMMBlockwiseOpSum(
|
||||
block_csr_array, B, E, C, has_idx, N, num_M_blocks, num_K_blocks,
|
||||
M_block_size, kernel);
|
||||
}
|
||||
|
||||
#ifdef DEBUG
|
||||
endTick = __rdtsc();
|
||||
LOG(INFO) << "stage3 ticks = " << (endTick - startTick);
|
||||
startTick = __rdtsc();
|
||||
#endif // DEBUG
|
||||
|
||||
SpMMFreeBlocks(
|
||||
block_csr_array, num_M_blocks, num_K_blocks, Op::use_lhs, Op::use_rhs);
|
||||
|
||||
#ifdef DEBUG
|
||||
endTick = __rdtsc();
|
||||
LOG(INFO) << "stage4 ticks = " << (endTick - startTick);
|
||||
#endif // DEBUG
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Optimized CPU kernel of SpMM-Sum on Csr format.
|
||||
* @param bcast Broadcast information.
|
||||
* @param csr The Csr matrix.
|
||||
* @param ufeat The feature on source nodes.
|
||||
* @param efeat The feature on edges.
|
||||
* @param out The result feature on destination nodes.
|
||||
* @note it uses libxsmm, blocking and dynamic thread scheduling.
|
||||
*/
|
||||
template <typename IdType, typename DType, typename Op>
|
||||
void SpMMSumCsrLibxsmm(
|
||||
const BcastOff &bcast, const CSRMatrix &csr, NDArray ufeat, NDArray efeat,
|
||||
NDArray out) {
|
||||
NDArray dummy;
|
||||
SpMMRedopCsrOpt<IdType, DType, Op, op::Add<DType>>(
|
||||
bcast, csr, ufeat, efeat, out, dummy, dummy);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Optimized CPU kernel of SpMM-Min/Max on Csr format.
|
||||
* @param bcast Broadcast information.
|
||||
* @param csr The Csr matrix.
|
||||
* @param ufeat The feature on source nodes.
|
||||
* @param efeat The feature on edges.
|
||||
* @param out The result feature on destination nodes.
|
||||
* @param argu Arg-Min/Max on source nodes.
|
||||
* @param arge Arg-Min/Max on edges.
|
||||
* @note it uses libxsmm, blocking and dynamic thread scheduling.
|
||||
*/
|
||||
template <typename IdType, typename DType, typename Op, typename Cmp>
|
||||
void SpMMCmpCsrLibxsmm(
|
||||
const BcastOff &bcast, const CSRMatrix &csr, NDArray ufeat, NDArray efeat,
|
||||
NDArray out, NDArray argu, NDArray arge) {
|
||||
SpMMRedopCsrOpt<IdType, DType, Op, Cmp>(
|
||||
bcast, csr, ufeat, efeat, out, argu, arge);
|
||||
}
|
||||
|
||||
} // namespace cpu
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // USE_LIBXSMM
|
||||
#endif // _WIN32
|
||||
|
||||
#endif // DGL_ARRAY_CPU_SPMM_BLOCKING_LIBXSMM_H_
|
||||
@@ -0,0 +1,231 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/traversal.cc
|
||||
* @brief Graph traversal implementation
|
||||
*/
|
||||
|
||||
#include "./traversal.h"
|
||||
|
||||
#include <dgl/graph_traversal.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <queue>
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
namespace {
|
||||
// A utility view class to wrap a vector into a queue.
|
||||
template <typename DType>
|
||||
struct VectorQueueWrapper {
|
||||
std::vector<DType>* vec;
|
||||
size_t head = 0;
|
||||
|
||||
explicit VectorQueueWrapper(std::vector<DType>* vec) : vec(vec) {}
|
||||
|
||||
void push(const DType& elem) { vec->push_back(elem); }
|
||||
|
||||
DType top() const { return vec->operator[](head); }
|
||||
|
||||
void pop() { ++head; }
|
||||
|
||||
bool empty() const { return head == vec->size(); }
|
||||
|
||||
size_t size() const { return vec->size() - head; }
|
||||
};
|
||||
|
||||
// Internal function to merge multiple traversal traces into one ndarray.
|
||||
// It is similar to zip the vectors together.
|
||||
template <typename DType>
|
||||
IdArray MergeMultipleTraversals(const std::vector<std::vector<DType>>& traces) {
|
||||
int64_t max_len = 0, total_len = 0;
|
||||
for (size_t i = 0; i < traces.size(); ++i) {
|
||||
const int64_t tracelen = traces[i].size();
|
||||
max_len = std::max(max_len, tracelen);
|
||||
total_len += traces[i].size();
|
||||
}
|
||||
IdArray ret = IdArray::Empty(
|
||||
{total_len}, DGLDataType{kDGLInt, sizeof(DType) * 8, 1},
|
||||
DGLContext{kDGLCPU, 0});
|
||||
DType* ret_data = static_cast<DType*>(ret->data);
|
||||
for (int64_t i = 0; i < max_len; ++i) {
|
||||
for (size_t j = 0; j < traces.size(); ++j) {
|
||||
const int64_t tracelen = traces[j].size();
|
||||
if (i >= tracelen) {
|
||||
continue;
|
||||
}
|
||||
*(ret_data++) = traces[j][i];
|
||||
}
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
// Internal function to compute sections if multiple traversal traces
|
||||
// are merged into one ndarray.
|
||||
template <typename DType>
|
||||
IdArray ComputeMergedSections(const std::vector<std::vector<DType>>& traces) {
|
||||
int64_t max_len = 0;
|
||||
for (size_t i = 0; i < traces.size(); ++i) {
|
||||
const int64_t tracelen = traces[i].size();
|
||||
max_len = std::max(max_len, tracelen);
|
||||
}
|
||||
IdArray ret = IdArray::Empty(
|
||||
{max_len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
|
||||
int64_t* ret_data = static_cast<int64_t*>(ret->data);
|
||||
for (int64_t i = 0; i < max_len; ++i) {
|
||||
int64_t sec_len = 0;
|
||||
for (size_t j = 0; j < traces.size(); ++j) {
|
||||
const int64_t tracelen = traces[j].size();
|
||||
if (i < tracelen) {
|
||||
++sec_len;
|
||||
}
|
||||
}
|
||||
*(ret_data++) = sec_len;
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
Frontiers BFSNodesFrontiers(const CSRMatrix& csr, IdArray source) {
|
||||
std::vector<IdType> ids;
|
||||
std::vector<int64_t> sections;
|
||||
VectorQueueWrapper<IdType> queue(&ids);
|
||||
auto visit = [&](const int64_t v) {};
|
||||
auto make_frontier = [&]() {
|
||||
if (!queue.empty()) {
|
||||
// do not push zero-length frontier
|
||||
sections.push_back(queue.size());
|
||||
}
|
||||
};
|
||||
BFSTraverseNodes<IdType>(csr, source, &queue, visit, make_frontier);
|
||||
|
||||
Frontiers front;
|
||||
front.ids = VecToIdArray(ids, sizeof(IdType) * 8);
|
||||
front.sections = VecToIdArray(sections, sizeof(int64_t) * 8);
|
||||
return front;
|
||||
}
|
||||
|
||||
template Frontiers BFSNodesFrontiers<kDGLCPU, int32_t>(
|
||||
const CSRMatrix&, IdArray);
|
||||
template Frontiers BFSNodesFrontiers<kDGLCPU, int64_t>(
|
||||
const CSRMatrix&, IdArray);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
Frontiers BFSEdgesFrontiers(const CSRMatrix& csr, IdArray source) {
|
||||
std::vector<IdType> ids;
|
||||
std::vector<int64_t> sections;
|
||||
// NOTE: std::queue has no top() method.
|
||||
std::vector<IdType> nodes;
|
||||
VectorQueueWrapper<IdType> queue(&nodes);
|
||||
auto visit = [&](const IdType e) { ids.push_back(e); };
|
||||
bool first_frontier = true;
|
||||
auto make_frontier = [&] {
|
||||
if (first_frontier) {
|
||||
first_frontier = false; // do not push the first section when doing edges
|
||||
} else if (!queue.empty()) {
|
||||
// do not push zero-length frontier
|
||||
sections.push_back(queue.size());
|
||||
}
|
||||
};
|
||||
BFSTraverseEdges<IdType>(csr, source, &queue, visit, make_frontier);
|
||||
|
||||
Frontiers front;
|
||||
front.ids = VecToIdArray(ids, sizeof(IdType) * 8);
|
||||
front.sections = VecToIdArray(sections, sizeof(int64_t) * 8);
|
||||
return front;
|
||||
}
|
||||
|
||||
template Frontiers BFSEdgesFrontiers<kDGLCPU, int32_t>(
|
||||
const CSRMatrix&, IdArray);
|
||||
template Frontiers BFSEdgesFrontiers<kDGLCPU, int64_t>(
|
||||
const CSRMatrix&, IdArray);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
Frontiers TopologicalNodesFrontiers(const CSRMatrix& csr) {
|
||||
std::vector<IdType> ids;
|
||||
std::vector<int64_t> sections;
|
||||
VectorQueueWrapper<IdType> queue(&ids);
|
||||
auto visit = [&](const uint64_t v) {};
|
||||
auto make_frontier = [&]() {
|
||||
if (!queue.empty()) {
|
||||
// do not push zero-length frontier
|
||||
sections.push_back(queue.size());
|
||||
}
|
||||
};
|
||||
TopologicalNodes<IdType>(csr, &queue, visit, make_frontier);
|
||||
|
||||
Frontiers front;
|
||||
front.ids = VecToIdArray(ids, sizeof(IdType) * 8);
|
||||
front.sections = VecToIdArray(sections, sizeof(int64_t) * 8);
|
||||
return front;
|
||||
}
|
||||
|
||||
template Frontiers TopologicalNodesFrontiers<kDGLCPU, int32_t>(
|
||||
const CSRMatrix&);
|
||||
template Frontiers TopologicalNodesFrontiers<kDGLCPU, int64_t>(
|
||||
const CSRMatrix&);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
Frontiers DGLDFSEdges(const CSRMatrix& csr, IdArray source) {
|
||||
const int64_t len = source->shape[0];
|
||||
const IdType* src_data = static_cast<IdType*>(source->data);
|
||||
std::vector<std::vector<IdType>> edges(len);
|
||||
|
||||
for (int64_t i = 0; i < len; ++i) {
|
||||
auto visit = [&](IdType e, int tag) { edges[i].push_back(e); };
|
||||
DFSLabeledEdges<IdType>(csr, src_data[i], false, false, visit);
|
||||
}
|
||||
|
||||
Frontiers front;
|
||||
front.ids = MergeMultipleTraversals(edges);
|
||||
front.sections = ComputeMergedSections(edges);
|
||||
return front;
|
||||
}
|
||||
|
||||
template Frontiers DGLDFSEdges<kDGLCPU, int32_t>(const CSRMatrix&, IdArray);
|
||||
template Frontiers DGLDFSEdges<kDGLCPU, int64_t>(const CSRMatrix&, IdArray);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
Frontiers DGLDFSLabeledEdges(
|
||||
const CSRMatrix& csr, IdArray source, const bool has_reverse_edge,
|
||||
const bool has_nontree_edge, const bool return_labels) {
|
||||
const int64_t len = source->shape[0];
|
||||
const IdType* src_data = static_cast<IdType*>(source->data);
|
||||
std::vector<std::vector<IdType>> edges(len);
|
||||
std::vector<std::vector<int64_t>> tags;
|
||||
|
||||
if (return_labels) {
|
||||
tags.resize(len);
|
||||
}
|
||||
|
||||
for (int64_t i = 0; i < len; ++i) {
|
||||
auto visit = [&](IdType e, int64_t tag) {
|
||||
edges[i].push_back(e);
|
||||
if (return_labels) {
|
||||
tags[i].push_back(tag);
|
||||
}
|
||||
};
|
||||
DFSLabeledEdges<IdType>(
|
||||
csr, src_data[i], has_reverse_edge, has_nontree_edge, visit);
|
||||
}
|
||||
|
||||
Frontiers front;
|
||||
front.ids = MergeMultipleTraversals(edges);
|
||||
front.sections = ComputeMergedSections(edges);
|
||||
if (return_labels) {
|
||||
front.tags = MergeMultipleTraversals(tags);
|
||||
}
|
||||
|
||||
return front;
|
||||
}
|
||||
|
||||
template Frontiers DGLDFSLabeledEdges<kDGLCPU, int32_t>(
|
||||
const CSRMatrix&, IdArray, const bool, const bool, const bool);
|
||||
template Frontiers DGLDFSLabeledEdges<kDGLCPU, int64_t>(
|
||||
const CSRMatrix&, IdArray, const bool, const bool, const bool);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,310 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/traversal.h
|
||||
* @brief Graph traversal routines.
|
||||
*
|
||||
* Traversal routines generate frontiers. Frontiers can be node frontiers or
|
||||
* edge frontiers depending on the traversal function. Each frontier is a list
|
||||
* of nodes/edges (specified by their ids). An optional tag can be specified for
|
||||
* each node/edge (represented by an int value).
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CPU_TRAVERSAL_H_
|
||||
#define DGL_ARRAY_CPU_TRAVERSAL_H_
|
||||
|
||||
#include <dgl/graph_interface.h>
|
||||
|
||||
#include <stack>
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
/**
|
||||
* @brief Traverse the graph in a breadth-first-search (BFS) order.
|
||||
*
|
||||
* The queue object must suffice following interface:
|
||||
* Members:
|
||||
* void push(IdType); // push one node
|
||||
* IdType top(); // get the first node
|
||||
* void pop(); // pop one node
|
||||
* bool empty(); // return true if the queue is empty
|
||||
* size_t size(); // return the size of the queue
|
||||
* For example, std::queue<IdType> is a valid queue type.
|
||||
*
|
||||
* The visit function must be compatible with following interface:
|
||||
* void (*visit)(IdType );
|
||||
*
|
||||
* The frontier function must be compatible with following interface:
|
||||
* void (*make_frontier)(void);
|
||||
*
|
||||
* @param graph The graph.
|
||||
* @param sources Source nodes.
|
||||
* @param reversed If true, BFS follows the in-edge direction
|
||||
* @param queue The queue used to do bfs.
|
||||
* @param visit The function to call when a node is visited.
|
||||
* @param make_frontier The function to indicate that a new froniter can be
|
||||
* made;
|
||||
*/
|
||||
template <
|
||||
typename IdType, typename Queue, typename VisitFn, typename FrontierFn>
|
||||
void BFSTraverseNodes(
|
||||
const CSRMatrix &csr, IdArray source, Queue *queue, VisitFn visit,
|
||||
FrontierFn make_frontier) {
|
||||
const int64_t len = source->shape[0];
|
||||
const IdType *src_data = static_cast<IdType *>(source->data);
|
||||
|
||||
const IdType *indptr_data = static_cast<IdType *>(csr.indptr->data);
|
||||
const IdType *indices_data = static_cast<IdType *>(csr.indices->data);
|
||||
const int64_t num_nodes = csr.num_rows;
|
||||
std::vector<bool> visited(num_nodes);
|
||||
for (int64_t i = 0; i < len; ++i) {
|
||||
const IdType u = src_data[i];
|
||||
visited[u] = true;
|
||||
visit(u);
|
||||
queue->push(u);
|
||||
}
|
||||
make_frontier();
|
||||
|
||||
while (!queue->empty()) {
|
||||
const size_t size = queue->size();
|
||||
for (size_t i = 0; i < size; ++i) {
|
||||
const IdType u = queue->top();
|
||||
queue->pop();
|
||||
for (auto idx = indptr_data[u]; idx < indptr_data[u + 1]; ++idx) {
|
||||
auto v = indices_data[idx];
|
||||
if (!visited[v]) {
|
||||
visited[v] = true;
|
||||
visit(v);
|
||||
queue->push(v);
|
||||
}
|
||||
}
|
||||
}
|
||||
make_frontier();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Traverse the graph in a breadth-first-search (BFS) order, returning
|
||||
* the edges of the BFS tree.
|
||||
*
|
||||
* The queue object must suffice following interface:
|
||||
* Members:
|
||||
* void push(IdType); // push one node
|
||||
* IdType top(); // get the first node
|
||||
* void pop(); // pop one node
|
||||
* bool empty(); // return true if the queue is empty
|
||||
* size_t size(); // return the size of the queue
|
||||
* For example, std::queue<IdType> is a valid queue type.
|
||||
*
|
||||
* The visit function must be compatible with following interface:
|
||||
* void (*visit)(IdType );
|
||||
*
|
||||
* The frontier function must be compatible with following interface:
|
||||
* void (*make_frontier)(void);
|
||||
*
|
||||
* @param graph The graph.
|
||||
* @param sources Source nodes.
|
||||
* @param reversed If true, BFS follows the in-edge direction
|
||||
* @param queue The queue used to do bfs.
|
||||
* @param visit The function to call when a node is visited.
|
||||
* The argument would be edge ID.
|
||||
* @param make_frontier The function to indicate that a new frontier can be
|
||||
* made;
|
||||
*/
|
||||
template <
|
||||
typename IdType, typename Queue, typename VisitFn, typename FrontierFn>
|
||||
void BFSTraverseEdges(
|
||||
const CSRMatrix &csr, IdArray source, Queue *queue, VisitFn visit,
|
||||
FrontierFn make_frontier) {
|
||||
const int64_t len = source->shape[0];
|
||||
const IdType *src_data = static_cast<IdType *>(source->data);
|
||||
|
||||
const IdType *indptr_data = static_cast<IdType *>(csr.indptr->data);
|
||||
const IdType *indices_data = static_cast<IdType *>(csr.indices->data);
|
||||
const IdType *eid_data = static_cast<IdType *>(csr.data->data);
|
||||
|
||||
const int64_t num_nodes = csr.num_rows;
|
||||
std::vector<bool> visited(num_nodes);
|
||||
for (int64_t i = 0; i < len; ++i) {
|
||||
const IdType u = src_data[i];
|
||||
visited[u] = true;
|
||||
queue->push(u);
|
||||
}
|
||||
make_frontier();
|
||||
|
||||
while (!queue->empty()) {
|
||||
const size_t size = queue->size();
|
||||
for (size_t i = 0; i < size; ++i) {
|
||||
const IdType u = queue->top();
|
||||
queue->pop();
|
||||
for (auto idx = indptr_data[u]; idx < indptr_data[u + 1]; ++idx) {
|
||||
auto e = eid_data ? eid_data[idx] : idx;
|
||||
const IdType v = indices_data[idx];
|
||||
if (!visited[v]) {
|
||||
visited[v] = true;
|
||||
visit(e);
|
||||
queue->push(v);
|
||||
}
|
||||
}
|
||||
}
|
||||
make_frontier();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Traverse the graph in topological order.
|
||||
*
|
||||
* The queue object must suffice following interface:
|
||||
* Members:
|
||||
* void push(IdType); // push one node
|
||||
* IdType top(); // get the first node
|
||||
* void pop(); // pop one node
|
||||
* bool empty(); // return true if the queue is empty
|
||||
* size_t size(); // return the size of the queue
|
||||
* For example, std::queue<IdType> is a valid queue type.
|
||||
*
|
||||
* The visit function must be compatible with following interface:
|
||||
* void (*visit)(IdType );
|
||||
*
|
||||
* The frontier function must be compatible with following interface:
|
||||
* void (*make_frontier)(void);
|
||||
*
|
||||
* @param graph The graph.
|
||||
* @param reversed If true, follows the in-edge direction
|
||||
* @param queue The queue used to do bfs.
|
||||
* @param visit The function to call when a node is visited.
|
||||
* @param make_frontier The function to indicate that a new froniter can be
|
||||
* made;
|
||||
*/
|
||||
template <
|
||||
typename IdType, typename Queue, typename VisitFn, typename FrontierFn>
|
||||
void TopologicalNodes(
|
||||
const CSRMatrix &csr, Queue *queue, VisitFn visit,
|
||||
FrontierFn make_frontier) {
|
||||
int64_t num_visited_nodes = 0;
|
||||
const IdType *indptr_data = static_cast<IdType *>(csr.indptr->data);
|
||||
const IdType *indices_data = static_cast<IdType *>(csr.indices->data);
|
||||
|
||||
const int64_t num_nodes = csr.num_rows;
|
||||
const int64_t num_edges = csr.indices->shape[0];
|
||||
std::vector<int64_t> degrees(num_nodes, 0);
|
||||
for (int64_t eid = 0; eid < num_edges; ++eid) {
|
||||
degrees[indices_data[eid]]++;
|
||||
}
|
||||
|
||||
for (int64_t vid = 0; vid < num_nodes; ++vid) {
|
||||
if (degrees[vid] == 0) {
|
||||
visit(vid);
|
||||
queue->push(static_cast<IdType>(vid));
|
||||
++num_visited_nodes;
|
||||
}
|
||||
}
|
||||
make_frontier();
|
||||
|
||||
while (!queue->empty()) {
|
||||
const size_t size = queue->size();
|
||||
for (size_t i = 0; i < size; ++i) {
|
||||
const IdType u = queue->top();
|
||||
queue->pop();
|
||||
for (auto idx = indptr_data[u]; idx < indptr_data[u + 1]; ++idx) {
|
||||
const IdType v = indices_data[idx];
|
||||
if (--(degrees[v]) == 0) {
|
||||
visit(v);
|
||||
queue->push(v);
|
||||
++num_visited_nodes;
|
||||
}
|
||||
}
|
||||
}
|
||||
make_frontier();
|
||||
}
|
||||
|
||||
if (num_visited_nodes != num_nodes) {
|
||||
LOG(FATAL)
|
||||
<< "Error in topological traversal: loop detected in the given graph.";
|
||||
}
|
||||
}
|
||||
|
||||
/** @brief Tags for ``DFSEdges``. */
|
||||
enum DFSEdgeTag {
|
||||
kForward = 0,
|
||||
kReverse,
|
||||
kNonTree,
|
||||
};
|
||||
/**
|
||||
* @brief Traverse the graph in a depth-first-search (DFS) order.
|
||||
*
|
||||
* The traversal visit edges in its DFS order. Edges have three tags:
|
||||
* FORWARD(0), REVERSE(1), NONTREE(2)
|
||||
*
|
||||
* A FORWARD edge is one in which `u` has been visisted but `v` has not.
|
||||
* A REVERSE edge is one in which both `u` and `v` have been visisted and the
|
||||
* edge is in the DFS tree. A NONTREE edge is one in which both `u` and `v` have
|
||||
* been visisted but the edge is NOT in the DFS tree.
|
||||
*
|
||||
* @param source Source node.
|
||||
* @param reversed If true, DFS follows the in-edge direction
|
||||
* @param has_reverse_edge If true, REVERSE edges are included
|
||||
* @param has_nontree_edge If true, NONTREE edges are included
|
||||
* @param visit The function to call when an edge is visited; the edge id and
|
||||
* its tag will be given as the arguments.
|
||||
*/
|
||||
template <typename IdType, typename VisitFn>
|
||||
void DFSLabeledEdges(
|
||||
const CSRMatrix &csr, IdType source, bool has_reverse_edge,
|
||||
bool has_nontree_edge, VisitFn visit) {
|
||||
const int64_t num_nodes = csr.num_rows;
|
||||
CHECK_GE(num_nodes, source)
|
||||
<< "source " << source << " is out of range [0," << num_nodes << "]";
|
||||
const IdType *indptr_data = static_cast<IdType *>(csr.indptr->data);
|
||||
const IdType *indices_data = static_cast<IdType *>(csr.indices->data);
|
||||
const IdType *eid_data = static_cast<IdType *>(csr.data->data);
|
||||
|
||||
if (indptr_data[source + 1] - indptr_data[source] == 0) {
|
||||
// no out-going edges from the source node
|
||||
return;
|
||||
}
|
||||
|
||||
typedef std::tuple<IdType, size_t, bool> StackEntry;
|
||||
std::stack<StackEntry> stack;
|
||||
std::vector<bool> visited(num_nodes);
|
||||
visited[source] = true;
|
||||
stack.push(std::make_tuple(source, 0, false));
|
||||
IdType u = 0;
|
||||
int64_t i = 0;
|
||||
bool on_tree = false;
|
||||
|
||||
while (!stack.empty()) {
|
||||
std::tie(u, i, on_tree) = stack.top();
|
||||
const IdType v = indices_data[indptr_data[u] + i];
|
||||
const IdType uv =
|
||||
eid_data ? eid_data[indptr_data[u] + i] : indptr_data[u] + i;
|
||||
if (visited[v]) {
|
||||
if (!on_tree && has_nontree_edge) {
|
||||
visit(uv, kNonTree);
|
||||
} else if (on_tree && has_reverse_edge) {
|
||||
visit(uv, kReverse);
|
||||
}
|
||||
stack.pop();
|
||||
// find next one.
|
||||
if (indptr_data[u] + i < indptr_data[u + 1] - 1) {
|
||||
stack.push(std::make_tuple(u, i + 1, false));
|
||||
}
|
||||
} else {
|
||||
visited[v] = true;
|
||||
std::get<2>(stack.top()) = true;
|
||||
visit(uv, kForward);
|
||||
// expand
|
||||
if (indptr_data[v] < indptr_data[v + 1]) {
|
||||
stack.push(std::make_tuple(v, 0, false));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CPU_TRAVERSAL_H_
|
||||
@@ -0,0 +1,57 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/array_cumsum.cu
|
||||
* @brief Array cumsum GPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray CumSum(IdArray array, bool prepend_zero) {
|
||||
const int64_t len = array.NumElements();
|
||||
if (len == 0)
|
||||
return !prepend_zero ? array
|
||||
: aten::Full(0, 1, array->dtype.bits, array->ctx);
|
||||
|
||||
auto device = runtime::DeviceAPI::Get(array->ctx);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const IdType* in_d = array.Ptr<IdType>();
|
||||
IdArray ret;
|
||||
IdType* out_d = nullptr;
|
||||
if (prepend_zero) {
|
||||
ret = aten::Full(0, len + 1, array->dtype.bits, array->ctx);
|
||||
out_d = ret.Ptr<IdType>() + 1;
|
||||
} else {
|
||||
ret = aten::NewIdArray(len, array->ctx, array->dtype.bits);
|
||||
out_d = ret.Ptr<IdType>();
|
||||
}
|
||||
// Allocate workspace
|
||||
size_t workspace_size = 0;
|
||||
CUDA_CALL(cub::DeviceScan::InclusiveSum(
|
||||
nullptr, workspace_size, in_d, out_d, len, stream));
|
||||
void* workspace = device->AllocWorkspace(array->ctx, workspace_size);
|
||||
|
||||
// Compute cumsum
|
||||
CUDA_CALL(cub::DeviceScan::InclusiveSum(
|
||||
workspace, workspace_size, in_d, out_d, len, stream));
|
||||
|
||||
device->FreeWorkspace(array->ctx, workspace);
|
||||
|
||||
return ret;
|
||||
}
|
||||
|
||||
template IdArray CumSum<kDGLCUDA, int32_t>(IdArray, bool);
|
||||
template IdArray CumSum<kDGLCUDA, int64_t>(IdArray, bool);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,98 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file array/cpu/array_index_select.cu
|
||||
* @brief Array index select GPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./array_index_select.cuh"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename DType, typename IdType>
|
||||
NDArray IndexSelect(NDArray array, IdArray index) {
|
||||
const int64_t arr_len = array->shape[0];
|
||||
const int64_t len = index->shape[0];
|
||||
int64_t num_feat = 1;
|
||||
std::vector<int64_t> shape{len};
|
||||
for (int d = 1; d < array->ndim; ++d) {
|
||||
num_feat *= array->shape[d];
|
||||
shape.emplace_back(array->shape[d]);
|
||||
}
|
||||
|
||||
// use index->ctx for pinned array
|
||||
NDArray ret = NDArray::Empty(shape, array->dtype, index->ctx);
|
||||
if (len == 0 || arr_len * num_feat == 0) return ret;
|
||||
DType* ret_data = static_cast<DType*>(ret->data);
|
||||
|
||||
const DType* array_data = static_cast<DType*>(cuda::GetDevicePointer(array));
|
||||
const IdType* idx_data = static_cast<IdType*>(index->data);
|
||||
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
if (num_feat == 1) {
|
||||
const int nt = cuda::FindNumThreads(len);
|
||||
const int nb = (len + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
IndexSelectSingleKernel, nb, nt, 0, stream, array_data, idx_data, len,
|
||||
arr_len, ret_data);
|
||||
} else {
|
||||
dim3 block(256, 1);
|
||||
while (static_cast<int64_t>(block.x) >= 2 * num_feat) {
|
||||
block.x /= 2;
|
||||
block.y *= 2;
|
||||
}
|
||||
const dim3 grid((len + block.y - 1) / block.y);
|
||||
CUDA_KERNEL_CALL(
|
||||
IndexSelectMultiKernel, grid, block, 0, stream, array_data, num_feat,
|
||||
idx_data, len, arr_len, ret_data);
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
template NDArray IndexSelect<kDGLCUDA, int32_t, int32_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelect<kDGLCUDA, int32_t, int64_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelect<kDGLCUDA, int64_t, int32_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelect<kDGLCUDA, int64_t, int64_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelect<kDGLCUDA, __half, int32_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelect<kDGLCUDA, __half, int64_t>(NDArray, IdArray);
|
||||
#if BF16_ENABLED
|
||||
template NDArray IndexSelect<kDGLCUDA, __nv_bfloat16, int32_t>(
|
||||
NDArray, IdArray);
|
||||
template NDArray IndexSelect<kDGLCUDA, __nv_bfloat16, int64_t>(
|
||||
NDArray, IdArray);
|
||||
#endif // BF16_ENABLED
|
||||
template NDArray IndexSelect<kDGLCUDA, float, int32_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelect<kDGLCUDA, float, int64_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelect<kDGLCUDA, double, int32_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelect<kDGLCUDA, double, int64_t>(NDArray, IdArray);
|
||||
|
||||
template <DGLDeviceType XPU, typename DType>
|
||||
DType IndexSelect(NDArray array, int64_t index) {
|
||||
auto device = runtime::DeviceAPI::Get(array->ctx);
|
||||
DType ret = static_cast<DType>(0.0f);
|
||||
device->CopyDataFromTo(
|
||||
static_cast<DType*>(array->data) + index, 0, &ret, 0, sizeof(DType),
|
||||
array->ctx, DGLContext{kDGLCPU, 0}, array->dtype);
|
||||
return ret;
|
||||
}
|
||||
|
||||
template int32_t IndexSelect<kDGLCUDA, int32_t>(NDArray array, int64_t index);
|
||||
template int64_t IndexSelect<kDGLCUDA, int64_t>(NDArray array, int64_t index);
|
||||
template uint32_t IndexSelect<kDGLCUDA, uint32_t>(NDArray array, int64_t index);
|
||||
template uint64_t IndexSelect<kDGLCUDA, uint64_t>(NDArray array, int64_t index);
|
||||
template __half IndexSelect<kDGLCUDA, __half>(NDArray array, int64_t index);
|
||||
#if BF16_ENABLED
|
||||
template __nv_bfloat16 IndexSelect<kDGLCUDA, __nv_bfloat16>(
|
||||
NDArray array, int64_t index);
|
||||
#endif // BF16_ENABLED
|
||||
template float IndexSelect<kDGLCUDA, float>(NDArray array, int64_t index);
|
||||
template double IndexSelect<kDGLCUDA, double>(NDArray array, int64_t index);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,87 @@
|
||||
/**
|
||||
* Copyright (c) 2021-2022 by Contributors
|
||||
* @file array/cuda/array_index_select.cuh
|
||||
* @brief Array index select GPU kernel implementation
|
||||
*/
|
||||
|
||||
#ifndef DGL_ARRAY_CUDA_ARRAY_INDEX_SELECT_CUH_
|
||||
#define DGL_ARRAY_CUDA_ARRAY_INDEX_SELECT_CUH_
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <typename DType, typename IdType>
|
||||
__global__ void IndexSelectSingleKernel(
|
||||
const DType* array, const IdType* index, const int64_t length,
|
||||
const int64_t arr_len, DType* out, const int64_t* perm = nullptr) {
|
||||
int64_t tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < length) {
|
||||
assert(index[tx] >= 0 && index[tx] < arr_len);
|
||||
const auto out_row = perm ? perm[tx] : tx;
|
||||
out[out_row] = array[index[tx]];
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename DType, typename IdType>
|
||||
__global__ void IndexSelectMultiKernel(
|
||||
const DType* const array, const int64_t num_feat, const IdType* const index,
|
||||
const int64_t length, const int64_t arr_len, DType* const out,
|
||||
const int64_t* perm = nullptr) {
|
||||
int64_t out_row_index = blockIdx.x * blockDim.y + threadIdx.y;
|
||||
|
||||
const int64_t stride = blockDim.y * gridDim.x;
|
||||
|
||||
while (out_row_index < length) {
|
||||
int64_t col = threadIdx.x;
|
||||
const int64_t in_row = index[out_row_index];
|
||||
assert(in_row >= 0 && in_row < arr_len);
|
||||
const auto out_row = perm ? perm[out_row_index] : out_row_index;
|
||||
while (col < num_feat) {
|
||||
out[out_row * num_feat + col] = array[in_row * num_feat + col];
|
||||
col += blockDim.x;
|
||||
}
|
||||
out_row_index += stride;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename DType, typename IdType>
|
||||
__global__ void IndexScatterSingleKernel(
|
||||
const DType* array, const IdType* index, const int64_t length,
|
||||
const int64_t arr_len, DType* out) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < length) {
|
||||
assert(index[tx] >= 0 && index[tx] < arr_len);
|
||||
out[index[tx]] = array[tx];
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename DType, typename IdType>
|
||||
__global__ void IndexScatterMultiKernel(
|
||||
const DType* const array, const int64_t num_feat, const IdType* const index,
|
||||
const int64_t length, const int64_t arr_len, DType* const out) {
|
||||
int64_t in_row = blockIdx.x * blockDim.y + threadIdx.y;
|
||||
|
||||
const int64_t stride = blockDim.y * gridDim.x;
|
||||
|
||||
while (in_row < length) {
|
||||
int64_t col = threadIdx.x;
|
||||
const int64_t out_row = index[in_row];
|
||||
assert(out_row >= 0 && out_row < arr_len);
|
||||
while (col < num_feat) {
|
||||
out[out_row * num_feat + col] = array[in_row * num_feat + col];
|
||||
col += blockDim.x;
|
||||
}
|
||||
in_row += stride;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CUDA_ARRAY_INDEX_SELECT_CUH_
|
||||
@@ -0,0 +1,71 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/array_nonzero.cc
|
||||
* @brief Array nonzero CPU implementation
|
||||
*/
|
||||
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <typename IdType>
|
||||
struct IsNonZeroIndex {
|
||||
explicit IsNonZeroIndex(const IdType* array) : array_(array) {}
|
||||
|
||||
__device__ bool operator()(const int64_t index) { return array_[index] != 0; }
|
||||
|
||||
const IdType* array_;
|
||||
};
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray NonZero(IdArray array) {
|
||||
const auto& ctx = array->ctx;
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
|
||||
const int64_t len = array->shape[0];
|
||||
IdArray ret = NewIdArray(len, ctx, 64);
|
||||
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
|
||||
const IdType* const in_data = static_cast<const IdType*>(array->data);
|
||||
int64_t* const out_data = static_cast<int64_t*>(ret->data);
|
||||
|
||||
IsNonZeroIndex<IdType> comp(in_data);
|
||||
cub::CountingInputIterator<int64_t> counter(0);
|
||||
|
||||
// room for cub to output on GPU
|
||||
int64_t* d_num_nonzeros =
|
||||
static_cast<int64_t*>(device->AllocWorkspace(ctx, sizeof(int64_t)));
|
||||
|
||||
size_t temp_size = 0;
|
||||
CUDA_CALL(cub::DeviceSelect::If(
|
||||
nullptr, temp_size, counter, out_data, d_num_nonzeros, len, comp,
|
||||
stream));
|
||||
void* temp = device->AllocWorkspace(ctx, temp_size);
|
||||
CUDA_CALL(cub::DeviceSelect::If(
|
||||
temp, temp_size, counter, out_data, d_num_nonzeros, len, comp, stream));
|
||||
device->FreeWorkspace(ctx, temp);
|
||||
|
||||
// copy number of selected elements from GPU to CPU
|
||||
int64_t num_nonzeros = cuda::GetCUDAScalar(device, ctx, d_num_nonzeros);
|
||||
device->FreeWorkspace(ctx, d_num_nonzeros);
|
||||
device->StreamSync(ctx, stream);
|
||||
|
||||
// truncate array to size
|
||||
return ret.CreateView({num_nonzeros}, ret->dtype, 0);
|
||||
}
|
||||
|
||||
template IdArray NonZero<kDGLCUDA, int32_t>(IdArray);
|
||||
template IdArray NonZero<kDGLCUDA, int64_t>(IdArray);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,441 @@
|
||||
/**
|
||||
* Copyright (c) 2020-2021 by Contributors
|
||||
* @file array/cuda/array_op_impl.cu
|
||||
* @brief Array operator GPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "../../runtime/cuda/cuda_hashtable.cuh"
|
||||
#include "../arith.h"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
using namespace runtime::cuda;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
///////////////////////////// BinaryElewise /////////////////////////////
|
||||
|
||||
template <typename IdType, typename Op>
|
||||
__global__ void _BinaryElewiseKernel(
|
||||
const IdType* lhs, const IdType* rhs, IdType* out, int64_t length) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < length) {
|
||||
out[tx] = Op::Call(lhs[tx], rhs[tx]);
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename Op>
|
||||
IdArray BinaryElewise(IdArray lhs, IdArray rhs) {
|
||||
const int64_t len = lhs->shape[0];
|
||||
IdArray ret = NewIdArray(lhs->shape[0], lhs->ctx, lhs->dtype.bits);
|
||||
const IdType* lhs_data = static_cast<IdType*>(lhs->data);
|
||||
const IdType* rhs_data = static_cast<IdType*>(rhs->data);
|
||||
IdType* ret_data = static_cast<IdType*>(ret->data);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
int nt = cuda::FindNumThreads(len);
|
||||
int nb = (len + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
(_BinaryElewiseKernel<IdType, Op>), nb, nt, 0, stream, lhs_data, rhs_data,
|
||||
ret_data, len);
|
||||
return ret;
|
||||
}
|
||||
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::Add>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::Sub>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::Mul>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::Div>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::Mod>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::GT>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::LT>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::GE>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::LE>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::EQ>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::NE>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::Add>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::Sub>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::Mul>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::Div>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::Mod>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::GT>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::LT>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::GE>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::LE>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::EQ>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::NE>(
|
||||
IdArray lhs, IdArray rhs);
|
||||
|
||||
template <typename IdType, typename Op>
|
||||
__global__ void _BinaryElewiseKernel(
|
||||
const IdType* lhs, IdType rhs, IdType* out, int64_t length) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < length) {
|
||||
out[tx] = Op::Call(lhs[tx], rhs);
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename Op>
|
||||
IdArray BinaryElewise(IdArray lhs, IdType rhs) {
|
||||
const int64_t len = lhs->shape[0];
|
||||
IdArray ret = NewIdArray(lhs->shape[0], lhs->ctx, lhs->dtype.bits);
|
||||
const IdType* lhs_data = static_cast<IdType*>(lhs->data);
|
||||
IdType* ret_data = static_cast<IdType*>(ret->data);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
int nt = cuda::FindNumThreads(len);
|
||||
int nb = (len + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
(_BinaryElewiseKernel<IdType, Op>), nb, nt, 0, stream, lhs_data, rhs,
|
||||
ret_data, len);
|
||||
return ret;
|
||||
}
|
||||
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::Add>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::Sub>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::Mul>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::Div>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::Mod>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::GT>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::LT>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::GE>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::LE>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::EQ>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::NE>(
|
||||
IdArray lhs, int32_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::Add>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::Sub>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::Mul>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::Div>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::Mod>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::GT>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::LT>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::GE>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::LE>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::EQ>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::NE>(
|
||||
IdArray lhs, int64_t rhs);
|
||||
|
||||
template <typename IdType, typename Op>
|
||||
__global__ void _BinaryElewiseKernel(
|
||||
IdType lhs, const IdType* rhs, IdType* out, int64_t length) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < length) {
|
||||
out[tx] = Op::Call(lhs, rhs[tx]);
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename Op>
|
||||
IdArray BinaryElewise(IdType lhs, IdArray rhs) {
|
||||
const int64_t len = rhs->shape[0];
|
||||
IdArray ret = NewIdArray(rhs->shape[0], rhs->ctx, rhs->dtype.bits);
|
||||
const IdType* rhs_data = static_cast<IdType*>(rhs->data);
|
||||
IdType* ret_data = static_cast<IdType*>(ret->data);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
int nt = cuda::FindNumThreads(len);
|
||||
int nb = (len + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
(_BinaryElewiseKernel<IdType, Op>), nb, nt, 0, stream, lhs, rhs_data,
|
||||
ret_data, len);
|
||||
return ret;
|
||||
}
|
||||
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::Add>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::Sub>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::Mul>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::Div>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::Mod>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::GT>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::LT>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::GE>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::LE>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::EQ>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int32_t, arith::NE>(
|
||||
int32_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::Add>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::Sub>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::Mul>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::Div>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::Mod>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::GT>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::LT>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::GE>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::LE>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::EQ>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
template IdArray BinaryElewise<kDGLCUDA, int64_t, arith::NE>(
|
||||
int64_t lhs, IdArray rhs);
|
||||
|
||||
template <typename IdType, typename Op>
|
||||
__global__ void _UnaryElewiseKernel(
|
||||
const IdType* lhs, IdType* out, int64_t length) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < length) {
|
||||
out[tx] = Op::Call(lhs[tx]);
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename Op>
|
||||
IdArray UnaryElewise(IdArray lhs) {
|
||||
const int64_t len = lhs->shape[0];
|
||||
IdArray ret = NewIdArray(lhs->shape[0], lhs->ctx, lhs->dtype.bits);
|
||||
const IdType* lhs_data = static_cast<IdType*>(lhs->data);
|
||||
IdType* ret_data = static_cast<IdType*>(ret->data);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
int nt = cuda::FindNumThreads(len);
|
||||
int nb = (len + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
(_UnaryElewiseKernel<IdType, Op>), nb, nt, 0, stream, lhs_data, ret_data,
|
||||
len);
|
||||
return ret;
|
||||
}
|
||||
|
||||
template IdArray UnaryElewise<kDGLCUDA, int32_t, arith::Neg>(IdArray lhs);
|
||||
template IdArray UnaryElewise<kDGLCUDA, int64_t, arith::Neg>(IdArray lhs);
|
||||
|
||||
///////////////////////////// Full /////////////////////////////
|
||||
|
||||
template <typename DType>
|
||||
__global__ void _FullKernel(DType* out, int64_t length, DType val) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < length) {
|
||||
out[tx] = val;
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename DType>
|
||||
NDArray Full(DType val, int64_t length, DGLContext ctx) {
|
||||
NDArray ret = NDArray::Empty({length}, DGLDataTypeTraits<DType>::dtype, ctx);
|
||||
DType* ret_data = static_cast<DType*>(ret->data);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
int nt = cuda::FindNumThreads(length);
|
||||
int nb = (length + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
(_FullKernel<DType>), nb, nt, 0, stream, ret_data, length, val);
|
||||
return ret;
|
||||
}
|
||||
|
||||
template IdArray Full<kDGLCUDA, int32_t>(
|
||||
int32_t val, int64_t length, DGLContext ctx);
|
||||
template IdArray Full<kDGLCUDA, int64_t>(
|
||||
int64_t val, int64_t length, DGLContext ctx);
|
||||
template IdArray Full<kDGLCUDA, __half>(
|
||||
__half val, int64_t length, DGLContext ctx);
|
||||
#if BF16_ENABLED
|
||||
template IdArray Full<kDGLCUDA, __nv_bfloat16>(
|
||||
__nv_bfloat16 val, int64_t length, DGLContext ctx);
|
||||
#endif // BF16_ENABLED
|
||||
template IdArray Full<kDGLCUDA, float>(
|
||||
float val, int64_t length, DGLContext ctx);
|
||||
template IdArray Full<kDGLCUDA, double>(
|
||||
double val, int64_t length, DGLContext ctx);
|
||||
|
||||
///////////////////////////// Range /////////////////////////////
|
||||
|
||||
template <typename IdType>
|
||||
__global__ void _RangeKernel(IdType* out, IdType low, IdType length) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < length) {
|
||||
out[tx] = low + tx;
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray Range(IdType low, IdType high, DGLContext ctx) {
|
||||
CHECK(high >= low) << "high must be bigger than low";
|
||||
const IdType length = high - low;
|
||||
IdArray ret = NewIdArray(length, ctx, sizeof(IdType) * 8);
|
||||
if (length == 0) return ret;
|
||||
IdType* ret_data = static_cast<IdType*>(ret->data);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
int nt = cuda::FindNumThreads(length);
|
||||
int nb = (length + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
(_RangeKernel<IdType>), nb, nt, 0, stream, ret_data, low, length);
|
||||
return ret;
|
||||
}
|
||||
|
||||
template IdArray Range<kDGLCUDA, int32_t>(int32_t, int32_t, DGLContext);
|
||||
template IdArray Range<kDGLCUDA, int64_t>(int64_t, int64_t, DGLContext);
|
||||
|
||||
///////////////////////////// Relabel_ //////////////////////////////
|
||||
|
||||
template <typename IdType>
|
||||
__global__ void _RelabelKernel(
|
||||
IdType* out, int64_t length, DeviceOrderedHashTable<IdType> table) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride_x = gridDim.x * blockDim.x;
|
||||
|
||||
while (tx < length) {
|
||||
out[tx] = table.Search(out[tx])->local;
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray Relabel_(const std::vector<IdArray>& arrays) {
|
||||
IdArray all_nodes = Concat(arrays);
|
||||
const int64_t total_length = all_nodes->shape[0];
|
||||
|
||||
if (total_length == 0) {
|
||||
return all_nodes;
|
||||
}
|
||||
|
||||
const auto& ctx = arrays[0]->ctx;
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
|
||||
// build node maps and get the induced nodes
|
||||
OrderedHashTable<IdType> node_map(total_length, ctx, stream);
|
||||
int64_t num_induced = 0;
|
||||
int64_t* num_induced_device =
|
||||
static_cast<int64_t*>(device->AllocWorkspace(ctx, sizeof(int64_t)));
|
||||
IdArray induced_nodes = NewIdArray(total_length, ctx, sizeof(IdType) * 8);
|
||||
|
||||
CUDA_CALL(cudaMemsetAsync(
|
||||
num_induced_device, 0, sizeof(*num_induced_device), stream));
|
||||
|
||||
node_map.FillWithDuplicates(
|
||||
all_nodes.Ptr<IdType>(), all_nodes->shape[0], induced_nodes.Ptr<IdType>(),
|
||||
num_induced_device, stream);
|
||||
// copy using the internal current stream
|
||||
device->CopyDataFromTo(
|
||||
num_induced_device, 0, &num_induced, 0, sizeof(num_induced), ctx,
|
||||
DGLContext{kDGLCPU, 0}, DGLDataType{kDGLInt, 64, 1});
|
||||
|
||||
device->StreamSync(ctx, stream);
|
||||
device->FreeWorkspace(ctx, num_induced_device);
|
||||
|
||||
// resize the induced nodes
|
||||
induced_nodes->shape[0] = num_induced;
|
||||
|
||||
// relabel
|
||||
const int nt = 128;
|
||||
for (IdArray arr : arrays) {
|
||||
const int64_t length = arr->shape[0];
|
||||
int nb = (length + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
(_RelabelKernel<IdType>), nb, nt, 0, stream, arr.Ptr<IdType>(), length,
|
||||
node_map.DeviceHandle());
|
||||
}
|
||||
|
||||
return induced_nodes;
|
||||
}
|
||||
|
||||
template IdArray Relabel_<kDGLCUDA, int32_t>(
|
||||
const std::vector<IdArray>& arrays);
|
||||
template IdArray Relabel_<kDGLCUDA, int64_t>(
|
||||
const std::vector<IdArray>& arrays);
|
||||
|
||||
///////////////////////////// AsNumBits /////////////////////////////
|
||||
|
||||
template <typename InType, typename OutType>
|
||||
__global__ void _CastKernel(const InType* in, OutType* out, size_t length) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < length) {
|
||||
out[tx] = in[tx];
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
IdArray AsNumBits(IdArray arr, uint8_t bits) {
|
||||
const std::vector<int64_t> shape(arr->shape, arr->shape + arr->ndim);
|
||||
IdArray ret = IdArray::Empty(shape, DGLDataType{kDGLInt, bits, 1}, arr->ctx);
|
||||
const int64_t length = ret.NumElements();
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
int nt = cuda::FindNumThreads(length);
|
||||
int nb = (length + nt - 1) / nt;
|
||||
if (bits == 32) {
|
||||
CUDA_KERNEL_CALL(
|
||||
(_CastKernel<IdType, int32_t>), nb, nt, 0, stream,
|
||||
static_cast<IdType*>(arr->data), static_cast<int32_t*>(ret->data),
|
||||
length);
|
||||
} else {
|
||||
CUDA_KERNEL_CALL(
|
||||
(_CastKernel<IdType, int64_t>), nb, nt, 0, stream,
|
||||
static_cast<IdType*>(arr->data), static_cast<int64_t*>(ret->data),
|
||||
length);
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
template IdArray AsNumBits<kDGLCUDA, int32_t>(IdArray arr, uint8_t bits);
|
||||
template IdArray AsNumBits<kDGLCUDA, int64_t>(IdArray arr, uint8_t bits);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,61 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file array/cuda/array_scatter.cu
|
||||
* @brief Array scatter GPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <typename DType, typename IdType>
|
||||
__global__ void _ScatterKernel(
|
||||
const IdType* index, const DType* value, int64_t length, DType* out) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < length) {
|
||||
out[index[tx]] = value[tx];
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename DType, typename IdType>
|
||||
void Scatter_(IdArray index, NDArray value, NDArray out) {
|
||||
const int64_t len = index->shape[0];
|
||||
const IdType* idx = index.Ptr<IdType>();
|
||||
const DType* val = value.Ptr<DType>();
|
||||
DType* outd = out.Ptr<DType>();
|
||||
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const int nt = cuda::FindNumThreads(len);
|
||||
const int nb = (len + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(_ScatterKernel, nb, nt, 0, stream, idx, val, len, outd);
|
||||
}
|
||||
|
||||
template void Scatter_<kDGLCUDA, int32_t, int32_t>(IdArray, NDArray, NDArray);
|
||||
template void Scatter_<kDGLCUDA, int64_t, int32_t>(IdArray, NDArray, NDArray);
|
||||
template void Scatter_<kDGLCUDA, __half, int32_t>(IdArray, NDArray, NDArray);
|
||||
#if BF16_ENABLED
|
||||
template void Scatter_<kDGLCUDA, __nv_bfloat16, int32_t>(
|
||||
IdArray, NDArray, NDArray);
|
||||
#endif // BF16_ENABLED
|
||||
template void Scatter_<kDGLCUDA, float, int32_t>(IdArray, NDArray, NDArray);
|
||||
template void Scatter_<kDGLCUDA, double, int32_t>(IdArray, NDArray, NDArray);
|
||||
template void Scatter_<kDGLCUDA, int32_t, int64_t>(IdArray, NDArray, NDArray);
|
||||
template void Scatter_<kDGLCUDA, int64_t, int64_t>(IdArray, NDArray, NDArray);
|
||||
template void Scatter_<kDGLCUDA, __half, int64_t>(IdArray, NDArray, NDArray);
|
||||
#if BF16_ENABLED
|
||||
template void Scatter_<kDGLCUDA, __nv_bfloat16, int64_t>(
|
||||
IdArray, NDArray, NDArray);
|
||||
#endif // BF16_ENABLED
|
||||
template void Scatter_<kDGLCUDA, float, int64_t>(IdArray, NDArray, NDArray);
|
||||
template void Scatter_<kDGLCUDA, double, int64_t>(IdArray, NDArray, NDArray);
|
||||
|
||||
}; // namespace impl
|
||||
}; // namespace aten
|
||||
}; // namespace dgl
|
||||
@@ -0,0 +1,61 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cpu/array_sort.cu
|
||||
* @brief Array sort GPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::pair<IdArray, IdArray> Sort(IdArray array, int num_bits) {
|
||||
const auto& ctx = array->ctx;
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
const int64_t nitems = array->shape[0];
|
||||
IdArray orig_idx = Range(0, nitems, 64, ctx);
|
||||
IdArray sorted_array = NewIdArray(nitems, ctx, array->dtype.bits);
|
||||
IdArray sorted_idx = NewIdArray(nitems, ctx, 64);
|
||||
|
||||
const IdType* keys_in = array.Ptr<IdType>();
|
||||
const int64_t* values_in = orig_idx.Ptr<int64_t>();
|
||||
IdType* keys_out = sorted_array.Ptr<IdType>();
|
||||
int64_t* values_out = sorted_idx.Ptr<int64_t>();
|
||||
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
if (num_bits == 0) {
|
||||
num_bits = sizeof(IdType) * 8;
|
||||
}
|
||||
|
||||
// Allocate workspace
|
||||
size_t workspace_size = 0;
|
||||
CUDA_CALL(cub::DeviceRadixSort::SortPairs(
|
||||
nullptr, workspace_size, keys_in, keys_out, values_in, values_out, nitems,
|
||||
0, num_bits, stream));
|
||||
void* workspace = device->AllocWorkspace(ctx, workspace_size);
|
||||
|
||||
// Compute
|
||||
CUDA_CALL(cub::DeviceRadixSort::SortPairs(
|
||||
workspace, workspace_size, keys_in, keys_out, values_in, values_out,
|
||||
nitems, 0, num_bits, stream));
|
||||
|
||||
device->FreeWorkspace(ctx, workspace);
|
||||
|
||||
return std::make_pair(sorted_array, sorted_idx);
|
||||
}
|
||||
|
||||
template std::pair<IdArray, IdArray> Sort<kDGLCUDA, int32_t>(
|
||||
IdArray, int num_bits);
|
||||
template std::pair<IdArray, IdArray> Sort<kDGLCUDA, int64_t>(
|
||||
IdArray, int num_bits);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,336 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file array/cuda/atomic.cuh
|
||||
* @brief Atomic functions
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CUDA_ATOMIC_CUH_
|
||||
#define DGL_ARRAY_CUDA_ATOMIC_CUH_
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <cassert>
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
|
||||
#include "bf16.cuh"
|
||||
#include "fp16.cuh"
|
||||
|
||||
#if __CUDA_ARCH__ >= 600
|
||||
#include <cuda_fp16.h>
|
||||
#endif
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace cuda {
|
||||
|
||||
// Type trait for selecting code type
|
||||
template <int Bytes>
|
||||
struct Code {};
|
||||
|
||||
template <>
|
||||
struct Code<2> {
|
||||
typedef unsigned short int Type; // NOLINT
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Code<4> {
|
||||
typedef unsigned int Type; // NOLINT
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Code<8> {
|
||||
typedef unsigned long long int Type; // NOLINT
|
||||
};
|
||||
|
||||
// Helper class for converting to/from atomicCAS compatible types.
|
||||
template <typename T>
|
||||
struct Cast {
|
||||
typedef typename Code<sizeof(T)>::Type Type;
|
||||
static __device__ __forceinline__ Type Encode(T val) {
|
||||
return static_cast<Type>(val);
|
||||
}
|
||||
static __device__ __forceinline__ T Decode(Type code) {
|
||||
return static_cast<T>(code);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Cast<half> {
|
||||
typedef Code<sizeof(half)>::Type Type;
|
||||
static __device__ __forceinline__ Type Encode(half val) {
|
||||
return __half_as_ushort(val);
|
||||
}
|
||||
static __device__ __forceinline__ half Decode(Type code) {
|
||||
return __ushort_as_half(code);
|
||||
}
|
||||
};
|
||||
|
||||
#if BF16_ENABLED
|
||||
template <>
|
||||
struct Cast<__nv_bfloat16> {
|
||||
typedef Code<sizeof(__nv_bfloat16)>::Type Type;
|
||||
static __device__ __forceinline__ Type Encode(__nv_bfloat16 val) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
return __bfloat16_as_ushort(val);
|
||||
#else
|
||||
printf(
|
||||
"Atomic operations are not supported for bfloat16 (BF16) "
|
||||
"on GPUs with compute capability less than 8.0.\n");
|
||||
__trap();
|
||||
return static_cast<Type>(0);
|
||||
#endif
|
||||
}
|
||||
static __device__ __forceinline__ __nv_bfloat16 Decode(Type code) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
return __ushort_as_bfloat16(code);
|
||||
#else
|
||||
printf(
|
||||
"Atomic operations are not supported for bfloat16 (BF16) "
|
||||
"on GPUs with compute capability less than 8.0.\n");
|
||||
__trap();
|
||||
return static_cast<__nv_bfloat16>(0.0f);
|
||||
#endif
|
||||
}
|
||||
};
|
||||
#endif // BF16_ENABLED
|
||||
|
||||
template <>
|
||||
struct Cast<float> {
|
||||
typedef Code<sizeof(float)>::Type Type;
|
||||
static __device__ __forceinline__ Type Encode(float val) {
|
||||
return __float_as_uint(val);
|
||||
}
|
||||
static __device__ __forceinline__ float Decode(Type code) {
|
||||
return __uint_as_float(code);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Cast<double> {
|
||||
typedef Code<sizeof(double)>::Type Type;
|
||||
static __device__ __forceinline__ Type Encode(double val) {
|
||||
return __double_as_longlong(val);
|
||||
}
|
||||
static __device__ __forceinline__ double Decode(Type code) {
|
||||
return __longlong_as_double(code);
|
||||
}
|
||||
};
|
||||
|
||||
static __device__ __forceinline__ unsigned short int atomicCASshort( // NOLINT
|
||||
unsigned short int* address, // NOLINT
|
||||
unsigned short int compare, // NOLINT
|
||||
unsigned short int val) { // NOLINT
|
||||
static_assert(CUDART_VERSION >= 10000, "Requires at least CUDA 10");
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__) >= 700)
|
||||
return atomicCAS(address, compare, val);
|
||||
#else
|
||||
(void)address;
|
||||
(void)compare;
|
||||
(void)val;
|
||||
printf(
|
||||
"Atomic operations are not supported for half precision (FP16) "
|
||||
"on this GPU.\n");
|
||||
__trap();
|
||||
return val;
|
||||
#endif // (defined(__CUDA_ARCH__) && (__CUDA_ARCH__) >= 700)
|
||||
}
|
||||
|
||||
#define DEFINE_ATOMIC(NAME) \
|
||||
template <typename T> \
|
||||
__device__ __forceinline__ T Atomic##NAME(T* addr, T val) { \
|
||||
typedef typename Cast<T>::Type CT; \
|
||||
CT* addr_as_ui = reinterpret_cast<CT*>(addr); \
|
||||
CT old = *addr_as_ui; \
|
||||
CT assumed = old; \
|
||||
do { \
|
||||
assumed = old; \
|
||||
old = atomicCAS( \
|
||||
addr_as_ui, assumed, \
|
||||
Cast<T>::Encode(OP(val, Cast<T>::Decode(old)))); \
|
||||
} while (assumed != old); \
|
||||
return Cast<T>::Decode(old); \
|
||||
}
|
||||
|
||||
#define DEFINE_ATOMIC_16BIT(NAME, dtype) \
|
||||
template <> \
|
||||
__device__ __forceinline__ dtype Atomic##NAME<dtype>( \
|
||||
dtype * addr, dtype val) { \
|
||||
typedef uint16_t CT; \
|
||||
CT* addr_as_ui = reinterpret_cast<CT*>(addr); \
|
||||
CT old = *addr_as_ui; \
|
||||
CT assumed = old; \
|
||||
do { \
|
||||
assumed = old; \
|
||||
old = atomicCASshort( \
|
||||
addr_as_ui, assumed, \
|
||||
Cast<dtype>::Encode(OP(val, Cast<dtype>::Decode(old)))); \
|
||||
} while (assumed != old); \
|
||||
return Cast<dtype>::Decode(old); \
|
||||
}
|
||||
|
||||
#define OP(a, b) max(a, b)
|
||||
DEFINE_ATOMIC(Max)
|
||||
DEFINE_ATOMIC_16BIT(Max, half)
|
||||
#if BF16_ENABLED
|
||||
DEFINE_ATOMIC_16BIT(Max, __nv_bfloat16)
|
||||
#endif // BF16_ENABLED
|
||||
#undef OP
|
||||
|
||||
#define OP(a, b) min(a, b)
|
||||
DEFINE_ATOMIC(Min)
|
||||
DEFINE_ATOMIC_16BIT(Min, half)
|
||||
#if BF16_ENABLED
|
||||
DEFINE_ATOMIC_16BIT(Min, __nv_bfloat16)
|
||||
#endif // BF16_ENABLED
|
||||
#undef OP
|
||||
|
||||
#define OP(a, b) a + b
|
||||
DEFINE_ATOMIC(Add)
|
||||
#undef OP
|
||||
|
||||
/**
|
||||
* @brief Performs an atomic compare-and-swap on 64 bit integers. That is,
|
||||
* it the word `old` at the memory location `address`, computes
|
||||
* `(old == compare ? val : old)` , and stores the result back to memory at
|
||||
* the same address.
|
||||
*
|
||||
* @param address The address to perform the atomic operation on.
|
||||
* @param compare The value to compare to.
|
||||
* @param val The new value to conditionally store.
|
||||
*
|
||||
* @return The old value at the address.
|
||||
*/
|
||||
inline __device__ int64_t
|
||||
AtomicCAS(int64_t* const address, const int64_t compare, const int64_t val) {
|
||||
// match the type of "::atomicCAS", so ignore lint warning
|
||||
using Type = unsigned long long int; // NOLINT
|
||||
|
||||
static_assert(sizeof(Type) == sizeof(*address), "Type width must match");
|
||||
|
||||
return atomicCAS(
|
||||
reinterpret_cast<Type*>(address), static_cast<Type>(compare),
|
||||
static_cast<Type>(val));
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Performs an atomic compare-and-swap on 32 bit integers. That is,
|
||||
* it the word `old` at the memory location `address`, computes
|
||||
* `(old == compare ? val : old)` , and stores the result back to memory at
|
||||
* the same address.
|
||||
*
|
||||
* @param address The address to perform the atomic operation on.
|
||||
* @param compare The value to compare to.
|
||||
* @param val The new value to conditionally store.
|
||||
*
|
||||
* @return The old value at the address.
|
||||
*/
|
||||
inline __device__ int32_t
|
||||
AtomicCAS(int32_t* const address, const int32_t compare, const int32_t val) {
|
||||
// match the type of "::atomicCAS", so ignore lint warning
|
||||
using Type = int; // NOLINT
|
||||
|
||||
static_assert(sizeof(Type) == sizeof(*address), "Type width must match");
|
||||
|
||||
return atomicCAS(
|
||||
reinterpret_cast<Type*>(address), static_cast<Type>(compare),
|
||||
static_cast<Type>(val));
|
||||
}
|
||||
|
||||
inline __device__ int64_t AtomicMax(int64_t* const address, const int64_t val) {
|
||||
// match the type of "::atomicCAS", so ignore lint warning
|
||||
using Type = unsigned long long int; // NOLINT
|
||||
|
||||
static_assert(sizeof(Type) == sizeof(*address), "Type width must match");
|
||||
|
||||
return atomicMax(reinterpret_cast<Type*>(address), static_cast<Type>(val));
|
||||
}
|
||||
|
||||
inline __device__ int32_t AtomicMax(int32_t* const address, const int32_t val) {
|
||||
// match the type of "::atomicCAS", so ignore lint warning
|
||||
using Type = int; // NOLINT
|
||||
|
||||
static_assert(sizeof(Type) == sizeof(*address), "Type width must match");
|
||||
|
||||
return atomicMax(reinterpret_cast<Type*>(address), static_cast<Type>(val));
|
||||
}
|
||||
|
||||
template <>
|
||||
__device__ __forceinline__ float AtomicAdd<float>(float* addr, float val) {
|
||||
#if __CUDA_ARCH__ >= 200
|
||||
return atomicAdd(addr, val);
|
||||
#else
|
||||
typedef float T;
|
||||
typedef typename Cast<T>::Type CT;
|
||||
CT* addr_as_ui = reinterpret_cast<CT*>(addr);
|
||||
CT old = *addr_as_ui;
|
||||
CT assumed = old;
|
||||
do {
|
||||
assumed = old;
|
||||
old = atomicCAS(
|
||||
addr_as_ui, assumed, Cast<T>::Encode(Cast<T>::Decode(old) + val));
|
||||
} while (assumed != old);
|
||||
return Cast<T>::Decode(old);
|
||||
#endif // __CUDA_ARCH__
|
||||
}
|
||||
|
||||
template <>
|
||||
__device__ __forceinline__ double AtomicAdd<double>(double* addr, double val) {
|
||||
#if __CUDA_ARCH__ >= 600
|
||||
return atomicAdd(addr, val);
|
||||
#else
|
||||
typedef double T;
|
||||
typedef typename Cast<T>::Type CT;
|
||||
CT* addr_as_ui = reinterpret_cast<CT*>(addr);
|
||||
CT old = *addr_as_ui;
|
||||
CT assumed = old;
|
||||
do {
|
||||
assumed = old;
|
||||
old = atomicCAS(
|
||||
addr_as_ui, assumed, Cast<T>::Encode(Cast<T>::Decode(old) + val));
|
||||
} while (assumed != old);
|
||||
return Cast<T>::Decode(old);
|
||||
#endif
|
||||
}
|
||||
|
||||
#if defined(CUDART_VERSION) && CUDART_VERSION >= 10000
|
||||
template <>
|
||||
__device__ __forceinline__ half AtomicAdd<half>(half* addr, half val) {
|
||||
// make sure we have half support
|
||||
#if __CUDA_ARCH__ >= 700
|
||||
return atomicAdd(addr, val);
|
||||
#else
|
||||
(void)addr;
|
||||
(void)val;
|
||||
printf(
|
||||
"Atomic operations are not supported for half precision (FP16) "
|
||||
"on this GPU.\n");
|
||||
__trap();
|
||||
return val;
|
||||
#endif // __CUDA_ARCH__ >= 700
|
||||
}
|
||||
#endif // defined(CUDART_VERSION) && CUDART_VERSION >= 10000
|
||||
|
||||
#if BF16_ENABLED
|
||||
template <>
|
||||
__device__ __forceinline__ __nv_bfloat16
|
||||
AtomicAdd<__nv_bfloat16>(__nv_bfloat16* addr, __nv_bfloat16 val) {
|
||||
// make sure we have bfloat16 support
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
return atomicAdd(addr, val);
|
||||
#else
|
||||
(void)addr;
|
||||
(void)val;
|
||||
printf(
|
||||
"Atomic operations are not supported for bfloat16 (BF16) "
|
||||
"on GPUs with compute capability less than 8.0.\n");
|
||||
__trap();
|
||||
return val;
|
||||
#endif // defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
}
|
||||
#endif // BF16_ENABLED
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CUDA_ATOMIC_CUH_
|
||||
@@ -0,0 +1,149 @@
|
||||
/**
|
||||
* Copyright (c) 2022 by Contributors
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
* @file array/cuda/bf16.cuh
|
||||
* @brief bfloat16 related functions.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CUDA_BF16_CUH_
|
||||
#define DGL_ARRAY_CUDA_BF16_CUH_
|
||||
|
||||
#if BF16_ENABLED
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
static __device__ __forceinline__ __nv_bfloat16
|
||||
max(__nv_bfloat16 a, __nv_bfloat16 b) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
return __hmax(a, b);
|
||||
#else
|
||||
return __nv_bfloat16(max(float(a), float(b))); // NOLINT
|
||||
#endif
|
||||
}
|
||||
|
||||
static __device__ __forceinline__ __nv_bfloat16
|
||||
min(__nv_bfloat16 a, __nv_bfloat16 b) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
return __hmin(a, b);
|
||||
#else
|
||||
return __nv_bfloat16(min(float(a), float(b))); // NOLINT
|
||||
#endif
|
||||
}
|
||||
|
||||
#ifdef __CUDACC__
|
||||
// Arithmetic BF16 operations for architecture >= 8.0 are already defined in
|
||||
// cuda_bf16.h
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)
|
||||
// CUDA 12.2 adds "emulated" support for older architectures.
|
||||
#if defined(CUDART_VERSION) && (CUDART_VERSION < 12020)
|
||||
__device__ __forceinline__ __nv_bfloat16
|
||||
operator+(const __nv_bfloat16& lh, const __nv_bfloat16& rh) {
|
||||
return __nv_bfloat16(float(lh) + float(rh)); // NOLINT
|
||||
}
|
||||
__device__ __forceinline__ __nv_bfloat16
|
||||
operator-(const __nv_bfloat16& lh, const __nv_bfloat16& rh) {
|
||||
return __nv_bfloat16(float(lh) - float(rh)); // NOLINT
|
||||
}
|
||||
__device__ __forceinline__ __nv_bfloat16
|
||||
operator*(const __nv_bfloat16& lh, const __nv_bfloat16& rh) {
|
||||
return __nv_bfloat16(float(lh) * float(rh)); // NOLINT
|
||||
}
|
||||
__device__ __forceinline__ __nv_bfloat16
|
||||
operator/(const __nv_bfloat16& lh, const __nv_bfloat16& rh) {
|
||||
return __nv_bfloat16(float(lh) / float(rh)); // NOLINT
|
||||
}
|
||||
|
||||
__device__ __forceinline__ __nv_bfloat16& operator+=(
|
||||
__nv_bfloat16& lh, const __nv_bfloat16& rh) { // NOLINT
|
||||
lh = __nv_bfloat16(float(lh) + float(rh)); // NOLINT
|
||||
return lh;
|
||||
}
|
||||
__device__ __forceinline__ __nv_bfloat16& operator-=(
|
||||
__nv_bfloat16& lh, const __nv_bfloat16& rh) { // NOLINT
|
||||
lh = __nv_bfloat16(float(lh) - float(rh)); // NOLINT
|
||||
return lh;
|
||||
}
|
||||
__device__ __forceinline__ __nv_bfloat16& operator*=(
|
||||
__nv_bfloat16& lh, const __nv_bfloat16& rh) { // NOLINT
|
||||
lh = __nv_bfloat16(float(lh) * float(rh)); // NOLINT
|
||||
return lh;
|
||||
}
|
||||
__device__ __forceinline__ __nv_bfloat16& operator/=(
|
||||
__nv_bfloat16& lh, const __nv_bfloat16& rh) { // NOLINT
|
||||
lh = __nv_bfloat16(float(lh) / float(rh)); // NOLINT
|
||||
return lh;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ __nv_bfloat16& operator++(
|
||||
__nv_bfloat16& h) { // NOLINT
|
||||
h = __nv_bfloat16(float(h) + 1.0f); // NOLINT
|
||||
return h;
|
||||
}
|
||||
__device__ __forceinline__ __nv_bfloat16& operator--(
|
||||
__nv_bfloat16& h) { // NOLINT
|
||||
h = __nv_bfloat16(float(h) - 1.0f); // NOLINT
|
||||
return h;
|
||||
}
|
||||
__device__ __forceinline__ __nv_bfloat16
|
||||
operator++(__nv_bfloat16& h, int) { // NOLINT
|
||||
__nv_bfloat16 ret = h;
|
||||
h = __nv_bfloat16(float(h) + 1.0f); // NOLINT
|
||||
return ret;
|
||||
}
|
||||
__device__ __forceinline__ __nv_bfloat16
|
||||
operator--(__nv_bfloat16& h, int) { // NOLINT
|
||||
__nv_bfloat16 ret = h;
|
||||
h = __nv_bfloat16(float(h) - 1.0f); // NOLINT
|
||||
return ret;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ __nv_bfloat16 operator+(const __nv_bfloat16& h) {
|
||||
return h;
|
||||
}
|
||||
__device__ __forceinline__ __nv_bfloat16 operator-(const __nv_bfloat16& h) {
|
||||
return __nv_bfloat16(-float(h)); // NOLINT
|
||||
}
|
||||
|
||||
__device__ __forceinline__ bool operator==(
|
||||
const __nv_bfloat16& lh, const __nv_bfloat16& rh) {
|
||||
return float(lh) == float(rh); // NOLINT
|
||||
}
|
||||
__device__ __forceinline__ bool operator!=(
|
||||
const __nv_bfloat16& lh, const __nv_bfloat16& rh) {
|
||||
return float(lh) != float(rh); // NOLINT
|
||||
}
|
||||
__device__ __forceinline__ bool operator>(
|
||||
const __nv_bfloat16& lh, const __nv_bfloat16& rh) {
|
||||
return float(lh) > float(rh); // NOLINT
|
||||
}
|
||||
__device__ __forceinline__ bool operator<(
|
||||
const __nv_bfloat16& lh, const __nv_bfloat16& rh) {
|
||||
return float(lh) < float(rh); // NOLINT
|
||||
}
|
||||
__device__ __forceinline__ bool operator>=(
|
||||
const __nv_bfloat16& lh, const __nv_bfloat16& rh) {
|
||||
return float(lh) >= float(rh); // NOLINT
|
||||
}
|
||||
__device__ __forceinline__ bool operator<=(
|
||||
const __nv_bfloat16& lh, const __nv_bfloat16& rh) {
|
||||
return float(lh) <= float(rh); // NOLINT
|
||||
}
|
||||
#endif // defined(CUDART_VERSION) && (CUDART_VERSION < 12020)
|
||||
#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)
|
||||
#endif // __CUDACC__
|
||||
|
||||
#endif // BF16_ENABLED
|
||||
|
||||
#endif // DGL_ARRAY_CUDA_BF16_CUH_
|
||||
@@ -0,0 +1,137 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/coo2csr.cc
|
||||
* @brief COO2CSR
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using runtime::NDArray;
|
||||
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix COOToCSR(COOMatrix coo) {
|
||||
LOG(FATAL) << "Unreachable code.";
|
||||
return {};
|
||||
}
|
||||
|
||||
template <>
|
||||
CSRMatrix COOToCSR<kDGLCUDA, int32_t>(COOMatrix coo) {
|
||||
auto* thr_entry = runtime::CUDAThreadEntry::ThreadLocal();
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
// allocate cusparse handle if needed
|
||||
if (!thr_entry->cusparse_handle) {
|
||||
CUSPARSE_CALL(cusparseCreate(&(thr_entry->cusparse_handle)));
|
||||
}
|
||||
CUSPARSE_CALL(cusparseSetStream(thr_entry->cusparse_handle, stream));
|
||||
|
||||
bool row_sorted = coo.row_sorted;
|
||||
bool col_sorted = coo.col_sorted;
|
||||
if (!row_sorted) {
|
||||
// we only need to sort the rows to perform conversion
|
||||
coo = COOSort(coo, false);
|
||||
col_sorted = coo.col_sorted;
|
||||
}
|
||||
|
||||
const int64_t nnz = coo.row->shape[0];
|
||||
CHECK_NO_OVERFLOW(coo.row->dtype, nnz);
|
||||
// TODO(minjie): Many of our current implementation assumes that CSR must have
|
||||
// a data array. This is a temporary workaround. Remove this after:
|
||||
// - The old immutable graph implementation is deprecated.
|
||||
// - The old binary reduce kernel is deprecated.
|
||||
if (!COOHasData(coo))
|
||||
coo.data = aten::Range(0, nnz, coo.row->dtype.bits, coo.row->ctx);
|
||||
|
||||
NDArray indptr =
|
||||
aten::NewIdArray(coo.num_rows + 1, coo.row->ctx, coo.row->dtype.bits);
|
||||
int32_t* indptr_ptr = static_cast<int32_t*>(indptr->data);
|
||||
CUSPARSE_CALL(cusparseXcoo2csr(
|
||||
thr_entry->cusparse_handle, coo.row.Ptr<int32_t>(), nnz, coo.num_rows,
|
||||
indptr_ptr, CUSPARSE_INDEX_BASE_ZERO));
|
||||
|
||||
return CSRMatrix(
|
||||
coo.num_rows, coo.num_cols, indptr, coo.col, coo.data, col_sorted);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Search for the insertion positions for needle in the hay.
|
||||
*
|
||||
* The hay is a list of sorted elements and the result is the insertion position
|
||||
* of each needle so that the insertion still gives sorted order.
|
||||
*
|
||||
* It essentially perform binary search to find upper bound for each needle
|
||||
* elements.
|
||||
*
|
||||
* For example:
|
||||
* hay = [0, 0, 1, 2, 2]
|
||||
* needle = [0, 1, 2, 3]
|
||||
* then,
|
||||
* out = [2, 3, 5, 5]
|
||||
*/
|
||||
template <typename IdType>
|
||||
__global__ void _SortedSearchKernelUpperBound(
|
||||
const IdType* hay, int64_t hay_size, const IdType* needles,
|
||||
int64_t num_needles, IdType* pos) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < num_needles) {
|
||||
const IdType ele = needles[tx];
|
||||
// binary search
|
||||
IdType lo = 0, hi = hay_size;
|
||||
while (lo < hi) {
|
||||
IdType mid = (lo + hi) >> 1;
|
||||
if (hay[mid] <= ele) {
|
||||
lo = mid + 1;
|
||||
} else {
|
||||
hi = mid;
|
||||
}
|
||||
}
|
||||
pos[tx] = lo;
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <>
|
||||
CSRMatrix COOToCSR<kDGLCUDA, int64_t>(COOMatrix coo) {
|
||||
const auto& ctx = coo.row->ctx;
|
||||
const auto nbits = coo.row->dtype.bits;
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
bool row_sorted = coo.row_sorted;
|
||||
bool col_sorted = coo.col_sorted;
|
||||
if (!row_sorted) {
|
||||
coo = COOSort(coo, false);
|
||||
col_sorted = coo.col_sorted;
|
||||
}
|
||||
|
||||
const int64_t nnz = coo.row->shape[0];
|
||||
// TODO(minjie): Many of our current implementation assumes that CSR must have
|
||||
// a data array. This is a temporary workaround. Remove this after:
|
||||
// - The old immutable graph implementation is deprecated.
|
||||
// - The old binary reduce kernel is deprecated.
|
||||
if (!COOHasData(coo))
|
||||
coo.data = aten::Range(0, nnz, coo.row->dtype.bits, coo.row->ctx);
|
||||
|
||||
IdArray rowids = Range(0, coo.num_rows, nbits, ctx);
|
||||
const int nt = cuda::FindNumThreads(coo.num_rows);
|
||||
const int nb = (coo.num_rows + nt - 1) / nt;
|
||||
IdArray indptr = Full(0, coo.num_rows + 1, nbits, ctx);
|
||||
CUDA_KERNEL_CALL(
|
||||
_SortedSearchKernelUpperBound, nb, nt, 0, stream, coo.row.Ptr<int64_t>(),
|
||||
nnz, rowids.Ptr<int64_t>(), coo.num_rows, indptr.Ptr<int64_t>() + 1);
|
||||
|
||||
return CSRMatrix(
|
||||
coo.num_rows, coo.num_cols, indptr, coo.col, coo.data, col_sorted);
|
||||
}
|
||||
|
||||
template CSRMatrix COOToCSR<kDGLCUDA, int32_t>(COOMatrix coo);
|
||||
template CSRMatrix COOToCSR<kDGLCUDA, int64_t>(COOMatrix coo);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,168 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/coo_sort.cc
|
||||
* @brief Sort COO index
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include "../../c_api_common.h"
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using runtime::NDArray;
|
||||
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
///////////////////////////// COOSort_ /////////////////////////////
|
||||
|
||||
/**
|
||||
* @brief Encode row and column IDs into a single scalar per edge.
|
||||
*
|
||||
* @tparam IdType The type to encode as.
|
||||
* @param row The row (src) IDs per edge.
|
||||
* @param col The column (dst) IDs per edge.
|
||||
* @param nnz The number of edges.
|
||||
* @param col_bits The number of bits used to encode the destination. The row
|
||||
* information is packed into the remaining bits.
|
||||
* @param key The encoded edges (output).
|
||||
*/
|
||||
template <typename IdType>
|
||||
__global__ void _COOEncodeEdgesKernel(
|
||||
const IdType* const row, const IdType* const col, const int64_t nnz,
|
||||
const int col_bits, IdType* const key) {
|
||||
int64_t tx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
|
||||
|
||||
if (tx < nnz) {
|
||||
key[tx] = row[tx] << col_bits | col[tx];
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Decode row and column IDs from the encoded edges.
|
||||
*
|
||||
* @tparam IdType The type the edges are encoded as.
|
||||
* @param key The encoded edges.
|
||||
* @param nnz The number of edges.
|
||||
* @param col_bits The number of bits used to store the column/dst ID.
|
||||
* @param row The row (src) IDs per edge (output).
|
||||
* @param col The col (dst) IDs per edge (output).
|
||||
*/
|
||||
template <typename IdType>
|
||||
__global__ void _COODecodeEdgesKernel(
|
||||
const IdType* const key, const int64_t nnz, const int col_bits,
|
||||
IdType* const row, IdType* const col) {
|
||||
int64_t tx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
|
||||
|
||||
if (tx < nnz) {
|
||||
const IdType k = key[tx];
|
||||
row[tx] = k >> col_bits;
|
||||
col[tx] = k & ((1 << col_bits) - 1);
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
void COOSort_(COOMatrix* coo, bool sort_column) {
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const int row_bits = cuda::_NumberOfBits(coo->num_rows);
|
||||
|
||||
const int64_t nnz = coo->row->shape[0];
|
||||
if (sort_column) {
|
||||
const int col_bits = cuda::_NumberOfBits(coo->num_cols);
|
||||
const int num_bits = row_bits + col_bits;
|
||||
|
||||
const int nt = 256;
|
||||
const int nb = (nnz + nt - 1) / nt;
|
||||
CHECK(static_cast<int64_t>(nb) * nt >= nnz);
|
||||
|
||||
IdArray pos = aten::NewIdArray(nnz, coo->row->ctx, coo->row->dtype.bits);
|
||||
|
||||
CUDA_KERNEL_CALL(
|
||||
_COOEncodeEdgesKernel, nb, nt, 0, stream, coo->row.Ptr<IdType>(),
|
||||
coo->col.Ptr<IdType>(), nnz, col_bits, pos.Ptr<IdType>());
|
||||
|
||||
auto sorted = Sort(pos, num_bits);
|
||||
|
||||
CUDA_KERNEL_CALL(
|
||||
_COODecodeEdgesKernel, nb, nt, 0, stream, sorted.first.Ptr<IdType>(),
|
||||
nnz, col_bits, coo->row.Ptr<IdType>(), coo->col.Ptr<IdType>());
|
||||
|
||||
if (aten::COOHasData(*coo))
|
||||
coo->data = IndexSelect(coo->data, sorted.second);
|
||||
else
|
||||
coo->data = AsNumBits(sorted.second, coo->row->dtype.bits);
|
||||
coo->row_sorted = coo->col_sorted = true;
|
||||
} else {
|
||||
const int num_bits = row_bits;
|
||||
|
||||
auto sorted = Sort(coo->row, num_bits);
|
||||
|
||||
coo->row = sorted.first;
|
||||
coo->col = IndexSelect(coo->col, sorted.second);
|
||||
|
||||
if (aten::COOHasData(*coo))
|
||||
coo->data = IndexSelect(coo->data, sorted.second);
|
||||
else
|
||||
coo->data = AsNumBits(sorted.second, coo->row->dtype.bits);
|
||||
coo->row_sorted = true;
|
||||
}
|
||||
}
|
||||
|
||||
template void COOSort_<kDGLCUDA, int32_t>(COOMatrix* coo, bool sort_column);
|
||||
template void COOSort_<kDGLCUDA, int64_t>(COOMatrix* coo, bool sort_column);
|
||||
|
||||
///////////////////////////// COOIsSorted /////////////////////////////
|
||||
|
||||
template <typename IdType>
|
||||
__global__ void _COOIsSortedKernel(
|
||||
const IdType* row, const IdType* col, int64_t nnz, int8_t* row_sorted,
|
||||
int8_t* col_sorted) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < nnz) {
|
||||
if (tx == 0) {
|
||||
row_sorted[0] = 1;
|
||||
col_sorted[0] = 1;
|
||||
} else {
|
||||
row_sorted[tx] = static_cast<int8_t>(row[tx - 1] <= row[tx]);
|
||||
col_sorted[tx] =
|
||||
static_cast<int8_t>(row[tx - 1] < row[tx] || col[tx - 1] <= col[tx]);
|
||||
}
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::pair<bool, bool> COOIsSorted(COOMatrix coo) {
|
||||
const int64_t nnz = coo.row->shape[0];
|
||||
const auto& ctx = coo.row->ctx;
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
// We allocate a workspace of 2*nnz bytes. It wastes a little bit memory but
|
||||
// should be fine.
|
||||
int8_t* row_flags = static_cast<int8_t*>(device->AllocWorkspace(ctx, nnz));
|
||||
int8_t* col_flags = static_cast<int8_t*>(device->AllocWorkspace(ctx, nnz));
|
||||
const int nt = cuda::FindNumThreads(nnz);
|
||||
const int nb = (nnz + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
_COOIsSortedKernel, nb, nt, 0, stream, coo.row.Ptr<IdType>(),
|
||||
coo.col.Ptr<IdType>(), nnz, row_flags, col_flags);
|
||||
|
||||
const bool row_sorted = cuda::AllTrue(row_flags, nnz, ctx);
|
||||
const bool col_sorted =
|
||||
row_sorted ? cuda::AllTrue(col_flags, nnz, ctx) : false;
|
||||
|
||||
device->FreeWorkspace(ctx, row_flags);
|
||||
device->FreeWorkspace(ctx, col_flags);
|
||||
|
||||
return {row_sorted, col_sorted};
|
||||
}
|
||||
|
||||
template std::pair<bool, bool> COOIsSorted<kDGLCUDA, int32_t>(COOMatrix coo);
|
||||
template std::pair<bool, bool> COOIsSorted<kDGLCUDA, int64_t>(COOMatrix coo);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,183 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/csr2coo.cc
|
||||
* @brief CSR2COO
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
#include <thrust/iterator/constant_iterator.h>
|
||||
#include <thrust/iterator/counting_iterator.h>
|
||||
#include <thrust/iterator/transform_iterator.h>
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using runtime::NDArray;
|
||||
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix CSRToCOO(CSRMatrix csr) {
|
||||
LOG(FATAL) << "Unreachable codes";
|
||||
return {};
|
||||
}
|
||||
|
||||
template <>
|
||||
COOMatrix CSRToCOO<kDGLCUDA, int32_t>(CSRMatrix csr) {
|
||||
auto* thr_entry = runtime::CUDAThreadEntry::ThreadLocal();
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
// allocate cusparse handle if needed
|
||||
if (!thr_entry->cusparse_handle) {
|
||||
CUSPARSE_CALL(cusparseCreate(&(thr_entry->cusparse_handle)));
|
||||
}
|
||||
CUSPARSE_CALL(cusparseSetStream(thr_entry->cusparse_handle, stream));
|
||||
|
||||
NDArray indptr = csr.indptr, indices = csr.indices, data = csr.data;
|
||||
const int32_t* indptr_ptr = static_cast<int32_t*>(indptr->data);
|
||||
NDArray row =
|
||||
aten::NewIdArray(indices->shape[0], indptr->ctx, indptr->dtype.bits);
|
||||
int32_t* row_ptr = static_cast<int32_t*>(row->data);
|
||||
|
||||
CUSPARSE_CALL(cusparseXcsr2coo(
|
||||
thr_entry->cusparse_handle, indptr_ptr, indices->shape[0], csr.num_rows,
|
||||
row_ptr, CUSPARSE_INDEX_BASE_ZERO));
|
||||
|
||||
return COOMatrix(
|
||||
csr.num_rows, csr.num_cols, row, indices, data, true, csr.sorted);
|
||||
}
|
||||
|
||||
struct RepeatIndex {
|
||||
template <typename IdType>
|
||||
__host__ __device__ auto operator()(IdType i) {
|
||||
return thrust::make_constant_iterator(i);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename IdType>
|
||||
struct OutputBufferIndexer {
|
||||
const IdType* indptr;
|
||||
IdType* buffer;
|
||||
__host__ __device__ auto operator()(IdType i) { return buffer + indptr[i]; }
|
||||
};
|
||||
|
||||
template <typename IdType>
|
||||
struct AdjacentDifference {
|
||||
const IdType* indptr;
|
||||
__host__ __device__ auto operator()(IdType i) {
|
||||
return indptr[i + 1] - indptr[i];
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
COOMatrix CSRToCOO<kDGLCUDA, int64_t>(CSRMatrix csr) {
|
||||
const auto& ctx = csr.indptr->ctx;
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
|
||||
const int64_t nnz = csr.indices->shape[0];
|
||||
const auto nbits = csr.indptr->dtype.bits;
|
||||
IdArray ret_row = NewIdArray(nnz, ctx, nbits);
|
||||
|
||||
runtime::CUDAWorkspaceAllocator allocator(csr.indptr->ctx);
|
||||
thrust::counting_iterator<int64_t> iota(0);
|
||||
|
||||
auto input_buffer = thrust::make_transform_iterator(iota, RepeatIndex{});
|
||||
auto output_buffer = thrust::make_transform_iterator(
|
||||
iota, OutputBufferIndexer<int64_t>{
|
||||
csr.indptr.Ptr<int64_t>(), ret_row.Ptr<int64_t>()});
|
||||
auto buffer_sizes = thrust::make_transform_iterator(
|
||||
iota, AdjacentDifference<int64_t>{csr.indptr.Ptr<int64_t>()});
|
||||
|
||||
constexpr int64_t max_copy_at_once = std::numeric_limits<int32_t>::max();
|
||||
for (int64_t i = 0; i < csr.num_rows; i += max_copy_at_once) {
|
||||
std::size_t temp_storage_bytes = 0;
|
||||
CUDA_CALL(cub::DeviceCopy::Batched(
|
||||
nullptr, temp_storage_bytes, input_buffer + i, output_buffer + i,
|
||||
buffer_sizes + i, std::min(csr.num_rows - i, max_copy_at_once),
|
||||
stream));
|
||||
|
||||
auto temp = allocator.alloc_unique<char>(temp_storage_bytes);
|
||||
|
||||
CUDA_CALL(cub::DeviceCopy::Batched(
|
||||
temp.get(), temp_storage_bytes, input_buffer + i, output_buffer + i,
|
||||
buffer_sizes + i, std::min(csr.num_rows - i, max_copy_at_once),
|
||||
stream));
|
||||
}
|
||||
|
||||
return COOMatrix(
|
||||
csr.num_rows, csr.num_cols, ret_row, csr.indices, csr.data, true,
|
||||
csr.sorted);
|
||||
}
|
||||
|
||||
template COOMatrix CSRToCOO<kDGLCUDA, int32_t>(CSRMatrix csr);
|
||||
template COOMatrix CSRToCOO<kDGLCUDA, int64_t>(CSRMatrix csr);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix CSRToCOODataAsOrder(CSRMatrix csr) {
|
||||
LOG(FATAL) << "Unreachable codes";
|
||||
return {};
|
||||
}
|
||||
|
||||
template <>
|
||||
COOMatrix CSRToCOODataAsOrder<kDGLCUDA, int32_t>(CSRMatrix csr) {
|
||||
COOMatrix coo = CSRToCOO<kDGLCUDA, int32_t>(csr);
|
||||
if (aten::IsNullArray(coo.data)) return coo;
|
||||
|
||||
auto* thr_entry = runtime::CUDAThreadEntry::ThreadLocal();
|
||||
auto device = runtime::DeviceAPI::Get(coo.row->ctx);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
// allocate cusparse handle if needed
|
||||
if (!thr_entry->cusparse_handle) {
|
||||
CUSPARSE_CALL(cusparseCreate(&(thr_entry->cusparse_handle)));
|
||||
}
|
||||
CUSPARSE_CALL(cusparseSetStream(thr_entry->cusparse_handle, stream));
|
||||
|
||||
NDArray row = coo.row, col = coo.col, data = coo.data;
|
||||
int32_t* row_ptr = static_cast<int32_t*>(row->data);
|
||||
int32_t* col_ptr = static_cast<int32_t*>(col->data);
|
||||
int32_t* data_ptr = static_cast<int32_t*>(data->data);
|
||||
|
||||
size_t workspace_size = 0;
|
||||
CUSPARSE_CALL(cusparseXcoosort_bufferSizeExt(
|
||||
thr_entry->cusparse_handle, coo.num_rows, coo.num_cols, row->shape[0],
|
||||
data_ptr, row_ptr, &workspace_size));
|
||||
void* workspace = device->AllocWorkspace(row->ctx, workspace_size);
|
||||
CUSPARSE_CALL(cusparseXcoosortByRow(
|
||||
thr_entry->cusparse_handle, coo.num_rows, coo.num_cols, row->shape[0],
|
||||
data_ptr, row_ptr, col_ptr, workspace));
|
||||
device->FreeWorkspace(row->ctx, workspace);
|
||||
|
||||
// The row and column field have already been reordered according
|
||||
// to data, thus the data field will be deprecated.
|
||||
coo.data = aten::NullArray();
|
||||
coo.row_sorted = false;
|
||||
coo.col_sorted = false;
|
||||
return coo;
|
||||
}
|
||||
|
||||
template <>
|
||||
COOMatrix CSRToCOODataAsOrder<kDGLCUDA, int64_t>(CSRMatrix csr) {
|
||||
COOMatrix coo = CSRToCOO<kDGLCUDA, int64_t>(csr);
|
||||
if (aten::IsNullArray(coo.data)) return coo;
|
||||
const auto& sorted = Sort(coo.data);
|
||||
|
||||
coo.row = IndexSelect(coo.row, sorted.second);
|
||||
coo.col = IndexSelect(coo.col, sorted.second);
|
||||
|
||||
// The row and column field have already been reordered according
|
||||
// to data, thus the data field will be deprecated.
|
||||
coo.data = aten::NullArray();
|
||||
coo.row_sorted = false;
|
||||
coo.col_sorted = false;
|
||||
return coo;
|
||||
}
|
||||
|
||||
template COOMatrix CSRToCOODataAsOrder<kDGLCUDA, int32_t>(CSRMatrix csr);
|
||||
template COOMatrix CSRToCOODataAsOrder<kDGLCUDA, int64_t>(CSRMatrix csr);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,100 @@
|
||||
/**
|
||||
* Copyright (c) 2021 by Contributors
|
||||
* @file array/cuda/csr_get_data.cu
|
||||
* @brief Retrieve entries of a CSR matrix
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include <numeric>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using runtime::NDArray;
|
||||
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename DType>
|
||||
NDArray CSRGetData(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, DType filler) {
|
||||
const int64_t rowlen = rows->shape[0];
|
||||
const int64_t collen = cols->shape[0];
|
||||
|
||||
CHECK((rowlen == collen) || (rowlen == 1) || (collen == 1))
|
||||
<< "Invalid row and col id array.";
|
||||
|
||||
const int64_t row_stride = (rowlen == 1 && collen != 1) ? 0 : 1;
|
||||
const int64_t col_stride = (collen == 1 && rowlen != 1) ? 0 : 1;
|
||||
|
||||
const int64_t rstlen = std::max(rowlen, collen);
|
||||
IdArray rst = NDArray::Empty({rstlen}, weights->dtype, rows->ctx);
|
||||
if (rstlen == 0) return rst;
|
||||
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const int nt = cuda::FindNumThreads(rstlen);
|
||||
const int nb = (rstlen + nt - 1) / nt;
|
||||
if (return_eids)
|
||||
BUG_IF_FAIL(DGLDataTypeTraits<DType>::dtype == rows->dtype)
|
||||
<< "DType does not match row's dtype.";
|
||||
|
||||
const IdType* indptr_data =
|
||||
static_cast<IdType*>(cuda::GetDevicePointer(csr.indptr));
|
||||
const IdType* indices_data =
|
||||
static_cast<IdType*>(cuda::GetDevicePointer(csr.indices));
|
||||
const IdType* data_data =
|
||||
CSRHasData(csr) ? static_cast<IdType*>(cuda::GetDevicePointer(csr.data))
|
||||
: nullptr;
|
||||
|
||||
// TODO(minjie): use binary search for sorted csr
|
||||
CUDA_KERNEL_CALL(
|
||||
cuda::_LinearSearchKernel, nb, nt, 0, stream, indptr_data, indices_data,
|
||||
data_data, rows.Ptr<IdType>(), cols.Ptr<IdType>(), row_stride, col_stride,
|
||||
rstlen, return_eids ? nullptr : weights.Ptr<DType>(), filler,
|
||||
rst.Ptr<DType>());
|
||||
return rst;
|
||||
}
|
||||
|
||||
template NDArray CSRGetData<kDGLCUDA, int32_t, __half>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, __half filler);
|
||||
template NDArray CSRGetData<kDGLCUDA, int64_t, __half>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, __half filler);
|
||||
#if BF16_ENABLED
|
||||
template NDArray CSRGetData<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, __nv_bfloat16 filler);
|
||||
template NDArray CSRGetData<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, __nv_bfloat16 filler);
|
||||
#endif // BF16_ENABLED
|
||||
template NDArray CSRGetData<kDGLCUDA, int32_t, float>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, float filler);
|
||||
template NDArray CSRGetData<kDGLCUDA, int64_t, float>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, float filler);
|
||||
template NDArray CSRGetData<kDGLCUDA, int32_t, double>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, double filler);
|
||||
template NDArray CSRGetData<kDGLCUDA, int64_t, double>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, double filler);
|
||||
|
||||
// For CSRGetData<XPU, IdType>(CSRMatrix, NDArray, NDArray)
|
||||
template NDArray CSRGetData<kDGLCUDA, int32_t, int32_t>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, int32_t filler);
|
||||
template NDArray CSRGetData<kDGLCUDA, int64_t, int64_t>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols, bool return_eids,
|
||||
NDArray weights, int64_t filler);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,332 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/csr_mm.cu
|
||||
* @brief SpSpMM/SpGEMM C APIs and definitions.
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/device_api.h>
|
||||
|
||||
#include <limits>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./cusparse_dispatcher.cuh"
|
||||
#include "./functor.cuh"
|
||||
namespace dgl {
|
||||
|
||||
using namespace dgl::runtime;
|
||||
|
||||
namespace aten {
|
||||
namespace cusparse {
|
||||
|
||||
#if CUDART_VERSION >= 12000
|
||||
|
||||
/** @brief Cusparse implementation of SpGEMM on Csr format for CUDA 12.0+ */
|
||||
template <typename DType, typename IdType>
|
||||
std::pair<CSRMatrix, NDArray> CusparseSpgemm(
|
||||
const CSRMatrix& A, const NDArray A_weights_array, const CSRMatrix& B,
|
||||
const NDArray B_weights_array) {
|
||||
// We use Spgemm (SpSpMM) to perform following operation:
|
||||
// C = A x B, where A, B and C are sparse matrices in csr format.
|
||||
const int nnzA = A.indices->shape[0];
|
||||
const int nnzB = B.indices->shape[0];
|
||||
const DType alpha = 1.0;
|
||||
const DType beta = 0.0;
|
||||
auto transA = CUSPARSE_OPERATION_NON_TRANSPOSE;
|
||||
auto transB = CUSPARSE_OPERATION_NON_TRANSPOSE;
|
||||
// device
|
||||
auto ctx = A.indptr->ctx;
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
auto* thr_entry = runtime::CUDAThreadEntry::ThreadLocal();
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const DType* A_weights = A_weights_array.Ptr<DType>();
|
||||
const DType* B_weights = B_weights_array.Ptr<DType>();
|
||||
// allocate cusparse handle if needed
|
||||
if (!thr_entry->cusparse_handle) {
|
||||
CUSPARSE_CALL(cusparseCreate(&(thr_entry->cusparse_handle)));
|
||||
}
|
||||
CUSPARSE_CALL(cusparseSetStream(thr_entry->cusparse_handle, stream));
|
||||
// all one data array
|
||||
cusparseSpMatDescr_t matA, matB, matC;
|
||||
IdArray dC_csrOffsets =
|
||||
IdArray::Empty({A.num_rows + 1}, A.indptr->dtype, A.indptr->ctx);
|
||||
IdType* dC_csrOffsets_data = dC_csrOffsets.Ptr<IdType>();
|
||||
constexpr auto idtype = cusparse_idtype<IdType>::value;
|
||||
constexpr auto dtype = cuda_dtype<DType>::value;
|
||||
// Create sparse matrix A, B and C in CSR format
|
||||
CUSPARSE_CALL(cusparseCreateCsr(
|
||||
&matA, A.num_rows, A.num_cols, nnzA, A.indptr.Ptr<IdType>(),
|
||||
A.indices.Ptr<IdType>(),
|
||||
// cusparseCreateCsr only accepts non-const pointers.
|
||||
const_cast<DType*>(A_weights), idtype, idtype, CUSPARSE_INDEX_BASE_ZERO,
|
||||
dtype));
|
||||
CUSPARSE_CALL(cusparseCreateCsr(
|
||||
&matB, B.num_rows, B.num_cols, nnzB, B.indptr.Ptr<IdType>(),
|
||||
B.indices.Ptr<IdType>(),
|
||||
// cusparseCreateCsr only accepts non-const pointers.
|
||||
const_cast<DType*>(B_weights), idtype, idtype, CUSPARSE_INDEX_BASE_ZERO,
|
||||
dtype));
|
||||
CUSPARSE_CALL(cusparseCreateCsr(
|
||||
&matC, A.num_rows, B.num_cols, 0, dC_csrOffsets_data, nullptr, nullptr,
|
||||
idtype, idtype, CUSPARSE_INDEX_BASE_ZERO, dtype));
|
||||
// SpGEMM Computation
|
||||
cusparseSpGEMMDescr_t spgemmDesc;
|
||||
cusparseSpGEMMAlg_t alg = CUSPARSE_SPGEMM_DEFAULT;
|
||||
|
||||
CUSPARSE_CALL(cusparseSpGEMM_createDescr(&spgemmDesc));
|
||||
size_t workspace_size1 = 0, workspace_size2 = 0, workspace_size3 = 0;
|
||||
// ask bufferSize1 bytes for external memory
|
||||
CUSPARSE_CALL(cusparseSpGEMM_workEstimation(
|
||||
thr_entry->cusparse_handle, transA, transB, &alpha, matA, matB, &beta,
|
||||
matC, dtype, alg, spgemmDesc, &workspace_size1, NULL));
|
||||
void* workspace1 = (device->AllocWorkspace(ctx, workspace_size1));
|
||||
// inspect the matrices A and B to understand the memory requiremnent
|
||||
cusparseStatus_t e = cusparseSpGEMM_workEstimation(
|
||||
thr_entry->cusparse_handle, transA, transB, &alpha, matA, matB, &beta,
|
||||
matC, dtype, alg, spgemmDesc, &workspace_size1, workspace1);
|
||||
// CUSPARSE_SPGEMM_DEFAULT not support getting num_prods > 2^31 -1
|
||||
// and throws insufficient memory error within workEstimation call
|
||||
if (e == CUSPARSE_STATUS_INSUFFICIENT_RESOURCES) {
|
||||
// fall back to ALG2 to estimate num_prods
|
||||
alg = CUSPARSE_SPGEMM_ALG2;
|
||||
device->FreeWorkspace(ctx, workspace1);
|
||||
// rerun cusparseSpGEMM_workEstimation
|
||||
CUSPARSE_CALL(cusparseSpGEMM_workEstimation(
|
||||
thr_entry->cusparse_handle, transA, transB, &alpha, matA, matB, &beta,
|
||||
matC, dtype, alg, spgemmDesc, &workspace_size1, NULL));
|
||||
workspace1 = (device->AllocWorkspace(ctx, workspace_size1));
|
||||
CUSPARSE_CALL(cusparseSpGEMM_workEstimation(
|
||||
thr_entry->cusparse_handle, transA, transB, &alpha, matA, matB, &beta,
|
||||
matC, dtype, alg, spgemmDesc, &workspace_size1, workspace1));
|
||||
} else {
|
||||
CHECK(e == CUSPARSE_STATUS_SUCCESS) << "CUSPARSE ERROR in SpGEMM: " << e;
|
||||
}
|
||||
|
||||
// get the number of intermediate products required for SpGEMM compute
|
||||
// num_prods indicates device memory consumption for SpGEMM if using ALG2/3
|
||||
int64_t num_prods;
|
||||
CUSPARSE_CALL(cusparseSpGEMM_getNumProducts(spgemmDesc, &num_prods));
|
||||
|
||||
// assume free GPU mem at least ~15G for below heuristics to work
|
||||
// user-defined medium problem size (below will use DEFAULT)
|
||||
int64_t MEDIUM_NUM_PRODUCTS = 400000000; // 400*1000*1000;
|
||||
// user-defined large problem size (above will use ALG3)
|
||||
int64_t LARGE_NUM_PRODUCTS = 800000000; // 800*1000*1000;
|
||||
|
||||
// switch to ALG2/ALG3 for medium & large problem size
|
||||
if (alg == CUSPARSE_SPGEMM_DEFAULT && num_prods > MEDIUM_NUM_PRODUCTS) {
|
||||
// use ALG3 for very large problem
|
||||
alg = num_prods > LARGE_NUM_PRODUCTS ? CUSPARSE_SPGEMM_ALG3
|
||||
: CUSPARSE_SPGEMM_ALG2;
|
||||
|
||||
device->FreeWorkspace(ctx, workspace1);
|
||||
// rerun cusparseSpGEMM_workEstimation
|
||||
CUSPARSE_CALL(cusparseSpGEMM_workEstimation(
|
||||
thr_entry->cusparse_handle, transA, transB, &alpha, matA, matB, &beta,
|
||||
matC, dtype, alg, spgemmDesc, &workspace_size1, NULL));
|
||||
workspace1 = (device->AllocWorkspace(ctx, workspace_size1));
|
||||
CUSPARSE_CALL(cusparseSpGEMM_workEstimation(
|
||||
thr_entry->cusparse_handle, transA, transB, &alpha, matA, matB, &beta,
|
||||
matC, dtype, alg, spgemmDesc, &workspace_size1, workspace1));
|
||||
} else if (alg == CUSPARSE_SPGEMM_ALG2 && num_prods > LARGE_NUM_PRODUCTS) {
|
||||
// no need to rerun cusparseSpGEMM_workEstimation between ALG2 and ALG3
|
||||
alg = CUSPARSE_SPGEMM_ALG3;
|
||||
}
|
||||
|
||||
if (alg == CUSPARSE_SPGEMM_ALG2 || alg == CUSPARSE_SPGEMM_ALG3) {
|
||||
// estimate memory for ALG2/ALG3; note chunk_fraction is only used by ALG3
|
||||
// reduce chunk_fraction if crash due to mem., but it trades off speed
|
||||
float chunk_fraction = num_prods < 4 * LARGE_NUM_PRODUCTS ? 0.15 : 0.05;
|
||||
CUSPARSE_CALL(cusparseSpGEMM_estimateMemory(
|
||||
thr_entry->cusparse_handle, transA, transB, &alpha, matA, matB, &beta,
|
||||
matC, dtype, alg, spgemmDesc, chunk_fraction, &workspace_size3, NULL,
|
||||
NULL));
|
||||
void* workspace3 = (device->AllocWorkspace(ctx, workspace_size3));
|
||||
CUSPARSE_CALL(cusparseSpGEMM_estimateMemory(
|
||||
thr_entry->cusparse_handle, transA, transB, &alpha, matA, matB, &beta,
|
||||
matC, dtype, alg, spgemmDesc, chunk_fraction, &workspace_size3,
|
||||
workspace3, &workspace_size2));
|
||||
device->FreeWorkspace(ctx, workspace3);
|
||||
} else {
|
||||
CUSPARSE_CALL(cusparseSpGEMM_compute(
|
||||
thr_entry->cusparse_handle, transA, transB, &alpha, matA, matB, &beta,
|
||||
matC, dtype, alg, spgemmDesc, &workspace_size2, NULL));
|
||||
}
|
||||
// ask bufferSize2 bytes for external memory
|
||||
void* workspace2 = device->AllocWorkspace(ctx, workspace_size2);
|
||||
// compute the intermediate product of A * B
|
||||
CUSPARSE_CALL(cusparseSpGEMM_compute(
|
||||
thr_entry->cusparse_handle, transA, transB, &alpha, matA, matB, &beta,
|
||||
matC, dtype, alg, spgemmDesc, &workspace_size2, workspace2));
|
||||
// get matrix C non-zero entries C_nnz1
|
||||
int64_t C_num_rows1, C_num_cols1, C_nnz1;
|
||||
CUSPARSE_CALL(
|
||||
cusparseSpMatGetSize(matC, &C_num_rows1, &C_num_cols1, &C_nnz1));
|
||||
IdArray dC_columns = IdArray::Empty({C_nnz1}, A.indptr->dtype, A.indptr->ctx);
|
||||
NDArray dC_weights =
|
||||
NDArray::Empty({C_nnz1}, A_weights_array->dtype, A.indptr->ctx);
|
||||
IdType* dC_columns_data = dC_columns.Ptr<IdType>();
|
||||
DType* dC_weights_data = dC_weights.Ptr<DType>();
|
||||
// update matC with the new pointers
|
||||
CUSPARSE_CALL(cusparseCsrSetPointers(
|
||||
matC, dC_csrOffsets_data, dC_columns_data, dC_weights_data));
|
||||
// copy the final products to the matrix C
|
||||
CUSPARSE_CALL(cusparseSpGEMM_copy(
|
||||
thr_entry->cusparse_handle, transA, transB, &alpha, matA, matB, &beta,
|
||||
matC, dtype, alg, spgemmDesc));
|
||||
|
||||
device->FreeWorkspace(ctx, workspace1);
|
||||
device->FreeWorkspace(ctx, workspace2);
|
||||
// destroy matrix/vector descriptors
|
||||
CUSPARSE_CALL(cusparseSpGEMM_destroyDescr(spgemmDesc));
|
||||
CUSPARSE_CALL(cusparseDestroySpMat(matA));
|
||||
CUSPARSE_CALL(cusparseDestroySpMat(matB));
|
||||
CUSPARSE_CALL(cusparseDestroySpMat(matC));
|
||||
return {
|
||||
CSRMatrix(
|
||||
A.num_rows, B.num_cols, dC_csrOffsets, dC_columns,
|
||||
NullArray(dC_csrOffsets->dtype, dC_csrOffsets->ctx)),
|
||||
dC_weights};
|
||||
}
|
||||
|
||||
#else // CUDART_VERSION < 12000
|
||||
|
||||
/** @brief Cusparse implementation of SpGEMM on Csr format for older CUDA
|
||||
* versions */
|
||||
template <typename DType, typename IdType>
|
||||
std::pair<CSRMatrix, NDArray> CusparseSpgemm(
|
||||
const CSRMatrix& A, const NDArray A_weights_array, const CSRMatrix& B,
|
||||
const NDArray B_weights_array) {
|
||||
int nnzC;
|
||||
csrgemm2Info_t info = nullptr;
|
||||
size_t workspace_size;
|
||||
const DType alpha = 1.;
|
||||
const int nnzA = A.indices->shape[0];
|
||||
const int nnzB = B.indices->shape[0];
|
||||
const int m = A.num_rows;
|
||||
const int n = A.num_cols;
|
||||
const int k = B.num_cols;
|
||||
auto ctx = A.indptr->ctx;
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
auto* thr_entry = runtime::CUDAThreadEntry::ThreadLocal();
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
auto idtype = A.indptr->dtype;
|
||||
auto dtype = A_weights_array->dtype;
|
||||
const DType* A_weights = A_weights_array.Ptr<DType>();
|
||||
const DType* B_weights = B_weights_array.Ptr<DType>();
|
||||
if (!thr_entry->cusparse_handle) {
|
||||
CUSPARSE_CALL(cusparseCreate(&(thr_entry->cusparse_handle)));
|
||||
}
|
||||
CUSPARSE_CALL(cusparseSetStream(thr_entry->cusparse_handle, stream));
|
||||
CUSPARSE_CALL(cusparseSetPointerMode(
|
||||
thr_entry->cusparse_handle, CUSPARSE_POINTER_MODE_HOST));
|
||||
|
||||
CUSPARSE_CALL(cusparseCreateCsrgemm2Info(&info));
|
||||
|
||||
cusparseMatDescr_t matA, matB, matC, matD;
|
||||
CUSPARSE_CALL(cusparseCreateMatDescr(&matA));
|
||||
CUSPARSE_CALL(cusparseCreateMatDescr(&matB));
|
||||
CUSPARSE_CALL(cusparseCreateMatDescr(&matC));
|
||||
CUSPARSE_CALL(cusparseCreateMatDescr(&matD)); // needed even if D is null
|
||||
|
||||
CUSPARSE_CALL(CSRGEMM<DType>::bufferSizeExt(
|
||||
thr_entry->cusparse_handle, m, n, k, &alpha, matA, nnzA,
|
||||
A.indptr.Ptr<IdType>(), A.indices.Ptr<IdType>(), matB, nnzB,
|
||||
B.indptr.Ptr<IdType>(), B.indices.Ptr<IdType>(), nullptr, matD, 0,
|
||||
nullptr, nullptr, info, &workspace_size));
|
||||
|
||||
void* workspace = device->AllocWorkspace(ctx, workspace_size);
|
||||
IdArray C_indptr = IdArray::Empty({m + 1}, idtype, ctx);
|
||||
CUSPARSE_CALL(CSRGEMM<DType>::nnz(
|
||||
thr_entry->cusparse_handle, m, n, k, matA, nnzA, A.indptr.Ptr<IdType>(),
|
||||
A.indices.Ptr<IdType>(), matB, nnzB, B.indptr.Ptr<IdType>(),
|
||||
B.indices.Ptr<IdType>(), matD, 0, nullptr, nullptr, matC,
|
||||
C_indptr.Ptr<IdType>(), &nnzC, info, workspace));
|
||||
|
||||
IdArray C_indices = IdArray::Empty({nnzC}, idtype, ctx);
|
||||
NDArray C_weights = NDArray::Empty({nnzC}, dtype, ctx);
|
||||
CUSPARSE_CALL(CSRGEMM<DType>::compute(
|
||||
thr_entry->cusparse_handle, m, n, k, &alpha, matA, nnzA, A_weights,
|
||||
A.indptr.Ptr<IdType>(), A.indices.Ptr<IdType>(), matB, nnzB, B_weights,
|
||||
B.indptr.Ptr<IdType>(), B.indices.Ptr<IdType>(), nullptr, matD, 0,
|
||||
nullptr, nullptr, nullptr, matC, C_weights.Ptr<DType>(),
|
||||
C_indptr.Ptr<IdType>(), C_indices.Ptr<IdType>(), info, workspace));
|
||||
|
||||
device->FreeWorkspace(ctx, workspace);
|
||||
CUSPARSE_CALL(cusparseDestroyCsrgemm2Info(info));
|
||||
CUSPARSE_CALL(cusparseDestroyMatDescr(matA));
|
||||
CUSPARSE_CALL(cusparseDestroyMatDescr(matB));
|
||||
CUSPARSE_CALL(cusparseDestroyMatDescr(matC));
|
||||
CUSPARSE_CALL(cusparseDestroyMatDescr(matD));
|
||||
|
||||
return {
|
||||
CSRMatrix(
|
||||
m, k, C_indptr, C_indices, NullArray(C_indptr->dtype, C_indptr->ctx)),
|
||||
C_weights};
|
||||
}
|
||||
|
||||
#endif // CUDART_VERSION >= 12000
|
||||
} // namespace cusparse
|
||||
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
std::pair<CSRMatrix, NDArray> CSRMM(
|
||||
const CSRMatrix& A, NDArray A_weights, const CSRMatrix& B,
|
||||
NDArray B_weights) {
|
||||
auto ctx = A.indptr->ctx;
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
CSRMatrix newA, newB;
|
||||
bool cast = false;
|
||||
|
||||
// Cast 64 bit indices to 32 bit.
|
||||
if (A.indptr->dtype.bits == 64) {
|
||||
newA = CSRMatrix(
|
||||
A.num_rows, A.num_cols, AsNumBits(A.indptr, 32),
|
||||
AsNumBits(A.indices, 32), AsNumBits(A.data, 32));
|
||||
newB = CSRMatrix(
|
||||
B.num_rows, B.num_cols, AsNumBits(B.indptr, 32),
|
||||
AsNumBits(B.indices, 32), AsNumBits(B.data, 32));
|
||||
cast = true;
|
||||
}
|
||||
|
||||
// Reorder weights if A or B has edge IDs
|
||||
NDArray newA_weights, newB_weights;
|
||||
if (CSRHasData(A)) newA_weights = IndexSelect(A_weights, A.data);
|
||||
if (CSRHasData(B)) newB_weights = IndexSelect(B_weights, B.data);
|
||||
|
||||
auto result = cusparse::CusparseSpgemm<DType, int32_t>(
|
||||
cast ? newA : A, CSRHasData(A) ? newA_weights : A_weights,
|
||||
cast ? newB : B, CSRHasData(B) ? newB_weights : B_weights);
|
||||
|
||||
// Cast 32 bit indices back to 64 bit if necessary
|
||||
if (cast) {
|
||||
CSRMatrix C = result.first;
|
||||
return {
|
||||
CSRMatrix(
|
||||
C.num_rows, C.num_cols, AsNumBits(C.indptr, 64),
|
||||
AsNumBits(C.indices, 64), AsNumBits(C.data, 64)),
|
||||
result.second};
|
||||
} else {
|
||||
return result;
|
||||
}
|
||||
}
|
||||
|
||||
template std::pair<CSRMatrix, NDArray> CSRMM<kDGLCUDA, int32_t, __half>(
|
||||
const CSRMatrix&, NDArray, const CSRMatrix&, NDArray);
|
||||
template std::pair<CSRMatrix, NDArray> CSRMM<kDGLCUDA, int64_t, __half>(
|
||||
const CSRMatrix&, NDArray, const CSRMatrix&, NDArray);
|
||||
#if BF16_ENABLED
|
||||
template std::pair<CSRMatrix, NDArray> CSRMM<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
const CSRMatrix&, NDArray, const CSRMatrix&, NDArray);
|
||||
template std::pair<CSRMatrix, NDArray> CSRMM<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
const CSRMatrix&, NDArray, const CSRMatrix&, NDArray);
|
||||
#endif // BF16_ENABLED
|
||||
template std::pair<CSRMatrix, NDArray> CSRMM<kDGLCUDA, int32_t, float>(
|
||||
const CSRMatrix&, NDArray, const CSRMatrix&, NDArray);
|
||||
template std::pair<CSRMatrix, NDArray> CSRMM<kDGLCUDA, int64_t, float>(
|
||||
const CSRMatrix&, NDArray, const CSRMatrix&, NDArray);
|
||||
template std::pair<CSRMatrix, NDArray> CSRMM<kDGLCUDA, int32_t, double>(
|
||||
const CSRMatrix&, NDArray, const CSRMatrix&, NDArray);
|
||||
template std::pair<CSRMatrix, NDArray> CSRMM<kDGLCUDA, int64_t, double>(
|
||||
const CSRMatrix&, NDArray, const CSRMatrix&, NDArray);
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,151 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/csr_sort.cc
|
||||
* @brief Sort CSR index
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using runtime::NDArray;
|
||||
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
/**
|
||||
* @brief Check whether each row is sorted.
|
||||
*/
|
||||
template <typename IdType>
|
||||
__global__ void _SegmentIsSorted(
|
||||
const IdType* indptr, const IdType* indices, int64_t num_rows,
|
||||
int8_t* flags) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < num_rows) {
|
||||
bool f = true;
|
||||
for (IdType i = indptr[tx] + 1; f && i < indptr[tx + 1]; ++i) {
|
||||
f = (indices[i - 1] <= indices[i]);
|
||||
}
|
||||
flags[tx] = static_cast<int8_t>(f);
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
bool CSRIsSorted(CSRMatrix csr) {
|
||||
const auto& ctx = csr.indptr->ctx;
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
// We allocate a workspace of num_rows bytes. It wastes a little bit memory
|
||||
// but should be fine.
|
||||
int8_t* flags =
|
||||
static_cast<int8_t*>(device->AllocWorkspace(ctx, csr.num_rows));
|
||||
const int nt = cuda::FindNumThreads(csr.num_rows);
|
||||
const int nb = (csr.num_rows + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
_SegmentIsSorted, nb, nt, 0, stream, csr.indptr.Ptr<IdType>(),
|
||||
csr.indices.Ptr<IdType>(), csr.num_rows, flags);
|
||||
bool ret = cuda::AllTrue(flags, csr.num_rows, ctx);
|
||||
device->FreeWorkspace(ctx, flags);
|
||||
return ret;
|
||||
}
|
||||
|
||||
template bool CSRIsSorted<kDGLCUDA, int32_t>(CSRMatrix csr);
|
||||
template bool CSRIsSorted<kDGLCUDA, int64_t>(CSRMatrix csr);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
void CSRSort_(CSRMatrix* csr) {
|
||||
LOG(FATAL) << "Unreachable codes";
|
||||
}
|
||||
|
||||
template <>
|
||||
void CSRSort_<kDGLCUDA, int32_t>(CSRMatrix* csr) {
|
||||
auto* thr_entry = runtime::CUDAThreadEntry::ThreadLocal();
|
||||
auto device = runtime::DeviceAPI::Get(csr->indptr->ctx);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
// allocate cusparse handle if needed
|
||||
if (!thr_entry->cusparse_handle) {
|
||||
CUSPARSE_CALL(cusparseCreate(&(thr_entry->cusparse_handle)));
|
||||
}
|
||||
CUSPARSE_CALL(cusparseSetStream(thr_entry->cusparse_handle, stream));
|
||||
|
||||
NDArray indptr = csr->indptr;
|
||||
NDArray indices = csr->indices;
|
||||
const auto& ctx = indptr->ctx;
|
||||
const int64_t nnz = indices->shape[0];
|
||||
if (!aten::CSRHasData(*csr))
|
||||
csr->data = aten::Range(0, nnz, indices->dtype.bits, ctx);
|
||||
NDArray data = csr->data;
|
||||
|
||||
size_t workspace_size = 0;
|
||||
CUSPARSE_CALL(cusparseXcsrsort_bufferSizeExt(
|
||||
thr_entry->cusparse_handle, csr->num_rows, csr->num_cols, nnz,
|
||||
indptr.Ptr<int32_t>(), indices.Ptr<int32_t>(), &workspace_size));
|
||||
void* workspace = device->AllocWorkspace(ctx, workspace_size);
|
||||
|
||||
cusparseMatDescr_t descr;
|
||||
CUSPARSE_CALL(cusparseCreateMatDescr(&descr));
|
||||
CUSPARSE_CALL(cusparseSetMatType(descr, CUSPARSE_MATRIX_TYPE_GENERAL));
|
||||
CUSPARSE_CALL(cusparseSetMatIndexBase(descr, CUSPARSE_INDEX_BASE_ZERO));
|
||||
CUSPARSE_CALL(cusparseXcsrsort(
|
||||
thr_entry->cusparse_handle, csr->num_rows, csr->num_cols, nnz, descr,
|
||||
indptr.Ptr<int32_t>(), indices.Ptr<int32_t>(), data.Ptr<int32_t>(),
|
||||
workspace));
|
||||
|
||||
csr->sorted = true;
|
||||
|
||||
// free resources
|
||||
CUSPARSE_CALL(cusparseDestroyMatDescr(descr));
|
||||
device->FreeWorkspace(ctx, workspace);
|
||||
}
|
||||
|
||||
template <>
|
||||
void CSRSort_<kDGLCUDA, int64_t>(CSRMatrix* csr) {
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
auto device = runtime::DeviceAPI::Get(csr->indptr->ctx);
|
||||
|
||||
const auto& ctx = csr->indptr->ctx;
|
||||
const int64_t nnz = csr->indices->shape[0];
|
||||
const auto nbits = csr->indptr->dtype.bits;
|
||||
if (!aten::CSRHasData(*csr)) csr->data = aten::Range(0, nnz, nbits, ctx);
|
||||
|
||||
IdArray new_indices = csr->indices.Clone();
|
||||
IdArray new_data = csr->data.Clone();
|
||||
|
||||
const int64_t* offsets = csr->indptr.Ptr<int64_t>();
|
||||
const int64_t* key_in = csr->indices.Ptr<int64_t>();
|
||||
int64_t* key_out = new_indices.Ptr<int64_t>();
|
||||
const int64_t* value_in = csr->data.Ptr<int64_t>();
|
||||
int64_t* value_out = new_data.Ptr<int64_t>();
|
||||
|
||||
// Allocate workspace
|
||||
size_t workspace_size = 0;
|
||||
CUDA_CALL(cub::DeviceSegmentedRadixSort::SortPairs(
|
||||
nullptr, workspace_size, key_in, key_out, value_in, value_out, nnz,
|
||||
csr->num_rows, offsets, offsets + 1, 0, sizeof(int64_t) * 8, stream));
|
||||
void* workspace = device->AllocWorkspace(ctx, workspace_size);
|
||||
|
||||
// Compute
|
||||
CUDA_CALL(cub::DeviceSegmentedRadixSort::SortPairs(
|
||||
workspace, workspace_size, key_in, key_out, value_in, value_out, nnz,
|
||||
csr->num_rows, offsets, offsets + 1, 0, sizeof(int64_t) * 8, stream));
|
||||
|
||||
csr->sorted = true;
|
||||
csr->indices = new_indices;
|
||||
csr->data = new_data;
|
||||
|
||||
// free resources
|
||||
device->FreeWorkspace(ctx, workspace);
|
||||
}
|
||||
|
||||
template void CSRSort_<kDGLCUDA, int32_t>(CSRMatrix* csr);
|
||||
template void CSRSort_<kDGLCUDA, int64_t>(CSRMatrix* csr);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,177 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/spmm.cu
|
||||
* @brief SpGEAM C APIs and definitions.
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/device_api.h>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./cusparse_dispatcher.cuh"
|
||||
#include "./functor.cuh"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using namespace dgl::runtime;
|
||||
|
||||
namespace aten {
|
||||
namespace cusparse {
|
||||
|
||||
/** Cusparse implementation of SpSum on Csr format. */
|
||||
template <typename DType, typename IdType>
|
||||
std::pair<CSRMatrix, NDArray> CusparseCsrgeam2(
|
||||
const CSRMatrix& A, const NDArray A_weights_array, const CSRMatrix& B,
|
||||
const NDArray B_weights_array) {
|
||||
const int m = A.num_rows;
|
||||
const int n = A.num_cols;
|
||||
const int nnzA = A.indices->shape[0];
|
||||
const int nnzB = B.indices->shape[0];
|
||||
int nnzC;
|
||||
const DType alpha = 1.0;
|
||||
const DType beta = 1.0;
|
||||
auto ctx = A.indptr->ctx;
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
auto* thr_entry = runtime::CUDAThreadEntry::ThreadLocal();
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const DType* A_weights = A_weights_array.Ptr<DType>();
|
||||
const DType* B_weights = B_weights_array.Ptr<DType>();
|
||||
// allocate cusparse handle if needed
|
||||
if (!thr_entry->cusparse_handle)
|
||||
CUSPARSE_CALL(cusparseCreate(&(thr_entry->cusparse_handle)));
|
||||
CUSPARSE_CALL(cusparseSetStream(thr_entry->cusparse_handle, stream));
|
||||
|
||||
cusparseMatDescr_t matA, matB, matC;
|
||||
CUSPARSE_CALL(cusparseCreateMatDescr(&matA));
|
||||
CUSPARSE_CALL(cusparseCreateMatDescr(&matB));
|
||||
CUSPARSE_CALL(cusparseCreateMatDescr(&matC));
|
||||
|
||||
cusparseSetPointerMode(
|
||||
thr_entry->cusparse_handle, CUSPARSE_POINTER_MODE_HOST);
|
||||
size_t workspace_size = 0;
|
||||
/* prepare output C */
|
||||
IdArray dC_csrOffsets = IdArray::Empty({m + 1}, A.indptr->dtype, ctx);
|
||||
IdType* dC_csrOffsets_data = dC_csrOffsets.Ptr<IdType>();
|
||||
IdArray dC_columns;
|
||||
NDArray dC_weights;
|
||||
IdType* dC_columns_data = dC_columns.Ptr<IdType>();
|
||||
DType* dC_weights_data = dC_weights.Ptr<DType>();
|
||||
/* prepare buffer */
|
||||
CUSPARSE_CALL(CSRGEAM<DType>::bufferSizeExt(
|
||||
thr_entry->cusparse_handle, m, n, &alpha, matA, nnzA, A_weights,
|
||||
A.indptr.Ptr<IdType>(), A.indices.Ptr<IdType>(), &beta, matB, nnzB,
|
||||
B_weights, B.indptr.Ptr<IdType>(), B.indices.Ptr<IdType>(), matC,
|
||||
dC_weights_data, dC_csrOffsets_data, dC_columns_data, &workspace_size));
|
||||
|
||||
void* workspace = device->AllocWorkspace(ctx, workspace_size);
|
||||
CUSPARSE_CALL(CSRGEAM<DType>::nnz(
|
||||
thr_entry->cusparse_handle, m, n, matA, nnzA, A.indptr.Ptr<IdType>(),
|
||||
A.indices.Ptr<IdType>(), matB, nnzB, B.indptr.Ptr<IdType>(),
|
||||
B.indices.Ptr<IdType>(), matC, dC_csrOffsets_data, &nnzC, workspace));
|
||||
|
||||
dC_columns = IdArray::Empty({nnzC}, A.indptr->dtype, ctx);
|
||||
dC_weights = NDArray::Empty({nnzC}, A_weights_array->dtype, ctx);
|
||||
dC_columns_data = dC_columns.Ptr<IdType>();
|
||||
dC_weights_data = dC_weights.Ptr<DType>();
|
||||
|
||||
CUSPARSE_CALL(CSRGEAM<DType>::compute(
|
||||
thr_entry->cusparse_handle, m, n, &alpha, matA, nnzA, A_weights,
|
||||
A.indptr.Ptr<IdType>(), A.indices.Ptr<IdType>(), &beta, matB, nnzB,
|
||||
B_weights, B.indptr.Ptr<IdType>(), B.indices.Ptr<IdType>(), matC,
|
||||
dC_weights_data, dC_csrOffsets_data, dC_columns_data, workspace));
|
||||
|
||||
device->FreeWorkspace(ctx, workspace);
|
||||
// destroy matrix/vector descriptors
|
||||
CUSPARSE_CALL(cusparseDestroyMatDescr(matA));
|
||||
CUSPARSE_CALL(cusparseDestroyMatDescr(matB));
|
||||
CUSPARSE_CALL(cusparseDestroyMatDescr(matC));
|
||||
return {
|
||||
CSRMatrix(
|
||||
A.num_rows, A.num_cols, dC_csrOffsets, dC_columns,
|
||||
NullArray(dC_csrOffsets->dtype, dC_csrOffsets->ctx), true),
|
||||
dC_weights};
|
||||
}
|
||||
} // namespace cusparse
|
||||
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
std::pair<CSRMatrix, NDArray> CSRSum(
|
||||
const std::vector<CSRMatrix>& As, const std::vector<NDArray>& A_weights) {
|
||||
const int64_t M = As[0].num_rows;
|
||||
const int64_t N = As[0].num_cols;
|
||||
const int64_t n = As.size();
|
||||
|
||||
// Cast 64 bit indices to 32 bit
|
||||
std::vector<CSRMatrix> newAs;
|
||||
newAs.reserve(n);
|
||||
bool cast = false;
|
||||
if (As[0].indptr->dtype.bits == 64) {
|
||||
for (int i = 0; i < n; ++i)
|
||||
newAs.emplace_back(
|
||||
As[i].num_rows, As[i].num_cols, AsNumBits(As[i].indptr, 32),
|
||||
AsNumBits(As[i].indices, 32), AsNumBits(As[i].data, 32));
|
||||
cast = true;
|
||||
} else {
|
||||
for (int i = 0; i < n; ++i) newAs.push_back(As[i]);
|
||||
}
|
||||
|
||||
// cuSPARSE csrgeam2 requires the CSR to be sorted.
|
||||
// TODO(BarclayII): ideally the sorted CSR should be cached but I'm not sure
|
||||
// how to do it.
|
||||
for (int i = 0; i < n; ++i) {
|
||||
if (!newAs[i].sorted) newAs[i] = CSRSort(newAs[i]);
|
||||
}
|
||||
|
||||
// Reorder weights if A[i] has edge IDs
|
||||
std::vector<NDArray> A_weights_reordered(n);
|
||||
for (int i = 0; i < n; ++i) {
|
||||
if (CSRHasData(newAs[i]))
|
||||
A_weights_reordered[i] = IndexSelect(A_weights[i], newAs[i].data);
|
||||
else
|
||||
A_weights_reordered[i] = A_weights[i];
|
||||
}
|
||||
|
||||
// Loop and sum
|
||||
auto result = std::make_pair(
|
||||
CSRMatrix(
|
||||
newAs[0].num_rows, newAs[0].num_cols, newAs[0].indptr,
|
||||
newAs[0].indices,
|
||||
NullArray(newAs[0].indptr->dtype, newAs[0].indptr->ctx)),
|
||||
A_weights_reordered[0]); // Weights already reordered so we don't need
|
||||
// As[0].data
|
||||
for (int64_t i = 1; i < n; ++i)
|
||||
result = cusparse::CusparseCsrgeam2<DType, int32_t>(
|
||||
result.first, result.second, newAs[i], A_weights_reordered[i]);
|
||||
|
||||
// Cast 32 bit indices back to 64 bit if necessary
|
||||
if (cast) {
|
||||
CSRMatrix C = result.first;
|
||||
return {
|
||||
CSRMatrix(
|
||||
C.num_rows, C.num_cols, AsNumBits(C.indptr, 64),
|
||||
AsNumBits(C.indices, 64), AsNumBits(C.data, 64), true),
|
||||
result.second};
|
||||
} else {
|
||||
return result;
|
||||
}
|
||||
}
|
||||
|
||||
template std::pair<CSRMatrix, NDArray> CSRSum<kDGLCUDA, int32_t, __half>(
|
||||
const std::vector<CSRMatrix>&, const std::vector<NDArray>&);
|
||||
template std::pair<CSRMatrix, NDArray> CSRSum<kDGLCUDA, int64_t, __half>(
|
||||
const std::vector<CSRMatrix>&, const std::vector<NDArray>&);
|
||||
#if BF16_ENABLED
|
||||
template std::pair<CSRMatrix, NDArray> CSRSum<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
const std::vector<CSRMatrix>&, const std::vector<NDArray>&);
|
||||
template std::pair<CSRMatrix, NDArray> CSRSum<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
const std::vector<CSRMatrix>&, const std::vector<NDArray>&);
|
||||
#endif // BF16_ENABLED
|
||||
template std::pair<CSRMatrix, NDArray> CSRSum<kDGLCUDA, int32_t, float>(
|
||||
const std::vector<CSRMatrix>&, const std::vector<NDArray>&);
|
||||
template std::pair<CSRMatrix, NDArray> CSRSum<kDGLCUDA, int64_t, float>(
|
||||
const std::vector<CSRMatrix>&, const std::vector<NDArray>&);
|
||||
template std::pair<CSRMatrix, NDArray> CSRSum<kDGLCUDA, int32_t, double>(
|
||||
const std::vector<CSRMatrix>&, const std::vector<NDArray>&);
|
||||
template std::pair<CSRMatrix, NDArray> CSRSum<kDGLCUDA, int64_t, double>(
|
||||
const std::vector<CSRMatrix>&, const std::vector<NDArray>&);
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,95 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/csr_transpose.cc
|
||||
* @brief CSR transpose (convert to CSC)
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using runtime::NDArray;
|
||||
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix CSRTranspose(CSRMatrix csr) {
|
||||
LOG(FATAL) << "Unreachable codes";
|
||||
return {};
|
||||
}
|
||||
|
||||
template <>
|
||||
CSRMatrix CSRTranspose<kDGLCUDA, int32_t>(CSRMatrix csr) {
|
||||
#if CUDART_VERSION < 12000
|
||||
auto* thr_entry = runtime::CUDAThreadEntry::ThreadLocal();
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
// allocate cusparse handle if needed
|
||||
if (!thr_entry->cusparse_handle) {
|
||||
CUSPARSE_CALL(cusparseCreate(&(thr_entry->cusparse_handle)));
|
||||
}
|
||||
CUSPARSE_CALL(cusparseSetStream(thr_entry->cusparse_handle, stream));
|
||||
|
||||
NDArray indptr = csr.indptr, indices = csr.indices, data = csr.data;
|
||||
const int64_t nnz = indices->shape[0];
|
||||
const auto& ctx = indptr->ctx;
|
||||
const auto bits = indptr->dtype.bits;
|
||||
if (aten::IsNullArray(data)) data = aten::Range(0, nnz, bits, ctx);
|
||||
const int32_t* indptr_ptr = static_cast<int32_t*>(indptr->data);
|
||||
const int32_t* indices_ptr = static_cast<int32_t*>(indices->data);
|
||||
const void* data_ptr = data->data;
|
||||
|
||||
// (BarclayII) csr2csc doesn't seem to clear the content of cscColPtr if nnz
|
||||
// == 0. We need to do it ourselves.
|
||||
NDArray t_indptr = aten::Full(0, csr.num_cols + 1, bits, ctx);
|
||||
NDArray t_indices = aten::NewIdArray(nnz, ctx, bits);
|
||||
NDArray t_data = aten::NewIdArray(nnz, ctx, bits);
|
||||
int32_t* t_indptr_ptr = static_cast<int32_t*>(t_indptr->data);
|
||||
int32_t* t_indices_ptr = static_cast<int32_t*>(t_indices->data);
|
||||
void* t_data_ptr = t_data->data;
|
||||
|
||||
#if CUDART_VERSION >= 10010
|
||||
auto device = runtime::DeviceAPI::Get(csr.indptr->ctx);
|
||||
// workspace
|
||||
size_t workspace_size;
|
||||
CUSPARSE_CALL(cusparseCsr2cscEx2_bufferSize(
|
||||
thr_entry->cusparse_handle, csr.num_rows, csr.num_cols, nnz, data_ptr,
|
||||
indptr_ptr, indices_ptr, t_data_ptr, t_indptr_ptr, t_indices_ptr,
|
||||
CUDA_R_32F, CUSPARSE_ACTION_NUMERIC, CUSPARSE_INDEX_BASE_ZERO,
|
||||
CUSPARSE_CSR2CSC_ALG1, // see cusparse doc for reference
|
||||
&workspace_size));
|
||||
void* workspace = device->AllocWorkspace(ctx, workspace_size);
|
||||
CUSPARSE_CALL(cusparseCsr2cscEx2(
|
||||
thr_entry->cusparse_handle, csr.num_rows, csr.num_cols, nnz, data_ptr,
|
||||
indptr_ptr, indices_ptr, t_data_ptr, t_indptr_ptr, t_indices_ptr,
|
||||
CUDA_R_32F, CUSPARSE_ACTION_NUMERIC, CUSPARSE_INDEX_BASE_ZERO,
|
||||
CUSPARSE_CSR2CSC_ALG1, // see cusparse doc for reference
|
||||
workspace));
|
||||
device->FreeWorkspace(ctx, workspace);
|
||||
#else
|
||||
CUSPARSE_CALL(cusparseScsr2csc(
|
||||
thr_entry->cusparse_handle, csr.num_rows, csr.num_cols, nnz,
|
||||
static_cast<const float*>(data_ptr), indptr_ptr, indices_ptr,
|
||||
static_cast<float*>(t_data_ptr), t_indices_ptr, t_indptr_ptr,
|
||||
CUSPARSE_ACTION_NUMERIC, CUSPARSE_INDEX_BASE_ZERO));
|
||||
#endif
|
||||
|
||||
return CSRMatrix(
|
||||
csr.num_cols, csr.num_rows, t_indptr, t_indices, t_data, false);
|
||||
#else
|
||||
return COOToCSR(COOTranspose(CSRToCOO(csr, false)));
|
||||
#endif
|
||||
}
|
||||
|
||||
template <>
|
||||
CSRMatrix CSRTranspose<kDGLCUDA, int64_t>(CSRMatrix csr) {
|
||||
return COOToCSR(COOTranspose(CSRToCOO(csr, false)));
|
||||
}
|
||||
|
||||
template CSRMatrix CSRTranspose<kDGLCUDA, int32_t>(CSRMatrix csr);
|
||||
template CSRMatrix CSRTranspose<kDGLCUDA, int64_t>(CSRMatrix csr);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,140 @@
|
||||
/**
|
||||
* Copyright (c) 2021 by Contributors
|
||||
* @file array/cuda/cuda_filter.cc
|
||||
* @brief Object for selecting items in a set, or selecting items not in a set.
|
||||
*/
|
||||
|
||||
#include <dgl/runtime/device_api.h>
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "../../runtime/cuda/cuda_hashtable.cuh"
|
||||
#include "../filter.h"
|
||||
|
||||
using namespace dgl::runtime::cuda;
|
||||
|
||||
namespace dgl {
|
||||
namespace array {
|
||||
|
||||
namespace {
|
||||
|
||||
template <typename IdType, bool include>
|
||||
__global__ void _IsInKernel(
|
||||
DeviceOrderedHashTable<IdType> table, const IdType* const array,
|
||||
const int64_t size, IdType* const mark) {
|
||||
const int64_t idx = threadIdx.x + blockDim.x * blockIdx.x;
|
||||
if (idx < size) {
|
||||
mark[idx] = table.Contains(array[idx]) ^ (!include);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename IdType>
|
||||
__global__ void _InsertKernel(
|
||||
const IdType* const prefix, const int64_t size, IdType* const result) {
|
||||
const int64_t idx = threadIdx.x + blockDim.x * blockIdx.x;
|
||||
if (idx < size) {
|
||||
if (prefix[idx] != prefix[idx + 1]) {
|
||||
result[prefix[idx]] = idx;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename IdType, bool include>
|
||||
IdArray _PerformFilter(const OrderedHashTable<IdType>& table, IdArray test) {
|
||||
const auto& ctx = test->ctx;
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
const int64_t size = test->shape[0];
|
||||
cudaStream_t cudaStream = runtime::getCurrentCUDAStream();
|
||||
|
||||
if (size == 0) {
|
||||
return test;
|
||||
}
|
||||
|
||||
// we need two arrays: 1) to act as a prefixsum
|
||||
// for the number of entries that will be inserted, and
|
||||
// 2) to collect the included items.
|
||||
IdType* prefix = static_cast<IdType*>(
|
||||
device->AllocWorkspace(ctx, sizeof(IdType) * (size + 1)));
|
||||
|
||||
// will resize down later
|
||||
IdArray result = aten::NewIdArray(size, ctx, sizeof(IdType) * 8);
|
||||
|
||||
// mark each index based on it's existence in the hashtable
|
||||
{
|
||||
const dim3 block(256);
|
||||
const dim3 grid((size + block.x - 1) / block.x);
|
||||
|
||||
CUDA_KERNEL_CALL(
|
||||
(_IsInKernel<IdType, include>), grid, block, 0, cudaStream,
|
||||
table.DeviceHandle(), static_cast<const IdType*>(test->data), size,
|
||||
prefix);
|
||||
}
|
||||
|
||||
// generate prefix-sum
|
||||
{
|
||||
size_t workspace_bytes;
|
||||
CUDA_CALL(cub::DeviceScan::ExclusiveSum(
|
||||
nullptr, workspace_bytes, static_cast<IdType*>(nullptr),
|
||||
static_cast<IdType*>(nullptr), size + 1, cudaStream));
|
||||
void* workspace = device->AllocWorkspace(ctx, workspace_bytes);
|
||||
|
||||
CUDA_CALL(cub::DeviceScan::ExclusiveSum(
|
||||
workspace, workspace_bytes, prefix, prefix, size + 1, cudaStream));
|
||||
device->FreeWorkspace(ctx, workspace);
|
||||
}
|
||||
|
||||
// copy number using the internal current stream;
|
||||
IdType num_unique;
|
||||
device->CopyDataFromTo(
|
||||
prefix + size, 0, &num_unique, 0, sizeof(num_unique), ctx,
|
||||
DGLContext{kDGLCPU, 0}, test->dtype);
|
||||
|
||||
// insert items into set
|
||||
{
|
||||
const dim3 block(256);
|
||||
const dim3 grid((size + block.x - 1) / block.x);
|
||||
|
||||
CUDA_KERNEL_CALL(
|
||||
_InsertKernel, grid, block, 0, cudaStream, prefix, size,
|
||||
static_cast<IdType*>(result->data));
|
||||
}
|
||||
device->FreeWorkspace(ctx, prefix);
|
||||
|
||||
return result.CreateView({num_unique}, result->dtype);
|
||||
}
|
||||
|
||||
template <typename IdType>
|
||||
class CudaFilterSet : public Filter {
|
||||
public:
|
||||
explicit CudaFilterSet(IdArray array)
|
||||
: table_(array->shape[0], array->ctx, runtime::getCurrentCUDAStream()) {
|
||||
cudaStream_t cudaStream = runtime::getCurrentCUDAStream();
|
||||
table_.FillWithUnique(
|
||||
static_cast<const IdType*>(array->data), array->shape[0], cudaStream);
|
||||
}
|
||||
|
||||
IdArray find_included_indices(IdArray test) override {
|
||||
return _PerformFilter<IdType, true>(table_, test);
|
||||
}
|
||||
|
||||
IdArray find_excluded_indices(IdArray test) override {
|
||||
return _PerformFilter<IdType, false>(table_, test);
|
||||
}
|
||||
|
||||
private:
|
||||
OrderedHashTable<IdType> table_;
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
FilterRef CreateSetFilter(IdArray set) {
|
||||
return FilterRef(std::make_shared<CudaFilterSet<IdType>>(set));
|
||||
}
|
||||
|
||||
template FilterRef CreateSetFilter<kDGLCUDA, int32_t>(IdArray set);
|
||||
template FilterRef CreateSetFilter<kDGLCUDA, int64_t>(IdArray set);
|
||||
|
||||
} // namespace array
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,238 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/dispatcher.cuh
|
||||
* @brief Templates to dispatch into different cuSPARSE routines based on the
|
||||
* type argument.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CUDA_CUSPARSE_DISPATCHER_CUH_
|
||||
#define DGL_ARRAY_CUDA_CUSPARSE_DISPATCHER_CUH_
|
||||
|
||||
#include <cusparse.h>
|
||||
#include <dgl/runtime/c_runtime_api.h>
|
||||
|
||||
#include "bf16.cuh"
|
||||
#include "fp16.cuh"
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
|
||||
/** @brief cusparseXcsrgemm dispatcher */
|
||||
template <typename DType>
|
||||
struct CSRGEMM {
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t bufferSizeExt(Args... args) {
|
||||
BUG_IF_FAIL(false) << "This piece of code should not be reached.";
|
||||
return static_cast<cusparseStatus_t>(0);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t nnz(Args... args) {
|
||||
return cusparseXcsrgemm2Nnz(args...);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t compute(Args... args) {
|
||||
BUG_IF_FAIL(false) << "This piece of code should not be reached.";
|
||||
return static_cast<cusparseStatus_t>(0);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct CSRGEMM<__half> {
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t bufferSizeExt(Args... args) {
|
||||
// TODO(ndickson): There is no cusparseHcsrgemm2_bufferSizeExt, so a
|
||||
// different implementation would be required.
|
||||
LOG(FATAL) << "CSRGEMM::bufferSizeExt does not support dtype half (FP16).";
|
||||
return static_cast<cusparseStatus_t>(0);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t nnz(Args... args) {
|
||||
return cusparseXcsrgemm2Nnz(args...);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t compute(Args... args) {
|
||||
// TODO(ndickson): There is no cusparseHcsrgemm2, so a different
|
||||
// implementation would be required.
|
||||
LOG(FATAL) << "CSRGEMM::compute does not support dtype half (FP16).";
|
||||
return static_cast<cusparseStatus_t>(0);
|
||||
}
|
||||
};
|
||||
|
||||
#if BF16_ENABLED
|
||||
template <>
|
||||
struct CSRGEMM<__nv_bfloat16> {
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t bufferSizeExt(Args... args) {
|
||||
// TODO(ndickson): There is no cusparseHcsrgemm2_bufferSizeExt, so a
|
||||
// different implementation would be required.
|
||||
LOG(FATAL)
|
||||
<< "CSRGEMM::bufferSizeExt does not support dtype bfloat16 (BF16).";
|
||||
return static_cast<cusparseStatus_t>(0);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t nnz(Args... args) {
|
||||
return cusparseXcsrgemm2Nnz(args...);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t compute(Args... args) {
|
||||
// TODO(ndickson): There is no cusparseHcsrgemm2, so a different
|
||||
// implementation would be required.
|
||||
LOG(FATAL) << "CSRGEMM::compute does not support dtype bfloat16 (BF16).";
|
||||
return static_cast<cusparseStatus_t>(0);
|
||||
}
|
||||
};
|
||||
#endif // BF16_ENABLED
|
||||
|
||||
template <>
|
||||
struct CSRGEMM<float> {
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t bufferSizeExt(Args... args) {
|
||||
return cusparseScsrgemm2_bufferSizeExt(args...);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t nnz(Args... args) {
|
||||
return cusparseXcsrgemm2Nnz(args...);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t compute(Args... args) {
|
||||
return cusparseScsrgemm2(args...);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct CSRGEMM<double> {
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t bufferSizeExt(Args... args) {
|
||||
return cusparseDcsrgemm2_bufferSizeExt(args...);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t nnz(Args... args) {
|
||||
return cusparseXcsrgemm2Nnz(args...);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t compute(Args... args) {
|
||||
return cusparseDcsrgemm2(args...);
|
||||
}
|
||||
};
|
||||
|
||||
/** @brief cusparseXcsrgeam dispatcher */
|
||||
template <typename DType>
|
||||
struct CSRGEAM {
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t bufferSizeExt(Args... args) {
|
||||
BUG_IF_FAIL(false) << "This piece of code should not be reached.";
|
||||
return static_cast<cusparseStatus_t>(0);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t nnz(Args... args) {
|
||||
return cusparseXcsrgeam2Nnz(args...);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t compute(Args... args) {
|
||||
BUG_IF_FAIL(false) << "This piece of code should not be reached.";
|
||||
return static_cast<cusparseStatus_t>(0);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct CSRGEAM<__half> {
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t bufferSizeExt(Args... args) {
|
||||
// TODO(ndickson): There is no cusparseHcsrgeam2_bufferSizeExt, so a
|
||||
// different implementation would be required.
|
||||
LOG(FATAL) << "CSRGEAM::bufferSizeExt does not support dtype half (FP16).";
|
||||
return static_cast<cusparseStatus_t>(0);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t nnz(Args... args) {
|
||||
return cusparseXcsrgeam2Nnz(args...);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t compute(Args... args) {
|
||||
// TODO(ndickson): There is no cusparseHcsrgeam2, so a different
|
||||
// implementation would be required.
|
||||
LOG(FATAL) << "CSRGEAM::compute does not support dtype half (FP16).";
|
||||
return static_cast<cusparseStatus_t>(0);
|
||||
}
|
||||
};
|
||||
|
||||
#if BF16_ENABLED
|
||||
template <>
|
||||
struct CSRGEAM<__nv_bfloat16> {
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t bufferSizeExt(Args... args) {
|
||||
// TODO(ndickson): There is no cusparseHcsrgeam2_bufferSizeExt, so a
|
||||
// different implementation would be required.
|
||||
LOG(FATAL)
|
||||
<< "CSRGEAM::bufferSizeExt does not support dtype bfloat16 (BF16).";
|
||||
return static_cast<cusparseStatus_t>(0);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t nnz(Args... args) {
|
||||
return cusparseXcsrgeam2Nnz(args...);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t compute(Args... args) {
|
||||
// TODO(ndickson): There is no cusparseHcsrgeam2, so a different
|
||||
// implementation would be required.
|
||||
LOG(FATAL) << "CSRGEAM::compute does not support dtype bfloat16 (BF16).";
|
||||
return static_cast<cusparseStatus_t>(0);
|
||||
}
|
||||
};
|
||||
#endif // BF16_ENABLED
|
||||
|
||||
template <>
|
||||
struct CSRGEAM<float> {
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t bufferSizeExt(Args... args) {
|
||||
return cusparseScsrgeam2_bufferSizeExt(args...);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t nnz(Args... args) {
|
||||
return cusparseXcsrgeam2Nnz(args...);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t compute(Args... args) {
|
||||
return cusparseScsrgeam2(args...);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct CSRGEAM<double> {
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t bufferSizeExt(Args... args) {
|
||||
return cusparseDcsrgeam2_bufferSizeExt(args...);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t nnz(Args... args) {
|
||||
return cusparseXcsrgeam2Nnz(args...);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
static inline cusparseStatus_t compute(Args... args) {
|
||||
return cusparseDcsrgeam2(args...);
|
||||
}
|
||||
};
|
||||
|
||||
}; // namespace aten
|
||||
}; // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CUDA_CUSPARSE_DISPATCHER_CUH_
|
||||
@@ -0,0 +1,185 @@
|
||||
/**
|
||||
* Copyright (c) 2022, NVIDIA CORPORATION.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
* @file array/gpu/disjoint_union.cu
|
||||
* @brief Disjoint union GPU implementation.
|
||||
*/
|
||||
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <typename IdType>
|
||||
__global__ void _DisjointUnionKernel(
|
||||
IdType** arrs, IdType* prefix, IdType* offset, IdType* out, int64_t n_arrs,
|
||||
int n_elms) {
|
||||
IdType tx = static_cast<IdType>(blockIdx.x) * blockDim.x + threadIdx.x;
|
||||
const int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < n_elms) {
|
||||
IdType i = dgl::cuda::_UpperBound(offset, n_arrs, tx) - 1;
|
||||
if (arrs[i] == NULL) {
|
||||
out[tx] = tx;
|
||||
} else {
|
||||
IdType j = tx - offset[i];
|
||||
out[tx] = arrs[i][j] + prefix[i];
|
||||
}
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::tuple<IdArray, IdArray, IdArray> _ComputePrefixSums(
|
||||
const std::vector<COOMatrix>& coos) {
|
||||
IdType n = coos.size(), nbits = coos[0].row->dtype.bits;
|
||||
IdArray n_rows = NewIdArray(n, CPU, nbits);
|
||||
IdArray n_cols = NewIdArray(n, CPU, nbits);
|
||||
IdArray n_elms = NewIdArray(n, CPU, nbits);
|
||||
|
||||
IdType* n_rows_data = n_rows.Ptr<IdType>();
|
||||
IdType* n_cols_data = n_cols.Ptr<IdType>();
|
||||
IdType* n_elms_data = n_elms.Ptr<IdType>();
|
||||
|
||||
dgl::runtime::parallel_for(0, coos.size(), [&](IdType b, IdType e) {
|
||||
for (IdType i = b; i < e; ++i) {
|
||||
n_rows_data[i] = coos[i].num_rows;
|
||||
n_cols_data[i] = coos[i].num_cols;
|
||||
n_elms_data[i] = coos[i].row->shape[0];
|
||||
}
|
||||
});
|
||||
|
||||
return std::make_tuple(
|
||||
CumSum(n_rows.CopyTo(coos[0].row->ctx), true),
|
||||
CumSum(n_cols.CopyTo(coos[0].row->ctx), true),
|
||||
CumSum(n_elms.CopyTo(coos[0].row->ctx), true));
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
void _Merge(
|
||||
IdType** arrs, IdType* prefix, IdType* offset, IdType* out, int64_t n_arrs,
|
||||
int n_elms, DGLContext ctx, DGLDataType dtype, cudaStream_t stream) {
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
int nt = 256;
|
||||
int nb = (n_elms + nt - 1) / nt;
|
||||
|
||||
IdType** arrs_dev = static_cast<IdType**>(
|
||||
device->AllocWorkspace(ctx, n_arrs * sizeof(IdType*)));
|
||||
|
||||
device->CopyDataFromTo(
|
||||
arrs, 0, arrs_dev, 0, sizeof(IdType*) * n_arrs, DGLContext{kDGLCPU, 0},
|
||||
ctx, dtype);
|
||||
|
||||
CUDA_KERNEL_CALL(
|
||||
_DisjointUnionKernel, nb, nt, 0, stream, arrs_dev, prefix, offset, out,
|
||||
n_arrs, n_elms);
|
||||
|
||||
device->FreeWorkspace(ctx, arrs_dev);
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix DisjointUnionCoo(const std::vector<COOMatrix>& coos) {
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
auto device = runtime::DeviceAPI::Get(coos[0].row->ctx);
|
||||
uint64_t src_offset = 0, dst_offset = 0;
|
||||
bool has_data = false;
|
||||
bool row_sorted = true;
|
||||
bool col_sorted = true;
|
||||
|
||||
// check if data index array
|
||||
for (size_t i = 0; i < coos.size(); ++i) {
|
||||
CHECK_SAME_DTYPE(coos[0].row, coos[i].row);
|
||||
CHECK_SAME_CONTEXT(coos[0].row, coos[i].row);
|
||||
has_data |= COOHasData(coos[i]);
|
||||
}
|
||||
|
||||
auto prefixes = _ComputePrefixSums<XPU, IdType>(coos);
|
||||
auto prefix_src = static_cast<IdType*>(std::get<0>(prefixes)->data);
|
||||
auto prefix_dst = static_cast<IdType*>(std::get<1>(prefixes)->data);
|
||||
auto prefix_elm = static_cast<IdType*>(std::get<2>(prefixes)->data);
|
||||
|
||||
std::unique_ptr<IdType*[]> rows(new IdType*[coos.size()]);
|
||||
std::unique_ptr<IdType*[]> cols(new IdType*[coos.size()]);
|
||||
std::unique_ptr<IdType*[]> data(new IdType*[coos.size()]);
|
||||
|
||||
for (size_t i = 0; i < coos.size(); i++) {
|
||||
row_sorted &= coos[i].row_sorted;
|
||||
col_sorted &= coos[i].col_sorted;
|
||||
rows[i] = coos[i].row.Ptr<IdType>();
|
||||
cols[i] = coos[i].col.Ptr<IdType>();
|
||||
data[i] = coos[i].data.Ptr<IdType>();
|
||||
}
|
||||
|
||||
auto ctx = coos[0].row->ctx;
|
||||
auto dtype = coos[0].row->dtype;
|
||||
|
||||
IdType n_elements = 0;
|
||||
device->CopyDataFromTo(
|
||||
&prefix_elm[coos.size()], 0, &n_elements, 0, sizeof(IdType),
|
||||
coos[0].row->ctx, DGLContext{kDGLCPU, 0}, coos[0].row->dtype);
|
||||
|
||||
device->CopyDataFromTo(
|
||||
&prefix_src[coos.size()], 0, &src_offset, 0, sizeof(IdType),
|
||||
coos[0].row->ctx, DGLContext{kDGLCPU, 0}, coos[0].row->dtype);
|
||||
|
||||
device->CopyDataFromTo(
|
||||
&prefix_dst[coos.size()], 0, &dst_offset, 0, sizeof(IdType),
|
||||
coos[0].row->ctx, DGLContext{kDGLCPU, 0}, coos[0].row->dtype);
|
||||
|
||||
// Union src array
|
||||
IdArray result_src =
|
||||
NewIdArray(n_elements, coos[0].row->ctx, coos[0].row->dtype.bits);
|
||||
_Merge<XPU, IdType>(
|
||||
rows.get(), prefix_src, prefix_elm, result_src.Ptr<IdType>(), coos.size(),
|
||||
n_elements, ctx, dtype, stream);
|
||||
|
||||
// Union dst array
|
||||
IdArray result_dst =
|
||||
NewIdArray(n_elements, coos[0].col->ctx, coos[0].col->dtype.bits);
|
||||
_Merge<XPU, IdType>(
|
||||
cols.get(), prefix_dst, prefix_elm, result_dst.Ptr<IdType>(), coos.size(),
|
||||
n_elements, ctx, dtype, stream);
|
||||
|
||||
// Union data array if exists and fetch number of elements
|
||||
IdArray result_dat = NullArray();
|
||||
if (has_data) {
|
||||
result_dat =
|
||||
NewIdArray(n_elements, coos[0].row->ctx, coos[0].row->dtype.bits);
|
||||
_Merge<XPU, IdType>(
|
||||
data.get(), prefix_elm, prefix_elm, result_dat.Ptr<IdType>(),
|
||||
coos.size(), n_elements, ctx, dtype, stream);
|
||||
}
|
||||
|
||||
return COOMatrix(
|
||||
src_offset, dst_offset, result_src, result_dst, result_dat, row_sorted,
|
||||
col_sorted);
|
||||
}
|
||||
|
||||
template COOMatrix DisjointUnionCoo<kDGLCUDA, int32_t>(
|
||||
const std::vector<COOMatrix>& coos);
|
||||
template COOMatrix DisjointUnionCoo<kDGLCUDA, int64_t>(
|
||||
const std::vector<COOMatrix>& coos);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,134 @@
|
||||
/**
|
||||
* Copyright (c) 2020-2022 by Contributors
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
* @file array/cuda/fp16.cuh
|
||||
* @brief float16 related functions.
|
||||
* @note this file is modified from TVM project:
|
||||
* https://github.com/apache/tvm/blob/e561007f0c330e3d14c2bc8a3ef40fb741db9004/src/target/source/literal/cuda_half_t.h.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CUDA_FP16_CUH_
|
||||
#define DGL_ARRAY_CUDA_FP16_CUH_
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
static __device__ __forceinline__ half max(half a, half b) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 530
|
||||
return __hgt(__half(a), __half(b)) ? a : b;
|
||||
#else
|
||||
return __half(max(float(a), float(b))); // NOLINT
|
||||
#endif
|
||||
}
|
||||
|
||||
static __device__ __forceinline__ half min(half a, half b) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 530
|
||||
return __hlt(__half(a), __half(b)) ? a : b;
|
||||
#else
|
||||
return __half(min(float(a), float(b))); // NOLINT
|
||||
#endif
|
||||
}
|
||||
|
||||
#ifdef __CUDACC__
|
||||
// Arithmetic FP16 operations for architecture >= 5.3 are already defined in
|
||||
// cuda_fp16.h
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 530)
|
||||
// CUDA 12.2 adds "emulated" support for older architectures.
|
||||
#if defined(CUDART_VERSION) && (CUDART_VERSION < 12020)
|
||||
__device__ __forceinline__ __half
|
||||
operator+(const __half& lh, const __half& rh) {
|
||||
return __half(float(lh) + float(rh)); // NOLINT
|
||||
}
|
||||
__device__ __forceinline__ __half
|
||||
operator-(const __half& lh, const __half& rh) {
|
||||
return __half(float(lh) - float(rh)); // NOLINT
|
||||
}
|
||||
__device__ __forceinline__ __half
|
||||
operator*(const __half& lh, const __half& rh) {
|
||||
return __half(float(lh) * float(rh)); // NOLINT
|
||||
}
|
||||
__device__ __forceinline__ __half
|
||||
operator/(const __half& lh, const __half& rh) {
|
||||
return __half(float(lh) / float(rh)); // NOLINT
|
||||
}
|
||||
|
||||
__device__ __forceinline__ __half& operator+=(
|
||||
__half& lh, const __half& rh) { // NOLINT
|
||||
lh = __half(float(lh) + float(rh)); // NOLINT
|
||||
return lh;
|
||||
}
|
||||
__device__ __forceinline__ __half& operator-=(
|
||||
__half& lh, const __half& rh) { // NOLINT
|
||||
lh = __half(float(lh) - float(rh)); // NOLINT
|
||||
return lh;
|
||||
}
|
||||
__device__ __forceinline__ __half& operator*=(
|
||||
__half& lh, const __half& rh) { // NOLINT
|
||||
lh = __half(float(lh) * float(rh)); // NOLINT
|
||||
return lh;
|
||||
}
|
||||
__device__ __forceinline__ __half& operator/=(
|
||||
__half& lh, const __half& rh) { // NOLINT
|
||||
lh = __half(float(lh) / float(rh)); // NOLINT
|
||||
return lh;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ __half& operator++(__half& h) { // NOLINT
|
||||
h = __half(float(h) + 1.0f); // NOLINT
|
||||
return h;
|
||||
}
|
||||
__device__ __forceinline__ __half& operator--(__half& h) { // NOLINT
|
||||
h = __half(float(h) - 1.0f); // NOLINT
|
||||
return h;
|
||||
}
|
||||
__device__ __forceinline__ __half operator++(__half& h, int) { // NOLINT
|
||||
__half ret = h;
|
||||
h = __half(float(h) + 1.0f); // NOLINT
|
||||
return ret;
|
||||
}
|
||||
__device__ __forceinline__ __half operator--(__half& h, int) { // NOLINT
|
||||
__half ret = h;
|
||||
h = __half(float(h) - 1.0f); // NOLINT
|
||||
return ret;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ __half operator+(const __half& h) { return h; }
|
||||
__device__ __forceinline__ __half operator-(const __half& h) {
|
||||
return __half(-float(h)); // NOLINT
|
||||
}
|
||||
|
||||
__device__ __forceinline__ bool operator==(const __half& lh, const __half& rh) {
|
||||
return float(lh) == float(rh); // NOLINT
|
||||
}
|
||||
__device__ __forceinline__ bool operator!=(const __half& lh, const __half& rh) {
|
||||
return float(lh) != float(rh); // NOLINT
|
||||
}
|
||||
__device__ __forceinline__ bool operator>(const __half& lh, const __half& rh) {
|
||||
return float(lh) > float(rh); // NOLINT
|
||||
}
|
||||
__device__ __forceinline__ bool operator<(const __half& lh, const __half& rh) {
|
||||
return float(lh) < float(rh); // NOLINT
|
||||
}
|
||||
__device__ __forceinline__ bool operator>=(const __half& lh, const __half& rh) {
|
||||
return float(lh) >= float(rh); // NOLINT
|
||||
}
|
||||
__device__ __forceinline__ bool operator<=(const __half& lh, const __half& rh) {
|
||||
return float(lh) <= float(rh); // NOLINT
|
||||
}
|
||||
#endif // defined(CUDART_VERSION) && (CUDART_VERSION < 12020)
|
||||
#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 530)
|
||||
#endif // __CUDACC__
|
||||
|
||||
#endif // DGL_ARRAY_CUDA_FP16_CUH_
|
||||
@@ -0,0 +1,456 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/functor.cuh
|
||||
* @brief Functors for template on CUDA
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CUDA_FUNCTOR_CUH_
|
||||
#define DGL_ARRAY_CUDA_FUNCTOR_CUH_
|
||||
|
||||
#include <cmath>
|
||||
#include <limits>
|
||||
|
||||
#include "./atomic.cuh"
|
||||
#include "./fp16.cuh"
|
||||
#include "bf16.cuh"
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace cuda {
|
||||
|
||||
/////////////////////////// CUDA binary operators //////////////////////////////
|
||||
namespace binary {
|
||||
template <typename DType>
|
||||
struct Add {
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = true;
|
||||
static constexpr bool reduce_last_dim = false;
|
||||
static __device__ __forceinline__ DType
|
||||
Call(const DType *lhs, const DType *rhs, int64_t len = 1) {
|
||||
return lhs[0] + rhs[0];
|
||||
}
|
||||
};
|
||||
template <typename DType>
|
||||
constexpr bool Add<DType>::use_lhs;
|
||||
template <typename DType>
|
||||
constexpr bool Add<DType>::use_rhs;
|
||||
template <typename DType>
|
||||
constexpr bool Add<DType>::reduce_last_dim;
|
||||
|
||||
template <typename DType>
|
||||
struct Sub {
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = true;
|
||||
static constexpr bool reduce_last_dim = false;
|
||||
static __device__ __forceinline__ DType
|
||||
Call(const DType *lhs, const DType *rhs, int64_t len = 1) {
|
||||
return lhs[0] - rhs[0];
|
||||
}
|
||||
};
|
||||
template <typename DType>
|
||||
constexpr bool Sub<DType>::use_lhs;
|
||||
template <typename DType>
|
||||
constexpr bool Sub<DType>::use_rhs;
|
||||
template <typename DType>
|
||||
constexpr bool Sub<DType>::reduce_last_dim;
|
||||
|
||||
template <typename DType>
|
||||
struct Mul {
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = true;
|
||||
static constexpr bool reduce_last_dim = false;
|
||||
static __device__ __forceinline__ DType
|
||||
Call(const DType *lhs, const DType *rhs, int64_t len = 1) {
|
||||
return lhs[0] * rhs[0];
|
||||
}
|
||||
};
|
||||
template <typename DType>
|
||||
constexpr bool Mul<DType>::use_lhs;
|
||||
template <typename DType>
|
||||
constexpr bool Mul<DType>::use_rhs;
|
||||
template <typename DType>
|
||||
constexpr bool Mul<DType>::reduce_last_dim;
|
||||
|
||||
template <typename DType>
|
||||
struct Div {
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = true;
|
||||
static constexpr bool reduce_last_dim = false;
|
||||
static __device__ __forceinline__ DType
|
||||
Call(const DType *lhs, const DType *rhs, int64_t len = 1) {
|
||||
return lhs[0] / rhs[0];
|
||||
}
|
||||
};
|
||||
template <typename DType>
|
||||
constexpr bool Div<DType>::use_lhs;
|
||||
template <typename DType>
|
||||
constexpr bool Div<DType>::use_rhs;
|
||||
template <typename DType>
|
||||
constexpr bool Div<DType>::reduce_last_dim;
|
||||
|
||||
template <typename DType>
|
||||
struct CopyLhs {
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = false;
|
||||
static constexpr bool reduce_last_dim = false;
|
||||
static __device__ __forceinline__ DType
|
||||
Call(const DType *lhs, const DType *rhs, int64_t len = 1) {
|
||||
return lhs[0];
|
||||
}
|
||||
};
|
||||
template <typename DType>
|
||||
constexpr bool CopyLhs<DType>::use_lhs;
|
||||
template <typename DType>
|
||||
constexpr bool CopyLhs<DType>::use_rhs;
|
||||
template <typename DType>
|
||||
constexpr bool CopyLhs<DType>::reduce_last_dim;
|
||||
|
||||
template <typename DType>
|
||||
struct CopyRhs {
|
||||
static constexpr bool use_lhs = false;
|
||||
static constexpr bool use_rhs = true;
|
||||
static constexpr bool reduce_last_dim = false;
|
||||
static __device__ __forceinline__ DType
|
||||
Call(const DType *lhs, const DType *rhs, int64_t len = 1) {
|
||||
return rhs[0];
|
||||
}
|
||||
};
|
||||
template <typename DType>
|
||||
constexpr bool CopyRhs<DType>::use_lhs;
|
||||
template <typename DType>
|
||||
constexpr bool CopyRhs<DType>::use_rhs;
|
||||
template <typename DType>
|
||||
constexpr bool CopyRhs<DType>::reduce_last_dim;
|
||||
|
||||
template <typename DType>
|
||||
struct Dot {
|
||||
static constexpr bool use_lhs = true;
|
||||
static constexpr bool use_rhs = true;
|
||||
static constexpr bool reduce_last_dim = true;
|
||||
static __device__ __forceinline__ DType
|
||||
Call(const DType *lhs, const DType *rhs, int64_t len = 1) {
|
||||
DType rst = static_cast<DType>(0.0f);
|
||||
for (int64_t i = 0; i < len; ++i) {
|
||||
rst += lhs[i] * rhs[i];
|
||||
}
|
||||
return rst;
|
||||
}
|
||||
};
|
||||
template <typename DType>
|
||||
constexpr bool Dot<DType>::use_lhs;
|
||||
template <typename DType>
|
||||
constexpr bool Dot<DType>::use_rhs;
|
||||
template <typename DType>
|
||||
constexpr bool Dot<DType>::reduce_last_dim;
|
||||
|
||||
} // end of namespace binary
|
||||
|
||||
/////////////////////////// CUDA reduce operators //////////////////////////////
|
||||
namespace reduce {
|
||||
template <typename Idx, typename DType, bool atomic>
|
||||
struct _Sum {
|
||||
static constexpr __host__ __device__ __forceinline__ DType zero() {
|
||||
return 0.;
|
||||
}
|
||||
static constexpr bool require_arg = false;
|
||||
static __device__ __forceinline__ void Call(
|
||||
DType *out_buf, Idx *arg_u_buf, Idx *arg_e_buf, DType val, Idx uid,
|
||||
Idx eid) {
|
||||
if (!atomic) {
|
||||
*out_buf += val;
|
||||
} else {
|
||||
cuda::AtomicAdd(out_buf, val);
|
||||
}
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
DType *out_buf, Idx *arg_buf, DType val, Idx id) {
|
||||
if (!atomic) {
|
||||
*out_buf += val;
|
||||
} else {
|
||||
cuda::AtomicAdd(out_buf, val);
|
||||
}
|
||||
}
|
||||
static __device__ __forceinline__ void CallArg(
|
||||
Idx fid, Idx *arg_u_buf, Idx *arg_e_buf, DType val, DType val_ref,
|
||||
Idx uid, Idx eid) {}
|
||||
};
|
||||
|
||||
template <typename Idx, typename DType, bool atomic = false>
|
||||
struct Sum : _Sum<Idx, DType, atomic> {};
|
||||
|
||||
template <typename Idx, bool atomic>
|
||||
struct Sum<Idx, __half, atomic> : _Sum<Idx, __half, atomic> {
|
||||
static constexpr __host__ __device__ __forceinline__ __half zero() {
|
||||
return __float2half_rn(0.);
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
__half *out_buf, Idx *arg_u_buf, Idx *arg_e_buf,
|
||||
__half val, Idx uid, Idx eid) {
|
||||
_Sum<Idx, __half, atomic>::Call(
|
||||
out_buf, arg_u_buf, arg_e_buf, val, uid, eid);
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
__half *out_buf, Idx *arg_buf, __half val, Idx id) {
|
||||
_Sum<Idx, __half, atomic>::Call(out_buf, arg_buf, val, id);
|
||||
}
|
||||
// sometimes we have to use float in reduction for better precision
|
||||
static __device__ __forceinline__ void Call(
|
||||
float *out_buf, Idx *arg_u_buf, Idx *arg_e_buf,
|
||||
__half val, Idx uid, Idx eid) {
|
||||
_Sum<Idx, float, atomic>::Call(out_buf, arg_u_buf, arg_e_buf,
|
||||
static_cast<float>(val), uid, eid);
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
float *out_buf, Idx *arg_buf, __half val, Idx id) {
|
||||
_Sum<Idx, float, atomic>::Call(out_buf, arg_buf,
|
||||
static_cast<float>(val), id);
|
||||
}
|
||||
};
|
||||
|
||||
#if BF16_ENABLED
|
||||
template <typename Idx, bool atomic>
|
||||
struct Sum<Idx, __nv_bfloat16, atomic> : _Sum<Idx, __nv_bfloat16, atomic> {
|
||||
static constexpr __host__ __device__ __forceinline__ __nv_bfloat16 zero() {
|
||||
return __float2bfloat16_rn(0.);
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
__nv_bfloat16 *out_buf, Idx *arg_u_buf, Idx *arg_e_buf,
|
||||
__nv_bfloat16 val, Idx uid, Idx eid) {
|
||||
_Sum<Idx, __nv_bfloat16, atomic>::Call(
|
||||
out_buf, arg_u_buf, arg_e_buf, val, uid, eid);
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
__nv_bfloat16 *out_buf, Idx *arg_buf, __nv_bfloat16 val, Idx id) {
|
||||
_Sum<Idx, __nv_bfloat16, atomic>::Call(out_buf, arg_buf, val, id);
|
||||
}
|
||||
// sometimes we have to use float in reduction for better precision
|
||||
static __device__ __forceinline__ void Call(
|
||||
float *out_buf, Idx *arg_u_buf, Idx *arg_e_buf,
|
||||
__nv_bfloat16 val, Idx uid, Idx eid) {
|
||||
_Sum<Idx, float, atomic>::Call(out_buf, arg_u_buf, arg_e_buf,
|
||||
static_cast<float>(val), uid, eid);
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
float *out_buf, Idx *arg_buf, __nv_bfloat16 val, Idx id) {
|
||||
_Sum<Idx, float, atomic>::Call(out_buf, arg_buf,
|
||||
static_cast<float>(val), id);
|
||||
}
|
||||
};
|
||||
#endif // BF16_ENABLED
|
||||
|
||||
template <typename Idx, typename DType, bool atomic>
|
||||
struct _Max {
|
||||
static constexpr __host__ __device__ __forceinline__ DType zero() {
|
||||
return -std::numeric_limits<DType>::infinity();
|
||||
}
|
||||
static constexpr bool require_arg = true;
|
||||
static __device__ __forceinline__ void Call(
|
||||
DType *out_buf, Idx *arg_u_buf, Idx *arg_e_buf, DType val, Idx uid,
|
||||
Idx eid) {
|
||||
if (!atomic) {
|
||||
if (*out_buf < val) {
|
||||
*out_buf = val;
|
||||
*arg_u_buf = uid;
|
||||
*arg_e_buf = eid;
|
||||
}
|
||||
} else {
|
||||
cuda::AtomicMax(out_buf, val);
|
||||
}
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
DType *out_buf, Idx *arg_buf, DType val, Idx id) {
|
||||
if (!atomic) {
|
||||
if (*out_buf < val) {
|
||||
*out_buf = val;
|
||||
*arg_buf = id;
|
||||
}
|
||||
} else {
|
||||
cuda::AtomicMax(out_buf, val);
|
||||
}
|
||||
}
|
||||
static __device__ __forceinline__ void CallArg(
|
||||
Idx fid, Idx *arg_u_buf, Idx *arg_e_buf, DType val, DType val_ref,
|
||||
Idx uid, Idx eid) {
|
||||
if (atomic) {
|
||||
if (val == val_ref) {
|
||||
if (arg_u_buf) arg_u_buf[fid] = uid;
|
||||
if (arg_e_buf) arg_e_buf[fid] = eid;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <typename Idx, typename DType, bool atomic = false>
|
||||
struct Max : _Max<Idx, DType, atomic> {};
|
||||
|
||||
template <typename Idx, bool atomic>
|
||||
struct Max<Idx, __half, atomic> : _Max<Idx, __half, atomic> {
|
||||
static constexpr __host__ __device__ __forceinline__ __half zero() {
|
||||
return __float2half_rn(-6.550400e+04f);
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
__half *out_buf, Idx *arg_u_buf, Idx *arg_e_buf,
|
||||
__half val, Idx uid, Idx eid) {
|
||||
_Max<Idx, __half, atomic>::Call(
|
||||
out_buf, arg_u_buf, arg_e_buf, val, uid, eid);
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
__half *out_buf, Idx *arg_buf, __half val, Idx id) {
|
||||
_Max<Idx, __half, atomic>::Call(out_buf, arg_buf, val, id);
|
||||
}
|
||||
// sometimes we have to use float in reduction for better precision
|
||||
static __device__ __forceinline__ void Call(
|
||||
float *out_buf, Idx *arg_u_buf, Idx *arg_e_buf,
|
||||
__half val, Idx uid, Idx eid) {
|
||||
_Max<Idx, float, atomic>::Call(out_buf, arg_u_buf, arg_e_buf,
|
||||
static_cast<float>(val), uid, eid);
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
float *out_buf, Idx *arg_buf, __half val, Idx id) {
|
||||
_Max<Idx, float, atomic>::Call(out_buf, arg_buf,
|
||||
static_cast<float>(val), id);
|
||||
}
|
||||
};
|
||||
|
||||
#if BF16_ENABLED
|
||||
template <typename Idx, bool atomic>
|
||||
struct Max<Idx, __nv_bfloat16, atomic> : _Max<Idx, __nv_bfloat16, atomic> {
|
||||
static constexpr __host__ __device__ __forceinline__ __nv_bfloat16 zero() {
|
||||
return __float2bfloat16_rn(-std::numeric_limits<float>::infinity());
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
__nv_bfloat16 *out_buf, Idx *arg_u_buf, Idx *arg_e_buf,
|
||||
__nv_bfloat16 val, Idx uid, Idx eid) {
|
||||
_Max<Idx, __nv_bfloat16, atomic>::Call(
|
||||
out_buf, arg_u_buf, arg_e_buf, val, uid, eid);
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
__nv_bfloat16 *out_buf, Idx *arg_buf, __nv_bfloat16 val, Idx id) {
|
||||
_Max<Idx, __nv_bfloat16, atomic>::Call(out_buf, arg_buf, val, id);
|
||||
}
|
||||
// sometimes we have to use float in reduction for better precision
|
||||
static __device__ __forceinline__ void Call(
|
||||
float *out_buf, Idx *arg_u_buf, Idx *arg_e_buf,
|
||||
__nv_bfloat16 val, Idx uid, Idx eid) {
|
||||
_Max<Idx, float, atomic>::Call(out_buf, arg_u_buf, arg_e_buf,
|
||||
static_cast<float>(val), uid, eid);
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
float *out_buf, Idx *arg_buf, __nv_bfloat16 val, Idx id) {
|
||||
_Max<Idx, float, atomic>::Call(out_buf, arg_buf,
|
||||
static_cast<float>(val), id);
|
||||
}
|
||||
};
|
||||
#endif // BF16_ENABLED
|
||||
|
||||
template <typename Idx, typename DType, bool atomic>
|
||||
struct _Min {
|
||||
static constexpr __host__ __device__ __forceinline__ DType zero() {
|
||||
return std::numeric_limits<DType>::infinity();
|
||||
}
|
||||
static constexpr bool require_arg = true;
|
||||
static __device__ __forceinline__ void Call(
|
||||
DType *out_buf, Idx *arg_u_buf, Idx *arg_e_buf, DType val, Idx uid,
|
||||
Idx eid) {
|
||||
if (!atomic) {
|
||||
if (*out_buf > val) {
|
||||
*out_buf = val;
|
||||
*arg_u_buf = uid;
|
||||
*arg_e_buf = eid;
|
||||
}
|
||||
} else {
|
||||
cuda::AtomicMin(out_buf, val);
|
||||
}
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
DType *out_buf, Idx *arg_buf, DType val, Idx id) {
|
||||
if (!atomic) {
|
||||
if (*out_buf > val) {
|
||||
*out_buf = val;
|
||||
*arg_buf = id;
|
||||
}
|
||||
} else {
|
||||
cuda::AtomicMin(out_buf, val);
|
||||
}
|
||||
}
|
||||
static __device__ __forceinline__ void CallArg(
|
||||
Idx fid, Idx *arg_u_buf, Idx *arg_e_buf, DType val, DType val_ref,
|
||||
Idx uid, Idx eid) {
|
||||
if (atomic) {
|
||||
if (val == val_ref) {
|
||||
if (arg_u_buf) arg_u_buf[fid] = uid;
|
||||
if (arg_e_buf) arg_e_buf[fid] = eid;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <typename Idx, typename DType, bool atomic = false>
|
||||
struct Min : _Min<Idx, DType, atomic> {};
|
||||
|
||||
template <typename Idx, bool atomic>
|
||||
struct Min<Idx, __half, atomic> : _Min<Idx, __half, atomic> {
|
||||
static constexpr __host__ __device__ __forceinline__ __half zero() {
|
||||
return __float2half_rn(6.550400e+04f);
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
__half *out_buf, Idx *arg_u_buf, Idx *arg_e_buf,
|
||||
__half val, Idx uid, Idx eid) {
|
||||
_Min<Idx, __half, atomic>::Call(
|
||||
out_buf, arg_u_buf, arg_e_buf, val, uid, eid);
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
__half *out_buf, Idx *arg_buf, __half val, Idx id) {
|
||||
_Min<Idx, __half, atomic>::Call(out_buf, arg_buf, val, id);
|
||||
}
|
||||
// sometimes we have to use float in reduction for better precision
|
||||
static __device__ __forceinline__ void Call(
|
||||
float *out_buf, Idx *arg_u_buf, Idx *arg_e_buf,
|
||||
__half val, Idx uid, Idx eid) {
|
||||
_Min<Idx, float, atomic>::Call(out_buf, arg_u_buf, arg_e_buf,
|
||||
static_cast<float>(val), uid, eid);
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
float *out_buf, Idx *arg_buf, __half val, Idx id) {
|
||||
_Min<Idx, float, atomic>::Call(out_buf, arg_buf,
|
||||
static_cast<float>(val), id);
|
||||
}
|
||||
};
|
||||
|
||||
#if BF16_ENABLED
|
||||
template <typename Idx, bool atomic>
|
||||
struct Min<Idx, __nv_bfloat16, atomic> : _Min<Idx, __nv_bfloat16, atomic> {
|
||||
static constexpr __host__ __device__ __forceinline__ __nv_bfloat16 zero() {
|
||||
return __float2bfloat16_rn(std::numeric_limits<float>::infinity());
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
__nv_bfloat16 *out_buf, Idx *arg_u_buf, Idx *arg_e_buf,
|
||||
__nv_bfloat16 val, Idx uid, Idx eid) {
|
||||
_Min<Idx, __nv_bfloat16, atomic>::Call(
|
||||
out_buf, arg_u_buf, arg_e_buf, val, uid, eid);
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
__nv_bfloat16 *out_buf, Idx *arg_buf, __nv_bfloat16 val, Idx id) {
|
||||
_Min<Idx, __nv_bfloat16, atomic>::Call(out_buf, arg_buf, val, id);
|
||||
}
|
||||
// sometimes we have to use float in reduction for better precision
|
||||
static __device__ __forceinline__ void Call(
|
||||
float *out_buf, Idx *arg_u_buf, Idx *arg_e_buf,
|
||||
__nv_bfloat16 val, Idx uid, Idx eid) {
|
||||
_Min<Idx, float, atomic>::Call(out_buf, arg_u_buf, arg_e_buf,
|
||||
static_cast<float>(val), uid, eid);
|
||||
}
|
||||
static __device__ __forceinline__ void Call(
|
||||
float *out_buf, Idx *arg_buf, __nv_bfloat16 val, Idx id) {
|
||||
_Min<Idx, float, atomic>::Call(out_buf, arg_buf,
|
||||
static_cast<float>(val), id);
|
||||
}
|
||||
};
|
||||
#endif // BF16_ENABLED
|
||||
|
||||
} // namespace reduce
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CUDA_FUNCTOR_CUH_
|
||||
@@ -0,0 +1,464 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/gather_mm.cu
|
||||
* @brief GatherMM C APIs and definitions.
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include <algorithm> // std::swap
|
||||
|
||||
#include "./atomic.cuh"
|
||||
#include "./functor.cuh"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
using namespace cuda;
|
||||
namespace aten {
|
||||
|
||||
namespace {
|
||||
|
||||
/** @brief Call cuBLAS GEMM API for dense matmul operation for float and double.
|
||||
*/
|
||||
template <typename DType>
|
||||
cublasStatus_t cublasGemm(
|
||||
cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
|
||||
int m, int n, int k, const DType* alpha, const DType* A, int lda,
|
||||
const DType* B, int ldb, const DType* beta, DType* C, int ldc) {
|
||||
LOG(INFO) << "Not supported dtype";
|
||||
return CUBLAS_STATUS_EXECUTION_FAILED;
|
||||
}
|
||||
|
||||
template <>
|
||||
cublasStatus_t cublasGemm<__half>(
|
||||
cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
|
||||
int m, int n, int k, const __half* alpha, const __half* A, int lda,
|
||||
const __half* B, int ldb, const __half* beta, __half* C, int ldc) {
|
||||
return cublasHgemm(
|
||||
handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc);
|
||||
}
|
||||
|
||||
#if BF16_ENABLED
|
||||
template <>
|
||||
cublasStatus_t cublasGemm<__nv_bfloat16>(
|
||||
cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
|
||||
int m, int n, int k, const __nv_bfloat16* alpha, const __nv_bfloat16* A,
|
||||
int lda, const __nv_bfloat16* B, int ldb, const __nv_bfloat16* beta,
|
||||
__nv_bfloat16* C, int ldc) {
|
||||
float alpha_float = __bfloat162float(*alpha);
|
||||
float beta_float = __bfloat162float(*beta);
|
||||
return cublasGemmEx(
|
||||
handle, transa, transb, m, n, k, &alpha_float, A, CUDA_R_16BF, lda, B,
|
||||
CUDA_R_16BF, ldb, &beta_float, C, CUDA_R_16BF, ldc, CUBLAS_COMPUTE_32F,
|
||||
CUBLAS_GEMM_DEFAULT_TENSOR_OP);
|
||||
}
|
||||
#endif // BF16_ENABLED
|
||||
|
||||
template <>
|
||||
cublasStatus_t cublasGemm<float>(
|
||||
cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
|
||||
int m, int n, int k, const float* alpha, const float* A, int lda,
|
||||
const float* B, int ldb, const float* beta, float* C, int ldc) {
|
||||
return cublasSgemm(
|
||||
handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc);
|
||||
}
|
||||
|
||||
template <>
|
||||
cublasStatus_t cublasGemm<double>(
|
||||
cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
|
||||
int m, int n, int k, const double* alpha, const double* A, int lda,
|
||||
const double* B, int ldb, const double* beta, double* C, int ldc) {
|
||||
return cublasDgemm(
|
||||
handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace cuda {
|
||||
|
||||
/**
|
||||
* @note Each row of A multiplies a segment of matrix of B of dimension in_len *
|
||||
* outlen. One warp is assigned to process one row of A. Each WARP sequentially
|
||||
* multiplies one element of A and a row of B to compute partial result of the
|
||||
* output. A is loaded in shared memory in a coalesced way. Output matrix is
|
||||
* loaded in registers. B should get benefit from L2 cache.
|
||||
*/
|
||||
template <typename Idx, typename DType>
|
||||
__global__ void GatherMMScatterKernel(
|
||||
const DType* __restrict__ A, const DType* __restrict__ B,
|
||||
DType* __restrict__ C, const Idx* __restrict__ idx_a,
|
||||
const Idx* __restrict__ idx_b, const Idx* __restrict__ idx_c,
|
||||
const int64_t num_rows, const int64_t in_len, const int64_t out_len) {
|
||||
unsigned int tId = threadIdx.x;
|
||||
unsigned int laneId = tId & 31;
|
||||
unsigned int gId = (blockIdx.x * blockDim.x + threadIdx.x);
|
||||
unsigned int warpId = gId >> 5;
|
||||
unsigned int row = warpId;
|
||||
if (row < num_rows) {
|
||||
const unsigned int local_row =
|
||||
row & 3; // hardcoded for TB size 128 (4 warps)
|
||||
const Idx cur_rowA = (idx_a) ? idx_a[row] : row;
|
||||
const Idx cur_rowB = (idx_b) ? idx_b[row] : row;
|
||||
const Idx cur_rowC = (idx_c) ? idx_c[row] : row;
|
||||
const Idx B_offset = cur_rowB * in_len * out_len;
|
||||
const int sh_a_tile = 64;
|
||||
__shared__ DType sh_A[4 * sh_a_tile];
|
||||
int a_tile = sh_a_tile;
|
||||
for (unsigned int k_start = 0; k_start < in_len; k_start += 64) {
|
||||
if ((in_len - k_start) < a_tile) a_tile = in_len - k_start;
|
||||
// Load A in shared mem in a coalesced way
|
||||
for (unsigned int l = laneId; l < a_tile; l += 32)
|
||||
sh_A[local_row * sh_a_tile + l] = A[cur_rowA * in_len + (k_start + l)];
|
||||
__syncwarp();
|
||||
|
||||
for (unsigned int outloop = 0; outloop < out_len; outloop += 32) {
|
||||
DType out_reg = static_cast<DType>(0.0f); // thread private
|
||||
const unsigned int l = laneId;
|
||||
if (l < out_len) {
|
||||
// iterate over elements of a row of A
|
||||
for (unsigned int i = 0; i < a_tile; i++) {
|
||||
const DType a_val = sh_A[local_row * sh_a_tile + i];
|
||||
// iterate over elements of a row of B in parallel
|
||||
out_reg +=
|
||||
a_val * B[B_offset + ((i + k_start) * out_len + (outloop + l))];
|
||||
}
|
||||
if (idx_c) {
|
||||
AtomicAdd(C + cur_rowC * out_len + (outloop + l), out_reg);
|
||||
} else {
|
||||
C[cur_rowC * out_len + (outloop + l)] += out_reg;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @note Output matrix is accumulated via atomic operations. Rest of the
|
||||
* strategies are similar to GatherMMKernel. One warp is assigned to process one
|
||||
* row of A. Each WARP sequentially multiplies one element of A and a row of B
|
||||
* to compute partial result of the output. A is loaded in shared memory in a
|
||||
* coalesced way. B should get benefit from L2 cache.
|
||||
*/
|
||||
template <typename Idx, typename DType>
|
||||
__global__ void GatherMMScatterKernel2(
|
||||
const DType* __restrict__ A, const DType* __restrict__ B,
|
||||
DType* __restrict__ C, const Idx* __restrict__ idx_a,
|
||||
const Idx* __restrict__ idx_b, const Idx* __restrict__ idx_c,
|
||||
const int64_t num_rows, const int64_t in_len, const int64_t out_len) {
|
||||
unsigned int tId = threadIdx.x;
|
||||
unsigned int laneId = tId & 31;
|
||||
unsigned int gId = (blockIdx.x * blockDim.x + threadIdx.x);
|
||||
unsigned int warpId = gId >> 5;
|
||||
unsigned int row = warpId;
|
||||
if (row < num_rows) {
|
||||
const unsigned int local_row =
|
||||
row & 3; // hardcoded for TB size 128 (4 warps)
|
||||
const Idx row_a = (idx_a) ? idx_a[row] : row;
|
||||
const Idx row_b = (idx_b) ? idx_b[row] : row;
|
||||
const Idx row_c = (idx_c) ? idx_c[row] : row;
|
||||
const Idx C_offset = row_c * in_len * out_len;
|
||||
const int sh_a_tile = 64;
|
||||
__shared__ DType sh_A[4 * sh_a_tile];
|
||||
int a_tile = sh_a_tile;
|
||||
for (unsigned int k_start = 0; k_start < in_len; k_start += 64) {
|
||||
if ((in_len - k_start) < a_tile) a_tile = in_len - k_start;
|
||||
/* Load A in shared mem in a coalesced way */
|
||||
for (unsigned int l = laneId; l < a_tile; l += 32)
|
||||
sh_A[local_row * sh_a_tile + l] = A[row_a * in_len + (k_start + l)];
|
||||
__syncwarp();
|
||||
|
||||
for (unsigned int outloop = 0; outloop < out_len; outloop += 32) {
|
||||
DType out_reg = static_cast<DType>(0.0f); // thread private
|
||||
const unsigned int l = laneId;
|
||||
if (l < out_len) {
|
||||
const DType b_val = B[row_b * out_len + (outloop + l)];
|
||||
/* iterate over elements of a row of A */
|
||||
for (unsigned int i = 0; i < a_tile; i++) {
|
||||
const DType a_val = sh_A[local_row * sh_a_tile + i];
|
||||
const Idx C_idx =
|
||||
C_offset + ((i + k_start) * out_len + (outloop + l));
|
||||
AtomicAdd(C + C_idx, a_val * b_val);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
|
||||
/**
|
||||
* @brief Implementation of Gather_mm operator. The input matrix A is
|
||||
* expected to be sorted according to relation type.
|
||||
* @param A The input dense matrix of dimension m x k
|
||||
* @param B The input dense matrix of dimension k x n
|
||||
* @param C The output dense matrix of dimension m x n
|
||||
* @param seglen_A The input vector of size R. Each element
|
||||
* is the length of segments of input ``A``
|
||||
* @param a_trans Matrix A to be transposed
|
||||
* @param b_trans Matrix B to be transposed
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SegmentMM(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans) {
|
||||
auto device = runtime::DeviceAPI::Get(A->ctx);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const DType* A_data = A.Ptr<DType>();
|
||||
const DType* B_data = B.Ptr<DType>();
|
||||
const IdType* seglen_A_data = seglen_A.Ptr<IdType>();
|
||||
DType* C_data = C.Ptr<DType>();
|
||||
int64_t A_offset = 0, B_offset = 0, C_offset = 0;
|
||||
int64_t m, n, k;
|
||||
int64_t num_rel = seglen_A.NumElements();
|
||||
DType alpha = 1., beta = 0.;
|
||||
|
||||
auto* thr_entry = runtime::CUDAThreadEntry::ThreadLocal();
|
||||
if (!thr_entry->cublas_handle)
|
||||
CUBLAS_CALL(cublasCreate(&(thr_entry->cublas_handle)));
|
||||
CUBLAS_CALL(cublasSetStream(thr_entry->cublas_handle, stream));
|
||||
|
||||
IdType m_offset = 0;
|
||||
for (IdType etype = 0; etype < num_rel; ++etype) {
|
||||
m = seglen_A_data[etype]; // rows of A
|
||||
CHECK_LE(m_offset + m, A->shape[0])
|
||||
<< "Segment index out of bound of A->shape[0].";
|
||||
n = B->shape[2]; // cols of B
|
||||
k = B->shape[1]; // cols of A == rows of B
|
||||
int ldb = n, lda = k, ldc = n;
|
||||
cublasOperation_t transB = CUBLAS_OP_N;
|
||||
cublasOperation_t transA = CUBLAS_OP_N;
|
||||
if (b_trans) {
|
||||
transB = CUBLAS_OP_T;
|
||||
ldb = n, lda = n, ldc = k;
|
||||
std::swap(n, k);
|
||||
}
|
||||
CUBLAS_CALL(cublasGemm<DType>(
|
||||
thr_entry->cublas_handle, transB, transA, n, m, k, &alpha,
|
||||
B_data + B_offset, ldb, A_data + A_offset, lda, &beta,
|
||||
C_data + C_offset, ldc));
|
||||
A_offset += m * k;
|
||||
B_offset += k * n;
|
||||
C_offset += m * n;
|
||||
m_offset += m;
|
||||
}
|
||||
}
|
||||
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SegmentMMBackwardB(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen) {
|
||||
auto device = runtime::DeviceAPI::Get(A->ctx);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const DType* A_data = A.Ptr<DType>();
|
||||
const DType* dC_data = dC.Ptr<DType>();
|
||||
const IdType* seglen_data = seglen.Ptr<IdType>();
|
||||
DType* dB_data = dB.Ptr<DType>();
|
||||
int64_t A_offset = 0, dC_offset = 0, dB_offset = 0;
|
||||
int64_t m, n, k;
|
||||
int64_t num_rel = seglen.NumElements();
|
||||
DType alpha = 1., beta = 0.;
|
||||
|
||||
auto* thr_entry = runtime::CUDAThreadEntry::ThreadLocal();
|
||||
if (!thr_entry->cublas_handle)
|
||||
CUBLAS_CALL(cublasCreate(&(thr_entry->cublas_handle)));
|
||||
CUBLAS_CALL(cublasSetStream(thr_entry->cublas_handle, stream));
|
||||
|
||||
IdType k_offset = 0;
|
||||
for (IdType etype = 0; etype < num_rel; ++etype) {
|
||||
m = dC->shape[1];
|
||||
n = A->shape[1];
|
||||
k = seglen_data[etype];
|
||||
CHECK_LE(k_offset + k, A->shape[0])
|
||||
<< "Segement index out of bound of A->shape[0].";
|
||||
int lddC = m, ldA = n, lddB = m;
|
||||
cublasOperation_t trans_dC = CUBLAS_OP_N;
|
||||
cublasOperation_t trans_A = CUBLAS_OP_T;
|
||||
CUBLAS_CALL(cublasGemm<DType>(
|
||||
thr_entry->cublas_handle, trans_dC, trans_A, m, n, k, &alpha,
|
||||
dC_data + dC_offset, lddC, A_data + A_offset, ldA, &beta,
|
||||
dB_data + dB_offset, lddB));
|
||||
dC_offset += m * k;
|
||||
A_offset += n * k;
|
||||
dB_offset += m * n;
|
||||
k_offset += k;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Implementation of Gather_mm operator. The input matrix A is
|
||||
* expected to be sorted according to relation type.
|
||||
* @param A The input dense matrix of dimension m x k
|
||||
* @param B The input dense matrix of dimension k x n
|
||||
* @param C The output dense matrix of dimension m x n
|
||||
* @param idx_a The input vector to gather left hand operand on
|
||||
* @param idx_b The input vector to gather right hand operand on
|
||||
*/
|
||||
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void GatherMM(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b) {
|
||||
auto device = runtime::DeviceAPI::Get(A->ctx);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
int64_t out_len = B->shape[2]; // cols of B
|
||||
int64_t in_len = A->shape[1]; // cols of A
|
||||
const int64_t tot_num_rows = A->shape[0];
|
||||
const int ntx = 128;
|
||||
const int warp_size = 32;
|
||||
const int nbx = ((tot_num_rows * warp_size + ntx - 1) / ntx);
|
||||
const dim3 nblks(nbx);
|
||||
const dim3 nthrs(ntx);
|
||||
CUDA_KERNEL_CALL(
|
||||
(cuda::GatherMMScatterKernel<IdType, DType>), nblks, nthrs, 0, stream,
|
||||
A.Ptr<DType>(), B.Ptr<DType>(), C.Ptr<DType>(), idx_a.Ptr<IdType>(),
|
||||
idx_b.Ptr<IdType>(), nullptr, tot_num_rows, in_len, out_len);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Implementation of Gather_mm operator. The input matrix A is
|
||||
* expected to be sorted according to relation type.
|
||||
* @param A The input dense matrix of dimension m x k
|
||||
* @param B The input dense matrix of dimension k x n
|
||||
* @param C The output dense matrix of dimension m x n
|
||||
* @param idx_a The input vector to gather left hand operand on
|
||||
* @param idx_b The input vector to gather right hand operand on
|
||||
* @param idx_c The input vector to gather output operand on
|
||||
* @param num_rel The number of idx types in idx_b
|
||||
* @param a_trans Matrix A to be transposed
|
||||
* @param b_trans Matrix B to be transposed
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void GatherMMScatter(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c) {
|
||||
auto device = runtime::DeviceAPI::Get(A->ctx);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const IdType* idx_c_data = idx_c.Ptr<IdType>();
|
||||
int64_t out_len = (B->ndim == 2) ? B->shape[1] : B->shape[2]; // cols of B
|
||||
int64_t in_len = A->shape[1]; // cols of A
|
||||
int64_t tot_num_rows = A->shape[0];
|
||||
const int ntx = 128;
|
||||
const int warp_size = 32;
|
||||
const int nbx = ((tot_num_rows * warp_size + ntx - 1) / ntx);
|
||||
const dim3 nblks(nbx);
|
||||
const dim3 nthrs(ntx);
|
||||
if (B->ndim == 3) {
|
||||
CUDA_KERNEL_CALL(
|
||||
(cuda::GatherMMScatterKernel<IdType, DType>), nblks, nthrs, 0, stream,
|
||||
A.Ptr<DType>(), B.Ptr<DType>(), C.Ptr<DType>(), idx_a.Ptr<IdType>(),
|
||||
idx_b.Ptr<IdType>(), idx_c.Ptr<IdType>(), tot_num_rows, in_len,
|
||||
out_len);
|
||||
} else {
|
||||
// Custom kernel for W_grad[idx_c[i]] = H^T[i] * C.grad[i]
|
||||
// This kernel accesses rows of A in a transposed way w/o explicitly
|
||||
// converting A
|
||||
CUDA_KERNEL_CALL(
|
||||
(cuda::GatherMMScatterKernel2<IdType, DType>), nblks, nthrs, 0, stream,
|
||||
A.Ptr<DType>(), B.Ptr<DType>(), C.Ptr<DType>(), idx_a.Ptr<IdType>(),
|
||||
idx_b.Ptr<IdType>(), idx_c.Ptr<IdType>(), tot_num_rows, in_len,
|
||||
out_len);
|
||||
}
|
||||
}
|
||||
|
||||
template void GatherMM<kDGLCUDA, int32_t, __half>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b);
|
||||
template void GatherMM<kDGLCUDA, int64_t, __half>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b);
|
||||
#if BF16_ENABLED
|
||||
template void GatherMM<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b);
|
||||
template void GatherMM<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b);
|
||||
#endif // BF16_ENABLED
|
||||
template void GatherMM<kDGLCUDA, int32_t, float>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b);
|
||||
template void GatherMM<kDGLCUDA, int64_t, float>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b);
|
||||
template void GatherMM<kDGLCUDA, int32_t, double>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b);
|
||||
template void GatherMM<kDGLCUDA, int64_t, double>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b);
|
||||
|
||||
template void GatherMMScatter<kDGLCUDA, int32_t, __half>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c);
|
||||
template void GatherMMScatter<kDGLCUDA, int64_t, __half>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c);
|
||||
#if BF16_ENABLED
|
||||
template void GatherMMScatter<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c);
|
||||
template void GatherMMScatter<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c);
|
||||
#endif // BF16_ENABLED
|
||||
template void GatherMMScatter<kDGLCUDA, int32_t, float>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c);
|
||||
template void GatherMMScatter<kDGLCUDA, int64_t, float>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c);
|
||||
template void GatherMMScatter<kDGLCUDA, int32_t, double>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c);
|
||||
template void GatherMMScatter<kDGLCUDA, int64_t, double>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c);
|
||||
|
||||
template void SegmentMM<kDGLCUDA, int32_t, __half>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans);
|
||||
template void SegmentMM<kDGLCUDA, int64_t, __half>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans);
|
||||
#if BF16_ENABLED
|
||||
template void SegmentMM<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans);
|
||||
template void SegmentMM<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans);
|
||||
#endif // BF16_ENABLED
|
||||
template void SegmentMM<kDGLCUDA, int32_t, float>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans);
|
||||
template void SegmentMM<kDGLCUDA, int64_t, float>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans);
|
||||
template void SegmentMM<kDGLCUDA, int32_t, double>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans);
|
||||
template void SegmentMM<kDGLCUDA, int64_t, double>(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans);
|
||||
|
||||
template void SegmentMMBackwardB<kDGLCUDA, int32_t, __half>(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen);
|
||||
template void SegmentMMBackwardB<kDGLCUDA, int64_t, __half>(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen);
|
||||
#if BF16_ENABLED
|
||||
template void SegmentMMBackwardB<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen);
|
||||
template void SegmentMMBackwardB<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen);
|
||||
#endif // BF16_ENABLED
|
||||
template void SegmentMMBackwardB<kDGLCUDA, int32_t, float>(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen);
|
||||
template void SegmentMMBackwardB<kDGLCUDA, int64_t, float>(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen);
|
||||
template void SegmentMMBackwardB<kDGLCUDA, int32_t, double>(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen);
|
||||
template void SegmentMMBackwardB<kDGLCUDA, int64_t, double>(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen);
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,144 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/ge_spmm.cuh
|
||||
* @brief GE-SpMM CUDA kernel function header.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CUDA_GE_SPMM_CUH_
|
||||
#define DGL_ARRAY_CUDA_GE_SPMM_CUH_
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./utils.h"
|
||||
#include "atomic.cuh"
|
||||
#include "macro.cuh"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using namespace cuda;
|
||||
|
||||
namespace aten {
|
||||
namespace cuda {
|
||||
|
||||
/**
|
||||
* @brief CUDA kernel of GE-SpMM on Csr.
|
||||
* @note GE-SpMM: https://arxiv.org/pdf/2007.03179.pdf
|
||||
* The grid dimension x and y are reordered for better performance.
|
||||
*/
|
||||
template <typename Idx, typename DType, typename BinaryOp>
|
||||
__global__ void GESpMMKernel(
|
||||
const DType* __restrict__ ufeat, const DType* __restrict__ efeat,
|
||||
DType* __restrict__ out, const Idx* __restrict__ indptr,
|
||||
const Idx* __restrict__ indices, const int64_t num_rows,
|
||||
const int64_t num_cols, const int64_t feat_len) {
|
||||
const Idx rid =
|
||||
blockIdx.x * blockDim.y + threadIdx.y; // over vertices dimension
|
||||
const Idx fid = (blockIdx.y * 64) + threadIdx.x; // over feature dimension
|
||||
|
||||
if (rid < num_rows && fid < feat_len) {
|
||||
const Idx low = __ldg(indptr + rid), high = __ldg(indptr + rid + 1);
|
||||
DType accum_0 = 0., accum_1 = 0.;
|
||||
|
||||
if (blockIdx.y != gridDim.y - 1) { // fid + 32 < feat_len
|
||||
for (Idx left = low; left < high; left += 32) {
|
||||
if (left + 32 <= high) {
|
||||
#pragma unroll
|
||||
for (Idx i = 0; i < 32; ++i) {
|
||||
const Idx eid = left + i;
|
||||
const Idx cid = __ldg(indices + eid);
|
||||
const Idx offset = feat_len * cid + fid;
|
||||
if (BinaryOp::use_rhs) {
|
||||
accum_0 += BinaryOp::Call(ufeat + offset, efeat + eid);
|
||||
accum_1 += BinaryOp::Call(ufeat + offset + 32, efeat + eid);
|
||||
} else {
|
||||
accum_0 += ufeat[offset];
|
||||
accum_1 += ufeat[offset + 32];
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for (Idx i = 0; left + i < high; ++i) {
|
||||
const Idx eid = left + i;
|
||||
const Idx cid = __ldg(indices + eid);
|
||||
const Idx offset = feat_len * cid + fid;
|
||||
if (BinaryOp::use_rhs) {
|
||||
accum_0 += BinaryOp::Call(ufeat + offset, efeat + eid);
|
||||
accum_1 += BinaryOp::Call(ufeat + offset + 32, efeat + eid);
|
||||
} else {
|
||||
accum_0 += ufeat[offset];
|
||||
accum_1 += ufeat[offset + 32];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
out[feat_len * rid + fid] = accum_0;
|
||||
out[feat_len * rid + fid + 32] = accum_1;
|
||||
}
|
||||
} else {
|
||||
const Idx fid_0 = fid < feat_len ? fid : 0,
|
||||
fid_1 = fid + 32 < feat_len ? fid + 32 : 0;
|
||||
for (int left = low; left < high; left += 32) {
|
||||
if (left + 32 <= high) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 32; ++i) {
|
||||
const Idx eid = left + i;
|
||||
const Idx cid = __ldg(indices + eid);
|
||||
const Idx offset = feat_len * cid;
|
||||
if (BinaryOp::use_rhs) {
|
||||
accum_0 += BinaryOp::Call(ufeat + offset + fid_0, efeat + eid);
|
||||
accum_1 += BinaryOp::Call(ufeat + offset + fid_1, efeat + eid);
|
||||
} else {
|
||||
accum_0 += ufeat[offset + fid_0];
|
||||
accum_1 += ufeat[offset + fid_1];
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for (int i = 0; i + left < high; ++i) {
|
||||
const Idx eid = left + i;
|
||||
const Idx cid = __ldg(indices + eid);
|
||||
const Idx offset = feat_len * cid;
|
||||
if (BinaryOp::use_rhs) {
|
||||
accum_0 += BinaryOp::Call(ufeat + offset + fid_0, efeat + eid);
|
||||
accum_1 += BinaryOp::Call(ufeat + offset + fid_1, efeat + eid);
|
||||
} else {
|
||||
accum_0 += ufeat[offset + fid_0];
|
||||
accum_1 += ufeat[offset + fid_1];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
out[feat_len * rid + fid] = accum_0;
|
||||
if (fid + 32 < feat_len) out[feat_len * rid + fid + 32] = accum_1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Idx, typename DType, typename BinaryOp>
|
||||
void GESpMMCsr(
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
int64_t feat_len) {
|
||||
const Idx* indptr = csr.indptr.Ptr<Idx>();
|
||||
const Idx* indices = csr.indices.Ptr<Idx>();
|
||||
const DType* ufeat_data = ufeat.Ptr<DType>();
|
||||
const DType* efeat_data = efeat.Ptr<DType>();
|
||||
DType* out_data = out.Ptr<DType>();
|
||||
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
|
||||
const int ntx = 32;
|
||||
const int nty = 32;
|
||||
const int nby = (feat_len + (ntx * 2) - 1) / (ntx * 2);
|
||||
const int nbx = (csr.num_rows + nty - 1) / nty;
|
||||
const dim3 nblks(nbx, nby);
|
||||
const dim3 nthrs(ntx, nty);
|
||||
const int sh_mem_size = 0;
|
||||
|
||||
CUDA_KERNEL_CALL(
|
||||
(GESpMMKernel<Idx, DType, BinaryOp>), nblks, nthrs, sh_mem_size, stream,
|
||||
ufeat_data, efeat_data, out_data, indptr, indices, csr.num_rows,
|
||||
csr.num_cols, feat_len);
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CUDA_GE_SPMM_CUH_
|
||||
@@ -0,0 +1,833 @@
|
||||
/*!
|
||||
* Copyright (c) 2022, NVIDIA Corporation
|
||||
* Copyright (c) 2022, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
|
||||
* All rights reserved.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*
|
||||
* @file array/cuda/labor_sampling.cu
|
||||
* @brief labor sampling
|
||||
*/
|
||||
|
||||
#include <dgl/aten/coo.h>
|
||||
#include <dgl/random.h>
|
||||
#include <dgl/runtime/device_api.h>
|
||||
#include <thrust/binary_search.h>
|
||||
#include <thrust/copy.h>
|
||||
#include <thrust/execution_policy.h>
|
||||
#include <thrust/gather.h>
|
||||
#include <thrust/iterator/constant_iterator.h>
|
||||
#include <thrust/iterator/counting_iterator.h>
|
||||
#include <thrust/iterator/zip_iterator.h>
|
||||
#include <thrust/reduce.h>
|
||||
#include <thrust/shuffle.h>
|
||||
#include <thrust/transform.h>
|
||||
#include <thrust/zip_function.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cub/cub.cuh> // NOLINT
|
||||
#include <limits>
|
||||
#include <numeric>
|
||||
#include <type_traits>
|
||||
#include <utility>
|
||||
|
||||
#include "../../array/cuda/utils.h"
|
||||
#include "../../random/continuous_seed.h"
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./functor.cuh"
|
||||
#include "./spmm.cuh"
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
using dgl::random::continuous_seed;
|
||||
|
||||
constexpr int BLOCK_SIZE = 128;
|
||||
constexpr int CTA_SIZE = 128;
|
||||
constexpr double eps = 0.0001;
|
||||
|
||||
namespace {
|
||||
|
||||
template <typename IdType>
|
||||
struct TransformOp {
|
||||
const IdType* idx_coo;
|
||||
const IdType* rows;
|
||||
const IdType* indptr;
|
||||
const IdType* subindptr;
|
||||
const IdType* indices;
|
||||
const IdType* data_arr;
|
||||
bool is_pinned;
|
||||
__host__ __device__ auto operator()(IdType idx) {
|
||||
const auto in_row = idx_coo[idx];
|
||||
const auto row = rows[in_row];
|
||||
const auto in_idx = indptr[in_row] + idx - subindptr[in_row];
|
||||
const auto u = indices[is_pinned ? idx : in_idx];
|
||||
const auto data = data_arr ? data_arr[in_idx] : in_idx;
|
||||
return thrust::make_tuple(row, u, data);
|
||||
}
|
||||
};
|
||||
|
||||
template <
|
||||
typename IdType, typename FloatType, typename probs_t, typename A_t,
|
||||
typename B_t>
|
||||
struct TransformOpImp {
|
||||
probs_t probs;
|
||||
A_t A;
|
||||
B_t B;
|
||||
const IdType* idx_coo;
|
||||
const IdType* rows;
|
||||
const FloatType* cs;
|
||||
const IdType* indptr;
|
||||
const IdType* subindptr;
|
||||
const IdType* indices;
|
||||
const IdType* data_arr;
|
||||
bool is_pinned;
|
||||
__host__ __device__ auto operator()(IdType idx) {
|
||||
const auto ps = probs[idx];
|
||||
const auto in_row = idx_coo[idx];
|
||||
const auto c = cs[in_row];
|
||||
const auto row = rows[in_row];
|
||||
const auto in_idx = indptr[in_row] + idx - subindptr[in_row];
|
||||
const auto u = indices[is_pinned ? idx : in_idx];
|
||||
const auto w = A[in_idx];
|
||||
const auto w2 = B[in_idx];
|
||||
const auto data = data_arr ? data_arr[in_idx] : in_idx;
|
||||
return thrust::make_tuple(
|
||||
in_row, row, u, data, w / min((FloatType)1, c * w2 * ps));
|
||||
}
|
||||
};
|
||||
|
||||
template <typename FloatType>
|
||||
struct StencilOp {
|
||||
const FloatType* cs;
|
||||
template <typename IdType>
|
||||
__host__ __device__ auto operator()(
|
||||
IdType in_row, FloatType ps, FloatType rnd) {
|
||||
return rnd <= cs[in_row] * ps;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename IdType, typename FloatType, typename ps_t, typename A_t>
|
||||
struct StencilOpFused {
|
||||
const continuous_seed seed;
|
||||
const IdType* idx_coo;
|
||||
const FloatType* cs;
|
||||
const ps_t probs;
|
||||
const A_t A;
|
||||
const IdType* subindptr;
|
||||
const IdType* indptr;
|
||||
const IdType* indices;
|
||||
const IdType* nids;
|
||||
bool is_pinned;
|
||||
__device__ auto operator()(IdType idx) {
|
||||
const auto in_row = idx_coo[idx];
|
||||
const auto ps = probs[idx];
|
||||
IdType rofs = idx - subindptr[in_row];
|
||||
const auto in_idx = indptr[in_row] + rofs;
|
||||
const auto u = indices[is_pinned ? idx : in_idx];
|
||||
const auto t = nids ? nids[u] : u; // t in the paper
|
||||
// rolled random number r_t is a function of the random_seed and t
|
||||
const float rnd = seed.uniform(t);
|
||||
return rnd <= cs[in_row] * A[in_idx] * ps;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename IdType, typename FloatType>
|
||||
struct TransformOpMean {
|
||||
const IdType* ds;
|
||||
const FloatType* ws;
|
||||
__host__ __device__ auto operator()(IdType idx, FloatType ps) {
|
||||
return ps * ds[idx] / ws[idx];
|
||||
}
|
||||
};
|
||||
|
||||
struct TransformOpMinWith1 {
|
||||
template <typename FloatType>
|
||||
__host__ __device__ auto operator()(FloatType x) {
|
||||
return min((FloatType)1, x);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename IdType>
|
||||
struct IndptrFunc {
|
||||
const IdType* indptr;
|
||||
const IdType* in_deg;
|
||||
__host__ __device__ auto operator()(IdType row) {
|
||||
return indptr[row] + (in_deg ? in_deg[row] : 0);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename FloatType>
|
||||
struct SquareFunc {
|
||||
__host__ __device__ auto operator()(FloatType x) {
|
||||
return thrust::make_tuple(x, x * x);
|
||||
}
|
||||
};
|
||||
|
||||
struct TupleSum {
|
||||
template <typename T>
|
||||
__host__ __device__ T operator()(const T& a, const T& b) const {
|
||||
return thrust::make_tuple(
|
||||
thrust::get<0>(a) + thrust::get<0>(b),
|
||||
thrust::get<1>(a) + thrust::get<1>(b));
|
||||
}
|
||||
};
|
||||
|
||||
template <typename IdType, typename FloatType>
|
||||
struct DegreeFunc {
|
||||
const IdType num_picks;
|
||||
const IdType* rows;
|
||||
const IdType* indptr;
|
||||
IdType* in_deg;
|
||||
IdType* inrow_indptr;
|
||||
FloatType* cs;
|
||||
__host__ __device__ auto operator()(IdType tIdx) {
|
||||
const auto out_row = rows[tIdx];
|
||||
const auto indptr_val = indptr[out_row];
|
||||
const auto d = indptr[out_row + 1] - indptr_val;
|
||||
in_deg[tIdx] = d;
|
||||
inrow_indptr[tIdx] = indptr_val;
|
||||
cs[tIdx] = num_picks / (FloatType)d;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename IdType, typename FloatType>
|
||||
__global__ void _CSRRowWiseOneHopExtractorKernel(
|
||||
const continuous_seed seed, const IdType hop_size,
|
||||
const IdType* const indptr, const IdType* const subindptr,
|
||||
const IdType* const indices, const IdType* const idx_coo,
|
||||
const IdType* const nids, const FloatType* const A, FloatType* const rands,
|
||||
IdType* const hop, FloatType* const A_l) {
|
||||
IdType tx = static_cast<IdType>(blockIdx.x) * blockDim.x + threadIdx.x;
|
||||
const int stride_x = gridDim.x * blockDim.x;
|
||||
|
||||
while (tx < hop_size) {
|
||||
IdType rpos = idx_coo[tx];
|
||||
IdType rofs = tx - subindptr[rpos];
|
||||
const auto in_idx = indptr[rpos] + rofs;
|
||||
const auto not_pinned = indices != hop;
|
||||
const auto u = indices[not_pinned ? in_idx : tx];
|
||||
if (not_pinned) hop[tx] = u;
|
||||
const auto t = nids ? nids[u] : u;
|
||||
if (A) A_l[tx] = A[in_idx];
|
||||
// rolled random number r_t is a function of the random_seed and t
|
||||
rands[tx] = (FloatType)seed.uniform(t);
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
constexpr int CACHE_LINE_SIZE = 128;
|
||||
|
||||
template <typename IdType>
|
||||
struct AlignmentFunc {
|
||||
static_assert(CACHE_LINE_SIZE % sizeof(IdType) == 0);
|
||||
const IdType* in_deg;
|
||||
const int64_t* perm;
|
||||
IdType num_rows;
|
||||
__host__ __device__ auto operator()(IdType row) {
|
||||
constexpr int num_elements = CACHE_LINE_SIZE / sizeof(IdType);
|
||||
return in_deg[perm ? perm[row % num_rows] : row] + num_elements - 1;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename IdType>
|
||||
__global__ void _CSRRowWiseOneHopExtractorAlignedKernel(
|
||||
const IdType hop_size, const IdType num_rows, const IdType* const indptr,
|
||||
const IdType* const subindptr, const IdType* const subindptr_aligned,
|
||||
const IdType* const indices, IdType* const hop, const int64_t* const perm) {
|
||||
IdType tx = static_cast<IdType>(blockIdx.x) * blockDim.x + threadIdx.x;
|
||||
const int stride_x = gridDim.x * blockDim.x;
|
||||
|
||||
while (tx < hop_size) {
|
||||
const IdType rpos_ =
|
||||
dgl::cuda::_UpperBound(subindptr_aligned, num_rows, tx) - 1;
|
||||
const IdType rpos = perm ? perm[rpos_] : rpos_;
|
||||
const auto out_row = subindptr[rpos];
|
||||
const auto d = subindptr[rpos + 1] - out_row;
|
||||
const int offset =
|
||||
((uint64_t)(indices + indptr[rpos] - subindptr_aligned[rpos_]) %
|
||||
CACHE_LINE_SIZE) /
|
||||
sizeof(IdType);
|
||||
const IdType rofs = tx - subindptr_aligned[rpos_] - offset;
|
||||
if (rofs >= 0 && rofs < d) {
|
||||
const auto in_idx = indptr[rpos] + rofs;
|
||||
assert((uint64_t)(indices + in_idx - tx) % CACHE_LINE_SIZE == 0);
|
||||
const auto u = indices[in_idx];
|
||||
hop[out_row + rofs] = u;
|
||||
}
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename IdType, typename FloatType, int BLOCK_CTAS, int TILE_SIZE>
|
||||
__global__ void _CSRRowWiseLayerSampleDegreeKernel(
|
||||
const IdType num_picks, const IdType num_rows, FloatType* const cs,
|
||||
const FloatType* const ds, const FloatType* const d2s,
|
||||
const IdType* const indptr, const FloatType* const probs,
|
||||
const FloatType* const A, const IdType* const subindptr) {
|
||||
typedef cub::BlockReduce<FloatType, BLOCK_SIZE> BlockReduce;
|
||||
__shared__ typename BlockReduce::TempStorage temp_storage;
|
||||
__shared__ FloatType var_1_bcast[BLOCK_CTAS];
|
||||
|
||||
// we assign one warp per row
|
||||
assert(blockDim.x == CTA_SIZE);
|
||||
assert(blockDim.y == BLOCK_CTAS);
|
||||
|
||||
IdType out_row = blockIdx.x * TILE_SIZE + threadIdx.y;
|
||||
const auto last_row =
|
||||
min(static_cast<IdType>(blockIdx.x + 1) * TILE_SIZE, num_rows);
|
||||
|
||||
constexpr FloatType ONE = 1;
|
||||
|
||||
while (out_row < last_row) {
|
||||
const auto in_row_start = indptr[out_row];
|
||||
const auto out_row_start = subindptr[out_row];
|
||||
|
||||
const IdType degree = subindptr[out_row + 1] - out_row_start;
|
||||
|
||||
if (degree > 0) {
|
||||
// stands for k in in arXiv:2210.13339, i.e. fanout
|
||||
const auto k = min(num_picks, degree);
|
||||
// slightly better than NS
|
||||
const FloatType d_ = ds ? ds[out_row] : degree;
|
||||
// stands for right handside of Equation (22) in arXiv:2210.13339
|
||||
FloatType var_target =
|
||||
d_ * d_ / k + (ds ? d2s[out_row] - d_ * d_ / degree : 0);
|
||||
|
||||
auto c = cs[out_row];
|
||||
const int num_valid = min(degree, (IdType)CTA_SIZE);
|
||||
// stands for left handside of Equation (22) in arXiv:2210.13339
|
||||
FloatType var_1;
|
||||
do {
|
||||
var_1 = 0;
|
||||
if (A) {
|
||||
for (int idx = threadIdx.x; idx < degree; idx += CTA_SIZE) {
|
||||
const auto w = A[in_row_start + idx];
|
||||
const auto ps = probs ? probs[out_row_start + idx] : w;
|
||||
var_1 += w > 0 ? w * w / min(ONE, c * ps) : 0;
|
||||
}
|
||||
} else {
|
||||
for (int idx = threadIdx.x; idx < degree; idx += CTA_SIZE) {
|
||||
const auto ps = probs[out_row_start + idx];
|
||||
var_1 += 1 / min(ONE, c * ps);
|
||||
}
|
||||
}
|
||||
var_1 = BlockReduce(temp_storage).Sum(var_1, num_valid);
|
||||
if (threadIdx.x == 0) var_1_bcast[threadIdx.y] = var_1;
|
||||
__syncthreads();
|
||||
var_1 = var_1_bcast[threadIdx.y];
|
||||
|
||||
c *= var_1 / var_target;
|
||||
} while (min(var_1, var_target) / max(var_1, var_target) < 1 - eps);
|
||||
|
||||
if (threadIdx.x == 0) cs[out_row] = c;
|
||||
}
|
||||
|
||||
out_row += BLOCK_CTAS;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
template <typename IdType>
|
||||
int log_size(const IdType size) {
|
||||
if (size <= 0) return 0;
|
||||
for (int i = 0; i < static_cast<int>(sizeof(IdType)) * 8; i++)
|
||||
if (((size - 1) >> i) == 0) return i;
|
||||
return sizeof(IdType) * 8;
|
||||
}
|
||||
|
||||
template <typename IdType, typename FloatType, typename exec_policy_t>
|
||||
void compute_importance_sampling_probabilities(
|
||||
CSRMatrix mat, const IdType hop_size, cudaStream_t stream,
|
||||
const continuous_seed seed, const IdType num_rows, const IdType* indptr,
|
||||
const IdType* subindptr, const IdType* indices, IdArray idx_coo_arr,
|
||||
const IdType* nids,
|
||||
FloatArray cs_arr, // holds the computed cs values, has size num_rows
|
||||
const bool weighted, const FloatType* A, const FloatType* ds,
|
||||
const FloatType* d2s, const IdType num_picks, DGLContext ctx,
|
||||
const runtime::CUDAWorkspaceAllocator& allocator,
|
||||
const exec_policy_t& exec_policy, const int importance_sampling,
|
||||
IdType* hop_1, // holds the contiguous one-hop neighborhood, has size |E|
|
||||
FloatType* rands, // holds the rolled random numbers r_t for each edge, has
|
||||
// size |E|
|
||||
FloatType* probs_found) { // holds the computed pi_t values for each edge,
|
||||
// has size |E|
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
auto idx_coo = idx_coo_arr.Ptr<IdType>();
|
||||
auto cs = cs_arr.Ptr<FloatType>();
|
||||
FloatArray A_l_arr = weighted
|
||||
? NewFloatArray(hop_size, ctx, sizeof(FloatType) * 8)
|
||||
: NullArray();
|
||||
auto A_l = A_l_arr.Ptr<FloatType>();
|
||||
|
||||
const int max_log_num_vertices = log_size(mat.num_cols);
|
||||
|
||||
{ // extracts the onehop neighborhood cols to a contiguous range into hop_1
|
||||
const dim3 block(BLOCK_SIZE);
|
||||
const dim3 grid((hop_size + BLOCK_SIZE - 1) / BLOCK_SIZE);
|
||||
CUDA_KERNEL_CALL(
|
||||
(_CSRRowWiseOneHopExtractorKernel<IdType, FloatType>), grid, block, 0,
|
||||
stream, seed, hop_size, indptr, subindptr, indices, idx_coo, nids,
|
||||
weighted ? A : nullptr, rands, hop_1, A_l);
|
||||
}
|
||||
int64_t hop_uniq_size = 0;
|
||||
IdArray hop_new_arr = NewIdArray(hop_size, ctx, sizeof(IdType) * 8);
|
||||
auto hop_new = hop_new_arr.Ptr<IdType>();
|
||||
auto hop_unique = allocator.alloc_unique<IdType>(hop_size);
|
||||
// After this block, hop_unique holds the unique set of one-hop neighborhood
|
||||
// and hop_new holds the relabeled hop_1, idx_coo already holds relabeled
|
||||
// destination. hop_unique[hop_new] == hop_1 holds
|
||||
{
|
||||
auto hop_2 = allocator.alloc_unique<IdType>(hop_size);
|
||||
auto hop_3 = allocator.alloc_unique<IdType>(hop_size);
|
||||
|
||||
device->CopyDataFromTo(
|
||||
hop_1, 0, hop_2.get(), 0, sizeof(IdType) * hop_size, ctx, ctx,
|
||||
mat.indptr->dtype);
|
||||
|
||||
cub::DoubleBuffer<IdType> hop_b(hop_2.get(), hop_3.get());
|
||||
|
||||
{
|
||||
std::size_t temp_storage_bytes = 0;
|
||||
CUDA_CALL(cub::DeviceRadixSort::SortKeys(
|
||||
nullptr, temp_storage_bytes, hop_b, hop_size, 0, max_log_num_vertices,
|
||||
stream));
|
||||
|
||||
auto temp = allocator.alloc_unique<char>(temp_storage_bytes);
|
||||
|
||||
CUDA_CALL(cub::DeviceRadixSort::SortKeys(
|
||||
temp.get(), temp_storage_bytes, hop_b, hop_size, 0,
|
||||
max_log_num_vertices, stream));
|
||||
}
|
||||
|
||||
auto hop_counts = allocator.alloc_unique<IdType>(hop_size + 1);
|
||||
auto hop_unique_size = allocator.alloc_unique<int64_t>(1);
|
||||
|
||||
{
|
||||
std::size_t temp_storage_bytes = 0;
|
||||
CUDA_CALL(cub::DeviceRunLengthEncode::Encode(
|
||||
nullptr, temp_storage_bytes, hop_b.Current(), hop_unique.get(),
|
||||
hop_counts.get(), hop_unique_size.get(), hop_size, stream));
|
||||
|
||||
auto temp = allocator.alloc_unique<char>(temp_storage_bytes);
|
||||
|
||||
CUDA_CALL(cub::DeviceRunLengthEncode::Encode(
|
||||
temp.get(), temp_storage_bytes, hop_b.Current(), hop_unique.get(),
|
||||
hop_counts.get(), hop_unique_size.get(), hop_size, stream));
|
||||
|
||||
device->CopyDataFromTo(
|
||||
hop_unique_size.get(), 0, &hop_uniq_size, 0, sizeof(hop_uniq_size),
|
||||
ctx, DGLContext{kDGLCPU, 0}, mat.indptr->dtype);
|
||||
}
|
||||
|
||||
thrust::lower_bound(
|
||||
exec_policy, hop_unique.get(), hop_unique.get() + hop_uniq_size, hop_1,
|
||||
hop_1 + hop_size, hop_new);
|
||||
}
|
||||
|
||||
// @todo Consider creating a CSC because the SpMV will be done multiple times.
|
||||
COOMatrix rmat(
|
||||
num_rows, hop_uniq_size, idx_coo_arr, hop_new_arr, NullArray(), true,
|
||||
mat.sorted);
|
||||
|
||||
BcastOff bcast_off;
|
||||
bcast_off.use_bcast = false;
|
||||
bcast_off.out_len = 1;
|
||||
bcast_off.lhs_len = 1;
|
||||
bcast_off.rhs_len = 1;
|
||||
|
||||
FloatArray probs_arr =
|
||||
NewFloatArray(hop_uniq_size, ctx, sizeof(FloatType) * 8);
|
||||
auto probs_1 = probs_arr.Ptr<FloatType>();
|
||||
FloatArray probs_arr_2 =
|
||||
NewFloatArray(hop_uniq_size, ctx, sizeof(FloatType) * 8);
|
||||
auto probs = probs_arr_2.Ptr<FloatType>();
|
||||
auto arg_u = NewIdArray(hop_uniq_size, ctx, sizeof(IdType) * 8);
|
||||
auto arg_e = NewIdArray(hop_size, ctx, sizeof(IdType) * 8);
|
||||
|
||||
double prev_ex_nodes = hop_uniq_size;
|
||||
|
||||
for (int iters = 0; iters < importance_sampling || importance_sampling < 0;
|
||||
iters++) {
|
||||
if (weighted && iters == 0) {
|
||||
cuda::SpMMCoo<
|
||||
IdType, FloatType, cuda::binary::Mul<FloatType>,
|
||||
cuda::reduce::Max<IdType, FloatType, true>>(
|
||||
bcast_off, rmat, cs_arr, A_l_arr, probs_arr_2, arg_u, arg_e);
|
||||
} else {
|
||||
cuda::SpMMCoo<
|
||||
IdType, FloatType, cuda::binary::CopyLhs<FloatType>,
|
||||
cuda::reduce::Max<IdType, FloatType, true>>(
|
||||
bcast_off, rmat, cs_arr, NullArray(), iters ? probs_arr : probs_arr_2,
|
||||
arg_u, arg_e);
|
||||
}
|
||||
|
||||
if (iters)
|
||||
thrust::transform(
|
||||
exec_policy, probs_1, probs_1 + hop_uniq_size, probs, probs,
|
||||
thrust::multiplies<FloatType>{});
|
||||
|
||||
thrust::gather(
|
||||
exec_policy, hop_new, hop_new + hop_size, probs, probs_found);
|
||||
|
||||
{
|
||||
constexpr int BLOCK_CTAS = BLOCK_SIZE / CTA_SIZE;
|
||||
// the number of rows each thread block will cover
|
||||
constexpr int TILE_SIZE = BLOCK_CTAS;
|
||||
const dim3 block(CTA_SIZE, BLOCK_CTAS);
|
||||
const dim3 grid((num_rows + TILE_SIZE - 1) / TILE_SIZE);
|
||||
CUDA_KERNEL_CALL(
|
||||
(_CSRRowWiseLayerSampleDegreeKernel<
|
||||
IdType, FloatType, BLOCK_CTAS, TILE_SIZE>),
|
||||
grid, block, 0, stream, (IdType)num_picks, num_rows, cs,
|
||||
weighted ? ds : nullptr, weighted ? d2s : nullptr, indptr,
|
||||
probs_found, A, subindptr);
|
||||
}
|
||||
|
||||
{
|
||||
auto probs_min_1 =
|
||||
thrust::make_transform_iterator(probs, TransformOpMinWith1{});
|
||||
const double cur_ex_nodes = thrust::reduce(
|
||||
exec_policy, probs_min_1, probs_min_1 + hop_uniq_size, 0.0);
|
||||
if (cur_ex_nodes / prev_ex_nodes >= 1 - eps) break;
|
||||
prev_ex_nodes = cur_ex_nodes;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/////////////////////////////// CSR ///////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename FloatType>
|
||||
std::pair<COOMatrix, FloatArray> CSRLaborSampling(
|
||||
CSRMatrix mat, IdArray rows_arr, const int64_t num_picks,
|
||||
FloatArray prob_arr, const int importance_sampling, IdArray random_seed_arr,
|
||||
float seed2_contribution, IdArray NIDs) {
|
||||
const bool weighted = !IsNullArray(prob_arr);
|
||||
|
||||
const auto& ctx = rows_arr->ctx;
|
||||
|
||||
runtime::CUDAWorkspaceAllocator allocator(ctx);
|
||||
|
||||
const auto stream = runtime::getCurrentCUDAStream();
|
||||
const auto exec_policy = thrust::cuda::par_nosync(allocator).on(stream);
|
||||
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
|
||||
const IdType num_rows = rows_arr->shape[0];
|
||||
IdType* const rows = rows_arr.Ptr<IdType>();
|
||||
IdType* const nids = IsNullArray(NIDs) ? nullptr : NIDs.Ptr<IdType>();
|
||||
FloatType* const A = prob_arr.Ptr<FloatType>();
|
||||
|
||||
IdType* const indptr_ = mat.indptr.Ptr<IdType>();
|
||||
IdType* const indices_ = mat.indices.Ptr<IdType>();
|
||||
IdType* const data = CSRHasData(mat) ? mat.data.Ptr<IdType>() : nullptr;
|
||||
|
||||
// Read indptr only once in case it is pinned and access is slow.
|
||||
auto indptr = allocator.alloc_unique<IdType>(num_rows);
|
||||
// compute in-degrees
|
||||
auto in_deg = allocator.alloc_unique<IdType>(num_rows + 1);
|
||||
// cs stands for c_s in arXiv:2210.13339
|
||||
FloatArray cs_arr = NewFloatArray(num_rows, ctx, sizeof(FloatType) * 8);
|
||||
auto cs = cs_arr.Ptr<FloatType>();
|
||||
// ds stands for A_{*s} in arXiv:2210.13339
|
||||
FloatArray ds_arr = weighted
|
||||
? NewFloatArray(num_rows, ctx, sizeof(FloatType) * 8)
|
||||
: NullArray();
|
||||
auto ds = ds_arr.Ptr<FloatType>();
|
||||
// d2s stands for (A^2)_{*s} in arXiv:2210.13339, ^2 is elementwise.
|
||||
FloatArray d2s_arr = weighted
|
||||
? NewFloatArray(num_rows, ctx, sizeof(FloatType) * 8)
|
||||
: NullArray();
|
||||
auto d2s = d2s_arr.Ptr<FloatType>();
|
||||
|
||||
thrust::counting_iterator<IdType> iota(0);
|
||||
thrust::for_each(
|
||||
exec_policy, iota, iota + num_rows,
|
||||
DegreeFunc<IdType, FloatType>{
|
||||
(IdType)num_picks, rows, indptr_, in_deg.get(), indptr.get(), cs});
|
||||
|
||||
if (weighted) {
|
||||
auto b_offsets = thrust::make_transform_iterator(
|
||||
iota, IndptrFunc<IdType>{indptr.get(), nullptr});
|
||||
auto e_offsets = thrust::make_transform_iterator(
|
||||
iota, IndptrFunc<IdType>{indptr.get(), in_deg.get()});
|
||||
|
||||
auto A_A2 = thrust::make_transform_iterator(A, SquareFunc<FloatType>{});
|
||||
auto ds_d2s = thrust::make_zip_iterator(ds, d2s);
|
||||
|
||||
size_t prefix_temp_size = 0;
|
||||
CUDA_CALL(cub::DeviceSegmentedReduce::Reduce(
|
||||
nullptr, prefix_temp_size, A_A2, ds_d2s, num_rows, b_offsets, e_offsets,
|
||||
TupleSum{}, thrust::make_tuple((FloatType)0, (FloatType)0), stream));
|
||||
auto temp = allocator.alloc_unique<char>(prefix_temp_size);
|
||||
CUDA_CALL(cub::DeviceSegmentedReduce::Reduce(
|
||||
temp.get(), prefix_temp_size, A_A2, ds_d2s, num_rows, b_offsets,
|
||||
e_offsets, TupleSum{}, thrust::make_tuple((FloatType)0, (FloatType)0),
|
||||
stream));
|
||||
}
|
||||
|
||||
// fill subindptr
|
||||
IdArray subindptr_arr = NewIdArray(num_rows + 1, ctx, sizeof(IdType) * 8);
|
||||
auto subindptr = subindptr_arr.Ptr<IdType>();
|
||||
|
||||
IdType hop_size;
|
||||
{
|
||||
size_t prefix_temp_size = 0;
|
||||
CUDA_CALL(cub::DeviceScan::ExclusiveSum(
|
||||
nullptr, prefix_temp_size, in_deg.get(), subindptr, num_rows + 1,
|
||||
stream));
|
||||
auto temp = allocator.alloc_unique<char>(prefix_temp_size);
|
||||
CUDA_CALL(cub::DeviceScan::ExclusiveSum(
|
||||
temp.get(), prefix_temp_size, in_deg.get(), subindptr, num_rows + 1,
|
||||
stream));
|
||||
|
||||
device->CopyDataFromTo(
|
||||
subindptr, num_rows * sizeof(hop_size), &hop_size, 0, sizeof(hop_size),
|
||||
ctx, DGLContext{kDGLCPU, 0}, mat.indptr->dtype);
|
||||
}
|
||||
IdArray hop_arr = NewIdArray(hop_size, ctx, sizeof(IdType) * 8);
|
||||
CSRMatrix smat(
|
||||
num_rows, mat.num_cols, subindptr_arr, hop_arr, NullArray(), mat.sorted);
|
||||
// @todo Consider fusing CSRToCOO into StencilOpFused kernel
|
||||
auto smatcoo = CSRToCOO(smat, false);
|
||||
|
||||
auto idx_coo_arr = smatcoo.row;
|
||||
auto idx_coo = idx_coo_arr.Ptr<IdType>();
|
||||
|
||||
auto hop_1 = hop_arr.Ptr<IdType>();
|
||||
const bool is_pinned = mat.indices.IsPinned();
|
||||
if (is_pinned) {
|
||||
const auto res = Sort(rows_arr, log_size(mat.num_rows));
|
||||
const int64_t* perm = static_cast<int64_t*>(res.second->data);
|
||||
|
||||
IdType hop_size; // Shadows the original one as this is temporary
|
||||
auto subindptr_aligned = allocator.alloc_unique<IdType>(num_rows + 1);
|
||||
{
|
||||
auto modified_in_deg = thrust::make_transform_iterator(
|
||||
iota, AlignmentFunc<IdType>{in_deg.get(), perm, num_rows});
|
||||
size_t prefix_temp_size = 0;
|
||||
CUDA_CALL(cub::DeviceScan::ExclusiveSum(
|
||||
nullptr, prefix_temp_size, modified_in_deg, subindptr_aligned.get(),
|
||||
num_rows + 1, stream));
|
||||
auto temp = allocator.alloc_unique<char>(prefix_temp_size);
|
||||
CUDA_CALL(cub::DeviceScan::ExclusiveSum(
|
||||
temp.get(), prefix_temp_size, modified_in_deg,
|
||||
subindptr_aligned.get(), num_rows + 1, stream));
|
||||
|
||||
device->CopyDataFromTo(
|
||||
subindptr_aligned.get(), num_rows * sizeof(hop_size), &hop_size, 0,
|
||||
sizeof(hop_size), ctx, DGLContext{kDGLCPU, 0}, mat.indptr->dtype);
|
||||
}
|
||||
const dim3 block(BLOCK_SIZE);
|
||||
const dim3 grid((hop_size + BLOCK_SIZE - 1) / BLOCK_SIZE);
|
||||
CUDA_KERNEL_CALL(
|
||||
(_CSRRowWiseOneHopExtractorAlignedKernel<IdType>), grid, block, 0,
|
||||
stream, hop_size, num_rows, indptr.get(), subindptr,
|
||||
subindptr_aligned.get(), indices_, hop_1, perm);
|
||||
}
|
||||
const auto indices = is_pinned ? hop_1 : indices_;
|
||||
|
||||
auto rands =
|
||||
allocator.alloc_unique<FloatType>(importance_sampling ? hop_size : 1);
|
||||
auto probs_found =
|
||||
allocator.alloc_unique<FloatType>(importance_sampling ? hop_size : 1);
|
||||
|
||||
if (weighted) {
|
||||
// Recompute c for weighted graphs.
|
||||
constexpr int BLOCK_CTAS = BLOCK_SIZE / CTA_SIZE;
|
||||
// the number of rows each thread block will cover
|
||||
constexpr int TILE_SIZE = BLOCK_CTAS;
|
||||
const dim3 block(CTA_SIZE, BLOCK_CTAS);
|
||||
const dim3 grid((num_rows + TILE_SIZE - 1) / TILE_SIZE);
|
||||
CUDA_KERNEL_CALL(
|
||||
(_CSRRowWiseLayerSampleDegreeKernel<
|
||||
IdType, FloatType, BLOCK_CTAS, TILE_SIZE>),
|
||||
grid, block, 0, stream, (IdType)num_picks, num_rows, cs, ds, d2s,
|
||||
indptr.get(), nullptr, A, subindptr);
|
||||
}
|
||||
|
||||
const continuous_seed random_seed =
|
||||
IsNullArray(random_seed_arr)
|
||||
? continuous_seed(RandomEngine::ThreadLocal()->RandInt(1000000000))
|
||||
: continuous_seed(random_seed_arr, seed2_contribution);
|
||||
|
||||
if (importance_sampling)
|
||||
compute_importance_sampling_probabilities<
|
||||
IdType, FloatType, decltype(exec_policy)>(
|
||||
mat, hop_size, stream, random_seed, num_rows, indptr.get(), subindptr,
|
||||
indices, idx_coo_arr, nids, cs_arr, weighted, A, ds, d2s,
|
||||
(IdType)num_picks, ctx, allocator, exec_policy, importance_sampling,
|
||||
hop_1, rands.get(), probs_found.get());
|
||||
|
||||
IdArray picked_row = NewIdArray(hop_size, ctx, sizeof(IdType) * 8);
|
||||
IdArray picked_col = NewIdArray(hop_size, ctx, sizeof(IdType) * 8);
|
||||
IdArray picked_idx = NewIdArray(hop_size, ctx, sizeof(IdType) * 8);
|
||||
FloatArray picked_imp =
|
||||
importance_sampling || weighted
|
||||
? NewFloatArray(hop_size, ctx, sizeof(FloatType) * 8)
|
||||
: NullArray();
|
||||
|
||||
IdType* const picked_row_data = picked_row.Ptr<IdType>();
|
||||
IdType* const picked_col_data = picked_col.Ptr<IdType>();
|
||||
IdType* const picked_idx_data = picked_idx.Ptr<IdType>();
|
||||
FloatType* const picked_imp_data = picked_imp.Ptr<FloatType>();
|
||||
|
||||
auto picked_inrow = allocator.alloc_unique<IdType>(
|
||||
importance_sampling || weighted ? hop_size : 1);
|
||||
|
||||
// Sample edges here
|
||||
IdType num_edges;
|
||||
{
|
||||
thrust::constant_iterator<FloatType> one(1);
|
||||
if (importance_sampling) {
|
||||
auto output = thrust::make_zip_iterator(
|
||||
picked_inrow.get(), picked_row_data, picked_col_data, picked_idx_data,
|
||||
picked_imp_data);
|
||||
if (weighted) {
|
||||
auto transformed_output = thrust::make_transform_output_iterator(
|
||||
output,
|
||||
TransformOpImp<
|
||||
IdType, FloatType, FloatType*, FloatType*, decltype(one)>{
|
||||
probs_found.get(), A, one, idx_coo, rows, cs, indptr.get(),
|
||||
subindptr, indices, data, is_pinned});
|
||||
auto stencil =
|
||||
thrust::make_zip_iterator(idx_coo, probs_found.get(), rands.get());
|
||||
num_edges =
|
||||
thrust::copy_if(
|
||||
exec_policy, iota, iota + hop_size, stencil, transformed_output,
|
||||
thrust::make_zip_function(StencilOp<FloatType>{cs})) -
|
||||
transformed_output;
|
||||
} else {
|
||||
auto transformed_output = thrust::make_transform_output_iterator(
|
||||
output,
|
||||
TransformOpImp<
|
||||
IdType, FloatType, FloatType*, decltype(one), decltype(one)>{
|
||||
probs_found.get(), one, one, idx_coo, rows, cs, indptr.get(),
|
||||
subindptr, indices, data, is_pinned});
|
||||
auto stencil =
|
||||
thrust::make_zip_iterator(idx_coo, probs_found.get(), rands.get());
|
||||
num_edges =
|
||||
thrust::copy_if(
|
||||
exec_policy, iota, iota + hop_size, stencil, transformed_output,
|
||||
thrust::make_zip_function(StencilOp<FloatType>{cs})) -
|
||||
transformed_output;
|
||||
}
|
||||
} else {
|
||||
if (weighted) {
|
||||
auto output = thrust::make_zip_iterator(
|
||||
picked_inrow.get(), picked_row_data, picked_col_data,
|
||||
picked_idx_data, picked_imp_data);
|
||||
auto transformed_output = thrust::make_transform_output_iterator(
|
||||
output,
|
||||
TransformOpImp<
|
||||
IdType, FloatType, decltype(one), FloatType*, FloatType*>{
|
||||
one, A, A, idx_coo, rows, cs, indptr.get(), subindptr, indices,
|
||||
data, is_pinned});
|
||||
const auto pred =
|
||||
StencilOpFused<IdType, FloatType, decltype(one), FloatType*>{
|
||||
random_seed, idx_coo, cs, one, A,
|
||||
subindptr, indptr.get(), indices, nids, is_pinned};
|
||||
num_edges = thrust::copy_if(
|
||||
exec_policy, iota, iota + hop_size, iota,
|
||||
transformed_output, pred) -
|
||||
transformed_output;
|
||||
} else {
|
||||
auto output = thrust::make_zip_iterator(
|
||||
picked_row_data, picked_col_data, picked_idx_data);
|
||||
auto transformed_output = thrust::make_transform_output_iterator(
|
||||
output, TransformOp<IdType>{
|
||||
idx_coo, rows, indptr.get(), subindptr, indices, data,
|
||||
is_pinned});
|
||||
const auto pred =
|
||||
StencilOpFused<IdType, FloatType, decltype(one), decltype(one)>{
|
||||
random_seed, idx_coo, cs, one, one,
|
||||
subindptr, indptr.get(), indices, nids, is_pinned};
|
||||
num_edges = thrust::copy_if(
|
||||
exec_policy, iota, iota + hop_size, iota,
|
||||
transformed_output, pred) -
|
||||
transformed_output;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Normalize edge weights here
|
||||
if (importance_sampling || weighted) {
|
||||
thrust::constant_iterator<IdType> one(1);
|
||||
// contains degree information
|
||||
auto ds = allocator.alloc_unique<IdType>(num_rows);
|
||||
// contains sum of edge weights
|
||||
auto ws = allocator.alloc_unique<FloatType>(num_rows);
|
||||
// contains degree information only for vertices with nonzero degree
|
||||
auto ds_2 = allocator.alloc_unique<IdType>(num_rows);
|
||||
// contains sum of edge weights only for vertices with nonzero degree
|
||||
auto ws_2 = allocator.alloc_unique<FloatType>(num_rows);
|
||||
auto output_ = thrust::make_zip_iterator(ds.get(), ws.get());
|
||||
// contains row ids only for vertices with nonzero degree
|
||||
auto keys = allocator.alloc_unique<IdType>(num_rows);
|
||||
auto input = thrust::make_zip_iterator(one, picked_imp_data);
|
||||
auto new_end = thrust::reduce_by_key(
|
||||
exec_policy, picked_inrow.get(), picked_inrow.get() + num_edges, input,
|
||||
keys.get(), output_, thrust::equal_to<IdType>{}, TupleSum{});
|
||||
{
|
||||
thrust::constant_iterator<IdType> zero_int(0);
|
||||
thrust::constant_iterator<FloatType> zero_float(0);
|
||||
auto input = thrust::make_zip_iterator(zero_int, zero_float);
|
||||
auto output = thrust::make_zip_iterator(ds_2.get(), ws_2.get());
|
||||
thrust::copy(exec_policy, input, input + num_rows, output);
|
||||
{
|
||||
const auto num_rows_2 = new_end.first - keys.get();
|
||||
thrust::scatter(
|
||||
exec_policy, output_, output_ + num_rows_2, keys.get(), output);
|
||||
}
|
||||
}
|
||||
{
|
||||
auto input =
|
||||
thrust::make_zip_iterator(picked_inrow.get(), picked_imp_data);
|
||||
auto transformed_input = thrust::make_transform_iterator(
|
||||
input, thrust::make_zip_function(TransformOpMean<IdType, FloatType>{
|
||||
ds_2.get(), ws_2.get()}));
|
||||
thrust::copy(
|
||||
exec_policy, transformed_input, transformed_input + num_edges,
|
||||
picked_imp_data);
|
||||
}
|
||||
}
|
||||
|
||||
picked_row = picked_row.CreateView({num_edges}, picked_row->dtype);
|
||||
picked_col = picked_col.CreateView({num_edges}, picked_col->dtype);
|
||||
picked_idx = picked_idx.CreateView({num_edges}, picked_idx->dtype);
|
||||
if (importance_sampling || weighted)
|
||||
picked_imp = picked_imp.CreateView({num_edges}, picked_imp->dtype);
|
||||
|
||||
return std::make_pair(
|
||||
COOMatrix(mat.num_rows, mat.num_cols, picked_row, picked_col, picked_idx),
|
||||
picked_imp);
|
||||
}
|
||||
|
||||
template std::pair<COOMatrix, FloatArray>
|
||||
CSRLaborSampling<kDGLCUDA, int32_t, float>(
|
||||
CSRMatrix, IdArray, int64_t, FloatArray, int, IdArray, float, IdArray);
|
||||
template std::pair<COOMatrix, FloatArray>
|
||||
CSRLaborSampling<kDGLCUDA, int64_t, float>(
|
||||
CSRMatrix, IdArray, int64_t, FloatArray, int, IdArray, float, IdArray);
|
||||
template std::pair<COOMatrix, FloatArray>
|
||||
CSRLaborSampling<kDGLCUDA, int32_t, double>(
|
||||
CSRMatrix, IdArray, int64_t, FloatArray, int, IdArray, float, IdArray);
|
||||
template std::pair<COOMatrix, FloatArray>
|
||||
CSRLaborSampling<kDGLCUDA, int64_t, double>(
|
||||
CSRMatrix, IdArray, int64_t, FloatArray, int, IdArray, float, IdArray);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,53 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/macro.cuh
|
||||
* @brief Macro to call SPMM/SDDMM cuda kernels.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CUDA_MACRO_CUH_
|
||||
#define DGL_ARRAY_CUDA_MACRO_CUH_
|
||||
|
||||
///////////////////////// Dispatchers //////////////////////////
|
||||
|
||||
/* Macro used for switching between broadcasting and non-broadcasting kernels.
|
||||
* It also copies the auxiliary information for calculating broadcasting offsets
|
||||
* to GPU.
|
||||
*/
|
||||
#define BCAST_IDX_CTX_SWITCH(BCAST, EDGE_MAP, CTX, LHS_OFF, RHS_OFF, ...) \
|
||||
do { \
|
||||
const BcastOff &info = (BCAST); \
|
||||
if (!info.use_bcast) { \
|
||||
constexpr bool UseBcast = false; \
|
||||
if ((EDGE_MAP)) { \
|
||||
constexpr bool UseIdx = true; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else { \
|
||||
constexpr bool UseIdx = false; \
|
||||
{ __VA_ARGS__ } \
|
||||
} \
|
||||
} else { \
|
||||
constexpr bool UseBcast = true; \
|
||||
const DGLContext ctx = (CTX); \
|
||||
const auto device = runtime::DeviceAPI::Get(ctx); \
|
||||
(LHS_OFF) = static_cast<int64_t *>(device->AllocWorkspace( \
|
||||
ctx, sizeof(int64_t) * info.lhs_offset.size())); \
|
||||
CUDA_CALL(cudaMemcpy( \
|
||||
(LHS_OFF), &info.lhs_offset[0], \
|
||||
sizeof(int64_t) * info.lhs_offset.size(), cudaMemcpyHostToDevice)); \
|
||||
(RHS_OFF) = static_cast<int64_t *>(device->AllocWorkspace( \
|
||||
ctx, sizeof(int64_t) * info.rhs_offset.size())); \
|
||||
CUDA_CALL(cudaMemcpy( \
|
||||
(RHS_OFF), &info.rhs_offset[0], \
|
||||
sizeof(int64_t) * info.rhs_offset.size(), cudaMemcpyHostToDevice)); \
|
||||
if ((EDGE_MAP)) { \
|
||||
constexpr bool UseIdx = true; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else { \
|
||||
constexpr bool UseIdx = false; \
|
||||
{ __VA_ARGS__ } \
|
||||
} \
|
||||
device->FreeWorkspace(ctx, (LHS_OFF)); \
|
||||
device->FreeWorkspace(ctx, (RHS_OFF)); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
#endif // DGL_ARRAY_CUDA_MACRO_CUH_
|
||||
@@ -0,0 +1,220 @@
|
||||
/**
|
||||
* Copyright (c) 2021 by Contributors
|
||||
* @file array/cuda/negative_sampling.cu
|
||||
* @brief rowwise sampling
|
||||
*/
|
||||
|
||||
#include <curand_kernel.h>
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/array_iterator.h>
|
||||
#include <dgl/random.h>
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./utils.h"
|
||||
|
||||
using namespace dgl::runtime;
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
namespace {
|
||||
|
||||
template <typename IdType>
|
||||
__global__ void _GlobalUniformNegativeSamplingKernel(
|
||||
const IdType* __restrict__ indptr, const IdType* __restrict__ indices,
|
||||
IdType* __restrict__ row, IdType* __restrict__ col, int64_t num_row,
|
||||
int64_t num_col, int64_t num_samples, int num_trials,
|
||||
bool exclude_self_loops, int32_t random_seed) {
|
||||
int64_t tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int stride_x = gridDim.x * blockDim.x;
|
||||
|
||||
curandStatePhilox4_32_10_t
|
||||
rng; // this allows generating 4 32-bit ints at a time
|
||||
curand_init(random_seed * gridDim.x + blockIdx.x, threadIdx.x, 0, &rng);
|
||||
|
||||
while (tx < num_samples) {
|
||||
for (int i = 0; i < num_trials; ++i) {
|
||||
uint4 result = curand4(&rng);
|
||||
// Turns out that result.x is always 0 with the above RNG.
|
||||
uint64_t y_hi = result.y >> 16;
|
||||
uint64_t y_lo = result.y & 0xFFFF;
|
||||
uint64_t z = static_cast<uint64_t>(result.z);
|
||||
uint64_t w = static_cast<uint64_t>(result.w);
|
||||
int64_t u = static_cast<int64_t>(((y_lo << 32L) | z) % num_row);
|
||||
int64_t v = static_cast<int64_t>(((y_hi << 32L) | w) % num_col);
|
||||
|
||||
if (exclude_self_loops && (u == v)) continue;
|
||||
|
||||
// binary search of v among indptr[u:u+1]
|
||||
int64_t b = indptr[u], e = indptr[u + 1] - 1;
|
||||
bool found = false;
|
||||
while (b <= e) {
|
||||
int64_t m = (b + e) / 2;
|
||||
if (indices[m] == v) {
|
||||
found = true;
|
||||
break;
|
||||
} else if (indices[m] < v) {
|
||||
b = m + 1;
|
||||
} else {
|
||||
e = m - 1;
|
||||
}
|
||||
}
|
||||
|
||||
if (!found) {
|
||||
row[tx] = u;
|
||||
col[tx] = v;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename DType>
|
||||
struct IsNotMinusOne {
|
||||
__device__ __forceinline__ bool operator()(const std::pair<DType, DType>& a) {
|
||||
return a.first != -1;
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Sort ordered pairs in ascending order, using \a tmp_major and \a
|
||||
* tmp_minor as temporary buffers, each with \a n elements.
|
||||
*/
|
||||
template <typename IdType>
|
||||
void SortOrderedPairs(
|
||||
runtime::DeviceAPI* device, DGLContext ctx, IdType* major, IdType* minor,
|
||||
IdType* tmp_major, IdType* tmp_minor, int64_t n, cudaStream_t stream) {
|
||||
// Sort ordered pairs in lexicographical order by two radix sorts since
|
||||
// cub's radix sorts are stable.
|
||||
// We need a 2*n auxiliary storage to store the results form the first radix
|
||||
// sort.
|
||||
size_t s1 = 0, s2 = 0;
|
||||
void* tmp1 = nullptr;
|
||||
void* tmp2 = nullptr;
|
||||
|
||||
// Radix sort by minor key first, reorder the major key in the progress.
|
||||
CUDA_CALL(cub::DeviceRadixSort::SortPairs(
|
||||
tmp1, s1, minor, tmp_minor, major, tmp_major, n, 0, sizeof(IdType) * 8,
|
||||
stream));
|
||||
tmp1 = device->AllocWorkspace(ctx, s1);
|
||||
CUDA_CALL(cub::DeviceRadixSort::SortPairs(
|
||||
tmp1, s1, minor, tmp_minor, major, tmp_major, n, 0, sizeof(IdType) * 8,
|
||||
stream));
|
||||
|
||||
// Radix sort by major key next.
|
||||
CUDA_CALL(cub::DeviceRadixSort::SortPairs(
|
||||
tmp2, s2, tmp_major, major, tmp_minor, minor, n, 0, sizeof(IdType) * 8,
|
||||
stream));
|
||||
tmp2 = (s2 > s1) ? device->AllocWorkspace(ctx, s2)
|
||||
: tmp1; // reuse buffer if s2 <= s1
|
||||
CUDA_CALL(cub::DeviceRadixSort::SortPairs(
|
||||
tmp2, s2, tmp_major, major, tmp_minor, minor, n, 0, sizeof(IdType) * 8,
|
||||
stream));
|
||||
|
||||
if (tmp1 != tmp2) device->FreeWorkspace(ctx, tmp2);
|
||||
device->FreeWorkspace(ctx, tmp1);
|
||||
}
|
||||
|
||||
}; // namespace
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::pair<IdArray, IdArray> CSRGlobalUniformNegativeSampling(
|
||||
const CSRMatrix& csr, int64_t num_samples, int num_trials,
|
||||
bool exclude_self_loops, bool replace, double redundancy) {
|
||||
auto ctx = csr.indptr->ctx;
|
||||
auto dtype = csr.indptr->dtype;
|
||||
const int64_t num_row = csr.num_rows;
|
||||
const int64_t num_col = csr.num_cols;
|
||||
const int64_t num_actual_samples =
|
||||
static_cast<int64_t>(num_samples * (1 + redundancy));
|
||||
IdArray row = Full<IdType>(-1, num_actual_samples, ctx);
|
||||
IdArray col = Full<IdType>(-1, num_actual_samples, ctx);
|
||||
IdArray out_row = IdArray::Empty({num_actual_samples}, dtype, ctx);
|
||||
IdArray out_col = IdArray::Empty({num_actual_samples}, dtype, ctx);
|
||||
IdType* row_data = row.Ptr<IdType>();
|
||||
IdType* col_data = col.Ptr<IdType>();
|
||||
IdType* out_row_data = out_row.Ptr<IdType>();
|
||||
IdType* out_col_data = out_col.Ptr<IdType>();
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const int nt = cuda::FindNumThreads(num_actual_samples);
|
||||
const int nb = (num_actual_samples + nt - 1) / nt;
|
||||
std::pair<IdArray, IdArray> result;
|
||||
int64_t num_out;
|
||||
|
||||
CUDA_KERNEL_CALL(
|
||||
_GlobalUniformNegativeSamplingKernel, nb, nt, 0, stream,
|
||||
csr.indptr.Ptr<IdType>(), csr.indices.Ptr<IdType>(), row_data, col_data,
|
||||
num_row, num_col, num_actual_samples, num_trials, exclude_self_loops,
|
||||
RandomEngine::ThreadLocal()->RandInt32());
|
||||
|
||||
size_t tmp_size = 0;
|
||||
int64_t* num_out_cuda =
|
||||
static_cast<int64_t*>(device->AllocWorkspace(ctx, sizeof(int64_t)));
|
||||
IsNotMinusOne<IdType> op;
|
||||
PairIterator<IdType> begin(row_data, col_data);
|
||||
PairIterator<IdType> out_begin(out_row_data, out_col_data);
|
||||
CUDA_CALL(cub::DeviceSelect::If(
|
||||
nullptr, tmp_size, begin, out_begin, num_out_cuda, num_actual_samples, op,
|
||||
stream));
|
||||
void* tmp = device->AllocWorkspace(ctx, tmp_size);
|
||||
CUDA_CALL(cub::DeviceSelect::If(
|
||||
tmp, tmp_size, begin, out_begin, num_out_cuda, num_actual_samples, op,
|
||||
stream));
|
||||
num_out = cuda::GetCUDAScalar(device, ctx, num_out_cuda);
|
||||
|
||||
if (!replace) {
|
||||
IdArray unique_row = IdArray::Empty({num_out}, dtype, ctx);
|
||||
IdArray unique_col = IdArray::Empty({num_out}, dtype, ctx);
|
||||
IdType* unique_row_data = unique_row.Ptr<IdType>();
|
||||
IdType* unique_col_data = unique_col.Ptr<IdType>();
|
||||
PairIterator<IdType> unique_begin(unique_row_data, unique_col_data);
|
||||
|
||||
SortOrderedPairs(
|
||||
device, ctx, out_row_data, out_col_data, unique_row_data,
|
||||
unique_col_data, num_out, stream);
|
||||
|
||||
size_t tmp_size_unique = 0;
|
||||
void* tmp_unique = nullptr;
|
||||
CUDA_CALL(cub::DeviceSelect::Unique(
|
||||
nullptr, tmp_size_unique, out_begin, unique_begin, num_out_cuda,
|
||||
num_out, stream));
|
||||
tmp_unique = (tmp_size_unique > tmp_size)
|
||||
? device->AllocWorkspace(ctx, tmp_size_unique)
|
||||
: tmp; // reuse buffer
|
||||
CUDA_CALL(cub::DeviceSelect::Unique(
|
||||
tmp_unique, tmp_size_unique, out_begin, unique_begin, num_out_cuda,
|
||||
num_out, stream));
|
||||
num_out = cuda::GetCUDAScalar(device, ctx, num_out_cuda);
|
||||
|
||||
num_out = std::min(num_samples, num_out);
|
||||
result = {
|
||||
unique_row.CreateView({num_out}, dtype),
|
||||
unique_col.CreateView({num_out}, dtype)};
|
||||
|
||||
if (tmp_unique != tmp) device->FreeWorkspace(ctx, tmp_unique);
|
||||
} else {
|
||||
num_out = std::min(num_samples, num_out);
|
||||
result = {
|
||||
out_row.CreateView({num_out}, dtype),
|
||||
out_col.CreateView({num_out}, dtype)};
|
||||
}
|
||||
|
||||
device->FreeWorkspace(ctx, tmp);
|
||||
device->FreeWorkspace(ctx, num_out_cuda);
|
||||
return result;
|
||||
}
|
||||
|
||||
template std::pair<IdArray, IdArray> CSRGlobalUniformNegativeSampling<
|
||||
kDGLCUDA, int32_t>(const CSRMatrix&, int64_t, int, bool, bool, double);
|
||||
template std::pair<IdArray, IdArray> CSRGlobalUniformNegativeSampling<
|
||||
kDGLCUDA, int64_t>(const CSRMatrix&, int64_t, int, bool, bool, double);
|
||||
|
||||
}; // namespace impl
|
||||
}; // namespace aten
|
||||
}; // namespace dgl
|
||||
@@ -0,0 +1,366 @@
|
||||
/**
|
||||
* Copyright (c) 2021 by Contributors
|
||||
* @file array/cuda/rowwise_sampling.cu
|
||||
* @brief uniform rowwise sampling
|
||||
*/
|
||||
|
||||
#include <curand_kernel.h>
|
||||
#include <dgl/random.h>
|
||||
#include <dgl/runtime/device_api.h>
|
||||
#include <dgl/runtime/tensordispatch.h>
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
#include <numeric>
|
||||
|
||||
#include "../../array/cuda/atomic.cuh"
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./utils.h"
|
||||
|
||||
using namespace dgl::cuda;
|
||||
using namespace dgl::aten::cuda;
|
||||
using TensorDispatcher = dgl::runtime::TensorDispatcher;
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int BLOCK_SIZE = 128;
|
||||
|
||||
/**
|
||||
* @brief Compute the size of each row in the sampled CSR, without replacement.
|
||||
*
|
||||
* @tparam IdType The type of node and edge indexes.
|
||||
* @param num_picks The number of non-zero entries to pick per row.
|
||||
* @param num_rows The number of rows to pick.
|
||||
* @param in_rows The set of rows to pick.
|
||||
* @param in_ptr The index where each row's edges start.
|
||||
* @param out_deg The size of each row in the sampled matrix, as indexed by
|
||||
* `in_rows` (output).
|
||||
*/
|
||||
template <typename IdType>
|
||||
__global__ void _CSRRowWiseSampleDegreeKernel(
|
||||
const int64_t num_picks, const int64_t num_rows,
|
||||
const IdType* const in_rows, const IdType* const in_ptr,
|
||||
IdType* const out_deg) {
|
||||
const int tIdx = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
if (tIdx < num_rows) {
|
||||
const int in_row = in_rows[tIdx];
|
||||
const int out_row = tIdx;
|
||||
out_deg[out_row] = min(
|
||||
static_cast<IdType>(num_picks), in_ptr[in_row + 1] - in_ptr[in_row]);
|
||||
|
||||
if (out_row == num_rows - 1) {
|
||||
// make the prefixsum work
|
||||
out_deg[num_rows] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Compute the size of each row in the sampled CSR, with replacement.
|
||||
*
|
||||
* @tparam IdType The type of node and edge indexes.
|
||||
* @param num_picks The number of non-zero entries to pick per row.
|
||||
* @param num_rows The number of rows to pick.
|
||||
* @param in_rows The set of rows to pick.
|
||||
* @param in_ptr The index where each row's edges start.
|
||||
* @param out_deg The size of each row in the sampled matrix, as indexed by
|
||||
* `in_rows` (output).
|
||||
*/
|
||||
template <typename IdType>
|
||||
__global__ void _CSRRowWiseSampleDegreeReplaceKernel(
|
||||
const int64_t num_picks, const int64_t num_rows,
|
||||
const IdType* const in_rows, const IdType* const in_ptr,
|
||||
IdType* const out_deg) {
|
||||
const int tIdx = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
if (tIdx < num_rows) {
|
||||
const int64_t in_row = in_rows[tIdx];
|
||||
const int64_t out_row = tIdx;
|
||||
|
||||
if (in_ptr[in_row + 1] - in_ptr[in_row] == 0) {
|
||||
out_deg[out_row] = 0;
|
||||
} else {
|
||||
out_deg[out_row] = static_cast<IdType>(num_picks);
|
||||
}
|
||||
|
||||
if (out_row == num_rows - 1) {
|
||||
// make the prefixsum work
|
||||
out_deg[num_rows] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Perform row-wise uniform sampling on a CSR matrix,
|
||||
* and generate a COO matrix, without replacement.
|
||||
*
|
||||
* @tparam IdType The ID type used for matrices.
|
||||
* @tparam TILE_SIZE The number of rows covered by each threadblock.
|
||||
* @param rand_seed The random seed to use.
|
||||
* @param num_picks The number of non-zeros to pick per row.
|
||||
* @param num_rows The number of rows to pick.
|
||||
* @param in_rows The set of rows to pick.
|
||||
* @param in_ptr The indptr array of the input CSR.
|
||||
* @param in_index The indices array of the input CSR.
|
||||
* @param data The data array of the input CSR.
|
||||
* @param out_ptr The offset to write each row to in the output COO.
|
||||
* @param out_rows The rows of the output COO (output).
|
||||
* @param out_cols The columns of the output COO (output).
|
||||
* @param out_idxs The data array of the output COO (output).
|
||||
*/
|
||||
template <typename IdType, int TILE_SIZE>
|
||||
__global__ void _CSRRowWiseSampleUniformKernel(
|
||||
const uint64_t rand_seed, const int64_t num_picks, const int64_t num_rows,
|
||||
const IdType* const in_rows, const IdType* const in_ptr,
|
||||
const IdType* const in_index, const IdType* const data,
|
||||
const IdType* const out_ptr, IdType* const out_rows, IdType* const out_cols,
|
||||
IdType* const out_idxs) {
|
||||
// we assign one warp per row
|
||||
assert(blockDim.x == BLOCK_SIZE);
|
||||
|
||||
int64_t out_row = blockIdx.x * TILE_SIZE;
|
||||
const int64_t last_row =
|
||||
min(static_cast<int64_t>(blockIdx.x + 1) * TILE_SIZE, num_rows);
|
||||
|
||||
curandStatePhilox4_32_10_t rng;
|
||||
curand_init(rand_seed * gridDim.x + blockIdx.x, threadIdx.x, 0, &rng);
|
||||
|
||||
while (out_row < last_row) {
|
||||
const int64_t row = in_rows[out_row];
|
||||
const int64_t in_row_start = in_ptr[row];
|
||||
const int64_t deg = in_ptr[row + 1] - in_row_start;
|
||||
const int64_t out_row_start = out_ptr[out_row];
|
||||
|
||||
if (deg <= num_picks) {
|
||||
// just copy row when there is not enough nodes to sample.
|
||||
for (int idx = threadIdx.x; idx < deg; idx += BLOCK_SIZE) {
|
||||
const IdType in_idx = in_row_start + idx;
|
||||
out_rows[out_row_start + idx] = row;
|
||||
out_cols[out_row_start + idx] = in_index[in_idx];
|
||||
out_idxs[out_row_start + idx] = data ? data[in_idx] : in_idx;
|
||||
}
|
||||
} else {
|
||||
// generate permutation list via reservoir algorithm
|
||||
for (int idx = threadIdx.x; idx < num_picks; idx += BLOCK_SIZE) {
|
||||
out_idxs[out_row_start + idx] = idx;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int idx = num_picks + threadIdx.x; idx < deg; idx += BLOCK_SIZE) {
|
||||
const int num = curand(&rng) % (idx + 1);
|
||||
if (num < num_picks) {
|
||||
// use max so as to achieve the replacement order the serial
|
||||
// algorithm would have
|
||||
AtomicMax(out_idxs + out_row_start + num, idx);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// copy permutation over
|
||||
for (int idx = threadIdx.x; idx < num_picks; idx += BLOCK_SIZE) {
|
||||
const IdType perm_idx = out_idxs[out_row_start + idx] + in_row_start;
|
||||
out_rows[out_row_start + idx] = row;
|
||||
out_cols[out_row_start + idx] = in_index[perm_idx];
|
||||
out_idxs[out_row_start + idx] = data ? data[perm_idx] : perm_idx;
|
||||
}
|
||||
}
|
||||
out_row += 1;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Perform row-wise uniform sampling on a CSR matrix,
|
||||
* and generate a COO matrix, with replacement.
|
||||
*
|
||||
* @tparam IdType The ID type used for matrices.
|
||||
* @tparam TILE_SIZE The number of rows covered by each threadblock.
|
||||
* @param rand_seed The random seed to use.
|
||||
* @param num_picks The number of non-zeros to pick per row.
|
||||
* @param num_rows The number of rows to pick.
|
||||
* @param in_rows The set of rows to pick.
|
||||
* @param in_ptr The indptr array of the input CSR.
|
||||
* @param in_index The indices array of the input CSR.
|
||||
* @param data The data array of the input CSR.
|
||||
* @param out_ptr The offset to write each row to in the output COO.
|
||||
* @param out_rows The rows of the output COO (output).
|
||||
* @param out_cols The columns of the output COO (output).
|
||||
* @param out_idxs The data array of the output COO (output).
|
||||
*/
|
||||
template <typename IdType, int TILE_SIZE>
|
||||
__global__ void _CSRRowWiseSampleUniformReplaceKernel(
|
||||
const uint64_t rand_seed, const int64_t num_picks, const int64_t num_rows,
|
||||
const IdType* const in_rows, const IdType* const in_ptr,
|
||||
const IdType* const in_index, const IdType* const data,
|
||||
const IdType* const out_ptr, IdType* const out_rows, IdType* const out_cols,
|
||||
IdType* const out_idxs) {
|
||||
// we assign one warp per row
|
||||
assert(blockDim.x == BLOCK_SIZE);
|
||||
|
||||
int64_t out_row = blockIdx.x * TILE_SIZE;
|
||||
const int64_t last_row =
|
||||
min(static_cast<int64_t>(blockIdx.x + 1) * TILE_SIZE, num_rows);
|
||||
|
||||
curandStatePhilox4_32_10_t rng;
|
||||
curand_init(rand_seed * gridDim.x + blockIdx.x, threadIdx.x, 0, &rng);
|
||||
|
||||
while (out_row < last_row) {
|
||||
const int64_t row = in_rows[out_row];
|
||||
const int64_t in_row_start = in_ptr[row];
|
||||
const int64_t out_row_start = out_ptr[out_row];
|
||||
const int64_t deg = in_ptr[row + 1] - in_row_start;
|
||||
|
||||
if (deg > 0) {
|
||||
// each thread then blindly copies in rows only if deg > 0.
|
||||
for (int idx = threadIdx.x; idx < num_picks; idx += BLOCK_SIZE) {
|
||||
const int64_t edge = curand(&rng) % deg;
|
||||
const int64_t out_idx = out_row_start + idx;
|
||||
out_rows[out_idx] = row;
|
||||
out_cols[out_idx] = in_index[in_row_start + edge];
|
||||
out_idxs[out_idx] =
|
||||
data ? data[in_row_start + edge] : in_row_start + edge;
|
||||
}
|
||||
}
|
||||
out_row += 1;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
///////////////////////////// CSR sampling //////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix _CSRRowWiseSamplingUniform(
|
||||
CSRMatrix mat, IdArray rows, const int64_t num_picks, const bool replace) {
|
||||
const auto& ctx = rows->ctx;
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
|
||||
const int64_t num_rows = rows->shape[0];
|
||||
const IdType* const slice_rows = static_cast<const IdType*>(rows->data);
|
||||
|
||||
IdArray picked_row =
|
||||
NewIdArray(num_rows * num_picks, ctx, sizeof(IdType) * 8);
|
||||
IdArray picked_col =
|
||||
NewIdArray(num_rows * num_picks, ctx, sizeof(IdType) * 8);
|
||||
IdArray picked_idx =
|
||||
NewIdArray(num_rows * num_picks, ctx, sizeof(IdType) * 8);
|
||||
IdType* const out_rows = static_cast<IdType*>(picked_row->data);
|
||||
IdType* const out_cols = static_cast<IdType*>(picked_col->data);
|
||||
IdType* const out_idxs = static_cast<IdType*>(picked_idx->data);
|
||||
|
||||
const IdType* in_ptr = static_cast<IdType*>(GetDevicePointer(mat.indptr));
|
||||
const IdType* in_cols = static_cast<IdType*>(GetDevicePointer(mat.indices));
|
||||
const IdType* data = CSRHasData(mat)
|
||||
? static_cast<IdType*>(GetDevicePointer(mat.data))
|
||||
: nullptr;
|
||||
|
||||
// compute degree
|
||||
IdType* out_deg = static_cast<IdType*>(
|
||||
device->AllocWorkspace(ctx, (num_rows + 1) * sizeof(IdType)));
|
||||
if (replace) {
|
||||
const dim3 block(512);
|
||||
const dim3 grid((num_rows + block.x - 1) / block.x);
|
||||
CUDA_KERNEL_CALL(
|
||||
_CSRRowWiseSampleDegreeReplaceKernel, grid, block, 0, stream, num_picks,
|
||||
num_rows, slice_rows, in_ptr, out_deg);
|
||||
} else {
|
||||
const dim3 block(512);
|
||||
const dim3 grid((num_rows + block.x - 1) / block.x);
|
||||
CUDA_KERNEL_CALL(
|
||||
_CSRRowWiseSampleDegreeKernel, grid, block, 0, stream, num_picks,
|
||||
num_rows, slice_rows, in_ptr, out_deg);
|
||||
}
|
||||
|
||||
// fill out_ptr
|
||||
IdType* out_ptr = static_cast<IdType*>(
|
||||
device->AllocWorkspace(ctx, (num_rows + 1) * sizeof(IdType)));
|
||||
size_t prefix_temp_size = 0;
|
||||
CUDA_CALL(cub::DeviceScan::ExclusiveSum(
|
||||
nullptr, prefix_temp_size, out_deg, out_ptr, num_rows + 1, stream));
|
||||
void* prefix_temp = device->AllocWorkspace(ctx, prefix_temp_size);
|
||||
CUDA_CALL(cub::DeviceScan::ExclusiveSum(
|
||||
prefix_temp, prefix_temp_size, out_deg, out_ptr, num_rows + 1, stream));
|
||||
device->FreeWorkspace(ctx, prefix_temp);
|
||||
device->FreeWorkspace(ctx, out_deg);
|
||||
|
||||
cudaEvent_t copyEvent;
|
||||
CUDA_CALL(cudaEventCreate(©Event));
|
||||
|
||||
NDArray new_len_tensor;
|
||||
if (TensorDispatcher::Global()->IsAvailable()) {
|
||||
new_len_tensor = NDArray::PinnedEmpty(
|
||||
{1}, DGLDataTypeTraits<IdType>::dtype, DGLContext{kDGLCPU, 0});
|
||||
} else {
|
||||
// use pageable memory, it will unecessarily block but be functional
|
||||
new_len_tensor = NDArray::Empty(
|
||||
{1}, DGLDataTypeTraits<IdType>::dtype, DGLContext{kDGLCPU, 0});
|
||||
}
|
||||
|
||||
// copy using the internal current stream
|
||||
CUDA_CALL(cudaMemcpyAsync(
|
||||
new_len_tensor->data, out_ptr + num_rows, sizeof(IdType),
|
||||
cudaMemcpyDeviceToHost, stream));
|
||||
CUDA_CALL(cudaEventRecord(copyEvent, stream));
|
||||
|
||||
const uint64_t random_seed = RandomEngine::ThreadLocal()->RandInt(1000000000);
|
||||
|
||||
// select edges
|
||||
// the number of rows each thread block will cover
|
||||
constexpr int TILE_SIZE = 128 / BLOCK_SIZE;
|
||||
if (replace) { // with replacement
|
||||
const dim3 block(BLOCK_SIZE);
|
||||
const dim3 grid((num_rows + TILE_SIZE - 1) / TILE_SIZE);
|
||||
CUDA_KERNEL_CALL(
|
||||
(_CSRRowWiseSampleUniformReplaceKernel<IdType, TILE_SIZE>), grid, block,
|
||||
0, stream, random_seed, num_picks, num_rows, slice_rows, in_ptr,
|
||||
in_cols, data, out_ptr, out_rows, out_cols, out_idxs);
|
||||
} else { // without replacement
|
||||
const dim3 block(BLOCK_SIZE);
|
||||
const dim3 grid((num_rows + TILE_SIZE - 1) / TILE_SIZE);
|
||||
CUDA_KERNEL_CALL(
|
||||
(_CSRRowWiseSampleUniformKernel<IdType, TILE_SIZE>), grid, block, 0,
|
||||
stream, random_seed, num_picks, num_rows, slice_rows, in_ptr, in_cols,
|
||||
data, out_ptr, out_rows, out_cols, out_idxs);
|
||||
}
|
||||
device->FreeWorkspace(ctx, out_ptr);
|
||||
|
||||
// wait for copying `new_len` to finish
|
||||
CUDA_CALL(cudaEventSynchronize(copyEvent));
|
||||
CUDA_CALL(cudaEventDestroy(copyEvent));
|
||||
|
||||
const IdType new_len = static_cast<const IdType*>(new_len_tensor->data)[0];
|
||||
picked_row = picked_row.CreateView({new_len}, picked_row->dtype);
|
||||
picked_col = picked_col.CreateView({new_len}, picked_col->dtype);
|
||||
picked_idx = picked_idx.CreateView({new_len}, picked_idx->dtype);
|
||||
|
||||
return COOMatrix(
|
||||
mat.num_rows, mat.num_cols, picked_row, picked_col, picked_idx);
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
COOMatrix CSRRowWiseSamplingUniform(
|
||||
CSRMatrix mat, IdArray rows, const int64_t num_picks, const bool replace) {
|
||||
if (num_picks == -1) {
|
||||
// Basically this is UnitGraph::InEdges().
|
||||
COOMatrix coo = CSRToCOO(CSRSliceRows(mat, rows), false);
|
||||
IdArray sliced_rows = IndexSelect(rows, coo.row);
|
||||
return COOMatrix(
|
||||
mat.num_rows, mat.num_cols, sliced_rows, coo.col, coo.data);
|
||||
} else {
|
||||
return _CSRRowWiseSamplingUniform<XPU, IdType>(
|
||||
mat, rows, num_picks, replace);
|
||||
}
|
||||
}
|
||||
|
||||
template COOMatrix CSRRowWiseSamplingUniform<kDGLCUDA, int32_t>(
|
||||
CSRMatrix, IdArray, int64_t, bool);
|
||||
template COOMatrix CSRRowWiseSamplingUniform<kDGLCUDA, int64_t>(
|
||||
CSRMatrix, IdArray, int64_t, bool);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,696 @@
|
||||
/**
|
||||
* Copyright (c) 2022 by Contributors
|
||||
* @file array/cuda/rowwise_sampling_prob.cu
|
||||
* @brief weighted rowwise sampling. The degree computing kernels and
|
||||
* host-side functions are partially borrowed from the uniform rowwise
|
||||
* sampling code rowwise_sampling.cu.
|
||||
* @author pengqirong (OPPO), dlasalle and Xin from Nvidia.
|
||||
*/
|
||||
#include <curand_kernel.h>
|
||||
#include <dgl/random.h>
|
||||
#include <dgl/runtime/device_api.h>
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
#include <numeric>
|
||||
|
||||
#include "../../array/cuda/atomic.cuh"
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./utils.h"
|
||||
|
||||
// require CUB 1.17 to use DeviceSegmentedSort
|
||||
static_assert(
|
||||
CUB_VERSION >= 101700, "Require CUB >= 1.17 to use DeviceSegmentedSort");
|
||||
|
||||
namespace dgl {
|
||||
using namespace cuda;
|
||||
using namespace aten::cuda;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int BLOCK_SIZE = 128;
|
||||
|
||||
/**
|
||||
* @brief Compute the size of each row in the sampled CSR, without replacement.
|
||||
* temp_deg is calculated for rows with deg > num_picks.
|
||||
* For these rows, we will calculate their A-Res values and sort them to get
|
||||
* top-num_picks.
|
||||
*
|
||||
* @tparam IdType The type of node and edge indexes.
|
||||
* @param num_picks The number of non-zero entries to pick per row.
|
||||
* @param num_rows The number of rows to pick.
|
||||
* @param in_rows The set of rows to pick.
|
||||
* @param in_ptr The index where each row's edges start.
|
||||
* @param out_deg The size of each row in the sampled matrix, as indexed by
|
||||
* `in_rows` (output).
|
||||
* @param temp_deg The size of each row in the input matrix, as indexed by
|
||||
* `in_rows` (output).
|
||||
*/
|
||||
template <typename IdType>
|
||||
__global__ void _CSRRowWiseSampleDegreeKernel(
|
||||
const int64_t num_picks, const int64_t num_rows,
|
||||
const IdType* const in_rows, const IdType* const in_ptr,
|
||||
IdType* const out_deg, IdType* const temp_deg) {
|
||||
const int64_t tIdx = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
if (tIdx < num_rows) {
|
||||
const int64_t in_row = in_rows[tIdx];
|
||||
const int64_t out_row = tIdx;
|
||||
const IdType deg = in_ptr[in_row + 1] - in_ptr[in_row];
|
||||
// temp_deg is used to generate ares_ptr
|
||||
temp_deg[out_row] = deg > static_cast<IdType>(num_picks) ? deg : 0;
|
||||
out_deg[out_row] = min(static_cast<IdType>(num_picks), deg);
|
||||
|
||||
if (out_row == num_rows - 1) {
|
||||
// make the prefixsum work
|
||||
out_deg[num_rows] = 0;
|
||||
temp_deg[num_rows] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Compute the size of each row in the sampled CSR, with replacement.
|
||||
* We need the actual in degree of each row to store CDF values.
|
||||
*
|
||||
* @tparam IdType The type of node and edge indexes.
|
||||
* @param num_picks The number of non-zero entries to pick per row.
|
||||
* @param num_rows The number of rows to pick.
|
||||
* @param in_rows The set of rows to pick.
|
||||
* @param in_ptr The index where each row's edges start.
|
||||
* @param out_deg The size of each row in the sampled matrix, as indexed by
|
||||
* `in_rows` (output).
|
||||
* @param temp_deg The size of each row in the input matrix, as indexed by
|
||||
* `in_rows` (output).
|
||||
*/
|
||||
template <typename IdType>
|
||||
__global__ void _CSRRowWiseSampleDegreeReplaceKernel(
|
||||
const int64_t num_picks, const int64_t num_rows,
|
||||
const IdType* const in_rows, const IdType* const in_ptr,
|
||||
IdType* const out_deg, IdType* const temp_deg) {
|
||||
const int64_t tIdx = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
if (tIdx < num_rows) {
|
||||
const int64_t in_row = in_rows[tIdx];
|
||||
const int64_t out_row = tIdx;
|
||||
const IdType deg = in_ptr[in_row + 1] - in_ptr[in_row];
|
||||
temp_deg[out_row] = deg;
|
||||
out_deg[out_row] = deg == 0 ? 0 : static_cast<IdType>(num_picks);
|
||||
|
||||
if (out_row == num_rows - 1) {
|
||||
// make the prefixsum work
|
||||
out_deg[num_rows] = 0;
|
||||
temp_deg[num_rows] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Equivalent to numpy expression: array[idx[off:off + len]]
|
||||
*
|
||||
* @tparam IdType The ID type used for indices.
|
||||
* @tparam FloatType The float type used for array values.
|
||||
* @param array The array to be selected.
|
||||
* @param idx_data The index mapping array.
|
||||
* @param index The index of value to be selected.
|
||||
* @param offset The offset to start.
|
||||
* @param out The selected value (output).
|
||||
*/
|
||||
template <typename IdType, typename FloatType>
|
||||
__device__ void _DoubleSlice(
|
||||
const FloatType* const array, const IdType* const idx_data,
|
||||
const IdType idx, const IdType offset, FloatType* const out) {
|
||||
if (idx_data) {
|
||||
*out = array[idx_data[offset + idx]];
|
||||
} else {
|
||||
*out = array[offset + idx];
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Compute A-Res value. A-Res value needs to be calculated only if deg
|
||||
* is greater than num_picks in weighted rowwise sampling without replacement.
|
||||
*
|
||||
* @tparam IdType The ID type used for matrices.
|
||||
* @tparam FloatType The Float type used for matrices.
|
||||
* @tparam TILE_SIZE The number of rows covered by each threadblock.
|
||||
* @param rand_seed The random seed to use.
|
||||
* @param num_picks The number of non-zeros to pick per row.
|
||||
* @param num_rows The number of rows to pick.
|
||||
* @param in_rows The set of rows to pick.
|
||||
* @param in_ptr The indptr array of the input CSR.
|
||||
* @param data The data array of the input CSR.
|
||||
* @param prob The probability array of the input CSR.
|
||||
* @param ares_ptr The offset to write each row to in the A-res array.
|
||||
* @param ares_idxs The A-Res value corresponding index array, the index of
|
||||
* input CSR (output).
|
||||
* @param ares The A-Res value array (output).
|
||||
* @author pengqirong (OPPO)
|
||||
*/
|
||||
template <typename IdType, typename FloatType, int TILE_SIZE>
|
||||
__global__ void _CSRAResValueKernel(
|
||||
const uint64_t rand_seed, const int64_t num_picks, const int64_t num_rows,
|
||||
const IdType* const in_rows, const IdType* const in_ptr,
|
||||
const IdType* const data, const FloatType* const prob,
|
||||
const IdType* const ares_ptr, IdType* const ares_idxs,
|
||||
FloatType* const ares) {
|
||||
int64_t out_row = blockIdx.x * TILE_SIZE;
|
||||
const int64_t last_row =
|
||||
min(static_cast<int64_t>(blockIdx.x + 1) * TILE_SIZE, num_rows);
|
||||
|
||||
curandStatePhilox4_32_10_t rng;
|
||||
curand_init(rand_seed * gridDim.x + blockIdx.x, threadIdx.x, 0, &rng);
|
||||
|
||||
while (out_row < last_row) {
|
||||
const int64_t row = in_rows[out_row];
|
||||
const int64_t in_row_start = in_ptr[row];
|
||||
const int64_t deg = in_ptr[row + 1] - in_row_start;
|
||||
// A-Res value needs to be calculated only if deg is greater than num_picks
|
||||
// in weighted rowwise sampling without replacement
|
||||
if (deg > num_picks) {
|
||||
const int64_t ares_row_start = ares_ptr[out_row];
|
||||
|
||||
for (int64_t idx = threadIdx.x; idx < deg; idx += BLOCK_SIZE) {
|
||||
const int64_t in_idx = in_row_start + idx;
|
||||
const int64_t ares_idx = ares_row_start + idx;
|
||||
FloatType item_prob;
|
||||
_DoubleSlice<IdType, FloatType>(
|
||||
prob, data, idx, in_row_start, &item_prob);
|
||||
// compute A-Res value
|
||||
ares[ares_idx] = static_cast<FloatType>(
|
||||
__powf(curand_uniform(&rng), 1.0f / item_prob));
|
||||
ares_idxs[ares_idx] = static_cast<IdType>(in_idx);
|
||||
}
|
||||
}
|
||||
out_row += 1;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Perform weighted row-wise sampling on a CSR matrix, and generate a COO
|
||||
* matrix, without replacement. After sorting, we select top-num_picks items.
|
||||
*
|
||||
* @tparam IdType The ID type used for matrices.
|
||||
* @tparam FloatType The Float type used for matrices.
|
||||
* @tparam TILE_SIZE The number of rows covered by each threadblock.
|
||||
* @param num_picks The number of non-zeros to pick per row.
|
||||
* @param num_rows The number of rows to pick.
|
||||
* @param in_rows The set of rows to pick.
|
||||
* @param in_ptr The indptr array of the input CSR.
|
||||
* @param in_cols The columns array of the input CSR.
|
||||
* @param data The data array of the input CSR.
|
||||
* @param out_ptr The offset to write each row to in the output COO.
|
||||
* @param ares_ptr The offset to write each row to in the ares array.
|
||||
* @param sort_ares_idxs The sorted A-Res value corresponding index array, the
|
||||
* index of input CSR.
|
||||
* @param out_rows The rows of the output COO (output).
|
||||
* @param out_cols The columns of the output COO (output).
|
||||
* @param out_idxs The data array of the output COO (output).
|
||||
* @author pengqirong (OPPO)
|
||||
*/
|
||||
template <typename IdType, typename FloatType, int TILE_SIZE>
|
||||
__global__ void _CSRRowWiseSampleKernel(
|
||||
const int64_t num_picks, const int64_t num_rows,
|
||||
const IdType* const in_rows, const IdType* const in_ptr,
|
||||
const IdType* const in_cols, const IdType* const data,
|
||||
const IdType* const out_ptr, const IdType* const ares_ptr,
|
||||
const IdType* const sort_ares_idxs, IdType* const out_rows,
|
||||
IdType* const out_cols, IdType* const out_idxs) {
|
||||
// we assign one warp per row
|
||||
assert(blockDim.x == BLOCK_SIZE);
|
||||
|
||||
int64_t out_row = blockIdx.x * TILE_SIZE;
|
||||
const int64_t last_row =
|
||||
min(static_cast<int64_t>(blockIdx.x + 1) * TILE_SIZE, num_rows);
|
||||
|
||||
while (out_row < last_row) {
|
||||
const int64_t row = in_rows[out_row];
|
||||
const int64_t in_row_start = in_ptr[row];
|
||||
const int64_t out_row_start = out_ptr[out_row];
|
||||
const int64_t deg = in_ptr[row + 1] - in_row_start;
|
||||
|
||||
if (deg > num_picks) {
|
||||
const int64_t ares_row_start = ares_ptr[out_row];
|
||||
for (int64_t idx = threadIdx.x; idx < num_picks; idx += BLOCK_SIZE) {
|
||||
// get in and out index, the in_idx is one of top num_picks A-Res value
|
||||
// corresponding index in input CSR.
|
||||
const int64_t out_idx = out_row_start + idx;
|
||||
const int64_t ares_idx = ares_row_start + idx;
|
||||
const int64_t in_idx = sort_ares_idxs[ares_idx];
|
||||
// copy permutation over
|
||||
out_rows[out_idx] = static_cast<IdType>(row);
|
||||
out_cols[out_idx] = in_cols[in_idx];
|
||||
out_idxs[out_idx] = static_cast<IdType>(data ? data[in_idx] : in_idx);
|
||||
}
|
||||
} else {
|
||||
for (int64_t idx = threadIdx.x; idx < deg; idx += BLOCK_SIZE) {
|
||||
// get in and out index
|
||||
const int64_t out_idx = out_row_start + idx;
|
||||
const int64_t in_idx = in_row_start + idx;
|
||||
// copy permutation over
|
||||
out_rows[out_idx] = static_cast<IdType>(row);
|
||||
out_cols[out_idx] = in_cols[in_idx];
|
||||
out_idxs[out_idx] = static_cast<IdType>(data ? data[in_idx] : in_idx);
|
||||
}
|
||||
}
|
||||
out_row += 1;
|
||||
}
|
||||
}
|
||||
|
||||
// A stateful callback functor that maintains a running prefix to be applied
|
||||
// during consecutive scan operations.
|
||||
template <typename FloatType>
|
||||
struct BlockPrefixCallbackOp {
|
||||
// Running prefix
|
||||
FloatType running_total;
|
||||
// Constructor
|
||||
__device__ BlockPrefixCallbackOp(FloatType running_total)
|
||||
: running_total(running_total) {}
|
||||
// Callback operator to be entered by the first warp of threads in the block.
|
||||
// Thread-0 is responsible for returning a value for seeding the block-wide
|
||||
// scan.
|
||||
__device__ FloatType operator()(FloatType block_aggregate) {
|
||||
FloatType old_prefix = running_total;
|
||||
running_total += block_aggregate;
|
||||
return old_prefix;
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Perform weighted row-wise sampling on a CSR matrix, and generate a COO
|
||||
* matrix, with replacement. We store the CDF (unnormalized) of all neighbors of
|
||||
* a row in global memory and use binary search to find inverse indices as
|
||||
* selected items.
|
||||
*
|
||||
* @tparam IdType The ID type used for matrices.
|
||||
* @tparam FloatType The Float type used for matrices.
|
||||
* @tparam TILE_SIZE The number of rows covered by each threadblock.
|
||||
* @param rand_seed The random seed to use.
|
||||
* @param num_picks The number of non-zeros to pick per row.
|
||||
* @param num_rows The number of rows to pick.
|
||||
* @param in_rows The set of rows to pick.
|
||||
* @param in_ptr The indptr array of the input CSR.
|
||||
* @param in_cols The columns array of the input CSR.
|
||||
* @param data The data array of the input CSR.
|
||||
* @param prob The probability array of the input CSR.
|
||||
* @param out_ptr The offset to write each row to in the output COO.
|
||||
* @param cdf_ptr The offset of each cdf segment.
|
||||
* @param cdf The global buffer to store cdf segments.
|
||||
* @param out_rows The rows of the output COO (output).
|
||||
* @param out_cols The columns of the output COO (output).
|
||||
* @param out_idxs The data array of the output COO (output).
|
||||
* @author pengqirong (OPPO)
|
||||
*/
|
||||
template <typename IdType, typename FloatType, int TILE_SIZE>
|
||||
__global__ void _CSRRowWiseSampleReplaceKernel(
|
||||
const uint64_t rand_seed, const int64_t num_picks, const int64_t num_rows,
|
||||
const IdType* const in_rows, const IdType* const in_ptr,
|
||||
const IdType* const in_cols, const IdType* const data,
|
||||
const FloatType* const prob, const IdType* const out_ptr,
|
||||
const IdType* const cdf_ptr, FloatType* const cdf, IdType* const out_rows,
|
||||
IdType* const out_cols, IdType* const out_idxs) {
|
||||
// we assign one warp per row
|
||||
assert(blockDim.x == BLOCK_SIZE);
|
||||
|
||||
int64_t out_row = blockIdx.x * TILE_SIZE;
|
||||
const int64_t last_row =
|
||||
min(static_cast<int64_t>(blockIdx.x + 1) * TILE_SIZE, num_rows);
|
||||
|
||||
curandStatePhilox4_32_10_t rng;
|
||||
curand_init(rand_seed * gridDim.x + blockIdx.x, threadIdx.x, 0, &rng);
|
||||
|
||||
while (out_row < last_row) {
|
||||
const int64_t row = in_rows[out_row];
|
||||
const int64_t in_row_start = in_ptr[row];
|
||||
const int64_t out_row_start = out_ptr[out_row];
|
||||
const int64_t cdf_row_start = cdf_ptr[out_row];
|
||||
const int64_t deg = in_ptr[row + 1] - in_row_start;
|
||||
const FloatType MIN_THREAD_DATA = static_cast<FloatType>(0.0f);
|
||||
|
||||
if (deg > 0) {
|
||||
// Specialize BlockScan for a 1D block of BLOCK_SIZE threads
|
||||
typedef cub::BlockScan<FloatType, BLOCK_SIZE> BlockScan;
|
||||
// Allocate shared memory for BlockScan
|
||||
__shared__ typename BlockScan::TempStorage temp_storage;
|
||||
// Initialize running total
|
||||
BlockPrefixCallbackOp<FloatType> prefix_op(MIN_THREAD_DATA);
|
||||
|
||||
int64_t max_iter = (1 + (deg - 1) / BLOCK_SIZE) * BLOCK_SIZE;
|
||||
// Have the block iterate over segments of items
|
||||
for (int64_t idx = threadIdx.x; idx < max_iter; idx += BLOCK_SIZE) {
|
||||
// Load a segment of consecutive items that are blocked across threads
|
||||
FloatType thread_data;
|
||||
if (idx < deg)
|
||||
_DoubleSlice<IdType, FloatType>(
|
||||
prob, data, idx, in_row_start, &thread_data);
|
||||
else
|
||||
thread_data = MIN_THREAD_DATA;
|
||||
thread_data = max(thread_data, MIN_THREAD_DATA);
|
||||
// Collectively compute the block-wide inclusive prefix sum
|
||||
BlockScan(temp_storage)
|
||||
.InclusiveSum(thread_data, thread_data, prefix_op);
|
||||
__syncthreads();
|
||||
|
||||
// Store scanned items to cdf array
|
||||
if (idx < deg) {
|
||||
cdf[cdf_row_start + idx] = thread_data;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int64_t idx = threadIdx.x; idx < num_picks; idx += BLOCK_SIZE) {
|
||||
// get random value
|
||||
FloatType sum = cdf[cdf_row_start + deg - 1];
|
||||
FloatType rand = static_cast<FloatType>(curand_uniform(&rng) * sum);
|
||||
// get the offset of the first value within cdf array which is greater
|
||||
// than random value.
|
||||
int64_t item = cub::UpperBound<FloatType*, int64_t, FloatType>(
|
||||
&cdf[cdf_row_start], deg, rand);
|
||||
item = min(item, deg - 1);
|
||||
// get in and out index
|
||||
const int64_t in_idx = in_row_start + item;
|
||||
const int64_t out_idx = out_row_start + idx;
|
||||
// copy permutation over
|
||||
out_rows[out_idx] = static_cast<IdType>(row);
|
||||
out_cols[out_idx] = in_cols[in_idx];
|
||||
out_idxs[out_idx] = static_cast<IdType>(data ? data[in_idx] : in_idx);
|
||||
}
|
||||
}
|
||||
out_row += 1;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename IdType, typename DType, typename BoolType>
|
||||
__global__ void _GenerateFlagsKernel(
|
||||
int64_t n, const IdType* idx, const DType* values, DType criteria,
|
||||
BoolType* output) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < n) {
|
||||
output[tx] = (values[idx ? idx[tx] : tx] != criteria);
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename DType, typename MaskGen>
|
||||
COOMatrix COOGeneralRemoveIf(const COOMatrix& coo, MaskGen maskgen) {
|
||||
using namespace dgl::cuda;
|
||||
|
||||
const auto idtype = coo.row->dtype;
|
||||
const auto ctx = coo.row->ctx;
|
||||
const int64_t nnz = coo.row->shape[0];
|
||||
const IdType* row = coo.row.Ptr<IdType>();
|
||||
const IdType* col = coo.col.Ptr<IdType>();
|
||||
const IdArray& eid =
|
||||
COOHasData(coo) ? coo.data : Range(0, nnz, sizeof(IdType) * 8, ctx);
|
||||
const IdType* data = coo.data.Ptr<IdType>();
|
||||
IdArray new_row = IdArray::Empty({nnz}, idtype, ctx);
|
||||
IdArray new_col = IdArray::Empty({nnz}, idtype, ctx);
|
||||
IdArray new_eid = IdArray::Empty({nnz}, idtype, ctx);
|
||||
IdType* new_row_data = new_row.Ptr<IdType>();
|
||||
IdType* new_col_data = new_col.Ptr<IdType>();
|
||||
IdType* new_eid_data = new_eid.Ptr<IdType>();
|
||||
auto stream = runtime::getCurrentCUDAStream();
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
|
||||
int8_t* flags = static_cast<int8_t*>(device->AllocWorkspace(ctx, nnz));
|
||||
int nt = dgl::cuda::FindNumThreads(nnz);
|
||||
int64_t nb = (nnz + nt - 1) / nt;
|
||||
|
||||
maskgen(nb, nt, stream, nnz, data, flags);
|
||||
|
||||
int64_t* rst =
|
||||
static_cast<int64_t*>(device->AllocWorkspace(ctx, sizeof(int64_t)));
|
||||
MaskSelect(device, ctx, row, flags, new_row_data, nnz, rst, stream);
|
||||
MaskSelect(device, ctx, col, flags, new_col_data, nnz, rst, stream);
|
||||
MaskSelect(device, ctx, data, flags, new_eid_data, nnz, rst, stream);
|
||||
|
||||
int64_t new_len = GetCUDAScalar(device, ctx, rst);
|
||||
|
||||
device->FreeWorkspace(ctx, flags);
|
||||
device->FreeWorkspace(ctx, rst);
|
||||
return COOMatrix(
|
||||
coo.num_rows, coo.num_cols, new_row.CreateView({new_len}, idtype, 0),
|
||||
new_col.CreateView({new_len}, idtype, 0),
|
||||
new_eid.CreateView({new_len}, idtype, 0));
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename DType>
|
||||
COOMatrix _COORemoveIf(
|
||||
const COOMatrix& coo, const NDArray& values, DType criteria) {
|
||||
const DType* val = values.Ptr<DType>();
|
||||
auto maskgen = [val, criteria](
|
||||
int nb, int nt, cudaStream_t stream, int64_t nnz,
|
||||
const IdType* data, int8_t* flags) {
|
||||
CUDA_KERNEL_CALL(
|
||||
(_GenerateFlagsKernel<IdType, DType, int8_t>), nb, nt, 0, stream, nnz,
|
||||
data, val, criteria, flags);
|
||||
};
|
||||
return COOGeneralRemoveIf<XPU, IdType, DType, decltype(maskgen)>(
|
||||
coo, maskgen);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
/////////////////////////////// CSR ///////////////////////////////
|
||||
|
||||
/**
|
||||
* @brief Perform weighted row-wise sampling on a CSR matrix, and generate a COO
|
||||
* matrix. Use CDF sampling algorithm for with replacement:
|
||||
* 1) Calculate the CDF of all neighbor's prob.
|
||||
* 2) For each [0, num_picks), generate a rand ~ U(0, 1). Use binary search to
|
||||
* find its index in the CDF array as a chosen item.
|
||||
* Use A-Res sampling algorithm for without replacement:
|
||||
* 1) For rows with deg > num_picks, calculate A-Res values for all neighbors.
|
||||
* 2) Sort the A-Res array and select top-num_picks as chosen items.
|
||||
*
|
||||
* @tparam XPU The device type used for matrices.
|
||||
* @tparam IdType The ID type used for matrices.
|
||||
* @tparam FloatType The Float type used for matrices.
|
||||
* @param mat The CSR matrix.
|
||||
* @param rows The set of rows to pick.
|
||||
* @param num_picks The number of non-zeros to pick per row.
|
||||
* @param prob The probability array of the input CSR.
|
||||
* @param replace Is replacement sampling?
|
||||
* @author pengqirong (OPPO), dlasalle and Xin from Nvidia.
|
||||
*/
|
||||
template <DGLDeviceType XPU, typename IdType, typename FloatType>
|
||||
COOMatrix _CSRRowWiseSampling(
|
||||
const CSRMatrix& mat, const IdArray& rows, int64_t num_picks,
|
||||
const FloatArray& prob, bool replace) {
|
||||
const auto& ctx = rows->ctx;
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
|
||||
const int64_t num_rows = rows->shape[0];
|
||||
const IdType* const slice_rows = static_cast<const IdType*>(rows->data);
|
||||
|
||||
IdArray picked_row =
|
||||
NewIdArray(num_rows * num_picks, ctx, sizeof(IdType) * 8);
|
||||
IdArray picked_col =
|
||||
NewIdArray(num_rows * num_picks, ctx, sizeof(IdType) * 8);
|
||||
IdArray picked_idx =
|
||||
NewIdArray(num_rows * num_picks, ctx, sizeof(IdType) * 8);
|
||||
IdType* const out_rows = static_cast<IdType*>(picked_row->data);
|
||||
IdType* const out_cols = static_cast<IdType*>(picked_col->data);
|
||||
IdType* const out_idxs = static_cast<IdType*>(picked_idx->data);
|
||||
|
||||
const IdType* in_ptr = static_cast<IdType*>(GetDevicePointer(mat.indptr));
|
||||
const IdType* in_cols = static_cast<IdType*>(GetDevicePointer(mat.indices));
|
||||
const IdType* data = CSRHasData(mat)
|
||||
? static_cast<IdType*>(GetDevicePointer(mat.data))
|
||||
: nullptr;
|
||||
const FloatType* prob_data = static_cast<FloatType*>(GetDevicePointer(prob));
|
||||
|
||||
// compute degree
|
||||
// out_deg: the size of each row in the sampled matrix
|
||||
// temp_deg: the size of each row we will manipulate in sampling
|
||||
// 1) for w/o replacement: in degree if it's greater than num_picks else 0
|
||||
// 2) for w/ replacement: in degree
|
||||
IdType* out_deg = static_cast<IdType*>(
|
||||
device->AllocWorkspace(ctx, (num_rows + 1) * sizeof(IdType)));
|
||||
IdType* temp_deg = static_cast<IdType*>(
|
||||
device->AllocWorkspace(ctx, (num_rows + 1) * sizeof(IdType)));
|
||||
if (replace) {
|
||||
const dim3 block(512);
|
||||
const dim3 grid((num_rows + block.x - 1) / block.x);
|
||||
CUDA_KERNEL_CALL(
|
||||
_CSRRowWiseSampleDegreeReplaceKernel, grid, block, 0, stream, num_picks,
|
||||
num_rows, slice_rows, in_ptr, out_deg, temp_deg);
|
||||
} else {
|
||||
const dim3 block(512);
|
||||
const dim3 grid((num_rows + block.x - 1) / block.x);
|
||||
CUDA_KERNEL_CALL(
|
||||
_CSRRowWiseSampleDegreeKernel, grid, block, 0, stream, num_picks,
|
||||
num_rows, slice_rows, in_ptr, out_deg, temp_deg);
|
||||
}
|
||||
|
||||
// fill temp_ptr
|
||||
IdType* temp_ptr = static_cast<IdType*>(
|
||||
device->AllocWorkspace(ctx, (num_rows + 1) * sizeof(IdType)));
|
||||
size_t prefix_temp_size = 0;
|
||||
CUDA_CALL(cub::DeviceScan::ExclusiveSum(
|
||||
nullptr, prefix_temp_size, temp_deg, temp_ptr, num_rows + 1, stream));
|
||||
void* prefix_temp = device->AllocWorkspace(ctx, prefix_temp_size);
|
||||
CUDA_CALL(cub::DeviceScan::ExclusiveSum(
|
||||
prefix_temp, prefix_temp_size, temp_deg, temp_ptr, num_rows + 1, stream));
|
||||
device->FreeWorkspace(ctx, prefix_temp);
|
||||
device->FreeWorkspace(ctx, temp_deg);
|
||||
|
||||
// TODO(Xin): The copy here is too small, and the overhead of creating
|
||||
// cuda events cannot be ignored. Just use synchronized copy.
|
||||
IdType temp_len;
|
||||
// copy using the internal current stream.
|
||||
device->CopyDataFromTo(
|
||||
temp_ptr, num_rows * sizeof(temp_len), &temp_len, 0, sizeof(temp_len),
|
||||
ctx, DGLContext{kDGLCPU, 0}, mat.indptr->dtype);
|
||||
device->StreamSync(ctx, stream);
|
||||
|
||||
// fill out_ptr
|
||||
IdType* out_ptr = static_cast<IdType*>(
|
||||
device->AllocWorkspace(ctx, (num_rows + 1) * sizeof(IdType)));
|
||||
prefix_temp_size = 0;
|
||||
CUDA_CALL(cub::DeviceScan::ExclusiveSum(
|
||||
nullptr, prefix_temp_size, out_deg, out_ptr, num_rows + 1, stream));
|
||||
prefix_temp = device->AllocWorkspace(ctx, prefix_temp_size);
|
||||
CUDA_CALL(cub::DeviceScan::ExclusiveSum(
|
||||
prefix_temp, prefix_temp_size, out_deg, out_ptr, num_rows + 1, stream));
|
||||
device->FreeWorkspace(ctx, prefix_temp);
|
||||
device->FreeWorkspace(ctx, out_deg);
|
||||
|
||||
cudaEvent_t copyEvent;
|
||||
CUDA_CALL(cudaEventCreate(©Event));
|
||||
// TODO(dlasalle): use pinned memory to overlap with the actual sampling, and
|
||||
// wait on a cudaevent
|
||||
IdType new_len;
|
||||
// copy using the internal current stream.
|
||||
device->CopyDataFromTo(
|
||||
out_ptr, num_rows * sizeof(new_len), &new_len, 0, sizeof(new_len), ctx,
|
||||
DGLContext{kDGLCPU, 0}, mat.indptr->dtype);
|
||||
CUDA_CALL(cudaEventRecord(copyEvent, stream));
|
||||
|
||||
// allocate workspace
|
||||
// 1) for w/ replacement, it's a global buffer to store cdf segments (one
|
||||
// segment for each row).
|
||||
// 2) for w/o replacement, it's used to store a-res segments (one segment for
|
||||
// each row with degree > num_picks)
|
||||
FloatType* temp = static_cast<FloatType*>(
|
||||
device->AllocWorkspace(ctx, temp_len * sizeof(FloatType)));
|
||||
|
||||
const uint64_t rand_seed = RandomEngine::ThreadLocal()->RandInt(1000000000);
|
||||
|
||||
// select edges
|
||||
// the number of rows each thread block will cover
|
||||
constexpr int TILE_SIZE = 128 / BLOCK_SIZE;
|
||||
if (replace) { // with replacement.
|
||||
const dim3 block(BLOCK_SIZE);
|
||||
const dim3 grid((num_rows + TILE_SIZE - 1) / TILE_SIZE);
|
||||
CUDA_KERNEL_CALL(
|
||||
(_CSRRowWiseSampleReplaceKernel<IdType, FloatType, TILE_SIZE>), grid,
|
||||
block, 0, stream, rand_seed, num_picks, num_rows, slice_rows, in_ptr,
|
||||
in_cols, data, prob_data, out_ptr, temp_ptr, temp, out_rows, out_cols,
|
||||
out_idxs);
|
||||
device->FreeWorkspace(ctx, temp);
|
||||
} else { // without replacement
|
||||
IdType* temp_idxs = static_cast<IdType*>(
|
||||
device->AllocWorkspace(ctx, (temp_len) * sizeof(IdType)));
|
||||
|
||||
// Compute A-Res value. A-Res value needs to be calculated only if deg
|
||||
// is greater than num_picks in weighted rowwise sampling without
|
||||
// replacement.
|
||||
const dim3 block(BLOCK_SIZE);
|
||||
const dim3 grid((num_rows + TILE_SIZE - 1) / TILE_SIZE);
|
||||
CUDA_KERNEL_CALL(
|
||||
(_CSRAResValueKernel<IdType, FloatType, TILE_SIZE>), grid, block, 0,
|
||||
stream, rand_seed, num_picks, num_rows, slice_rows, in_ptr, data,
|
||||
prob_data, temp_ptr, temp_idxs, temp);
|
||||
|
||||
// sort A-Res value array.
|
||||
FloatType* sort_temp = static_cast<FloatType*>(
|
||||
device->AllocWorkspace(ctx, temp_len * sizeof(FloatType)));
|
||||
IdType* sort_temp_idxs = static_cast<IdType*>(
|
||||
device->AllocWorkspace(ctx, temp_len * sizeof(IdType)));
|
||||
|
||||
cub::DoubleBuffer<FloatType> sort_keys(temp, sort_temp);
|
||||
cub::DoubleBuffer<IdType> sort_values(temp_idxs, sort_temp_idxs);
|
||||
|
||||
void* d_temp_storage = nullptr;
|
||||
size_t temp_storage_bytes = 0;
|
||||
CUDA_CALL(cub::DeviceSegmentedSort::SortPairsDescending(
|
||||
d_temp_storage, temp_storage_bytes, sort_keys, sort_values, temp_len,
|
||||
num_rows, temp_ptr, temp_ptr + 1, stream));
|
||||
d_temp_storage = device->AllocWorkspace(ctx, temp_storage_bytes);
|
||||
CUDA_CALL(cub::DeviceSegmentedSort::SortPairsDescending(
|
||||
d_temp_storage, temp_storage_bytes, sort_keys, sort_values, temp_len,
|
||||
num_rows, temp_ptr, temp_ptr + 1, stream));
|
||||
device->FreeWorkspace(ctx, d_temp_storage);
|
||||
device->FreeWorkspace(ctx, temp);
|
||||
device->FreeWorkspace(ctx, temp_idxs);
|
||||
device->FreeWorkspace(ctx, sort_temp);
|
||||
device->FreeWorkspace(ctx, sort_temp_idxs);
|
||||
|
||||
// select tok-num_picks as results
|
||||
CUDA_KERNEL_CALL(
|
||||
(_CSRRowWiseSampleKernel<IdType, FloatType, TILE_SIZE>), grid, block, 0,
|
||||
stream, num_picks, num_rows, slice_rows, in_ptr, in_cols, data, out_ptr,
|
||||
temp_ptr, sort_values.Current(), out_rows, out_cols, out_idxs);
|
||||
}
|
||||
|
||||
device->FreeWorkspace(ctx, temp_ptr);
|
||||
device->FreeWorkspace(ctx, out_ptr);
|
||||
|
||||
// wait for copying `new_len` to finish
|
||||
CUDA_CALL(cudaEventSynchronize(copyEvent));
|
||||
CUDA_CALL(cudaEventDestroy(copyEvent));
|
||||
|
||||
picked_row = picked_row.CreateView({new_len}, picked_row->dtype);
|
||||
picked_col = picked_col.CreateView({new_len}, picked_col->dtype);
|
||||
picked_idx = picked_idx.CreateView({new_len}, picked_idx->dtype);
|
||||
|
||||
return COOMatrix(
|
||||
mat.num_rows, mat.num_cols, picked_row, picked_col, picked_idx);
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType, typename DType>
|
||||
COOMatrix CSRRowWiseSampling(
|
||||
CSRMatrix mat, IdArray rows, int64_t num_picks, FloatArray prob,
|
||||
bool replace) {
|
||||
COOMatrix result;
|
||||
if (num_picks == -1) {
|
||||
// Basically this is UnitGraph::InEdges().
|
||||
COOMatrix coo = CSRToCOO(CSRSliceRows(mat, rows), false);
|
||||
IdArray sliced_rows = IndexSelect(rows, coo.row);
|
||||
result =
|
||||
COOMatrix(mat.num_rows, mat.num_cols, sliced_rows, coo.col, coo.data);
|
||||
} else {
|
||||
result = _CSRRowWiseSampling<XPU, IdType, DType>(
|
||||
mat, rows, num_picks, prob, replace);
|
||||
}
|
||||
// NOTE(BarclayII): I'm removing the entries with zero probability after
|
||||
// sampling. Is there a better way?
|
||||
return _COORemoveIf<XPU, IdType, DType>(result, prob, static_cast<DType>(0));
|
||||
}
|
||||
|
||||
template COOMatrix CSRRowWiseSampling<kDGLCUDA, int32_t, float>(
|
||||
CSRMatrix, IdArray, int64_t, FloatArray, bool);
|
||||
template COOMatrix CSRRowWiseSampling<kDGLCUDA, int64_t, float>(
|
||||
CSRMatrix, IdArray, int64_t, FloatArray, bool);
|
||||
template COOMatrix CSRRowWiseSampling<kDGLCUDA, int32_t, double>(
|
||||
CSRMatrix, IdArray, int64_t, FloatArray, bool);
|
||||
template COOMatrix CSRRowWiseSampling<kDGLCUDA, int64_t, double>(
|
||||
CSRMatrix, IdArray, int64_t, FloatArray, bool);
|
||||
// These are not being called, but we instantiate them anyway to prevent missing
|
||||
// symbols in Debug build
|
||||
template COOMatrix CSRRowWiseSampling<kDGLCUDA, int32_t, int8_t>(
|
||||
CSRMatrix, IdArray, int64_t, FloatArray, bool);
|
||||
template COOMatrix CSRRowWiseSampling<kDGLCUDA, int64_t, int8_t>(
|
||||
CSRMatrix, IdArray, int64_t, FloatArray, bool);
|
||||
template COOMatrix CSRRowWiseSampling<kDGLCUDA, int32_t, uint8_t>(
|
||||
CSRMatrix, IdArray, int64_t, FloatArray, bool);
|
||||
template COOMatrix CSRRowWiseSampling<kDGLCUDA, int64_t, uint8_t>(
|
||||
CSRMatrix, IdArray, int64_t, FloatArray, bool);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,99 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/sddmm.cu
|
||||
* @brief SDDMM C APIs and definitions.
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include "./functor.cuh"
|
||||
#include "./sddmm.cuh"
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
|
||||
/**
|
||||
* @brief CUDA implementation of g-SDDMM on Csr format.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SDDMMCsr(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target) {
|
||||
SWITCH_OP(op, Op, {
|
||||
SWITCH_TARGET(lhs_target, rhs_target, LhsTarget, RhsTarget, {
|
||||
cuda::SDDMMCsr<IdType, DType, Op, LhsTarget, RhsTarget>(
|
||||
bcast, csr, lhs, rhs, out);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA implementation of g-SDDMM on Coo format.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SDDMMCoo(
|
||||
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target) {
|
||||
SWITCH_OP(op, Op, {
|
||||
SWITCH_TARGET(lhs_target, rhs_target, LhsTarget, RhsTarget, {
|
||||
cuda::SDDMMCoo<IdType, DType, Op, LhsTarget, RhsTarget>(
|
||||
bcast, coo, lhs, rhs, out);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
template void SDDMMCsr<kDGLCUDA, int32_t, __half>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCsr<kDGLCUDA, int64_t, __half>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
#if BF16_ENABLED
|
||||
template void SDDMMCsr<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCsr<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
#endif // BF16_ENABLED
|
||||
template void SDDMMCsr<kDGLCUDA, int32_t, float>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCsr<kDGLCUDA, int64_t, float>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCsr<kDGLCUDA, int32_t, double>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCsr<kDGLCUDA, int64_t, double>(
|
||||
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
|
||||
template void SDDMMCoo<kDGLCUDA, int32_t, __half>(
|
||||
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCoo<kDGLCUDA, int64_t, __half>(
|
||||
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
#if BF16_ENABLED
|
||||
template void SDDMMCoo<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCoo<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
#endif // BF16_ENABLED
|
||||
template void SDDMMCoo<kDGLCUDA, int32_t, float>(
|
||||
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCoo<kDGLCUDA, int64_t, float>(
|
||||
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCoo<kDGLCUDA, int32_t, double>(
|
||||
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
template void SDDMMCoo<kDGLCUDA, int64_t, double>(
|
||||
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,368 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/sddmm.cuh
|
||||
* @brief SDDMM CUDA kernel function header.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CUDA_SDDMM_CUH_
|
||||
#define DGL_ARRAY_CUDA_SDDMM_CUH_
|
||||
|
||||
#include <dgl/bcast.h>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "../selector.h"
|
||||
#include "./functor.cuh"
|
||||
#include "./utils.h"
|
||||
#include "atomic.cuh"
|
||||
#include "bf16.cuh"
|
||||
#include "fp16.cuh"
|
||||
#include "functor.cuh"
|
||||
#include "macro.cuh"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using namespace cuda;
|
||||
|
||||
namespace aten {
|
||||
namespace cuda {
|
||||
|
||||
#define SWITCH_OP(op, Op, ...) \
|
||||
do { \
|
||||
if ((op) == "add") { \
|
||||
typedef cuda::binary::Add<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "sub") { \
|
||||
typedef cuda::binary::Sub<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "mul") { \
|
||||
typedef cuda::binary::Mul<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "div") { \
|
||||
typedef cuda::binary::Div<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "copy_lhs") { \
|
||||
typedef cuda::binary::CopyLhs<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "copy_rhs") { \
|
||||
typedef cuda::binary::CopyRhs<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "dot") { \
|
||||
typedef cuda::binary::Dot<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else { \
|
||||
LOG(FATAL) << "Unsupported SpMM/SDDMM binary operator: " << op; \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
#define SWITCH_RHS(rhs_target, RhsTarget, ...) \
|
||||
do { \
|
||||
if ((rhs_target) == 0) { \
|
||||
constexpr int RhsTarget = 0; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((rhs_target) == 1) { \
|
||||
constexpr int RhsTarget = 1; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((rhs_target) == 2) { \
|
||||
constexpr int RhsTarget = 2; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else { \
|
||||
LOG(INFO) << "Invalid rhs target: " << (rhs_target); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
#define SWITCH_TARGET(lhs_target, rhs_target, LhsTarget, RhsTarget, ...) \
|
||||
do { \
|
||||
if ((lhs_target) == 0) { \
|
||||
constexpr int LhsTarget = 0; \
|
||||
SWITCH_RHS(rhs_target, RhsTarget, __VA_ARGS__); \
|
||||
} else if ((lhs_target) == 1) { \
|
||||
constexpr int LhsTarget = 1; \
|
||||
SWITCH_RHS(rhs_target, RhsTarget, __VA_ARGS__); \
|
||||
} else if ((lhs_target) == 2) { \
|
||||
constexpr int LhsTarget = 2; \
|
||||
SWITCH_RHS(rhs_target, RhsTarget, __VA_ARGS__); \
|
||||
} else { \
|
||||
LOG(INFO) << "Invalid lhs target: " << (lhs_target); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
constexpr unsigned int full_mask = 0xffffffff;
|
||||
|
||||
/**
|
||||
* @brief CUDA kernel of g-SDDMM on Coo format.
|
||||
* @note it uses edge parallel strategy, different threadblocks (on y-axis)
|
||||
* is responsible for the computation on different edges. Threadblocks
|
||||
* on the x-axis are responsible for the computation on different
|
||||
* positions in feature dimension.
|
||||
*/
|
||||
template <
|
||||
typename Idx, typename DType, typename BinaryOp, bool UseBcast = false,
|
||||
bool UseIdx = false, int LhsTarget = 0, int RhsTarget = 2>
|
||||
__global__ void SDDMMCooKernel(
|
||||
const DType* __restrict__ lhs, const DType* __restrict__ rhs,
|
||||
DType* __restrict__ out, const Idx* __restrict__ row,
|
||||
const Idx* __restrict__ col, const Idx* __restrict__ edge_map, int64_t N,
|
||||
int64_t M, int64_t E, int64_t reduce_size,
|
||||
const int64_t* __restrict__ lhs_off, const int64_t* __restrict__ rhs_off,
|
||||
int64_t lhs_len, int64_t rhs_len, int64_t out_len) {
|
||||
// SDDMM with COO.
|
||||
Idx ty = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
const Idx stride_y = blockDim.y * gridDim.y;
|
||||
while (ty < E) {
|
||||
const Idx src = _ldg(row + ty);
|
||||
const Idx dst = _ldg(col + ty);
|
||||
const Idx eid = UseIdx ? _ldg(edge_map + ty) : ty;
|
||||
const DType* lhsoff =
|
||||
BinaryOp::use_lhs
|
||||
? (lhs + Selector<LhsTarget>::Call(src, eid, dst) * lhs_len)
|
||||
: nullptr;
|
||||
const DType* rhsoff =
|
||||
BinaryOp::use_rhs
|
||||
? (rhs + Selector<RhsTarget>::Call(src, eid, dst) * rhs_len)
|
||||
: nullptr;
|
||||
DType* outoff = out + eid * out_len;
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int stride_x = blockDim.x * gridDim.x;
|
||||
while (tx < out_len) {
|
||||
const Idx lhs_add = UseBcast ? lhs_off[tx] : tx;
|
||||
const Idx rhs_add = UseBcast ? rhs_off[tx] : tx;
|
||||
DType val = BinaryOp::Call(
|
||||
lhsoff + lhs_add * reduce_size, rhsoff + rhs_add * reduce_size,
|
||||
reduce_size);
|
||||
outoff[tx] = val;
|
||||
tx += stride_x;
|
||||
}
|
||||
ty += stride_y;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA kernel of SDDMM-dot on Coo format, accelerated with tree
|
||||
* reduction.
|
||||
* @note it uses edge parallel strategy, different threadblocks (on y-axis)
|
||||
* is responsible for the computation on different edges. Threadblocks
|
||||
* on the x-axis are responsible for the computation on different
|
||||
* positions in feature dimension.
|
||||
*/
|
||||
template <
|
||||
typename Idx, typename DType, bool UseBcast = false, bool UseIdx = false,
|
||||
int LhsTarget = 0, int RhsTarget = 2>
|
||||
__global__ void SDDMMCooTreeReduceKernel(
|
||||
const DType* __restrict__ lhs, const DType* __restrict__ rhs,
|
||||
DType* __restrict__ out, const Idx* __restrict__ row,
|
||||
const Idx* __restrict__ col, const Idx* __restrict__ edge_map, int64_t N,
|
||||
int64_t M, int64_t E, int64_t reduce_size,
|
||||
const int64_t* __restrict__ lhs_off, const int64_t* __restrict__ rhs_off,
|
||||
int64_t lhs_len, int64_t rhs_len, int64_t out_len) {
|
||||
Idx ty = blockIdx.x * blockDim.y + threadIdx.y;
|
||||
if (ty < E) {
|
||||
const Idx src = _ldg(row + ty);
|
||||
const Idx dst = _ldg(col + ty);
|
||||
const Idx eid = UseIdx ? _ldg(edge_map + ty) : ty;
|
||||
const DType* lhsoff =
|
||||
lhs + Selector<LhsTarget>::Call(src, eid, dst) * lhs_len;
|
||||
const DType* rhsoff =
|
||||
rhs + Selector<RhsTarget>::Call(src, eid, dst) * rhs_len;
|
||||
DType* outoff = out + eid * out_len;
|
||||
int tx = threadIdx.x; // tx < 32
|
||||
for (int i = blockIdx.y; i < out_len;
|
||||
i += gridDim.y) { // over output feature dimension
|
||||
const Idx lhs_add = UseBcast ? __ldg(lhs_off + i) : i;
|
||||
const Idx rhs_add = UseBcast ? __ldg(rhs_off + i) : i;
|
||||
DType val = reduce::Sum<Idx, DType>::zero();
|
||||
for (int j = tx; j < reduce_size; j += 64) {
|
||||
val += lhsoff[lhs_add * reduce_size + j] *
|
||||
rhsoff[rhs_add * reduce_size + j];
|
||||
if (j + 32 < reduce_size)
|
||||
val += lhsoff[lhs_add * reduce_size + j + 32] *
|
||||
rhsoff[rhs_add * reduce_size + j + 32];
|
||||
}
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset /= 2)
|
||||
val += __shfl_down_sync(full_mask, val, offset);
|
||||
if (tx == 0) outoff[i] = val;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Binary search the row_offsets to find the source node of the edge id.
|
||||
template <typename Idx>
|
||||
__device__ __forceinline__ Idx
|
||||
BinarySearchSrc(const Idx* array, Idx length, Idx eid) {
|
||||
Idx lo = 0, hi = length - 1;
|
||||
while (lo < hi) {
|
||||
Idx mid = (lo + hi) >> 1;
|
||||
if (_ldg(array + mid) <= eid) {
|
||||
lo = mid + 1;
|
||||
} else {
|
||||
hi = mid;
|
||||
}
|
||||
}
|
||||
// INVARIANT: lo == hi
|
||||
if (_ldg(array + hi) == eid) {
|
||||
return hi;
|
||||
} else {
|
||||
return hi - 1;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA kernel of g-SDDMM on Csr format.
|
||||
* @note it uses edge parallel strategy, different threadblocks (on y-axis)
|
||||
* is responsible for the computation on different edges. Threadblocks
|
||||
* on the x-axis are responsible for the computation on different
|
||||
* positions in feature dimension. To efficiently find the source node idx and
|
||||
* destination node index of an given edge on Csr format, it uses binary search
|
||||
* (time complexity O(log N)).
|
||||
*/
|
||||
template <
|
||||
typename Idx, typename DType, typename BinaryOp, bool UseBcast = false,
|
||||
bool UseIdx = false, int LhsTarget = 0, int RhsTarget = 2>
|
||||
__global__ void SDDMMCsrKernel(
|
||||
const DType* __restrict__ lhs, const DType* __restrict__ rhs,
|
||||
DType* __restrict__ out, const Idx* __restrict__ indptr,
|
||||
const Idx* __restrict__ indices, const Idx* __restrict__ edge_map,
|
||||
int64_t N, int64_t M, int64_t E, int64_t reduce_size,
|
||||
const int64_t* __restrict__ lhs_off, const int64_t* __restrict__ rhs_off,
|
||||
int64_t lhs_len, int64_t rhs_len, int64_t out_len) {
|
||||
// SDDMM with Csr.
|
||||
Idx ty = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
const Idx stride_y = blockDim.y * gridDim.y;
|
||||
while (ty < E) {
|
||||
const Idx src = BinarySearchSrc<Idx>(indptr, N + 1, ty);
|
||||
const Idx dst = _ldg(indices + ty);
|
||||
const Idx eid = UseIdx ? _ldg(edge_map + ty) : ty;
|
||||
int64_t tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int64_t stride_x = blockDim.x * gridDim.x;
|
||||
const DType* lhsoff =
|
||||
BinaryOp::use_lhs
|
||||
? (lhs + Selector<LhsTarget>::Call(src, eid, dst) * lhs_len)
|
||||
: nullptr;
|
||||
const DType* rhsoff =
|
||||
BinaryOp::use_rhs
|
||||
? (rhs + Selector<RhsTarget>::Call(src, eid, dst) * rhs_len)
|
||||
: nullptr;
|
||||
DType* outoff = out + eid * out_len;
|
||||
while (tx < out_len) {
|
||||
const Idx lhs_add = UseBcast ? lhs_off[tx] : tx;
|
||||
const Idx rhs_add = UseBcast ? rhs_off[tx] : tx;
|
||||
DType val = BinaryOp::Call(
|
||||
lhsoff + lhs_add * reduce_size, rhsoff + rhs_add * reduce_size,
|
||||
reduce_size);
|
||||
outoff[tx] = val;
|
||||
tx += stride_x;
|
||||
}
|
||||
ty += stride_y;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA implementation of g-SDDMM on Coo format.
|
||||
* @param bcast Broadcast information.
|
||||
* @param coo The Coo matrix.
|
||||
* @param lhs The left hand side operand feature.
|
||||
* @param rhs The right hand size operand feature.
|
||||
* @param out The result feature on edges.
|
||||
*/
|
||||
template <
|
||||
typename Idx, typename DType, typename Op, int LhsTarget = 0,
|
||||
int RhsTarget = 2>
|
||||
void SDDMMCoo(
|
||||
const BcastOff& bcast, const COOMatrix& coo, NDArray lhs, NDArray rhs,
|
||||
NDArray out) {
|
||||
const Idx* row = coo.row.Ptr<Idx>();
|
||||
const Idx* col = coo.col.Ptr<Idx>();
|
||||
const Idx* edge_map = coo.data.Ptr<Idx>();
|
||||
const DType* lhs_data = lhs.Ptr<DType>();
|
||||
const DType* rhs_data = rhs.Ptr<DType>();
|
||||
DType* out_data = out.Ptr<DType>();
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
|
||||
int64_t *lhs_off = nullptr, *rhs_off = nullptr;
|
||||
int64_t len = bcast.out_len, lhs_len = bcast.lhs_len, rhs_len = bcast.rhs_len;
|
||||
int64_t reduce_dim = bcast.reduce_size;
|
||||
|
||||
const int64_t nnz = coo.row->shape[0];
|
||||
const bool use_idx = !IsNullArray(coo.data);
|
||||
|
||||
if (std::is_same<Op, binary::Dot<DType> >::value && reduce_dim >= 32) {
|
||||
const int ntx = 32; // on feature dimension
|
||||
const int nty = 8; // on out dimension
|
||||
const int nbx = (nnz + nty - 1) / nty;
|
||||
const int nby = FindNumBlocks<'y'>(len);
|
||||
const dim3 nblks(nbx, nby);
|
||||
const dim3 nthrs(ntx, nty);
|
||||
BCAST_IDX_CTX_SWITCH(bcast, use_idx, out->ctx, lhs_off, rhs_off, {
|
||||
CUDA_KERNEL_CALL(
|
||||
(SDDMMCooTreeReduceKernel<
|
||||
Idx, DType, UseBcast, UseIdx, LhsTarget, RhsTarget>),
|
||||
nblks, nthrs, 0, stream, lhs_data, rhs_data, out_data, row, col,
|
||||
edge_map, coo.num_rows, coo.num_cols, nnz, reduce_dim, lhs_off,
|
||||
rhs_off, lhs_len, rhs_len, len);
|
||||
});
|
||||
} else {
|
||||
const int ntx = FindNumThreads(len);
|
||||
const int nty = CUDA_MAX_NUM_THREADS / ntx;
|
||||
const int nbx = (len + ntx - 1) / ntx;
|
||||
const int nby = FindNumBlocks<'y'>((nnz + nty - 1) / nty);
|
||||
const dim3 nblks(nbx, nby);
|
||||
const dim3 nthrs(ntx, nty);
|
||||
BCAST_IDX_CTX_SWITCH(bcast, use_idx, out->ctx, lhs_off, rhs_off, {
|
||||
CUDA_KERNEL_CALL(
|
||||
(SDDMMCooKernel<
|
||||
Idx, DType, Op, UseBcast, UseIdx, LhsTarget, RhsTarget>),
|
||||
nblks, nthrs, 0, stream, lhs_data, rhs_data, out_data, row, col,
|
||||
edge_map, coo.num_rows, coo.num_cols, nnz, reduce_dim, lhs_off,
|
||||
rhs_off, lhs_len, rhs_len, len);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA implementation of g-SDDMM on Csr format.
|
||||
* @param bcast Broadcast information.
|
||||
* @param csr The Csr matrix.
|
||||
* @param lhs The left hand side operand feature.
|
||||
* @param rhs The right hand size operand feature.
|
||||
* @param out The result feature on edges.
|
||||
*/
|
||||
template <
|
||||
typename Idx, typename DType, typename Op, int LhsTarget = 0,
|
||||
int RhsTarget = 2>
|
||||
void SDDMMCsr(
|
||||
const BcastOff& bcast, const CSRMatrix& csr, NDArray lhs, NDArray rhs,
|
||||
NDArray out) {
|
||||
const Idx* indptr = csr.indptr.Ptr<Idx>();
|
||||
const Idx* indices = csr.indices.Ptr<Idx>();
|
||||
const Idx* edge_map = csr.data.Ptr<Idx>();
|
||||
const DType* lhs_data = lhs.Ptr<DType>();
|
||||
const DType* rhs_data = rhs.Ptr<DType>();
|
||||
DType* out_data = out.Ptr<DType>();
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
int64_t N = csr.num_rows, M = csr.num_cols, E = csr.indices->shape[0];
|
||||
|
||||
int64_t *lhs_off = nullptr, *rhs_off = nullptr;
|
||||
int64_t len = bcast.out_len, lhs_len = bcast.lhs_len, rhs_len = bcast.rhs_len;
|
||||
int64_t reduce_dim = bcast.reduce_size;
|
||||
|
||||
const int ntx = FindNumThreads(len);
|
||||
const int nty = CUDA_MAX_NUM_THREADS / ntx;
|
||||
const int nbx = (len + ntx - 1) / ntx;
|
||||
const int nby = FindNumBlocks<'y'>((E + nty - 1) / nty);
|
||||
const dim3 nblks(nbx, nby);
|
||||
const dim3 nthrs(ntx, nty);
|
||||
const bool use_idx = !IsNullArray(csr.data);
|
||||
|
||||
BCAST_IDX_CTX_SWITCH(bcast, use_idx, out->ctx, lhs_off, rhs_off, {
|
||||
CUDA_KERNEL_CALL(
|
||||
(SDDMMCsrKernel<
|
||||
Idx, DType, Op, UseBcast, UseIdx, LhsTarget, RhsTarget>),
|
||||
nblks, nthrs, 0, stream, lhs_data, rhs_data, out_data, indptr, indices,
|
||||
edge_map, N, M, E, reduce_dim, lhs_off, rhs_off, lhs_len, rhs_len, len);
|
||||
});
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CUDA_SDDMM_CUH_
|
||||
@@ -0,0 +1,91 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/sddmm.cu
|
||||
* @brief SDDMM C APIs and definitions.
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include "./sddmm.cuh"
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
|
||||
/**
|
||||
* @brief CUDA implementation of g-SDDMM on heterograph using
|
||||
Csr format.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SDDMMCooHetero(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& vec_lhs,
|
||||
const std::vector<NDArray>& vec_rhs, std::vector<NDArray> vec_out,
|
||||
int lhs_target, int rhs_target, const std::vector<dgl_type_t>& lhs_eid,
|
||||
const std::vector<dgl_type_t>& rhs_eid) {
|
||||
SWITCH_OP(op, Op, {
|
||||
SWITCH_TARGET(lhs_target, rhs_target, LhsTarget, RhsTarget, {
|
||||
/* Call SDDMM CUDA kernel for each relation type sequentially */
|
||||
for (dgl_type_t etype = 0; etype < lhs_eid.size(); ++etype) {
|
||||
COOMatrix coo = vec_coo[etype];
|
||||
NDArray lhs = vec_lhs[lhs_eid[etype]];
|
||||
NDArray rhs = vec_rhs[rhs_eid[etype]];
|
||||
NDArray out = vec_out[etype];
|
||||
cuda::SDDMMCoo<IdType, DType, Op, LhsTarget, RhsTarget>(
|
||||
bcast, coo, lhs, rhs, out);
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
template void SDDMMCooHetero<kDGLCUDA, int32_t, __half>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCooHetero<kDGLCUDA, int64_t, __half>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
#if BF16_ENABLED
|
||||
template void SDDMMCooHetero<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCooHetero<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
#endif // BF16_ENABLED
|
||||
template void SDDMMCooHetero<kDGLCUDA, int32_t, float>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCooHetero<kDGLCUDA, int64_t, float>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCooHetero<kDGLCUDA, int32_t, double>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCooHetero<kDGLCUDA, int64_t, double>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,90 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/sddmm.cu
|
||||
* @brief SDDMM C APIs and definitions.
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include "./sddmm.cuh"
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
|
||||
/**
|
||||
* @brief CUDA implementation of g-SDDMM on heterograph using Csr format.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SDDMMCsrHetero(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& vec_lhs,
|
||||
const std::vector<NDArray>& vec_rhs, std::vector<NDArray> vec_out,
|
||||
int lhs_target, int rhs_target, const std::vector<dgl_type_t>& lhs_eid,
|
||||
const std::vector<dgl_type_t>& rhs_eid) {
|
||||
SWITCH_OP(op, Op, {
|
||||
SWITCH_TARGET(lhs_target, rhs_target, LhsTarget, RhsTarget, {
|
||||
/* Call SDDMM CUDA kernel for each relation type sequentially */
|
||||
for (dgl_type_t etype = 0; etype < lhs_eid.size(); ++etype) {
|
||||
CSRMatrix csr = vec_csr[etype];
|
||||
NDArray lhs = vec_lhs[lhs_eid[etype]];
|
||||
NDArray rhs = vec_rhs[rhs_eid[etype]];
|
||||
NDArray out = vec_out[etype];
|
||||
cuda::SDDMMCsr<IdType, DType, Op, LhsTarget, RhsTarget>(
|
||||
bcast, csr, lhs, rhs, out);
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
template void SDDMMCsrHetero<kDGLCUDA, int32_t, __half>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCsrHetero<kDGLCUDA, int64_t, __half>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
#if BF16_ENABLED
|
||||
template void SDDMMCsrHetero<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCsrHetero<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
#endif // BF16_ENABLED
|
||||
template void SDDMMCsrHetero<kDGLCUDA, int32_t, float>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCsrHetero<kDGLCUDA, int64_t, float>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCsrHetero<kDGLCUDA, int32_t, double>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
template void SDDMMCsrHetero<kDGLCUDA, int64_t, double>(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
|
||||
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target, const std::vector<dgl_type_t>& in_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,157 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/segment_reduce.cu
|
||||
* @brief Segment reduce C APIs and definitions.
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/base_heterograph.h>
|
||||
|
||||
#include "./functor.cuh"
|
||||
#include "./segment_reduce.cuh"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using namespace cuda;
|
||||
|
||||
namespace aten {
|
||||
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SegmentReduce(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg) {
|
||||
if (op == "sum") {
|
||||
cuda::SegmentReduce<IdType, DType, cuda::reduce::Sum<IdType, DType>>(
|
||||
feat, offsets, out, arg);
|
||||
} else if (op == "max") {
|
||||
cuda::SegmentReduce<IdType, DType, cuda::reduce::Max<IdType, DType>>(
|
||||
feat, offsets, out, arg);
|
||||
} else if (op == "min") {
|
||||
cuda::SegmentReduce<IdType, DType, cuda::reduce::Min<IdType, DType>>(
|
||||
feat, offsets, out, arg);
|
||||
} else {
|
||||
LOG(FATAL) << "Not implemented";
|
||||
}
|
||||
}
|
||||
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void ScatterAdd(NDArray feat, NDArray idx, NDArray out) {
|
||||
cuda::ScatterAdd<IdType, DType>(feat, idx, out);
|
||||
}
|
||||
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void UpdateGradMinMax_hetero(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out) {
|
||||
cuda::UpdateGradMinMax_hetero<IdType, DType>(
|
||||
g, op, feat, idx, idx_etype, out);
|
||||
}
|
||||
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void BackwardSegmentCmp(NDArray feat, NDArray arg, NDArray out) {
|
||||
cuda::BackwardSegmentCmp<IdType, DType>(feat, arg, out);
|
||||
}
|
||||
|
||||
template void SegmentReduce<kDGLCUDA, int32_t, __half>(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg);
|
||||
template void SegmentReduce<kDGLCUDA, int64_t, __half>(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg);
|
||||
#if BF16_ENABLED
|
||||
template void SegmentReduce<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg);
|
||||
template void SegmentReduce<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg);
|
||||
#endif // BF16_ENABLED
|
||||
template void SegmentReduce<kDGLCUDA, int32_t, float>(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg);
|
||||
template void SegmentReduce<kDGLCUDA, int64_t, float>(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg);
|
||||
template void SegmentReduce<kDGLCUDA, int32_t, double>(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg);
|
||||
template void SegmentReduce<kDGLCUDA, int64_t, double>(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg);
|
||||
|
||||
template void ScatterAdd<kDGLCUDA, int32_t, __half>(
|
||||
NDArray feat, NDArray idx, NDArray out);
|
||||
template void ScatterAdd<kDGLCUDA, int64_t, __half>(
|
||||
NDArray feat, NDArray idx, NDArray out);
|
||||
#if BF16_ENABLED
|
||||
template void ScatterAdd<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
NDArray feat, NDArray idx, NDArray out);
|
||||
template void ScatterAdd<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
NDArray feat, NDArray idx, NDArray out);
|
||||
#endif // BF16_ENABLED
|
||||
template void ScatterAdd<kDGLCUDA, int32_t, float>(
|
||||
NDArray feat, NDArray idx, NDArray out);
|
||||
template void ScatterAdd<kDGLCUDA, int64_t, float>(
|
||||
NDArray feat, NDArray idx, NDArray out);
|
||||
template void ScatterAdd<kDGLCUDA, int32_t, double>(
|
||||
NDArray feat, NDArray idx, NDArray out);
|
||||
template void ScatterAdd<kDGLCUDA, int64_t, double>(
|
||||
NDArray feat, NDArray idx, NDArray out);
|
||||
|
||||
template void UpdateGradMinMax_hetero<kDGLCUDA, int32_t, __half>(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out);
|
||||
template void UpdateGradMinMax_hetero<kDGLCUDA, int64_t, __half>(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out);
|
||||
#if BF16_ENABLED
|
||||
template void UpdateGradMinMax_hetero<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out);
|
||||
template void UpdateGradMinMax_hetero<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out);
|
||||
#endif // BF16_ENABLED
|
||||
template void UpdateGradMinMax_hetero<kDGLCUDA, int32_t, float>(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out);
|
||||
template void UpdateGradMinMax_hetero<kDGLCUDA, int64_t, float>(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out);
|
||||
template void UpdateGradMinMax_hetero<kDGLCUDA, int32_t, double>(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out);
|
||||
template void UpdateGradMinMax_hetero<kDGLCUDA, int64_t, double>(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out);
|
||||
|
||||
template void BackwardSegmentCmp<kDGLCUDA, int32_t, __half>(
|
||||
NDArray feat, NDArray arg, NDArray out);
|
||||
template void BackwardSegmentCmp<kDGLCUDA, int64_t, __half>(
|
||||
NDArray feat, NDArray arg, NDArray out);
|
||||
#if BF16_ENABLED
|
||||
template void BackwardSegmentCmp<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
NDArray feat, NDArray arg, NDArray out);
|
||||
template void BackwardSegmentCmp<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
NDArray feat, NDArray arg, NDArray out);
|
||||
#endif // BF16_ENABLED
|
||||
template void BackwardSegmentCmp<kDGLCUDA, int32_t, float>(
|
||||
NDArray feat, NDArray arg, NDArray out);
|
||||
template void BackwardSegmentCmp<kDGLCUDA, int64_t, float>(
|
||||
NDArray feat, NDArray arg, NDArray out);
|
||||
template void BackwardSegmentCmp<kDGLCUDA, int32_t, double>(
|
||||
NDArray feat, NDArray arg, NDArray out);
|
||||
template void BackwardSegmentCmp<kDGLCUDA, int64_t, double>(
|
||||
NDArray feat, NDArray arg, NDArray out);
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,262 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/segment_reduce.cuh
|
||||
* @brief Segment reduce kernel function header.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CUDA_SEGMENT_REDUCE_CUH_
|
||||
#define DGL_ARRAY_CUDA_SEGMENT_REDUCE_CUH_
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./atomic.cuh"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using namespace cuda;
|
||||
using namespace runtime;
|
||||
|
||||
namespace aten {
|
||||
namespace cuda {
|
||||
|
||||
/**
|
||||
* @brief CUDA kernel of segment reduce.
|
||||
* @note each blockthread is responsible for aggregation on a row
|
||||
* in the result tensor.
|
||||
*/
|
||||
template <typename IdType, typename DType, typename ReduceOp>
|
||||
__global__ void SegmentReduceKernel(
|
||||
const DType* feat, const IdType* offsets, DType* out, IdType* arg,
|
||||
int64_t n, int64_t dim) {
|
||||
for (int row = blockIdx.x; row < n; row += gridDim.x) {
|
||||
int col = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
while (col < dim) {
|
||||
typename accum_dtype<DType>::type local_accum = ReduceOp::zero();
|
||||
IdType local_arg = -1;
|
||||
for (IdType i = offsets[row]; i < offsets[row + 1]; ++i) {
|
||||
ReduceOp::Call(&local_accum, &local_arg, feat[i * dim + col], i);
|
||||
}
|
||||
out[row * dim + col] = static_cast<DType>(local_accum);
|
||||
if (ReduceOp::require_arg) arg[row * dim + col] = local_arg;
|
||||
col += gridDim.y * blockDim.x;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA kernel of scatter add.
|
||||
* @note each blockthread is responsible for adding a row in feature tensor
|
||||
* to a target row in output tensor.
|
||||
*/
|
||||
template <typename IdType, typename DType>
|
||||
__global__ void ScatterAddKernel(
|
||||
const DType* feat, const IdType* idx, DType* out, int64_t n, int64_t dim) {
|
||||
for (int row = blockIdx.x; row < n; row += gridDim.x) {
|
||||
const int write_row = idx[row];
|
||||
int col = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
while (col < dim) {
|
||||
cuda::AtomicAdd(out + write_row * dim + col, feat[row * dim + col]);
|
||||
col += gridDim.y * blockDim.x;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA kernel to update gradients for reduce op max/min
|
||||
* @note each WARP (group of 32 threads) is responsible for adding a row in
|
||||
* feature tensor to a target row in output tensor.
|
||||
*/
|
||||
|
||||
template <typename IdType, typename DType>
|
||||
__global__ void UpdateGradMinMaxHeteroKernel(
|
||||
const DType* feat, const IdType* idx, const IdType* idx_type, DType* out,
|
||||
int64_t n, int64_t dim, int type) {
|
||||
unsigned int tId = threadIdx.x;
|
||||
unsigned int laneId = tId & 31;
|
||||
unsigned int gId = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
unsigned int warpId = gId >> 5;
|
||||
unsigned int warp_size = 32;
|
||||
unsigned int row = warpId;
|
||||
|
||||
while (row < n) {
|
||||
for (unsigned int col = laneId; col < dim; col += warp_size) {
|
||||
if (type == idx_type[row * dim + col]) {
|
||||
const int write_row = idx[row * dim + col];
|
||||
cuda::AtomicAdd(out + write_row * dim + col, feat[row * dim + col]);
|
||||
}
|
||||
}
|
||||
row += blockDim.x * gridDim.x;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA kernel of backward phase in segment min/max.
|
||||
* @note each blockthread is responsible for writing a row in the
|
||||
* result gradient tensor by lookup the ArgMin/Max for index information.
|
||||
*/
|
||||
template <typename IdType, typename DType>
|
||||
__global__ void BackwardSegmentCmpKernel(
|
||||
const DType* feat, const IdType* arg, DType* out, int64_t n, int64_t dim) {
|
||||
for (int row = blockIdx.x; row < n; row += gridDim.x) {
|
||||
int col = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
while (col < dim) {
|
||||
int write_row = arg[row * dim + col];
|
||||
if (write_row >= 0) {
|
||||
out[write_row * dim + col] = feat[row * dim + col];
|
||||
}
|
||||
col += gridDim.y * blockDim.x;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA implementation of forward phase of Segment Reduce.
|
||||
* @param feat The input tensor.
|
||||
* @param offsets The offsets tensor.
|
||||
* @param out The output tensor.
|
||||
* @param arg An auxiliary tensor storing ArgMax/Min information,
|
||||
*/
|
||||
template <typename IdType, typename DType, typename ReduceOp>
|
||||
void SegmentReduce(NDArray feat, NDArray offsets, NDArray out, NDArray arg) {
|
||||
const DType* feat_data = feat.Ptr<DType>();
|
||||
const IdType* offsets_data = offsets.Ptr<IdType>();
|
||||
DType* out_data = out.Ptr<DType>();
|
||||
IdType* arg_data = arg.Ptr<IdType>();
|
||||
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
int64_t n = out->shape[0];
|
||||
int64_t dim = 1;
|
||||
for (int i = 1; i < out->ndim; ++i) dim *= out->shape[i];
|
||||
|
||||
const int nbx = FindNumBlocks<'x'>(n);
|
||||
const int ntx = FindNumThreads(dim);
|
||||
const int nby = FindNumBlocks<'y'>((dim + ntx - 1) / ntx);
|
||||
const int nty = 1;
|
||||
const dim3 nblks(nbx, nby);
|
||||
const dim3 nthrs(ntx, nty);
|
||||
// TODO(zihao): try cub's DeviceSegmentedReduce and compare the performance.
|
||||
CUDA_KERNEL_CALL(
|
||||
(SegmentReduceKernel<IdType, DType, ReduceOp>), nblks, nthrs, 0, stream,
|
||||
feat_data, offsets_data, out_data, arg_data, n, dim);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA implementation of Scatter Add (on first dimension).
|
||||
* @note math equation: out[idx[i], *] += feat[i, *]
|
||||
* @param feat The input tensor.
|
||||
* @param idx The indices tensor.
|
||||
* @param out The output tensor.
|
||||
*/
|
||||
template <typename IdType, typename DType>
|
||||
void ScatterAdd(NDArray feat, NDArray idx, NDArray out) {
|
||||
const DType* feat_data = feat.Ptr<DType>();
|
||||
const IdType* idx_data = idx.Ptr<IdType>();
|
||||
DType* out_data = out.Ptr<DType>();
|
||||
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
int64_t n = feat->shape[0];
|
||||
int64_t dim = 1;
|
||||
for (int i = 1; i < out->ndim; ++i) dim *= out->shape[i];
|
||||
|
||||
const int nbx = FindNumBlocks<'x'>(n);
|
||||
const int ntx = FindNumThreads(dim);
|
||||
const int nby = FindNumBlocks<'y'>((dim + ntx - 1) / ntx);
|
||||
const int nty = 1;
|
||||
const dim3 nblks(nbx, nby);
|
||||
const dim3 nthrs(ntx, nty);
|
||||
CUDA_KERNEL_CALL(
|
||||
(ScatterAddKernel<IdType, DType>), nblks, nthrs, 0, stream, feat_data,
|
||||
idx_data, out_data, n, dim);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA implementation to update gradients for reduce op max/min
|
||||
* @param graph The input heterogeneous graph.
|
||||
* @param op The binary operator, could be `copy_u`, `copy_e'.
|
||||
* @param list_feat List of the input tensors.
|
||||
* @param list_idx List of the indices tensors.
|
||||
* @param list_idx_etype List of the node- or edge-type tensors.
|
||||
* @param list_out List of the output tensors.
|
||||
*/
|
||||
template <typename IdType, typename DType>
|
||||
void UpdateGradMinMax_hetero(
|
||||
const HeteroGraphPtr& graph, const std::string& op,
|
||||
const std::vector<NDArray>& list_feat, const std::vector<NDArray>& list_idx,
|
||||
const std::vector<NDArray>& list_idx_types,
|
||||
std::vector<NDArray>* list_out) {
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
if (op == "copy_lhs" || op == "copy_rhs") {
|
||||
std::vector<std::vector<dgl_id_t>> src_dst_ntypes(
|
||||
graph->NumVertexTypes(), std::vector<dgl_id_t>());
|
||||
for (dgl_type_t etype = 0; etype < graph->NumEdgeTypes(); ++etype) {
|
||||
auto pair = graph->meta_graph()->FindEdge(etype);
|
||||
const dgl_id_t dst_ntype = pair.first; // graph is reversed
|
||||
const dgl_id_t src_ntype = pair.second;
|
||||
auto same_src_dst_ntype = std::find(
|
||||
std::begin(src_dst_ntypes[dst_ntype]),
|
||||
std::end(src_dst_ntypes[dst_ntype]), src_ntype);
|
||||
// if op is "copy_lhs", relation type with same src and dst node type will
|
||||
// be updated once
|
||||
if (op == "copy_lhs" &&
|
||||
same_src_dst_ntype != std::end(src_dst_ntypes[dst_ntype]))
|
||||
continue;
|
||||
src_dst_ntypes[dst_ntype].push_back(src_ntype);
|
||||
const DType* feat_data = list_feat[dst_ntype].Ptr<DType>();
|
||||
const IdType* idx_data = list_idx[dst_ntype].Ptr<IdType>();
|
||||
const IdType* idx_type_data = list_idx_types[dst_ntype].Ptr<IdType>();
|
||||
int type = (op == "copy_lhs") ? src_ntype : etype;
|
||||
DType* out_data = (*list_out)[type].Ptr<DType>();
|
||||
int dim = 1;
|
||||
for (int i = 1; i < (*list_out)[type]->ndim; ++i)
|
||||
dim *= (*list_out)[type]->shape[i];
|
||||
int n = list_feat[dst_ntype]->shape[0];
|
||||
const int th_per_row = 32;
|
||||
const int ntx = 128;
|
||||
const int nbx = FindNumBlocks<'x'>((n * th_per_row + ntx - 1) / ntx);
|
||||
const dim3 nblks(nbx);
|
||||
const dim3 nthrs(ntx);
|
||||
CUDA_KERNEL_CALL(
|
||||
(UpdateGradMinMaxHeteroKernel<IdType, DType>), nblks, nthrs, 0,
|
||||
stream, feat_data, idx_data, idx_type_data, out_data, n, dim, type);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA implementation of backward phase of Segment Reduce with Min/Max
|
||||
* reducer.
|
||||
* @note math equation: out[arg[i, k], k] = feat[i, k]
|
||||
* @param feat The input
|
||||
* tensor.
|
||||
* @param arg The ArgMin/Max information, used for indexing.
|
||||
* @param out The output tensor.
|
||||
*/
|
||||
template <typename IdType, typename DType>
|
||||
void BackwardSegmentCmp(NDArray feat, NDArray arg, NDArray out) {
|
||||
const DType* feat_data = feat.Ptr<DType>();
|
||||
const IdType* arg_data = arg.Ptr<IdType>();
|
||||
DType* out_data = out.Ptr<DType>();
|
||||
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
int64_t n = feat->shape[0];
|
||||
int64_t dim = 1;
|
||||
for (int i = 1; i < out->ndim; ++i) dim *= out->shape[i];
|
||||
|
||||
const int nbx = FindNumBlocks<'x'>(n);
|
||||
const int ntx = FindNumThreads(dim);
|
||||
const int nby = FindNumBlocks<'y'>((dim + ntx - 1) / ntx);
|
||||
const int nty = 1;
|
||||
const dim3 nblks(nbx, nby);
|
||||
const dim3 nthrs(ntx, nty);
|
||||
CUDA_KERNEL_CALL(
|
||||
(BackwardSegmentCmpKernel<IdType, DType>), nblks, nthrs, 0, stream,
|
||||
feat_data, arg_data, out_data, n, dim);
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CUDA_SEGMENT_REDUCE_CUH_
|
||||
@@ -0,0 +1,139 @@
|
||||
/**
|
||||
* Copyright (c) 2021 by contributors.
|
||||
* @file array/cuda/spmat_op_impl_coo.cu
|
||||
* @brief COO operator GPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include <numeric>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./atomic.cuh"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using runtime::NDArray;
|
||||
using namespace cuda;
|
||||
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <typename IdType>
|
||||
__device__ void _warpReduce(volatile IdType* sdata, IdType tid) {
|
||||
sdata[tid] += sdata[tid + 32];
|
||||
sdata[tid] += sdata[tid + 16];
|
||||
sdata[tid] += sdata[tid + 8];
|
||||
sdata[tid] += sdata[tid + 4];
|
||||
sdata[tid] += sdata[tid + 2];
|
||||
sdata[tid] += sdata[tid + 1];
|
||||
}
|
||||
|
||||
template <typename IdType>
|
||||
__global__ void _COOGetRowNNZKernel(
|
||||
const IdType* __restrict__ row_indices, IdType* __restrict__ glb_cnt,
|
||||
const int64_t row_query, IdType nnz) {
|
||||
__shared__ IdType local_cnt[1024];
|
||||
IdType tx = threadIdx.x;
|
||||
IdType bx = blockIdx.x;
|
||||
local_cnt[tx] = 0;
|
||||
IdType start = bx * blockDim.x;
|
||||
while (start < nnz) {
|
||||
if (start + tx < nnz)
|
||||
local_cnt[tx] = (row_indices[start + tx] == row_query);
|
||||
__syncthreads();
|
||||
if (tx < 512) {
|
||||
local_cnt[tx] += local_cnt[tx + 512];
|
||||
__syncthreads();
|
||||
}
|
||||
if (tx < 256) {
|
||||
local_cnt[tx] += local_cnt[tx + 256];
|
||||
__syncthreads();
|
||||
}
|
||||
if (tx < 128) {
|
||||
local_cnt[tx] += local_cnt[tx + 128];
|
||||
__syncthreads();
|
||||
}
|
||||
if (tx < 64) {
|
||||
local_cnt[tx] += local_cnt[tx + 64];
|
||||
__syncthreads();
|
||||
}
|
||||
if (tx < 32) {
|
||||
_warpReduce(local_cnt, tx);
|
||||
}
|
||||
if (tx == 0) {
|
||||
cuda::AtomicAdd(glb_cnt, local_cnt[tx]);
|
||||
}
|
||||
start += blockDim.x * gridDim.x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
int64_t COOGetRowNNZ(COOMatrix coo, int64_t row) {
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const auto& ctx = coo.row->ctx;
|
||||
IdType nnz = coo.row->shape[0];
|
||||
IdType nt = 1024;
|
||||
IdType nb = dgl::cuda::FindNumBlocks<'x'>((nnz + nt - 1) / nt);
|
||||
NDArray rst = NDArray::Empty({1}, coo.row->dtype, coo.row->ctx);
|
||||
_Fill(rst.Ptr<IdType>(), 1, IdType(0));
|
||||
CUDA_KERNEL_CALL(
|
||||
_COOGetRowNNZKernel, nb, nt, 0, stream, coo.row.Ptr<IdType>(),
|
||||
rst.Ptr<IdType>(), row, nnz);
|
||||
rst = rst.CopyTo(DGLContext{kDGLCPU, 0});
|
||||
return *rst.Ptr<IdType>();
|
||||
}
|
||||
|
||||
template int64_t COOGetRowNNZ<kDGLCUDA, int32_t>(COOMatrix, int64_t);
|
||||
template int64_t COOGetRowNNZ<kDGLCUDA, int64_t>(COOMatrix, int64_t);
|
||||
|
||||
template <typename IdType>
|
||||
__global__ void _COOGetAllRowNNZKernel(
|
||||
const IdType* __restrict__ row_indices, IdType* __restrict__ glb_cnts,
|
||||
IdType nnz) {
|
||||
IdType eid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
while (eid < nnz) {
|
||||
IdType row = row_indices[eid];
|
||||
cuda::AtomicAdd(glb_cnts + row, IdType(1));
|
||||
eid += blockDim.x * gridDim.x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
NDArray COOGetRowNNZ(COOMatrix coo, NDArray rows) {
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const auto& ctx = coo.row->ctx;
|
||||
IdType nnz = coo.row->shape[0];
|
||||
IdType num_rows = coo.num_rows;
|
||||
IdType num_queries = rows->shape[0];
|
||||
if (num_queries == 1) {
|
||||
auto rows_cpu = rows.CopyTo(DGLContext{kDGLCPU, 0});
|
||||
int64_t row = *rows_cpu.Ptr<IdType>();
|
||||
IdType nt = 1024;
|
||||
IdType nb = dgl::cuda::FindNumBlocks<'x'>((nnz + nt - 1) / nt);
|
||||
NDArray rst = NDArray::Empty({1}, coo.row->dtype, coo.row->ctx);
|
||||
_Fill(rst.Ptr<IdType>(), 1, IdType(0));
|
||||
CUDA_KERNEL_CALL(
|
||||
_COOGetRowNNZKernel, nb, nt, 0, stream, coo.row.Ptr<IdType>(),
|
||||
rst.Ptr<IdType>(), row, nnz);
|
||||
return rst;
|
||||
} else {
|
||||
IdType nt = 1024;
|
||||
IdType nb = dgl::cuda::FindNumBlocks<'x'>((nnz + nt - 1) / nt);
|
||||
NDArray in_degrees = NDArray::Empty({num_rows}, rows->dtype, rows->ctx);
|
||||
_Fill(in_degrees.Ptr<IdType>(), num_rows, IdType(0));
|
||||
CUDA_KERNEL_CALL(
|
||||
_COOGetAllRowNNZKernel, nb, nt, 0, stream, coo.row.Ptr<IdType>(),
|
||||
in_degrees.Ptr<IdType>(), nnz);
|
||||
return IndexSelect(in_degrees, rows);
|
||||
}
|
||||
}
|
||||
|
||||
template NDArray COOGetRowNNZ<kDGLCUDA, int32_t>(COOMatrix, NDArray);
|
||||
template NDArray COOGetRowNNZ<kDGLCUDA, int64_t>(COOMatrix, NDArray);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,654 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/spmat_op_impl_csr.cu
|
||||
* @brief CSR operator CPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
#include <thrust/execution_policy.h>
|
||||
#include <thrust/for_each.h>
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
#include <numeric>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./atomic.cuh"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using runtime::NDArray;
|
||||
using namespace cuda;
|
||||
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
///////////////////////////// CSRIsNonZero /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
bool CSRIsNonZero(CSRMatrix csr, int64_t row, int64_t col) {
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const auto& ctx = csr.indptr->ctx;
|
||||
IdArray rows = aten::VecToIdArray<int64_t>({row}, sizeof(IdType) * 8, ctx);
|
||||
IdArray cols = aten::VecToIdArray<int64_t>({col}, sizeof(IdType) * 8, ctx);
|
||||
rows = rows.CopyTo(ctx);
|
||||
cols = cols.CopyTo(ctx);
|
||||
IdArray out = aten::NewIdArray(1, ctx, sizeof(IdType) * 8);
|
||||
const IdType* data = nullptr;
|
||||
// TODO(minjie): use binary search for sorted csr
|
||||
CUDA_KERNEL_CALL(
|
||||
dgl::cuda::_LinearSearchKernel, 1, 1, 0, stream, csr.indptr.Ptr<IdType>(),
|
||||
csr.indices.Ptr<IdType>(), data, rows.Ptr<IdType>(), cols.Ptr<IdType>(),
|
||||
1, 1, 1, static_cast<IdType*>(nullptr), static_cast<IdType>(-1),
|
||||
out.Ptr<IdType>());
|
||||
out = out.CopyTo(DGLContext{kDGLCPU, 0});
|
||||
return *out.Ptr<IdType>() != -1;
|
||||
}
|
||||
|
||||
template bool CSRIsNonZero<kDGLCUDA, int32_t>(CSRMatrix, int64_t, int64_t);
|
||||
template bool CSRIsNonZero<kDGLCUDA, int64_t>(CSRMatrix, int64_t, int64_t);
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
NDArray CSRIsNonZero(CSRMatrix csr, NDArray row, NDArray col) {
|
||||
const auto rowlen = row->shape[0];
|
||||
const auto collen = col->shape[0];
|
||||
const auto rstlen = std::max(rowlen, collen);
|
||||
NDArray rst = NDArray::Empty({rstlen}, row->dtype, row->ctx);
|
||||
if (rstlen == 0) return rst;
|
||||
const int64_t row_stride = (rowlen == 1 && collen != 1) ? 0 : 1;
|
||||
const int64_t col_stride = (collen == 1 && rowlen != 1) ? 0 : 1;
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const int nt = dgl::cuda::FindNumThreads(rstlen);
|
||||
const int nb = (rstlen + nt - 1) / nt;
|
||||
const IdType* data = nullptr;
|
||||
const IdType* indptr_data =
|
||||
static_cast<IdType*>(GetDevicePointer(csr.indptr));
|
||||
const IdType* indices_data =
|
||||
static_cast<IdType*>(GetDevicePointer(csr.indices));
|
||||
// TODO(minjie): use binary search for sorted csr
|
||||
CUDA_KERNEL_CALL(
|
||||
dgl::cuda::_LinearSearchKernel, nb, nt, 0, stream, indptr_data,
|
||||
indices_data, data, row.Ptr<IdType>(), col.Ptr<IdType>(), row_stride,
|
||||
col_stride, rstlen, static_cast<IdType*>(nullptr),
|
||||
static_cast<IdType>(-1), rst.Ptr<IdType>());
|
||||
return rst != -1;
|
||||
}
|
||||
|
||||
template NDArray CSRIsNonZero<kDGLCUDA, int32_t>(CSRMatrix, NDArray, NDArray);
|
||||
template NDArray CSRIsNonZero<kDGLCUDA, int64_t>(CSRMatrix, NDArray, NDArray);
|
||||
|
||||
///////////////////////////// CSRHasDuplicate /////////////////////////////
|
||||
|
||||
/**
|
||||
* @brief Check whether each row does not have any duplicate entries.
|
||||
* Assume the CSR is sorted.
|
||||
*/
|
||||
template <typename IdType>
|
||||
__global__ void _SegmentHasNoDuplicate(
|
||||
const IdType* indptr, const IdType* indices, int64_t num_rows,
|
||||
int8_t* flags) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < num_rows) {
|
||||
bool f = true;
|
||||
for (IdType i = indptr[tx] + 1; f && i < indptr[tx + 1]; ++i) {
|
||||
f = (indices[i - 1] != indices[i]);
|
||||
}
|
||||
flags[tx] = static_cast<int8_t>(f);
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
bool CSRHasDuplicate(CSRMatrix csr) {
|
||||
if (!csr.sorted) csr = CSRSort(csr);
|
||||
const auto& ctx = csr.indptr->ctx;
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
// We allocate a workspace of num_rows bytes. It wastes a little bit memory
|
||||
// but should be fine.
|
||||
int8_t* flags =
|
||||
static_cast<int8_t*>(device->AllocWorkspace(ctx, csr.num_rows));
|
||||
const int nt = dgl::cuda::FindNumThreads(csr.num_rows);
|
||||
const int nb = (csr.num_rows + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
_SegmentHasNoDuplicate, nb, nt, 0, stream, csr.indptr.Ptr<IdType>(),
|
||||
csr.indices.Ptr<IdType>(), csr.num_rows, flags);
|
||||
bool ret = dgl::cuda::AllTrue(flags, csr.num_rows, ctx);
|
||||
device->FreeWorkspace(ctx, flags);
|
||||
return !ret;
|
||||
}
|
||||
|
||||
template bool CSRHasDuplicate<kDGLCUDA, int32_t>(CSRMatrix csr);
|
||||
template bool CSRHasDuplicate<kDGLCUDA, int64_t>(CSRMatrix csr);
|
||||
|
||||
///////////////////////////// CSRGetRowNNZ /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
int64_t CSRGetRowNNZ(CSRMatrix csr, int64_t row) {
|
||||
const IdType cur = aten::IndexSelect<IdType>(csr.indptr, row);
|
||||
const IdType next = aten::IndexSelect<IdType>(csr.indptr, row + 1);
|
||||
return next - cur;
|
||||
}
|
||||
|
||||
template int64_t CSRGetRowNNZ<kDGLCUDA, int32_t>(CSRMatrix, int64_t);
|
||||
template int64_t CSRGetRowNNZ<kDGLCUDA, int64_t>(CSRMatrix, int64_t);
|
||||
|
||||
template <typename IdType>
|
||||
__global__ void _CSRGetRowNNZKernel(
|
||||
const IdType* vid, const IdType* indptr, IdType* out, int64_t length) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < length) {
|
||||
const IdType vv = vid[tx];
|
||||
out[tx] = indptr[vv + 1] - indptr[vv];
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
NDArray CSRGetRowNNZ(CSRMatrix csr, NDArray rows) {
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const auto len = rows->shape[0];
|
||||
const IdType* vid_data = rows.Ptr<IdType>();
|
||||
const IdType* indptr_data =
|
||||
static_cast<IdType*>(GetDevicePointer(csr.indptr));
|
||||
NDArray rst = NDArray::Empty({len}, rows->dtype, rows->ctx);
|
||||
IdType* rst_data = static_cast<IdType*>(rst->data);
|
||||
const int nt = dgl::cuda::FindNumThreads(len);
|
||||
const int nb = (len + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
_CSRGetRowNNZKernel, nb, nt, 0, stream, vid_data, indptr_data, rst_data,
|
||||
len);
|
||||
return rst;
|
||||
}
|
||||
|
||||
template NDArray CSRGetRowNNZ<kDGLCUDA, int32_t>(CSRMatrix, NDArray);
|
||||
template NDArray CSRGetRowNNZ<kDGLCUDA, int64_t>(CSRMatrix, NDArray);
|
||||
|
||||
////////////////////////// CSRGetRowColumnIndices //////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
NDArray CSRGetRowColumnIndices(CSRMatrix csr, int64_t row) {
|
||||
const int64_t len = impl::CSRGetRowNNZ<XPU, IdType>(csr, row);
|
||||
const int64_t offset =
|
||||
aten::IndexSelect<IdType>(csr.indptr, row) * sizeof(IdType);
|
||||
return csr.indices.CreateView({len}, csr.indices->dtype, offset);
|
||||
}
|
||||
|
||||
template NDArray CSRGetRowColumnIndices<kDGLCUDA, int32_t>(CSRMatrix, int64_t);
|
||||
template NDArray CSRGetRowColumnIndices<kDGLCUDA, int64_t>(CSRMatrix, int64_t);
|
||||
|
||||
///////////////////////////// CSRGetRowData /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
NDArray CSRGetRowData(CSRMatrix csr, int64_t row) {
|
||||
const int64_t len = impl::CSRGetRowNNZ<XPU, IdType>(csr, row);
|
||||
const int64_t offset =
|
||||
aten::IndexSelect<IdType>(csr.indptr, row) * sizeof(IdType);
|
||||
if (aten::CSRHasData(csr))
|
||||
return csr.data.CreateView({len}, csr.data->dtype, offset);
|
||||
else
|
||||
return aten::Range(
|
||||
offset, offset + len, csr.indptr->dtype.bits, csr.indptr->ctx);
|
||||
}
|
||||
|
||||
template NDArray CSRGetRowData<kDGLCUDA, int32_t>(CSRMatrix, int64_t);
|
||||
template NDArray CSRGetRowData<kDGLCUDA, int64_t>(CSRMatrix, int64_t);
|
||||
|
||||
///////////////////////////// CSRSliceRows /////////////////////////////
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix CSRSliceRows(CSRMatrix csr, int64_t start, int64_t end) {
|
||||
const int64_t num_rows = end - start;
|
||||
const IdType st_pos = aten::IndexSelect<IdType>(csr.indptr, start);
|
||||
const IdType ed_pos = aten::IndexSelect<IdType>(csr.indptr, end);
|
||||
const IdType nnz = ed_pos - st_pos;
|
||||
IdArray ret_indptr = aten::IndexSelect(csr.indptr, start, end + 1) - st_pos;
|
||||
// indices and data can be view arrays
|
||||
IdArray ret_indices = csr.indices.CreateView(
|
||||
{nnz}, csr.indices->dtype, st_pos * sizeof(IdType));
|
||||
IdArray ret_data;
|
||||
if (CSRHasData(csr))
|
||||
ret_data =
|
||||
csr.data.CreateView({nnz}, csr.data->dtype, st_pos * sizeof(IdType));
|
||||
else
|
||||
ret_data =
|
||||
aten::Range(st_pos, ed_pos, csr.indptr->dtype.bits, csr.indptr->ctx);
|
||||
return CSRMatrix(
|
||||
num_rows, csr.num_cols, ret_indptr, ret_indices, ret_data, csr.sorted);
|
||||
}
|
||||
|
||||
template CSRMatrix CSRSliceRows<kDGLCUDA, int32_t>(CSRMatrix, int64_t, int64_t);
|
||||
template CSRMatrix CSRSliceRows<kDGLCUDA, int64_t>(CSRMatrix, int64_t, int64_t);
|
||||
|
||||
/**
|
||||
* @brief Copy data segment to output buffers
|
||||
*
|
||||
* For the i^th row r = row[i], copy the data from indptr[r] ~ indptr[r+1]
|
||||
* to the out_data from out_indptr[i] ~ out_indptr[i+1]
|
||||
*
|
||||
* If the provided `data` array is nullptr, write the read index to the
|
||||
* out_data.
|
||||
*
|
||||
*/
|
||||
template <typename IdType, typename DType>
|
||||
__global__ void _SegmentCopyKernel(
|
||||
const IdType* indptr, const DType* data, const IdType* row, int64_t length,
|
||||
int64_t n_row, const IdType* out_indptr, DType* out_data) {
|
||||
IdType tx = static_cast<IdType>(blockIdx.x) * blockDim.x + threadIdx.x;
|
||||
const int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < length) {
|
||||
IdType rpos = dgl::cuda::_UpperBound(out_indptr, n_row, tx) - 1;
|
||||
IdType rofs = tx - out_indptr[rpos];
|
||||
const IdType u = row[rpos];
|
||||
out_data[tx] = data ? data[indptr[u] + rofs] : indptr[u] + rofs;
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix CSRSliceRows(CSRMatrix csr, NDArray rows) {
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const int64_t len = rows->shape[0];
|
||||
IdArray ret_indptr = aten::CumSum(aten::CSRGetRowNNZ(csr, rows), true);
|
||||
const int64_t nnz = aten::IndexSelect<IdType>(ret_indptr, len);
|
||||
|
||||
const int nt = 256; // for better GPU usage of small invocations
|
||||
const int nb = (nnz + nt - 1) / nt;
|
||||
|
||||
// Copy indices.
|
||||
IdArray ret_indices = NDArray::Empty({nnz}, csr.indptr->dtype, rows->ctx);
|
||||
|
||||
const IdType* indptr_data =
|
||||
static_cast<IdType*>(GetDevicePointer(csr.indptr));
|
||||
const IdType* indices_data =
|
||||
static_cast<IdType*>(GetDevicePointer(csr.indices));
|
||||
const IdType* data_data =
|
||||
CSRHasData(csr) ? static_cast<IdType*>(GetDevicePointer(csr.data))
|
||||
: nullptr;
|
||||
|
||||
CUDA_KERNEL_CALL(
|
||||
_SegmentCopyKernel, nb, nt, 0, stream, indptr_data, indices_data,
|
||||
rows.Ptr<IdType>(), nnz, len, ret_indptr.Ptr<IdType>(),
|
||||
ret_indices.Ptr<IdType>());
|
||||
// Copy data.
|
||||
IdArray ret_data = NDArray::Empty({nnz}, csr.indptr->dtype, rows->ctx);
|
||||
CUDA_KERNEL_CALL(
|
||||
_SegmentCopyKernel, nb, nt, 0, stream, indptr_data, data_data,
|
||||
rows.Ptr<IdType>(), nnz, len, ret_indptr.Ptr<IdType>(),
|
||||
ret_data.Ptr<IdType>());
|
||||
return CSRMatrix(
|
||||
len, csr.num_cols, ret_indptr, ret_indices, ret_data, csr.sorted);
|
||||
}
|
||||
|
||||
template CSRMatrix CSRSliceRows<kDGLCUDA, int32_t>(CSRMatrix, NDArray);
|
||||
template CSRMatrix CSRSliceRows<kDGLCUDA, int64_t>(CSRMatrix, NDArray);
|
||||
|
||||
///////////////////////////// CSRGetDataAndIndices /////////////////////////////
|
||||
|
||||
/**
|
||||
* @brief Generate a 0-1 mask for each index that hits the provided (row, col)
|
||||
* index.
|
||||
*
|
||||
* Examples:
|
||||
* Given a CSR matrix (with duplicate entries) as follows:
|
||||
* [[0, 1, 2, 0, 0],
|
||||
* [1, 0, 0, 0, 0],
|
||||
* [0, 0, 1, 1, 0],
|
||||
* [0, 0, 0, 0, 0]]
|
||||
* Given rows: [0, 1], cols: [0, 2, 3]
|
||||
* The result mask is: [0, 1, 1, 1, 0, 0]
|
||||
*/
|
||||
template <typename IdType>
|
||||
__global__ void _SegmentMaskKernel(
|
||||
const IdType* indptr, const IdType* indices, const IdType* row,
|
||||
const IdType* col, int64_t row_stride, int64_t col_stride, int64_t length,
|
||||
IdType* mask) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < length) {
|
||||
int rpos = tx * row_stride, cpos = tx * col_stride;
|
||||
const IdType r = row[rpos], c = col[cpos];
|
||||
for (IdType i = indptr[r]; i < indptr[r + 1]; ++i) {
|
||||
if (indices[i] == c) {
|
||||
mask[i] = 1;
|
||||
}
|
||||
}
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Search for the insertion positions for needle in the hay.
|
||||
*
|
||||
* The hay is a list of sorted elements and the result is the insertion position
|
||||
* of each needle so that the insertion still gives sorted order.
|
||||
*
|
||||
* It essentially perform binary search to find lower bound for each needle
|
||||
* elements. Require the largest elements in the hay is larger than the given
|
||||
* needle elements. Commonly used in searching for row IDs of a given set of
|
||||
* coordinates.
|
||||
*/
|
||||
template <typename IdType>
|
||||
__global__ void _SortedSearchKernel(
|
||||
const IdType* hay, int64_t hay_size, const IdType* needles,
|
||||
int64_t num_needles, IdType* pos) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < num_needles) {
|
||||
const IdType ele = needles[tx];
|
||||
// binary search
|
||||
IdType lo = 0, hi = hay_size - 1;
|
||||
while (lo < hi) {
|
||||
IdType mid = (lo + hi) >> 1;
|
||||
if (hay[mid] <= ele) {
|
||||
lo = mid + 1;
|
||||
} else {
|
||||
hi = mid;
|
||||
}
|
||||
}
|
||||
pos[tx] = (hay[hi] == ele) ? hi : hi - 1;
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
std::vector<NDArray> CSRGetDataAndIndices(
|
||||
CSRMatrix csr, NDArray row, NDArray col) {
|
||||
const auto rowlen = row->shape[0];
|
||||
const auto collen = col->shape[0];
|
||||
const auto len = std::max(rowlen, collen);
|
||||
if (len == 0) return {NullArray(), NullArray(), NullArray()};
|
||||
|
||||
const auto& ctx = row->ctx;
|
||||
const auto nbits = row->dtype.bits;
|
||||
const int64_t nnz = csr.indices->shape[0];
|
||||
const int64_t row_stride = (rowlen == 1 && collen != 1) ? 0 : 1;
|
||||
const int64_t col_stride = (collen == 1 && rowlen != 1) ? 0 : 1;
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
|
||||
const IdType* indptr_data =
|
||||
static_cast<IdType*>(GetDevicePointer(csr.indptr));
|
||||
const IdType* indices_data =
|
||||
static_cast<IdType*>(GetDevicePointer(csr.indices));
|
||||
|
||||
// Generate a 0-1 mask for matched (row, col) positions.
|
||||
IdArray mask = Full(0, nnz, nbits, ctx);
|
||||
const int nt = dgl::cuda::FindNumThreads(len);
|
||||
const int nb = (len + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
_SegmentMaskKernel, nb, nt, 0, stream, indptr_data, indices_data,
|
||||
row.Ptr<IdType>(), col.Ptr<IdType>(), row_stride, col_stride, len,
|
||||
mask.Ptr<IdType>());
|
||||
|
||||
IdArray idx = AsNumBits(NonZero(mask), nbits);
|
||||
if (idx->shape[0] == 0)
|
||||
// No data. Return three empty arrays.
|
||||
return {idx, idx, idx};
|
||||
|
||||
// Search for row index
|
||||
IdArray ret_row = NewIdArray(idx->shape[0], ctx, nbits);
|
||||
const int nt2 = dgl::cuda::FindNumThreads(idx->shape[0]);
|
||||
const int nb2 = (idx->shape[0] + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
_SortedSearchKernel, nb2, nt2, 0, stream, indptr_data, csr.num_rows,
|
||||
idx.Ptr<IdType>(), idx->shape[0], ret_row.Ptr<IdType>());
|
||||
|
||||
// Column & data can be obtained by index select.
|
||||
IdArray ret_col = IndexSelect(csr.indices, idx);
|
||||
IdArray ret_data = CSRHasData(csr) ? IndexSelect(csr.data, idx) : idx;
|
||||
return {ret_row, ret_col, ret_data};
|
||||
}
|
||||
|
||||
template std::vector<NDArray> CSRGetDataAndIndices<kDGLCUDA, int32_t>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols);
|
||||
template std::vector<NDArray> CSRGetDataAndIndices<kDGLCUDA, int64_t>(
|
||||
CSRMatrix csr, NDArray rows, NDArray cols);
|
||||
|
||||
///////////////////////////// CSRSliceMatrix /////////////////////////////
|
||||
|
||||
int64_t _UpPower(int64_t numel) {
|
||||
uint64_t ret = 1 << static_cast<uint64_t>(std::log2(numel) + 1);
|
||||
return ret;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Thomas Wang's 32 bit Mix Function.
|
||||
* Source link: https://gist.github.com/badboy/6267743
|
||||
*/
|
||||
__device__ inline uint32_t _Hash32Shift(uint32_t key) {
|
||||
key = ~key + (key << 15);
|
||||
key = key ^ (key >> 12);
|
||||
key = key + (key << 2);
|
||||
key = key ^ (key >> 4);
|
||||
key = key * 2057;
|
||||
key = key ^ (key >> 16);
|
||||
return key;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Thomas Wang's 64 bit Mix Function.
|
||||
* Source link: https://gist.github.com/badboy/6267743
|
||||
*/
|
||||
__device__ inline uint64_t _Hash64Shift(uint64_t key) {
|
||||
key = (~key) + (key << 21);
|
||||
key = key ^ (key >> 24);
|
||||
key = (key + (key << 3)) + (key << 8);
|
||||
key = key ^ (key >> 14);
|
||||
key = (key + (key << 2)) + (key << 4);
|
||||
key = key ^ (key >> 28);
|
||||
key = key + (key << 31);
|
||||
return key;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief A hashmap designed for CSRSliceMatrix, similar in function to set. For
|
||||
* performance, it can only be created and called in the cuda kernel.
|
||||
*/
|
||||
template <typename IdType>
|
||||
struct NodeQueryHashmap {
|
||||
__device__ inline NodeQueryHashmap(IdType* Kptr, size_t numel)
|
||||
: kptr_(Kptr), capacity_(numel) {}
|
||||
|
||||
/**
|
||||
* @brief Insert a key. It must be called by cuda threads.
|
||||
*
|
||||
* @param key The key to be inserted.
|
||||
*/
|
||||
__device__ inline void Insert(IdType key) {
|
||||
uint32_t delta = 1;
|
||||
uint32_t pos = Hash(key);
|
||||
IdType prev = dgl::aten::cuda::AtomicCAS(&kptr_[pos], kEmptyKey_, key);
|
||||
while (prev != key && prev != kEmptyKey_) {
|
||||
pos = Hash(pos + delta);
|
||||
delta += 1;
|
||||
prev = dgl::aten::cuda::AtomicCAS(&kptr_[pos], kEmptyKey_, key);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Check whether a key exists within the hashtable. It must be called
|
||||
* by cuda threads.
|
||||
*
|
||||
* @param key The key to check for.
|
||||
* @return True if the key exists in the hashtable.
|
||||
*/
|
||||
__device__ inline bool Query(IdType key) {
|
||||
uint32_t delta = 1;
|
||||
uint32_t pos = Hash(key);
|
||||
while (true) {
|
||||
if (kptr_[pos] == key) return true;
|
||||
if (kptr_[pos] == kEmptyKey_) return false;
|
||||
pos = Hash(pos + delta);
|
||||
delta += 1;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
__device__ inline uint32_t Hash(int32_t key) {
|
||||
return _Hash32Shift(key) & (capacity_ - 1);
|
||||
}
|
||||
|
||||
__device__ inline uint32_t Hash(uint32_t key) {
|
||||
return _Hash32Shift(key) & (capacity_ - 1);
|
||||
}
|
||||
|
||||
__device__ inline uint32_t Hash(int64_t key) {
|
||||
return static_cast<uint32_t>(_Hash64Shift(key)) & (capacity_ - 1);
|
||||
}
|
||||
|
||||
__device__ inline uint32_t Hash(uint64_t key) {
|
||||
return static_cast<uint32_t>(_Hash64Shift(key)) & (capacity_ - 1);
|
||||
}
|
||||
|
||||
IdType kEmptyKey_{-1};
|
||||
IdType* kptr_;
|
||||
uint32_t capacity_{0};
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Generate a 0-1 mask for each index whose column is in the provided
|
||||
* hashmap. It also counts the number of masked values per row.
|
||||
*
|
||||
* @tparam IdType The ID type used for matrices.
|
||||
* @tparam WARP_SIZE The number of cuda threads in a cuda warp.
|
||||
* @tparam BLOCK_WARPS The number of warps in a cuda block.
|
||||
* @tparam TILE_SIZE The number of rows covered by each threadblock.
|
||||
*/
|
||||
template <typename IdType, int WARP_SIZE, int BLOCK_WARPS, int TILE_SIZE>
|
||||
__global__ void _SegmentMaskColKernel(
|
||||
const IdType* indptr, const IdType* indices, int64_t num_rows,
|
||||
IdType* hashmap_buffer, int64_t buffer_size, IdType* mask, IdType* count) {
|
||||
assert(blockDim.x == WARP_SIZE);
|
||||
assert(blockDim.y == BLOCK_WARPS);
|
||||
|
||||
int warp_id = threadIdx.y;
|
||||
int laneid = threadIdx.x;
|
||||
IdType out_row = blockIdx.x * TILE_SIZE + threadIdx.y;
|
||||
IdType last_row =
|
||||
min(static_cast<IdType>((blockIdx.x + 1) * TILE_SIZE),
|
||||
static_cast<IdType>(num_rows));
|
||||
|
||||
NodeQueryHashmap<IdType> hashmap(hashmap_buffer, buffer_size);
|
||||
typedef cub::WarpReduce<IdType> WarpReduce;
|
||||
__shared__ typename WarpReduce::TempStorage temp_storage[BLOCK_WARPS];
|
||||
|
||||
while (out_row < last_row) {
|
||||
IdType local_count = 0;
|
||||
IdType in_row_start = indptr[out_row];
|
||||
IdType in_row_end = indptr[out_row + 1];
|
||||
for (int idx = in_row_start + laneid; idx < in_row_end; idx += WARP_SIZE) {
|
||||
bool is_in = hashmap.Query(indices[idx]);
|
||||
if (is_in) {
|
||||
local_count += 1;
|
||||
mask[idx] = 1;
|
||||
}
|
||||
}
|
||||
IdType reduce_count = WarpReduce(temp_storage[warp_id]).Sum(local_count);
|
||||
if (laneid == 0) {
|
||||
count[out_row] = reduce_count;
|
||||
}
|
||||
out_row += BLOCK_WARPS;
|
||||
}
|
||||
}
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
CSRMatrix CSRSliceMatrix(
|
||||
CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols) {
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const auto& ctx = rows->ctx;
|
||||
const auto& dtype = rows->dtype;
|
||||
const auto nbits = dtype.bits;
|
||||
const int64_t new_nrows = rows->shape[0];
|
||||
const int64_t new_ncols = cols->shape[0];
|
||||
|
||||
if (new_nrows == 0 || new_ncols == 0)
|
||||
return CSRMatrix(
|
||||
new_nrows, new_ncols, Full(0, new_nrows + 1, nbits, ctx),
|
||||
NullArray(dtype, ctx), NullArray(dtype, ctx));
|
||||
|
||||
// First slice rows
|
||||
csr = CSRSliceRows(csr, rows);
|
||||
|
||||
if (csr.indices->shape[0] == 0)
|
||||
return CSRMatrix(
|
||||
new_nrows, new_ncols, Full(0, new_nrows + 1, nbits, ctx),
|
||||
NullArray(dtype, ctx), NullArray(dtype, ctx));
|
||||
|
||||
// Generate a 0-1 mask for matched (row, col) positions.
|
||||
IdArray mask = Full(0, csr.indices->shape[0], nbits, ctx);
|
||||
// A count for how many masked values per row.
|
||||
IdArray count = NewIdArray(csr.num_rows, ctx, nbits);
|
||||
CUDA_CALL(
|
||||
cudaMemset(count.Ptr<IdType>(), 0, sizeof(IdType) * (csr.num_rows)));
|
||||
|
||||
// Generate a NodeQueryHashmap buffer. The key of the hashmap is col.
|
||||
// For performance, the load factor of the hashmap is in (0.25, 0.5);
|
||||
// Because num_cols is usually less than 1 Million (on GPU), the
|
||||
// memory overhead is not significant (less than 31MB) at a low load factor.
|
||||
int64_t buffer_size = _UpPower(new_ncols) * 2;
|
||||
IdArray hashmap_buffer = Full(-1, buffer_size, nbits, ctx);
|
||||
|
||||
using it = thrust::counting_iterator<int64_t>;
|
||||
runtime::CUDAWorkspaceAllocator allocator(ctx);
|
||||
const auto exec_policy = thrust::cuda::par_nosync(allocator).on(stream);
|
||||
thrust::for_each(
|
||||
exec_policy, it(0), it(new_ncols),
|
||||
[key = cols.Ptr<IdType>(), buffer = hashmap_buffer.Ptr<IdType>(),
|
||||
buffer_size] __device__(int64_t i) {
|
||||
NodeQueryHashmap<IdType> hashmap(buffer, buffer_size);
|
||||
hashmap.Insert(key[i]);
|
||||
});
|
||||
|
||||
const IdType* indptr_data =
|
||||
static_cast<IdType*>(GetDevicePointer(csr.indptr));
|
||||
const IdType* indices_data =
|
||||
static_cast<IdType*>(GetDevicePointer(csr.indices));
|
||||
|
||||
// Execute SegmentMaskColKernel
|
||||
const int64_t num_rows = csr.num_rows;
|
||||
constexpr int WARP_SIZE = 32;
|
||||
// With a simple fine-tuning, TILE_SIZE=16 gives a good performance.
|
||||
constexpr int TILE_SIZE = 16;
|
||||
constexpr int BLOCK_WARPS = CUDA_MAX_NUM_THREADS / WARP_SIZE;
|
||||
IdType nb =
|
||||
dgl::cuda::FindNumBlocks<'x'>((num_rows + TILE_SIZE - 1) / TILE_SIZE);
|
||||
const dim3 nthrs(WARP_SIZE, BLOCK_WARPS);
|
||||
const dim3 nblks(nb);
|
||||
CUDA_KERNEL_CALL(
|
||||
(_SegmentMaskColKernel<IdType, WARP_SIZE, BLOCK_WARPS, TILE_SIZE>), nblks,
|
||||
nthrs, 0, stream, indptr_data, indices_data, num_rows,
|
||||
hashmap_buffer.Ptr<IdType>(), buffer_size, mask.Ptr<IdType>(),
|
||||
count.Ptr<IdType>());
|
||||
|
||||
IdArray idx = AsNumBits(NonZero(mask), nbits);
|
||||
if (idx->shape[0] == 0)
|
||||
return CSRMatrix(
|
||||
new_nrows, new_ncols, Full(0, new_nrows + 1, nbits, ctx),
|
||||
NullArray(dtype, ctx), NullArray(dtype, ctx));
|
||||
|
||||
// Indptr needs to be adjusted according to the new nnz per row.
|
||||
IdArray ret_indptr = CumSum(count, true);
|
||||
|
||||
// Column & data can be obtained by index select.
|
||||
IdArray ret_col = IndexSelect(csr.indices, idx);
|
||||
IdArray ret_data = CSRHasData(csr) ? IndexSelect(csr.data, idx) : idx;
|
||||
|
||||
// Relabel column
|
||||
IdArray col_hash = NewIdArray(csr.num_cols, ctx, nbits);
|
||||
Scatter_(cols, Range(0, cols->shape[0], nbits, ctx), col_hash);
|
||||
ret_col = IndexSelect(col_hash, ret_col);
|
||||
|
||||
return CSRMatrix(new_nrows, new_ncols, ret_indptr, ret_col, ret_data);
|
||||
}
|
||||
|
||||
template CSRMatrix CSRSliceMatrix<kDGLCUDA, int32_t>(
|
||||
CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols);
|
||||
template CSRMatrix CSRSliceMatrix<kDGLCUDA, int64_t>(
|
||||
CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,179 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/spmm.cu
|
||||
* @brief SPMM C APIs and definitions.
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include <cstdlib>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./functor.cuh"
|
||||
#include "./ge_spmm.cuh"
|
||||
#include "./spmm.cuh"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using namespace cuda;
|
||||
|
||||
namespace aten {
|
||||
|
||||
/**
|
||||
* @brief CUDA implementation of g-SpMM on Csr format.
|
||||
* @note use cusparse if the reduce operator is `sum` and there is
|
||||
* no broadcast, use dgl's kernel in other cases.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SpMMCsr(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux) {
|
||||
bool is_scalar_efeat = efeat.NumElements() == csr.indices->shape[0];
|
||||
bool use_efeat = op != "copy_lhs";
|
||||
bool use_deterministic_alg_only = false;
|
||||
if (NULL != std::getenv("USE_DETERMINISTIC_ALG"))
|
||||
use_deterministic_alg_only = true;
|
||||
|
||||
if (reduce == "sum") {
|
||||
bool more_nnz = (csr.indices->shape[0] > csr.num_rows * csr.num_cols);
|
||||
if (op == "copy_lhs" && cusparse_available<DType, IdType>(more_nnz)) {
|
||||
// cusparse
|
||||
int64_t x_length = 1;
|
||||
for (int i = 1; i < ufeat->ndim; ++i) x_length *= ufeat->shape[i];
|
||||
CusparseCsrmm2<DType, IdType>(
|
||||
ufeat->ctx, csr, static_cast<DType*>(ufeat->data), nullptr,
|
||||
static_cast<DType*>(out->data), x_length, use_deterministic_alg_only);
|
||||
} else if (
|
||||
op == "mul" && is_scalar_efeat &&
|
||||
cusparse_available<DType, IdType>(more_nnz)) {
|
||||
// cusparse
|
||||
int64_t x_length = 1;
|
||||
for (int i = 1; i < ufeat->ndim; ++i) x_length *= ufeat->shape[i];
|
||||
if (!IsNullArray(csr.data)) {
|
||||
efeat = IndexSelect(efeat, csr.data);
|
||||
}
|
||||
CusparseCsrmm2<DType, IdType>(
|
||||
ufeat->ctx, csr, static_cast<DType*>(ufeat->data),
|
||||
static_cast<DType*>(efeat->data), static_cast<DType*>(out->data),
|
||||
x_length, use_deterministic_alg_only);
|
||||
} else { // general kernel
|
||||
SWITCH_OP(op, Op, {
|
||||
cuda::SpMMCsr<IdType, DType, Op, cuda::reduce::Sum<IdType, DType> >(
|
||||
bcast, csr, ufeat, efeat, out, NullArray(), NullArray());
|
||||
});
|
||||
}
|
||||
} else if (reduce == "max") {
|
||||
SWITCH_OP(op, Op, {
|
||||
cuda::SpMMCsr<IdType, DType, Op, cuda::reduce::Max<IdType, DType> >(
|
||||
bcast, csr, ufeat, efeat, out, out_aux[0], out_aux[1]);
|
||||
});
|
||||
} else if (reduce == "min") {
|
||||
SWITCH_OP(op, Op, {
|
||||
cuda::SpMMCsr<IdType, DType, Op, cuda::reduce::Min<IdType, DType> >(
|
||||
bcast, csr, ufeat, efeat, out, out_aux[0], out_aux[1]);
|
||||
});
|
||||
} else {
|
||||
LOG(FATAL) << "Not implemented";
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA implementation of g-SpMM on Coo format.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SpMMCoo(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux) {
|
||||
if (reduce == "sum") {
|
||||
SWITCH_OP(op, Op, {
|
||||
cuda::SpMMCoo<IdType, DType, Op, cuda::reduce::Sum<IdType, DType, true> >(
|
||||
bcast, coo, ufeat, efeat, out, NullArray(), NullArray());
|
||||
});
|
||||
} else if (reduce == "max") {
|
||||
SWITCH_OP(op, Op, {
|
||||
cuda::SpMMCoo<IdType, DType, Op, cuda::reduce::Max<IdType, DType, true> >(
|
||||
bcast, coo, ufeat, efeat, out, out_aux[0], out_aux[1]);
|
||||
});
|
||||
} else if (reduce == "min") {
|
||||
SWITCH_OP(op, Op, {
|
||||
cuda::SpMMCoo<IdType, DType, Op, cuda::reduce::Min<IdType, DType, true> >(
|
||||
bcast, coo, ufeat, efeat, out, out_aux[0], out_aux[1]);
|
||||
});
|
||||
} else {
|
||||
LOG(FATAL) << "Not implemented";
|
||||
}
|
||||
}
|
||||
|
||||
template void SpMMCsr<kDGLCUDA, int32_t, __half>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCsr<kDGLCUDA, int64_t, __half>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
#if BF16_ENABLED
|
||||
template void SpMMCsr<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCsr<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
#endif // BF16_ENABLED
|
||||
template void SpMMCsr<kDGLCUDA, int32_t, float>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCsr<kDGLCUDA, int64_t, float>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCsr<kDGLCUDA, int32_t, double>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCsr<kDGLCUDA, int64_t, double>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
|
||||
template void SpMMCoo<kDGLCUDA, int32_t, __half>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCoo<kDGLCUDA, int64_t, __half>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
#if BF16_ENABLED
|
||||
template void SpMMCoo<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCoo<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
#endif // BF16_ENABLED
|
||||
template void SpMMCoo<kDGLCUDA, int32_t, float>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCoo<kDGLCUDA, int64_t, float>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCoo<kDGLCUDA, int32_t, double>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
template void SpMMCoo<kDGLCUDA, int64_t, double>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,802 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/spmm.cuh
|
||||
* @brief SPMM CUDA kernel function header.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CUDA_SPMM_CUH_
|
||||
#define DGL_ARRAY_CUDA_SPMM_CUH_
|
||||
|
||||
#include <dgl/bcast.h>
|
||||
|
||||
#include <limits>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./utils.h"
|
||||
#include "atomic.cuh"
|
||||
#include "bf16.cuh"
|
||||
#include "fp16.cuh"
|
||||
#include "macro.cuh"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using namespace cuda;
|
||||
|
||||
namespace aten {
|
||||
|
||||
/**
|
||||
* @brief Determine whether cusparse SpMM function is applicable.
|
||||
*/
|
||||
template <typename DType, typename IdType>
|
||||
inline bool cusparse_available(bool more_nnz_than_matrix_size) {
|
||||
#if CUDART_VERSION < 11000
|
||||
if (std::is_same<IdType, int>::value &&
|
||||
(std::is_same<DType, float>::value || std::is_same<DType, double>::value))
|
||||
return true;
|
||||
return false;
|
||||
#else
|
||||
if (std::is_same<DType, __half>::value ||
|
||||
std::is_same<DType, __nv_bfloat16>::value)
|
||||
return false; // cusparse's SpMM on fp16 is slow, temporally disabled.
|
||||
// If the CSR matrix has more NNZ than matrix size, we should not use
|
||||
// cuSPARSE 11.1.
|
||||
return !more_nnz_than_matrix_size;
|
||||
#endif
|
||||
}
|
||||
|
||||
namespace {
|
||||
|
||||
/** @brief Call cuBLAS geam API for transpose operation for float and double. */
|
||||
template <typename DType>
|
||||
cublasStatus_t Xgeam(
|
||||
cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
|
||||
int m, int n, const DType* alpha, const DType* A, int lda,
|
||||
const DType* beta, const DType* B, int ldb, DType* C, int ldc) {
|
||||
LOG(FATAL) << "Not supported dtype";
|
||||
return CUBLAS_STATUS_EXECUTION_FAILED;
|
||||
}
|
||||
|
||||
template <>
|
||||
cublasStatus_t Xgeam<__half>(
|
||||
cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
|
||||
int m, int n, const __half* alpha, const __half* A, int lda,
|
||||
const __half* beta, const __half* B, int ldb, __half* C, int ldc) {
|
||||
// TODO(ndickson): There is no cublasHgeam, so a different
|
||||
// implementation would be required.
|
||||
LOG(FATAL) << "Xgeam does not support dtype half (FP16)";
|
||||
return CUBLAS_STATUS_EXECUTION_FAILED;
|
||||
}
|
||||
|
||||
#if BF16_ENABLED
|
||||
template <>
|
||||
cublasStatus_t Xgeam<__nv_bfloat16>(
|
||||
cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
|
||||
int m, int n, const __nv_bfloat16* alpha, const __nv_bfloat16* A, int lda,
|
||||
const __nv_bfloat16* beta, const __nv_bfloat16* B, int ldb,
|
||||
__nv_bfloat16* C, int ldc) {
|
||||
// TODO(ndickson): There is no cublasHgeam, so a different
|
||||
// implementation would be required.
|
||||
LOG(FATAL) << "Xgeam does not support dtype bfloat16 (BF16)";
|
||||
return CUBLAS_STATUS_EXECUTION_FAILED;
|
||||
}
|
||||
#endif // BF16_ENABLED
|
||||
|
||||
template <>
|
||||
cublasStatus_t Xgeam<float>(
|
||||
cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
|
||||
int m, int n, const float* alpha, const float* A, int lda,
|
||||
const float* beta, const float* B, int ldb, float* C, int ldc) {
|
||||
return cublasSgeam(
|
||||
handle, transa, transb, m, n, alpha, A, lda, beta, B, ldb, C, ldc);
|
||||
}
|
||||
|
||||
template <>
|
||||
cublasStatus_t Xgeam<double>(
|
||||
cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
|
||||
int m, int n, const double* alpha, const double* A, int lda,
|
||||
const double* beta, const double* B, int ldb, double* C, int ldc) {
|
||||
return cublasDgeam(
|
||||
handle, transa, transb, m, n, alpha, A, lda, beta, B, ldb, C, ldc);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Transpose operator kernel implementation.
|
||||
* @note not efficient but it's not a bottleneck, used for float16 dtype.
|
||||
*/
|
||||
template <typename DType>
|
||||
__global__ void _TransposeKernel(
|
||||
const DType* __restrict__ in, DType* __restrict__ out, int n, int m) {
|
||||
int i = blockIdx.x;
|
||||
for (int j = threadIdx.x; j < m; j += blockDim.x)
|
||||
out[i * m + j] = in[j * n + i];
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Tranpose the input matrix.
|
||||
* @param row number of rows of input matrix.
|
||||
* @param col number of columns of input matrix.
|
||||
*/
|
||||
template <typename DType>
|
||||
void _Transpose(const DType* in, DType* out, int row, int col) {
|
||||
DType alpha = 1., beta = 0.;
|
||||
auto* thr_entry = runtime::CUDAThreadEntry::ThreadLocal();
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
if (!thr_entry->cublas_handle)
|
||||
CUBLAS_CALL(cublasCreate(&(thr_entry->cublas_handle)));
|
||||
CUBLAS_CALL(cublasSetStream(thr_entry->cublas_handle, stream));
|
||||
CUBLAS_CALL(Xgeam<DType>(
|
||||
thr_entry->cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, row, col, &alpha, in,
|
||||
col, &beta, nullptr, row, out, row));
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Tranpose the input matrix for data type half.
|
||||
* @note cuBLAS has no geam API for half data type, fallback to our kernel.
|
||||
*/
|
||||
template <>
|
||||
void _Transpose<__half>(const __half* in, __half* out, int row, int col) {
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
int nt = FindNumThreads(row);
|
||||
int nb = col;
|
||||
CUDA_KERNEL_CALL(_TransposeKernel, nb, nt, 0, stream, in, out, col, row);
|
||||
}
|
||||
|
||||
#if BF16_ENABLED
|
||||
/**
|
||||
* @brief Tranpose the input matrix for data type half.
|
||||
* @note cuBLAS has no geam API for bf16 data type, fallback to our kernel.
|
||||
*/
|
||||
template <>
|
||||
void _Transpose<__nv_bfloat16>(
|
||||
const __nv_bfloat16* in, __nv_bfloat16* out, int row, int col) {
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
int nt = FindNumThreads(row);
|
||||
int nb = col;
|
||||
CUDA_KERNEL_CALL(_TransposeKernel, nb, nt, 0, stream, in, out, col, row);
|
||||
}
|
||||
#endif // BF16_ENABLED
|
||||
|
||||
#if CUDART_VERSION < 11000
|
||||
template <typename DType>
|
||||
cusparseStatus_t Xcsrmm2(
|
||||
cusparseHandle_t handle, cusparseOperation_t transA,
|
||||
cusparseOperation_t transB, int m, int n, int k, int nnz,
|
||||
const DType* alpha, const cusparseMatDescr_t descrA, const DType* csrValA,
|
||||
const int* csrRowPtrA, const int* csrColIndA, const DType* B, int ldb,
|
||||
const DType* beta, DType* C, int ldc) {
|
||||
LOG(INFO) << "Not supported dtype";
|
||||
return CUSPARSE_STATUS_EXECUTION_FAILED;
|
||||
}
|
||||
|
||||
template <>
|
||||
cusparseStatus_t Xcsrmm2<float>(
|
||||
cusparseHandle_t handle, cusparseOperation_t transA,
|
||||
cusparseOperation_t transB, int m, int n, int k, int nnz,
|
||||
const float* alpha, const cusparseMatDescr_t descrA, const float* csrValA,
|
||||
const int* csrRowPtrA, const int* csrColIndA, const float* B, int ldb,
|
||||
const float* beta, float* C, int ldc) {
|
||||
return cusparseScsrmm2(
|
||||
handle, transA, transB, m, n, k, nnz, alpha, descrA, csrValA, csrRowPtrA,
|
||||
csrColIndA, B, ldb, beta, C, ldc);
|
||||
}
|
||||
|
||||
template <>
|
||||
cusparseStatus_t Xcsrmm2<double>(
|
||||
cusparseHandle_t handle, cusparseOperation_t transA,
|
||||
cusparseOperation_t transB, int m, int n, int k, int nnz,
|
||||
const double* alpha, const cusparseMatDescr_t descrA, const double* csrValA,
|
||||
const int* csrRowPtrA, const int* csrColIndA, const double* B, int ldb,
|
||||
const double* beta, double* C, int ldc) {
|
||||
return cusparseDcsrmm2(
|
||||
handle, transA, transB, m, n, k, nnz, alpha, descrA, csrValA, csrRowPtrA,
|
||||
csrColIndA, B, ldb, beta, C, ldc);
|
||||
}
|
||||
#endif
|
||||
|
||||
/** Cusparse implementation of SpMM on Csr format. */
|
||||
template <typename DType, typename IdType>
|
||||
void CusparseCsrmm2(
|
||||
const DGLContext& ctx, const CSRMatrix& csr, const DType* B_data,
|
||||
const DType* A_data, DType* C_data, int x_length,
|
||||
bool use_deterministic_alg_only = false) {
|
||||
// We use csrmm2 to perform following operation:
|
||||
// C = A x B, where A is a sparse matrix in csr format, B is the dense matrix
|
||||
// for node feature tensor. However, since cusparse only supports
|
||||
// column-major, while our tensor is stored in row-major, the actual
|
||||
// computation is: C = trans(A x trans(B)). Currently, we use cublasXgeam to
|
||||
// implement transposition and allocate intermediate workspace memory for
|
||||
// this.
|
||||
const int m = csr.num_rows;
|
||||
const int n = x_length;
|
||||
const int k = csr.num_cols;
|
||||
const int nnz = csr.indices->shape[0];
|
||||
const DType alpha = 1.0;
|
||||
const DType beta = 0.0;
|
||||
// device
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
auto* thr_entry = runtime::CUDAThreadEntry::ThreadLocal();
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
// allocate cusparse handle if needed
|
||||
if (!thr_entry->cusparse_handle) {
|
||||
CUSPARSE_CALL(cusparseCreate(&(thr_entry->cusparse_handle)));
|
||||
}
|
||||
CUSPARSE_CALL(cusparseSetStream(thr_entry->cusparse_handle, stream));
|
||||
// all one data array
|
||||
DType* valptr = nullptr;
|
||||
if (!A_data) {
|
||||
valptr =
|
||||
static_cast<DType*>(device->AllocWorkspace(ctx, nnz * sizeof(DType)));
|
||||
_Fill(valptr, nnz, static_cast<DType>(1.));
|
||||
}
|
||||
#if CUDART_VERSION >= 11000
|
||||
cusparseSpMatDescr_t matA;
|
||||
cusparseDnMatDescr_t matB, matC;
|
||||
constexpr auto dtype = cuda_dtype<DType>::value;
|
||||
constexpr auto idtype = cusparse_idtype<IdType>::value;
|
||||
CUSPARSE_CALL(cusparseCreateCsr(
|
||||
&matA, m, k, nnz, static_cast<IdType*>(csr.indptr->data),
|
||||
static_cast<IdType*>(csr.indices->data),
|
||||
const_cast<DType*>(valptr ? valptr : A_data), idtype, idtype,
|
||||
CUSPARSE_INDEX_BASE_ZERO, dtype));
|
||||
CUSPARSE_CALL(cusparseCreateDnMat(
|
||||
&matB, k, n, n, const_cast<DType*>(B_data), dtype, CUSPARSE_ORDER_ROW));
|
||||
CUSPARSE_CALL(
|
||||
cusparseCreateDnMat(&matC, m, n, n, C_data, dtype, CUSPARSE_ORDER_ROW));
|
||||
|
||||
auto transA = CUSPARSE_OPERATION_NON_TRANSPOSE;
|
||||
auto transB = CUSPARSE_OPERATION_NON_TRANSPOSE;
|
||||
size_t workspace_size;
|
||||
cusparseSpMMAlg_t spmm_alg = use_deterministic_alg_only
|
||||
? CUSPARSE_SPMM_CSR_ALG3
|
||||
: CUSPARSE_SPMM_CSR_ALG2;
|
||||
CUSPARSE_CALL(cusparseSpMM_bufferSize(
|
||||
thr_entry->cusparse_handle, transA, transB, &alpha, matA, matB, &beta,
|
||||
matC, dtype, spmm_alg, &workspace_size));
|
||||
void* workspace = device->AllocWorkspace(ctx, workspace_size);
|
||||
CUSPARSE_CALL(cusparseSpMM(
|
||||
thr_entry->cusparse_handle, transA, transB, &alpha, matA, matB, &beta,
|
||||
matC, dtype, spmm_alg, workspace));
|
||||
device->FreeWorkspace(ctx, workspace);
|
||||
|
||||
CUSPARSE_CALL(cusparseDestroySpMat(matA));
|
||||
CUSPARSE_CALL(cusparseDestroyDnMat(matB));
|
||||
CUSPARSE_CALL(cusparseDestroyDnMat(matC));
|
||||
#else
|
||||
// allocate matrix for temporary transposed output
|
||||
DType* trans_out =
|
||||
static_cast<DType*>(device->AllocWorkspace(ctx, m * n * sizeof(DType)));
|
||||
|
||||
cusparseMatDescr_t descr;
|
||||
CUSPARSE_CALL(cusparseCreateMatDescr(&descr));
|
||||
CUSPARSE_CALL(cusparseSetMatType(descr, CUSPARSE_MATRIX_TYPE_GENERAL));
|
||||
CUSPARSE_CALL(cusparseSetMatIndexBase(descr, CUSPARSE_INDEX_BASE_ZERO));
|
||||
CUSPARSE_CALL(Xcsrmm2<DType>(
|
||||
thr_entry->cusparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE,
|
||||
CUSPARSE_OPERATION_TRANSPOSE, m, n, k, nnz, &alpha, descr,
|
||||
(valptr) ? valptr : A_data, static_cast<int32_t*>(csr.indptr->data),
|
||||
static_cast<int32_t*>(csr.indices->data), B_data, n, &beta, trans_out,
|
||||
m));
|
||||
CUSPARSE_CALL(cusparseDestroyMatDescr(descr));
|
||||
// transpose the output matrix
|
||||
_Transpose(trans_out, C_data, n, m);
|
||||
device->FreeWorkspace(ctx, trans_out);
|
||||
#endif
|
||||
if (valptr) device->FreeWorkspace(ctx, valptr);
|
||||
}
|
||||
|
||||
/** Cusparse implementation of SpMM on Csr format. */
|
||||
template <typename DType, typename IdType>
|
||||
void CusparseCsrmm2Hetero(
|
||||
const DGLContext& ctx, const CSRMatrix& csr, const DType* B_data,
|
||||
const DType* A_data, DType* C_data, int64_t x_length, cudaStream_t strm_id,
|
||||
bool use_deterministic_alg_only = false) {
|
||||
// We use csrmm2 to perform following operation:
|
||||
// C = A x B, where A is a sparse matrix in csr format, B is the dense matrix
|
||||
// for node feature tensor. However, since cusparse only supports
|
||||
// column-major, while our tensor is stored in row-major, the actual
|
||||
// computation is: C = trans(A x trans(B)). Currently, we use cublasXgeam to
|
||||
// implement transposition and allocate intermediate workspace memory for
|
||||
// this.
|
||||
int int_maxlimit = std::numeric_limits<int>::max();
|
||||
CHECK_GE(int_maxlimit, (csr.num_rows));
|
||||
CHECK_GE(int_maxlimit, csr.num_cols);
|
||||
CHECK_GE(int_maxlimit, csr.indices->shape[0]);
|
||||
const int m = csr.num_rows;
|
||||
const int n = x_length;
|
||||
const int k = csr.num_cols;
|
||||
const int nnz = csr.indices->shape[0];
|
||||
const DType alpha = 1.0;
|
||||
const DType beta = 1.0;
|
||||
// device
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
auto* thr_entry = runtime::CUDAThreadEntry::ThreadLocal();
|
||||
// allocate cusparse handle if needed
|
||||
if (!thr_entry->cusparse_handle) {
|
||||
CUSPARSE_CALL(cusparseCreate(&(thr_entry->cusparse_handle)));
|
||||
}
|
||||
CUSPARSE_CALL(cusparseSetStream(thr_entry->cusparse_handle, strm_id));
|
||||
// all one data array
|
||||
DType* valptr = nullptr;
|
||||
if (!A_data) {
|
||||
valptr =
|
||||
static_cast<DType*>(device->AllocWorkspace(ctx, nnz * sizeof(DType)));
|
||||
_Fill(valptr, nnz, static_cast<DType>(1.));
|
||||
}
|
||||
#if CUDART_VERSION >= 11000
|
||||
cusparseSpMatDescr_t matA;
|
||||
cusparseDnMatDescr_t matB, matC;
|
||||
constexpr auto dtype = cuda_dtype<DType>::value;
|
||||
constexpr auto idtype = cusparse_idtype<IdType>::value;
|
||||
CUSPARSE_CALL(cusparseCreateCsr(
|
||||
&matA, m, k, nnz, static_cast<IdType*>(csr.indptr->data),
|
||||
static_cast<IdType*>(csr.indices->data),
|
||||
const_cast<DType*>(valptr ? valptr : A_data), idtype, idtype,
|
||||
CUSPARSE_INDEX_BASE_ZERO, dtype));
|
||||
CUSPARSE_CALL(cusparseCreateDnMat(
|
||||
&matB, k, n, n, const_cast<DType*>(B_data), dtype, CUSPARSE_ORDER_ROW));
|
||||
CUSPARSE_CALL(
|
||||
cusparseCreateDnMat(&matC, m, n, n, C_data, dtype, CUSPARSE_ORDER_ROW));
|
||||
|
||||
auto transA = CUSPARSE_OPERATION_NON_TRANSPOSE;
|
||||
auto transB = CUSPARSE_OPERATION_NON_TRANSPOSE;
|
||||
size_t workspace_size;
|
||||
cusparseSpMMAlg_t spmm_alg = use_deterministic_alg_only
|
||||
? CUSPARSE_SPMM_CSR_ALG3
|
||||
: CUSPARSE_SPMM_CSR_ALG2;
|
||||
CUSPARSE_CALL(cusparseSpMM_bufferSize(
|
||||
thr_entry->cusparse_handle, transA, transB, &alpha, matA, matB, &beta,
|
||||
matC, dtype, spmm_alg, &workspace_size));
|
||||
void* workspace = device->AllocWorkspace(ctx, workspace_size);
|
||||
CUSPARSE_CALL(cusparseSpMM(
|
||||
thr_entry->cusparse_handle, transA, transB, &alpha, matA, matB, &beta,
|
||||
matC, dtype, spmm_alg, workspace));
|
||||
device->FreeWorkspace(ctx, workspace);
|
||||
|
||||
CUSPARSE_CALL(cusparseDestroySpMat(matA));
|
||||
CUSPARSE_CALL(cusparseDestroyDnMat(matB));
|
||||
CUSPARSE_CALL(cusparseDestroyDnMat(matC));
|
||||
#else
|
||||
cusparseMatDescr_t descr;
|
||||
CUSPARSE_CALL(cusparseCreateMatDescr(&descr));
|
||||
CUSPARSE_CALL(cusparseSetMatType(descr, CUSPARSE_MATRIX_TYPE_GENERAL));
|
||||
CUSPARSE_CALL(cusparseSetMatIndexBase(descr, CUSPARSE_INDEX_BASE_ZERO));
|
||||
CHECK_EQ(sizeof(IdType), sizeof(int32_t));
|
||||
CUSPARSE_CALL(Xcsrmm2<DType>(
|
||||
thr_entry->cusparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE,
|
||||
CUSPARSE_OPERATION_TRANSPOSE, m, n, k, nnz, &alpha, descr,
|
||||
(valptr) ? valptr : A_data, static_cast<int32_t*>(csr.indptr->data),
|
||||
static_cast<int32_t*>(csr.indices->data), B_data, n, &beta, C_data, m));
|
||||
CUSPARSE_CALL(cusparseDestroyMatDescr(descr));
|
||||
#endif
|
||||
if (valptr) device->FreeWorkspace(ctx, valptr);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
#define SWITCH_OP(op, Op, ...) \
|
||||
do { \
|
||||
if ((op) == "add") { \
|
||||
typedef cuda::binary::Add<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "sub") { \
|
||||
typedef cuda::binary::Sub<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "mul") { \
|
||||
typedef cuda::binary::Mul<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "div") { \
|
||||
typedef cuda::binary::Div<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "copy_lhs") { \
|
||||
typedef cuda::binary::CopyLhs<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else if ((op) == "copy_rhs") { \
|
||||
typedef cuda::binary::CopyRhs<DType> Op; \
|
||||
{ __VA_ARGS__ } \
|
||||
} else { \
|
||||
LOG(FATAL) << "Unsupported SpMM binary operator: " << op; \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
namespace cuda {
|
||||
|
||||
/**
|
||||
* @brief CUDA kernel of g-SpMM on Coo format.
|
||||
* @note it uses edge parallel strategy, different threadblocks (on y-axis)
|
||||
* is responsible for the computation on different edges. Threadblocks
|
||||
* on the x-axis are responsible for the computation on different
|
||||
* positions in feature dimension. To avoid possible data hazards, it uses
|
||||
* atomic operators for reduction.
|
||||
*/
|
||||
template <
|
||||
typename Idx, typename DType, typename BinaryOp, typename ReduceOp,
|
||||
bool UseBcast = false, bool UseIdx = false>
|
||||
__global__ void SpMMCooKernel(
|
||||
const DType* __restrict__ ufeat, const DType* __restrict__ efeat,
|
||||
DType* __restrict__ out, Idx* __restrict__ arg_u, Idx* __restrict__ arg_e,
|
||||
const Idx* __restrict__ row, const Idx* __restrict__ col,
|
||||
const Idx* __restrict__ edge_map, int64_t N, int64_t M, int64_t E,
|
||||
const int64_t* __restrict__ ubcast_off,
|
||||
const int64_t* __restrict__ ebcast_off, int64_t ufeat_len,
|
||||
int64_t efeat_len, int64_t out_len) {
|
||||
// SPMM with COO.
|
||||
Idx ty = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
const Idx stride_y = blockDim.y * gridDim.y;
|
||||
while (ty < E) {
|
||||
const Idx src = _ldg(row + ty);
|
||||
const Idx dst = _ldg(col + ty);
|
||||
const Idx eid = UseIdx ? _ldg(edge_map + ty) : ty;
|
||||
int64_t tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int64_t stride_x = blockDim.x * gridDim.x;
|
||||
const DType* uoff = BinaryOp::use_lhs ? (ufeat + src * ufeat_len) : nullptr;
|
||||
const DType* eoff = BinaryOp::use_rhs ? (efeat + eid * efeat_len) : nullptr;
|
||||
DType* outoff = out + dst * out_len;
|
||||
while (tx < out_len) {
|
||||
const int64_t lhs_add = UseBcast ? ubcast_off[tx] : tx;
|
||||
const int64_t rhs_add = UseBcast ? ebcast_off[tx] : tx;
|
||||
DType val = BinaryOp::Call(uoff + lhs_add, eoff + rhs_add);
|
||||
Idx* arguoff = nullptr; // arguoff is not used in SpMMCoo.
|
||||
Idx* argeoff = nullptr; // argeoff is not used in SpMMCoo.
|
||||
ReduceOp::Call(outoff + tx, arguoff, argeoff, val, src, eid);
|
||||
tx += stride_x;
|
||||
}
|
||||
ty += stride_y;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA kernel to compute argu and arge in g-SpMM on Coo format.
|
||||
* @note it uses edge parallel strategy, different threadblocks (on y-axis)
|
||||
* is responsible for the computation on different edges. Threadblocks
|
||||
* on the x-axis are responsible for the computation on different
|
||||
* positions in feature dimension.
|
||||
*/
|
||||
template <
|
||||
typename Idx, typename DType, typename BinaryOp, typename ReduceOp,
|
||||
bool UseBcast = false, bool UseIdx = false>
|
||||
__global__ void ArgSpMMCooKernel(
|
||||
const DType* __restrict__ ufeat, const DType* __restrict__ efeat,
|
||||
DType* __restrict__ out, Idx* __restrict__ arg_u, Idx* __restrict__ arg_e,
|
||||
const Idx* __restrict__ row, const Idx* __restrict__ col,
|
||||
const Idx* __restrict__ edge_map, int64_t N, int64_t M, int64_t E,
|
||||
const int64_t* __restrict__ ubcast_off,
|
||||
const int64_t* __restrict__ ebcast_off, int64_t ufeat_len,
|
||||
int64_t efeat_len, int64_t out_len) {
|
||||
// SPMM with COO arg max/min.
|
||||
Idx ty = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
const Idx stride_y = blockDim.y * gridDim.y;
|
||||
while (ty < E) {
|
||||
const Idx src = _ldg(row + ty);
|
||||
const Idx dst = _ldg(col + ty);
|
||||
const Idx eid = UseIdx ? _ldg(edge_map + ty) : ty;
|
||||
int64_t tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int64_t stride_x = blockDim.x * gridDim.x;
|
||||
const DType* uoff = BinaryOp::use_lhs ? (ufeat + src * ufeat_len) : nullptr;
|
||||
const DType* eoff = BinaryOp::use_rhs ? (efeat + eid * efeat_len) : nullptr;
|
||||
const DType* outoff = out + dst * out_len;
|
||||
Idx* arguoff = BinaryOp::use_lhs ? (arg_u + dst * out_len) : nullptr;
|
||||
Idx* argeoff = BinaryOp::use_rhs ? (arg_e + dst * out_len) : nullptr;
|
||||
while (tx < out_len) {
|
||||
int64_t lhs_add = UseBcast ? ubcast_off[tx] : tx;
|
||||
int64_t rhs_add = UseBcast ? ebcast_off[tx] : tx;
|
||||
DType val = BinaryOp::Call(uoff + lhs_add, eoff + rhs_add);
|
||||
ReduceOp::CallArg(tx, arguoff, argeoff, val, outoff[tx], src, eid);
|
||||
tx += stride_x;
|
||||
}
|
||||
ty += stride_y;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA kernel of g-SpMM on Csr format.
|
||||
* @note it uses node parallel strategy, different threadblocks (on y-axis)
|
||||
* is responsible for the computation on different destination nodes.
|
||||
* Threadblocks on the x-axis are responsible for the computation on
|
||||
* different positions in feature dimension.
|
||||
*/
|
||||
template <
|
||||
typename Idx, typename DType, typename BinaryOp, typename ReduceOp,
|
||||
bool UseBcast = false, bool UseIdx = false>
|
||||
__global__ void SpMMCsrKernel(
|
||||
const DType* __restrict__ ufeat, const DType* __restrict__ efeat,
|
||||
DType* __restrict__ out, Idx* __restrict__ arg_u, Idx* __restrict__ arg_e,
|
||||
const Idx* __restrict__ indptr, const Idx* __restrict__ indices,
|
||||
const Idx* __restrict__ edge_map, int64_t num_rows, int64_t num_cols,
|
||||
const int64_t* __restrict__ ubcast_off,
|
||||
const int64_t* __restrict__ ebcast_off, int64_t ufeat_len,
|
||||
int64_t efeat_len, int64_t out_len) {
|
||||
// SPMM with CSR.
|
||||
int ty = blockIdx.x * blockDim.y + threadIdx.y;
|
||||
const Idx stride_y = blockDim.y * gridDim.x;
|
||||
const int stride_x = blockDim.x * gridDim.y;
|
||||
while (ty < num_rows) {
|
||||
int tx = blockIdx.y * blockDim.x + threadIdx.x;
|
||||
while (tx < out_len) {
|
||||
typename accum_dtype<DType>::type local_accum = ReduceOp::zero();
|
||||
Idx local_argu = 0, local_arge = 0;
|
||||
const int lhs_add = UseBcast ? ubcast_off[tx] : tx;
|
||||
const int rhs_add = UseBcast ? ebcast_off[tx] : tx;
|
||||
for (Idx i = indptr[ty]; i < indptr[ty + 1]; ++i) {
|
||||
const Idx eid = UseIdx ? _ldg(edge_map + i) : i;
|
||||
const Idx cid = _ldg(indices + i);
|
||||
const DType* uoff =
|
||||
BinaryOp::use_lhs ? (ufeat + cid * ufeat_len) : nullptr;
|
||||
const DType* eoff =
|
||||
BinaryOp::use_rhs ? (efeat + eid * efeat_len) : nullptr;
|
||||
DType out = BinaryOp::Call(uoff + lhs_add, eoff + rhs_add);
|
||||
ReduceOp::Call(&local_accum, &local_argu, &local_arge, out, cid, eid);
|
||||
}
|
||||
// The use of += is to compute cross-type reducing on heterogeneous graph
|
||||
// when reduce op is `sum`.
|
||||
// C = SpMM(SpA, B) + C
|
||||
// Separate kernel `SpMMCmpCsrHeteroKernel` is used for max- and
|
||||
// min-reducer. It does not affect the output on homogeneous graph as
|
||||
// `out` is initialized to zero.
|
||||
out[ty * out_len + tx] += static_cast<DType>(local_accum);
|
||||
if (ReduceOp::require_arg && BinaryOp::use_lhs)
|
||||
arg_u[ty * out_len + tx] = local_argu;
|
||||
if (ReduceOp::require_arg && BinaryOp::use_rhs)
|
||||
arg_e[ty * out_len + tx] = local_arge;
|
||||
tx += stride_x;
|
||||
}
|
||||
ty += stride_y;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA kernel of SpMM-Min/Max on Csr format.
|
||||
* @note it uses node parallel strategy, different threadblocks (on y-axis)
|
||||
* is responsible for the computation on different destination nodes.
|
||||
* Threadblocks on the x-axis are responsible for the computation on
|
||||
* different positions in feature dimension.
|
||||
*/
|
||||
template <
|
||||
typename Idx, typename DType, typename BinaryOp, typename ReduceOp,
|
||||
bool UseBcast = false, bool UseIdx = false>
|
||||
__global__ void SpMMCmpCsrHeteroKernel(
|
||||
const DType* __restrict__ ufeat, const DType* __restrict__ efeat,
|
||||
DType* __restrict__ out, Idx* __restrict__ arg_u, Idx* __restrict__ arg_e,
|
||||
Idx* __restrict__ arg_u_ntype, Idx* __restrict__ arg_e_etype,
|
||||
const Idx* __restrict__ indptr, const Idx* __restrict__ indices,
|
||||
const Idx* __restrict__ edge_map, int64_t num_rows, int64_t num_cols,
|
||||
const int64_t* __restrict__ ubcast_off,
|
||||
const int64_t* __restrict__ ebcast_off, int64_t ufeat_len,
|
||||
int64_t efeat_len, int64_t out_len, const int src_type, const int etype) {
|
||||
// SPMM with CSR.
|
||||
int ty = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
const Idx stride_y = blockDim.y * gridDim.y;
|
||||
const int stride_x = blockDim.x * gridDim.x;
|
||||
while (ty < num_rows) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
while (tx < out_len) {
|
||||
using accum_type = typename accum_dtype<DType>::type;
|
||||
accum_type local_accum =
|
||||
static_cast<accum_type>(out[ty * out_len + tx]); // ReduceOp::zero();
|
||||
Idx local_argu = 0, local_arge = 0;
|
||||
const int lhs_add = UseBcast ? ubcast_off[tx] : tx;
|
||||
const int rhs_add = UseBcast ? ebcast_off[tx] : tx;
|
||||
for (Idx i = indptr[ty]; i < indptr[ty + 1]; ++i) {
|
||||
const Idx eid = UseIdx ? _ldg(edge_map + i) : i;
|
||||
const Idx cid = _ldg(indices + i);
|
||||
const DType* uoff =
|
||||
BinaryOp::use_lhs ? (ufeat + cid * ufeat_len) : nullptr;
|
||||
const DType* eoff =
|
||||
BinaryOp::use_rhs ? (efeat + eid * efeat_len) : nullptr;
|
||||
DType tmp_out = BinaryOp::Call(uoff + lhs_add, eoff + rhs_add);
|
||||
ReduceOp::Call(
|
||||
&local_accum, &local_argu, &local_arge, tmp_out, cid, eid);
|
||||
}
|
||||
// Update output only when max/min values are different that original
|
||||
// output
|
||||
DType new_out = static_cast<DType>(local_accum);
|
||||
if (out[ty * out_len + tx] != new_out) {
|
||||
out[ty * out_len + tx] = new_out;
|
||||
if (ReduceOp::require_arg && BinaryOp::use_lhs) {
|
||||
arg_u[ty * out_len + tx] = local_argu;
|
||||
arg_u_ntype[ty * out_len + tx] = src_type;
|
||||
}
|
||||
if (ReduceOp::require_arg && BinaryOp::use_rhs) {
|
||||
arg_e[ty * out_len + tx] = local_arge;
|
||||
arg_e_etype[ty * out_len + tx] = etype;
|
||||
}
|
||||
}
|
||||
tx += stride_x;
|
||||
}
|
||||
ty += stride_y;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA implementation of g-SpMM on Coo format.
|
||||
* @param bcast Broadcast information.
|
||||
* @param coo The Coo matrix.
|
||||
* @param ufeat The feature on source nodes.
|
||||
* @param efeat The feature on edges.
|
||||
* @param out The result feature on destination nodes.
|
||||
* @param argu Arg-Min/Max on source nodes, which refers the source node indices
|
||||
* correspond to the minimum/maximum values of reduction result on
|
||||
* destination nodes. It's useful in computing gradients of Min/Max
|
||||
* reducer.
|
||||
* @param arge Arg-Min/Max on edges. which refers the source node indices
|
||||
* correspond to the minimum/maximum values of reduction result on
|
||||
* destination nodes. It's useful in computing gradients of Min/Max
|
||||
* reducer.
|
||||
*/
|
||||
template <typename Idx, typename DType, typename BinaryOp, typename ReduceOp>
|
||||
void SpMMCoo(
|
||||
const BcastOff& bcast, const COOMatrix& coo, NDArray ufeat, NDArray efeat,
|
||||
NDArray out, NDArray argu, NDArray arge) {
|
||||
/**
|
||||
* TODO(Xin): Disable half precision for SpMMCoo due to the round-off error.
|
||||
* We should use fp32 for the accumulation but it's hard to modify the
|
||||
* current implementation.
|
||||
*/
|
||||
#if BF16_ENABLED
|
||||
if (std::is_same<DType, __half>::value ||
|
||||
std::is_same<DType, __nv_bfloat16>::value)
|
||||
#else
|
||||
if (std::is_same<DType, __half>::value)
|
||||
#endif // BF16_ENABLED
|
||||
LOG(FATAL) << "SpMMCoo doesn't support half precision fow now. "
|
||||
<< "Please use SpMMCsr instead by allowing the graph "
|
||||
<< "materialize CSR/CSC formats.";
|
||||
const Idx *row = coo.row.Ptr<Idx>(), *col = coo.col.Ptr<Idx>(),
|
||||
*edge_map = coo.data.Ptr<Idx>();
|
||||
const DType *ufeat_data = ufeat.Ptr<DType>(),
|
||||
*efeat_data = efeat.Ptr<DType>();
|
||||
DType* out_data = out.Ptr<DType>();
|
||||
Idx *argu_data = argu.Ptr<Idx>(), *arge_data = arge.Ptr<Idx>();
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const int64_t N = coo.num_rows, M = coo.num_cols, E = coo.row->shape[0];
|
||||
|
||||
int64_t *ubcast_off = nullptr, *ebcast_off = nullptr;
|
||||
int64_t len = bcast.out_len, lhs_len = bcast.lhs_len, rhs_len = bcast.rhs_len;
|
||||
|
||||
int64_t out_size = out.NumElements();
|
||||
const int nt = FindNumThreads(out_size);
|
||||
const int nb = (out_size + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
_FillKernel, nb, nt, 0, stream, out_data, out_size, ReduceOp::zero());
|
||||
|
||||
const int ntx = FindNumThreads(len);
|
||||
const int nty = CUDA_MAX_NUM_THREADS / ntx;
|
||||
const int nbx = (len + ntx - 1) / ntx;
|
||||
const int nby = FindNumBlocks<'y'>((E + nty - 1) / nty);
|
||||
const dim3 nblks(nbx, nby);
|
||||
const dim3 nthrs(ntx, nty);
|
||||
const bool use_idx = !IsNullArray(coo.data);
|
||||
|
||||
BCAST_IDX_CTX_SWITCH(bcast, use_idx, ufeat->ctx, ubcast_off, ebcast_off, {
|
||||
CUDA_KERNEL_CALL(
|
||||
(SpMMCooKernel<Idx, DType, BinaryOp, ReduceOp, UseBcast, UseIdx>),
|
||||
nblks, nthrs, 0, stream, ufeat_data, efeat_data, out_data, argu_data,
|
||||
arge_data, row, col, edge_map, N, M, E, ubcast_off, ebcast_off, lhs_len,
|
||||
rhs_len, len);
|
||||
if (ReduceOp::require_arg) {
|
||||
CUDA_KERNEL_CALL(
|
||||
(ArgSpMMCooKernel<Idx, DType, BinaryOp, ReduceOp, UseBcast, UseIdx>),
|
||||
nblks, nthrs, 0, stream, ufeat_data, efeat_data, out_data, argu_data,
|
||||
arge_data, row, col, edge_map, N, M, E, ubcast_off, ebcast_off,
|
||||
lhs_len, rhs_len, len);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA implementation of g-SpMM on Csr format.
|
||||
* @param bcast Broadcast information.
|
||||
* @param csr The Csr matrix.
|
||||
* @param ufeat The feature on source nodes.
|
||||
* @param efeat The feature on edges.
|
||||
* @param out The result feature on destination nodes.
|
||||
* @param argu Arg-Min/Max on source nodes, which refers the source node indices
|
||||
* correspond to the minimum/maximum values of reduction result on
|
||||
* destination nodes. It's useful in computing gradients of Min/Max
|
||||
* reducer.
|
||||
* @param arge Arg-Min/Max on edges. which refers the source node indices
|
||||
* correspond to the minimum/maximum values of reduction result on
|
||||
* destination nodes. It's useful in computing gradients of Min/Max
|
||||
* reducer.
|
||||
*/
|
||||
template <typename Idx, typename DType, typename BinaryOp, typename ReduceOp>
|
||||
void SpMMCsr(
|
||||
const BcastOff& bcast, const CSRMatrix& csr, NDArray ufeat, NDArray efeat,
|
||||
NDArray out, NDArray argu, NDArray arge) {
|
||||
const Idx* indptr = csr.indptr.Ptr<Idx>();
|
||||
const Idx* indices = csr.indices.Ptr<Idx>();
|
||||
const Idx* edge_map = csr.data.Ptr<Idx>();
|
||||
const DType* ufeat_data = ufeat.Ptr<DType>();
|
||||
const DType* efeat_data = efeat.Ptr<DType>();
|
||||
DType* out_data = out.Ptr<DType>();
|
||||
Idx* argu_data = argu.Ptr<Idx>();
|
||||
Idx* arge_data = arge.Ptr<Idx>();
|
||||
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
|
||||
int64_t *ubcast_off = nullptr, *ebcast_off = nullptr;
|
||||
int64_t len = bcast.out_len, lhs_len = bcast.lhs_len, rhs_len = bcast.rhs_len;
|
||||
const int ntx = FindNumThreads(len);
|
||||
const int nty = CUDA_MAX_NUM_THREADS / ntx;
|
||||
const int nby = (len + ntx - 1) / ntx;
|
||||
const int nbx = FindNumBlocks<'x'>((csr.num_rows + nty - 1) / nty);
|
||||
const dim3 nblks(nbx, nby);
|
||||
const dim3 nthrs(ntx, nty);
|
||||
const bool use_idx = !IsNullArray(csr.data);
|
||||
|
||||
BCAST_IDX_CTX_SWITCH(
|
||||
bcast, use_idx, ufeat->ctx, ubcast_off, ebcast_off,
|
||||
{CUDA_KERNEL_CALL(
|
||||
(SpMMCsrKernel<Idx, DType, BinaryOp, ReduceOp, UseBcast, UseIdx>),
|
||||
nblks, nthrs, 0, stream, ufeat_data, efeat_data, out_data, argu_data,
|
||||
arge_data, indptr, indices, edge_map, csr.num_rows, csr.num_cols,
|
||||
ubcast_off, ebcast_off, lhs_len, rhs_len, len)});
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief CUDA kernel of SpMM-Min/Max on Csr format on heterogeneous graph.
|
||||
* @param bcast Broadcast information.
|
||||
* @param csr The Csr matrix.
|
||||
* @param ufeat The feature on source nodes.
|
||||
* @param efeat The feature on edges.
|
||||
* @param out The result feature on destination nodes.
|
||||
* @param argu Arg-Min/Max on source nodes, which refers the source node indices
|
||||
* correspond to the minimum/maximum values of reduction result on
|
||||
* destination nodes. It's useful in computing gradients of Min/Max
|
||||
* reducer.
|
||||
* @param arge Arg-Min/Max on edges. which refers the source node indices
|
||||
* correspond to the minimum/maximum values of reduction result on
|
||||
* destination nodes. It's useful in computing gradients of Min/Max
|
||||
* reducer.
|
||||
* @param argu_ntype Node type of the arg-Min/Max on source nodes, which refers
|
||||
* the source node types correspond to the minimum/maximum values of reduction
|
||||
* result on destination nodes. It's useful in computing gradients of Min/Max
|
||||
* reducer.
|
||||
* @param arge_etype Edge-type of the arg-Min/Max on edges. which refers the
|
||||
* source node indices correspond to the minimum/maximum values of reduction
|
||||
* result on destination nodes. It's useful in computing gradients of Min/Max
|
||||
* reducer.
|
||||
* @param src_type Node type of the source nodes of an etype
|
||||
* @param etype Edge type
|
||||
*/
|
||||
template <typename Idx, typename DType, typename BinaryOp, typename ReduceOp>
|
||||
void SpMMCmpCsrHetero(
|
||||
const BcastOff& bcast, const CSRMatrix& csr, NDArray ufeat, NDArray efeat,
|
||||
NDArray out, NDArray argu, NDArray arge, NDArray argu_ntype,
|
||||
NDArray arge_etype, const int src_type, const int etype) {
|
||||
const Idx* indptr = csr.indptr.Ptr<Idx>();
|
||||
const Idx* indices = csr.indices.Ptr<Idx>();
|
||||
const Idx* edge_map = csr.data.Ptr<Idx>();
|
||||
const DType* ufeat_data = ufeat.Ptr<DType>();
|
||||
const DType* efeat_data = efeat.Ptr<DType>();
|
||||
DType* out_data = out.Ptr<DType>();
|
||||
Idx* argu_data = argu.Ptr<Idx>();
|
||||
Idx* arge_data = arge.Ptr<Idx>();
|
||||
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
|
||||
int64_t *ubcast_off = nullptr, *ebcast_off = nullptr;
|
||||
int64_t len = bcast.out_len, lhs_len = bcast.lhs_len, rhs_len = bcast.rhs_len;
|
||||
const int ntx = FindNumThreads(len);
|
||||
const int nty = CUDA_MAX_NUM_THREADS / ntx;
|
||||
const int nbx = (len + ntx - 1) / ntx;
|
||||
const int nby = FindNumBlocks<'y'>((csr.num_rows + nty - 1) / nty);
|
||||
const dim3 nblks(nbx, nby);
|
||||
const dim3 nthrs(ntx, nty);
|
||||
const bool use_idx = !IsNullArray(csr.data);
|
||||
|
||||
BCAST_IDX_CTX_SWITCH(
|
||||
bcast, use_idx, ufeat->ctx, ubcast_off, ebcast_off,
|
||||
{CUDA_KERNEL_CALL(
|
||||
(SpMMCmpCsrHeteroKernel<
|
||||
Idx, DType, BinaryOp, ReduceOp, UseBcast, UseIdx>),
|
||||
nblks, nthrs, 0, stream, ufeat_data, efeat_data, out_data, argu_data,
|
||||
arge_data, static_cast<Idx*>(argu_ntype->data),
|
||||
static_cast<Idx*>(arge_etype->data), indptr, indices, edge_map,
|
||||
csr.num_rows, csr.num_cols, ubcast_off, ebcast_off, lhs_len, rhs_len,
|
||||
len, src_type, etype)});
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CUDA_SPMM_CUH_
|
||||
@@ -0,0 +1,262 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/spmm.cu
|
||||
* @brief SPMM C APIs and definitions.
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include <cstdlib>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./functor.cuh"
|
||||
#include "./ge_spmm.cuh"
|
||||
#include "./spmm.cuh"
|
||||
|
||||
namespace dgl {
|
||||
|
||||
using namespace cuda;
|
||||
|
||||
namespace aten {
|
||||
|
||||
/**
|
||||
* @brief CUDA implementation of g-SpMM on Csr format.
|
||||
* @note use cusparse if the reduce operator is `sum` and there is
|
||||
* no broadcast, use dgl's kernel in other cases.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SpMMCsrHetero(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr,
|
||||
const std::vector<NDArray>& vec_ufeat,
|
||||
const std::vector<NDArray>& vec_efeat, std::vector<NDArray>* vec_out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_ntids, // ufeat node type id
|
||||
const std::vector<dgl_type_t>& out_ntids) { // output node type id
|
||||
bool is_scalar_efeat =
|
||||
vec_efeat[0].NumElements() == vec_csr[0].indices->shape[0];
|
||||
bool use_efeat = op != "copy_lhs";
|
||||
auto device = runtime::DeviceAPI::Get(vec_csr[0].indptr->ctx);
|
||||
std::vector<DType*> trans_out((*vec_out).size(), NULL);
|
||||
bool use_deterministic_alg_only = false;
|
||||
if (NULL != std::getenv("USE_DETERMINISTIC_ALG"))
|
||||
use_deterministic_alg_only = true;
|
||||
|
||||
bool use_legacy_cusparsemm =
|
||||
(CUDART_VERSION < 11000) && (reduce == "sum") &&
|
||||
// legacy cuSPARSE does not care about NNZ, hence the argument "false".
|
||||
((op == "copy_lhs" && cusparse_available<DType, IdType>(false)) ||
|
||||
(op == "mul" && is_scalar_efeat &&
|
||||
cusparse_available<DType, IdType>(false)));
|
||||
// Create temporary output buffer to store non-transposed output
|
||||
if (use_legacy_cusparsemm) {
|
||||
for (dgl_type_t ntype = 0; ntype < (*vec_out).size(); ++ntype) {
|
||||
const int m = (*vec_out)[ntype]->shape[0];
|
||||
const int n = (*vec_out)[ntype]->shape[1];
|
||||
if (m == 0) continue;
|
||||
DType* out = static_cast<DType*>(device->AllocWorkspace(
|
||||
vec_csr[0].indptr->ctx, m * n * sizeof(DType)));
|
||||
CUDA_CALL(cudaMemset(out, 0, m * n * sizeof(DType)));
|
||||
trans_out[ntype] = out;
|
||||
}
|
||||
}
|
||||
// Check shape of ufeat for all relation type and compute feature size
|
||||
int64_t x_length = 1;
|
||||
for (dgl_type_t etype = 0; etype < (ufeat_ntids.size() - 1); ++etype) {
|
||||
NDArray ufeat = vec_ufeat[ufeat_ntids[etype]];
|
||||
NDArray next_ufeat = vec_ufeat[ufeat_ntids[etype + 1]];
|
||||
CHECK_EQ(ufeat->ndim, next_ufeat->ndim)
|
||||
<< "Input features have different shapes";
|
||||
for (int i = 1; i < ufeat->ndim; ++i) {
|
||||
if (ufeat->shape[i] != next_ufeat->shape[i]) {
|
||||
if (ufeat->shape[i] == 1 || next_ufeat->shape[i] == 1)
|
||||
LOG(FATAL) << "Homogenized message passing on heterogeneous graphs "
|
||||
"does not support "
|
||||
<< "automatic broadcasting. Please manually broadcast it "
|
||||
"before calling "
|
||||
<< "message passing functions.";
|
||||
else
|
||||
LOG(FATAL) << "Input features have different shapes.";
|
||||
return;
|
||||
}
|
||||
|
||||
if (etype == 0) x_length *= ufeat->shape[i];
|
||||
}
|
||||
}
|
||||
// TODO(Israt): Can python do the following initializations while creating the
|
||||
// tensors?
|
||||
if (reduce == "max" || reduce == "min") {
|
||||
const int64_t dim = bcast.out_len;
|
||||
std::vector<bool> updated((*vec_out).size(), false);
|
||||
for (dgl_type_t etype = 0; etype < ufeat_ntids.size(); ++etype) {
|
||||
DType* out_off = (*vec_out)[out_ntids[etype]].Ptr<DType>();
|
||||
if (reduce == "max")
|
||||
_Fill(
|
||||
out_off, vec_csr[etype].num_rows * dim,
|
||||
cuda::reduce::Max<IdType, DType>::zero());
|
||||
else // min
|
||||
_Fill(
|
||||
out_off, vec_csr[etype].num_rows * dim,
|
||||
cuda::reduce::Min<IdType, DType>::zero());
|
||||
const dgl_type_t dst_id = out_ntids[etype];
|
||||
if (!updated[dst_id]) {
|
||||
updated[dst_id] = true;
|
||||
if (op == "copy_lhs") {
|
||||
IdType* argu_ntype = (*out_aux)[2][dst_id].Ptr<IdType>();
|
||||
_Fill(
|
||||
argu_ntype, vec_csr[etype].num_rows * dim,
|
||||
static_cast<IdType>(-1));
|
||||
}
|
||||
if (op == "copy_rhs") {
|
||||
IdType* arge_etype = (*out_aux)[3][dst_id].Ptr<IdType>();
|
||||
_Fill(
|
||||
arge_etype, vec_csr[etype].num_rows * dim,
|
||||
static_cast<IdType>(-1));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
for (dgl_type_t etype = 0; etype < ufeat_ntids.size(); ++etype) {
|
||||
const dgl_type_t src_id = ufeat_ntids[etype];
|
||||
const dgl_type_t dst_id = out_ntids[etype];
|
||||
CSRMatrix csr = vec_csr[etype];
|
||||
if (reduce == "sum") {
|
||||
bool more_nnz = (csr.indices->shape[0] > csr.num_rows * csr.num_cols);
|
||||
/* Call SpMM for each relation type */
|
||||
if (op == "copy_lhs" &&
|
||||
cusparse_available<DType, IdType>(more_nnz)) { // cusparse
|
||||
/* If CUDA is less than 11.0, put the output in trans_out for later
|
||||
* transposition */
|
||||
DType* out = (CUDART_VERSION < 11000)
|
||||
? trans_out[dst_id]
|
||||
: static_cast<DType*>((*vec_out)[dst_id]->data);
|
||||
CusparseCsrmm2Hetero<DType, IdType>(
|
||||
csr.indptr->ctx, csr, static_cast<DType*>(vec_ufeat[src_id]->data),
|
||||
nullptr, out, x_length, stream, use_deterministic_alg_only);
|
||||
} else if (
|
||||
op == "mul" && is_scalar_efeat &&
|
||||
cusparse_available<DType, IdType>(more_nnz)) { // cusparse
|
||||
NDArray efeat = vec_efeat[etype];
|
||||
if (!IsNullArray(csr.data)) efeat = IndexSelect(efeat, csr.data);
|
||||
CusparseCsrmm2Hetero<DType, IdType>(
|
||||
csr.indptr->ctx, csr, static_cast<DType*>(vec_ufeat[src_id]->data),
|
||||
static_cast<DType*>(efeat->data),
|
||||
// TODO(Israt): Change (*vec_out) to trans_out to support CUDA
|
||||
// version < 11
|
||||
static_cast<DType*>((*vec_out)[dst_id]->data), x_length, stream,
|
||||
use_deterministic_alg_only);
|
||||
} else { // general kernel
|
||||
NDArray ufeat =
|
||||
(vec_ufeat.size() == 0) ? NullArray() : vec_ufeat[src_id];
|
||||
NDArray efeat =
|
||||
(vec_efeat.size() == 0) ? NullArray() : vec_efeat[etype];
|
||||
SWITCH_OP(op, Op, {
|
||||
cuda::SpMMCsr<IdType, DType, Op, cuda::reduce::Sum<IdType, DType>>(
|
||||
bcast, csr, ufeat, efeat, (*vec_out)[dst_id], NullArray(),
|
||||
NullArray());
|
||||
});
|
||||
}
|
||||
} else if (reduce == "max") {
|
||||
SWITCH_OP(op, Op, {
|
||||
NDArray ufeat =
|
||||
(vec_ufeat.size() == 0) ? NullArray() : vec_ufeat[src_id];
|
||||
NDArray efeat =
|
||||
(vec_efeat.size() == 0) ? NullArray() : vec_efeat[etype];
|
||||
cuda::SpMMCmpCsrHetero<
|
||||
IdType, DType, Op, cuda::reduce::Max<IdType, DType>>(
|
||||
bcast, csr, ufeat, efeat, (*vec_out)[dst_id], (*out_aux)[0][dst_id],
|
||||
(*out_aux)[1][dst_id], (*out_aux)[2][dst_id], (*out_aux)[3][dst_id],
|
||||
src_id, etype);
|
||||
});
|
||||
} else if (reduce == "min") {
|
||||
SWITCH_OP(op, Op, {
|
||||
NDArray ufeat =
|
||||
(vec_ufeat.size() == 0) ? NullArray() : vec_ufeat[src_id];
|
||||
NDArray efeat =
|
||||
(vec_efeat.size() == 0) ? NullArray() : vec_efeat[etype];
|
||||
cuda::SpMMCmpCsrHetero<
|
||||
IdType, DType, Op, cuda::reduce::Min<IdType, DType>>(
|
||||
bcast, csr, ufeat, efeat, (*vec_out)[dst_id], (*out_aux)[0][dst_id],
|
||||
(*out_aux)[1][dst_id], (*out_aux)[2][dst_id], (*out_aux)[3][dst_id],
|
||||
src_id, etype);
|
||||
});
|
||||
} else {
|
||||
LOG(FATAL) << "Not implemented";
|
||||
}
|
||||
}
|
||||
|
||||
if (use_legacy_cusparsemm) {
|
||||
// transpose output
|
||||
for (dgl_type_t ntype = 0; ntype < (*vec_out).size(); ++ntype) {
|
||||
const int m = (*vec_out)[ntype]->shape[0];
|
||||
const int n = (*vec_out)[ntype]->shape[1];
|
||||
if (m == 0) continue;
|
||||
DType* C_data = static_cast<DType*>((*vec_out)[ntype]->data);
|
||||
_Transpose(trans_out[ntype], C_data, n, m);
|
||||
device->FreeWorkspace(vec_csr[0].indptr->ctx, trans_out[ntype]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template void SpMMCsrHetero<kDGLCUDA, int32_t, __half>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
|
||||
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_ntids,
|
||||
const std::vector<dgl_type_t>& out_ntids);
|
||||
template void SpMMCsrHetero<kDGLCUDA, int64_t, __half>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
|
||||
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_ntids,
|
||||
const std::vector<dgl_type_t>& out_ntids);
|
||||
#if BF16_ENABLED
|
||||
template void SpMMCsrHetero<kDGLCUDA, int32_t, __nv_bfloat16>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
|
||||
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_ntids,
|
||||
const std::vector<dgl_type_t>& out_ntids);
|
||||
template void SpMMCsrHetero<kDGLCUDA, int64_t, __nv_bfloat16>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
|
||||
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_ntids,
|
||||
const std::vector<dgl_type_t>& out_ntids);
|
||||
#endif // BF16_ENABLED
|
||||
template void SpMMCsrHetero<kDGLCUDA, int32_t, float>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
|
||||
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_ntids,
|
||||
const std::vector<dgl_type_t>& out_ntids);
|
||||
template void SpMMCsrHetero<kDGLCUDA, int64_t, float>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
|
||||
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_ntids,
|
||||
const std::vector<dgl_type_t>& out_ntids);
|
||||
template void SpMMCsrHetero<kDGLCUDA, int32_t, double>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
|
||||
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_ntids,
|
||||
const std::vector<dgl_type_t>& out_ntids);
|
||||
template void SpMMCsrHetero<kDGLCUDA, int64_t, double>(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
|
||||
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_ntids,
|
||||
const std::vector<dgl_type_t>& out_ntids);
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,33 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/utils.cu
|
||||
* @brief Utilities for CUDA kernels.
|
||||
*/
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
#include "./utils.h"
|
||||
|
||||
namespace dgl {
|
||||
namespace cuda {
|
||||
|
||||
bool AllTrue(int8_t* flags, int64_t length, const DGLContext& ctx) {
|
||||
auto device = runtime::DeviceAPI::Get(ctx);
|
||||
int8_t* rst = static_cast<int8_t*>(device->AllocWorkspace(ctx, 1));
|
||||
// Call CUB's reduction
|
||||
size_t workspace_size = 0;
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
CUDA_CALL(cub::DeviceReduce::Min(
|
||||
nullptr, workspace_size, flags, rst, length, stream));
|
||||
void* workspace = device->AllocWorkspace(ctx, workspace_size);
|
||||
CUDA_CALL(cub::DeviceReduce::Min(
|
||||
workspace, workspace_size, flags, rst, length, stream));
|
||||
int8_t cpu_rst = GetCUDAScalar(device, ctx, rst);
|
||||
device->FreeWorkspace(ctx, workspace);
|
||||
device->FreeWorkspace(ctx, rst);
|
||||
return cpu_rst == 1;
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,301 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/cuda/utils.h
|
||||
* @brief Utilities for CUDA kernels.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_CUDA_UTILS_H_
|
||||
#define DGL_ARRAY_CUDA_UTILS_H_
|
||||
|
||||
#include <dgl/runtime/c_runtime_api.h>
|
||||
#include <dgl/runtime/device_api.h>
|
||||
#include <dgl/runtime/ndarray.h>
|
||||
#include <dmlc/logging.h>
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
#include <type_traits>
|
||||
|
||||
#include "../../runtime/cuda/cuda_common.h"
|
||||
|
||||
namespace dgl {
|
||||
namespace cuda {
|
||||
|
||||
#define CUDA_MAX_NUM_BLOCKS_X 0x7FFFFFFF
|
||||
#define CUDA_MAX_NUM_BLOCKS_Y 0xFFFF
|
||||
#define CUDA_MAX_NUM_BLOCKS_Z 0xFFFF
|
||||
// The max number of threads per block
|
||||
#define CUDA_MAX_NUM_THREADS 256
|
||||
|
||||
/** @brief Calculate the number of threads needed given the dimension length.
|
||||
*
|
||||
* It finds the biggest number that is smaller than min(dim, max_nthrs)
|
||||
* and is also power of two.
|
||||
*/
|
||||
inline int FindNumThreads(int dim, int max_nthrs = CUDA_MAX_NUM_THREADS) {
|
||||
CHECK_GE(dim, 0);
|
||||
if (dim == 0) return 1;
|
||||
int ret = max_nthrs;
|
||||
while (ret > dim) {
|
||||
ret = ret >> 1;
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
int _NumberOfBits(const T& range) {
|
||||
if (range <= 1) {
|
||||
// ranges of 0 or 1 require no bits to store
|
||||
return 0;
|
||||
}
|
||||
|
||||
int bits = 1;
|
||||
const auto urange = static_cast<std::make_unsigned_t<T>>(range);
|
||||
while (bits < static_cast<int>(sizeof(T) * 8) && (1ull << bits) < urange) {
|
||||
++bits;
|
||||
}
|
||||
|
||||
if (bits < static_cast<int>(sizeof(T) * 8)) {
|
||||
CHECK_EQ((range - 1) >> bits, 0);
|
||||
}
|
||||
CHECK_NE((range - 1) >> (bits - 1), 0);
|
||||
|
||||
return bits;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Find number of blocks is smaller than nblks and max_nblks
|
||||
* on the given axis ('x', 'y' or 'z').
|
||||
*/
|
||||
template <char axis>
|
||||
inline int FindNumBlocks(int nblks, int max_nblks = -1) {
|
||||
int default_max_nblks = -1;
|
||||
switch (axis) {
|
||||
case 'x':
|
||||
default_max_nblks = CUDA_MAX_NUM_BLOCKS_X;
|
||||
break;
|
||||
case 'y':
|
||||
default_max_nblks = CUDA_MAX_NUM_BLOCKS_Y;
|
||||
break;
|
||||
case 'z':
|
||||
default_max_nblks = CUDA_MAX_NUM_BLOCKS_Z;
|
||||
break;
|
||||
default:
|
||||
LOG(FATAL) << "Axis " << axis << " not recognized";
|
||||
break;
|
||||
}
|
||||
if (max_nblks == -1) max_nblks = default_max_nblks;
|
||||
CHECK_NE(nblks, 0);
|
||||
if (nblks < max_nblks) return nblks;
|
||||
return max_nblks;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ __forceinline__ T _ldg(T* addr) {
|
||||
#if __CUDA_ARCH__ >= 350
|
||||
return __ldg(addr);
|
||||
#else
|
||||
return *addr;
|
||||
#endif
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Return true if the given bool flag array is all true.
|
||||
* The input bool array is in int8_t type so it is aligned with byte address.
|
||||
*
|
||||
* @param flags The bool array.
|
||||
* @param length The length.
|
||||
* @param ctx Device context.
|
||||
* @return True if all the flags are true.
|
||||
*/
|
||||
bool AllTrue(int8_t* flags, int64_t length, const DGLContext& ctx);
|
||||
|
||||
/**
|
||||
* @brief CUDA Kernel of filling the vector started from ptr of size length
|
||||
* with val.
|
||||
* @note internal use only.
|
||||
*/
|
||||
template <typename DType>
|
||||
__global__ void _FillKernel(DType* ptr, size_t length, DType val) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < length) {
|
||||
ptr[tx] = val;
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
/** @brief Fill the vector started from ptr of size length with val */
|
||||
template <typename DType>
|
||||
void _Fill(DType* ptr, size_t length, DType val) {
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
int nt = FindNumThreads(length);
|
||||
int nb =
|
||||
(length + nt - 1) / nt; // on x-axis, no need to worry about upperbound.
|
||||
CUDA_KERNEL_CALL(cuda::_FillKernel, nb, nt, 0, stream, ptr, length, val);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Search adjacency list linearly for each (row, col) pair and
|
||||
* write the data under the matched position in the indices array to the output.
|
||||
*
|
||||
* If there is no match, the value in \c filler is written.
|
||||
* If there are multiple matches, only the first match is written.
|
||||
* If the given data array is null, write the matched position to the output.
|
||||
*/
|
||||
template <typename IdType, typename DType>
|
||||
__global__ void _LinearSearchKernel(
|
||||
const IdType* indptr, const IdType* indices, const IdType* data,
|
||||
const IdType* row, const IdType* col, int64_t row_stride,
|
||||
int64_t col_stride, int64_t length, const DType* weights, DType filler,
|
||||
DType* out) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < length) {
|
||||
int rpos = tx * row_stride, cpos = tx * col_stride;
|
||||
IdType v = -1;
|
||||
const IdType r = row[rpos], c = col[cpos];
|
||||
for (IdType i = indptr[r]; i < indptr[r + 1]; ++i) {
|
||||
if (indices[i] == c) {
|
||||
v = data ? data[i] : i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (v == -1) {
|
||||
out[tx] = filler;
|
||||
} else {
|
||||
// The casts here are to be able to handle DType being __half.
|
||||
// GCC treats int64_t as a distinct type from long long, so
|
||||
// without the explcit cast to long long, it errors out saying
|
||||
// that the implicit cast results in an ambiguous choice of
|
||||
// constructor for __half.
|
||||
// The using statement is to avoid a linter error about using
|
||||
// long or long long.
|
||||
using LongLong = long long; // NOLINT
|
||||
out[tx] = weights ? weights[v] : DType(LongLong(v));
|
||||
}
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
|
||||
#if BF16_ENABLED
|
||||
/**
|
||||
* @brief Specialization for bf16 because conversion from long long to bfloat16
|
||||
* doesn't exist before SM80.
|
||||
*/
|
||||
template <typename IdType>
|
||||
__global__ void _LinearSearchKernel(
|
||||
const IdType* indptr, const IdType* indices, const IdType* data,
|
||||
const IdType* row, const IdType* col, int64_t row_stride,
|
||||
int64_t col_stride, int64_t length, const __nv_bfloat16* weights,
|
||||
__nv_bfloat16 filler, __nv_bfloat16* out) {
|
||||
int tx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int stride_x = gridDim.x * blockDim.x;
|
||||
while (tx < length) {
|
||||
int rpos = tx * row_stride, cpos = tx * col_stride;
|
||||
IdType v = -1;
|
||||
const IdType r = row[rpos], c = col[cpos];
|
||||
for (IdType i = indptr[r]; i < indptr[r + 1]; ++i) {
|
||||
if (indices[i] == c) {
|
||||
v = data ? data[i] : i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (v == -1) {
|
||||
out[tx] = filler;
|
||||
} else {
|
||||
// If the result is saved in bf16, it should be fine to convert it to
|
||||
// float first
|
||||
out[tx] = weights ? weights[v] : __nv_bfloat16(static_cast<float>(v));
|
||||
}
|
||||
tx += stride_x;
|
||||
}
|
||||
}
|
||||
#endif // BF16_ENABLED
|
||||
|
||||
template <typename DType>
|
||||
inline DType GetCUDAScalar(
|
||||
runtime::DeviceAPI* device_api, DGLContext ctx, const DType* cuda_ptr) {
|
||||
DType result;
|
||||
device_api->CopyDataFromTo(
|
||||
cuda_ptr, 0, &result, 0, sizeof(result), ctx, DGLContext{kDGLCPU, 0},
|
||||
DGLDataTypeTraits<DType>::dtype);
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Given a sorted array and a value this function returns the index
|
||||
* of the first element which compares greater than value.
|
||||
*
|
||||
* This function assumes 0-based index
|
||||
* @param A: ascending sorted array
|
||||
* @param n: size of the A
|
||||
* @param x: value to search in A
|
||||
* @return index, i, of the first element st. A[i]>x. If x>=A[n-1] returns n.
|
||||
* if x<A[0] then it returns 0.
|
||||
*/
|
||||
template <typename IdType>
|
||||
__device__ IdType _UpperBound(const IdType* A, int64_t n, IdType x) {
|
||||
IdType l = 0, r = n, m = 0;
|
||||
while (l < r) {
|
||||
m = l + (r - l) / 2;
|
||||
if (x >= A[m]) {
|
||||
l = m + 1;
|
||||
} else {
|
||||
r = m;
|
||||
}
|
||||
}
|
||||
return l;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Given a sorted array and a value this function returns the index
|
||||
* of the element who is equal to val. If not exist returns n+1
|
||||
*
|
||||
* This function assumes 0-based index
|
||||
* @param A: ascending sorted array
|
||||
* @param n: size of the A
|
||||
* @param x: value to search in A
|
||||
* @return index, i, st. A[i]==x. If such an index not exists returns 'n'.
|
||||
*/
|
||||
template <typename IdType>
|
||||
__device__ IdType _BinarySearch(const IdType* A, int64_t n, IdType x) {
|
||||
IdType l = 0, r = n - 1, m = 0;
|
||||
while (l <= r) {
|
||||
m = l + (r - l) / 2;
|
||||
if (A[m] == x) {
|
||||
return m;
|
||||
}
|
||||
if (A[m] < x) {
|
||||
l = m + 1;
|
||||
} else {
|
||||
r = m - 1;
|
||||
}
|
||||
}
|
||||
return n; // not found
|
||||
}
|
||||
|
||||
template <typename DType, typename BoolType>
|
||||
void MaskSelect(
|
||||
runtime::DeviceAPI* device, const DGLContext& ctx, const DType* input,
|
||||
const BoolType* mask, DType* output, int64_t n, int64_t* rst,
|
||||
cudaStream_t stream) {
|
||||
size_t workspace_size = 0;
|
||||
CUDA_CALL(cub::DeviceSelect::Flagged(
|
||||
nullptr, workspace_size, input, mask, output, rst, n, stream));
|
||||
void* workspace = device->AllocWorkspace(ctx, workspace_size);
|
||||
CUDA_CALL(cub::DeviceSelect::Flagged(
|
||||
workspace, workspace_size, input, mask, output, rst, n, stream));
|
||||
device->FreeWorkspace(ctx, workspace);
|
||||
}
|
||||
|
||||
inline void* GetDevicePointer(runtime::NDArray array) {
|
||||
void* ptr = array->data;
|
||||
if (array.IsPinned()) {
|
||||
CUDA_CALL(cudaHostGetDevicePointer(&ptr, ptr, 0));
|
||||
}
|
||||
return ptr;
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CUDA_UTILS_H_
|
||||
@@ -0,0 +1,131 @@
|
||||
/**
|
||||
* Copyright (c) 2019-2022 by Contributors
|
||||
* @file array/cuda/uvm/array_index_select_uvm.cu
|
||||
* @brief Array index select GPU implementation
|
||||
*/
|
||||
#include <dgl/array.h>
|
||||
|
||||
#include "../../../runtime/cuda/cuda_common.h"
|
||||
#include "../array_index_select.cuh"
|
||||
#include "../utils.h"
|
||||
#include "./array_index_select_uvm.cuh"
|
||||
|
||||
namespace dgl {
|
||||
using runtime::NDArray;
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
template <typename DType, typename IdType>
|
||||
NDArray IndexSelectCPUFromGPU(NDArray array, IdArray index) {
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const int64_t arr_len = array->shape[0];
|
||||
const int64_t len = index->shape[0];
|
||||
int64_t num_feat = 1;
|
||||
std::vector<int64_t> shape{len};
|
||||
|
||||
CHECK(array.IsPinned());
|
||||
const DType* array_data = static_cast<DType*>(cuda::GetDevicePointer(array));
|
||||
CHECK_EQ(index->ctx.device_type, kDGLCUDA);
|
||||
|
||||
for (int d = 1; d < array->ndim; ++d) {
|
||||
num_feat *= array->shape[d];
|
||||
shape.emplace_back(array->shape[d]);
|
||||
}
|
||||
|
||||
NDArray ret = NDArray::Empty(shape, array->dtype, index->ctx);
|
||||
if (len == 0 || arr_len * num_feat == 0) return ret;
|
||||
DType* ret_data = static_cast<DType*>(ret->data);
|
||||
|
||||
auto res = Sort(index, cuda::_NumberOfBits(arr_len));
|
||||
const IdType* idx_data = static_cast<IdType*>(res.first->data);
|
||||
const int64_t* perm_data = static_cast<int64_t*>(res.second->data);
|
||||
|
||||
if (num_feat == 1) {
|
||||
const int nt = cuda::FindNumThreads(len);
|
||||
const int nb = (len + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
IndexSelectSingleKernel, nb, nt, 0, stream, array_data, idx_data, len,
|
||||
arr_len, ret_data, perm_data);
|
||||
} else {
|
||||
dim3 block(256, 1);
|
||||
while (static_cast<int64_t>(block.x) >= 2 * num_feat) {
|
||||
block.x /= 2;
|
||||
block.y *= 2;
|
||||
}
|
||||
const dim3 grid((len + block.y - 1) / block.y);
|
||||
if (num_feat * sizeof(DType) < 2 * CACHE_LINE_SIZE) {
|
||||
CUDA_KERNEL_CALL(
|
||||
IndexSelectMultiKernel, grid, block, 0, stream, array_data, num_feat,
|
||||
idx_data, len, arr_len, ret_data, perm_data);
|
||||
} else {
|
||||
CUDA_KERNEL_CALL(
|
||||
IndexSelectMultiKernelAligned, grid, block, 0, stream, array_data,
|
||||
num_feat, idx_data, len, arr_len, ret_data, perm_data);
|
||||
}
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
// floating point types are treated as their equal width integer types
|
||||
template NDArray IndexSelectCPUFromGPU<int8_t, int32_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelectCPUFromGPU<int8_t, int64_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelectCPUFromGPU<int16_t, int32_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelectCPUFromGPU<int16_t, int64_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelectCPUFromGPU<int32_t, int32_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelectCPUFromGPU<int32_t, int64_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelectCPUFromGPU<int64_t, int32_t>(NDArray, IdArray);
|
||||
template NDArray IndexSelectCPUFromGPU<int64_t, int64_t>(NDArray, IdArray);
|
||||
|
||||
template <typename DType, typename IdType>
|
||||
void IndexScatterGPUToCPU(NDArray dest, IdArray index, NDArray source) {
|
||||
cudaStream_t stream = runtime::getCurrentCUDAStream();
|
||||
const DType* source_data = static_cast<DType*>(source->data);
|
||||
const IdType* idx_data = static_cast<IdType*>(index->data);
|
||||
const int64_t arr_len = dest->shape[0];
|
||||
const int64_t len = index->shape[0];
|
||||
int64_t num_feat = 1;
|
||||
std::vector<int64_t> shape{len};
|
||||
|
||||
CHECK(dest.IsPinned());
|
||||
DType* dest_data = static_cast<DType*>(cuda::GetDevicePointer(dest));
|
||||
CHECK_EQ(index->ctx.device_type, kDGLCUDA);
|
||||
CHECK_EQ(source->ctx.device_type, kDGLCUDA);
|
||||
|
||||
for (int d = 1; d < source->ndim; ++d) {
|
||||
num_feat *= source->shape[d];
|
||||
}
|
||||
|
||||
if (len == 0) return;
|
||||
|
||||
if (num_feat == 1) {
|
||||
const int nt = cuda::FindNumThreads(len);
|
||||
const int nb = (len + nt - 1) / nt;
|
||||
CUDA_KERNEL_CALL(
|
||||
IndexScatterSingleKernel, nb, nt, 0, stream, source_data, idx_data, len,
|
||||
arr_len, dest_data);
|
||||
} else {
|
||||
dim3 block(256, 1);
|
||||
while (static_cast<int64_t>(block.x) >= 2 * num_feat) {
|
||||
block.x /= 2;
|
||||
block.y *= 2;
|
||||
}
|
||||
const dim3 grid((len + block.y - 1) / block.y);
|
||||
CUDA_KERNEL_CALL(
|
||||
IndexScatterMultiKernel, grid, block, 0, stream, source_data, num_feat,
|
||||
idx_data, len, arr_len, dest_data);
|
||||
}
|
||||
}
|
||||
|
||||
// floating point types are treated as their equal width integer types
|
||||
template void IndexScatterGPUToCPU<int8_t, int32_t>(NDArray, IdArray, NDArray);
|
||||
template void IndexScatterGPUToCPU<int8_t, int64_t>(NDArray, IdArray, NDArray);
|
||||
template void IndexScatterGPUToCPU<int16_t, int32_t>(NDArray, IdArray, NDArray);
|
||||
template void IndexScatterGPUToCPU<int16_t, int64_t>(NDArray, IdArray, NDArray);
|
||||
template void IndexScatterGPUToCPU<int32_t, int32_t>(NDArray, IdArray, NDArray);
|
||||
template void IndexScatterGPUToCPU<int32_t, int64_t>(NDArray, IdArray, NDArray);
|
||||
template void IndexScatterGPUToCPU<int64_t, int32_t>(NDArray, IdArray, NDArray);
|
||||
template void IndexScatterGPUToCPU<int64_t, int64_t>(NDArray, IdArray, NDArray);
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,52 @@
|
||||
/**
|
||||
* Copyright (c) 2021 by Contributors
|
||||
* @file array/cpu/array_index_select_uvm.cuh
|
||||
* @brief Array index select GPU kernel implementation
|
||||
*/
|
||||
|
||||
#ifndef DGL_ARRAY_CUDA_UVM_ARRAY_INDEX_SELECT_UVM_CUH_
|
||||
#define DGL_ARRAY_CUDA_UVM_ARRAY_INDEX_SELECT_UVM_CUH_
|
||||
|
||||
#define CACHE_LINE_SIZE 128
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace impl {
|
||||
|
||||
/**
|
||||
* This is a cross-device access version of IndexSelectMultiKernel.
|
||||
* Since the memory access over PCIe is more sensitive to the
|
||||
* data access aligment (cacheline), we need a separate version here.
|
||||
*/
|
||||
template <typename DType, typename IdType>
|
||||
__global__ void IndexSelectMultiKernelAligned(
|
||||
const DType* const array, const int64_t num_feat, const IdType* const index,
|
||||
const int64_t length, const int64_t arr_len, DType* const out,
|
||||
const int64_t* perm = nullptr) {
|
||||
int64_t out_row_index = blockIdx.x * blockDim.y + threadIdx.y;
|
||||
|
||||
const int64_t stride = blockDim.y * gridDim.x;
|
||||
|
||||
while (out_row_index < length) {
|
||||
int64_t col = threadIdx.x;
|
||||
const int64_t in_row = index[out_row_index];
|
||||
assert(in_row >= 0 && in_row < arr_len);
|
||||
const int64_t idx_offset =
|
||||
((uint64_t)(&array[in_row * num_feat]) % CACHE_LINE_SIZE) /
|
||||
sizeof(DType);
|
||||
col = col - idx_offset;
|
||||
const auto out_row = perm ? perm[out_row_index] : out_row_index;
|
||||
while (col < num_feat) {
|
||||
if (col >= 0)
|
||||
out[out_row * num_feat + col] = array[in_row * num_feat + col];
|
||||
col += blockDim.x;
|
||||
}
|
||||
out_row_index += stride;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace impl
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_CUDA_UVM_ARRAY_INDEX_SELECT_UVM_CUH_
|
||||
@@ -0,0 +1,54 @@
|
||||
/**
|
||||
* Copyright (c) 2021 by Contributors
|
||||
* @file array/filter.cc
|
||||
* @brief Object for selecting items in a set, or selecting items not in a set.
|
||||
*/
|
||||
|
||||
#include "./filter.h"
|
||||
|
||||
#include <dgl/packed_func_ext.h>
|
||||
#include <dgl/runtime/packed_func.h>
|
||||
#include <dgl/runtime/registry.h>
|
||||
|
||||
namespace dgl {
|
||||
namespace array {
|
||||
|
||||
using namespace dgl::runtime;
|
||||
|
||||
template <DGLDeviceType XPU, typename IdType>
|
||||
FilterRef CreateSetFilter(IdArray set);
|
||||
|
||||
DGL_REGISTER_GLOBAL("utils.filter._CAPI_DGLFilterCreateFromSet")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
IdArray array = args[0];
|
||||
auto ctx = array->ctx;
|
||||
// TODO(nv-dlasalle): Implement CPU version.
|
||||
if (ctx.device_type == kDGLCUDA) {
|
||||
#ifdef DGL_USE_CUDA
|
||||
ATEN_ID_TYPE_SWITCH(array->dtype, IdType, {
|
||||
*rv = CreateSetFilter<kDGLCUDA, IdType>(array);
|
||||
});
|
||||
#else
|
||||
LOG(FATAL) << "GPU support not compiled.";
|
||||
#endif
|
||||
} else {
|
||||
LOG(FATAL) << "CPU support not yet implemented.";
|
||||
}
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("utils.filter._CAPI_DGLFilterFindIncludedIndices")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
FilterRef filter = args[0];
|
||||
IdArray array = args[1];
|
||||
*rv = filter->find_included_indices(array);
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("utils.filter._CAPI_DGLFilterFindExcludedIndices")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
FilterRef filter = args[0];
|
||||
IdArray array = args[1];
|
||||
*rv = filter->find_excluded_indices(array);
|
||||
});
|
||||
|
||||
} // namespace array
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,49 @@
|
||||
/**
|
||||
* Copyright (c) 2021 by Contributors
|
||||
* @file array/filter.h
|
||||
* @brief Object for selecting items in a set, or selecting items not in a set.
|
||||
*/
|
||||
|
||||
#ifndef DGL_ARRAY_FILTER_H_
|
||||
#define DGL_ARRAY_FILTER_H_
|
||||
|
||||
#include <dgl/array.h>
|
||||
#include <dgl/runtime/object.h>
|
||||
|
||||
namespace dgl {
|
||||
namespace array {
|
||||
|
||||
class Filter : public runtime::Object {
|
||||
public:
|
||||
static constexpr const char* _type_key = "array.Filter";
|
||||
DGL_DECLARE_OBJECT_TYPE_INFO(Filter, Object);
|
||||
|
||||
/**
|
||||
* @brief From the test set of items, get the index of those which are
|
||||
* included by this filter.
|
||||
*
|
||||
* @param test The set of items to check for.
|
||||
*
|
||||
* @return The indices of the items from `test` that are selected by
|
||||
* this filter.
|
||||
*/
|
||||
virtual IdArray find_included_indices(IdArray test) = 0;
|
||||
|
||||
/**
|
||||
* @brief From the test set of items, get the indices of those which are
|
||||
* excluded by this filter.
|
||||
*
|
||||
* @param test The set of items to check for.
|
||||
*
|
||||
* @return The indices of the items from `test` that are not selected by this
|
||||
* filter.
|
||||
*/
|
||||
virtual IdArray find_excluded_indices(IdArray test) = 0;
|
||||
};
|
||||
|
||||
DGL_DEFINE_OBJECT_REF(FilterRef, Filter);
|
||||
|
||||
} // namespace array
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_FILTER_H_
|
||||
@@ -0,0 +1,803 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/kernel.cc
|
||||
* @brief New kernels
|
||||
*/
|
||||
#include <dgl/base_heterograph.h>
|
||||
#include <dgl/packed_func_ext.h>
|
||||
|
||||
#include "../c_api_common.h"
|
||||
#include "./check.h"
|
||||
#include "kernel_decl.h"
|
||||
|
||||
using namespace dgl::runtime;
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
namespace {} // namespace
|
||||
|
||||
/** @brief Generalized Sparse Matrix-Matrix Multiplication. */
|
||||
void SpMM(
|
||||
const std::string& op, const std::string& reduce, HeteroGraphPtr graph,
|
||||
NDArray ufeat, NDArray efeat, NDArray out, std::vector<NDArray> out_aux) {
|
||||
// TODO(zihao): format tuning
|
||||
SparseFormat format = graph->SelectFormat(0, CSC_CODE);
|
||||
const auto& bcast = CalcBcastOff(op, ufeat, efeat);
|
||||
|
||||
ATEN_XPU_SWITCH_CUDA(graph->Context().device_type, XPU, "SpMM", {
|
||||
ATEN_ID_TYPE_SWITCH(graph->DataType(), IdType, {
|
||||
ATEN_FLOAT_TYPE_SWITCH_16BITS(out->dtype, Dtype, XPU, "Feature data", {
|
||||
if (format == SparseFormat::kCSC) {
|
||||
SpMMCsr<XPU, IdType, Dtype>(
|
||||
op, reduce, bcast, graph->GetCSCMatrix(0), ufeat, efeat, out,
|
||||
out_aux);
|
||||
} else if (format == SparseFormat::kCOO) {
|
||||
SpMMCoo<XPU, IdType, Dtype>(
|
||||
op, reduce, bcast, graph->GetCOOMatrix(0), ufeat, efeat, out,
|
||||
out_aux);
|
||||
} else {
|
||||
LOG(FATAL) << "SpMM only supports CSC and COO formats";
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** @brief Generalized segmented dense Matrix-Matrix Multiplication. */
|
||||
void SegmentMM(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
|
||||
bool A_trans, bool B_trans) {
|
||||
CHECK_EQ(A->ndim, 2) << "segment_mm expects a 2D tensor for the first input.";
|
||||
CHECK_EQ(B->ndim, 3)
|
||||
<< "segment_mm expects a 3D tensor for the second input.";
|
||||
CHECK(!A_trans);
|
||||
if (B_trans) {
|
||||
CHECK_EQ(A->shape[1], B->shape[2])
|
||||
<< "segment_mm expects A.shape[1] == B.shape[2] when B_trans=True";
|
||||
} else {
|
||||
CHECK_EQ(A->shape[1], B->shape[1])
|
||||
<< "segment_mm expects A.shape[1] == B.shape[1]";
|
||||
}
|
||||
CHECK_EQ(B->shape[0], seglen_A.NumElements())
|
||||
<< "segment_mm expects len(seglen_A) == B.shape[0]";
|
||||
CHECK_EQ(seglen_A->ctx.device_type, kDGLCPU)
|
||||
<< "segment_mm expects seglen_A to be on CPU.";
|
||||
CHECK(A->ctx == B->ctx)
|
||||
<< "segment_mm expects A and B to be of the same device";
|
||||
ATEN_XPU_SWITCH_CUDA(A->ctx.device_type, XPU, "SegmentMM", {
|
||||
ATEN_ID_TYPE_SWITCH(seglen_A->dtype, IdType, {
|
||||
ATEN_FLOAT_TYPE_SWITCH_16BITS(A->dtype, Dtype, XPU, "Feature data", {
|
||||
SegmentMM<XPU, IdType, Dtype>(A, B, C, seglen_A, A_trans, B_trans);
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
void SegmentMMBackwardB(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen) {
|
||||
CHECK_EQ(A->ndim, 2) << "segment_mm_backward operator expects a 2D tensor "
|
||||
"for the first input.";
|
||||
CHECK_EQ(dC->ndim, 2) << "segment_mm_backward operator expects a 2D tensor "
|
||||
"for the second input.";
|
||||
CHECK_EQ(seglen->ctx.device_type, kDGLCPU)
|
||||
<< "segment_mm expects seglen to be on CPU.";
|
||||
ATEN_XPU_SWITCH_CUDA(A->ctx.device_type, XPU, "SegmentMMBackwardB", {
|
||||
ATEN_ID_TYPE_SWITCH(seglen->dtype, IdType, {
|
||||
ATEN_FLOAT_TYPE_SWITCH_16BITS(A->dtype, Dtype, XPU, "Feature data", {
|
||||
SegmentMMBackwardB<XPU, IdType, Dtype>(A, dC, dB, seglen);
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** @brief Generalized Dense Matrix-Matrix Multiplication according to relation
|
||||
* types. */
|
||||
void GatherMM(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b) {
|
||||
CHECK_EQ(A->ndim, 2)
|
||||
<< "gather_mm operator expects a 2D tensor for the first input.";
|
||||
CHECK_EQ(B->ndim, 3)
|
||||
<< "gather_mm operator expects a 3D tensor for the second input.";
|
||||
CHECK(A->ctx == B->ctx)
|
||||
<< "gather_mm expects all arguments to be on the same device.";
|
||||
if (aten::IsNullArray(idx_a)) {
|
||||
CHECK_EQ(A->shape[0], idx_b->shape[0])
|
||||
<< "gather_mm expects len(idx_b) == A.shape[0] when idx_a is None.";
|
||||
CHECK(A->ctx == idx_b->ctx)
|
||||
<< "gather_mm expects all arguments to be on the same device.";
|
||||
} else if (aten::IsNullArray(idx_b)) {
|
||||
CHECK_EQ(B->shape[0], idx_a->shape[0])
|
||||
<< "gather_mm expects len(idx_a) == B.shape[0] when idx_b is None.";
|
||||
CHECK(A->ctx == idx_a->ctx)
|
||||
<< "gather_mm expects all arguments to be on the same device.";
|
||||
} else {
|
||||
CHECK_EQ(idx_a->shape[0], idx_b->shape[0])
|
||||
<< "gather_mm expects len(idx_a) == len(idx_b) when both idx_a and "
|
||||
"idx_b are given.";
|
||||
CHECK(A->ctx == idx_a->ctx && A->ctx == idx_b->ctx)
|
||||
<< "gather_mm expects all arguments to be on the same device.";
|
||||
}
|
||||
const auto idtype = aten::IsNullArray(idx_a) ? idx_b->dtype : idx_a->dtype;
|
||||
ATEN_XPU_SWITCH_CUDA(A->ctx.device_type, XPU, "GatherMM", {
|
||||
ATEN_ID_TYPE_SWITCH(idtype, IdType, {
|
||||
ATEN_FLOAT_TYPE_SWITCH_16BITS(A->dtype, Dtype, XPU, "Feature data", {
|
||||
GatherMM<XPU, IdType, Dtype>(A, B, C, idx_a, idx_b);
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** @brief Generalized Dense Matrix-Matrix Multiplication according to relation
|
||||
* types. */
|
||||
void GatherMMScatter(
|
||||
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c) {
|
||||
CHECK_EQ(A->ndim, 2)
|
||||
<< "gather_mm_scatter expects a 2D tensor for the first input.";
|
||||
CHECK(A->ctx == B->ctx)
|
||||
<< "gather_mm_scatter expects all arguments to be on the same device.";
|
||||
if (!aten::IsNullArray(idx_c))
|
||||
CHECK(A->ctx == idx_c->ctx)
|
||||
<< "gather_mm_scatter expects all arguments to be on the same device.";
|
||||
if (aten::IsNullArray(idx_a) && !aten::IsNullArray(idx_b)) {
|
||||
CHECK_EQ(A->shape[0], idx_b->shape[0])
|
||||
<< "gather_mm_scatter expects len(idx_b) == A.shape[0] when idx_a is "
|
||||
"None.";
|
||||
CHECK(A->ctx == idx_b->ctx)
|
||||
<< "gather_mm_scatter expects all arguments to be on the same device.";
|
||||
} else if (aten::IsNullArray(idx_b) && !aten::IsNullArray(idx_a)) {
|
||||
CHECK_EQ(B->shape[0], idx_a->shape[0])
|
||||
<< "gather_mm_scatter expects len(idx_a) == B.shape[0] when idx_b is "
|
||||
"None.";
|
||||
CHECK(A->ctx == idx_a->ctx)
|
||||
<< "gather_mm_scatter expects all arguments to be on the same device.";
|
||||
} else if (!aten::IsNullArray(idx_b) && !aten::IsNullArray(idx_a)) {
|
||||
CHECK_EQ(idx_a->shape[0], idx_b->shape[0])
|
||||
<< "gather_mm_scatter expects len(idx_a) == len(idx_b) "
|
||||
<< "when both idx_a and idx_b are given.";
|
||||
CHECK(A->ctx == idx_a->ctx && A->ctx == idx_b->ctx)
|
||||
<< "gather_mm_scatter expects all arguments to be on the same device.";
|
||||
}
|
||||
ATEN_XPU_SWITCH_CUDA(A->ctx.device_type, XPU, "GatherMM", {
|
||||
ATEN_ID_TYPE_SWITCH(idx_c->dtype, IdType, {
|
||||
ATEN_FLOAT_TYPE_SWITCH_16BITS(A->dtype, Dtype, XPU, "Feature data", {
|
||||
GatherMMScatter<XPU, IdType, Dtype>(A, B, C, idx_a, idx_b, idx_c);
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** @brief Generalized Sparse Matrix-Matrix Multiplication with hetero-graph
|
||||
* support. */
|
||||
void SpMMHetero(
|
||||
const std::string& op, const std::string& reduce, HeteroGraphPtr graph,
|
||||
const std::vector<NDArray>& ufeat_vec,
|
||||
const std::vector<NDArray>& efeat_vec, std::vector<NDArray>* out,
|
||||
std::vector<std::vector<NDArray>>* out_aux) {
|
||||
SparseFormat format = graph->SelectFormat(0, CSC_CODE);
|
||||
|
||||
std::vector<CSRMatrix> vec_graph;
|
||||
std::vector<dgl_type_t> ufeat_eid;
|
||||
std::vector<dgl_type_t> efeat_eid;
|
||||
std::vector<dgl_type_t> out_eid;
|
||||
auto pair = graph->meta_graph()->FindEdge(0); // first etype
|
||||
NDArray ufeat_etype0 =
|
||||
(ufeat_vec.size() == 0) ? NullArray() : ufeat_vec[pair.first];
|
||||
NDArray efeat_etype0 = (efeat_vec.size() == 0) ? NullArray() : efeat_vec[0];
|
||||
for (dgl_type_t etype = 0; etype < graph->NumEdgeTypes(); ++etype) {
|
||||
vec_graph.push_back(graph->GetCSCMatrix(etype));
|
||||
auto pair = graph->meta_graph()->FindEdge(etype);
|
||||
ufeat_eid.push_back(pair.first);
|
||||
efeat_eid.push_back(etype);
|
||||
out_eid.push_back(pair.second);
|
||||
if (ufeat_etype0->shape[1] != ufeat_vec[pair.first]->shape[1])
|
||||
LOG(FATAL) << "Column width of the input node features of all etypes "
|
||||
"must be same.";
|
||||
if (efeat_etype0->shape[1] != efeat_vec[etype]->shape[1])
|
||||
LOG(FATAL) << "Column width of the input edge features of all etypes "
|
||||
"must be same.";
|
||||
}
|
||||
const auto& bcast = CalcBcastOff(op, ufeat_etype0, efeat_etype0);
|
||||
|
||||
ATEN_XPU_SWITCH_CUDA(graph->Context().device_type, XPU, "SpMM", {
|
||||
ATEN_ID_TYPE_SWITCH(
|
||||
graph->DataType(), IdType, {
|
||||
ATEN_FLOAT_TYPE_SWITCH_16BITS(
|
||||
(*out)[out_eid[0]]->dtype, Dtype, XPU, "Feature data", {
|
||||
if (format == SparseFormat::kCSC) {
|
||||
SpMMCsrHetero<XPU, IdType, Dtype>(
|
||||
op, reduce, bcast, vec_graph, ufeat_vec, efeat_vec, out,
|
||||
out_aux, ufeat_eid, out_eid);
|
||||
} else {
|
||||
// TODO(Israt): Add support for COO format
|
||||
LOG(FATAL)
|
||||
<< "SpMM only supports CSC format for graphs with number "
|
||||
<< "of relation types > 1";
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** @brief Generalized Sampled Dense-Dense Matrix Multiplication. */
|
||||
void SDDMM(
|
||||
const std::string& op, HeteroGraphPtr graph, NDArray lhs, NDArray rhs,
|
||||
NDArray out, int lhs_target, int rhs_target) {
|
||||
// TODO(zihao): format tuning
|
||||
SparseFormat format = graph->SelectFormat(0, COO_CODE);
|
||||
const auto& bcast = CalcBcastOff(op, lhs, rhs);
|
||||
|
||||
ATEN_XPU_SWITCH_CUDA(graph->Context().device_type, XPU, "SDDMM", {
|
||||
ATEN_ID_TYPE_SWITCH(graph->DataType(), IdType, {
|
||||
ATEN_FLOAT_TYPE_SWITCH_16BITS(out->dtype, Dtype, XPU, "Feature data", {
|
||||
if (format == SparseFormat::kCSR) {
|
||||
SDDMMCsr<XPU, IdType, Dtype>(
|
||||
op, bcast, graph->GetCSRMatrix(0), lhs, rhs, out, lhs_target,
|
||||
rhs_target);
|
||||
} else if (format == SparseFormat::kCOO) {
|
||||
SDDMMCoo<XPU, IdType, Dtype>(
|
||||
op, bcast, graph->GetCOOMatrix(0), lhs, rhs, out, lhs_target,
|
||||
rhs_target);
|
||||
} else {
|
||||
LOG(FATAL) << "SDDMM only supports CSR and COO formats";
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Find the src/dst/etype id based on the target 'u', 'v' or 'e'.
|
||||
*
|
||||
* @param graph The input graph.
|
||||
* @param target 'u', 'v' or 'e'. The target of the lhs or rhs data of an etype.
|
||||
* @param etype Relation type of the input graph.
|
||||
*/
|
||||
int get_typeid_by_target(HeteroGraphPtr graph, int target, dgl_type_t etype) {
|
||||
auto pair = graph->meta_graph()->FindEdge(etype);
|
||||
if (target == 0) return pair.first;
|
||||
if (target == 2) return pair.second;
|
||||
return etype;
|
||||
}
|
||||
|
||||
/** @brief Generalized Sampled Dense-Dense Matrix Multiplication. */
|
||||
void SDDMMHetero(
|
||||
const std::string& op, HeteroGraphPtr graph, std::vector<NDArray> lhs,
|
||||
std::vector<NDArray> rhs, std::vector<NDArray> out, int lhs_target,
|
||||
int rhs_target) {
|
||||
SparseFormat format = graph->SelectFormat(0, COO_CODE);
|
||||
|
||||
std::vector<dgl_type_t> lhs_eid;
|
||||
std::vector<dgl_type_t> rhs_eid;
|
||||
for (dgl_type_t etype = 0; etype < graph->NumEdgeTypes(); ++etype) {
|
||||
lhs_eid.push_back(get_typeid_by_target(graph, lhs_target, etype));
|
||||
rhs_eid.push_back(get_typeid_by_target(graph, rhs_target, etype));
|
||||
}
|
||||
const auto& bcast = CalcBcastOff(op, lhs[lhs_eid[0]], rhs[rhs_eid[0]]);
|
||||
|
||||
ATEN_XPU_SWITCH_CUDA(graph->Context().device_type, XPU, "SDDMM", {
|
||||
ATEN_ID_TYPE_SWITCH(graph->DataType(), IdType, {
|
||||
ATEN_FLOAT_TYPE_SWITCH_16BITS(
|
||||
out[rhs_eid[0]]->dtype, Dtype, XPU, "Feature data", {
|
||||
if (format == SparseFormat::kCSR) {
|
||||
std::vector<CSRMatrix> vec_csr;
|
||||
for (dgl_type_t etype = 0; etype < graph->NumEdgeTypes();
|
||||
++etype) {
|
||||
vec_csr.push_back(graph->GetCSRMatrix(etype));
|
||||
}
|
||||
SDDMMCsrHetero<XPU, IdType, Dtype>(
|
||||
op, bcast, vec_csr, lhs, rhs, out, lhs_target, rhs_target,
|
||||
lhs_eid, rhs_eid);
|
||||
} else if (format == SparseFormat::kCOO) {
|
||||
std::vector<COOMatrix> vec_coo;
|
||||
for (dgl_type_t etype = 0; etype < graph->NumEdgeTypes();
|
||||
++etype) {
|
||||
vec_coo.push_back(graph->GetCOOMatrix(etype));
|
||||
}
|
||||
SDDMMCooHetero<XPU, IdType, Dtype>(
|
||||
op, bcast, vec_coo, lhs, rhs, out, lhs_target, rhs_target,
|
||||
lhs_eid, rhs_eid);
|
||||
} else {
|
||||
LOG(FATAL) << "SDDMM only supports CSR and COO formats";
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** @brief Generalized Edge_softmax op for forward */
|
||||
void Edge_softmax_forward(
|
||||
const std::string& op, HeteroGraphPtr graph, NDArray ufeat, NDArray efeat,
|
||||
NDArray out) {
|
||||
// TODO(zhejiang): add gpu op for edge_softmax
|
||||
const auto& bcast = CalcBcastOff(op, ufeat, efeat);
|
||||
|
||||
ATEN_XPU_SWITCH(graph->Context().device_type, XPU, "edge_softmax", {
|
||||
ATEN_ID_TYPE_SWITCH(graph->DataType(), IdType, {
|
||||
ATEN_FLOAT_TYPE_SWITCH_16BITS(
|
||||
out->dtype, Dtype, XPU, "edge_softmax out data", {
|
||||
Edge_softmax_csr_forward<XPU, IdType, Dtype>(
|
||||
op, bcast, graph->GetCSCMatrix(0), ufeat, efeat, out);
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** @brief Generalized Edge_softmax op for backward */
|
||||
void Edge_softmax_backward(
|
||||
const std::string& op, HeteroGraphPtr graph, NDArray out, NDArray sds,
|
||||
NDArray back_out, NDArray ufeat) {
|
||||
// TODO(zhejiang): add gpu op for edge_softmax
|
||||
const auto& bcast = CalcBcastOff(op, ufeat, sds);
|
||||
|
||||
ATEN_XPU_SWITCH(graph->Context().device_type, XPU, "edge_softmax_back", {
|
||||
ATEN_ID_TYPE_SWITCH(graph->DataType(), IdType, {
|
||||
ATEN_FLOAT_TYPE_SWITCH_16BITS(
|
||||
out->dtype, Dtype, XPU, "edge_softmax out data_back", {
|
||||
Edge_softmax_csr_backward<XPU, IdType, Dtype>(
|
||||
op, bcast, graph->GetCSCMatrix(0), out, sds, back_out);
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
NDArray GetEdgeMapping(HeteroGraphRef graph) {
|
||||
SparseFormat format = graph->SelectFormat(0, CSC_CODE);
|
||||
if (format == SparseFormat::kCSC) {
|
||||
return graph.sptr()->GetCSCMatrix(0).data;
|
||||
} else {
|
||||
return NullArray();
|
||||
}
|
||||
}
|
||||
|
||||
/** @brief Segment reduce dispatch function. */
|
||||
void SegmentReduceDispatch(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg) {
|
||||
ATEN_XPU_SWITCH_CUDA(feat->ctx.device_type, XPU, "SegmentReduce", {
|
||||
ATEN_ID_TYPE_SWITCH(offsets->dtype, IdType, {
|
||||
ATEN_FLOAT_TYPE_SWITCH_16BITS(feat->dtype, Dtype, XPU, "Feature data", {
|
||||
SegmentReduce<XPU, IdType, Dtype>(op, feat, offsets, out, arg);
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** @brief Scatter Add (on first dimension) dispatch function. */
|
||||
void ScatterAddDispatch(NDArray feat, NDArray idx, NDArray out) {
|
||||
ATEN_XPU_SWITCH_CUDA(feat->ctx.device_type, XPU, "ScatterAdd", {
|
||||
ATEN_ID_TYPE_SWITCH(idx->dtype, IdType, {
|
||||
ATEN_FLOAT_TYPE_SWITCH_16BITS(feat->dtype, Dtype, XPU, "Feature data", {
|
||||
ScatterAdd<XPU, IdType, Dtype>(feat, idx, out);
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** @brief Update gradients (reduce op max/min) dispatch function on
|
||||
* heterogeneous graph. */
|
||||
void UpdateGradMinMaxDispatchHetero(
|
||||
const HeteroGraphPtr& graph, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out) {
|
||||
auto pair = graph->meta_graph()->FindEdge(0); // checking the first etype
|
||||
auto src_id = pair.first;
|
||||
ATEN_XPU_SWITCH_CUDA(feat[src_id]->ctx.device_type, XPU, "ScatterAdd", {
|
||||
ATEN_ID_TYPE_SWITCH(idx[src_id]->dtype, IdType, {
|
||||
ATEN_FLOAT_TYPE_SWITCH_16BITS(
|
||||
feat[src_id]->dtype, Dtype, XPU, "Feature data", {
|
||||
UpdateGradMinMax_hetero<XPU, IdType, Dtype>(
|
||||
graph, op, feat, idx, idx_etype, out);
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** @brief Backward segment cmp dispatch function.*/
|
||||
void BackwardSegmentCmpDispatch(NDArray feat, NDArray arg, NDArray out) {
|
||||
ATEN_XPU_SWITCH_CUDA(feat->ctx.device_type, XPU, "BackwardSegmentCmp", {
|
||||
ATEN_ID_TYPE_SWITCH(arg->dtype, IdType, {
|
||||
ATEN_FLOAT_TYPE_SWITCH_16BITS(feat->dtype, Dtype, XPU, "Feature data", {
|
||||
BackwardSegmentCmp<XPU, IdType, Dtype>(feat, arg, out);
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
std::pair<CSRMatrix, NDArray> CSRMM(
|
||||
CSRMatrix A, NDArray A_weights, CSRMatrix B, NDArray B_weights) {
|
||||
CHECK_EQ(A.num_cols, B.num_rows)
|
||||
<< "The number of nodes of destination node type of the first graph must "
|
||||
"be the "
|
||||
"same as the number of nodes of source node type of the second graph.";
|
||||
CheckCtx(
|
||||
A.indptr->ctx, {A_weights, B_weights},
|
||||
{"A's edge weights", "B's edge weights"});
|
||||
CHECK_EQ(A.indptr->ctx, B.indptr->ctx) << "Device of two graphs must match.";
|
||||
CHECK_EQ(A.indptr->dtype, B.indptr->dtype)
|
||||
<< "ID types of two graphs must match.";
|
||||
CHECK_EQ(A_weights->dtype, B_weights->dtype)
|
||||
<< "Data types of two edge weights must match.";
|
||||
|
||||
std::pair<CSRMatrix, NDArray> ret;
|
||||
ATEN_XPU_SWITCH_CUDA(A.indptr->ctx.device_type, XPU, "CSRMM", {
|
||||
ATEN_ID_TYPE_SWITCH(A.indptr->dtype, IdType, {
|
||||
ATEN_FLOAT_TYPE_SWITCH(A_weights->dtype, DType, "Edge weights", {
|
||||
ret = CSRMM<XPU, IdType, DType>(A, A_weights, B, B_weights);
|
||||
});
|
||||
});
|
||||
});
|
||||
return ret;
|
||||
}
|
||||
|
||||
std::pair<CSRMatrix, NDArray> CSRSum(
|
||||
const std::vector<CSRMatrix>& A, const std::vector<NDArray>& A_weights) {
|
||||
CHECK(A.size() > 0) << "The list of graphs must not be empty.";
|
||||
CHECK_EQ(A.size(), A_weights.size())
|
||||
<< "The list of edge weights must have the same length as the list of "
|
||||
"graphs.";
|
||||
const auto ctx = A[0].indptr->ctx;
|
||||
const auto idtype = A[0].indptr->dtype;
|
||||
const auto dtype = A_weights[0]->dtype;
|
||||
const auto num_rows = A[0].num_rows;
|
||||
const auto num_cols = A[0].num_cols;
|
||||
for (size_t i = 0; i < A.size(); ++i) {
|
||||
CHECK_EQ(A[i].indptr->ctx, ctx)
|
||||
<< "The devices of all graphs must be equal.";
|
||||
CHECK_EQ(A[i].indptr->dtype, idtype)
|
||||
<< "The ID types of all graphs must be equal.";
|
||||
CHECK_EQ(A[i].indices->shape[0], A_weights[i]->shape[0])
|
||||
<< "Shape of edge weights does not match the number of edges.";
|
||||
CHECK_EQ(A_weights[i]->ctx, ctx) << "The devices of edge weights must be "
|
||||
"the same as that of the graphs.";
|
||||
CHECK_EQ(A_weights[i]->dtype, dtype)
|
||||
<< "The data types of all edge weights must be equal.";
|
||||
CHECK_EQ(A[i].num_rows, num_rows)
|
||||
<< "Graphs must have the same number of nodes.";
|
||||
CHECK_EQ(A[i].num_cols, num_cols)
|
||||
<< "Graphs must have the same number of nodes.";
|
||||
}
|
||||
|
||||
std::pair<CSRMatrix, NDArray> ret;
|
||||
ATEN_XPU_SWITCH_CUDA(ctx.device_type, XPU, "CSRSum", {
|
||||
ATEN_ID_TYPE_SWITCH(idtype, IdType, {
|
||||
ATEN_FLOAT_TYPE_SWITCH(dtype, DType, "Edge weights", {
|
||||
ret = CSRSum<XPU, IdType, DType>(A, A_weights);
|
||||
});
|
||||
});
|
||||
});
|
||||
return ret;
|
||||
}
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLKernelSpMM")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
HeteroGraphRef graph = args[0];
|
||||
const std::string op = args[1];
|
||||
const std::string reduce_op = args[2];
|
||||
NDArray U = args[3];
|
||||
NDArray E = args[4];
|
||||
NDArray V = args[5];
|
||||
NDArray ArgU = args[6];
|
||||
NDArray ArgE = args[7];
|
||||
CheckCtx(
|
||||
graph->Context(), {U, E, V, ArgU, ArgE},
|
||||
{"U_data", "E_data", "out", "Arg_U", "Arg_E"});
|
||||
CheckContiguous(
|
||||
{U, E, V, ArgU, ArgE}, {"U_data", "E_data", "out", "Arg_U", "Arg_E"});
|
||||
CHECK_EQ(graph->NumEdgeTypes(), 1);
|
||||
auto pair =
|
||||
graph->meta_graph()->FindEdge(0); // only one etype in the graph.
|
||||
const dgl_type_t src_vtype = pair.first;
|
||||
const dgl_type_t dst_vtype = pair.second;
|
||||
CheckShape(
|
||||
{graph->NumVertices(src_vtype), graph->NumEdges(0),
|
||||
graph->NumVertices(dst_vtype)},
|
||||
{0, 1, 2, 2, 2}, {U, E, V, ArgU, ArgE},
|
||||
{"U_data", "E_data", "out", "Arg_U", "Arg_E"});
|
||||
SpMM(op, reduce_op, graph.sptr(), U, E, V, {ArgU, ArgE});
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLKernelGATHERMM")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
NDArray A = args[0];
|
||||
NDArray B = args[1];
|
||||
NDArray C = args[2];
|
||||
NDArray idx_a = args[3];
|
||||
NDArray idx_b = args[4];
|
||||
GatherMM(A, B, C, idx_a, idx_b);
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLKernelGATHERMMSCATTER")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
NDArray A = args[0];
|
||||
NDArray B = args[1];
|
||||
NDArray C = args[2];
|
||||
NDArray idx_a = args[3];
|
||||
NDArray idx_b = args[4];
|
||||
NDArray idx_c = args[5];
|
||||
GatherMMScatter(A, B, C, idx_a, idx_b, idx_c);
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLKernelSEGMENTMM")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
NDArray A = args[0];
|
||||
NDArray B = args[1];
|
||||
NDArray C = args[2];
|
||||
NDArray seglen_A = args[3];
|
||||
bool A_trans = args[4];
|
||||
bool B_trans = args[5];
|
||||
SegmentMM(A, B, C, seglen_A, A_trans, B_trans);
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLKernelSEGMENTMMBackwardB")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
NDArray A = args[0];
|
||||
NDArray dC = args[1];
|
||||
NDArray dB = args[2];
|
||||
NDArray seglen = args[3];
|
||||
SegmentMMBackwardB(A, dC, dB, seglen);
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLKernelEdge_softmax_forward")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
HeteroGraphRef graph = args[0];
|
||||
const std::string op = args[1];
|
||||
NDArray U = args[2];
|
||||
NDArray E = args[3];
|
||||
NDArray V = args[4];
|
||||
Edge_softmax_forward(op, graph.sptr(), U, E, V);
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLKernelEdge_softmax_backward")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
HeteroGraphRef graph = args[0];
|
||||
const std::string op = args[1];
|
||||
NDArray out = args[2];
|
||||
NDArray sds = args[3];
|
||||
NDArray back_out = args[4];
|
||||
NDArray ufeat = args[5];
|
||||
Edge_softmax_backward(op, graph.sptr(), out, sds, back_out, ufeat);
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLKernelSpMMHetero")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
HeteroGraphRef graph = args[0];
|
||||
const std::string op = args[1];
|
||||
const std::string reduce_op = args[2];
|
||||
List<Value> list_U = args[3];
|
||||
List<Value> list_E = args[4];
|
||||
List<Value> list_V = args[5];
|
||||
List<Value> list_ArgU = args[6];
|
||||
List<Value> list_ArgE = args[7];
|
||||
List<Value> list_ArgU_ntype = args[8];
|
||||
List<Value> list_ArgE_etype = args[9];
|
||||
std::vector<std::vector<NDArray>> Arg_vec; // ArgU + ArgE
|
||||
for (int i = 0; i < 4; ++i) { // ArgU + ArgE + ArgU_ntype + ArgE_etype
|
||||
Arg_vec.push_back(std::vector<NDArray>());
|
||||
}
|
||||
std::vector<NDArray> U_vec = ListValueToVector<NDArray>(list_U);
|
||||
std::vector<NDArray> V_vec = ListValueToVector<NDArray>(list_V);
|
||||
std::vector<NDArray> E_vec = ListValueToVector<NDArray>(list_E);
|
||||
Arg_vec[0] = ListValueToVector<NDArray>(list_ArgU);
|
||||
Arg_vec[1] = ListValueToVector<NDArray>(list_ArgE);
|
||||
Arg_vec[2] = ListValueToVector<NDArray>(list_ArgU_ntype);
|
||||
Arg_vec[3] = ListValueToVector<NDArray>(list_ArgE_etype);
|
||||
for (dgl_type_t etype = 0; etype < graph->NumEdgeTypes(); ++etype) {
|
||||
auto pair = graph->meta_graph()->FindEdge(etype);
|
||||
const dgl_id_t src_id = pair.first;
|
||||
const dgl_id_t dst_id = pair.second;
|
||||
NDArray U = (U_vec.size() == 0) ? NullArray() : U_vec[src_id];
|
||||
NDArray E = (E_vec.size() == 0) ? NullArray() : E_vec[etype];
|
||||
CheckCtx(
|
||||
graph->Context(),
|
||||
{U, E, V_vec[dst_id], Arg_vec[0][dst_id], Arg_vec[1][dst_id]},
|
||||
{"U_data", "E_data", "out", "Arg_U", "Arg_E"});
|
||||
CheckContiguous(
|
||||
{U, E, V_vec[dst_id], Arg_vec[0][dst_id], Arg_vec[1][dst_id]},
|
||||
{"U_data", "E_data", "out", "Arg_U", "Arg_E"});
|
||||
}
|
||||
SpMMHetero(op, reduce_op, graph.sptr(), U_vec, E_vec, &V_vec, &Arg_vec);
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLKernelSDDMM")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
HeteroGraphRef graph = args[0];
|
||||
const std::string op = args[1];
|
||||
NDArray lhs = args[2];
|
||||
NDArray rhs = args[3];
|
||||
NDArray out = args[4];
|
||||
int lhs_target = args[5];
|
||||
int rhs_target = args[6];
|
||||
CheckCtx(graph->Context(), {lhs, rhs, out}, {"lhs", "rhs", "out"});
|
||||
CheckContiguous({lhs, rhs, out}, {"lhs", "rhs", "out"});
|
||||
CHECK_EQ(graph->NumEdgeTypes(), 1);
|
||||
auto pair =
|
||||
graph->meta_graph()->FindEdge(0); // only one etype in the graph.
|
||||
const dgl_type_t src_vtype = pair.first;
|
||||
const dgl_type_t dst_vtype = pair.second;
|
||||
|
||||
CheckShape(
|
||||
{graph->NumVertices(src_vtype), graph->NumEdges(0),
|
||||
graph->NumVertices(dst_vtype)},
|
||||
{lhs_target, rhs_target, 1}, {lhs, rhs, out},
|
||||
{"U_data", "E_data", "V_data"});
|
||||
SDDMM(op, graph.sptr(), lhs, rhs, out, lhs_target, rhs_target);
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLKernelSDDMMHetero")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
HeteroGraphRef graph = args[0];
|
||||
const std::string op = args[1];
|
||||
List<Value> list_lhs = args[2];
|
||||
List<Value> list_rhs = args[3];
|
||||
List<Value> list_out = args[4];
|
||||
int lhs_target = args[5];
|
||||
int rhs_target = args[6];
|
||||
std::vector<NDArray> vec_lhs;
|
||||
std::vector<NDArray> vec_rhs;
|
||||
std::vector<NDArray> vec_out;
|
||||
|
||||
vec_lhs.reserve(list_lhs.size());
|
||||
vec_rhs.reserve(list_rhs.size());
|
||||
vec_out.reserve(list_out.size());
|
||||
|
||||
for (Value val : list_lhs) {
|
||||
vec_lhs.push_back(val->data);
|
||||
}
|
||||
for (Value val : list_rhs) {
|
||||
vec_rhs.push_back(val->data);
|
||||
}
|
||||
for (Value val : list_out) {
|
||||
vec_out.push_back(val->data);
|
||||
}
|
||||
SDDMMHetero(
|
||||
op, graph.sptr(), vec_lhs, vec_rhs, vec_out, lhs_target, rhs_target);
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLKernelSegmentReduce")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
const std::string op = args[0];
|
||||
NDArray feat = args[1];
|
||||
NDArray offsets = args[2];
|
||||
NDArray out = args[3];
|
||||
NDArray arg = args[4];
|
||||
CheckCtx(feat->ctx, {feat, offsets, out}, {"feat", "offsets", "out"});
|
||||
CheckContiguous({feat, offsets, out}, {"feat", "offsets", "out"});
|
||||
SegmentReduceDispatch(op, feat, offsets, out, arg);
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLKernelScatterAdd")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
NDArray feat = args[0];
|
||||
NDArray idx = args[1];
|
||||
NDArray out = args[2];
|
||||
CheckCtx(feat->ctx, {feat, idx, out}, {"feat", "idx", "out"});
|
||||
CheckContiguous({feat, idx, out}, {"feat", "idx", "out"});
|
||||
ScatterAddDispatch(feat, idx, out);
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLKernelUpdateGradMinMaxHetero")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
HeteroGraphRef graph = args[0];
|
||||
const std::string op = args[1];
|
||||
List<Value> list_feat = args[2];
|
||||
List<Value> list_idx = args[3];
|
||||
List<Value> list_idx_etype = args[4];
|
||||
List<Value> list_out = args[5];
|
||||
std::vector<NDArray> vec_feat = ListValueToVector<NDArray>(list_feat);
|
||||
std::vector<NDArray> vec_idx = ListValueToVector<NDArray>(list_idx);
|
||||
std::vector<NDArray> vec_idx_etype =
|
||||
ListValueToVector<NDArray>(list_idx_etype);
|
||||
std::vector<NDArray> vec_out = ListValueToVector<NDArray>(list_out);
|
||||
// CheckCtx(feat->ctx, {feat, idx, out}, {"feat", "idx", "out"});
|
||||
// CheckContiguous({feat, idx, out}, {"feat", "idx", "out"});
|
||||
UpdateGradMinMaxDispatchHetero(
|
||||
graph.sptr(), op, vec_feat, vec_idx, vec_idx_etype, &vec_out);
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLKernelBwdSegmentCmp")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
NDArray feat = args[0];
|
||||
NDArray arg = args[1];
|
||||
NDArray out = args[2];
|
||||
CheckCtx(feat->ctx, {feat, arg, out}, {"feat", "arg", "out"});
|
||||
CheckContiguous({feat, arg, out}, {"feat", "arg", "out"});
|
||||
BackwardSegmentCmpDispatch(feat, arg, out);
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLKernelGetEdgeMapping")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
HeteroGraphRef graph = args[0];
|
||||
*rv = GetEdgeMapping(graph);
|
||||
});
|
||||
|
||||
/**
|
||||
* @brief Sparse matrix multiplication with graph interface.
|
||||
*
|
||||
* @param A_ref The left operand.
|
||||
* @param A_weights The edge weights of graph A.
|
||||
* @param B_ref The right operand.
|
||||
* @param B_weights The edge weights of graph B.
|
||||
* @param num_vtypes The number of vertex types of the graph to be returned.
|
||||
* @return A pair consisting of the new graph as well as its edge weights.
|
||||
*/
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLCSRMM")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
const HeteroGraphRef A_ref = args[0];
|
||||
NDArray A_weights = args[1];
|
||||
const HeteroGraphRef B_ref = args[2];
|
||||
NDArray B_weights = args[3];
|
||||
int num_vtypes = args[4];
|
||||
|
||||
const HeteroGraphPtr A = A_ref.sptr();
|
||||
const HeteroGraphPtr B = B_ref.sptr();
|
||||
CHECK_EQ(A->NumEdgeTypes(), 1)
|
||||
<< "The first graph must have only one edge type.";
|
||||
CHECK_EQ(B->NumEdgeTypes(), 1)
|
||||
<< "The second graph must have only one edge type.";
|
||||
const auto A_csr = A->GetCSRMatrix(0);
|
||||
const auto B_csr = B->GetCSRMatrix(0);
|
||||
auto result = CSRMM(A_csr, A_weights, B_csr, B_weights);
|
||||
|
||||
List<ObjectRef> ret;
|
||||
ret.push_back(
|
||||
HeteroGraphRef(CreateFromCSR(num_vtypes, result.first, ALL_CODE)));
|
||||
ret.push_back(Value(MakeValue(result.second)));
|
||||
*rv = ret;
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLCSRSum")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
List<HeteroGraphRef> A_refs = args[0];
|
||||
List<Value> A_weights = args[1];
|
||||
|
||||
std::vector<NDArray> weights = ListValueToVector<NDArray>(A_weights);
|
||||
std::vector<CSRMatrix> mats;
|
||||
mats.reserve(A_refs.size());
|
||||
int num_vtypes = 0;
|
||||
for (auto A_ref : A_refs) {
|
||||
const HeteroGraphPtr A = A_ref.sptr();
|
||||
CHECK_EQ(A->NumEdgeTypes(), 1)
|
||||
<< "Graphs must have only one edge type.";
|
||||
mats.push_back(A->GetCSRMatrix(0));
|
||||
if (num_vtypes == 0) num_vtypes = A->NumVertexTypes();
|
||||
}
|
||||
auto result = CSRSum(mats, weights);
|
||||
|
||||
List<ObjectRef> ret;
|
||||
ret.push_back(
|
||||
HeteroGraphRef(CreateFromCSR(num_vtypes, result.first, ALL_CODE)));
|
||||
ret.push_back(Value(MakeValue(result.second)));
|
||||
*rv = ret;
|
||||
});
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLCSRMask")
|
||||
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
const HeteroGraphRef A_ref = args[0];
|
||||
NDArray A_weights = args[1];
|
||||
const HeteroGraphRef B_ref = args[2];
|
||||
|
||||
const HeteroGraphPtr A = A_ref.sptr();
|
||||
const HeteroGraphPtr B = B_ref.sptr();
|
||||
CHECK_EQ(A->NumEdgeTypes(), 1)
|
||||
<< "Both graphs must have only one edge type.";
|
||||
CHECK_EQ(B->NumEdgeTypes(), 1)
|
||||
<< "Both graphs must have only one edge type.";
|
||||
const CSRMatrix& A_csr = A->GetCSRMatrix(0);
|
||||
const COOMatrix& B_coo = B->GetCOOMatrix(0);
|
||||
CHECK_EQ(A_csr.num_rows, B_coo.num_rows)
|
||||
<< "Both graphs must have the same number of nodes.";
|
||||
CHECK_EQ(A_csr.num_cols, B_coo.num_cols)
|
||||
<< "Both graphs must have the same number of nodes.";
|
||||
|
||||
NDArray result;
|
||||
ATEN_FLOAT_TYPE_SWITCH(A_weights->dtype, DType, "Edge weights", {
|
||||
result =
|
||||
aten::CSRGetData<DType>(A_csr, B_coo.row, B_coo.col, A_weights, 0.);
|
||||
});
|
||||
*rv = result;
|
||||
});
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,196 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/kernel_decl.h
|
||||
* @brief Sparse matrix format-specific operator declarations.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_KERNEL_DECL_H_
|
||||
#define DGL_ARRAY_KERNEL_DECL_H_
|
||||
|
||||
#include <dgl/base_heterograph.h>
|
||||
#include <dgl/bcast.h>
|
||||
#include <dgl/runtime/ndarray.h>
|
||||
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
|
||||
/**
|
||||
* @brief Generalized Sparse Matrix Dense Matrix Multiplication on Csr format.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SpMMCsr(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const aten::CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
|
||||
/**
|
||||
* @brief Generalized Sparse Matrix Dense Matrix Multiplication on Csr format
|
||||
* with heterograph support.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SpMMCsrHetero(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
|
||||
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
|
||||
std::vector<std::vector<NDArray>>* out_aux,
|
||||
const std::vector<dgl_type_t>& ufeat_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
/**
|
||||
* @brief Generalized Sparse Matrix Dense Matrix Multiplication on Coo format.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SpMMCoo(
|
||||
const std::string& op, const std::string& reduce, const BcastOff& bcast,
|
||||
const aten::COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
|
||||
std::vector<NDArray> out_aux);
|
||||
|
||||
/**
|
||||
* @brief Generalized Sampled Dense-Dense Matrix Multiplication on Csr format.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SDDMMCsr(
|
||||
const std::string& op, const BcastOff& bcast, const aten::CSRMatrix& csr,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
/**
|
||||
* @brief Generalized Sampled Dense-Dense Matrix Multiplication on Csr format
|
||||
* with heterograph support.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SDDMMCsrHetero(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& vec_lhs,
|
||||
const std::vector<NDArray>& vec_rhs, std::vector<NDArray> vec_out,
|
||||
int lhs_target, int rhs_target, const std::vector<dgl_type_t>& ufeat_eid,
|
||||
const std::vector<dgl_type_t>& out_eid);
|
||||
|
||||
/**
|
||||
* @brief Generalized Sampled Dense-Dense Matrix Multiplication on Coo format.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SDDMMCoo(
|
||||
const std::string& op, const BcastOff& bcast, const aten::COOMatrix& coo,
|
||||
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
|
||||
|
||||
/**
|
||||
* @brief Generalized Sampled Dense-Dense Matrix Multiplication on Coo format
|
||||
* with heterograph support.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SDDMMCooHetero(
|
||||
const std::string& op, const BcastOff& bcast,
|
||||
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& vec_lhs,
|
||||
const std::vector<NDArray>& vec_rhs, std::vector<NDArray> vec_out,
|
||||
int lhs_target, int rhs_target, const std::vector<dgl_type_t>& lhs_eid,
|
||||
const std::vector<dgl_type_t>& rhs_eid);
|
||||
|
||||
/**
|
||||
* @brief Generalized Dense Matrix-Matrix Multiplication according to relation
|
||||
* types.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void GatherMM(
|
||||
const NDArray A, const NDArray B, NDArray out, const NDArray idx_a,
|
||||
const NDArray idx_b);
|
||||
|
||||
/**
|
||||
* @brief Generalized Dense Matrix-Matrix Multiplication according to relation
|
||||
* types.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void GatherMMScatter(
|
||||
const NDArray A, const NDArray B, NDArray out, const NDArray idx_a,
|
||||
const NDArray idx_b, const NDArray idx_c);
|
||||
|
||||
/**
|
||||
* @brief Generalized segmented dense Matrix-Matrix Multiplication.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SegmentMM(
|
||||
const NDArray A, const NDArray B, NDArray out, const NDArray seglen_A,
|
||||
bool a_trans, bool b_trans);
|
||||
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SegmentMMBackwardB(
|
||||
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen);
|
||||
|
||||
/**
|
||||
* @brief Segment reduce.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void SegmentReduce(
|
||||
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
|
||||
NDArray arg);
|
||||
|
||||
/**
|
||||
* @brief Scatter Add on first dimension.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void ScatterAdd(NDArray feat, NDArray idx, NDArray out);
|
||||
|
||||
/**
|
||||
* @brief Update gradients for reduce operator max and min on first dimension.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void UpdateGradMinMax_hetero(
|
||||
const HeteroGraphPtr& g, const std::string& op,
|
||||
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
|
||||
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out);
|
||||
|
||||
/**
|
||||
* @brief Backward function of segment cmp.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void BackwardSegmentCmp(NDArray feat, NDArray arg, NDArray out);
|
||||
|
||||
/**
|
||||
* @brief Sparse-sparse matrix multiplication
|
||||
*
|
||||
* @param A The left operand.
|
||||
* @param A_weights The weights of matrix as a 1D tensor.
|
||||
* @param B The right operand.
|
||||
* @param B_weights The weights of matrix as a 1D tensor.
|
||||
*
|
||||
* @note GPU implementation will cast the indices to 32 bit.
|
||||
* @note The zero entries in the result are not removed.
|
||||
* @note The CSR matrix should not have duplicate entries.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
std::pair<CSRMatrix, NDArray> CSRMM(
|
||||
const CSRMatrix& A, NDArray A_weights, const CSRMatrix& B,
|
||||
NDArray B_weights);
|
||||
|
||||
/**
|
||||
* @brief Sparse-sparse matrix summation.
|
||||
*
|
||||
* @param A The sparse matrices with the same size.
|
||||
* @param A_weights The weights of each sparse matrix as a 1D tensor.
|
||||
*
|
||||
* @note GPU implementation will cast the indices to 32 bit.
|
||||
* @note The zero entries in the result are not removed.
|
||||
* @note The CSR matrix should not have duplicate entries.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
std::pair<CSRMatrix, NDArray> CSRSum(
|
||||
const std::vector<CSRMatrix>& A, const std::vector<NDArray>& A_weights);
|
||||
|
||||
/**
|
||||
* @brief Edge_softmax_csr forward function on Csr format.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void Edge_softmax_csr_forward(
|
||||
const std::string& op, const BcastOff& bcast, const aten::CSRMatrix& csr,
|
||||
NDArray ufeat, NDArray efeat, NDArray out);
|
||||
/**
|
||||
* @brief Edge_softmax_csr backward function on Csr format.
|
||||
*/
|
||||
template <int XPU, typename IdType, typename DType>
|
||||
void Edge_softmax_csr_backward(
|
||||
const std::string& op, const BcastOff& bcast, const aten::CSRMatrix& csr,
|
||||
NDArray ufeat, NDArray efeat, NDArray out);
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_KERNEL_DECL_H_
|
||||
@@ -0,0 +1,660 @@
|
||||
/**
|
||||
* Copyright (c) 2021 Intel Corporation
|
||||
*
|
||||
* @file distgnn/partition/main_Libra.py
|
||||
* @brief Libra - Vertex-cut based graph partitioner for distirbuted training
|
||||
* @author Vasimuddin Md <vasimuddin.md@intel.com>,
|
||||
* Guixiang Ma <guixiang.ma@intel.com>
|
||||
* Sanchit Misra <sanchit.misra@intel.com>,
|
||||
* Ramanarayan Mohanty <ramanarayan.mohanty@intel.com>,
|
||||
* Sasikanth Avancha <sasikanth.avancha@intel.com>
|
||||
* Nesreen K. Ahmed <nesreen.k.ahmed@intel.com>
|
||||
*/
|
||||
|
||||
#include <dgl/base_heterograph.h>
|
||||
#include <dgl/packed_func_ext.h>
|
||||
#include <dgl/random.h>
|
||||
#include <dgl/runtime/parallel_for.h>
|
||||
#include <dmlc/omp.h>
|
||||
#include <stdint.h>
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "../c_api_common.h"
|
||||
#include "./check.h"
|
||||
#include "kernel_decl.h"
|
||||
|
||||
using namespace dgl::runtime;
|
||||
|
||||
namespace dgl {
|
||||
namespace aten {
|
||||
|
||||
template <typename IdType>
|
||||
int32_t Ver2partition(IdType in_val, int64_t *node_map, int32_t num_parts) {
|
||||
int32_t pos = 0;
|
||||
for (int32_t p = 0; p < num_parts; p++) {
|
||||
if (in_val < node_map[p]) return pos;
|
||||
pos = pos + 1;
|
||||
}
|
||||
LOG(FATAL) << "Error: Unexpected output in Ver2partition!";
|
||||
return -1;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Identifies the lead loaded partition/community for a given edge
|
||||
* assignment.
|
||||
*/
|
||||
int32_t LeastLoad(int64_t *community_edges, int32_t nc) {
|
||||
std::vector<int> loc;
|
||||
int32_t min = 1e9;
|
||||
for (int32_t i = 0; i < nc; i++) {
|
||||
if (community_edges[i] < min) {
|
||||
min = community_edges[i];
|
||||
}
|
||||
}
|
||||
for (int32_t i = 0; i < nc; i++) {
|
||||
if (community_edges[i] == min) {
|
||||
loc.push_back(i);
|
||||
}
|
||||
}
|
||||
|
||||
int32_t r = RandomEngine::ThreadLocal()->RandInt(loc.size());
|
||||
CHECK(loc[r] < nc);
|
||||
return loc[r];
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Libra - vertexcut based graph partitioning.
|
||||
* It takes list of edges from input DGL graph and distributed them among nc
|
||||
* partitions During edge distribution, Libra assign a given edge to a partition
|
||||
* based on the end vertices, in doing so, it tries to minimized the splitting
|
||||
* of the graph vertices. In case of conflict Libra assigns an edge to the least
|
||||
* loaded partition/community.
|
||||
* @param[in] nc Number of partitions/communities
|
||||
* @param[in] node_degree per node degree
|
||||
* @param[in] edgenum_unassigned node degree
|
||||
* @param[out] community_weights weight of the created partitions
|
||||
* @param[in] u src nodes
|
||||
* @param[in] v dst nodes
|
||||
* @param[out] w weight per edge
|
||||
* @param[out] out partition assignment of the edges
|
||||
* @param[in] N_n number of nodes in the input graph
|
||||
* @param[in] N_e number of edges in the input graph
|
||||
* @param[in] prefix output/partition storage location
|
||||
*/
|
||||
template <typename IdType, typename IdType2>
|
||||
void LibraVertexCut(
|
||||
int32_t nc, NDArray node_degree, NDArray edgenum_unassigned,
|
||||
NDArray community_weights, NDArray u, NDArray v, NDArray w, NDArray out,
|
||||
int64_t N_n, int64_t N_e, const std::string &prefix) {
|
||||
int32_t *out_ptr = out.Ptr<int32_t>();
|
||||
IdType2 *node_degree_ptr = node_degree.Ptr<IdType2>();
|
||||
IdType2 *edgenum_unassigned_ptr = edgenum_unassigned.Ptr<IdType2>();
|
||||
IdType *u_ptr = u.Ptr<IdType>();
|
||||
IdType *v_ptr = v.Ptr<IdType>();
|
||||
int64_t *w_ptr = w.Ptr<int64_t>();
|
||||
int64_t *community_weights_ptr = community_weights.Ptr<int64_t>();
|
||||
|
||||
std::vector<std::vector<int32_t> > node_assignments(N_n);
|
||||
std::vector<IdType2> replication_list;
|
||||
// local allocations
|
||||
int64_t *community_edges = new int64_t[nc]();
|
||||
int64_t *cache = new int64_t[nc]();
|
||||
|
||||
int64_t meter = static_cast<int>(N_e / 100);
|
||||
for (int64_t i = 0; i < N_e; i++) {
|
||||
IdType u = u_ptr[i]; // edge end vertex 1
|
||||
IdType v = v_ptr[i]; // edge end vertex 2
|
||||
int64_t w = w_ptr[i]; // edge weight
|
||||
|
||||
CHECK(u < N_n);
|
||||
CHECK(v < N_n);
|
||||
|
||||
if (i % meter == 0) {
|
||||
fprintf(stderr, ".");
|
||||
fflush(0);
|
||||
}
|
||||
|
||||
if (node_assignments[u].size() == 0 && node_assignments[v].size() == 0) {
|
||||
int32_t c = LeastLoad(community_edges, nc);
|
||||
out_ptr[i] = c;
|
||||
CHECK_LT(c, nc);
|
||||
|
||||
community_edges[c]++;
|
||||
community_weights_ptr[c] = community_weights_ptr[c] + w;
|
||||
node_assignments[u].push_back(c);
|
||||
if (u != v) node_assignments[v].push_back(c);
|
||||
|
||||
CHECK(node_assignments[u].size() <= static_cast<size_t>(nc))
|
||||
<< "[bug] 1. generated splits (u) are greater than nc!";
|
||||
CHECK(node_assignments[v].size() <= static_cast<size_t>(nc))
|
||||
<< "[bug] 1. generated splits (v) are greater than nc!";
|
||||
edgenum_unassigned_ptr[u]--;
|
||||
edgenum_unassigned_ptr[v]--;
|
||||
} else if (
|
||||
node_assignments[u].size() != 0 && node_assignments[v].size() == 0) {
|
||||
for (uint32_t j = 0; j < node_assignments[u].size(); j++) {
|
||||
int32_t cind = node_assignments[u][j];
|
||||
cache[j] = community_edges[cind];
|
||||
}
|
||||
int32_t cindex = LeastLoad(cache, node_assignments[u].size());
|
||||
int32_t c = node_assignments[u][cindex];
|
||||
out_ptr[i] = c;
|
||||
community_edges[c]++;
|
||||
community_weights_ptr[c] = community_weights_ptr[c] + w;
|
||||
|
||||
node_assignments[v].push_back(c);
|
||||
CHECK(node_assignments[v].size() <= static_cast<size_t>(nc))
|
||||
<< "[bug] 2. generated splits (v) are greater than nc!";
|
||||
edgenum_unassigned_ptr[u]--;
|
||||
edgenum_unassigned_ptr[v]--;
|
||||
} else if (
|
||||
node_assignments[v].size() != 0 && node_assignments[u].size() == 0) {
|
||||
for (uint32_t j = 0; j < node_assignments[v].size(); j++) {
|
||||
int32_t cind = node_assignments[v][j];
|
||||
cache[j] = community_edges[cind];
|
||||
}
|
||||
int32_t cindex = LeastLoad(cache, node_assignments[v].size());
|
||||
int32_t c = node_assignments[v][cindex];
|
||||
CHECK(c < nc) << "[bug] 2. partition greater than nc !!";
|
||||
out_ptr[i] = c;
|
||||
|
||||
community_edges[c]++;
|
||||
community_weights_ptr[c] = community_weights_ptr[c] + w;
|
||||
|
||||
node_assignments[u].push_back(c);
|
||||
CHECK(node_assignments[u].size() <= static_cast<size_t>(nc))
|
||||
<< "[bug] 3. generated splits (u) are greater than nc!";
|
||||
edgenum_unassigned_ptr[u]--;
|
||||
edgenum_unassigned_ptr[v]--;
|
||||
} else {
|
||||
std::vector<int> setv(nc), intersetv;
|
||||
for (int32_t j = 0; j < nc; j++) setv[j] = 0;
|
||||
int32_t interset = 0;
|
||||
|
||||
CHECK(node_assignments[u].size() <= static_cast<size_t>(nc))
|
||||
<< "[bug] 4. generated splits (u) are greater than nc!";
|
||||
CHECK(node_assignments[v].size() <= static_cast<size_t>(nc))
|
||||
<< "[bug] 4. generated splits (v) are greater than nc!";
|
||||
for (size_t j = 0; j < node_assignments[v].size(); j++) {
|
||||
CHECK(node_assignments[v][j] < nc)
|
||||
<< "[bug] 4. Part assigned (v) greater than nc!";
|
||||
setv[node_assignments[v][j]]++;
|
||||
}
|
||||
|
||||
for (size_t j = 0; j < node_assignments[u].size(); j++) {
|
||||
CHECK(node_assignments[u][j] < nc)
|
||||
<< "[bug] 4. Part assigned (u) greater than nc!";
|
||||
setv[node_assignments[u][j]]++;
|
||||
}
|
||||
|
||||
for (int32_t j = 0; j < nc; j++) {
|
||||
CHECK(setv[j] <= 2) << "[bug] 4. unexpected computed value !!!";
|
||||
if (setv[j] == 2) {
|
||||
interset++;
|
||||
intersetv.push_back(j);
|
||||
}
|
||||
}
|
||||
if (interset) {
|
||||
for (size_t j = 0; j < intersetv.size(); j++) {
|
||||
int32_t cind = intersetv[j];
|
||||
cache[j] = community_edges[cind];
|
||||
}
|
||||
int32_t cindex = LeastLoad(cache, intersetv.size());
|
||||
int32_t c = intersetv[cindex];
|
||||
CHECK(c < nc) << "[bug] 4. partition greater than nc !!";
|
||||
out_ptr[i] = c;
|
||||
community_edges[c]++;
|
||||
community_weights_ptr[c] = community_weights_ptr[c] + w;
|
||||
edgenum_unassigned_ptr[u]--;
|
||||
edgenum_unassigned_ptr[v]--;
|
||||
} else {
|
||||
if (node_degree_ptr[u] < node_degree_ptr[v]) {
|
||||
for (uint32_t j = 0; j < node_assignments[u].size(); j++) {
|
||||
int32_t cind = node_assignments[u][j];
|
||||
cache[j] = community_edges[cind];
|
||||
}
|
||||
int32_t cindex = LeastLoad(cache, node_assignments[u].size());
|
||||
int32_t c = node_assignments[u][cindex];
|
||||
CHECK(c < nc) << "[bug] 5. partition greater than nc !!";
|
||||
out_ptr[i] = c;
|
||||
community_edges[c]++;
|
||||
community_weights_ptr[c] = community_weights_ptr[c] + w;
|
||||
|
||||
for (uint32_t j = 0; j < node_assignments[v].size(); j++) {
|
||||
CHECK(node_assignments[v][j] != c)
|
||||
<< "[bug] 5. duplicate partition (v) assignment !!";
|
||||
}
|
||||
|
||||
node_assignments[v].push_back(c);
|
||||
CHECK(node_assignments[v].size() <= static_cast<size_t>(nc))
|
||||
<< "[bug] 5. generated splits (v) greater than nc!!";
|
||||
replication_list.push_back(v);
|
||||
edgenum_unassigned_ptr[u]--;
|
||||
edgenum_unassigned_ptr[v]--;
|
||||
} else {
|
||||
for (uint32_t j = 0; j < node_assignments[v].size(); j++) {
|
||||
int32_t cind = node_assignments[v][j];
|
||||
cache[j] = community_edges[cind];
|
||||
}
|
||||
int32_t cindex = LeastLoad(cache, node_assignments[v].size());
|
||||
int32_t c = node_assignments[v][cindex];
|
||||
CHECK(c < nc) << "[bug] 6. partition greater than nc !!";
|
||||
out_ptr[i] = c;
|
||||
community_edges[c]++;
|
||||
community_weights_ptr[c] = community_weights_ptr[c] + w;
|
||||
for (uint32_t j = 0; j < node_assignments[u].size(); j++) {
|
||||
CHECK(node_assignments[u][j] != c)
|
||||
<< "[bug] 6. duplicate partition (u) assignment !!";
|
||||
}
|
||||
if (u != v) node_assignments[u].push_back(c);
|
||||
|
||||
CHECK(node_assignments[u].size() <= static_cast<size_t>(nc))
|
||||
<< "[bug] 6. generated splits (u) greater than nc!!";
|
||||
replication_list.push_back(u);
|
||||
edgenum_unassigned_ptr[u]--;
|
||||
edgenum_unassigned_ptr[v]--;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
delete cache;
|
||||
|
||||
for (int64_t c = 0; c < nc; c++) {
|
||||
std::string path = prefix + "/community" + std::to_string(c) + ".txt";
|
||||
|
||||
FILE *fp = fopen(path.c_str(), "w");
|
||||
CHECK_NE(fp, static_cast<FILE *>(NULL))
|
||||
<< "Error: can not open file: " << path.c_str();
|
||||
|
||||
for (int64_t i = 0; i < N_e; i++) {
|
||||
if (out_ptr[i] == c)
|
||||
fprintf(
|
||||
fp, "%ld,%ld,%ld\n", static_cast<int64_t>(u_ptr[i]),
|
||||
static_cast<int64_t>(v_ptr[i]), w_ptr[i]);
|
||||
}
|
||||
fclose(fp);
|
||||
}
|
||||
|
||||
std::string path = prefix + "/replicationlist.csv";
|
||||
FILE *fp = fopen(path.c_str(), "w");
|
||||
CHECK_NE(fp, static_cast<FILE *>(NULL))
|
||||
<< "Error: can not open file: " << path.c_str();
|
||||
|
||||
fprintf(fp, "## The Indices of Nodes that are replicated :: Header");
|
||||
printf("\nTotal replication: %ld\n", replication_list.size());
|
||||
|
||||
for (uint64_t i = 0; i < replication_list.size(); i++)
|
||||
fprintf(fp, "%ld\n", static_cast<int64_t>(replication_list[i]));
|
||||
|
||||
printf("Community weights:\n");
|
||||
for (int64_t c = 0; c < nc; c++) printf("%ld ", community_weights_ptr[c]);
|
||||
printf("\n");
|
||||
|
||||
printf("Community edges:\n");
|
||||
for (int64_t c = 0; c < nc; c++) printf("%ld ", community_edges[c]);
|
||||
printf("\n");
|
||||
|
||||
delete[] community_edges;
|
||||
fclose(fp);
|
||||
}
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLLibraVertexCut")
|
||||
.set_body([](DGLArgs args, DGLRetValue *rv) {
|
||||
int32_t nc = args[0];
|
||||
NDArray node_degree = args[1];
|
||||
NDArray edgenum_unassigned = args[2];
|
||||
NDArray community_weights = args[3];
|
||||
NDArray u = args[4];
|
||||
NDArray v = args[5];
|
||||
NDArray w = args[6];
|
||||
NDArray out = args[7];
|
||||
int64_t N = args[8];
|
||||
int64_t N_e = args[9];
|
||||
std::string prefix = args[10];
|
||||
|
||||
ATEN_ID_TYPE_SWITCH(node_degree->dtype, IdType2, {
|
||||
ATEN_ID_TYPE_SWITCH(u->dtype, IdType, {
|
||||
LibraVertexCut<IdType, IdType2>(
|
||||
nc, node_degree, edgenum_unassigned, community_weights, u, v, w,
|
||||
out, N, N_e, prefix);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* @brief
|
||||
* 1. Builds dictionary (ldt) for assigning local node IDs to nodes in the
|
||||
* partitions
|
||||
* 2. Builds dictionary (gdt) for storing copies (local ID) of split nodes
|
||||
* These dictionaries will be used in the subsequesnt stages to setup
|
||||
* tracking of split nodes copies across the partition, setting up partition
|
||||
* `ndata` dictionaries.
|
||||
* @param[out] a local src node ID of an edge in a partition
|
||||
* @param[out] b local dst node ID of an edge in a partition
|
||||
* @param[-] indices temporary memory, keeps track of global node ID to local
|
||||
* node ID in a partition
|
||||
* @param[out] ldt_key per partition dict for storing global and local node IDs
|
||||
* (consecutive)
|
||||
* @param[out] gdt_key global dict for storing number of local nodes (or split
|
||||
* nodes) for a given global node ID
|
||||
* @param[out] gdt_value global dict, stores local node IDs (due to split)
|
||||
* across partitions for a given global node ID
|
||||
* @param[out] node_map keeps track of range of local node IDs (consecutive)
|
||||
* given to the nodes in the partitions
|
||||
* @param[in, out] offset start of the range of local node IDs for this
|
||||
* partition
|
||||
* @param[in] nc number of partitions/communities
|
||||
* @param[in] c current partition number
|
||||
* @param[in] fsize size of pre-allocated
|
||||
* memory tensor
|
||||
* @param[in] prefix input Libra partition file location
|
||||
*/
|
||||
List<Value> Libra2dglBuildDict(
|
||||
NDArray a, NDArray b, NDArray indices, NDArray ldt_key, NDArray gdt_key,
|
||||
NDArray gdt_value, NDArray node_map, NDArray offset, int32_t nc, int32_t c,
|
||||
int64_t fsize, const std::string &prefix) {
|
||||
int64_t *indices_ptr = indices.Ptr<int64_t>(); // 1D temp array
|
||||
int64_t *ldt_key_ptr =
|
||||
ldt_key.Ptr<int64_t>(); // 1D local nodes <-> global nodes
|
||||
int64_t *gdt_key_ptr = gdt_key.Ptr<int64_t>(); // 1D #split copies per node
|
||||
int64_t *gdt_value_ptr = gdt_value.Ptr<int64_t>(); // 2D tensor
|
||||
int64_t *node_map_ptr = node_map.Ptr<int64_t>(); // 1D tensor
|
||||
int64_t *offset_ptr = offset.Ptr<int64_t>(); // 1D tensor
|
||||
int32_t width = nc;
|
||||
|
||||
int64_t *a_ptr = a.Ptr<int64_t>(); // stores local src and dst node ID,
|
||||
int64_t *b_ptr = b.Ptr<int64_t>(); // to create the partition graph
|
||||
|
||||
int64_t N_n = indices->shape[0];
|
||||
int64_t num_nodes = ldt_key->shape[0];
|
||||
|
||||
for (int64_t i = 0; i < N_n; i++) {
|
||||
indices_ptr[i] = -100;
|
||||
}
|
||||
|
||||
int64_t pos = 0;
|
||||
int64_t edge = 0;
|
||||
std::string path = prefix + "/community" + std::to_string(c) + ".txt";
|
||||
FILE *fp = fopen(path.c_str(), "r");
|
||||
CHECK_NE(fp, static_cast<FILE *>(NULL))
|
||||
<< "Error: can not open file: " << path.c_str();
|
||||
|
||||
while (!feof(fp) && edge < fsize) {
|
||||
int64_t u, v;
|
||||
float w;
|
||||
CHECK_EQ(
|
||||
fscanf(fp, "%ld,%ld,%f\n", &u, &v, &w),
|
||||
3); // reading an edge - the src and dst global node IDs
|
||||
|
||||
if (indices_ptr[u] ==
|
||||
-100) { // if already not assigned a local node ID, local node ID is
|
||||
ldt_key_ptr[pos] = u; // already assigned for this global node ID
|
||||
CHECK(pos < num_nodes); // Sanity check
|
||||
indices_ptr[u] =
|
||||
pos++; // consecutive local node ID for a given global node ID
|
||||
}
|
||||
if (indices_ptr[v] == -100) { // if already not assigned a local node ID
|
||||
ldt_key_ptr[pos] = v;
|
||||
CHECK(pos < num_nodes); // Sanity check
|
||||
indices_ptr[v] = pos++;
|
||||
}
|
||||
a_ptr[edge] = indices_ptr[u]; // new local ID for an edge
|
||||
b_ptr[edge++] = indices_ptr[v]; // new local ID for an edge
|
||||
}
|
||||
CHECK(edge <= fsize)
|
||||
<< "[Bug] memory allocated for #edges per partition is not enough.";
|
||||
fclose(fp);
|
||||
|
||||
List<Value> ret;
|
||||
ret.push_back(Value(
|
||||
MakeValue(pos))); // returns total number of nodes in this partition
|
||||
ret.push_back(Value(
|
||||
MakeValue(edge))); // returns total number of edges in this partition
|
||||
|
||||
for (int64_t i = 0; i < pos; i++) {
|
||||
int64_t u = ldt_key_ptr[i]; // global node ID
|
||||
// int64_t v = indices_ptr[u];
|
||||
int64_t v = i; // local node ID
|
||||
int64_t *ind =
|
||||
&gdt_key_ptr[u]; // global dict, total number of local node IDs (an
|
||||
// offset) as of now for a given global node ID
|
||||
int64_t *ptr = gdt_value_ptr + u * width;
|
||||
ptr[*ind] =
|
||||
offset_ptr[0] + v; // stores a local node ID for the global node ID
|
||||
(*ind)++;
|
||||
CHECK_NE(v, -100);
|
||||
CHECK(*ind <= nc);
|
||||
}
|
||||
node_map_ptr[c] =
|
||||
offset_ptr[0] +
|
||||
pos; // since local node IDs for a partition are consecutive,
|
||||
// we maintain the range of local node IDs like this
|
||||
offset_ptr[0] += pos;
|
||||
|
||||
return ret;
|
||||
}
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLLibra2dglBuildDict")
|
||||
.set_body([](DGLArgs args, DGLRetValue *rv) {
|
||||
NDArray a = args[0];
|
||||
NDArray b = args[1];
|
||||
NDArray indices = args[2];
|
||||
NDArray ldt_key = args[3];
|
||||
NDArray gdt_key = args[4];
|
||||
NDArray gdt_value = args[5];
|
||||
NDArray node_map = args[6];
|
||||
NDArray offset = args[7];
|
||||
int32_t nc = args[8];
|
||||
int32_t c = args[9];
|
||||
int64_t fsize = args[10];
|
||||
std::string prefix = args[11];
|
||||
List<Value> ret = Libra2dglBuildDict(
|
||||
a, b, indices, ldt_key, gdt_key, gdt_value, node_map, offset, nc, c,
|
||||
fsize, prefix);
|
||||
*rv = ret;
|
||||
});
|
||||
|
||||
/**
|
||||
* @brief sets up the 1-level tree among the clones of the split-nodes.
|
||||
* @param[in] gdt_key global dict for assigning consecutive node IDs to nodes
|
||||
* across all the partitions
|
||||
* @param[in] gdt_value global dict for assigning consecutive node IDs to nodes
|
||||
* across all the partition
|
||||
* @param[out] lrtensor keeps the root node ID of 1-level tree
|
||||
* @param[in] nc number of partitions/communities
|
||||
* @param[in] Nn number of nodes in the input graph
|
||||
*/
|
||||
void Libra2dglSetLR(
|
||||
NDArray gdt_key, NDArray gdt_value, NDArray lrtensor, int32_t nc,
|
||||
int64_t Nn) {
|
||||
int64_t *gdt_key_ptr = gdt_key.Ptr<int64_t>(); // 1D tensor
|
||||
int64_t *gdt_value_ptr = gdt_value.Ptr<int64_t>(); // 2D tensor
|
||||
int64_t *lrtensor_ptr = lrtensor.Ptr<int64_t>(); // 1D tensor
|
||||
|
||||
int32_t width = nc;
|
||||
int64_t cnt = 0;
|
||||
int64_t avg_split_copy = 0, scnt = 0;
|
||||
|
||||
for (int64_t i = 0; i < Nn; i++) {
|
||||
if (gdt_key_ptr[i] <= 0) {
|
||||
cnt++;
|
||||
} else {
|
||||
int32_t val = RandomEngine::ThreadLocal()->RandInt(gdt_key_ptr[i]);
|
||||
CHECK(val >= 0 && val < gdt_key_ptr[i]);
|
||||
CHECK(gdt_key_ptr[i] <= nc);
|
||||
|
||||
int64_t *ptr = gdt_value_ptr + i * width;
|
||||
lrtensor_ptr[i] = ptr[val];
|
||||
}
|
||||
if (gdt_key_ptr[i] > 1) {
|
||||
avg_split_copy += gdt_key_ptr[i];
|
||||
scnt++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLLibra2dglSetLR")
|
||||
.set_body([](DGLArgs args, DGLRetValue *rv) {
|
||||
NDArray gdt_key = args[0];
|
||||
NDArray gdt_value = args[1];
|
||||
NDArray lrtensor = args[2];
|
||||
int32_t nc = args[3];
|
||||
int64_t Nn = args[4];
|
||||
|
||||
Libra2dglSetLR(gdt_key, gdt_value, lrtensor, nc, Nn);
|
||||
});
|
||||
|
||||
/**
|
||||
* @brief For each node in a partition, it creates a list of remote clone IDs;
|
||||
* also, for each node in a partition, it gathers the data (feats, label,
|
||||
* trian, test) from input graph.
|
||||
* @param[out] feat node features in current partition c.
|
||||
* @param[in] gfeat input graph node features.
|
||||
* @param[out] adj list of node IDs of remote clones.
|
||||
* @param[out] inner_nodes marks whether a node is split or not.
|
||||
* @param[in] ldt_key per partition dict for tracking global to local node IDs
|
||||
* @param[out] gdt_key global dict for storing number of local nodes (or split
|
||||
* nodes) for a given global node ID
|
||||
* @param[out] gdt_value global
|
||||
* dict, stores local node IDs (due to split) across partitions for
|
||||
* a given global node ID.
|
||||
* @param[in] node_map keeps track of range of local node IDs (consecutive)
|
||||
* given to the nodes in the partitions.
|
||||
* @param[out] lr 1-level tree marking for local split nodes.
|
||||
* @param[in] lrtensor global (all the partitions) 1-level tree.
|
||||
* @param[in] num_nodes number of nodes in current partition.
|
||||
* @param[in] nc number of partitions/communities.
|
||||
* @param[in] c current partition/community.
|
||||
* @param[in] feat_size node feature vector size.
|
||||
* @param[out] labels local (for this partition) labels.
|
||||
* @param[out] trainm local (for this partition) training nodes.
|
||||
* @param[out] testm local (for this partition) testing nodes.
|
||||
* @param[out] valm local (for this partition) validation nodes.
|
||||
* @param[in] glabels global (input graph) labels.
|
||||
* @param[in] gtrainm glabal (input graph) training nodes.
|
||||
* @param[in] gtestm glabal (input graph) testing nodes.
|
||||
* @param[in] gvalm glabal (input graph) validation nodes.
|
||||
* @param[out] Nn number of nodes in the input graph.
|
||||
*/
|
||||
template <typename IdType, typename IdType2, typename DType>
|
||||
void Libra2dglBuildAdjlist(
|
||||
NDArray feat, NDArray gfeat, NDArray adj, NDArray inner_node,
|
||||
NDArray ldt_key, NDArray gdt_key, NDArray gdt_value, NDArray node_map,
|
||||
NDArray lr, NDArray lrtensor, int64_t num_nodes, int32_t nc, int32_t c,
|
||||
int32_t feat_size, NDArray labels, NDArray trainm, NDArray testm,
|
||||
NDArray valm, NDArray glabels, NDArray gtrainm, NDArray gtestm,
|
||||
NDArray gvalm, int64_t Nn) {
|
||||
DType *feat_ptr = feat.Ptr<DType>(); // 2D tensor
|
||||
DType *gfeat_ptr = gfeat.Ptr<DType>(); // 2D tensor
|
||||
int64_t *adj_ptr = adj.Ptr<int64_t>(); // 2D tensor
|
||||
int32_t *inner_node_ptr = inner_node.Ptr<int32_t>();
|
||||
int64_t *ldt_key_ptr = ldt_key.Ptr<int64_t>();
|
||||
int64_t *gdt_key_ptr = gdt_key.Ptr<int64_t>();
|
||||
int64_t *gdt_value_ptr = gdt_value.Ptr<int64_t>(); // 2D tensor
|
||||
int64_t *node_map_ptr = node_map.Ptr<int64_t>();
|
||||
int64_t *lr_ptr = lr.Ptr<int64_t>();
|
||||
int64_t *lrtensor_ptr = lrtensor.Ptr<int64_t>();
|
||||
int32_t width = nc - 1;
|
||||
|
||||
runtime::parallel_for(0, num_nodes, [&](int64_t s, int64_t e) {
|
||||
for (int64_t i = s; i < e; i++) {
|
||||
int64_t k = ldt_key_ptr[i];
|
||||
int64_t v = i;
|
||||
int64_t ind = gdt_key_ptr[k];
|
||||
|
||||
int64_t *adj_ptr_ptr = adj_ptr + v * width;
|
||||
if (ind == 1) {
|
||||
for (int32_t j = 0; j < width; j++) adj_ptr_ptr[j] = -1;
|
||||
inner_node_ptr[i] = 1;
|
||||
lr_ptr[i] = -200;
|
||||
} else {
|
||||
lr_ptr[i] = lrtensor_ptr[k];
|
||||
int64_t *ptr = gdt_value_ptr + k * nc;
|
||||
int64_t pos = 0;
|
||||
CHECK(ind <= nc);
|
||||
int32_t flg = 0;
|
||||
for (int64_t j = 0; j < ind; j++) {
|
||||
if (ptr[j] == lr_ptr[i]) flg = 1;
|
||||
if (c != Ver2partition<int64_t>(ptr[j], node_map_ptr, nc))
|
||||
adj_ptr_ptr[pos++] = ptr[j];
|
||||
}
|
||||
CHECK_EQ(flg, 1);
|
||||
CHECK(pos == ind - 1);
|
||||
for (; pos < width; pos++) adj_ptr_ptr[pos] = -1;
|
||||
inner_node_ptr[i] = 0;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// gather
|
||||
runtime::parallel_for(0, num_nodes, [&](int64_t s, int64_t e) {
|
||||
for (int64_t i = s; i < e; i++) {
|
||||
int64_t k = ldt_key_ptr[i];
|
||||
int64_t ind = i * feat_size;
|
||||
DType *optr = gfeat_ptr + ind;
|
||||
DType *iptr = feat_ptr + k * feat_size;
|
||||
|
||||
for (int32_t j = 0; j < feat_size; j++) optr[j] = iptr[j];
|
||||
}
|
||||
|
||||
IdType *labels_ptr = labels.Ptr<IdType>();
|
||||
IdType *glabels_ptr = glabels.Ptr<IdType>();
|
||||
IdType2 *trainm_ptr = trainm.Ptr<IdType2>();
|
||||
IdType2 *gtrainm_ptr = gtrainm.Ptr<IdType2>();
|
||||
IdType2 *testm_ptr = testm.Ptr<IdType2>();
|
||||
IdType2 *gtestm_ptr = gtestm.Ptr<IdType2>();
|
||||
IdType2 *valm_ptr = valm.Ptr<IdType2>();
|
||||
IdType2 *gvalm_ptr = gvalm.Ptr<IdType2>();
|
||||
|
||||
for (int64_t i = 0; i < num_nodes; i++) {
|
||||
int64_t k = ldt_key_ptr[i];
|
||||
CHECK(k >= 0 && k < Nn);
|
||||
glabels_ptr[i] = labels_ptr[k];
|
||||
gtrainm_ptr[i] = trainm_ptr[k];
|
||||
gtestm_ptr[i] = testm_ptr[k];
|
||||
gvalm_ptr[i] = valm_ptr[k];
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
DGL_REGISTER_GLOBAL("sparse._CAPI_DGLLibra2dglBuildAdjlist")
|
||||
.set_body([](DGLArgs args, DGLRetValue *rv) {
|
||||
NDArray feat = args[0];
|
||||
NDArray gfeat = args[1];
|
||||
NDArray adj = args[2];
|
||||
NDArray inner_node = args[3];
|
||||
NDArray ldt_key = args[4];
|
||||
NDArray gdt_key = args[5];
|
||||
NDArray gdt_value = args[6];
|
||||
NDArray node_map = args[7];
|
||||
NDArray lr = args[8];
|
||||
NDArray lrtensor = args[9];
|
||||
int64_t num_nodes = args[10];
|
||||
int32_t nc = args[11];
|
||||
int32_t c = args[12];
|
||||
int32_t feat_size = args[13];
|
||||
NDArray labels = args[14];
|
||||
NDArray trainm = args[15];
|
||||
NDArray testm = args[16];
|
||||
NDArray valm = args[17];
|
||||
NDArray glabels = args[18];
|
||||
NDArray gtrainm = args[19];
|
||||
NDArray gtestm = args[20];
|
||||
NDArray gvalm = args[21];
|
||||
int64_t Nn = args[22];
|
||||
|
||||
ATEN_FLOAT_TYPE_SWITCH(feat->dtype, DType, "Features", {
|
||||
ATEN_ID_TYPE_SWITCH(trainm->dtype, IdType2, {
|
||||
ATEN_ID_BITS_SWITCH((glabels->dtype).bits, IdType, {
|
||||
Libra2dglBuildAdjlist<IdType, IdType2, DType>(
|
||||
feat, gfeat, adj, inner_node, ldt_key, gdt_key, gdt_value,
|
||||
node_map, lr, lrtensor, num_nodes, nc, c, feat_size, labels,
|
||||
trainm, testm, valm, glabels, gtrainm, gtestm, gvalm, Nn);
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
} // namespace aten
|
||||
} // namespace dgl
|
||||
@@ -0,0 +1,59 @@
|
||||
/**
|
||||
* Copyright (c) 2020 by Contributors
|
||||
* @file array/selector.h
|
||||
* @brief Selector functions to select among src/edge/dst attributes.
|
||||
*/
|
||||
#ifndef DGL_ARRAY_SELECTOR_H_
|
||||
#define DGL_ARRAY_SELECTOR_H_
|
||||
|
||||
#include <dmlc/logging.h>
|
||||
|
||||
namespace dgl {
|
||||
|
||||
namespace {
|
||||
|
||||
#ifdef __CUDACC__
|
||||
#define DGLDEVICE __device__
|
||||
#define DGLINLINE __forceinline__
|
||||
#else
|
||||
#define DGLDEVICE
|
||||
#define DGLINLINE inline
|
||||
#endif // __CUDACC__
|
||||
|
||||
} // namespace
|
||||
|
||||
/**
|
||||
* @brief Select among src/edge/dst feature/idx.
|
||||
* @note the integer argument target specifies which target
|
||||
* to choose, 0: src, 1: edge, 2: dst.
|
||||
*/
|
||||
template <int target>
|
||||
struct Selector {
|
||||
template <typename T>
|
||||
static DGLDEVICE DGLINLINE T Call(T src, T edge, T dst) {
|
||||
LOG(INFO) << "Target " << target << " not recognized.";
|
||||
return src;
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
template <typename T>
|
||||
DGLDEVICE DGLINLINE T Selector<0>::Call(T src, T edge, T dst) {
|
||||
return src;
|
||||
}
|
||||
|
||||
template <>
|
||||
template <typename T>
|
||||
DGLDEVICE DGLINLINE T Selector<1>::Call(T src, T edge, T dst) {
|
||||
return edge;
|
||||
}
|
||||
|
||||
template <>
|
||||
template <typename T>
|
||||
DGLDEVICE DGLINLINE T Selector<2>::Call(T src, T edge, T dst) {
|
||||
return dst;
|
||||
}
|
||||
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_ARRAY_SELECTOR_H_
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user