لهجة 'sdy'

تحدّد لهجة Shardy (SDY) تمثيلًا لتقسيم مصفوفة كثيفة يستند إلى المحاور ومكونات إضافية لواجهة برمجة التطبيقات لإرفاق عمليات التقسيم بالمصفوفات الكثيفة.

العمليات

sdy.all_gather (sdy::AllGatherOp)

تنفيذ عملية جمع شاملة على طول المحاور

البنية:

operation ::= `sdy.all_gather` $gathering_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

تجمع هذه الدالة أجزاء من مصفوفة تانسور على طول المحاور المحدّدة في gathering_axes.

gathering_axes هي قائمة بقوائم المحاور. القائمة الخارجية أكبر من أبعاد مصفوفة السلاسل الزمنية. تحدّد كل قائمة داخلية المحاور التي يجب إجراء جمعٍ منفصل على البُعد المعنيّ وفقًا لها. وسيتم تطبيقها على تقسيم الم Operand (tensor) للحصول على تقسيم النتيجة (out_sharding).

يُرجى العلم أنّه لا يتم استخدام out_sharding لتحديد تقسيم النتيجة. بدلاً من ذلك، يتم تحديد تقسيم النتيجة من خلال تقسيم المُشغِّل وgathering_axes، ويجب أن يتطابق out_sharding مع هذا التقسيم المستنتج.

مثال:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b", "c"}, {}, {"d"}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.all_gather [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a"}, {}, {}\]> : tensor<8x8x8xf32>

القيود:

  • يجب أن تستوفي القيود الواردة في Sdy_CollectiveOpInterface.
  • يجب أن تستوفي العناصر في gathering_axes القيود الواردة في AxisRefListAttr.
  • يؤدي تطبيق gathering_axes على تقسيم الم Operand إلى الحصول على out_sharding.

السمات: SameOperandsAndResultType

الواجهات: InferTypeOpInterface وSdy_CollectiveOpInterface

السمات:

السمةنوع MLIRالوصف
gathering_axes::mlir::sdy::ListOfAxisRefListsAttrقائمة قوائم مراجع المحاور
out_sharding::mlir::sdy::TensorShardingAttrتقسيم المتّجه

المَعلمات:

المُعامل الوصف
tensor مصفوفة من قيم أي نوع

النتائج:

النتيجة الوصف
result مصفوفة من قيم أي نوع

sdy.all_reduce (sdy::AllReduceOp)

تنفيذ عملية تواصل شاملة على طول المحاور

البنية:

operation ::= `sdy.all_reduce` $reduction_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

تُقلِّل أجزاء من مصفوفة لاتّجاهات محدّدة في reduction_axes. لا يُعدّ ترتيب reduction_axes مهمًا للنتيجة، ولكن يمكنه التأثير في ترتيب مجموعات النُسخ المطابقة.

القيود:

  • يجب أن تستوفي القيود الواردة في Sdy_CollectiveOpInterface.
  • يجب أن تستوفي reduction_axes القيود الواردة في AxisRefListAttr.
  • يجب ألّا تتداخل reduction_axes مع محاور تقسيم الم Operand

السمات: SameOperandsAndResultType

الواجهات: CollectiveOpInterface وInferTypeOpInterface

السمات:

السمةنوع MLIRالوصف
reduction_axes::mlir::sdy::AxisRefListAttrقائمة بمراجع المحاور
out_sharding::mlir::sdy::TensorShardingAttrتقسيم المتّجه

المَعلمات:

المُعامل الوصف
tensor مصفوفة من قيم أي نوع

النتائج:

النتيجة الوصف
result مصفوفة من قيم أي نوع

sdy.all_slice (sdy::AllSliceOp)

تنفيذ عملية شريحة ديناميكية على طول المحاور

البنية:

operation ::= `sdy.all_slice` $slicing_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

تقطّع أجزاء من مصفوفة تانسور على طول المحاور المحدّدة في slicing_axes. هناك تناظر جبري بين sdy.all_slice وsdy.all_gather.

slicing_axes هي قائمة بقوائم المحاور. القائمة الخارجية أكبر من أبعاد مصفوفة السلاسل المتعدّدة. تحدّد كل قائمة داخلية المحاور التي يجب تنفيذ شريحة عليها في السمة المعنية. سيتم تطبيقه على تقسيم الم Operand (tensor) للحصول على تقسيم النتيجة (out_sharding).

يُرجى العلم أنّه لا يتم استخدام out_sharding لتحديد تقسيم النتيجة. بدلاً من ذلك، يتم تحديد تقسيم النتيجة من خلال تقسيم المَعلمة وslicing_axes، ويجب أن تتطابق out_sharding مع هذا التقسيم المستنتج.

