Bits y empaquetado: meter cuatro canales en un u32

Un shader que mueve mucha memoria gasta más tiempo esperando datos que calculando, y la forma más directa de acelerarlo es que los datos ocupen menos. WGSL trae para eso una familia de funciones de empaquetado que traducen entre vectores de flotantes y enteros de 32 bits en una sola instrucción de hardware, y un juego completo de operaciones de bits que casi nadie usa hasta que le hace falta.

Los operadores de bits

Los seis de siempre: &, |, ^, ~, << y >>. Funcionan sobre i32 y u32 y sobre sus vectores, componente a componente.

La diferencia entre los dos tipos está en el desplazamiento a la derecha: sobre u32 es lógico e introduce ceros, y sobre i32 es aritmético y propaga el bit de signo. El desplazamiento a la izquierda es igual en los dos.

let banderas : u32 = 0b1011u;
let tiene = (banderas & 0b0010u) != 0u;    // comprobar un bit
let puesto = banderas | 0b0100u;           // poner un bit
let quitado = banderas & ~0b0001u;         // quitar un bit
let cambiado = banderas ^ 0b1000u;         // invertir un bit

El operando de desplazamiento tiene que ser u32, y desplazar por un número de bits mayor o igual que la anchura del tipo no es una operación con resultado útil; si el compilador puede ver que la constante se pasa de rango, lo rechaza directamente.

Un uso que aparece constantemente es empaquetar varias banderas y un índice en el mismo u32 para no gastar cuatro campos en una struct:

// Bits 0..15: indice de material. Bits 16..23: banderas.
fn material(empaquetado : u32) -> u32 { return empaquetado & 0xffffu; }
fn esTransparente(empaquetado : u32) -> bool { return (empaquetado & (1u << 16u)) != 0u; }

Las funciones de bits

extractBits e insertBits sustituyen las máscaras y los desplazamientos escritos a mano, y se compilan a la instrucción de campo de bits que tienen casi todas las GPUs:

// Los mismos dos campos de antes, con las funciones.
let indice = extractBits(empaquetado, 0u, 16u);
let bandera = extractBits(empaquetado, 16u, 8u);

countOneBits es la operación con la que se cuentan elementos en máscaras de cobertura y en algoritmos de compactación. reverseBits aparece en la FFT y en el ordenado por bits de Morton. Y firstLeadingBit es el logaritmo entero en base dos, que es lo que se usa para calcular niveles de mip y profundidades de árbol.

Empaquetar valores normalizados

Aquí está lo verdaderamente útil. Estas funciones convierten entre vectores de flotantes y un u32, aplicando la cuantización de un formato normalizado, y en la mayoría del hardware son una sola instrucción.

Cada una tiene su unpack correspondiente que devuelve el vector original con el error de cuantización. La disposición es siempre la misma: el componente 0 va en los bits menos significativos.

let color = vec4f(1.0, 0.5, 0.25, 1.0);
let empaquetado = pack4x8unorm(color);      // 0xff4080ff aproximadamente
let recuperado = unpack4x8unorm(empaquetado);

Las conversiones son las de siempre: unorm acota a cero y uno y multiplica por 255; snorm acota a menos uno y uno y multiplica por 127. La pérdida es real y hay que dimensionarla antes de usarla: 8 bits dan 256 niveles, que es suficiente para un color final y claramente insuficiente para una normal o para una posición.

pack2x16float es distinto de los otros cuatro: no cuantiza a un rango normalizado, sino que convierte dos f32 a dos medios flotantes con su rango y su precisión. Es el que quieres cuando el valor no está acotado.

Con la extensión packed_4x8_integer_dot_product aparecen además pack4xI8, pack4xU8, sus variantes con acotado, sus unpack, y las dos operaciones que dan nombre a la extensión, dot4I8Packed y dot4U8Packed, que calculan un producto escalar de cuatro enteros de 8 bits empaquetados en un solo u32. Es la instrucción que usan las redes neuronales cuantizadas.

Los casos reales

Un G-buffer compacto. En vez de tres targets de rgba16float, uno de rgba32uint con todo empaquetado:

struct Gbuffer {
  @location(0) datos : vec4u,
};

@fragment
fn fs(entrada : Entrada) -> Gbuffer {
  var s : Gbuffer;
  s.datos.x = pack4x8unorm(vec4f(albedo, 0.0));       // color base
  s.datos.y = pack2x16snorm(normalOctaedrica);        // normal en 32 bits
  s.datos.z = pack4x8unorm(vec4f(rugosidad, metalico, oclusion, 0.0));
  s.datos.w = bitcast<u32>(profundidadLineal);
  return s;
}

Vértices comprimidos leídos desde un storage buffer. Una normal por vértice en snorm8x4 ocupa 4 bytes en vez de 12, y se desempaqueta con una instrucción en el vertex shader.

Estado de simulación. Un contador y tres banderas en el mismo u32 que se manipula con operaciones atómicas, en vez de cuatro campos que necesitarían cuatro atómicas distintas.