اپراتورهای عصبی فوریه (FNO) ابزارهای قدرتمندی برای یادگیری اپراتورهای حل معادلات دیفرانسیل جزئی هستند، اما فاقد بهینهسازیهای آگاه از معماری هستند. لایه فوریه آنها FFT، فیلتر کردن، GEMM، صفرگذاری و iFFT را به عنوان مراحل جداگانه اجرا میکند که منجر به اجرای چند هسته و ترافیک بیش از حد حافظه سراسری میشود. الگوی محاسباتی FFT -> GEMM -> iFFT در رابطه با ادغام هسته GPU و بهینهسازی طرحبندی حافظه، توجه کافی را دریافت نکرده است. روشهای فعلی مانند Quantum ESPRESSO، Octopus و CP2K تماسهای جداگانهای با روالهای FFT و BLAS برقرار میکنند. با این حال، آنها دارای سه محدودیت هستند: استفاده جزئی از فرکانس با عملیات کپی حافظه اضافی، فقدان قابلیتهای فیلتر کردن فرکانس بومی در cuFFT و تراکنشهای بیش از حد حافظه بین مراحل پردازش.
FNO یک خط لوله را پیادهسازی میکند که با یک FFT پیشرو روی نقشههای ویژگی ورودی شروع میشود، فیلتر طیفی را اعمال میکند و خروجی را از طریق FFT معکوس بازسازی میکند. این فرآیند مستلزم مراحل برش دامنه فرکانس و صفرگذاری است که فریمورکهای فعلی مانند PyTorch به دلیل محدودیتهای cuFFT در پشتیبانی بومی از برش ورودی/خروجی، به عنوان هستههای کپی حافظه جداگانه اجرا میکنند. کتابخانههای FFT پیشرو مانند cuFFT و VkFFT فاقد قابلیتهای برش داده داخلی هستند. FFTهای دوبعدی سنتی هر دو مرحله 1D-FFT را در امتداد ابعاد فضایی اعمال میکنند، اما FNO وزنهای طیفی را در سراسر بعد کانال اعمال میکند، که نشاندهنده فرصتی برای جدا کردن مراحل FFT با نگه داشتن اولین FFT یکبعدی در امتداد محورهای فضایی و تفسیر مجدد مرحله دوم FFT در امتداد بعد پنهان است.
محققان دانشگاه کالیفرنیا، ریورساید، ایالات متحده آمریکا، TurboFNO را پیشنهاد کردهاند، اولین هسته GPU کاملاً ادغامشده FFT-GEMM-iFFT با بهینهسازیهای FFT داخلی. این رویکرد با توسعه هستههای FFT و GEMM از ابتدا آغاز میشود که به عملکردی قابل مقایسه یا سریعتر از cuBLAS و cuFFT منبع بسته دست مییابند. یک نوع FFT برای ادغام موثر بارهای کاری FFT و GEMM معرفی شده است که در آن یک بلوک رشته واحد روی بعد پنهان تکرار میشود، و با حلقه k در GEMM همسو میشود. علاوه بر این، دو الگوی چرخشی حافظه مشترک برای دستیابی به 100٪ استفاده از بانک حافظه هنگام ارسال خروجی FFT به GEMM و فعال کردن iFFT برای بازیابی نتایج GEMM به طور مستقیم از حافظه مشترک طراحی شدهاند.
TurboFNO پیادهسازیهای بهینهشده هستههای FFT و CGEMM را برای فعال کردن ادغام موثر و بهینهسازیهای FFT داخلی ادغام میکند. استراتژی ادغام هسته در TurboFNO از طریق سه سطح پیشرفت میکند: ادغام FFT-GEMM، ادغام GEMM-iFFT و ادغام کامل FFT-GEMM-iFFT. هر مرحله شامل تراز کردن گردش کار FFT با GEMM، حل ناهماهنگیهای طرحبندی داده و حذف تضادهای بانک حافظه مشترک است. تکنیکهای کلیدی شامل اصلاح طرحبندی خروجی FFT برای مطابقت با فرمت ورودی GEMM، اعمال چرخاندن رشته برای دسترسی بدون تضاد به حافظه مشترک و ادغام FFT معکوس به عنوان یک مرحله پایانی CGEMM برای دور زدن نوشتن حافظه سراسری میانی و افزایش مکانیابی حافظه است.
TurboFNO عملکرد عالی را در ارزیابیهای 1D و 2D FNO نشان میدهد. در تستهای 1D FNO، گردش کار بهینهشده FFT-CGEMM-iFFT به افزایش سرعت تا 100٪ نسبت به PyTorch دست مییابد و به طور متوسط 50٪ بهبود مییابد. این دستاوردها از هرس FFT ناشی میشود که محاسبات را 25٪ تا 67.5٪ کاهش میدهد. هسته کاملاً ادغامشده FFT-CGEMM-iFFT تا 150٪ سرعت بیشتری نسبت به PyTorch ارائه میدهد و 10٪ تا 20٪ بهبود اضافی نسبت به استراتژیهای ادغام جزئی ارائه میدهد. به طور مشابه، در 2D FNO، گردش کار بهینهشده بهتر از PyTorch با افزایش سرعت متوسط بالای 50٪ و حداکثر بهبودهای 100٪ عمل میکند. هسته کاملاً ادغامشده 2D به 50٪ -105٪ سرعت بیشتری نسبت به PyTorch بدون کاهش عملکرد دست مییابد، علیرغم سربار اضافی تراز کردن طرحبندی بار کاری FFT با جریان داده CGEMM.
در این مقاله، محققان TurboFNO را معرفی کردند، اولین هسته GPU کاملاً ادغامشده که FFT، CGEMM و iFFT را برای تسریع اپراتورهای عصبی فوریه ادغام میکند. آنها یک سری بهینهسازیهای آگاه از معماری را برای غلبه بر ناکارآمدیها در پیادهسازیهای FNO معمولی، مانند اجرای بیش از حد هسته و ترافیک حافظه سراسری، توسعه دادند. اینها شامل یک هسته FFT سفارشی با فیلتر کردن فرکانس داخلی و صفرگذاری، یک نوع FFT سازگار با GEMM که رفتار حلقه k را تقلید میکند، و استراتژیهای چرخشی حافظه مشترک که استفاده از بانک را از 25٪ به 100٪ بهبود میبخشد، میشود. TurboFNO به افزایش سرعت تا 150٪ دست مییابد و به طور متوسط 67٪ سود عملکرد را در تمام تنظیمات آزمایششده حفظ میکند.