الديناميكية في 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.

هناك نتيجتان أساسيتان لتعدّد الأشكال:

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

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

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

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

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

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

تصدير البرامج التي تتضمّن سمات ديناميكية

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

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

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

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

void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
                                           TypeRange refinedTypes);

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

بشكلٍ فردي، تكون عمليات التمرير التي تميل إلى أن تكون مفيدة لتحسين الشكل هي:

stablehlo-refine-arguments لاستبدال وسيطات الإدخال بأنواع موترات محددة.
stablehlo-refine-shapes لنشر معلومات شكل وسيط الإدخال الجديد في البرنامج بأكمله.
stablehlo-canonicalize-dynamism لاستبدال العمليات الديناميكية ببدائلها الثابتة.
stablehlo-check-shape-assertions للتحقّق من عمليات الاستدعاء المخصّصة الخاصة بتأكيدات الشكل وإزالتها.

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

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

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

نموذج 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 مع رمز IR صالح بشكل ديناميكي، ما سيؤدي إلى بث الثابت لمطابقة شكل الإدخال الديناميكي على النحو التالي:

// 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>
}

الديناميكية في StableHLO تنظيم صفحاتك في مجموعات يمكنك حفظ المحتوى وتصنيفه حسب إعداداتك المفضّلة.