Tratar la dispersión como una propiedad en lugar de un detalle de implementación (tedioso) permite que el dispersor del compilador1 genere código disperso automáticamente. [Bik96] fue pionero en el concepto de álgebra lineal en MT1, y [Kjolstad17, Kjolstad20] lo formalizó para el álgebra de tensores en el proyecto del compilador de álgebra de tensores dispersos(TACO).
El dialecto SparseTensor admite todos los atributos, tipos, operaciones y pases que se requieren para que los tipos de tensores dispersos sean ciudadanos de primera clase dentro de la infraestructura del compilador de MLIR. El dialecto forma un puente entre las operaciones de alto nivel en tipos de tensores dispersos y las operaciones de nivel inferior en los esquemas de almacenamiento disperso reales que constan de posiciones, coordenadas y valores. La compatibilidad de nivel inferior puede consistir en código completamente generado o proporcionarse a través de una pequeña biblioteca de compatibilidad de tiempo de ejecución dispersa.
La implementación de MLIR [Biketal22] sigue de cerca la "teoría de la iteración dispersa" que constituye la base de TACO. El compilador tiene una regla de reescritura para cada expresión de tensor en el dialecto de Linalg (notación de índice de tensor de MLIR): la dispersión de cada nivel de un tensor se indica con los tipos de nivel (p.ej., denso, comprimido, singleton) y una especificación del orden en los niveles (consulta [Chou18] para obtener análisis detallados y posibles extensiones de estos tipos de nivel). Para cada expresión de tensor en la entrada, el compilador primero construye un grafo de iteración ordenado topológicamente que refleja el orden de las coordenadas con respecto a los niveles de cada tensor. Esto garantiza que todos los tensores se visiten en el orden natural de nivel y coordenada. A continuación, se construyen las celosías de iteración para cada índice en orden topológico. Cada punto de la celosía de iteración consta de una conjunción de coordenadas de tensor y una subexpresión de tensor que se debe evaluar para esa conjunción. Dentro de la retícula, los puntos de iteración se ordenan según la forma en que se agotan las coordenadas. Por lo tanto, estas celosías de iteración impulsan la generación de código disperso real, una asignación uno a uno relativamente sencilla de las celosías de iteración a combinaciones de bucles for, bucles while y sentencias if. Los resultados de tensores dispersos que se materializan sin inicializar se controlan con inserciones directas, si todos los bucles paralelos son los más externos, o con inserciones que pasan indirectamente por una expansión del patrón de acceso unidimensional (también conocido como espacio de trabajo), cuando es posible [Gustavson72,Bik96,Kjolstad19].
[Bik96] Aart J.C. Bik. Compiler Support for Sparse Matrix Computations. Tesis doctoral, Universidad de Leiden, mayo de 1996.
[Biketal22] Aart J.C. Bik, Penporn Koanantakool, Tatiana Shpeisman, Nicolas Vasilache, Bixia Zheng y Fredrik Kjolstad. Compatibilidad del compilador para los cálculos de tensores dispersos en MLIR. ACM Transactions on Architecture and Code Optimization, junio de 2022. Consulta https://dl.acm.org/doi/10.1145/3544559.
[Chou18] Stephen Chou, Fredrik Berg Kjolstad y Saman Amarasinghe. Abstracción de formato para compiladores de álgebra de tensores dispersos. Proceedings of the ACM on Programming Languages, octubre de 2018.
[Chou20] Stephen Chou, Fredrik Berg Kjolstad y Saman Amarasinghe. Automatic Generation of Efficient Sparse Tensor Format Conversion Routines. Proceedings of the 41st ACM SIGPLAN Conference on Programming Language Design and Implementation, junio de 2020.
[Gustavson72] Fred G. Gustavson. Algunas técnicas básicas para resolver sistemas dispersos de ecuaciones lineales. En Sparse Matrices and Their Applications, páginas 41 a 52. Plenum Press, Nueva York, 1972.
[Kjolstad17] Fredrik Berg Kjolstad, Shoaib Ashraf Kamil, Stephen Chou, David Lugato y Saman Amarasinghe. Es el compilador de álgebra de tensores. Proceedings of the ACM on Programming Languages, octubre de 2017.
[Kjolstad19] Fredrik Berg Kjolstad, Peter Ahrens, Shoaib Ashraf Kamil y Saman Amarasinghe. Tensor Algebra Compilation with Workspaces, Proceedings of the IEEE/ACM International Symposium on Code Generation and Optimization, 2019.
[Kjolstad20] Fredrik Berg Kjolstad. Compilación de álgebra de tensores dispersos. Tesis de doctorado, MIT, febrero de 2020.
¿Te gustaría colaborar? Estos son algunos buenos primeros problemas.
-
Ten en cuenta que ahora preferimos el término "esparcificador" en lugar de la terminología también comúnmente utilizada de "compilador disperso" para referirnos a ese pase y dejar en claro que el pase de esparcificador no es un compilador independiente, sino que debe ser una parte integral de cualquier canalización de compilador que se cree con la infraestructura del compilador de MLIR. ↩