Cuda атомарность и условные переходы

Я пытаюсь написать CUDA версию кода serial как часть реализации периодического граничного условия в алгоритме молекулярной динамики. Идея состоит в том, что крошечную часть частиц, которые имеют нестандартные позиции, нужно вернуть обратно, используя один из двух ways, с ограничением на количество раз, которое я использую первым способом.

По сути, это сводится к следующему MWE. У меня есть массив x[N], где N большой, и следующий код serial.

#include <cstdlib>

int main()
{
  int N =30000;
  double x[30000];
  int Nmax = 10, count = 0;

  for(int i = 0; i < N; i++)
    x[i] = 1.0*(rand()%3);

  for(int i = 0; i < N; i++)
   {
      if(x[i] > 2.9)
        {
          if(count < Nmax)
            {
              x[i] += 0.1; //first way
              count++;
            }
          else
            x[i] -= 0.2; //second way
        }
    }
}

Пожалуйста, предположим, что x[i] > 2.9 только для небольшой части (около 12-15) из 30000 элементов x[i].

Обратите внимание, что последовательность i не важна, т. е. нет необходимости иметь 10 наименьшее i для использования x[i] += 0.1, что делает алгоритм потенциально распараллеливаемым. Я подумал о следующей CUDA версии MWE, которая компилируется с nvcc -arch sm_35 main.cu, где main.cu читается как

#include <cstdlib>

__global__ void PeriodicCondition(double *x, int *N, int *Nmax, int *count)
{
  int i = threadIdx.x+blockIdx.x*blockDim.x;
  if(i < N[0])
    {
      if(x[i] > 2.9)
        {
           if(count[0] < Nmax[0]) //===============(line a)
             {
               x[i] += 0.1; //first way
               atomicAdd(&count[0],1); //========(line b)
             }
           else
             x[i] -= 0.2; //second way
        }
    }
}

int main()
{
  int N = 30000;
  double x[30000];
  int Nmax = 10, count = 0;

  srand(128512);
  for(int i = 0; i < N; i++)
    x[i] = 1.0*(rand()%3);

  double *xD;
  cudaMalloc( (void**) &xD, N*sizeof(double) );
  cudaMemcpy( xD, &x, N*sizeof(double),cudaMemcpyHostToDevice );

  int *countD;
  cudaMalloc( (void**) &countD, sizeof(int) );
  cudaMemcpy( countD, &count, sizeof(int),cudaMemcpyHostToDevice );

  int *ND;
  cudaMalloc( (void**) &ND, sizeof(int) );
  cudaMemcpy( ND, &N, sizeof(int),cudaMemcpyHostToDevice );

  int *NmaxD;
  cudaMalloc( (void**) &NmaxD, sizeof(int) );
  cudaMemcpy( NmaxD, &Nmax, sizeof(int),cudaMemcpyHostToDevice );

  PeriodicCondition<<<938,32>>>(xD, ND, NmaxD, countD);

  cudaFree(NmaxD);
  cudaFree(ND);
  cudaFree(countD);
  cudaFree(xD);

}

Конечно, это неверно, потому что условие if для (line a) использует переменную, обновленную в (line b), которая может быть устаревшей. Это чем-то похоже на флаг изменения Cuda atomics, однако я не уверен, что и как использовать критические разделы помогут.

Есть ли способ убедиться, что count[0] обновлен, когда каждый поток проверяет условие if на (line a), не делая код слишком последовательным?