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();
הפונקציה מחזירה את כלל הפיצול של הפעולה.