يصف هذا المستند تحليل فهرسة HLO، الذي يتيح لك حساب خرائط الفهرسة لعمليات HLO بشكل رمزي. خريطة الفهرسة هي دالة تحدد مؤشرات أحد الموترين بمؤشرات أخرى، مثل مؤشرات مخرجات تعليمات HLO إلى مؤشرات مدخلات تعليمات HLO أو العكس.
مثال
لبثّ من tensor<20xf32> إلى tensor<10x20x30xf32>
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
خريطة الفهرسة من الإخراج إلى الإدخال هي \((i, j, k) \mapsto (j)\) $i \in [0, 10]\(, \)j \in [0, 20]\( and \)k \in [0, 30]$.
الحافز
تستخدم وحدة معالجة الرسومات XLA عدة حلول مخصصة لشرح مخططات الدمج واستخدام المعاملات والتقسيم إلى أجزاء (مزيد من التفاصيل أدناه). الهدف من تحليل الفهرسة هو توفير مكون قابل لإعادة الاستخدام لحالات الاستخدام هذه. يستند تحليل الفهرسة إلى البنية الأساسية لخريطة أفين في MLIR ويضيف دلالات HLO.
تجانس
يصبح التفكير في اندماج الذاكرة ممكنًا في الحالات غير البسيطة عندما نعرف العناصر/الشرائح من المدخلات التي تتم قراءتها لحساب عنصر من المخرجات.
استخدام خاص
يشير الاستخدام الخاص في XLA إلى المقدار الذي يتم استخدامه لكل إدخال في التعليمات بافتراض استخدام مخرجاته بالكامل. حاليًا، لا يتم احتساب الاستخدام أيضًا لحالة عامة. يتيح تحليل الفهرسة حساب الاستخدام بدقة.
التجزئة إلى أقسام عمودية:
المربّع/الشريحة عبارة عن مجموعة فرعية مستطيلة الشكل للغاية من متوتر يحدده الإزاحة والأحجام والخطوات. نشر المربّع هو طريقة لحساب مَعلمات المربّعات للمنتج أو مستهلك العملية باستخدام مَعلمات التجانب في العملية نفسها. هناك مكتبة بالفعل تقوم بتنفيذ ذلك لـ softmax وdot. يمكن جعل انتشار المربّع أكثر عمومية وقوة إذا تم التعبير عنه من خلال خرائط الفهرسة.
الدالة والنطاق
خريطة الفهرسة هي دالة \(\boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\) تربط فهرسًا متعددًا \(\boldsymbol{d}\) لوحدة مغناطيسية \(A\) عناصر أو نطاقات متدرج \(B\). تشير المعلمة \(\boldsymbol{s}\) إلى نطاقات الفهارس المتوفرة في الموتر \(B\)، ولكن ليس في الموتر \(A\).
على سبيل المثال، إذا كان لدينا تخفيض من tensor<2x4x8x16xf32> إلى
tensor<4x8xf32>، تكون خريطة الفهرسة من الإخراج الثنائي الأبعاد إلى الإدخال رباعي الأبعاد هي
\((d_0, d_1) \mapsto (s_0, d_0, d_1, s_1)\)، حيث \(d_i\) هي مَعلَمات
السمة التي تتوافق مع مؤشرات موصّل الإخراج. ترمّز المعلَمات \(s_j\)
قيمًا متعددة، أي لحساب \((d_0, d_1)\) عنصر من الناتج، نحتاج
إلى \((s_0, d_0, d_1, s_1)\) عناصر من المدخلات، حيث \(s_0 \in [0, 2)\) و\(s_1 \in [0, 16)\).
ويمكن إنشاء هذا التعيين من سمات تعليمات HLO أو تعيينات التعليمات غير المُدمَجة للحصول على فهرسة من أجل الدمج. يحتوي التعيين أيضًا على مجال يحدد عناصر المثلث الموجود في التعيين.
\[ \begin{eqnarray} \boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\; &s.t.& \\ \boldsymbol{lb}_d &\leq& \boldsymbol{d} \leq \boldsymbol{ub}_d \\ \boldsymbol{lb}_s &\leq& \boldsymbol{s} \leq \boldsymbol{ub}_s \\ \boldsymbol{lb}_g &\leq& \boldsymbol{g}(\boldsymbol{d}, \boldsymbol{s}) \leq \boldsymbol{ub}_g \end{eqnarray} \]
وبما أننا نريد تقليل إعادة الحساب، فإننا بحاجة إلى مكتبة للعمليات الحسابية الرمزية. تعتمد دالة XLA على MLIR، لذلك نستخدم mlir::AffineMAP بدلاً من كتابة مكتبة حسابية رمزية.
يبدو أنّ AffineMap العادي
(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)
هناك نوعان من المَعلمات بسهولة في "AffineMap"، هما: السمات والرموز التي يمكننا استخدامها \(\boldsymbol d\) و \(\boldsymbol s\) على التوالي.
لا تحتوي السمة AffineMap على أي بيانات وصفية حول نطاقات الأبعاد، لذا علينا أن نقدّم هذه البيانات بأنفسنا.
struct Range {
int64_t lower_bound;
int64_t upper_bound;
};
struct IndexingMap {
mlir::AffineMap affine_map;
std::vector<Range> dimension_ranges;
std::vector<Range> symbol_ranges;
llvm::DenseMap<mlir::AffineExpr, Range> expr_ranges;
};
تعمل dim_ranges على ترميز قيود المربّع الشاملة لمَعلمات الأبعاد \(\boldsymbol{d}\) في خريطة الفهرسة، والتي تتزامن عادةً مع شكل منظّم الإخراج للعمليات مثل تبديل الموضع والتقليل والعنصر والنقطة، ولكن هناك بعض الاستثناءات مثل HloConcatenateInstruction.
تشفّر symbol_ranges القيم المحتملة التي \(\boldsymbol {s}\) يمكن أن تتخذها المعلَمات.
لنلقِ نظرة على مثال لفهم معنى كل ما سبق.
فهرسة الخرائط للعمليات غير المدمجة
عناصر
بالنسبة إلى العمليات القائمة على العناصر، فإن خريطة الفهرسة هي هوية.
p0 = f32[10, 20] parameter(0)
p1 = f32[10, 20] parameter(1)
add = f32[10, 20] add(p0, p1)
المخرجات المدخلة للخرائط:
- الناتج -> الإدخال_0: \((d_0, d_1) \mapsto (d_0, d_1)\) لرمز $\boldsymbol{d} \in [0,9] \tفترة [0، 19]\(, i.e. \)\boldsymbol{d} \in {\rm Dom}(output)$
- المخرجات -> الإدخال_1: \((d_0, d_1) \mapsto (d_0, d_1)\) لـ $\boldsymbol{d} \in {\rm Dom} (output)$
الإدخال في خرائط الإخراج
- الإدخال_i -> الناتج: \((d_0, d_1) \mapsto (d_0, d_1)\) لـ $\boldsymbol{d} \in {\rm Dom}(output)$
إعلان الرسائل
يعني البث أنّه ستتم إزالة بعض الأبعاد عندما نعيّن الناتج في المدخلات ونضيفها عندما نربط المدخلات بالمخرجات.
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
المخرجات المدخلة للخريطة:
- المخرجات -> الإدخال: \((d_0, d_1, d_2) \mapsto (d_1)\) لـ $\boldsymbol{d} \in {\rm Dom}(output)$
إدخال خريطة الإخراج
- إدخال -> إخراج: \((d_0) \mapsto (s_0, d_1, s_1)\) لرمز $\boldsymbol{d} \in {\rm Dom}(output)\( and \)\boldsymbol{s} \in [0, 9] \tفترة [0، 29]$.
لاحظ أن لدينا الآن \(\boldsymbol s\) على الجانب الأيمن لتعيين الإدخال إلى الإخراج. هذه هي الرموز التي تمثل نطاقات القيم. على سبيل المثال، في هذه الحالة تحديدًا، يتم تعيين كل عنصر من عناصر الإدخال ذي الفهرس \(d_0\) إلى شريحة 10×1×30 من الناتج.
كونستانت ويوتا
وهي لا تتضمن أي معلَمات إدخال، وبالتالي لا حاجة إلى احتساب الفهرسة.
تبديل الموضع
خريطة الفهرسة لنقل الموضع هي تبديل لأبعاد الإدخال/الإخراج.
p0 = f32[3, 12288, 6, 128] parameter(0)
transpose = f32[3, 6, 128, 12288] transpose(p0), dimensions={0, 2, 3, 1}
المخرجات المدخلة للخريطة:
- المخرجات -> الإدخال: \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_3, d_1, d_2)\) لـ \(\boldsymbol{d} \in {\rm Dom}(output)\)
إدخال خريطة الإخراج:
- الإدخال -> المخرجات: \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_2, d_3, d_1)\) لـ \(\boldsymbol{d} \in {\rm Dom}(input)\)
تسديدة عكسية
تؤدي خريطة الفهرسة للعكس إلى تغيير الأبعاد التي تم إرجاعها إلى $upper_bound(d_i) - d_i$:
p0 = f32[1, 17, 9, 9] parameter(0)
reverse = f32[1, 17, 9, 9] reverse(p0), dimensions={1, 2}
المخرجات المدخلة للخريطة:
- المخرجات -> الإدخال: $(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(output)$
إدخال خريطة الإخراج:
- الإدخال -> الإخراج: $(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(input)$
(Variadic)reset
يحتوي الاختزال المتنوّع على عدة مدخلات والعديد من وحدات الإدخال، وتضيف الخريطة من الناتج إلى الإدخال الأبعاد المخفَّضة. لذا، فهي تعمل كعكس البث بشكل ما.
p0 = f32[256,10] parameter(0)
p0_init = f32[] constant(-inf)
p1 = s32[256,10] parameter(1)
p1_init = s32[] constant(0)
reduce = (f32[10], s32[10]) reduce(p0, p1, p0_init, p1_init),
dimensions={0}, to_apply=min
المخرجات المدخلة للخرائط:
- المخرجات -> enter_j: \((d_0) \mapsto (s_0, d_0)\) لـ $\boldsymbol{d} \in {\rm Dom}(output)\( and \)\boldsymbol{s} \in [0, 9]$
- المخرجات -> init_j: \((d_0) \mapsto ()\) لـ $\boldsymbol{d} \in {\rm Dom}(output)$
الإدخال في خرائط الإخراج:
- إدخال_i -> مخرجات_j: \((d_0, d_1) \mapsto (d_1)\) لـ $\boldsymbol{d} \in {\rm Dom}(input)$
- init_i -> generate_j: \(() \mapsto (s_0)\) لـ \(\boldsymbol{s} \in [0, 9]\)
لمدة \(i, j = 0, \ldots, INPUT\\_COUNT\).
شريحة
تؤدي الفهرسة من الإخراج إلى المدخل إلى الحصول على خريطة فهرسة تدريجية صالحة لكل عنصر من عناصر الناتج. يقتصر التعيين من المدخل إلى الإخراج على نطاق تدريجي من العناصر في الإدخال.
p0 = f32[10, 20, 50] parameter(0)
slice = f32[5, 3, 25] slice(f32[10, 20, 50] p0),
slice={[5:10:1], [3:20:7], [0:50:2]}
المخرجات المدخلة للخريطة:
- المخرجات -> الإدخال: \((d_0, d_1, d_2) \mapsto (d_0 + 5, 7d_1 + 3, 2d_2)\) لـ \(\boldsymbol{d} \in {\rm Dom}(output)\)
إدخال خريطة الإخراج:
- الإدخال -> المخرجات: \((d_0, d_1, d_2) \mapsto (d_0, d_1 / 7, d_2 / 2)\) لـ \(\boldsymbol{d} \in [5, 9] \times [3, 19] \times [0, 49]\) مع الخطوات $[1, 7, 2]$.
يُحدَّد اللون لاحقًا: فهرسة الإدخال إلى الإخراج
إعادة التشكيل
إعادة الأشكال تأتي بنكهات مختلفة.
تصغير الشكل
هذه إعادة تشكيل "خطي" من N-D إلى 1D.
p0 = f32[4,8] parameter(0)
reshape = f32[32] reshape(p0)
المخرجات المدخلة للخريطة:
- المخرجات -> الإدخال: \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) لـ $\boldsymbol{d} \in {\rm Dom}(output)$
إدخال خريطة الإخراج:
- الإدخال -> المخرجات: \((d_0, d_1) \mapsto (8 d_0 + d_1)\) لـ $\boldsymbol{d} \in {\rm Dom}(input)$.
توسيع الشكل
هذه عملية عكسية لـ "تصغير الشكل"، تعيد تشكيل إدخال أحادي الأبعاد في مخرجات N-D.
p0 = f32[32] parameter(0)
reshape = f32[4, 8] reshape(p0)
المخرجات المدخلة للخريطة:
- المخرجات -> الإدخال: \((d_0, d_1) \mapsto (8 d_0 + d_1)\) لـ $\boldsymbol{d} \in {\rm Dom}(output)$
إدخال خريطة الإخراج:
- الإدخال -> المخرجات: \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) لـ $\boldsymbol{d} \in {\rm Dom}(input)$.
إعادة تشكيل عامة
هذه هي عمليات إعادة الشكل التي لا يمكن تمثيلها كشكل توسيع أو تصغير فردي. يمكن تمثيلها فقط كتركيبة من شكلَين أو أكثر من أشكال التوسيع أو التصغير.
المثال 1: إزالة الخط الخطي
p0 = f32[4,8] parameter(0)
reshape = f32[2, 4, 4] reshape(p0)
يمكن تمثيل إعادة الشكل هذه على شكل تركيبة لشكل تصغير
من tensor<4x8xf32> إلى tensor<32xf32> ثم توسيع نطاق الشكل إلى
tensor<2x4x4xf32>.
المخرجات المدخلة للخريطة:
- المخرجات -> الإدخال: $(d_0, d_1, d_2) \mapsto (2d_0 + (4d_1 + d_2) / 8, 4d_1 + d_2) \mod 8)$
لمدة \(\boldsymbol{d} \in {\rm Dom}(output)\)
إدخال خريطة الإخراج:
- الإدخال -> المخرجات: $(d_0, d_1) \mapsto ((8d_0 + d_1) / 16, ((8d_0 + d_1) \mod 16) / 4, d_1 \mod 4)$
لمدة \(\boldsymbol{d} \in {\rm Dom}(input)\).
المثال 2: الأشكال الفرعية الموسَّعة والمصغرة
p0 = f32[4, 8, 12] parameter(0)
reshape = f32[32, 3, 4] reshape(p0)
يمكن تمثيل إعادة الشكل هذه كتركيبة من إعادة شكلين. يعمل العنصر الأول على تصغير البُعد الخارجي من tensor<4x8x12xf32> إلى tensor<32x12xf32>، في حين يؤدي البُعد الثاني إلى توسيع البُعد الأعمق tensor<32x12xf32> إلى "tensor<32x3x4xf32>".
المخرجات المدخلة للخريطة:
- المخرجات -> الإدخال: \((d_0, d_1, d_2) \mapsto (d_0 / 8, d_0 \mod 8, 4d_1 + d_2)\) لـ \(\boldsymbol{d} \in {\rm Dom}(output)\)
إدخال خريطة الإخراج:
- الإدخال -> المخرجات: \((d_0, d_1, d_2) \mapsto (8d_0 + d_1, d_2 / 4, d_2 \mod 4)\) لـ \(\boldsymbol{d} \in {\rm Dom}(input)\).
إرسال Bitcast
يمكن تمثيل إجراء تحويل بت كبديل كسلسلة من عمليات تبديل الموضع وإعادة الشكل. وبالتالي، فإن خرائط الفهرسة هي مجرد تركيبة من خرائط الفهرسة لهذه التسلسل.
سلسلة
يتم تحديد تعيين مخرجات إلى إدخال لـ concat لجميع المدخلات، ولكن مع النطاقات غير المتداخلة، أي أنه سيتم استخدام إدخال واحد فقط في كل مرة.
p0 = f32[3,50] parameter(0)
p1 = f32[3,30] parameter(1)
concat = f32[3,80] concatenate(f32[3,50] p0, f32[3,30] p1),
dimensions={1}
المخرجات المدخلة للخريطة:
- المخرجات -> الإدخال 1:
\((d_0, d_1) \mapsto (d_0, d_1)\) من أجل \(\boldsymbol{d} \in [0, 2] \times [0, 49]\)
- المخرجات -> الإدخال 2:
\((d_0, d_1) \mapsto (d_0, d_1 - 50)\) لرمز $\boldsymbol{d} \في [0، 2] \tفترة [50، 79]$
إدخالات خريطة الإخراج:
- الإدخال 1 -> إخراج: \((d_0, d_1) \mapsto (d_0, d_1)\) لـ $\boldsymbol{d} \in {\rm Dom}(input_1)$.
- الإدخال 2 -> إخراج: \((d_0, d_1) \mapsto (d_0, d_1 + 50)\) لـ $\boldsymbol{d} \in {\rm Dom}(input_2)$.
نقطة (تم تنفيذ إخراج إلى الإدخال
تتشابه خرائط الفهرسة الخاصة بالنقطة كثيرًا مع خرائط الاختزال.
p0 = f32[4, 128, 256] parameter(0)
p1 = f32[4, 256, 64] parameter(1)
dot = f32[4, 128, 64] dot(p0, p1),
lhs_batch_dims={0}, rhs_batch_dims={0},
lhs_contracting_dims={2}, rhs_contracting_dims={1}
المخرجات التي يتم إدخالها على الخرائط:
- المخرجات -> enter_1: \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) لـ \(\boldsymbol{d} \in {\rm Dom}(output)\) و \(\boldsymbol{s} \in [0, 255]\)
- المخرجات -> enter_2: \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_2)\) لـ \(\boldsymbol{d} \in {\rm Dom}(output)\) و \(\boldsymbol{s} \in [0, 255]\)
مدخلات خرائط الإخراج:
- الإدخال_1 -> المخرجات: \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) لـ \(\boldsymbol{d} \in {\rm Dom}(input_1)\) و \(\boldsymbol{s} \in [0, 63]\)
- الإدخال_2 -> المخرجات: \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_1)\) لـ \(\boldsymbol{d} \in {\rm Dom}(input_2)\) و \(\boldsymbol{s} \in [0, 127]\)
تقليل المهلة (يُحدَّد اللون لاحقًا)
وسادة (يُحدَّد اللون لاحقًا)
فهرسة الخرائط لميزة الدمج
خريطة الفهرسة الخاصة بعملية الاندماج هي تركيبة من خرائط الفهرسة لكل عملية في المجموعة. قد تتم قراءة بعض الإدخالات عدة مرات بأنماط وصول مختلفة.
إدخال واحد والعديد من خرائط الفهرسة
إليك مثال على \(p_0 + p_0^T\)
f {
p0 = f32[1000, 1000] parameter(0)
transpose_p0 = f32[1000, 1000]{0, 1} transpose(p0), dimensions={1, 0}
ROOT a0 = f32[1000, 1000] add(p0, transpose_p0)
}
ستكون خرائط فهرسة إخراج الإخراج إلى الإدخال p0 هي $(d_0, d_1) \mapsto (d_0,
d_1)\( and \)(d_0, d_1) \mapsto (d_1, d_0)$. ويعني ذلك أنّه لحساب عنصر واحد من الناتج، قد نحتاج إلى قراءة معلَمة الإدخال مرتين.
إدخال واحد، خريطة فهرسة تمت إزالة تكرارها
هناك حالات تكون فيها خرائط الفهرسة متطابقة، على الرغم من أنّها غير واضحة على الفور.
f {
p0 = f32[20, 10, 50] parameter(0)
lhs_transpose_1 = f32[10, 20, 50] transpose(p0), dimensions={1, 0, 2}
lhs_e = f32[10, 20, 50] exponential(lhs_transpose_1)
lhs_transpose_2 = f32[10, 50, 20] transpose(lhs_e), dimensions={0, 2, 1}
rhs_transpose_1 = f32[50, 10, 20] transpose(p0), dimensions={2, 1, 0}
rhs_log = f32[50, 10, 20] exponential(rhs_transpose_1)
rhs_transpose_2 = f32[10, 50, 20] transpose(rhs_log), dimensions={1, 0, 2}
ROOT add = f32[10, 50, 20] add(lhs_transpose_2, rhs_transpose_2)
}
خريطة فهرسة إخراج إلى الإدخال لـ p0 في هذه الحالة هي $(d_0, d_1, d_2)
\mapsto (d_2, d_0, d_1)$ فقط.
Softmax
ربط فهرسة الإخراج إلى الإدخال لـ parameter 0 لـ softmax:
- \((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\)
- \((d_0, d_1, d_2)[s_0] \mapsto (d_0, d_1, s_0)\)
لـ \(\boldsymbol{d} \in {\rm Dom}(output)\) و \(\boldsymbol{s} \in [0, 124]\) إلى البُعد الداخلي للإدخال.
أداة تبسيط الفهرسة
لا يمكن لأداة تبسيط عملية التحميل التلقائية للسمة mlir::AffineMap
أيّ افتراضات حول نطاقات الأبعاد أو الرموز. لذلك، لا يمكن
تبسيط التعبيرات باستخدام mod وdivبفعالية.
يمكننا الاستفادة من المعرفة حول الحدود الدنيا والعلوية للتعبيرات الفرعية في الخرائط الارتباطية لتبسيطها أكثر.
يمكن للمبسط إعادة كتابة التعبيرات التالية.
- \((d_0, d_1) \mapsto (d_0 + d1 / 16, d1 \mod 16)\) بالنسبة إلى $\boldsymbol{d} \في [0, 6] \tمدة [0، 14]\( becomes \)(d_0، d_1) \mapsto (d_0، d_1)$
- $(d_0, d_1, d_2) \mapsto ((100d_0 + 10d_1 + d_2) /100, ((100d_0 + 10d_1 + d_2) \mod 100) / 10, d_2 \mod 10 ,$_2,\( for \)d_i)\( for \)d_i)\( becomes \)
- $(d_0, d_1, d_2) \mapsto ((16d_0 + 4d_1 + d_2) /8, (16d_0 + 4d_1 + d_2) \mod 8)\( for \)d_i \في [0، 9]\( becomes \)(d_0, d_1، /d_1+، d_2, +d_2) (16d_0 + d_2)
- \((d_0, d_1) \mapsto (-(-11d_0 - d_1 + 109) / 11 + 9)\) لـ $\boldsymbol{d} \في [0, 9] \tميات [0، 10]\( becomes \)(d_0، d_1) \mapsto (d_0)$.
يتيح لنا مبسّط فهرسة الخريطة معرفة أن بعض أشكال إعادة التشكيل المتسلسلة في HLO تلغي بعضها البعض.
p0 = f32[10, 10, 10] parameter(0)
reshape1 = f32[50, 20] reshape(p0)
reshape2 = f32[10, 10, 10] reshape(reshape1)
بعد إنشاء خرائط الفهرسة وتبسيطها سنحصل على
\((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\).