ניב 'sdy'

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

תפעול

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.

תכונות: IsolatedFromAbove, RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

מאפיינים:

אופרטנדים:

תוצאות:

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, ‏ ShardableDataFlowOpInterface

מאפיינים:

אופרטנדים:

תוצאות:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

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

תחביר:

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

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

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

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

מאפיינים: AlwaysSpeculatableImplTrait, Elementwise, 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, ‏ Elementwise, ‏ 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)

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

אם לפיצול יש מידות פתוחות וצירים לא מוגבלים, פירוש הדבר הוא שאפשר לפצל את החיתוך עוד יותר לאורך הממדים הפתוחים.

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

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

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

ממשקים: InferTypeOpInterface

מאפיינים:

אופרטנדים:

תוצאות:

sdy.sharding_group (sdy::ShardingGroupOp)

פעולה של קבוצת חלוקה

תחביר:

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

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

מאפיינים:

אופרטנדים:

מאפיינים

AxisRefAttr

התייחסות לציר מלא או לציר משנה מפוצל

תחביר:

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

פרמטרים:

DimMappingAttr

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

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

פרמטרים:

DimensionShardingAttr

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

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

פרמטרים:

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

פרמטרים:

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
  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 שתואמים בין אופרטנדים לתוצאות.

השדה is_custom_rule מתאר אם זהו כלל שהוגדר על ידי משתמש לפעולה stablehlo.custom_call. הכלי לחלוקת המחיצות לא יודע איך לחלק את הפעולות האלה, ולכן המשתמש צריך להסביר לו איך. כשמדובר בכלל בהתאמה אישית, הוא תמיד נשמר/לעולם לא מוסר. הערך is_custom_rule יכול להיות True רק לפעולות 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. לכן, המאפיין sub-axis-info מכיל את שני המספרים האלה, והוא מסומן באופן הבא: (m)k עבור m לפני הגודל ו-k בגודל.

פרמטרים:

TensorMappingAttr

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

תחביר:

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

פרמטרים:

TensorShardingAttr

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

תחביר:

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

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

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

פרמטרים:

TensorShardingPerValueAttr

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

תחביר:

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

פרמטרים:

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

PropagationDirection

כיוון ההפצה 'טיפוסים בני מנייה (enum)'

מקרים: