از HLO تا Thunks

HLO پیش بهینه سازی

ما با HLO پیش بهینه سازی شروع می کنیم. HLO پیش بهینه‌سازی شامل عملیات‌هایی نیست که برای XLA داخلی در نظر گرفته می‌شوند، به‌عنوان مثال fusion یا bitcast . Ops در این مرحله طرحی ندارد یا اگر داشته باشد نادیده گرفته می شود. HLO پیش بهینه‌سازی معمولاً توسط چارچوب‌های سطح بالاتری مانند Tensorflow و JAX تولید می‌شود. هنگام استفاده از پرچم XLA -xla_dump_to ، HLO پیش‌بهینه‌سازی به فایلی با پسوند نام فایل "before_optimizations.txt" ریخته می‌شود.

بهینه سازی ماژول HLO

خط لوله XLA:GPU، HLO پیش بهینه‌سازی شده را با اجرای دنباله‌ای از گذرها به HLO بهینه‌شده تبدیل می‌کند. پاس ها را می توان از نظر معنایی با هم گروه بندی کرد و به ترتیب زیر اجرا کرد:

شاردینگ پاس های مرتبط

Shardy Partitioner یا Shardy SPMD.

بهینه سازی می گذرد.

این می تواند هم مجوزهای قانونی و هم مجوزهای ساده سازی را شامل شود.

بهینه سازی جمعی می گذرد.

مشابه پاس‌های بهینه‌سازی ، اما بر عملیات‌های جمعی تمرکز دارد.

پاس های تکلیف چیدمان

به هر عملیات HLO یک طرح اختصاص داده می شود که بخشی از شکل دستورالعمل است. چیدمان نحوه قرارگیری تانسور به صورت فیزیکی در حافظه را کنترل می کند.

نمونه ای از شکل با طرح بندی:

f32[10,20,30]{2,0,1}

بعد از نوع عنصر، ابعاد منطقی شکل و به دنبال آن جایگشت چیدمان به ترتیب جزئی به بزرگ وجود دارد. در این مثال، جزئی ترین بعد 30، دومین بعد جزئی 10 و بعد عمده 20 است.

هدف از تخصیص چیدمان، به حداقل رساندن تعداد جابجایی های فیزیکی است که با استفاده از یک استراتژی حریصانه مورد نیاز است. با محدودیت‌های طرح‌بندی خاصی شروع می‌شود (مثلاً کتابخانه‌های CuDNN/cuBLAS ابعاد متوالی را انتظار دارند) و طرح‌بندی «پایین» و سپس «بالا» نمودار HLO را منتشر می‌کند. در پایان انتشار طرح، برخی از دستورالعمل ها ممکن است دارای طرح بندی متناقضی باشند، یکی از یک عملوند منتشر شده است، یکی از یک کاربر منتشر شده است. برای حل این تضاد، یک دستورالعمل HLO copy درج می شود که طرح را از طرح عملوند به طرح بندی دستورالعمل تغییر می دهد.

نرمال سازی چیدمان می گذرد

با توجه به اینکه تشخیص شکل فیزیکی تا حدودی دشوار است، نرمال‌سازی طرح سعی می‌کند شکل را بازنویسی کند به طوری که از طرح‌بندی پیش‌فرض {rank-1, rank-2, …, 0} استفاده می‌کند. در مثال بالا، شکل نرمال شده f32[20,10,30]{2,1,0} خواهد بود. عملیات کپی که طرح‌بندی را تغییر می‌دهند به ترکیبی از transpose + bitcast بازنویسی می‌شوند. با توجه به اینکه در حال حاضر نمی‌توانیم همه عملیات‌ها را عادی کنیم، هنوز برخی از عملیات‌ها وجود دارند که ممکن است طرح‌بندی‌های غیر پیش‌فرض داشته باشند، به ویژه gather و dot . در مرزهای بین عملیات نرمال شده و عملیات غیر عادی، عملیات bitcast وجود دارد که نشان‌دهنده یک transpose است، یعنی یک جابجایی با یک چیدمان اختصاص داده شده که آن را از نظر فیزیکی بدون عملیات تبدیل می‌کند.

