Einführung

Wenn Sparsity als Eigenschaft und nicht als (mühsames) Implementierungsdetail behandelt wird, kann der Sparsifier des Compilers1 automatisch spärlichen Code generieren. Das Konzept wurde für die lineare Algebra von [Bik96] in MT1 entwickelt und von [Kjolstad17,Kjolstad20] im Sparse Tensor Algebra Compiler (TACO) für die Tensoralgebra formalisiert.

Der SparseTensor-Dialekt unterstützt alle Attribute, Typen, Operationen und Durchläufe, die erforderlich sind, um SparseTensor-Typen zu erstklassigen Elementen in der MLIR-Compilerinfrastruktur zu machen. Der Dialekt bildet eine Brücke zwischen High-Level-Vorgängen für Typen von dünnbesetzten Tensoren und Low-Level-Vorgängen für die tatsächlichen dünnbesetzten Speicherschemas, die aus Positionen, Koordinaten und Werten bestehen. Die Unterstützung auf niedrigerer Ebene kann aus vollständig generiertem Code bestehen oder durch eine kleine, spärliche Laufzeit-Supportbibliothek bereitgestellt werden.

Die MLIR-Implementierung [Biketal22] folgt eng der „Theorie der spärlichen Iteration“, die die Grundlage von TACO bildet. Der Compiler hat eine Umschreibregel für jeden Tensorausdruck im Linalg-Dialekt (Tensorindexnotation von MLIR): Die Sparsity jeder Ebene eines Tensors wird durch die Ebenentypen (z.B. dicht, komprimiert, Singleton) und eine Spezifikation der Reihenfolge auf den Ebenen angegeben (siehe [Chou18] für eine ausführliche Erläuterung und mögliche Erweiterungen dieser Ebenentypen). Für jeden Tensorausdruck in der Eingabe erstellt der Compiler zuerst einen topologisch sortierten Iterationsgraphen, der die Reihenfolge der Koordinaten in Bezug auf die Ebenen der einzelnen Tensoren widerspiegelt. So wird sichergestellt, dass alle Tensoren in der natürlichen Reihenfolge der Ebenenkoordinaten besucht werden. Als Nächstes werden für jeden Index in topologischer Reihenfolge Iterations-Lattices erstellt. Jeder Iterationsgitterpunkt besteht aus einer Konjunktion von Tensorkoordinaten und einem Tensor-(Unter-)Ausdruck, der für diese Konjunktion ausgewertet werden muss. Innerhalb des Gitters werden die Iterationspunkte entsprechend der Reihenfolge angeordnet, in der die Koordinaten durchlaufen werden. Diese Iterationsgitter steuern die eigentliche Generierung von spärlichem Code, eine relativ einfache Eins-zu-eins-Zuordnung von Iterationsgittern zu Kombinationen von For-Schleifen, While-Schleifen und If-Anweisungen. Ausgaben für dünnbesetzte Tensoren, die nicht initialisiert werden, werden entweder durch direkte Einfügungen verarbeitet, wenn alle parallelen Schleifen die äußersten sind, oder durch Einfügungen, die indirekt über eine 1-dimensionale Erweiterung des Zugriffsmusters (auch als Arbeitsbereich bezeichnet) erfolgen, sofern dies möglich ist [Gustavson72,Bik96,Kjolstad19].

  • [Bik96] Aart J.C. Bik. Compiler-Unterstützung für Berechnungen mit dünnbesetzten Matrizen. Doktorarbeit, Universität Leiden, Mai 1996.

  • [Biketal22] Aart J.C. Bik, Penporn Koanantakool, Tatiana Shpeisman, Nicolas Vasilache, Bixia Zheng und Fredrik Kjolstad. Compilerunterstützung für Berechnungen mit spärlichen Tensoren in MLIR. ACM Transactions on Architecture and Code Optimization, Juni 2022. Weitere Informationen: https://dl.acm.org/doi/10.1145/3544559

  • [Chou18] Stephen Chou, Fredrik Berg Kjolstad und Saman Amarasinghe. Formatabstraktion für Compiler für spärliche Tensoralgebra. Proceedings of the ACM on Programming Languages, Oktober 2018.

  • [Chou20] Stephen Chou, Fredrik Berg Kjolstad und Saman Amarasinghe. Automatische Generierung von effizienten Routinen zur Konvertierung in das Sparse-Tensor-Format. Proceedings of the 41st ACM SIGPLAN Conference on Programming Language Design and Implementation, Juni 2020.

  • [Gustavson72] Fred G. Gustavson. Einige grundlegende Techniken zum Lösen dünnbesetzter linearer Gleichungssysteme. In Sparse Matrices and Their Applications, Seiten 41–52. Plenum Press, New York, 1972.

  • [Kjolstad17] Fredrik Berg Kjolstad, Shoaib Ashraf Kamil, Stephen Chou, David Lugato und Saman Amarasinghe. Der Tensor Algebra Compiler. Proceedings of the ACM on Programming Languages, Oktober 2017.

  • [Kjolstad19] Fredrik Berg Kjolstad, Peter Ahrens, Shoaib Ashraf Kamil und Saman Amarasinghe. Tensor Algebra Compilation with Workspaces, Proceedings of the IEEE/ACM International Symposium on Code Generation and Optimization, 2019.

  • [Kjolstad20] Fredrik Berg Kjolstad. Kompilierung der Algebra für dünnbesetzte Tensoren. Dissertation, MIT, Februar 2020.

Möchten Sie mitmachen? Hier sind einige gute erste Probleme.

  1. Wir bevorzugen jetzt den Begriff „Sparsifier“ gegenüber dem ebenfalls häufig verwendeten Begriff „Sparse Compiler“, um auf einen solchen Pass zu verweisen. So soll deutlich gemacht werden, dass der Sparsifier-Pass kein separater Compiler ist, sondern ein integraler Bestandteil jeder Compiler-Pipeline sein sollte, die mit der MLIR-Compiler-Infrastruktur erstellt wird. ↩