HPC-kääntäjien tulevaisuus: Miten NektarIR ratkaisee heterogeenisen raudan haasteet

Jos olet seurannut high-performance computingin maailmaa viime aikoina, olet todennäköisesti huomannut huolestuttavan kehityssuunnan: laitteistokirjo on sirpaloitunut pahasti. Nykyaikaiset supertietokoneet ja jopa tehokkaat työasemat hyödyntävät heterogeenisiä arkkitehtuureja – perinteisten CPU-prosessorien rinnalla on GPU-kiihdyttimiä, mahdollisesti FPGA-piirejä ja eri valmistajien erikoiskiihdyttimiä. Jokainen näistä puhuu eri kieltä optimoinnin suhteen, ja oman koodin tehokas ajaminen kaikilla alustoilla voi tuntua mahdottomalta tehtävältä.

Tähän ongelmaan NektarIR tarjoaa ratkaisun, ja on virkistävää nähdä jonkun tarttuvan siihen systemaattisesti sen sijaan että vain kirjoitettaisiin lisää käsinoptimoitua assembly-koodia.

Käännösputken vallankumous

Perinteiset kääntäjät kuten GCC tai LLVM toimivat melko korkealla tasolla – ne muuttavat lähdekoodin optimoituksi konekieleksi. Mutta kun työskennellään toimialakohtaisten operaatioiden parissa, olipa kyse sitten finite element -menetelmistä tai koneoppimisen ytimekkyistä, syntyy valtava semanttinen kuilu sen välille, mitä toimialan asiantuntija ajattelee (differentiaalioperaattorit, alkio-matriisit) ja mitä rauta oikeasti ymmärtää (rekisterit, muistihierarkiat, SIMD-kanavat).

MLIR, Googlen kehittämä ja nykyään LLVM:n kulmakivi, esitti nerokkaan ajatuksen: miksi ei tukisi useita eri tasojen väliaikaisia esityksiä sen sijaan että olisi yksi monoliittinen ratkaisu? Jokainen taso voidaan virittää tietylle abstraktiotasolle. Voit määritellä omia "murteita" (dialects), jotka vangitsevat suoraan toimialakohtaiset käsitteet, ja sitten laskea ne asteittain kohti laitteistokohtaisia käskyjä samalla kun optimointeja tehdään jokaisella tasolla.

NektarIR soveltaa tätä filosofiaa spektraalisiin/hp-elementtimenetelmiin virtausdynamiikan laskennassa. Jos et ole aiemmin törmännyt spektraalisiin elementtimenetelmiin, ajattele niitä kehittyneenä numeerisena tekniikkana osittaisdifferentiaaliyhtälöiden ratkaisemiseen – ne ovat erityisen suosittuja ilmailu- ja säämallinnuksessa, koska ne saavuttavat korkean tarkkuuden suhteellisen karkealla hilarakenteella.

Miksi elementtioperaatiot ovat erilaisia

Tässä on se seikka, joka tekee finite element -menetelmistä kiinnostavia käännöksen kannalta: suorituskyvyn määrää yllättävän pieni joukko yleisiä operaattoreita, joita yhdistellään eri tavoin monimutkaisten differentiaalioperaattoreiden luomiseksi. Puhumme operaatioista kuten kantafunktioiden evaluointi, matriisi-vektoritulokset ja kvadratuurilaskenta – ne näyttävät yksinkertaisilta mutta ovat äärimmäisen herkkiä muistiakteerausmalleille ja välimuistin käytölle.

Perinteisissä lähestymistavoissa kehittäjät kirjoittivat optimoituja ytimekkyksiä jokaiselle arkkitehtuurille ja ylläpitivät rinnakkaisia versioita CPU- ja GPU-alustoille. Tuloksena on ylläpitokauhu, jossa yksittäinen algoritminen muutos voi vaatia kymmenien toimittajakohtaisten toteutusten päivittämistä.

NektarIR kiertää tämän ongelman luomalla oman MLIR-murteensa, joka esittää suoraan nämä elementtioperaatiot. Ajattele sitä sanastona, joka on suunniteltu nimenomaan spektraalisten elementtilaskelmien kuvaamiseen. Kun tämä toimialakohtainen kieli on kääntäjäinfrastruktuurin ytimessä, voidaan kirjoittaa optimointeja, jotka ymmärtävät mitä oikeasti ollaan tekemässä – eivät vain raakoja operaatioita.

Asteittaisen laskemisen taika

