ניב 'sdy'

דיאלקט Shardy (SDY)

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

יומן גרסאות: ‫0.0.1: הוספת צירים לא מצומצמים ל-TensorShardingAttr.

תפעול

‫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>

Constraints:

• חייב לעמוד במגבלות שמפורטות ב-Sdy_CollectiveOpInterface.
• הרכיבים ב-gathering_axes צריכים לעמוד באילוצים שמפורטים ב-AxisRefListAttr.
• החלת gathering_axes על חלוקת האופרנד לשברירים נותנת out_sharding.

מאפיינים: SameOperandsAndResultType

ממשקים: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

מאפיינים:

אופרנדים:

תוצאות:

‫sdy.all_reduce (sdy::AllReduceOp)

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

תחביר:

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

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

Constraints:

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

מאפיינים: SameOperandsAndResultType

ממשקים: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

מאפיינים:

אופרנדים:

תוצאות:

‫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>

Constraints:

• חייב לעמוד במגבלות שמפורטות ב-Sdy_CollectiveOpInterface.
• הרכיבים ב-slicing_axes צריכים לעמוד באילוצים שמפורטים ב-AxisRefListAttr.
• החלת slicing_axes על חלוקת האופרנד לשברירים נותנת out_sharding.

מאפיינים: SameOperandsAndResultType

ממשקים: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

מאפיינים:

אופרנדים:

תוצאות:

‫sdy.all_to_all (sdy::AllToAllOp)

מבצע תקשורת בין כל הצירים

תחביר:

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

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

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

