array_index_select.cc 2.09 KB
Newer Older
1
/**
2
 *  Copyright (c) 2019 by Contributors
3
4
 * @file array/cpu/array_index_select.cc
 * @brief Array index select CPU implementation
5
6
7
8
9
10
11
12
 */
#include <dgl/array.h>

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

13
template <DGLDeviceType XPU, typename DType, typename IdType>
14
NDArray IndexSelect(NDArray array, IdArray index) {
15
16
17
  CHECK_EQ(array->shape[0], array.NumElements())
      << "Only support tensor"
      << " whose first dimension equals number of elements, e.g. (5,), (5, 1)";
18

19
20
21
22
23
24
25
26
27
28
29
30
31
  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;
}

32
33
34
35
36
37
38
39
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);
40

41
template <DGLDeviceType XPU, typename DType>
42
DType IndexSelect(NDArray array, int64_t index) {
43
44
45
46
  const DType* data = static_cast<DType*>(array->data);
  return data[index];
}

47
48
49
50
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);
51

52
53
54
}  // namespace impl
}  // namespace aten
}  // namespace dgl