‫'mpmd' Dialect

תפעול

mpmd.assign (mpmd::AssignOp)

פעולת הקצאה

תחביר:

operation ::= `mpmd.assign` attr-dict $tensor `:` functional-type(operands, results)

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

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

שם ה-mesh של סוג התוצאה צריך להתאים ל-mesh בטופולוגיה, והסוג הגלובלי שלו צריך להיות זהה לסוג האופרנד.

המקור של פעולת ההקצאה הוא המקור של הרשת, למשל named_computation,‏ mesh inference וכו'.

מאפיינים: AlwaysSpeculatableImplTrait, ‏ HasParent<::mlir::func::FuncOp, ForOp>

ממשקים: ConditionallySpeculatable, ‏ NoMemoryEffect (MemoryEffectOpInterface)

השפעות: MemoryEffects::Effect{}

מאפיינים:

אופרנדים:

תוצאות:

mpmd.broadcast (mpmd::BroadcastOp)

פעולת שידור

תחביר:

operation ::= `mpmd.broadcast` attr-dict $tensor `:` type($tensor)

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

מאפיינים: AlwaysSpeculatableImplTrait, ‏ SameOperandsAndResultType

ממשקים: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

השפעות: MemoryEffects::Effect{}

אופרנדים:

תוצאות:

‫mpmd.call (mpmd::CallOp)

פונקציית התקשרות ספציפית ל-MPMD

תחביר:

operation ::= `mpmd.call` $callee `(` $tensors `)` attr-dict `:` functional-type(operands, results)

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

ממשקים: ArgAndResultAttrsOpInterface, CallOpInterface, SymbolUserOpInterface

מאפיינים:

אופרנדים:

תוצאות:

‫mpmd.for (mpmd::ForOp)

למפעילים

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

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

אופרטור for יכול לקבל ולהחזיר כל סוג, אבל ה-TypeID של הסוגים האלה חייב להיות זהה – למשל, כל סוגי הטנסור או כל סוגי ה-MPMD mesh וכו'. כך אפשר להשתמש באופרטור ברמות שונות, ולשתף את ההטמעה והטרנספורמציות.

מאפיינים: HLO_PairwiseSameOperandAndResultType, ‏ RecursiveMemoryEffects, ‏ RecursivelySpeculatableImplTrait, ‏ SingleBlockImplicitTerminator<ReturnOp>, ‏ SingleBlock

ממשקי API: ConditionallySpeculatable, LoopLikeOpInterface, OpAsmOpInterface, ShardableDataFlowOpInterface

מאפיינים:

אופרנדים:

תוצאות:

mpmd.fragment (mpmd::FragmentOp)

פעולת קטע

הפונקציה מקצה חישוב, כלומר בלוק של פעולות, לרשת ספציפית בטופולוגיית MPMD, שמיועדת להרצה כקטע תוכנית SPMD נפרד.

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

שם הרשת של הפריט צריך להיות זהה לשם של רשת בטופולוגיה.

הפרגמנט כולל רשימה של מקורות, כלומר מטא-נתונים עם מידע על חישובים מקוריים בשם named_computations שיצרו את הפרגמנט הזה, ו-staged_id שמוגדר iff אם זה פרגמנט שהוגדר על ידי משתמש, כלומר יש לו רשימה לא ריקה של מקורות. הפרמטר האופציונלי in_shardings מציין את הפיצול של ארגומנטים של בלוקים של פרגמנט, שתואמים לאופרנדים. הפרמטר האופציונלי out_shardings מציין את החלוקה של התוצאות.

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

מאפיינים: HasParent<::mlir::func::FuncOp, ForOp>, IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

ממשקים: ConditionallySpeculatable, ‏ ShardableDataFlowOpInterface

מאפיינים:

אופרנדים:

תוצאות:

‫mpmd.fragment_call (mpmd::FragmentCallOp)

פעולת קריאה של קטע

מייצג קריאה לפונקציה שמכילה גוף של מקטע MPMD, כלומר חישוב שמוקצה לרשת ספציפית בטופולוגיה של MPMD, שמיועד לביצוע כמקטע תוכנית SPMD נפרד.

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

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

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

דוגמה:

%2 = mpmd.fragment_call<mesh="m1",> origin=[] @my_fragment(%0, %1) :
  (m<esh>_tensor..., m<esh>_te>nsor...) - m<esh>_tensor...

מאפיינים: HasParent<::mlir::func::FuncOp>, ‏ MemRefsNormalizable

ממשקים: ArgAndResultAttrsOpInterface, CallOpInterface, SymbolUserOpInterface

מאפיינים:

אופרנדים:

תוצאות:

‫mpmd.named_computation (mpmd::NamedComputationOp)

פעולה בהיקף בעל שם

מקבץ חישוב, כלומר בלוק של פעולות, ונותן לו שם ומספר טרנספוזיציות באמצעות המאפיין UserOrigin. אפשר להשתמש ב-NamedComputation כדי להקצות רשת לחישוב ב-MPMD או לאופטימיזציות.

המספר של פעולות הטרנספוזיציה (ברירת מחדל=0) מציין אם החישוב שצוין נוצר על ידי מספר מסוים של טרנספורמציות של טרנספוזיציה של JAX AD.

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

מאפיינים: IsolatedFromAbove, ‏ RecursiveMemoryEffects, ‏ RecursivelySpeculatableImplTrait, ‏ SingleBlockImplicitTerminator<ReturnOp>, ‏ SingleBlock

ממשקים: ConditionallySpeculatable

מאפיינים:

אופרנדים:

תוצאות:

mpmd.named_tensor (mpmd::NamedTensorOp)

הקצאת טנסור לרשת

תחביר:

operation ::= `mpmd.named_tensor` $tensor `name````=```$name attr-dict `:` type($result)

פעולת זהות שמשייכת את התוצאה של הטנזור לשם נתון. אפשר להשתמש ב-NamedTensor הזה כדי להקצות רשת לטנזור ב-MPMD.

מאפיינים: AlwaysSpeculatableImplTrait, ‏ SameOperandsAndResultType

ממשקים: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

השפעות: MemoryEffects::Effect{}

מאפיינים:

אופרנדים:

תוצאות:

‫mpmd.reduce (mpmd::ReduceOp)

פעולת צמצום בין רשתות

תחביר:

operation ::= `mpmd.reduce` `` $reduction attr-dict $tensors `:` functional-type(operands, results)

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

מאפיינים: AlwaysSpeculatableImplTrait, ‏ SameOperandsAndResultType

ממשקים: ConditionallySpeculatable, ‏ NoMemoryEffect (MemoryEffectOpInterface)

השפעות: MemoryEffects::Effect{}

מאפיינים:

אופרנדים:

תוצאות:

mpmd.return (mpmd::ReturnOp)

הפעולה mpmd.return מסיימת את האזורים שמצורפים לפעולות מבוססות אזור של mpmd. היא מקבלת מספר משתנה של ארגומנטים: רשימה של ערכים שהסוגים שלהם יכולים להיות כל סוג (אבל מאותו סוג, למשל AnyTensor) ולכן אפשר לעשות בה שימוש חוזר ברמות שונות של מחסנית ה-IR של MPMD.

תחביר:

operation ::= `mpmd.return` attr-dict $results (`:` type($results)^)?

מאפיינים: AlwaysSpeculatableImplTrait, ‏ Terminator

ממשקים: ConditionallySpeculatable, ‏ NoMemoryEffect (MemoryEffectOpInterface)

השפעות: MemoryEffects::Effect{}

אופרנדים:

mpmd.transfer (mpmd::TransferOp)

פעולת העברה

תחביר:

operation ::= `mpmd.transfer` attr-dict $tensor `:` functional-type(operands, results)

העברה של טנסור מבוזר מרשת אחת לרשת אחרת.

השמות של ה-mesh של האופרנד וסוגי התוצאות צריכים להתאים ל-mesh בטופולוגיה, והסוגים הגלובליים שלהם צריכים להיות זהים.

מאפיינים: AlwaysSpeculatableImplTrait, ‏ HasParent<::mlir::func::FuncOp>

ממשקים: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), ShardingRuleOpInterface

השפעות: MemoryEffects::Effect{}

אופרנדים:

תוצאות:

mpmd.unassign (mpmd::UnassignOp)

פעולת ביטול ההקצאה

תחביר:

operation ::= `mpmd.unassign` attr-dict $tensor `:` functional-type(operands, results)

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

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

שם הרשת של סוג האופרנד צריך להתאים לרשת בטופולוגיה, והסוג הגלובלי שלו צריך להיות זהה לסוג התוצאה.

מאפיינים: AlwaysSpeculatableImplTrait, ‏ HasParent<::mlir::func::FuncOp, ForOp>, ‏ InferTensorType

ממשקי API: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

השפעות: MemoryEffects::Effect{}

מאפיינים:

אופרנדים:

תוצאות:

מאפיינים

MeshWithOriginsAttr

התאימו את התוכן למקורות שלו.

תחביר:

#mpmd.mesh_with_origins<
  ::llvm::StringRef,   # mesh_name
  ::llvm::ArrayRef<OriginAttr>   # origins
>

פרמטרים:

MeshesWithOriginsAttr

רשימה של רשתות עם המקורות שלהן.

תחביר:

#mpmd.meshes_with_origins<
  ::llvm::ArrayRef<MeshWithOriginsAttr>   # value
>

פרמטרים:

NamedMeshAttr

זוג עם שם ורשת.

תחביר:

#mpmd.named_mesh<
  ::llvm::StringRef,   # name
  sdy::MeshAttr   # mesh
>

פרמטרים:

OriginAttr

מקור הקצאת הרשת.

תחביר:

#mpmd.origin<
  ::llvm::StringRef   # origin_label
>

המקור של הקצאת רשת.

‫origin_label היא תווית קריאה (לבני אדם) למקור. הוא מיועד לשימוש לצורך ניפוי באגים.

פרמטרים:

ReductionAttr

מציין ירידה.

תחביר:

#mpmd.reduction<
  `none` | `add` | `max` | `min` | `mul` | `or` | `and`   # reduction_type
>

פרמטרים:

TopologyAttr

טופולוגיה של רשתות עם שמות.

תחביר:

#mpmd.topology<
  ::llvm::ArrayRef<NamedMeshAttr>   # meshes
>

פרמטרים:

UserOriginAttr

מקור החישוב שצוין על ידי המשתמש.

תחביר:

#mpmd.user_origin<
  ::mlir::StringAttr,   # userName
  int64_t   # transposeCount
>

פרמטרים:

סוגים

MeshTensorType

סוג טנסור של רשת

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

פרמטרים:

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

ReductionType

מאפיין סוג ההפחתה

מקרים:

EdgeNodeType

Edge node type enum

מקרים:

PropagationDirection

ספירה (enum) של כיוון ההפצה

מקרים:

ReductionOp

Enum של פעולת צמצום

מקרים: