ترجمت واجهة Cloud Translation API‏ هذه الصفحة.
Switch to English

MLIR CodeGen لـ XLA

يعمل XLA على HloInstruction ويقوم بإجراء العديد من التحسينات على هذا التمثيل ، ومشاركة الكثير من هذه بين الأجهزة المستهدفة. عند نقطة ما يتم حساب جدول خطي ويتم تعيين ذاكرة التخزين المؤقت لكل قيمة بشكل ثابت. يعمل برنامج التشفير المحدد للجهاز من خلال اجتياز هذا التسلسل واستدعاء "بواعث" لإنشاء تمثيل مناسب للجهاز (على سبيل المثال ، وظيفة LLVM واحدة لكل حساب XLA على وحدة المعالجة المركزية ، أو سلسلة من "thunks" مغلفة لعمليات GPU وربما توليد PTX عندما تستهدف GPU).

كخطوة مرحلية ، نحن حاليًا بصدد اعتراض العملية مباشرة بعد اكتمال XLA لمرحلة تعيين المخزن المؤقت وإصدار وحدة MLIR في اللهجة lhlo بدلاً من ذلك. من هناك نقوم بتنفيذ برنامج التشفير باستخدام مكونات MLIR (Linalg و affine و GPU لهجة بشكل أساسي) اعتمادًا على الجهاز.

فيما يلي خطة التسجيل لترحيل XLA / GPU بشكل lhlo باستخدام lhlo كمدخل تشفير.

مهام

مضيف جهاز
نمط الإدخال HloInstruction * (المهمة 1) HloInstruction * (المهمة 1)
تنسيق الإخراج xla :: Thunk (المهمة 2) LLVM IR (المهمة 3)
  • تقوم المهمة 1 بتغيير كل من تنسيق إدخال المضيف والجهاز من HloInstruction * إلى LHLO.
  • تغير المهمة 2 تنسيق الإخراج للمضيف من thunks إلى "بعض منصات الهبوط للمضيف" (انظر أدناه).
  • تقوم المهمة 3 بترحيل إخراج الجهاز من LLVM IR إلى شكل من أشكال MLIR. إنه اختياري لهذا المشروع ، وانظر قسم "Migrating Device LLVM IR" للحصول على التفاصيل.

يعطي هذا المشروع الأولوية لتمكين نماذج قابلة للتشغيل من البداية إلى النهاية مع بواعث LHLO قدر الإمكان. وهذا يعني أن قائمة الترتيب التالية للأهداف حسب الأولوية:

  • اجعل XLA / GPU قابلاً للتشغيل مع بواعث LHLO ، مع وجود Thunks وبواعث الحالية غير المعدلة.
  • تخلص من الإشارات إلى HloInstruction * في LHLO ، كل حالة على حدة:
    • قم بتبديل باعث قديم إلى باعث قائم على MLIR (مثل Linalg) ، أو
    • ترجمة الباعث الحالي ميكانيكيًا لأخذ تمثيل MLIR (الترحيل إلى Standard with GPU Dialect).

ترحيل Thunks (المهمة 2)

xla :: gpu :: Thunk هي بنية بيانات:

  • يمكن استدعاؤه من المضيف (xla :: gpu :: Thunk :: ExecuteOnStream ()).
  • يحمل بيانات مختلفة في فئاته الفرعية.
  • يتفاعل مع BufferAllocation :: Slice و StreamExecutor.
  • تطلق النواة
  • مكالمات في جميع مكتبات وقت التشغيل.

تشمل تكلفة ذلك:

  • تمثيل بيانات التكوين الخاصة بالعمليات (مثل تكوينات الالتفاف).
  • ترحيل شكل المرجع وأشكال المعامل.
  • تمثل شجرة الأشواك (بينما ، الشرط ، إلخ).

عمل الترحيل مستقل عن ترحيل LHLO / باعث. في ظل الموارد المحدودة ، تم تحديد أولوياتها خلف ترحيل LHLO / باعث.

