الأشكال والتخطيط

بنية عملية XLA

إليك مثالاً على هدف رفيع المستوى:

add.936 = bf16[8,1,1280,16384]{3,2,0,1:T(8,128)(2,1)}
          add(exponential.183, broadcast.3115)

يتكوّن ذلك من المكوّنات التالية:

• اسم العملية: add.936
  • هذا هو الاسم الفريد للعملية.
• الشكل: bf16[8,1,1280,16384]
  • هذا هو شكل الناتج الخاص بالعملية، حيث يكون نوع البيانات هنا هو bf16 والشكل هو [8,1,1280,16384].
• التنسيق (مع التجانب): 3,2,0,1:T(8,128)(2,1)
  • يصف هذا كيفية تخزين المصفوفة في الذاكرة. يشير 3,2,0,1 إلى ترتيب المحاور في الذاكرة (أي الترتيب حسب الأعمدة، أو حسب الصفوف، وما إلى ذلك)، ويشير T(8,128)(2,1) إلى التجانب والحشو المستخدَمَين.
  • يكون التنسيق اختياريًا. في حال عدم تحديد ذلك، لن يتم تقسيم الصورة إلى مربّعات، وسيتم افتراض ترتيب الأبعاد من الأكبر إلى الأصغر.
• العملية: add
  • العملية التي يتم تنفيذها في هذا المثال، القيمة هي Add، وهي مذكورة أيضًا في اسم العملية.
• الوسيطات: exponential.183، broadcast.3115
  • تتلقّى هذه العملية وسيطتَين يتم تحديدهما باستخدام اسمَيهما الفريدَين.

لنأخذ مثالاً آخر، وهو عملية دمج:

%fusion.3 = bf16[32,32,4096]{2,1,0:T(8,128)(2,1)S(1)}
            fusion(bf16[32,32,8192]{2,1,0:T(8,128)(2,1)S(1)} %fusion.32),
            kind=kCustom, calls=%all-reduce-scatter.3

بالإضافة إلى المكوّنات الموضّحة سابقًا، يتضمّن ذلك ما يلي:

• السمتان kind وcalls
  • تقدّم هذه السمة معلومات إضافية عن العملية التي يتم تنفيذها، وهي في هذه الحالة: الدمج.
• موقع الذاكرة (معرّف مساحة الذاكرة): S(1)
  • يشير ذلك إلى مساحة/موقع الذاكرة حيث يتم تخزين الصفيف. يشير S(1) هنا إلى أنّ هذا الصفيف موجود في VMEM (على وحدة TPU).
• تفاصيل الشكل والتنسيق لوسيطة الإدخال %fusion.32

توضّح الأقسام التالية "الأشكال" والتصميم ومعرّفات مساحة الذاكرة. يمكنك الاطّلاع على مزيد من المعلومات حول التقسيم إلى مربّعات في مقالة التنسيق المقسّم إلى مربّعات.

الأشكال

يصف ملف XLA ShapeProto proto (xla_data.proto) عدد السمات وحجمها ونوع البيانات لمصفوفة متعددة الأبعاد (array باختصار).

المصطلحات والرموز والاصطلاحات

• عدد الأبعاد الصحيح للصفيفة هو عدد الأبعاد التي يكون حجمها أكبر من 1.

• يتم ترقيم السمات من 0 إلى N-1 لمصفوفة ذات N سمات. حجم السمة هو عدد صحيح غير سالب. على وجه الخصوص، يكون الحجم 0 صالحًا. أرقام السمات هي تصنيفات اختيارية لتسهيل الاستخدام. لا يشير ترتيب أرقام الأبعاد هذه إلى ترتيب معيّن بين المحورين الرئيسي والثانوي في تصميم الشكل. يتم تحديد التنسيق من خلال LayoutProto.

• حسب الاصطلاح، يتم إدراج السمات بترتيب تصاعدي حسب رقم السمة. على سبيل المثال، بالنسبة إلى صفيف ثلاثي الأبعاد بحجم [A x B x C]، يكون حجم البُعد 0 هو A، وحجم البُعد 1 هو B، وحجم البُعد 2 هو C.

  تتيح بعض الأدوات المساعدة في XLA أيضًا استخدام الفهرسة السلبية المشابهة للغة Python: البُعد -1 هو البُعد الأخير (أي ما يعادل N-1 لصفيف N أبعاد). على سبيل المثال، بالنسبة إلى الصفيف الثلاثي الأبعاد الموضّح أعلاه، يكون حجم البُعد -1 هو C، وحجم البُعد -2 هو B، وهكذا.

• غالبًا ما تحتوي المصفوفات الثنائية والثلاثية والرباعية الأبعاد على أحرف معيّنة مرتبطة بالأبعاد. على سبيل المثال، بالنسبة إلى صفيفة ثنائية الأبعاد:

  • السمة 0: y
  • السمة 1: x

  بالنسبة إلى مصفوفة ثلاثية الأبعاد:

  • السمة 0: z
  • السمة 1: y
  • السمة 2: x

  بالنسبة إلى مصفوفة رباعية الأبعاد:

  • السمة 0: p
  • السمة 1: z
  • السمة 2: y
  • السمة 3: x