הפעולה all-to-all תוחל על חלוקת האופרנד (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>

Constraints:

• חייב לעמוד במגבלות שמפורטות ב-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, SymbolUserOpInterface

מאפיינים:

אופרנדים:

תוצאות:

‫sdy.collective_permute (sdy::CollectivePermuteOp)

מבצע תקשורת של החלפת צירים באמצעות החלפה מחזורית קולקטיבית

תחביר:

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

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

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

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

בדוגמה שלמטה, גודל הטנזור המפוצל הוא tensor<1x4x2xf32>, והוא נשמר על ידי הפעולה collective permute.

דוגמה:

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>

Constraints:

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

מאפיינים: SameOperandsAndResultType

ממשקים: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

מאפיינים:

אופרנדים:

תוצאות:

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

בזמן הפעולה הזה יש 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, ‏ SymbolUserOpInterface

מאפיינים:

אופרנדים:

תוצאות:

‫sdy.func_data_flow_edge (sdy::FuncDataFlowEdgeOp)

פעולת קצה של זרימת נתונים של קלט/פלט של פונקציה

תחביר:

operation ::= `sdy.func_data_flow_edge` $operand attr-dict `:` type($result)

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

מאפיינים: SameOperandsAndResultType

ממשקים: InferTypeOpInterface, ‏ SymbolUserOpInterface

אופרנדים:

תוצאות:

‫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.

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

Constraints:

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

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

ממשקים: ShardableDataFlowOpInterface, ‏ SymbolUserOpInterface

מאפיינים:

אופרנדים:

תוצאות:

‫sdy.mesh (sdy::MeshOp)

רשת בעלת שם

תחביר:

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

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

מאפיינים: HasParent<ModuleOp>, ‏ SymbolName

ממשקים: 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

ממשקי API: ConditionallySpeculatable, InferTypeOpInterface, ShardableDataFlowOpInterface, SymbolUserOpInterface

מאפיינים:

אופרנדים:

תוצאות:

‫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.reduce_scatter (sdy::ReduceScatterOp)

מבצעת תקשורת של צמצום פיזור לאורך צירים

תחביר:

operation ::= `sdy.reduce_scatter` ($reduction_op^)? $reduce_scatter_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

מצמצמת נתחים של טנסור לאורך צירים שצוינו ב-reduce_scatter_axes, ואז מפזרת את התוצאה לאורך אותם צירים. הפעולה הזו היא בעצם שילוב של sdy.all_reduce ואחריו sdy.all_slice באותו reduce_scatter_axes.

Constraints:

• חייב לעמוד במגבלות שמפורטות ב-Sdy_CollectiveOpInterface.
• הרכיבים ב-reduce_scatter_axes צריכים לעמוד באילוצים שמפורטים ב-AxisRefListAttr.
• החלת reduce_scatter_axes על חלוקת האופרנד מקבלת out_sharding.

מאפיינים: SameOperandsAndResultType

ממשקים: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

מאפיינים:

אופרנדים:

תוצאות:

‫sdy.replicated_to_unreduced (sdy::ReplicatedToUnreducedOp)

העברה של צירים משוכפלים באופן מרומז או מפורש לצירים לא מצומצמים.

תחביר:

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

הערך axes צריך להיות משוכפל באופן מרומז או מפורש באופרנד. הפעולה הזו גורמת לכך שהם לא יצומצמו בתוצאה. הקשר בינינו הוא:

all-reduce(replicated-to-unreduced(x, axes), axes) = x

דוגמה:

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

Constraints:

• חייב לעמוד במגבלות שמפורטות ב-Sdy_CollectiveOpInterface.
• הפרמטר axes חייב לעמוד באילוצים שמפורטים בפרמטר AxisRefListAttr.
• צריך למיין את axes ביחס לרשת.
• השדות axes לא ריקים.
• הפיצול של הקלט והפלט חייב להיות זהה.
• ‫axes צריך להיות משוכפל באופן מרומז או מפורש בפיצול האופרנד.
• inUnreducedAxes + axes = outUnreducedAxes.

מאפיינים: SameOperandsAndResultType

ממשקים: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

מאפיינים:

אופרנדים:

תוצאות:

‫sdy.reshard (sdy::ReshardOp)

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

תחביר:

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

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

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

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

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

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

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

מאפיינים:

אופרנדים:

תוצאות:

‫sdy.return (sdy::ReturnOp)

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

תחביר:

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

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

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

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

אופרנדים:

‫sdy.sharded_to_unreduced (sdy::ShardedToUnreducedOp)

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

תחביר:

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

צריך להשתמש ב-axes כדי לפצל את האופרנד. הפעולה הזו מבטלת את הצמצום שלהם בתוצאה. היחסים בינינו הם:

‫all-gather(x, axes) = all-reduce(sharded-to-unreduced(x, axes), axes), where all-gather, sharded-to-unreduced, all-reduce are applied on the same axes.

דוגמה:

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

Constraints:

• חייב לעמוד במגבלות שמפורטות ב-Sdy_CollectiveOpInterface.
• הרכיבים ב-axes צריכים לעמוד באילוצים שמפורטים ב-AxisRefListAttr.
• החלת axes על חלוקת האופרנד לשברירים נותנת out_sharding.

מאפיינים: SameOperandsAndResultType

ממשקים: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

מאפיינים:

אופרנדים:

תוצאות:

‫sdy.sharding_constraint (sdy::ShardingConstraintOp)

מגבילה טנסור לפי הפיצול שצוין

תחביר:

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

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

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

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

• אין לו שימושים (תלוי) – כלומר, ה-sharding המצורף הוא הדרך שבה צריך לבצע sharding של טנסור הקלט עצמו.
• ‫Have uses – כלומר, הפיצול המצורף הוא הפיצול שבו צריך לפצל את השימושים באילוץ הפיצול, בעוד ששימושים אחרים בטנזור הקלט עשויים להיות בעלי פיצול שונה (אם אין שימושים אחרים בטנזור הקלט, ההתנהגות זהה למקרה של no uses).

מאפיינים: SameOperandsAndResultType

ממשקים: InferTypeOpInterface, ‏ SymbolUserOpInterface

מאפיינים:

אופרנדים:

תוצאות:

sdy.sharding_group (sdy::ShardingGroupOp)

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

תחביר:

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

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

ממשקים: InferTypeOpInterface

מאפיינים:

אופרנדים:

מאפיינים

AllToAllParamAttr

פרמטר all-to-all

תחביר:

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

טופל שמכיל את הצירים ואת מאפייני המקור או היעד שרוצים לבצע עליהם את ההמרה.

פרמטרים:

AllToAllParamListAttr

רשימה של כל הפרמטרים מסוג all-to-all

תחביר:

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

פרמטרים:

AxisRefAttr

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

תחביר:

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

Constraints:

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

פרמטרים:

AxisRefListAttr

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

תחביר:

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

Constraints:

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

פרמטרים:

AxisToPropagationDetailsAttr

פרטים על זרימת הקצוות של ההתפשטות בציר ובמקור ספציפיים.

תחביר:

#sdy.axis_to_propagation_details<
  ::mlir::sdy::AxisRefAttr,   # axis_name
  ::mlir::sdy::EdgeValueRefAttr,   # source
  ::llvm::ArrayRef<EdgeValueRefAttr>   # targets
>

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

פרמטרים:

DimMappingAttr

רשימה של אינדקסים של גורמים למאפיין

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

Constraints:

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

פרמטרים:

DimensionShardingAttr

חלוקת מאפיינים

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

Constraints:

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

פרמטרים:

EdgeValueRefAttr

הפניה לאינדקס מסוים של קצה ערך מהסוג type.

תחביר:

#sdy.edge_value_ref<
  `operand` | `result`,   # type
  int64_t   # index
>

פרמטרים:

ListOfAxisRefListsAttr

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

תחביר:

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

פרמטרים:

ManualAxesAttr

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

תחביר:

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

פרמטרים:

MeshAttr

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

תחביר:

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

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

אם רשימת הצירים ריקה

• אם לא מספקים את device_ids, מדובר ברשת ריקה.
• אם מציינים את device_ids, הוא חייב להיות מספר שלם לא שלילי יחיד. אנחנו קוראים לו maximal-sharding mesh.

אם מספקים את רשימת הצירים

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

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

• רשת ריקה מייצגת רשת placeholder שאפשר להחליף במהלך ההפצה: <[]>
• רשת מש עם רשימה ריקה של צירים ומזהה מכשיר יחיד לא שלילי, שהיא רשת מש עם פיצול מקסימלי: <[], device_ids=[3]>
• רשת עם שני צירים ומזהי מכשירים מרומזים iota(6): <["a"=2, "b"=3]>
• רשת עם שני צירים ומזהי מכשירים מפורשים שמציינים את סדר המכשירים: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

Constraints:

• הערכים ברכיב device_ids לא יכולים להיות שליליים.
• אם הערך של axes ריק, הגודל של device_ids יכול להיות 0 (רשת ריקה) או 1 (רשת עם חלוקה מקסימלית).
• אם axes לא ריק,
  • אסור שיהיו שמות כפולים לרכיבים במאפיין 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 גם אם אי אפשר לבצע חלוקה לשברים שלהם. הסיבה העיקרית לכך היא שלמות, כי להרבה פעולות כמו פעולות נקודתיות יש מימדים בגודל אחד שתואמים בין אופרנדים ותוצאות.

סוגי גורמים:

• ‫reduction_factors מכיל את האינדקסים של הגורמים שצריך לצמצם, כמו המאפיינים של החוזה בפעולת נקודה. הגורמים האלה יכולים להיות באופרנדים אבל לא בתוצאות.
• ‫need_replication_factors: מכיל את האינדקסים של גורמים שנדרש שכפול מלא שלהם, כמו המאפיין הממוין בפעולת מיון.
• השדה permutation_factors מכיל את האינדקסים של גורמים שנדרש להם collective-permute אם הם מפולגים, כמו ממדי הריפוד בפעולת ריפוד.
• כל שאר הגורמים נחשבים לגורמים שמועברים דרך המערכת, כלומר גורמים שלא נדרשת תקשורת לגביהם אם הם מחולקים באופן זהה בכל הטנסורים שממופים אליהם.

‫blocked_propagation_factors מכיל את הגורמים שאסור להפיץ לאורך הפיצולים. היא אורתוגונלית לסוגי הגורמים. כלומר, גורם חסימה יכול להיות כל אחד מסוגי הגורמים.

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

Constraints:

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

פרמטרים:

PropagationEdgesAttr

מטא-נתונים של קצה ההפצה לכל שלבי ההפצה.

תחביר:

#sdy.propagation_edges<
  ::llvm::ArrayRef<PropagationOneStepAttr>   # value
>

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

פרמטרים:

PropagationOneStepAttr

מטא-נתונים של הפצה לכל שלב.

תחביר:

#sdy.propagation_one_step<
  int64_t,   # step_index
  ::llvm::ArrayRef<AxisToPropagationDetailsAttr>   # axis_entries
>

פרטי ההתפשטות של כל הצירים בשלב התפשטות יחיד.

פרמטרים:

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 אחרי השינוי.

Constraints:

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

פרמטרים:

TensorMappingAttr

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

תחביר:

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

Constraints:

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

פרמטרים:

TensorShardingAttr

Tensor sharding

תחביר:

#sdy.sharding<
  ::mlir::Attribute,   # mesh_or_ref
  ::llvm::ArrayRef<DimensionShardingAttr>,   # dim_shardings
  ::llvm::ArrayRef<AxisRefAttr>,   # replicated_axes
  ::llvm::ArrayRef<AxisRefAttr>,   # unreduced_axes
  `sum` | `max` | `min`   # reduction_op
>

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

שימו לב: אם לא מציינים מאפיין חלוקה ל-shards בטנזור, המשמעות היא חלוקה ל-shards של טנזור פתוח לחלוטין.

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

ל-sharding יכולים להיות צירים לא מצומצמים (שמצוינים על ידי unreduced_axes), כלומר הטנזור לא מצומצם לאורך הצירים האלה. לדוגמה, אם המימד של כיווץ מטריצות (matmul) מפולח לאורך ציר x גם בצד שמאל (lhs) וגם בצד ימין (rhs), התוצאה לא מצומצמת לאורך x. הפעלת all-reduce על הטנזור לאורך הצירים שלא צומצמו תגרום לשכפול הטנזור לאורך הצירים האלה. עם זאת, לא חייבים לצמצם טנזור עם צירים לא מצומצמים באופן מיידי, והוא יכול להישאר לא מצומצם כשמעבירים אותו לפעולות לינאריות כמו stablehlo.add (כל עוד גם lhs וגם rhs לא מצומצמים), ואז לצמצם אותו. אנחנו מניחים שסוג ההפחתה הוא סכום, אבל יכול להיות שבעתיד נתמוך בהפחתות אחרות.

Constraints:

• הרכיבים ב-dim_shardings צריכים לעמוד באילוצים שמפורטים ב-DimensionShardingAttr.
• הרכיבים ב-replicated_axes צריכים לעמוד באילוצים שמפורטים ב-AxisRefListAttr.
• הרכיבים ב-unreduced_axes צריכים לעמוד באילוצים שמפורטים ב-AxisRefListAttr.
• אם סוג הטנזור המתאים הוא לא ShapedType, הפיצול צריך להיות בדרגה 0 וללא צירים משוכפלים.
• אם מדובר ב-ShapedType:
  • ל-Tensor צריך להיות דרגה.
  • מספר חלוקות המאפיינים שווה לדרגה של הטנזור.
  • הממדים של גודל 0 לא מחולקים.
• אין הפניות כפולות לצירים או לצירים משניים שחופפים זה לזה ב-dim_shardings, ב-replicated_axes וב-unreduced_axes.
• הפריטים ב-replicated_axes וב-unreduced_axes מסודרים לפי mesh_or_ref (ראו AxisRefAttr::getMeshComparator).

פרמטרים:

TensorShardingPerValueAttr

Tensor sharding per operand/result of an op

תחביר:

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

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

Constraints:

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

פרמטרים:

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

EdgeNodeType

Edge node type enum

מקרים:

PropagationDirection

ספירה (enum) של כיוון ההפצה

מקרים:

ReductionOp

Enum של פעולת צמצום

מקרים: