find patient groups

groups.cu

#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#include <thrust/sort.h>

#define N 10

__global__
static void find_groups(int *locs, int *sorted, int num)
{
    int bid = blockIdx.y * gridDim.x + blockIdx.x;
    int tid = bid * blockDim.x + threadIdx.x;

    if (tid < num) {
        int curr = sorted[tid];
        if (tid == 0 || curr != sorted[tid - 1]) locs[curr] = tid;
    }

}

int main()
{
    int h_P0[N] = {0, 0, 1, 2, 1, 1, 0, 2, 0, 0};
    int h_P1[N] = {0, 1, 1, 2, 1, 2, 0, 2, 1, 0};

    thrust::host_vector<int> th_P0(h_P0, h_P0 + N);
    thrust::host_vector<int> th_P1(h_P1, h_P1 + N);

    thrust::device_vector<int> td_P0 = th_P0;
    thrust::device_vector<int> td_P1 = th_P1;

    thrust::device_vector<int> td_S0(N);
    thrust::device_vector<int> td_S1(N);

    thrust::sequence(td_S0.begin(), td_S0.end());
    thrust::sequence(td_S1.begin(), td_S1.end());

    thrust::stable_sort_by_key(td_P0.begin(), td_P0.end(), td_S0.begin());
    thrust::stable_sort_by_key(td_P1.begin(), td_P1.end(), td_S1.begin());

    thrust::device_vector<int> td_l0(3);
    thrust::device_vector<int> td_l1(3);

    int threads =  256;
    int blocks_x = (N + 256) / 256;
    int blocks_y = (blocks_x + 65535) / 65535;
    dim3 blocks(blocks_x, blocks_y);

    int *d_l0 = thrust::raw_pointer_cast(td_l0.data());
    int *d_l1 = thrust::raw_pointer_cast(td_l1.data());
    int *d_P0 = thrust::raw_pointer_cast(td_P0.data());
    int *d_P1 = thrust::raw_pointer_cast(td_P1.data());

    find_groups<<<blocks, threads>>>(d_l0, d_P0, N);
    find_groups<<<blocks, threads>>>(d_l1, d_P1, N);
    printf("Sorted\n");
    thrust::host_vector<int> th_S0 = td_S0;
    thrust::host_vector<int> th_S1 = td_S1;
    printf("t\t t+1\n");
    for (int i = 0;  i < N; i++) {
        printf("%d\t %d\n", th_S0[i], th_S1[i]);
    }

    thrust::host_vector<int> th_l0 = td_l0;
    thrust::host_vector<int> th_l1 = td_l1;

    char groups[4] = "SIR";
    printf("Groups \t t \t t + 1\n");
    for (int i = 0; i < 3; i++) {
        int t0_st = th_l0[i];
        int t0_en = i == 2 ? N : th_l0[i + 1];
        int t1_st = th_l1[i];
        int t1_en = i == 2 ? N  : th_l1[i + 1];
        printf("%c \t [%d-%d]\t [%d-%d]\n",
               groups[i], t0_st, t0_en - 1, t1_st, t1_en - 1);
    }
    return 0;
}