NVIDIA Omniverse NuRec egy neurális rekonstrukciós pipeline, amely többforrású szenzoradatokból (kamerák, lidar) épít nagy hűségű, háromdimenziós reprezentációkat valós környezetről. Az eszközt autonóm jármű (AV) és robotikai platformok által rögzített dinamikus jelenetek rekonstruálására használják, hogy azok szimulációra kész digitális környezetekké váljanak, melyek az NVIDIA Omniverse és kapcsolódó szimulációs munkafolyamatokban renderelhetők, lejátszhatók és elemezhetők.
Ezek a rekonstrukciók fontos szerepet játszanak a fizikai mesterséges intelligencia és autonóm rendszerek fejlesztésében: a mérnökök egy valós vezetési vagy robotikai forgatókönyvet rögzíthetnek, rekonstruálhatják a környezetet, majd megvizsgálhatják vagy újrajátszhatják a jelenetet a modellviselkedés megértéséhez, észlelési eredmények validálásához, szintetikus nézőpontok generálásához vagy tanítószett előállításához.
A NuRec nagy realizmust kombinál GPU-gyorsított rendereléssel és neurális rendering technikákkal, például Gaussian splattinggel, de ez jelentős számítási költséggel jár: nagy adatmennyiségek, komplex PyTorch-alapú tanulási hurkok és speciális CUDA-kernelek terhelik a GPU-kat.
Célok és kiinduló állapot
A teljesítmény kritikus a NuRec munkafolyamatoknál, mert a rekonstrukciós átfutási idő közvetlenül befolyásolja a mérnöki iterációs sebességet. Tipikus felhasználási eset például egy érdekes vagy problémás AV-felvétel gyors rekonstruálása a hibakeresés gyorsítása érdekében; ha a rekonstrukció órákat vesz igénybe, az jelentősen lassítja a munkát. A kezdeti állapotban viszonylag rövid felvételek rekonstruálása is órákat vehetett igénybe a jelenettől és konfigurációtól függően. A csapat hosszú távú célja ambiciózus: valós idejű rekonstrukció, ahol egy 30 másodperces felvételt körülbelül 30 másodperc alatt lehet rekonstruálni.
A teljesítményjavítás nem csak magához a rekonstruáláshoz fontos: a rekonstruált jelenetek renderelési munkafolyamatai (pl. megerősítő tanuláshoz, szintetikus adatgeneráláshoz vagy nagyszabású szimulációkhoz) rengeteg képkockát hozhatnak létre, így akár kisebb sebességnövekedés is jelentős megtakarítást eredményezhet GPU-idő és infrastruktúra-költség tekintetében.
Profilozás Nsight Systems-szel — hol a szűk keresztmetszet?
A NuRec optimalizálásához NVIDIA profilozó és optimalizáló eszközöket használtak, elsősorban az NVIDIA Nsight Systems és az NVIDIA Nsight Compute eszközöket. Az első lépés egy Nsight Systems profil lefuttatása volt a kiindulási állapot megértéséhez.
A PyTorch-ban található NVTX annotációkat használva a csapat egyetlen forward iterációra szűkített (forward pass) profilokat készített. Az első feltételezés az volt, hogy a renderelő kernel viszi a futásidő nagy részét, azonban a CUDA hardver idősora azt mutatta, hogy a GPU gyakran alulhasznosított vagy egyáltalán nem használt — sok apró kernel volt jelen, és hiányzott a folyamatos GPU-terhelés.
További NVTX annotációk bevezetésével kiderült, hogy a collect_gaussian_parameters függvény sok időt vett el a renderelés kezdete előtt, és többször is meghívódott egy forward pass során. Mélyebbre ásva az interpolate függvény tűnt időigényesnek (~4.148 ms), amely sok kis kernelt és memóriaműveletet hívott meg, elnyomva a GPU hatékony kihasználását.
Kis kernel fuzionálása és szinkronizációk eltávolítása
Az interpolate alatt található kódot átvizsgálva a csapat a kis kernelfeladatokat összefűzve egyetlen nagyobb kernelbe tömörítette. Ennek eredményeként az interpolate függvény futása 4.184 ms-ról 83.81 µs-ra csökkent — ez közel 50× gyorsulás.
Ezt követően hosszú cudaStreamSynchronize hívások kerültek azonosításra, amelyek megakadályozták, hogy a CPU gyorsan elindítson több kernel-láncot, miközben a GPU foglalt volt. Egy szinkronizáció eltávolítása után mindig másik pont vált bottleneckké, így a csapat iteratívan eltávolította a fölösleges szinkronizációkat, amíg a CPU képes lett hatékonyan sorba állítani a munkát a GPU ideje alatt. Ennek hatására a korábban apró kernelhívások már nem a CPU indítási idejéhez voltak kötve, és a GPU-terhelés kompaktabbá, folyamatosabbá vált.
Hot-kernellek azonosítása és renderBackward optimalizálása Nsight Compute-tal
Az Nsight Systems segített megtalálni a forrást, majd az Nsight Compute segítségével egyes kernellek részletes profilozása következett. A renderBackward kernel mind kamerás, mind lidar feldolgozásnál használatos volt és a profillal kiderült, hogy csak körülbelül 15% foglaltságot (occupancy) ért el. Továbbá a kernel viselkedése és erőforrásigénye jelentősen eltér attól függően, hogy lidar vagy kamera adaton fut.
A hosszabb három renderBackward példa lidar adathoz tartozott, míg a másik három kamerához. Mindkét esetben azonban statikusan 167 regisztert allokáltak szálanként. Az Nsight Compute összehasonlítása megmutatta, hogy bár a megosztott memória-használat nagy volt mindkét esetben, a kamera esetén a shared memory hozzáférések száma mintegy 75%-kal kevesebb volt, mégis ugyanannyi memóriát foglalt a blokk.
Ennek megfelelően a csapat úgy döntött, hogy szétválasztja a kernel-verziókat attól függően, hogy lidar vagy kamera adatot dolgoznak-e fel. Mindkét verzió esetén finomhangolták a regiszterhasználatot a launch_bounds minősítővel és a blokkonként allokált shared memory mennyiségét. A cudaFuncSetCacheConfig futási API-t használva a preferenciát a nagyobb shared memory és kisebb L1 cache irányába állították.
A tuning után a lidar és kamera kernel regiszterigénye 167-ről rendre 64-re és 128-ra csökkent, és mindkettő hatékonyan futott körülbelül a korábbi megosztott memória felével. Ennek eredményeként az occupancy ~15%-ról 30–50%-ra nőtt, és a futási idő jelentősen javult: a leghosszabb lidar kernel futási ideje 31 ms-ról 18 ms-ra csökkent.
Hátralevő feladatok: long-tail és munkaterhelet egyenlőtlenség
A szerzők megjegyzik, hogy további javítási lehetőségek vannak: az egyik jelenlegi probléma a kernel hosszú farok (long-tail) hatása, amely munkaterheleti egyenlőtlenségre utal. Az Nsight Compute PM Sampling része azt mutatja, hogy a kernel első felében átlagosan 32 aktív warp van, majd idővel fokozatosan csökken az aktivitás, és az utolsó néhány milliszekundumban kevesebb mint egy aktív warp van ciklusonként. Ideálisan minden warp aktív maradna a teljes kernelidő alatt.
Hol kezdje el az érdeklődő?
A teljesítmény-analízis és optimalizálás iteratív folyamat: profil futtatása, probléma azonosítása, javítás, majd újraprofilozás. Az NVIDIA Nsight Systems és Nsight Compute eszközök megkönnyítik ezt a munkát: mindkettő ingyenesen elérhető. A szerzők arra buzdítanak, hogy töltse le az eszközöket és próbálja ki saját esetén, és ha kérdések vagy megosztanivaló tapasztalat van, írjon a NVIDIA Developer Forums-ra.
Elismerések
Külön köszönet Francois Trudelnek, Joey Lainak és Rodolfo Limának, akik hozzájárultak a munkához.
