من HLO إلى Thunks

يوضّح هذا المستند رحلة وحدة High Level Optimizer (HLO) في XLA من حالتها الأولية إلى ملف تنفيذي نهائي. في بعض الأحيان، سنحذف كلمة "وحدة" ونشير إلى هذا المقياس باسم "مقياس HLO" فقط.

مخطط HLO إلى thunks

Pre-optimization 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 عمليات تمثّل عملية تبديل، أي عملية تبديل مع تحديد تخطيط يجعلها عملية غير مجدية فعليًا.

يؤدي توحيد تنسيق التخطيط أيضًا إلى توضيح بعض عمليات التبديل الضمني، وهو أمر مهم لأنّ إنشاء الرموز البرمجية يمكنه التعامل مع عمليات التبديل الواضحة باستخدام أداة إصدار مخصّصة. على سبيل المثال، يُسمح من الناحية الفنية بأن يكون لعملية إعادة التشكيل تخطيط مادي مختلف بين المعامِل والنتيجة (على سبيل المثال، بسبب اختلاف الترتيب). تعمل عملية ReshapeDecomposer التي يتم تنفيذها كجزء من عمليات تسوية التنسيق على تحويل عملية إعادة التشكيل إلى تسلسل من transpose وإعادة التشكيل bitcast وtranspose.

عمليات تحسين تحديد تخطيط المشاركات

أهم عمليات الدمج هنا هي عمليات دمج Triton (عمليات دمج GEMM + عمليات دمج Softmax/Layernorm) أو عمليات إعادة الكتابة إلى استدعاءات المكتبة. يتم أيضًا تنفيذ عملية الضبط التلقائي في هذه الخطوة، حيث يختار XLA بين مصادر إخراج مختلفة، ويحدّد أفضل خوارزمية لعمليات الالتفاف أو النقاط، ويجد أفضل تقسيم إلى مربعات لعمليات الدمج التي يتعامل معها مصدر إخراج Triton وما إلى ذلك.

بطاقات الدمج

عمليتان الدمج الرئيسيتان هما PriorityFusion وMulti-Output.

في PriorityFusion، ننشئ عمليات دمج موجّهة بنموذج التكلفة. عند الدمج، نسمح بتكرار العمليات مع عدة مستخدمين إذا كان يمكن دمج العملية مع جميع المستخدمين. وسنسمح أيضًا بتوسيع عمليات دمج Triton Softmax الحالية إذا كان ذلك ممكنًا.

Multi-Output fusion هي عملية دمج منفصلة تتيح دمج العمليات/عمليات الدمج التي تتشارك في معامل. يمكنه أيضًا دمج المعامِلات/عمليات دمج المعامِلات في المستخدمين بدون تكرار من خلال إضافة مخرجات إضافية، وبالتالي يمكن إعادة توجيه المستخدمين الآخرين للعملية المراد دمجها إلى هذه المخرجات. يجب أن يكون هذا المرور حريصًا على عدم إدخال دورات في الرسم البياني HLO.

بعد دمج عمليات الإخراج المتعددة، يتم تنفيذ عملية إزالة التعبيرات الفرعية المشتركة (تمرير HloCSE)، ما قد يؤدي إلى دمج العمليات المكررة سابقًا معًا إذا انتهى بها الأمر في عملية الدمج نفسها.

عدة عمليات دمج بعد الدمج

العديد من عمليات التمرير ذات الصلة بالمجموعات (مثل تحويلها إلى غير متزامن أو فرض ترتيب نسبي معيّن للمجموعات)

أخيرًا، ننفّذ CopyInsertion حيث تتم إضافة النُسخ لضمان عدم استبدال العمليات الموضعية للبيانات التي لا تزال مطلوبة في مكان آخر.

