El problema fundamental que TIRx aborda es la tensión inherente entre la abstracción de alto nivel en los compiladores de kernels de Machine Learning y la necesidad de control explícito de bajo nivel para exprimir el máximo rendimiento en hardware acelerador de vanguardia. A medida que el hardware de IA evoluciona rápidamente con nuevas instrucciones, jerarquías de memoria y patrones de cooperación, los compiladores de alto nivel como Triton, que ocultan detalles de asignación de hilos y movimiento de memoria, se encuentran bajo presión. Estos sistemas son excelentes para patrones de kernels maduros, pero fallan cuando los ingenieros necesitan interactuar directamente con las características más recientes del hardware.

TIRx se posiciona como una capa intermedia, ofreciendo un DSL que permite a los expertos controlar directamente la orquestación, la sincronización y el uso de intrínsecos de hardware, mientras que expone primitivas de tile recurrentes al compilador para su reutilización y optimización. Esta filosofía de diseño permite que el stack de software evolucione con el hardware, adaptándose a nuevas generaciones de aceleradores sin requerir una reescritura completa del compilador. La relevancia actual radica en la explosión de nuevos diseños de hardware de IA y la necesidad de optimizar kernels para modelos de lenguaje grandes (LLMs) y otras cargas de trabajo intensivas en computación, donde cada porcentaje de rendimiento cuenta.

Arquitectura del Sistema

La arquitectura de TIRx se centra en un DSL hardware-nativo y un compilador ligero que opera sobre Apache TVM. El modelo de programación se basa en tres ingredientes clave: Execution Scope, Tensor Layout y Tile Primitive Dispatch. El Execution Scope define quién ejecuta una operación y con qué granularidad (ej. nivel de hilo o warpgroup), utilizando construcciones de control de flujo y namespaces de primitivas. El Tensor Layout describe la ubicación física de un tensor lógico en la jerarquía de memoria (global, compartida, registros, memoria de tensor, SRAM del acelerador) y cómo sus elementos se distribuyen (sharding, replication, offset) a través de ejes de hardware nombrados (laneid, warpid, etc.). Este layout es un contrato de almacenamiento, no una interfaz de particionamiento de trabajo.

El Tile Primitive Dispatch es el componente que traduce una llamada a una primitiva de tile en IR nativo. Utiliza los layouts de los operandos, el execution scope y el backend objetivo (o una sugerencia explícita) para seleccionar la implementación correspondiente. Por ejemplo, una copia de memoria global a compartida podría resolverse a una operación TMA, o una multiplicación de matrices a WGMMA o una instrucción de array sistólico. Este dispatch genera los bucles y cálculos de direccionamiento necesarios para aplicar la instrucción a todo el tile. El compilador de TIRx es ligero porque, después del dispatch de primitivas, el programa ya es IR de kernel nativo y puede traducirse directamente a código de backend (ej. CUDA C++/PTX), evitando pases de optimización pesados en la ruta crítica para expresar nuevos kernels. Las optimizaciones avanzadas (especialización automática de warps, inferencia de layout, búsqueda de schedules) se tratan como capas opcionales que pueden construirse sobre este camino de bajada directo.

Flujo de Compilación de Kernel TIRx

  1. 1 Código Fuente TIRx Programa con IR hardware-nativo y llamadas a primitivas de tile.
  2. 2 Parser Análisis sintáctico del código fuente.
  3. 3 Primitive Dispatch Resuelve llamadas a primitivas de tile a IR nativo (bucles, cálculos de direc...
  4. 4 Generación de Código Backend Traduce el IR nativo a código específico del hardware (ej. CUDA C++/PTX).
  5. 5 Ejecución en Hardware El kernel compilado se ejecuta en GPU/acelerador.
CapaTecnologíaJustificación
orchestration TIRx DSL Define la estructura de pipeline, sincronización, asignación de roles y colocación de memoria directamente en el código fuente, permitiendo control explícito de bajo nivel. vs Triton, TileLang
compute Tile Primitives Expone operaciones recurrentes (ej. movimiento de datos, multiplicación de matrices) al compilador para su reutilización y dispatch optimizado a intrínsecos de hardware. vs Funciones de biblioteca opacas, Implementaciones manuales repetitivas
storage Tensor Layouts (Storage-First Interface) Describe cómo los datos de un tensor lógico se mapean a coordenadas físicas en la memoria del hardware, incluyendo sharding, replicación y offset, informando al dispatch de primitivas. vs CuTe (layout como partición de trabajo), Layouts genéricos sin ejes nombrados
orchestration Apache TVM Infraestructura de compilador subyacente sobre la cual TIRx está construido, proporcionando la base para la representación intermedia (IR) y las utilidades del compilador. vs LLVM, Compiladores propietarios

Trade-offs

Ganancias
  • ▲ Control explícito de bajo nivel para hardware de vanguardia
  • ▲ Rendimiento competitivo con baselines de estado del arte
  • ▲ Extensibilidad para nuevo hardware sin reescritura del compilador
  • ▲ Composabilidad de tareas para megakernels
  • ▲ Soporte para programación de kernels por agentes (Agentic Kernel Programming)
Costes
  • △ Mayor esfuerzo de ingeniería para escribir kernels a mano
  • △ Curva de aprendizaje más pronunciada que DSLs de alto nivel 
T.Layout([
    T.D.Shard(T.axis.Warp, 4, 16),
    T.D.Shard(T.axis.Lane, 4, 4),
    T.D.Shard(T.axis.Reg, 2, 1),
    T.R.Replica(T.axis.Warp, 2, 4),
    T.O.Offset(T.axis.Warp, 5)
])
Ejemplo de cómo se define un layout de tensor en TIRx usando la API de Python, especificando sharding, replicación y offset en ejes de hardware nombrados.

Fundamentos Teóricos

La necesidad de un control de bajo nivel en la programación de GPUs y aceleradores se remonta a los primeros días de la computación paralela y de alto rendimiento. Conceptos como la gestión explícita de la memoria compartida y la sincronización entre hilos son fundamentales en la programación CUDA y OpenCL, reflejando principios de concurrencia y paralelismo. La idea de un DSL para la generación de código optimizado para hardware específico tiene raíces en la investigación de compiladores y lenguajes específicos de dominio (DSLs) para optimizar el rendimiento, un campo explorado por autores como Aiken y sus colegas en el contexto de lenguajes para computación científica.

La distinción entre un 'layout como contrato de almacenamiento' y un 'layout como interfaz de particionamiento de trabajo' en TIRx, en contraste con sistemas como CuTe, refleja una decisión de diseño fundamental sobre la separación de preocupaciones. Mientras CuTe (parte de la biblioteca CUTLASS de NVIDIA) utiliza layouts para derivar la partición de trabajo, TIRx los usa para informar al dispatch de primitivas, lo que es consistente con la filosofía de 'mantener la base pequeña y explícita'. Este enfoque puede verse como una aplicación de principios de diseño de compiladores donde la representación intermedia (IR) se mantiene lo más cercana posible al hardware para facilitar la generación de código de alto rendimiento, mientras que las abstracciones de más alto nivel se construyen sobre esta base. La capacidad de TIRx para soportar la generación de kernels por agentes (Agentic Kernel Programming) se alinea con la investigación en compiladores auto-optimizables y programación generativa, donde los sistemas buscan explorar espacios de diseño para encontrar implementaciones óptimas, un concepto que ha sido explorado en trabajos como el de Chen et al. (2018) sobre TVM.