الديناميكية في StableHLO

تم توضيح الحالة الحالية للديناميكية بشكل أكثر رسمية في مستند RFC حول الديناميكية، وتقدّم هذه الصفحة نظرة عامة على مستوى عالٍ حول مستند RFC وتناقش واجهات برمجة التطبيقات والأدوات المهمة للتفاعل مع البرامج الديناميكية.

المصطلحات المتعلّقة بالتأثير الديناميكي ونظرة عامة على الدعم

أولاً، سنتناول بعض المصطلحات التي ستظهر في هذا المستند، بالإضافة إلى مقدمة موجزة عن استخدامها في StableHLO:

السمات الديناميكية

تشير السمات الديناميكية إلى أيّ سمة لا يُعرَف حجمها. في StableHLO، نمثّل السمات الديناميكية باستخدام ?، أي tensor<16x?xf32>.

الديناميكية المحدودة

يشير الثبات الديناميكي إلى سمة ديناميكية لها قيمة لها حدّ صعِد معلوم. يكون هذا مفيدًا بشكل عام لتوفير مساحة متروكة للمتسخ أثناء التنفيذ. في StableHLO، نمثل الديناميكا المشروطة باستخدام #stablehlo.bounds كترميز تنسور، أي متّصِل بالرتبة 2 مع بُعد ديناميكي واحد مرتبط بالرقم 16 ويمكن تمثيل الآخر بدون حدود بالشكل tensor<?x?xf32, #stablehlo.bounds<16, ?>>.

يمكن لـ StableHLO تمثيل الديناميكية المحدودة، ولكن هناك دعم محدود للإطار العمل، والذي ينبع من TensorFlow، مع بعض الدعم في PyTorch/XLA.

ديناميكية غير محدودة

يشير السلوك الديناميكي غير المحدود، كما يوحي الاسم، إلى سمة ديناميكية ليس لها حدود معروفة للحجم. هذا النوع من الديناميكية شائع جدًا في StableHLO، مع دعم JAX وPyTorch/XLA وTF، ويتم استخدامه غالبًا لتصدير النماذج ذات حجم الدفعة الديناميكي أو طول التسلسل.

وفي StableHLO، يمكننا ببساطة إزالة ترميز الحدود لهذا الشكل من الديناميكية، أي tensor<?x?xf32>.

أشكال متعددة للأشكال

الشكل المتعدّد الأشكال هو مصطلح اكتسبناه من JAX.

هناك نتيجتان رئيسيتان لتشكيل التعددية:

1. تعود كل الديناميكية في البرنامج إلى وسيطات الإدخال.
2. ترتبط كلّ الديناميكية بأشكال مصفوفة تينسور فقط، أي أنّها لا تعتمد على البيانات.

باستخدام هاتين القاعدتَين، بعد معرفة الأشكال الثابتة للبرنامج، يمكننا أخذ برنامج ديناميكي وتحسينه بالكامل إلى برنامج ثابت لأجل الترجمة (راجِع "عمليات المترجم لتحسين البرامج الديناميكية").

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

الديناميكية المستندة إلى البيانات

يشير السلوك الديناميكي المستند إلى البيانات إلى أحجام السمات الديناميكية التي تتعلّق بالبيانات داخل مصفوفة تانسور. المثال الأساسي هو دالة nonzeros التي تُرجِع فهارس جميع العناصر التي تكون 0 في قيمة مصفوفة. لا يمكن معرفة الشكل بدون تقييم البيانات، ولكن يمكن غالبًا تجميعه باستخدام الديناميكية المحدودة، مع إنفاق ذاكرة إضافية على حجم مصفوفة السواتل المحتملة للإخراج.

يمكن وضع نماذج لكثير من العمليات الديناميكية التي تعتمد على البيانات باستخدام الديناميكية المحدودة، حيث يتم تحديد حدّ أقصى لحجم مصفوفة تينسور، وستنفذ الأجهزة بشكل عام هذا من خلال إضافة عناصر إلى مصفوفة تينسور. تتوفّر حاليًا بعض الإمكانات للاستفادة من الأداء الديناميكي المستنِد إلى البيانات في PyTorch/XLA وTensorFlow، ولكن لا تتتبّع JAX حاليًا العمليات التي منجرِ إلى الأداء الديناميكي المستنِد إلى البيانات.

تصدير البرامج بأبعاد ديناميكية

اطّلِع على الأدلة التعليمية حول StableHLO للحصول على معلومات عن كيفية تصدير البرامج التي تتضمّن أحجام دفعات أو أطوال تسلسلات ديناميكية:

• البرنامج التعليمي حول JAX > التصدير باستخدام حجم حزمة ديناميكي
• دليل PyTorch/XLA التعليمي > التصدير باستخدام حجم الدفعة الديناميكية

أدوات تجميع البيانات لتحسين البرامج الديناميكية

إزالة مسار الإرسال لمرور المحتوى الديناميكي

