La generación de código eficiente para arquitecturas heterogéneas y cargas de trabajo de alto rendimiento (HPC, Machine Learning) es un problema fundamental en la computación. Los compiladores tradicionales a menudo pierden información semántica valiosa al bajar rápidamente a representaciones de bajo nivel (como CFG o bucles genéricos), lo que dificulta optimizaciones avanzadas como el tiling, la fusión de bucles o la delegación a bibliotecas optimizadas. Esto lleva a una "brecha de abstracción" que los compiladores deben reconstruir, a menudo de forma heurística y frágil.
Linalg aborda este problema manteniendo la información semántica de alto nivel directamente en su Representación Intermedia (IR) a través de operaciones declarativas. En lugar de depender de análisis complejos para inferir propiedades de las operaciones, Linalg codifica estas propiedades en la IR misma. Esto permite que las transformaciones sean "correctas por construcción" y simplifica la lógica del compilador, facilitando la aplicación de optimizaciones que de otro modo serían prohibitivamente complejas o propensas a errores. La clave es un "lowering progresivo" que reduce el potencial de una función (número de instrucciones de bajo nivel) gradualmente, sin perder la intención del programador.
Arquitectura del Sistema
Linalg se implementa como un dialecto dentro del marco MLIR (Multi-Level Intermediate Representation). MLIR proporciona una infraestructura extensible para definir IRs en múltiples niveles de abstracción, lo que Linalg aprovecha para su enfoque de "lowering progresivo". Las operaciones de Linalg (ops) se definen declarativamente, utilizando las capacidades de MLIR como Regions, OpInterfaces, ODS (Operation Definition Specification) y Declarative Rewrite Rules (DRR).
El dialecto Linalg define varias "formas" de operaciones: Generic (la más general, puede codificar cualquier operación de bucle perfectamente anidado), Category (como linalg.contract o linalg.elementwise, que agrupan semánticas similares), Named (operaciones específicas como linalg.matmul o linalg.add) y Composite (operaciones no perfectamente anidadas como linalg.softmax que se expanden a otras operaciones). Estas formas son interconvertibles (excepto Composite), permitiendo que las transformaciones operen en la representación más conveniente. Por ejemplo, un linalg.matmul puede ser "tiled" y seguir siendo un linalg.matmul a un nivel más bajo, encapsulando un linalg.matmul más pequeño dentro de bucles. Esto contrasta con enfoques que bajarían a bucles genéricos y perderían la semántica de matmul.
Las transformaciones en Linalg se basan en reglas de reescritura locales (DRR) y se diseñan para preservar la información de alto nivel. Por ejemplo, el tiling se enfoca en el espacio de datos (creando "views" en memrefs strided de MLIR) en lugar del espacio de iteración, lo que simplifica las expresiones de dirección y evita la complejidad de los compiladores poliédricos tradicionales. Linalg se complementa con el dialecto Vector de MLIR para operaciones vectoriales estructuradas, siguiendo principios de diseño similares.
| Capa | Tecnología | Justificación |
|---|---|---|
| orchestration | MLIR | Infraestructura extensible para definir y manipular representaciones intermedias (IR) en múltiples niveles de abstracción, permitiendo el diseño de dialectos como Linalg. vs LLVM (directamente), GCC (directamente) |
| data-processing | Linalg Dialect | Define operaciones declarativas de alto nivel para álgebra lineal y computación tensorial, facilitando la optimización y generación de código para HPC y ML. vs ONNX, XLA HLO, HalideIR, TVM Relay, Tensor Comprehensions |
| data-processing | Vector Dialect (MLIR) | Complementa a Linalg proporcionando operaciones estructuradas en vectores, siguiendo principios de diseño similares para optimización vectorial. |
| data-processing | Affine Dialect (MLIR) | Proporciona abstracciones poliédricas más cercanas a la representación SSA, sirviendo como objetivo de lowering para Linalg y permitiendo transformaciones de nivel medio. vs GRAPHITE (GCC), Polly (LLVM) |
Trade-offs
Ganancias
- ▲ Simplicidad y corrección de transformaciones
- ▲ Preservación de información semántica de alto nivel
- ▲ Interoperabilidad con bibliotecas de alto rendimiento (BLAS, DNN)
- ▲ Reducción de la proliferación de operaciones en la IR
- ▲ Facilidad para la co-diseño frontend-compilador
Costes
- △ Curva de aprendizaje inicial para el diseño de dialectos en MLIR
- △ Posiblemente, mayor verbosidad en la IR de alto nivel comparada con IRs de bajo nivel
Fundamentos Teóricos
El diseño de Linalg se inspira en décadas de investigación en compiladores y lenguajes específicos de dominio (DSLs). Conceptos como la "preservación de información" y el "lowering progresivo" se oponen a la tendencia de muchos sistemas de bajar rápidamente a representaciones de bajo nivel como el SSA (Static Single Assignment) o CFG (Control Flow Graph), que son fundamentales en compiladores como LLVM. La idea de que "Algorithms + Data Structures = Programs" (Wirth, 1976) se reinterpreta en Linalg para enfatizar la fuerte relación entre el flujo de control y las estructuras de datos, influyendo en cómo se modelan las operaciones y transformaciones.
Linalg toma lecciones de compiladores poliédricos (como GRAPHITE en GCC o Polly en LLVM) y el dialecto Affine de MLIR, pero busca evitar sus complejidades, como la dificultad de integrar el modelo matemático poliédrico con SSA o la complejidad exponencial de algunos algoritmos. También se nutre de la experiencia de sistemas como Halide, TVM, XLA, ONNX y Tensor Comprehensions, adoptando sus fortalezas (como la especificación declarativa de operaciones y transformaciones) y mitigando sus debilidades (como la proliferación de operaciones o la pérdida prematura de información). La noción de "raising" en compiladores poliédricos (reconstruir información de alto nivel a partir de IR de bajo nivel) es un problema que Linalg busca evitar al mantener la información declarativamente desde el principio.