لدينا العديد من الخيارات حول كيفية خفض جزء جانب المضيف من LHLO:

  • TFRT
    • (Pro) أغلفة CUDA و HIP رائعة للاستخدام.
    • (Pro) من السهل تنفيذ مكالمات المكتبة (cuDNN ، cuBLAS ، cuFFT ، إلخ) ، حيث يتم تفسير عمليات TFRT بواسطة كود C ++.
    • (Con) الجانب المضيف قيد التطوير ولم يتم اختباره.
  • كود وحدة المعالجة المركزية Jitted
    • (برو) قدرة منخفضة كبيرة. قم بإنشاء عدد قليل من الحلقات والظروف وانتهى الأمر.
    • (Con) GPUDialect لم يقم بعد بتخصيص سلاسل / تيارات / عدم التزامن / تخصيص الجهاز.
    • (Con) دعم وقت تشغيل CUDA / HIP هو الحد الأدنى (مسار مجموعة الأدوات ، الإصدار ، التحميل الديناميكي ، إلخ).
  • وقت تشغيل XLA الحالي (التفسير)

القرار: اعتماد TFRT ، ولكن أيضًا دعم كود وحدة المعالجة المركزية النفاثة في TFRT.

ترحيل الجهاز LLVM IR (المهمة 3)

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

ElementalIrEmitter فريد من نوعه في ذلك:

  • يتم مشاركة جزء كبير من الرمز بين XLA / GPU ووحدة المعالجة المركزية.
  • إنه يمثل جزءًا كبيرًا من العمليات التي تتم مشاهدتها في النماذج ، بما في ذلك جميع العمليات الحكيمة للعناصر.
  • تعتمد معظم عمليات الاندماج على عنصر ElemalIrEmitter فقط.
  • إنه بسيط من الناحية الهيكلية ، حيث يصف DAG تبعية البيانات بين عناصر المرجع وعناصر المعامل.
  • غالبًا ما تكون محمولة وعالية المستوى (على سبيل المثال على عكس GPU kReduce و GPU kCopy).
  • دعم الشكل الديناميكي سهل على الأقل لعناصر العنصر.

الآن ، لجميع العمليات ، المنبعثة أو غير المنبعثة ، هناك العديد من النكهات للحالة النهائية لكل عملية XLA:

  1. يبقى رمز الجهاز كـ LLVM IR.
  2. إعادة تشكيل الباعث القديم ليكون مثل LHLO -> MLIR LLVM لهجة:
    • (التكلفة) سيتم التخلص منها إذا أردنا في النهاية الانتقال إلى الإصدار القياسي.
    • (فائدة) إنه سهل وميكانيكي. يمكن القيام به في فترة قصيرة.
    • (فائدة) لا تفيد أكثر مقارنة بـ (1).
  3. الباعثات القديمة لإعادة البناء لتكون مثل LHLO -> MLIR GPU + Standard + Loops:
    • (التكلفة) رفع أجهزة البث الموجودة إلى مستوى قياسي يقدم بعض التحديات. يجب تحويل المؤشرات و GEP إلى MemRefs و SubViews. ضمان اكتمال amdgpu أمر آخر.
    • (التكلفة) يعتمد XLA / GPU بشكل كبير على بيانات تعريف LLVM:
      • range مؤشرات الكتلة / الخيط.
      • align ، dereferenceable ، invariant.load ، alias.scope ، noalias لتحميل / مخازن.
      • llvm.loop.unroll.disable ، llvm.loop.unroll.full ، llvm.loop.vectorize.enable للحلقات المتسلسلة.
    • (فائدة) يمكن أن تكون طويلة الأمد. أكثر سهولة.
  4. إعادة باعث باعثات قديمة ليكون LHLO -> Linalg ، وكتابة بواعث Linalg جديدة
    • (التكلفة) هذه حالة بحالة. مقارنة بالخيارات السابقة ، يجب أن يمر تنفيذ جديد يتطابق مع أداء XLA من خلال العلامة المرجعية <-> لتحسين سير العمل ، والتي يمكن أن تكون تكلفة كبيرة لبعض العمليات.
    • (فائدة) مكدس موحد ؛ دعم المجتمع ؛ قابلية التنقل؛ المزيد من إمكانات التحسين.

