در حال توسعه یک باطن جدید برای XLA

این راهنما برای مهندسان سیستمی است که می‌خواهند XLA برنامه‌هایی را که سخت‌افزار آنها را به طور موثر هدف قرار می‌دهند، خروجی بگیرد. راهنما گام به گام نیست و دانش LLVM ، Bazel و XLA را فرض می کند.

XLA یک رابط انتزاعی ارائه می دهد که یک معماری یا شتاب دهنده جدید می تواند برای ایجاد یک backend برای اجرای برنامه های ML خروجی توسط XLA پیاده سازی کند. هدف‌گیری مجدد XLA باید به طور قابل‌توجهی ساده‌تر و مقیاس‌پذیرتر از پیاده‌سازی هر عملیات موجود از یک فریمورک فرانت‌اند مانند PyTorch یا TensorFlow برای سخت‌افزار جدید باشد.

اکثر پیاده سازی ها در یکی از سناریوهای زیر قرار می گیرند:

  1. معماری CPU موجود هنوز به طور رسمی توسط XLA پشتیبانی نمی شود، با یا بدون پشتیبان LLVM موجود.
  2. سخت افزار غیر شبیه CPU با پشتیبان LLVM موجود.
  3. سخت افزار غیر شبیه CPU بدون پشتوانه LLVM موجود.

سناریو 1: معماری CPU موجود هنوز به طور رسمی توسط XLA پشتیبانی نشده است

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

اگر فروشنده سخت‌افزار یک Backend LLVM برای سخت‌افزار خود دارد، ساده است که باطن را با LLVM ساخته شده با XLA پیوند دهید. در حالت JIT، پشتیبان XLA CPU کدی را برای CPU میزبان منتشر می کند. برای کامپایل زود هنگام، xla::AotCompilationOptions می تواند یک LLVM سه گانه برای پیکربندی معماری هدف ارائه دهد.

اگر هیچ باطن LLVM موجود وجود نداشته باشد، اما نوع دیگری از تولیدکننده کد وجود داشته باشد، باید بتوان از بیشتر پشتیبان CPU موجود استفاده مجدد کرد.

سناریو 2: سخت افزار غیر شبیه به CPU با پشتیبان LLVM موجود

می‌توان یک xla::Compiler جدید را روی کلاس‌های xla::CPUCompiler و xla::GPUCompiler مدل‌سازی کرد، زیرا از قبل LLVM IR منتشر می‌کنند. بسته به ماهیت سخت افزار، ممکن است بسیاری از جنبه های نسل LLVM IR باید تغییر کند، اما بسیاری از کدها را می توان با پشتیبان های موجود به اشتراک گذاشت.

یک مثال خوب برای دنبال کردن، باطن GPU XLA است. باطن GPU یک ISA غیر شبیه CPU را هدف قرار می دهد و بنابراین برخی از جنبه های تولید کد آن منحصر به دامنه GPU است. انواع دیگر سخت افزار، به عنوان مثال DSPهایی مانند Hexagon (که دارای یک backend بالادستی LLVM است)، می توانند از بخش هایی از منطق انتشار LLVM IR استفاده مجدد کنند، اما سایر قسمت ها منحصر به فرد خواهند بود.

سناریو 3: سخت افزار غیر شبیه CPU بدون پشتوانه LLVM موجود

اگر امکان استفاده از LLVM وجود ندارد، بهترین گزینه این است که یک Backend جدید برای XLA برای سخت افزار مورد نظر پیاده سازی کنید. این گزینه به بیشترین تلاش نیاز دارد. کلاس هایی که باید پیاده سازی شوند به شرح زیر است:

  • StreamExecutor : برای بسیاری از دستگاه ها به همه روش های StreamExecutor نیاز نیست. برای جزئیات بیشتر به اجرای StreamExecutor موجود مراجعه کنید.
  • xla::Compiler : این کلاس کامپایل یک محاسبات HLO را در یک xla::Executable کپسوله می کند.
  • xla::Executable : این کلاس برای راه اندازی یک محاسبات کامپایل شده بر روی پلتفرم استفاده می شود.
  • xla::TransferManager : این کلاس به backend ها امکان می دهد مکانیسم های مخصوص پلتفرم را برای ساخت داده های واقعی XLA از دسته های حافظه دستگاه ارائه کنند. به عبارت دیگر، به کپسوله کردن انتقال داده ها از میزبان به دستگاه و برگشت کمک می کند.