Operaciones atómicas en WGSL

Una barrera ordena el tiempo; una atómica hace indivisible una operación. Son herramientas distintas para problemas distintos, y confundirlas produce código que sincroniza sin necesidad o que cuenta mal sin darse cuenta. WGSL ofrece un catálogo pequeño y deliberadamente austero: dos tipos, once funciones, dos address spaces. Con eso se construyen contadores, histogramas, listas de anexado, reservas de espacio y todos los cerrojos que de verdad se pueden usar en una GPU, que son menos de los que parece.

El catálogo y sus límites

Un tipo atómico se escribe atomic con el tipo entre ángulos, y solo admite dos: i32 y u32. No hay atómicas de f32, ni de f16, ni de vectores. Y solo se puede declarar en dos address spaces: storage con acceso read_write, y workgroup. En un uniform buffer es ilegal, en private es ilegal, y en un storage de solo lectura también.

// En storage: compartido por todo el dispatch.
struct Contadores {
  vivas:  atomic<u32>,
  libres: atomic<u32>,
};
@group(0) @binding(0) var<storage, read_write> cont: Contadores;
@group(0) @binding(1) var<storage, read_write> histo: array<atomic<u32>, 256>;

// En workgroup: compartido solo por el grupo.
var<workgroup> histoLocal: array<atomic<u32>, 256>;
var<workgroup> maxLocal: atomic<u32>;

Todas las funciones toman como primer argumento un puntero a la variable atómica, que se obtiene con el operador de dirección. Las que modifican devuelven el valor antiguo, no el nuevo, y ese detalle es la clave de casi todos los patrones útiles.

Que atomicAdd devuelva el valor antiguo es lo que convierte un contador en un repartidor de índices. Cada invocación que llama recibe un número distinto, y eso resuelve el problema de “quiero anexar un elemento a una lista y no sé cuántos van a anexar los demás”:

// Anexar de forma segura a una lista compartida.
fn anexar(v: u32) {
  let ranura = atomicAdd(&cont.vivas, 1u);   // reserva exclusiva
  if (ranura < arrayLength(&lista)) {
    lista[ranura] = v;
  }
}

El guardia contra la longitud es imprescindible: si hay más aportaciones que capacidad, el contador sigue creciendo y las últimas escrituras se saldrían del buffer. Además hay que recordar saturar el contador antes de usarlo como número de elementos, porque puede haber quedado por encima de la capacidad.

atomicSub tiene una trampa concreta con u32: restar por debajo de cero no da un error, envuelve y produce un número enorme. Un pool de libres del que se saca sin comprobar acaba con un contador de cuatro mil millones. La forma segura es reservar con atomicSub, comprobar el valor antiguo devuelto, y devolver la reserva si no había existencias:

fn sacarLibre() -> i32 {
  let antes = atomicSub(&cont.libres, 1u);
  if (antes == 0u) {
    atomicAdd(&cont.libres, 1u);   // deshacer: no habia
    return -1;
  }
  return i32(antes - 1u);
}

Comparación e intercambio: sumar flotantes

atomicCompareExchangeWeak es la operación universal: con ella se construye cualquier operación atómica que no exista de fábrica. Compara el contenido con cmp y, si coincide, lo sustituye por v. Devuelve un struct con dos campos: old_value, el contenido antes de la operación, y exchanged, un booleano que dice si el cambio se produjo.

El “weak” del nombre significa que puede fallar espuriamente aunque los valores coincidan, así que siempre se usa dentro de un bucle que reintenta.

La aplicación más útil es la suma de flotantes, que WGSL no ofrece. Se reinterpretan los bits del flotante como u32, se hace la suma en el mundo de los flotantes y se intenta escribir el resultado; si otro se adelantó, se reintenta con su valor.

@group(0) @binding(0) var<storage, read_write> acumulador: atomic<u32>;

