test_partition.cc 6.66 KB
Newer Older
1
#include <gtest/gtest.h>
2

3
#include "../../src/partition/ndarray_partition.h"
4
#include "./common.h"
5
6
7
8

using namespace dgl;
using namespace dgl::partition;

9
template <DGLDeviceType XPU, typename IdType>
10
void _TestRemainder_GeneratePermutation() {
11
12
  const int64_t size = 160000;
  const int num_parts = 7;
13
  NDArrayPartitionRef part = CreatePartitionRemainderBased(size, num_parts);
14

15
16
  IdArray idxs =
      aten::Range(0, size / 10, sizeof(IdType) * 8, DGLContext{XPU, 0});
17

18
  std::pair<IdArray, IdArray> result = part->GeneratePermutation(idxs);
19
20

  // first part of result should be the permutation
21
  IdArray perm = result.first.CopyTo(DGLContext{kDGLCPU, 0});
22
23
  ASSERT_TRUE(perm.Ptr<IdType>() != nullptr);
  ASSERT_EQ(perm->shape[0], idxs->shape[0]);
24
  const IdType* const perm_cpu = static_cast<const IdType*>(perm->data);
25
26

  // second part of result should be the counts
27
  IdArray counts = result.second.CopyTo(DGLContext{kDGLCPU, 0});
28
29
  ASSERT_TRUE(counts.Ptr<int64_t>() != nullptr);
  ASSERT_EQ(counts->shape[0], num_parts);
30
  const int64_t* const counts_cpu = static_cast<const int64_t*>(counts->data);
31

32
  std::vector<int64_t> prefix(num_parts + 1, 0);
33
  for (int p = 0; p < num_parts; ++p) {
34
    prefix[p + 1] = prefix[p] + counts_cpu[p];
35
36
37
38
  }
  ASSERT_EQ(prefix.back(), idxs->shape[0]);

  // copy original indexes to cpu
39
  idxs = idxs.CopyTo(DGLContext{kDGLCPU, 0});
40
  const IdType* const idxs_cpu = static_cast<const IdType*>(idxs->data);
41
42

  for (int p = 0; p < num_parts; ++p) {
43
    for (int64_t i = prefix[p]; i < prefix[p + 1]; ++i) {
44
45
46
47
48
      EXPECT_EQ(idxs_cpu[perm_cpu[i]] % num_parts, p);
    }
  }
}

49
template <DGLDeviceType XPU, typename IdType>
50
51
52
void _TestRemainder_MapToX() {
  const int64_t size = 160000;
  const int num_parts = 7;
53
  NDArrayPartitionRef part = CreatePartitionRemainderBased(size, num_parts);
54
55

  for (int part_id = 0; part_id < num_parts; ++part_id) {
56
57
    IdArray local = aten::Range(
        0, part->PartSize(part_id), sizeof(IdType) * 8, DGLContext{XPU, 0});
58
59
60
61
62
63
    IdArray global = part->MapToGlobal(local, part_id);
    IdArray act_local = part->MapToLocal(global).CopyTo(CPU);

    // every global index should have the same remainder as the part id
    ASSERT_EQ(global->shape[0], local->shape[0]);
    global = global.CopyTo(CPU);
64
    for (int64_t i = 0; i < global->shape[0]; ++i) {
65
66
67
      EXPECT_EQ(Ptr<IdType>(global)[i] % num_parts, part_id)
          << "i=" << i << ", num_parts=" << num_parts
          << ", part_id=" << part_id;
68
69
70
71
72
    }

    // the remapped local indices to should match the original
    local = local.CopyTo(CPU);
    ASSERT_EQ(local->shape[0], act_local->shape[0]);
73
    for (int64_t i = 0; i < act_local->shape[0]; ++i) {
74
75
76
77
78
      EXPECT_EQ(Ptr<IdType>(local)[i], Ptr<IdType>(act_local)[i]);
    }
  }
}

79
80
TEST(PartitionTest, TestRemainderPartition) {
#ifdef DGL_USE_CUDA
81
82
  _TestRemainder_GeneratePermutation<kDGLCUDA, int32_t>();
  _TestRemainder_GeneratePermutation<kDGLCUDA, int64_t>();
83

84
85
  _TestRemainder_MapToX<kDGLCUDA, int32_t>();
  _TestRemainder_MapToX<kDGLCUDA, int64_t>();
86
87
88
89
#endif
  // CPU is not implemented
}

90
91
template <typename INDEX, typename RANGE>
int _FindPart(const INDEX idx, const RANGE* const range, const int num_parts) {
92
  for (int i = 0; i < num_parts; ++i) {
93
    if (range[i + 1] > idx) {
94
95
96
97
98
99
100
      return i;
    }
  }

  return -1;
}