نرمال‌سازی چیدمان همچنین باعث می‌شود که برخی از جابجایی‌های ضمنی صریح باشند که مهم است زیرا کدژن می‌تواند جابجایی‌های صریح را با یک امیتر اختصاصی انجام دهد. به عنوان مثال، یک تغییر شکل از نظر فنی مجاز است که طرح فیزیکی متفاوتی بین عملوند و نتیجه داشته باشد (مثلاً به دلیل رتبه متفاوت). پاس ReshapeDecomposer که به عنوان بخشی از پاس‌های عادی‌سازی طرح اجرا می‌شود، یک تغییر شکل را به دنباله‌ای از transpose ، reshape bitcast و transpose تبدیل می‌کند.

پاس های بهینه سازی تخصیص طرح بندی پست

مهم‌ترین پاس‌ها در اینجا ترکیب‌های Triton (تلفیقی GEMM + Softmax/Layernorm fusions) یا بازنویسی در فراخوان‌های کتابخانه هستند. اما همچنین Autotuning در این مرحله اجرا می‌شود، جایی که ما بهترین الگوریتم را برای کانولوشن‌ها یا نقطه‌ها، یا بهترین کاشی کاری برای نقاطی که توسط فرستنده قدیمی Triton GEMM مدیریت می‌شود، انتخاب می‌کنیم، یا اینکه آیا باید از Triton یا Cublas برای یک همجوشی نقطه‌ای خاص استفاده کنیم.

فیوژن می گذرد

دو پاس اصلی 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=.* (یا یک عبارت منظم خاص برای محدود کردن آن به پاس‌های خاص) استفاده کنید.

برنامه ریزی

یک HloModule بدون زمان‌بندی هنوز درجاتی از آزادی در جهت پردازش عملیات‌ها دارد. اساساً هر مرتب‌سازی توپولوژیکی با توجه به رابطه عملوند/نتیجه و وابستگی‌های کنترل خوب است. زمان بندی نظم خاصی را اجرا می کند. این بر میزان حافظه مورد نیاز تأثیر می‌گذارد، زیرا تا زمانی که همه خوانندگان آن بافر پردازش نشده باشند، نمی‌توانیم از یک بافر دوباره استفاده کنیم. در مرحله اولیه، الگوریتم‌های زمان‌بندی مختلف را امتحان می‌کنیم و برنامه‌ای را انتخاب می‌کنیم که حداکثر مصرف حافظه را به حداقل برساند.

در ادامه، ما مجوز LatencyHidingScheduler را اجرا می کنیم که سعی می کند همپوشانی محاسباتی-ارتباطات را به حداکثر برساند، اما ممکن است دوباره استفاده از حافظه را افزایش دهد.

پس از زمان‌بندی، HloRematerialization اجرا می‌کنیم که سعی می‌کند در صورتی که پیک مصرف حافظه بیشتر از مقدار حافظه موجود باشد، استفاده از حافظه را کاهش دهد. این به قیمت کارایی است، به عنوان مثال ممکن است برخی از فیوژن ها تقسیم شوند و برخی عملیات ممکن است تکرار شوند تا طول عمر بافر کمتری داشته باشند. اگر مادی‌سازی مجدد اتفاق می‌افتد، به طور بالقوه منطقی به نظر می‌رسد که آیا راه‌هایی در سمت مدل برای کاهش مقدار حافظه مورد نیاز وجود دارد (مثلاً اندازه‌های دسته‌ای کوچکتر).

Thunks و CommandBuffers

TBD

BufferAssignment

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

1. HloDataflowAnalysis HloValues ​​(بافرهای اساسا منطقی) را به دستورالعمل ها اختصاص می دهد. برای عملیات در محل، HloValue یک عملوند را می توان دوباره استفاده کرد. یک op ممکن است بیش از یک HloValue تعریف کند (مثلاً با یک شکل نتیجه تاپل).

2. HloAliasAnalysis تلاش می کند تا بافرها را برای عملیات aliasing ترکیب کند و یک نقشه برداری از HloValue به HloBuffer محاسبه می کند.

3. BufferAssignment نگاشت HloBuffers را برای برش های بافر در داخل یک بافر بزرگ محاسبه می کند، به گونه ای که همان برش بافر برای HloBuffers مختلف با طول عمر همپوشانی استفاده نمی شود. برای عملیات هایی که ممکن است نام مستعار داشته باشند، مشکلی ندارد که یک همپوشانی جزئی وجود داشته باشد (زمان پایان یک HloBuffer ممکن است با زمان شروع HloBuffer دیگر مطابقت داشته باشد). هنگام استفاده از flag -xla_dump_to ، برخی از اطلاعات مربوط به تخصیص بافر به فایلی با پسوند نام "after_optimizations-buffer-assignment.txt" ریخته می شود.