ניב 'sdy'

הדיאלקט Shardy‏ (SDY) מגדיר ייצוג של חלוקת טינסורים לפי ציר ורכיבי API נוספים כדי לצרף חלוקות לטינסורים.

תפעול

sdy.all_gather (sdy::AllGatherOp)

ביצוע תקשורת all-gather לאורך צירים

תחביר:

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

איסוף קטעי טינסור לאורך צירים שצוינו ב-gathering_axes.

gathering_axes היא רשימה של רשימות של צירים. הרשימה החיצונית חורגת מהמאפיינים של הטנזור. כל רשימה פנימית מציינת את הצירים שבהם צריך לבצע איסוף נפרד של המאפיין הרלוונטי. הוא יחול על חלוקת המשנה של המשתנה (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 על חלוקת המשתנה לחלקים מקבלת את הערך out_sharding.

מאפיינים: SameOperandsAndResultType

ממשקים: InferTypeOpInterface, ‏ Sdy_CollectiveOpInterface

מאפיינים:

אופרטנדים:

תוצאות:

sdy.all_reduce (sdy::AllReduceOp)

ביצוע תקשורת all-reduce לאורך צירים

תחביר:

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 לא יכול לחפוף לצירים של חלוקת המשנה של המשתנה.

מאפיינים: SameOperandsAndResultType

ממשקים: CollectiveOpInterface, ‏ InferTypeOpInterface

מאפיינים:

אופרטנדים:

תוצאות:

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 היא רשימה של רשימות של צירים. הרשימה החיצונית חורגת מהמאפיינים של הטנזור. כל רשימה פנימית מציינת את הצירים שבהם צריך לבצע חיתוך של המאפיין הרלוונטי. הוא יחול על חלוקת המשנה של המשתנה (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 על חלוקת המשתנה לחלקים מקבלת את הערך out_sharding.

מאפיינים: SameOperandsAndResultType

ממשקים: CollectiveOpInterface, ‏ InferTypeOpInterface

מאפיינים:

אופרטנדים:

תוצאות:

sdy.all_to_all (sdy::AllToAllOp)

ביצוע תקשורת 'הכול לכולם' לאורך צירים

תחביר:

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

לכל קבוצה של ערכי טור (tuple) מסוג (axes, src_dim, tgt_dim) ברשימת הפרמטרים, הפעולה הזו חותכת קטעי טינסור לאורך המאפיין tgt_dim והצירים שצוינו ב-axes, מפזרת את הקטעים האלה לאורך הצירים ומארגנת אותם ברצף לאורך המאפיין src_dim.

הפעולה הזו היא למעשה שילוב של all-gather לאורך src_dim ו-axes, ואחריו all-slice לאורך tgt_dim ו-axes. כלומר, סיומת של המאפיין src_dim של חלוקת הצירים בטנסור הקלט מצורפת למאפיין tgt_dim של חלוקת הצירים בטנסור הפלט.

הפעולה 'הכול אל הכול' תחול על חלוקת המשנה של המשתנה (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 בחלוקת המזהה של הפעולה מקבלת את הערך out_sharding.

מאפיינים: SameOperandsAndResultType

ממשקים: InferTypeOpInterface, ‏ Sdy_CollectiveOpInterface

מאפיינים:

אופרטנדים:

תוצאות:

sdy.collective_permute (sdy::CollectivePermuteOp)

ביצוע תקשורת של תמורה קבוצתית כדי להחליף צירים

תחביר:

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

שליחת מקטע של הטנסור הקלט מכל מכשיר למכשיר אחר כדי לשנות את הסדר של הצירים שמחלקים את הטנסור או להחליף אותם.

טרנספורמציה של חלוקה קולקטיבית יכולה לשנות את חלוקת המשנה של הקלט כך שכל מאפיין חייב להיות מחולק כמו קודם, כלומר הוא חייב להיות מחולק לפי צירים שמכפלת הגדלים שלהם תואמת לזו של הצירים שבהם התאמת הטנזור בעבר.

אפשר להשתמש באפשרות הזו כדי לשנות את הסדר של צירים במאפיין יחיד או במאפיינים שונים, ולהחליף צירים מחולקים לשברים בצירים עם רפליקה.

בדוגמה הבאה, הגודל של הטנסור המפוצל הוא 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 חייבת להיות זהה למכפלה של חלוקת המאפיין התואם של הפונקציה.

מאפיינים: SameOperandsAndResultType

ממשקים: CollectiveOpInterface, ‏ InferTypeOpInterface

מאפיינים:

אופרטנדים:

תוצאות:

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

מאפיינים:

תוצאות:

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

לפעולה הזו של 'while' יש 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

מאפיינים:

אופרטנדים:

תוצאות:

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. ההעברה תתבצע דרך הגוף בכל צירים חופשיים – אלה שלא נמצאים ברשימה manual_axes.

אילוצים:

• הרכיבים ב-in_shardings וב-out_shardings חייבים לעמוד באילוצים שמפורטים ב-TensorShardingAttr.
• מספר הקלט/הפלט של הטנסורים הגלובליים והמקומיים של אזור הפעולה חייב להיות זהה.
• הצירים הידניים חייבים להופיע לפני כל ציר חופשי בכל חלוקה של מאפיין (dim) לפלחים.
• לא ניתן להוסיף מרווחים לצירים ידניים. כלומר, גודל המאפיין חייב להיות ניתן לחלוקה בגודל של הצירים הידניים התואמים.
• הצורות הגלובליות והמקומיות של הארגומנטים/התוצאות של אזורי הפעולה חייבות להיות זהות.
• לא מתבצעת חלוקה של צירים ידניים.

מאפיינים: IsolatedFromAbove, RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

ממשקים: ShardableDataFlowOpInterface

מאפיינים:

אופרטנדים:

תוצאות:

sdy.mesh (sdy::MeshOp)

רשת בעלת שם

תחביר:

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

הגדרת רשת חדשה בעלת שם. לכל רשתות העמודות במודול צריך להיות אותו מספר מכשירים (למעט רשתות עמודות עם device_id יחיד). הרשת היא פעולת Symbol שמופיעה ב-SymbolTable של המודול, וניתן להפנות אליה באמצעות name שלה.

מאפיינים: HasParent<ModuleOp>

ממשקים: Symbol

מאפיינים:

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

מאפיינים:

אופרטנדים:

תוצאות:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

פעולת מחסום להפצה

תחביר:

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

הפעולה הזו פועלת כמו פעולת זהות, ומפיקה את אותו ערך שהיא קיבלה כקלט. אבל מבחינת ההפצה, היא תאפשר לה לעבור דרכו רק בכיוון מסוים.

כך אפשר למנוע את ההפצה של חלוקות בין השימושים בתוצאה של פעולת המחסום לבין המשתנה שלה.

• FORWARD פירושו שחלוקות יכולות לזרום רק מהאופרטנד לתוצאה.
• BACKWARD פירושו שחלוקות יכולות לזרום מהתוצאה לפעולה בלבד.
• הערך NONE מציין שאין חלוקה לפלחים שיכולה להופיע בפעולה הזו.
• לא ניתן לציין את BOTH, כי הפעולה הזו תהיה מיותרת.

מאפיינים: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

ממשקים: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

השפעות: MemoryEffects::Effect{}

מאפיינים:

אופרטנדים:

תוצאות:

sdy.reshard (sdy::ReshardOp)

חלוקה מחדש של טנסור לחלוקה אחרת

תחביר:

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

חלוקה מחדש של הטנסור הקלט לפי החלוקה שצוינה, ששונה מהחלוקה הקיימת של הטנסור הקלט.

שתי הפונקציות ShardingConstraintOp ו-ReshardOp מצרפות חלוקה ל-tensor. משך החיים שלהם הוא:

1. לפני ההפצה של חלוקת המשנה, המשתמשים מוסיפים את ShardingConstraintOp.
2. ההפצה של הפיצול צורכת את ShardingConstraintOp. אין ‎ShardingConstraintOp בתוצאות של ההפצה של חלוקת המחיצות. במקום זאת, אפשר להוסיף את ReshardOp במקרה הצורך.
3. חלוקה למחיצות ממירה את ReshardOp לפעולה קולקטיבית (או לפעולה של זהות). לא אמורה להיות פעולת ReshardOp בתוצאות של מחלק המחיצות.

// TODO(b/331680067). מוסיפים תבנית קנוניזציה כדי להסיר פעולות רישות מיותרות.

מאפיינים: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

ממשקים: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

השפעות: MemoryEffects::Effect{}

מאפיינים:

אופרטנדים:

תוצאות:

sdy.return (sdy::ReturnOp)

הפעולה sdy.return מסיימת את האזורים שמצורפים לפעולות מבוססות-אזור של sdy וכל פעולה מבוססת-אזור אחרת של Shardy. היא משתנה (variadic): כארגומנטים היא מקבלת רשימה של ערכים שיכולים להיות מכל סוג (אבל מאותו סוג, למשל AnyTensor), ולכן אפשר לעשות בה שימוש חוזר ברמות שונות של סטאק ה-IR של Shardy.

תחביר:

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

מאפיינים: AlwaysSpeculatableImplTrait, Terminator

ממשקים: ConditionallySpeculatable, ‏ NoMemoryEffect (MemoryEffectOpInterface)

השפעות: MemoryEffects::Effect{}

אופרטנדים:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

הגבלת טינסור לחלוקה (sharding) שצוינה

תחביר:

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

הוספת חלוקה לענפים (shard) לטרנספורמציה ביניים (למשל, התוצאה של matmul) כדי לציין איך צריך לפצל את הטרנספורמציה הזו או קבוצת משנה של השימושים שלה.

אם לחלוקה יש מאפיינים פתוחים ויש לה צירים ללא אילוצים, סימן שאפשר לפצל את הטנזור עוד יותר לפי המאפיינים הפתוחים.

הפעולה הזו יכולה:

• אין להם שימוש (dangling) – כלומר, חלוקת המשנה המצורפת היא האופן שבו צריך לפצל את הטנסור של הקלט עצמו.
• יש שימושים – כלומר, חלוקת המשנה המצורפת היא האופן שבו צריך לפצל את השימושים של אופרטורי האילוצים של חלוקת המשנה, בעוד שלשימושים אחרים של הטנזור הקלט עשויה להיות חלוקת משנה שונה (אם אין שימושים אחרים לטנזור הקלט, ההתנהגות זהה לתרחיש 'אין שימושים').

מאפיינים: SameOperandsAndResultType

ממשקים: InferTypeOpInterface

מאפיינים:

אופרטנדים:

תוצאות:

sdy.sharding_group (sdy::ShardingGroupOp)

הגבלה של הטנסורים בקבוצה כך שיהיה להם אותו חלוקה לפלחים.

תחביר:

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

הפעולה הזו מספקת ממשק להקצאת טינסורים לקבוצות חלוקה (קבוצות של טינסורים שיחולו עליהן אכיפה של חלוקות זהות). במהלך ההעברה, ברגע שרכיב אחד של הקבוצה מחולק למקטעים, כל שאר המרכיבים מחולקים למקטעים באותו אופן. הפעולה הזו מקבלת את מזהה הקבוצה של הארגומנט ולא מחזירה תוצאה, אלא משנה את הייצוג הפנימי של קבוצת הפיצול כדי להוסיף את הטנסור של הקלט לקבוצה עם המזהה הנתון.

ממשקים: InferTypeOpInterface

מאפיינים:

אופרטנדים:

מאפיינים

AllToAllParamAttr

פרמטר 'הכול לכולם'

תחביר:

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

קבוצת ערכים (tuple) שמכילה את הצירים ואת מאפייני המקור/היעד שעבורם רוצים לבצע את ההתאמה 'הכול אל הכול'.

פרמטרים:

AlltoAllParamListAttr

רשימת פרמטרים של 'הכול לכולם'

תחביר:

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

פרמטרים:

AxisRefAttr

הפניה לציר מלא או לציר משנה מפוצל

תחביר:

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

אילוצים:

• הערך name חייב להופיע במגבלה MeshAttr.
• אם השדה sub_axis_info קיים, הוא חייב לעמוד במגבלות של SubAxisInfoAttr.

פרמטרים:

AxisRefListAttr

רשימת הפניות לצירים

תחביר:

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

אילוצים:

• הרכיבים ב-value חייבים לעמוד באילוצים של AxisRefAttr.
• אין ערכי axis-ref כפולים או צירים משניים שחופפים זה לזה.
• אין שני ערכי axis-ref סמוכים שהם צירים משניים עוקבים של אותו ציר מלא, כלומר אפשר למזג אותם לציר משני אחד או לציר המלא.

פרמטרים:

DimMappingAttr

רשימת מדדי הגורמים של מאפיין

רשימה ריקה מציינת שמדובר במיפוי null (הניתוח או ההדפסה מתבצעים באמצעות *), כלומר המאפיין לא ממופה לאף גורם.

אילוצים:

• יש לפחות אינדקס גורם אחד.
• אינדקסי הגורמים חייבים להיות בטווח [0, $factor_sizes).
• אם יש כמה גורמים, אף אחד מהם לא יכול להיות בגודל 1.
• אין אינדקסים כפולים של גורמים.

פרמטרים:

DimensionShardingAttr

חלוקה של מאפיינים לפלחים

רשימה של שמות צירים לחלוקה של מאפיין טינסור, מגדול לקטן, משתנה בוליאני שמציין אם אפשר לפצל את המאפיין עוד יותר ומספר שלם אופציונלי שמציין את העדיפות של חלוקת המאפיין, שתישמר במהלך ההפצה של החלוקה. רמות העדיפות מגיעות מהערות של חלוקת משתמשים, וערך נמוך יותר מציין רמת עדיפות גבוהה יותר. אם העדיפות לא מופיעה בהערה, המערכת תשתמש בעדיפות הגבוהה ביותר.

אילוצים:

• הרכיבים ב-axes חייבים לעמוד באילוצים שמפורטים ב-AxisRefListAttr.
• אם יש עדיפות לחלוקה של מאפיין:
  • העדיפות חייבת להיות גדולה מ-0 או שווה לה.
  • אם המאפיין סגור, יש לו לפחות ציר אחד.

פרמטרים:

ListOfAxisRefListsAttr

רשימת רשימות של ערכי עזר לצירים

תחביר:

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

פרמטרים:

ManualAxesAttr

רשימת צירים שבהם פעולת החישוב הידנית היא ידנית

תחביר:

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

פרמטרים:

MeshAttr

רשת של צירים ורשימת מכשירים

תחביר:

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

רשת היא רשימה של צירים ורשימה אופציונלית של מזהי מכשירים שמציינת את סדר המכשירים.

אם רשימת הצירים ריקה, למערך יש ציר משתמע ללא שם בגודל 1. במקרה כזה, אם לא מציינים רשימת מזהי מכשירים, רשימת מזהי המכשירים המשתמעת היא [0]. אם מציינים רשימת מזהי מכשירים, היא חייבת להכיל מספר שלם אחד של ערך לא שלילי כלשהו. אנחנו מכנים את המצב הזה 'מקסימום חלוקה למקטעים'.

בכל המקרים שבהם לא מתבצעת חלוקה מקסימלית, אם מצוין רשימת מזהי מכשירים, המכפלה של גדלי הצירים צריכה להתאים למספר המכשירים. אם לא צוינה רשימה של מזהי מכשירים, רשימת מזהי המכשירים המשתמעת היא iota(product(axes)). כדי לפשט את התהליך, אסור גם לציין רשימת מזהי מכשירים זהה ל-iota(product(axes)); במקרה כזה, אסור לציין רשימת מזהי מכשירים.

ריכזנו כאן כמה דוגמאות לרשתות:

• רשת ריקה מייצגת רשת placeholder שאפשר להחליף במהלך ההעברה: <[]>
• רשת עם ציר ללא שם ומזהה מכשיר מפורש, שמשמשת בדרך כלל לייצוג חלוקה מקסימלית: <[], 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)).

פרמטרים:

MeshAxisAttr

ציר בעל שם במערך

תחביר:

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

פרמטרים:

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 גם אם אי אפשר לפצל אותם, בעיקר כדי לשמור על השלמות. לפעולות רבות, כמו פעולות נקודה-לנקודה, יש מאפיינים בגודל 1 שתואמים בין אופרטנדים לתוצאות.

סוגי גורמים:

• המאפיין reduction_factors מכיל את המדדים של הגורמים שצריך לצמצם, כמו המאפיינים המתכווצים בפעולה של נקודה (dot).
• השדה need_replication_factors מכיל את האינדקסים של הגורמים שדורשים שכפול מלא, כמו המאפיין הממוין בפעולת מיון.
• השדה permutation_factors מכיל את האינדקסים של הגורמים שדורשים טרנספוזיציה קולקטיבית אם הם מחולקים למקטעים, כמו מאפייני המילוי בפעולה של מילוי.
• כל שאר הגורמים נחשבים לגורמים של העברה, כלומר גורמים שלא דורשים תקשורת אם הם מחולקים לאותה דרך בכל הטנסורים שממופים אליהם.

blocked_propagation_factors מכיל את הגורמים שאסור להפיץ ביניהם חלוקות. הוא אורתוגונלי לסוגים של הגורמים. כלומר, גורם של העברה חסומה יכול להיות כל אחד מסוגי הגורמים.

השדה is_custom_rule מתאר אם זהו כלל שהוגדר על ידי משתמש. המשתמשים יכולים להגדיר כללי חלוקה לפלחים לשיחות בהתאמה אישית, או לשנות את כללי החלוקה לפלחים שהוגדרו מראש לפעולות הרגילות. כלל מותאם אישית תמיד נשמר ולא מוסר אף פעם.

אילוצים:

• מספר המיפויים של אופרטורים/תוצאות חייב להתאים למספר האופרטורים/התוצאות של הפעולה.
• יש לפחות מיפוי אחד (לא ניתן ליצור כלל לפעולה ללא אופרטורים/תוצאות).
• הדירוג של כל TensorMappingAttr תואם לדירוג של סוג הטנזור התואם.
• לכל קבוצה של גורמים (reduction_factors,‏ need_replication_factors, ‏ permutation_factors):
  • הרכיבים חייבים להיות בטווח [0, $factor_sizes].
  • אין אינדקסים כפולים של גורמים בכל קבוצה ובין הקבוצות.

פרמטרים:

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. לכן, המאפיין sub-axis-info מכיל את שני המספרים האלה, והוא מסומן באופן הבא: (m)k עבור m לפני הגודל ו-k בגודל.

אילוצים:

• הערך של pre-size הוא לפחות 1.
• הערך של size גדול מ-1.
• pre-size חייב לחלק את הגודל של הציר המלא, כלומר גם pre-size וגם size צריכים לחלק את הגודל של הציר המלא, והציר המשני לא יכול לחרוג מהציר המלא.
• הגודל של ציר המשנה לא שווה לגודל של הציר המלא התואם. במקרה כזה, צריך להשתמש בציר המלא במקום זאת.

פרמטרים:

TensorMappingAttr

מיפויים של גורמים לכל מאפיין של טינסור.

תחביר:

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

אילוצים:

• הרכיבים ב-dim_mappings חייבים לעמוד באילוצים ב-DimMappingAttr.
• אין אינדקסים כפולים של גורמים במאפיינים שונים.

פרמטרים:

TensorShardingAttr

חלוקה של טינסורים

תחביר:

#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).

פרמטרים:

TensorShardingPerValueAttr

חלוקה של טינסורים לפי אופרטור/תוצאה של פעולה

תחביר:

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

רשימה של TensorShardingAttr, אחת לכל אופרנד/תוצאה של פעולה.

אילוצים:

• הרכיבים ב-shardings חייבים לעמוד באילוצים של TensorShardingAttr.

פרמטרים:

טיפוסים בני מנייה (enum)

PropagationDirection

Propagation direction enum

מקרים: