הגדרות של OpInterface

ShardableDataFlowOpInterface (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::ArrayRef<mlir::sdy::TensorShardingAttr> getBlockArgumentEdgeOwnerShardings();

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

setBlockArgumentEdgeOwnerSharding

void setBlockArgumentEdgeOwnerSharding(unsigned index, mlir::sdy::TensorShardingAttr sharding);

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

setBlockArgumentEdgeOwnerShardings

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

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

getOpResultEdgeOwnerShardings

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

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

setOpResultEdgeOwnerSharding

void setOpResultEdgeOwnerSharding(unsigned index, mlir::sdy::TensorShardingAttr sharding);

הגדרת sharding של הבעלים של קצה התוצאה של הפעולה באמצעות index הנתון.

setOpResultEdgeOwnerShardings

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

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

getBlockArgumentEdgeOwners

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

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

getOpResultEdgeOwners

mlir::ResultRange getOpResultEdgeOwners();

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

getEdgeSources

mlir::SmallVector<mlir::Value> getEdgeSources(mlir::Value target);

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

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‏ (ShardingRuleOpInterface)

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

שיטות:

getShardingRule

mlir::sdy::OpShardingRuleAttr getShardingRule();

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