في نهاية عملية التحسين، يتم تفريغ HLO المحسَّن إذا تم استخدام العلامة -xla_dump_to في ملف يتضمّن لاحقة اسم الملف "after_optimizations.txt". إذا أردت تفريغ HLO بعد عمليات المرور الوسيطة التي تغيّر HloModule فعليًا، يمكنك استخدام العلامة -xla_dump_hlo_pass_re=.* (أو تعبير عادي محدّد لحصرها في عمليات مرور معيّنة).

الجدولة

يظلّ بإمكان وحدة HLO بدون جدول زمني التمتّع ببعض الحرية في ترتيب معالجة العمليات. ويكون أي ترتيب طوبولوجي يحترم علاقات المعامل/النتيجة والتبعيات الخاصة بالتحكّم صالحًا. تحدّد الجدولة الترتيب المحدّد الذي سيتم استخدامه. المشكلة الرئيسية في هذه المرحلة هي الحد الأقصى لاستهلاك الذاكرة الذي يعتمد على مدة بقاء الموترات. في خطوة أولية، نجرب خوارزميات مختلفة لتحديد الجدول الزمني ونختار الجدول الذي من المفترض أن يقلّل من الحد الأقصى لاستهلاك الذاكرة. يُرجى العِلم أنّنا لا نتعامل في هذه المرحلة مع مخازن مؤقتة فعلية (سيحدث ذلك في "تخصيص المخزن المؤقت") ونحاكي استخدام الذاكرة.

بعد ذلك، يتم تنفيذ LatencyHidingScheduler عدة مرات لمحاولة زيادة التداخل بين العمليات الحسابية وعمليات الاتصال إلى أقصى حد. ولكن قد يؤدي ذلك إلى زيادة استخدام الذاكرة مرة أخرى.

أخيرًا، في حال كان الحد الأقصى لاستهلاك الذاكرة أعلى من مقدار الذاكرة المتوفّرة لدينا، ننفّذ HloRematerialization. تحاول هذه الخطوة تقليل استخدام الذاكرة على حساب الأداء، مثلاً، قد يتم تقسيم بعض عمليات الدمج وقد يتم تكرار بعض العمليات للحصول على مدة أقصر لتخزين البيانات مؤقتًا. في حال حدوث إعادة إنشاء، قد يكون من المفيد البحث عن طرق لتقليل متطلبات الذاكرة من جهة النموذج (مثل استخدام أحجام دفعات أصغر).

تحديد المخزن المؤقت

قبل الانتقال مباشرةً إلى LLVM IR، ننفّذ عمليات تمرير تخصيص المخزن المؤقت التي ستخصّص شرائح المخزن المؤقت لكل تعليمات في الرسم البياني HLO. يتم تنفيذ عملية تخصيص المخزن المؤقت على عدة خطوات:

  1. HloDataflowAnalysis تعيّن HloValues (وهي في الأساس مخازن مؤقتة منطقية) للتعليمات. بالنسبة إلى العمليات الموضعية، يمكن إعادة استخدام HloValue لأحد المعامِلات. يمكن أن يحدّد عامل التشغيل أكثر من HloValue واحد (على سبيل المثال، مع شكل نتيجة صفوف).

  2. تحاول HloAliasAnalysis دمج المخازن المؤقتة لعمليات التسمية المستعارة، وتحسب عملية ربط من HloValue إلى HloBuffer.

  3. تحسب BufferAssignment عملية ربط بين HloBuffers وشرائح المخزن المؤقت داخل مخزن مؤقت كبير بطريقة لا يتم فيها استخدام شريحة المخزن المؤقت نفسها مع HloBuffers مختلفة ذات فترات صلاحية متداخلة. بالنسبة إلى العمليات التي قد تستخدم أسماء مستعارة، لا بأس بحدوث تداخل طفيف (قد يتزامن وقت انتهاء إحدى العمليات HloBuffer مع وقت بدء العملية الأخرى HloBuffer). عند استخدام العلامة -xla_dump_to، يتم تفريغ بعض المعلومات حول عملية تخصيص المخزن المؤقت في ملف يحمل لاحقة الاسم "after_optimizations-buffer-assignment.txt".