مثال:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}, {}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.all_slice [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a", "b", "c"}, {}, {"d"}\]> : tensor<8x8x8xf32>

القيود:

  • يجب أن تستوفي العناصر في slicing_axes القيود الواردة في AxisRefListAttr.
  • يجب أن تستوفي القيود الواردة في Sdy_CollectiveOpInterface.
  • يؤدي تطبيق slicing_axes على تقسيم الم Operand إلى الحصول على out_sharding.

السمات: SameOperandsAndResultType

الواجهات: CollectiveOpInterface وInferTypeOpInterface

السمات:

السمةنوع MLIRالوصف
slicing_axes::mlir::sdy::ListOfAxisRefListsAttrقائمة قوائم مراجع المحاور
out_sharding::mlir::sdy::TensorShardingAttrتقسيم المتّجه

المَعلمات:

المُعامل الوصف
tensor مصفوفة من قيم أي نوع

النتائج:

النتيجة الوصف
result مصفوفة من قيم أي نوع

sdy.all_to_all (sdy::AllToAllOp)

تنفيذ عملية تواصل بين جميع العناصر على طول المحاور

البنية:

operation ::= `sdy.all_to_all` $params $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

لكلّ مجموعة (الأجسام، سمة_المصدر، سمة_الهدف) في قائمة المَعلمات، تقطع هذه العملية أجزاءً من مصفوفة كثيفة على طول السمة tgt_dim والأجسام المحدّدة في axes، وتوزّع هذه الأجزاء على طول المحاور، وتسلسلها على طول السمة src_dim.

هذه العملية هي في الأساس عبارة عن تركيبة من جمع الكل على طول src_dim وaxes، متبوعة بقسمة الكل على طول tgt_dim وaxes، أي تتم إضافة لاحقة لسمة تقسيم المحاور src_dim على مصفوفة الإدخال إلى سمة تقسيم المحاور tgt_dim على مصفوفة الإخراج.

سيتم تطبيق الإجراء "الكلّ مع الكلّ" على تقسيم الم Operand (tensor) للحصول على تقسيم النتيجة (out_sharding).

يُرجى العلم أنّه لا يتم استخدام out_sharding لتحديد تقسيم النتيجة. بدلاً من ذلك، يتم تحديد تقسيم النتيجة حسب تقسيم المَعلمة وsrc_dim وtgt_dim وaxes، ويجب أن يتطابقout_sharding مع هذا التقسيم المستنتَج.

مثال:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b"}, {"c"}, {}, {}\]>]>} : tensor<8x8x4x4x32>
%2 = sdy.all_to_all [{"b"}: 0->2, {"c"}: 1->3] %1 out_sharding=<@mesh, [{"a"}, {}, {"b"}, {"c"}\]> : tensor<8x8x4x4x32>

القيود:

  • يجب أن تستوفي القيود الواردة في Sdy_CollectiveOpInterface.
  • يجب ألا تكون قائمة المَعلمات فارغة.
  • لكل مَعلمة في params:
    • يجب أن تستوفي العناصر في axes قيود AxisRefAttr.
    • يجب أن يكون src_dim وtgt_dim سمتَين صالحتَين (غير سالبتَين وأقل من ترتيب المتجه).
    • يجب أن تكون أيّ src_dim أو tgt_dim فريدة في جميع المَعلمات.
    • يجب ترتيب src_dim تصاعديًا في جميع المَعلمات.
  • يؤدي نقل axes من src_dim إلى tgt_dim في تقسيم الم Operand إلى الحصول على out_sharding.

السمات: SameOperandsAndResultType

الواجهات: InferTypeOpInterface وSdy_CollectiveOpInterface

السمات:

السمةنوع MLIRالوصف
params::mlir::sdy::AlltoAllParamListAttrقائمة بالمَعلمات التي تُطبَّق على جميع العناصر
out_sharding::mlir::sdy::TensorShardingAttrتقسيم المتّجه

المَعلمات:

المُعامل الوصف
tensor مصفوفة من قيم أي نوع

النتائج:

النتيجة الوصف
result مصفوفة من قيم أي نوع

sdy.collective_permute (sdy::CollectivePermuteOp)

تنفيذ عملية تواصل تبديل جماعي لاستبدال المحاور

البنية:

operation ::= `sdy.collective_permute` $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

تُرسِل هذه الوظيفة جزءًا من مصفوفة الإدخال من كل جهاز إلى آخر لإعادة ترتيب/استبدال المحاور التي تقسم المصفوفة.

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

ويُعدّ ذلك مفيدًا لإعادة ترتيب المحاور في سمة واحدة أو على مستوى سمات مختلفة، وتبديل المحاور المجزّأة بالمحاور المكرّرة.

في المثال أدناه، حجم مصفوفة Tensor المجزّأة هو tensor<1x4x2xf32>، ويظلّ محفوظًا من خلال الترتيب المجمّع.

مثال:

sdy.mesh @mesh = <["a"=2, "b"=2, "c"=4, "d"=2, "e"=2, "f"=2]>
%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "c"}, {"f"}, {"d", "e"}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.collective_permute %1 out_sharding=<@mesh, [{"c":(1)2, "b", "f"}, {"a"}, {"e", "d"}\]> : tensor<8x8x8xf32>

القيود:

  • يجب أن تستوفي القيود الواردة في Sdy_CollectiveOpInterface.
  • إذا كانت شبكة تقسيم الإدخال والإخراج مختلفة، يجب أن تملك تلك الشبكات المحاور نفسها بالضبط وترتيبًا مختلفًا لمعرّفات الأجهزة.
  • لكل سمة، يجب أن يتطابق حاصل ضرب أحجام محور التجزئة في out_sharding مع حاصل ضرب تجزئة سمة الم Operand المقابلة.

السمات: SameOperandsAndResultType

الواجهات: CollectiveOpInterface وInferTypeOpInterface

السمات:

السمةنوع MLIRالوصف
out_sharding::mlir::sdy::TensorShardingAttrتقسيم المتّجه

المَعلمات:

المُعامل الوصف
tensor مصفوفة من قيم أي نوع

النتائج:

النتيجة الوصف
result مصفوفة من قيم أي نوع

sdy.constant (sdy::ConstantOp)

العملية الثابتة

تُنشئ مصفوفة output من عدد ثابت value.

راجِع: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#constant

مثال:

%output = sdy.constant dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>

السمات: AlwaysSpeculatableImplTrait

الواجهات: ConditionallySpeculatable وInferTypeOpInterface وNoMemoryEffect (MemoryEffectOpInterface)

التأثيرات: MemoryEffects::Effect{}