• تتلقّى الدوال في XLA API التي تتضمّن سمات هذه السمات بترتيب تصاعدي حسب رقم السمة. يتطابق ذلك مع الترتيب المستخدَم عند تمرير السمات كـ initializer_list، مثل:

  ShapeUtil::MakeShape(F32, {A, B, C, D})

  سيتم إنشاء شكل يتألف مصفوفة حجم أبعاده من التسلسل [A, B, C, D].

التنسيق

يوضّح ملف LayoutProto الأولي طريقة تمثيل الصفيف في الذاكرة. ويتضمّن الحقول التالية:

message LayoutProto {
  repeated int64 minor_to_major;
  int64 tail_padding_alignment_in_elements;
  ...
}

ترتيب السمات من الأقل أهمية إلى الأكثر أهمية

الحقل المطلوب الوحيد هو minor_to_major. يصف هذا الحقل ترتيب السمات من الأقل إلى الأكثر أهمية ضمن شكل معيّن. تمثّل القيم في minor_to_major ترتيبًا لأبعاد المصفوفة (من 0 إلى N-1 لمصفوفة ذات N بعد)، حيث تكون القيمة الأولى هي البُعد الأقل أهمية وصولاً إلى القيمة الأخيرة التي تمثّل البُعد الأكثر أهمية. البُعد الأقل أهمية هو البُعد الذي يتغيّر بشكل أسرع عند الانتقال بين عناصر المصفوفة المرتبة في الذاكرة الخطية.

على سبيل المثال، ضع في اعتبارك مصفوفة ثنائية الأبعاد بالحجم [2 x 3]:

a b c
d e f

في هذا المثال، القيمة 0 هي الحجم 2، والقيمة 1 هي الحجم 3. إذا كان الحقل minor_to_major في التنسيق هو [0, 1]، يكون المقياس 0 هو المقياس الأقل أهمية، ويكون المقياس 1 هو المقياس الأكثر أهمية. يتوافق ذلك مع التنسيق التالي في الذاكرة الخطية:

a d b e c f

إنّ ترتيب الأبعاد من الأصغر إلى الأكبر من 0 إلى N-1 يشبه الترتيب حسب الأعمدة (بالنسبة إلى الأبعاد الثنائية). بافتراض ترتيب رتيب للأبعاد، يمكننا الإشارة إلى هذا التنسيق في الرمز ببساطة على أنّه "البعد 0 هو البعد الأصغر".

من ناحية أخرى، إذا كان الحقل minor_to_major في التنسيق هو [1, 0]، سيكون التنسيق في الذاكرة الخطية كما يلي:

a b c d e f

إنّ ترتيب الأبعاد من الأصغر إلى الأكبر، أي من N-1 إلى 0، بالنسبة إلى مصفوفة N ذات الأبعاد، يشبه الترتيب حسب الصف (بالنسبة إلى المصفوفات الثنائية الأبعاد). بافتراض ترتيب رتيب للأبعاد، يمكننا الإشارة إلى هذا التصميم في الرمز ببساطة على أنّه "البعد 0 هو البعد الرئيسي".

الترتيب التلقائي من الأصغر إلى الأكبر

التنسيق التلقائي للأشكال التي تم إنشاؤها حديثًا هو "ترتيب الأبعاد هو من الأكبر إلى الأصغر" (أي [N-1, ..., 0]).

مساحة متروكة

يحدّد الحقل tail_padding_alignment_in_elements محاذاة مصفوفة tiled من حيث عدد العناصر. بعد تطبيق التجانب، ستتم إضافة عناصر مبطّنة في نهاية التصميم إلى أن يصبح العدد الإجمالي للعناصر مضاعفًا لهذه القيمة.

الفهرسة في المصفوفات

تقدّم الفئة IndexUtil في index_util.h أدوات مساعدة للتحويل بين الفهارس المتعددة الأبعاد والفهارس الخطية مع توفّر شكل وتصميم. تتضمّن الفهارس المتعدّدة الأبعاد int64 فهرسًا لكل سمة. الفهارس الخطية هي قيمة int64 واحدة يتم استخدامها كفهرس في المخزن المؤقت الذي يحتوي على الصفيف. راجِع shape_util.h وlayout_util.h في الدليل نفسه للاطّلاع على الأدوات المساعدة التي تسهّل إنشاء الأشكال والتصاميم ومعالجتها.

معرّفات "مساحة الذاكرة"

في HLO، يمكن إضافة تعليق توضيحي إلى كل مصفوفة باستخدام معرّف مساحة تخزين، ويُكتب على النحو التالي: S(n).

• يشير S(0) (غالبًا ما يتم حذفه) إلى ذاكرة النطاق الترددي العالي (HBM) للجهاز.
• يشير S(1) إلى الذاكرة الافتراضية على الجهاز (VMEM).
• تمثّل S(2) وS(3) وما إلى ذلك مساحات ذاكرة إضافية خاصة بالجهاز.
• يشير S(5) إلى ذاكرة المضيف.