Task: sph_cell_index
Category: Scientific simulation
Difficulty: 2
Tags: sph,spatial_hashing,sorting,neighbor_search,cell_linked_list
Source: DualSPHysics: https://github.com/DualSPHysics/DualSPHysics
Description: Constructs SPH cell-linked-list index ranges while preserving ordering and tie-breaking behavior.

Prompt template (YAML):
# SPH Cell-Linked List Construction -- Prompt Template

task_description: |
  Construct a cell-linked list data structure for Smoothed Particle Hydrodynamics
  (SPH) neighbor search. Given N particle positions in 3D and a uniform grid:

  1. Compute the cell index for each particle: cell_id = floor(pos/cell_size)
  2. Sort particles by cell index
  3. Build begin/end arrays: for each cell, store the index range [begin, end)
     of particles in the sorted order that belong to that cell

  This is the spatial indexing step from DualSPHysics (Crespo et al. 2015),
  specifically the KerCalcBeginEndCell kernel from JCellDivGpu_ker.cu.

  Based on: DualSPHysics JCellDivGpu_ker.cu, KerCalcBeginEndCell.
  Reference: Crespo et al. "DualSPHysics: Open-source parallel CFD solver based on
  SPH", Computer Physics Communications, 2015.

interface: |
  ```c
  // Single end-to-end call. Do H2D + kernel + D2H and synchronize before returning.
  // Called multiple times (warmup + timed); must be idempotent.
  extern "C" void solution_compute(
      int N,
      const float* xs,
      const float* ys,
      const float* zs,
      float cell_size,
      int grid_nx,
      int grid_ny,
      int grid_nz,
      int num_cells,
      int* sorted_indices,
      int* cell_begin,
      int* cell_end
  );
  ```

input_size_notes: |
  | Size   | N (particles) |
  |--------|--------------|
  | small  | 50,000       |
  | medium | 500,000      |
  | large  | 5,000,000    |

  - Grid dimensions: ~100x50x50 cells (domain 4x2x2, cell_size=0.04)
  - num_cells = grid_nx * grid_ny * grid_nz (~250,000)
  - Cell index: cx + cy*grid_nx + cz*grid_nx*grid_ny
  - Particles are uniformly distributed (each cell has ~N/num_cells particles)

output_contract: |
  - sorted_indices[i]: the original particle index of the i-th particle in cell-sorted order
  - cell_begin[c]: index into sorted_indices where cell c starts (-1 if empty)
  - cell_end[c]: index into sorted_indices where cell c ends (-1 if empty)
  - For non-empty cell c: particles sorted_indices[cell_begin[c]..cell_end[c]-1] all have cell_id == c
  - Correctness mode: exact (integer outputs)
  - solution_compute must be idempotent and synchronize before returning

algorithm_background: |
  **Cell-Linked List for SPH Neighbor Search**

  SPH simulations need to find all particle pairs within a cutoff radius (2h).
  A cell-linked list divides space into uniform cells of size >= 2h, so
  neighbors can only be in the same cell or 26 adjacent cells.

  Construction steps:
  1. **Cell assignment**: For each particle p with position (x,y,z):
     ```
     cx = floor(x / cell_size), cy = floor(y / cell_size), cz = floor(z / cell_size)
     cell_id = cx + cy * grid_nx + cz * grid_nx * grid_ny
     ```

  2. **Sort by cell_id**: Reorder particles so all particles in the same cell
     are contiguous. This is typically done with a radix sort on GPU.

  3. **Build begin/end arrays** (KerCalcBeginEndCell in DualSPHysics):
     Scan the sorted cell_id array. When cell_id changes between positions i-1
     and i, set cell_end[old_cell] = i and cell_begin[new_cell] = i.

  The begin/end arrays allow O(1) lookup of all particles in any cell.

