Compactar los visibles: contador atómico o scan

Al salir del culling tienes un array<u32> de cien mil posiciones donde la mayoría son ceros. Eso no se puede dibujar: una llamada indirecta necesita un número de instancias y una lista contigua de índices, no un array disperso con huecos. Comprimir ese array es un problema clásico del cómputo paralelo llamado stream compaction, tiene exactamente dos soluciones prácticas, y la diferencia entre ellas no es de rendimiento sino de si tu fotograma es reproducible.

De flags a lista

El array que dejaron los pases de culling tiene esta pinta:

indice:  0  1  2  3  4  5  6  7  8  9
flag:    1  0  0  1  1  0  1  0  0  1

Y lo que necesita el vertex shader es la otra cosa:

posicion: 0  1  2  3  4
indice:   0  3  4  6  9      instanceCount = 5

La razón es directa. Una llamada de dibujo instanciada asigna a cada instancia un @builtin(instance_index) que va de cero a instanceCount - 1, contiguo y sin huecos. Ese número es lo único que el shader recibe para saber qué está dibujando. Si le pasaras el array de flags, la instancia número 1 no tendría forma de averiguar que le corresponde el objeto 3 sin recorrer el array desde el principio, que es precisamente el trabajo secuencial que estamos intentando evitar.

Con la lista densa, el vertex shader hace una indirección de un solo acceso: let objeto = visibles[instance_index];. Y el número de elementos de esa lista es, literalmente, el instanceCount que hay que escribir en el buffer indirecto.

El problema es que las cien mil invocaciones que escriben la lista tienen que ponerse de acuerdo sobre quién ocupa qué posición, y cada una solo conoce su propio flag. Ahí está toda la dificultad.

El contador atómico

La solución obvia y la más rápida de escribir: una variable compartida que empieza en cero, y cada invocación con flag a uno reserva un hueco incrementándola.

@group(0) @binding(0) var<storage, read>       visible  : array<u32>;
@group(0) @binding(1) var<storage, read_write> indices  : array<u32>;
@group(0) @binding(2) var<storage, read_write> contador : atomic<u32>;

@compute @workgroup_size(64)
fn compactar(@builtin(global_invocation_id) gid : vec3<u32>) {
  let i = gid.x;
  if (i >= arrayLength(&visible)) {
    return;
  }
  if (visible[i] == 0u) {
    return;
  }
  // atomicAdd devuelve el valor ANTERIOR al incremento: ese es tu hueco.
  let hueco = atomicAdd(&contador, 1u);
  indices[hueco] = i;
}

Diez líneas y funciona. atomicAdd garantiza que dos invocaciones nunca reciben el mismo valor de retorno, por muchas que compitan y en cualquier orden que lleguen. El contador hay que ponerlo a cero al principio de cada fotograma, con un queue.writeBuffer de cuatro bytes o con un kernel trivial de una invocación.

En teoría esto debería ser un desastre de contención: cien mil invocaciones peleando por la misma dirección de memoria. En la práctica va sorprendentemente bien, porque las GPU modernas resuelven los atómicos en la caché de nivel dos sin subir a memoria y porque el compilador suele agregar las operaciones dentro de un grupo SIMD antes de emitir una sola operación al exterior. Aun así, cuando la fracción de visibles es alta merece la pena hacer esa agregación tú mismo y de forma explícita, con un contador por workgroup:

var<workgroup> localCont : atomic<u32>;
var<workgroup> baseGlobal : u32;

@compute @workgroup_size(64)
fn compactarAgregado(
  @builtin(global_invocation_id)  gid : vec3<u32>,
  @builtin(local_invocation_index) li : u32,
) {
  if (li == 0u) {
    atomicStore(&localCont, 0u);
  }
  workgroupBarrier();

  let i = gid.x;
  let pasa = i < arrayLength(&visible) && visible[i] == 1u;

  var localIdx = 0u;
  if (pasa) {
    localIdx = atomicAdd(&localCont, 1u);   // atomico en memoria compartida
  }
  workgroupBarrier();

  // Una sola operacion atomica global por workgroup en vez de una por invocacion.
  if (li == 0u) {
    baseGlobal = atomicAdd(&contador, atomicLoad(&localCont));
  }
  workgroupBarrier();

  if (pasa) {
    indices[baseGlobal + localIdx] = i;
  }
}

De 64 operaciones atómicas globales por workgroup a una. Los atómicos en var<workgroup> viven en la memoria compartida de la unidad de cómputo, que es dos órdenes de magnitud más rápida que la global.

Lo que este código no te da es el orden. Los workgroups se planifican en el orden que decide el hardware, que depende de qué unidades quedan libres, que depende de lo que estuviera ejecutándose antes. Dos fotogramas con la misma escena, la misma cámara y los mismos flags producen listas con los mismos elementos en distinto orden.

El prefix sum exclusivo

La alternativa da el mismo resultado con el orden original preservado. La idea: si supieras, para cada posición i, cuántos visibles hay antes de ella, ese número sería su hueco. Y esa cantidad es exactamente el prefix sum exclusivo del array de flags.

flag:      1  0  0  1  1  0  1  0  0  1
exclusivo: 0  1  1  1  2  3  3  4  4  4      total = 5

Cada invocación lee su flag y su prefijo y escribe sin hablar con nadie. El último elemento del scan más el último flag es el total.

Calcularlo en paralelo es el problema no trivial. Para un bloque que cabe en un workgroup, el algoritmo de Hillis-Steele lo resuelve en log2(n) pasos usando memoria compartida:

