Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
array_pack.cc 3.28 KB
Newer Older
Quan (Andy) Gan's avatar
[Feature] Random Walk for 0.5 (#1209)  
Quan (Andy) Gan committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83 
/*!
 *  Copyright (c) 2019 by Contributors
 * \file array/cpu/array_index_select.cc
 * \brief Array index select CPU implementation
 */
#include <dgl/array.h>
#include <tuple>
#include <utility>

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

template<DLDeviceType 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);

#pragma omp parallel for
  for (int64_t i = 0; i < rows; ++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<kDLCPU, int32_t, int32_t>(NDArray, IdArray);
template std::pair<NDArray, IdArray> ConcatSlices<kDLCPU, int64_t, int32_t>(NDArray, IdArray);
template std::pair<NDArray, IdArray> ConcatSlices<kDLCPU, float, int32_t>(NDArray, IdArray);
template std::pair<NDArray, IdArray> ConcatSlices<kDLCPU, double, int32_t>(NDArray, IdArray);
template std::pair<NDArray, IdArray> ConcatSlices<kDLCPU, int32_t, int64_t>(NDArray, IdArray);
template std::pair<NDArray, IdArray> ConcatSlices<kDLCPU, int64_t, int64_t>(NDArray, IdArray);
template std::pair<NDArray, IdArray> ConcatSlices<kDLCPU, float, int64_t>(NDArray, IdArray);
template std::pair<NDArray, IdArray> ConcatSlices<kDLCPU, double, int64_t>(NDArray, IdArray);

template<DLDeviceType 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);
#pragma omp parallel for
  for (int64_t i = 0; i < rows; ++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<kDLCPU, int32_t>(NDArray, int32_t);
template std::tuple<NDArray, IdArray, IdArray> Pack<kDLCPU, int64_t>(NDArray, int64_t);
template std::tuple<NDArray, IdArray, IdArray> Pack<kDLCPU, float>(NDArray, float);
template std::tuple<NDArray, IdArray, IdArray> Pack<kDLCPU, double>(NDArray, double);

};  // namespace impl

};  // namespace aten

};  // namespace dgl