hints_l2: |
  - Step 1 (cell assignment) is embarrassingly parallel: one thread per particle
  - Step 2 (sort) is the bottleneck: use CUB DeviceRadixSort or Thrust sort_by_key
  - Step 3 (begin/end) is parallel: each thread checks if its cell_id differs from predecessor
  - Use shared memory in Step 3 (as DualSPHysics does) for the boundary detection
  - Pre-allocate all device memory inside solution_compute

hints_l1: |
  ### 1. Cell Assignment Kernel
  ```cuda
  __global__ void KerComputeCellId(int N, const float* xs, const float* ys,
      const float* zs, float cell_size, int nx, int ny, unsigned int* cell_ids,
      unsigned int* indices) {
      int i = blockIdx.x * blockDim.x + threadIdx.x;
      if (i >= N) return;
      int cx = (int)floorf(xs[i] / cell_size);
      int cy = (int)floorf(ys[i] / cell_size);
      int cz = (int)floorf(zs[i] / cell_size);
      cell_ids[i] = cx + cy * nx + cz * nx * ny;
      indices[i] = i;
  }
  ```

  ### 2. Sort
  Use CUB DeviceRadixSort::SortPairs to sort (cell_ids, indices) by cell_ids.
  This gives O(N) radix sort performance on GPU.

  ### 3. Begin/End Kernel (matches DualSPHysics KerCalcBeginEndCell)
  ```cuda
  __global__ void KerCalcBeginEndCell(int N, const unsigned int* cellpart,
      int* begcell, int* endcell) {
      extern __shared__ unsigned int scell[];
      int i = blockIdx.x * blockDim.x + threadIdx.x;
      if (i < N) {
          unsigned int cel = cellpart[i];
          scell[threadIdx.x + 1] = cel;
          if (i > 0 && threadIdx.x == 0) scell[0] = cellpart[i - 1];
      }
      __syncthreads();
      if (i < N) {
          unsigned int cel = cellpart[i];
          if (i == 0 || cel != scell[threadIdx.x]) {
              begcell[cel] = i;
              if (i > 0) endcell[scell[threadIdx.x]] = i;
          }
          if (i == N - 1) endcell[cel] = i + 1;
      }
  }
  ```

  ### 4. Launch Configuration
  - 256 threads per block for all kernels
  - Shared memory: (blockDim.x + 1) * sizeof(unsigned int) for begin/end kernel


CPU reference implementation (C):
// cpu_reference.c -- SPH cell-linked list construction (CPU baseline)
//
// Faithfully ported from DualSPHysics JCellDivGpu_ker.cu
//   KerCalcBeginEndCell + sorting logic
//
// For each particle: compute cell_id = floor(pos/cell_size) linearized.
// Sort particles by cell_id. Build begin/end arrays for each cell.
// NO file I/O, NO main(). All I/O handled by task_io_cpu.c.
//
// Reference: Crespo et al. "DualSPHysics: Open-source parallel CFD solver
// based on SPH", Computer Physics Communications, 2015.

#include <math.h>
#include <stdlib.h>
#include <string.h>

// ===== Module-level static state =====
static int    g_N;
static const float* g_xs;
static const float* g_ys;
static const float* g_zs;
static float  g_cell_size;
static int    g_grid_nx;
static int    g_grid_ny;
static int    g_grid_nz;

// ===== Helper: compute linear cell index from position =====
// Matches DualSPHysics cell computation: cellpart[p] = cx + cy*nx + cz*nx*ny

static int ComputeCellIndex(float x, float y, float z,
                            float cell_size, int nx, int ny)
{
    int cx = (int)floorf(x / cell_size);
    int cy = (int)floorf(y / cell_size);
    int cz = (int)floorf(z / cell_size);
    /* Clamp to valid range */
    if (cx < 0) cx = 0;
    if (cy < 0) cy = 0;
    if (cz < 0) cz = 0;
    if (cx >= nx) cx = nx - 1;
    if (cy >= ny) cy = ny - 1;
    return cx + cy * nx + cz * nx * ny;
}

// ===== Sort helper: comparison for qsort =====

typedef struct {
    int cell_id;
    int original_index;
} CellParticlePair;

