An OpenCL kernel for neighbor search

sph_neighbor.cl

#define NNMAX 16
float R2(float4 p)
{
  return p.x*p.x + p.y*p.y + p.z*p.z;
}

__kernel void sph_neighbor(__global float4 *pos, 
			   __global float *size,
			   __global int *next,
			   __global int *more, 
			   __global int *res,
			   int n)
{
  unsigned int gid = get_global_id(0);
  if (gid >= n) return;

  float4 pi = pos[gid];  // x,  y,  z, mj
  float  hi = size[gid];
  float  hihi = hi*hi;

  int list[NNMAX];
  int nn = 0;
  int full = 0;
  for(int i = 0; i < NNMAX; i++) list[i] = -1;

  int cur = n;
  while(cur != -1) {
    float4 pj = pos[cur];
    float4 dx = pi - pj;
    float  r2 = R2(dx);
    float  hj = size[cur];

    if (cur < n) {
      // particle
      if (r2 <= hihi && full == 0) {
	list[nn++] = cur;
       	if (nn == NNMAX) full = 1;
      }
      cur = next[cur];	
    } else {
      // cell
      float dum = hi + hj;
      if (r2 > dum*dum) {
	cur = next[cur];	
      } else {
	cur = more[cur];
      }
    }
  }

  if (full == 1) list[0] = -2;

  for(int i = 0; i < NNMAX; i++) {
    res[gid*NNMAX + i] = list[i];
  }
}