Todellinen voima piilee laskentaputkessa (lowering pipeline). NektarIR:ssä lähdetään korkean tason esityksestä, joka näyttää hyvin samankaltaiselta kuin miten matemaatikko kuvailisi elementtioperaattoria. Kun esitys siirtyy läpi peräkkäisten vaiheiden, se muuttuu yhä konkreettisemmaksi:

1. Toimialakohtainen murre vangitsee matemaattisen rakenteen
2. Bufferisaatio käsittelee muistin allokaatiota ja datan siirtoa
3. Laitteistokohtainen laskenta mukauttaa operaatiot CPU- tai GPU-suoritusmalleihin
4. Lopullinen koodigenerointi tuottaa optimoituja konekäskyjä

Jokaisella tasolla voidaan soveltaa kyseiselle abstraktiotasolle spesifisiä optimointeja. Tämä on huomattavasti tehokkaampaa kuin yritys jälkikäteen selvittää korkean tason tarkoitus alhaisen tason operaatioista – jotain, mihin perinteiset kääntäjät kamppailevat toimialakohtaisen koodin kanssa.

Todellisia tuloksia

NektarIR-tiimi demonstroi lähestymistapaansa vertaamalla sitä vakiintuneeseen Nektar++ spektraaliseen/hp-elementtikehykseen. Tulokset osoittavat, että heidän just-in-time -käännösmenetelmänsä saavuttaa kilpailukykyisen tai paremman suorituskyvyn sekä CPU- että GPU-arkkitehtuureilla. Tämä on vaikuttavaa, kun ottaa huomioon että he työskentelevät yhdestä yhtenäisestä koodikannasta käsin sen sijaan että jokaiselle alustalle olisi erikseen optimoituja toteutuksia.

Tämä on erityisen merkittävää, koska CFD-simulaatiot pyörivät usein päiviä tai viikkoja supertietokoneilla. Jopa pienet prosentuaaliset parannukset ytimekkyiden suorituskyvyssä muuttuvat merkittäviksi ajan ja energian säästöiksi skaalautuessa.

Mitä tämä tarkoittaa kehittäjille

Saatat ajatella, että "tämä on mielenkiintoista, mutta en kirjoita CFD-ratkaisuja." Reilu pointti – mutta taustalla oleva tekniikka on laajalti sovellettavissa. Toimialakohtaisten murteiden määrittely, asteittainen laskenta abstraktiotasojen läpi ja kohdennettujen optimointien soveltaminen jokaisella tasolla on käyttökelpoista paljon finite element -menetelmiä laajemmin.

Näemme jo MLIR:n käyttöä TensorFlow'n ML-kuormissa, IREE:n sulautetussa AI-inferenssissä ja erilaisissa tutkimusprojekteissa kvanttilaskennassa ja signaalinkäsittelyssä. Kun erikoisrauta yleistyy, kyky kirjoittaa koodia kerran ja kääntää se tehokkaasti eri kohteisiin muuttuu yhä arvokkaammaksi.

NektarIR-projekti osoittaa, että olemme siirtymässä aikakauteen, jossa kääntäjät voidaan todella suunnitella yhdessä kohdealueidensa kanssa. Sen sijaan että pakotettaisiin toimialan asiantuntijoita oppimaan kääntäjätekniikkaa tai odotettaisiin kääntäjäkehittäjien ymmärtävän spektraalisten elementtimenetelmien vivahteita, voimme rakentaa työkaluja jotka siltaavat tämän kuilun tyylikkäästi.

HPC-yhteisölle tämä edustaa tietä eteenpäin nykytilanteesta, jossa sovellusten siirtäminen uusille arkkitehtuureille vaatii sankarillisia insinööriponnistuksia. Kehittäjille yleisesti tämä antaa esinäkymän siihen, miten käännöskirjallisuus voi kehittyä käsittelemään yhä heterogeenisempää laskennan tulevaisuutta.

Koodi on saatavilla kokeilua varten, ja lähestymistapa on yleistettävissä niin että samankaltaisia tekniikoita voitaisiin soveltaa muihinkin laskenta-intensiivisiin toimialoihin. Simuloidaanpa sitten ilmastomallinnusta tai yritetään puristada lisää suorituskykyä GPU-kiihdytetystä sovelluksesta, NektarIR:n taustalla olevat ideat kannattaa ymmärtää.

Haluatko jakaa kokemuksiasi heterogeenisen raudan haasteista omissa projekteissasi? Mikä on näkemyksesi toimialakohtaisista käännösmenetelmistä? Kommentoi alle – kuulemme mielellämme ajatuksiasi.