Reutilización de datos y tiling

Puedes optimizar el patrón de acceso hasta que cada transacción aproveche el cien por cien del bus y aun así estar moviendo veinte veces más bytes de los necesarios, simplemente porque traes el mismo valor una y otra vez. Solo hay un principio que reduce el tráfico total en lugar de hacerlo más eficiente: traer un dato una vez y usarlo muchas. La memoria compartida del workgroup existe exactamente para eso, y la multiplicación de matrices es el caso donde la aritmética sale tan limpia que se puede seguir con lápiz.

Traer una vez, usar muchas

Multiplica dos matrices de N por N. La versión directa asigna una invocación a cada elemento del resultado, y esa invocación lee una fila completa de A y una columna completa de B: 2N valores para producir un número. Con N² resultados, el total son 2N³ lecturas de memoria global.

Pon números. Con N igual a 1024 y f32:

Cien veces más tiempo moviendo bytes que multiplicando. Y el dato interesante es que cada valor de A se lee N veces —una por cada columna del resultado— y cada valor de B también. Las matrices completas ocupan 4 MiB cada una: estás moviendo 8,6 GB para leer 8 MiB de datos. El mil por uno de repetición es todo el problema.

El tiling lo ataca de frente. Divide el resultado en bloques de T por T. Cada bloque necesita una banda de T filas de A y una banda de T columnas de B, que se recorren en trozos de T por T. Cada trozo se carga una sola vez en memoria compartida y desde ahí lo leen las T invocaciones que lo necesitan. El tráfico global pasa de 2N³ a 2N³/T lecturas: el factor de ahorro es exactamente el lado del bloque.

Con T igual a 16, los 8,6 GB se convierten en 537 MB y los 21,5 ms en 1,34 ms. Sigue habiendo seis veces más tiempo de memoria que de cálculo —el kernel continúa limitado por memoria— pero ha ganado un factor de dieciséis sin cambiar una sola operación aritmética.

El kernel completo

const T: u32 = 16u;   // tiene que coincidir con @workgroup_size

struct Dims { M: u32, N: u32, K: u32, relleno: u32 };

@group(0) @binding(0) var<uniform> d: Dims;
@group(0) @binding(1) var<storage, read>       A: array<f32>;   // M por K
@group(0) @binding(2) var<storage, read>       B: array<f32>;   // K por N
@group(0) @binding(3) var<storage, read_write> C: array<f32>;   // M por N

var<workgroup> tileA: array<array<f32, T>, T>;
var<workgroup> tileB: array<array<f32, T>, T>;

@compute @workgroup_size(16, 16)
fn gemm(@builtin(workgroup_id)        wid: vec3u,
        @builtin(local_invocation_id) lid: vec3u) {

  let fila = wid.y * T + lid.y;
  let col  = wid.x * T + lid.x;
  let pasos = (d.K + T - 1u) / T;

  var acc = 0.0;

  for (var p: u32 = 0u; p < pasos; p = p + 1u) {
    let aCol  = p * T + lid.x;
    let bFila = p * T + lid.y;

    // Carga cooperativa: cada invocacion trae un elemento de cada tile.
    // Los hilos con lid.x consecutivo leen posiciones consecutivas de A y de B:
    // las dos cargas son coalescentes.
    if (fila < d.M && aCol < d.K) {
      tileA[lid.y][lid.x] = A[fila * d.K + aCol];
    } else {
      tileA[lid.y][lid.x] = 0.0;      // el cero no altera el producto
    }
    if (bFila < d.K && col < d.N) {
      tileB[lid.y][lid.x] = B[bFila * d.N + col];
    } else {
      tileB[lid.y][lid.x] = 0.0;
    }

    workgroupBarrier();   // nadie calcula hasta que el tile este completo

    for (var k: u32 = 0u; k < T; k = k + 1u) {
      acc = fma(tileA[lid.y][k], tileB[k][lid.x], acc);
    }

    workgroupBarrier();   // nadie sobrescribe el tile hasta que todos hayan leido
  }

  if (fila < d.M && col < d.N) {
    C[fila * d.N + col] = acc;
  }
}

Del lado de JavaScript el dispatch es una cuadrícula de bloques, y maxComputeWorkgroupsPerDimension vale 65535, lo que con bloques de 16 da matrices de hasta un millón de columnas: no es el límite que te va a molestar.

pass.dispatchWorkgroups(Math.ceil(N / 16), Math.ceil(M / 16));

Hay cuatro decisiones en ese kernel que conviene que sepas defender.

