wehub-resource-sync
75 lines
2.4 KiB
C++
75 lines
2.4 KiB
C++
/**
|
|
* Copyright (c) 2019-2022 by Contributors
|
|
* @file array/uvm_array.cc
|
|
* @brief DGL array utilities implementation
|
|
*/
|
|
#include <dgl/array.h>
|
|
|
|
#include <sstream>
|
|
|
|
#include "../c_api_common.h"
|
|
#include "./uvm_array_op.h"
|
|
|
|
using namespace dgl::runtime;
|
|
|
|
namespace dgl {
|
|
namespace aten {
|
|
|
|
NDArray IndexSelectCPUFromGPU(NDArray array, IdArray index) {
|
|
#ifdef DGL_USE_CUDA
|
|
CHECK(array.IsPinned()) << "Input array must be in pinned memory.";
|
|
CHECK_EQ(index->ctx.device_type, kDGLCUDA) << "Index must be on the GPU.";
|
|
CHECK_GE(array->ndim, 1) << "Input array must have at least 1 dimension.";
|
|
CHECK_EQ(index->ndim, 1) << "Index must be a 1D array.";
|
|
|
|
ATEN_DTYPE_BITS_ONLY_SWITCH(array->dtype, DType, "values", {
|
|
ATEN_ID_TYPE_SWITCH(index->dtype, IdType, {
|
|
return impl::IndexSelectCPUFromGPU<DType, IdType>(array, index);
|
|
});
|
|
});
|
|
#endif
|
|
LOG(FATAL) << "IndexSelectCPUFromGPU requires CUDA.";
|
|
// Should be unreachable
|
|
return NDArray{};
|
|
}
|
|
|
|
void IndexScatterGPUToCPU(NDArray dest, IdArray index, NDArray source) {
|
|
#ifdef DGL_USE_CUDA
|
|
CHECK(dest.IsPinned()) << "Destination array must be in pinned memory.";
|
|
CHECK_EQ(index->ctx.device_type, kDGLCUDA) << "Index must be on the GPU.";
|
|
CHECK_EQ(source->ctx.device_type, kDGLCUDA)
|
|
<< "Source array must be on the GPU.";
|
|
CHECK_EQ(dest->dtype, source->dtype) << "Destination array and source "
|
|
"array must have the same dtype.";
|
|
CHECK_GE(dest->ndim, 1)
|
|
<< "Destination array must have at least 1 dimension.";
|
|
CHECK_EQ(index->ndim, 1) << "Index must be a 1D array.";
|
|
|
|
ATEN_DTYPE_BITS_ONLY_SWITCH(source->dtype, DType, "values", {
|
|
ATEN_ID_TYPE_SWITCH(index->dtype, IdType, {
|
|
impl::IndexScatterGPUToCPU<DType, IdType>(dest, index, source);
|
|
});
|
|
});
|
|
#else
|
|
LOG(FATAL) << "IndexScatterGPUToCPU requires CUDA.";
|
|
#endif
|
|
}
|
|
|
|
DGL_REGISTER_GLOBAL("ndarray.uvm._CAPI_DGLIndexSelectCPUFromGPU")
|
|
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
|
NDArray array = args[0];
|
|
IdArray index = args[1];
|
|
*rv = IndexSelectCPUFromGPU(array, index);
|
|
});
|
|
|
|
DGL_REGISTER_GLOBAL("ndarray.uvm._CAPI_DGLIndexScatterGPUToCPU")
|
|
.set_body([](DGLArgs args, DGLRetValue* rv) {
|
|
NDArray dest = args[0];
|
|
IdArray index = args[1];
|
|
NDArray source = args[2];
|
|
IndexScatterGPUToCPU(dest, index, source);
|
|
});
|
|
|
|
} // namespace aten
|
|
} // namespace dgl
|