-sdy-close-shardings

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

-sdy-drop-sharding-rules

הסרת OpShardingRuleAttr מכל התפעולים הרשומים.

-sdy-insert-explicit-reshards

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

חלוקה תואמת למשנה (shard) היא למעשה חלוקה שמאפשרת לפעולה לקבל את אופרטנדים המחולקים וליצור תוצאה מחולקת ללא צורך בתקשורת של חלוקה מחדש (שימו לב: יכול להיות שהפעולה עדיין תצטרך תקשורת, כמו all-reduce או halo-swaps).

אחרי ההעברה, יכול להיות שחלק מהפעולות עדיין יכללו חלוקות לא תואמות.

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

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

דוגמה להבהרה:

קלט:

mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"y"},{"x"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
stablehlo.dot %lhs, %rhs {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
  : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>

פלט:

sdy.mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"x"}, {"y"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
%0 = sdy.reshard %rhs <@mesh, \[{"y"}, {}\]> : tensor<32x16xf32>
stablehlo.dot %lhs, %0 {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
  : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>

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

-sdy-remove-sharding-groups

הסרת ShardingGroupOps אחרי העברה.

-sdy-sharding-constraint-to-reshard

המרת ShardingConstraintOp ל-ReshardOp.

-sdy-sink-data-flow-edges

מזין את כל DataFlowEdgeOp לקלט שלו.

העברת חלוקת המשנה של כל DataFlowEdgeOp לקלט שלו (יעד הבסיס של הקצה), והחלפת הפעולה (op) בקלט שלו.

TODO(tomnatan): consider moving the sharding to all targets that can have a sharding attached.

-sdy-update-non-divisible-input-output-shardings

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

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