THCPropose.cuh 1.28 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
#ifndef THC_PROPOSE_INC
#define THC_PROPOSE_INC

rusty1s's avatar
rusty1s committed
4
5
6
#include "common.cuh"

__global__ void proposeKernel(int64_t *color, int64_t *prop, int64_t *row, int64_t *col,
rusty1s's avatar
rusty1s committed
7
                              int64_t *degree, int64_t *cumDegree, ptrdiff_t nNodes) {
rusty1s's avatar
rusty1s committed
8
  KERNEL_LOOP(i, nNodes) {
rusty1s's avatar
rusty1s committed
9
10
    if (color[i] != -1) { continue; }  // Only visit blue nodes.
    ptrdiff_t c; bool isDead = true;
rusty1s's avatar
rusty1s committed
11
    for (ptrdiff_t e = cumDegree[i] - degree[i]; e < cumDegree[i]; e++) {
rusty1s's avatar
rusty1s committed
12
      c = col[e];
rusty1s's avatar
rusty1s committed
13
14
      if (isDead && color[c] < 0) { isDead = false; }  // Unmatched neighbor found.
      if (color[c] == -2) { prop[i] = c; break; }  // Propose to first red neighbor.
rusty1s's avatar
rusty1s committed
15
    }
rusty1s's avatar
rusty1s committed
16
    if (isDead) { color[i] = i; }  // Mark node as dead.
rusty1s's avatar
rusty1s committed
17
18
19
  }
}

rusty1s's avatar
rusty1s committed
20
void THCTensor_propose(THCState *state, THCudaLongTensor *color, THCudaLongTensor *prop,
rusty1s's avatar
rusty1s committed
21
22
23
24
25
                       THCudaLongTensor *row, THCudaLongTensor *col, THCudaLongTensor *degree,
                       THCudaLongTensor *cumDegree) {
  KERNEL_RUN(proposeKernel, THCudaLongTensor_nElement(state, color),
             THCudaLongTensor_data(state, color), THCudaLongTensor_data(state, prop),
             THCudaLongTensor_data(state, row), THCudaLongTensor_data(state, col),
rusty1s's avatar
rusty1s committed
26
             THCudaLongTensor_data(state, degree), THCudaLongTensor_data(state, cumDegree));
rusty1s's avatar
rusty1s committed
27
}
rusty1s's avatar
rusty1s committed
28
29

#endif  // THC_PROPOSE_INC