الاستنتاجات:

  • لا تذهب إلى (2). (1) أو (3) أفضل من (2) فقط. (2) تكلف أكثر من (1) ، لأنها تتطلب الكثير من إعادة الهيكلة الميكانيكية. مع (1) لا يزال بإمكاننا تحقيق هدف تمكين XLA من التقاط بواعث MLIR. هذا عن طريق عمل LHLO -> LLVM IR -> تشغيل بواعث الجهاز القديم.
  • تعمل عناصر ElementalIrEmitter على (4) ، ولكن ليس بشكل متزايد. لا توجد طريقة للقيام بذلك عن طريق المرجع ، لأن جميع العمليات المنبعثة من العناصر مرتبطة بنفس الرسم البياني. يمكن أن يكون هذا العمل أيضًا بمثابة نقطة توحيد للعديد من القوى المستمرة (xla / service / mlir_gpu ، مولد النواة ، Linalg).
  • جميع العمليات الأخرى تذهب لـ (1). كهدف تمدد ، قد يتم ترحيلهم إلى (3) أو (4).

تحديد الأولويات

في حين أن جميع المهام الثلاث المذكورة أعلاه متوازية ، إلا أنه في ظل الموارد المحدودة يجب إجراء تسلسل لها. يركز تحديد الأولويات على النتائج المرئية لإكمال كل مهمة.

الأولوية هي: Task1 (LHLO for emitation emersers)> Task 2 (Thunks)> Task 3 (MLIR emitters).

بنهاية المهمة 1 ، يمكن لمستخدمي XLA توليد LHLO (مثل مولد النواة) وتنفيذها. لن يكون تنسيق الترجمة MLIR قابل للتسلسل.

بنهاية المهمة 2 ، ينخفض ​​LHLO إلى MLIR المناسب القابل للتسلسل. هذا يتيح تجميع دون اتصال.

بنهاية المهمة 3 ، تعتمد جميع أجهزة إرسال XLA على MLIR في تنفيذها.

التصميم التفصيلي

الخطوة 1: (المهمة 1) أكمل LHLO واجعل أجهزة البث القديمة تأخذ LHLO

تجعل هذه الخطوة جميع بواعث XLA / GPU الحالية تتفاعل مع عمليات MLIR. هذه الخطوة هي إعادة هيكلة خالصة و NFC.

هذه الخطوة ميكانيكية في الغالب ، ولكن من الجدير ملاحظة الاختلافات التالية بين HloComputation غير المتداخلة و LHLO:

  • لكل HloInstruction وصول مباشر إلى معاملاته (DAG لتدفق البيانات). على العكس من ذلك ، يمكن لكل مرجع LHLO الوصول فقط إلى مخازن المعامل الخاصة به (ثنائي بين العمليات والمخازن المؤقتة). يجب أن تخضع عمليات LHLO لسلاسل استخدام-الوصول للوصول إلى عمليات المعامل الخاصة بهم.
  • لا تصل البواعث القديمة غير المتداخلة بشكل تجريبي تقريبًا إلى معاملاتها. الاستثناء الوحيد هو kReduce.
  • لا تصل أجهزة البث القديمة غير المتداخلة إلى BufferAssignment إلا للحصول على الشرائح ، وليس للوصول إلى هياكل البيانات الإضافية مثل dataflow_analysis () أو alias_analysis (). يقوم llvm_ir ببناء تحليل الاسم المستعار الخاص به () بناءً على معلومات الشريحة.

الاستنتاج هو أن LHLO يجب أن يكون مناسبًا دون أي متاعب كبيرة.

الخطوة 2: دعم التنميط (اختياري)

هذه الخطوة مطلوبة فقط إذا بدأنا في تجاهل بعض منطق XLA Thunk (انظر الخطوة التالية).

قبل تشغيل أي بواعث قائمة على MLIR فعليًا ، نحتاج إلى ملف تعريف لبواعث تعتمد على MLIR.

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

هناك ما يقرب من ثلاثة أساليب لدعم التنميط في MLIR:

  • قم بتشغيل منشئ ملفات التعريف من طرف إلى طرف
  • أضف مرجع ملف شخصي لكل عملية تشغيل في LHLO ، باستخدام ملف تعريف محقون.

