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

MLIR CodeGen لـ XLA

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

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

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

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

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

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

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

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

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

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

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

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

ElementalIrEmitter فريد من نوعه في:

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

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

  1. يبقى رمز الجهاز مثل LLVM IR.
  2. إعادة تشكيل الباعث القديم ليكون مثل LHLO -> لهجة MLIR LLVM:
    • (التكلفة) سيكون عملاً بعيدًا إذا أردنا الانتقال في النهاية إلى المعيار.
    • (فائدة) إنه سهل وميكانيكي. يمكن القيام به في فترة قصيرة.
    • (فائدة) لا يستفيد أكثر من (1).
  3. تكون بواعث إعادة البناء القديمة مثل LHLO -> MLIR GPU + Standard + Loops:
    • (التكلفة) يقدم رفع البواعث الحالية إلى المعيار بعض التحديات. يجب تحويل المؤشرات و GEPs إلى 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 -> تشغيل بواعث الأجهزة القديمة.
  • تذهب العمليات الأولية للمرسل إلى (4) ، ولكن ليس بشكل تدريجي. لا توجد طريقة للقيام بذلك المرجع من خلال المرجع ، لأن جميع العمليات المنبعثة بشكل أساسي مرتبطة بنفس الرسم البياني. يمكن أن يكون هذا العمل أيضًا بمثابة نقطة توحيد للعديد من القوى الجارية (مولد النواة ، Linalg).
  • جميع العمليات الأخرى تذهب إلى (1). كهدف ممتد ، قد يتم ترحيلهم إلى (3) أو (4).

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

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

ترتيب الأولويات هو: Task1 (LHLO للبواعث القديمة)> المهمة 2 (Thunks)> المهمة 3 (بواعث MLIR).

بحلول نهاية المهمة 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 يبني alias_analysis () الخاص به بناءً على معلومات الشريحة.

الاستنتاج هو أن 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

كملاحظة ، هناك ما يقرب من ثلاثة أنواع من 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 من أجل الكثير ، لذلك يبقى المقدار الفعلي للعمل الذي يتعين القيام به غير معروف.