הגדרות של OpInterface

CollectiveOpInterface‏ (Sdy_CollectiveOpInterface)

ממשק לכל הפעולות הקבוצתיות. תיבת encapsulation של פונקציות get/set נפוצות למאפיין outSharding.

אילוצים:

  • לביטוי המבצע חייב להיות חלוקה לפלחים, או שהפונקציה allowMissingInputSharding() מחזירה את הערך true.
  • out_sharding תקין ביחס לסוג המתאים.
  • אם הערך המוחזר של allowDifferentMeshes() הוא false, חלוקת המשנה של המשתנה והתוצאה חייבת להיות באותה רשת.
  • דירוג זהה לחלוקת המשנה של המשתנה והתוצאה.

שיטות:

getOutSharding

::mlir::sdy::TensorShardingAttr getOutSharding();

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

setOutShardingAttr

void setOutShardingAttr(::mlir::sdy::TensorShardingAttr sharding);

הגדרת חלוקת הענפים של הטנסור הפלט של הפעולה הקולקטיבית.

getTensor

::mlir::TypedValue<::mlir::TensorType> getTensor();

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

getType

::mlir::Type getType();

אחזור הסוג של תוצאת הפעולה הקולקטיבית.

allowDifferentMeshes

bool allowDifferentMeshes();

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

allowMissingInputSharding

bool allowMissingInputSharding();

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

ShardableDataFlowOpInterface‏ (Sdy_ShardableDataFlowOpInterface)

ממשק אופרטורי שמאפשר ל-shardy להפיץ חלוקות באמצעות צמתים של זרימת נתונים של אופרטורים שמרחיבים את הממשק הזה.

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

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

לדוגמה:

  y_1, ..., y_n = custom_op (x_1, ..., x_n)
                  ((body_arg_1,..., body_arg_n) {
                    ...
                    return return_value_1, ..., return_value_n
                  })

ל-custom_op יש שני סוגים של צמתים של זרימת נתונים, n צמתים כל אחד בין return_value_i (מקורות) לבין y_i (יעדים) ו-n צמתים כל אחד בין x_i(מקורות) לבין body_arg_i(יעדים). במקרה כזה, בעלי הקצוות הם אותם יעדים.

דוגמה למבצע עם כמה מטרות:

  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.

שיטות:

getBlockArgumentEdgeOwnerShardings

mlir::SmallVector<mlir::sdy::TensorShardingAttr> getBlockArgumentEdgeOwnerShardings();

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

setBlockArgumentEdgeOwnerShardings

void setBlockArgumentEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);

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

getOpResultEdgeOwnerShardings

mlir::SmallVector<mlir::sdy::TensorShardingAttr> getOpResultEdgeOwnerShardings();

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

setOpResultEdgeOwnerShardings

void setOpResultEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);

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

transformTargetSharding

mlir::sdy::TensorShardingAttr transformTargetSharding(mlir::Value target, mlir::sdy::TensorShardingAttr sharding, mlir::sdy::DataFlowShardingTransformType transformType);

טרנספורמציה של sharding של היעד בהתאם ל-transformType

מידע נוסף זמין בכתובת DataFlowShardingTransformType.

getBlockArgumentEdgeOwners

mlir::ArrayRef<mlir::BlockArgument> getBlockArgumentEdgeOwners();

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

getOpResultEdgeOwners

mlir::ResultRange getOpResultEdgeOwners();

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

getEdgeSources

mlir::SmallVector<mlir::OpOperand*> getEdgeSources(mlir::Value owner);

הפונקציה מקבלת את מקורות הקצה של זרימת הנתונים לפי הקצה owner.

getEdgeOwnerFromTarget

mlir::Value getEdgeOwnerFromTarget(mlir::Value target);

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

getEdgeOwnerFromSource

mlir::Value getEdgeOwnerFromSource(mlir::OpOperand&source);

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

getNonEdgeOwnerTargets

mlir::SmallVector<mlir::Value> getNonEdgeOwnerTargets(mlir::Value owner);

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

ShardingRuleOpInterface‏ (Sdy_ShardingRuleOpInterface)

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

שיטות:

getShardingRule

mlir::sdy::OpShardingRuleAttr getShardingRule();

הפונקציה מחזירה את כלל הפיצול של הפעולה.