Thunks

بعد تحسين رسم بياني HLO وجدولته، يتم تحويله إلى تسلسل خطي من أجزاء صغيرة من الرمز البرمجي لتنفيذه على نظام خلفي معيّن (وحدة المعالجة المركزية أو وحدة معالجة الرسومات).

في XLA، Thunk هو تجريد لوحدة عمل مستقلة ينفّذها وقت التشغيل. قد يكون ذلك إطلاقًا لنواة مجمّعة أو عملية معيّنة أو استدعاء مكتبة أو بنية للتحكّم في التدفق أو اتصال جماعي أو غير ذلك. يمثّل تسلسل Thunk الملف التنفيذي الكامل لخادم خلفي معيّن.

Thunk Emission

تُعرف عملية تحويل عملية حسابية مجدوَلة في HLO إلى تسلسل ثنك باسم "إصدار ثنك". ويتم التعامل مع ذلك من خلال فئة مصدر مخصّصة في كل نظام خلفي.

بالنسبة إلى الخلفية لوحدة معالجة الرسومات، يتم التعامل مع ذلك من خلال IrEmitterUnnested. تكرِّر EmitHloComputation قائمة تعليمات HLO المُجدوَلة في عملية حسابية، ثم ترسلها إلى طريقة Emit... متخصّصة (مثل EmitFusion وEmitConvolutionThunk وEmitWhile). ينشئ كل من هذه الطرق كائنات Thunk المناسبة ويضيفها إلى تسلسل Thunk.

بالنسبة إلى وحدة المعالجة المركزية (CPU) الخلفية، تؤدي ThunkEmitter هذا الدور ويتم تنظيمها بطريقة مشابهة. يتم تضمين ThunkSequence النهائي في CpuExecutable.

يُرجى العِلم أنّ كل تعليمات في عملية احتساب إدخال وحدة HLO قد تتوافق مع صفر (kTuple، kConstant، ..) أو ثنك واحد أو عدة ثنكات (مثل تعليمات الترتيب) في تسلسل الثنكات النهائي.

مخازن الأوامر المؤقتة: تحسين التنفيذ على وحدة معالجة الرسومات

تتيح أجهزة وحدات معالجة الرسومات الحديثة تسجيل سلسلة من عمليات وحدة معالجة الرسومات (عمليات تشغيل النواة، ونسخ الذاكرة، وما إلى ذلك) مرة واحدة ثم إعادة تشغيل السلسلة عدة مرات بأقل قدر من الحمل الزائد لوحدة المعالجة المركزية. هذا التحسين مهم جدًا للأداء، لا سيما لأحمال العمل التي تتضمّن العديد من النواة الصغيرة التي يتم تشغيلها بسرعة. تستخدم XLA مخزن الأوامر المؤقت كطبقة تجريدية لـ CUDA Graphs أو HIP Graphs. يتم تحديد الواجهة الأساسية في GpuCommandBuffer.

يتم تمثيل مخزن مؤقت للأوامر في تسلسل ثنك من خلال CommandBufferThunk.

لا ينتج المُصدر هذا الثنك مباشرةً من تعليمات HLO. بدلاً من ذلك، يتم ذلك من خلال CommandBufferConversionPass الذي يتم تشغيله على ThunkSequence نفسه.

تحدّد عملية التحويل هذه التسلسلات الفرعية المتجاورة من أجزاء التعليمات البرمجية المتوافقة (مثل سلسلة من KernelThunk وGemmThunk). ثم يستبدل التسلسل الفرعي الذي تم العثور عليه بعلامة CommandBufferThunk واحدة. يحتوي Thunk الجديد على منطق thunk الأصلي في شكل قائمة من عناصر CommandBufferCmd خفيفة الوزن. عندما يتم تنفيذ CommandBufferThunk لأول مرة على بث معيّن لوحدة معالجة الرسومات، يتم "تسجيل" تسلسل الأوامر في مخزن مؤقت لأوامر الأجهزة. وفي جميع عمليات التنفيذ اللاحقة، يتم ببساطة إصدار أمر واحد إلى وحدة معالجة الرسومات (GPU) من أجل "إعادة تشغيل" التسلسل المسجّل. ويؤدي ذلك إلى تجنُّب الحمل الزائد لوحدة المعالجة المركزية الناتج عن تشغيل كل نواة على حدة.

