این سند، مسیر یک ماژول بهینه‌سازی سطح بالا (HLO) XLA را از حالت اولیه تا یک فایل اجرایی نهایی تشریح می‌کند. گاهی اوقات ما «ماژول» را حذف می‌کنیم و فقط به آن «HLO» می‌گوییم.
HLO پیش از بهینه‌سازی
ما با ماژول HLO پیش از بهینه‌سازی شروع می‌کنیم. HLO پیش از بهینه‌سازی شامل عملیات ( ops ) که برای XLA داخلی در نظر گرفته می‌شوند، مانند fusion یا bitcast نمی‌شود. عملیات‌ها در این مرحله طرح‌بندی ندارند، یا اگر داشته باشند، نادیده گرفته می‌شوند. HLO پیش از بهینه‌سازی معمولاً توسط چارچوب‌های سطح بالاتر مانند TensorFlow و JAX تولید می‌شود. هنگام استفاده از پرچم XLA -xla_dump_to ، HLO پیش از بهینه‌سازی در فایلی با پسوند نام فایل "before_optimizations.txt" ذخیره می‌شود.
ماژول HLO را بهینه کنید
خط لوله XLA:GPU با اجرای دنباله‌ای از مراحل، HLO پیش‌بهینه‌سازی را به HLO بهینه تبدیل می‌کند. مراحل را می‌توان از نظر معنایی گروه‌بندی کرد و به ترتیب زیر اجرا کرد:
پاس‌های مرتبط با شاردینگ
این شامل گذرگاه‌هایی مانند Shardy Partitioner یا گذرگاه‌هایی برای شاردینگ SPMD می‌شود.
مراحل بهینه‌سازی
این می‌تواند شامل مجوزهای قانونی‌سازی و مجوزهای ساده‌سازی باشد.
مسیرهای بهینه‌سازی جمعی
مشابه مراحل بهینه‌سازی است، اما بر عملیات جمعی تمرکز دارد.
مراحل انتساب طرح‌بندی
به هر عملیات HLO یک طرح‌بندی اختصاص داده می‌شود که بخشی از شکل دستورالعمل است. این طرح‌بندی نحوه قرارگیری فیزیکی تانسور در حافظه را کنترل می‌کند.
مثالی از یک شکل با طرح‌بندی:
f32[10,20,30]{2,0,1}
بعد از نوع عنصر، ابعاد منطقی شکل قرار دارند و به دنبال آن جایگشت طرح‌بندی به ترتیب از کوچک به بزرگ آمده است. در این مثال، کوچکترین بُعد 30، دومین کوچکترین بُعد 10 و بزرگترین بُعد 20 است.
هدف از تخصیص طرح‌بندی، به حداقل رساندن تعداد جابجایی‌های فیزیکی مورد نیاز با استفاده از یک استراتژی حریصانه است. این کار با محدودیت‌های طرح‌بندی خاصی شروع می‌شود (مثلاً کتابخانه‌های cuDNN/cuBLAS انتظار ابعاد متوالی را دارند) و طرح‌بندی‌ها را "پایین" و سپس "بالا" در نمودار HLO منتشر می‌کند. در پایان انتشار طرح‌بندی، برخی از دستورالعمل‌ها ممکن است طرح‌بندی‌های متناقضی داشته باشند، یکی از یک عملوند منتشر می‌شود، دیگری از یک کاربر منتشر می‌شود. برای حل این تضاد، یک دستورالعمل copy HLO درج می‌شود که طرح‌بندی را از طرح‌بندی عملوند به طرح‌بندی دستورالعمل تغییر می‌دهد.
مراحل نرمال‌سازی طرح‌بندی
با توجه به اینکه تشخیص شکل فیزیکی تا حدودی دشوار است، نرمال‌سازی طرح‌بندی تلاش می‌کند تا شکل را طوری بازنویسی کند که از طرح‌بندی پیش‌فرض {rank-1, rank-2, …, 0} استفاده کند. در مثال بالا، شکل نرمال‌شده f32[20,10,30]{2,1,0} خواهد بود. عملیات کپی که طرح‌بندی‌ها را تغییر می‌دهند، به صورت ترکیبی از transpose و bitcast بازنویسی می‌شوند. با توجه به اینکه در حال حاضر نمی‌توانیم همه عملیات‌ها را نرمال‌سازی کنیم، هنوز برخی عملیات‌ها وجود دارند که ممکن است طرح‌بندی‌های غیر پیش‌فرض داشته باشند، به‌ویژه gather و dot . در مرزهای بین عملیات‌های نرمال‌شده و عملیات‌های غیر نرمال‌شده، عملیات‌های bitcast وجود خواهند داشت که نشان‌دهنده یک transpose هستند، یعنی یک transpose با طرح‌بندی اختصاص داده شده که آن را از نظر فیزیکی بدون عملیات می‌کند.
نرمال‌سازی طرح‌بندی همچنین برخی از انتقال‌های ضمنی را صریح می‌کند که مهم است زیرا codegen می‌تواند انتقال‌های صریح را با یک ساطع‌کننده اختصاصی مدیریت کند. به عنوان مثال، یک تغییر شکل از نظر فنی مجاز است که طرح فیزیکی متفاوتی بین عملوند و نتیجه داشته باشد (مثلاً به دلیل رتبه متفاوت). مسیر ReshapeDecomposer که به عنوان بخشی از مسیرهای نرمال‌سازی طرح‌بندی اجرا می‌شود، یک تغییر شکل را به دنباله‌ای از transpose ، reshape bitcast و transpose تبدیل می‌کند.
بهینه‌سازی تخصیص طرح‌بندی پست‌ها (post layout) با موفقیت انجام شد.
مهم‌ترین مراحل در اینجا، ادغام‌های تریتون (ترکیب‌های GEMM + ادغام‌های Softmax/Layernorm) یا بازنویسی‌ها برای فراخوانی‌های کتابخانه هستند. تنظیم خودکار نیز در این مرحله اجرا می‌شود، که در آن XLA بین ساطع‌کننده‌های مختلف یکی را انتخاب می‌کند، بهترین الگوریتم را برای پیچش‌ها یا نقاط انتخاب می‌کند، بهترین کاشی‌کاری را برای ادغام‌های مدیریت‌شده توسط ساطع‌کننده تریتون پیدا می‌کند و غیره.
عبورهای فیوژن
دو مرحله اصلی عبارتند از PriorityFusion و Multi-Output fusion.
در PriorityFusion ، ما ادغام‌ها را بر اساس مدل هزینه تشکیل می‌دهیم. هنگام ادغام، در صورتی که بتوان عملیات را در همه کاربران ادغام کرد، اجازه تکرار عملیات با چندین کاربر را می‌دهیم. همچنین در صورت امکان، اجازه گسترش ادغام‌های Triton Softmax موجود را می‌دهیم.
ادغام Multi-Output یک مسیر جداگانه است که امکان ادغام عملیات/ترکیب‌هایی را که یک عملوند مشترک دارند، فراهم می‌کند. همچنین می‌تواند با اضافه کردن خروجی(های) اضافی، عملیات/ترکیب‌های عملوند را بدون تکرار در کاربران ادغام کند، بنابراین سایر کاربران عملیاتی که قرار است ادغام شوند، می‌توانند به این خروجی‌ها هدایت شوند. این مسیر باید با دقت انجام شود تا چرخه‌هایی را در نمودار HLO وارد نکند.
پس از ادغام چند خروجی، حذف زیرعبارت‌های مشترک ( HloCSE pass) اجرا می‌شود و اگر عملیات‌های تکراری قبلی در یک ادغام قرار گرفته باشند، آنها را دوباره با هم ادغام می‌کند.
چندین پاس پس از همجوشی
چندین مرحله مربوط به کالکتیوها (مانند تبدیل آنها به ناهمگام، یا اعمال ترتیب نسبی خاصی از کالکتیوها).
در نهایت، CopyInsertion اجرا می‌کنیم که در آن کپی‌ها اضافه می‌شوند تا اطمینان حاصل شود که عملیات درجا، داده‌هایی را که هنوز در جای دیگر مورد نیاز هستند، بازنویسی نمی‌کنند.
در پایان بهینه‌سازی، HLO بهینه‌سازی‌شده در صورت استفاده از پرچم -xla_dump_to در فایلی که پسوند نام فایل آن "after_optimizations.txt" است، ذخیره می‌شود. اگر می‌خواهید HLO را پس از گذرهای میانی که در واقع HloModule را تغییر می‌دهند، ذخیره کنید، می‌توانید از پرچم -xla_dump_hlo_pass_re=.* (یا یک عبارت منظم خاص برای محدود کردن آن به گذرهای خاص) استفاده کنید.
زمان‌بندی
یک ماژول HLO بدون زمان‌بندی، هنوز هم در ترتیب پردازش عملیات‌ها، درجه‌ای از آزادی دارد. هر نوع توپولوژیکی که به روابط عملوند/نتیجه و وابستگی‌های کنترلی احترام بگذارد، معتبر است. زمان‌بندی تعیین می‌کند که از چه ترتیب خاصی استفاده شود. نگرانی اصلی در این مرحله، حداکثر مصرف حافظه است که به طول عمر تانسورها بستگی دارد. در مرحله اولیه، الگوریتم‌های زمان‌بندی مختلف را امتحان می‌کنیم و زمان‌بندی را انتخاب می‌کنیم که باید حداکثر مصرف حافظه را به حداقل برساند. توجه داشته باشید که در این مرحله هنوز با بافرهای فیزیکی کار نمی‌کنیم (این اتفاق در "تخصیص بافر" رخ می‌دهد) و استفاده از حافظه را شبیه‌سازی می‌کنیم.
سپس LatencyHidingScheduler اجرا می‌شود و سعی می‌کند همپوشانی محاسبات-ارتباطات را به حداکثر برساند. اما این ممکن است دوباره باعث افزایش استفاده از حافظه شود.
در نهایت، در صورتی که اوج مصرف حافظه بیشتر از میزان حافظه موجود باشد، HloRematerialization اجرا می‌کنیم. این مرحله تلاش می‌کند تا مصرف حافظه را به قیمت کاهش عملکرد کاهش دهد، زیرا به عنوان مثال ممکن است برخی از ادغام‌ها تقسیم شوند و برخی از عملیات‌ها برای داشتن طول عمر بافر کوتاه‌تر، کپی شوند. اگر rematerialization رخ دهد، بررسی راه‌هایی برای کاهش نیاز به حافظه در سمت مدل (مثلاً استفاده از اندازه‌های دسته‌ای کوچکتر) می‌تواند مفید باشد.
تخصیص بافر
درست قبل از اینکه به LLVM IR برسیم، مراحل تخصیص بافر را اجرا می‌کنیم که برش‌های بافر را به هر دستورالعمل در نمودار HLO اختصاص می‌دهد. تخصیص بافر در چند مرحله انجام می‌شود:
HloDataflowAnalysisHloValues(اساساً بافرهای منطقی) را به دستورالعمل‌ها اختصاص می‌دهد. برای عملیات درجا،HloValueیک عملوند می‌تواند دوباره استفاده شود. یک عملیات ممکن است بیش از یکHloValueتعریف کند (مثلاً با شکل نتیجه‌ی چندتایی).HloAliasAnalysisتلاش می‌کند تا بافرها را برای عملیات نامگذاری مستعار ترکیب کند و نگاشتی ازHloValueبهHloBufferمحاسبه کند.BufferAssignmentنگاشتHloBuffersرا به برش‌های بافر درون یک بافر بزرگ محاسبه می‌کند به گونه‌ای که از همان برش بافر برایHloBuffersمختلف با طول عمر همپوشانی استفاده نشود. برای عملیاتی که ممکن است نام مستعار داشته باشند، اشکالی ندارد که کمی همپوشانی وجود داشته باشد (زمان پایان یکHloBufferممکن است با زمان شروعHloBufferدیگر همزمان باشد). هنگام استفاده از پرچم-xla_dump_to، برخی از اطلاعات مربوط به انتساب بافر در فایلی با پسوند نام "after_optimizations-buffer-assignment.txt" ذخیره می‌شود.
تانکس
پس از اینکه یک گراف HLO بهینه و زمان‌بندی شد، به یک توالی خطی از thunkها برای یک backend خاص (CPU یا GPU) تبدیل می‌شود.
در XLA، یک Thunk انتزاعی از یک واحد کاری مستقل است که زمان اجرا آن را اجرا می‌کند. این می‌تواند یک راه‌اندازی هسته کامپایل شده، عملیات خاص، فراخوانی کتابخانه، ساختار جریان کنترل، ارتباط جمعی و غیره باشد. یک توالی Thunk نشان دهنده کل فایل اجرایی برای یک backend خاص است.
انتشار تان
فرآیند تبدیل یک محاسبه HLO زمان‌بندی‌شده به یک توالی thunk، "انتشار thunk" نامیده می‌شود. این کار توسط یک کلاس emitter اختصاصی در هر backend انجام می‌شود.
برای GPU Backend، این کار توسط IrEmitterUnnested انجام می‌شود. EmitHloComputation لیست زمان‌بندی‌شده‌ی دستورالعمل‌های HLO را در یک محاسبه تکرار می‌کند و به یک متد تخصصی Emit... (مثلاً EmitFusion ، EmitConvolutionThunk ، EmitWhile ) ارسال می‌کند. هر یک از این متدها شیء(های) Thunk مناسب را می‌سازد و آنها را به دنباله thunk اضافه می‌کند.
برای CPU Backend، ThunkEmitter این نقش را انجام می‌دهد و به شیوه‌ای مشابه سازماندهی شده است. Final ThunkSequence در CpuExecutable تعبیه شده است.
توجه داشته باشید که هر دستورالعمل در محاسبه ورودی یک ماژول HLO ممکن است با هیچ ( kTuple ، kConstant ، ..)، یک یا چندین (برای مثال دستورالعمل مرتب‌سازی) thunk در توالی thunk نهایی مطابقت داشته باشد.
بافرهای فرمان: بهینه‌سازی اجرا روی پردازنده گرافیکی
سخت‌افزار GPU مدرن امکان ضبط یک توالی از عملیات GPU (اجرای هسته، کپی حافظه و غیره) را یک بار و سپس پخش مجدد این توالی چندین بار با حداقل سربار CPU فراهم می‌کند. این یک بهینه‌سازی عملکرد حیاتی است، به خصوص برای بارهای کاری با هسته‌های کوچک و سریع. XLA از Command Buffer به عنوان انتزاعی از نمودارهای CUDA یا نمودارهای HIP استفاده می‌کند. رابط اصلی در GpuCommandBuffer تعریف شده است.
یک بافر فرمان توسط CommandBufferThunk در یک توالی thunk نمایش داده می‌شود.
امیتر این thunk را مستقیماً از دستورالعمل‌های HLO تولید نمی‌کند. در عوض، این کار توسط CommandBufferConversionPass که روی خود ThunkSequence اجرا می‌شود، انجام می‌شود.
این گذرگاه، زیردنباله‌های پیوسته از thunkهای سازگار (مثلاً مجموعه‌ای از KernelThunk و GemmThunk ها) را شناسایی می‌کند. سپس زیردنباله یافت‌شده را با یک CommandBufferThunk واحد جایگزین می‌کند. thunk جدید، منطق thunkهای اصلی را به عنوان فهرستی از اشیاء سبک CommandBufferCmd کپسوله‌سازی می‌کند. هنگامی که یک CommandBufferThunk برای اولین بار در یک جریان GPU مشخص اجرا می‌شود، توالی دستورات خود را در یک بافر دستور سخت‌افزاری "ضبط" می‌کند. در تمام اجراهای بعدی، به سادگی یک دستور واحد به GPU صادر می‌کند تا توالی ضبط‌شده را "بازپخش" کند. این کار از سربار CPU برای راه‌اندازی هر هسته جداگانه جلوگیری می‌کند.
قابل اجرا
محصول نهایی خط لوله کامپایل XLA یک فایل اجرایی مستقل و مختص پلتفرم است. این شیء تمام اطلاعات مورد نیاز برای اجرای برنامه کامپایل شده روی دستگاه هدف، مانند CPU یا GPU، را کپسوله می‌کند. این پل ارتباطی بین کامپایلر و محیط اجرا است. محیط‌های اجرای مدرن مانند PJRT از انتزاع‌های سطح بالاتر استفاده می‌کنند (به PjRtExecutable مراجعه کنید)، اما در نهایت یک فایل اجرایی مختص به backend را در بر می‌گیرند.
یک Executable شامل چندین بخش کلیدی از اطلاعات تولید شده در طول کامپایل است. در حالی که محتوای دقیق آن بسته به backend متفاوت است، اما عموماً شامل موارد زیر است:
کد کامپایل‌شده: این کد ماشین سطح پایین است که روی دستگاه اجرا می‌شود. برای CPUها، این معمولاً یک یا چند فایل شیء است. برای GPUها، این کد دستگاه کامپایل‌شده با فرمت PTX یا HSACO است که در زمان اجرا روی GPU بارگذاری می‌شود.
طرح اجرا (ThunkSequence): هسته منطق زمان اجرا. این یک توالی خطی از اشیاء Thunk است. هر thunk نشان دهنده یک واحد کاری واحد است، مانند راه اندازی یک هسته، فراخوانی یک تابع کتابخانه (مثلاً cuBLAS) یا مدیریت جریان کنترل. زمان اجرا با تکرار از طریق این توالی، برنامه را اجرا می‌کند.
طرح‌بندی حافظه (اختصاص بافر): این بخش حیاتی از فراداده، که توسط BufferAssigner تولید می‌شود، طرح‌بندی کامل حافظه برای محاسبه را توصیف می‌کند. این بخش، اندازه هر بافر و نحوه تخصیص و استفاده مجدد از حافظه برای پارامترها، خروجی‌ها و مقادیر موقت را مشخص می‌کند. زمان اجرا از این بخش برای تخصیص حافظه دستگاه و ارسال اشاره‌گرهای صحیح به هر thunk استفاده می‌کند.
(اختیاری) ماژول HLO: برای اشکال‌زدایی و پروفایل‌بندی، فایل اجرایی اغلب ارجاعی به HloModule نهایی و بهینه‌شده‌ای که از آن کامپایل شده است را نگه می‌دارد.
ایجاد فایل اجرایی نهایی توسط کامپایلر برای هر backend خاص هماهنگ می‌شود. متد RunBackend در پیاده‌سازی کامپایلر، آخرین مرحله در فرآیند کامپایل است که تمام مصنوعات کامپایل شده را در یک شیء اجرایی بسته‌بندی می‌کند. GpuCompiler و CpuCompiler به ترتیب GPU و CPU را هدف قرار می‌دهند.
وقتی کاربری دستور Execute... را روی یک فایل اجرایی فراخوانی می‌کند، محیط اجرا از BufferAssignment برای تخصیص حافظه استفاده می‌کند و سپس ThunkSequence برای اجرای عملیات روی دستگاه با استفاده از کد کامپایل شده فراخوانی می‌کند.