array_index_select.cc

/**
 *  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>

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);
  runtime::parallel_for(
      0,
      len,
      1000,  // Thread scheduling overhead is bigger with tiny grain size.
      [idx_data, arr_len, ret_data, array_data] (size_t begin, size_t end) {
        for (size_t i = begin; i < end; ++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