Construir el índice de la rejilla con un prefix sum

Tienes cuántas partículas hay en cada celda. Lo que necesitas es dónde empieza cada celda dentro de un array donde las partículas están agrupadas por celda. Convertir lo primero en lo segundo es exactamente un scan exclusivo, el algoritmo del nivel anterior, aplicado a un array de dos millones de contadores. Y con el array de inicios, colocar cada partícula es una atómica por partícula. Tres fases, cuatro conceptos que ya conoces, y una estructura que responde “quién está cerca de mí” en tiempo constante.

Del conteo al inicio

Si conteos[c] es el número de partículas de la celda c, entonces el scan exclusivo de ese array da, en la posición c, cuántas partículas hay en todas las celdas anteriores. En un array donde las partículas están ordenadas por celda, ese número es precisamente el índice donde empieza la celda c.

conteos  = [3, 0, 5, 2, 1, 4]
inicios  = [0, 3, 3, 8, 10, 11]     scan exclusivo
finales  = [3, 3, 8, 10, 11, 15]    inicios[c] + conteos[c]

La celda 1 está vacía y su inicio coincide con su final, que es exactamente lo que hace que el bucle de vecinos no itere ninguna vez sobre ella sin necesidad de una comprobación especial.

El final no hace falta guardarlo aparte: es inicios[c + 1], siempre que el array de inicios tenga una posición extra al final con el total. Esa posición extra es gratis y ahorra un array entero de dos millones de posiciones, así que conviene reservarla desde el principio.

El scan es el scan de varios bloques sin ninguna modificación, sobre u32 en vez de f32. Con una rejilla de 128 al cubo son 2.097.152 elementos, que con bloques de 512 necesitan tres niveles: 4096 bloques, luego 8, luego 1. Son seis dispatches de scan, y su coste conjunto en una GPU de escritorio ronda los 200 microsegundos.

Ahí hay una decisión de diseño que a menudo se pasa por alto: el coste del scan depende del número de celdas, no del número de partículas. Con cien mil partículas y una rejilla de 128 al cubo, el scan de dos millones de celdas cuesta más que todo lo demás junto. Si tu número de partículas es bajo, una rejilla más gruesa —o una tabla hash con una fracción de las entradas— sale mucho mejor.

La dispersión

Con los inicios calculados, cada partícula necesita saber en qué posición dentro de su celda va. Eso se resuelve con un cursor atómico por celda, inicializado a cero:

@group(0) @binding(0) var<uniform> p: Params;
@group(0) @binding(1) var<storage, read>       celdaDe:  array<u32>;   // del conteo
@group(0) @binding(2) var<storage, read>       inicios:  array<u32>;   // del scan
@group(0) @binding(3) var<storage, read_write> cursor:   array<atomic<u32>>;
@group(0) @binding(4) var<storage, read_write> ordenado: array<u32>;   // indices

@compute @workgroup_size(64)
fn dispersar(@builtin(global_invocation_id) gid: vec3u) {
  let i = gid.x;
  if (i >= p.conteo) { return; }

  let c = celdaDe[i];
  let rango = atomicAdd(&cursor[c], 1u);
  ordenado[inicios[c] + rango] = i;
}

El array cursor se pone a cero con clearBuffer antes del dispatch. Y hay un truco que ahorra un array de dos millones de posiciones: reutilizar el array de conteos como cursor, ya que después del scan su contenido no hace falta. Se limpia y se usa.

El resultado, ordenado, es un array de índices de partícula agrupados por celda. La celda c ocupa el tramo [inicios[c], inicios[c+1]), y recorrerla es un bucle sobre ese tramo.

// Recorrer las particulas de una celda.
let ini = inicios[c];
let fin = inicios[c + 1u];
for (var k: u32 = ini; k < fin; k = k + 1u) {
  let j = ordenado[k];            // indice de la particula vecina
  // ...
}

La cadena completa

Reunido todo, la construcción de la rejilla es esta secuencia, encolada en un único compute pass:

Ocho dispatches y dos limpiezas para un millón de partículas y dos millones de celdas. En una GPU de gama media eso ronda el medio milisegundo, que dentro de un presupuesto de dieciséis es perfectamente asumible y deja el grueso del frame para la interacción, que es lo caro.

Fíjate en que las fases 1, 2 y 4 tienen que ser dispatches separados. La 2 necesita todos los conteos y la 4 necesita todos los inicios, y esas dependencias cruzan la frontera entre workgroups. No es una elección de estilo: es la única estructura posible.

Las dos limpiezas se pueden encolar antes de abrir el compute pass, porque clearBuffer es una operación del encoder y no de un pass. El orden relativo respecto a los dispatches lo garantiza el encoder.

El orden dentro de la celda

El atomicAdd del cursor reparte ranuras en el orden en que las invocaciones llegan a la atómica, que depende de la planificación. Dos ejecuciones con las mismas posiciones producen la misma agrupación por celda pero un orden distinto dentro de cada celda.

Para una búsqueda de vecinos eso da exactamente igual: la suma de las fuerzas de los vecinos es la misma sea cual sea el orden en que se recorren, salvo por el último bit de la suma en coma flotante. Y si esa diferencia de un bit se realimenta durante diez mil pasos, dos ejecuciones divergen visiblemente. Es el mismo fenómeno del que hablábamos con las atómicas de flotante, aquí por la puerta de atrás.

Si necesitas reproducibilidad exacta —para un test de regresión, para grabar y reproducir, para una simulación que tiene que dar el mismo resultado en dos máquinas— hay que sustituir el cursor atómico por un ranking estable, exactamente el que usa la fase de dispersión de un radix sort: ordenar el bloque localmente por celda con particiones de bits y calcular el rango dentro del bloque con un scan local. Cuesta bastante más y solo hay que pagarlo cuando de verdad hace falta.

Una alternativa intermedia que a veces basta: usar el índice de partícula como criterio de desempate en un segundo pase que ordene cada celda. Si las celdas tienen pocas decenas de elementos, ordenarlas con una red pequeña en memoria compartida es barato. Pero antes de llegar ahí conviene preguntarse si el problema real es el orden o si es que la simulación es caótica de todas formas, cosa que un fluido lo es por naturaleza.