قابل للتنفيذ

المنتج النهائي لعملية تجميع XLA هو ملف تنفيذي مستقل خاص بالمنصة. يحتوي هذا العنصر على جميع المعلومات اللازمة لتشغيل البرنامج المجمّع على جهاز مستهدف، مثل وحدة المعالجة المركزية أو وحدة معالجة الرسومات. وهي تمثّل حلقة الوصل بين المترجم والوقت الفعلي لتنفيذ البرنامج. تستخدم بيئات التشغيل الحديثة، مثل PJRT، تجريدات ذات مستوى أعلى قليلاً (راجِع PjRtExecutable)، ولكنها في النهاية تتضمّن ملفًا تنفيذيًا خاصًا بالخادم الخلفي.

يحتوي Executable على عدة معلومات أساسية يتم إنشاؤها أثناء عملية التجميع. على الرغم من أنّ المحتوى الدقيق يختلف حسب الخلفية، إلا أنّه يشمل بشكل عام ما يلي:

  • الرمز البرمجي المجمَّع: هو رمز الجهاز المنخفض المستوى الذي سيتم تشغيله على الجهاز. بالنسبة إلى وحدات المعالجة المركزية، يكون ذلك عادةً ملفًا واحدًا أو أكثر من ملفات الكائنات. بالنسبة إلى وحدات معالجة الرسومات، يكون هذا هو رمز الجهاز المجمَّع بتنسيق PTX أو HSACO، والذي يتم تحميله على وحدة معالجة الرسومات في وقت التشغيل.

  • خطة التنفيذ (ThunkSequence): هي جوهر منطق وقت التشغيل. هذه سلسلة خطية من عناصر Thunk. يمثّل كل ثنك وحدة عمل واحدة، مثل تشغيل نواة أو استدعاء دالة مكتبة (مثل cuBLAS) أو معالجة تدفق التحكّم. ينفِّذ وقت التشغيل البرنامج من خلال تكرار هذا التسلسل.

  • تنسيق الذاكرة (BufferAssignment): هذه البيانات الوصفية مهمة جدًا، وهي من إنتاج BufferAssigner، وتصف تنسيق الذاكرة الكامل للحساب. تحدّد هذه السمة حجم كل مخزن مؤقت وكيفية تخصيص الذاكرة وإعادة استخدامها للمَعلمات والنتائج والقيم المؤقتة. يستخدم وقت التشغيل ذلك لتخصيص ذاكرة الجهاز وتمرير المؤشرات الصحيحة إلى كل ثنك.

  • (اختياري) وحدة HLO: لأغراض تصحيح الأخطاء وإنشاء الملفات الشخصية، يحتفظ الملف التنفيذي غالبًا بمرجع إلى HloModule النهائي والمحسّن الذي تم تجميعه منه.

يتم تنسيق عملية إنشاء الملف التنفيذي النهائي بواسطة المترجم لكل نظام خلفي محدد. الطريقة RunBackend في تنفيذ برنامج التجميع هي الخطوة الأخيرة في عملية التجميع، وهي تحزّم جميع العناصر المجمّعة في عنصر قابل للتنفيذ. يستهدف GpuCompiler وCpuCompiler وحدة معالجة الرسومات ووحدة المعالجة المركزية على التوالي.

عندما يستدعي المستخدم Execute... على ملف تنفيذي، يستخدم وقت التشغيل BufferAssignment لتخصيص الذاكرة، ثم يستدعي ThunkSequence لتشغيل العمليات على الجهاز باستخدام الرمز البرمجي المترجَم.