السمات:

السمةنوع MLIRالوصف
value::mlir::ElementsAttrسمة متّجه/متجهّد ثابت

النتائج:

النتيجة الوصف
output مصفوفة ذات شكل ثابت من أي نوع من القيم

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

عملية معالجة في حافة تدفق البيانات

البنية:

operation ::= `sdy.data_flow_edge` $input (`sharding````=``` $sharding^)? attr-dict `:` type($result)

تحدِّد حافة تدفق البيانات لبعض العمليات X جسرًا بين مجموعة من المصادر (كلّ منها إما عامل تشغيل في X أو عامل تشغيل في عنصر إنهاء السلسلة في X) ومجموعة من الاستهدافات (كلّ منها إما نتيجة X أو مَعلمة سلسلة في X)، بحيث يجب تقسيم جميع المصادر والاستهدافات بالطريقة نفسها.

يمكن أن تحتوي العملية على حواف متعددة لتدفّق البيانات تكون متعامدة مع بعضها.

على سبيل المثال:

  y_0, ..., y_n = while (x_0, ..., x_n)
                  ((pred_arg_0,... , pred_arg_n) { ... })
                  ((body_arg_0,..., body_arg_n) {
                    ...
                    return return_value_0, ..., return_value_n
                  })

يحتوي هذا الإجراء الذي يستغرق وقتًا طويلاً على n من حواف تدفق البيانات، وتكون حواف تدفق البيانات i بين المصادر x_i وreturn_value_i والاستهدافات y_i وpred_arg_i body_arg_i.

يأخذ العنصر sdy.data_flow_edge مالك الحافة كمدخل (يمكن أن يكون أيًا من الاستهدافات، ولكن يُفضّل أن يكون نتيجة عملية بدلاً من دالّة الحظر)، والذي يجب ألا يكون له أي استخدامات أخرى. هذه العملية ليست عملية نقية لأنّه يمكنها استخدام مدخلات لم تكن لها أي استخدامات في الأصل.

يحتوي sdy.data_flow_edge أيضًا على تقسيم اختياري لجميع استهدافات الحافة، ويجب تعديل هذا التقسيم بدلاً من تقسيم الاستهدافات (إذا كان يمكن إرفاقه) أثناء النشر. يكون ذلك مفيدًا عندما يحتوي أحد الإجراءات على العديد من الحواف، لأنّه من الأفضل إجراء ما يلي:

  • يتم نشرها من خلال كل حافة على حدة.
  • تعديل تقسيم كلّ حافة بشكل منفصل بدلاً من جميع الاستهدافات في آنٍ واحد (على سبيل المثال، تتضمّن عملية TensorShardingPerValueAttr واحدًا غير قابل للتغيير لتقسيم نتيجة ).
  • إضافة كلّ حافة إلى قائمة العمل بشكلٍ منفصل عند اختلاف تجزئة مصدر

ستؤدي عملية النشر إلى نشر الأجزاء بين جميع مصادر وأهداف sdy.data_flow_edge كما لو كانت عملية عادية بالمصادر كمُعاملين والأهداف كنتائج، وهوية sdy.op_sharding_rule. وهذا يعني أنّه تتم المعالجة من المصادر إلى الأهداف، ومن الأهداف إلى المصادر.

لا نسمح بتحديد إدخال sdy.data_flow_edge من خلال عملية SdyDialect، لذا يمكننا افتراض أنّه تم تحديده من خلال عملية تحتوي على سمة sdy.sharding غير مسجَّلة.

السمات: SameOperandsAndResultType

الواجهات: InferTypeOpInterface

السمات:

السمةنوع MLIRالوصف
sharding::mlir::sdy::TensorShardingAttrتقسيم المتّجه

المَعلمات:

المُعامل الوصف
input تكون من أي نوع من القيم

النتائج:

النتيجة الوصف
result تكون من أي نوع من القيم

sdy.manual_computation (sdy::ManualComputationOp)

عملية التوازي على أجهزة متعددة باستخدام المجموعات اليدوية

البنية:

operation ::= `sdy.manual_computation` `(`operands`)`
              `in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)
              `out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)
              `manual_axes````=```$manual_axes
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:`
              functional-type(operands, results)

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

الجسم محلي بالنسبة إلى المحاور اليدوية. سيتمّ النشر من خلال الجسم على أيّ محاور مجانية، أيّ المحاور غير المُدرَجة في قائمة manual_axes.

القيود:

  • يجب أن تستوفي العناصر في in_shardings وout_shardings القيود الواردة في TensorShardingAttr.
  • يجب أن يتطابق عدد مدخلات/مخرجات مصفوفة تينسور العالمية والمحلية لمنطقة الحساب.
  • يجب أن تأتي المحاور اليدوية قبل أي محاور مجانية في كل عملية تقسيم سمة.
  • لا يمكن أن تضيف المحاور اليدوية مساحة فارغة. على وجه التحديد، يجب أن يكون حجم السمة قابلاً للقسمة على حجم المحاور اليدوية المقابل.
  • يجب أن تتطابق الأشكال العالمية والمحلية لوسيطات/نتائج مناطق العمليات.
  • لا يتم تقسيم أيّ محاور يدوية.

السمات: IsolatedFromAbove وRecursiveMemoryEffects وSingleBlockImplicitTerminator<ReturnOp> وSingleBlock

الواجهات: ShardableDataFlowOpInterface

السمات:

السمةنوع MLIRالوصف
in_shardings::mlir::sdy::TensorShardingPerValueAttrتقسيم المتّجه لكلّ عامل تشغيل/نتيجة لعملية تشغيل
out_shardings::mlir::sdy::TensorShardingPerValueAttrتقسيم المتّجه لكلّ عامل تشغيل/نتيجة لعملية تشغيل
manual_axes::mlir::sdy::ManualAxesAttrقائمة بالأعمدة التي يكون فيها ManualComputationOp يدويًا

المَعلمات:

المُعامل الوصف
tensors متغيّر عددي من مصفوفة مرتبة من قيم أي نوع

النتائج:

النتيجة الوصف
results متغيّر عددي من مصفوفة مرتبة من قيم أي نوع

sdy.mesh (sdy::MeshOp)

الشبكة المُسمّاة

البنية:

operation ::= `sdy.mesh` $sym_name `=` $mesh attr-dict

تحدِّد شبكة جديدة مُسمّاة. يجب أن تتضمّن جميع الشبكات في الوحدة العدد نفسه للأجهزة (باستثناء الشبكات التي تحتوي على معرّف جهاز واحد). الشبكة هي عملية Symbol تظهر في SymbolTable الوحدة ويمكن الإشارة إليها من خلال name.

السمات: HasParent<ModuleOp>

الواجهات: Symbol

السمات:

السمةنوع MLIRالوصف
sym_name::mlir::StringAttrسمة سلسلة نصية
mesh::mlir::sdy::MeshAttrشبكة المحاور وقائمة بالأجهزة

sdy.named_computation (sdy::NamedComputationOp)

عملية حسابية مُسمّاة

البنية:

operation ::= `sdy.named_computation` `<`$name`>` `` `(` $operands `)`
              (`in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)^)?
              (`out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)^)?
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:` functional-type($operands, results)

تُجمِّع عملية حسابية، أي مجموعة من العمليات، وتمنحها اسمًا. ستتم عملية النشر داخل المنطقة أو خارجها كما لو كان كل المحتوى مضمّنًا.

ويمكن استخدام هذا الإجراء للتعامل مع التوسّع من خلال تعليمات الاتصال إلى الدوالّ الأخرى. على أي مستخدمين من Shardy كتابة عملية استيراد/تصدير تؤدي إلى تحويل عمليات الاتصال إلى عمليات sdy.named_computation، مع تكرار/نسخ نص الدالة التي يتم استدعاؤها في نص named_computation.

يجب أن يكون نوع وسيطات كلّ عنصر وقيمها المعروضة في المنطقة هو النوع نفسه لنوع المَعلمات ونوع النتائج لعملية التشغيل.

مثال:

%1 = sdy.named_computation<"foo">(%0) (%arg1: tensor<16x32xf32>) {
  sdy.return %arg1 : tensor<16x32xf32>
} : (tensor<16x32xf32>) -> tensor<16x32xf32>

السمات: IsolatedFromAbove وRecursiveMemoryEffects وRecursivelySpeculatableImplTrait وSingleBlockImplicitTerminator<ReturnOp> وSingleBlock

الواجهات: ConditionallySpeculatable وInferTypeOpInterface وShardableDataFlowOpInterface

السمات:

السمةنوع MLIRالوصف
name::mlir::StringAttrسمة سلسلة نصية
in_shardings::mlir::sdy::TensorShardingPerValueAttrتقسيم المتّجه لكلّ عامل تشغيل/نتيجة لعملية تشغيل
out_shardings::mlir::sdy::TensorShardingPerValueAttrتقسيم المتّجه لكلّ عامل تشغيل/نتيجة لعملية تشغيل

المَعلمات:

المُعامل الوصف
operands متغيّر عددي من أي نوع

النتائج:

النتيجة الوصف
«unnamed» متغيّر عددي من أي نوع

sdy.propagation_barrier (sdy::PropagationBarrierOp)

عملية حاجز الانتشار

البنية:

operation ::= `sdy.propagation_barrier` $input `allowed_direction````=```$allowed_direction attr-dict `:` type($input)

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

ويمنع ذلك من نشر الأجزاء بين استخدامات نتيجة عملية الحاجز وعمله.

  • FORWARD يعني أنّه لا يمكن تقسيم البيانات إلا من الم Operand إلى النتيجة.
  • BACKWARD يعني أنّه لا يمكن تقسيم البيانات إلا من النتيجة إلى الم Operand.
  • تعني NONE أنّه لا يمكن نشر أي تقسيم من خلال هذه العملية.
  • لا يمكن تحديد BOTH، لأنّ هذا الإجراء سيكون زائداً.

السمات: AlwaysSpeculatableImplTrait، SameOperandsAndResultType

الواجهات: ConditionallySpeculatable وInferTypeOpInterface وNoMemoryEffect (MemoryEffectOpInterface)

التأثيرات: MemoryEffects::Effect{}

السمات:

السمةنوع MLIRالوصف
allowed_direction::mlir::sdy::PropagationDirectionAttrتعداد اتجاه الانتشار

المَعلمات:

المُعامل الوصف
input مصفوفة مرتبة لأي نوع من القيم

النتائج:

النتيجة الوصف
result مصفوفة مرتبة لأي نوع من القيم

sdy.reshard (sdy::ReshardOp)

إعادة تقسيم مصفوفة كثيفة إلى تقسيم مختلف

البنية:

operation ::= `sdy.reshard` $input $sharding attr-dict `:` type($result)

تعيد تقسيم مصفوفة الإدخال باستخدام التقسيم المحدّد، والذي يختلف عن التقسيم الحالي لمصفوفة الإدخال.

يربط كلّ من ShardingConstraintOp وReshardOp عملية تقسيم إلى مصفوفة. مدة استخدامها هي:

  1. قبل نشر التجزئة، يضيف المستخدمون ShardingConstraintOp.
  2. تستهلك عملية نشر التقسيم إلى أجزاء ShardingConstraintOp. لا يتوفّر سوى ShardingConstraintOp في نتائج نشر التجزئة. بدلاً من ذلك، قد تتم إضافة ReshardOp إذا لزم الأمر.
  3. يحوّل أداة التقسيم عملية ReshardOp إلى عملية جماعية (أو عملية الهوية). يجب ألا يكون هناك ReshardOp في نتائج أداة التقسيم.

// TODO(b/331680067). أضِف نمطًا لتوحيد البيانات لإزالة عمليات // reshard المتكرّرة.

السمات: AlwaysSpeculatableImplTrait، SameOperandsAndResultType

الواجهات: ConditionallySpeculatable وInferTypeOpInterface وNoMemoryEffect (MemoryEffectOpInterface)

التأثيرات: MemoryEffects::Effect{}

السمات:

السمةنوع MLIRالوصف
sharding::mlir::sdy::TensorShardingAttrتقسيم المتّجه

المَعلمات:

المُعامل الوصف
input مصفوفة من قيم أي نوع

النتائج:

النتيجة الوصف
result مصفوفة من قيم أي نوع

sdy.return (sdy::ReturnOp)

تُنهي عملية sdy.return المناطق المرتبطة بعملياتsdy المستندة إلى المنطقة وأي عمليات أخرى مستندة إلى منطقة Shardy. وهو متغيّر: يأخذ كوسيطات قائمة بالقيم التي يمكن أن تكون أنواعها أيًا (ولكن من النوع نفسه، مثل AnyTensor)، وبالتالي يمكن إعادة استخدامه على مستويات مختلفة من حِزمة Shardy IR.

البنية:

operation ::= `sdy.return` attr-dict ($results^ `:` type($results))?

السمات: AlwaysSpeculatableImplTrait، Terminator

الواجهات: ConditionallySpeculatable وNoMemoryEffect (MemoryEffectOpInterface)

التأثيرات: MemoryEffects::Effect{}

المَعلمات:

المُعامل الوصف
results متغيّر عددي من أي نوع

sdy.sharding_constraint (sdy::ShardingConstraintOp)

تقييد مصفوفة تينسور بالشريحة المحدّدة

البنية:

operation ::= `sdy.sharding_constraint` $input $sharding attr-dict `:` type($result)

تُرفِق هذه الدالة عملية تقسيم إلى مصفوفة وسيطة (مثل نتيجة matmul) للإشارة إلى أنّ هذه هي الطريقة التي يجب بها تقسيم هذه المصفوفة أو مجموعة فرعية من استخداماتها.

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

يمكن أن يؤدي هذا الإجراء إلى:

  • لا يكون لها أي استخدامات (غير مرتبطة) - ما يعني أنّ التقسيم المرفق هو الطريقة التي يجب بها تقسيم ملف tensor للدخل نفسه.
  • أن يكون لها استخدامات: ويعني ذلك أنّ التقسيم المرفق هو الطريقة التي يجب تقسيم استخدامات عملية قيد التقسيم بها، في حين أنّ استخدامات ناقلات المصفوفات المدخلة الأخرى قد يكون لها تقسيم مختلف (إذا لم تكن ناقلات المصفوفات المدخلة لها استخدامات أخرى، يكون السلوك هو نفسه في حالة عدم توفّر أي استخدامات).

السمات: SameOperandsAndResultType

الواجهات: InferTypeOpInterface

السمات:

السمةنوع MLIRالوصف
sharding::mlir::sdy::TensorShardingAttrتقسيم المتّجه

المَعلمات:

المُعامل الوصف
input مصفوفة من قيم أي نوع

النتائج:

النتيجة الوصف
result مصفوفة من قيم أي نوع

sdy.sharding_group (sdy::ShardingGroupOp)

تقييد مصفوفات السلاسل في المجموعة لاستخدام التقسيم نفسه

البنية:

operation ::= `sdy.sharding_group` $input `group_id````=```$group_id attr-dict `:` type($input)

يوفّر هذا الإجراء واجهة لتحديد مصفوفات تانسور إلى مجموعات تقسيم ( مجموعات من مصفوفات تانسور التي سيتم فرض تقسيمها بشكلٍ متطابق). أثناء النشر، بعد تقسيم عنصر مجموعة واحد، سيتم تقسيم كل العناصر الأخرى بالطريقة نفسها تمامًا. تأخذ هذه العملية معرّف مجموعة الوسيطات ولا تعرِض أي نتيجة، ولكنّها تعدّل بدلاً من ذلك التمثيل الداخلي لمجموعة التجزئة لإضافة مصفوفة الإدخال إلى المجموعة التي تحمل المعرّف المحدَّد.

الواجهات: InferTypeOpInterface

السمات:

السمةنوع MLIRالوصف
group_id::mlir::IntegerAttrسمة عدد صحيح بدون علامة بسعة 64 بت

المَعلمات:

المُعامل الوصف
input مصفوفة مرتبة لأي نوع من القيم

السمات

AllToAllParamAttr

مَعلمة الكلّ للكلّ

البنية:

#sdy.all_to_all_param<
  ::llvm::ArrayRef<AxisRefAttr>,   # axes
  int64_t,   # src_dim
  int64_t   # tgt_dim
>

مجموعة تحتوي على المحاور وسمات المصدر/الهدف لإجراء عملية مقارنة بين كل العناصر

المَعلمات:

المَعلمة نوع C++ الوصف
المحاور ::llvm::ArrayRef<AxisRefAttr> المحاور التي يتم إجراء عمليات الربط التام بينها
src_dim int64_t فهرس سمة المصدر
tgt_dim int64_t فهرس السمة المستهدَفة

AlltoAllParamListAttr

قائمة بالمَعلمات التي تُطبَّق على جميع العناصر

البنية:

#sdy.all_to_all_param_list<
  ::llvm::ArrayRef<AllToAllParamAttr>   # value
>

المَعلمات:

المَعلمة نوع C++ الوصف
القيمة ::llvm::ArrayRef<AllToAllParamAttr>

AxisRefAttr

إشارة إلى محور كامل أو محور فرعي مُقسَّم

البنية:

#sdy.axis_ref<
  ::llvm::StringRef,   # name
  SubAxisInfoAttr   # sub_axis_info
>

القيود:

  • يجب أن تكون name متوفّرة في الحدّ MeshAttr.
  • إذا كان sub_axis_info متوفّرًا، يجب أن يستوفي القيود المفروضة على SubAxisInfoAttr.

المَعلمات:

المَعلمة نوع C++ الوصف
الاسم ::llvm::StringRef اسم هذا المحور
sub_axis_info SubAxisInfoAttr معلومات إضافية إذا كان هذا محورًا فرعيًا

AxisRefListAttr

قائمة إحالات المحاور

البنية:

#sdy.axis_ref_list<
  ::llvm::ArrayRef<AxisRefAttr>   # value
>

القيود:

  • يجب أن تستوفي العناصر في value قيود AxisRefAttr.
  • لا تتوفّر إشارات محور أو محاور فرعية مكرّرة تتداخل مع بعضها.
  • لا يكون أيّ مرجعَي محور مجاورَين محورَين فرعيَّين متتاليين للمحور الكامل نفسه، أي أنّه يمكن دمجهما في محور فرعي واحد أو في المحور الكامل.

المَعلمات:

المَعلمة نوع C++ الوصف
القيمة ::llvm::ArrayRef<AxisRefAttr>

DimMappingAttr

قائمة مؤشرات العوامل لسمة معيّنة

تشير القائمة الفارغة إلى أنّ هذا ربط فارغ (يتم تحليله/طباعته باستخدام *)، أي أنّه لم يتم ربط السمة بأي عوامل.

القيود:

  • أن يكون هناك فهرس عامل واحد على الأقل
  • يجب أن تكون فهارس العوامل ضمن النطاق [0، $factor_sizes).
  • إذا كانت هناك عوامل متعدّدة، لا يمكن أن يكون لأيّ منها حجم 1.
  • عدم توفّر مؤشرات عوامل مكرّرة

المَعلمات:

المَعلمة نوع C++ الوصف
factor_indices ::llvm::ArrayRef<int64_t> العوامل التي تمّ ربط هذه السمة بها

DimensionShardingAttr

تقسيم السمات

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

القيود:

  • يجب أن تستوفي العناصر في axes القيود الواردة في AxisRefListAttr.
  • إذا كان تجزئة السمة لها أولوية:
    • تكون الأولوية أكبر من أو تساوي 0.
    • تحتوي السمة على محور واحد على الأقل إذا كانت مغلقة.

المَعلمات:

المَعلمة نوع C++ الوصف
المحاور ::llvm::ArrayRef<AxisRefAttr> إشارات المحاور
is_closed bool ما إذا كان لا يمكن تقسيم هذه السمة إلى أقسام أخرى
الحملة std::optional<int64_t> الأولوية المستخدَمة أثناء النشر المستنِد إلى أولوية المستخدم

ListOfAxisRefListsAttr

قائمة قوائم مراجع المحاور

البنية:

#sdy.list_of_axis_ref_lists<
  ::llvm::ArrayRef<AxisRefListAttr>   # value
>

المَعلمات:

المَعلمة نوع C++ الوصف
القيمة ::llvm::ArrayRef<AxisRefListAttr>

ManualAxesAttr

قائمة بالأعمدة التي يكون فيها ManualComputationOp يدويًا

البنية:

#sdy.manual_axes<
  ::llvm::ArrayRef<StringAttr>   # value
>

المَعلمات:

المَعلمة نوع C++ الوصف
القيمة ::llvm::ArrayRef<StringAttr>

MeshAttr

شبكة المحاور وقائمة الأجهزة

البنية:

#sdy.mesh<
  ::llvm::ArrayRef<MeshAxisAttr>,   # axes
  ::llvm::ArrayRef<int64_t>   # device_ids
>

الشبكة هي قائمة بالأعمدة وقائمة اختيارية بمعرّفات الأجهزة التي تحدّد ترتيب الأجهزة.

إذا كانت قائمة المحاور فارغة، تحتوي الشبكة على محور غير مُسمّى ضمنيًا بحجم 1. في هذه الحالة، إذا لم يتم تقديم قائمة أرقام تعريف الأجهزة، تكون قائمة أرقام تعريف الأجهزة الضمنية هي [0]، وإذا تم تقديم قائمة أرقام تعريف الأجهزة، يجب أن تحتوي على عدد صحيح واحد من أي قيمة غير سالبة. ونُطلق على ذلك اسم حالة التجزئة القصوى.

في جميع حالات عدم استخدام التجزئة القصوى، إذا تم تحديد قائمة بأرقام تعريف الأجهزة، يجب أن يتطابق حاصل ضرب أحجام المحاور مع عدد الأجهزة. في حال عدم تحديد قائمة معرّفات الأجهزة، تكون قائمة معرّفات الأجهزة الضمنية هي iota(product(axes)). من أجل البساطة، لا نسمح أيضًا بتحديد قائمة أرقام تعريف الأجهزة التي هي نفسها iota(product(axes));، وفي هذه الحالة، يجب عدم تحديد قائمة أرقام تعريف الأجهزة.

في ما يلي بعض الأمثلة على الشبكات:

  • تمثّل الشبكة الفارغة شبكة عنصر نائب يمكن استبدالها أثناء الانتشار: <[]>
  • شبكة تتضمّن محورًا بدون اسم ورقم تعريف جهاز صريح، يُستخدَم عادةً لتمثيل الحد الأقصى لتقسيم البيانات: <[], device_ids=[3]>
  • شبكة تتضمّن محورَين ومعرّفات أجهزة ضمنية iota(6):‏ <["a"=2, "b"=3]>
  • شبكة تتضمّن محورَين ومعرّفات أجهزة صريحة تحدّد ترتيب الأجهزة: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

القيود:

  • يجب ألّا تحتوي العناصر في axes على أسماء مكرّرة.
  • في حال تحديد device_ids:
    • يجب أن يتطابق حاصل ضرب أحجام المحاور مع عدد الأجهزة.
    • ويجب أن تكون جميع عناصرها غير سالبة.
    • يجب ألا يساوي device_ids iota(product(axis_sizes)).
    • يجب أن يكون الحقل device_ids الذي تم ترتيبه iota(product(axis_sizes)).

المَعلمات:

المَعلمة نوع C++ الوصف
المحاور ::llvm::ArrayRef<MeshAxisAttr> محاور الشبكة
device_ids ::llvm::ArrayRef<int64_t> ترتيب الأجهزة الصريح أو الحد الأقصى لمعرّف الجهاز

MeshAxisAttr

محور مُعنوَن في شبكة

البنية:

#sdy.mesh_axis<
  ::llvm::StringRef,   # name
  int64_t   # size
>

المَعلمات:

المَعلمة نوع C++ الوصف
الاسم ::llvm::StringRef الاسم
الحجم int64_t حجم هذا المحور

OpShardingRuleAttr

يحدِّد كيفية تقسيم عملية معيّنة.

البنية:

#sdy.op_sharding_rule<
  ::llvm::ArrayRef<int64_t>,   # factor_sizes
  ::llvm::ArrayRef<TensorMappingAttr>,   # operand_mappings
  ::llvm::ArrayRef<TensorMappingAttr>,   # result_mappings
  ::llvm::ArrayRef<int64_t>,   # reduction_factors
  ::llvm::ArrayRef<int64_t>,   # need_replication_factors
  ::llvm::ArrayRef<int64_t>,   # permutation_factors
  ::llvm::ArrayRef<int64_t>,   # blocked_propagation_factors
  bool   # is_custom_rule
>

تحدّد قاعدة التجزئة كيفية تقسيم عملية وفقًا للخصائص المختلفة في العملية، مثل أيّ سمات أو شكل المُعامِلات أو شكل النتائج، وما إلى ذلك. على سبيل المثال:

%0 = stablehlo.add %arg0, %arg1 {
    sdy.sharding_rule = #sdy.op_sharding_rule<
        ([i, j],[i, j])->([i, j])
        {i=8, j=8}>
} : tensor<8x8xf32>
%1 = stablehlo.dot_general %arg2, %arg3, contracting_dims = [1] x [0] {
  sdy.sharding_rule = #sdy.op_sharding_rule<
      ([i, k],[k, j])->([i, j])
      {i=8, j=16, k=8}>
}: (tensor<8x8xf32>, tensor<8x16xf32>) -> tensor<8x16xf32>

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

أنواع العوامل:

  • يحتوي reduction_factors على مؤشرات العوامل التي تتطلّب التخفيض، مثل الأبعاد المتعاقدة في عملية النقطة.
  • يحتوي العمود need_replication_factors على مؤشرات العوامل التي تتطلّب تكرارًا كاملاً، مثل السمة التي تم ترتيبها في عملية ترتيب.
  • يحتوي permutation_factors على فهارس العوامل التي تتطلّب تبديلاً جماعيًا إذا كانت مجزّأة، مثل سمات الحشو في عملية الحشو.
  • تُعتبر جميع العوامل الأخرى عوامل تمرير، أي عوامل لا تتطلّب أي تواصل إذا تم تقسيمها بالطريقة نفسها على جميع المتسلسلات التي تم ربطها بها.

يحتوي blocked_propagation_factors على العوامل التي لا يُسمح بنشر الأقسام على أساسها. وهي متعامدة مع أنواع العوامل. على وجه التحديد، يمكن أن يكون عامل الانتشار المحظور أيًا من أنواع العوامل.

يصف الحقل is_custom_rule ما إذا كانت هذه قاعدة حدّدها مستخدم. يمكن للمستخدمين تحديد قواعد التجزئة لمكالماتهم المخصّصة أو استبدال قواعد التجزئة المحدّدة مسبقًا للعمليات العادية. يتم الحفاظ على القاعدة المخصّصة أو عدم إزالتها مطلقًا.

القيود:

  • يجب أن يتطابق عدد عمليات الربط بين الم Operand/النتيجة مع عدد الم Operand/النتائج لعملية التشغيل.
  • تتوفر عملية ربط واحدة على الأقل (لا يمكن أن تتضمّن قاعدة عملية لا تحتوي على عوامل تشغيل/نتائج).
  • يتطابق ترتيب كل TensorMappingAttr مع ترتيب نوع المكثّف المقابل.
  • لكل مجموعة من العوامل (reduction_factors، need_replication_factors، permutation_factors):
    • يجب أن تكون العناصر ضمن النطاق [0, $factor_sizes].
    • عدم توفّر مؤشرات عوامل مكرّرة ضمن كل مجموعة وعبر المجموعات

المَعلمات:

المَعلمة نوع C++ الوصف
factor_sizes ::llvm::ArrayRef<int64_t> أحجام جميع العوامل في هذه القاعدة
operand_mappings ::llvm::ArrayRef<TensorMappingAttr> عمليات ربط المعاملات
result_mappings ::llvm::ArrayRef<TensorMappingAttr> عمليات ربط النتائج
reduction_factors ::llvm::ArrayRef<int64_t> العوامل التي تتطلّب خفضًا
need_replication_factors ::llvm::ArrayRef<int64_t> العوامل التي تتطلّب التكرار الكامل
permutation_factors ::llvm::ArrayRef<int64_t> العوامل التي تتطلّب استخدام دالة collective-permute
blocked_propagation_factors ::llvm::ArrayRef<int64_t> العوامل التي لا يتم نشر الأجزاء على أساسها
is_custom_rule bool ما إذا كانت القاعدة مخصّصة لطلب stablehlo.custom_call

SubAxisInfoAttr

معلومات عن كيفية اشتقاق هذا المحور الفرعي من المحور الكامل

البنية:

#sdy.sub_axis_info<
  int64_t,   # pre_size
  int64_t   # size
>

عند تقسيم محور كامل إلى n محور فرعي، تتم إعادة تشكيل المحور إلى [k_1,...,k_n]، ويمكن التعبير عن المحور الفرعي i من خلال حاصل ضرب كل مقاسات المحور على يساره m=prod(k_1,...,k_(i-1)) (المعروفة أيضًا باسم المقاس المُسبَق) والمقاس k_i. وبالتالي، تحتوي سمة معلومات-المحور-الفرعي على هذين الرقمين ويتم تمثيلها على النحو التالي: (m)k للحجم المُسبَق m والحجم k.

القيود:

  • pre-size هو 1 على الأقل.
  • size أكبر من 1.
  • يجب أن يقسم pre-size حجم المحور الكامل، أي أن يقسم كل من pre-size وsize حجم المحور الكامل، ولا يتجاوز المحور الفرعي حجم المحور الكامل.
  • لا يساوي حجم المحور الفرعي حجم المحور الكامل المناظر له، وفي هذه الحالة يجب استخدام المحور الكامل بدلاً من ذلك.

المَعلمات:

المَعلمة نوع C++ الوصف
pre_size int64_t حاصل ضرب أحجام المحاور الفرعية على يمين هذا المحور الفرعي
الحجم int64_t حجم هذا المحور الفرعي

TensorMappingAttr

عمليات ربط العوامل لكل سمة من سمات مصفوفة تانسور

البنية:

#sdy.tensor_mapping<
  ::llvm::ArrayRef<DimMappingAttr>   # dim_mappings
>

القيود:

  • يجب أن تستوفي العناصر في dim_mappings القيود الواردة في DimMappingAttr.
  • عدم توفّر فهارس مسبّبات مكرّرة في جميع السمات

المَعلمات:

المَعلمة نوع C++ الوصف
dim_mappings ::llvm::ArrayRef<DimMappingAttr> عمليات ربط السمات

TensorShardingAttr

تقسيم مصفوفات Tensor

البنية:

#sdy.sharding<
  ::mlir::Attribute,   # mesh_or_ref
  ::llvm::ArrayRef<DimensionShardingAttr>,   # dim_shardings
  ::llvm::ArrayRef<AxisRefAttr>   # replicated_axes
>

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

يمكن تحديد الشبكة المرتبطة بهذا التجزئة إما باستخدام اسم رمز يشير إلى رمز MeshOp ملائم أو MeshAttr مضمّن.

القيود:

  • يجب أن تستوفي العناصر في dim_shardings القيود الواردة في DimensionShardingAttr.
  • يجب أن تستوفي العناصر في replicated_axes القيود الواردة في AxisRefListAttr.
  • إذا لم يكن نوع مصفوفة العمود المقابلة هو ShapedType، يجب أن يكون للتقسيم ترتيب 0 وألا تتضمّن أيّ محاور مكرّرة.
  • يجب أن يكون للمركم ترتيب.
  • يساوي عدد تقسيمات السمات رتبة المصفوفة.
  • لا يتم تقسيم السمات التي يبلغ حجمها 0.
  • يتم ترتيب العناصر في replicated_axes حسب mesh_or_ref (راجِع AxisRefAttr::getMeshComparator).

المَعلمات:

المَعلمة نوع C++ الوصف
mesh_or_ref ::mlir::Attribute سمة الشبكة أو سمة مرجع رمز الشبكة المسطّحة
dim_shardings ::llvm::ArrayRef<DimensionShardingAttr> تقسيم السمات
replicated_axes ::llvm::ArrayRef<AxisRefAttr> إشارات المحاور

TensorShardingPerValueAttr

تقسيم مصفوفة متجاوبة حسب الم Operand/نتيجة عملية الحساب

البنية:

#sdy.sharding_per_value<
  ::llvm::ArrayRef<TensorShardingAttr>   # shardings
>

قائمة بعناصر TensorShardingAttr، واحدة لكل عامل أو نتيجة لعملية

القيود:

  • يجب أن تستوفي العناصر في shardings قيود TensorShardingAttr.

المَعلمات:

المَعلمة نوع C++ الوصف
عمليات التقسيم ::llvm::ArrayRef<TensorShardingAttr> التجزئة حسب القيمة

عمليات التعداد

PropagationDirection

قائمة بفئات اتجاه النشر

الحالات:

الرمز القيمة سلسلة
لم يتم اختيار لون 0 لم يتم اختيار لون
FORWARD 1 FORWARD
للخلف 2 للخلف
BOTH 3 BOTH