array_index_select.cu 3.87 KB
Newer Older
1
2
3
4
5
6
/*!
 *  Copyright (c) 2019 by Contributors
 * \file array/cpu/array_index_select.cu
 * \brief Array index select GPU implementation
 */
#include <dgl/array.h>
7

8
#include "../../runtime/cuda/cuda_common.h"
9
#include "./array_index_select.cuh"
10
#include "./utils.h"
11
12
13
14
15
16

namespace dgl {
using runtime::NDArray;
namespace aten {
namespace impl {

17
template <DGLDeviceType XPU, typename DType, typename IdType>
18
NDArray IndexSelect(NDArray array, IdArray index) {
19
  cudaStream_t stream = runtime::getCurrentCUDAStream();
20
21
22
23
  const DType* array_data = array.Ptr<DType>();
  if (array.IsPinned()) {
    CUDA_CALL(cudaHostGetDevicePointer(&array_data, array.Ptr<DType>(), 0));
  }
24
25
26
  const IdType* idx_data = static_cast<IdType*>(index->data);
  const int64_t arr_len = array->shape[0];
  const int64_t len = index->shape[0];
27
28
29
30
31
32
33
  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]);
  }

34
  // use index->ctx for pinned array
35
  NDArray ret = NDArray::Empty(shape, array->dtype, index->ctx);
36
  if (len == 0) return ret;
37
  DType* ret_data = static_cast<DType*>(ret->data);
38
39

  if (num_feat == 1) {
40
41
42
43
44
    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);
45
  } else {
46
47
48
49
50
51
52
53
54
    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);
55
  }
56
57
58
  return ret;
}

59
60
61
62
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);
63
#ifdef USE_FP16
64
65
template NDArray IndexSelect<kDGLCUDA, __half, int32_t>(NDArray, IdArray);
template NDArray IndexSelect<kDGLCUDA, __half, int64_t>(NDArray, IdArray);
66
#endif
67
68
69
70
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);
71

72
template <DGLDeviceType XPU, typename DType>
73
DType IndexSelect(NDArray array, int64_t index) {
74
  auto device = runtime::DeviceAPI::Get(array->ctx);
75
76
77
78
79
80
81
82
#ifdef USE_FP16
  // The initialization constructor for __half is apparently a device-
  // only function in some setups, but the current function, IndexSelect,
  // isn't run on the device, so it doesn't have access to that constructor.
  using SafeDType = typename std::conditional<
      std::is_same<DType, __half>::value, uint16_t, DType>::type;
  SafeDType ret = 0;
#else
83
  DType ret = 0;
84
#endif
85
  device->CopyDataFromTo(
86
87
88
      static_cast<DType*>(array->data) + index, 0,
      reinterpret_cast<DType*>(&ret), 0, sizeof(DType), array->ctx,
      DGLContext{kDGLCPU, 0}, array->dtype);
89
  return reinterpret_cast<DType&>(ret);
90
91
}

92
93
94
95
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);
96
#ifdef USE_FP16
97
template __half IndexSelect<kDGLCUDA, __half>(NDArray array, int64_t index);
98
#endif
99
100
template float IndexSelect<kDGLCUDA, float>(NDArray array, int64_t index);
template double IndexSelect<kDGLCUDA, double>(NDArray array, int64_t index);
101
102
103
104

}  // namespace impl
}  // namespace aten
}  // namespace dgl