روز جامعه ML 9 نوامبر است! برای به روز رسانی از TensorFlow، JAX به ما بپیوندید، و بیشتر بیشتر بدانید

معماری XLA

چرا ما XLA را ساختیم؟

ما چندین هدف برای XLA برای کار با TensorFlow داشتیم:

  • بهبود سرعت اجرا برای کاهش زمان اجرای Ops های کوتاه مدت ، زیرگراف ها را برای از بین بردن سربار از زمان اجرا TensorFlow ، عملیات فیوز خط لوله برای کاهش سربار حافظه ، جمع آوری کنید و به شکل های شناخته شده تنسور اختصاص دهید تا انتشار دائمی تهاجمی را فراهم کند.

  • استفاده از حافظه را بهبود ببخشید. تجزیه و تحلیل و زمانبندی استفاده از حافظه ، در اصل حذف بسیاری از بافرهای ذخیره سازی متوسط.

  • اعتماد به Ops سفارشی را کاهش دهید. با بهبود عملکرد Ops های سطح پایین که به صورت خودکار ذوب می شوند ، نیاز به بسیاری از Opsهای سفارشی را برطرف کنید تا با عملکردهای سفارشی که با دست ذوب شده اند مطابقت داشته باشد.

  • ردپای موبایل را کاهش دهید. با تدوین زودهنگام زیرشاخه و انتشار یک جفت فایل شی / هدر که می تواند مستقیماً به برنامه دیگری مرتبط شود ، زمان اجرا TensorFlow را حذف کنید. این نتایج می تواند رد پای استنباط موبایل را با چند مرتبه اندازه کاهش دهد.

  • قابلیت حمل را بهبود ببخشید. نوشتن یک باطن جدید برای سخت افزار جدید نسبتاً آسان است ، در آن زمان بخش بزرگی از برنامه های TensorFlow روی آن سخت افزار اصلاح نشده اجرا می شوند. این در تضاد با رویکرد تخصصی کردن Ops های یکپارچه منفرد برای سخت افزار جدید است که برای استفاده از آن Ops نیاز به بازنویسی برنامه های TensorFlow دارد.

XLA چگونه کار می کند؟

زبان ورودی به XLA "HLO IR" یا فقط HLO (عملیات سطح بالا) نامیده می شود. معناشناسی HLO در صفحه Operation Semantics شرح داده شده است. راحت تر است که HLO را به عنوان یک IR کامپایلر در نظر بگیریم.

XLA نمودارهای ("محاسبات") تعریف شده در HLO را می گیرد و آنها را در دستورالعمل های ماشین برای معماری های مختلف جمع می کند. XLA مدولار است به این معنا که اسلات در یک باطن جایگزین برای هدف قرار دادن برخی از معماری های جدید HW آسان است . باطن پردازنده برای x64 و ARM64 و همچنین باطن GPU NVIDIA در درخت منبع TensorFlow وجود دارد.

نمودار زیر روند تدوین در XLA را نشان می دهد:

XLA دارای چندین بهینه سازی و گذر از تجزیه و تحلیل است که مستقل از هدف هستند ، مانند CSE ، همجوشی عملیات مستقل از هدف و تحلیل بافر برای اختصاص حافظه زمان اجرا برای محاسبه.

پس از مرحله مستقل از هدف ، XLA محاسبه HLO را به یک backend ارسال می کند. باطن می تواند بهینه سازی های بیشتری در سطح HLO را انجام دهد ، این بار با توجه به اطلاعات خاص مورد نیاز و نیازها. به عنوان مثال ، باطن GPU XLA ممکن است همجوشی عملیاتی را که به طور خاص برای مدل برنامه نویسی GPU مفید است انجام دهد و نحوه تقسیم محاسبات به جریان ها را تعیین کند. در این مرحله ، باطن ها همچنین ممکن است برخی از عملکردها یا ترکیبات آنها را برای تماس کتابخانه ای بهینه شده مطابقت دهند.

مرحله بعدی تولید کد خاص هدف است. باطن پردازنده و پردازنده گرافیکی همراه با XLA از LLVM برای IR سطح پایین ، بهینه سازی و تولید کد استفاده می کنند. این باطن ها LLVM IR لازم برای نمایش محاسبات XLA HLO را به روشی کارآمد منتشر می کنند و سپس LLVM را فراخوانی می کنند تا کد بومی را از این IR LLVM منتشر کنند.

GPU backend از GPU های NVIDIA از طریق LLVM NVPTX پشتیبانی می کند. backend پردازنده از چندین ISA پردازنده پشتیبانی می کند.