El estado de las partículas en storage buffers

Un sistema de partículas en GPU es, antes que nada, una decisión sobre la disposición de la memoria. El shader que integra el movimiento se ejecuta un millón de veces por frame y está limitado por ancho de banda de principio a fin: la aritmética de una integración son quince instrucciones y las lecturas son treinta y dos bytes. Elegir bien cómo están dispuestos esos bytes cambia el rendimiento por un factor que la aritmética no puede compensar de ninguna manera.

La estructura y su alineación

Las reglas de alineación de WGSL castigan la ingenuidad. Un vec3f tiene tamaño 12 y alineación 16, así que dos vec3f seguidos ocupan 32 bytes con 8 desperdiciados. Colocar un f32 detrás de cada vec3f rellena el hueco y no cuesta nada:

struct Particula {
  pos:  vec3f,   // offset  0, ocupa 12
  vida: f32,     // offset 12, rellena el hueco de alineacion
  vel:  vec3f,   // offset 16, ocupa 12
  masa: f32,     // offset 28
};                // tamano 32, alineacion 16

Treinta y dos bytes por partícula, sin un solo byte perdido. La misma estructura con vida y masa al final ocuparía 48, un 50 por ciento más de tráfico para los mismos datos.

Del lado de JavaScript, la escritura tiene que respetar exactamente esa disposición. Una vista Float32Array sobre un ArrayBuffer con paso de 8 flotantes hace el trabajo:

const N = 1_000_000;
const FLOTANTES = 8;                       // 32 bytes
const inicial = new Float32Array(N * FLOTANTES);
for (let i = 0; i < N; i++) {
  const o = i * FLOTANTES;
  inicial[o + 0] = (Math.random() - 0.5) * 10;   // pos.x
  inicial[o + 1] = (Math.random() - 0.5) * 10;   // pos.y
  inicial[o + 2] = (Math.random() - 0.5) * 10;   // pos.z
  inicial[o + 3] = 1 + Math.random() * 3;        // vida
  inicial[o + 4] = 0;                            // vel.x
  inicial[o + 5] = 0;                            // vel.y
  inicial[o + 6] = 0;                            // vel.z
  inicial[o + 7] = 0.5 + Math.random();          // masa
}

const bufParticulas = device.createBuffer({
  size: inicial.byteLength,
  usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
  mappedAtCreation: true,
});
new Float32Array(bufParticulas.getMappedRange()).set(inicial);
bufParticulas.unmap();

mappedAtCreation evita una copia intermedia: el buffer se crea ya mapeado en memoria accesible y se escribe directamente. Para 32 megabytes de estado inicial la diferencia con writeBuffer es apreciable, y en el arranque cuenta.

Array de estructuras o estructura de arrays

La disposición anterior es un array de estructuras: cada partícula ocupa 32 bytes contiguos. La alternativa es la estructura de arrays: un buffer con todas las posiciones, otro con todas las velocidades.

// Estructura de arrays.
@group(0) @binding(1) var<storage, read_write> pos: array<vec4f>;
@group(0) @binding(2) var<storage, read_write> vel: array<vec4f>;

La diferencia se mide contando bytes útiles por transacción de memoria. Un warp de 32 invocaciones que lee p[i].pos en la disposición de estructuras toca 32 posiciones separadas por 32 bytes: el hardware trae 1024 bytes y usa 384. Con la estructura de arrays, las 32 lecturas de pos[i] son 512 bytes contiguos y se usan todos.

De ahí sale una regla que no admite atajos: el array de estructuras gana cuando el kernel toca toda la estructura, la estructura de arrays gana cuando toca un subconjunto.

El integrador toca posición, velocidad y masa: array de estructuras. El shader de render solo necesita posición y color: estructura de arrays. Como los dos existen en el mismo sistema, la solución habitual es híbrida, separando por frecuencia de uso más que por campo:

// Caliente: lo que cambia cada frame y lee el integrador.
struct Cinematica { pos: vec3f, vida: f32, vel: vec3f, masa: f32 };
@group(0) @binding(1) var<storage, read_write> cine: array<Cinematica>;

// Frio: lo que se escribe al nacer y solo lee el render.
struct Aspecto { color: vec4f, tam: f32, semilla: f32, _r: vec2f };
@group(0) @binding(2) var<storage, read> aspecto: array<Aspecto>;