هناك بعض عمليات التحسين المفيدة لتحسين الأشكال، ومن المفيد أنّها كلها مجمّعة في مسار معالجة createStablehloRemoveDynamismPipeline:

void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
                                           TypeRange refinedTypes);

البطاقات الفردية لتحسين الأداء الديناميكي

في ما يلي الخطوات التي تكون مفيدة بشكلٍ فردي لتحسين الشكل:

• stablehlo-refine-arguments لاستبدال مَعلمات الإدخال بأنواع مصفوفات كثيفة محدّدة
• stablehlo-refine-shapes لنشر مَعلمة الإدخال الجديدة لتشكيل المعلومات في البرنامج بأكمله.
• stablehlo-canonicalize-dynamism لاستبدال ops الديناميكية بنُسخها الثابتة

اطّلِع على المستندات المرتبطة للحصول على المعلومات والأمثلة المحدّثة.

مثال: ما مدى فائدة الأسلوب الديناميكي، وكيف يمكنني استخدامه؟

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

نموذج add_one الثابت

سنستخدم نموذج add_one البسيط التالي لتوضيح ذلك:

def add_one(x):
  return x + 1

عند تتبُّع المسار باستخدام tensor<4xf32>، سنحصل على برنامج StableHLO التالي:

// File: add_one.mlir
func.func @add_one(%arg0: tensor<4xf32>) -> tensor<4xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<4xf32>
  %0 = stablehlo.add %arg0, %cst : tensor<4xf32>
  return %0 : tensor<4xf32>
}

لن يعمل هذا النموذج إلا مع وسيطات الإدخال التي تتضمّن الشكل tensor<4xf32>. إذا غيّرنا حجم الحزمة أو طول التسلسل، سنحتاج إلى إعادة تتبُّع رمز المصدر وإعادة خفضه إلى StableHLO، ولا يمكننا ضمان أنّه لا يزال بإمكاننا الوصول إلى رمز المصدر.

نموذج add_one الديناميكي

وهنا يأتي دور الديناميكية متعددة الأشكال للأشكال. بدلاً من ذلك، يمكن لـ JAX و PyTorch/XLA إصدار نموذج add_one باستخدام رمز برمجي صالح ديناميكيًا والذي سيُرسِل الثابت لمطابقة شكل الإدخال الديناميكي على النحو التالي:

// File: add_one_dynamic.mlir
func.func public @main(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  %cst = stablehlo.constant dense<1.0> : tensor<f32>
  %0 = stablehlo.get_dimension_size %arg0, dim = 0 : (tensor<?xf32>) -> tensor<i32>
  %1 = stablehlo.reshape %0 : (tensor<i32>) -> tensor<1xi32>
  %2 = stablehlo.dynamic_broadcast_in_dim %cst, %1, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
  %3 = stablehlo.add %arg0, %2 : tensor<?xf32>
  return %3 : tensor<?xf32>
}

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

تحسين النموذج الديناميكي

على سبيل المثال، يمكن أن يؤدي ترتيب البطاقات التالية إلى تحسين هذا البرنامج بالكامل:

stablehlo-opt add_one_dynamic.mlir \
  --stablehlo-refine-arguments='types=tensor<16xf32>' \
  --stablehlo-refine-shapes \
  --stablehlo-canonicalize-dynamism

وإليك كيفية تحويل البرنامج بشكل تدريجي:

// After stablehlo-refine-arguments: Inputs updated, shapes not propagated
func.func public @main(%arg0: tensor<16xf32>) -> tensor<?xf32> {
  %c = stablehlo.constant dense<16> : tensor<1xi64>
  %0 = stablehlo.custom_call @stablehlo.shape_refinement_operand_wrapper(%arg0, %c) {indices_of_shape_operands = dense<1> : tensor<1xi64>} : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
  ...
  %3 = stablehlo.dynamic_broadcast_in_dim %cst, %2, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
  %4 = stablehlo.add %0, %3 : tensor<?xf32>
  return %4 : tensor<?xf32>
}

// After stablehlo-refine-shapes: Shapes propagated, dynamic ops still exist
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
  %c = stablehlo.constant dense<16> : tensor<1xi32>
  %0 = stablehlo.dynamic_broadcast_in_dim %cst, %c, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<16xf32>
  %1 = stablehlo.add %arg0, %0 : tensor<16xf32>
  return %1 : tensor<16xf32>
}

// After stablehlo-canonicalize-dynamism: Dynamic ops replaced with static ops
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
  %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<16xf32>
  %1 = stablehlo.add %arg0, %0 : tensor<16xf32>
  return %1 : tensor<16xf32>
}

// (Bonus) Use ` --stablehlo-aggressive-simplification` pass to canonicalize the
// constant broadcast, leaving us with the original static program in this case.
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<16xf32>
  %0 = stablehlo.add %arg0, %cst : tensor<16xf32>
  return %0 : tensor<16xf32>
}