API de fragmentos de WMMA

Para programar directamente los Tensor Cores, utilice WMMA, la API de multiplicación de matrices con acumulación por warp del espacio de nombres nvcuda::wmma. 🧩

WMMA funciona a nivel de warp: los 32 hilos de un warp trabajan juntos en un bloque. Debe pensar en bloques, no en hilos individuales.

Un fragmento es el tipo de datos de WMMA que contiene la parte de un warp de un bloque de matriz. Cada hilo posee internamente algunos de sus elementos.

Declare fragmentos etiquetados como matrix_a, matrix_b o accumulator. La etiqueta indica a WMMA cómo participa el bloque en la multiplicación con acumulación.

Los fragmentos utilizan tamaños de bloque establecidos, como 16 por 16 por 16. Elija una forma compatible; el hardware solo acepta esas combinaciones.

Primero, llame a load_matrix_sync para transferir un bloque de la memoria a un fragmento. Esta carga distribuye los datos entre el warp automáticamente.

Antes de sumar, ponga a cero el bloque de resultados con fill_fragment. Un acumulador limpio garantiza que las sumas comiencen desde un valor conocido.

La llamada principal es mma_sync. Ejecuta D = A por B más C en los fragmentos cargados utilizando directamente los Tensor Cores.

Por último, store_matrix_sync escribe el fragmento acumulador de nuevo en la memoria. Este almacenamiento reúne las partes de cada hilo en un único bloque.

El sufijo _sync recuerda que cada llamada de WMMA es sincrónica respecto al warp. Los 32 lanes deben alcanzarla juntos o el comportamiento será indefinido.

Las cargas y los almacenamientos necesitan la dimensión líder, el salto entre filas, además de una disposición por filas o por columnas para leer correctamente la memoria.

Ha recorrido el flujo de WMMA: declarar un fragmento, cargar los bloques, borrar el acumulador, llamar a mma_sync y, después, almacenar. Cada paso es sincrónico respecto al warp. 🙌