const TAM : u32 = 256u;

var<workgroup> tmp : array<u32, TAM>;   // 256 * 4 = 1024 bytes

@group(0) @binding(0) var<storage, read>       visible : array<u32>;
@group(0) @binding(1) var<storage, read_write> prefijo : array<u32>;
@group(0) @binding(2) var<storage, read_write> sumas   : array<u32>;

@compute @workgroup_size(256)
fn scanBloque(
  @builtin(global_invocation_id)   gid : vec3<u32>,
  @builtin(local_invocation_index) li  : u32,
  @builtin(workgroup_id)           wid : vec3<u32>,
) {
  let n = arrayLength(&visible);
  let i = gid.x;

  var v = 0u;
  if (i < n) {
    v = visible[i];
  }
  tmp[li] = v;
  workgroupBarrier();

  // Hillis-Steele: scan inclusivo en log2(256) = 8 pasos.
  // Leer en registro, barrera, escribir, barrera. Las barreras estan fuera
  // del if porque workgroupBarrier exige flujo de control uniforme.
  for (var paso = 1u; paso < TAM; paso = paso << 1u) {
    var suma = tmp[li];
    if (li >= paso) {
      suma = suma + tmp[li - paso];
    }
    workgroupBarrier();
    tmp[li] = suma;
    workgroupBarrier();
  }

  // De inclusivo a exclusivo: restar el propio valor.
  if (i < n) {
    prefijo[i] = tmp[li] - v;
  }

  // La ultima invocacion publica la suma total del bloque.
  if (li == TAM - 1u) {
    sumas[wid.x] = tmp[li];
  }
}

Eso resuelve 256 elementos. Para cien mil hacen falta tres pasos, y el patrón es el mismo que verías en cualquier implementación seria de scan:

Paso 1. scanBloque sobre los 100.000 elementos con 391 workgroups. Deja prefijo correcto dentro de cada bloque y un array sumas de 391 totales de bloque.

Paso 2. Un scan exclusivo sobre esos 391 valores. Como 391 es mayor que 256, tampoco cabe en un bloque, así que se aplica el mismo esquema recursivamente: dos bloques y un scan de dos elementos que ya es trivial. En la práctica se implementa con la misma tubería llamada dos veces. El resultado es un array de 391 desplazamientos: cuántos visibles hay en todos los bloques anteriores al bloque b.

Paso 3. Sumar a cada elemento el desplazamiento de su bloque.

@group(0) @binding(0) var<storage, read_write> prefijo        : array<u32>;
@group(0) @binding(1) var<storage, read>       desplazamientos : array<u32>;

@compute @workgroup_size(256)
fn sumarDesplazamiento(
  @builtin(global_invocation_id) gid : vec3<u32>,
  @builtin(workgroup_id)         wid : vec3<u32>,
) {
  let i = gid.x;
  if (i >= arrayLength(&prefijo)) {
    return;
  }
  prefijo[i] = prefijo[i] + desplazamientos[wid.x];
}

Y la dispersión final, que es donde se materializa la lista:

@compute @workgroup_size(64)
fn dispersar(@builtin(global_invocation_id) gid : vec3<u32>) {
  let i = gid.x;
  if (i >= arrayLength(&visible)) {
    return;
  }
  if (visible[i] == 1u) {
    indices[prefijo[i]] = i;     // el hueco ya estaba decidido, sin atomicos
  }
}

Tres o cuatro dispatches y ni una sola operación atómica. Cada invocación escribe en una posición que nadie más va a tocar, calculada de forma puramente determinista.

Cuál elegir

La tabla, con criterio y sin matices de más:

Usa el contador atómico cuando el orden de la lista sea irrelevante: un lote de instancias del mismo material y la misma malla donde lo único que importa es cuántas hay. Es la mayoría de los casos, y escribirlo cuesta cinco minutos.

Usa el scan cuando el orden importe, y hay dos situaciones donde importa de verdad. La primera es cuando el array de instancias viene preordenado por material, por pipeline o por distancia a la cámara, y quieres que la lista compactada herede ese orden. Preordenar por material minimiza los cambios de estado; preordenar de delante hacia atrás maximiza el rechazo temprano de profundidad, que en una escena con complejidad de profundidad alta vale más que cualquier otra optimización de fragmento. El contador atómico destruye ese orden y con él el beneficio entero.

La segunda es el determinismo, y esta es la que la gente descubre tarde. Un efecto temporal —antialiasing acumulativo, oclusión ambiental con historial, reproyección— asocia el resultado de un fotograma con el del anterior. Si esa asociación pasa por el índice de instancia, y el índice de instancia lo asigna un atómico, la asociación es basura. También lo notarás al depurar: si dos capturas de la misma escena producen listas distintas, no puedes comparar buffers, no puedes bisecar un cambio y no puedes escribir un test de regresión que compruebe un hash del buffer de salida. El scan te devuelve la propiedad de que la misma entrada produce el mismo buffer bit a bit, y esa propiedad vale más de lo que parece hasta el día que la necesitas.

Hay un punto medio que funciona bien: contador atómico con agregación por workgroup, que es el segundo kernel de arriba. Los bloques salen en orden arbitrario entre sí, pero dentro de cada bloque los elementos quedan contiguos y en orden. Si tu criterio de ordenación es por material y has ordenado el buffer de instancias por material, los objetos de un mismo material caen en bloques contiguos, así que el resultado conserva la agrupación aunque no el orden global. Coste: un dispatch y treinta líneas.