En traitant la parcimonie comme une propriété plutôt que comme un détail d'implémentation (fastidieux), le sparsifier du compilateur1 génère automatiquement du code parcimonieux. Le concept a été développé pour l'algèbre linéaire par [Bik96] dans MT1 et formalisé pour l'algèbre tensorielle par [Kjolstad17,Kjolstad20] dans le projet TACO(Sparse Tensor Algebra Compiler).
Le dialecte SparseTensor accepte tous les attributs, types, opérations et passes nécessaires pour faire des types SparseTensor des citoyens de première classe dans l'infrastructure du compilateur MLIR. Le dialecte sert de pont entre les opérations de haut niveau sur les types de Tensor creux et les opérations de niveau inférieur sur les schémas de stockage creux réels, qui se composent de positions, de coordonnées et de valeurs. Le support de niveau inférieur peut consister en un code entièrement généré ou être fourni au moyen d'une petite bibliothèque de support d'exécution clairsemée.
L'implémentation MLIR [Biketal22] suit de près la "théorie de l'itération creuse" qui constitue la base de TACO. Le compilateur dispose d'une règle de réécriture pour chaque expression de Tensor dans le dialecte Linalg (notation d'index de Tensor de MLIR) : la parcimonie de chaque niveau d'un Tensor est indiquée par les types de niveau (par exemple, dense, compressé, singleton) et une spécification de l'ordre sur les niveaux (voir [Chou18] pour une discussion approfondie et les extensions possibles de ces types de niveau). Pour chaque expression Tensor dans l'entrée, le compilateur construit d'abord un graphique d'itération trié topologiquement qui reflète l'ordre des coordonnées par rapport aux niveaux de chaque Tensor. Cela garantit que tous les Tensors sont visités dans l'ordre naturel des niveaux et des coordonnées. Ensuite, des treillis d'itération sont construits pour chaque index dans l'ordre topologique. Chaque point de la grille d'itération se compose d'une conjonction de coordonnées de Tensor et d'une sous-expression Tensor qui doit être évaluée pour cette conjonction. Dans le treillis, les points d'itération sont ordonnés en fonction de la manière dont les coordonnées sont épuisées. Ainsi, ces treillis d'itération pilotent la génération de code creux, qui est un mappage un-à-un relativement simple des treillis d'itération vers des combinaisons de boucles for, de boucles while et d'instructions if. Les sorties de tenseurs creux qui se matérialisent non initialisées sont gérées soit par des insertions directes, si toutes les boucles parallèles sont les plus externes, soit par des insertions qui passent indirectement par une expansion de modèle d'accès unidimensionnel (également appelée espace de travail), lorsque cela est possible [Gustavson72,Bik96,Kjolstad19].
[Bik96] Aart J.C. Bik. Compatibilité du compilateur avec les calculs de matrices creuses. Thèse de doctorat, Université de Leyde, mai 1996.
[Biketal22] Aart J.C. Bik, Penporn Koanantakool, Tatiana Shpeisman, Nicolas Vasilache, Bixia Zheng et Fredrik Kjolstad. Compatibilité du compilateur avec les calculs de tenseurs creux dans MLIR. ACM Transactions on Architecture and Code Optimization, juin 2022. Voir : https://dl.acm.org/doi/10.1145/3544559
[Chou18] Stephen Chou, Fredrik Berg Kjolstad et Saman Amarasinghe. Abstraction de format pour les compilateurs d'algèbre de tenseurs creux. Proceedings of the ACM on Programming Languages, octobre 2018.
[Chou20] Stephen Chou, Fredrik Berg Kjolstad et Saman Amarasinghe. Génération automatique de routines de conversion de format de tenseur creux efficace. Proceedings of the 41st ACM SIGPLAN Conference on Programming Language Design and Implementation, juin 2020.
[Gustavson72] Fred G. Gustavson. Quelques techniques de base pour résoudre des systèmes d'équations linéaires creux. Dans Sparse Matrices and Their Applications, pages 41 à 52. Plenum Press, New York, 1972.
[Kjolstad17] Fredrik Berg Kjolstad, Shoaib Ashraf Kamil, Stephen Chou, David Lugato et Saman Amarasinghe. Compilateur d'algèbre tensorielle. Proceedings of the ACM on Programming Languages, octobre 2017.
[Kjolstad19] Fredrik Berg Kjolstad, Peter Ahrens, Shoaib Ashraf Kamil et Saman Amarasinghe. Tensor Algebra Compilation with Workspaces, Proceedings of the IEEE/ACM International Symposium on Code Generation and Optimization, 2019.
[Kjolstad20] Fredrik Berg Kjolstad. Compilation d'algèbre de tenseurs creux. Thèse de doctorat, MIT, février 2020.
Vous souhaitez contribuer ? Voici quelques problèmes simples.
-
Veuillez noter que nous préférons désormais le terme "sparsifier" à celui de "compilateur sparse", également couramment utilisé, pour désigner un tel pass. Cela permet de préciser que le pass de sparsification n'est pas un compilateur distinct, mais doit faire partie intégrante de tout pipeline de compilation créé avec l'infrastructure de compilation MLIR. ↩