XLA (Accelerated Linear Algebra) یک کامپایلر یادگیری ماشینی (ML) است که جبر خطی را بهینه می‌کند و باعث بهبود سرعت اجرا و استفاده از حافظه می‌شود. این صفحه یک نمای کلی از اهداف و معماری کامپایلر XLA ارائه می دهد.
اهداف
امروزه XLA از چندین فریم ورک ML (از جمله PyTorch، TensorFlow و JAX) پشتیبانی می‌کند و بخشی از پروژه OpenXLA است - یک اکوسیستم از فناوری‌های کامپایلر منبع باز برای ML که به طور مشترک توسط سازمان‌های سخت‌افزار و نرم‌افزار ML پیشرو توسعه داده شده است. قبل از ایجاد پروژه OpenXLA، XLA در داخل پروژه TensorFlow توسعه داده شد، اما اهداف اساسی یکسان است:
بهبود سرعت اجرا زیرگراف ها را برای کاهش زمان اجرای عملیات کوتاه مدت و حذف سربار از زمان اجرا، ترکیب عملیات خط لوله برای کاهش سربار حافظه، و تخصصی کردن اشکال تانسور شناخته شده برای امکان انتشار ثابت تهاجمی تر، جمع آوری کنید.
بهبود استفاده از حافظه تجزیه و تحلیل و زمان بندی استفاده از حافظه، حذف بسیاری از بافرهای ذخیره سازی متوسط.
اتکا به عملیات سفارشی را کاهش دهید. با بهبود عملکرد عملیات‌های سطح پایین با ذوب خودکار برای مطابقت با عملکردهای سفارشی که در ابتدا با دست ترکیب می‌شدند، نیاز به بسیاری از عملیات‌های سفارشی را از بین ببرید.
بهبود قابلیت حمل نوشتن یک backend جدید برای سخت افزار جدید را نسبتاً آسان کنید، به طوری که بخش بزرگی از مدل های ML می توانند بدون تغییر روی آن سخت افزار اجرا شوند. این در تضاد با رویکرد تخصصی کردن عملیات یکپارچه منفرد برای سخت‌افزار جدید است، که نیاز به بازنویسی مدل‌ها برای استفاده از آن عملیات دارد.
چگونه کار می کند
کامپایلر XLA نمودارهای مدل را از چارچوب های ML تعریف شده در StableHLO می گیرد و آنها را در دستورالعمل های ماشین برای معماری های مختلف کامپایل می کند. StableHLO یک مجموعه عملیات نسخه شده (HLO = عملیات سطح بالا) را تعریف می کند که یک لایه قابل حمل بین چارچوب های ML و کامپایلر فراهم می کند.
به طور کلی، فرآیند کامپایل که نمودار مدل را به یک فایل اجرایی بهینه شده برای هدف تبدیل می کند، شامل مراحل زیر است:
XLA چندین بهینه سازی داخلی و پاس های تحلیلی را روی نمودار StableHLO انجام می دهد که مستقل از هدف هستند، مانند CSE ، ترکیب عملیات مستقل از هدف، و تجزیه و تحلیل بافر برای تخصیص حافظه زمان اجرا برای محاسبات. در طول این مرحله بهینه سازی، XLA همچنین گویش StableHLO را به یک گویش HLO داخلی تبدیل می کند.
XLA محاسبات HLO را برای بهینه سازی های بیشتر در سطح HLO، این بار با در نظر گرفتن اطلاعات و نیازهای خاص هدف، به یک باطن ارسال می کند. به عنوان مثال، باطن GPU ممکن است ترکیبات عملیاتی را انجام دهد که به طور خاص برای مدل برنامه نویسی GPU مفید است و نحوه تقسیم کردن محاسبات را به جریان ها تعیین می کند. در این مرحله، پشتیبان‌ها همچنین ممکن است عملیات خاصی یا ترکیبی از آنها را با فراخوان‌های بهینه‌شده کتابخانه مطابقت دهند.
سپس backend تولید کد خاص هدف را انجام می دهد. پشتیبان‌های CPU و GPU همراه با XLA از LLVM برای IR سطح پایین، بهینه‌سازی و تولید کد استفاده می‌کنند. این backendها LLVM IR لازم برای نمایش محاسبات HLO را به شیوه ای کارآمد منتشر می کنند و سپس LLVM را فراخوانی می کنند تا کد بومی را از این LLVM IR منتشر کند.
در این فرآیند، کامپایلر XLA ماژولار است به این معنا که به راحتی می توان آن را در یک باطن جایگزین برای هدف قرار دادن معماری جدید HW قرار داد. پشتیبان GPU در حال حاضر از پردازنده های گرافیکی NVIDIA از طریق باطن LLVM NVPTX پشتیبانی می کند. باطن CPU از چندین ISA CPU پشتیبانی می کند.