fn atomicAddF32(p: ptr<storage, atomic<u32>, read_write>, v: f32) {
  var viejo = atomicLoad(p);
  loop {
    let nuevo = bitcast<u32>(bitcast<f32>(viejo) + v);
    let r = atomicCompareExchangeWeak(p, viejo, nuevo);
    if (r.exchanged) { break; }
    viejo = r.old_value;      // otro se adelanto: reintentar con su valor
  }
}

Es correcto y funciona, con dos advertencias. Cuesta bastante más que un atomicAdd nativo, sobre todo con contención alta, porque cada colisión obliga a repetir la vuelta. Y el resultado depende del orden de llegada, con lo que no es reproducible bit a bit entre ejecuciones. Para totales que se muestran, irrelevante; para tests de regresión o simulaciones realimentadas, un problema serio.

Un truco que evita el bucle cuando los valores son positivos: los flotantes IEEE 754 no negativos comparan igual como enteros sin signo que como flotantes, porque el bit de signo es cero y el exponente está sesgado. Así que un máximo de flotantes positivos se puede hacer con atomicMax sobre el patrón de bits directamente.

// Solo valido si v >= 0. Para negativos hay que invertir los bits.
fn atomicMaxPositivo(p: ptr<storage, atomic<u32>, read_write>, v: f32) {
  atomicMax(p, bitcast<u32>(v));
}

El patrón de dos niveles

La contención es el enemigo de las atómicas. Si 65536 invocaciones incrementan el mismo contador, el hardware serializa esas 65536 operaciones: el kernel deja de ser paralelo en ese punto y el tiempo pasa a ser proporcional al número de invocaciones.

La solución es agregar primero en memoria compartida, donde las atómicas son mucho más rápidas y la contención está limitada al workgroup, y hacer una sola aportación global por grupo. Un histograma lo ilustra perfectamente.

const TAM:  u32 = 256u;
const BINS: u32 = 256u;

@group(0) @binding(0) var<uniform> params: Params;                  // params.conteo
@group(0) @binding(1) var<storage, read> datos: array<u32>;
@group(0) @binding(2) var<storage, read_write> histo: array<atomic<u32>, BINS>;

var<workgroup> local: array<atomic<u32>, BINS>;

@compute @workgroup_size(TAM)
fn main(@builtin(global_invocation_id) gid: vec3u,
        @builtin(local_invocation_index) li: u32,
        @builtin(num_workgroups) nwg: vec3u) {

  // La memoria workgroup arranca a cero, pero lo dejamos explicito:
  // si este kernel se reutilizara en un bucle, haria falta igualmente.
  var b = li;
  loop {
    if (b >= BINS) { break; }
    atomicStore(&local[b], 0u);
    b += TAM;
  }
  workgroupBarrier();

  // Acumular en memoria compartida, con zancada para no perder coalescencia.
  let zancada = nwg.x * TAM;
  var i = gid.x;
  loop {
    if (i >= params.conteo) { break; }
    atomicAdd(&local[datos[i] & (BINS - 1u)], 1u);
    i += zancada;
  }
  workgroupBarrier();

  // Una sola aportacion global por bin y por workgroup.
  b = li;
  loop {
    if (b >= BINS) { break; }
    let c = atomicLoad(&local[b]);
    if (c > 0u) { atomicAdd(&histo[b], c); }
    b += TAM;
  }
}

Con un millón de elementos y 256 workgroups, la versión ingenua hace un millón de atómicas globales; esta hace un millón de atómicas locales, que son órdenes de magnitud más baratas, y como mucho 65536 globales. La diferencia medida suele estar entre cinco y veinte veces.

La condición if (c > 0u) no es cosmética: en datos sesgados, la mayoría de los bins de un workgroup están vacíos y saltarse la atómica global ahorra la mayor parte del tráfico restante.

Este mismo patrón —agregar en el workgroup, aportar una vez al global— es el esqueleto de la reducción y de contar partículas por celda. Es probablemente la idea con mejor relación entre sencillez y rendimiento de todo el cómputo en GPU.