نهج "من طرف إلى طرف" شفاف لـ MLIR ، لكنه يعاني من نفس المشكلة التي تجعل XLA لا تستخدمه في المقام الأول: مكالمات المكتبة التي تم جمعها بواسطة المحلل (nvprof / ...) لا يمكن بسهولة أن تتصل بـ HLO العمليات. على سبيل المثال ، يطلق cuDNN عدة نواة لكل HLO ، ومن الصعب معرفة أي نواة تتوافق مع HLO.

يتطلب نهج "المحلل المحقون":

  • LHLO لاتخاذ ملف تعريف كمعلمة.
  • إدراج profile.start / profile.end قبل وبعد كل مرجع.
  • تمريرة من هذا يخفض الملف الشخصي. {start، end} إلى تنفيذ C ++.

لا يمكن إجراء التنميط الدقيق بسهولة للعمليات التي تم إنشاؤها بواسطة MLIR ، حيث:

  • لا يحتوي MLIR على موقت ، ولا يعتمد على TFRT / StreamExecutor.
  • لا يقوم MLIR باستدعاء وظائف C بسهولة مع معلمات معقدة.

الخطوة 3: (المهمة 2) ترحيل Thunks

كملاحظة ، هناك ما يقرب من ثلاثة أنواع من السوط:

  • KernelThunk ، الذي يطلق نواة.
  • أذرع التحكم في التدفق ، والتي لها منطق تدفق التحكم في المضيف (مشروط ، بينما ، من أجل ، التسلسل) وحبات جسم الإطلاق.
  • thunks المكتبة: cuDNN ، cuBLAS ، cuFFT ، NCCL ، إلخ.

الخطة هي:

  • اجعل Thunks (de) قابلاً للتسلسل.
  • ساعد في تحسين TFRT إلى حالة حيث يمكنها دعم هذه الدلالات.
  • مع تحسن الحالة ، قم بترحيل الأفخاذ بشكل تدريجي.

يتم ترتيب عناصر العمل هذه جزئيًا فقط. يجب تقييم أمر التنفيذ الفعلي / التوازي الهندسي أثناء سيره.

الخطوة 4: (المهمة 3) تم ترحيل ElementalIrEmitter

بمجرد أن يصبح التنميط جاهزًا ، يمكننا إكمال وضبط جميع بواعث ElementalIrEmitter في MLIR. ثم نقوم بتشغيلها بشكل افتراضي ، بافتراض أن جميع هذه البواعث القائمة على MLIR تستخدم دفقًا واحدًا.

لاحظ أنه من المفيد ترحيل ElementalIrEmitter الخاص بـ XLA / CPU أيضًا ، نظرًا لأنها تشترك في جزء كبير من التعليمات البرمجية.

مع الانتهاء من جميع المقاييس وصيد الأداء (TODO: تحديد تكافؤ الأداء) ، نقوم بتشغيل باعث العناصر الجديد القائم على MLIR ، وحذف ElementalIrEmitter القديم.

توفر هذه الخطوة أيضًا انتقالات سهلة للانصهار (عمليات متداخلة) للترحيل اللاحق.

الخطوة 5: دعم متعدد الدفق أو إسقاط

لا يمكننا حذف بعض الباعثات حتى ندعمها في MLIR ، أو نسقط الميزة. إنه قدر كبير نسبيًا من العمل في MLIR وكمية صغيرة من الكسب لـ XLA. يجب علينا التحقيق في المستخدمين الحاليين لمستخدمي XLA / GPU متعدد البث ، ومحاولة حذف هذه الميزة إذا كان ذلك معقولًا.

الخطوة 6: (المهمة 3) عمليات الجهاز المُرحّلة

تقوم هذه الخطوة بترحيل جميع العمليات غير المتداخلة ، ثم يمكننا حذف جميع البواعث غير المتداخلة.

هذا يستدعي إعادة كتابة / إعادة بناء kCopy و kReduce. لقد تم بالفعل عمل kReduce من أجل الكثير ، لذلك يبقى أن نرى الكمية الفعلية من العمل التي يجب القيام بها.