El integrador lee y escribe 32 bytes por partícula y no toca el aspecto en absoluto. El render lee 32 bytes de aspecto y 16 de cinemática. Ninguno de los dos arrastra datos que no usa.

Una precaución con vec3f en arrays de storage: array<vec3f> tiene paso 16, no 12, porque el elemento hereda la alineación del tipo. Si escribes desde JavaScript un array empaquetado de tres flotantes por elemento, no coincidirá. O usas vec4f y aceptas el cuarto componente como carga útil, o usas tres arrays de f32.

Doble buffer o en sitio

La pregunta es sencilla de contestar y casi todo el mundo la responde mal por exceso de prudencia: hace falta doble buffer si y solo si una invocación lee el estado de otra partícula.

Un integrador de partículas independientes —gravedad, arrastre, un campo de fuerzas que solo depende de la posición propia— lee y escribe únicamente su propia partícula. No hay carrera posible, porque no hay dos invocaciones tocando la misma dirección. Se actualiza en sitio, con un solo buffer.

// En sitio: correcto para particulas independientes.
@compute @workgroup_size(64)
fn integrar(@builtin(global_invocation_id) gid: vec3u) {
  let i = gid.x;
  if (i >= p.conteo) { return; }
  var q = cine[i];                       // copia local
  q.vel = q.vel + aceleracion(q) * p.dt;
  q.pos = q.pos + q.vel * p.dt;
  cine[i] = q;
}

Un integrador con interacción —repulsión entre vecinos, atracción n-cuerpos, cualquier cosa que lea cine[j] con j distinto de i— sí necesita doble buffer, porque si no, unas invocaciones leerían el estado nuevo de sus vecinos y otras el viejo, según el orden de planificación. El resultado no sería simplemente impreciso: sería distinto en cada ejecución.

// Ping-pong: dos buffers y dos bind groups que se alternan.
const uso = GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST;
const bufA = device.createBuffer({ size: bytes, usage: uso });
const bufB = device.createBuffer({ size: bytes, usage: uso });

const bgAB = device.createBindGroup({ layout, entries: [
  { binding: 0, resource: { buffer: uniforms } },
  { binding: 1, resource: { buffer: bufA } },   // lee
  { binding: 2, resource: { buffer: bufB } },   // escribe
]});
const bgBA = device.createBindGroup({ layout, entries: [
  { binding: 0, resource: { buffer: uniforms } },
  { binding: 1, resource: { buffer: bufB } },
  { binding: 2, resource: { buffer: bufA } },
]});

let paridad = 0;
function frame() {
  const pass = enc.beginComputePass();
  pass.setPipeline(pipeIntegrar);
  pass.setBindGroup(0, paridad === 0 ? bgAB : bgBA);
  pass.dispatchWorkgroups(Math.ceil(N / 64));
  pass.end();
  paridad ^= 1;
}

El coste del doble buffer no es solo la memoria: es que cada frame lee un array completo y escribe otro, en vez de leer y escribir el mismo, con lo que la caché no ayuda nada. Con 32 megabytes de estado son 64 megabytes de tráfico por frame en cualquier caso, pero la versión en sitio tiene mucha más probabilidad de encontrar el dato caliente.

Cuántas partículas caben

Dos límites acotan el tamaño. maxStorageBufferBindingSize vale 134.217.728 bytes por defecto, o sea 128 MiB, y maxBufferSize vale 268.435.456, o sea 256 MiB. El que manda es el primero, porque es el que limita lo que se puede enlazar de una vez.

Con 32 bytes por partícula, 128 MiB dan 4.194.304 partículas en un solo binding. Con la disposición híbrida, donde cinemática y aspecto van en buffers distintos, el mismo límite se aplica a cada uno por separado y el techo efectivo sube.

Muchos dispositivos ofrecen más y se puede pedir, con la disciplina habitual de leer del adaptador y quedarse con el mínimo:

const pedido = Math.min(512 * 1024 * 1024, adapter.limits.maxStorageBufferBindingSize);
const device = await adapter.requestDevice({
  requiredLimits: { maxStorageBufferBindingSize: pedido,
                    maxBufferSize: Math.min(pedido, adapter.limits.maxBufferSize) },
});

Si necesitas más de lo que da un binding, la salida es trocear el array en varios buffers y lanzar un dispatch por trozo, con el desplazamiento en un uniform. Es incómodo pero no cuesta rendimiento, porque los dispatches son secuenciales y cada uno satura el chip igual.