Las dos barreras no son una duplicada. La primera resuelve el peligro de escritura seguida de lectura: si una invocación empieza a multiplicar antes de que sus compañeras hayan terminado de rellenar el tile, lee basura. La segunda resuelve el de lectura seguida de escritura: si una invocación rápida empieza la iteración siguiente y sobrescribe tileA mientras otra todavía está en el bucle interior, corrompe el cálculo de la otra. Olvidar la segunda es el error clásico, y es traicionero porque un workgroup pequeño que quepa en un solo grupo SIMD avanza en bloqueo de paso y da el resultado correcto por accidente; con 256 invocaciones repartidas en ocho grupos SIMD ya no, y falla en una máquina y no en otra.

La comprobación de límites carga ceros en vez de salir. La tentación es escribir if (fila >= d.M) { return; } al principio. No se puede: workgroupBarrier() exige control de flujo uniforme dentro del workgroup, y si una parte de las invocaciones ha salido, las que quedan esperan en la barrera a compañeras que nunca van a llegar. El comportamiento formal es indefinido y el práctico es un cuelgue o un resultado sucio. Cargar cero en las posiciones fuera de rango es correcto porque el cero es el elemento neutro del producto acumulado, y todas las invocaciones llegan a todas las barreras.

Las dos cargas son coalescentes. En A[fila * d.K + aCol], las invocaciones con lid.x consecutivo leen direcciones consecutivas; en B[bFila * d.N + col], también. Si hubieras escrito el tile de A transpuesto al cargarlo, la carga dejaría de ser coalescente y perderías por un lado lo ganado por el otro.

El bucle interior lee la memoria compartida de dos formas distintas. tileB[k][lid.x] recorre posiciones consecutivas al variar lid.x, que es el caso ideal. tileA[lid.y][k] da la misma dirección para todas las invocaciones que comparten lid.y, y una lectura de la misma dirección por muchos carriles se resuelve con una difusión, que es gratis. Ninguno de los dos accesos serializa.

El presupuesto: 16384 bytes y 256 invocaciones

maxComputeWorkgroupStorageSize vale 16384 bytes por defecto y maxComputeInvocationsPerWorkgroup vale 256. Haz la cuenta explícita antes de elegir el lado del bloque, porque los dos límites muerden a la vez y el que muerde primero no es el que la gente espera:

Un bloque de 16 por 16 de f32 ocupa 1024 bytes por matriz y las dos matrices caben de sobra. Un bloque de 64 por 64 ocupa exactamente los 16384 bytes del límite él solo, así que ni siquiera hay sitio para el segundo, y además necesitaría 4096 invocaciones por workgroup, dieciséis veces por encima del máximo. La conclusión es que con un hilo por elemento del bloque, 16 por 16 es el techo, y no por la memoria compartida sino por el número de invocaciones.

Los 2048 bytes del caso bueno tienen otra virtud: son tan pocos que la memoria compartida deja de ser el recurso que limita cuántos workgroups residen a la vez en una unidad de ejecución. Si la unidad dispone de 64 KB de memoria compartida —cifra típica, dependiente del hardware—, con 2048 bytes por grupo caben treinta y dos, y el cuello pasa a estar en los registros o en el número máximo de invocaciones residentes.

El mismo principio fuera de las matrices

El desenfoque separable. Un núcleo gaussiano de k por k cuesta k² muestras por píxel. Como el gaussiano es separable, se puede hacer en dos pasadas de k muestras cada una, una horizontal y otra vertical, con un resultado idéntico salvo redondeo. El ahorro crece con el radio:

El precio es una textura intermedia: a 1080p en RGBA16F son 16,6 MB que se escriben y se vuelven a leer, unos 83 microsegundos a 400 GB/s. Por debajo de k igual a 5 no se amortiza; a partir de 9 gana siempre. Y ojo con la condición: la separabilidad es una propiedad del núcleo, no una técnica. El gaussiano y la caja son separables; un filtro bilateral no lo es, y las versiones separables que circulan son aproximaciones que dejan artefactos en los bordes.

El halo compartido. En una convolución con un núcleo de radio R, un bloque de 16 por 16 píxeles necesita leer un rectángulo de (16 + 2R) por (16 + 2R) para producir sus 256 salidas. Con R igual a 4 son 576 lecturas para 256 resultados: 2,25 lecturas por salida, frente a las 81 de la versión directa. El halo tiene un sobrecoste propio que conviene cuantificar, porque decide el tamaño del bloque:

Los bloques grandes amortizan mejor el borde. Es la razón real por la que los kernels de convolución tienden a bloques de 32 por 32 con varias salidas por hilo, y no a 8 por 8.

La regla de decisión, que resume el nivel entero en una frase: si un dato lo va a leer más de un hilo del mismo workgroup, cárgalo a memoria compartida; si no, no. Un factor de reutilización de uno significa que has añadido una escritura, una barrera y una lectura para sustituir una lectura de memoria global que además habría sido coalescente. Con factor dos apenas empatas. A partir de cuatro, gana claro.