static int compare_cell_particle(const void* a, const void* b)
{
    const CellParticlePair* pa = (const CellParticlePair*)a;
    const CellParticlePair* pb = (const CellParticlePair*)b;
    if (pa->cell_id < pb->cell_id) return -1;
    if (pa->cell_id > pb->cell_id) return 1;
    /* Stable sort tie-break by original index */
    if (pa->original_index < pb->original_index) return -1;
    if (pa->original_index > pb->original_index) return 1;
    return 0;
}

// ===== Ported from DualSPHysics KerCalcBeginEndCell (CPU version) =====
// Original: __global__ void KerCalcBeginEndCell(unsigned n, unsigned pini,
//   const unsigned* cellpart, int2* begcell)
// CPU adaptation: sequential scan over sorted cell_ids to find begin/end.

static void KerCalcBeginEndCell(int N, const int* sorted_cell_ids,
                                int num_cells, int* cell_begin, int* cell_end)
{
    int i;
    /* Initialize all cells as empty */
    for (i = 0; i < num_cells; i++) {
        cell_begin[i] = -1;
        cell_end[i]   = -1;
    }

    if (N == 0) return;

    /* First particle starts a cell */
    cell_begin[sorted_cell_ids[0]] = 0;

    for (i = 1; i < N; i++) {
        if (sorted_cell_ids[i] != sorted_cell_ids[i - 1]) {
            /* Previous cell ends here */
            cell_end[sorted_cell_ids[i - 1]] = i;
            /* New cell begins here */
            cell_begin[sorted_cell_ids[i]] = i;
        }
    }

    /* Last particle's cell ends at N */
    cell_end[sorted_cell_ids[N - 1]] = N;
}

// ===== Public interface =====

static void _orbench_old_init(int N,
                   const float* xs, const float* ys, const float* zs,
                   float cell_size, int grid_nx, int grid_ny, int grid_nz)
{
    g_N = N;
    g_xs = xs;
    g_ys = ys;
    g_zs = zs;
    g_cell_size = cell_size;
    g_grid_nx = grid_nx;
    g_grid_ny = grid_ny;
    g_grid_nz = grid_nz;
}

static void _orbench_old_compute(int N, int num_cells,
                      int* sorted_indices, int* cell_begin, int* cell_end)
{
    int i;

    /* Step 1: Compute cell index for each particle */
    CellParticlePair* pairs = (CellParticlePair*)malloc(N * sizeof(CellParticlePair));
    for (i = 0; i < N; i++) {
        pairs[i].cell_id = ComputeCellIndex(g_xs[i], g_ys[i], g_zs[i],
                                            g_cell_size, g_grid_nx, g_grid_ny);
        pairs[i].original_index = i;
    }

    /* Step 2: Sort by cell_id (matches DualSPHysics SortDataParticles) */
    qsort(pairs, N, sizeof(CellParticlePair), compare_cell_particle);

    /* Step 3: Extract sorted indices and sorted cell_ids */
    int* sorted_cell_ids = (int*)malloc(N * sizeof(int));
    for (i = 0; i < N; i++) {
        sorted_indices[i] = pairs[i].original_index;
        sorted_cell_ids[i] = pairs[i].cell_id;
    }

    /* Step 4: Build begin/end cell arrays (KerCalcBeginEndCell) */
    KerCalcBeginEndCell(N, sorted_cell_ids, num_cells, cell_begin, cell_end);

    free(sorted_cell_ids);
    free(pairs);
}

// ── Unified compute_only wrapper (auto-migrated) ──
void solution_compute(int N, const float* xs, const float* ys, const float* zs, float cell_size, int grid_nx, int grid_ny, int grid_nz, int num_cells, int* sorted_indices, int* cell_begin, int* cell_end) {
    _orbench_old_init(N, xs, ys, zs, cell_size, grid_nx, grid_ny, grid_nz);
    _orbench_old_compute(N, num_cells, sorted_indices, cell_begin, cell_end);
}