Construir un compilador de GPU desde cero: de PyTorch a CUDA mediante IR multinivel (Emmy)

Aviso: el contenido descargado está truncado. El texto se corta a mitad de la sección de autotuning, así que este resumen cubre solo lo que aparece en la descarga y no incluye las conclusiones ni los resultados de los benchmarks.

El artículo es un tutorial publicado en el newsletter Kernel Space, firmado por Dmitry Trifonov. Sigue una operación de PyTorch hasta CUDA mediante lowering progresivo, usando un compilador llamado Emmy (versión 0.3.22). El autor cuenta que en 2018 empezó, junto a un amigo, a construir la pila de lowering del Apple Neural Engine sin experiencia previa en compiladores. Según él, construir compiladores de ML sigue pareciendo un trabajo de pioneros, con muchos marcos, ideas y trucos. Considera que la pila de IR y el lowering progresivo son de lo más asentado del campo.

Para reproducirlo hace falta Python 3.12+ y la instalación de emmy-ml con PyTorch 2.14.1. Se puede inspeccionar el IR sin GPU. Para los benchmarks se necesita una GPU NVIDIA y nvcc. El ejemplo conductor es RMSNorm: calcula la media de los cuadrados de la fila, suma epsilon, aplica la raíz cuadrada inversa y escala cada elemento por ese valor y por su peso aprendido.

El pipeline tiene seis etapas, que el artículo recorre una a una.

Torch IR. El frontend captura el módulo con el mecanismo de exportación de PyTorch. En esta etapa RMSNorm sigue siendo una sola operación en el grafo torch.fx, y los pesos del checkpoint aparecen como load("weight").

Tensor IR. Unifica los frontends. Cada framework guarda sus operaciones con convenciones distintas. Por ejemplo, una capa densa de PyTorch almacena el peso transpuesto respecto a la de otros frameworks. Tensor IR hace explícitos la indexación y las reducciones mediante llamadas index_map, que cubren transposiciones y broadcasts sin aritmética adicional. Así, un nn.Linear y un matmul directo acaban describiendo el mismo producto de forma [4,3,2] y la misma suma sobre el eje de entrada. RMSNorm se descompone aquí en pasos aritméticos: cuadrado, suma, división por el recuento, suma de epsilon, rsqrt y dos multiplicaciones.

Loop IR. Convierte cada nodo en bucles escalares y fusiona las regiones que pueda unir legalmente. Si se lanzara un kernel CUDA por nodo, habría que escribir y releer tensores intermedios en memoria. El ejemplo exp(-x) muestra dos bucles que se fusionan en uno que lee, niega, aplica exp y escribe, sin tensor intermedio. RMSNorm queda en un único kernel, y la estadística permanece dentro del kernel en lugar de pasar por memoria global.

Normalización. Dentro de Loop IR elimina trabajo duplicado y pone programas equivalentes en forma canónica. En el ejemplo (x*x)+(x*x), sin normalización hay cuatro cargas y dos multiplicaciones. Con ella queda una carga y una multiplicación. Además, x+x+y e y+(x+x) producen el mismo cuerpo porque la suma es conmutativa. La identidad de un kernel depende del cuerpo normalizado, de las entradas y salidas tipadas y de las decisiones del compilador. Por eso se puede reutilizar un schedule eficiente entre formas equivalentes.

Tile IR. Es la primera etapa con decisiones conscientes de la GPU. Empareja un árbol de Fold, independiente del schedule, con decisiones de planificación separadas, siguiendo la separación algoritmo/schedule de Halide. Un Fold expresa un map, una reducción o un scan mediante operandos, una función por elemento (lift), un estado inicial y una operación de combinación opcionales, y un observador opcional. En el ejemplo de RMSNorm, un Fold interno reduce los cuadrados de la fila. Otro convierte la suma en una escala. El externo produce cada elemento de salida. El schedule elegido hace que 256 workers cooperen en la reducción de cada fila. Emmy puede explorar otros schedules para el mismo árbol antes de bajarlo a Kernel IR y CUDA.

Códec de schedules. El matmul sirve para ver cómo interactúan varias decisiones: cuántos warps cubren la salida, cuántas sumas parciales guarda cada warp en registros, cómo se divide la reducción y cómo se mueven las entradas por memoria compartida. El códec registra cada combinación completa como una fila canónica, con campos como WORK, TILE, REDUCE, STAGE y RASTER. El artículo compara dos filas para un matmul FP16 de 128×128, ambas con la misma instrucción de tensor core. Una mayor tesela en registros reutiliza más cada valor cargado de A y B y mejora la relación cómputo/tráfico de memoria, pero consume registros y puede dejar menos warps residentes. El autor señala que hacen falta mediciones para saber cuál gana. El CUDA emitido incluye copias TMA con barreras, cargas de fragmentos con ldmatrix y la instrucción mma de tensor core, con doble buffer. Emmy recorre los puntos de decisión y combina solo opciones compatibles. El autor sugiere que un agente LLM podría proponer filas a través de la misma interfaz y que el compilador comprobaría su legalidad.

Partir un kernel. Aunque Loop IR fusiona todo lo legal, a veces conviene volver a dividir, y Tile IR tiene la información de hardware para decidirlo. En el ejemplo de dos matmuls encadenados, H=AB y O=HC, la versión fusionada recalcula cada elemento de H para cada una de las 16 columnas de salida. Con PLACE=cut, cada elemento de H se calcula una vez. A cambio, aparece un intermedio de 4 KiB y un segundo lanzamiento de kernel, pero cada matmul puede recibir un schedule distinto.

Autotuning. La última sección, que se corta en la descarga, empieza con RMSNorm como ejemplo corto de ajuste. Su espacio de schedules es pequeño, se convierte en un solo kernel y se puede muestrear rápido. El autor adelanta que las ganancias suelen ser pequeñas y que matmul y attention ofrecen más rango, pero tardan mucho más en ajustarse. El texto recibido no incluye los resultados.