101
template <DGLDeviceType XPU, typename IdType>
102
103
104
void _TestRange_GeneratePermutation() {
  const int64_t size = 160000;
  const int num_parts = 7;
105
106
  IdArray range = aten::NewIdArray(
      num_parts + 1, DGLContext{kDGLCPU, 0}, sizeof(IdType) * 8);
107
  for (int i = 0; i < num_parts; ++i) {
108
    range.Ptr<IdType>()[i] = (size / num_parts) * i;
109
110
111
  }
  range.Ptr<IdType>()[num_parts] = size;
  NDArrayPartitionRef part = CreatePartitionRangeBased(
112
      size, num_parts, range.CopyTo(DGLContext{XPU, 0}));
113

114
115
  IdArray idxs =
      aten::Range(0, size / 10, sizeof(IdType) * 8, DGLContext{XPU, 0});
116
117
118
119

  std::pair<IdArray, IdArray> result = part->GeneratePermutation(idxs);

  // first part of result should be the permutation
120
  IdArray perm = result.first.CopyTo(DGLContext{kDGLCPU, 0});
121
122
  ASSERT_TRUE(perm.Ptr<IdType>() != nullptr);
  ASSERT_EQ(perm->shape[0], idxs->shape[0]);
123
  const IdType* const perm_cpu = static_cast<const IdType*>(perm->data);
124
125

  // second part of result should be the counts
126
  IdArray counts = result.second.CopyTo(DGLContext{kDGLCPU, 0});
127
128
  ASSERT_TRUE(counts.Ptr<int64_t>() != nullptr);
  ASSERT_EQ(counts->shape[0], num_parts);
129
  const int64_t* const counts_cpu = static_cast<const int64_t*>(counts->data);
130

131
  std::vector<int64_t> prefix(num_parts + 1, 0);
132
  for (int p = 0; p < num_parts; ++p) {
133
    prefix[p + 1] = prefix[p] + counts_cpu[p];
134
135
136
137
  }
  ASSERT_EQ(prefix.back(), idxs->shape[0]);

  // copy original indexes to cpu
138
  idxs = idxs.CopyTo(DGLContext{kDGLCPU, 0});
139
  const IdType* const idxs_cpu = static_cast<const IdType*>(idxs->data);
140
141

  for (int p = 0; p < num_parts; ++p) {
142
143
144
    for (int64_t i = prefix[p]; i < prefix[p + 1]; ++i) {
      EXPECT_EQ(
          _FindPart(idxs_cpu[perm_cpu[i]], range.Ptr<IdType>(), num_parts), p);
145
146
147
148
    }
  }
}

149
template <DGLDeviceType XPU, typename IdType>
150
151
152
void _TestRange_MapToX() {
  const int64_t size = 160000;
  const int num_parts = 7;
153
154
  IdArray range = aten::NewIdArray(
      num_parts + 1, DGLContext{kDGLCPU, 0}, sizeof(IdType) * 8);
155
  for (int i = 0; i < num_parts; ++i) {
156
    Ptr<IdType>(range)[i] = (size / num_parts) * i;
157
158
159
  }
  range.Ptr<IdType>()[num_parts] = size;
  NDArrayPartitionRef part = CreatePartitionRangeBased(
160
      size, num_parts, range.CopyTo(DGLContext{XPU, 0}));
161
162

  for (int part_id = 0; part_id < num_parts; ++part_id) {
163
164
    IdArray local = aten::Range(
        0, part->PartSize(part_id), sizeof(IdType) * 8, DGLContext{XPU, 0});
165
166
167
168
169
    IdArray global = part->MapToGlobal(local, part_id);
    IdArray act_local = part->MapToLocal(global).CopyTo(CPU);

    ASSERT_EQ(global->shape[0], local->shape[0]);
    global = global.CopyTo(CPU);
170
    for (int64_t i = 0; i < global->shape[0]; ++i) {
171
172
173
174
175
      EXPECT_EQ(
          _FindPart(Ptr<IdType>(global)[i], Ptr<IdType>(range), num_parts),
          part_id)
          << "i=" << i << ", num_parts=" << num_parts << ", part_id=" << part_id
          << ", shape=" << global->shape[0];
176
177
178
179
180
    }

    // the remapped local indices to should match the original
    local = local.CopyTo(CPU);
    ASSERT_EQ(local->shape[0], act_local->shape[0]);
181
    for (int64_t i = 0; i < act_local->shape[0]; ++i) {
182
183
184
185
186
187
188
      EXPECT_EQ(Ptr<IdType>(local)[i], Ptr<IdType>(act_local)[i]);
    }
  }
}

TEST(PartitionTest, TestRangePartition) {
#ifdef DGL_USE_CUDA
189
190
  _TestRange_GeneratePermutation<kDGLCUDA, int32_t>();
  _TestRange_GeneratePermutation<kDGLCUDA, int64_t>();
191

192
193
  _TestRange_MapToX<kDGLCUDA, int32_t>();
  _TestRange_MapToX<kDGLCUDA, int64_t>();
194
195
196
#endif
  // CPU is not implemented
}