Wprowadzenie

Traktowanie rzadkości jako właściwości zamiast (żmudnego) szczegółu implementacji pozwala kompilatorowi1 automatycznie generować rzadki kod. Koncepcja ta została wprowadzona w algebrze liniowej przez [Bik96] w MT1, a następnie sformalizowana w algebrze tensorowej przez [Kjolstad17,Kjolstad20] w ramach projektu kompilatora algebry tensorów rzadkich(TACO).

Dialekt SparseTensor obsługuje wszystkie atrybuty, typy, operacje i przekształcenia, które są wymagane, aby typy tensorów rzadkich były traktowane jako elementy podstawowe w infrastrukturze kompilatora MLIR. Dialekt ten stanowi pomost między operacjami wysokiego poziomu na typach tensorów rzadkich a operacjami niższego poziomu na rzeczywistych schematach przechowywania rzadkiego, które składają się z pozycji, współrzędnych i wartości. Obsługa na niższym poziomie może obejmować w pełni wygenerowany kod lub być zapewniana za pomocą małej, rzadkiej biblioteki obsługi środowiska wykonawczego.

Implementacja MLIR [Biketal22] jest ściśle zgodna z „teorią iteracji rzadkich”, która stanowi podstawę TACO. Kompilator ma regułę przepisywania dla każdego wyrażenia tensora w dialekcie Linalg (notacja indeksu tensora MLIR): rzadkość każdego poziomu tensora jest wskazywana przez typy poziomów (np. gęsty, skompresowany, pojedynczy) i specyfikację kolejności na poziomach (szczegółowe omówienie i możliwe rozszerzenia tych typów poziomów znajdziesz w [Chou18]). W przypadku każdego wyrażenia tensora w danych wejściowych kompilator najpierw tworzy posortowany topologicznie wykres iteracji, który odzwierciedla kolejność współrzędnych względem poziomów każdego tensora. Dzięki temu wszystkie tensory są odwiedzane w naturalnej kolejności poziomów i współrzędnych. Następnie dla każdego indeksu w porządku topologicznym konstruowane są kraty iteracji. Każdy punkt siatki iteracji składa się z koniunkcji współrzędnych tensora i wyrażenia tensorowego (podwyrażenia), które należy obliczyć dla tej koniunkcji. W siatce punkty iteracji są uporządkowane zgodnie z kolejnością wyczerpywania się współrzędnych. W ten sposób te siatki iteracji generują rzeczywisty kod rzadki, który jest stosunkowo prostym mapowaniem jeden do jednego z siatek iteracji na kombinacje pętli for, pętli while i instrukcji if. Wyjścia tensora rzadkiego, które materializują niezainicjowane wartości, są obsługiwane przez bezpośrednie wstawienia, jeśli wszystkie pętle równoległe są najbardziej zewnętrzne, lub przez wstawienia, które pośrednio przechodzą przez 1-wymiarowe rozszerzenie wzorca dostępu (tzw. przestrzeń roboczą), jeśli jest to możliwe [Gustavson72,Bik96,Kjolstad19].

  • [Bik96] Aart J.C. Bik. Compiler Support for Sparse Matrix Computations. Praca doktorska, Uniwersytet w Lejdzie, maj 1996 r.

  • [Biketal22] Aart J.C. Bik, Penporn Koanantakool, Tatiana Shpeisman, Nicolas Vasilache, Bixia Zheng i Fredrik Kjolstad. Obsługa kompilatora obliczeń na tensorach rzadkich w MLIR. ACM Transactions on Architecture and Code Optimization, czerwiec 2022 r. Więcej informacji: https://dl.acm.org/doi/10.1145/3544559

  • [Chou18] Stephen Chou, Fredrik Berg Kjolstad i Saman Amarasinghe. Format Abstraction for Sparse Tensor Algebra Compilers. Proceedings of the ACM on Programming Languages, październik 2018 r.

  • [Chou20] Stephen Chou, Fredrik Berg Kjolstad i Saman Amarasinghe. Automatyczne generowanie wydajnych procedur konwersji formatu tensora rzadkiego. Proceedings of the 41st ACM SIGPLAN Conference on Programming Language Design and Implementation, czerwiec 2020 r.

  • [Gustavson72] Fred G. Gustavson. Podstawowe techniki rozwiązywania rzadkich układów równań liniowych. W książce Sparse Matrices and Their Applications, strony 41–52. Plenum Press, Nowy Jork, 1972.

  • [Kjolstad17] Fredrik Berg Kjolstad, Shoaib Ashraf Kamil, Stephen Chou, David Lugato i Saman Amarasinghe. Kompilator algebry tensorowej. Proceedings of the ACM on Programming Languages, październik 2017 r.

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

  • [Kjolstad20] Fredrik Berg Kjolstad. Kompilacja algebry tensorów rzadkich. Praca doktorska, MIT, luty 2020 r.

Chcesz pomóc? Oto kilka dobrych pierwszych zadań.

  1. Zwróć uwagę, że zamiast terminu „kompilator rzadki” (ang. sparse compiler), który jest również powszechnie używany, wolimy teraz używać terminu „sparsifier” (ang. sparsifier), aby wyraźnie zaznaczyć, że ten etap nie jest osobnym kompilatorem, ale powinien być integralną częścią każdego potoku kompilatora zbudowanego przy użyciu infrastruktury kompilatora MLIR. ↩