Výkonné výpočty s mnohem menším kódem

Mnoho společností investuje značné prostředky do najímání talentů pro tvorbu vysoce výkonných knihoven kódu, které jsou základem moderních systémů umělé inteligence. NVIDIA například vyvinula některé z nejpokročilejších knihoven pro vysoce výkonné výpočty (HPC), čímž si vytvořila konkurenční výhodu, kterou se ostatním těžko daří překonat.

Ale co kdyby pár studentů dokázalo během několika měsíců konkurovat špičkovým knihovnám HPC s pouhými několika stovkami řádků kódu místo desítek či stovek tisíc?

Právě to ukázali výzkumníci z MIT Computer Science and Artificial Intelligence Laboratory (CSAIL) s novým programovacím jazykem nazvaným Exo 2.

Exo 2 patří do nové kategorie programovacích jazyků, které profesor MIT Jonathan Ragan-Kelley nazývá „jazyky s uživatelsky definovatelným plánováním“ (USL). Namísto spoléhání se na to, že neprůhledný kompilátor automaticky vygeneruje nejrychlejší možný kód, USL dávají programátory do řídící pozice, čímž jim umožňuje psát „plány“, které explicitně řídí, jak kompilátor generuje kód. To umožňuje inženýrům zabývajícím se výkonem transformovat jednoduché programy, které specifikují, co se má počítat, na komplexní programy, které dělají totéž jako původní specifikace, ale mnohem, mnohem rychleji.

Jedním z omezení stávajících USL (jako je původní Exo) je jejich relativně pevná sada operací plánování, což ztěžuje opětovné použití kódu plánování napříč různými „jádry“ (jednotlivé komponenty ve vysoce výkonné knihovně).

Naopak Exo 2 umožňuje uživatelům definovat nové operace plánování externě ke kompilátoru, čímž usnadňuje vytváření znovu použitelných knihoven plánování. Hlavní autorka Yuka Ikarashi, doktorandka MIT v oboru elektrotechniky a informatiky a spolupracovnice CSAIL, říká, že Exo 2 může snížit celkový kód plánování o faktor 100 a dosáhnout výkonu srovnatelného s implementacemi špičkové úrovně na několika různých platformách, včetně Basic Linear Algebra Subprograms (BLAS), které pohánějí mnoho aplikací strojového učení. To z něj dělá atraktivní možnost pro inženýry v oblasti HPC zaměřené na optimalizaci jader napříč různými operacemi, datovými typy a cílovými architekturami.

„Je to přístup zdola nahoru k automatizaci, spíše než provádění vyhledávání ML/AI v kódu s vysokým výkonem,“ říká Ikarashi. „To znamená, že inženýři zabývající se výkonem a implementátoři hardwaru si mohou napsat vlastní knihovnu plánování, která představuje sadu optimalizačních technik, které se mají aplikovat na jejich hardware pro dosažení špičkového výkonu.“

Jednou z hlavních výhod Exo 2 je, že snižuje množství kódovacího úsilí potřebného v daném okamžiku opětovným použitím kódu plánování napříč aplikacemi a hardwarovými cíli. Výzkumníci implementovali knihovnu plánování s přibližně 2 000 řádky kódu v Exo 2, která zapouzdřuje znovu použitelné optimalizace, které jsou specifické pro lineární algebru a cílové zařízení (zrychlovače hardwaru AVX512, AVX2, Neon a Gemmini). Tato knihovna konsoliduje úsilí o plánování napříč více než 80 vysoce výkonnými jádry s až tuctem řádků kódu, čímž dosahuje výkonu srovnatelného s, nebo lepšího než MKL, OpenBLAS, BLIS a Halide.

Exo 2 zahrnuje nový mechanismus nazvaný „kurzory“, který poskytuje takzvanou „stabilní referenci“ pro ukazování na objektový kód během celého procesu plánování. Ikarashi říká, že stabilní reference je nezbytná pro uživatele, aby zapouzdřili plány do funkce knihovny, protože to činí kód plánování nezávislým na transformacích objektového kódu.

„Věříme, že USL by měly být navrženy tak, aby byly uživatelsky rozšiřitelné, spíše než aby měly pevnou sadu operací,“ říká Ikarashi. „Tímto způsobem se jazyk může rozrůstat tak, aby podporoval velké projekty prostřednictvím implementace knihoven, které vyhovují různým požadavkům na optimalizaci a aplikačním oborům.“

Konstrukce Exo 2 umožňuje inženýrům zabývajícím se výkonem soustředit se na strategie optimalizace na vysoké úrovni a zároveň zajistit, aby základní objektový kód zůstal funkčně ekvivalentní díky použití bezpečných primitiv. V budoucnu tým doufá, že rozšíří podporu Exo 2 pro různé typy hardwarových zrychlovačů, jako jsou GPU. Několik probíhajících projektů si klade za cíl zlepšit samotnou analýzu kompilátoru, a to z hlediska správnosti, doby kompilace a výrazové síly.

Ikarashi a Ragan-Kelley spoluautoři článku s postgraduálními studenty Kevinem Qianem a Samirem Droubim, Alexem Reinkingem z Adobe a bývalým postdoktorandem CSAIL Gilbertem Bernsteinem, nyní profesorem na University of Washington. Tento výzkum byl částečně financován agenturou U.S. Defense Advanced Research Projects Agency (DARPA) a U.S. National Science Foundation, zatímco první autorka byla také podporována stipendii Masason, Funai a Quad.